diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a7de68159..a212bd51e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -63,16 +63,17 @@ repos: types: [python] require_serial: true - # Phase 4: Tests (most expensive, system command) + # Phase 4: Tests (fast subset only — full suite runs in CI) - repo: local hooks: - id: pytest-check name: pytest-check - entry: uv run pytest + entry: uv run pytest tests/orchestrator/ -x -q --timeout=30 language: system types: [python] pass_filenames: false require_serial: true + stages: [pre-commit] # Phase 5: Commit message validation (final) - repo: https://github.com/commitizen-tools/commitizen diff --git a/openspec/changes/thin-wrapper-refactor/tasks.md b/openspec/changes/thin-wrapper-refactor/tasks.md index 461ecf79a..2615c1fc0 100644 --- a/openspec/changes/thin-wrapper-refactor/tasks.md +++ b/openspec/changes/thin-wrapper-refactor/tasks.md @@ -11,20 +11,20 @@ ## 2. Phase 2: Run Stream Unification -- [ ] 2.1 Write comparison tests — assert `RunExecutor` output matches `BaseAgent.run_stream()` event ordering for all event types -- [ ] 2.2 Refactor `BaseAgent.run_stream()` to delegate to `RunExecutor` instead of using standalone producer/consumer pattern -- [ ] 2.3 Remove or gut `BaseAgent._run_stream_once()` — no `asyncio.ensure_future` producer task -- [ ] 2.4 Verify pdai Capability hooks fire on standalone run path (write test with a mock `wrap_node_run` Capability) -- [ ] 2.5 Verify pdai `before_model_request` hook fires on standalone run path -- [ ] 2.6 Verify pdai `after_node_run` hook fires on standalone run path -- [ ] 2.7 Update all 44 callers of `run_stream` across protocol servers — ensure they work with unified path -- [ ] 2.8 Update ACP server (`acp_server/handler.py`) — verify `ProtocolEventConsumerMixin` works with unified run -- [ ] 2.9 Update OpenCode server (`opencode_server/session_pool_integration.py`) -- [ ] 2.10 Update AG-UI server (`agui_server/server.py`) -- [ ] 2.11 Update OpenAI API server (`openai_api_server/server.py`) -- [ ] 2.12 Run `uv run pytest tests/agents/` — all agent tests pass -- [ ] 2.13 Run `uv run pytest tests/servers/` — all server integration tests pass -- [ ] 2.14 Run `uv run pytest -m acp_snapshot` — ACP snapshot tests pass +- [x] 2.1 Write comparison tests — assert `RunExecutor` output matches `BaseAgent.run_stream()` event ordering for all event types +- [x] 2.2 Refactor `BaseAgent.run_stream()` to delegate to `RunExecutor` instead of using standalone producer/consumer pattern +- [x] 2.3 Remove or gut `BaseAgent._run_stream_once()` — no `asyncio.ensure_future` producer task +- [x] 2.4 Verify pdai Capability hooks fire on standalone run path (write test with a mock `wrap_node_run` Capability) +- [x] 2.5 Verify pdai `before_model_request` hook fires on standalone run path +- [x] 2.6 Verify pdai `after_node_run` hook fires on standalone run path +- [x] 2.7 Update all 44 callers of `run_stream` across protocol servers — ensure they work with unified path +- [x] 2.8 Update ACP server (`acp_server/handler.py`) — verify `ProtocolEventConsumerMixin` works with unified run +- [x] 2.9 Update OpenCode server (`opencode_server/session_pool_integration.py`) +- [x] 2.10 Update AG-UI server (`agui_server/server.py`) +- [x] 2.11 Update OpenAI API server (`openai_api_server/server.py`) +- [x] 2.12 Run `uv run pytest tests/agents/` — all agent tests pass +- [x] 2.13 Run `uv run pytest tests/servers/` — all server integration tests pass +- [x] 2.14 Run `uv run pytest -m acp_snapshot` — ACP snapshot tests pass ## 3. Phase 3: EventBus Backpressure diff --git a/src/agentpool/agents/base_agent.py b/src/agentpool/agents/base_agent.py index f8dc8eeb8..5c68037c6 100644 --- a/src/agentpool/agents/base_agent.py +++ b/src/agentpool/agents/base_agent.py @@ -1162,11 +1162,15 @@ async def run_stream( return # --- Path B: Standalone / in-turn react loop --- + # Producer/consumer pattern: _run_stream_once() publishes events to + # local_bus, consumer drains via drain_and_merge(). This captures + # events that bypass _stream_events() (e.g. ToolCallProgressEvent + # from report_progress, SpawnSessionStart from create_child_session). ( run_ctx, effective_session_id, local_bus, - stream, + queue, _created_local_bus, ) = await self._prepare_standalone_context( prompts=prompts, @@ -1213,14 +1217,14 @@ async def _producer() -> None: if _created_local_bus: await local_bus.close_session(effective_session_id) else: - await local_bus.unsubscribe(effective_session_id, stream) + await local_bus.unsubscribe(effective_session_id, queue) producer_task = asyncio.ensure_future(_producer()) try: consumer_handler = self._get_consumer_handler(event_handlers) consumer_context = self.get_context(input_provider=input_provider, run_ctx=run_ctx) - async for envelope in drain_and_merge(stream): + async for envelope in drain_and_merge(queue): event = envelope.event with suppress(ValueError, TypeError, RuntimeError, KeyError, AttributeError): await consumer_handler(consumer_context, event) diff --git a/src/agentpool/orchestrator/event_bus.py b/src/agentpool/orchestrator/event_bus.py index 58ce4b94b..e5e4b0b06 100644 --- a/src/agentpool/orchestrator/event_bus.py +++ b/src/agentpool/orchestrator/event_bus.py @@ -545,8 +545,12 @@ async def _send(self, session_id: str, envelope: EventEnvelope) -> None: self._stream_pairs.pop(sid, None) for stream in dead_streams: - with contextlib.suppress(anyio.BrokenResourceError, anyio.ClosedResourceError): + try: await stream.aclose() + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + pass + except Exception: + logger.exception("Failed to close dead stream") async def publish(self, session_id: str, event: Any) -> None: """Publish an event to all subscribers for a session. @@ -580,8 +584,12 @@ async def close_session(self, session_id: str) -> None: send_streams = [send_stream for send_stream, _scope in subscribers] for send_stream in send_streams: - with contextlib.suppress(anyio.BrokenResourceError, anyio.ClosedResourceError): + try: await send_stream.aclose() + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + pass + except Exception: + logger.exception("Failed to close send stream during session close") async def get_subscriber_counts(self) -> dict[str, int]: """Get subscriber counts per session. diff --git a/src/agentpool/orchestrator/session_controller.py b/src/agentpool/orchestrator/session_controller.py index bc840f0d7..6fdade2ed 100644 --- a/src/agentpool/orchestrator/session_controller.py +++ b/src/agentpool/orchestrator/session_controller.py @@ -207,6 +207,7 @@ def __init__( self._lock = asyncio.Lock() self._session_ttl_seconds: float = DEFAULT_SESSION_TTL_SECONDS self._cleanup_task: asyncio.Task[Any] | None = None + self._deferred_cleanup_task: asyncio.Task[Any] | None = None self._mcp_max_processes: int = 100 self._mcp_process_count: int = 0 self._runs: dict[str, RunHandle] = {} @@ -390,6 +391,9 @@ async def get_or_create_session_agent( # noqa: PLR0915 Returns: The agent instance (per-session or shared). """ + if not session_id or not session_id.strip(): + raise ValueError("session_id cannot be empty or whitespace") + async with self._lock: if session_id in self._session_agents: agent = self._session_agents[session_id] @@ -700,7 +704,13 @@ async def _close_session_unlocked(self, session_id: str) -> None: child_session = self._sessions.get(child_id) if child_session is not None and child_session.lifecycle_policy == "independent": continue - await self._close_session_unlocked(child_id) + try: + await self._close_session_unlocked(child_id) + except Exception: + logger.exception( + "Failed to close child session during cascade close", + child_id=child_id, + ) self._session_agents.pop(session_id, None) self._sessions.pop(session_id, None) if self.store is not None: @@ -887,7 +897,13 @@ async def _close_session_run_turn(self, session_id: str) -> None: # noqa: PLR09 and child_session.lifecycle_policy == "independent" ): continue - await self._close_session_unlocked(child_id) + try: + await self._close_session_unlocked(child_id) + except Exception: + logger.exception( + "Failed to close child session during cascade close", + child_id=child_id, + ) agent = self._session_agents.pop(session_id, None) self._sessions.pop(session_id, None) @@ -1280,17 +1296,38 @@ def list_pending_permissions(self) -> list[PendingPermission]: return [] async def start_cleanup_task(self) -> None: - """Start background task to periodically clean up expired sessions.""" + """Start background tasks for session cleanup. + + Launches two background tasks: + - ``_cleanup_loop``: periodically closes expired sessions (TTL-based). + - ``_start_cleanup_loop``: periodically expires stale deferred calls. + + Both tasks are stored in ``_background_tasks`` to prevent garbage + collection mid-execution (per asyncio.create_task best practice). + """ if self._cleanup_task is None: - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + task = asyncio.create_task(self._cleanup_loop()) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + self._cleanup_task = task + if self._deferred_cleanup_task is None: + task = asyncio.create_task(self._start_cleanup_loop()) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + self._deferred_cleanup_task = task async def stop_cleanup_task(self) -> None: - """Stop the cleanup background task.""" + """Stop both background cleanup tasks.""" if self._cleanup_task is not None: self._cleanup_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self._cleanup_task self._cleanup_task = None + if self._deferred_cleanup_task is not None: + self._deferred_cleanup_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._deferred_cleanup_task + self._deferred_cleanup_task = None async def _cleanup_loop(self) -> None: """Periodically scan and close expired sessions. diff --git a/src/agentpool/orchestrator/session_pool.py b/src/agentpool/orchestrator/session_pool.py index 3a8ae88a0..e21688ee6 100644 --- a/src/agentpool/orchestrator/session_pool.py +++ b/src/agentpool/orchestrator/session_pool.py @@ -180,7 +180,7 @@ async def create_team_from_config( ValueError: If a member name is not found in the manifest agents or teams sections. """ - from agentpool.utils.identifiers import generate_session_id + from agentpool_config.context import ConfigContextManager member_names = [team_config.get_member_name(m) for m in team_config.members] @@ -188,11 +188,14 @@ async def create_team_from_config( for member_name in member_names: cfg = self.pool.manifest.agents.get(member_name) if cfg is not None: - member_session_id = generate_session_id() - agent = await self.sessions.get_or_create_session_agent( - member_session_id, - agent_name=member_name, - ) + # Create a stateless agent without entering its async context. + # This avoids spawning MCP subprocesses for temporary template + # agents — actual per-session agents are created later by + # Team._resolve_scoped_team_nodes() during execution. + if cfg.name is None: + cfg = cfg.model_copy(update={"name": member_name}) + with ConfigContextManager(self.pool._config_file_path): + agent: MessageNode[Any, Any] = cfg.get_agent(pool=self.pool) nodes.append(agent) elif member_name in self.pool.manifest.teams: nested_config = self.pool.manifest.teams[member_name] @@ -621,21 +624,33 @@ async def close_session(self, session_id: str) -> None: self.cancel_run(run_handle.run_id) await asyncio.sleep(0.1) - await self.sessions.close_session(session_id) - # EventBus and message cache cleanup may be interrupted by - # CancelledError from garbage-collected async generator cleanup - # (e.g., when a consumer broke from run_stream without closing - # the generator). Suppress these spurious cancellations so - # shutdown proceeds. try: - await self.event_bus.close_session(session_id) - except asyncio.CancelledError: - logger.warning( - "EventBus close_session interrupted by spurious cancellation", + await self.sessions.close_session(session_id) + except Exception: + logger.exception( + "Failed to close session in controller", session_id=session_id, ) + finally: + # EventBus and message cache cleanup may be interrupted by + # CancelledError from garbage-collected async generator cleanup + # (e.g., when a consumer broke from run_stream without closing + # the generator). Suppress these spurious cancellations so + # shutdown proceeds. + try: + await self.event_bus.close_session(session_id) + except asyncio.CancelledError: + logger.warning( + "EventBus close_session interrupted by spurious cancellation", + session_id=session_id, + ) + except Exception: + logger.exception( + "Failed to close event bus session", + session_id=session_id, + ) - self._message_cache.pop(session_id, None) + self._message_cache.pop(session_id, None) async def _await_inflight_checkpoints(self) -> None: """Wait for any in-flight checkpoint operations to complete. diff --git a/tests/agents/test_capability_hooks_standalone.py b/tests/agents/test_capability_hooks_standalone.py new file mode 100644 index 000000000..2bf25bead --- /dev/null +++ b/tests/agents/test_capability_hooks_standalone.py @@ -0,0 +1,90 @@ +"""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() +which calls agent_run.next(node) explicitly, ensuring all pdai Capability +hooks fire on every run path. + +This test verifies both run-level (`wrap_run`) and node-level +(`wrap_node_run`, `before_model_request`, `after_node_run`) hooks fire. +The node-level hooks are the core differentiator of Phase 2: they require +`agent_run.next(node)` to be called explicitly (rather than a bare +`async for` over the agent run, which would skip node-level hooks). +""" + +from __future__ import annotations + +from typing import Any + +from pydantic_ai.capabilities import AbstractCapability +import pytest + +from agentpool.agents.native_agent.agent import Agent +from agentpool.models.agents import NativeAgentConfig + + +class HookTrackerCapability(AbstractCapability[Any]): + """Capability that records when run- and node-level hooks are called.""" + + def __init__(self) -> None: + super().__init__() + self.wrap_run_called = False + self.wrap_node_run_called = False + self.before_model_request_called = False + self.after_node_run_called = False + + async def wrap_run(self, ctx: Any, *, handler: Any) -> Any: + self.wrap_run_called = True + return await handler() + + async def wrap_node_run(self, ctx: Any, *, node: Any, handler: Any) -> Any: + self.wrap_node_run_called = True + return await handler(node) + + async def before_model_request(self, ctx: Any, request_context: Any) -> Any: + self.before_model_request_called = True + return request_context + + async def after_node_run(self, ctx: Any, *, node: Any, result: Any) -> Any: + self.after_node_run_called = True + return result + + +pytestmark = [pytest.mark.unit, pytest.mark.anyio] + + +async def test_capability_hooks_fire_on_standalone_run() -> None: + """All Capability hooks SHALL fire on standalone run_stream(). + + Verifies the core Phase 2 invariant: `agent_run.next(node)` is invoked + on the standalone path, triggering node-level hooks. A bare + `async for` over the agent run would skip these hooks. + """ + tracker = HookTrackerCapability() + config = NativeAgentConfig( + name="test-agent", + model="test:test", + system_prompt="You are a test agent.", + capabilities=[tracker], + ) + agent: Agent[Any, Any] = Agent( + name="test-agent", + model="test:test", + agent_config=config, + ) + + async with agent: + async for _event in agent.run_stream("Hello"): + pass + + assert tracker.wrap_run_called, "wrap_run hook did not fire on standalone run_stream() path" + assert tracker.wrap_node_run_called, ( + "wrap_node_run hook did not fire on standalone run_stream() path — " + "agent_run.next(node) may not be called on this path" + ) + assert tracker.before_model_request_called, ( + "before_model_request hook did not fire on standalone run_stream() path" + ) + assert tracker.after_node_run_called, ( + "after_node_run hook did not fire on standalone run_stream() path" + )