From 2807cfd61ead026d30fa8cc3fe88b74eaef4642e Mon Sep 17 00:00:00 2001 From: Ziang Xie Date: Sat, 2 May 2026 01:42:56 +0800 Subject: [PATCH] fix: async plan isolation, blocking API call, and CORS for production - plan_tool.py: Replace module-level _current_plan list with per-session _session_plans dict to prevent cross-session plan corruption when multiple concurrent sessions use the plan tool. - hf_repo_git_tool.py: Wrap get_repo_discussions in _async_call to avoid blocking the event loop with a synchronous HTTP request. - backend/main.py: Add SPACE_HOST env var to CORS allow_origins so the production HF Spaces frontend is not rejected by the middleware. - terminal_display.py / main.py: Thread session through plan display functions for session-aware plan retrieval. --- agent/main.py | 2 +- agent/tools/hf_repo_git_tool.py | 13 ++++++++----- agent/tools/plan_tool.py | 23 +++++++++++++++-------- agent/utils/terminal_display.py | 9 +++++---- backend/main.py | 21 ++++++++++++++------- 5 files changed, 43 insertions(+), 25 deletions(-) diff --git a/agent/main.py b/agent/main.py index 606aaf8e..735b51c4 100644 --- a/agent/main.py +++ b/agent/main.py @@ -347,8 +347,8 @@ def _cancel_event(): shimmer.stop() stream_buf.discard() print_turn_complete() - print_plan() session = session_holder[0] if session_holder else None + print_plan(session=session) if session is not None: await session.send_deferred_turn_complete_notification(event) turn_complete_event.set() diff --git a/agent/tools/hf_repo_git_tool.py b/agent/tools/hf_repo_git_tool.py index d7e2323a..21fed435 100644 --- a/agent/tools/hf_repo_git_tool.py +++ b/agent/tools/hf_repo_git_tool.py @@ -285,11 +285,14 @@ async def _list_prs(self, args: Dict[str, Any]) -> ToolResult: repo_type = args.get("repo_type", "model") status = args.get("status", "all") # open, closed, all - discussions = list(self.api.get_repo_discussions( - repo_id=repo_id, - repo_type=repo_type, - discussion_status=status if status != "all" else None, - )) + def _fetch(): + return list(self.api.get_repo_discussions( + repo_id=repo_id, + repo_type=repo_type, + discussion_status=status if status != "all" else None, + )) + + discussions = await _async_call(_fetch) if not discussions: return {"formatted": f"No discussions in {repo_id}", "totalResults": 0, "resultsShared": 0} diff --git a/agent/tools/plan_tool.py b/agent/tools/plan_tool.py index a923d53c..de619343 100644 --- a/agent/tools/plan_tool.py +++ b/agent/tools/plan_tool.py @@ -5,8 +5,9 @@ from .types import ToolResult -# In-memory storage for the current plan (raw structure from agent) -_current_plan: List[Dict[str, str]] = [] +# Per-session plan storage to avoid cross-session corruption +_session_plans: Dict[str, List[Dict[str, str]]] = {} +_last_plan_session_id: str | None = None class PlanTool: @@ -26,7 +27,7 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: Returns: ToolResult with formatted output """ - global _current_plan + global _session_plans, _last_plan_session_id todos = params.get("todos", []) @@ -54,8 +55,10 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: "isError": True, } - # Store the raw todos structure in memory - _current_plan = todos + # Store per-session to prevent cross-session plan corruption + session_id = self.session.session_id if self.session else "__no_session__" + _session_plans[session_id] = todos + _last_plan_session_id = session_id # Emit plan update event if session is available if self.session: @@ -76,9 +79,13 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: } -def get_current_plan() -> List[Dict[str, str]]: - """Get the current plan (raw structure).""" - return _current_plan +def get_current_plan(session_id: str | None = None) -> List[Dict[str, str]]: + """Get the current plan for a session (raw structure).""" + if session_id: + return _session_plans.get(session_id, []) + if _last_plan_session_id: + return _session_plans.get(_last_plan_session_id, []) + return [] # Tool specification diff --git a/agent/utils/terminal_display.py b/agent/utils/terminal_display.py index f2b73301..25d2dbe9 100644 --- a/agent/utils/terminal_display.py +++ b/agent/utils/terminal_display.py @@ -437,11 +437,12 @@ def print_help() -> None: # ── Plan display ─────────────────────────────────────────────────────── -def format_plan_display() -> str: +def format_plan_display(session=None) -> str: """Format the current plan for display.""" from agent.tools.plan_tool import get_current_plan - plan = get_current_plan() + session_id = session.session_id if session else None + plan = get_current_plan(session_id=session_id) if not plan: return "" @@ -462,8 +463,8 @@ def format_plan_display() -> str: return "\n".join(lines) -def print_plan() -> None: - plan_str = format_plan_display() +def print_plan(session=None) -> None: + plan_str = format_plan_display(session=session) if plan_str: _console.print(plan_str) diff --git a/backend/main.py b/backend/main.py index f6bc64d1..9b3f873e 100644 --- a/backend/main.py +++ b/backend/main.py @@ -70,15 +70,22 @@ async def lifespan(app: FastAPI): lifespan=lifespan, ) -# CORS middleware for development +# CORS middleware +allow_origins = [ + "http://localhost:5173", # Vite dev server + "http://localhost:3000", + "http://127.0.0.1:5173", + "http://127.0.0.1:3000", +] + +# Add production HF Spaces URL when deployed +space_host = os.environ.get("SPACE_HOST") +if space_host: + allow_origins.append(f"https://{space_host}") + app.add_middleware( CORSMiddleware, - allow_origins=[ - "http://localhost:5173", # Vite dev server - "http://localhost:3000", - "http://127.0.0.1:5173", - "http://127.0.0.1:3000", - ], + allow_origins=allow_origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"],