fix(api): separate run control from stream lifetime
This commit is contained in:
teknium1
Teknium
parent
1da89a5f3d
commit
8f18fa104f
2 changed files with 188 additions and 43 deletions
|
|
@ -882,9 +882,13 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
self._run_streams: Dict[str, "asyncio.Queue[Optional[Dict]]"] = {}
|
||||
# Creation timestamps for orphaned-run TTL sweep
|
||||
self._run_streams_created: Dict[str, float] = {}
|
||||
# Runs with a connected SSE consumer; their queue is actively draining.
|
||||
self._run_stream_subscribers: set[str] = set()
|
||||
# Active run agent/task references for stop support
|
||||
self._active_run_agents: Dict[str, Any] = {}
|
||||
self._active_run_tasks: Dict[str, "asyncio.Task"] = {}
|
||||
# Stop is cooperative: the executor thread may outlive the HTTP request.
|
||||
self._stopping_run_ids: set[str] = set()
|
||||
# Pollable run status for dashboards and external control-plane UIs.
|
||||
self._run_statuses: Dict[str, Dict[str, Any]] = {}
|
||||
# Active approval session key for each run_id. The approval core
|
||||
|
|
@ -4304,6 +4308,8 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
def _text_cb(delta: Optional[str]) -> None:
|
||||
if delta is None:
|
||||
return
|
||||
if run_id not in self._run_streams:
|
||||
return
|
||||
try:
|
||||
loop.call_soon_threadsafe(q.put_nowait, {
|
||||
"event": "message.delta",
|
||||
|
|
@ -4328,6 +4334,18 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
async def _run_and_close():
|
||||
try:
|
||||
self._set_run_status(run_id, "running")
|
||||
if run_id in self._stopping_run_ids:
|
||||
q.put_nowait({
|
||||
"event": "run.cancelled",
|
||||
"run_id": run_id,
|
||||
"timestamp": time.time(),
|
||||
})
|
||||
self._set_run_status(
|
||||
run_id,
|
||||
"cancelled",
|
||||
last_event="run.cancelled",
|
||||
)
|
||||
return
|
||||
agent = self._create_agent(
|
||||
ephemeral_system_prompt=ephemeral_system_prompt,
|
||||
session_id=session_id,
|
||||
|
|
@ -4412,10 +4430,21 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
return r, u
|
||||
|
||||
result, usage = await asyncio.get_running_loop().run_in_executor(None, _run_sync)
|
||||
if run_id in self._stopping_run_ids:
|
||||
q.put_nowait({
|
||||
"event": "run.cancelled",
|
||||
"run_id": run_id,
|
||||
"timestamp": time.time(),
|
||||
})
|
||||
self._set_run_status(
|
||||
run_id,
|
||||
"cancelled",
|
||||
last_event="run.cancelled",
|
||||
)
|
||||
# Check for structured failure (non-retryable client errors like
|
||||
# 401/400 return failed=True instead of raising, so the except
|
||||
# block below never fires — issue #15561).
|
||||
if isinstance(result, dict) and result.get("failed"):
|
||||
elif isinstance(result, dict) and result.get("failed"):
|
||||
error_msg = _redact_api_error_text(result.get("error") or "agent run failed")
|
||||
q.put_nowait({
|
||||
"event": "run.failed",
|
||||
|
|
@ -4497,6 +4526,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
self._active_run_agents.pop(run_id, None)
|
||||
self._active_run_tasks.pop(run_id, None)
|
||||
self._run_approval_sessions.pop(run_id, None)
|
||||
self._stopping_run_ids.discard(run_id)
|
||||
|
||||
task = asyncio.create_task(_run_and_close())
|
||||
self._active_run_tasks[run_id] = task
|
||||
|
|
@ -4548,6 +4578,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
return web.json_response(_openai_error(f"Run not found: {run_id}", code="run_not_found"), status=404)
|
||||
|
||||
q = self._run_streams[run_id]
|
||||
self._run_stream_subscribers.add(run_id)
|
||||
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
|
|
@ -4575,6 +4606,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
except Exception as exc:
|
||||
logger.debug("[api_server] SSE stream error for run %s: %s", run_id, exc)
|
||||
finally:
|
||||
self._run_stream_subscribers.discard(run_id)
|
||||
self._run_streams.pop(run_id, None)
|
||||
self._run_streams_created.pop(run_id, None)
|
||||
|
||||
|
|
@ -4683,6 +4715,7 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
return web.json_response(_openai_error(f"Run not found: {run_id}", code="run_not_found"), status=404)
|
||||
|
||||
self._set_run_status(run_id, "stopping", last_event="run.stopping")
|
||||
self._stopping_run_ids.add(run_id)
|
||||
|
||||
if agent is not None:
|
||||
try:
|
||||
|
|
@ -4690,41 +4723,29 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
if task is not None and not task.done():
|
||||
task.cancel()
|
||||
# Bounded wait: run_conversation() executes in the default
|
||||
# executor thread which task.cancel() cannot preempt — we rely on
|
||||
# agent.interrupt() above to break the loop. Cap the wait so a
|
||||
# slow/unresponsive interrupt can't hang this handler.
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(task), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"[api_server] stop for run %s timed out after 5s; "
|
||||
"agent may still be finishing the current step",
|
||||
run_id,
|
||||
)
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
return web.json_response({"run_id": run_id, "status": "stopping"})
|
||||
|
||||
async def _sweep_orphaned_runs(self) -> None:
|
||||
"""Periodically clean up expired transport state for inactive runs."""
|
||||
"""Periodically expire transport buffers and terminal status records."""
|
||||
while True:
|
||||
await asyncio.sleep(60)
|
||||
self._sweep_orphaned_runs_once(time.time())
|
||||
|
||||
def _sweep_orphaned_runs_once(self, now: Optional[float] = None) -> None:
|
||||
"""Expire old SSE buffers without treating transport age as run age."""
|
||||
if now is None:
|
||||
now = time.time()
|
||||
stale = [
|
||||
run_id
|
||||
for run_id, created_at in list(self._run_streams_created.items())
|
||||
if now - created_at > self._RUN_STREAM_TTL
|
||||
and (
|
||||
(task := self._active_run_tasks.get(run_id)) is None
|
||||
or task.done()
|
||||
)
|
||||
]
|
||||
for run_id in stale:
|
||||
logger.debug("[api_server] sweeping orphaned run %s", run_id)
|
||||
stale = [
|
||||
run_id
|
||||
for run_id, created_at in list(self._run_streams_created.items())
|
||||
if now - created_at > self._RUN_STREAM_TTL
|
||||
and run_id not in self._run_stream_subscribers
|
||||
]
|
||||
for run_id in stale:
|
||||
logger.debug("[api_server] sweeping expired run transport %s", run_id)
|
||||
task = self._active_run_tasks.get(run_id)
|
||||
task_done = task is None or task.done()
|
||||
if task_done:
|
||||
try:
|
||||
from tools.approval import unregister_gateway_notify
|
||||
|
||||
|
|
@ -4733,20 +4754,24 @@ class APIServerAdapter(BasePlatformAdapter):
|
|||
unregister_gateway_notify(approval_session_key)
|
||||
except Exception:
|
||||
pass
|
||||
self._run_streams.pop(run_id, None)
|
||||
self._run_streams_created.pop(run_id, None)
|
||||
# The transport TTL always bounds buffering. Live control state is
|
||||
# independent and survives until the executor-backed task returns.
|
||||
self._run_streams.pop(run_id, None)
|
||||
self._run_streams_created.pop(run_id, None)
|
||||
if task_done:
|
||||
self._active_run_agents.pop(run_id, None)
|
||||
self._active_run_tasks.pop(run_id, None)
|
||||
self._run_approval_sessions.pop(run_id, None)
|
||||
self._stopping_run_ids.discard(run_id)
|
||||
|
||||
stale_statuses = [
|
||||
run_id
|
||||
for run_id, status in list(self._run_statuses.items())
|
||||
if status.get("status") in {"completed", "failed", "cancelled"}
|
||||
and now - float(status.get("updated_at", 0) or 0) > self._RUN_STATUS_TTL
|
||||
]
|
||||
for run_id in stale_statuses:
|
||||
self._run_statuses.pop(run_id, None)
|
||||
stale_statuses = [
|
||||
run_id
|
||||
for run_id, status in list(self._run_statuses.items())
|
||||
if status.get("status") in {"completed", "failed", "cancelled"}
|
||||
and now - float(status.get("updated_at", 0) or 0) > self._RUN_STATUS_TTL
|
||||
]
|
||||
for run_id in stale_statuses:
|
||||
self._run_statuses.pop(run_id, None)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# BasePlatformAdapter interface
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue