From d119659d1799d9c77a482fd94f450bc50f9ac245 Mon Sep 17 00:00:00 2001 From: rogeriomoura Date: Mon, 15 Jun 2026 08:04:07 -0500 Subject: [PATCH] feat(subagents): add invoke_agents batch tool for deterministic parallel fan-out Introduce a single `invoke_agents` tool that fans a list of sub-agent tasks out internally via `asyncio.gather`, instead of relying on the model/provider to honour `parallel_tool_calls` across N separate `invoke_agent` calls. Why this beats N x invoke_agent: - Deterministic parallelism: Code Puppy owns the concurrency, not the provider. Guaranteed to actually run in parallel. - One reviewable tool call instead of N scattered ones. - Order-preserved results (zip(tasks, results)). - Failure isolation: return_exceptions=True + normalization means one bad sub-agent becomes a structured `error` field, siblings unaffected. Supporting changes: - messaging/bus.py: move session context into a ContextVar with set/reset token semantics so concurrently-running sub-agents stay isolated and never clobber a sibling's session id. Export reset_session_context. - SubAgentTask gains an optional per-task `model_name` override (omit for the agent's configured model), closing the "parallel + custom models" gap. A BeforeValidator tolerates providers that double-encode the `tasks` arg as a JSON string. - config.get_allow_parallel_tool_calls(): opt-in flag to permit parallel tool calls outside yolo mode; model_factory honours it. - run_stats.py: render the batch (list) result as an OK/FAIL summary. - Register invoke_agents in the tool registry and the code-puppy agent. Tests: per-task model override plumbing, blank-normalization contract, ContextVar session isolation, and batch failure isolation. 129 passing. --- code_puppy/agents/agent_code_puppy.py | 1 + code_puppy/agents/run_stats.py | 28 ++- code_puppy/config.py | 19 ++ code_puppy/messaging/__init__.py | 2 + code_puppy/messaging/bus.py | 63 ++++- code_puppy/model_factory.py | 14 +- code_puppy/tools/__init__.py | 2 + code_puppy/tools/agent_tools.py | 2 + code_puppy/tools/subagent_invocation.py | 149 +++++++++++- tests/messaging/test_bus.py | 11 +- tests/test_agent_tools_coverage.py | 302 +++++++++++++++++++++--- tests/test_messaging_bus.py | 5 +- tests/test_messaging_init.py | 1 + 13 files changed, 536 insertions(+), 63 deletions(-) diff --git a/code_puppy/agents/agent_code_puppy.py b/code_puppy/agents/agent_code_puppy.py index 4f1eb9381..6cce8f2b9 100644 --- a/code_puppy/agents/agent_code_puppy.py +++ b/code_puppy/agents/agent_code_puppy.py @@ -25,6 +25,7 @@ def get_available_tools(self) -> list[str]: return [ "list_agents", "invoke_agent", + "invoke_agents", "list_files", "read_file", "grep", diff --git a/code_puppy/agents/run_stats.py b/code_puppy/agents/run_stats.py index 02d98a37a..c555c2323 100644 --- a/code_puppy/agents/run_stats.py +++ b/code_puppy/agents/run_stats.py @@ -276,7 +276,9 @@ def format_conversation_stats( # Tool names whose response has already been rendered inline (streamed or # via display_non_streamed_result) and should NOT be dumped again. -_ALREADY_RENDERED_TOOLS = frozenset({"invoke_agent", "invoke_agent_with_model"}) +_ALREADY_RENDERED_TOOLS = frozenset( + {"invoke_agent", "invoke_agent_with_model", "invoke_agents"} +) # Tools whose output is already rendered by the rich_renderer via MessageBus # messages (FileContentMessage, DiffMessage, GrepResultMessage, etc.). @@ -415,9 +417,31 @@ def _render_high_mode_agent_result( result: Any, dur_str: str, ) -> None: - """Compact one-liner for invoke_agent / invoke_agent_with_model results.""" + """Compact one-liner for invoke_agent(s) results. + + ``invoke_agent`` / ``invoke_agent_with_model`` return a single + ``AgentInvokeOutput``; ``invoke_agents`` returns a *list* of them. Handle + both so the batch fan-out gets a useful OK/FAIL summary instead of "?". + """ from rich.markup import escape as _esc + # Batch result (invoke_agents): summarise per-agent success/failure. + if isinstance(result, (list, tuple)): + total = len(result) + ok = sum(1 for r in result if getattr(r, "error", None) is None) + failed = total - ok + status = ( + f"[green]{ok} OK[/green]" + if failed == 0 + else f"[green]{ok} OK[/green] / [red]{failed} FAIL[/red]" + ) + names = ", ".join(getattr(r, "agent_name", None) or "?" for r in result) + console.print( + f"[dim] \u21a9 {_esc(tool_name)} returned ({dur_str} ms): " + f"{total} sub-agent(s) {status} \u2014 {_esc(names)}[/dim]" + ) + return + agent_name = getattr(result, "agent_name", None) or "?" error = getattr(result, "error", None) response = getattr(result, "response", None) diff --git a/code_puppy/config.py b/code_puppy/config.py index bb643fdff..864848076 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -1198,6 +1198,25 @@ def get_yolo_mode(): return True +def get_allow_parallel_tool_calls(): + """ + Checks puppy.cfg for 'allow_parallel_tool_calls' (case-insensitive value). + Defaults to False if not set. + + When off (default), parallel tool calls are disabled outside yolo mode so + the user can review each call sequentially. Turn this on to let the model + fan out tool calls (including parallel sub-agent invocations) even when + yolo mode is off. Note: the ``invoke_agents`` batch tool fans out + internally via one reviewable tool call and does NOT require this flag. + Allowed values for ON: 1, '1', 'true', 'yes', 'on' (case-insensitive). + """ + true_vals = {"1", "true", "yes", "on"} + cfg_val = get_value("allow_parallel_tool_calls") + if cfg_val is not None: + return str(cfg_val).strip().lower() in true_vals + return False + + def get_safety_permission_level(): """ Checks puppy.cfg for 'safety_permission_level' (case-insensitive in value only). diff --git a/code_puppy/messaging/__init__.py b/code_puppy/messaging/__init__.py index 64af8fef3..119e7c07f 100644 --- a/code_puppy/messaging/__init__.py +++ b/code_puppy/messaging/__init__.py @@ -47,6 +47,7 @@ get_message_bus, get_session_context, reset_message_bus, + reset_session_context, set_session_context, ) from .bus import emit as bus_emit # Convenience functions (new API versions) @@ -247,6 +248,7 @@ "reset_message_bus", # Session context "set_session_context", + "reset_session_context", "get_session_context", # New API convenience functions (prefixed to avoid collision) "bus_emit", diff --git a/code_puppy/messaging/bus.py b/code_puppy/messaging/bus.py index f53206ad8..c81b18dec 100644 --- a/code_puppy/messaging/bus.py +++ b/code_puppy/messaging/bus.py @@ -35,6 +35,7 @@ import asyncio import queue import threading +from contextvars import ContextVar, Token from typing import Any, Dict, List, Optional, Tuple from uuid import uuid4 @@ -57,6 +58,16 @@ UserInputRequest, ) +# Session context lives in a ContextVar (NOT a plain shared attribute) so that +# concurrently-running sub-agents each get an isolated copy. ``asyncio`` copies +# the current context when a Task is created, so every fanned-out sub-agent +# invocation sees its own session id and can't clobber a sibling's. This mirrors +# the per-session browser isolation in ``browser_manager`` — same pattern, one +# source of truth for "who am I right now". +_session_context_var: ContextVar[Optional[str]] = ContextVar( + "code_puppy_session_id", default=None +) + class MessageBus: """Central coordinator for bidirectional Agent <-> UI communication. @@ -89,8 +100,8 @@ def __init__(self, maxsize: int = 1000) -> None: # Request/Response correlation: prompt_id → Future (for async usage) self._pending_requests: Dict[str, asyncio.Future[Any]] = {} - # Session context for multi-agent tracking - self._current_session_id: Optional[str] = None + # Session context for multi-agent tracking lives in a module-level + # ContextVar (see top of file) so parallel sub-agents stay isolated. # ========================================================================= # Outgoing Messages (Agent → UI) @@ -106,10 +117,13 @@ def emit(self, message: AnyMessage) -> None: Args: message: The message to emit. """ - # Auto-tag message with current session if not already set + # Auto-tag message with current session if not already set. The session + # id comes from the ContextVar so it's correct per-task even when many + # sub-agents run concurrently. + current_session_id = _session_context_var.get() with self._lock: - if message.session_id is None and self._current_session_id is not None: - message.session_id = self._current_session_id + if message.session_id is None and current_session_id is not None: + message.session_id = current_session_id if not self._has_active_renderer: self._startup_buffer.append(message) @@ -181,17 +195,32 @@ def emit_shell_line(self, line: str, stream: str = "stdout") -> None: # Session Context (Multi-Agent Tracking) # ========================================================================= - def set_session_context(self, session_id: Optional[str]) -> None: + def set_session_context(self, session_id: Optional[str]) -> Token: """Set the current session context for auto-tagging messages. When set, all messages emitted via emit() will be automatically tagged with this session_id unless they already have one set. + Returns a ``Token`` that callers MUST hand back to + ``reset_session_context`` (typically in a ``finally``) so the previous + value is restored without the racey "save the old value, restore it + later" dance that breaks under concurrency. + Args: session_id: The session ID to tag messages with, or None to clear. + + Returns: + A contextvars ``Token`` for restoring the prior context. """ - with self._lock: - self._current_session_id = session_id + return _session_context_var.set(session_id) + + def reset_session_context(self, token: Token) -> None: + """Restore the session context to the value captured by ``token``. + + Args: + token: The ``Token`` returned by a prior ``set_session_context``. + """ + _session_context_var.reset(token) def get_session_context(self) -> Optional[str]: """Get the current session context. @@ -199,8 +228,7 @@ def get_session_context(self) -> Optional[str]: Returns: The current session_id, or None if not set. """ - with self._lock: - return self._current_session_id + return _session_context_var.get() # ========================================================================= # User Input Requests (Agent waits for UI response) @@ -594,9 +622,17 @@ def emit_shell_line(line: str, stream: str = "stdout") -> None: get_message_bus().emit_shell_line(line, stream) -def set_session_context(session_id: Optional[str]) -> None: - """Set the session context on the global bus.""" - get_message_bus().set_session_context(session_id) +def set_session_context(session_id: Optional[str]) -> Token: + """Set the session context on the global bus. + + Returns a ``Token`` to pass to ``reset_session_context`` when done. + """ + return get_message_bus().set_session_context(session_id) + + +def reset_session_context(token: Token) -> None: + """Restore the session context using a token from ``set_session_context``.""" + get_message_bus().reset_session_context(token) def get_session_context() -> Optional[str]: @@ -624,5 +660,6 @@ def get_session_context() -> Optional[str]: "emit_shell_line", # Session context "set_session_context", + "reset_session_context", "get_session_context", ] diff --git a/code_puppy/model_factory.py b/code_puppy/model_factory.py index d4b12f9a5..d312022f8 100644 --- a/code_puppy/model_factory.py +++ b/code_puppy/model_factory.py @@ -24,7 +24,12 @@ from . import callbacks from .claude_cache_client import ClaudeCacheAsyncClient, patch_anthropic_client_messages -from .config import EXTRA_MODELS_FILE, get_value, get_yolo_mode +from .config import ( + EXTRA_MODELS_FILE, + get_allow_parallel_tool_calls, + get_value, + get_yolo_mode, +) from .http_utils import create_async_client, get_cert_bundle_path, get_http2 from .provider_identity import ( make_anthropic_provider, @@ -168,8 +173,11 @@ def make_model_settings( effective_settings = get_effective_model_settings(model_name) model_settings_dict.update(effective_settings) - # Disable parallel tool calls when yolo_mode is off (sequential so user can review each call) - if not get_yolo_mode(): + # Disable parallel tool calls when yolo_mode is off so the user can review + # each call sequentially -- UNLESS they've explicitly opted into parallel + # fan-out via allow_parallel_tool_calls (needed for Claude-Code-style + # parallel sub-agent spawning outside yolo mode). + if not get_yolo_mode() and not get_allow_parallel_tool_calls(): model_settings_dict["parallel_tool_calls"] = False # Default to clear_thinking=False for GLM-4.7 and GLM-5 models (preserved thinking) diff --git a/code_puppy/tools/__init__.py b/code_puppy/tools/__init__.py index be4298a9e..810189221 100644 --- a/code_puppy/tools/__init__.py +++ b/code_puppy/tools/__init__.py @@ -4,6 +4,7 @@ from code_puppy.tools.subagent_invocation import ( register_invoke_agent, register_invoke_agent_with_model, + register_invoke_agents, ) from code_puppy.tools.ask_user_question import register_ask_user_question @@ -94,6 +95,7 @@ "list_agents": register_list_agents, "invoke_agent": register_invoke_agent, "invoke_agent_with_model": register_invoke_agent_with_model, + "invoke_agents": register_invoke_agents, "list_available_models": register_list_available_models, # File Operations "list_files": register_list_files, diff --git a/code_puppy/tools/agent_tools.py b/code_puppy/tools/agent_tools.py index ef6a8a8d5..4cc2e536d 100644 --- a/code_puppy/tools/agent_tools.py +++ b/code_puppy/tools/agent_tools.py @@ -273,6 +273,7 @@ def list_agents(context: RunContext) -> ListAgentsOutput: _active_subagent_tasks, register_invoke_agent, register_invoke_agent_with_model, + register_invoke_agents, ) __all__ = [ @@ -290,6 +291,7 @@ def list_agents(context: RunContext) -> ListAgentsOutput: "_validate_session_id", "register_invoke_agent", "register_invoke_agent_with_model", + "register_invoke_agents", "register_list_agents", "set_session_context", ] diff --git a/code_puppy/tools/subagent_invocation.py b/code_puppy/tools/subagent_invocation.py index ec298ead5..45b9ebe42 100644 --- a/code_puppy/tools/subagent_invocation.py +++ b/code_puppy/tools/subagent_invocation.py @@ -2,11 +2,13 @@ import asyncio import inspect +import json import traceback from contextlib import AsyncExitStack from functools import partial -from typing import Set +from typing import Annotated, Set +from pydantic import BaseModel, BeforeValidator, Field from pydantic_ai import Agent, RunContext, UsageLimits from code_puppy.callbacks import ( @@ -22,7 +24,7 @@ emit_info, emit_success, get_message_bus, - get_session_context, + reset_session_context, set_session_context, ) from code_puppy.tools.agent_tools import ( @@ -113,9 +115,11 @@ async def _invoke_agent_impl( ) ) - # Save current session context and set the new one for this sub-agent - previous_session_id = get_session_context() - set_session_context(session_id) + # Set the session context for this sub-agent. ``set_session_context`` + # returns a ContextVar token; we restore via ``reset_session_context`` in + # the ``finally`` below. This is reentrant + parallel-safe — no clobbering + # between concurrently-running sibling sub-agents. + session_token = set_session_context(session_id) # Set browser session for browser tools (qa-kitten, etc.) # This allows parallel agent invocations to each have their own browser @@ -404,8 +408,8 @@ async def _invoke_agent_impl( ) finally: - # Restore the previous session context - set_session_context(previous_session_id) + # Restore the previous session context via the token + reset_session_context(session_token) # Reset browser session context from code_puppy.tools.browser.browser_manager import ( _browser_session_var, @@ -504,3 +508,134 @@ async def invoke_agent_with_model( ) return invoke_agent_with_model + + +class SubAgentTask(BaseModel): + """A single sub-agent invocation within a parallel batch.""" + + agent_name: str = Field(description="Name of the sub-agent to invoke.") + prompt: str = Field(description="Task prompt for this sub-agent.") + session_id: str | None = Field( + default=None, + description="Optional kebab-case session id for continuing memory.", + ) + model_name: str | None = Field( + default=None, + description=( + "Optional per-task model alias override for this invocation only. " + "Omit to use the sub-agent's configured model (the common case)." + ), + ) + + def resolved_model_name(self) -> str | None: + """Return the override model, treating blank strings as 'no override'. + + Mirrors the normalization in ``invoke_agent_with_model`` so a whitespace + alias never leaks through as a bogus override. + """ + if self.model_name is None: + return None + normalized = self.model_name.strip() + return normalized or None + + +def _coerce_tasks(value): + """Tolerate a JSON-string ``tasks`` arg from providers that double-encode it. + + pydantic v2 validates tool args with ``validate_python``, which will NOT + parse a JSON *string* into a *list* (only ``validate_json`` does). Some + LLM providers serialise nested array/object tool arguments as a JSON + string, so ``tasks`` arrives as ``'[{...}]'`` rather than ``[{...}]`` and + pydantic raises ``list_type`` before any sub-agent runs. Parse the string + back into structured data here; pass real lists straight through. + """ + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError: + # Let pydantic raise its normal, informative validation error. + return value + return value + + +# Tool-arg type for the batch fan-out. The BeforeValidator makes us robust to +# providers that hand ``tasks`` over as a double-encoded JSON string. +TasksArg = Annotated[list[SubAgentTask], BeforeValidator(_coerce_tasks)] + + +def register_invoke_agents(agent): + """Register the parallel fan-out ``invoke_agents`` tool. + + Unlike emitting several ``invoke_agent`` calls in one turn (which depends on + the model/provider honouring ``parallel_tool_calls``), this is a SINGLE + reviewable tool call that fans the tasks out internally with + ``asyncio.gather``. Each task runs in its own asyncio Task, so the + ContextVar-based session/browser isolation keeps them from clobbering each + other. This is the DRY, deterministic way to spawn a pack of puppies. + """ + + async def invoke_agents( + context: RunContext, + tasks: TasksArg, + ) -> list[AgentInvokeOutput]: + """Invoke multiple sub-agents concurrently and collect every result. + + Use this when you have independent tasks that can run in parallel + (e.g. "review these three files"). Each sub-agent runs with its own + isolated session, and results come back in the same order as ``tasks``. + + Args: + tasks: List of sub-agent tasks, each with agent_name, prompt, an + optional session_id, and an optional per-task model_name + override (omit model_name to use the sub-agent's configured + model). This closes the "parallel + custom models" gap so the + batch path is no longer limited to each agent's default model. + + Returns: + list[AgentInvokeOutput]: One result per task, order-preserved. + A failure in one task is captured in that task's ``error`` field + and never cancels its siblings. + """ + if not tasks: + return [] + + emit_info( + f"Spawning {len(tasks)} sub-agent(s) in parallel: " + + ", ".join(t.agent_name for t in tasks), + message_group=generate_group_id("invoke_agent", "batch"), + ) + + coros = [ + _invoke_agent_impl( + context=context, + agent_name=task.agent_name, + prompt=task.prompt, + session_id=task.session_id, + model_name=task.resolved_model_name(), + ) + for task in tasks + ] + + raw_results = await asyncio.gather(*coros, return_exceptions=True) + + # Normalise: an exception escaping _invoke_agent_impl (it mostly + # swallows its own, but belt-and-braces) becomes a structured error + # so one bad puppy doesn't poison the whole batch's return type. + results: list[AgentInvokeOutput] = [] + for task, outcome in zip(tasks, raw_results): + if isinstance(outcome, BaseException): + results.append( + AgentInvokeOutput( + response=None, + agent_name=task.agent_name, + session_id=task.session_id, + model_name=task.resolved_model_name(), + error=f"Error invoking agent '{task.agent_name}': {outcome}", + ) + ) + else: + results.append(outcome) + + return results + + return agent.tool(invoke_agents) diff --git a/tests/messaging/test_bus.py b/tests/messaging/test_bus.py index 0930261a3..ed804eb79 100644 --- a/tests/messaging/test_bus.py +++ b/tests/messaging/test_bus.py @@ -36,7 +36,14 @@ @pytest.fixture def bus(): - return MessageBus(maxsize=10) + # Session context now lives in a process-global ContextVar (so parallel + # sub-agents stay isolated), which means it can leak across tests. Reset + # it before and after each test to keep them hermetic. + set_session_context(None) + try: + yield MessageBus(maxsize=10) + finally: + set_session_context(None) # ========================================================================= @@ -517,6 +524,8 @@ def delayed_put(): def test_global_session_context(): reset_message_bus() + # Context lives in a process-global ContextVar; clear any leak first. + set_session_context(None) assert get_session_context() is None set_session_context("s1") assert get_session_context() == "s1" diff --git a/tests/test_agent_tools_coverage.py b/tests/test_agent_tools_coverage.py index 3c4013f1d..a568914b4 100644 --- a/tests/test_agent_tools_coverage.py +++ b/tests/test_agent_tools_coverage.py @@ -410,11 +410,8 @@ async def test_invoke_agent_model_not_found_error(self): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error") as mock_emit_error, patch( "code_puppy.agents.agent_manager.load_agent", @@ -482,13 +479,10 @@ def temporary_override(model_name): patch("code_puppy.tools.subagent_invocation.get_message_bus") ) stack.enter_context( - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ) + patch("code_puppy.tools.subagent_invocation.set_session_context") ) stack.enter_context( - patch("code_puppy.tools.subagent_invocation.set_session_context") + patch("code_puppy.tools.subagent_invocation.reset_session_context") ) stack.enter_context(patch("code_puppy.tools.subagent_invocation.emit_info")) stack.enter_context( @@ -619,11 +613,8 @@ def temporary_override(model_name): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error"), patch( "code_puppy.agents.agent_manager.load_agent", @@ -669,7 +660,9 @@ async def test_invoke_agent_session_context_restored_on_error(self): mock_agent_config.get_model_name.return_value = "test-model" mock_agent_config.get_system_prompt.return_value = "Test" - set_context_calls = [] + # Token-based contextvar API: set returns a token, reset restores it. + sentinel_token = object() + reset_calls = [] with ( patch( @@ -678,12 +671,12 @@ async def test_invoke_agent_session_context_restored_on_error(self): ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="original-parent", + "code_puppy.tools.subagent_invocation.set_session_context", + return_value=sentinel_token, ), patch( - "code_puppy.tools.subagent_invocation.set_session_context", - side_effect=lambda x: set_context_calls.append(x), + "code_puppy.tools.subagent_invocation.reset_session_context", + side_effect=lambda tok: reset_calls.append(tok), ), patch("code_puppy.tools.subagent_invocation.emit_error"), patch( @@ -715,8 +708,9 @@ async def test_invoke_agent_session_context_restored_on_error(self): # Should have error assert result.error is not None - # Session context should still be restored - assert "original-parent" in set_context_calls + # Session context must still be torn down via the returned token, + # even when the run errors out before completing. + assert reset_calls == [sentinel_token] class TestInvokeAgentPartialSessionSaveOnCrash: @@ -761,11 +755,8 @@ async def test_partial_history_saved_when_run_crashes(self): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error"), patch("code_puppy.tools.subagent_invocation.emit_info"), patch( @@ -821,11 +812,8 @@ async def test_no_save_when_no_progress_beyond_loaded_history(self): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error"), patch("code_puppy.tools.subagent_invocation.emit_info"), patch( @@ -872,11 +860,8 @@ async def test_save_failure_does_not_mask_original_error(self): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error"), patch("code_puppy.tools.subagent_invocation.emit_info"), patch( @@ -924,11 +909,8 @@ async def test_load_agent_itself_crashes_no_save_attempted(self): return_value="test-group", ), patch("code_puppy.tools.subagent_invocation.get_message_bus") as mock_bus, - patch( - "code_puppy.tools.subagent_invocation.get_session_context", - return_value="parent", - ), patch("code_puppy.tools.subagent_invocation.set_session_context"), + patch("code_puppy.tools.subagent_invocation.reset_session_context"), patch("code_puppy.tools.subagent_invocation.emit_error"), patch( "code_puppy.agents.agent_manager.load_agent", @@ -1133,3 +1115,251 @@ def capture_tool(func): # Verify emit_info was called (at least for banner) assert mock_emit_info.called + + +class TestRegisterInvokeAgentsBatch: + """Tests for the parallel fan-out ``invoke_agents`` batch tool.""" + + def _get_registered_invoke_agents(self): + from code_puppy.tools.subagent_invocation import register_invoke_agents + + captured = {} + + class FakeAgent: + def tool(self, fn): + captured["fn"] = fn + return fn + + register_invoke_agents(FakeAgent()) + return captured["fn"] + + @pytest.mark.asyncio + async def test_empty_task_list_returns_empty(self): + invoke_agents = self._get_registered_invoke_agents() + result = await invoke_agents(MagicMock(), []) + assert result == [] + + @pytest.mark.asyncio + async def test_fans_out_concurrently_and_preserves_order(self): + from code_puppy.tools.subagent_invocation import SubAgentTask + + invoke_agents = self._get_registered_invoke_agents() + completion_order = [] + + async def fake_impl( + context, agent_name, prompt, session_id=None, model_name=None + ): + # "slow" sleeps so a later-listed "fast" can overtake it, + # proving the calls actually run concurrently. + import asyncio + + await asyncio.sleep(0.05 if agent_name == "slow" else 0.0) + completion_order.append(agent_name) + return AgentInvokeOutput( + response=f"{agent_name}-done", + agent_name=agent_name, + session_id=session_id, + model_name="m", + ) + + tasks = [ + SubAgentTask(agent_name="slow", prompt="p1"), + SubAgentTask(agent_name="fast", prompt="p2"), + ] + + with ( + patch( + "code_puppy.tools.subagent_invocation._invoke_agent_impl", + side_effect=fake_impl, + ), + patch("code_puppy.tools.subagent_invocation.emit_info"), + ): + results = await invoke_agents(MagicMock(), tasks) + + # Concurrency: fast finished before slow despite being listed second. + assert completion_order[0] == "fast" + # Order preservation: results match the input order, not completion. + assert [r.agent_name for r in results] == ["slow", "fast"] + assert results[0].response == "slow-done" + + @pytest.mark.asyncio + async def test_one_failure_does_not_poison_siblings(self): + from code_puppy.tools.subagent_invocation import SubAgentTask + + invoke_agents = self._get_registered_invoke_agents() + + async def fake_impl( + context, agent_name, prompt, session_id=None, model_name=None + ): + if agent_name == "boom": + raise RuntimeError("kaboom") + return AgentInvokeOutput( + response=f"{agent_name}-done", + agent_name=agent_name, + session_id=session_id, + model_name="m", + ) + + tasks = [ + SubAgentTask(agent_name="ok", prompt="p1"), + SubAgentTask(agent_name="boom", prompt="p2"), + ] + + with ( + patch( + "code_puppy.tools.subagent_invocation._invoke_agent_impl", + side_effect=fake_impl, + ), + patch("code_puppy.tools.subagent_invocation.emit_info"), + ): + results = await invoke_agents(MagicMock(), tasks) + + assert results[0].response == "ok-done" + assert results[1].response is None + assert results[1].error is not None + assert "kaboom" in results[1].error + + @pytest.mark.asyncio + async def test_per_task_model_name_is_passed_through(self): + """Each task's model_name override must reach ``_invoke_agent_impl``. + + This closes the "parallel + custom models" gap: the batch path used to + hardcode ``model_name=None``, so per-task model selection was + impossible without falling back to individual ``invoke_agent_with_model`` + calls. + """ + from code_puppy.tools.subagent_invocation import SubAgentTask + + invoke_agents = self._get_registered_invoke_agents() + seen = {} + + async def fake_impl( + context, agent_name, prompt, session_id=None, model_name=None + ): + seen[agent_name] = model_name + return AgentInvokeOutput( + response=f"{agent_name}-done", + agent_name=agent_name, + session_id=session_id, + model_name=model_name, + ) + + tasks = [ + SubAgentTask(agent_name="custom", prompt="p1", model_name="gpt-fancy"), + # Blank override normalizes to None (no bogus override leaks). + SubAgentTask(agent_name="blank", prompt="p2", model_name=" "), + # Omitted entirely -> default behaviour preserved (None). + SubAgentTask(agent_name="default", prompt="p3"), + ] + + with ( + patch( + "code_puppy.tools.subagent_invocation._invoke_agent_impl", + side_effect=fake_impl, + ), + patch("code_puppy.tools.subagent_invocation.emit_info"), + ): + results = await invoke_agents(MagicMock(), tasks) + + assert seen == { + "custom": "gpt-fancy", + "blank": None, + "default": None, + } + assert results[0].model_name == "gpt-fancy" + + def test_resolved_model_name_normalizes_blanks(self): + """Whitespace-only overrides collapse to None, mirroring with_model.""" + from code_puppy.tools.subagent_invocation import SubAgentTask + + assert SubAgentTask(agent_name="a", prompt="p").resolved_model_name() is None + assert ( + SubAgentTask( + agent_name="a", prompt="p", model_name=" " + ).resolved_model_name() + is None + ) + assert ( + SubAgentTask( + agent_name="a", prompt="p", model_name=" gpt-x " + ).resolved_model_name() + == "gpt-x" + ) + + def test_tasks_arg_coerces_json_string(self): + """Regression: providers that double-encode ``tasks`` as a JSON string. + + pydantic's ``validate_python`` refuses to parse a JSON string into a + list (``list_type`` error). The ``BeforeValidator`` on ``TasksArg`` + must coerce it so the live batch tool stops faceplanting. This is the + serialization boundary the object-based tests above never exercised. + """ + import json + + from pydantic import TypeAdapter + + from code_puppy.tools.subagent_invocation import SubAgentTask, TasksArg + + ta = TypeAdapter(TasksArg) + payload = [{"agent_name": "a", "prompt": "p"}] + + # Real list still passes straight through. + from_list = ta.validate_python(payload) + # Double-encoded JSON string now coerces instead of raising list_type. + from_string = ta.validate_python(json.dumps(payload)) + + assert from_list == from_string + assert from_string == [SubAgentTask(agent_name="a", prompt="p")] + + def test_tasks_arg_rejects_non_json_string(self): + """Genuine garbage strings must still raise, not silently swallow.""" + from pydantic import TypeAdapter, ValidationError + + from code_puppy.tools.subagent_invocation import TasksArg + + ta = TypeAdapter(TasksArg) + with pytest.raises(ValidationError): + ta.validate_python("definitely not json {") + + @pytest.mark.asyncio + async def test_invoke_agents_accepts_json_string_end_to_end(self): + """The tool itself works when handed a validated JSON-string payload.""" + import json + + from pydantic import TypeAdapter + + from code_puppy.tools.subagent_invocation import TasksArg + + invoke_agents = self._get_registered_invoke_agents() + + async def fake_impl( + context, agent_name, prompt, session_id=None, model_name=None + ): + return AgentInvokeOutput( + response=f"{agent_name}-done", + agent_name=agent_name, + session_id=session_id, + model_name="m", + ) + + # Simulate the provider boundary: tasks arrive as a JSON string, then + # get validated into SubAgentTask objects before the tool body runs. + raw = json.dumps( + [ + {"agent_name": "one", "prompt": "p1"}, + {"agent_name": "two", "prompt": "p2"}, + ] + ) + tasks = TypeAdapter(TasksArg).validate_python(raw) + + with ( + patch( + "code_puppy.tools.subagent_invocation._invoke_agent_impl", + side_effect=fake_impl, + ), + patch("code_puppy.tools.subagent_invocation.emit_info"), + ): + results = await invoke_agents(MagicMock(), tasks) + + assert [r.agent_name for r in results] == ["one", "two"] + assert results[0].response == "one-done" diff --git a/tests/test_messaging_bus.py b/tests/test_messaging_bus.py index 6ee0c6282..324a5f983 100644 --- a/tests/test_messaging_bus.py +++ b/tests/test_messaging_bus.py @@ -34,7 +34,10 @@ def test_initialization_default(self): assert bus._maxsize == 1000 assert isinstance(bus._outgoing, queue.Queue) assert isinstance(bus._incoming, queue.Queue) - assert bus._current_session_id is None + # Session context now lives in a process-global ContextVar, not a + # per-instance attribute. A fresh bus should report no active session. + bus.set_session_context(None) + assert bus.get_session_context() is None assert not bus._has_active_renderer assert bus._startup_buffer == [] diff --git a/tests/test_messaging_init.py b/tests/test_messaging_init.py index 8344a6db0..0f79efd6e 100644 --- a/tests/test_messaging_init.py +++ b/tests/test_messaging_init.py @@ -174,6 +174,7 @@ def test_expected_export_count(self): "patch_markdown_headings", # Session management "set_session_context", + "reset_session_context", "get_session_context", # Shell output rendering "ShellLineMessage",