fix(agent): address session and streaming concurrency bugs

This commit is contained in:
hamb1y
2026-05-28 22:54:46 +08:00
committed by Xubin Ren
parent 1a4ae8994d
commit 0df60416ba
10 changed files with 250 additions and 39 deletions
+10 -6
View File
@@ -16,6 +16,7 @@ There is **no** sub-agent orchestrator and **no** special WebSocket ``agent_ui``
from __future__ import annotations
from contextvars import ContextVar
from datetime import datetime
from typing import TYPE_CHECKING, Any
@@ -45,15 +46,19 @@ class _GoalToolsMixin(ContextAware):
def __init__(self, sessions: SessionManager, bus: Any | None = None) -> None:
self._sessions = sessions
self._bus = bus
self._request_ctx: RequestContext | None = None
self._request_ctx: ContextVar[RequestContext | None] = ContextVar(
f"{self.__class__.__name__}_request_ctx",
default=None,
)
def set_context(self, ctx: RequestContext) -> None:
self._request_ctx = ctx
self._request_ctx.set(ctx)
def _session(self):
if self._request_ctx is None:
request_ctx = self._request_ctx.get()
if request_ctx is None:
return None
key = self._request_ctx.session_key
key = request_ctx.session_key
if not key:
return None
return self._sessions.get_or_create(key)
@@ -61,7 +66,7 @@ class _GoalToolsMixin(ContextAware):
async def _publish_goal_state_ws(self, metadata: dict[str, Any]) -> None:
"""Fan-out authoritative goal snapshot for this WebSocket chat only."""
bus = self._bus
rc = self._request_ctx
rc = self._request_ctx.get()
if bus is None or rc is None or rc.channel != "websocket":
return
cid = (rc.chat_id or "").strip()
@@ -224,4 +229,3 @@ class CompleteGoalTool(Tool, _GoalToolsMixin):
if tail:
return f"Goal marked complete ({ended}). Recap:\n{tail}"
return f"Goal marked complete ({ended})."