diff --git a/.omo/plans/fix-mcp-session-lifecycle.md b/.omo/plans/fix-mcp-session-lifecycle.md deleted file mode 100644 index 55de89956..000000000 --- a/.omo/plans/fix-mcp-session-lifecycle.md +++ /dev/null @@ -1,407 +0,0 @@ -# fix-mcp-session-lifecycle - Work Plan - -## TL;DR (For humans) - -**What you'll get:** MCP resources (toolsets, transports, ACP connections) are properly cleaned up when sessions close or WebSocket connections drop, eliminating stale toolset references on session resume. This is Phase 1 of a 2-phase MCP lifecycle redesign. - -**Why this approach:** The root cause is that session-scoped MCP state is scattered across 4 objects (`_toolset_cache`, `_session_connection_pool`, `_mcp_snapshot`, `AcpMcpConnectionManager._connections`) with no coordinated cleanup. The fix centralizes session-scoped state into `_SessionContext` on `MCPManager`, adds a per-session `asyncio.Lock` for concurrency-safe idempotent `cleanup_session()`, and wires cleanup into all 3 close paths (ACPSession.close, SessionController.close_session, WebSocket disconnect hook). All 8 design decisions were made through 3 rounds of Oracle + Gemini review. - -**What it will NOT do:** No changes to MessageNode base class, AgentPool registry, MCPResourceProvider model, or config API. No Phase 2 features (per-agent MCPManager removal, allow/block lists, pool-level MCP consolidation). No ACP v2 migration. - -**Effort:** Large -**Risk:** Medium — touches 7 source files across 3 subsystems (MCP manager, ACP server, session orchestrator), but all design decisions are made and codebase state is verified -**Decisions to sanity-check:** D4 (as_capability(session_id) API change — only caller is get_agentlet), D5 (reverse index at manager level instead of changing _session_streams key type), D6 (two-layer close-then-recreate in resume_session) - -Your next move: approve to start work, or run a high-accuracy review first. Full execution detail follows below. - ---- - -> TL;DR (machine): Large effort, Medium risk, 33 todos across 7 waves — fix MCP session lifecycle by centralizing session-scoped state in _SessionContext, wiring cleanup into all close paths, and fixing resume_session. - -## Scope -### Must have -- `_SessionContext` dataclass on `MCPManager` with per-session `toolset_cache`, `connection_pool`, `snapshot`, `acp_connection_ids`, and `_cleanup_lock` -- `MCPManager.get_or_create_session()`, `update_session_snapshot()`, `add_acp_transport()`, `cleanup_session()` methods -- `as_capability(session_id: str | None = None)` simplified API replacing `as_capability(snapshot=..., session_pool=...)` -- `_session_connections: dict[str, set[tuple[str, int]]]` reverse index on `AcpMcpConnectionManager` with `register_session_connection()` and `cleanup_session()` methods -- `has_active_sessions()` on `AcpMcpConnection` -- `cleanup_session()` wired into `ACPSession.close()` and `SessionController._close_session_run_turn()` -- `_session_id` stored on `Agent` and propagated to `as_capability()` call in `get_agentlet()` -- `resume_session()` close-then-recreate via two-layer cleanup (SessionController.close_session + ACPSession.close) -- `on_disconnect` callback in `_handle_websocket_client()` + `close_all_sessions_for_connection()` on `ACPSessionManager` -- All existing bug-documenting tests flipped to fix-verifying tests -- Full unit + integration test coverage for all new code paths - -### Must NOT have (guardrails, anti-slop, scope boundaries) -- NO changes to `MessageNode` base class (`messagenode.py`) -- NO changes to `AgentPool` registry (`pool.py`) -- NO removal of per-agent MCPManager (Phase 2) -- NO config API changes — no new YAML fields, no new public config models -- NO changes to `MCPResourceProvider` model or `ResourceProvider` base class -- NO Phase 2 features (allow/block lists, pool-level MCP consolidation, skill MCP dual path consolidation) -- NO ACP v2 protocol migration -- NO `getattr`/`hasattr` — full type safety with match/case or isinstance -- NO TODOs left in code - -### Metis gap resolutions (folded into todos below) -1. **GAP-1 (Critical)**: `AcpMcpConnection.register_session()` currently returns `SessionStreamPair`, NOT an int key. **Resolution**: Modify `register_session()` to return `tuple[SessionStreamPair, int]` — the pair AND the internal `_next_session_key`. Callers store the int key in `acp_connection_ids`. Affected: T1, T3, T6, T7, T8. -2. **GAP-3 (Critical)**: `AgentSideConnection` has no `connection_id` for `_connection_sessions` lookup. **Resolution**: Generate a UUID4 string for each WebSocket connection at accept time, store it on the `AgentSideConnection` instance (add a `connection_id: str` attribute set in `_handle_websocket_client()`), and use that as the key in `_connection_sessions`. The `on_disconnect` callback receives the `AgentSideConnection` and reads `.connection_id`. Affected: T24, T25, T26. -3. **GAP-4 (High)**: `_session_id` storage location undefined. **Resolution**: Use `run_ctx.session_id` in `get_agentlet()` (already available via `AgentRunContext`) instead of storing `self._session_id` on Agent. This avoids duplicating state. If `run_ctx` is None or `run_ctx.session_id` is None, fall back to `session_id=None` (global-only capabilities). Affected: T12, T16. -4. **GAP-5 (High)**: `connect_acp_mcp_server()` signature must gain `session_id: str` parameter. **Resolution**: Change signature to `connect_acp_mcp_server(self, server: AcpMcpServer, session_id: str) -> str`. Update call site at `session.py:478`. Affected: T8. -5. **GAP-7 (Medium)**: `_make_capability()` is a closure inside `as_capability()` and can't access per-session cache. **Resolution**: Pass `toolset_cache: dict[str, Any]` as a parameter to `_make_capability()` instead of accessing `self._toolset_cache` directly. When `cache=True`, pass `self._toolset_cache`; when `cache=False`, pass `ctx.toolset_cache`. Affected: T10. -6. **GAP-11 (High)**: Race condition — `cleanup_session()` can pop context while `as_capability()` reads it. **Resolution**: `as_capability()` acquires `ctx._cleanup_lock` before reading the session context (shared lock via `asyncio.Lock` — but `asyncio.Lock` is exclusive, not shared). Alternative: `as_capability()` catches `KeyError` on `_session_contexts` lookup and falls back to global-only. **Chosen**: Catch `KeyError` fallback approach — simpler, no lock contention. Affected: T10. -7. **GAP-12 (High)**: `AcpMcpConnectionManager.cleanup_session()` has no lock. **Resolution**: Add `_cleanup_lock: asyncio.Lock` to `AcpMcpConnectionManager.__init__`. `cleanup_session()` acquires it. Double-cleanup from `MCPManager.cleanup_session()` + direct call is idempotent (pop from `_session_connections` returns None on second call). Affected: T7. -8. **GAP-15 (High)**: Task 7.5 manual test violates zero-user-intervention. **Resolution**: Replace with automated integration test using mock ACP client. Test creates a mock WebSocket connection, sends ACP messages, simulates disconnect, reconnects, and verifies MCP tools work. Affected: T33. -9. **GAP-14 (Medium)**: Resume after close — closed session re-opening. **Resolution**: `SessionController.close_session()` marks session as closed in store (line 830-832) but `_get_or_create_session_locked()` creates fresh `SessionState` if not in `_sessions` dict (which was popped at line 933). The store's "closed" flag is informational — a new `SessionState` is created. This is validated by the existing resume flow. No change needed, but T20 acceptance criteria must verify this explicitly. - -## Verification strategy -> Zero human intervention - all verification is agent-executed. -- Test decision: tests-after (implementation first, tests in same todo) + pytest -- Evidence: .omo/evidence/task--fix-mcp-session-lifecycle. -- Lint: `uv run ruff check src/` — zero errors -- Types: `uv run --no-group docs mypy src/` — zero errors on changed files -- Tests: `uv run pytest tests/mcp_server/ tests/agentpool_server/acp_server/ -v` -- Unit marker: `uv run pytest -m unit` - -## Execution strategy -### Parallel execution waves -> Target 5-8 todos per wave. Fewer than 3 (except the final) means under-splitting. - -### Dependency matrix -| Todo | Depends on | Blocks | Can parallelize with | -| --- | --- | --- | --- | -| T1 (_SessionContext + _session_contexts) | — | T2, T3, T4, T5, T10, T15 | T6, T7, T8, T9 | -| T2 (get_or_create_session + update_session_snapshot) | T1 | T4, T10, T16 | T3, T6, T7, T8, T9 | -| T3 (add_acp_transport) | T1 | T5, T8 | T2, T6, T7, T8, T9 | -| T4 (cleanup_session with lock) | T1, T2 | T15, T17, T19, T22 | T5, T6, T7, T8, T9 | -| T5 (session context unit tests) | T1-T4 | — | T6, T7, T8, T9 | -| T6 (_session_connections + register_session_connection) | — | T7, T8, T15 | T1-T5, T9 | -| T7 (AcpMcpConnectionManager.cleanup_session + has_active_sessions) | T6 | T15, T17, T19, T22 | T1-T5, T8, T9 | -| T8 (Wire register_session_connection into connect_acp_mcp_server) | T3, T6 | T15 | T1-T5, T7, T9 | -| T9 (ACP session connection unit tests) | T6, T7, T8 | — | T1-T5 | -| T10 (as_capability new signature + _make_capability + _process_snapshot) | T1, T2 | T11, T12, T13, T14 | — | -| T11 (session-scoped vs global routing) | T10 | T12 | T13, T14 | -| T12 (Update get_agentlet call site) | T10, T11 | T15, T16 | T13, T14 | -| T13 (Update test_mcpmanager_caching.py) | T10 | — | T14 | -| T14 (Flip test_stale_mcp_connection.py to fix-verifying) | T10 | — | T13 | -| T15 (Wire cleanup into ACPSession.close + SessionController) | T4, T7, T8, T12 | T17, T18, T19, T20, T22 | T16 | -| T16 (get_or_create_session_agent + _session_id on Agent) | T2, T12 | T17, T18 | T15 | -| T17 (Integration: create→run→close→verify empty) | T15, T16 | — | T18, T19 | -| T18 (Integration: close→recreate same ID→fresh resources) | T15, T16 | — | T17, T19 | -| T19 (Test: concurrent cleanup_session calls) | T4, T15 | T22 | T17, T18 | -| T20 (Fix resume_session close-then-recreate) | T15, T16 | T21, T22, T23 | — | -| T21 (Test: resume→old closed→fresh MCP) | T20 | — | T22, T23 | -| T22 (Test: resume after WebSocket reconnect) | T15, T19, T20 | — | T21, T23 | -| T23 (Test: resume with active run→RunHandle cancelled) | T20 | — | T21, T22 | -| T24 (on_disconnect param + ConnectionClosed hook) | T15 | T25, T26, T27 | — | -| T25 (_connection_sessions + close_all_sessions_for_connection) | T15, T24 | T26, T28 | — | -| T26 (Wire on_disconnect in server setup) | T24, T25 | T27, T28 | — | -| T27 (Tests: disconnect closes + other connections unaffected) | T25, T26 | — | T28 | -| T28 (Test: disconnect during active run→RunHandle cancelled) | T25, T26 | — | T27 | -| T29-T33 (End-to-end verification) | ALL | — | — | - -## Todos -> Implementation + Test = ONE todo. Never separate. - -### Wave 1: P1a — MCPManager Session Tracking (foundation) - -- [x] 1. Add `_SessionContext` dataclass and `_session_contexts` dict to MCPManager - What to do / Must NOT do: Create a `@dataclass` named `_SessionContext` with fields: `connection_pool: SessionConnectionPool`, `toolset_cache: dict[str, Any]`, `snapshot: McpConfigSnapshot | None`, `acp_connection_ids: list[tuple[str, int]]`, `_cleanup_lock: asyncio.Lock`. Add `_session_contexts: dict[str, _SessionContext]` to `MCPManager.__init__` (after `_toolset_cache` at line 147). Import `SessionConnectionPool` from `agentpool.mcp_server.session_pool`, `McpConfigSnapshot` from `agentpool.mcp_server.config_snapshot`. **Metis GAP-1 resolution**: `AcpMcpConnection.register_session()` currently returns `SessionStreamPair` only — it must be modified (in T7) to return `tuple[SessionStreamPair, int]` so the int key can be stored in `acp_connection_ids`. Must NOT remove or rename existing `_toolset_cache` (D3: retained for global configs). - Parallelization: Wave 1 | Blocked by: — | Blocks: T2, T3, T4, T5 - References: `src/agentpool/mcp_server/manager.py:115` (MCPManager class), `manager.py:123-147` (__init__ fields), `manager.py:147` (_toolset_cache line), `src/agentpool/mcp_server/session_pool.py` (SessionConnectionPool class with `cleanup(timeout=5.0)` method and `copy_pre_created_transports()`), `src/agentpool/mcp_server/config_snapshot.py` (McpConfigSnapshot frozen dataclass with `pool_configs`, `agent_configs`, `session_configs`, `skill_configs` fields and `global_configs`/`session_scoped_configs` properties) - Acceptance criteria: `uv run python -c "from agentpool.mcp_server.manager import MCPManager, _SessionContext; print(_SessionContext.__dataclass_fields__.keys())"` prints fields including `connection_pool`, `toolset_cache`, `snapshot`, `acp_connection_ids`, `_cleanup_lock`. `uv run ruff check src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) `uv run python -c "import asyncio; from agentpool.mcp_server.manager import _SessionContext; ctx = _SessionContext(connection_pool=None, toolset_cache={}, snapshot=None, acp_connection_ids=[], _cleanup_lock=asyncio.Lock()); print(ctx)"` runs without error. (failure) Verify `MCPManager()` has `_session_contexts` attribute initialized as empty dict. Evidence: `.omo/evidence/task-1-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(mcp): add _SessionContext dataclass and _session_contexts to MCPManager - -- [x] 2. Implement `get_or_create_session()` and `update_session_snapshot()` on MCPManager - What to do / Must NOT do: Add `get_or_create_session(self, session_id: str) -> _SessionContext` — if `session_id` not in `_session_contexts`, create new `_SessionContext` with fresh `SessionConnectionPool()`, empty `toolset_cache`, `snapshot=None`, empty `acp_connection_ids`, new `asyncio.Lock()`. Return existing if present. Add `update_session_snapshot(self, session_id: str, snapshot: McpConfigSnapshot) -> None` — calls `get_or_create_session(session_id)` then sets `.snapshot = snapshot`. Must NOT raise if session already exists (idempotent). - Parallelization: Wave 1 | Blocked by: T1 | Blocks: T4, T10, T16 - References: `src/agentpool/mcp_server/manager.py:115` (class), `session_pool.py` (SessionConnectionPool constructor — check if it takes args), `config_snapshot.py` (McpConfigSnapshot type) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_get_or_create_session" -v` passes (test added in T5). `uv run ruff check src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) Create manager, call `get_or_create_session("s1")` twice, verify same object returned. (failure) Call `update_session_snapshot` on non-existent session, verify it creates the context. Evidence: `.omo/evidence/task-2-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(mcp): implement get_or_create_session and update_session_snapshot - -- [x] 3. Implement `add_acp_transport()` on MCPManager - What to do / Must NOT do: Add `add_acp_transport(self, session_id: str, client_id: str, transport: ClientTransport, connection_id: str, session_key: int) -> None` — gets session context via `get_or_create_session(session_id)`, adds transport to `ctx.connection_pool` (check SessionConnectionPool's API for adding transports — it uses `(client_id, skill_name)` keys), appends `(connection_id, session_key)` to `ctx.acp_connection_ids`. Must NOT create duplicate entries if called twice with same args. - Parallelization: Wave 1 | Blocked by: T1 | Blocks: T5, T8 - References: `src/agentpool/mcp_server/manager.py:115` (class), `session_pool.py` (SessionConnectionPool — check how transports are stored, keyed by `(client_id, skill_name)`), `src/acp/client/protocol.py` (ClientTransport type) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_add_acp_transport" -v` passes (test added in T5). `uv run ruff check src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) Add transport, verify it appears in `ctx.connection_pool` and `ctx.acp_connection_ids`. (failure) Add transport to non-existent session, verify session context is created. Evidence: `.omo/evidence/task-3-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(mcp): implement add_acp_transport for session-scoped ACP tracking - -- [x] 4. Implement `cleanup_session()` on MCPManager with per-session lock - What to do / Must NOT do: Add `async cleanup_session(self, session_id: str) -> None`. Acquire `ctx._cleanup_lock` (from `get_or_create_session`). In try block: (1) clear `ctx.toolset_cache` dict, (2) call `await ctx.connection_pool.cleanup()` (with try/except to log but not re-raise), (3) delegate to ACP cleanup — if `self._acp_mcp_manager` is not None, call `await self._acp_mcp_manager.cleanup_session(session_id)` (with try/except to log). In finally block: always `self._session_contexts.pop(session_id, None)`. Must NOT re-raise exceptions from intermediate steps. Must NOT skip the pop in finally. D8: the lock makes concurrent calls idempotent — second caller blocks on lock, then finds session already popped. - Parallelization: Wave 1 | Blocked by: T1, T2 | Blocks: T15, T17, T19, T22 - References: `src/agentpool/mcp_server/manager.py:115` (class), `manager.py:275` (disconnect_all — pattern for clearing toolset cache), `manager.py:438` (cleanup — pattern for exit_stack closing), `session_pool.py` (SessionConnectionPool.cleanup(timeout=5.0) method), `src/agentpool_server/acp_server/acp_mcp_manager.py:253` (AcpMcpConnectionManager — will have cleanup_session() after T7). Note: MCPManager may need an `_acp_mcp_manager: AcpMcpConnectionManager | None = None` field to delegate ACP cleanup — check if it already has a reference, if not add one. - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_cleanup_session" -v` passes (test added in T5). `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_concurrent_cleanup" -v` passes. `uv run ruff check src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) Create session, add resources, call `cleanup_session()`, verify `_session_contexts` is empty. (failure) Call `cleanup_session()` twice concurrently (asyncio.gather), verify no errors and second call is no-op. Evidence: `.omo/evidence/task-4-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(mcp): implement cleanup_session with per-session asyncio.Lock - -- [x] 5. Unit tests for MCPManager session context lifecycle - What to do / Must NOT do: Create `tests/mcp_server/test_session_lifecycle.py` with tests: (1) `test_get_or_create_session_creates_and_returns_same` — two calls return same object, (2) `test_get_or_create_session_creates_fresh_for_different_ids` — different session_ids get different contexts, (3) `test_update_session_snapshot_stores_snapshot` — snapshot is stored correctly, (4) `test_add_acp_transport_stores_transport_and_ids` — transport and (connection_id, session_key) are stored, (5) `test_cleanup_session_clears_all_resources` — after cleanup, `_session_contexts` is empty, (6) `test_cleanup_session_is_idempotent` — double-call is no-op, (7) `test_concurrent_cleanup_session_no_error` — asyncio.gather of two cleanup calls. Use `@pytest.mark.unit`. Use `pytest.fixture` for MCPManager instance. Must NOT use `getattr`/`hasattr` — use direct attribute access with type annotations. - Parallelization: Wave 1 | Blocked by: T1-T4 | Blocks: — - References: `tests/mcp_server/test_mcpmanager_caching.py` (existing test patterns), `tests/mcp_server/test_session_pool.py` (SessionConnectionPool test patterns), `tests/conftest.py` (fixtures, TestModel, observability disabled) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -v` — all 7 tests pass. `uv run pytest -m unit tests/mcp_server/test_session_lifecycle.py -v` — all pass with unit marker. - QA scenarios: (happy) All 7 tests pass. (failure) Intentionally break cleanup (remove pop from finally), verify test 6 and 7 fail. Evidence: `.omo/evidence/task-5-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): add session context lifecycle unit tests - -### Wave 2: P1c — AcpMcpConnectionManager Session Tracking (parallel with Wave 1) - -- [x] 6. Add `_session_connections` dict and `register_session_connection()` to AcpMcpConnectionManager - What to do / Must NOT do: Add `_session_connections: dict[str, set[tuple[str, int]]]` to `AcpMcpConnectionManager.__init__` (after `_connections` at line 259). Maps `session_id` → set of `(connection_id, session_key)` tuples. Add `register_session_connection(self, session_id: str, connection_id: str, session_key: int) -> None` — adds `(connection_id, session_key)` to the session's set, creating the set if missing. Must NOT modify `AcpMcpConnection._session_streams` (D5: reverse index at manager level, not changing int keys). - Parallelization: Wave 2 | Blocked by: — | Blocks: T7, T8, T15 - References: `src/agentpool_server/acp_server/acp_mcp_manager.py:253` (AcpMcpConnectionManager class), `acp_mcp_manager.py:259` (_connections dict), `acp_mcp_manager.py:34` (AcpMcpConnection class), `acp_mcp_manager.py:50` (_session_streams dict with int keys), `acp_mcp_manager.py:52` (_next_session_key int), `acp_mcp_manager.py:78` (register_session returns SessionStreamPair) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py -k "test_register_session_connection" -v` passes (test added in T9). `uv run ruff check src/agentpool_server/acp_server/acp_mcp_manager.py` passes. - QA scenarios: (happy) Register a connection, verify it appears in `_session_connections`. (failure) Register same tuple twice, verify set deduplicates. Evidence: `.omo/evidence/task-6-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): add session connection tracking to AcpMcpConnectionManager - -- [x] 7. Implement `cleanup_session()` and `has_active_sessions()` on AcpMcpConnectionManager/AcpMcpConnection - What to do / Must NOT do: Add `has_active_sessions(self) -> bool` to `AcpMcpConnection` (line 34) — returns `len(self._session_streams) > 0`. **Metis GAP-1 resolution**: Modify `register_session()` at `acp_mcp_manager.py:78` to return `tuple[SessionStreamPair, int]` instead of just `SessionStreamPair` — return `(pair, key)` where `key` is the internal `_next_session_key`. Add `async cleanup_session(self, session_id: str) -> None` to `AcpMcpConnectionManager` (line 253) — **Metis GAP-12 resolution**: acquire `_cleanup_lock` (new `asyncio.Lock` added to `__init__`) before proceeding. Pop `session_id` from `_session_connections`, for each `(connection_id, session_key)` tuple: look up `AcpMcpConnection` via `self._connections[connection_id]`, look up `SessionStreamPair` via `conn._session_streams[session_key]` (note: _session_streams uses int keys, session_key is int from the modified `register_session()`), call `conn.unregister_session(pair)`, after processing all tuples for a connection check `conn.has_active_sessions()` — if False, optionally remove the connection (check existing `remove_connection()` logic at line ~290 for cleanup pattern). Must NOT change `_session_streams` key type from int to str (D5). - Parallelization: Wave 2 | Blocked by: T6 | Blocks: T15, T17, T19, T22 - References: `src/agentpool_server/acp_server/acp_mcp_manager.py:253` (class), `acp_mcp_manager.py:67` (close method), `acp_mcp_manager.py:78` (register_session), `acp_mcp_manager.py:98` (unregister_session takes SessionStreamPair), `acp_mcp_manager.py:50` (_session_streams dict), `acp_mcp_manager.py:224` (broadcast_to_sessions — pattern for iterating sessions) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py -k "test_cleanup_session" -v` passes (test added in T9). `uv run ruff check src/agentpool_server/acp_server/acp_mcp_manager.py` passes. - QA scenarios: (happy) Register 2 sessions on same connection, cleanup one, verify connection still has 1 active session. (failure) Cleanup all sessions, verify connection is removed or has `has_active_sessions() == False`. Evidence: `.omo/evidence/task-7-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): implement cleanup_session and has_active_sessions - -- [x] 8. Wire `register_session_connection()` into `connect_acp_mcp_server()` - What to do / Must NOT do: **Metis GAP-5 resolution**: Change `connect_acp_mcp_server()` signature at `acp_agent.py:823` from `connect_acp_mcp_server(self, server: AcpMcpServer) -> str` to `connect_acp_mcp_server(self, server: AcpMcpServer, session_id: str) -> str`. Update call site at `session.py:478` to pass `self.session_id`. After calling `AcpMcpConnection.register_session()` (which now returns `tuple[SessionStreamPair, int]` per T7), extract the `session_key` (int) from the return. Call `self._mcp_manager.register_session_connection(session_id, connection_id, session_key)` and `agent.mcp.add_acp_transport(session_id, client_id, transport, connection_id, session_key)` (or equivalent MCPManager method from T3). Must NOT change the `SessionStreamPair` return type. - Parallelization: Wave 2 | Blocked by: T3, T6 | Blocks: T15 - References: `src/agentpool_server/acp_server/acp_agent.py:823` (connect_acp_mcp_server), `acp_agent.py:846` (disconnect_acp_mcp_server), `acp_agent.py:238` (_mcp_manager field), `acp_agent.py:263` (_mcp_manager init), `acp_mcp_manager.py:78` (register_session returns SessionStreamPair — check if key is stored on the pair or accessible), `src/agentpool_server/acp_server/session.py:165` (self.agent is BaseAgent) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py -k "test_connect_registers_session" -v` passes (test added in T9). `uv run ruff check src/agentpool_server/acp_server/acp_agent.py` passes. - QA scenarios: (happy) Connect ACP MCP server, verify `register_session_connection()` was called with correct session_id and connection_id. (failure) Connect without session_id, verify graceful handling (no crash). Evidence: `.omo/evidence/task-8-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): wire register_session_connection into connect_acp_mcp_server - -- [x] 9. Unit tests for AcpMcpConnectionManager session connection tracking - What to do / Must NOT do: Create `tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py` with tests: (1) `test_register_session_connection_adds_to_set`, (2) `test_register_deduplicates_same_tuple`, (3) `test_cleanup_session_unregisters_streams`, (4) `test_cleanup_preserves_shared_connection`, (5) `test_cleanup_removes_connection_with_no_sessions`, (6) `test_has_active_sessions_true_when_streams_exist`, (7) `test_has_active_sessions_false_when_empty`, (8) `test_connect_acp_mcp_server_registers_session` (integration with T8). Use `@pytest.mark.unit` for 1-7, `@pytest.mark.integration` for 8. Must NOT use `getattr`/`hasattr`. - Parallelization: Wave 2 | Blocked by: T6, T7, T8 | Blocks: — - References: `tests/agentpool_server/acp_server/test_acp_mcp_manager.py` (existing test patterns), `tests/agentpool_server/acp_server/test_acp_mcp_agent_integration.py` (integration test patterns) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py -v` — all 8 tests pass. `uv run pytest -m unit tests/agentpool_server/acp_server/test_acp_mcp_session_cleanup.py -v` — 7 unit tests pass. - QA scenarios: (happy) All 8 tests pass. (failure) Remove `has_active_sessions()` check from cleanup, verify test 4 (shared connection preservation) fails. Evidence: `.omo/evidence/task-9-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): add session connection cleanup unit tests - -### Wave 3: P1b — as_capability Session-Aware API (depends on Wave 1) - -- [x] 10. Change `as_capability()` signature and modify `_make_capability()` and `_process_snapshot()` - What to do / Must NOT do: Change `as_capability(self, snapshot: McpConfigSnapshot | None = None, session_pool: SessionConnectionPool | None = None)` at `manager.py:301` to `as_capability(self, session_id: str | None = None) -> AggregatingCapability` (or whatever the return type is — check current signature). When `session_id` is provided: look up `_SessionContext` via `get_or_create_session(session_id)`, use its `snapshot`, `connection_pool`, and `toolset_cache`. **Metis GAP-11 resolution**: Wrap the `_session_contexts` lookup in try/except `KeyError` — if the session context was popped by concurrent `cleanup_session()`, fall back to global-only capabilities (log a warning). This avoids the race condition without lock contention. When `session_id` is None: process only global configs from `self.servers` (backward compat, use `_toolset_cache`). **Metis GAP-7 resolution**: Modify `_make_capability(self, server, transport)` at line 374 to accept `toolset_cache: dict[str, Any]` parameter (instead of accessing `self._toolset_cache` directly) — when processing global configs, pass `self._toolset_cache`; when processing session-scoped configs, pass `ctx.toolset_cache`. Modify `_process_snapshot(self, snap)` at line 396 to pass the correct `toolset_cache` for session-scoped vs global configs. Must NOT remove `_toolset_cache` (D3: retained for global configs). Must NOT break `session_id=None` backward compat path. - Parallelization: Wave 3 | Blocked by: T1, T2 | Blocks: T11, T12, T13, T14 - References: `src/agentpool/mcp_server/manager.py:301` (as_capability current signature), `manager.py:374` (_make_capability), `manager.py:396` (_process_snapshot), `manager.py:147` (_toolset_cache), `config_snapshot.py` (McpConfigSnapshot.global_configs and .session_scoped_configs properties), `session_pool.py` (SessionConnectionPool) - Acceptance criteria: `uv run pytest tests/mcp_server/test_manager_capability.py -v` — all existing 19 tests pass (may need updates in T13). `uv run ruff check src/agentpool/mcp_server/manager.py` passes. `uv run --no-group docs mypy src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) Call `as_capability(session_id="s1")` with a session context that has a snapshot, verify session-scoped configs use per-session cache. (failure) Call `as_capability(session_id=None)`, verify only global configs are processed. Evidence: `.omo/evidence/task-10-fix-mcp-session-lifecycle.txt` - Commit: Y | refactor(mcp): change as_capability to session_id-based API - -- [x] 11. Implement session-scoped vs global config routing in `as_capability()` - What to do / Must NOT do: Inside `as_capability(session_id)`: if `session_id` is not None and `ctx.snapshot` is not None, call `_process_snapshot(ctx.snapshot, cache=False, toolset_cache=ctx.toolset_cache, connection_pool=ctx.connection_pool)` for session-scoped configs and `_process_snapshot(ctx.snapshot, cache=True)` for global configs. If `session_id` is None, process `self.servers` global configs with `_toolset_cache` as before. Must NOT mix session-scoped toolsets into `_toolset_cache`. - Parallelization: Wave 3 | Blocked by: T10 | Blocks: T12 - References: `src/agentpool/mcp_server/manager.py:301` (as_capability), `manager.py:396` (_process_snapshot), `config_snapshot.py` (global_configs property returns pool+agent configs, session_scoped_configs returns session+skill configs) - Acceptance criteria: `uv run pytest tests/mcp_server/test_manager_capability.py -k "session" -v` passes (tests updated in T13). `uv run ruff check src/agentpool/mcp_server/manager.py` passes. - QA scenarios: (happy) Session-scoped config produces toolset in `ctx.toolset_cache`, NOT in `_toolset_cache`. (failure) Global config produces toolset in `_toolset_cache`, NOT in `ctx.toolset_cache`. Evidence: `.omo/evidence/task-11-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(mcp): route session-scoped configs to per-session cache - -- [x] 12. Update `get_agentlet()` call site in `agent.py` - What to do / Must NOT do: At `agent.py:901-903`, change `mcp_capabilities = await self.mcp.as_capability(snapshot=self._mcp_snapshot, session_pool=self._session_connection_pool)` to `mcp_capabilities = await self.mcp.as_capability(session_id=run_ctx.session_id if run_ctx else None)`. **Metis GAP-4 resolution**: Use `run_ctx.session_id` (already available via `AgentRunContext` parameter in `get_agentlet()`) instead of storing `self._session_id` on Agent. This avoids duplicating state. If `run_ctx` is None or `run_ctx.session_id` is None, pass `session_id=None` (global-only capabilities). Remove the direct setting of `self._mcp_snapshot` and `self._session_connection_pool` on the agent if they are now managed through `MCPManager.get_or_create_session()` and `update_session_snapshot()`. However, keep `self._mcp_snapshot` and `self._session_connection_pool` fields for backward compat if other code reads them — check all references. Must NOT remove `_mcp_snapshot` or `_session_connection_pool` field declarations if other code references them. - Parallelization: Wave 3 | Blocked by: T10, T11 | Blocks: T15, T16 - References: `src/agentpool/agents/native_agent/agent.py:901-903` (as_capability call), `agent.py:333-334` (_mcp_snapshot and _session_connection_pool declarations), `src/agentpool/orchestrator/session_controller.py:504-505` (child agent sets _mcp_snapshot and _session_connection_pool), `session_controller.py:586-587` (main agent sets _mcp_snapshot and _session_connection_pool) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -v` — existing ACP tests pass. `uv run ruff check src/agentpool/agents/native_agent/agent.py` passes. `uv run --no-group docs mypy src/agentpool/agents/native_agent/agent.py` passes. - QA scenarios: (happy) Agentlet creation calls `as_capability(session_id=...)` with correct session_id. (failure) Call `as_capability(session_id=None)`, verify it returns global-only capabilities. Evidence: `.omo/evidence/task-12-fix-mcp-session-lifecycle.txt` - Commit: Y | refactor(agent): update get_agentlet to use as_capability(session_id) - -- [x] 13. Update existing tests in `test_mcpmanager_caching.py` - What to do / Must NOT do: Update all 6 tests in `tests/mcp_server/test_mcpmanager_caching.py` to use new `as_capability(session_id=...)` API instead of `as_capability(snapshot=..., session_pool=...)`. Tests: toolset cache sharing, client_id keying, aggregating provider, no dedup hack, engineer/librarian scoping. For tests that verify cache sharing behavior, update to test per-session cache isolation instead. Must NOT delete tests — update them to verify the new behavior. - Parallelization: Wave 3 | Blocked by: T10 | Blocks: — - References: `tests/mcp_server/test_mcpmanager_caching.py` (6 existing tests), `tests/mcp_server/test_manager_capability.py` (19 existing tests — may also need updates) - Acceptance criteria: `uv run pytest tests/mcp_server/test_mcpmanager_caching.py -v` — all 6 updated tests pass. `uv run pytest tests/mcp_server/test_manager_capability.py -v` — all 19 tests pass. - QA scenarios: (happy) All tests pass with new API. (failure) Revert API change, verify tests fail with old signature. Evidence: `.omo/evidence/task-13-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): update caching tests for session_id API - -- [x] 14. Flip `test_stale_mcp_connection.py` tests from bug-documenting to fix-verifying - What to do / Must NOT do: Update all 5 tests in `tests/mcp_server/test_stale_mcp_connection.py`: (1) `test_session_resume_returns_stale_toolset_from_cache` → rename to `test_session_resume_returns_fresh_toolset` and assert session 2 gets a DIFFERENT toolset object, (2) `test_acp_client_id_is_deterministic` → keep as-is (still valid), (3) `test_session_pool_provides_fresh_transport` → keep as-is (still valid), (4) `test_multiple_acp_servers_all_go_stale` → rename to `test_multiple_acp_servers_get_fresh_toolsets` and assert freshness, (5) `test_disconnect_all_clears_cache_but_not_called_on_resume` → rename to `test_cleanup_session_clears_per_session_cache` and verify cleanup works. Add `try/finally` or `@pytest.fixture` teardown for resource cleanup — current tests skip cleanup on assertion failure. Must NOT keep assertions that verify the bug exists. - Parallelization: Wave 3 | Blocked by: T10 | Blocks: — - References: `tests/mcp_server/test_stale_mcp_connection.py` (5 existing tests documenting the bug) - Acceptance criteria: `uv run pytest tests/mcp_server/test_stale_mcp_connection.py -v` — all 5 updated tests pass. Tests verify the FIX, not the bug. - QA scenarios: (happy) Session 2 gets fresh toolset after session 1 is cleaned up. (failure) Remove cleanup_session call, verify test 1 and 4 fail (stale toolset returned). Evidence: `.omo/evidence/task-14-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): flip stale connection tests to verify fix - -### Wave 4: P1d — Wire cleanup_session into Close Paths (depends on Waves 1+2) - -- [x] 15. Wire `cleanup_session()` into `ACPSession.close()` and `SessionController._close_session_run_turn()` - What to do / Must NOT do: (1) In `session.py:795-823` (`ACPSession.close()`), add `await self.agent.mcp.cleanup_session(self.session_id)` BEFORE existing env/signal/prompt cleanup (before `acp_env.__aexit__()`). Check that `self.agent` has `.mcp` attribute and `.session_id` is accessible — `self.session_id` should be on ACPSession (check `session.py` for the field name, it may be `self._session_id` or similar). (2) In `session_controller.py:835-949` (`_close_session_run_turn()`), add `await agent.mcp.cleanup_session(session_id)` BEFORE `agent.__aexit__()` call (before line 941). Must NOT call cleanup_session AFTER `agent.__aexit__()` (agent context may be torn down). Must NOT skip cleanup if `is_per_session_agent=False` — the shared MCPManager still has session-scoped contexts that need cleanup. - Parallelization: Wave 4 | Blocked by: T4, T7, T8, T12 | Blocks: T17, T18, T19, T20, T22 - References: `src/agentpool_server/acp_server/session.py:795-823` (ACPSession.close), `session.py:165` (self.agent is BaseAgent), `src/agentpool/orchestrator/session_controller.py:835-949` (_close_session_run_turn), `session_controller.py:941-947` (agent.__aexit__ call with is_per_session_agent check) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -v` passes. `uv run pytest tests/agentpool_server/acp_server/ -v` — existing tests pass. `uv run ruff check src/agentpool_server/acp_server/session.py src/agentpool/orchestrator/session_controller.py` passes. - QA scenarios: (happy) Close session, verify `cleanup_session()` was called and `_session_contexts` is empty. (failure) Close session with active run, verify RunHandle is cancelled with timeout before cleanup. Evidence: `.omo/evidence/task-15-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(session): wire cleanup_session into ACPSession.close and SessionController - -- [x] 16. Wire `get_or_create_session()` in `get_or_create_session_agent()` and update MCP snapshot setup - What to do / Must NOT do: **Metis GAP-4 resolution**: Do NOT add `self._session_id` to Agent — `get_agentlet()` uses `run_ctx.session_id` instead (see T12). (1) In `session_controller.py:397-672` (`get_or_create_session_agent()`), when creating a new agent: call `agent.mcp.get_or_create_session(session_id)` to create the session context, and call `agent.mcp.update_session_snapshot(session_id, snapshot)` if a snapshot is available (replace the direct `agent._mcp_snapshot = ...` and `agent._session_connection_pool = ...` setting at lines 504-505 and 586-587). **Metis GAP-6 resolution**: This must be done on ALL 3 agent creation paths: (a) child session (line 444-546), (b) main native (line 548-601), (c) non-native (line 603-657). For child sessions, the MCPManager is the parent's/pool's shared one — calling `get_or_create_session` on it is correct (session_ids are unique). Must NOT remove the `_mcp_snapshot` and `_session_connection_pool` field declarations if other code reads them — but do redirect the setting through MCPManager. - Parallelization: Wave 4 | Blocked by: T2, T12 | Blocks: T17, T18 - References: `src/agentpool/agents/native_agent/agent.py:333-334` (field declarations), `src/agentpool/orchestrator/session_controller.py:397-672` (get_or_create_session_agent), `session_controller.py:504-505` (child agent MCP setup), `session_controller.py:586-587` (main agent MCP setup) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -v` — existing tests pass. `uv run ruff check src/agentpool/agents/native_agent/agent.py src/agentpool/orchestrator/session_controller.py` passes. `uv run --no-group docs mypy src/agentpool/agents/native_agent/agent.py` passes. - QA scenarios: (happy) Create session agent, verify `_session_id` is set and `_session_contexts` has the session. (failure) Create agent without session_id, verify `as_capability(session_id=None)` works. Evidence: `.omo/evidence/task-16-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(agent): add _session_id and wire get_or_create_session in SessionController - -- [x] 17. Integration test: create session → run turn → close → verify empty contexts - What to do / Must NOT do: Create integration test in `tests/mcp_server/test_session_lifecycle.py` (or a new `tests/integration/test_session_cleanup.py`): create an AgentPool with a native agent that has MCP servers, create a session, run a turn (use TestModel), close the session, verify `agent.mcp._session_contexts` is empty and `agent.mcp._toolset_cache` has no session-scoped entries. Use `@pytest.mark.integration`. Must NOT use real model calls — use TestModel from pydantic-ai. - Parallelization: Wave 4 | Blocked by: T15, T16 | Blocks: — - References: `tests/conftest.py` (TestModel setup, observability disabled), `tests/mcp_server/test_mcp_provider_lifecycle.py` (integration test patterns) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_integration_create_run_close" -v` passes. - QA scenarios: (happy) After close, `_session_contexts` is empty. (failure) Remove cleanup call from close path, verify test fails (context still present). Evidence: `.omo/evidence/task-17-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): integration test for session create→run→close lifecycle - -- [x] 18. Integration test: close → recreate same ID → verify fresh MCP resources - What to do / Must NOT do: Create integration test: create session "s1", run turn, close session, create new session "s1" (same ID), verify the new session has fresh MCP resources (different toolset objects, fresh connection pool). Use `@pytest.mark.integration`. Must NOT reuse the old session object. - Parallelization: Wave 4 | Blocked by: T15, T16 | Blocks: — - References: Same as T17 - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_integration_close_recreate_fresh" -v` passes. - QA scenarios: (happy) New session "s1" has fresh resources, different from old session "s1". (failure) Remove cleanup, verify old resources leak into new session. Evidence: `.omo/evidence/task-18-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): integration test for session close→recreate freshness - -- [x] 19. Test: concurrent `cleanup_session()` calls (WebSocket disconnect + SessionController) - What to do / Must NOT do: Create test that simulates concurrent cleanup: spawn `asyncio.gather(agent.mcp.cleanup_session("s1"), agent.mcp.cleanup_session("s1"))`. Verify no errors, no double-cleanup, `_session_contexts` is empty. Use `@pytest.mark.unit`. Must NOT use real WebSocket connections — mock the disconnect trigger. - Parallelization: Wave 4 | Blocked by: T4, T15 | Blocks: T22 - References: `tests/mcp_server/test_session_lifecycle.py` (existing test patterns from T5) - Acceptance criteria: `uv run pytest tests/mcp_server/test_session_lifecycle.py -k "test_concurrent_cleanup_from_two_paths" -v` passes. - QA scenarios: (happy) Both calls complete without error, only one does actual cleanup. (failure) Remove lock from cleanup_session, verify race condition or double-cleanup error. Evidence: `.omo/evidence/task-19-fix-mcp-session-lifecycle.txt` - Commit: Y | test(mcp): concurrent cleanup_session from WebSocket and SessionController - -### Wave 5: P1e — Fix resume_session Early-Return (depends on Wave 4) - -- [x] 20. Remove early-return and implement close-then-recreate in `resume_session()` - What to do / Must NOT do: In `session_manager.py:243-249`, remove the early-return that returns stale session when `session_id in self._acp_sessions`. Replace with: (1) if session exists, call `SessionController.close_session(session_id)` first (handles RunHandle lifecycle with 10s timeout + cancel, calls `agent.mcp.cleanup_session()` via T15, calls `agent.__aexit__()`), (2) then call `ACPSession.close()` for ACP-specific cleanup (acp_env, signals, prompts — also calls `cleanup_session()` via T15, but idempotent via D8 lock), (3) remove from `_acp_sessions`, (4) proceed to create fresh session. Fallback: if `SessionController` is unavailable (tests), call `ACPSession.close()` only. Must NOT skip the `SessionController.close_session()` call when it's available — it handles active runs. Must NOT skip `ACPSession.close()` — it handles ACP-specific state. - Parallelization: Wave 5 | Blocked by: T15, T16 | Blocks: T21, T22, T23 - References: `src/agentpool_server/acp_server/session_manager.py:243-249` (early-return to remove), `session_manager.py:45` (_acp_sessions dict), `session_manager.py:371-391` (close_all_sessions pattern), `src/agentpool/orchestrator/session_controller.py:951-966` (close_session one-liner) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -v` — existing tests pass. `uv run ruff check src/agentpool_server/acp_server/session_manager.py` passes. `uv run --no-group docs mypy src/agentpool_server/acp_server/session_manager.py` passes. - QA scenarios: (happy) Resume existing session, verify old session is closed and new session has fresh resources. (failure) Resume with active run, verify RunHandle is cancelled with timeout. Evidence: `.omo/evidence/task-20-fix-mcp-session-lifecycle.txt` - Commit: Y | fix(acp): resume_session close-then-recreate instead of early-return - -- [x] 21. Test: resume → verify old session closed → fresh MCP resources - What to do / Must NOT do: Create test: create session, run turn, resume same session, verify old session was closed (check `_acp_sessions` had old entry removed and re-added), verify new session has fresh MCP resources (different toolset objects). Use `@pytest.mark.integration`. - Parallelization: Wave 5 | Blocked by: T20 | Blocks: — - References: `tests/agentpool_server/acp_server/test_acp_mcp_agent_integration.py` (integration test patterns) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "test_resume_closes_old_session" -v` passes. - QA scenarios: (happy) Resumed session has fresh MCP resources. (failure) Revert early-return, verify test fails (stale resources). Evidence: `.omo/evidence/task-21-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): resume_session closes old and creates fresh - -- [x] 22. Test: resume after WebSocket reconnect → fresh ACP connections - What to do / Must NOT do: Create test: create session with ACP MCP server, simulate WebSocket disconnect, reconnect, resume session, verify fresh ACP connections are created and no stale connection references remain. Use `@pytest.mark.integration`. Must NOT use real WebSocket — mock the connection/disconnect. - Parallelization: Wave 5 | Blocked by: T15, T19, T20 | Blocks: — - References: `tests/agentpool_server/acp_server/test_acp_mcp_agent_integration.py` - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "test_resume_after_reconnect" -v` passes. - QA scenarios: (happy) After reconnect+resume, ACP connections are fresh. (failure) Don't close old session on resume, verify stale connections persist. Evidence: `.omo/evidence/task-22-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): resume after WebSocket reconnect creates fresh connections - -- [x] 23. Test: resume with active run → RunHandle cancelled with timeout - What to do / Must NOT do: Create test: create session, start a long-running turn, resume same session while run is active, verify RunHandle is cancelled with timeout before cleanup proceeds. Use `@pytest.mark.integration`. Must NOT block forever — use `asyncio.wait_for` in test with 30s timeout. - Parallelization: Wave 5 | Blocked by: T20 | Blocks: — - References: `src/agentpool/orchestrator/session_controller.py:835-949` (_close_session_run_turn with 10s timeout) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "test_resume_with_active_run" -v` passes. - QA scenarios: (happy) RunHandle cancelled, cleanup proceeds, new session created. (failure) Remove timeout from close_session, verify test hangs (would need timeout). Evidence: `.omo/evidence/task-23-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): resume with active run cancels RunHandle - -### Wave 6: P1f — WebSocket Disconnect Hook (depends on Waves 4+2) - -- [x] 24. Add `on_disconnect` parameter to `_handle_websocket_client()` and call in `ConnectionClosed` handler - What to do / Must NOT do: Add `on_disconnect: Callable[[AgentSideConnection], Awaitable[None]] | None = None` parameter to `_handle_websocket_client()` at `transports.py:355`. **Metis GAP-3 resolution**: Generate a UUID4 string for each WebSocket connection at accept time, store it as `conn.connection_id: str` attribute on the `AgentSideConnection` instance (set right after creation at line 376). In the `ConnectionClosed` exception handler (line 412), call `await on_disconnect(conn)` BEFORE `conn.close()` in the finally block (line 414). The callback reads `conn.connection_id` to look up sessions. If `on_disconnect` is None, skip the call (backward compat). Must NOT make `on_disconnect` a required parameter. Must NOT call `on_disconnect` after `conn.close()`. - Parallelization: Wave 6 | Blocked by: T15 | Blocks: T25, T26, T27 - References: `src/acp/transports.py:355-428` (_handle_websocket_client), `transports.py:412` (ConnectionClosed catch), `transports.py:414-428` (finally block) - Acceptance criteria: `uv run pytest tests/ -k "websocket" -v` — existing WebSocket tests pass. `uv run ruff check src/acp/transports.py` passes. `uv run --no-group docs mypy src/acp/transports.py` passes. - QA scenarios: (happy) Disconnect triggers `on_disconnect` callback with connection object. (failure) `on_disconnect=None`, verify no callback called and existing behavior unchanged. Evidence: `.omo/evidence/task-24-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): add on_disconnect callback to websocket handler - -- [x] 25. Add `_connection_sessions` to ACPSessionManager and implement `close_all_sessions_for_connection()` - What to do / Must NOT do: (1) Add `_connection_sessions: dict[str, set[str]]` (connection_id → session_ids) to `ACPSessionManager.__init__` (after `_acp_sessions` at line 45). **Metis GAP-3 resolution**: The `connection_id` is the UUID4 string generated and stored on `AgentSideConnection.connection_id` (from T24). Populate `_connection_sessions` when sessions are created/resumed — add `session_id` to `_connection_sessions[connection_id]` set. The `connection_id` must be passed from the `Client` object or from the `AgentSideConnection` when creating the session. Check `ACPSessionManager.create_session()` to see how `client: Client` is received and how to access the underlying connection's `connection_id`. (2) Implement `async close_all_sessions_for_connection(self, connection_id: str) -> None` — iterates sessions for the connection. For each session: call `SessionController.close_session(session_id)` first (RunHandle lifecycle with timeout + cancel), then call `ACPSession.close()` for ACP-specific cleanup. Both must be called — SessionController handles RunHandle + agent lifecycle, ACPSession.close() handles ACP-specific state. Remove the connection entry from `_connection_sessions` after all sessions are closed. Must NOT skip SessionController.close_session() when available. Must NOT skip ACPSession.close(). - Parallelization: Wave 6 | Blocked by: T15, T24 | Blocks: T26, T28 - References: `src/agentpool_server/acp_server/session_manager.py:45` (_acp_sessions), `session_manager.py:371-391` (close_all_sessions pattern), `src/agentpool/orchestrator/session_controller.py:951-966` (close_session) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "close_all_sessions_for_connection" -v` passes. `uv run ruff check src/agentpool_server/acp_server/session_manager.py` passes. - QA scenarios: (happy) Disconnect connection, all sessions for that connection are closed. (failure) Disconnect, verify sessions on other connections are NOT affected. Evidence: `.omo/evidence/task-25-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): implement close_all_sessions_for_connection - -- [x] 26. Wire `on_disconnect` callback in server setup - What to do / Must NOT do: In the server setup that creates `_handle_websocket_client()` call (search for where `_handle_websocket_client` is called — likely in `ACPWebSocketTransport` or a server module), pass a callback that calls `ACPSessionManager.close_all_sessions_for_connection(connection_id)`. The callback needs access to the `ACPSessionManager` instance and the `connection_id` — check how the connection_id is determined at the call site. Must NOT create a circular dependency between transports.py and session_manager.py — use a callback, not a direct import. - Parallelization: Wave 6 | Blocked by: T24, T25 | Blocks: T27, T28 - References: Search for `_handle_websocket_client` call sites in `src/acp/` and `src/agentpool_server/acp_server/`. Check `src/acp/transports.py` for `ACPWebSocketTransport` class. - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -v` — existing tests pass. `uv run ruff check src/` passes on changed files. - QA scenarios: (happy) WebSocket disconnect triggers `close_all_sessions_for_connection()`. (failure) Callback not wired, verify disconnect doesn't close sessions. Evidence: `.omo/evidence/task-26-fix-mcp-session-lifecycle.txt` - Commit: Y | feat(acp): wire on_disconnect to close_all_sessions_for_connection - -- [x] 27. Tests: WebSocket disconnect closes sessions + other connections unaffected - What to do / Must NOT do: Create tests: (1) `test_websocket_disconnect_closes_all_sessions` — create 2 sessions on same connection, disconnect, verify both closed via `cleanup_session()`, (2) `test_websocket_disconnect_preserves_other_connections` — create sessions on 2 connections, disconnect one, verify only that connection's sessions are closed. Use `@pytest.mark.integration`. Must NOT use real WebSocket — mock connection/disconnect. - Parallelization: Wave 6 | Blocked by: T25, T26 | Blocks: — - References: `tests/agentpool_server/acp_server/test_acp_mcp_agent_integration.py` - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "websocket_disconnect" -v` — both tests pass. - QA scenarios: (happy) Disconnect closes all sessions for that connection. (failure) Don't wire callback, verify sessions remain open. Evidence: `.omo/evidence/task-27-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): websocket disconnect closes sessions and preserves others - -- [x] 28. Test: WebSocket disconnect during active run → RunHandle cancelled with timeout - What to do / Must NOT do: Create test: create session, start long-running turn, simulate WebSocket disconnect, verify RunHandle is cancelled with timeout before cleanup proceeds. Use `@pytest.mark.integration`. Must NOT block forever — use `asyncio.wait_for` in test with 30s timeout. - Parallelization: Wave 6 | Blocked by: T25, T26 | Blocks: — - References: `src/agentpool/orchestrator/session_controller.py:835-949` (_close_session_run_turn with 10s timeout) - Acceptance criteria: `uv run pytest tests/agentpool_server/acp_server/ -k "websocket_disconnect_during_run" -v` passes. - QA scenarios: (happy) RunHandle cancelled, cleanup proceeds. (failure) Remove timeout, verify test hangs. Evidence: `.omo/evidence/task-28-fix-mcp-session-lifecycle.txt` - Commit: Y | test(acp): websocket disconnect during active run cancels RunHandle - -### Wave 7: End-to-End Verification - -- [x] 29. Run full test suite for MCP and ACP server - What to do / Must NOT do: Run `uv run pytest tests/mcp_server/ tests/agentpool_server/acp_server/ -v` and verify all tests pass. Capture full output. Must NOT mark any test as `xfail` or `skip` to make it pass. - Parallelization: Wave 7 | Blocked by: ALL | Blocks: — - References: All test files mentioned in previous todos - Acceptance criteria: `uv run pytest tests/mcp_server/ tests/agentpool_server/acp_server/ -v` — 0 failures, 0 errors. - QA scenarios: (happy) All tests pass. (failure) Any test fails — fix before proceeding. Evidence: `.omo/evidence/task-29-fix-mcp-session-lifecycle.txt` - Commit: N - -- [x] 30. Run unit test suite - What to do / Must NOT do: Run `uv run pytest -m unit` and verify all unit tests pass. Must NOT include slow or integration tests. - Parallelization: Wave 7 | Blocked by: ALL | Blocks: — - References: All test files - Acceptance criteria: `uv run pytest -m unit` — 0 failures, 0 errors. - QA scenarios: (happy) All unit tests pass. (failure) Any unit test fails — fix before proceeding. Evidence: `.omo/evidence/task-30-fix-mcp-session-lifecycle.txt` - Commit: N - -- [x] 31. Ruff lint check - What to do / Must NOT do: Run `uv run ruff check src/` and verify zero errors. Must NOT add `# noqa` comments to suppress errors. - Parallelization: Wave 7 | Blocked by: ALL | Blocks: — - References: All changed source files - Acceptance criteria: `uv run ruff check src/` — 0 errors. - QA scenarios: (happy) Zero lint errors. (failure) Any lint error — fix before proceeding. Evidence: `.omo/evidence/task-31-fix-mcp-session-lifecycle.txt` - Commit: N - -- [x] 32. Mypy type check - What to do / Must NOT do: Run `uv run --no-group docs mypy src/` and verify zero errors on changed files. Must NOT use `# type: ignore` to suppress errors (use proper type annotations). - Parallelization: Wave 7 | Blocked by: ALL | Blocks: — - References: All changed source files - Acceptance criteria: `uv run --no-group docs mypy src/` — 0 errors on changed files. - QA scenarios: (happy) Zero type errors. (failure) Any type error — fix before proceeding. Evidence: `.omo/evidence/task-32-fix-mcp-session-lifecycle.txt` - Commit: N - -- [x] 33. Automated end-to-end ACP test (replaces manual test per Metis GAP-15) - What to do / Must NOT do: Create automated integration test in `tests/agentpool_server/acp_server/test_e2e_session_lifecycle.py`: (1) Start ACP server in-process with a config that has MCP servers (use TestModel), (2) Create a mock WebSocket client that connects, (3) Create session, (4) Use MCP tool (mock), (5) Simulate WebSocket disconnect (close the mock connection), (6) Reconnect with new mock client, (7) Resume session, (8) Verify MCP tools work with fresh connections (assert toolset objects are different from pre-disconnect). Use `@pytest.mark.integration` and `@pytest.mark.slow`. Must NOT require a real ACP client or real model API key. - Parallelization: Wave 7 | Blocked by: ALL | Blocks: — - References: `agentpool serve-acp config.yml` command, example configs in `site/examples/*/config.yml` - Acceptance criteria: All 8 steps complete successfully. MCP tools work after reconnect+resume. - QA scenarios: (happy) Full flow works, MCP tools functional after resume. (failure) MCP tools fail after resume — indicates stale resources. Evidence: `.omo/evidence/task-33-fix-mcp-session-lifecycle.txt` - Commit: N - -## Final verification wave -> Runs in parallel after ALL todos. ALL must APPROVE. Surface results and wait for the user's explicit okay before declaring complete. -- [x] F1. Plan compliance audit — verify every task in `openspec/changes/fix-mcp-session-lifecycle/tasks.md` is implemented and checked off. Compare task-by-task. -- [x] F2. Code quality review — `uv run ruff check src/` and `uv run --no-group docs mypy src/` both pass with zero errors. Review changed code for `getattr`/`hasattr` usage (forbidden), missing type annotations, TODOs left in code. -- [x] F3. Real manual QA — run the manual ACP test from T33: connect → session → MCP tool → disconnect → reconnect → resume → verify MCP tools work. Capture output as evidence. -- [x] F4. Scope fidelity — verify NO changes to `MessageNode`, `AgentPool` registry, `MCPResourceProvider` model, or config API. Verify NO Phase 2 features were introduced. Verify all 5 stale-mcp tests are now fix-verifying (not bug-documenting). - -## Commit strategy -- One commit per todo that has `Commit: Y` (28 commits) -- Todos with `Commit: N` (T29-T33 verification) are verification-only, no commits -- Commit message format: `(): ` matching repo style -- Types: `feat`, `fix`, `refactor`, `test` -- Scopes: `mcp`, `acp`, `agent`, `session` -- Branch: `fix-mcp-session-lifecycle` (already created as worktree) - -## Success criteria -1. `uv run pytest tests/mcp_server/ tests/agentpool_server/acp_server/ -v` — 0 failures -2. `uv run pytest -m unit` — 0 failures -3. `uv run ruff check src/` — 0 errors -4. `uv run --no-group docs mypy src/` — 0 errors on changed files -5. Manual ACP test (T33) — MCP tools work after WebSocket disconnect + reconnect + resume -6. All 5 tests in `test_stale_mcp_connection.py` verify the fix (not the bug) -7. `_session_contexts` is empty after session close on all close paths (ACPSession.close, SessionController.close_session, WebSocket disconnect) -8. `resume_session()` creates fresh session, not returning stale one diff --git a/openspec/changes/m4-multi-config/tasks.md b/openspec/changes/m4-multi-config/tasks.md index ac1935f22..b4c9f36a5 100644 --- a/openspec/changes/m4-multi-config/tasks.md +++ b/openspec/changes/m4-multi-config/tasks.md @@ -119,6 +119,20 @@ - [ ] 14.5 Implement `AgentPool.from_registry(registry: ConfigRegistry, config_id: str = "default") -> AgentPool` classmethod — retrieves config from registry, initializes infrastructure from HostConfig, constructs HostContext with `config_id`, compiles agents via AgentFactory - [ ] 14.6 Verify `async with AgentPool("config.yml")` still works: internally creates ConfigRegistry, registers file with `config_id="default"`, delegates to `from_registry()` - [ ] 14.7 Write unit tests: `from_registry` produces correct HostContext with config_id, file-path constructor preserves all behavior, HostContext model_cache is per-Host scoped, HostContext reconstruction preserves config_id/tenant_id +- [ ] 14.8 Remove pool-level `input_provider` fallback in `NodeContext.get_input_provider()` (step 3 of 4 in `src/agentpool/messaging/context.py:60-61`). This fallback accesses `self.pool._input_provider` (private field on `AgentPool`), which is already deprecated — `BaseAgent._input_provider` warns to "Use SessionState.input_provider instead". The session-level provider (step 2) and ContextVar fallback (step 4) already cover all cases. After removal, `NodeContext.pool` is only used for `prompt_manager` (which `HostContext` already has), enabling the `NodeContext.pool` → `NodeContext.host` migration in task 14.9. + - Verify: `grep -n 'pool\._input_provider' src/agentpool/messaging/context.py` returns 0 + - Verify: `uv run pytest tests/ -x` passes (no test relies on pool-level input_provider fallback) +- [ ] 14.9 Migrate `NodeContext.pool: AgentPool | None` → `NodeContext.host: HostContext | None` in `src/agentpool/messaging/context.py`: + - Update `get_input_provider()` to use `self.host.input_provider` instead of `self.pool._input_provider` + - Update `prompt_manager` property to use `self.host.prompt_manager` instead of `self.pool.prompt_manager` + - Update `TeamContext(pool=...)` → `TeamContext(host=...)` in `base_team.py:get_context()` + - Update `AgentContext(pool=...)` → `AgentContext(host=...)` in `base_agent.py:get_context()` + - Verify: `grep -rn 'NodeContext.*pool' src/` returns 0 (field renamed) + - Verify: `grep -rn '\.pool\b' src/agentpool/messaging/context.py` returns 0 +- [ ] 14.10 Move `get_skill_instructions_for_node()` from `AgentPool` to `SkillsManager` — update `base_team.py:_load_skill_instructions()` to use `self.host_context.skills_registry.get_skill_instructions_for_node()` instead of `self._agent_pool.get_skill_instructions_for_node()`. Also move `skill_provider` property if needed. + - Verify: `grep -rn 'get_skill_instructions_for_node' src/agentpool/delegation/` uses `skills_registry` not `_agent_pool` +- [ ] 14.11 Audit remaining `_agent_pool` references in protocol servers (`acp_server/acp_agent.py:254,1114`, `opencode_server/state.py:119`, `opencode_server/routes/agent_routes.py:162`) and migrate to `host_context` accessors where possible. References that need the agent registry (for `get_agent()` / `register()`) should use `AgentRegistry` interface from `AgentContext.agent_registry`. Defer any `from_callback(agent_pool=)` refactoring to M5. + - Verify: `grep -rn '\._agent_pool' src/agentpool_server/` returns 0 (or only M5-deferred `from_callback` sites) ## 15. Hot Reload: Triggers and Turn-Level Snapshot diff --git a/src/acp/agent/acp_agent_api.py b/src/acp/agent/acp_agent_api.py index b947c2192..e7120a3bf 100644 --- a/src/acp/agent/acp_agent_api.py +++ b/src/acp/agent/acp_agent_api.py @@ -3,7 +3,7 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable from acp.schema import ( AuthenticateRequest, @@ -22,7 +22,7 @@ if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import AsyncIterator, Sequence from acp.agent.protocol import Agent from acp.schema import ( @@ -35,6 +35,7 @@ NewSessionResponse, PromptResponse, ResumeSessionResponse, + SessionUpdate, SetSessionConfigOptionResponse, SetSessionModelResponse, SetSessionModeResponse, @@ -42,19 +43,62 @@ from acp.schema.mcp import McpServer +@runtime_checkable +class _SessionStateProtocol(Protocol): + """Protocol for session state objects that ACPAgentAPI can poll for updates.""" + + def pop_update(self) -> SessionUpdate | None: ... + def clear(self) -> None: ... + + +@runtime_checkable +class _UpdateEventProtocol(Protocol): + """Protocol for update events that ACPAgentAPI can wait on.""" + + async def wait_with_timeout(self, timeout: float | None = None) -> bool: ... + def clear(self) -> None: ... + + class ACPAgentAPI: """Thin wrapper for client-to-agent ACP interactions. Avoids manual instantiation of request/notification objects. + + When optional ``state`` and ``update_event`` are provided, the instance + also satisfies the :class:`~agentpool.agents.acp_agent.turn.ACPClientProtocol` + protocol by implementing :meth:`stream_events` and :meth:`get_messages`. """ - def __init__(self, connection: Agent) -> None: + def __init__( + self, + connection: Agent, + *, + state: _SessionStateProtocol | None = None, + update_event: _UpdateEventProtocol | None = None, + ) -> None: """Initialize agent API helper. Args: connection: The Agent protocol connection (e.g., ClientSideConnection) + state: Optional session state for polling updates (enables stream_events) + update_event: Optional event signaled when new updates arrive """ self.connection = connection + self._state: _SessionStateProtocol | None = state + self._update_event: _UpdateEventProtocol | None = update_event + self._consumed_updates: list[SessionUpdate] = [] + + def _attach_state( + self, + state: _SessionStateProtocol, + update_event: _UpdateEventProtocol, + ) -> None: + """Attach state and update event after construction. + + Allows deferred wiring when state/event are created after the API. + """ + self._state = state + self._update_event = update_event async def initialize( self, @@ -224,3 +268,45 @@ async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any async def ext_notification(self, method: str, params: dict[str, Any]) -> None: """Send an extension notification to the agent.""" await self.connection.ext_notification(method, params) + + async def stream_events( + self, + response: PromptResponse, + ) -> AsyncIterator[SessionUpdate]: + """Yield raw ACP session updates from the state queue. + + Polls :meth:`_SessionStateProtocol.pop_update` in a loop, waiting + up to 50 ms between drain cycles for new updates to arrive via + ``_update_event``. Once a full drain cycle produces no updates, + the iterator ends. + + Updates are also collected in ``_consumed_updates`` so that + :meth:`get_messages` can return them after streaming completes. + + Args: + response: The prompt response (unused — updates come from state) + """ + self._consumed_updates.clear() + if self._state is None or self._update_event is None: + return + while True: + try: + await self._update_event.wait_with_timeout(0.05) + self._update_event.clear() + except TimeoutError: + pass + drained_any = False + while (update := self._state.pop_update()) is not None: + self._consumed_updates.append(update) + yield update + drained_any = True + if not drained_any: + break + + async def get_messages(self, session_id: str) -> list[SessionUpdate]: + """Return all session updates consumed during :meth:`stream_events`. + + Args: + session_id: The ACP session ID (unused — updates are already collected) + """ + return list(self._consumed_updates) diff --git a/src/agentpool/agents/acp_agent/acp_agent.py b/src/agentpool/agents/acp_agent/acp_agent.py index 096ec9af6..b1c7e853a 100644 --- a/src/agentpool/agents/acp_agent/acp_agent.py +++ b/src/agentpool/agents/acp_agent/acp_agent.py @@ -34,24 +34,23 @@ from datetime import datetime import os from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Self, cast +from typing import TYPE_CHECKING, Any, ClassVar, Self import uuid import anyio from pydantic import HttpUrl from pydantic_ai import ( - ModelRequest, - ModelResponse, TextPart, + ThinkingPart, + ToolCallPart, ToolReturnPart, UserContent, - UserPromptPart, ) from acp import InitializeRequest from acp.agent import ACPAgentAPI from agentpool.agents.acp_agent.session_state import ACPSessionState -from agentpool.agents.acp_agent.turn import ACPTurn +from agentpool.agents.acp_agent.turn import ACPClientProtocol, ACPTurn from agentpool.agents.base_agent import BaseAgent from agentpool.agents.events import ( RunStartedEvent, @@ -67,7 +66,6 @@ ) from agentpool.log import get_logger from agentpool.messaging import ChatMessage -from agentpool.orchestrator.core import EventEnvelope from agentpool.utils.subprocess_utils import SubprocessError, run_with_process_monitor from agentpool.utils.token_breakdown import calculate_usage_from_parts @@ -79,7 +77,6 @@ from anyio.abc import Process from evented_config import EventConfig from exxec import ExecutionEnvironment - from pydantic_ai import ThinkingPart, ToolCallPart, UserContent from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.messages import ModelMessage from slashed import BaseCommand @@ -90,7 +87,6 @@ from acp.schema.capabilities import AgentCapabilities from acp.schema.mcp import McpServer from agentpool.agents.acp_agent.client_handler import ACPClientHandler - from agentpool.agents.acp_agent.turn import ACPClientProtocol from agentpool.agents.context import AgentRunContext from agentpool.agents.events import RichAgentStreamEvent from agentpool.agents.modes import ModeCategory @@ -350,6 +346,10 @@ async def _initialize(self) -> None: output_stream=self._process.stdout, ) self._api = ACPAgentAPI(self._connection) + self._api._attach_state( + state=self._state, + update_event=self._client_handler._update_event, + ) init_response = await self._connection.initialize(self._init_request) self._agent_info = init_response.agent_info self._caps = init_response.agent_capabilities @@ -426,12 +426,8 @@ async def _stream_events( # noqa: PLR0915 wait_for_connections: bool | None = None, store_history: bool = True, ) -> AsyncIterator[RichAgentStreamEvent[str]]: - from agentpool.agents.acp_agent.acp_converters import ( - convert_to_acp_content, - to_finish_reason, - ) + from agentpool.agents.acp_agent.acp_converters import to_finish_reason - # Update input provider if provided if input_provider is not None and self._client_handler: self._client_handler._input_provider = input_provider if not self._api or not self._sdk_session_id or not self._state: @@ -439,11 +435,6 @@ async def _stream_events( # noqa: PLR0915 run_id = str(uuid.uuid4()) self._state.clear() - model_messages: list[ModelResponse | ModelRequest] = [] - initial_request = ModelRequest(parts=[UserPromptPart(content=prompts)]) - model_messages.append(initial_request) - current_response_parts: list[TextPart | ThinkingPart | ToolCallPart] = [] - text_chunks: list[str] = [] assert session_id is not None yield RunStartedEvent( @@ -452,7 +443,7 @@ async def _stream_events( # noqa: PLR0915 agent_name=self.name, parent_session_id=parent_session_id, ) - final_blocks = convert_to_acp_content(prompts) + # Handle ephemeral execution (fork session if store_history=False) acp_session_id = self._sdk_session_id if not store_history and self._sdk_session_id: @@ -460,98 +451,59 @@ async def _stream_events( # noqa: PLR0915 fork_response = await self._api.fork_session(self._sdk_session_id, cwd) acp_session_id = fork_response.session_id self.log.debug("Forked session", parent=self._sdk_session_id, fork=acp_session_id) - self.log.debug("Starting streaming prompt", num_blocks=len(final_blocks)) - prompt_task = asyncio.create_task(self._api.prompt(acp_session_id, final_blocks)) - self._prompt_task = prompt_task - - async def poll_acp_events() -> AsyncIterator[RichAgentStreamEvent[str]]: - """Poll raw updates from ACP state, convert to events, until prompt completes.""" - from agentpool.agents.acp_agent.acp_converters import acp_to_native_event - - assert self._state - while not prompt_task.done(): - if self._client_handler: - try: - await self._client_handler._update_event.wait_with_timeout(0.05) - self._client_handler._update_event.clear() - except TimeoutError: - pass - while (update := self._state.pop_update()) is not None: - if native_event := acp_to_native_event(update): - yield native_event - while (update := self._state.pop_update()) is not None: - if native_event := acp_to_native_event(update): - yield native_event + + # Create ACPTurn and delegate to execute() + assert self._api is not None + _acp_client: ACPClientProtocol = self._api + turn = ACPTurn( + acp_client=_acp_client, + prompts=prompts, + run_ctx=run_ctx, + session_id=acp_session_id, + agent_name=self.name, + hooks=self.hooks, + env=self.env, + ) tool_metadata: dict[str, dict[str, Any]] = {} + text_chunks: list[str] = [] + current_response_parts: list[TextPart | ThinkingPart | ToolCallPart] = [] try: agent_ctx = self.get_context(run_ctx=run_ctx, input_provider=input_provider) async with self._tool_bridge.set_run_context(agent_ctx, prompt=prompts): - send_stream, receive_stream = anyio.create_memory_object_stream( - max_buffer_size=1000 - ) - - async def _forward_acp_events() -> None: - try: - async for event in poll_acp_events(): - try: - await send_stream.send(event) - except (anyio.ClosedResourceError, anyio.BrokenResourceError): - return - finally: - await send_stream.aclose() - - # Do NOT subscribe to run_ctx.event_bus here: in standalone mode - # the producer publishes _stream_events() output back into the - # same local EventBus, creating a self-echo infinite loop. - _bg_tasks: set[asyncio.Task[Any]] = set() - task_a = asyncio.create_task(_forward_acp_events()) - _bg_tasks.add(task_a) - task_a.add_done_callback(_bg_tasks.discard) - - try: - async for raw_event in receive_stream: - event = ( - raw_event.event if isinstance(raw_event, EventEnvelope) else raw_event - ) - if isinstance(event, ToolResultMetadataEvent): - tool_metadata[event.tool_call_id] = event.metadata - continue - if run_ctx.cancelled: - self.log.info("Stream cancelled by user") - break - if isinstance(event, ToolCallCompleteEvent): - enriched_event = event - if not enriched_event.agent_name: - enriched_event = replace(enriched_event, agent_name=self.name) - if ( - enriched_event.metadata is None - and enriched_event.tool_call_id in tool_metadata - ): - enriched_event = replace( - enriched_event, - metadata=tool_metadata[enriched_event.tool_call_id], - ) - output_event = enriched_event - else: - output_event = event - part = event_to_part(output_event) - if isinstance(part, TextPart): - text_chunks.append(part.content) - if part and not isinstance(part, ToolReturnPart): - current_response_parts.append(part) - yield output_event - finally: - for t in list(_bg_tasks): - t.cancel() - for t in list(_bg_tasks): - try: - await t - except asyncio.CancelledError: - pass - except Exception: - self.log.exception("Error during background task cleanup") + async for event in turn.execute(): + # Capture ToolResultMetadataEvent for enrichment (not yielded) + if isinstance(event, ToolResultMetadataEvent): + tool_metadata[event.tool_call_id] = event.metadata + continue + # Check cancellation + if run_ctx.cancelled: + self.log.info("Stream cancelled by user") + break + # Enrich ToolCallCompleteEvent with metadata and agent_name + output_event: RichAgentStreamEvent[str] = event + if isinstance(event, ToolCallCompleteEvent): + enriched = event + if not enriched.agent_name: + enriched = replace(enriched, agent_name=self.name) + if enriched.metadata is None and enriched.tool_call_id in tool_metadata: + enriched = replace( + enriched, + metadata=tool_metadata[enriched.tool_call_id], + ) + output_event = enriched + # Don't yield StreamCompleteEvent from ACPTurn — + # we build our own with usage/cost and finish_reason + if isinstance(output_event, StreamCompleteEvent): + break + # Track parts for usage calculation and text for cancellation + part = event_to_part(output_event) + if isinstance(part, TextPart): + text_chunks.append(part.content) + if part and not isinstance(part, ToolReturnPart): + current_response_parts.append(part) + yield output_event except asyncio.CancelledError: self.log.info("Stream cancelled via task cancellation") run_ctx.cancelled = True @@ -565,7 +517,7 @@ async def _forward_acp_events() -> None: session_id=session_id, parent_id=user_msg.message_id, model_name=self.model_name, - messages=model_messages, + messages=[], metadata={}, finish_reason="stop", ) @@ -573,19 +525,19 @@ async def _forward_acp_events() -> None: self._prompt_task = None return - response = await prompt_task - finish_reason = to_finish_reason(response.stop_reason) - if current_response_parts: - model_messages.append( - ModelResponse( - parts=current_response_parts, - finish_reason=finish_reason, - model_name=self.model_name, - provider_name=self._provider_type, - ) - ) + # Build enriched StreamCompleteEvent with finish_reason and usage/cost + final_message = turn.final_message + finish_reason = ( + to_finish_reason(turn._prompt_response.stop_reason) + if turn._prompt_response is not None + else "stop" + ) - text_content = "".join(text_chunks) + text_content = ( + final_message.content + if isinstance(final_message.content, str) + else "".join(text_chunks) + ) usage, cost_info = await calculate_usage_from_parts( input_parts=prompts, response_parts=current_response_parts, @@ -602,7 +554,7 @@ async def _forward_acp_events() -> None: session_id=session_id, parent_id=user_msg.message_id, model_name=self.model_name, - messages=model_messages, + messages=turn.message_history, metadata={}, finish_reason=finish_reason, usage=usage, @@ -645,16 +597,12 @@ def create_turn( Returns: An ACPTurn instance for single-cycle execution. """ - # TODO: ACPAgentAPI does not implement ACPClientProtocol fully — - # it lacks stream_events() and get_messages(). At runtime this will raise - # AttributeError when ACPTurn.execute() calls those methods. An adapter - # wrapping ACPAgentAPI with async futures / notification registry is needed - # for full integration. + assert self._api is not None + _acp_client: ACPClientProtocol = self._api return ACPTurn( - acp_client=cast("ACPClientProtocol", self._api), - prompts=prompts, # type: ignore[arg-type] + acp_client=_acp_client, + prompts=prompts, run_ctx=run_ctx, - message_history=message_history, session_id=self._sdk_session_id or run_ctx.session_id, agent_name=self.name, hooks=self.hooks, diff --git a/src/agentpool/agents/acp_agent/turn.py b/src/agentpool/agents/acp_agent/turn.py index 2a5de9288..b6549cf8d 100644 --- a/src/agentpool/agents/acp_agent/turn.py +++ b/src/agentpool/agents/acp_agent/turn.py @@ -10,9 +10,12 @@ import asyncio import time -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Protocol, runtime_checkable from uuid import uuid4 +from pydantic import ValidationError + +from acp.exceptions import RequestError from agentpool.agents.events import ( RunErrorEvent, StreamCompleteEvent, @@ -24,7 +27,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Sequence from typing import Any - from pydantic_ai import ModelMessage + from pydantic_ai import ModelMessage, UserContent from acp.schema import ContentBlock, PromptResponse, SessionUpdate from agentpool.agents.context import AgentRunContext @@ -33,6 +36,7 @@ from agentpool.messaging import ChatMessage +@runtime_checkable class ACPClientProtocol(Protocol): """Protocol defining the ACP client interface expected by ACPTurn. @@ -102,9 +106,8 @@ class ACPTurn(HookAwareTurn, Turn): def __init__( self, acp_client: ACPClientProtocol, - prompts: list[str], + prompts: list[UserContent], run_ctx: AgentRunContext, - message_history: list[ModelMessage], session_id: str, agent_name: str | None = None, hooks: AgentHooks | None = None, @@ -118,6 +121,7 @@ def __init__( self._agent_name = agent_name self._hooks = hooks self._agent_env = env + self._prompt_response: PromptResponse | None = None @property def _hook_env(self) -> Any | None: @@ -169,15 +173,16 @@ async def execute(self) -> AsyncGenerator[RichAgentStreamEvent[Any]]: # noqa: P return # Convert all user prompts to ACP ContentBlock list. # Join all prompts instead of taking only the last one. - full_prompt = "\n\n".join(self._prompts) if self._prompts else "" + full_prompt = "\n\n".join(str(p) for p in self._prompts) if self._prompts else "" content = convert_to_acp_content([full_prompt]) # --- Phase 1: Send prompt --- try: response = await self._acp_client.prompt(self._session_id, content) + self._prompt_response = response except asyncio.CancelledError: raise - except Exception as exc: # noqa: BLE001 + except (RequestError, ConnectionError, RuntimeError, ValidationError) as exc: yield RunErrorEvent( message=str(exc), run_id=run_id, @@ -218,7 +223,7 @@ async def execute(self) -> AsyncGenerator[RichAgentStreamEvent[Any]]: # noqa: P yield native_event except asyncio.CancelledError: raise - except Exception as exc: # noqa: BLE001 + except (RequestError, ConnectionError, RuntimeError, ValueError) as exc: yield RunErrorEvent( message=str(exc), run_id=run_id, @@ -231,7 +236,7 @@ async def execute(self) -> AsyncGenerator[RichAgentStreamEvent[Any]]: # noqa: P raw_updates = await self._acp_client.get_messages(self._session_id) except asyncio.CancelledError: raise - except Exception as exc: # noqa: BLE001 + except (RequestError, ConnectionError, ValidationError) as exc: yield RunErrorEvent( message=str(exc), run_id=run_id, diff --git a/src/agentpool/agents/base_agent.py b/src/agentpool/agents/base_agent.py index 5c0d2049e..26758a6e3 100644 --- a/src/agentpool/agents/base_agent.py +++ b/src/agentpool/agents/base_agent.py @@ -11,7 +11,7 @@ import os from pathlib import Path import re -from typing import TYPE_CHECKING, Any, ClassVar, Literal, assert_never, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, assert_never, overload import warnings from anyenv import MultiEventHandler, method_spawner @@ -1028,11 +1028,11 @@ def queue_prompt(self, *prompts: PromptCompatible, session_id: str | None = None processed in a continuation loop without exiting the stream. This allows tools or external code to schedule follow-up work. - !!! warning "Deprecated for pooled native agents" - Use ``agent_pool.session_pool.followup()`` instead. + !!! warning "Deprecated for pooled agents" + Use ``host_context.session_pool.followup()`` instead. - For non-native agents and standalone native agents, the existing - injection_manager-based path remains unchanged. + For standalone agents (no session pool), the injection_manager-based + path is used as a fallback. Args: *prompts: Prompts to queue (same format as run/run_stream) @@ -1046,11 +1046,11 @@ async def my_tool(ctx: AgentContext) -> str: """ run_ctx = self.get_active_run_context(session_id=session_id) - # Pooled native agents: delegate to session_pool.followup(). + # Pooled agents: delegate to session_pool.followup(). ctx = self.host_context - if self.AGENT_TYPE == "native" and ctx is not None and ctx.session_pool is not None: + if ctx is not None and ctx.session_pool is not None: warnings.warn( - "queue_prompt() is deprecated for pooled native agents. " + "queue_prompt() is deprecated for pooled agents. " "Use host_context.session_pool.followup() instead.", DeprecationWarning, stacklevel=2, @@ -1065,9 +1065,8 @@ async def my_tool(ctx: AgentContext) -> str: session_pool.followup(effective_session_id, combined) ) return - # Standalone native agents: fall through to legacy path - # Legacy path for non-native agents and standalone native agents + # Standalone agents: use injection_manager if run_ctx is not None and run_ctx.injection_manager is not None: combined = "\n".join(str(p) for p in prompts) run_ctx.injection_manager.inject(combined) @@ -1080,11 +1079,11 @@ def inject_prompt(self, message: str, session_id: str | None = None) -> None: iteration completes, the message is automatically queued for the next iteration. - !!! warning "Deprecated for pooled native agents" - Use ``agent_pool.session_pool.steer()`` instead. + !!! warning "Deprecated for pooled agents" + Use ``host_context.session_pool.steer()`` instead. - For non-native agents and standalone native agents, the existing - injection_manager-based path remains unchanged. + For standalone agents (no session pool), the injection_manager-based + path is used as a fallback. Args: message: Message to inject @@ -1100,10 +1099,10 @@ async def my_tool(ctx: AgentContext) -> str: run_ctx = self.get_active_run_context(session_id=session_id) ctx = self.host_context - # Pooled native agents: delegate to session_pool.steer(). - if self.AGENT_TYPE == "native" and ctx is not None and ctx.session_pool is not None: + # Pooled agents: delegate to session_pool.steer(). + if ctx is not None and ctx.session_pool is not None: warnings.warn( - "inject_prompt() is deprecated for pooled native agents. " + "inject_prompt() is deprecated for pooled agents. " "Use host_context.session_pool.steer() instead.", DeprecationWarning, stacklevel=2, @@ -1127,51 +1126,10 @@ async def my_tool(ctx: AgentContext) -> str: session_pool.steer(most_recent.session_id, message) ) return - # Standalone native agents: fall through to legacy path - # Legacy path for non-native agents and standalone native agents - # CRITICAL: Check run_ctx.completed to avoid injecting into a turn that - # has already finished (e.g., after end_turn). If the turn is complete, - # the message would be stuck in injection_manager.pending forever. - # In that case, delegate to SessionPool for auto-resume. + # Standalone agents: use injection_manager if run_ctx is not None and not run_ctx.completed and run_ctx.injection_manager is not None: run_ctx.injection_manager.inject(message) - return - - # No active run context — delegate to SessionPool for auto-resume - effective_session_id = session_id or (run_ctx.session_id if run_ctx else None) - if ctx is not None and effective_session_id is not None: - _session_pool = ctx.session_pool - if _session_pool is None: - return - # Fire-and-forget: delegate to SessionPool for auto-resume. - # Use task_manager to prevent GC of the task mid-execution. - self.task_manager.fire_and_forget( - _session_pool.inject_prompt(effective_session_id, message) - ) - return - - # FALLBACK for shared agents: effective_session_id is None but session_pool exists. - # This handles the case where BackgroundTaskProvider calls inject_prompt - # after background task completion when the agent has no fixed session_id. - if ctx is not None: - _session_pool = ctx.session_pool - if _session_pool is not None: - sessions = _session_pool.sessions.find_sessions_by_agent_name(self.name) - if sessions: - most_recent = max(sessions, key=lambda s: s.last_active_at) - self.task_manager.fire_and_forget( - _session_pool.receive_request( - most_recent.session_id, message, priority="asap" - ) - ) - return - - # No pool or session_id available — log warning - self.log.warning( - "inject_prompt called but no active run context or session pool available", - agent_name=self.name, - ) def has_pending_injections(self, session_id: str | None = None) -> bool: """Check if there are pending injections. @@ -1544,10 +1502,6 @@ async def _run_stream_once( """ from agentpool.messaging import ChatMessage - # Clear hooks_fired so the new turn's hooks can fire - # even if the previous turn already fired them. - run_ctx.hooks_fired.clear() - # Convert prompts to standard UserContent format converted_prompts = await convert_prompts(prompts) # Prepend any staged content @@ -1587,31 +1541,6 @@ async def _run_stream_once( conversation.add_chat_messages([user_msg]) try: - # Execute pre-turn hooks (guarded against double-firing with HookAwareTurn) - # Native agents fire hooks via HookAwareTurn in NativeTurn.execute(); - # only ACP standalone path still uses this old hook firing. - if self.AGENT_TYPE != "native" and self.hooks and "pre_turn" not in run_ctx.hooks_fired: - run_ctx.hooks_fired.add("pre_turn") - pre_turn_result = await self.hooks.run_pre_turn_hooks( - agent_name=self.name, - prompt=user_msg.content - if isinstance(user_msg.content, str) - else str(user_msg.content), - session_id=session_id, - ) - if pre_turn_result.get("decision") == "deny": - run_ctx.cancelled = True - cancel_msg = ChatMessage( - content="", - role="assistant", - name=self.name, - session_id=session_id, - ) - yield StreamCompleteEvent( - message=cast("ChatMessage[TResult]", cancel_msg), cancelled=True - ) - return - async for event in self._stream_events( run_ctx, [*pending_parts, *converted_prompts], @@ -1648,27 +1577,6 @@ async def _run_stream_once( # TaskGroup cancellation from interrupting hooks/routing/persistence if final_message is not None: with anyio.CancelScope(shield=True): - # Execute post-turn hooks (guarded against double-firing with HookAwareTurn) - # Native agents fire hooks via HookAwareTurn in NativeTurn.execute(); - # only ACP standalone path still uses this old hook firing. - if ( - self.AGENT_TYPE != "native" - and self.hooks - and "post_turn" not in run_ctx.hooks_fired - ): - run_ctx.hooks_fired.add("post_turn") - prompt_str = ( - user_msg.content - if isinstance(user_msg.content, str) - else str(user_msg.content) - ) - await self.hooks.run_post_turn_hooks( - agent_name=self.name, - prompt=prompt_str, - result=final_message.content, - session_id=session_id, - ) - # Emit signal (always - for event handlers). # Skip when run_ctx was provided by SessionPool (Path A); # the Path A wrapper in run_stream() handles emission in that case. diff --git a/src/agentpool/agents/context.py b/src/agentpool/agents/context.py index eea60bc4c..9e530ac2f 100644 --- a/src/agentpool/agents/context.py +++ b/src/agentpool/agents/context.py @@ -94,14 +94,6 @@ class AgentRunContext: cancelled: bool = False """Whether the run has been cancelled.""" - hooks_fired: set[str] = field(default_factory=set) - """Tracks which hook events have fired this turn to prevent double-firing. - - Cleared at the start of each turn by ``RunHandle.start()`` and - ``_run_stream_once()``. Entries are event names like ``"pre_turn"``, - ``"post_turn"``, ``"pre_tool_use:{tool_call_id}"``. - """ - run_id: str = field(default_factory=lambda: uuid.uuid4().hex) """Unique identifier for this run.""" diff --git a/src/agentpool/agents/native_agent/agent.py b/src/agentpool/agents/native_agent/agent.py index 319489ad7..154963963 100644 --- a/src/agentpool/agents/native_agent/agent.py +++ b/src/agentpool/agents/native_agent/agent.py @@ -75,8 +75,7 @@ ) from agentpool.delegation import AgentPool from agentpool.hooks import AgentHooks - from agentpool.mcp_server.config_snapshot import McpConfigEntry, McpConfigSnapshot - from agentpool.mcp_server.session_pool import SessionConnectionPool + from agentpool.mcp_server.config_snapshot import McpConfigEntry from agentpool.messaging import MessageNode from agentpool.models.agents import NativeAgentConfig, ToolMode from agentpool.orchestrator.turn import Turn @@ -331,11 +330,6 @@ def __init__( # noqa: PLR0915 self._providers = list(providers) if providers else None # model discovery self._direct_history_processors = list(history_processors) if history_processors else None self._resolved_history_processors: list[Callable[..., Any]] | None = None - # MCP lifecycle snapshot — set externally (e.g. by SessionController) to - # enable snapshot-aware capability building in get_agentlet(). - # When None, get_agentlet() falls back to the legacy get_capabilities() path. - self._mcp_snapshot: McpConfigSnapshot | None = None - self._session_connection_pool: SessionConnectionPool | None = None self._extra_capabilities: list[Any] = capabilities or [] def _build_pool_configs(self) -> tuple[McpConfigEntry, ...]: @@ -911,9 +905,9 @@ async def get_agentlet[AgentOutputType]( # noqa: PLR0915 tool_capabilities.extend(mcp_capabilities) # 5. Skill capabilities — from pool-scoped instances created during __aenter__. # Each SkillManagerCap provides tools and MCP servers. - ctx = self.host_context - if ctx is not None and ctx.pool is not None: - pool_capabilities = ctx.pool.skill_capabilities + pool = self._agent_pool + if pool is not None: + pool_capabilities = pool.skill_capabilities if pool_capabilities: from agentpool.capabilities.skill_manager_cap import SkillManagerCap diff --git a/src/agentpool/capabilities/mcp_server_cap.py b/src/agentpool/capabilities/mcp_server_cap.py index b02eac27c..38b97b571 100644 --- a/src/agentpool/capabilities/mcp_server_cap.py +++ b/src/agentpool/capabilities/mcp_server_cap.py @@ -160,6 +160,10 @@ async def _ensure_client(self) -> MCPClient: # Success — set up change notification callbacks. async def _on_tools_changed() -> None: + # Cross-layer wiring: this ChangeEvent(kind="tools_changed") is + # consumed by the OpenCode server's _watch_mcp_tool_changes task + # (server.py) which converts it to McpToolsChangedEvent and + # broadcasts it as an SSE event to connected clients. event = ChangeEvent( capability_name=self._name, kind="tools_changed", diff --git a/src/agentpool/delegation/base_team.py b/src/agentpool/delegation/base_team.py index 337a9723a..5a69c05f8 100644 --- a/src/agentpool/delegation/base_team.py +++ b/src/agentpool/delegation/base_team.py @@ -408,7 +408,10 @@ def get_context( shared_pool: AgentPool | None = None for agent in self.iter_agents(): - pool = agent.host_context.pool if agent.host_context else None + # TODO(m4): Migrate to `agent.host_context` once NodeContext.pool + # is replaced with NodeContext.host: HostContext | None. + # Requires adding `input_provider` to HostContext (see M4 task 14.8). + pool = agent._agent_pool if pool: pool_id = id(pool) if pool_id not in pool_ids: @@ -772,7 +775,7 @@ async def _load_member_skill_instructions( self, member_skills: dict[str, list[str]], ) -> dict[str, str]: - if not member_skills or self.host_context is None: + if not member_skills or self._agent_pool is None: return {} result: dict[str, str] = {} @@ -787,8 +790,7 @@ async def _load_member_skill_instructions( return result async def _load_skill_instructions(self, skill_name: str, member_name: str) -> str: - ctx = self.host_context - pool = ctx.pool if ctx is not None else None + pool = self._agent_pool if pool is None or pool.skill_provider is None: from agentpool.skills.exceptions import SkillNotFoundError diff --git a/src/agentpool/delegation/pool.py b/src/agentpool/delegation/pool.py index ff7e002d2..b12723569 100644 --- a/src/agentpool/delegation/pool.py +++ b/src/agentpool/delegation/pool.py @@ -253,7 +253,6 @@ def get_context(self) -> HostContext: session_pool=self._session_pool, config_file_path=self._config_file_path, main_agent_name=self._safe_main_agent_name(), - pool=self, extension_registry=self._extension_registry, ) return self._host_context diff --git a/src/agentpool/host/context.py b/src/agentpool/host/context.py index 303dc6322..acfbd212c 100644 --- a/src/agentpool/host/context.py +++ b/src/agentpool/host/context.py @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from agentpool.host.stubs import CapabilityCache, ModelCache, ModelRegistry @@ -20,7 +20,6 @@ from upathtools import UPath from agentpool.capabilities.extension_registry import ExtensionRegistry - from agentpool.delegation.pool import AgentPool from agentpool.mcp_server.manager import MCPManager from agentpool.models.manifest import AgentsManifest from agentpool.orchestrator import SessionPool @@ -83,5 +82,4 @@ class HostContext: model_registry: ModelRegistry = field(default_factory=ModelRegistry) model_cache: ModelCache = field(default_factory=ModelCache) main_agent_name: str | None = None - pool: AgentPool[Any] | None = None extension_registry: ExtensionRegistry | None = None diff --git a/src/agentpool/lifecycle/__init__.py b/src/agentpool/lifecycle/__init__.py index 2bfbca961..0fe63cf1a 100644 --- a/src/agentpool/lifecycle/__init__.py +++ b/src/agentpool/lifecycle/__init__.py @@ -38,6 +38,7 @@ Feedback, Prompt, ResumeResult, + RunOutcome, RunState, ToolExecutionRecord, ) @@ -60,6 +61,7 @@ "ProtocolChannel", "ProtocolTrigger", "ResumeResult", + "RunOutcome", "RunState", "ScheduledTrigger", "SnapshotStore", diff --git a/src/agentpool/lifecycle/comm_channel.py b/src/agentpool/lifecycle/comm_channel.py index a99913f0c..b38139f43 100644 --- a/src/agentpool/lifecycle/comm_channel.py +++ b/src/agentpool/lifecycle/comm_channel.py @@ -95,6 +95,25 @@ def queue(self) -> asyncio.Queue[Any]: """The internal event queue, accessible for RunLoop draining.""" return self._queue + @property + def publishes_to_event_bus(self) -> bool: + """DirectChannel does not publish to the EventBus. + + Returns: + Always ``False``. + """ + return False + + def set_replaying(self, flag: bool) -> None: + """Set the replaying flag. + + When ``True``, journaling is skipped during ``publish()``. + + Args: + flag: ``True`` to enable replaying mode, ``False`` to disable. + """ + self._replaying = flag + def attach(self, run_loop: Any) -> None: """Store a reference to the RunLoop. @@ -150,6 +169,20 @@ def recv(self) -> Feedback | None: """ return None + def deliver_feedback(self, feedback: Feedback) -> bool: + """Reject feedback (unidirectional channel). + + DirectChannel does not support feedback delivery. Returns + ``False`` so the caller can fall back to the queue-based path. + + Args: + feedback: Ignored. + + Returns: + Always ``False``. + """ + return False + def close(self) -> None: """Drain the queue and mark the channel as closed. @@ -197,6 +230,25 @@ def __init__( self._run_loop: Any = None self._state: RunState | None = None + def set_replaying(self, flag: bool) -> None: + """Set the replaying flag. + + When ``True``, journaling is skipped during ``publish()``. + + Args: + flag: ``True`` to enable replaying mode, ``False`` to disable. + """ + self._replaying = flag + + @property + def publishes_to_event_bus(self) -> bool: + """ProtocolChannel publishes to the EventBus internally. + + Returns: + Always ``True``. + """ + return True + def attach(self, run_loop: Any) -> None: """Store a reference to the RunLoop for feedback routing. @@ -260,15 +312,19 @@ def recv(self) -> Feedback | None: except asyncio.QueueEmpty: return None - def deliver_feedback(self, feedback: Feedback) -> None: + def deliver_feedback(self, feedback: Feedback) -> bool: """Enqueue feedback from SessionController. This is how steer/followup messages arrive at the RunLoop. Args: feedback: The feedback to enqueue. + + Returns: + Always ``True`` (ProtocolChannel supports feedback delivery). """ self._feedback_queue.put_nowait(feedback) + return True def close(self) -> None: """Clean up the feedback queue and mark as closed. diff --git a/src/agentpool/lifecycle/protocols.py b/src/agentpool/lifecycle/protocols.py index 81f9581a4..dfa9e88a7 100644 --- a/src/agentpool/lifecycle/protocols.py +++ b/src/agentpool/lifecycle/protocols.py @@ -210,7 +210,33 @@ class CommChannel(Protocol): (append for deltas, upsert for entity-state events). """ - _replaying: bool + def set_replaying(self, flag: bool) -> None: + """Set the replaying flag. + + When ``True``, the channel skips journaling on ``publish()``. + Used during crash recovery replay to avoid duplicate entries. + + Args: + flag: ``True`` to enable replaying mode, ``False`` to disable. + """ + ... + + @property + def publishes_to_event_bus(self) -> bool: + """Whether this channel publishes events to the EventBus internally. + + ``ProtocolChannel`` publishes to the EventBus inside its own + ``publish()`` method, so the RunLoop must NOT also call + ``event_bus.publish()`` directly to avoid double-publishing. + + ``DirectChannel`` does not publish to the EventBus, so the + direct call in the RunLoop is required. + + Returns: + ``True`` if the channel publishes to the EventBus + internally, ``False`` otherwise. + """ + ... def attach(self, run_loop: Any) -> None: """Store a reference to the RunLoop, enabling the feedback loop. @@ -252,6 +278,22 @@ def recv(self) -> Feedback | None: """ ... + def deliver_feedback(self, feedback: Feedback) -> bool: + """Deliver feedback (steer/followup) to the RunLoop. + + Bidirectional channels (``ProtocolChannel``) enqueue the + feedback and return ``True``. Unidirectional channels + (``DirectChannel``) do not support feedback and return + ``False`` so the caller can fall back to the queue-based path. + + Args: + feedback: The feedback message to deliver. + + Returns: + ``True`` if the feedback was handled, ``False`` otherwise. + """ + ... + def close(self) -> None: """Release all resources held by the CommChannel.""" ... diff --git a/src/agentpool/lifecycle/types.py b/src/agentpool/lifecycle/types.py index a116a7888..1ad6239e0 100644 --- a/src/agentpool/lifecycle/types.py +++ b/src/agentpool/lifecycle/types.py @@ -25,6 +25,23 @@ class RunState(Enum): DONE = "done" +class RunOutcome(Enum): + """Terminal outcome for a completed RunLoop. + + Set on ``RunHandle.outcome`` when the run reaches ``RunState.DONE``. + ``None`` means the run has not yet terminated (still ``IDLE`` or + ``RUNNING``) or was closed without a specific outcome. + + - ``COMPLETED`` — Run finished normally. + - ``FAILED`` — Run finished with an error. + - ``CHECKPOINTED`` — Run state persisted for later resumption. + """ + + COMPLETED = "completed" + FAILED = "failed" + CHECKPOINTED = "checkpointed" + + @dataclass class Prompt: """Incoming prompt delivered to the RunLoop by a TriggerSource. @@ -128,6 +145,7 @@ class EventEnvelope: "Feedback", "Prompt", "ResumeResult", + "RunOutcome", "RunState", "ToolExecutionRecord", ] diff --git a/src/agentpool/mcp_server/manager.py b/src/agentpool/mcp_server/manager.py index 6efa76e6a..7790ba278 100644 --- a/src/agentpool/mcp_server/manager.py +++ b/src/agentpool/mcp_server/manager.py @@ -117,7 +117,7 @@ async def _process_tool_call( @dataclass -class _SessionContext: +class McpSessionContext: """Per-session MCP state container for the :class:`MCPManager`. Holds all session-scoped MCP resources so that each session has its own @@ -177,7 +177,7 @@ def __init__( self._accessible_roots = accessible_roots self._global_pool = GlobalConnectionPool() self._toolset_cache: dict[str, Any] = {} - self._session_contexts: dict[str, _SessionContext] = {} + self._session_contexts: dict[str, McpSessionContext] = {} self._acp_mcp_manager: AcpMcpConnectionManager | None = None def add_server_config(self, cfg: MCPServerConfig | str) -> None: @@ -185,10 +185,10 @@ def add_server_config(self, cfg: MCPServerConfig | str) -> None: resolved = BaseMCPServerConfig.from_string(cfg) if isinstance(cfg, str) else cfg self.servers.append(resolved) - def get_or_create_session(self, session_id: str) -> _SessionContext: + def get_or_create_session(self, session_id: str) -> McpSessionContext: """Get or create the per-session MCP context for ``session_id``. - If no context exists for ``session_id``, a new ``_SessionContext`` + If no context exists for ``session_id``, a new ``McpSessionContext`` is created with a fresh ``SessionConnectionPool``, empty toolset cache, no snapshot, and an empty ACP connection list. Subsequent calls with the same ``session_id`` return the same object. @@ -197,19 +197,19 @@ def get_or_create_session(self, session_id: str) -> _SessionContext: session_id: Unique identifier for the session. Returns: - The ``_SessionContext`` for this session. + The ``McpSessionContext`` for this session. """ from agentpool.mcp_server.session_pool import SessionConnectionPool ctx = self._session_contexts.get(session_id) if ctx is None: - ctx = _SessionContext( + ctx = McpSessionContext( connection_pool=SessionConnectionPool(session_id=session_id), ) self._session_contexts[session_id] = ctx return ctx - def get_session_context(self, session_id: str) -> _SessionContext | None: + def get_session_context(self, session_id: str) -> McpSessionContext | None: """Get the session context for ``session_id`` without creating one. Returns ``None`` if no context exists for the session. @@ -234,6 +234,29 @@ def update_session_snapshot( ctx = self.get_or_create_session(session_id) ctx.snapshot = snapshot + async def add_transport( + self, + session_id: str, + client_id: str, + transport: ClientTransport, + skill_name: str | None = None, + ) -> None: + """Add a pre-created transport to the session's connection pool. + + Delegates to the internal :class:`SessionConnectionPool` for the + given session. If no session context exists yet, one is created + via ``get_or_create_session()``. + + Args: + session_id: Unique identifier for the session. + client_id: Client identifier for the MCP server. + transport: Pre-created fastmcp ``ClientTransport``. + skill_name: Optional skill name for skill-scoped MCP isolation. + """ + ctx = self.get_or_create_session(session_id) + if ctx.connection_pool is not None: + await ctx.connection_pool.add_transport(client_id, transport, skill_name) + def __repr__(self) -> str: return f"MCPManager(name={self.name!r}, servers={len(self.servers)})" @@ -397,7 +420,7 @@ async def get_capabilities( # noqa: PLR0915 skipped in global configs since pydantic-ai does not support ACP directly. Disabled servers are also skipped. - When ``session_id`` is provided, the session's ``_SessionContext`` is + When ``session_id`` is provided, the session's ``McpSessionContext`` is looked up via ``get_or_create_session()``. If a snapshot is stored on the context, configs are partitioned: diff --git a/src/agentpool/orchestrator/__init__.py b/src/agentpool/orchestrator/__init__.py index 8ff7a45b3..c462d30ec 100644 --- a/src/agentpool/orchestrator/__init__.py +++ b/src/agentpool/orchestrator/__init__.py @@ -4,7 +4,7 @@ - event_bus: EventBus, EventEnvelope, drain_and_merge - session_controller: SessionController, SessionState, exceptions - session_pool: SessionPool -- run: RunHandle, RunStatus +- run: RunHandle - metrics: MetricsCollector, SessionPoolMetrics - runtime_registry: RuntimeAgentRegistry """ @@ -18,7 +18,7 @@ drain_and_merge, ) from agentpool.orchestrator.metrics import MetricsCollector, SessionPoolMetrics -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.runtime_registry import RuntimeAgentRegistry from agentpool.orchestrator.session_controller import ( DEFAULT_SESSION_TTL_SECONDS, @@ -43,7 +43,6 @@ "EventEnvelope", "MetricsCollector", "RunHandle", - "RunStatus", "RuntimeAgentRegistry", "SessionBusyError", "SessionController", diff --git a/src/agentpool/orchestrator/core.py b/src/agentpool/orchestrator/core.py index df4480a8e..b3448dea2 100644 --- a/src/agentpool/orchestrator/core.py +++ b/src/agentpool/orchestrator/core.py @@ -25,7 +25,7 @@ _rebind, # noqa: F401 drain_and_merge, ) -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.session_controller import ( DEFAULT_SESSION_TTL_SECONDS, CheckpointMismatchError, @@ -50,7 +50,6 @@ "EventBus", "EventEnvelope", "RunHandle", - "RunStatus", "SessionBusyError", "SessionClosedError", "SessionController", diff --git a/src/agentpool/orchestrator/run.py b/src/agentpool/orchestrator/run.py index 4af820eed..6f99a02ec 100644 --- a/src/agentpool/orchestrator/run.py +++ b/src/agentpool/orchestrator/run.py @@ -5,7 +5,6 @@ import asyncio import contextlib from dataclasses import dataclass, field -from enum import Enum, auto from typing import TYPE_CHECKING, Any, Self import uuid @@ -26,6 +25,7 @@ Journal, MemoryJournal, MemorySnapshotStore, + RunOutcome, RunState, SnapshotStore, TriggerSource, @@ -127,28 +127,6 @@ def inject_cancelled_tool_results(messages: list[ModelMessage]) -> list[ModelMes return result -class RunStatus(Enum): - """Lifecycle states for an agent run. - - Values: - pending: RunHandle created but not yet started. - running: Actively executing. - completed: Finished normally. - failed: Finished with an error. - checkpointed: Run state persisted for later resumption. - idle: RunHandle created but no active turn. - done: RunHandle closed or cancelled. - """ - - pending = auto() - running = auto() - completed = auto() - failed = auto() - checkpointed = auto() - idle = auto() - done = auto() - - @dataclass class RunHandle: """Ephemeral runtime handle for a single agent run. @@ -168,7 +146,9 @@ class RunHandle: run_id: Unique identifier for this run. session_id: Session this run belongs to. agent_type: Type of agent running (e.g. ``"native"``, ``"claude"``). - status: Legacy lifecycle state (used by old code paths). + outcome: Terminal outcome (``RunOutcome.COMPLETED``, ``FAILED``, + ``CHECKPOINTED``) set when the run reaches ``RunState.DONE``. + ``None`` while the run is active or was closed without outcome. agent: The agent instance driving turns. event_bus: Event bus for publishing stream events. session: Per-session state containing the turn lock. @@ -177,7 +157,6 @@ class RunHandle: _cleanup_callback: Optional callback invoked with run_id during cleanup. active_agent_run: Reference to PydanticAI AgentRun, set by NativeTurn during execution and cleared in ``finally``. - _status: New primary lifecycle state (idle/running/done). _closing: Flag indicating :meth:`close` has been called. _idle_event: asyncio.Event that is set when idle (for wake-up). _message_queue: Queued prompts for the next turn. @@ -190,7 +169,7 @@ class RunHandle: run_id: str session_id: str agent_type: str - status: RunStatus = RunStatus.pending + outcome: RunOutcome | None = None agent: BaseAgent[Any, Any] | None = None event_bus: EventBus | None = None session: SessionState | None = None @@ -199,7 +178,6 @@ class RunHandle: _cleanup_callback: Callable[[str], None] | None = None active_agent_run: AgentRun[Any, Any] | None = None _cancel_fn: Callable[[], None] | None = None - _status: RunStatus = RunStatus.idle _closing: bool = False _closed: bool = False _idle_event: asyncio.Event = field(default_factory=_create_set_event) @@ -208,6 +186,15 @@ class RunHandle: _turn_complete_event: asyncio.Event = field(default_factory=asyncio.Event) _turn_was_cancelled: bool = False _interrupt_task: asyncio.Task[None] | None = None + _current_turn: Any = None + """The current Turn being executed. Set by ``_execute_turn()``, read by + ``_handle_turn_result()`` and ``_drain_events()``.""" + _current_turn_id: str | None = None + """The current turn ID. Set by ``_execute_turn()``, read by + ``_drain_events()``.""" + _current_turn_failed: bool = False + """Whether the current turn failed. Set by ``_execute_turn()``, read by + ``_handle_turn_result()``.""" # ------------------------------------------------------------------ # Lifecycle dimensions (M2) @@ -251,8 +238,9 @@ def __post_init__(self) -> None: """Initialize default lifecycle dimensions. Any dimension left as ``None`` is populated with the default - in-process implementation. The journal is injected into the - CommChannel to ensure the channel can persist events. + in-process implementation. Both ``DirectChannel`` and + ``ProtocolChannel`` receive the journal via their constructor, + so no post-hoc journal injection is needed. """ if self._journal is None: self._journal = MemoryJournal() @@ -260,14 +248,6 @@ def __post_init__(self) -> None: self._snapshot_store = MemorySnapshotStore() if self._comm_channel is None: self._comm_channel = DirectChannel(self._journal) - else: - # Inject journal into existing CommChannel if not already set. - try: - existing_journal: Journal | None = self._comm_channel._journal # type: ignore[attr-defined] - except AttributeError: - existing_journal = None - if existing_journal is None: - self._comm_channel._journal = self._journal # type: ignore[attr-defined] if self._event_transport is None: self._event_transport = InProcessTransport() if self._trigger_source is None: @@ -282,26 +262,6 @@ def is_running(self) -> bool: """ return self._run_state == RunState.RUNNING - @property - def _channel_publishes_to_event_bus(self) -> bool: - """Whether the CommChannel publishes events to the EventBus itself. - - ``ProtocolChannel`` calls ``event_bus.publish(session_id, event)`` - inside its own ``publish()`` method. When this is the case, - ``start()`` must NOT also call ``event_bus.publish()`` directly - to avoid double-publishing. - - ``DirectChannel`` does not publish to the EventBus, so the - direct call in ``start()`` is required. - - Returns: - ``True`` if the CommChannel publishes to the EventBus - internally, ``False`` otherwise. - """ - from agentpool.lifecycle.comm_channel import ProtocolChannel - - return isinstance(self._comm_channel, ProtocolChannel) - @property def recovered_tool_executions(self) -> list[Any]: """Tool executions from the interrupted Turn, for idempotent retry. @@ -398,7 +358,7 @@ def _inject_agent_context(self) -> None: # New session-level lifecycle # ------------------------------------------------------------------ - async def start(self, initial_prompt: str) -> AsyncGenerator[RichAgentStreamEvent[Any]]: # noqa: PLR0915 + async def start(self, initial_prompt: str) -> AsyncGenerator[RichAgentStreamEvent[Any]]: """Start the idle/wake/turn loop as an async generator. Yields :class:`RichAgentStreamEvent` tokens from each turn's @@ -430,23 +390,83 @@ async def start(self, initial_prompt: str) -> AsyncGenerator[RichAgentStreamEven # CancelNotification, native _iteration_task cancel). self._cancel_fn = self._create_cancel_fn() - # --- Lifecycle dimensions (M2) --- - # Crash recovery: check journal for prior state. - assert self._journal is not None # set by __post_init__ - assert self._snapshot_store is not None # set by __post_init__ - assert self._comm_channel is not None # set by __post_init__ - assert self._trigger_source is not None # set by __post_init__ + recovered_prompt = await self._handle_recovery() + + try: + async with session.turn_lock: + # For "retry" recovery, prepend the recovered prompt. + if recovered_prompt is not None: + current_prompts: list[str] = [recovered_prompt] + else: + current_prompts = [initial_prompt] + while not self._closing: + if not current_prompts: + current_prompts = await self._idle_loop() + if not current_prompts: + continue + + async for event in self._execute_turn( + agent, + event_bus, + session, + current_prompts, + ): + yield event + + action = await self._handle_turn_result(event_bus) + if action == "continue": + current_prompts = [] # Prevent re-execution of cancelled prompt + continue + if action == "break": + break + + current_prompts = await self._drain_events() + + finally: + self._closed = True + # Lifecycle state transition: → DONE. + with contextlib.suppress(Exception): + await self._transition(RunState.DONE) + self._turn_complete_event.set() + self.complete_event.set() + # Close lifecycle dimensions. + with contextlib.suppress(Exception): + if self._trigger_source is not None: + self._trigger_source.close() + with contextlib.suppress(Exception): + if self._comm_channel is not None: + self._comm_channel.close() + with contextlib.suppress(Exception): + if self._event_transport is not None: + self._event_transport.close() + + async def _handle_recovery(self) -> str | None: + """Perform crash recovery and subscribe lifecycle dimensions. + + Checks the journal for prior state. If an in-flight Turn is + detected, replays journaled events and applies the recovery + strategy (``"retry"`` or ``"mark_interrupted"``). Saves the + initial snapshot for fresh starts. Subscribes the trigger + source and CommChannel to this handle. + + Returns: + The recovered prompt for the ``"retry"`` strategy, or ``None``. + """ + assert self._journal is not None + assert self._snapshot_store is not None + assert self._comm_channel is not None + assert self._trigger_source is not None recovered_prompt: str | None = None resume_result = self._journal.resume(self._snapshot_store) if resume_result is not None: if resume_result.is_inflight: # In-flight crash recovery: replay journaled events. - self._comm_channel._replaying = True + self._comm_channel.set_replaying(True) try: for event in resume_result.events: await self._comm_channel.publish(event) finally: - self._comm_channel._replaying = False + self._comm_channel.set_replaying(False) self._recovered_inflight_turn_id = resume_result.inflight_turn_id # Apply recovery strategy. if self._recover_strategy == "retry": @@ -480,321 +500,345 @@ async def start(self, initial_prompt: str) -> AsyncGenerator[RichAgentStreamEven # Subscribe dimensions. self._trigger_source.subscribe(self) self._comm_channel.attach(self) + return recovered_prompt - try: - async with session.turn_lock: - # For "retry" recovery, prepend the recovered prompt. - if recovered_prompt is not None: - current_prompts: list[str] = [recovered_prompt] - else: - current_prompts = [initial_prompt] - while not self._closing: - if not current_prompts: - self._status = RunStatus.idle - self._idle_event.clear() - # Drain CommChannel feedback queue (ProtocolChannel) - # BEFORE deciding to block. Feedback may have been - # enqueued by steer/followup via deliver_feedback() - # while the loop was running (e.g., during cancel). - # Without this, the loop would block on - # _idle_event.wait() even though feedback is already - # available in the CommChannel. - if self._comm_channel is not None: - while True: - fb = self._comm_channel.recv() - if fb is None: - break - self._message_queue.append(fb.content) - # Check if messages were queued during cancel/cleanup - # before blocking. Without this, messages routed through - # _message_queue by the cancel path would deadlock: cancel() - # sets _idle_event, but clear() above removes it, and wait() - # blocks forever with no one to re-set it. - if not self._message_queue: - await self._idle_event.wait() - if self._closing: - # Drain any feedback from CommChannel - # before checking for pending messages. - if self._comm_channel is not None: - while True: - fb = self._comm_channel.recv() - if fb is None: - break - self._message_queue.append(fb.content) - # Process pending messages before exiting - # so close() with queued followups are - # handled as final Turns. - if not self._message_queue: - break - # Drain CommChannel feedback queue again after - # waking from idle (feedback may have arrived - # during the wait). - if self._comm_channel is not None: - while True: - fb = self._comm_channel.recv() - if fb is None: - break - self._message_queue.append(fb.content) - current_prompts = list(self._message_queue) - self._message_queue.clear() - if not current_prompts: - continue + async def _idle_loop(self) -> list[str]: + """Wait for idle, drain feedback, and collect prompts for next turn. - self._status = RunStatus.running - # Lifecycle state transition: IDLE → RUNNING. - await self._transition(RunState.RUNNING) - # Generate a unique turn_id for this Turn. - turn_id = str(uuid.uuid4()) - self.run_ctx.turn_id = turn_id - # Reset per-turn state: clear the completion event - # and clear any stale cancelled flag from a prior turn. - self._turn_complete_event.clear() - self._turn_was_cancelled = False - if self.run_ctx.cancelled: - self.run_ctx.cancelled = False - # Clear hooks_fired so the new turn's hooks can fire - # even if the previous turn already fired them. - self.run_ctx.hooks_fired.clear() - # Construct per-turn AgentContext and inject as deps - # so capabilities (SubagentCapability, etc.) can access - # the delegation service, resource sources, and host. - self._inject_agent_context() - turn = agent.create_turn( - prompts=current_prompts, # type: ignore[arg-type] - run_ctx=self.run_ctx, - message_history=self._message_history, - ) - # Publish RunStartedEvent before turn.execute() so - # consumers know a new turn is starting. This was - # previously yielded by NativeTurn.execute() itself, - # causing duplicate events when RunHandle.start() - # also published turn events. We publish to the event - # bus without yielding to avoid inflating the event - # count seen by generator consumers. - run_started = RunStartedEvent( - run_id=self.run_id, - session_id=self.session_id, - agent_name=self.agent_type, - parent_session_id=session.parent_session_id - if session is not None - else None, - ) - if not self._channel_publishes_to_event_bus: - await event_bus.publish(self.session_id, run_started) - await self._comm_channel.publish(run_started) - - # Set _current_input_provider ContextVar so MCP - # elicitation can access it during turn execution. - # Only set() without reset(): start() runs inside an - # asyncio.Task which copies the parent Context, so - # set() only affects this task's private context copy. - # When the task ends the context is discarded. Calling - # reset() is unnecessary and can raise ValueError when - # the async generator is GC-collected in a different - # Context (race between task cancellation and generator - # suspension at a yield point). - if session.input_provider is not None: - from agentpool.mcp_server.manager import _current_input_provider - - _current_input_provider.set(session.input_provider) - - # Save user prompt to agent conversation before execution. - # This ensures user messages are preserved even if the turn - # fails or is cancelled (mirroring _run_stream_once() behavior). - agent.conversation.add_chat_messages([ - ChatMessage( - content="\n".join(current_prompts), - role="user", - name=agent.name, - session_id=self.session_id, - ), - ]) + Clears the idle event, drains CommChannel feedback and message + queue, then blocks on the idle event if no prompts are available. + After waking, drains feedback again and returns the collected + prompts. - # Pre-turn snapshot: save prompt and turn_id for - # crash recovery. If the process crashes during - # turn.execute(), this snapshot allows the "retry" - # strategy to recover the prompt and turn_id. - self._snapshot_store.save({ - "state": RunState.RUNNING.value, - "run_id": self.run_id, - "turn_id": turn_id, - "prompt": "\n".join(current_prompts), - }) - - turn_failed = False - try: - async for event in turn.execute(): - if not self._channel_publishes_to_event_bus: - await event_bus.publish(self.session_id, event) - await self._comm_channel.publish(event) - # Save assistant final message to conversation BEFORE - # yielding. The _consume_run caller closes the generator - # immediately after receiving StreamCompleteEvent, which - # prevents any code after `yield event` from executing. - if isinstance(event, StreamCompleteEvent) and event.message is not None: - agent.conversation.add_chat_messages( - [event.message], - extend_last=True, - ) - yield event - if isinstance(event, RunErrorEvent): - turn_failed = True - break - if isinstance(event, StreamCompleteEvent): - break - except Exception as e: # noqa: BLE001 - turn_failed = True - error_event = RunErrorEvent( - message=str(e), - run_id=self.run_id, - agent_name=self.agent_type, - ) - if not self._channel_publishes_to_event_bus: - await event_bus.publish(self.session_id, error_event) - await self._comm_channel.publish(error_event) - yield error_event - - if self.run_ctx.cancelled: - # Turn was cancelled — publish RunFailedEvent, set turn - # complete, clear prompts, and continue to idle for next turn. - # RunFailedEvent must be published BEFORE _turn_complete_event - # so the event converter can emit - # TurnCompleteUpdate(stop_reason="cancelled"). - cancelled_event = RunFailedEvent( - run_id=self.run_id, - session_id=self.session_id, - exception=RuntimeError("Run cancelled"), - ) - if not self._channel_publishes_to_event_bus: - await event_bus.publish(self.session_id, cancelled_event) - await self._comm_channel.publish(cancelled_event) - # Capture cancelled state BEFORE setting _turn_complete_event. - # handle_prompt() checks run_handle.cancelled after waking from - # _turn_complete_event.wait(). But the loop may reset cancelled=False - # before handle_prompt() gets scheduled (e.g., when steer messages - # are queued). _turn_was_cancelled preserves the state for observation. - self._turn_was_cancelled = True - self._turn_complete_event.set() - # Route queued steer messages through _message_queue - # instead of directly into current_prompts. This forces - # the loop through idle, preserving cancelled=True for - # handle_prompt() to observe before the next turn resets it. - if self.run_ctx.queued_steer_messages: - self._message_queue.extend(self.run_ctx.queued_steer_messages) - self.run_ctx.queued_steer_messages.clear() - current_prompts = [] # Prevent re-execution of cancelled prompt - # Preserve the cancelled turn's message history so the - # next turn sees the partial conversation context. - # Without this, `continue` skips line 300 and the - # next turn starts with stale _message_history. - if not turn_failed: - with contextlib.suppress(RuntimeError): - self._message_history = turn.message_history - # Do NOT reset cancelled here — handle_prompt() needs to - # observe it. It will be reset at the start of the next turn. - continue + Returns: + List of prompts for the next turn. Empty list if closing + with no pending messages. + """ + self._idle_event.clear() + # Drain CommChannel feedback queue (ProtocolChannel) BEFORE + # deciding to block. Feedback may have been enqueued by + # steer/followup via deliver_feedback() while the loop was + # running (e.g., during cancel). Without this, the loop would + # block on _idle_event.wait() even though feedback is already + # available in the CommChannel. + if self._comm_channel is not None: + while True: + fb = self._comm_channel.recv() + if fb is None: + break + self._message_queue.append(fb.content) + # Check if messages were queued during cancel/cleanup before + # blocking. Without this, messages routed through _message_queue + # by the cancel path would deadlock: cancel() sets _idle_event, + # but clear() above removes it, and wait() blocks forever with + # no one to re-set it. + if not self._message_queue: + await self._idle_event.wait() + if self._closing: + # Drain any feedback from CommChannel before checking + # for pending messages. + if self._comm_channel is not None: + while True: + fb = self._comm_channel.recv() + if fb is None: + break + self._message_queue.append(fb.content) + # Process pending messages before exiting so close() + # with queued followups are handled as final Turns. + if not self._message_queue: + return [] + # Drain CommChannel feedback queue again after waking from + # idle (feedback may have arrived during the wait). + if self._comm_channel is not None: + while True: + fb = self._comm_channel.recv() + if fb is None: + break + self._message_queue.append(fb.content) + prompts = list(self._message_queue) + self._message_queue.clear() + return prompts + + async def _execute_turn( + self, + agent: BaseAgent[Any, Any], + event_bus: EventBus, + session: SessionState, + current_prompts: list[str], + ) -> AsyncGenerator[RichAgentStreamEvent[Any]]: + """Execute a single turn and yield stream events. + + Creates a Turn from the current prompts, publishes + ``RunStartedEvent``, saves the user prompt to conversation + history, takes a pre-turn snapshot, then executes the Turn + and yields each event. On exception, publishes and yields a + ``RunErrorEvent``. + + Stores the Turn, turn_id, and turn_failed flag on ``self`` for + downstream sub-methods (``_handle_turn_result``, + ``_drain_events``). - if turn_failed: - break + Args: + agent: The agent driving the turn. + event_bus: The event bus for publishing events. + session: The per-session state. + current_prompts: Prompts for this turn. + """ + assert self._comm_channel is not None + assert self._snapshot_store is not None + # Lifecycle state transition: IDLE -> RUNNING. + await self._transition(RunState.RUNNING) + # Generate a unique turn_id for this Turn. + turn_id = str(uuid.uuid4()) + self.run_ctx.turn_id = turn_id + # Reset per-turn state: clear the completion event and clear + # any stale cancelled flag from a prior turn. + self._turn_complete_event.clear() + self._turn_was_cancelled = False + if self.run_ctx.cancelled: + self.run_ctx.cancelled = False + # Construct per-turn AgentContext and inject as deps so + # capabilities (SubagentCapability, etc.) can access the + # delegation service, resource sources, and host. + self._inject_agent_context() + turn = agent.create_turn( + prompts=current_prompts, # type: ignore[arg-type] + run_ctx=self.run_ctx, + message_history=self._message_history, + ) + # Publish RunStartedEvent before turn.execute() so consumers + # know a new turn is starting. + run_started = RunStartedEvent( + run_id=self.run_id, + session_id=self.session_id, + agent_name=self.agent_type, + parent_session_id=session.parent_session_id if session is not None else None, + ) + if not self._comm_channel.publishes_to_event_bus: + await event_bus.publish(self.session_id, run_started) + await self._comm_channel.publish(run_started) + # Set _current_input_provider ContextVar so MCP elicitation can + # access it during turn execution. Only set() without reset(): + # start() runs inside an asyncio.Task which copies the parent + # Context, so set() only affects this task's private context + # copy. When the task ends the context is discarded. + if session.input_provider is not None: + from agentpool.mcp_server.manager import _current_input_provider + + _current_input_provider.set(session.input_provider) + # Save user prompt to agent conversation before execution. + # This ensures user messages are preserved even if the turn + # fails or is cancelled. + agent.conversation.add_chat_messages([ + ChatMessage( + content="\n".join(current_prompts), + role="user", + name=agent.name, + session_id=self.session_id, + ), + ]) + # Pre-turn snapshot: save prompt and turn_id for crash + # recovery. If the process crashes during turn.execute(), + # this snapshot allows the "retry" strategy to recover. + self._snapshot_store.save({ + "state": RunState.RUNNING.value, + "run_id": self.run_id, + "turn_id": turn_id, + "prompt": "\n".join(current_prompts), + }) + # Store turn state for downstream sub-methods. + self._current_turn = turn + self._current_turn_id = turn_id + self._current_turn_failed = False + turn_failed = False + try: + async for event in turn.execute(): + if not self._comm_channel.publishes_to_event_bus: + await event_bus.publish(self.session_id, event) + await self._comm_channel.publish(event) + # Save assistant final message to conversation BEFORE + # yielding. The _consume_run caller closes the generator + # immediately after receiving StreamCompleteEvent, which + # prevents any code after `yield event` from executing. + if isinstance(event, StreamCompleteEvent) and event.message is not None: + agent.conversation.add_chat_messages( + [event.message], + extend_last=True, + ) + yield event + if isinstance(event, RunErrorEvent): + turn_failed = True + break + if isinstance(event, StreamCompleteEvent): + break + except Exception as e: # noqa: BLE001 + turn_failed = True + error_event = RunErrorEvent( + message=str(e), + run_id=self.run_id, + agent_name=self.agent_type, + ) + if not self._comm_channel.publishes_to_event_bus: + await event_bus.publish(self.session_id, error_event) + await self._comm_channel.publish(error_event) + yield error_event + finally: + self._current_turn_failed = turn_failed - if not turn_failed: - with contextlib.suppress(RuntimeError): - self._message_history = turn.message_history + async def _handle_turn_result(self, event_bus: EventBus) -> str: + """Handle cancel and error outcomes after turn execution. - # Lifecycle state transition: RUNNING → IDLE. - await self._transition(RunState.IDLE) + If the turn was cancelled, publishes ``RunFailedEvent``, routes + queued steer messages, transitions to IDLE, and returns + ``"continue"``. If the turn failed, returns ``"break"``. + Otherwise saves message history and returns ``"proceed"``. - # Snapshot at turn boundary (after state transition). - self._snapshot_store.save( - { - "state": self._run_state.value, - "run_id": self.run_id, - "turn_id": turn_id, - }, - ) - # Save turn result for idempotency. - with contextlib.suppress(RuntimeError): - final_msg = turn.final_message - self._snapshot_store.save_turn_result(turn_id, final_msg) - - # Between turns: wait for background child tasks to complete, - # then collect their steer messages as prompts for next turn. - child_events_timed_out = False - if self.run_ctx.child_done_events: - try: - async with asyncio.timeout(30): - await asyncio.gather(*[ - e.wait() for e in list(self.run_ctx.child_done_events.values()) - ]) - except TimeoutError: - child_events_timed_out = True - logger.warning( - "Timeout waiting for child_done_events", - run_id=self.run_id, - pending=len(self.run_ctx.child_done_events), - ) - - # Collect queued steer messages from completed children - # as prompts for the next turn. - if self.run_ctx.queued_steer_messages: - self._message_queue.extend(self.run_ctx.queued_steer_messages) - self.run_ctx.queued_steer_messages.clear() - - if child_events_timed_out: - # On timeout, clear ALL child_done_events since we are - # proceeding regardless. New child tasks may have been - # registered during the wait, but we cannot wait further. - self.run_ctx.child_done_events.clear() - else: - # Only remove completed events; new child tasks may have - # been registered between gather() and here. - completed_keys = [ - k for k, e in list(self.run_ctx.child_done_events.items()) if e.is_set() - ] - for k in completed_keys: - del self.run_ctx.child_done_events[k] - - # Drain CommChannel feedback queue (ProtocolChannel). - # Feedback may have been enqueued by steer/followup - # via deliver_feedback() during the Turn. Steer - # feedback is prioritized as next-turn prompts. - feedback_steer: list[str] = [] - if self._comm_channel is not None: - while True: - fb = self._comm_channel.recv() - if fb is None: - break - if fb.is_steer: - feedback_steer.append(fb.content) - else: - self._message_queue.append(fb.content) - - current_prompts = feedback_steer + list(self._message_queue) - self._message_queue.clear() - - # Signal that this turn has completed normally. - self._turn_was_cancelled = False - self._turn_complete_event.set() - - self._status = RunStatus.done - finally: - self._status = RunStatus.done - self._closed = True - # Lifecycle state transition: → DONE. - with contextlib.suppress(Exception): - await self._transition(RunState.DONE) + Args: + event_bus: The event bus for publishing events. + + Returns: + ``"continue"`` (cancel path), ``"break"`` (failure), or + ``"proceed"`` (normal completion). + """ + turn = self._current_turn + turn_failed = self._current_turn_failed + if turn is None: + return "break" + + assert self._comm_channel is not None + + if self.run_ctx.cancelled: + # Turn was cancelled -- publish RunFailedEvent, set turn + # complete, clear prompts, and continue to idle for next + # turn. RunFailedEvent must be published BEFORE + # _turn_complete_event so the event converter can emit + # TurnCompleteUpdate(stop_reason="cancelled"). + cancelled_event = RunFailedEvent( + run_id=self.run_id, + session_id=self.session_id, + exception=RuntimeError("Run cancelled"), + ) + if not self._comm_channel.publishes_to_event_bus: + await event_bus.publish(self.session_id, cancelled_event) + await self._comm_channel.publish(cancelled_event) + # Capture cancelled state BEFORE setting _turn_complete_event. + # handle_prompt() checks run_handle.cancelled after waking + # from _turn_complete_event.wait(). But the loop may reset + # cancelled=False before handle_prompt() gets scheduled. + # _turn_was_cancelled preserves the state for observation. + self._turn_was_cancelled = True self._turn_complete_event.set() - self.complete_event.set() - # Close lifecycle dimensions. - with contextlib.suppress(Exception): - if self._trigger_source is not None: - self._trigger_source.close() - with contextlib.suppress(Exception): - if self._comm_channel is not None: - self._comm_channel.close() - with contextlib.suppress(Exception): - if self._event_transport is not None: - self._event_transport.close() + # Route queued steer messages through _message_queue instead + # of directly into current_prompts. This forces the loop + # through idle, preserving cancelled=True for handle_prompt() + # to observe before the next turn resets it. + if self.run_ctx.queued_steer_messages: + self._message_queue.extend(self.run_ctx.queued_steer_messages) + self.run_ctx.queued_steer_messages.clear() + # Prevent re-execution of cancelled prompt. + # Preserve the cancelled turn's message history so the next + # turn sees the partial conversation context. + if not turn_failed: + with contextlib.suppress(RuntimeError): + self._message_history = turn.message_history + # Do NOT reset cancelled here -- handle_prompt() needs to + # observe it. It will be reset at the start of the next turn. + # Lifecycle state transition: RUNNING -> IDLE (cancel path). + await self._transition(RunState.IDLE) + return "continue" + + if turn_failed: + return "break" + + with contextlib.suppress(RuntimeError): + self._message_history = turn.message_history + return "proceed" + + async def _drain_events(self) -> list[str]: + """Post-turn snapshot, child event collection, and feedback drain. + + Transitions to IDLE, saves a turn-boundary snapshot, saves the + turn result for idempotency, waits for background child tasks, + collects queued steer messages, drains CommChannel feedback, + and signals turn completion. + + Returns: + Prompts for the next turn (steer feedback + queued messages). + """ + turn = self._current_turn + turn_id = self._current_turn_id + assert turn is not None + assert turn_id is not None + assert self._snapshot_store is not None + + # Lifecycle state transition: RUNNING -> IDLE. + await self._transition(RunState.IDLE) + # Snapshot at turn boundary (after state transition). + self._snapshot_store.save( + { + "state": self._run_state.value, + "run_id": self.run_id, + "turn_id": turn_id, + }, + ) + # Save turn result for idempotency. + with contextlib.suppress(RuntimeError): + final_msg = turn.final_message + self._snapshot_store.save_turn_result(turn_id, final_msg) + # Between turns: wait for background child tasks to complete, + # then collect their steer messages as prompts for next turn. + child_events_timed_out = False + if self.run_ctx.child_done_events: + try: + async with asyncio.timeout(30): + await asyncio.gather(*[ + e.wait() for e in list(self.run_ctx.child_done_events.values()) + ]) + except TimeoutError: + child_events_timed_out = True + logger.warning( + "Timeout waiting for child_done_events", + run_id=self.run_id, + pending=len(self.run_ctx.child_done_events), + ) + # Collect queued steer messages from completed children as + # prompts for the next turn. + if self.run_ctx.queued_steer_messages: + self._message_queue.extend(self.run_ctx.queued_steer_messages) + self.run_ctx.queued_steer_messages.clear() + if child_events_timed_out: + # On timeout, clear ALL child_done_events since we are + # proceeding regardless. New child tasks may have been + # registered during the wait, but we cannot wait further. + self.run_ctx.child_done_events.clear() + else: + # Only remove completed events; new child tasks may have + # been registered between gather() and here. + completed_keys = [ + k for k, e in list(self.run_ctx.child_done_events.items()) if e.is_set() + ] + for k in completed_keys: + del self.run_ctx.child_done_events[k] + # Drain CommChannel feedback queue (ProtocolChannel). Feedback + # may have been enqueued by steer/followup via deliver_feedback() + # during the Turn. Steer feedback is prioritized as next-turn + # prompts. + feedback_steer: list[str] = [] + if self._comm_channel is not None: + while True: + fb = self._comm_channel.recv() + if fb is None: + break + if fb.is_steer: + feedback_steer.append(fb.content) + else: + self._message_queue.append(fb.content) + prompts = feedback_steer + list(self._message_queue) + self._message_queue.clear() + # Signal that this turn has completed normally. + self._turn_was_cancelled = False + self._turn_complete_event.set() + return prompts def steer(self, message: str) -> bool: """Inject a steer message into the active turn or wake idle handle. @@ -819,15 +863,10 @@ def steer(self, message: str) -> bool: if self._closing: return False - # Try ProtocolChannel feedback path (deliver_feedback exists). + # Try CommChannel feedback path (ProtocolChannel returns True). if self._comm_channel is not None: - try: - deliver: Any | None = self._comm_channel.deliver_feedback # type: ignore[attr-defined] - except AttributeError: - deliver = None - if deliver is not None: - feedback = Feedback(content=message, is_steer=True) - self._comm_channel.deliver_feedback(feedback) # type: ignore[attr-defined] + feedback = Feedback(content=message, is_steer=True) + if self._comm_channel.deliver_feedback(feedback): # Always set _idle_event when delivering via ProtocolChannel. # If the loop is running, the event is cleared when entering # idle, and the loop then drains CommChannel feedback. If the @@ -838,12 +877,12 @@ def steer(self, message: str) -> bool: return True # Fallback: DirectChannel path (existing logic). - if self._status == RunStatus.idle: + if self._run_state == RunState.IDLE: self._message_queue.append(message) self._idle_event.set() return True - if self._status == RunStatus.running: + if self._run_state == RunState.RUNNING: agent_run = self.active_agent_run if agent_run is not None: agent_run.enqueue(message, priority="asap") @@ -866,22 +905,17 @@ def followup(self, message: str) -> bool: if self._closing: return False - # Try ProtocolChannel feedback path (deliver_feedback exists). + # Try CommChannel feedback path (ProtocolChannel returns True). if self._comm_channel is not None: - try: - deliver: Any | None = self._comm_channel.deliver_feedback # type: ignore[attr-defined] - except AttributeError: - deliver = None - if deliver is not None: - feedback = Feedback(content=message, is_steer=False) - self._comm_channel.deliver_feedback(feedback) # type: ignore[attr-defined] + feedback = Feedback(content=message, is_steer=False) + if self._comm_channel.deliver_feedback(feedback): # Always set _idle_event (see steer() for rationale). self._idle_event.set() return True # Fallback: DirectChannel path (existing logic). self._message_queue.append(message) - if self._status == RunStatus.idle: + if self._run_state == RunState.IDLE: self._idle_event.set() return True @@ -956,17 +990,19 @@ def _start_task(self, task: asyncio.Task[Any] | None = None) -> None: Args: task: The asyncio.Task driving this run, if any. """ - self.status = RunStatus.running + self._run_state = RunState.RUNNING self.run_ctx.current_task = task def complete(self) -> None: """Transition the run to completed and trigger cleanup.""" - self.status = RunStatus.completed + self._run_state = RunState.DONE + self.outcome = RunOutcome.COMPLETED self._cleanup_run() def checkpoint(self) -> None: """Transition the run to checkpointed and trigger cleanup.""" - self.status = RunStatus.checkpointed + self._run_state = RunState.DONE + self.outcome = RunOutcome.CHECKPOINTED self._cleanup_run() def fail( @@ -981,7 +1017,8 @@ def fail( exception: Optional exception that caused the failure. event_bus: Optional event bus to publish RunFailedEvent on. """ - self.status = RunStatus.failed + self._run_state = RunState.DONE + self.outcome = RunOutcome.FAILED if exception is not None: self.run_ctx.cancelled = True if event_bus is not None: diff --git a/src/agentpool/orchestrator/session_controller.py b/src/agentpool/orchestrator/session_controller.py index f055a509f..9bbca1d2e 100644 --- a/src/agentpool/orchestrator/session_controller.py +++ b/src/agentpool/orchestrator/session_controller.py @@ -22,8 +22,9 @@ RunFailedEvent, StreamCompleteEvent, ) +from agentpool.lifecycle import RunState from agentpool.log import get_logger -from agentpool.orchestrator.run import RunHandle, RunStatus, inject_cancelled_tool_results +from agentpool.orchestrator.run import RunHandle, inject_cancelled_tool_results from agentpool.orchestrator.runtime_registry import RuntimeAgentRegistry from agentpool.sessions.models import PendingDeferredCall, SessionData from agentpool_server.opencode_server.models.session_info import SessionInfo @@ -481,14 +482,6 @@ async def get_or_create_session_agent( logger.info("Created session agent", session_id=session_id, agent_name=agent_name) return agent - available_manifest = list(self.pool.manifest.agents.keys()) - available_runtime = self._runtime_registry.names() - msg = ( - f"Agent config not found: {agent_name!r}. " - f"Available in manifest: {available_manifest}. " - f"Available in runtime registry: {available_runtime}." - ) - raise RuntimeError(msg) def list_sessions(self) -> list[SessionInfo]: """List all active sessions. @@ -1055,11 +1048,7 @@ async def receive_request( # or terminal run, clear it and start a new run. if session.current_run_id is not None: existing_run = self._runs.get(session.current_run_id) - if existing_run is None or existing_run._status in ( - RunStatus.failed, - RunStatus.completed, - RunStatus.done, - ): + if existing_run is None or existing_run._run_state == RunState.DONE: session.current_run_id = None if session.current_run_id is None: return self._start_run_handle(session, agent, session_id, content_str, deps=deps) diff --git a/src/agentpool/orchestrator/session_pool.py b/src/agentpool/orchestrator/session_pool.py index 43ff67f3a..180378681 100644 --- a/src/agentpool/orchestrator/session_pool.py +++ b/src/agentpool/orchestrator/session_pool.py @@ -20,7 +20,7 @@ ) from agentpool.log import get_logger from agentpool.orchestrator.event_bus import EventBus -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.session_controller import ( CheckpointMismatchError, SessionBusyError, @@ -1064,7 +1064,7 @@ async def receive_request( @property def active_runs(self) -> list[RunHandle]: """Get all currently active (running) RunHandles.""" - return [rh for rh in self.sessions._runs.values() if rh.status == RunStatus.running] + return [rh for rh in self.sessions._runs.values() if rh.is_running] def get_run(self, run_id: str) -> RunHandle | None: """Get a RunHandle by ID. diff --git a/src/agentpool/orchestrator/turn.py b/src/agentpool/orchestrator/turn.py index b90e179b5..d4f036284 100644 --- a/src/agentpool/orchestrator/turn.py +++ b/src/agentpool/orchestrator/turn.py @@ -91,17 +91,35 @@ class ACPTurn(HookAwareTurn, Turn): ... - Implement the three abstract properties: :attr:`_hook_env`, :attr:`_hook_agent_name`, and :attr:`_hook_prompt`. - All methods are no-ops when ``self._hooks`` is ``None``. The - ``hooks_fired`` set on :attr:`_run_ctx` prevents double-firing when both - the old (capability-based) and new (mixin-based) code paths are active - during migration. + All methods are no-ops when ``self._hooks`` is ``None``. + + Tool execution logging idempotency is tracked via :attr:`_logged_tools`, + a per-Turn-instance set. A new Turn is created for each turn, so the set + does not need cross-turn reset. """ _hooks: AgentHooks | None """Hooks container, set by host class ``__init__``. ``None`` = no hooks.""" _run_ctx: AgentRunContext - """Per-run context, set by host class ``__init__``. Provides ``hooks_fired``.""" + """Per-run context, set by host class ``__init__``.""" + + _logged_tools: set[str] + """Per-Turn set of tool log keys already logged to the journal. + + This Turn instance is not reused across turns — _logged_tools does not + need reset. If Turn instances are ever reused (e.g., for retry), add a + ``reset()`` method to clear this set. + """ + + def __init__(self) -> None: + """Initialize the HookAwareTurn mixin. + + Host classes must call ``super().__init__()`` in their own + ``__init__`` to ensure ``_logged_tools`` is initialized. + """ + super().__init__() + self._logged_tools: set[str] = set() @property @abstractmethod @@ -134,9 +152,6 @@ async def _fire_pre_turn_hooks(self) -> HookResult | None: """ if self._hooks is None: return None - if "pre_turn" in self._run_ctx.hooks_fired: - return None - self._run_ctx.hooks_fired.add("pre_turn") return await self._hooks.run_pre_turn_hooks( agent_name=self._hook_agent_name, prompt=self._hook_prompt, @@ -164,9 +179,6 @@ async def _fire_post_turn_hooks( """ if self._hooks is None: return None - if "post_turn" in self._run_ctx.hooks_fired: - return None - self._run_ctx.hooks_fired.add("post_turn") return await self._hooks.run_post_turn_hooks( agent_name=self._hook_agent_name, prompt=self._hook_prompt, @@ -195,13 +207,6 @@ async def _fire_pre_tool_hooks( """ if self._hooks is None: return None - if tool_call_id is not None: - guard_key = f"pre_tool_use:{tool_call_id}" - else: - guard_key = f"pre_tool_use:{tool_name}" - if guard_key in self._run_ctx.hooks_fired: - return None - self._run_ctx.hooks_fired.add(guard_key) return await self._hooks.run_pre_tool_hooks( agent_name=self._hook_agent_name, tool_name=tool_name, @@ -220,9 +225,8 @@ def _log_tool_execution( """Log a tool execution record to the journal for crash recovery. Creates a :class:`ToolExecutionRecord` and stores it via - ``journal.log_tool_execution()``. Uses a separate guard key - (``"tool_log:{tool_call_id}"``) to prevent double-logging when - both old and new code paths are active. + ``journal.log_tool_execution()``. Uses :attr:`_logged_tools` to + prevent double-logging within a single Turn instance. Skips silently if any of the following are missing: - ``run_ctx._run_handle`` (not running inside a pooled session) @@ -237,14 +241,16 @@ def _log_tool_execution( """ from agentpool.lifecycle.types import ToolExecutionRecord - # Separate guard to prevent double-logging. + # Per-Turn idempotency guard using _logged_tools set. + # This Turn instance is not reused across turns — _logged_tools + # does not need reset. if tool_call_id is not None: - log_guard = f"tool_log:{tool_call_id}" + log_key = f"tool_log:{tool_call_id}" else: - log_guard = f"tool_log:{tool_name}" - if log_guard in self._run_ctx.hooks_fired: + log_key = f"tool_log:{tool_name}" + if log_key in self._logged_tools: return - self._run_ctx.hooks_fired.add(log_guard) + self._logged_tools.add(log_key) run_handle = self._run_ctx._run_handle if run_handle is None: @@ -293,13 +299,6 @@ async def _fire_post_tool_hooks( if self._hooks is None: return None - if tool_call_id is not None: - guard_key = f"post_tool_use:{tool_call_id}" - else: - guard_key = f"post_tool_use:{tool_name}" - if guard_key in self._run_ctx.hooks_fired: - return None - self._run_ctx.hooks_fired.add(guard_key) return await self._hooks.run_post_tool_hooks( agent_name=self._hook_agent_name, tool_name=tool_name, diff --git a/src/agentpool_server/acp_server/acp_agent.py b/src/agentpool_server/acp_server/acp_agent.py index 553eb2fea..b6d06dd90 100644 --- a/src/agentpool_server/acp_server/acp_agent.py +++ b/src/agentpool_server/acp_server/acp_agent.py @@ -251,17 +251,19 @@ def __post_init__(self) -> None: self.client_capabilities: ClientCapabilities | None = None self.client_info: Implementation | None = None ctx = self.host_context - if ctx is None or ctx.pool is None: + pool = self.default_agent._agent_pool + if pool is None: msg = "Default agent has no associated pool" raise RuntimeError(msg) if self.session_manager is None: - self.session_manager = ACPSessionManager(pool=ctx.pool) + self.session_manager = ACPSessionManager(pool=pool) self.tasks = TaskManager() self._initialized = False self._sessions_cache: ListSessionsResponse | None = None self._sessions_cache_time: float = 0.0 # Connect to title generation signal to notify clients of session updates - ctx.storage.metadata_generated.connect(self._on_metadata_generated) + if ctx is not None: + ctx.storage.metadata_generated.connect(self._on_metadata_generated) # Initialize MCP-over-ACP connection manager self._mcp_manager = AcpMcpConnectionManager() @@ -981,7 +983,7 @@ async def _swap_session_agent(self, session_id: str, new_agent_name: str) -> dic raise RequestError.invalid_params(msg) # Block swap during active prompt - if hasattr(session, "_task_lock") and session._task_lock.locked(): + if session.is_busy: msg = {"session_id": session_id, "reason": "Prompt active"} raise RequestError.invalid_params(msg) @@ -1109,7 +1111,8 @@ async def swap_pool(self, config_path: str, agent_name: str | None = None) -> li # 6. Update cached agent config from new pool ctx = new_agent.host_context - if ctx is None or ctx.pool is None: + new_pool = new_agent._agent_pool + if ctx is None or new_pool is None: msg = "New agent has no associated pool" raise RuntimeError(msg) @@ -1133,7 +1136,7 @@ async def swap_pool(self, config_path: str, agent_name: str | None = None) -> li # 7. Update default_agent reference and pool self.default_agent = new_agent - self.session_manager._pool = ctx.pool + self.session_manager._pool = new_pool # 8. Invalidate sessions cache self._sessions_cache = None diff --git a/src/agentpool_server/acp_server/session.py b/src/agentpool_server/acp_server/session.py index bb6fe7e04..f65045a9f 100644 --- a/src/agentpool_server/acp_server/session.py +++ b/src/agentpool_server/acp_server/session.py @@ -32,6 +32,7 @@ from agentpool.commands.base import NodeCommand from agentpool.log import get_logger from agentpool.mcp_server.config_snapshot import McpConfigEntry, McpConfigSnapshot +from agentpool_config.commands import BaseCommandConfig from agentpool_server.acp_server.converters import ( convert_acp_mcp_server_to_config, from_acp_content, @@ -272,6 +273,15 @@ async def permission_callback( self.log.info("Created ACP session", current_agent=self.agent.name) + @property + def is_busy(self) -> bool: + """Whether the session is currently processing a prompt. + + Returns: + True if the task lock is held (active prompt processing). + """ + return self._task_lock.locked() + def _register_manifest_commands(self) -> None: """Register global commands from manifest to command_store. @@ -303,7 +313,7 @@ def _register_manifest_commands(self) -> None: "Failed to register manifest command", name=cmd_name, config_type=type(cmd_config).__name__ - if hasattr(cmd_config, "type") + if isinstance(cmd_config, BaseCommandConfig) else "unknown", ) @@ -396,9 +406,10 @@ async def initialize_mcp_servers(self) -> None: """Initialize MCP servers if any are configured. Session-level MCP servers are converted to :class:`McpConfigEntry` - objects and merged into the agent's ``_mcp_snapshot``. For - ACP-transport servers, the transport is created and stored in the - agent's ``_session_connection_pool`` so that snapshot-aware + objects and merged into the session's MCP config snapshot via + :meth:`MCPManager.update_session_snapshot`. For ACP-transport + servers, the transport is registered via + :meth:`MCPManager.add_acp_transport` so that snapshot-aware capability building can reuse it. """ if not self.mcp_servers: @@ -433,16 +444,10 @@ async def _init_server(server: McpServer) -> None: transport = AcpMcpTransport(conn, timeout=600.0) if isinstance(self.agent, Agent): - # Store on the agent's session connection pool - # if one is available (legacy path). - if self.agent._session_connection_pool is not None: - await self.agent._session_connection_pool.add_transport( - cfg.client_id, transport - ) - # Always register the ACP transport on the - # MCPManager's session context so that - # get_capabilities() can find it and child sessions - # can inherit it via copy_pre_created_transports(). + # Register the ACP transport on the MCPManager's + # session context so that get_capabilities() can + # find it and child sessions can inherit it via + # copy_pre_created_transports(). await self.agent.mcp.add_acp_transport( self.session_id, cfg.client_id, @@ -479,7 +484,8 @@ async def _init_server(server: McpServer) -> None: # Merge new session configs into the agent's MCP snapshot, deduplicating # by client_id so that re-initialisation does not duplicate entries. if entries and isinstance(self.agent, Agent): - existing = self.agent._mcp_snapshot + ctx = self.agent.mcp.get_session_context(self.session_id) + existing = ctx.snapshot if ctx is not None else None existing_session = existing.session_configs if existing is not None else () seen_ids: set[str] = {e.server_config.client_id for e in existing_session} merged: list[McpConfigEntry] = list(existing_session) @@ -488,7 +494,6 @@ async def _init_server(server: McpServer) -> None: merged.append(entry) seen_ids.add(entry.server_config.client_id) new_snapshot = (existing or McpConfigSnapshot()).with_session_configs(tuple(merged)) - self.agent._mcp_snapshot = new_snapshot # Sync the updated snapshot to the MCPManager's session context # so that get_capabilities(session_id) can discover ACP MCP configs # and child sessions can inherit them via copy_pre_created_transports(). diff --git a/src/agentpool_server/opencode_server/event_processor.py b/src/agentpool_server/opencode_server/event_processor.py index 9fb8ccb1b..88fba1c3e 100644 --- a/src/agentpool_server/opencode_server/event_processor.py +++ b/src/agentpool_server/opencode_server/event_processor.py @@ -42,9 +42,14 @@ PartDeltaEvent, PartUpdatedEvent, SessionErrorEvent, + SessionStatusEvent, TokenCache, Tokens, ) + +# Cross-layer import: McpToolsChangedEvent is an OpenCode SSE event that +# EventProcessor creates from core-layer ChangeEvent(kind="tools_changed"). +from agentpool_server.opencode_server.models.events import McpToolsChangedEvent from agentpool_server.opencode_server.models.parts import ( ReasoningPart, StepFinishPart, @@ -71,6 +76,7 @@ ) from agentpool_server.opencode_server.models.events import Event from agentpool_server.opencode_server.models.parts import ToolState + from agentpool_server.opencode_server.models.session import SessionStatusType logger = get_logger(__name__) @@ -88,6 +94,23 @@ class EventProcessor: def __init__(self) -> None: """Initialize the event processor.""" + @staticmethod + def create_mcp_tools_changed_event(server: str) -> McpToolsChangedEvent: + """Create an McpToolsChangedEvent for tool list refresh notification. + + Called by the server's ``_watch_mcp_tool_changes`` task when a + ``ChangeEvent(kind="tools_changed")`` is received from + ``McpServerCap.on_change()``. The resulting event is broadcast + to connected OpenCode clients so they can refresh their tool lists. + + Args: + server: Name of the MCP server whose tools changed. + + Returns: + ``McpToolsChangedEvent`` ready for broadcasting. + """ + return McpToolsChangedEvent.create(server=server) + async def process( self, event: RichAgentStreamEvent[Any], @@ -186,9 +209,14 @@ async def process( for e in self._process_tool_complete(ctx, tool_call_id, result, event_metadata): yield e - case StreamCompleteEvent(message=msg) if msg: + case StreamCompleteEvent(message=msg, cancelled=cancelled) if msg: for e in self._process_stream_complete(ctx, msg): yield e + status: SessionStatusType = "cancelled" if cancelled else "idle" + yield SessionStatusEvent.create( + session_id=ctx.session_id, + status_type=status, + ) case RunErrorEvent() as run_error_event: yield SessionErrorEvent.create( diff --git a/src/agentpool_server/opencode_server/models/events.py b/src/agentpool_server/opencode_server/models/events.py index 00fea41a0..fa5506bf7 100644 --- a/src/agentpool_server/opencode_server/models/events.py +++ b/src/agentpool_server/opencode_server/models/events.py @@ -848,10 +848,14 @@ class McpToolsChangedProperties(OpenCodeBaseModel): class McpToolsChangedEvent(OpenCodeBaseModel): - """MCP tools changed event - emitted when an MCP server's tool list changes. + """MCP tools changed event — emitted when an MCP server's tool list changes. - TODO: Hook into MCP SDK's ToolListChangedNotification to emit this event. - OpenCode only emits this from the notification handler, not on connect/disconnect. + Wired via: ``McpServerCap._on_tools_changed()`` → ``ChangeEvent(kind="tools_changed")`` + → ``ExtensionRegistry.merge_change_streams()`` → ``server._watch_mcp_tool_changes()`` + → ``EventProcessor.create_mcp_tools_changed_event()`` → ``state.broadcast_event()``. + + The ``ChangeEvent`` (core capability layer) is converted to this OpenCode SSE event + by the server layer, keeping the event type in OpenCode server models only. """ type: Literal["mcp.tools.changed"] = Field(default="mcp.tools.changed", init=False) diff --git a/src/agentpool_server/opencode_server/models/session.py b/src/agentpool_server/opencode_server/models/session.py index 0a894e9ba..f2084f72d 100644 --- a/src/agentpool_server/opencode_server/models/session.py +++ b/src/agentpool_server/opencode_server/models/session.py @@ -13,7 +13,7 @@ ) -SessionStatusType = Literal["idle", "busy", "retry"] +SessionStatusType = Literal["idle", "busy", "retry", "cancelled"] TodoStatus = Literal["pending", "in_progress", "completed"] TodoPriority = Literal["high", "medium", "low"] diff --git a/src/agentpool_server/opencode_server/routes/agent_routes.py b/src/agentpool_server/opencode_server/routes/agent_routes.py index 2d42b30a3..4e6723d12 100644 --- a/src/agentpool_server/opencode_server/routes/agent_routes.py +++ b/src/agentpool_server/opencode_server/routes/agent_routes.py @@ -159,7 +159,8 @@ async def list_skills(state: StateDep) -> list[SkillInfo]: # 1. Get MCP provider skills from skill resolver first # These will be overridden by local skills if names conflict - skill_resolver = ctx.pool.skill_resolver if ctx.pool is not None else None + pool = state.agent._agent_pool + skill_resolver = pool.skill_resolver if pool is not None else None if skill_resolver is not None: try: for provider_name in skill_resolver.list_providers(): diff --git a/src/agentpool_server/opencode_server/routes/message_routes.py b/src/agentpool_server/opencode_server/routes/message_routes.py index 4eeb52782..d656598b3 100644 --- a/src/agentpool_server/opencode_server/routes/message_routes.py +++ b/src/agentpool_server/opencode_server/routes/message_routes.py @@ -8,8 +8,8 @@ from fastapi import APIRouter, HTTPException, Query, status +from agentpool.lifecycle import RunOutcome from agentpool.log import get_logger -from agentpool.orchestrator.run import RunStatus from agentpool.utils import identifiers as identifier from agentpool.utils.time_utils import now_ms from agentpool_server.opencode_server.converters import ( @@ -583,7 +583,7 @@ async def _feed_adapter() -> None: await session_pool.event_bus.unsubscribe(session_id, event_stream) # Finalize based on run outcome - if run_handle.status != RunStatus.failed: + if run_handle.outcome != RunOutcome.FAILED: for oc_event in adapter.finalize(): await state.broadcast_event(oc_event) diff --git a/src/agentpool_server/opencode_server/server.py b/src/agentpool_server/opencode_server/server.py index ea03933e4..6d7eedde6 100644 --- a/src/agentpool_server/opencode_server/server.py +++ b/src/agentpool_server/opencode_server/server.py @@ -234,6 +234,37 @@ async def _watch_skill_changes() -> None: state._skill_change_task = asyncio.create_task(_watch_skill_changes()) + # Watch for MCP tool list changes and broadcast McpToolsChangedEvent. + # McpServerCap.on_change() yields ChangeEvent(kind="tools_changed") + # when an MCP server's tool list changes (via notifications/tools/list_changed). + # This watcher converts that to McpToolsChangedEvent and broadcasts it + # so connected OpenCode clients can refresh their tool lists. + from agentpool_server.opencode_server.event_processor import EventProcessor + + processor = EventProcessor() + + async def _watch_mcp_tool_changes() -> None: + """Watch for MCP tool change events and broadcast notifications.""" + stream = extension_registry.merge_change_streams(Scope(level=ScopeLevel.POOL)) + if stream is None: + return + async for event in stream: + if event.kind != "tools_changed": + continue + logger.info( + "MCP tools changed, broadcasting notification", + server=event.capability_name, + ) + try: + oc_event = processor.create_mcp_tools_changed_event( + server=event.capability_name, + ) + await state.broadcast_event(oc_event) + except Exception: + logger.exception("Failed to broadcast McpToolsChangedEvent") + + state._mcp_tool_change_task = asyncio.create_task(_watch_mcp_tool_changes()) + # Set up todo change callback to broadcast events async def on_todo_change(tracker: TodoTracker) -> None: """Broadcast todo updates to all active sessions.""" @@ -386,6 +417,16 @@ async def check_for_updates() -> None: except Exception: logger.exception("Error during skill change task cleanup") state._skill_change_task = None + # Cancel MCP tool change watcher + if state._mcp_tool_change_task is not None: + state._mcp_tool_change_task.cancel() + try: + await state._mcp_tool_change_task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Error during MCP tool change task cleanup") + state._mcp_tool_change_task = None # Then clean up background tasks await state.cleanup_tasks() # Then tear down watchers and shared infrastructure diff --git a/src/agentpool_server/opencode_server/session_pool_integration.py b/src/agentpool_server/opencode_server/session_pool_integration.py index 37f1e1ea6..f48601b88 100644 --- a/src/agentpool_server/opencode_server/session_pool_integration.py +++ b/src/agentpool_server/opencode_server/session_pool_integration.py @@ -19,8 +19,8 @@ SpawnSessionStart, StreamCompleteEvent, ) +from agentpool.lifecycle import RunState from agentpool.log import get_logger -from agentpool.orchestrator.run import RunStatus from agentpool.utils import identifiers as identifier from agentpool.utils.time_utils import now_ms from agentpool_server.mixins import ProtocolEventConsumerMixin @@ -822,9 +822,9 @@ async def get_session_status(self, session_id: str) -> SessionStatus | None: run_id = session.current_run_id if run_id is not None: run_handle = self.session_pool.sessions._runs.get(run_id) - if run_handle is not None and run_handle.status in ( - RunStatus.pending, - RunStatus.running, + if run_handle is not None and run_handle._run_state in ( + RunState.IDLE, + RunState.RUNNING, ): return SessionStatus(type="busy") diff --git a/src/agentpool_server/opencode_server/state.py b/src/agentpool_server/opencode_server/state.py index e3fecbbad..ce3131465 100644 --- a/src/agentpool_server/opencode_server/state.py +++ b/src/agentpool_server/opencode_server/state.py @@ -89,6 +89,7 @@ class ServerState: skill_bridge: Any = field(default=None) command_store: CommandStore | None = field(default=None) _skill_change_task: Any = field(default=None, repr=False) + _mcp_tool_change_task: Any = field(default=None, repr=False) session_pool_integration: Any = field(default=None) session_controller: SessionController | None = field(default=None) event_bridge: Any = field(default=None, repr=False) @@ -115,8 +116,7 @@ def __post_init__(self) -> None: # Cache non-session-scoped dependencies directly so they remain # accessible even after the shared ``self.agent`` is removed in a # later migration step. - _ctx = self.agent.host_context - self._pool: AgentPool[Any] | None = _ctx.pool if _ctx is not None else None + self._pool: AgentPool[Any] | None = self.agent._agent_pool self._storage: StorageManager | None = self.agent.storage # Create a standalone execution environment for shell commands. diff --git a/src/agentpool_server/opencode_server/stream_adapter.py b/src/agentpool_server/opencode_server/stream_adapter.py index 93386367a..0b01c54f6 100644 --- a/src/agentpool_server/opencode_server/stream_adapter.py +++ b/src/agentpool_server/opencode_server/stream_adapter.py @@ -210,21 +210,6 @@ async def convert_event(self, event: RichAgentStreamEvent[Any]) -> AsyncIterator self._step_finish_emitted = True yield oc_event - async def _handle_event(self, event: RichAgentStreamEvent[Any]) -> AsyncIterator[Event]: - """Backward-compatible event handler that delegates to EventProcessor. - - This method is deprecated but kept for tests that directly call it. - Use :meth:`convert_event` or :meth:`process_stream` instead for new code. - - Args: - event: The agent stream event to process. - - Yields: - OpenCode Event objects for broadcasting. - """ - async for oc_event in self.processor.process(event, self.main_context): - yield oc_event - def finalize(self) -> Iterator[Event]: """Yield final events after the stream has ended. diff --git a/tests/agentpool_server/acp_server/test_e2e_session_lifecycle.py b/tests/agentpool_server/acp_server/test_e2e_session_lifecycle.py index c304a51c8..4850f657f 100644 --- a/tests/agentpool_server/acp_server/test_e2e_session_lifecycle.py +++ b/tests/agentpool_server/acp_server/test_e2e_session_lifecycle.py @@ -173,7 +173,7 @@ async def test_e2e_session_lifecycle() -> None: # noqa: PLR0915 # Phase 7: VERIFY - Fresh resources, no stale references # ================================================================ - # 7a: New _SessionContext is a different object + # 7a: New McpSessionContext is a different object assert new_ctx is not old_ctx # 7b: New toolset_cache is fresh (different object, empty) diff --git a/tests/agentpool_server/acp_server/test_resume_active_run.py b/tests/agentpool_server/acp_server/test_resume_active_run.py index c3317cba4..e9ec6e179 100644 --- a/tests/agentpool_server/acp_server/test_resume_active_run.py +++ b/tests/agentpool_server/acp_server/test_resume_active_run.py @@ -15,7 +15,8 @@ import pytest -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.lifecycle import RunState +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.session_controller import SessionController, SessionState @@ -51,7 +52,7 @@ async def test_resume_with_active_run() -> None: agent_type="native", ) # Simulate a running run: status is running, complete_event NOT set. - run_handle.status = RunStatus.running + run_handle._run_state = RunState.RUNNING controller._runs[run_id] = run_handle # close_session should complete within 30s (internal 2s timeout + cleanup). diff --git a/tests/agentpool_server/acp_server/test_resume_reconnect.py b/tests/agentpool_server/acp_server/test_resume_reconnect.py index 1c6c4ec44..4e283f26c 100644 --- a/tests/agentpool_server/acp_server/test_resume_reconnect.py +++ b/tests/agentpool_server/acp_server/test_resume_reconnect.py @@ -11,7 +11,7 @@ 4. Verify old ACP connections are cleaned up (``_session_connections`` no longer has the session_id, the ``AcpMcpConnection`` is removed). 5. Simulate reconnect + resume by calling ``get_or_create_session`` again - (creates a fresh ``_SessionContext``) and registering a new ACP + (creates a fresh ``McpSessionContext``) and registering a new ACP connection with a new ``connection_id``. 6. Verify the new connection is fresh — different object, different ``connection_id``, no stale references from the pre-disconnect session. @@ -114,7 +114,7 @@ async def test_resume_after_reconnect() -> None: # noqa: PLR0915 assert mcp_manager.get_session_context(session_id) is None # --- Step 5: Reconnect + resume — create fresh session context --- - # get_or_create_session creates a NEW _SessionContext with fresh state + # get_or_create_session creates a NEW McpSessionContext with fresh state new_ctx = mcp_manager.get_or_create_session(session_id) # Create a fresh ACP connection (new connection_id, new session_key) @@ -135,7 +135,7 @@ async def test_resume_after_reconnect() -> None: # noqa: PLR0915 # --- Step 6: Verify fresh connections, no stale references --- - # 6a: New _SessionContext is a different object from the old one + # 6a: New McpSessionContext is a different object from the old one assert new_ctx is not old_ctx # 6b: New context has fresh toolset_cache (empty, different object) diff --git a/tests/agentpool_server/acp_server/test_resume_session_lifecycle.py b/tests/agentpool_server/acp_server/test_resume_session_lifecycle.py index e22d19407..5a6bbafcb 100644 --- a/tests/agentpool_server/acp_server/test_resume_session_lifecycle.py +++ b/tests/agentpool_server/acp_server/test_resume_session_lifecycle.py @@ -2,7 +2,7 @@ Verifies that resume_session() closes the old session (removing it from _acp_sessions) and creates a fresh session with new MCP resources -(different _SessionContext in MCPManager's session context). +(different McpSessionContext in MCPManager's session context). """ from __future__ import annotations @@ -28,7 +28,7 @@ async def test_resume_closes_old_session() -> None: 4. Call resume_session() with the same session_id. 5. Verify old session was closed (popped from _acp_sessions, close called). 6. Verify new session is a different object in _acp_sessions. - 7. Verify MCPManager's session context has a fresh _SessionContext + 7. Verify MCPManager's session context has a fresh McpSessionContext (different object, different toolset_cache, different connection_pool). """ session_id = "test-resume-lifecycle-1" diff --git a/tests/agentpool_server/acp_server/test_websocket_disconnect_during_run.py b/tests/agentpool_server/acp_server/test_websocket_disconnect_during_run.py index c6b3892ce..39a806e25 100644 --- a/tests/agentpool_server/acp_server/test_websocket_disconnect_during_run.py +++ b/tests/agentpool_server/acp_server/test_websocket_disconnect_during_run.py @@ -21,7 +21,8 @@ import pytest -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.lifecycle import RunState +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.session_controller import SessionController, SessionState from agentpool_server.acp_server.session_manager import ACPSessionManager @@ -58,7 +59,7 @@ async def test_websocket_disconnect_during_run() -> None: session_id=session_id, agent_type="native", ) - run_handle.status = RunStatus.running + run_handle._run_state = RunState.RUNNING controller._runs[run_id] = run_handle mock_acp_session = Mock() diff --git a/tests/agents/acp_agent/test_acp_turn_hooks.py b/tests/agents/acp_agent/test_acp_turn_hooks.py index 13e28f953..eb9768f84 100644 --- a/tests/agents/acp_agent/test_acp_turn_hooks.py +++ b/tests/agents/acp_agent/test_acp_turn_hooks.py @@ -126,7 +126,6 @@ def _make_turn( acp_client=client, # type: ignore[arg-type] prompts=["do something"], run_ctx=_make_run_ctx(), - message_history=[], session_id="test-acp-session", agent_name="test-acp-agent", hooks=hooks, diff --git a/tests/agents/acp_agent/test_create_turn.py b/tests/agents/acp_agent/test_create_turn.py index 4ccbc9de0..b28af0f41 100644 --- a/tests/agents/acp_agent/test_create_turn.py +++ b/tests/agents/acp_agent/test_create_turn.py @@ -7,8 +7,11 @@ import pytest from acp import InitializeRequest +from acp.agent.acp_agent_api import ACPAgentAPI from agentpool.agents.acp_agent import ACPAgent -from agentpool.agents.acp_agent.turn import ACPTurn +from agentpool.agents.acp_agent.client_handler import TimeoutableEvent +from agentpool.agents.acp_agent.session_state import ACPSessionState +from agentpool.agents.acp_agent.turn import ACPClientProtocol, ACPTurn from agentpool.agents.context import AgentRunContext @@ -28,3 +31,23 @@ def test_acp_agent_create_turn_returns_acp_turn() -> None: ) assert isinstance(turn, ACPTurn) + + +@pytest.mark.unit +def test_acp_agent_api_satisfies_acp_client_protocol() -> None: + """ACPAgentAPI with state+event satisfies ACPClientProtocol (no cast needed).""" + connection = MagicMock() + state = ACPSessionState(session_id="test-session") + update_event = TimeoutableEvent() + api = ACPAgentAPI(connection, state=state, update_event=update_event) + + assert isinstance(api, ACPClientProtocol) + + +@pytest.mark.unit +def test_acp_agent_api_satisfies_protocol_without_state() -> None: + """ACPAgentAPI satisfies ACPClientProtocol even without state (methods exist).""" + connection = MagicMock() + api = ACPAgentAPI(connection) + + assert isinstance(api, ACPClientProtocol) diff --git a/tests/agents/acp_agent/test_turn.py b/tests/agents/acp_agent/test_turn.py index 28f4dac9e..32e972c4a 100644 --- a/tests/agents/acp_agent/test_turn.py +++ b/tests/agents/acp_agent/test_turn.py @@ -85,7 +85,6 @@ async def test_acp_turn_prompt_stream_complete_cycle() -> None: acp_client=client, prompts=["Say hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -119,7 +118,6 @@ async def test_acp_turn_prompt_error_yields_run_error_event() -> None: acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -147,7 +145,6 @@ async def test_acp_turn_stream_error_yields_run_error_event() -> None: acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -173,7 +170,6 @@ async def test_acp_turn_message_history_and_final_message_after_execute() -> Non acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -203,7 +199,6 @@ async def test_acp_turn_properties_raise_before_execute() -> None: acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -224,7 +219,6 @@ async def test_acp_turn_cancelled_error_propagates() -> None: acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -246,7 +240,6 @@ async def test_acp_turn_cancelled_error_during_stream_propagates() -> None: acp_client=client, prompts=["Hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -268,7 +261,6 @@ async def test_acp_turn_empty_prompts_uses_empty_string() -> None: acp_client=client, prompts=[], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) diff --git a/tests/agents/acp_agent/test_turn_integration.py b/tests/agents/acp_agent/test_turn_integration.py index 4c26ff6ee..e830e1853 100644 --- a/tests/agents/acp_agent/test_turn_integration.py +++ b/tests/agents/acp_agent/test_turn_integration.py @@ -21,8 +21,9 @@ PartDeltaEvent, StreamCompleteEvent, ) +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.turn import Turn @@ -95,7 +96,6 @@ async def test_acp_turn_full_cycle_with_mock_client() -> None: acp_client=client, prompts=["Say hello"], run_ctx=run_ctx, - message_history=[], session_id="test-session", ) @@ -134,7 +134,7 @@ async def test_acp_turn_full_cycle_with_mock_client() -> None: class _BlockingTurn(Turn): """Stub Turn that blocks on a release event before completing. - Used to keep _status == RunStatus.running long enough to call steer(). + Used to keep _status == RunState.RUNNING long enough to call steer(). """ def __init__(self, *, release_event: asyncio.Event) -> None: @@ -186,7 +186,7 @@ async def _consume() -> None: await asyncio.sleep(0.05) # Turn should be running (blocked on release_event inside _BlockingTurn) - assert handle._status == RunStatus.running + assert handle._run_state == RunState.RUNNING # ACP path does not set active_agent_run (only NativeTurn does) assert handle.active_agent_run is None @@ -202,7 +202,7 @@ async def _consume() -> None: await asyncio.sleep(0.05) await consumer_task - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE # --------------------------------------------------------------------------- diff --git a/tests/agents/native_agent/test_inject_prompt_cross_task.py b/tests/agents/native_agent/test_inject_prompt_cross_task.py index e495a10b8..63e5de023 100644 --- a/tests/agents/native_agent/test_inject_prompt_cross_task.py +++ b/tests/agents/native_agent/test_inject_prompt_cross_task.py @@ -300,7 +300,7 @@ async def test_inject_prompt_triggers_continuation(slow_agent: Agent[None]) -> N """inject_prompt from a different task should cause run_stream to continue. The run_stream() loop checks for pending injections - after each _run_stream_once iteration. If inject_prompt() successfully + after each _stream_events iteration. If inject_prompt() successfully delivers to the injection manager (via SessionPool fallback), and the injection gets flushed, the loop should run another iteration. """ diff --git a/tests/agents/native_agent/test_native_turn_hooks.py b/tests/agents/native_agent/test_native_turn_hooks.py index c1c17c480..9e2814faa 100644 --- a/tests/agents/native_agent/test_native_turn_hooks.py +++ b/tests/agents/native_agent/test_native_turn_hooks.py @@ -166,10 +166,10 @@ async def test_tool_hooks_not_fired_by_hook_aware_turn_for_native() -> None: ``_fire_pre_tool_hooks`` / ``_fire_post_tool_hooks`` methods. The mixin methods are never called by ``NativeTurn.execute()``. - We verify this by checking that ``hooks_fired`` does NOT contain the - ``pre_tool_use:*`` or ``post_tool_use:*`` guard keys that HookAwareTurn - would set. The tool hooks themselves DO fire (via the capability), but - the mixin's guard mechanism is not used. + We verify this by checking that ``_logged_tools`` does NOT contain any + tool log keys. If the mixin's ``_fire_post_tool_hooks`` were called, + it would have invoked ``_log_tool_execution`` which adds keys to + ``_logged_tools``. """ _reset_calls() hooks = AgentHooks( @@ -201,30 +201,37 @@ def simple_tool() -> str: _ = [event async for event in turn.execute()] - # pre_turn and post_turn fire via HookAwareTurn (guard keys present) - assert "pre_turn" in run_ctx.hooks_fired - assert "post_turn" in run_ctx.hooks_fired + # pre_turn and post_turn fire via HookAwareTurn + event_names = [name for name, _ in hook_calls] + assert "pre_turn" in event_names + assert "post_turn" in event_names # Tool hooks fire via pydantic-ai Hooks capability, but HookAwareTurn's - # guard keys are NOT set because NativeTurn.execute() never calls - # _fire_pre_tool_hooks() / _fire_post_tool_hooks(). - tool_guard_keys = [ - k for k in run_ctx.hooks_fired if k.startswith(("pre_tool_use:", "post_tool_use:")) - ] - assert len(tool_guard_keys) == 0, ( - f"HookAwareTurn should not set tool guard keys for native agents, " - f"but found: {tool_guard_keys}" + # _log_tool_execution is NOT called because NativeTurn.execute() never + # calls _fire_pre_tool_hooks() / _fire_post_tool_hooks(). + assert len(turn._logged_tools) == 0, ( + f"HookAwareTurn should not log tool executions for native agents, " + f"but found: {turn._logged_tools}" ) # --------------------------------------------------------------------------- -# Test: hooks_fired guard prevents double-firing through old code path +# Test: hooks fire correctly without double-fire guard # --------------------------------------------------------------------------- @pytest.mark.unit async def test_hooks_fired_prevents_double_firing_via_old_path() -> None: - """Given hooks_fired already has 'pre_turn', a second fire returns None.""" + """Given the removal of hooks_fired guard, hooks fire on each call. + + The old ``hooks_fired`` double-fire guard was removed after T3 eliminated + the ACP standalone path that caused double-firing. With only the + ``Turn.execute()`` path active, hooks fire exactly once per turn. + + This test verifies that calling ``_fire_pre_turn_hooks`` after + ``execute()`` completes fires hooks again (no guard to block them). + This is the expected behavior — the guard is no longer needed. + """ _reset_calls() hooks = AgentHooks(pre_turn=[_make_recorder("pre_turn")]) agent = Agent( @@ -245,10 +252,10 @@ async def test_hooks_fired_prevents_double_firing_via_old_path() -> None: # Execute the turn (fires hooks once) _ = [event async for event in turn.execute()] - # Attempt to fire again — should be no-op + # Without the hooks_fired guard, calling again fires again. result = await turn._fire_pre_turn_hooks() - assert result is None + assert result is not None - # Only one pre_turn call despite the second attempt + # Two pre_turn calls: one from execute(), one from manual call pre_turn_count = sum(1 for name, _ in hook_calls if name == "pre_turn") - assert pre_turn_count == 1 + assert pre_turn_count == 2 diff --git a/tests/agents/test_base_agent_run_v2.py b/tests/agents/test_base_agent_run_v2.py index d6b44397c..ee0a24d74 100644 --- a/tests/agents/test_base_agent_run_v2.py +++ b/tests/agents/test_base_agent_run_v2.py @@ -14,9 +14,10 @@ from agentpool.agents.context import AgentRunContext from agentpool.agents.events import RunErrorEvent, RunStartedEvent, StreamCompleteEvent from agentpool.agents.native_agent.agent import Agent +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventBus -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.turn import Turn @@ -101,7 +102,7 @@ async def test_create_run_returns_run_handle_without_executing() -> None: ) assert isinstance(handle, RunHandle) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert handle._closing is False @@ -319,7 +320,7 @@ def test_no_duplicate_stream_complete_in_run_once() -> None: ) -def test_no_duplicate_stream_complete_in_run_stream_once() -> None: +def test_no_duplicate_stream_complete_in_stream_events() -> None: """_stream_events must not publish StreamCompleteEvent after turn.execute(). NativeTurn.execute() already yields StreamCompleteEvent as its terminal diff --git a/tests/agents/test_capability_hooks_standalone.py b/tests/agents/test_capability_hooks_standalone.py index 2bf25bead..9d5784435 100644 --- a/tests/agents/test_capability_hooks_standalone.py +++ b/tests/agents/test_capability_hooks_standalone.py @@ -1,7 +1,7 @@ """Tests verifying pdai Capability hooks fire on standalone run path. Phase 2 of thin-wrapper refactor: BaseAgent.run_stream() now delegates -directly to _run_stream_once() → _stream_events() → NativeTurn.execute() +directly to _stream_events() → NativeTurn.execute() which calls agent_run.next(node) explicitly, ensuring all pdai Capability hooks fire on every run path. diff --git a/tests/agents/test_create_turn.py b/tests/agents/test_create_turn.py index f02808fd0..f37a1c9e4 100644 --- a/tests/agents/test_create_turn.py +++ b/tests/agents/test_create_turn.py @@ -62,18 +62,43 @@ def test_acp_turn_joins_all_prompts_not_just_last() -> None: assert "third prompt" in new_result -def test_acp_adapter_has_todo_comment() -> None: - """ACP agent adapter gap must be documented with TODO, not just NOTE. +def test_acp_adapter_satisfies_protocol() -> None: + """ACP agent adapter must satisfy ACPClientProtocol (no TODO/cast needed). - The TODO comment must describe the required infrastructure - (async futures / notification registry) to prevent runtime crashes. + T1 built stream_events() and get_messages() on ACPAgentAPI, so the + create_turn method passes self._api directly without cast. This test + verifies the adapter is complete by checking that ACPAgentAPI has the + required methods and that create_turn does not use cast. """ + import inspect + + from acp.agent.acp_agent_api import ACPAgentAPI + from agentpool.agents.acp_agent.turn import ACPClientProtocol + + # ACPAgentAPI must have stream_events and get_messages methods + assert hasattr(ACPAgentAPI, "stream_events"), "ACPAgentAPI must have stream_events()" + assert hasattr(ACPAgentAPI, "get_messages"), "ACPAgentAPI must have get_messages()" + + # create_turn must not contain cast or TODO import agentpool.agents.acp_agent.acp_agent as acp_module source = inspect.getsource(acp_module.ACPAgent.create_turn) - assert "TODO" in source, "ACP adapter gap must be documented with TODO comment, not just NOTE" - assert "AttributeError" in source or "adapter" in source.lower(), ( - "TODO comment must describe the gap and required infrastructure" + assert "cast" not in source, "create_turn must not use cast — adapter is complete" + assert "TODO" not in source, "create_turn must not have TODO — adapter is complete" + + # ACPAgentAPI must satisfy ACPClientProtocol at runtime + from unittest.mock import MagicMock + + from agentpool.agents.acp_agent.client_handler import TimeoutableEvent + from agentpool.agents.acp_agent.session_state import ACPSessionState + + api = ACPAgentAPI( + MagicMock(), + state=ACPSessionState(session_id="test"), + update_event=TimeoutableEvent(), + ) + assert isinstance(api, ACPClientProtocol), ( + "ACPAgentAPI must satisfy ACPClientProtocol without cast" ) diff --git a/tests/agents/test_run_stream_direct_gating.py b/tests/agents/test_run_stream_direct_gating.py index e450bccc0..41bcac928 100644 --- a/tests/agents/test_run_stream_direct_gating.py +++ b/tests/agents/test_run_stream_direct_gating.py @@ -22,7 +22,7 @@ # --------------------------------------------------------------------------- -# Test helper: minimal BaseAgent subclass that spies on _run_stream_once +# Test helper: minimal BaseAgent subclass that spies on _stream_events # --------------------------------------------------------------------------- @@ -31,10 +31,10 @@ class _FakeEvent: class _GatingTestAgent(BaseAgent[None, str]): - """Test agent that tracks _run_stream_once calls. + """Test agent that tracks _stream_events calls. Subclasses override AGENT_TYPE to test native vs non-native gating. - ``_run_stream_once`` records every call and queues an extra prompt on the first + ``_stream_events`` records every call and queues an extra prompt on the first invocation, allowing the test to distinguish single-call (native) from multi-call (non-native) behaviour. """ @@ -63,9 +63,14 @@ async def _stream_events( user_msg: Any = None, **kwargs: Any, ) -> AsyncIterator[RichAgentStreamEvent[str]]: - # Not called because _run_stream_once is overridden below. - return - yield # pragma: no cover (make generator) + # Spy: record prompts passed to _stream_events. + self._call_log.append(tuple(prompts)) + # On the very first call, mark that an extra prompt was queued so + # the test can distinguish single-call (native) from multi-call + # (non-native) behaviour. + if not self._has_queued_extra: + self._has_queued_extra = True + yield _FakeEvent() # type: ignore[return-value] def create_turn( self, @@ -99,22 +104,6 @@ async def list_sessions( async def load_session(self, session_id: str) -> Any: return None - # -- spied method -------------------------------------------------------- - - async def _run_stream_once( - self, - run_ctx: AgentRunContext, - *prompts: Any, - **kwargs: Any, - ) -> AsyncIterator[RichAgentStreamEvent[str]]: - self._call_log.append(prompts) - # On the very first call, mark that an extra prompt was queued so - # the test can distinguish single-call (native) from multi-call - # (non-native) behaviour. - if not self._has_queued_extra: - self._has_queued_extra = True - yield _FakeEvent() # type: ignore[return-value] - class _NativeTestAgent(_GatingTestAgent): """Agent with AGENT_TYPE = 'native' (skips manual loop).""" @@ -137,7 +126,7 @@ async def test_native_agent_skips_manual_loop() -> None: """Native AGENT_TYPE should cause run_stream() to skip the while loop. When AGENT_TYPE == 'native', the extra prompt queued during - _run_stream_once must NOT be processed -- the method uses a simple + _stream_events must NOT be processed -- the method uses a simple ``async for`` and exits without re-checking the injection queue. """ call_log: list[tuple[Any, ...]] = [] @@ -146,9 +135,9 @@ async def test_native_agent_skips_manual_loop() -> None: events: list[object] = [] events.extend([event async for event in agent.run_stream("test prompt")]) - # Native path: _run_stream_once is called exactly once + # Native path: _stream_events is called exactly once assert len(call_log) == 1, ( - f"Expected 1 call to _run_stream_once for native agent, got {len(call_log)}" + f"Expected 1 call to _stream_events for native agent, got {len(call_log)}" ) # The queued extra prompt should still be in the injection manager assert agent._has_queued_extra, "Extra prompt should have been queued" diff --git a/tests/delegation/test_pool_get_context.py b/tests/delegation/test_pool_get_context.py index 37e66ab86..db1408fe4 100644 --- a/tests/delegation/test_pool_get_context.py +++ b/tests/delegation/test_pool_get_context.py @@ -44,14 +44,6 @@ def test_get_context_is_cached(pool: AgentPool[None]) -> None: assert ctx1 is ctx2 -@pytest.mark.unit -def test_get_context_pool_back_reference(pool: AgentPool[None]) -> None: - """HostContext.pool back-reference points to the originating pool.""" - ctx = pool.get_context() - - assert ctx.pool is pool - - @pytest.mark.unit def test_factory_returns_agent_factory(pool: AgentPool[None]) -> None: """_factory property returns an AgentFactory instance.""" diff --git a/tests/delegation/test_team_member_skills.py b/tests/delegation/test_team_member_skills.py index df49d9676..d6e50f342 100644 --- a/tests/delegation/test_team_member_skills.py +++ b/tests/delegation/test_team_member_skills.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -from unittest.mock import MagicMock import pytest @@ -15,12 +14,6 @@ class _Pool: async def get_skill_instructions_for_node(self, skill_name: str, node_name: str) -> str: return f"# {skill_name}\nUse this skill for {node_name}." - def get_context(self) -> MagicMock: - """Return a HostContext-like mock for host_context property.""" - ctx = MagicMock() - ctx.pool = self # pool points back to _Pool for skill_provider + methods - return ctx - def test_team_loads_member_skills_from_pool_provider() -> None: team = BaseTeam([], mode="parallel", name="review_team") diff --git a/tests/hooks/test_hook_aware_turn.py b/tests/hooks/test_hook_aware_turn.py index b2d8b6207..fc4e309ed 100644 --- a/tests/hooks/test_hook_aware_turn.py +++ b/tests/hooks/test_hook_aware_turn.py @@ -1,8 +1,10 @@ """Unit tests for HookAwareTurn mixin in isolation. -Tests the mixin's hook firing logic, guard key deduplication, and no-op -behavior when hooks are None. Uses a minimal host class that implements -the three required abstract properties. +Tests the mixin's hook firing logic and no-op behavior when hooks are None. +The hooks_fired double-fire guard was removed in T4 — hooks now fire on +every call without dedup. Tool execution logging idempotency is handled +by the per-Turn ``_logged_tools`` set. Uses a minimal host class that +implements the three required abstract properties. """ from __future__ import annotations @@ -58,6 +60,7 @@ def __init__( agent_name: str = "test-agent", prompt: str = "hello", ) -> None: + super().__init__() self._hooks = hooks self._run_ctx = run_ctx self._agent_name = agent_name @@ -117,13 +120,13 @@ async def test_all_four_hooks_fire_in_order() -> None: # --------------------------------------------------------------------------- -# Test: hooks_fired prevents double-firing +# Test: hooks fire on every call (hooks_fired guard removed in T4) # --------------------------------------------------------------------------- @pytest.mark.unit -async def test_pre_turn_hooks_dedup_on_double_call() -> None: - """Given two calls to _fire_pre_turn_hooks, only the first fires.""" +async def test_pre_turn_hooks_fire_on_double_call() -> None: + """Given two calls to _fire_pre_turn_hooks, both fire (no dedup guard).""" _reset_calls() hooks = AgentHooks(pre_turn=[_make_recording_hook("pre_turn")]) host = _make_host(hooks=hooks) @@ -132,22 +135,21 @@ async def test_pre_turn_hooks_dedup_on_double_call() -> None: result2 = await host._fire_pre_turn_hooks() assert result1 is not None - assert result2 is None - assert len(hook_calls) == 1 + assert result2 is not None + assert len(hook_calls) == 2 @pytest.mark.unit -async def test_post_turn_hooks_dedup_on_double_call() -> None: - """Given two calls to _fire_post_turn_hooks, only the first fires.""" +async def test_post_turn_hooks_fire_on_double_call() -> None: + """Given two calls to _fire_post_turn_hooks, both fire (no dedup guard).""" _reset_calls() hooks = AgentHooks(post_turn=[_make_recording_hook("post_turn")]) host = _make_host(hooks=hooks) await host._fire_post_turn_hooks(None) - result2 = await host._fire_post_turn_hooks(None) + await host._fire_post_turn_hooks(None) - assert result2 is None - assert len(hook_calls) == 1 + assert len(hook_calls) == 2 # --------------------------------------------------------------------------- @@ -193,8 +195,8 @@ async def test_post_tool_noop_when_hooks_none() -> None: @pytest.mark.unit -async def test_pre_turn_fires_and_sets_guard_key() -> None: - """Given a pre_turn hook, firing it adds 'pre_turn' to hooks_fired.""" +async def test_pre_turn_fires_and_records_call() -> None: + """Given a pre_turn hook, firing it records the call.""" _reset_calls() hooks = AgentHooks(pre_turn=[_make_recording_hook("pre_turn")]) run_ctx = AgentRunContext(session_id="test-session") @@ -202,8 +204,8 @@ async def test_pre_turn_fires_and_sets_guard_key() -> None: await host._fire_pre_turn_hooks() - assert "pre_turn" in run_ctx.hooks_fired assert len(hook_calls) == 1 + assert hook_calls[0][0] == "pre_turn" # --------------------------------------------------------------------------- @@ -238,12 +240,12 @@ async def _raise_and_fire() -> None: # --------------------------------------------------------------------------- -# Test: tool_call_id-scoped guard keys prevent cross-call double-firing +# Test: different tool_call_id allows both hooks to fire (no dedup guard) # --------------------------------------------------------------------------- @pytest.mark.unit -async def test_tool_guard_key_uses_tool_call_id() -> None: +async def test_tool_hooks_fire_for_different_call_ids() -> None: """Given same tool_name but different tool_call_id, both pre_tool hooks fire.""" _reset_calls() hooks = AgentHooks(pre_tool_use=[_make_recording_hook("pre_tool_use")]) @@ -254,13 +256,11 @@ async def test_tool_guard_key_uses_tool_call_id() -> None: await host._fire_pre_tool_hooks("same_tool", {"x": 2}, "call-B") assert len(hook_calls) == 2 - assert "pre_tool_use:call-A" in run_ctx.hooks_fired - assert "pre_tool_use:call-B" in run_ctx.hooks_fired @pytest.mark.unit -async def test_tool_guard_key_dedup_same_call_id() -> None: - """Given same tool_call_id twice, second pre_tool hook is skipped.""" +async def test_tool_hooks_fire_on_same_call_id() -> None: + """Given same tool_call_id twice, both pre_tool hooks fire (no dedup guard).""" _reset_calls() hooks = AgentHooks(pre_tool_use=[_make_recording_hook("pre_tool_use")]) run_ctx = AgentRunContext(session_id="test-session") @@ -269,12 +269,12 @@ async def test_tool_guard_key_dedup_same_call_id() -> None: await host._fire_pre_tool_hooks("tool", {}, "call-X") result2 = await host._fire_pre_tool_hooks("tool", {}, "call-X") - assert result2 is None - assert len(hook_calls) == 1 + assert result2 is not None + assert len(hook_calls) == 2 @pytest.mark.unit -async def test_post_tool_guard_key_uses_tool_call_id() -> None: +async def test_post_tool_hooks_fire_for_different_call_ids() -> None: """Given same tool_name but different tool_call_id, both post_tool hooks fire.""" _reset_calls() hooks = AgentHooks(post_tool_use=[_make_recording_hook("post_tool_use")]) @@ -285,8 +285,6 @@ async def test_post_tool_guard_key_uses_tool_call_id() -> None: await host._fire_post_tool_hooks("tool", {}, "out2", 2.0, "call-B") assert len(hook_calls) == 2 - assert "post_tool_use:call-A" in run_ctx.hooks_fired - assert "post_tool_use:call-B" in run_ctx.hooks_fired # --------------------------------------------------------------------------- @@ -305,28 +303,3 @@ async def test_pre_turn_deny_returns_deny_result() -> None: assert result is not None assert result.get("decision") == "deny" - - -# --------------------------------------------------------------------------- -# Test: hooks_fired cleared between turns -# --------------------------------------------------------------------------- - - -@pytest.mark.unit -async def test_hooks_fired_cleared_between_turns() -> None: - """Given hooks_fired populated in turn 1, clearing it allows turn 2 to fire.""" - _reset_calls() - hooks = AgentHooks(pre_turn=[_make_recording_hook("pre_turn")]) - run_ctx = AgentRunContext(session_id="test-session") - host = _make_host(hooks=hooks, run_ctx=run_ctx) - - # Turn 1 - await host._fire_pre_turn_hooks() - assert len(hook_calls) == 1 - - # Simulate RunHandle.start() clearing hooks_fired - run_ctx.hooks_fired.clear() - - # Turn 2 (same run_ctx — new turn) - await host._fire_pre_turn_hooks() - assert len(hook_calls) == 2 diff --git a/tests/hooks/test_hook_smoke_matrix.py b/tests/hooks/test_hook_smoke_matrix.py index bd6e8fcc1..9750c8c3f 100644 --- a/tests/hooks/test_hook_smoke_matrix.py +++ b/tests/hooks/test_hook_smoke_matrix.py @@ -219,7 +219,6 @@ async def _run_acp_standalone(hook_type: HookType) -> None: acp_client=client, # type: ignore[arg-type] prompts=["do something"], run_ctx=run_ctx, - message_history=[], session_id="smoke-acp-session", agent_name="smoke-acp-agent", hooks=hooks, diff --git a/tests/host/test_context.py b/tests/host/test_context.py index aee5a395c..d7a1085f1 100644 --- a/tests/host/test_context.py +++ b/tests/host/test_context.py @@ -1,7 +1,7 @@ """Unit tests for HostContext frozen dataclass. Covers immutability, construction with required fields, -default factory values, and pool back-reference default. +and default factory values. """ from __future__ import annotations @@ -73,12 +73,6 @@ def test_default_factory_fields_have_correct_defaults(): assert ctx.tenant_id is None -def test_pool_back_reference_defaults_to_none(): - """Given no pool argument, when constructing HostContext, then pool is None.""" - ctx = _make_context() - assert ctx.pool is None - - def test_session_pool_defaults_to_none(): """Given session_pool=None, when constructing HostContext, then session_pool is None.""" ctx = _make_context(session_pool=None) diff --git a/tests/host/test_factory.py b/tests/host/test_factory.py index 44ae10957..71753d943 100644 --- a/tests/host/test_factory.py +++ b/tests/host/test_factory.py @@ -26,7 +26,6 @@ def _make_host_context( ) -> Any: """Build a mock HostContext with common defaults.""" ctx = MagicMock() - ctx.pool = pool if pool is not None else MagicMock() ctx.config_file_path = config_file_path ctx.skills_tools_provider = skills_tools_provider ctx.mcp = mcp if mcp is not None else MagicMock() diff --git a/tests/integration/__snapshots__/test_acp_streaming.ambr b/tests/integration/__snapshots__/test_acp_streaming.ambr new file mode 100644 index 000000000..aa9575aa0 --- /dev/null +++ b/tests/integration/__snapshots__/test_acp_streaming.ambr @@ -0,0 +1,91 @@ +# serializer version: 1 +# name: test_acp_streaming_event_sequence_with_tool_metadata[asyncio] + list([ + dict({ + 'agent_name': 'test-acp-agent', + 'event_kind': 'run_started', + 'parent_session_id': None, + }), + dict({ + 'content': 'Hello', + 'delta_type': 'TextPartDelta', + 'event_kind': 'part_delta', + 'index': 0, + }), + dict({ + 'event_kind': 'tool_call_start', + 'kind': 'other', + 'raw_input': dict({ + 'path': '/test/file.py', + }), + 'title': 'Read file', + 'tool_call_id': 'tc-read-001', + 'tool_name': 'Read file', + }), + dict({ + 'agent_name': 'test-acp-agent', + 'event_kind': 'tool_call_complete', + 'metadata': dict({ + 'diff': dict({ + 'new': 'b', + 'old': 'a', + 'path': '/test/file.py', + }), + }), + 'tool_call_id': 'tc-read-001', + 'tool_input': dict({ + }), + 'tool_name': 'Read file', + 'tool_result': 'file contents here', + }), + dict({ + 'content': ' world', + 'delta_type': 'TextPartDelta', + 'event_kind': 'part_delta', + 'index': 0, + }), + dict({ + 'cancelled': False, + 'content': 'Hello world', + 'event_kind': 'stream_complete', + 'finish_reason': 'stop', + 'name': 'test-acp-agent', + 'role': 'assistant', + }), + ]) +# --- +# name: test_acp_streaming_text_only_sequence[asyncio] + list([ + dict({ + 'agent_name': 'test-acp-agent', + 'event_kind': 'run_started', + 'parent_session_id': None, + }), + dict({ + 'content': 'Hello', + 'delta_type': 'TextPartDelta', + 'event_kind': 'part_delta', + 'index': 0, + }), + dict({ + 'content': ' ', + 'delta_type': 'TextPartDelta', + 'event_kind': 'part_delta', + 'index': 0, + }), + dict({ + 'content': 'world', + 'delta_type': 'TextPartDelta', + 'event_kind': 'part_delta', + 'index': 0, + }), + dict({ + 'cancelled': False, + 'content': 'Hello world', + 'event_kind': 'stream_complete', + 'finish_reason': 'stop', + 'name': 'test-acp-agent', + 'role': 'assistant', + }), + ]) +# --- diff --git a/tests/integration/test_acp_streaming.py b/tests/integration/test_acp_streaming.py new file mode 100644 index 000000000..b1812f3e7 --- /dev/null +++ b/tests/integration/test_acp_streaming.py @@ -0,0 +1,486 @@ +"""ACP streaming snapshot baseline (V10). + +Captures the event sequence from ACP standalone streaming via +``ACPAgent._stream_events()``, including ``ToolCallCompleteEvent`` +enriched with ``ToolResultMetadataEvent`` metadata. + +This is a regression baseline for the pre-M4 protocol cleanup. The test +uses **syrupy** for snapshot comparison and is marked ``acp_snapshot`` +so it is excluded from the default test run (run explicitly with +``-m acp_snapshot`` or ``--snapshot-update``). + +Key event sequence captured: + 1. ``RunStartedEvent`` + 2. ``PartDeltaEvent`` (text chunk "Hello") + 3. ``ToolCallStartEvent`` (tool call "Read file") + 4. ``ToolCallCompleteEvent`` (enriched with metadata from + ``ToolResultMetadataEvent``) + 5. ``PartDeltaEvent`` (text chunk " world") + 6. ``StreamCompleteEvent`` +""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock, patch + +import pytest + +from acp import InitializeRequest +from acp.schema import ( + AgentMessageChunk, + PromptResponse, + TextContentBlock, + ToolCallProgress, + ToolCallStart, +) +from agentpool.agents.acp_agent import ACPAgent +from agentpool.agents.acp_agent.session_state import ACPSessionState +from agentpool.agents.context import AgentRunContext +from agentpool.agents.events import ( + PartDeltaEvent, + RunStartedEvent, + StreamCompleteEvent, + ToolCallCompleteEvent, + ToolCallStartEvent, + ToolResultMetadataEvent, +) +from agentpool.messaging import ChatMessage + + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from syrupy import SnapshotAssertion + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _serialize_event(event: Any) -> dict[str, Any]: # noqa: PLR0911 + """Serialize a RichAgentStreamEvent to a dict suitable for snapshot. + + Strips non-deterministic fields (uuids, timestamps) and converts + nested dataclass/Pydantic objects to plain dicts. + """ + if isinstance(event, RunStartedEvent): + return { + "event_kind": event.event_kind, + "agent_name": event.agent_name, + "parent_session_id": event.parent_session_id, + } + if isinstance(event, PartDeltaEvent): + delta = event.delta + # TextPartDelta or ThinkingPartDelta + delta_type = type(delta).__name__ + content = "" + if hasattr(delta, "content_delta"): + content = delta.content_delta + return { + "event_kind": "part_delta", + "delta_type": delta_type, + "content": content, + "index": event.index, + } + if isinstance(event, ToolCallStartEvent): + return { + "event_kind": event.event_kind, + "tool_call_id": event.tool_call_id, + "tool_name": event.tool_name, + "title": event.title, + "kind": event.kind, + "raw_input": event.raw_input, + } + if isinstance(event, ToolCallCompleteEvent): + return { + "event_kind": event.event_kind, + "tool_call_id": event.tool_call_id, + "tool_name": event.tool_name, + "tool_input": event.tool_input, + "tool_result": event.tool_result, + "agent_name": event.agent_name, + "metadata": event.metadata, + } + if isinstance(event, ToolResultMetadataEvent): + return { + "event_kind": event.event_kind, + "tool_call_id": event.tool_call_id, + "metadata": event.metadata, + } + if isinstance(event, StreamCompleteEvent): + msg = event.message + return { + "event_kind": event.event_kind, + "cancelled": event.cancelled, + "content": msg.content, + "role": msg.role, + "name": msg.name, + "finish_reason": msg.finish_reason, + } + # Fallback: type name + str + return {"event_kind": getattr(event, "event_kind", "unknown"), "type": type(event).__name__} + + +def _text_chunk(text: str) -> AgentMessageChunk: + """Create an ACP AgentMessageChunk with text content.""" + return AgentMessageChunk(content=TextContentBlock(text=text)) + + +def _tool_call_start( + tool_call_id: str, + title: str, + raw_input: dict[str, Any] | None = None, +) -> ToolCallStart: + """Create an ACP ToolCallStart update.""" + return ToolCallStart( + tool_call_id=tool_call_id, + title=title, + raw_input=raw_input or {}, + ) + + +def _tool_call_progress_completed( + tool_call_id: str, + title: str, + raw_output: str, +) -> ToolCallProgress: + """Create an ACP ToolCallProgress with status=completed.""" + return ToolCallProgress( + tool_call_id=tool_call_id, + status="completed", + title=title, + raw_output=raw_output, + ) + + +# --------------------------------------------------------------------------- +# Mock objects +# --------------------------------------------------------------------------- + + +class _MockUpdateEvent: + """Minimal mock of TimeoutableEvent for polling. + + Yields control to the event loop so that the prompt_task can execute + concurrently with the polling loop. + """ + + def __init__(self) -> None: + self._set = False + + def set(self) -> None: + self._set = True + + def clear(self) -> None: + self._set = False + + async def wait_with_timeout(self, timeout: float | None = None) -> bool: + # Yield control to the event loop so prompt_task can run. + # A tiny sleep prevents busy-looping while allowing concurrent task execution. + await asyncio.sleep(0.001) + return True + + +class _MockToolBridge: + """Mock ToolManagerBridge that does nothing in set_run_context.""" + + @asynccontextmanager + async def set_run_context( + self, + context: Any, + prompt: Any = None, + ) -> AsyncIterator[_MockToolBridge]: + yield self + + +class _MockACPClientHandler: + """Mock ACPClientHandler with just the _update_event needed.""" + + def __init__(self) -> None: + self._update_event = _MockUpdateEvent() + self._input_provider = None + + +class _MockAPI: + """Mock ACPAgentAPI that returns a PromptResponse after a short delay. + + Implements stream_events() and get_messages() with the same polling + logic as ACPAgentAPI so ACPTurn.execute() can use it as an + ACPClientProtocol. + """ + + def __init__( + self, + delay: float = 0.01, + state: Any | None = None, + update_event: Any | None = None, + ) -> None: + self._delay = delay + self._state = state + self._update_event = update_event + self._consumed_updates: list[Any] = [] + + async def prompt(self, session_id: str, content: list[Any]) -> PromptResponse: + await asyncio.sleep(self._delay) + return PromptResponse(stop_reason="end_turn") + + async def fork_session(self, session_id: str, cwd: str) -> Any: + raise NotImplementedError + + async def stream_events(self, response: Any) -> AsyncIterator[Any]: + """Poll state queue for updates, same as ACPAgentAPI.stream_events().""" + self._consumed_updates.clear() + if self._state is None or self._update_event is None: + return + while True: + try: + await self._update_event.wait_with_timeout(0.05) + self._update_event.clear() + except TimeoutError: + pass + drained_any = False + while (update := self._state.pop_update()) is not None: + self._consumed_updates.append(update) + yield update + drained_any = True + if not drained_any: + break + + async def get_messages(self, session_id: str) -> list[Any]: + return list(self._consumed_updates) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +class _PrePopulatedSessionState(ACPSessionState): + """Session state that re-populates updates after clear(). + + ``_stream_events`` calls ``self._state.clear()`` at the start of each + turn, which wipes the updates deque. This subclass overrides + ``clear()`` to re-add the pre-populated updates so they survive the + clear and are available for the polling loop. + """ + + def __init__(self, updates: list[Any], session_id: str = "test-session-id") -> None: + super().__init__(session_id=session_id) + self._initial_updates = list(updates) + + def clear(self) -> None: + """Clear and re-populate with the pre-populated updates.""" + self.updates.clear() + for update in self._initial_updates: + self.updates.append(update) + + +def _make_acp_agent_with_mocks( + updates: list[Any], +) -> ACPAgent[Any]: + """Create an ACPAgent with mocked internals for _stream_events testing. + + Pre-populates the session state with the given updates (ACP SessionUpdate + objects and optionally ToolResultMetadataEvent for metadata enrichment). + """ + init_request = MagicMock(spec=InitializeRequest) + agent = ACPAgent(command="test-cmd", init_request=init_request, name="test-acp-agent") + + # Set up mocked state with pre-populated updates that survive clear() + state = _PrePopulatedSessionState(updates=updates) + agent._state = state + + # Set up mocked client handler + agent._client_handler = _MockACPClientHandler() + + # Set up mocked API with state/event for stream_events() and get_messages() + agent._api = _MockAPI( + state=state, + update_event=agent._client_handler._update_event, + ) + + # Set up mocked session ID + agent._sdk_session_id = "test-session-id" + + # Set up mocked tool bridge + agent._tool_bridge = _MockToolBridge() + + return agent + + +def _make_run_ctx() -> AgentRunContext: + """Create a minimal AgentRunContext for testing.""" + return AgentRunContext(session_id="test-session-id") + + +# --------------------------------------------------------------------------- +# Patched acp_to_native_event +# --------------------------------------------------------------------------- + + +_original_acp_to_native_event: Any = None + + +def _patched_acp_to_native_event(update: Any) -> Any: + """Patched converter that passes through ToolResultMetadataEvent. + + The real ``acp_to_native_event`` only handles ACP ``SessionUpdate`` + types. This wrapper allows ``ToolResultMetadataEvent`` instances + (which are native events, not ACP updates) to pass through the + polling loop so that ``_stream_events`` can exercise the metadata + enrichment code path. + """ + if isinstance(update, ToolResultMetadataEvent): + return update + return _original_acp_to_native_event(update) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +@pytest.mark.acp_snapshot +async def test_acp_streaming_event_sequence_with_tool_metadata( + snapshot: SnapshotAssertion, +) -> None: + """Capture full event sequence from ACP standalone streaming. + + The sequence includes: + - RunStartedEvent + - PartDeltaEvent (text "Hello") + - ToolCallStartEvent (read file) + - ToolResultMetadataEvent (sidechannel, consumed for enrichment) + - ToolCallCompleteEvent (enriched with metadata) + - PartDeltaEvent (text " world") + - StreamCompleteEvent + """ + tool_call_id = "tc-read-001" + tool_metadata = {"diff": {"path": "/test/file.py", "old": "a", "new": "b"}} + + updates: list[Any] = [ + _text_chunk("Hello"), + _tool_call_start( + tool_call_id=tool_call_id, + title="Read file", + raw_input={"path": "/test/file.py"}, + ), + # ToolResultMetadataEvent is a native event, not an ACP update. + # It passes through the patched converter and gets consumed by + # _stream_events for metadata enrichment. + ToolResultMetadataEvent( + tool_call_id=tool_call_id, + metadata=tool_metadata, + ), + _tool_call_progress_completed( + tool_call_id=tool_call_id, + title="Read file", + raw_output="file contents here", + ), + _text_chunk(" world"), + ] + + agent = _make_acp_agent_with_mocks(updates) + run_ctx = _make_run_ctx() + + user_msg = ChatMessage[str]( + content="test prompt", + role="user", + message_id="msg-001", + session_id="test-session-id", + ) + + # Patch acp_to_native_event to allow ToolResultMetadataEvent passthrough + from agentpool.agents.acp_agent import acp_converters as converters_mod + + global _original_acp_to_native_event # noqa: PLW0603 + _original_acp_to_native_event = converters_mod.acp_to_native_event + + with patch.object(converters_mod, "acp_to_native_event", _patched_acp_to_native_event): + events: list[Any] = [ + event + async for event in agent._stream_events( + run_ctx, + ["test prompt"], + user_msg=user_msg, + message_history=[], + effective_parent_id=None, + session_id="test-session-id", + ) + ] + + # Serialize events for snapshot comparison + serialized = [_serialize_event(e) for e in events] + + # Verify the event sequence matches the snapshot + assert serialized == snapshot + + # Verify key properties of the event sequence + # 1. First event is RunStartedEvent + assert isinstance(events[0], RunStartedEvent) + assert events[0].agent_name == "test-acp-agent" + + # 2. ToolResultMetadataEvent is NOT in the output (consumed for enrichment) + metadata_events = [e for e in events if isinstance(e, ToolResultMetadataEvent)] + assert len(metadata_events) == 0 + + # 3. ToolCallCompleteEvent is enriched with metadata + complete_events = [e for e in events if isinstance(e, ToolCallCompleteEvent)] + assert len(complete_events) == 1 + assert complete_events[0].metadata == tool_metadata + assert complete_events[0].agent_name == "test-acp-agent" + + # 4. Last event is StreamCompleteEvent + assert isinstance(events[-1], StreamCompleteEvent) + assert events[-1].message.content == "Hello world" + assert events[-1].message.finish_reason == "stop" + + +@pytest.mark.anyio +@pytest.mark.acp_snapshot +async def test_acp_streaming_text_only_sequence( + snapshot: SnapshotAssertion, +) -> None: + """Capture event sequence for a simple text-only ACP stream. + + No tool calls, just text chunks followed by StreamCompleteEvent. + """ + updates: list[Any] = [ + _text_chunk("Hello"), + _text_chunk(" "), + _text_chunk("world"), + ] + + agent = _make_acp_agent_with_mocks(updates) + run_ctx = _make_run_ctx() + + user_msg = ChatMessage[str]( + content="test prompt", + role="user", + message_id="msg-002", + session_id="test-session-id", + ) + + events: list[Any] = [] + async for event in agent._stream_events( + run_ctx, + ["test prompt"], + user_msg=user_msg, + message_history=[], + effective_parent_id=None, + session_id="test-session-id", + ): + events.append(event) # noqa: PERF401 + + serialized = [_serialize_event(e) for e in events] + assert serialized == snapshot + + # Verify basic structure + assert isinstance(events[0], RunStartedEvent) + assert isinstance(events[-1], StreamCompleteEvent) + assert events[-1].message.content == "Hello world" diff --git a/tests/lifecycle/test_crash_recovery.py b/tests/lifecycle/test_crash_recovery.py index c1cc2180f..8a0627250 100644 --- a/tests/lifecycle/test_crash_recovery.py +++ b/tests/lifecycle/test_crash_recovery.py @@ -604,6 +604,7 @@ class _TestableHookAwareTurn(HookAwareTurn): """Concrete HookAwareTurn for testing _log_tool_execution.""" def __init__(self, run_ctx: AgentRunContext) -> None: + super().__init__() self._hooks = None self._run_ctx = run_ctx diff --git a/tests/lifecycle/test_run_loop.py b/tests/lifecycle/test_run_loop.py index 2ae118af8..c20b5ffd9 100644 --- a/tests/lifecycle/test_run_loop.py +++ b/tests/lifecycle/test_run_loop.py @@ -21,6 +21,7 @@ ) from agentpool.lifecycle import ( DirectChannel, + Feedback, ImmediateTrigger, InProcessTransport, MemoryJournal, @@ -29,7 +30,7 @@ ) from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventBus -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.turn import Turn @@ -493,7 +494,7 @@ async def test_crash_recovery_normal_resume() -> None: await consumer # Should have completed normally. - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE # --------------------------------------------------------------------------- @@ -693,7 +694,7 @@ async def _spy_publish(event: Any) -> None: async def test_steer_when_idle_direct_channel() -> None: """steer() when IDLE with DirectChannel appends to _message_queue and sets _idle_event.""" handle = _make_run_handle() - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE result = handle.steer("steer message") @@ -706,7 +707,7 @@ async def test_steer_when_idle_direct_channel() -> None: async def test_steer_when_running_direct_channel_with_agent_run() -> None: """steer() when RUNNING with DirectChannel injects via active_agent_run.enqueue().""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING mock_agent_run = MagicMock() mock_agent_run.enqueue = MagicMock() @@ -723,7 +724,7 @@ async def test_steer_when_running_direct_channel_with_agent_run() -> None: async def test_steer_when_running_direct_channel_without_agent_run() -> None: """steer() when RUNNING without active_agent_run queues to queued_steer_messages.""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING handle.active_agent_run = None result = handle.steer("steer message") @@ -736,7 +737,7 @@ async def test_steer_when_running_direct_channel_without_agent_run() -> None: async def test_followup_during_active_turn() -> None: """followup() during RUNNING appends to _message_queue without interrupting.""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING handle._idle_event.clear() result = handle.followup("followup message") @@ -751,7 +752,7 @@ async def test_followup_during_active_turn() -> None: async def test_followup_when_idle() -> None: """followup() when IDLE appends to _message_queue and sets _idle_event.""" handle = _make_run_handle() - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE result = handle.followup("followup message") @@ -764,7 +765,7 @@ async def test_followup_when_idle() -> None: async def test_close_while_idle_transitions_to_done() -> None: """close() while idle schedules transition to RunState.DONE.""" handle = _make_run_handle() - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE handle._run_state = RunState.IDLE handle.close() @@ -780,7 +781,7 @@ async def test_close_while_idle_transitions_to_done() -> None: async def test_close_twice_is_noop() -> None: """close() called twice: second call is a no-op.""" handle = _make_run_handle() - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE handle.close() first_closing = handle._closing @@ -891,7 +892,7 @@ async def test_close_with_pending_messages_processed() -> None: # Two turns should have been created: initial + followup. assert agent.create_turn.call_count == 2 - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE @pytest.mark.unit @@ -903,7 +904,7 @@ async def test_steer_with_protocol_channel_routes_via_deliver_feedback() -> None event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING result = handle.steer("protocol steer") @@ -924,7 +925,7 @@ async def test_followup_with_protocol_channel_routes_via_deliver_feedback() -> N event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING result = handle.followup("protocol followup") @@ -944,7 +945,7 @@ async def test_steer_protocol_channel_when_idle_sets_idle_event() -> None: event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE handle._idle_event.clear() result = handle.steer("idle steer") @@ -1011,7 +1012,7 @@ async def test_multi_turn_with_protocol_trigger() -> None: await consumer assert agent.create_turn.call_count == 2 - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE @pytest.mark.unit @@ -1036,16 +1037,14 @@ async def test_idempotency_skip_when_turn_result_exists() -> None: @pytest.mark.unit async def test_steer_direct_channel_does_not_use_deliver_feedback() -> None: - """steer() with DirectChannel does NOT call deliver_feedback (it doesn't exist).""" + """steer() with DirectChannel returns False from deliver_feedback (falls through to queue).""" handle = _make_run_handle() - # DirectChannel does not have deliver_feedback. - try: - _ = handle._comm_channel.deliver_feedback # type: ignore[union-attr] - raise AssertionError("DirectChannel should not have deliver_feedback") - except AttributeError: - pass # Expected - - handle._status = RunStatus.idle + # DirectChannel.deliver_feedback returns False (no-op). + assert handle._comm_channel is not None + feedback_result = handle._comm_channel.deliver_feedback(Feedback(content="test", is_steer=True)) + assert feedback_result is False + + handle._run_state = RunState.IDLE result = handle.steer("direct steer") assert result is True @@ -1056,7 +1055,7 @@ async def test_steer_direct_channel_does_not_use_deliver_feedback() -> None: async def test_followup_direct_channel_does_not_use_deliver_feedback() -> None: """followup() with DirectChannel does NOT call deliver_feedback.""" handle = _make_run_handle() - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE result = handle.followup("direct followup") @@ -1073,7 +1072,7 @@ async def test_steer_protocol_channel_does_not_touch_message_queue() -> None: event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING result = handle.steer("protocol steer") @@ -1093,7 +1092,7 @@ async def test_followup_protocol_channel_does_not_touch_message_queue() -> None: event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING result = handle.followup("protocol followup") @@ -1147,7 +1146,7 @@ async def test_close_while_running_lets_turn_finish() -> None: await asyncio.sleep(0.05) await consumer - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE assert handle._closed is True @@ -1210,7 +1209,7 @@ async def test_followup_protocol_channel_wakes_idle_loop() -> None: event_bus = EventBus() channel = ProtocolChannel(journal, event_bus, "test-session") handle = _make_run_handle(_comm_channel=channel) - handle._status = RunStatus.idle + handle._run_state = RunState.IDLE handle._idle_event.clear() result = handle.followup("idle followup") diff --git a/tests/lifecycle/test_session_migration.py b/tests/lifecycle/test_session_migration.py index 1e5b60cac..35cdafa6d 100644 --- a/tests/lifecycle/test_session_migration.py +++ b/tests/lifecycle/test_session_migration.py @@ -318,7 +318,7 @@ async def test_protocol_channel_publishes_to_event_bus() -> None: async def test_no_double_publish_with_protocol_channel() -> None: """When ProtocolChannel is the comm_channel, events are not double-published to EventBus. - The _channel_publishes_to_event_bus property returns True for + The publishes_to_event_bus property returns True for ProtocolChannel, so start() skips the direct event_bus.publish() call and lets ProtocolChannel.publish() handle EventBus delivery. """ @@ -346,7 +346,7 @@ async def test_no_double_publish_with_protocol_channel() -> None: ) # The property should return True. - assert run_handle._channel_publishes_to_event_bus is True + assert run_handle._comm_channel.publishes_to_event_bus is True # Cleanup run_handle.close() @@ -354,7 +354,7 @@ async def test_no_double_publish_with_protocol_channel() -> None: @pytest.mark.unit async def test_double_publish_with_direct_channel() -> None: - """With DirectChannel, _channel_publishes_to_event_bus is False. + """With DirectChannel, publishes_to_event_bus is False. DirectChannel does not publish to EventBus, so start() must call event_bus.publish() directly. @@ -381,7 +381,7 @@ async def test_double_publish_with_direct_channel() -> None: ) # The property should return False. - assert run_handle._channel_publishes_to_event_bus is False + assert run_handle._comm_channel.publishes_to_event_bus is False # Cleanup run_handle.close() diff --git a/tests/lifecycle/test_types.py b/tests/lifecycle/test_types.py index 2286fe82d..3d89c1c8d 100644 --- a/tests/lifecycle/test_types.py +++ b/tests/lifecycle/test_types.py @@ -288,7 +288,11 @@ def clear(self) -> None: ... class _DummyCommChannel: """Minimal CommChannel implementation for isinstance testing.""" - _replaying: bool = False + def set_replaying(self, flag: bool) -> None: ... + + @property + def publishes_to_event_bus(self) -> bool: + return False def attach(self, run_loop: Any) -> None: ... @@ -299,6 +303,9 @@ async def publish(self, event: Any) -> None: ... def recv(self) -> Feedback | None: return None + def deliver_feedback(self, feedback: Feedback) -> bool: + return False + def close(self) -> None: ... diff --git a/tests/mcp_server/test_e2e_session_controller.py b/tests/mcp_server/test_e2e_session_controller.py index f6afe8105..8ab5f6e70 100644 --- a/tests/mcp_server/test_e2e_session_controller.py +++ b/tests/mcp_server/test_e2e_session_controller.py @@ -5,7 +5,7 @@ per-session MCP state correctly. Tests: - G1 - Full create-session chain populates MCP _SessionContext with snapshot + G1 - Full create-session chain populates MCP McpSessionContext with snapshot G5 - close_session cleans agent.mcp session context and exits agent context G6 - resume_session creates real ACPSession with _acp_mcp_manager wired G13 - MCPManager.cleanup() cleans all session contexts @@ -64,11 +64,11 @@ def _make_mock_client() -> MagicMock: @pytest.mark.integration @pytest.mark.asyncio async def test_full_create_session_chain_populates_mcp_session_context() -> None: - """SessionController.get_or_create_session_agent() populates MCPManager _SessionContext. + """SessionController.get_or_create_session_agent() populates MCPManager McpSessionContext. Verifies that the real SessionController, when given a NativeAgentConfig, calls agent.mcp.get_or_create_session() and agent.mcp.update_session_snapshot() - so that the agent's MCPManager has a _SessionContext with a non-None snapshot. + so that the agent's MCPManager has a McpSessionContext with a non-None snapshot. """ from agentpool.agents.native_agent import Agent from agentpool.models.agents import NativeAgentConfig @@ -119,9 +119,9 @@ async def _mock_create_session_agent(**kwargs: object) -> Agent: session_id, agent_name="test_agent" ) - # Assert: agent's MCPManager has a _SessionContext with non-None snapshot + # Assert: agent's MCPManager has a McpSessionContext with non-None snapshot ctx = result_agent.mcp.get_session_context(session_id) - assert ctx is not None, "get_or_create_session_agent() must create a _SessionContext" + assert ctx is not None, "get_or_create_session_agent() must create a McpSessionContext" assert ctx.snapshot is not None, ( "get_or_create_session_agent() must call update_session_snapshot()" ) diff --git a/tests/mcp_server/test_review_fixes.py b/tests/mcp_server/test_review_fixes.py index 007a3ea9a..c983239fa 100644 --- a/tests/mcp_server/test_review_fixes.py +++ b/tests/mcp_server/test_review_fixes.py @@ -42,7 +42,7 @@ async def test_get_capabilities_does_not_recreate_cleaned_session() -> None: """get_capabilities(session_id=...) must not recreate a cleaned-up context. Bug: ``get_capabilities(session_id=...)`` calls ``get_or_create_session()`` - which recreates an empty ``_SessionContext`` if cleanup already popped it. + which recreates an empty ``McpSessionContext`` if cleanup already popped it. This is a memory leak — the dead context lingers forever with an empty snapshot, and the ``KeyError`` fallback code is dead. diff --git a/tests/mcp_server/test_review_fixes_r3.py b/tests/mcp_server/test_review_fixes_r3.py index dab113a53..8703893d5 100644 --- a/tests/mcp_server/test_review_fixes_r3.py +++ b/tests/mcp_server/test_review_fixes_r3.py @@ -3,7 +3,7 @@ Fix #3/#4 (REAL BUG): get_or_create_session_agent() calls parent_agent.mcp.get_or_create_session(parent_session_id) to read the parent's snapshot. If the parent session was already cleaned up, -get_or_create_session() CREATES a new phantom _SessionContext that +get_or_create_session() CREATES a new phantom McpSessionContext that will NEVER be cleaned up — memory leak. Fix #2 (IMPROVEMENT): on_disconnect callback in _handle_websocket_client() @@ -40,7 +40,7 @@ async def test_get_or_create_session_agent_does_not_recreate_cleaned_parent_sess Bug: The method calls parent_agent.mcp.get_or_create_session(parent_id) to read the parent's snapshot. If the parent was already cleaned up, - get_or_create_session() creates a phantom _SessionContext that leaks. + get_or_create_session() creates a phantom McpSessionContext that leaks. Steps: 1. Create a real Agent with a real MCPManager. @@ -135,7 +135,7 @@ async def test_get_or_create_session_agent_reads_parent_snapshot_without_leaking """get_or_create_session_agent() reads parent snapshot without recreating context. When the parent session is still active (not cleaned up), the method - should read the existing _SessionContext, not create a new one. + should read the existing McpSessionContext, not create a new one. This test passes both before and after the fix — it guards against regressions where the parent context is accidentally recreated. diff --git a/tests/mcp_server/test_session_close_integration.py b/tests/mcp_server/test_session_close_integration.py index df4964475..25a7676c7 100644 --- a/tests/mcp_server/test_session_close_integration.py +++ b/tests/mcp_server/test_session_close_integration.py @@ -58,7 +58,7 @@ async def test_integration_create_run_close() -> None: async def test_integration_close_recreate_fresh() -> None: """After closing a session and recreating with the same ID, the new context is fresh. - Verifies that ``cleanup_session`` fully removes the old ``_SessionContext`` + Verifies that ``cleanup_session`` fully removes the old ``McpSessionContext`` and a subsequent ``get_or_create_session`` with the same session ID creates a brand-new context with fresh resource objects (different toolset_cache, different connection_pool, and a ``None`` snapshot). @@ -89,7 +89,7 @@ async def test_integration_close_recreate_fresh() -> None: new_ctx = manager.get_or_create_session(session_id) assert session_id in manager._session_contexts - # 6. The new _SessionContext is a different object + # 6. The new McpSessionContext is a different object assert new_ctx is not original_ctx # 7. Fresh toolset_cache — different object, empty diff --git a/tests/mcp_server/test_session_lifecycle.py b/tests/mcp_server/test_session_lifecycle.py index 7fbb09738..ee89f56b9 100644 --- a/tests/mcp_server/test_session_lifecycle.py +++ b/tests/mcp_server/test_session_lifecycle.py @@ -1,6 +1,6 @@ """Unit tests for MCPManager session context lifecycle. -Covers ``_SessionContext`` creation, snapshot storage, ACP transport +Covers ``McpSessionContext`` creation, snapshot storage, ACP transport registration, and cleanup (including idempotency and concurrency safety). """ @@ -13,7 +13,7 @@ import pytest from agentpool.mcp_server.config_snapshot import McpConfigSnapshot -from agentpool.mcp_server.manager import MCPManager, _SessionContext +from agentpool.mcp_server.manager import MCPManager, McpSessionContext # --------------------------------------------------------------------------- @@ -53,12 +53,12 @@ def manager() -> MCPManager: def test_get_or_create_session_creates_and_returns_same( manager: MCPManager, ) -> None: - """Two calls with the same session_id return the same _SessionContext.""" + """Two calls with the same session_id return the same McpSessionContext.""" ctx1 = manager.get_or_create_session("sess-1") ctx2 = manager.get_or_create_session("sess-1") assert ctx1 is ctx2 - assert isinstance(ctx1, _SessionContext) + assert isinstance(ctx1, McpSessionContext) assert manager.get_session_context("sess-1") is not None @@ -71,7 +71,7 @@ def test_get_or_create_session_creates_and_returns_same( def test_get_or_create_session_creates_fresh_for_different_ids( manager: MCPManager, ) -> None: - """Different session_ids produce distinct _SessionContext objects.""" + """Different session_ids produce distinct McpSessionContext objects.""" ctx_a = manager.get_or_create_session("sess-a") ctx_b = manager.get_or_create_session("sess-b") diff --git a/tests/mcp_server/test_stale_mcp_connection.py b/tests/mcp_server/test_stale_mcp_connection.py index 56fe12cc1..c334758bb 100644 --- a/tests/mcp_server/test_stale_mcp_connection.py +++ b/tests/mcp_server/test_stale_mcp_connection.py @@ -8,7 +8,7 @@ the dead transport from the previous WebSocket session). After T10-T12, ``get_capabilities()`` accepts ``session_id`` and routes -session-scoped configs through per-session ``_SessionContext`` objects +session-scoped configs through per-session ``McpSessionContext`` objects with their own ``toolset_cache``. ``cleanup_session()`` clears the per-session cache, ensuring the next session gets a fresh toolset. @@ -81,7 +81,7 @@ async def test_session_resume_returns_fresh_toolset() -> None: """get_capabilities() returns a fresh toolset after session cleanup+recreate. After T10-T12, session-scoped toolsets are cached on the per-session - ``_SessionContext.toolset_cache`` rather than the global + ``McpSessionContext.toolset_cache`` rather than the global ``MCPManager._toolset_cache``. When ``cleanup_session()`` is called between sessions, the per-session context (including its toolset cache) is removed, so the next session gets a brand-new toolset with diff --git a/tests/orchestrator/test_cancel_e2e.py b/tests/orchestrator/test_cancel_e2e.py index 1173158a8..ed014f360 100644 --- a/tests/orchestrator/test_cancel_e2e.py +++ b/tests/orchestrator/test_cancel_e2e.py @@ -20,9 +20,9 @@ StreamCompleteEvent, ToolCallStartEvent, ) +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventEnvelope, SessionPool -from agentpool.orchestrator.run import RunStatus from agentpool.orchestrator.turn import Turn @@ -274,8 +274,8 @@ async def test_cancel_then_new_prompt_full_flow( # If second_handle is None, the existing RunHandle was steered (1:1 model). # Verify the first handle is not stuck in a running state - assert first_handle._status in (RunStatus.idle, RunStatus.done), ( - f"First RunHandle should be idle or done, got: {first_handle._status}" + assert first_handle._run_state in (RunState.IDLE, RunState.DONE), ( + f"First RunHandle should be idle or done, got: {first_handle._run_state}" ) # Cleanup: close the RunHandle first so the start() loop exits and @@ -421,8 +421,8 @@ async def test_double_cancel(mock_pool: MagicMock) -> None: assert RunFailedEvent in pre_types, f"Expected RunFailedEvent, got: {pre_types}" # RunHandle should be idle or done (not running) - assert first_handle._status in (RunStatus.idle, RunStatus.done), ( - f"RunHandle should be idle/done after double cancel, got: {first_handle._status}" + assert first_handle._run_state in (RunState.IDLE, RunState.DONE), ( + f"RunHandle should be idle/done after double cancel, got: {first_handle._run_state}" ) # Send a new prompt — should not hang @@ -603,8 +603,8 @@ async def test_cancel_during_tool_execution(mock_pool: MagicMock) -> None: ) # RunHandle should be idle or done - assert first_handle._status in (RunStatus.idle, RunStatus.done), ( - f"RunHandle should be idle/done after cancel, got: {first_handle._status}" + assert first_handle._run_state in (RunState.IDLE, RunState.DONE), ( + f"RunHandle should be idle/done after cancel, got: {first_handle._run_state}" ) first_handle.close() @@ -765,8 +765,8 @@ async def test_runhandle_dies_in_idle_loop(mock_pool: MagicMock) -> None: assert crash_handle.complete_event.is_set(), ( "complete_event should be set by finally block after error" ) - assert crash_handle._status == RunStatus.done, ( - f"RunHandle should be done after error, got: {crash_handle._status}" + assert crash_handle._run_state == RunState.DONE, ( + f"RunHandle should be done after error, got: {crash_handle._run_state}" ) # Verify _cleanup_run cleared current_run_id diff --git a/tests/orchestrator/test_child_done_events.py b/tests/orchestrator/test_child_done_events.py index 560def575..c2fc10de3 100644 --- a/tests/orchestrator/test_child_done_events.py +++ b/tests/orchestrator/test_child_done_events.py @@ -27,9 +27,10 @@ from agentpool import Agent from agentpool.agents.context import AgentRunContext from agentpool.agents.events import StreamCompleteEvent +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventBus, SessionState -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.turn import Turn @@ -135,7 +136,7 @@ async def test_empty_child_done_events_no_wait() -> None: assert len(events) == 1 assert isinstance(events[0], StreamCompleteEvent) assert run_ctx.child_done_events == {} - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE @pytest.mark.unit @@ -336,7 +337,7 @@ def test_child_done_events_items_wrapped_with_list() -> None: """ import agentpool.orchestrator.run as run_module - source = inspect.getsource(run_module.RunHandle.start) + source = inspect.getsource(run_module.RunHandle._drain_events) # Check that items() is wrapped with list() assert "list(self.run_ctx.child_done_events.items())" in source, ( "child_done_events.items() must be wrapped with list() for concurrent safety" @@ -351,7 +352,7 @@ def test_child_done_events_values_wrapped_with_list() -> None: """ import agentpool.orchestrator.run as run_module - source = inspect.getsource(run_module.RunHandle.start) + source = inspect.getsource(run_module.RunHandle._drain_events) assert "list(self.run_ctx.child_done_events.values())" in source, ( "child_done_events.values() must be wrapped with list() for concurrent safety" ) diff --git a/tests/orchestrator/test_receive_request_acp.py b/tests/orchestrator/test_receive_request_acp.py index 0a1ba7876..bb64081fb 100644 --- a/tests/orchestrator/test_receive_request_acp.py +++ b/tests/orchestrator/test_receive_request_acp.py @@ -13,8 +13,9 @@ import pytest from agentpool.agents.acp_agent import ACPAgent +from agentpool.lifecycle import RunState from agentpool.orchestrator.core import EventBus, SessionController -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle pytestmark = pytest.mark.unit @@ -215,7 +216,7 @@ async def test_cancel_then_receive_request_starts_new_run( agent_type="acp", event_bus=event_bus, ) - existing_run._status = RunStatus.failed + existing_run._run_state = RunState.DONE controller._runs["failed-run-id"] = existing_run controller.get_session("sess-cancel").current_run_id = "failed-run-id" # type: ignore[union-attr] diff --git a/tests/orchestrator/test_run_handle.py b/tests/orchestrator/test_run_handle.py index 2aee7ac8d..ad729ddf5 100644 --- a/tests/orchestrator/test_run_handle.py +++ b/tests/orchestrator/test_run_handle.py @@ -28,9 +28,10 @@ RunStartedEvent, StreamCompleteEvent, ) +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventBus, SessionState -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle from agentpool.orchestrator.turn import Turn @@ -154,14 +155,14 @@ async def _consume() -> None: await asyncio.sleep(0.05) # After consuming the single turn, handle should be idle - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE # Close to unblock the idle wait handle.close() await asyncio.sleep(0.05) await consumer_task - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE assert len(events) == 1 assert isinstance(events[0], StreamCompleteEvent) assert handle._message_history == ["msg1"] @@ -192,7 +193,7 @@ async def _consume() -> None: await asyncio.sleep(0.05) # Handle should be idle after first turn - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert not handle._idle_event.is_set() # cleared when entering idle # Steer while idle @@ -207,7 +208,7 @@ async def _consume() -> None: await asyncio.sleep(0.05) await consumer_task - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE # Two turns should have executed assert agent.create_turn.call_count == 2 @@ -229,7 +230,7 @@ async def _consume() -> None: consumer_task = asyncio.create_task(_consume()) await asyncio.sleep(0.05) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE result = handle.followup("followup message") assert result is True @@ -261,7 +262,7 @@ async def _consume() -> None: consumer_task = asyncio.create_task(_consume()) await asyncio.sleep(0.05) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert not handle._closing handle.close() @@ -271,7 +272,7 @@ async def _consume() -> None: await asyncio.sleep(0.05) await consumer_task - assert handle._status == RunStatus.done + assert handle._run_state == RunState.DONE @pytest.mark.unit @@ -298,7 +299,7 @@ async def test_followup_returns_false_when_closing() -> None: async def test_steer_while_running_with_agent_run() -> None: """Given a running RunHandle with active_agent_run, steer() enqueues.""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING mock_agent_run = MagicMock() handle.active_agent_run = mock_agent_run @@ -311,7 +312,7 @@ async def test_steer_while_running_with_agent_run() -> None: async def test_steer_while_running_without_agent_run() -> None: """Given a running RunHandle without active_agent_run, steer() queues to run_ctx.""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING handle.active_agent_run = None result = handle.steer("queue me") @@ -363,7 +364,7 @@ async def _consume() -> None: async def test_followup_while_running_does_not_set_idle_event() -> None: """Given a running RunHandle, followup() queues but does not set idle event.""" handle = _make_run_handle() - handle._status = RunStatus.running + handle._run_state = RunState.RUNNING handle._idle_event.clear() result = handle.followup("queued") @@ -376,7 +377,7 @@ async def test_followup_while_running_does_not_set_idle_event() -> None: async def test_initial_status_is_idle() -> None: """Given a freshly created RunHandle, _status is idle and _idle_event is set.""" handle = RunHandle(run_id="r", session_id="s", agent_type="native") - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert handle._idle_event.is_set() assert handle._closing is False assert handle._message_queue == [] @@ -498,7 +499,7 @@ async def _consume() -> None: consumer_task = asyncio.create_task(_consume()) await asyncio.sleep(0.05) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE # Queue two followups assert handle.followup("first followup") is True @@ -545,7 +546,7 @@ async def test_close_is_idempotent() -> None: async def test_steer_returns_false_when_done_status() -> None: """Given a RunHandle with _status=done (post-close), steer() returns False.""" handle = _make_run_handle() - handle._status = RunStatus.done + handle._run_state = RunState.DONE handle._closing = True result = handle.steer("message") @@ -556,7 +557,7 @@ async def test_steer_returns_false_when_done_status() -> None: async def test_followup_returns_false_when_done_status() -> None: """Given a RunHandle with _status=done (post-close), followup() returns False.""" handle = _make_run_handle() - handle._status = RunStatus.done + handle._run_state = RunState.DONE handle._closing = True result = handle.followup("message") @@ -1032,7 +1033,7 @@ async def test_cancel_returns_to_idle() -> None: handle.cancel() await asyncio.sleep(0.1) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert handle._turn_complete_event.is_set() handle.close() @@ -1070,7 +1071,7 @@ async def _consume() -> None: assert any(isinstance(e, RunFailedEvent) for e in published) # Returns to idle - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE handle.close() await asyncio.sleep(0.05) @@ -1135,7 +1136,7 @@ async def execute(self): # type: ignore[override] await asyncio.sleep(0.05) # After first turn: idle, event set, was cleared at start - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert handle._turn_complete_event.is_set() assert turn1.captured_start is False @@ -1144,7 +1145,7 @@ async def execute(self): # type: ignore[override] await asyncio.sleep(0.1) # After second turn: idle, event set, was cleared at start (was set between turns) - assert handle._status == RunStatus.idle + assert handle._run_state == RunState.IDLE assert handle._turn_complete_event.is_set() assert turn2.captured_start is False diff --git a/tests/orchestrator/test_run_lifecycle.py b/tests/orchestrator/test_run_lifecycle.py index 370580233..38f424a20 100644 --- a/tests/orchestrator/test_run_lifecycle.py +++ b/tests/orchestrator/test_run_lifecycle.py @@ -18,9 +18,10 @@ from agentpool import Agent from agentpool.agents.base_agent import _current_run_ctx_var from agentpool.agents.context import AgentRunContext +from agentpool.lifecycle import RunOutcome, RunState from agentpool.orchestrator.core import SessionPool from agentpool.orchestrator.metrics import MetricsCollector -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.orchestrator.run import RunHandle pytestmark = [pytest.mark.unit, pytest.mark.anyio] @@ -37,7 +38,7 @@ def test_run_handle_defaults() -> None: assert handle.run_id == "r1" assert handle.session_id == "s1" assert handle.agent_type == "native" - assert handle.status == RunStatus.pending + assert handle._run_state == RunState.IDLE assert handle.run_ctx.current_task is None assert not handle.complete_event.is_set() @@ -48,7 +49,7 @@ async def test_start_transitions_to_running() -> None: handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") task: asyncio.Task[Any] = asyncio.create_task(asyncio.sleep(0)) handle._start_task(task) - assert handle.status == RunStatus.running + assert handle.is_running assert handle.run_ctx.current_task is task await task @@ -57,7 +58,7 @@ def test_start_without_task() -> None: """start() works when no task is provided.""" handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle._start_task() - assert handle.status == RunStatus.running + assert handle.is_running assert handle.run_ctx.current_task is None @@ -65,7 +66,8 @@ def test_complete_transitions_and_sets_event() -> None: """complete() transitions to completed and sets complete_event.""" handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle.complete() - assert handle.status == RunStatus.completed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.COMPLETED assert handle.complete_event.is_set() @@ -93,7 +95,8 @@ def test_fail_transitions_and_sets_event() -> None: """fail() transitions to failed and sets complete_event.""" handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle.fail() - assert handle.status == RunStatus.failed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.FAILED assert handle.complete_event.is_set() @@ -102,7 +105,8 @@ def test_fail_with_exception_sets_cancelled() -> None: handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") exc = RuntimeError("boom") handle.fail(exc) - assert handle.status == RunStatus.failed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.FAILED assert handle.run_ctx.cancelled is True @@ -135,7 +139,7 @@ async def test_cancel_sets_cancelled_flag() -> None: handle.cancel() assert handle.run_ctx.cancelled is True # Status should remain running; cleanup is deferred - assert handle.status == RunStatus.running + assert handle.is_running assert not handle.complete_event.is_set() task.cancel() @@ -238,7 +242,7 @@ async def test_get_metrics_counts_native_vs_non_native( run_id="run-1", session_id="sess-native", agent_type="native", - status=RunStatus.running, + _run_state=RunState.RUNNING, ) session_pool.sessions._runs["run-1"] = handle_native @@ -248,7 +252,7 @@ async def test_get_metrics_counts_native_vs_non_native( run_id="run-2", session_id="sess-non-native", agent_type="non-native", - status=RunStatus.running, + _run_state=RunState.RUNNING, ) session_pool.sessions._runs["run-2"] = handle_non_native @@ -277,8 +281,8 @@ def ctxvar_agent() -> Agent[None]: @pytest.mark.unit @pytest.mark.asyncio -async def test_contextvar_set_during_run_stream_once(ctxvar_agent: Agent[None]) -> None: - """_current_run_ctx_var must be non-None during _run_stream_once and None after.""" +async def test_contextvar_set_during_stream_events(ctxvar_agent: Agent[None]) -> None: + """_current_run_ctx_var must be non-None during _stream_events and None after.""" # Before stream starts assert _current_run_ctx_var.get() is None @@ -286,11 +290,11 @@ async def test_contextvar_set_during_run_stream_once(ctxvar_agent: Agent[None]) # Fully consume the stream so the generator's finally block runs naturally async for _event in ctxvar_agent.run_stream("Test prompt"): - # During the stream _run_stream_once is active + # During the stream _stream_events is active if captured_ctx is None: captured_ctx = _current_run_ctx_var.get() assert captured_ctx is not None, ( - "_current_run_ctx_var must be set during _run_stream_once" + "_current_run_ctx_var must be set during _stream_events" ) assert isinstance(captured_ctx, AgentRunContext) diff --git a/tests/orchestrator/test_run_status.py b/tests/orchestrator/test_run_status.py index 187613f3a..f72b17d17 100644 --- a/tests/orchestrator/test_run_status.py +++ b/tests/orchestrator/test_run_status.py @@ -1,59 +1,59 @@ -"""Tests for the RunStatus enum in agentpool.orchestrator.run.""" +"""Tests for the RunState and RunOutcome enums in agentpool.lifecycle.types.""" from __future__ import annotations from enum import Enum -import pytest +from agentpool.lifecycle import RunOutcome, RunState -from agentpool.orchestrator.run import RunStatus - -@pytest.mark.unit -def test_run_status_enum_values() -> None: - """Given the RunStatus enum, it should define exactly 7 lifecycle states.""" - expected_values: set[str] = { - "pending", - "running", - "completed", - "failed", - "checkpointed", - "idle", - "done", - } - actual_values: set[str] = {m.name for m in RunStatus} +def test_run_state_defines_exactly_3_states() -> None: + """Given the RunState enum, it should define exactly 3 lifecycle states.""" + actual_values: set[str] = {m.name for m in RunState} + expected_values = {"IDLE", "RUNNING", "DONE"} assert actual_values == expected_values -@pytest.mark.unit -def test_run_status_is_enum() -> None: - """Given RunStatus, it should be a proper Enum subclass.""" - assert issubclass(RunStatus, Enum) +def test_run_state_is_enum() -> None: + """Given RunState, it should be a proper Enum subclass.""" + assert issubclass(RunState, Enum) + + +def test_run_state_distinctness() -> None: + """Given RunState, all states should be distinct.""" + assert RunState.IDLE is not RunState.DONE + assert RunState.IDLE is not RunState.RUNNING + assert RunState.RUNNING is not RunState.DONE + + +def test_run_state_value_strings() -> None: + """Given RunState, the value strings should match the expected names.""" + assert RunState.IDLE.value == "idle" + assert RunState.RUNNING.value == "running" + assert RunState.DONE.value == "done" + + +def test_run_outcome_defines_exactly_3_outcomes() -> None: + """Given the RunOutcome enum, it should define exactly 3 terminal outcomes.""" + actual_values: set[str] = {m.name for m in RunOutcome} + expected_values = {"COMPLETED", "FAILED", "CHECKPOINTED"} + assert actual_values == expected_values -@pytest.mark.unit -def test_run_status_idle_and_done_are_distinct() -> None: - """Given RunStatus has idle and done, they should be distinct from all existing states.""" - assert RunStatus.idle is not RunStatus.done - assert RunStatus.idle is not RunStatus.pending - assert RunStatus.idle is not RunStatus.running - assert RunStatus.idle is not RunStatus.completed - assert RunStatus.idle is not RunStatus.failed - assert RunStatus.idle is not RunStatus.checkpointed - assert RunStatus.done is not RunStatus.pending - assert RunStatus.done is not RunStatus.running - assert RunStatus.done is not RunStatus.completed - assert RunStatus.done is not RunStatus.failed - assert RunStatus.done is not RunStatus.checkpointed +def test_run_outcome_is_enum() -> None: + """Given RunOutcome, it should be a proper Enum subclass.""" + assert issubclass(RunOutcome, Enum) -@pytest.mark.unit -def test_run_status_idle_name() -> None: - """Given the idle member, its name should be 'idle'.""" - assert RunStatus.idle.name == "idle" +def test_run_outcome_distinctness() -> None: + """Given RunOutcome, all outcomes should be distinct.""" + assert RunOutcome.COMPLETED is not RunOutcome.FAILED + assert RunOutcome.COMPLETED is not RunOutcome.CHECKPOINTED + assert RunOutcome.FAILED is not RunOutcome.CHECKPOINTED -@pytest.mark.unit -def test_run_status_done_name() -> None: - """Given the done member, its name should be 'done'.""" - assert RunStatus.done.name == "done" +def test_run_outcome_value_strings() -> None: + """Given RunOutcome, the value strings should match the expected names.""" + assert RunOutcome.COMPLETED.value == "completed" + assert RunOutcome.FAILED.value == "failed" + assert RunOutcome.CHECKPOINTED.value == "checkpointed" diff --git a/tests/orchestrator/test_runhandle_checkpoint.py b/tests/orchestrator/test_runhandle_checkpoint.py index 22af90169..accac1197 100644 --- a/tests/orchestrator/test_runhandle_checkpoint.py +++ b/tests/orchestrator/test_runhandle_checkpoint.py @@ -1,7 +1,7 @@ """Tests for RunHandle checkpoint-aware status. Verifies that: -1. ``RunStatus.checkpointed`` exists in the enum. +1. ``RunOutcome.CHECKPOINTED`` exists in the enum. 2. ``RunHandle.checkpoint()`` transitions status and sets ``complete_event``. 3. ``RunFailedEvent`` is NOT emitted on checkpoint transition. 4. Resume creates a fresh ``RunHandle``. @@ -13,30 +13,30 @@ import pytest -from agentpool.orchestrator.run import RunHandle, RunStatus +from agentpool.lifecycle import RunOutcome, RunState +from agentpool.orchestrator.run import RunHandle pytestmark = [pytest.mark.unit, pytest.mark.anyio] # ============================================================================ -# RunStatus.checkpointed existence +# RunOutcome.CHECKPOINTED existence # ============================================================================ def test_checkpointed_status_exists() -> None: - """RunStatus.checkpointed must be a member of the enum.""" - assert hasattr(RunStatus, "checkpointed") - assert RunStatus.checkpointed is not None - assert isinstance(RunStatus.checkpointed, RunStatus) + """RunOutcome.CHECKPOINTED must be a member of the enum.""" + assert RunOutcome.CHECKPOINTED is not None + assert isinstance(RunOutcome.CHECKPOINTED, RunOutcome) def test_checkpointed_is_distinct() -> None: - """RunStatus.checkpointed must differ from existing states.""" - assert RunStatus.checkpointed != RunStatus.pending - assert RunStatus.checkpointed != RunStatus.running - assert RunStatus.checkpointed != RunStatus.completed - assert RunStatus.checkpointed != RunStatus.failed + """RunOutcome.CHECKPOINTED must differ from existing states.""" + assert RunOutcome.CHECKPOINTED != RunState.IDLE + assert RunOutcome.CHECKPOINTED != RunState.RUNNING + assert RunOutcome.CHECKPOINTED != RunOutcome.COMPLETED + assert RunOutcome.CHECKPOINTED != RunOutcome.FAILED # ============================================================================ @@ -47,7 +47,6 @@ def test_checkpointed_is_distinct() -> None: def test_checkpoint_method_exists() -> None: """RunHandle must have a checkpoint() method.""" handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") - assert hasattr(handle, "checkpoint") assert callable(handle.checkpoint) @@ -55,9 +54,10 @@ def test_checkpoint_transitions_from_running() -> None: """checkpoint() transitions from running to checkpointed.""" handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle._start_task() - assert handle.status == RunStatus.running + assert handle.is_running handle.checkpoint() - assert handle.status == RunStatus.checkpointed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.CHECKPOINTED def test_checkpoint_sets_complete_event() -> None: @@ -105,7 +105,8 @@ def test_checkpoint_does_not_emit_run_failed_event() -> None: # RunHandle.checkpoint() does not accept event_bus parameter, # so RunFailedEvent cannot be emitted - assert handle.status == RunStatus.checkpointed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.CHECKPOINTED assert handle.complete_event.is_set() @@ -136,12 +137,13 @@ def test_resume_creates_fresh_run_handle() -> None: old_handle = RunHandle(run_id="old-run", session_id="s1", agent_type="native") old_handle._start_task() old_handle.checkpoint() - assert old_handle.status == RunStatus.checkpointed + assert old_handle._run_state == RunState.DONE + assert old_handle.outcome == RunOutcome.CHECKPOINTED # Simulate resume: create a completely new handle new_handle = RunHandle(run_id="new-run", session_id="s1", agent_type="native") new_handle._start_task() - assert new_handle.status == RunStatus.running + assert new_handle.is_running assert new_handle.run_id != old_handle.run_id @@ -154,8 +156,8 @@ async def test_session_controller_skips_fail_on_checkpointed() -> None: """SessionController must skip fail() when RunHandle is checkpointed. This tests the guard in ``_run_turn_unlocked`` that checks - ``run_handle.status not in (RunStatus.completed, RunStatus.failed, - RunStatus.checkpointed)`` before calling ``run_handle.fail()``. + ``run_handle.outcome not in (RunOutcome.COMPLETED, RunOutcome.FAILED, + RunOutcome.CHECKPOINTED)`` before calling ``run_handle.fail()``. """ from agentpool import AgentsManifest from agentpool.delegation import AgentPool @@ -171,18 +173,19 @@ async def test_session_controller_skips_fail_on_checkpointed() -> None: handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle._start_task() handle.checkpoint() - assert handle.status == RunStatus.checkpointed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.CHECKPOINTED # Simulate the guard in _run_turn_unlocked's except block # (line ~1354-1358 in core.py): - # if run_handle.status not in (RunStatus.completed, RunStatus.failed): + # if run_handle.outcome not in (RunOutcome.COMPLETED, RunOutcome.FAILED): # run_handle.fail(...) - should_skip = handle.status in (RunStatus.completed, RunStatus.failed) - # With RunStatus.checkpointed added to the exclusion, fail() should NOT be called - should_fail = handle.status not in ( - RunStatus.completed, - RunStatus.failed, - RunStatus.checkpointed, + should_skip = handle.outcome in (RunOutcome.COMPLETED, RunOutcome.FAILED) + # With RunOutcome.CHECKPOINTED added to the exclusion, fail() should NOT be called + should_fail = handle.outcome not in ( + RunOutcome.COMPLETED, + RunOutcome.FAILED, + RunOutcome.CHECKPOINTED, ) assert not should_fail, "checkpointed runs must not transition to failed in except" assert not should_skip, "checkpointed status should be excluded from fail path" @@ -192,36 +195,39 @@ async def test_run_loop_finally_skips_complete_on_checkpointed() -> None: """Run loop must NOT call complete() when RunHandle is checkpointed. This tests the guard in ``_run_turn_unlocked``'s finally block that checks - ``run_handle.status not in (RunStatus.completed, RunStatus.failed)`` + ``run_handle.outcome not in (RunOutcome.COMPLETED, RunOutcome.FAILED)`` before calling ``run_handle.complete()``. """ # The finally block logic is: - # if run_handle.status not in (RunStatus.completed, RunStatus.failed): + # if run_handle.outcome not in (RunOutcome.COMPLETED, RunOutcome.FAILED): # run_handle.complete() # # When checkpointed, this guard should ALSO skip complete(): - # if run_handle.status not in (RunStatus.completed, RunStatus.failed, RunStatus.checkpointed): + # if run_handle.outcome not in ( + # RunOutcome.COMPLETED, RunOutcome.FAILED, RunOutcome.CHECKPOINTED, + # ): # if run_ctx.checkpointed: # run_handle.checkpoint() # else: # run_handle.complete() # # After checkpoint() was already called above, the guard must be: - assert RunStatus.checkpointed not in (RunStatus.completed, RunStatus.failed) + assert RunOutcome.CHECKPOINTED not in (RunOutcome.COMPLETED, RunOutcome.FAILED) # The finally block must NOT call complete() when status is checkpointed handle = RunHandle(run_id="r1", session_id="s1", agent_type="native") handle._start_task() handle.checkpoint() - assert handle.status == RunStatus.checkpointed + assert handle._run_state == RunState.DONE + assert handle.outcome == RunOutcome.CHECKPOINTED # If the guard only checks (completed, failed), it would call complete() # and change the status. Verify the guard must include checkpointed: - guard_ok = handle.status in (RunStatus.completed, RunStatus.failed) - guard_with_checkpointed = handle.status in ( - RunStatus.completed, - RunStatus.failed, - RunStatus.checkpointed, + guard_ok = handle.outcome in (RunOutcome.COMPLETED, RunOutcome.FAILED) + guard_with_checkpointed = handle.outcome in ( + RunOutcome.COMPLETED, + RunOutcome.FAILED, + RunOutcome.CHECKPOINTED, ) assert not guard_ok, "guard without checkpointed would incorrectly fall through" assert guard_with_checkpointed, "guard must include checkpointed to skip complete()" diff --git a/tests/orchestrator/test_session_lifecycle.py b/tests/orchestrator/test_session_lifecycle.py index bf9a4b8bb..6f6f9c366 100644 --- a/tests/orchestrator/test_session_lifecycle.py +++ b/tests/orchestrator/test_session_lifecycle.py @@ -54,7 +54,7 @@ async def run_stream( session_id: str | None = None, **kwargs: Any, ) -> AsyncIterator[Any]: - """Mock run_stream that delegates to _stream_impl via _run_stream_once.""" + """Mock run_stream that delegates to _stream_impl via _stream_events.""" if self._stream_impl is None: raise RuntimeError("No stream impl set") run_ctx = MagicMock() @@ -66,7 +66,7 @@ async def run_stream( # Yield at least one event so the run doesn't hang yield RunStartedEvent(session_id=session_id or "", run_id="run-mock") - async def _run_stream_once( + async def _stream_events( self, run_ctx: AgentRunContext, *prompts: Any, diff --git a/tests/orchestrator/test_session_pool.py b/tests/orchestrator/test_session_pool.py index e745b7d95..408e2878d 100644 --- a/tests/orchestrator/test_session_pool.py +++ b/tests/orchestrator/test_session_pool.py @@ -358,7 +358,8 @@ async def test_cache_copy_messages_invalidates_target_cache( # RunHandle delegation tests (feature-flag gated) # --------------------------------------------------------------------------- -from agentpool.orchestrator.run import RunHandle, RunStatus # noqa: E402 +from agentpool.lifecycle import RunState # noqa: E402 +from agentpool.orchestrator.run import RunHandle # noqa: E402 def _make_mock_agent() -> MagicMock: @@ -399,7 +400,7 @@ def _setup_active_run( run_handle.steer = MagicMock(return_value=True) run_handle.followup = MagicMock(return_value=True) run_handle.run_id = "test-run-id" - run_handle.status = RunStatus.running + run_handle._run_state = RunState.RUNNING session_pool.sessions._runs["test-run-id"] = run_handle session = session_pool.sessions.get_session(session_id) assert session is not None diff --git a/tests/orchestrator/test_session_pool_hooks.py b/tests/orchestrator/test_session_pool_hooks.py index 316e32196..d9b6e765c 100644 --- a/tests/orchestrator/test_session_pool_hooks.py +++ b/tests/orchestrator/test_session_pool_hooks.py @@ -5,7 +5,8 @@ The test verifies: 1. Hooks fire when going through RunHandle.start() (the SessionPool path) -2. hooks_fired is cleared between turns so turn 2 hooks still fire +2. Hooks fire in turn 2 without needing to clear any per-turn state + (new Turn instances have fresh _logged_tools sets) """ from __future__ import annotations @@ -129,13 +130,11 @@ async def _consume() -> None: async def test_hooks_fired_cleared_between_turns() -> None: """Given two sequential turns, turn 2 hooks still fire. - RunHandle.start() clears hooks_fired at the start of each turn. - Without this, turn 1's 'pre_turn' guard key would block turn 2's - pre_turn from firing. - - This test runs turn 1 through RunHandle.start(), then manually - creates a second turn with the same run_ctx (after clearing - hooks_fired) to verify hooks fire again. + Previously, ``hooks_fired`` on ``AgentRunContext`` needed to be cleared + between turns to prevent turn 1's guard keys from blocking turn 2. With + the ``hooks_fired`` guard removed (replaced by per-Turn ``_logged_tools`` + set), a new Turn instance is created for each turn with a fresh set, + so no explicit clearing is needed. """ _reset_calls() hooks = AgentHooks( @@ -167,10 +166,8 @@ async def _consume() -> None: pre_turn_count = sum(1 for name, _ in hook_calls if name == "pre_turn") assert pre_turn_count == 1, "pre_turn must fire in turn 1" - # Simulate RunHandle.start() clearing hooks_fired for turn 2 - run_ctx.hooks_fired.clear() - - # Turn 2: Create a new turn manually with the same run_ctx + # Turn 2: Create a new turn manually with the same run_ctx. + # No hooks_fired.clear() needed — new Turn has fresh _logged_tools. from agentpool.agents.native_agent.turn import NativeTurn turn = NativeTurn( @@ -183,7 +180,7 @@ async def _consume() -> None: _ = [event async for event in turn.execute()] pre_turn_count = sum(1 for name, _ in hook_calls if name == "pre_turn") - assert pre_turn_count == 2, "pre_turn must fire again in turn 2 after clearing hooks_fired" + assert pre_turn_count == 2, "pre_turn must fire again in turn 2" # --------------------------------------------------------------------------- diff --git a/tests/orchestrator/test_session_pool_input_provider.py b/tests/orchestrator/test_session_pool_input_provider.py index 640162014..90eebd212 100644 --- a/tests/orchestrator/test_session_pool_input_provider.py +++ b/tests/orchestrator/test_session_pool_input_provider.py @@ -61,7 +61,7 @@ def session_pool(mock_pool: MagicMock) -> SessionPool: @pytest.fixture def mock_agent() -> MagicMock: - """Return a mocked BaseAgent that captures kwargs passed to _run_stream_once.""" + """Return a mocked BaseAgent that captures kwargs passed to _stream_events.""" agent = MagicMock() agent.get_active_run_context.return_value = None agent.AGENT_TYPE = "native" @@ -78,7 +78,7 @@ async def _fake_stream( msg = ChatMessage(content="test response", role="assistant") yield StreamCompleteEvent(message=msg) - agent._run_stream_once = _fake_stream + agent._stream_events = _fake_stream return agent @@ -222,7 +222,7 @@ async def _capturing_get_agent( ) @pytest.mark.anyio - async def test_run_turn_passes_input_provider_to_run_stream_once( + async def test_run_turn_passes_input_provider_to_stream_events( self, session_pool: SessionPool, mock_pool: MagicMock, diff --git a/tests/servers/acp_server/test_acp_cancel_then_prompt.py b/tests/servers/acp_server/test_acp_cancel_then_prompt.py index fdce762c6..7756703d4 100644 --- a/tests/servers/acp_server/test_acp_cancel_then_prompt.py +++ b/tests/servers/acp_server/test_acp_cancel_then_prompt.py @@ -20,9 +20,9 @@ import pytest from agentpool.agents.events import RunFailedEvent, RunStartedEvent, StreamCompleteEvent +from agentpool.lifecycle import RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventEnvelope, SessionPool -from agentpool.orchestrator.run import RunStatus from agentpool.orchestrator.turn import Turn @@ -244,8 +244,8 @@ async def test_acp_cancel_then_prompt_no_hang( "New RunHandle should be a different instance if old one was cleaned up" ) - assert first_handle._status in (RunStatus.idle, RunStatus.done), ( - f"First RunHandle should be idle or done, got: {first_handle._status}" + assert first_handle._run_state in (RunState.IDLE, RunState.DONE), ( + f"First RunHandle should be idle or done, got: {first_handle._run_state}" ) # Cleanup: close the RunHandle first so the start() loop exits and @@ -345,8 +345,8 @@ def _create_turn( ) # RunHandle should be idle, waiting for the next prompt - assert first_handle._status == RunStatus.idle, ( - f"Expected RunHandle status idle after cancel, got {first_handle._status}" + assert first_handle._run_state == RunState.IDLE, ( + f"Expected RunHandle status idle after cancel, got {first_handle._run_state}" ) # No new events should have been published between cancel and idle diff --git a/tests/servers/acp_server/test_agent_role.py b/tests/servers/acp_server/test_agent_role.py index 39b063a81..64fe5febc 100644 --- a/tests/servers/acp_server/test_agent_role.py +++ b/tests/servers/acp_server/test_agent_role.py @@ -98,6 +98,7 @@ async def test_role_swap_success(self, mock_acp_agent): session._task_lock = MagicMock() session._task_lock.locked.return_value = False session.switch_active_agent = AsyncMock() + session.is_busy = False acp_agent.session_manager._acp_sessions = {"sess_1": session} acp_agent.session_manager.get_session = lambda sid: session @@ -127,6 +128,7 @@ async def test_role_swap_invalid_agent(self, mock_acp_agent): session.agent.name = "default" session._task_lock = MagicMock() session._task_lock.locked.return_value = False + session.is_busy = False session.switch_active_agent = AsyncMock(side_effect=ValueError("Agent not found")) acp_agent.session_manager.get_session = lambda sid: session @@ -142,6 +144,7 @@ async def test_swap_no_history_inheritance(self, mock_acp_agent): session.agent.name = "default" session._task_lock = MagicMock() session._task_lock.locked.return_value = False + session.is_busy = False session.switch_active_agent = AsyncMock() acp_agent.session_manager.get_session = lambda sid: session diff --git a/tests/servers/opencode_server/conftest.py b/tests/servers/opencode_server/conftest.py index d78755771..9f0411a37 100644 --- a/tests/servers/opencode_server/conftest.py +++ b/tests/servers/opencode_server/conftest.py @@ -308,10 +308,10 @@ def mock_agent(mock_env: Mock, mock_pool: Mock, storage_manager: StorageManager) agent._input_provider = None agent.run = AsyncMock(return_value=Mock(data="test response")) agent.agent_pool = mock_pool - # host_context is accessed by ServerState.__post_init__ instead of agent_pool - # state.py resolves _pool via _ctx.pool, so mock_pool.pool must return itself + # host_context is accessed by ServerState.__post_init__ for manifest etc. + # state.py resolves _pool via agent._agent_pool, so set it directly. + agent._agent_pool = mock_pool agent.host_context = mock_pool - mock_pool.pool = mock_pool # Real storage manager (accessed via state.storage -> agent.storage) agent.storage = storage_manager diff --git a/tests/servers/opencode_server/test_auto_resume_message_redflag.py b/tests/servers/opencode_server/test_auto_resume_message_redflag.py index 27170f060..87c6f5fad 100644 --- a/tests/servers/opencode_server/test_auto_resume_message_redflag.py +++ b/tests/servers/opencode_server/test_auto_resume_message_redflag.py @@ -66,7 +66,7 @@ def mock_agent_pool() -> Mock: pool.manifest.agents = {} pool._config_file_path = None - async def _mock_run_stream_once(*args: Any, **kwargs: Any) -> Any: + async def _mock_stream_events(*args: Any, **kwargs: Any) -> Any: """Yield a minimal run event sequence for testing.""" session_id = kwargs.get("session_id", "unknown") run_id = "run-mock-001" @@ -76,7 +76,7 @@ async def _mock_run_stream_once(*args: Any, **kwargs: Any) -> Any: ) mock_agent = Mock() - mock_agent._run_stream_once = _mock_run_stream_once + mock_agent._stream_events = _mock_stream_events mock_agent._input_provider = None mock_agent.conversation = Mock() mock_agent.conversation.add_chat_messages = Mock() diff --git a/tests/servers/opencode_server/test_cancelled_message.py b/tests/servers/opencode_server/test_cancelled_message.py index 003b27e25..fa0050229 100644 --- a/tests/servers/opencode_server/test_cancelled_message.py +++ b/tests/servers/opencode_server/test_cancelled_message.py @@ -114,10 +114,10 @@ def cancellable_mock_agent(): agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool # Set up SessionPool mock for new architecture - from agentpool.orchestrator.run import RunStatus + from agentpool.lifecycle import RunState session_pool = Mock() session_pool.sessions = Mock() @@ -130,7 +130,7 @@ def cancellable_mock_agent(): session_pool.sessions.store = None # Create a RunHandle that raises CancelledError when waiting run_handle = Mock() - run_handle.status = RunStatus.running + run_handle._run_state = RunState.RUNNING run_handle.complete_event = Mock() run_handle.complete_event.wait = AsyncMock(side_effect=asyncio.CancelledError) session_pool.receive_request = AsyncMock(return_value=run_handle) @@ -448,7 +448,7 @@ async def test_cancelled_message_preserves_conversation_history( """After cancellation, the agent's in-memory conversation must include the aborted response. The agent's `conversation.chat_messages` is what gets sent to the LLM as - conversation history. When a run is cancelled, `_run_stream_once` adds the + conversation history. When a run is cancelled, `_stream_events` adds the user message but never adds the assistant response (the post-processing code at base_agent.py:857-858 is skipped due to the exception). @@ -472,9 +472,9 @@ async def test_cancelled_message_preserves_conversation_history( session_id, sample_message_request, state, user_msg_id, user_msg_with_parts ) - # In the real flow, _run_stream_once adds the user message to conversation + # In the real flow, _stream_events adds the user message to conversation # (base_agent.py:784), and our CancelledError handler adds the aborted - # assistant message. Since CancellableAgentMock doesn't run _run_stream_once, + # assistant message. Since CancellableAgentMock doesn't run _stream_events, # only our handler's addition is reflected. What matters is that the # aborted assistant message IS present in the conversation. final_count = len(state.agent.conversation.chat_messages) diff --git a/tests/servers/opencode_server/test_concurrent_messages.py b/tests/servers/opencode_server/test_concurrent_messages.py index ed2097446..b1612bb6b 100644 --- a/tests/servers/opencode_server/test_concurrent_messages.py +++ b/tests/servers/opencode_server/test_concurrent_messages.py @@ -149,12 +149,13 @@ async def save_session(session_data: Any) -> None: # Mock receive_request to actually call agent.run_stream and publish events async def _mock_receive_request(*, session_id, content, priority, input_provider): - from agentpool.orchestrator.run import RunHandle, RunStatus + from agentpool.lifecycle import RunOutcome, RunState + from agentpool.orchestrator.run import RunHandle handle = Mock(spec=RunHandle) handle.run_id = "test-run" handle.session_id = session_id - handle.status = RunStatus.running + handle._run_state = RunState.RUNNING complete_event = asyncio.Event() handle.complete_event = complete_event @@ -163,9 +164,11 @@ async def _do_run(): stream = agent.run_stream(content, session_id=session_id) async for event in stream: await event_bus.publish(session_id, event) - handle.status = RunStatus.completed + handle._run_state = RunState.DONE + handle.outcome = RunOutcome.COMPLETED except Exception: # noqa: BLE001 - handle.status = RunStatus.failed + handle._run_state = RunState.DONE + handle.outcome = RunOutcome.FAILED finally: complete_event.set() @@ -179,7 +182,7 @@ async def _do_run(): agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool # Set up env mock env = Mock() diff --git a/tests/servers/opencode_server/test_ensure_session.py b/tests/servers/opencode_server/test_ensure_session.py index 53e6ccf3a..880ce979a 100644 --- a/tests/servers/opencode_server/test_ensure_session.py +++ b/tests/servers/opencode_server/test_ensure_session.py @@ -24,7 +24,7 @@ def create_mock_agent() -> MagicMock: agent.name = "test_agent" agent.session_id = "original_session_id" agent.host_context = MagicMock() - agent.host_context.pool = agent.host_context # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.host_context # state.py resolves _pool via agent._agent_pool agent.host_context.manifest.config_file_path = "test_config.yml" agent.host_context.storage.save_session = AsyncMock() agent.host_context.storage.load_session = AsyncMock(return_value=None) diff --git a/tests/servers/opencode_server/test_ensure_session_durable.py b/tests/servers/opencode_server/test_ensure_session_durable.py index 2a686916f..bef690722 100644 --- a/tests/servers/opencode_server/test_ensure_session_durable.py +++ b/tests/servers/opencode_server/test_ensure_session_durable.py @@ -36,7 +36,7 @@ def create_mock_agent() -> MagicMock: agent.name = "test_agent" agent.session_id = "original_session_id" agent.host_context = MagicMock() - agent.host_context.pool = agent.host_context # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.host_context # state.py resolves _pool via agent._agent_pool agent.host_context.manifest.config_file_path = "test_config.yml" agent.host_context.storage.save_session = AsyncMock() agent.host_context.storage.load_session = AsyncMock(return_value=None) diff --git a/tests/servers/opencode_server/test_ensure_session_store_first.py b/tests/servers/opencode_server/test_ensure_session_store_first.py index ca254ffd6..7162bb2a3 100644 --- a/tests/servers/opencode_server/test_ensure_session_store_first.py +++ b/tests/servers/opencode_server/test_ensure_session_store_first.py @@ -32,7 +32,7 @@ def create_mock_agent() -> MagicMock: agent.name = "test_agent" agent.session_id = "original_session_id" agent.host_context = MagicMock() - agent.host_context.pool = agent.host_context # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.host_context # state.py resolves _pool via agent._agent_pool agent.host_context.manifest.config_file_path = "test_config.yml" agent.host_context.storage.save_session = AsyncMock() agent.host_context.storage.load_session = AsyncMock(return_value=None) diff --git a/tests/servers/opencode_server/test_event_pipeline_e2e.py b/tests/servers/opencode_server/test_event_pipeline_e2e.py index ffc9a448c..9e2944441 100644 --- a/tests/servers/opencode_server/test_event_pipeline_e2e.py +++ b/tests/servers/opencode_server/test_event_pipeline_e2e.py @@ -100,7 +100,7 @@ def server_state(tmp_path: Any) -> ServerState: agent.name = "test-agent" agent.storage = Mock() agent.host_context = Mock() - agent.host_context.pool = agent.host_context # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.host_context # state.py resolves _pool via agent._agent_pool agent.host_context.session_pool = Mock() agent.host_context.session_pool.event_bus = Mock() # will be overridden state = ServerState(working_dir=str(tmp_path), agent=agent) diff --git a/tests/servers/opencode_server/test_event_processor.py b/tests/servers/opencode_server/test_event_processor.py index 418e3b8fe..c6c9c964b 100644 --- a/tests/servers/opencode_server/test_event_processor.py +++ b/tests/servers/opencode_server/test_event_processor.py @@ -26,6 +26,7 @@ MessageWithParts, PartDeltaEvent, PartUpdatedEvent, + SessionStatusEvent, TextPart, ) @@ -189,3 +190,114 @@ async def test_process_text_delta_without_start(server_state: ServerState) -> No first_part = assistant_msg.parts[0] assert isinstance(first_part, TextPart) assert first_part.text == "Some text" + + +# ============================================================================= +# StreamCompleteEvent Cancellation Tests +# ============================================================================= + + +def _make_stream_complete_context(server_state: ServerState) -> EventProcessorContext: + """Create a minimal EventProcessorContext for StreamCompleteEvent tests.""" + assistant_msg = MessageWithParts.assistant( + message_id="msg-1", + session_id="test-session", + time=MessageTime(created=0), + agent_name="test-agent", + model_id="test-model", + parent_id="parent-1", + provider_id="agentpool", + path=MessagePath(cwd="/tmp", root="/tmp"), + ) + return EventProcessorContext( + session_id="test-session", + assistant_msg_id="msg-1", + assistant_msg=assistant_msg, + state=server_state, + working_dir="/tmp", + ) + + +@pytest.mark.asyncio +async def test_stream_complete_emits_idle_status(server_state: ServerState) -> None: + """Test that a completed StreamCompleteEvent emits SessionStatusEvent(idle).""" + from agentpool.agents.events import StreamCompleteEvent + from agentpool.messaging import ChatMessage + + processor = EventProcessor() + ctx = _make_stream_complete_context(server_state) + + msg = ChatMessage[str](content="done", role="assistant") + event = StreamCompleteEvent(message=msg, cancelled=False) + + events = [e async for e in processor.process(event, ctx)] + + status_events = [e for e in events if isinstance(e, SessionStatusEvent)] + assert len(status_events) == 1 + assert status_events[0].properties.status.type == "idle" + + +@pytest.mark.asyncio +async def test_stream_complete_emits_cancelled_status(server_state: ServerState) -> None: + """Test that a cancelled StreamCompleteEvent emits SessionStatusEvent(cancelled).""" + from agentpool.agents.events import StreamCompleteEvent + from agentpool.messaging import ChatMessage + + processor = EventProcessor() + ctx = _make_stream_complete_context(server_state) + + msg = ChatMessage[str](content="partial", role="assistant") + event = StreamCompleteEvent(message=msg, cancelled=True) + + events = [e async for e in processor.process(event, ctx)] + + status_events = [e for e in events if isinstance(e, SessionStatusEvent)] + assert len(status_events) == 1 + assert status_events[0].properties.status.type == "cancelled" + + +# ============================================================================= +# McpToolsChangedEvent Tests +# ============================================================================= + + +def test_create_mcp_tools_changed_event() -> None: + """EventProcessor.create_mcp_tools_changed_event creates correct event.""" + from agentpool_server.opencode_server.models.events import McpToolsChangedEvent + + processor = EventProcessor() + event = processor.create_mcp_tools_changed_event(server="my_mcp_server") + + assert isinstance(event, McpToolsChangedEvent) + assert event.type == "mcp.tools.changed" + assert event.properties.server == "my_mcp_server" + + +@pytest.mark.anyio +async def test_mcp_tools_changed_event_from_change_event() -> None: + """Full wiring: ChangeEvent(kind='tools_changed') → McpToolsChangedEvent. + + Simulates the flow: + 1. McpServerCap.on_change() yields ChangeEvent(kind="tools_changed") + 2. EventProcessor.create_mcp_tools_changed_event() converts to McpToolsChangedEvent + """ + from agentpool.capabilities.change_event import ChangeEvent + from agentpool_server.opencode_server.models.events import McpToolsChangedEvent + + processor = EventProcessor() + + # Simulate a ChangeEvent from McpServerCap._on_tools_changed() + change_event = ChangeEvent( + capability_name="my_mcp_server", + kind="tools_changed", + source_uri="mcp://my_mcp_server", + ) + + # The server's _watch_mcp_tool_changes task would do this conversion + oc_event = processor.create_mcp_tools_changed_event( + server=change_event.capability_name, + ) + + assert isinstance(oc_event, McpToolsChangedEvent) + assert oc_event.properties.server == "my_mcp_server" + assert oc_event.type == "mcp.tools.changed" diff --git a/tests/servers/opencode_server/test_global_compat_routes.py b/tests/servers/opencode_server/test_global_compat_routes.py index a24a829d5..26e4b9466 100644 --- a/tests/servers/opencode_server/test_global_compat_routes.py +++ b/tests/servers/opencode_server/test_global_compat_routes.py @@ -55,7 +55,7 @@ def _server_state(tmp_path: Path) -> ServerState: agent._input_provider = None agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool agent.storage = storage_manager agent.get_available_models = AsyncMock(return_value=[]) diff --git a/tests/servers/opencode_server/test_opencode_model_switching.py b/tests/servers/opencode_server/test_opencode_model_switching.py index 53d7fcc32..9d5dd48a2 100644 --- a/tests/servers/opencode_server/test_opencode_model_switching.py +++ b/tests/servers/opencode_server/test_opencode_model_switching.py @@ -423,7 +423,7 @@ def _make_mock_state_with_session_agent( """Create a ServerState wired so get_or_create_session_agent returns per-session mocks.""" from unittest.mock import AsyncMock, Mock - from agentpool.orchestrator.run import RunStatus + from agentpool.lifecycle import RunOutcome, RunState from agentpool.utils.time_utils import now_ms from agentpool_server.opencode_server.state import ServerState @@ -465,7 +465,8 @@ async def _get_or_create_session_agent( # RunHandle that completes immediately run_handle = Mock() - run_handle.status = RunStatus.completed + run_handle._run_state = RunState.DONE + run_handle.outcome = RunOutcome.COMPLETED run_handle.complete_event = Mock() run_handle.complete_event.wait = AsyncMock(return_value=None) session_pool.receive_request = AsyncMock(return_value=run_handle) @@ -480,7 +481,7 @@ async def _get_or_create_session_agent( pool.session_pool = session_pool shared_agent.agent_pool = pool shared_agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + shared_agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool state = ServerState( working_dir=str(tmp_project_dir), diff --git a/tests/servers/opencode_server/test_question_abort_regression.py b/tests/servers/opencode_server/test_question_abort_regression.py index 516f16cba..6a6616435 100644 --- a/tests/servers/opencode_server/test_question_abort_regression.py +++ b/tests/servers/opencode_server/test_question_abort_regression.py @@ -265,11 +265,11 @@ async def _mock_receive_request( priority: str = "when_idle", input_provider: Any = None, ) -> Any: - from agentpool.orchestrator.run import RunStatus + from agentpool.lifecycle import RunOutcome, RunState complete_event = asyncio.Event() run_handle = Mock() - run_handle.status = RunStatus.running + run_handle._run_state = RunState.RUNNING run_handle.complete_event = complete_event async def _background_run(): @@ -277,9 +277,11 @@ async def _background_run(): stream = agent.run_stream(content, session_id=session_id) async for _ in stream: pass - run_handle.status = RunStatus.completed + run_handle._run_state = RunState.DONE + run_handle.outcome = RunOutcome.COMPLETED except Exception: # noqa: BLE001 - run_handle.status = RunStatus.failed + run_handle._run_state = RunState.DONE + run_handle.outcome = RunOutcome.FAILED finally: complete_event.set() @@ -308,7 +310,7 @@ def aborted_mock_agent(tmp_project_dir): agent = RunAbortedAgentMock() agent.agent_pool = _make_pool_mock(agent) agent.host_context = agent.agent_pool - agent.agent_pool.pool = agent.agent_pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.agent_pool # state.py resolves _pool via agent._agent_pool agent.env = _make_env_mock(str(tmp_project_dir)) agent.storage = agent.agent_pool.storage return agent @@ -320,7 +322,7 @@ def blocking_mock_agent(tmp_project_dir): agent = BlockingOnQuestionAgentMock() agent.agent_pool = _make_pool_mock(agent) agent.host_context = agent.agent_pool - agent.agent_pool.pool = agent.agent_pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.agent_pool # state.py resolves _pool via agent._agent_pool agent.env = _make_env_mock(str(tmp_project_dir)) agent.storage = agent.agent_pool.storage return agent @@ -343,9 +345,9 @@ def blocking_real_question_state(tmp_project_dir): placeholder_agent = RunAbortedAgentMock() placeholder_agent.agent_pool = _make_pool_mock(placeholder_agent) placeholder_agent.host_context = placeholder_agent.agent_pool - placeholder_agent.agent_pool.pool = ( + placeholder_agent._agent_pool = ( placeholder_agent.agent_pool - ) # state.py resolves _pool via _ctx.pool + ) # state.py resolves _pool via agent._agent_pool placeholder_agent.env = _make_env_mock(str(tmp_project_dir)) placeholder_agent.storage = placeholder_agent.agent_pool.storage state = ServerState(working_dir=str(tmp_project_dir), agent=placeholder_agent) @@ -372,7 +374,7 @@ def _cancel_all(): real_agent = BlockingOnRealQuestionAgentMock(state) real_agent.agent_pool = _make_pool_mock(real_agent) real_agent.host_context = real_agent.agent_pool - real_agent.agent_pool.pool = real_agent.agent_pool # state.py resolves _pool via _ctx.pool + real_agent._agent_pool = real_agent.agent_pool # state.py resolves _pool via agent._agent_pool real_agent.env = _make_env_mock(str(tmp_project_dir)) real_agent.storage = real_agent.agent_pool.storage state.agent = real_agent diff --git a/tests/servers/opencode_server/test_reasoning.py b/tests/servers/opencode_server/test_reasoning.py index 9a8166193..3fa15d492 100644 --- a/tests/servers/opencode_server/test_reasoning.py +++ b/tests/servers/opencode_server/test_reasoning.py @@ -1,4 +1,4 @@ -"""Tests for reasoning/thinking part behavior in OpenCode stream adapter.""" +"""Tests for reasoning/thinking part behavior in OpenCode event processor.""" from unittest.mock import MagicMock @@ -12,51 +12,68 @@ ) import pytest -from agentpool_server.opencode_server.models import PartDeltaEvent, PartUpdatedEvent +from agentpool_server.opencode_server.event_processor import EventProcessor +from agentpool_server.opencode_server.event_processor_context import ( + EventProcessorContext, +) +from agentpool_server.opencode_server.models import PartUpdatedEvent from agentpool_server.opencode_server.models.events import ( - PartDeltaEventProperties, PartUpdatedEventProperties, ) from agentpool_server.opencode_server.models.parts import ( ReasoningPart, TextPart as OpenCodeTextPart, ) -from agentpool_server.opencode_server.stream_adapter import OpenCodeStreamAdapter + + +def _make_processor_and_ctx( + mock_msg: MagicMock | None = None, +) -> tuple[EventProcessor, EventProcessorContext]: + """Create an EventProcessor and EventProcessorContext for testing.""" + if mock_msg is None: + mock_msg = MagicMock() + mock_msg.parts = [] + mock_state = MagicMock() + processor = EventProcessor() + ctx = EventProcessorContext( + session_id="test-session", + assistant_msg_id="msg-1", + assistant_msg=mock_msg, + state=mock_state, + working_dir=".", + ) + return processor, ctx + + +async def _process_event( + processor: EventProcessor, + ctx: EventProcessorContext, + event: object, +) -> list[object]: + """Process a single event and return the resulting events.""" + return [e async for e in processor.process(event, ctx)] # type: ignore[arg-type] @pytest.mark.asyncio async def test_thinking_events_create_reasoning_part(): """Verify ThinkingPart/ThinkingPartDelta events create ReasoningPart.""" - # Create a mock MessageWithParts mock_msg = MagicMock() mock_msg.parts = [] - mock_state = MagicMock() + processor, ctx = _make_processor_and_ctx(mock_msg) - adapter = OpenCodeStreamAdapter( - state=mock_state, - session_id="test-session", - assistant_msg_id="msg-1", - assistant_msg=mock_msg, - working_dir=".", + events = await _process_event( + processor, ctx, PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) ) - - # Use the adapter's _handle_event method directly - events = [ - e - async for e in adapter._handle_event( - PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) - ) - ] - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=" more...")) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=" more...")), ) - ]) + ) # Assert reasoning part was created and content accumulated - # Check both PartUpdatedEvent (creation) and PartDeltaEvent (delta updates) reasoning_parts = [] for e in events: if isinstance(e, PartUpdatedEvent): @@ -65,16 +82,11 @@ async def test_thinking_events_create_reasoning_part(): props.part, ReasoningPart ): reasoning_parts.append(props.part) - elif isinstance(e, PartDeltaEvent): - props = e.properties - if isinstance(props, PartDeltaEventProperties) and props.field == "text": - # Delta events don't have the full part, check adapter context - pass - # Also verify accumulation via adapter's main_context - assert adapter.main_context.reasoning_part is not None, "ReasoningPart should be created" - assert "Thinking..." in adapter.main_context.reasoning_part.text - assert " more..." in adapter.main_context.reasoning_part.text + # Verify accumulation via context + assert ctx.reasoning_part is not None, "ReasoningPart should be created" + assert "Thinking..." in ctx.reasoning_part.text + assert " more..." in ctx.reasoning_part.text @pytest.mark.asyncio @@ -84,72 +96,68 @@ async def test_multi_turn_thinking_creates_separate_parts(): This tests the fix for: "Multi-turn conversation thinking displayed in single block" Each thinking phase should be its own Part with its own ID. """ - # Create a mock MessageWithParts mock_msg = MagicMock() mock_msg.parts = [] - mock_state = MagicMock() + processor, ctx = _make_processor_and_ctx(mock_msg) - adapter = OpenCodeStreamAdapter( - state=mock_state, - session_id="test-session", - assistant_msg_id="msg-1", - assistant_msg=mock_msg, - working_dir=".", - ) - - events = [] + events: list[object] = [] # Simulate multi-turn conversation with thinking in each turn: # Turn 1: Thinking -> Text - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=0, part=ThinkingPart(content="First thinking...")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=0, part=ThinkingPart(content="First thinking...")) ) - ]) - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=" more thinking")) + ) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent( + index=0, delta=ThinkingPartDelta(content_delta=" more thinking") + ), ) - ]) + ) # End of thinking - text response starts - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=1, part=TextPart(content="First response")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=1, part=TextPart(content="First response")) ) - ]) + ) # Turn 2: Thinking -> Text (new turn, should be separate Part) - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=2, part=ThinkingPart(content="Second turn thinking...")) + events.extend( + await _process_event( + processor, + ctx, + PartStartEvent(index=2, part=ThinkingPart(content="Second turn thinking...")), ) - ]) - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=2, delta=ThinkingPartDelta(content_delta=" more")) + ) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent(index=2, delta=ThinkingPartDelta(content_delta=" more")), ) - ]) + ) # End of thinking - text response starts - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=3, delta=TextPartDelta(content_delta="Second response")) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent(index=3, delta=TextPartDelta(content_delta="Second response")), ) - ]) + ) # Turn 3: Thinking (should be third separate Part) - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=4, part=ThinkingPart(content="Third turn thinking...")) + events.extend( + await _process_event( + processor, + ctx, + PartStartEvent(index=4, part=ThinkingPart(content="Third turn thinking...")), ) - ]) + ) # Extract ReasoningParts from events reasoning_parts = [] @@ -174,7 +182,7 @@ async def test_multi_turn_thinking_creates_separate_parts(): # Assertions # We need to check that there are 3 unique reasoning phases (unique Part IDs) # Each thinking start creates a new Part, and deltas update the same Part - unique_reasoning_parts = {} + unique_reasoning_parts: dict[str, ReasoningPart] = {} for p in reasoning_parts: unique_reasoning_parts[p.id] = p @@ -222,44 +230,37 @@ async def test_single_thinking_phase_accumulates_correctly(): mock_msg = MagicMock() mock_msg.parts = [] - mock_state = MagicMock() + processor, ctx = _make_processor_and_ctx(mock_msg) - adapter = OpenCodeStreamAdapter( - state=mock_state, - session_id="test-session", - assistant_msg_id="msg-1", - assistant_msg=mock_msg, - working_dir=".", - ) - - events = [] + events: list[object] = [] # Single thinking phase with multiple deltas - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=0, part=ThinkingPart(content="Start ")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=0, part=ThinkingPart(content="Start ")) ) - ]) - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta="middle ")) + ) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta="middle ")), ) - ]) - events.extend([ - e - async for e in adapter._handle_event( - PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta="end")) + ) + events.extend( + await _process_event( + processor, + ctx, + PydanticPartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta="end")), ) - ]) + ) - # Verify accumulation via adapter context (not just events) + # Verify accumulation via context (not just events) # PartDeltaEvent yields deltas, not full parts, so check context directly - assert adapter.main_context.reasoning_part is not None, "ReasoningPart should be created" + assert ctx.reasoning_part is not None, "ReasoningPart should be created" # The content should be accumulated in the context - final_content = adapter.main_context.reasoning_part.text + final_content = ctx.reasoning_part.text expected = "Start middle end" assert final_content == expected, f"Expected '{expected}', got '{final_content}'" @@ -271,33 +272,23 @@ async def test_reasoning_part_gets_end_time_when_text_starts(): mock_msg.parts = [] mock_msg.update_part = MagicMock() - mock_state = MagicMock() + processor, ctx = _make_processor_and_ctx(mock_msg) - adapter = OpenCodeStreamAdapter( - state=mock_state, - session_id="test-session", - assistant_msg_id="msg-1", - assistant_msg=mock_msg, - working_dir=".", - ) - - events = [] + events: list[object] = [] # Thinking phase - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) ) - ]) + ) # Text starts - this should close out the reasoning part - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=1, part=TextPart(content="Response")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=1, part=TextPart(content="Response")) ) - ]) + ) # The reasoning part should have been updated with an end time reasoning_final_events = [ @@ -315,7 +306,7 @@ async def test_reasoning_part_gets_end_time_when_text_starts(): ) # The context should have cleared the reasoning_part reference - assert adapter.main_context.reasoning_part is None, ( + assert ctx.reasoning_part is None, ( "reasoning_part should be cleared from context after text starts" ) @@ -327,32 +318,22 @@ async def test_reasoning_part_gets_end_time_on_stream_complete(): mock_msg.parts = [] mock_msg.update_part = MagicMock() - mock_state = MagicMock() - - adapter = OpenCodeStreamAdapter( - state=mock_state, - session_id="test-session", - assistant_msg_id="msg-1", - assistant_msg=mock_msg, - working_dir=".", - ) + processor, ctx = _make_processor_and_ctx(mock_msg) - events = [] + events: list[object] = [] # Thinking phase only (no text follows) - events.extend([ - e - async for e in adapter._handle_event( - PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) + events.extend( + await _process_event( + processor, ctx, PartStartEvent(index=0, part=ThinkingPart(content="Thinking...")) ) - ]) + ) # Stream completes without any text starting - # Simulate what happens when _process_stream_complete is called from agentpool.messaging import ChatMessage chat_msg = ChatMessage(content="", role="assistant") - events.extend(list(adapter.processor._process_stream_complete(adapter.main_context, chat_msg))) + events.extend(list(processor._process_stream_complete(ctx, chat_msg))) # The reasoning part should have been finalized with an end time reasoning_final_events = [ @@ -370,6 +351,6 @@ async def test_reasoning_part_gets_end_time_on_stream_complete(): ) # The context should have cleared the reasoning_part reference - assert adapter.main_context.reasoning_part is None, ( + assert ctx.reasoning_part is None, ( "reasoning_part should be cleared from context after stream complete" ) diff --git a/tests/servers/opencode_server/test_route_discovery.py b/tests/servers/opencode_server/test_route_discovery.py index 0b7ed2e24..de7cd9579 100644 --- a/tests/servers/opencode_server/test_route_discovery.py +++ b/tests/servers/opencode_server/test_route_discovery.py @@ -63,7 +63,7 @@ def _server_state(tmp_path: Path) -> ServerState: agent._input_provider = None agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool agent.storage = storage_manager agent.get_available_models = AsyncMock(return_value=[]) agent.get_mcp_server_info = AsyncMock(return_value={}) diff --git a/tests/servers/opencode_server/test_server_lifecycle.py b/tests/servers/opencode_server/test_server_lifecycle.py index ac8b3ec1e..068b12453 100644 --- a/tests/servers/opencode_server/test_server_lifecycle.py +++ b/tests/servers/opencode_server/test_server_lifecycle.py @@ -52,7 +52,7 @@ def shared_agent(mock_env: Mock, mock_pool: Mock) -> Mock: agent._input_provider = None agent.agent_pool = mock_pool agent.host_context = mock_pool - mock_pool.pool = mock_pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = mock_pool # state.py resolves _pool via agent._agent_pool agent.storage = None return agent diff --git a/tests/servers/opencode_server/test_session_integration.py b/tests/servers/opencode_server/test_session_integration.py index d3abaab23..b829b553c 100644 --- a/tests/servers/opencode_server/test_session_integration.py +++ b/tests/servers/opencode_server/test_session_integration.py @@ -26,8 +26,8 @@ if TYPE_CHECKING: from collections.abc import AsyncIterator +from agentpool.lifecycle import RunState from agentpool.orchestrator.core import RunHandle, SessionPool -from agentpool.orchestrator.run import RunStatus from agentpool.sessions.models import SessionData from agentpool_server.opencode_server.input_provider import OpenCodeInputProvider from agentpool_server.opencode_server.state import ServerState @@ -611,7 +611,7 @@ def _create_turn_with_ctx(prompts: Any, run_ctx: Any, message_history: Any) -> A # Give the background task time to start and transition to running await asyncio.sleep(0.05) - assert run_handle._status == RunStatus.running + assert run_handle._run_state == RunState.RUNNING await integration.abort_session("test-session-011") diff --git a/tests/servers/opencode_server/test_session_scoped_consumer.py b/tests/servers/opencode_server/test_session_scoped_consumer.py index 4dee3bd8f..af2ecaaf1 100644 --- a/tests/servers/opencode_server/test_session_scoped_consumer.py +++ b/tests/servers/opencode_server/test_session_scoped_consumer.py @@ -40,7 +40,7 @@ def mock_agent_pool() -> Mock: pool.storage.load_session = AsyncMock(return_value=None) pool.storage.save_session = AsyncMock(return_value=None) - async def _mock_run_stream_once(*args: Any, **kwargs: Any) -> Any: + async def _mock_stream_events(*args: Any, **kwargs: Any) -> Any: """Yield a minimal run event sequence for testing.""" session_id = kwargs.get("session_id", "unknown") run_id = "run-mock-001" @@ -50,7 +50,7 @@ async def _mock_run_stream_once(*args: Any, **kwargs: Any) -> Any: ) mock_agent = Mock() - mock_agent._run_stream_once = _mock_run_stream_once + mock_agent._stream_events = _mock_stream_events mock_agent._input_provider = None mock_agent.conversation = Mock() mock_agent.conversation.add_chat_messages = Mock() @@ -92,7 +92,7 @@ def server_state(tmp_path: Any, mock_agent_pool: Mock) -> ServerState: agent.name = "test-agent" agent.storage = Mock() agent.host_context = mock_agent_pool - mock_agent_pool.pool = mock_agent_pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = mock_agent_pool # state.py resolves _pool via agent._agent_pool agent.env = Mock() return ServerState(working_dir=str(tmp_path), agent=agent) diff --git a/tests/servers/opencode_server/test_session_storage_load.py b/tests/servers/opencode_server/test_session_storage_load.py index a3392f55b..5a61c1aac 100644 --- a/tests/servers/opencode_server/test_session_storage_load.py +++ b/tests/servers/opencode_server/test_session_storage_load.py @@ -55,7 +55,7 @@ def mock_state_and_broadcast( agent.load_session = AsyncMock(return_value=session_data) # Explicitly set agent_pool so the SessionPool cold-load path is skipped agent.host_context = Mock() - agent.host_context.pool = agent.host_context # state.py resolves _pool via _ctx.pool + agent._agent_pool = agent.host_context # state.py resolves _pool via agent._agent_pool agent.host_context.session_pool = None # agent.conversation.chat_messages must be iterable (empty for cold load test) diff --git a/tests/servers/opencode_server/test_skill_autocomplete.py b/tests/servers/opencode_server/test_skill_autocomplete.py index 63f479902..358098d9b 100644 --- a/tests/servers/opencode_server/test_skill_autocomplete.py +++ b/tests/servers/opencode_server/test_skill_autocomplete.py @@ -125,7 +125,7 @@ async def test_command_endpoint_skill_provider_fallback( mock_resolver.list_providers = MagicMock(return_value=["test-provider"]) mock_resolver.get_provider = MagicMock(return_value=mock_provider) - mock_agent.host_context.pool.skill_resolver = mock_resolver # type: ignore[attr-defined] + mock_agent._agent_pool.skill_resolver = mock_resolver # type: ignore[attr-defined] mock_agent.host_context.skill_provider = None # type: ignore[attr-defined] # Mock empty MCP prompts @@ -231,7 +231,7 @@ async def test_command_endpoint_skill_bridge_with_provider_for_virtual_skills( mock_resolved.load_instructions = MagicMock(return_value="Virtual instructions with $1") mock_resolver = MagicMock() mock_resolver.resolve = AsyncMock(return_value=mock_resolved) - mock_agent.host_context.pool.skill_resolver = mock_resolver # type: ignore[attr-defined] + mock_agent._agent_pool.skill_resolver = mock_resolver # type: ignore[attr-defined] mock_agent.host_context.skill_provider = None # type: ignore[attr-defined] # Mock empty MCP prompts diff --git a/tests/servers/opencode_server/test_stream_adapter_event_feed.py b/tests/servers/opencode_server/test_stream_adapter_event_feed.py index abb40a21e..bddeaadba 100644 --- a/tests/servers/opencode_server/test_stream_adapter_event_feed.py +++ b/tests/servers/opencode_server/test_stream_adapter_event_feed.py @@ -15,9 +15,9 @@ import pytest from agentpool.agents.events import StreamCompleteEvent +from agentpool.lifecycle import RunOutcome, RunState from agentpool.messaging import ChatMessage from agentpool.orchestrator.core import EventBus -from agentpool.orchestrator.run import RunStatus from agentpool.utils import identifiers as identifier from agentpool.utils.time_utils import now_ms from agentpool_server.opencode_server.models import ( @@ -115,14 +115,15 @@ def mock_agent_with_event_bus(tmp_project_dir): # RunHandle whose complete_event we control from the test run_handle = Mock() - run_handle.status = RunStatus.completed + run_handle._run_state = RunState.DONE + run_handle.outcome = RunOutcome.COMPLETED run_handle.complete_event = asyncio.Event() session_pool.receive_request = AsyncMock(return_value=run_handle) pool.session_pool = session_pool agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool return agent, run_handle, event_bus diff --git a/tests/servers/opencode_server/test_subagent_completion_red_flags.py b/tests/servers/opencode_server/test_subagent_completion_red_flags.py index 9db617ae0..1cd3368ee 100644 --- a/tests/servers/opencode_server/test_subagent_completion_red_flags.py +++ b/tests/servers/opencode_server/test_subagent_completion_red_flags.py @@ -113,6 +113,6 @@ async def test_background_task_inject_prompt_wakes_lead_agent( ) # Verify the fallback path for shared agents (no fixed session_id) - assert "agent_pool" in source, ( - "inject_prompt must check agent_pool as fallback for shared agents" + assert "find_sessions_by_agent_name" in source, ( + "inject_prompt must use find_sessions_by_agent_name as fallback for shared agents" ) diff --git a/tests/servers/opencode_server/test_title_generation_nonblocking.py b/tests/servers/opencode_server/test_title_generation_nonblocking.py index 2986038d1..23cc3aa82 100644 --- a/tests/servers/opencode_server/test_title_generation_nonblocking.py +++ b/tests/servers/opencode_server/test_title_generation_nonblocking.py @@ -61,7 +61,7 @@ def _make_state(tmp_path: Any) -> ServerState: agent.agent_pool = pool agent.host_context = pool - pool.pool = pool # state.py resolves _pool via _ctx.pool + agent._agent_pool = pool # state.py resolves _pool via agent._agent_pool agent.storage = storage_mgr env = Mock()