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/src/octop/infra/agents/manager.py b/src/octop/infra/agents/manager.py index 0f07f47c..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 @@ -71,6 +72,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 +335,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 +406,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 +573,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 +674,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: @@ -1339,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)", @@ -1421,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/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..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, @@ -147,11 +148,32 @@ 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 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) + 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. @@ -182,14 +204,24 @@ 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() task = asyncio.create_task( self._run_loop(agent_id), name=f"proactive_care_{agent_id}", @@ -199,7 +231,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/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/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 bf8e57c8..454844a0 100644 --- a/tests/support/app.py +++ b/tests/support/app.py @@ -54,14 +54,21 @@ 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( 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() 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() diff --git a/tests/unit/proactive/test_scheduler.py b/tests/unit/proactive/test_scheduler.py index 935ac3b7..c865fb9a 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,63 @@ 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) + + +@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 == {}