From ae723c70d3799bfa2244b5cd1a5efdd4e92c4714 Mon Sep 17 00:00:00 2001 From: Bosheng0422 Date: Mon, 31 Aug 2026 20:39:20 +0800 Subject: [PATCH 1/4] fix: schedule proactive care for agents that never saved config Proactive care defaults to on, but the scheduler only listed persisted enabled rows, so most experts never entered the push loop. Co-authored-by: Cursor --- src/octop/infra/agents/manager.py | 10 +++++ .../infra/db/repos/proactive_care_config.py | 41 +++++++++++------- src/octop/infra/proactive/scheduler.py | 25 +++++++++-- src/octop/infra/server.py | 1 + tests/unit/proactive/test_scheduler.py | 42 +++++++++++++++++++ 5 files changed, 101 insertions(+), 18 deletions(-) diff --git a/src/octop/infra/agents/manager.py b/src/octop/infra/agents/manager.py index 0f07f47c..5ebb897b 100644 --- a/src/octop/infra/agents/manager.py +++ b/src/octop/infra/agents/manager.py @@ -71,6 +71,7 @@ from octop.infra.cron.manager import CronManager from octop.infra.db.repos.agents import AgentRow from octop.infra.db.services import RepoBundle + from octop.infra.proactive.scheduler import ProactiveCareScheduler from octop.infra.utils.paths import PathLayout logger = logging.getLogger(__name__) @@ -333,6 +334,7 @@ def __init__( self._expert_catalog = expert_catalog self._plugin_manager = plugin_manager self._cron_manager: CronManager | None = None + self._proactive_scheduler: ProactiveCareScheduler | None = None self._team_processor: Any | None = None self._harness_manager: HarnessAgentManager | None = None self._lock = asyncio.Lock() @@ -403,6 +405,10 @@ def set_cron_manager(self, cron_manager: CronManager) -> None: """Attach the process-wide CronManager (must be set before boot()).""" self._cron_manager = cron_manager + def set_proactive_scheduler(self, scheduler: ProactiveCareScheduler) -> None: + """Attach the process-wide proactive-care scheduler (optional; used after create/delete).""" + self._proactive_scheduler = scheduler + def set_team_processor(self, team_processor: Any | None) -> None: """Attach harness TeamProcessor (GlobalProcessor); required before boot().""" self._team_processor = team_processor @@ -566,6 +572,8 @@ async def create(self, spec: AgentCreateSpec, *, defer_bootstrap: bool = False) self._repos.audit_repo.write( actor=ACTOR_SYSTEM, action="agent.create", target=agent_id, payload=spec.name ) + if self._proactive_scheduler is not None: + self._proactive_scheduler.ensure_scheduled(agent_id) return row def _preserve_system_files_path(self, agent_id: str, cfg: dict[str, Any]) -> dict[str, Any]: @@ -665,6 +673,8 @@ async def delete(self, agent_id: str) -> None: except OSError: logger.exception("rmtree failed for %s; agent removed from DB anyway", workspace_dir) self._repos.agent_repo.delete(agent_id) + if self._proactive_scheduler is not None: + self._proactive_scheduler.cancel(agent_id) self._repos.audit_repo.write(actor=ACTOR_SYSTEM, action="agent.delete", target=agent_id) async def start(self, agent_id: str) -> None: diff --git a/src/octop/infra/db/repos/proactive_care_config.py b/src/octop/infra/db/repos/proactive_care_config.py index 3214d1ca..c54a6945 100644 --- a/src/octop/infra/db/repos/proactive_care_config.py +++ b/src/octop/infra/db/repos/proactive_care_config.py @@ -92,25 +92,36 @@ def upsert(self, config: ProactiveCareConfig) -> None: ) def list_enabled(self) -> list[ProactiveCareConfig]: - """List the configurations of all agents that have proactive care push enabled. + """List proactive-care configs for every enabled agent that should be scheduled. + + Agents without a ``proactive_care_config`` row use the same defaults as + ``get()`` (enabled=True). An explicit ``enabled=0`` row opts out. + Disabled agents (``agents.enabled = 0``) are excluded. Returns: The list of enabled ProactiveCareConfig. """ with self._db.connect() as conn: rows = conn.execute( - "SELECT agent_id, enabled, active_hours_start, active_hours_end, " - "min_interval_hours, max_interval_hours " - "FROM proactive_care_config WHERE enabled = 1", + "SELECT a.agent_id, c.enabled, c.active_hours_start, " + "c.active_hours_end, c.min_interval_hours, c.max_interval_hours " + "FROM agents a " + "LEFT JOIN proactive_care_config c ON c.agent_id = a.agent_id " + "WHERE a.enabled = 1 AND (c.agent_id IS NULL OR c.enabled = 1)", ).fetchall() - return [ - ProactiveCareConfig( - agent_id=row[0], - enabled=bool(row[1]), - active_hours_start=row[2], - active_hours_end=row[3], - min_interval_hours=row[4], - max_interval_hours=row[5], - ) - for row in rows - ] + configs: list[ProactiveCareConfig] = [] + for row in rows: + if row[1] is None: + configs.append(ProactiveCareConfig(agent_id=row[0])) + else: + configs.append( + ProactiveCareConfig( + agent_id=row[0], + enabled=bool(row[1]), + active_hours_start=row[2], + active_hours_end=row[3], + min_interval_hours=row[4], + max_interval_hours=row[5], + ) + ) + return configs diff --git a/src/octop/infra/proactive/scheduler.py b/src/octop/infra/proactive/scheduler.py index 6c12cd60..24365a0a 100644 --- a/src/octop/infra/proactive/scheduler.py +++ b/src/octop/infra/proactive/scheduler.py @@ -147,11 +147,26 @@ def replace_persistence( self._care_service.replace_care_push_repo(care_push_repo) async def start_all(self) -> None: - """At system startup, register random scheduling tasks for all agents with enabled=true.""" + """At system startup, register random scheduling tasks for all agents with enabled=true. + + Includes agents that never saved a config row: proactive care defaults to on. + """ configs = self._config_repo.list_enabled() logger.info("ProactiveCareScheduler: started, found %d enabled agents", len(configs)) for cfg in configs: - self._schedule(cfg.agent_id) + self.ensure_scheduled(cfg.agent_id) + + def ensure_scheduled(self, agent_id: str) -> None: + """Start the loop if this agent is enabled and not already scheduled.""" + if agent_id in self._tasks: + return + cfg = self._config_repo.get(agent_id) + if cfg.enabled: + self._schedule(agent_id) + logger.info( + "ProactiveCareScheduler: agent=%s scheduled (default enabled unless opted out)", + agent_id, + ) def reschedule(self, agent_id: str) -> None: """Cancel the current schedule and re-arrange the next trigger time with new config. @@ -190,6 +205,9 @@ async def shutdown(self) -> None: def _schedule(self, agent_id: str) -> None: """Create a scheduling task for an agent.""" + existing = self._tasks.get(agent_id) + if existing is not None and not existing.done(): + existing.cancel() task = asyncio.create_task( self._run_loop(agent_id), name=f"proactive_care_{agent_id}", @@ -199,7 +217,8 @@ def _schedule(self, agent_id: str) -> None: def _on_task_done(self, agent_id: str, task: asyncio.Task[None]) -> None: """Task-completion callback that handles exceptions.""" - self._tasks.pop(agent_id, None) + if self._tasks.get(agent_id) is task: + self._tasks.pop(agent_id, None) if task.cancelled(): return exc = task.exception() diff --git a/src/octop/infra/server.py b/src/octop/infra/server.py index 916d2235..cd9e0d49 100644 --- a/src/octop/infra/server.py +++ b/src/octop/infra/server.py @@ -274,6 +274,7 @@ async def _boot_runtime(self, config: OctopConfig) -> None: config_repo=self.services.repos.proactive_care_config_repo, session_repo=self.services.repos.session_repo, ) + registry.set_proactive_scheduler(proactive_scheduler) await registry.boot() await gateway.refresh_media_backends() diff --git a/tests/unit/proactive/test_scheduler.py b/tests/unit/proactive/test_scheduler.py index 935ac3b7..ea228b79 100644 --- a/tests/unit/proactive/test_scheduler.py +++ b/tests/unit/proactive/test_scheduler.py @@ -231,6 +231,27 @@ def test_config_repo_list_enabled(db: SqlitePool, config_repo: ProactiveCareConf assert aid2 not in enabled_ids +def test_config_repo_list_enabled_includes_default_on_agents( + db: SqlitePool, config_repo: ProactiveCareConfigRepo +): + """Agents with no config row default to enabled and must be scheduled.""" + uid = UserRepo(db).create(username="carol", password_hash="h", role="admin") + aid_default = new_ulid() + aid_off = new_ulid() + AgentRepo(db).create(agent_id=aid_default, user_id=uid, name="default-on") + AgentRepo(db).create(agent_id=aid_off, user_id=uid, name="opted-out") + config_repo.upsert(ProactiveCareConfig(agent_id=aid_off, enabled=False)) + + enabled = config_repo.list_enabled() + enabled_ids = {c.agent_id for c in enabled} + assert aid_default in enabled_ids + assert aid_off not in enabled_ids + default_cfg = next(c for c in enabled if c.agent_id == aid_default) + assert default_cfg.enabled is True + assert default_cfg.min_interval_hours == 5 + assert default_cfg.max_interval_hours == 24 + + # --------------------------------------------------------------------------- # ProactiveCareScheduler tests # --------------------------------------------------------------------------- @@ -314,3 +335,24 @@ async def test_scheduler_start_all_no_enabled( ) await scheduler.start_all() assert len(scheduler._tasks) == 0 + + +@pytest.mark.asyncio +async def test_scheduler_start_all_includes_default_on_agent( + agent_id: str, + config_repo: ProactiveCareConfigRepo, + db: SqlitePool, +): + """start_all should schedule agents that never saved a proactive-care row.""" + from octop.infra.db.repos.sessions import SessionRepo + + care_service = AsyncMock() + scheduler = ProactiveCareScheduler( + care_service=care_service, + config_repo=config_repo, + session_repo=SessionRepo(db), + ) + await scheduler.start_all() + assert agent_id in scheduler._tasks + scheduler.cancel(agent_id) + await asyncio.sleep(0.01) From 45ebc3e507d835d0fd6de81bf7faae37f47f78c3 Mon Sep 17 00:00:00 2001 From: Bosheng0422 Date: Wed, 2 Sep 2026 01:56:35 +0800 Subject: [PATCH 2/4] fix: attach gateway connector tools without rebuilding the agent prepare_chat_mcp rebuilt the agent whenever a builtin MCP server was missing from the live runtime. For gateway-mode connectors that rebuild is pure loss: they carry no HTTP transport, their tools are built in-process from stored credentials, and _post_start_agent runs the very same injection at the end anyway. Meanwhile the rebuild drops the harness instance and with it the checkpointer pool an in-flight turn is still writing to. Split the missing names on connector mode: gateway ones refresh their credentials and inject into the running agent, the rest keep the existing reload path. Co-authored-by: Cursor --- src/octop/infra/agents/manager.py | 55 ++++++++++++++++++++++++ src/octop/infra/connectors/builder.py | 17 ++++++++ tests/unit/agents/test_mcp_tool_cache.py | 55 ++++++++++++++++++++++++ 3 files changed, 127 insertions(+) diff --git a/src/octop/infra/agents/manager.py b/src/octop/infra/agents/manager.py index 5ebb897b..7c6fdac3 100644 --- a/src/octop/infra/agents/manager.py +++ b/src/octop/infra/agents/manager.py @@ -53,6 +53,7 @@ ) from octop.infra.connectors.builder import ( build_mcp_server_configs_for_user, + gateway_mcp_server_names, inject_missing_gateway_tools, ) from octop.infra.connectors.service import ConnectorService @@ -1349,6 +1350,23 @@ async def prepare_chat_mcp( continue agent.config.mcp_server_configs[name] = dict(spec) + gateway_missing: list[str] = [] + if builtin_missing and uid is not None: + gateway_names = gateway_mcp_server_names( + connector_repo=self._repos.connector_repo, + user_id=uid, + ) + gateway_missing = [n for n in builtin_missing if n in gateway_names] + builtin_missing = [n for n in builtin_missing if n not in gateway_names] + + if gateway_missing and uid is not None: + await self._attach_gateway_tools( + agent, + agent_id=agent_id, + user_id=uid, + server_names=gateway_missing, + ) + if builtin_missing: logger.info( "Reloading agent %s MCP tools (builtin_missing=%s)", @@ -1431,6 +1449,43 @@ async def prepare_chat_mcp( ) return still_missing + async def _attach_gateway_tools( + self, + agent: HarnessAgent, + *, + agent_id: str, + user_id: int, + server_names: list[str], + ) -> None: + """Add gateway connector tools to the running agent without rebuilding it. + + A rebuild would drop the harness instance (and with it the checkpointer + pool an in-flight turn still writes to) only to run the very same + in-process injection at the end of ``_post_start_agent``. + """ + repo = self._repos.connector_repo + for inst in repo.list_by_user(user_id): + if inst.status != "active" or inst.mcp_server_name not in server_names: + continue + try: + await self._connector_svc.ensure_fresh_credentials(inst.instance_id, inst.kind) + except Exception: + logger.exception( + "prepare_chat_mcp agent=%s: credential refresh failed for %s", + agent_id, + inst.mcp_server_name, + ) + inject_missing_gateway_tools( + agent, + svc=self._connector_svc, + connector_repo=repo, + user_id=user_id, + agent_id=agent_id, + mcp_server_configs=agent.config.mcp_server_configs, + ) + for name in server_names: + agent.config.mcp_server_configs.setdefault(name, {}) + # ------------------------------------------------------------------ # Settings persistence — push global policy into harness runtime # ------------------------------------------------------------------ diff --git a/src/octop/infra/connectors/builder.py b/src/octop/infra/connectors/builder.py index a7521c8e..d20bebb9 100644 --- a/src/octop/infra/connectors/builder.py +++ b/src/octop/infra/connectors/builder.py @@ -529,6 +529,23 @@ def build_mcp_server_configs_for_user( return configs +def gateway_mcp_server_names(*, connector_repo: Any, user_id: int) -> set[str]: + """MCP server names of *user_id*'s active gateway-mode connector instances. + + Gateway connectors carry no HTTP transport: their tools are built in-process + from stored credentials, so callers can attach them to a live agent instead + of rebuilding it. + """ + names: set[str] = set() + for inst in connector_repo.list_by_user(user_id): + if inst.status != "active": + continue + entry = get_catalog_entry(inst.kind) + if entry is not None and entry.mcp_mode == "gateway": + names.add(inst.mcp_server_name) + return names + + def inject_missing_gateway_tools( agent: Any, *, diff --git a/tests/unit/agents/test_mcp_tool_cache.py b/tests/unit/agents/test_mcp_tool_cache.py index efd6c0c0..52bc8f51 100644 --- a/tests/unit/agents/test_mcp_tool_cache.py +++ b/tests/unit/agents/test_mcp_tool_cache.py @@ -195,6 +195,61 @@ async def test_prepare_chat_mcp_injects_custom_from_cache() -> None: assert isinstance(injected[0], StructuredTool) +@pytest.mark.asyncio +async def test_prepare_chat_mcp_attaches_gateway_tools_without_reload() -> None: + """Gateway connectors inject in-process; rebuilding would drop the runtime. + + A rebuild ends in the very same injection, but first closes the harness + agent (and the checkpointer pool an in-flight turn still writes to) and + holds ``AgentManager._lock`` across workspace I/O. + """ + from octop.infra.agents.manager import AgentManager + from octop.infra.connectors.builder import mcp_server_name + + instance_id = "01TESTGATEWAYINSTANCE0001" + server = mcp_server_name("qq-mail", instance_id) + + inst = MagicMock() + inst.instance_id = instance_id + inst.kind = "qq-mail" + inst.status = "active" + inst.mcp_server_name = server + + agent = MagicMock() + agent._mcp_tools = [] + agent._mcp_tool_name_set = frozenset() + agent.config.mcp_server_configs = {} + + def _inject(tools: list[Any]) -> None: + agent._mcp_tools = [*tools, *agent._mcp_tools] + agent._mcp_tool_name_set = frozenset(str(getattr(t, "name", "")) for t in agent._mcp_tools) + + agent.inject_mcp_tools = MagicMock(side_effect=_inject) + + mgr = object.__new__(AgentManager) + mgr._repos = MagicMock() + mgr._repos.connector_repo.list_by_user.return_value = [inst] + mgr._connector_svc = MagicMock() + mgr._connector_svc.custom_harness_configs.return_value = {} + mgr._connector_svc.decrypt.return_value = {"email": "a@qq.com", "password": "code"} + mgr._connector_svc.ensure_fresh_credentials = AsyncMock(return_value={}) + mgr.get_agent = MagicMock(return_value=agent) + mgr.get_row = MagicMock(return_value=MagicMock()) + mgr._connector_uid_for = MagicMock(return_value=7) + mgr.reload_connectors = AsyncMock() + + failed = await mgr.prepare_chat_mcp("A1", [server], connector_user_id=7) + + # Asserted before the tool checks so a reintroduced rebuild reports itself + # rather than surfacing as "tools missing". + mgr.reload_connectors.assert_not_awaited() + assert failed == [] + assert agent.inject_mcp_tools.call_count == 1 + assert all(name.startswith(f"{server}_") for name in agent._mcp_tool_name_set) + assert server in agent.config.mcp_server_configs + mgr._connector_svc.ensure_fresh_credentials.assert_awaited_once_with(instance_id, "qq-mail") + + def test_wrap_tools_for_shared_use() -> None: lock = asyncio.Lock() inner = MagicMock() From 225a27a780d749ae3d51513086f0681b6e59b2e6 Mon Sep 17 00:00:00 2001 From: Bosheng0422 Date: Wed, 2 Sep 2026 01:57:03 +0800 Subject: [PATCH 3/4] fix: await cancellation on scheduler shutdown, and keep its sleeps out of tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ProactiveCareScheduler.shutdown() cancelled each task and returned without awaiting it, so a loop parked in a multi-hour sleep could outlive the shutdown that was supposed to end it. It now gathers the cancelled tasks before returning. Proactive care defaults to on, so creating an agent starts one of those sleeps. pytest-asyncio waits for leftover tasks before fixture teardown, which hangs the suite. octop_client shuts the scheduler down and calls the new suspend() so nothing reschedules; an autouse fixture covers tests that boot OctopServer another way. The scheduler's own unit tests opt out. suspend() exists for the tests. The alternative — every test remembering to tear the scheduler down — is the arrangement that produced the hang. Co-authored-by: Cursor --- src/octop/infra/proactive/scheduler.py | 22 ++++++++++++--- tests/conftest.py | 26 +++++++++++++++++ tests/support/app.py | 6 ++++ tests/unit/proactive/test_scheduler.py | 39 ++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 4 deletions(-) diff --git a/src/octop/infra/proactive/scheduler.py b/src/octop/infra/proactive/scheduler.py index 24365a0a..8d8e5273 100644 --- a/src/octop/infra/proactive/scheduler.py +++ b/src/octop/infra/proactive/scheduler.py @@ -133,6 +133,7 @@ def __init__( self._session_repo = session_repo # agent_id -> asyncio.Task self._tasks: dict[str, asyncio.Task[None]] = {} + self._suspended = False def replace_persistence( self, @@ -156,8 +157,14 @@ async def start_all(self) -> None: for cfg in configs: self.ensure_scheduled(cfg.agent_id) + def suspend(self) -> None: + """Stop creating new loops. Tests call this so leftover sleeps do not hang pytest.""" + self._suspended = True + def ensure_scheduled(self, agent_id: str) -> None: """Start the loop if this agent is enabled and not already scheduled.""" + if self._suspended: + return if agent_id in self._tasks: return cfg = self._config_repo.get(agent_id) @@ -197,14 +204,21 @@ def cancel(self, agent_id: str) -> None: task.cancel() async def shutdown(self) -> None: - """Shut down all scheduling tasks.""" - agent_ids = list(self._tasks.keys()) - for agent_id in agent_ids: - self.cancel(agent_id) + """Cancel every loop and wait so a long sleep cannot outlive the process.""" + pending: list[asyncio.Task[None]] = [] + for agent_id in list(self._tasks.keys()): + task = self._tasks.pop(agent_id, None) + if task and not task.done(): + task.cancel() + pending.append(task) + if pending: + await asyncio.gather(*pending, return_exceptions=True) logger.info("ProactiveCareScheduler: all scheduled tasks shut down") def _schedule(self, agent_id: str) -> None: """Create a scheduling task for an agent.""" + if self._suspended: + return existing = self._tasks.get(agent_id) if existing is not None and not existing.done(): existing.cancel() diff --git a/tests/conftest.py b/tests/conftest.py index 8d37cb87..bfd0bf9b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ from collections.abc import Iterator from pathlib import Path +from typing import Any import pytest @@ -43,6 +44,31 @@ def repo_root() -> Path: return _REPO_ROOT +# Creating an agent now starts a multi-hour asyncio.sleep (proactive care +# defaults to on). pytest-asyncio waits for leftover tasks before fixture +# teardown, so any test that boots OctopServer without going through +# ``octop_client`` would hang the suite. Scheduler unit tests opt out. +_PROACTIVE_SCHEDULER_TESTS = "tests/unit/proactive/test_scheduler.py" + + +@pytest.fixture(autouse=True) +def _suspend_proactive_care_loops( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if _module_path(request.node.nodeid) == _PROACTIVE_SCHEDULER_TESTS: + return + from octop.infra.proactive.scheduler import ProactiveCareScheduler + + monkeypatch.setattr(ProactiveCareScheduler, "ensure_scheduled", lambda self, _id: None) + + async def _start_all(self: Any) -> None: + return None + + monkeypatch.setattr(ProactiveCareScheduler, "start_all", _start_all) + monkeypatch.setattr(ProactiveCareScheduler, "_schedule", lambda self, _id: None) + + @pytest.fixture(autouse=True) def _isolated_user_home( tmp_path_factory: pytest.TempPathFactory, diff --git a/tests/support/app.py b/tests/support/app.py index bf8e57c8..c6e061be 100644 --- a/tests/support/app.py +++ b/tests/support/app.py @@ -54,6 +54,12 @@ async def octop_client( await srv.start() if bind_database and not srv.database_bound: await ensure_control_plane_bound(srv) + # Creating an agent now starts a random-interval sleep (default ON). + # pytest-asyncio waits for leftover tasks before fixture teardown, so + # those sleeps hang the suite. Production shutdown still cancels them. + if srv.app_runtime is not None: + await srv.app_runtime.proactive_scheduler.shutdown() + srv.app_runtime.proactive_scheduler.suspend() app = build_app(srv) try: async with httpx.AsyncClient( diff --git a/tests/unit/proactive/test_scheduler.py b/tests/unit/proactive/test_scheduler.py index ea228b79..c865fb9a 100644 --- a/tests/unit/proactive/test_scheduler.py +++ b/tests/unit/proactive/test_scheduler.py @@ -356,3 +356,42 @@ async def test_scheduler_start_all_includes_default_on_agent( assert agent_id in scheduler._tasks scheduler.cancel(agent_id) await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_long_sleep_quickly( + agent_id: str, + config_repo: ProactiveCareConfigRepo, + db: SqlitePool, +) -> None: + """A default-on loop sleeps for hours; shutdown must not wait that out.""" + from octop.infra.db.repos.sessions import SessionRepo + + scheduler = ProactiveCareScheduler( + care_service=AsyncMock(), + config_repo=config_repo, + session_repo=SessionRepo(db), + ) + scheduler.ensure_scheduled(agent_id) + assert agent_id in scheduler._tasks + await asyncio.wait_for(scheduler.shutdown(), timeout=2) + assert scheduler._tasks == {} + + +@pytest.mark.asyncio +async def test_suspend_skips_new_schedules( + agent_id: str, + config_repo: ProactiveCareConfigRepo, + db: SqlitePool, +) -> None: + from octop.infra.db.repos.sessions import SessionRepo + + scheduler = ProactiveCareScheduler( + care_service=AsyncMock(), + config_repo=config_repo, + session_repo=SessionRepo(db), + ) + scheduler.suspend() + scheduler.ensure_scheduled(agent_id) + await scheduler.start_all() + assert scheduler._tasks == {} From 1a432f282682a04a9ac1fccef834a68852108b47 Mon Sep 17 00:00:00 2001 From: Bosheng0422 Date: Wed, 2 Sep 2026 02:02:01 +0800 Subject: [PATCH 4/4] refactor: drive websockets on one event loop and drop the cross-loop workaround MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit starlette's sync TestClient runs the app on its own anyio portal loop. That left the process with two loops over one OctopServer, so primitives bound to a loop and shared by both — AgentManager._lock — deadlocked or raised "bound to a different event loop". The tests worked around it by running the ws session in a worker thread. tests.support.http now speaks ASGI directly, so the handler, the gateway workers and the test all share one loop, the way uvicorn runs it in production. That removes the reason the production handlers marshalled every outbound frame across loops with run_coroutine_threadsafe + wrap_future. The comment on that code named the cause outright — "this handler may run on a different loop (e.g. starlette's TestClient portal)" — so it was test-shaped machinery sitting on the path of every frame the dashboard receives. Both handlers now await the send directly. Co-authored-by: Cursor --- src/octop/api/routers/chat/notify_ws.py | 10 +- src/octop/api/routers/chat/ws.py | 12 +- tests/integration/test_chat_ws.py | 275 +++++---------------- tests/integration/test_e2e_golden_path.py | 36 +-- tests/integration/test_notifications_ws.py | 78 ++---- tests/support/app.py | 5 +- tests/support/http.py | 160 +++++++++--- 7 files changed, 228 insertions(+), 348 deletions(-) diff --git a/src/octop/api/routers/chat/notify_ws.py b/src/octop/api/routers/chat/notify_ws.py index 2b3f9cec..39d2df36 100644 --- a/src/octop/api/routers/chat/notify_ws.py +++ b/src/octop/api/routers/chat/notify_ws.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import contextlib import json import logging @@ -46,20 +45,13 @@ async def dashboard_notifications_ws( connection_id = uuid.uuid4().hex await websocket.accept() - # Harness/gateway workers run on the server loop; this handler may run on - # another (e.g. Starlette TestClient). Marshal frames onto the socket loop. - ws_loop = asyncio.get_running_loop() - - async def _emit_frame(frame: dict[str, Any]) -> None: + async def send_frame(frame: dict[str, Any]) -> None: if websocket.application_state != WebSocketState.CONNECTED: return await websocket.send_text( json.dumps(frame, ensure_ascii=False, default=json_chunk_default), ) - async def send_frame(frame: dict[str, Any]) -> None: - await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(_emit_frame(frame), ws_loop)) - hub.register(connection_id, send_frame, user_id=user.id) try: diff --git a/src/octop/api/routers/chat/ws.py b/src/octop/api/routers/chat/ws.py index 02fcb790..910c3fda 100644 --- a/src/octop/api/routers/chat/ws.py +++ b/src/octop/api/routers/chat/ws.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import contextlib import json import logging @@ -68,22 +67,13 @@ async def dashboard_chat_ws( connection_id = uuid.uuid4().hex await websocket.accept() - # The harness/gateway workers run on the server event loop, but this - # handler may run on a different loop (e.g. starlette's TestClient portal). - # Marshal outbound frames onto the loop that owns this WebSocket so - # cross-loop sends don't deadlock. - ws_loop = asyncio.get_running_loop() - - async def _emit_frame(frame: dict[str, Any]) -> None: + async def send_frame(frame: dict[str, Any]) -> None: if websocket.application_state != WebSocketState.CONNECTED: return await websocket.send_text( json.dumps(frame, ensure_ascii=False, default=json_chunk_default), ) - async def send_frame(frame: dict[str, Any]) -> None: - await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(_emit_frame(frame), ws_loop)) - hub.register(connection_id, send_frame) try: diff --git a/tests/integration/test_chat_ws.py b/tests/integration/test_chat_ws.py index cac5a6df..495f3d15 100644 --- a/tests/integration/test_chat_ws.py +++ b/tests/integration/test_chat_ws.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio -import json from collections.abc import AsyncIterator from pathlib import Path from typing import Any @@ -11,7 +10,6 @@ import httpx import pytest -from starlette.testclient import TestClient from starlette.websockets import WebSocketDisconnect from tests.support.app import octop_client @@ -23,7 +21,11 @@ seed_openai_provider, ) from tests.support.fakes import FakeHarnessAgent -from tests.support.http import ws_token +from tests.support.http import chat_ws_path, ws_connect + + +def _chat_ws(c: httpx.AsyncClient, aid: str, auth: dict[str, str]) -> Any: + return ws_connect(c._octop_app, chat_ws_path(aid, auth)) # type: ignore[attr-defined] @pytest.fixture @@ -43,80 +45,38 @@ async def env(tmp_octop_home: Path) -> AsyncIterator[Any]: yield c, srv, fake, users["alice"], users["bob"], aid -def _consume_ws_turn_sync( - app: object, - aid: str, - token: str, - body: dict[str, Any], -) -> list[dict[str, Any]]: - chunks: list[dict[str, Any]] = [] - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws: - ws.send_json(body) - while True: - raw = ws.receive_text() - chunk = json.loads(raw) - chunks.append(chunk) - if chunk.get("type") in ("done", "error"): - break - return chunks - - -def _disconnect_ws_turn_sync( - app: object, - aid: str, - token: str, -) -> None: - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws: - ws.send_json({"type": "user_turn", "text": "cancel me"}) - ws.receive_text() - - -def _turn_then_rebind_sync( - app: object, +async def _turn_then_rebind( + c: httpx.AsyncClient, aid: str, - token: str, + auth: dict[str, str], thread_id: str, - gate_release: Any, + gate_release: asyncio.Event, ) -> list[dict[str, Any]]: """Start a turn on conn A, drop it after the first token, re-subscribe on B. The fake stream is gated between first and second token so B can subscribe while the turn is still active, then receive the remaining chunks. """ - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_a: - ws_a.send_json( + async with _chat_ws(c, aid, auth) as ws_a: + await ws_a.send_json( { "type": "user_turn", "text": "slow please", "thread_id": thread_id, } ) - first = json.loads(ws_a.receive_text()) + first = await ws_a.receive_json() assert first.get("type") == "token" assert first.get("content") == "first" # A is closed; turn is blocked on gate_release (still active, not cancelled). frames: list[dict[str, Any]] = [] - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_b: - ws_b.send_json({"type": "subscribe", "thread_id": thread_id}) - status = json.loads(ws_b.receive_text()) - frames.append(status) + async with _chat_ws(c, aid, auth) as ws_b: + await ws_b.send_json({"type": "subscribe", "thread_id": thread_id}) + frames.append(await ws_b.receive_json()) # Release the slow stream only after B is subscribed. gate_release.set() - for _ in range(50): - raw = ws_b.receive_text() - frame = json.loads(raw) - frames.append(frame) - if frame.get("type") in ("done", "error"): - break + frames.extend(await ws_b.drain_turn()) return frames @@ -138,14 +98,7 @@ async def slow_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: cancel_spy = MagicMock(wraps=srv.app_runtime.agent_registry.cancel_stream) srv.app_runtime.agent_registry.cancel_stream = cancel_spy - frames = await asyncio.to_thread( - _turn_then_rebind_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid, - gate, - ) + frames = await _turn_then_rebind(c, aid, alice_auth, tid, gate) cancel_spy.assert_not_called() assert frames[0] == {"type": "turn_status", "thread_id": tid, "active": True} @@ -154,63 +107,46 @@ async def slow_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: assert frames[-1].get("type") == "done" -def _turn_with_second_subscriber_sync( - app: object, +async def _turn_with_second_subscriber( + c: httpx.AsyncClient, aid: str, - token: str, + auth: dict[str, str], thread_id: str, - gate_release: Any, + gate_release: asyncio.Event, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """Keep conn A open, subscribe conn B mid-turn; both receive remaining chunks.""" a_frames: list[dict[str, Any]] = [] b_frames: list[dict[str, Any]] = [] - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_a: - ws_a.send_json( + async with _chat_ws(c, aid, auth) as ws_a: + await ws_a.send_json( { "type": "user_turn", "text": "slow please", "thread_id": thread_id, } ) - a_frames.append(json.loads(ws_a.receive_text())) - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_b: - ws_b.send_json({"type": "subscribe", "thread_id": thread_id}) - b_frames.append(json.loads(ws_b.receive_text())) + a_frames.append(await ws_a.receive_json()) + async with _chat_ws(c, aid, auth) as ws_b: + await ws_b.send_json({"type": "subscribe", "thread_id": thread_id}) + b_frames.append(await ws_b.receive_json()) gate_release.set() - for _ in range(50): - frame = json.loads(ws_a.receive_text()) - a_frames.append(frame) - if frame.get("type") in ("done", "error"): - break - for _ in range(50): - frame = json.loads(ws_b.receive_text()) - b_frames.append(frame) - if frame.get("type") in ("done", "error"): - break + rest_a, rest_b = await asyncio.gather(ws_a.drain_turn(), ws_b.drain_turn()) + a_frames.extend(rest_a) + b_frames.extend(rest_b) return a_frames, b_frames -def _two_threads_turn_sync( - app: object, +async def _two_threads_turn( + c: httpx.AsyncClient, aid: str, - token: str, + auth: dict[str, str], tid_a: str, tid_b: str, *, text_a: str = "one", text_b: str = "two", ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Run two live turns on different threads from one TestClient thread. - - Each sync ``TestClient`` spins a blocking anyio portal; driving two portals - from separate ``asyncio.to_thread`` workers deadlocks the server's - ``run_coroutine_threadsafe`` outbound path under pytest-asyncio. Interleaving - receives on one thread still exercises concurrent per-thread routing. - """ + """Run two live turns on different threads over two concurrent sockets.""" body_a: dict[str, Any] = { "type": "user_turn", "text": text_a, @@ -223,31 +159,10 @@ def _two_threads_turn_sync( "messages": [{"role": "user", "content": text_b}], "thread_id": tid_b, } - a_frames: list[dict[str, Any]] = [] - b_frames: list[dict[str, Any]] = [] - with ( - TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_a, - TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws_b, - ): - ws_a.send_json(body_a) - ws_b.send_json(body_b) - pending = {"a": True, "b": True} - while pending["a"] or pending["b"]: - if pending["a"]: - frame = json.loads(ws_a.receive_text()) - a_frames.append(frame) - if frame.get("type") in ("done", "error"): - pending["a"] = False - if pending["b"]: - frame = json.loads(ws_b.receive_text()) - b_frames.append(frame) - if frame.get("type") in ("done", "error"): - pending["b"] = False - return a_frames, b_frames + async with _chat_ws(c, aid, auth) as ws_a, _chat_ws(c, aid, auth) as ws_b: + await ws_a.send_json(body_a) + await ws_b.send_json(body_b) + return await asyncio.gather(ws_a.drain_turn(), ws_b.drain_turn()) # type: ignore[return-value] async def test_ws_concurrent_subscribers_both_receive_later_chunks(env: Any) -> None: @@ -266,14 +181,7 @@ async def slow_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: agent.stream = slow_stream - a_frames, b_frames = await asyncio.to_thread( - _turn_with_second_subscriber_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid, - gate, - ) + a_frames, b_frames = await _turn_with_second_subscriber(c, aid, alice_auth, tid, gate) assert a_frames[0].get("content") == "first" assert [f.get("content") for f in a_frames if f.get("type") == "token"] == ["first", "second"] @@ -298,14 +206,7 @@ async def tagged_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any] agent.stream = tagged_stream - a_frames, b_frames = await asyncio.to_thread( - _two_threads_turn_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid_a, - tid_b, - ) + a_frames, b_frames = await _two_threads_turn(c, aid, alice_auth, tid_a, tid_b) a_tokens = [f.get("content") for f in a_frames if f.get("type") == "token"] b_tokens = [f.get("content") for f in b_frames if f.get("type") == "token"] @@ -334,15 +235,9 @@ async def _consume_ws_turn( if extra: body.update(extra) - # Run the blocking starlette TestClient session in a worker thread so the - # server's event loop stays free to process the turn (see ws_chat_turn). - return await asyncio.to_thread( - _consume_ws_turn_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(auth), - body, - ) + async with _chat_ws(c, aid, auth) as ws: + await ws.send_json(body) + return await ws.drain_turn() async def test_ws_emits_chunks_then_done(env: Any) -> None: @@ -367,28 +262,23 @@ async def slow_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: cancel_spy = MagicMock(wraps=original_cancel) srv.app_runtime.agent_registry.cancel_stream = cancel_spy - await asyncio.to_thread( - _disconnect_ws_turn_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - ) + async with _chat_ws(c, aid, alice_auth) as ws: + await ws.send_json({"type": "user_turn", "text": "cancel me"}) + await ws.receive_json() await asyncio.sleep(0.05) cancel_spy.assert_not_called() -def _subscribe_ws_sync( - app: object, +async def _subscribe_ws( + c: httpx.AsyncClient, aid: str, - token: str, + auth: dict[str, str], thread_id: str, ) -> dict[str, Any]: - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws: - ws.send_json({"type": "subscribe", "thread_id": thread_id}) - return json.loads(ws.receive_text()) + async with _chat_ws(c, aid, auth) as ws: + await ws.send_json({"type": "subscribe", "thread_id": thread_id}) + return await ws.receive_json() async def test_ws_subscribe_turn_status_idle(env: Any) -> None: @@ -397,13 +287,7 @@ async def test_ws_subscribe_turn_status_idle(env: Any) -> None: assert create.status_code == 201 tid = create.json()["thread_id"] - frame = await asyncio.to_thread( - _subscribe_ws_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid, - ) + frame = await _subscribe_ws(c, aid, alice_auth, tid) assert frame == {"type": "turn_status", "thread_id": tid, "active": False} @@ -420,30 +304,10 @@ async def test_ws_subscribe_rejects_another_users_thread(env: Any) -> None: assert response.status_code == 201, response.text tid = response.json()["thread_id"] - frame = await asyncio.to_thread( - _subscribe_ws_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid, - ) + frame = await _subscribe_ws(c, aid, alice_auth, tid) assert frame == {"type": "error", "message": f"thread {tid!r} not found"} -def _cancel_ws_turn_sync( - app: object, - aid: str, - token: str, - thread_id: str, -) -> None: - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={token}" - ) as ws: - ws.send_json({"type": "user_turn", "text": "cancel me", "thread_id": thread_id}) - ws.receive_text() # first token - ws.send_json({"type": "cancel", "thread_id": thread_id}) - - async def test_ws_cancel_frame_cancels_active_turn(env: Any) -> None: c, srv, _fake, alice_auth, _bob_auth, aid = env agent = srv.app_runtime.agent_registry.get_agent(aid) @@ -459,13 +323,10 @@ async def slow_stream(request: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: cancel_spy = MagicMock(wraps=original_cancel) srv.app_runtime.agent_registry.cancel_stream = cancel_spy - await asyncio.to_thread( - _cancel_ws_turn_sync, - c._octop_app, # type: ignore[attr-defined] - aid, - ws_token(alice_auth), - tid, - ) + async with _chat_ws(c, aid, alice_auth) as ws: + await ws.send_json({"type": "user_turn", "text": "cancel me", "thread_id": tid}) + await ws.receive_json() # first token + await ws.send_json({"type": "cancel", "thread_id": tid}) for _ in range(40): if cancel_spy.called: @@ -487,24 +348,16 @@ async def test_ws_emits_error_frame_on_exception(env: Any) -> None: async def test_ws_bad_agent_rejected(env: Any) -> None: c, _srv, _fake, alice_auth, _bob_auth, _aid = env - with ( - pytest.raises(WebSocketDisconnect), - TestClient(c._octop_app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/01HMISSING0000000000000000/chat/ws?token={ws_token(alice_auth)}" - ), - ): - pass + with pytest.raises(WebSocketDisconnect): + async with _chat_ws(c, "01HMISSING0000000000000000", alice_auth): + pass async def test_ws_cross_user_rejected(env: Any) -> None: c, _srv, _fake, _admin_auth, bob_auth, aid = env - with ( - pytest.raises(WebSocketDisconnect), - TestClient(c._octop_app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{aid}/chat/ws?token={ws_token(bob_auth)}" - ), - ): - pass + with pytest.raises(WebSocketDisconnect): + async with _chat_ws(c, aid, bob_auth): + pass async def test_ws_accepts_skills_and_model(env: Any) -> None: diff --git a/tests/integration/test_e2e_golden_path.py b/tests/integration/test_e2e_golden_path.py index 4cd5a2b4..a223005d 100644 --- a/tests/integration/test_e2e_golden_path.py +++ b/tests/integration/test_e2e_golden_path.py @@ -2,15 +2,12 @@ from __future__ import annotations -import asyncio -import json from typing import Any import pytest -from starlette.testclient import TestClient from tests.support.auth import bootstrap_admin -from tests.support.http import ws_token +from tests.support.http import ws_chat_turn @pytest.fixture @@ -18,37 +15,10 @@ async def env(env_fake_harness): yield env_fake_harness -def _ws_turn_sync(app: object, agent_id: str, token: str, text: str) -> list[dict[str, Any]]: - chunks: list[dict[str, Any]] = [] - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{agent_id}/chat/ws?token={token}" - ) as ws: - ws.send_json( - { - "type": "user_turn", - "text": text, - "messages": [{"role": "user", "content": text}], - } - ) - while True: - raw = ws.receive_text() - chunk = json.loads(raw) - chunks.append(chunk) - if chunk.get("type") in ("done", "error"): - break - return chunks - - async def _ws_turn( - client: object, agent_id: str, auth: dict[str, str], *, text: str = "Hello" + client: Any, agent_id: str, auth: dict[str, str], *, text: str = "Hello" ) -> list[dict[str, Any]]: - return await asyncio.to_thread( - _ws_turn_sync, - client._octop_app, - agent_id, - ws_token(auth), - text, # type: ignore[attr-defined] - ) + return await ws_chat_turn(client, agent_id, auth, text=text) async def test_full_golden_path(env: Any) -> None: diff --git a/tests/integration/test_notifications_ws.py b/tests/integration/test_notifications_ws.py index c62c1335..81c83b7f 100644 --- a/tests/integration/test_notifications_ws.py +++ b/tests/integration/test_notifications_ws.py @@ -2,20 +2,16 @@ from __future__ import annotations -import asyncio -import json -import threading from collections.abc import AsyncIterator from pathlib import Path from typing import Any import pytest -from starlette.testclient import TestClient from starlette.websockets import WebSocketDisconnect from tests.support.app import octop_client from tests.support.auth import auth_header, bootstrap_admin, ensure_users -from tests.support.http import ws_token +from tests.support.http import ws_connect, ws_token @pytest.fixture @@ -27,69 +23,47 @@ async def env(tmp_octop_home: Path) -> AsyncIterator[Any]: yield c, srv, users["alice"], users["bob"] +def _notifications_ws(c: Any, auth: dict[str, str] | None = None) -> Any: + query = f"?token={ws_token(auth)}" if auth else "" + return ws_connect(c._octop_app, f"/api/notifications/ws{query}") + + async def test_notifications_ws_ping_pong(env: Any) -> None: c, _srv, alice_auth, _bob_auth = env - with TestClient(c._octop_app).websocket_connect( # type: ignore[attr-defined] - f"/api/notifications/ws?token={ws_token(alice_auth)}" - ) as ws: - ws.send_json({"type": "ping"}) - assert json.loads(ws.receive_text()) == {"type": "pong"} + async with _notifications_ws(c, alice_auth) as ws: + await ws.send_json({"type": "ping"}) + assert await ws.receive_json() == {"type": "pong"} async def test_notifications_ws_missing_token_rejected(env: Any) -> None: c, _srv, _alice_auth, _bob_auth = env - with ( - pytest.raises(WebSocketDisconnect), - TestClient(c._octop_app).websocket_connect("/api/notifications/ws"), # type: ignore[attr-defined] - ): - pass + with pytest.raises(WebSocketDisconnect): + async with _notifications_ws(c): + pass async def test_notifications_ws_receives_user_push(env: Any) -> None: c, srv, alice_auth, _bob_auth = env user = srv.user_manager.get("alice") assert user is not None - token = ws_token(alice_auth) - connected = threading.Event() - frames: list[dict[str, object]] = [] - errors: list[BaseException] = [] - def _session() -> None: - try: - with TestClient(c._octop_app).websocket_connect( # type: ignore[attr-defined] - f"/api/notifications/ws?token={token}" - ) as ws: - ws.send_json({"type": "ping"}) - pong = json.loads(ws.receive_text()) - if pong != {"type": "pong"}: - raise AssertionError(f"unexpected pong: {pong!r}") - connected.set() - frames.append(json.loads(ws.receive_text())) - except BaseException as exc: - errors.append(exc) - connected.set() - - session_task = asyncio.create_task(asyncio.to_thread(_session)) - assert await asyncio.to_thread(connected.wait, 5) - if errors: - raise errors[0] - await srv.app_runtime.gateway.ws_hub.push_to_user( - user.id, - { - "type": "dashboard_push", - "agent_id": "a1", - "thread_id": "thr_1", - "text": "记得喝水", - "agent_name": "助手", - }, - ) - await asyncio.wait_for(session_task, timeout=5) - assert frames == [ - { + async with _notifications_ws(c, alice_auth) as ws: + await ws.send_json({"type": "ping"}) + assert await ws.receive_json() == {"type": "pong"} + await srv.app_runtime.gateway.ws_hub.push_to_user( + user.id, + { + "type": "dashboard_push", + "agent_id": "a1", + "thread_id": "thr_1", + "text": "记得喝水", + "agent_name": "助手", + }, + ) + assert await ws.receive_json() == { "type": "dashboard_push", "agent_id": "a1", "thread_id": "thr_1", "text": "记得喝水", "agent_name": "助手", } - ] diff --git a/tests/support/app.py b/tests/support/app.py index c6e061be..454844a0 100644 --- a/tests/support/app.py +++ b/tests/support/app.py @@ -66,8 +66,9 @@ async def octop_client( transport=httpx.ASGITransport(app=app), base_url="http://testserver", ) as client: - # Expose the ASGI app so tests can open WebSocket sessions via - # starlette's sync TestClient (httpx removed AsyncClient.websocket_connect). + # Expose the ASGI app so tests can open WebSocket sessions on + # this same event loop (tests.support.http.ws_connect); httpx + # removed AsyncClient.websocket_connect. client._octop_app = app # type: ignore[attr-defined] yield client, srv finally: diff --git a/tests/support/http.py b/tests/support/http.py index ff17e76b..1ade227e 100644 --- a/tests/support/http.py +++ b/tests/support/http.py @@ -3,39 +3,144 @@ from __future__ import annotations import asyncio +import contextlib import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any +from urllib.parse import unquote, urlsplit import httpx -from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +# A single-loop session should never stall; the cap keeps a regression from +# wedging the whole suite instead of failing this one test. +_WS_RECEIVE_TIMEOUT_S = 20.0 def ws_token(auth: dict[str, str]) -> str: return auth["Authorization"].split(" ", 1)[1] -def _consume_ws_chunks( - app: object, agent_id: str, token: str, body: dict[str, object] -) -> list[dict[str, object]]: - """Run the (blocking) starlette ``TestClient`` ws session in a worker thread. +class ASGIWebSocketSession: + """Drive an ASGI websocket endpoint on the caller's event loop. - starlette's sync ``TestClient.websocket_connect`` blocks the calling event - loop for the whole ``with`` block. The gateway workers that process the - turn run on that same loop, so blocking it would deadlock. Running the - session in a worker thread keeps the server loop free to stream chunks - back to the socket. + starlette's sync ``TestClient`` runs the app on its own anyio portal loop. + That gives the process two loops over one ``OctopServer``, and loop-bound + primitives shared by both (``AgentManager._lock``) then deadlock or raise + "bound to a different event loop". Speaking ASGI directly keeps the + handler, the gateway workers and the test on the same loop, exactly like + uvicorn in production. """ - chunks: list[dict[str, object]] = [] - with TestClient(app).websocket_connect( # type: ignore[attr-defined] - f"/api/agents/{agent_id}/chat/ws?token={token}", - ) as ws: - ws.send_json(body) + + def __init__(self, app: Any, path: str) -> None: + self._app = app + self._path = path + self._from_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._to_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._task: asyncio.Task[None] | None = None + + def _scope(self) -> dict[str, Any]: + url = urlsplit(self._path) + return { + "type": "websocket", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "scheme": "ws", + "path": unquote(url.path), + "raw_path": url.path.encode("ascii"), + "query_string": url.query.encode("ascii"), + "root_path": "", + "headers": [(b"host", b"testserver")], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "subprotocols": [], + "state": {}, + "app": self._app, + } + + async def _next_event(self, timeout: float) -> dict[str, Any]: + assert self._task is not None + getter: asyncio.Future[dict[str, Any]] = asyncio.ensure_future(self._from_app.get()) + done, _pending = await asyncio.wait( + {getter, self._task}, + timeout=timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if getter in done: + return getter.result() + getter.cancel() + with contextlib.suppress(asyncio.CancelledError): + await getter + if self._task in done: + # Surface a handler crash instead of a bare timeout. + self._task.result() + raise AssertionError("websocket handler returned without sending an event") + raise TimeoutError(f"no websocket event after {timeout:.0f}s") + + async def connect(self, timeout: float = _WS_RECEIVE_TIMEOUT_S) -> None: + self._task = asyncio.create_task( + self._app(self._scope(), self._to_app.get, self._from_app.put), + name=f"asgi-ws{self._path}", + ) + await self._to_app.put({"type": "websocket.connect"}) + event = await self._next_event(timeout) + if event["type"] == "websocket.close": + raise WebSocketDisconnect(event.get("code", 1000), event.get("reason")) + if event["type"] != "websocket.accept": + raise AssertionError(f"unexpected handshake event: {event['type']}") + + async def send_json(self, payload: Any) -> None: + await self._to_app.put( + {"type": "websocket.receive", "text": json.dumps(payload, ensure_ascii=False)}, + ) + + async def receive_json(self, *, timeout: float = _WS_RECEIVE_TIMEOUT_S) -> dict[str, Any]: + event = await self._next_event(timeout) + if event["type"] == "websocket.close": + raise WebSocketDisconnect(event.get("code", 1000), event.get("reason")) + raw = event.get("text") + if raw is None: + raw = bytes(event.get("bytes") or b"").decode("utf-8") + parsed: dict[str, Any] = json.loads(raw) + return parsed + + async def drain_turn(self, *, timeout: float = _WS_RECEIVE_TIMEOUT_S) -> list[dict[str, Any]]: + """Collect frames until the turn reports ``done`` or ``error``.""" + chunks: list[dict[str, Any]] = [] while True: - raw = ws.receive_text() - chunk = json.loads(raw) + chunk = await self.receive_json(timeout=timeout) chunks.append(chunk) if chunk.get("type") in ("done", "error"): - break - return chunks + return chunks + + async def close(self) -> None: + if self._task is None: + return + task = self._task + self._task = None + await self._to_app.put({"type": "websocket.disconnect", "code": 1000}) + try: + await asyncio.wait_for(task, timeout=_WS_RECEIVE_TIMEOUT_S) + except (asyncio.CancelledError, TimeoutError): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +@asynccontextmanager +async def ws_connect(app: Any, path: str) -> AsyncIterator[ASGIWebSocketSession]: + """Open an in-process websocket session against *app* on the current loop.""" + session = ASGIWebSocketSession(app, path) + try: + await session.connect() + yield session + finally: + await session.close() + + +def chat_ws_path(agent_id: str, auth: dict[str, str]) -> str: + return f"/api/agents/{agent_id}/chat/ws?token={ws_token(auth)}" async def ws_chat_turn( @@ -45,20 +150,15 @@ async def ws_chat_turn( *, mcp_servers: list[str] | None = None, text: str = "hi", -) -> list[dict[str, object]]: - body: dict[str, object] = { +) -> list[dict[str, Any]]: + body: dict[str, Any] = { "type": "user_turn", "text": text, "messages": [{"role": "user", "content": text}], } if mcp_servers is not None: body["mcp_servers"] = mcp_servers - # httpx dropped AsyncClient.websocket_connect; use starlette's sync TestClient - # in a worker thread (see _consume_ws_chunks). - return await asyncio.to_thread( - _consume_ws_chunks, - client._octop_app, # type: ignore[attr-defined] - agent_id, - ws_token(auth), - body, - ) + app = client._octop_app # type: ignore[attr-defined] + async with ws_connect(app, chat_ws_path(agent_id, auth)) as ws: + await ws.send_json(body) + return await ws.drain_turn()