diff --git a/docs/concepts.md b/docs/concepts.md index 41601c2..70571e9 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -130,7 +130,18 @@ Speed metrics are available directly on each `AssistantMessage` via `request_sta ## Session -The `session()` method returns the agent configured as an async context manager. Sessions handle: +Each call to `session()` returns a detached runtime with its own active tools, +provider and logger lifecycle, run state, skills, cache state, and outputs. Exact +built-in providers and loggers support overlapping sessions. Custom lifecycle +objects and subclasses of built-ins remain reusable across sequential sessions, +but the same configured object cannot be used by overlapping sessions; configure +a distinct instance for each concurrent owner. + +A provider-free `Agent.run()` can still be called directly. It is always an +independent root run, even when called from inside an active session, and does +not inherit that session's files, skills, warnings, or finish validation. + +Sessions handle: - Tool lifecycle (setup and teardown of ToolProviders) - File uploads to execution environment diff --git a/docs/extending/code_backends.md b/docs/extending/code_backends.md index f131890..c36d596 100644 --- a/docs/extending/code_backends.md +++ b/docs/extending/code_backends.md @@ -20,6 +20,12 @@ The base class provides: - `allowed_commands` - Optional regex patterns to restrict commands - File upload/download utilities +Custom backends and subclasses of built-in backends own their configured runtime +state. They can be reused sequentially, but overlapping sessions must use +separate backend instances. Stirrup rejects overlap on the same configured +instance before entering resources. Exact built-in backends are reconstructed +privately for each session. + ## Minimal Implementation ```python diff --git a/docs/extending/loggers.md b/docs/extending/loggers.md index f95b4e3..874bd15 100644 --- a/docs/extending/loggers.md +++ b/docs/extending/loggers.md @@ -228,6 +228,17 @@ async with agent.session() as session: print(metrics_logger.get_summary()) ``` +A custom logger instance is reused directly and supports sequential sessions. +The same configured logger cannot serve overlapping sessions, even when it is +attached to different `Agent` objects; the second session fails before it enters +any logger or provider resources. Use one logger instance per concurrent +owner, and share an explicitly synchronized metrics or event sink between those +instances when aggregation is required. Subclasses of the built-in logger follow +the same sequential-only rule. + +The exact built-in console and Slack loggers receive private per-session +lifecycle state automatically. + ## Combining Loggers ```python diff --git a/docs/extending/tools.md b/docs/extending/tools.md index cc52f83..c9d1229 100644 --- a/docs/extending/tools.md +++ b/docs/extending/tools.md @@ -53,6 +53,11 @@ class WeatherToolProvider: await self._client.aclose() ``` +Custom providers are reusable sequentially. For overlapping sessions, construct +one provider instance per concurrent owner; sharing the same configured object +is rejected before resources are entered. This also applies to subclasses of +built-in providers. + ## Tools with State ```python diff --git a/docs/guides/caching.md b/docs/guides/caching.md index 6a50512..595a6d2 100644 --- a/docs/guides/caching.md +++ b/docs/guides/caching.md @@ -76,5 +76,9 @@ cache_manager.clear_cache("abc123def456") ## Notes - Cache key is computed from the initial prompt—same prompt = same cache +- Cache files are not shared transactionally. Within one process, while + `cache_on_interrupt=True`, a concurrent root run with the same task prompt is + rejected before cache I/O. Different prompts can run concurrently; use + `cache_on_interrupt=False` when identical tasks must overlap without caching. - Caches are stored locally in `~/.cache/stirrup/` - Caches are automatically cleared on successful completion (by default) diff --git a/docs/guides/sub-agents.md b/docs/guides/sub-agents.md index 8ff152b..b40c124 100644 --- a/docs/guides/sub-agents.md +++ b/docs/guides/sub-agents.md @@ -63,6 +63,13 @@ agent.to_tool( Returns a `Tool[SubAgentParams, SubAgentMetadata]`. +The returned tool must execute within its owning parent's active session. Use +`async with parent.session() as session:` and run the parent through `session`; +a direct `await parent.run(...)` cannot delegate to sub-agents. Without an +active parent session, delegation returns an unsuccessful tool result instead +of treating the child as a root agent or saving its output to the current +working directory. + ### SubAgentParams When the parent calls the sub-agent tool, it provides: @@ -317,6 +324,9 @@ parent = Agent( - Sub-agent runs are synchronous (parent waits for completion) - All sub-agent messages are returned to parent (may use context) - File transfer only works with code execution environments +- A custom `ViewImageToolProvider` configured with the sub-agent's backend cannot + use `share_parent_exec_env=True`; disable sharing or use the exact built-in + view-image provider. ## Next Steps diff --git a/docs/guides/tool-providers.md b/docs/guides/tool-providers.md index 536c92e..d36c297 100644 --- a/docs/guides/tool-providers.md +++ b/docs/guides/tool-providers.md @@ -137,6 +137,23 @@ async with agent.session() as session: # Providers are cleaned up automatically ``` +A custom provider instance owns its lifecycle directly. It can be reused by +sequential sessions, including sessions created by different `Agent` objects. +Overlapping sessions cannot use the same configured custom provider instance; +Stirrup rejects the second session before entering any resources. Create a +separate provider instance for each concurrent owner. The same rule applies to +subclasses of built-in providers. + +Exact built-in providers are reconstructed privately for each session, so one +`Agent` can run overlapping sessions without sharing their clients, sandboxes, +temporary directories, or exit stacks. When a custom `ViewImageToolProvider` +subclass explicitly references a configured built-in code backend, the dependent +provider/backend pair follows the custom-provider rule instead: it is reusable +sequentially, and overlap on either configured object is rejected before entry. +Such a pair cannot be used by a `share_parent_exec_env=True` sub-agent because +sharing substitutes the parent's backend; disable sharing or use the exact +built-in `ViewImageToolProvider`. + ## Built-in ToolProviders Stirrup includes several ToolProviders: @@ -173,6 +190,13 @@ agent = Agent( ## Error Handling +Session teardown attempts every provider and logger cleanup and shields that +work from surrounding AnyIO cancellation. Custom-resource ownership is released +only after cleanup finishes or raises. If setup or the session body already +failed, cleanup failures are attached to that primary exception rather than +replacing it. A cleanup failure is raised directly only when there was no +earlier failure. + Handle setup/cleanup errors gracefully: ```python diff --git a/pyproject.toml b/pyproject.toml index adb60a9..4865484 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "openai>=1.0.0", "pillow>=10.4.0", "pydantic>=2.0.0", - "rich>=13.0.0", + "rich>=14.1.0", "tenacity>=5.0.0", "trafilatura>=1.9.0", ] @@ -62,7 +62,7 @@ dev = [ "mkdocstrings>=0.30.1", "mkdocstrings-python>=1.19.0", "pytest>=9.0.0", - "rich>=13.0.0", + "rich>=14.1.0", "ruff>=0.14.4", "ty>=0.0.1a32", ] diff --git a/src/stirrup/core/agent.py b/src/stirrup/core/agent.py index 71e937c..7a7ce59 100644 --- a/src/stirrup/core/agent.py +++ b/src/stirrup/core/agent.py @@ -1,17 +1,18 @@ -# Context var for passing parent depth to sub-agent executors import contextvars +import copy import glob as glob_module import inspect import logging import re import signal +import threading from contextlib import AsyncExitStack from dataclasses import dataclass, field from itertools import chain, takewhile from pathlib import Path from time import perf_counter from types import TracebackType -from typing import Annotated, Any, Self, overload +from typing import Annotated, Any, overload import anyio from pydantic import BaseModel, Field, ValidationError @@ -50,7 +51,10 @@ from stirrup.tools.finish import FinishParams as SimpleFinishParams from stirrup.utils.logging import AgentLogger, AgentLoggerBase -_PARENT_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar("parent_depth", default=0) +_ACTIVE_SESSION_RESERVATIONS: dict[tuple[str, int | str], object] = {} +_SESSION_RESERVATIONS_LOCK = threading.Lock() +_ACTIVE_SIGINT_SESSION_COUNT = 0 +_ORIGINAL_SIGINT_HANDLER: Any = None logger = logging.getLogger(__name__) @@ -76,6 +80,7 @@ class SessionState: """ exit_stack: AsyncExitStack + owner: "SessionAgent[Any, Any] | None" = None exec_env: CodeExecToolProvider | None = None output_dir: str | None = None # String path (contextual: local for root, in parent env for subagent) parent_exec_env: CodeExecToolProvider | None = None @@ -86,7 +91,7 @@ class SessionState: logger: AgentLoggerBase | None = None # Logger for pause/resume during user input -_SESSION_STATE: contextvars.ContextVar[SessionState] = contextvars.ContextVar("session_state") +_SESSION_STATE: contextvars.ContextVar[SessionState | None] = contextvars.ContextVar("session_state", default=None) __all__ = [ "Agent", @@ -97,6 +102,216 @@ class SessionState: LOGGER = logging.getLogger(__name__) +def _handle_interrupt_signal(_signum: int, _frame: object) -> None: + """Turn SIGINT into an exception so active sessions can cache on exit.""" + raise KeyboardInterrupt("Agent interrupted - state will be cached") + + +def _install_interrupt_handler() -> bool: + """Install one process handler for main-thread root sessions.""" + global _ACTIVE_SIGINT_SESSION_COUNT, _ORIGINAL_SIGINT_HANDLER + + if threading.current_thread() is not threading.main_thread(): + return False + if _ACTIVE_SIGINT_SESSION_COUNT == 0: + _ORIGINAL_SIGINT_HANDLER = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, _handle_interrupt_signal) + _ACTIVE_SIGINT_SESSION_COUNT += 1 + return True + + +def _release_interrupt_handler() -> None: + """Restore the process handler after the final root session exits.""" + global _ACTIVE_SIGINT_SESSION_COUNT, _ORIGINAL_SIGINT_HANDLER + + _ACTIVE_SIGINT_SESSION_COUNT -= 1 + if _ACTIVE_SIGINT_SESSION_COUNT == 0: + signal.signal(signal.SIGINT, _ORIGINAL_SIGINT_HANDLER) + _ORIGINAL_SIGINT_HANDLER = None + + +def _reserve_custom_session_resources(resources: list[object]) -> None: + """Claim custom lifecycle objects before any session resource is entered.""" + with _SESSION_RESERVATIONS_LOCK: + for resource in resources: + if ("custom", id(resource)) in _ACTIVE_SESSION_RESERVATIONS: + raise RuntimeError( + f"Overlapping sessions cannot share configured {type(resource).__name__}. " + "Custom providers, custom loggers, and subclasses of built-ins support sequential sessions only." + ) + for resource in resources: + _ACTIVE_SESSION_RESERVATIONS[("custom", id(resource))] = resource + + +def _release_custom_session_resources(resources: list[object]) -> None: + with _SESSION_RESERVATIONS_LOCK: + for resource in resources: + _ACTIVE_SESSION_RESERVATIONS.pop(("custom", id(resource)), None) + + +def _reserve_cached_root_run(task_hash: str, owner: object) -> None: + """Keep one cache-writing root run active for a task hash.""" + reservation = ("cache", task_hash) + with _SESSION_RESERVATIONS_LOCK: + active_owner = _ACTIVE_SESSION_RESERVATIONS.get(reservation) + if active_owner is not None: + raise RuntimeError( + "Concurrent root runs with cache_on_interrupt=True cannot use the same task cache. " + "Wait for the active run to finish or set cache_on_interrupt=False." + ) + _ACTIVE_SESSION_RESERVATIONS[reservation] = owner + + +def _release_cached_root_runs(task_hashes: set[str], owner: object) -> None: + with _SESSION_RESERVATIONS_LOCK: + for task_hash in task_hashes: + reservation = ("cache", task_hash) + if _ACTIVE_SESSION_RESERVATIONS.get(reservation) is owner: + _ACTIVE_SESSION_RESERVATIONS.pop(reservation) + + +def _record_session_cleanup_failure( + primary_exception: BaseException | None, + phase: str, + cleanup_failure: BaseException, +) -> BaseException: + """Keep the first failure and record later cleanup failures on it.""" + if primary_exception is None: + return cleanup_failure + primary_exception.add_note( + f"{phase} failed during session cleanup: {type(cleanup_failure).__name__}: {cleanup_failure}" + ) + return primary_exception + + +def _detach_builtin_provider(provider: ToolProvider, replacements: dict[int, ToolProvider]) -> ToolProvider | None: + """Return a fresh runtime owner for an exact built-in provider.""" + provider_type = type(provider) + + from stirrup.tools.view_image import ViewImageToolProvider + from stirrup.tools.web import WebToolProvider + + if isinstance(provider, LocalCodeExecToolProvider) and provider_type is LocalCodeExecToolProvider: + detached = copy.copy(provider) + detached._allowed_commands = ( # noqa: SLF001 + list(provider._allowed_commands) if provider._allowed_commands is not None else None # noqa: SLF001 + ) + detached._compiled_allowed = ( # noqa: SLF001 + list(provider._compiled_allowed) if provider._compiled_allowed is not None else None # noqa: SLF001 + ) + detached._temp_dir = None # noqa: SLF001 + return detached + if isinstance(provider, WebToolProvider) and provider_type is WebToolProvider: + detached = copy.copy(provider) + detached._client = None # noqa: SLF001 + return detached + if isinstance(provider, ViewImageToolProvider) and provider_type is ViewImageToolProvider: + detached = copy.copy(provider) + configured_exec_env = provider._exec_env # noqa: SLF001 + if configured_exec_env is not None: + replacement = replacements.get(id(configured_exec_env)) + detached._exec_env = ( # noqa: SLF001 + replacement if isinstance(replacement, CodeExecToolProvider) else configured_exec_env + ) + return detached + + if provider_type.__module__ == "stirrup.tools.code_backends.docker": + from stirrup.tools.code_backends.docker import DockerCodeExecToolProvider + + if isinstance(provider, DockerCodeExecToolProvider) and provider_type is DockerCodeExecToolProvider: + detached = copy.copy(provider) + detached._allowed_commands = ( # noqa: SLF001 + list(provider._allowed_commands) if provider._allowed_commands is not None else None # noqa: SLF001 + ) + detached._compiled_allowed = ( # noqa: SLF001 + list(provider._compiled_allowed) if provider._compiled_allowed is not None else None # noqa: SLF001 + ) + detached._env_vars = ( # noqa: SLF001 + list(provider._env_vars) if provider._env_vars is not None else None # noqa: SLF001 + ) + detached._temp_dir = None # noqa: SLF001 + detached._client = None # noqa: SLF001 + detached._container = None # noqa: SLF001 + return detached + elif provider_type.__module__ == "stirrup.tools.code_backends.e2b": + from stirrup.tools.code_backends.e2b import E2BCodeExecToolProvider + + if isinstance(provider, E2BCodeExecToolProvider) and provider_type is E2BCodeExecToolProvider: + detached = copy.copy(provider) + detached._allowed_commands = ( # noqa: SLF001 + list(provider._allowed_commands) if provider._allowed_commands is not None else None # noqa: SLF001 + ) + detached._compiled_allowed = ( # noqa: SLF001 + list(provider._compiled_allowed) if provider._compiled_allowed is not None else None # noqa: SLF001 + ) + detached._sandbox_kwargs = dict(provider._sandbox_kwargs) # noqa: SLF001 + for key in ("envs", "metadata"): + value = detached._sandbox_kwargs.get(key) # noqa: SLF001 + if isinstance(value, dict): + detached._sandbox_kwargs[key] = dict(value) # noqa: SLF001 + detached._sbx = None # noqa: SLF001 + return detached + elif provider_type.__module__ == "stirrup.tools.mcp": + from stirrup.tools.mcp import MCPToolProvider + + if isinstance(provider, MCPToolProvider) and provider_type is MCPToolProvider: + detached = copy.copy(provider) + detached._config = provider._config.model_copy(deep=True) # noqa: SLF001 + detached._server_names = ( # noqa: SLF001 + list(provider._server_names) if provider._server_names is not None else None # noqa: SLF001 + ) + detached._servers = {} # noqa: SLF001 + detached._tools = {} # noqa: SLF001 + detached._exit_stack = None # noqa: SLF001 + return detached + elif provider_type.__module__ == "stirrup.tools.browser_use": + from stirrup.tools.browser_use import BrowserUseToolProvider + + if isinstance(provider, BrowserUseToolProvider) and provider_type is BrowserUseToolProvider: + detached = copy.copy(provider) + detached._extra_args = ( # noqa: SLF001 + list(provider._extra_args) if provider._extra_args is not None else None # noqa: SLF001 + ) + detached._session = None # noqa: SLF001 + return detached + + return None + + +def _reset_logger_state(session_logger: AgentLoggerBase) -> None: + session_logger.name = "agent" + session_logger.model = None + session_logger.max_turns = None + session_logger.depth = 0 + session_logger.finish_params = None + session_logger.run_metadata = None + session_logger.output_dir = None + + +def _detach_logger(configured_logger: AgentLoggerBase) -> tuple[AgentLoggerBase, bool]: + """Detach exact built-in logger state; identify custom sequential owners.""" + logger_type = type(configured_logger) + if isinstance(configured_logger, AgentLogger) and logger_type is AgentLogger: + detached = copy.copy(configured_logger) + _reset_logger_state(detached) + detached._current_step = 0 # noqa: SLF001 + detached._tool_calls = 0 # noqa: SLF001 + detached._input_tokens = 0 # noqa: SLF001 + detached._output_tokens = 0 # noqa: SLF001 + detached._live = None # noqa: SLF001 + return detached, False + + if logger_type.__module__ == "stirrup.integrations.slack.slack": + from stirrup.integrations.slack.slack import SlackLogger + + if isinstance(configured_logger, SlackLogger) and logger_type is SlackLogger: + detached = copy.copy(configured_logger) + _reset_logger_state(detached) + return detached, False + + return configured_logger, True + + def _num_turns_remaining_msg(number_of_turns_remaining: int) -> TurnWarningMessage: """Create a user message warning the agent about remaining turns before max_turns is reached.""" if number_of_turns_remaining == 1: @@ -261,6 +476,14 @@ class Agent[FinishParams: BaseModel, FinishMeta]: finish_params, history, metadata = await session.run("Your task here") """ + _parent_session_state: SessionState | None + _session_state: SessionState | None + _session_state_token: contextvars.Token[SessionState | None] | None + _interrupt_handler_installed: bool + _custom_session_resources: list[object] + _custom_resources_reserved: bool + _reserved_cache_task_hashes: set[str] + @overload def __init__( self: "Agent[SimpleFinishParams, ToolUseCountMetadata]", @@ -371,6 +594,7 @@ def __init__( provides better performance (no file copying) and allows the subagent to see all files in the parent's environment. Only effective when the agent is used as a subagent via to_tool(). + Custom ViewImageToolProvider/backend pairs are unsupported. logger: Optional logger instance. If None, creates AgentLogger() internally. """ @@ -427,6 +651,7 @@ def __init__( # Cache state for resumption (set during run(), used in __aexit__ for caching on interrupt) self._current_task_hash: str | None = None self._current_run_state: CacheState | None = None + self._reserved_cache_task_hashes: set[str] = set() @property def name(self) -> str: @@ -473,8 +698,8 @@ def session( resume: bool = False, clear_cache_on_success: bool = True, cache_on_interrupt: bool = True, - ) -> Self: - """Configure a session and return self for use as async context manager. + ) -> "SessionAgent[FinishParams, FinishMeta]": + """Create a detached runtime for use as an async context manager. Args: output_dir: Directory to save output files from finish_params.paths @@ -498,35 +723,31 @@ def session( cache_on_interrupt: If True (default), set up a SIGINT handler to cache state on Ctrl+C. Set to False when running agents in threads or subprocesses where signal handlers cannot - be registered from non-main threads. + be registered from non-main threads. Concurrent root runs + in one process must use different task prompts while caching + is enabled. Returns: - Self, for use with `async with agent.session(...) as session:` + A fresh SessionAgent for use with `async with agent.session(...) as session:` Example: async with agent.session(output_dir="./output", input_files="data/*.csv") as session: result = await session.run("Analyze the CSV files") Note: - Multiple concurrent sessions from the same Agent instance are supported. - Each session maintains isolated state via ContextVar. + Exact built-in providers and loggers support overlapping sessions. + Custom lifecycle objects support sequential sessions only. """ - self._pending_output_dir = Path(output_dir) if output_dir else None - self._pending_input_files = input_files - self._pending_skills_dir = Path(skills_dir) if skills_dir else None - self._resume = resume - self._clear_cache_on_success = clear_cache_on_success - self._cache_on_interrupt = cache_on_interrupt - return self - - def _handle_interrupt(self, _signum: int, _frame: object) -> None: - """Handle SIGINT to ensure caching before exit. - - Converts the signal to a KeyboardInterrupt exception so that __aexit__ - is properly called and can cache the state before cleanup. - """ - raise KeyboardInterrupt("Agent interrupted - state will be cached") + return SessionAgent.from_agent( + self, + output_dir=output_dir, + input_files=input_files, + skills_dir=skills_dir, + resume=resume, + clear_cache_on_success=clear_cache_on_success, + cache_on_interrupt=cache_on_interrupt, + ) def _resolve_input_files(self, input_files: str | Path | list[str | Path]) -> list[Path]: """Resolve input file paths, expanding globs and normalizing to Path objects. @@ -737,35 +958,47 @@ def _validate_subagent_code_exec_requirements(self) -> None: # cell_contents can raise ValueError if empty - ignore async def __aenter__(self) -> "SessionAgent[FinishParams, FinishMeta]": - """Enter session context: set up tools, logging, and resources. - - Returns a SessionAgent wrapping this agent's state, providing access to - both static tools and ToolProvider-created tools. The SessionAgent shares - this agent's __dict__, so state changes are visible to __aexit__. - """ + """Enter this detached session and initialize its resources.""" + if not isinstance(self, SessionAgent): + raise RuntimeError("Use `async with agent.session(...)` to enter an Agent session") + if self._session_state is not None: + raise RuntimeError("Agent session is already active") + + _reserve_custom_session_resources(self._custom_session_resources) + self._custom_resources_reserved = True exit_stack = AsyncExitStack() - await exit_stack.__aenter__() - - # Get parent state if exists (for subagent file transfer) - parent_state = _SESSION_STATE.get(None) - - current_depth = _PARENT_DEPTH.get() + try: + await exit_stack.__aenter__() + except BaseException: + _release_custom_session_resources(self._custom_session_resources) + self._custom_resources_reserved = False + raise - # Create session state and store in ContextVar + parent_state = self._parent_session_state + current_depth = parent_state.depth + 1 if parent_state is not None else 0 state = SessionState( exit_stack=exit_stack, + owner=self, output_dir=str(self._pending_output_dir) if self._pending_output_dir else None, parent_exec_env=parent_state.exec_env if parent_state else None, depth=current_depth, logger=self._logger, ) - _SESSION_STATE.set(state) + self._session_state = state + self._session_state_token = _SESSION_STATE.set(state) + logger_entered = False try: # === TWO-PASS TOOL INITIALIZATION === # First pass initializes CodeExecToolProvider so that dependent tools # (like ViewImageToolProvider) can access state.exec_env in second pass. active_tools: list[Tool] = [] + code_exec_providers = [tool for tool in self._tools if isinstance(tool, CodeExecToolProvider)] + if len(code_exec_providers) > 1: + raise ValueError( + f"Agent can only have one CodeExecToolProvider, found {len(code_exec_providers)}: " + f"{[type(provider).__name__ for provider in code_exec_providers]}" + ) # Check if we should share parent's exec_env (subagent with share_parent_exec_env=True) should_share_exec_env = ( @@ -793,14 +1026,7 @@ async def __aenter__(self) -> "SessionAgent[FinishParams, FinishMeta]": code_exec_tool = state.exec_env.get_code_exec_tool() active_tools.append(code_exec_tool) else: - # OWNED EXEC ENV: Initialize our own CodeExecToolProvider (at most one allowed) - code_exec_providers = [t for t in self._tools if isinstance(t, CodeExecToolProvider)] - if len(code_exec_providers) > 1: - raise ValueError( - f"Agent can only have one CodeExecToolProvider, found {len(code_exec_providers)}: " - f"{[type(p).__name__ for p in code_exec_providers]}" - ) - + # OWNED EXEC ENV: Initialize our own CodeExecToolProvider. if code_exec_providers: provider = code_exec_providers[0] result = await exit_stack.enter_async_context(provider) @@ -916,8 +1142,8 @@ async def __aenter__(self) -> "SessionAgent[FinishParams, FinishMeta]": state.skills_metadata = load_skills_metadata(skills_path) logger.debug("[%s __aenter__] Loaded %d skills", self._name, len(state.skills_metadata)) self._pending_skills_dir = None # Clear pending state - elif parent_state and parent_state.skills_metadata: - # Sub-agent: inherit skills from parent + elif parent_state is not None and parent_state.skills_metadata: + # Sub-agent: inherit skills from its explicit parent state.skills_metadata = parent_state.skills_metadata logger.debug("[%s __aenter__] Inherited %d skills from parent", self._name, len(state.skills_metadata)) # Transfer skills directory from parent's exec_env to sub-agent's exec_env @@ -929,18 +1155,42 @@ async def __aenter__(self) -> "SessionAgent[FinishParams, FinishMeta]": self._logger.name = self._name self._logger.model = self._client.model_slug self._logger.max_turns = self._max_turns - # depth is already set (0 for main agent, passed in for sub-agents) + self._logger.depth = current_depth + self._logger.finish_params = None + self._logger.run_metadata = None + self._logger.output_dir = None self._logger.__enter__() + logger_entered = True - # Set up signal handler for graceful caching on interrupt (root agent only) if current_depth == 0 and self._cache_on_interrupt: - self._original_sigint = signal.getsignal(signal.SIGINT) - signal.signal(signal.SIGINT, self._handle_interrupt) + self._interrupt_handler_installed = _install_interrupt_handler() - return SessionAgent.from_agent(self) + return self - except Exception: - await exit_stack.__aexit__(None, None, None) + except BaseException as exc: + with anyio.CancelScope(shield=True): + if logger_entered: + try: + self._logger.__exit__(type(exc), exc, exc.__traceback__) + except BaseException as cleanup_failure: + _record_session_cleanup_failure(exc, "Logger", cleanup_failure) + try: + await exit_stack.__aexit__(type(exc), exc, exc.__traceback__) + except BaseException as cleanup_failure: + _record_session_cleanup_failure(exc, "Tool provider", cleanup_failure) + try: + if self._session_state_token is not None: + _SESSION_STATE.reset(self._session_state_token) + self._session_state_token = None + except BaseException as cleanup_failure: + _record_session_cleanup_failure(exc, "Ambient session context", cleanup_failure) + self._session_state = None + try: + if self._custom_resources_reserved: + _release_custom_session_resources(self._custom_session_resources) + self._custom_resources_reserved = False + except BaseException as cleanup_failure: + _record_session_cleanup_failure(exc, "Custom resource reservation", cleanup_failure) raise async def __aexit__( @@ -955,8 +1205,21 @@ async def __aexit__( - Root agent (depth=0): Saves files to local filesystem output_dir - Subagent (depth>0): Transfers files to parent's exec env at output_dir path """ - state = _SESSION_STATE.get() + state = self._session_state + if state is None: + raise RuntimeError("Agent session is not active") + + with anyio.CancelScope(shield=True): + await self._exit_session(state, exc_type, exc_val) + async def _exit_session( + self, + state: SessionState, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + ) -> None: + """Finalize and release a session inside its shielded teardown scope.""" + primary_exception = exc_val try: # Cache state on non-success exit (only at root level) should_cache = ( @@ -988,17 +1251,24 @@ async def __aexit__( if self._current_task_hash is None or self._current_run_state is None: raise ValueError("Cache state is unexpectedly None after should_cache check") - # Temporarily block SIGINT during cache save to prevent interruption - original_handler = signal.getsignal(signal.SIGINT) - signal.signal(signal.SIGINT, signal.SIG_IGN) - try: + if threading.current_thread() is threading.main_thread(): + # Prevent a second SIGINT from interrupting the cache write. + original_handler = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + cache_manager.save_state( + self._current_task_hash, + self._current_run_state, + exec_env_dir, + ) + finally: + signal.signal(signal.SIGINT, original_handler) + else: cache_manager.save_state( self._current_task_hash, self._current_run_state, exec_env_dir, ) - finally: - signal.signal(signal.SIGINT, original_handler) self._logger.info(f"Cached state for task {self._current_task_hash}") # Save files from finish_params.paths based on depth if state.output_dir and self._last_finish_params and state.exec_env: @@ -1064,27 +1334,83 @@ async def __aexit__( "Files will not be transferred.", state.depth, ) - finally: - # Restore original signal handler (root agent only) - if hasattr(self, "_original_sigint"): - signal.signal(signal.SIGINT, self._original_sigint) - del self._original_sigint + except BaseException as finalization_failure: + primary_exception = _record_session_cleanup_failure( + primary_exception, "Session finalization", finalization_failure + ) + + try: + if self._interrupt_handler_installed: + _release_interrupt_handler() + self._interrupt_handler_installed = False + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure(primary_exception, "Interrupt handler", cleanup_failure) - # Exit logger context + try: self._logger.finish_params = self._last_finish_params self._logger.run_metadata = self._last_run_metadata self._logger.output_dir = str(state.output_dir) if state.output_dir else None - self._logger.__exit__(exc_type, exc_val, exc_tb) + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure( + primary_exception, "Logger final state", cleanup_failure + ) + try: + current_exc_type = type(primary_exception) if primary_exception is not None else None + current_exc_tb = primary_exception.__traceback__ if primary_exception is not None else None + self._logger.__exit__(current_exc_type, primary_exception, current_exc_tb) + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure(primary_exception, "Logger", cleanup_failure) - # Reset active tools to static-only (remove session-initialized tools) - self._active_tools = {} - for tool in self._tools: - if isinstance(tool, Tool): - self._active_tools[tool.name] = tool + try: + self._active_tools = {tool.name: tool for tool in self._tools if isinstance(tool, Tool)} self._active_tools.update(self._finish_tools) + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure(primary_exception, "Active tool reset", cleanup_failure) + + try: + current_exc_type = type(primary_exception) if primary_exception is not None else None + current_exc_tb = primary_exception.__traceback__ if primary_exception is not None else None + await state.exit_stack.__aexit__(current_exc_type, primary_exception, current_exc_tb) + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure(primary_exception, "Tool provider", cleanup_failure) + + try: + if self._session_state_token is not None: + _SESSION_STATE.reset(self._session_state_token) + self._session_state_token = None + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure( + primary_exception, "Ambient session context", cleanup_failure + ) + self._session_state = None + + try: + if self._custom_resources_reserved: + _release_custom_session_resources(self._custom_session_resources) + self._custom_resources_reserved = False + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure( + primary_exception, "Custom resource reservation", cleanup_failure + ) + + try: + if self._reserved_cache_task_hashes: + _release_cached_root_runs(self._reserved_cache_task_hashes, self) + self._reserved_cache_task_hashes.clear() + except BaseException as cleanup_failure: + primary_exception = _record_session_cleanup_failure( + primary_exception, "Task cache reservation", cleanup_failure + ) + + if exc_val is None and primary_exception is not None: + raise primary_exception - # Cleanup all async resources - await state.exit_stack.__aexit__(exc_type, exc_val, exc_tb) + def _require_own_active_session_context(self) -> None: + """Reject use of a detached session outside its ambient ownership context.""" + if isinstance(self, SessionAgent): + session_state = self._session_state + if session_state is None or _SESSION_STATE.get() is not session_state: + raise RuntimeError("SessionAgent requires its own active session context") async def run_tool(self, tool_call: ToolCall, run_metadata: dict[str, list[Any]]) -> ToolMessage: """Execute a single tool call with error handling for invalid JSON/arguments. @@ -1092,6 +1418,7 @@ async def run_tool(self, tool_call: ToolCall, run_metadata: dict[str, list[Any]] Returns a ToolMessage containing either the tool output or an error description. Metadata from the tool result is stored in the provided run_metadata dict. """ + self._require_own_active_session_context() tool = self._active_tools.get(tool_call.name) result: ToolResult args_valid = True @@ -1108,19 +1435,14 @@ async def run_tool(self, tool_call: ToolCall, run_metadata: dict[str, list[Any]] args = tool_call.arguments if tool_call.arguments and tool_call.arguments.strip() else "{}" params = tool.parameters.model_validate_json(args) - # Set parent depth for sub-agent tools to read - prev_depth = _PARENT_DEPTH.set(self._logger.depth) - try: - if inspect.iscoroutinefunction(tool.executor): - result = await tool.executor(params) # ty: ignore[invalid-await] - elif self._run_sync_in_thread: - # ty: ignore - type checker doesn't understand iscoroutinefunction narrowing - result = await anyio.to_thread.run_sync(tool.executor, params) # ty: ignore[unresolved-attribute] - else: - # ty: ignore - iscoroutinefunction check above ensures this is sync - result = tool.executor(params) # ty: ignore[invalid-assignment] - finally: - _PARENT_DEPTH.reset(prev_depth) + if inspect.iscoroutinefunction(tool.executor): + result = await tool.executor(params) # ty: ignore[invalid-await] + elif self._run_sync_in_thread: + # ty: ignore - type checker doesn't understand iscoroutinefunction narrowing + result = await anyio.to_thread.run_sync(tool.executor, params) # ty: ignore[unresolved-attribute] + else: + # ty: ignore - iscoroutinefunction check above ensures this is sync + result = tool.executor(params) # ty: ignore[invalid-assignment] # Store metadata if present if result.metadata is not None: @@ -1334,17 +1656,64 @@ async def run( """ - # Guard: fail if ToolProviders are attached but we're not in a session - if self._has_tool_providers and not isinstance(self, SessionAgent): - provider_names = [type(t).__name__ for t in self._tools if isinstance(t, ToolProvider)] + self._require_own_active_session_context() + if not isinstance(self, SessionAgent) and self._has_tool_providers: + provider_names = [type(tool).__name__ for tool in self._tools if isinstance(tool, ToolProvider)] raise RuntimeError( f"Agent.run() called without a session, but the agent has ToolProviders " f"that require session initialization: {provider_names}. " f"Use `async with agent.session(...) as session: await session.run(...)` instead." ) - # Compute task hash for caching/resume + if isinstance(self, SessionAgent): + return await self._run(init_msgs, depth=depth) + + ambient_session_token = _SESSION_STATE.set(None) + reservation_owner = object() task_hash = compute_task_hash(init_msgs) + cache_reserved = False + primary_exception: BaseException | None = None + try: + if self._cache_on_interrupt: + _reserve_cached_root_run(task_hash, reservation_owner) + cache_reserved = True + return await self._run(init_msgs, depth=depth, task_hash=task_hash) + except BaseException as exc: + primary_exception = exc + raise + finally: + try: + if cache_reserved: + _release_cached_root_runs({task_hash}, reservation_owner) + except BaseException as cleanup_failure: + if primary_exception is None: + raise + primary_exception.add_note( + "Task cache reservation release failed after Agent.run(): " + f"{type(cleanup_failure).__name__}: {cleanup_failure}" + ) + finally: + _SESSION_STATE.reset(ambient_session_token) + + async def _run( + self, + init_msgs: str | list[ChatMessage], + *, + depth: int | None = None, + task_hash: str | None = None, + ) -> tuple[FinishParams | None, list[list[ChatMessage]], dict[str, Any]]: + """Execute a run after its session ownership and ambient context are established.""" + task_hash = task_hash or compute_task_hash(init_msgs) + state = _SESSION_STATE.get() + if ( + isinstance(self, SessionAgent) + and self._cache_on_interrupt + and state is not None + and state.depth == 0 + and task_hash not in self._reserved_cache_task_hashes + ): + _reserve_cached_root_run(task_hash, self) + self._reserved_cache_task_hashes.add(task_hash) self._current_task_hash = task_hash # Initialize cache manager @@ -1355,6 +1724,8 @@ async def run( # Try to resume from cache if requested if self._resume: state = _SESSION_STATE.get() + if state is None: + raise RuntimeError("Cannot resume an Agent run without an active session") cached = cache_manager.load_state(task_hash) if cached: # Restore files to exec env @@ -1524,161 +1895,205 @@ def to_tool( """ agent = self # Capture self for closure + sub_agent_tool: Tool[SubAgentParams, SubAgentMetadata] async def sub_agent_executor(params: SubAgentParams) -> ToolResult[SubAgentMetadata]: - """Execute the sub-agent with the given task. - - Sub-agents enter their own full session to ensure: - 1. Tool isolation - each agent only sees its own tools (fixes recursive sub-agent bug) - 2. Proper ToolProvider lifecycle - sub-agent's ToolProviders are initialized - 3. Correct logging - logger context is entered for proper output formatting - """ - # Get parent's depth and calculate subagent depth - parent_depth = _PARENT_DEPTH.get() - sub_agent_depth = parent_depth + 1 - - # Save parent's session state so we can restore it after subagent completes - # This ensures sibling subagents see the parent's state, not a previous sibling's stale state + """Execute this agent as a child of the active session.""" parent_session_state = _SESSION_STATE.get(None) - logger.debug( - "[%s] PRE-SESSION: _SESSION_STATE=%s, exec_env=%s, exec_env._temp_dir=%s", - agent.name, - id(parent_session_state) if parent_session_state else None, - type(parent_session_state.exec_env).__name__ - if parent_session_state and parent_session_state.exec_env - else None, - getattr(parent_session_state.exec_env, "_temp_dir", "N/A") - if parent_session_state and parent_session_state.exec_env - else None, - ) - - # Set _PARENT_DEPTH to subagent's depth BEFORE entering session - # so that __aenter__ reads the correct depth for SessionState.depth - prev_depth = _PARENT_DEPTH.set(sub_agent_depth) + if ( + parent_session_state is None + or parent_session_state.owner is None + or parent_session_state.owner.tools.get(agent.name) is not sub_agent_tool + ): + return ToolResult( + content=( + f"\nSub-agent tool '{agent.name}' requires an active parent " + "Agent session. Use `async with parent.session() as session:` before running the parent." + "\n" + ), + success=False, + metadata=SubAgentMetadata(message_history=[], run_metadata={}), + ) try: init_msgs: list[ChatMessage] = [] if system_prompt: init_msgs.append(SystemMessage(content=system_prompt)) init_msgs.append(UserMessage(content=params.task)) - # Sub-agent enters its own full session for tool isolation and proper lifecycle - # output_dir is a path within the parent's exec env (not local filesystem) - # Files are transferred to parent's env at __aexit__ via save_output_files(dest_env=parent) - async with agent.session( - output_dir=".", # Path in parent's exec env - input_files=list(params.input_files) if params.input_files else None, # ty: ignore[invalid-argument-type] - ) as agent_session: - # Override logger depth for proper indentation in console output - agent_session._logger.depth = sub_agent_depth # noqa: SLF001 - + agent_session = SessionAgent.from_agent( + agent, + output_dir=".", + input_files=list(params.input_files) if params.input_files else None, + parent_session_state=parent_session_state, + ) + async with agent_session: finish_params, msg_history, run_metadata = await agent_session.run(init_msgs) - # Extract the last assistant message with actual content (not just tool calls) - last_assistant_msg: AssistantMessage | None = None - for msg_group in reversed(msg_history): - for msg in reversed(msg_group): - if isinstance(msg, AssistantMessage) and msg.content: - last_assistant_msg = msg - break - if last_assistant_msg: + last_assistant_msg: AssistantMessage | None = None + for msg_group in reversed(msg_history): + for msg in reversed(msg_group): + if isinstance(msg, AssistantMessage) and msg.content: + last_assistant_msg = msg break - - # Build content from the assistant message and/or finish params - content_parts: list[str] = [] - - if last_assistant_msg and last_assistant_msg.content: - content = last_assistant_msg.content - if isinstance(content, list): - content = "\n".join(str(block) for block in content) - content_parts.append(content) - - # Include finish params if available (they often contain the actual result) - if finish_params is not None: - finish_dict = finish_params.model_dump() - if finish_dict: - content_parts.append(f"Finish params: {finish_dict}") - - # Report files transferred to parent's exec env (set in __aexit__) - transferred_paths = agent_session._transferred_paths # noqa: SLF001 - if transferred_paths: - content_parts.append(f"Files available in your environment: {transferred_paths}") - - if not content_parts: - result_content = "\nNo assistant message or finish params found\n" - else: - content = "\n".join(content_parts) - result_content = ( - f"" - f"\n{content}" - f"\n{finish_params is not None}" - f"\n" - ) - - # Create subagent metadata with token usage, message history, and run metadata - sub_metadata = SubAgentMetadata( - message_history=msg_history, - run_metadata=run_metadata, + if last_assistant_msg: + break + + content_parts: list[str] = [] + if last_assistant_msg and last_assistant_msg.content: + content = last_assistant_msg.content + if isinstance(content, list): + content = "\n".join(str(block) for block in content) + content_parts.append(content) + + if finish_params is not None: + finish_dict = finish_params.model_dump() + if finish_dict: + content_parts.append(f"Finish params: {finish_dict}") + + if agent_session._transferred_paths: # noqa: SLF001 + content_parts.append( + f"Files available in your environment: {agent_session._transferred_paths}" # noqa: SLF001 ) - return ToolResult(content=result_content, metadata=sub_metadata) + if not content_parts: + result_content = "\nNo assistant message or finish params found\n" + else: + content = "\n".join(content_parts) + result_content = ( + f"" + f"\n{content}" + f"\n{finish_params is not None}" + f"\n" + ) - except Exception as e: - # On error, return empty metadata - error_metadata = SubAgentMetadata( - message_history=[], - run_metadata={}, + return ToolResult( + content=result_content, + metadata=SubAgentMetadata(message_history=msg_history, run_metadata=run_metadata), ) + except Exception as exc: return ToolResult( - content=f"\n{e!s}\n", + content=f"\n{exc!s}\n", success=False, - metadata=error_metadata, + metadata=SubAgentMetadata(message_history=[], run_metadata={}), ) - finally: - # DEBUG: Log SESSION_STATE after subagent session - post_session_state = _SESSION_STATE.get(None) - logger.debug( - "[%s] POST-SESSION: _SESSION_STATE=%s, exec_env=%s, exec_env._temp_dir=%s", - agent.name, - id(post_session_state) if post_session_state else None, - type(post_session_state.exec_env).__name__ - if post_session_state and post_session_state.exec_env - else None, - getattr(post_session_state.exec_env, "_temp_dir", "N/A") - if post_session_state and post_session_state.exec_env - else None, - ) - - # Restore parent's depth - _PARENT_DEPTH.reset(prev_depth) - # Restore parent's session state so next sibling subagent sees it - if parent_session_state is not None: - _SESSION_STATE.set(parent_session_state) - return Tool[SubAgentParams, SubAgentMetadata]( + sub_agent_tool = Tool[SubAgentParams, SubAgentMetadata]( name=self._name, description=description, parameters=SubAgentParams, executor=sub_agent_executor, # ty: ignore[invalid-argument-type] ) + return sub_agent_tool class SessionAgent[FinishParams: BaseModel, FinishMeta](Agent[FinishParams, FinishMeta]): - """Agent running inside an active session with full tool access. + """Detached runtime returned by ``Agent.session()``. - Returned by ``async with agent.session(...) as session``. A SessionAgent - shares its ``__dict__`` with the parent Agent, so it has the same - configuration and state. The key difference is that ToolProvider-created - tools have been initialized and are available in ``_active_tools``. - - ``Agent.run()`` raises ``RuntimeError`` when ToolProviders are present - but the caller is not a SessionAgent. This makes the invalid state - (calling ``run()`` with uninitialized ToolProviders) a type-level - distinction rather than just a runtime flag. + Agent configuration, the LLM client, and static tools are shared. Active + tools, provider and logger lifecycle, run results, cache state, skills, and + output state belong to this session. """ @classmethod - def from_agent(cls, agent: Agent[FinishParams, FinishMeta]) -> "SessionAgent[FinishParams, FinishMeta]": - """Create a SessionAgent sharing the given Agent's ``__dict__``.""" - sa: SessionAgent[FinishParams, FinishMeta] = object.__new__(cls) - sa.__dict__ = agent.__dict__ - return sa + def from_agent( + cls, + agent: Agent[FinishParams, FinishMeta], + *, + output_dir: Path | str | None = None, + input_files: str | Path | list[str | Path] | None = None, + skills_dir: Path | str | None = None, + resume: bool = False, + clear_cache_on_success: bool = True, + cache_on_interrupt: bool = True, + parent_session_state: SessionState | None = None, + ) -> "SessionAgent[FinishParams, FinishMeta]": + """Shallow-copy configuration and replace all mutable runtime owners.""" + from stirrup.tools.view_image import ViewImageToolProvider + + configured_tools = agent._tools # noqa: SLF001 + configured_providers = {id(tool): tool for tool in configured_tools if isinstance(tool, ToolProvider)} + sequential_provider_ids = { + id(tool._exec_env) # noqa: SLF001 + for tool in configured_providers.values() + if isinstance(tool, ViewImageToolProvider) + and type(tool) is not ViewImageToolProvider + and tool._exec_env is not None # noqa: SLF001 + and id(tool._exec_env) in configured_providers # noqa: SLF001 + } + + replacements: dict[int, ToolProvider] = {} + should_share_exec_env = ( + agent._share_parent_exec_env # noqa: SLF001 + and parent_session_state is not None + and parent_session_state.exec_env is not None + ) + if should_share_exec_env and sequential_provider_ids: + raise ValueError( + f"Subagent '{agent._name}' cannot use share_parent_exec_env=True with a custom " # noqa: SLF001 + "ViewImageToolProvider configured for its own code backend. Disable environment sharing " + "or use the exact built-in ViewImageToolProvider." + ) + + for configured_tool in agent._tools: # noqa: SLF001 + if not isinstance(configured_tool, ToolProvider): + continue + if type(configured_tool) is ViewImageToolProvider: + continue + if id(configured_tool) in sequential_provider_ids: + replacements[id(configured_tool)] = configured_tool + continue + if should_share_exec_env and isinstance(configured_tool, CodeExecToolProvider): + replacements[id(configured_tool)] = parent_session_state.exec_env # type: ignore[assignment] + continue + replacements[id(configured_tool)] = ( + _detach_builtin_provider(configured_tool, replacements) or configured_tool + ) + + session_tools: list[Tool | ToolProvider] = [] + custom_resources: list[object] = [] + for configured_tool in agent._tools: # noqa: SLF001 + if not isinstance(configured_tool, ToolProvider): + session_tools.append(configured_tool) + continue + session_tool = replacements.get(id(configured_tool)) + if session_tool is None: + session_tool = _detach_builtin_provider(configured_tool, replacements) or configured_tool + replacements[id(configured_tool)] = session_tool + session_tools.append(session_tool) + if session_tool is configured_tool: + custom_resources.append(configured_tool) + + session_logger, is_custom_logger = _detach_logger(agent._logger) # noqa: SLF001 + if is_custom_logger: + custom_resources.append(session_logger) + custom_resources = list({id(resource): resource for resource in custom_resources}.values()) + + session: SessionAgent[FinishParams, FinishMeta] = object.__new__(cls) + session.__dict__ = agent.__dict__.copy() + session._tools = session_tools + session._finish_tools = dict(agent._finish_tools) # noqa: SLF001 + session._active_tools = {tool.name: tool for tool in session_tools if isinstance(tool, Tool)} + session._active_tools.update(session._finish_tools) + session._has_tool_providers = any(isinstance(tool, ToolProvider) for tool in session_tools) + session._logger = session_logger + + session._pending_output_dir = Path(output_dir) if output_dir else None + session._pending_input_files = list(input_files) if isinstance(input_files, list) else input_files + session._pending_skills_dir = Path(skills_dir) if skills_dir else None + session._resume = resume + session._clear_cache_on_success = clear_cache_on_success + session._cache_on_interrupt = cache_on_interrupt + + session._last_finish_params = None + session._last_run_metadata = {} + session._transferred_paths = [] + session._current_task_hash = None + session._current_run_state = None + session._parent_session_state = parent_session_state + session._session_state = None + session._session_state_token = None + session._interrupt_handler_installed = False + session._custom_session_resources = custom_resources + session._custom_resources_reserved = False + session._reserved_cache_task_hashes = set() + return session diff --git a/src/stirrup/integrations/slack/slack.py b/src/stirrup/integrations/slack/slack.py index e42b988..009bd2b 100644 --- a/src/stirrup/integrations/slack/slack.py +++ b/src/stirrup/integrations/slack/slack.py @@ -348,10 +348,10 @@ def build_agent(self, config: SlackAgentConfig, model_override: str | None = Non client = copy.copy(client) client._model = model_override # type: ignore[attr-defined] # noqa: SLF001 - # Deep copy tools so each concurrent run gets its own ToolProvider instances - # (ToolProviders have lifecycle state that can't be shared across sessions). + # Keep configured coordination objects shared. Agent sessions privately + # detach exact built-in lifecycle state and guard custom objects from overlap. tools: list[Tool | ToolProvider] = ( - copy.deepcopy(config.tools) + list(config.tools) if config.tools is not None else [ DockerCodeExecToolProvider.from_dockerfile(_DEFAULT_DOCKERFILE), diff --git a/tests/test_agent_sessions.py b/tests/test_agent_sessions.py new file mode 100644 index 0000000..df52a93 --- /dev/null +++ b/tests/test_agent_sessions.py @@ -0,0 +1,1273 @@ +"""Behavioral contracts for isolated Agent sessions.""" + +import json +import signal +from io import StringIO +from pathlib import Path +from typing import Any, Self + +import anyio +import pytest +from PIL import Image +from pydantic import BaseModel +from rich.console import Console + +from stirrup.constants import DEFAULT_FINISH_TOOL_NAME +from stirrup.core.agent import Agent +from stirrup.core.cache import CacheManager, CacheState, compute_task_hash +from stirrup.core.models import ( + AssistantMessage, + ChatMessage, + LLMClient, + SubAgentMetadata, + SystemMessage, + TokenUsage, + Tool, + ToolCall, + ToolMessage, + ToolProvider, + ToolResult, + UserMessage, +) +from stirrup.tools.code_backends.local import LocalCodeExecToolProvider +from stirrup.tools.finish import SIMPLE_FINISH_TOOL +from stirrup.tools.view_image import ViewImageToolProvider +from stirrup.utils.logging import AgentLogger, AgentLoggerBase + + +class ConcurrentPromptClient(LLMClient): + """Finish two overlapping runs from the prompt visible to each call.""" + + def __init__(self) -> None: + self._arrived = {"first prompt": anyio.Event(), "second prompt": anyio.Event()} + + @property + def model_slug(self) -> str: + return "concurrent-prompt-model" + + @property + def max_tokens(self) -> int: + return 100_000 + + async def generate(self, messages: list[ChatMessage], tools: dict[str, Tool]) -> AssistantMessage: + assert DEFAULT_FINISH_TOOL_NAME in tools + prompt = next( + str(message.content) + for message in reversed(messages) + if isinstance(message, UserMessage) and str(message.content) in self._arrived + ) + self._arrived[prompt].set() + other_prompt = "second prompt" if prompt == "first prompt" else "first prompt" + await self._arrived[other_prompt].wait() + return _finish_response(prompt, f"finish-{prompt}") + + +class ConcurrentIdenticalClient(LLMClient): + """Finish only after two calls for the same prompt overlap.""" + + def __init__(self) -> None: + self.call_count = 0 + self.all_arrived = anyio.Event() + + @property + def model_slug(self) -> str: + return "concurrent-identical-model" + + @property + def max_tokens(self) -> int: + return 100_000 + + async def generate(self, messages: list[ChatMessage], tools: dict[str, Tool]) -> AssistantMessage: + del messages + assert DEFAULT_FINISH_TOOL_NAME in tools + self.call_count += 1 + if self.call_count == 2: + self.all_arrived.set() + await self.all_arrived.wait() + return _finish_response("identical", f"finish-identical-{self.call_count}") + + +class BlockingFinishClient(LLMClient): + """Hold a run active until its cache boundary has been inspected.""" + + def __init__(self) -> None: + self.started = anyio.Event() + self.release = anyio.Event() + + @property + def model_slug(self) -> str: + return "blocking-finish-model" + + @property + def max_tokens(self) -> int: + return 100_000 + + async def generate(self, messages: list[ChatMessage], tools: dict[str, Tool]) -> AssistantMessage: + del messages + assert DEFAULT_FINISH_TOOL_NAME in tools + self.started.set() + await self.release.wait() + return _finish_response("finished", "finish-blocked") + + +class ConcurrentFileClient(LLMClient): + """Copy each overlapping session's input into its configured output.""" + + def __init__(self) -> None: + self._arrived = {"first file": anyio.Event(), "second file": anyio.Event()} + + @property + def model_slug(self) -> str: + return "concurrent-file-model" + + @property + def max_tokens(self) -> int: + return 100_000 + + async def generate(self, messages: list[ChatMessage], tools: dict[str, Tool]) -> AssistantMessage: + assert "code_exec" in tools + prompt = next( + str(message.content) + for message in reversed(messages) + if isinstance(message, UserMessage) and str(message.content) in self._arrived + ) + if not any(isinstance(message, ToolMessage) and message.name == "code_exec" for message in messages): + self._arrived[prompt].set() + other_prompt = "second file" if prompt == "first file" else "first file" + await self._arrived[other_prompt].wait() + return AssistantMessage( + content=f"copying {prompt}", + tool_calls=[ + ToolCall( + name="code_exec", + arguments='{"cmd": "cp input.txt result.txt"}', + tool_call_id=f"copy-{prompt}", + ) + ], + token_usage=TokenUsage(), + ) + return AssistantMessage( + content=f"finished {prompt}", + tool_calls=[ + ToolCall( + name=DEFAULT_FINISH_TOOL_NAME, + arguments=json.dumps({"reason": prompt, "paths": ["result.txt"]}), + tool_call_id=f"finish-{prompt}", + ) + ], + token_usage=TokenUsage(), + ) + + +class ScriptedClient(LLMClient): + def __init__(self, responses: list[AssistantMessage | Exception]) -> None: + self.responses = responses + self.messages_seen: list[list[ChatMessage]] = [] + + @property + def model_slug(self) -> str: + return "scripted-model" + + @property + def max_tokens(self) -> int: + return 100_000 + + async def generate(self, messages: list[ChatMessage], tools: dict[str, Tool]) -> AssistantMessage: + del tools + self.messages_seen.append(list(messages)) + response = self.responses.pop(0) + if isinstance(response, Exception): + raise response + return response + + +class CustomProvider(ToolProvider): + """Stateful custom provider used to verify sequential-only ownership.""" + + def __init__(self) -> None: + self.enter_count = 0 + self.exit_count = 0 + self.is_active = False + + async def __aenter__(self) -> Tool: + self.enter_count += 1 + self.is_active = True + + def inspect(_params: BaseModel) -> ToolResult: + return ToolResult(content=str(self.enter_count), success=self.is_active) + + return Tool(name="inspect_custom", description="Inspect custom state", executor=inspect) + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object, + ) -> None: + del exc_type, exc_val, exc_tb + self.exit_count += 1 + self.is_active = False + + +class BlockingCleanupProvider(CustomProvider): + """Custom provider whose asynchronous cleanup is externally released.""" + + def __init__(self) -> None: + super().__init__() + self.cleanup_started = anyio.Event() + self.cleanup_allowed = anyio.Event() + self.cleanup_finished = anyio.Event() + self.first_exit_exception: BaseException | None = None + self.first_cleanup_error = RuntimeError("cancelled cleanup failed") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object, + ) -> None: + if self.enter_count == 1: + self.first_exit_exception = exc_val + self.cleanup_started.set() + await self.cleanup_allowed.wait() + await super().__aexit__(exc_type, exc_val, exc_tb) + self.cleanup_finished.set() + if self.exit_count == 1: + raise self.first_cleanup_error + + +class CleanupFailingProvider(CustomProvider): + """Custom provider that records cleanup before failing it.""" + + def __init__(self) -> None: + super().__init__() + self.exit_error = RuntimeError("provider cleanup failed") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object, + ) -> None: + await super().__aexit__(exc_type, exc_val, exc_tb) + raise self.exit_error + + +class CountingViewImageProvider(ViewImageToolProvider): + """Custom view-image provider used to verify dependent backend ownership.""" + + def __init__(self, exec_env: LocalCodeExecToolProvider) -> None: + super().__init__(exec_env) + self.enter_count = 0 + + async def __aenter__(self) -> Tool: + self.enter_count += 1 + return await super().__aenter__() + + +class CustomLogger(AgentLoggerBase): + """Minimal stateful logger used to verify sequential-only ownership.""" + + def __init__(self) -> None: + self.name = "agent" + self.model: str | None = None + self.max_turns: int | None = None + self.depth = 0 + self.finish_params: BaseModel | None = None + self.run_metadata: dict[str, list[Any]] | None = None + self.output_dir: str | None = None + self.enter_count = 0 + self.exit_count = 0 + + def __enter__(self) -> Self: + self.enter_count += 1 + return self + + def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object) -> None: + del exc_type, exc_val, exc_tb + self.exit_count += 1 + + def on_step(self, step: int, tool_calls: int = 0, input_tokens: int = 0, output_tokens: int = 0) -> None: + del step, tool_calls, input_tokens, output_tokens + + def assistant_message(self, turn: int, max_turns: int, assistant_message: AssistantMessage) -> None: + del turn, max_turns, assistant_message + + def user_message(self, user_message: UserMessage) -> None: + del user_message + + def task_message(self, task: str | list[Any]) -> None: + del task + + def tool_result(self, tool_message: ToolMessage) -> None: + del tool_message + + def context_summarization_start(self, pct_used: float, cutoff: float) -> None: + del pct_used, cutoff + + def context_summarization_complete(self, summary: str, bridge: str) -> None: + del summary, bridge + + def debug(self, message: str, *args: object) -> None: + del message, args + + def info(self, message: str, *args: object) -> None: + del message, args + + def warning(self, message: str, *args: object) -> None: + del message, args + + def error(self, message: str, *args: object) -> None: + del message, args + + +class CleanupFailingLogger(CustomLogger): + """Custom logger that records cleanup before failing it.""" + + def __init__(self) -> None: + super().__init__() + self.exit_error = RuntimeError("logger cleanup failed") + + def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object) -> None: + super().__exit__(exc_type, exc_val, exc_tb) + raise self.exit_error + + +class WarningRecordingLogger(AgentLogger): + """Built-in logger subclass that records warnings without rendering them.""" + + def __init__(self) -> None: + super().__init__(show_spinner=False) + self.warnings_seen: list[str] = [] + + def warnings_message(self, warnings: list[str]) -> None: + self.warnings_seen.extend(warnings) + + +def _finish_response(reason: str, tool_call_id: str) -> AssistantMessage: + return AssistantMessage( + content=f"finished: {reason}", + tool_calls=[ + ToolCall( + name=DEFAULT_FINISH_TOOL_NAME, + arguments=json.dumps({"reason": reason, "paths": []}), + tool_call_id=tool_call_id, + ) + ], + token_usage=TokenUsage(), + ) + + +def _write_skill(skills_dir: Path, name: str) -> None: + skill_dir = skills_dir / name + skill_dir.mkdir(parents=True) + (skill_dir / "SKILL.md").write_text( + f"---\nname: {name}\ndescription: Available to delegated work.\n---\n\n# {name}\n" + ) + + +def _quiet_logger() -> AgentLogger: + return AgentLogger(show_spinner=False) + + +async def test_concurrent_exact_builtin_sessions_isolate_inputs_outputs_and_run_state(tmp_path: Path) -> None: + agent = Agent( + client=ConcurrentFileClient(), + name="file-agent", + max_turns=2, + tools=[LocalCodeExecToolProvider()], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + inputs: dict[str, Path] = {} + outputs: dict[str, Path] = {} + for prompt, content in (("first file", "FIRST-DATA"), ("second file", "SECOND-DATA")): + input_path = tmp_path / prompt / "input.txt" + input_path.parent.mkdir() + input_path.write_text(content) + inputs[prompt] = input_path + outputs[prompt] = tmp_path / f"{prompt}-output" + + async def run_prompt(prompt: str) -> None: + async with agent.session( + input_files=inputs[prompt], output_dir=outputs[prompt], cache_on_interrupt=False + ) as session: + finish_params, history, _ = await session.run(prompt) + assert finish_params is not None + assert finish_params.reason == prompt + assert all( + other not in str(message.content) + for turn in history + for message in turn + for other in ({"first file", "second file"} - {prompt}) + ) + + async with anyio.create_task_group() as task_group: + task_group.start_soon(run_prompt, "first file") + task_group.start_soon(run_prompt, "second file") + + assert (outputs["first file"] / "result.txt").read_text() == "FIRST-DATA" + assert (outputs["second file"] / "result.txt").read_text() == "SECOND-DATA" + + +async def test_concurrent_sessions_with_default_loggers_can_render_while_overlapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + output = StringIO() + monkeypatch.setattr( + "stirrup.utils.logging.console", + Console(file=output, force_terminal=False, width=120), + ) + agent = Agent(client=ScriptedClient([]), name="concurrent-logger", tools=[]) + first_entered = anyio.Event() + second_entered = anyio.Event() + + async def open_first_session() -> None: + async with agent.session(cache_on_interrupt=False): + first_entered.set() + await second_entered.wait() + + async def open_second_session() -> None: + await first_entered.wait() + async with agent.session(cache_on_interrupt=False): + second_entered.set() + + async with anyio.create_task_group() as task_group: + task_group.start_soon(open_first_session) + task_group.start_soon(open_second_session) + + assert output.getvalue().count("concurrent-logger") >= 2 + + +async def test_cached_root_runs_with_different_task_hashes_remain_concurrent( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("stirrup.core.cache.DEFAULT_CACHE_DIR", tmp_path) + agent = Agent( + client=ConcurrentPromptClient(), + name="different-cached-tasks", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + reasons: set[str] = set() + + async def run_prompt(prompt: str) -> None: + async with agent.session() as session: + finish_params, _, _ = await session.run(prompt) + assert finish_params is not None + reasons.add(finish_params.reason) + + async with anyio.create_task_group() as task_group: + task_group.start_soon(run_prompt, "first prompt") + task_group.start_soon(run_prompt, "second prompt") + + assert reasons == {"first prompt", "second prompt"} + + +async def test_identical_root_runs_remain_concurrent_when_interrupt_caching_is_disabled() -> None: + client = ConcurrentIdenticalClient() + agent = Agent( + client=client, + name="identical-uncached-tasks", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + + async def run_prompt() -> None: + async with agent.session(cache_on_interrupt=False) as session: + finish_params, _, _ = await session.run("same prompt") + assert finish_params is not None + + async with anyio.create_task_group() as task_group: + task_group.start_soon(run_prompt) + task_group.start_soon(run_prompt) + + assert client.call_count == 2 + + +async def test_concurrent_identical_cached_root_run_is_rejected_before_cache_io( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("stirrup.core.cache.DEFAULT_CACHE_DIR", tmp_path) + client = BlockingFinishClient() + agent = Agent( + client=client, + name="identical-cached-tasks", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + load_count = 0 + original_load_state = CacheManager.load_state + + def count_load(cache_manager: CacheManager, task_hash: str) -> CacheState | None: + nonlocal load_count + load_count += 1 + return original_load_state(cache_manager, task_hash) + + monkeypatch.setattr(CacheManager, "load_state", count_load) + + async def run_first() -> None: + async with agent.session(resume=True) as session: + await session.run("same cached prompt") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(run_first) + await client.started.wait() + + with pytest.raises(RuntimeError, match="cannot use the same task cache"): + async with agent.session(resume=True) as duplicate: + await duplicate.run("same cached prompt") + assert load_count == 1 + client.release.set() + + async with agent.session(resume=True) as reusable: + await reusable.run("same cached prompt") + + assert load_count == 2 + + +async def test_successful_run_keeps_task_cache_reserved_until_failed_session_teardown( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("stirrup.core.cache.DEFAULT_CACHE_DIR", tmp_path) + body_error = RuntimeError("body failed after run") + client = ScriptedClient( + [ + _finish_response("first", "finish-first"), + _finish_response("after teardown", "finish-after-teardown"), + ] + ) + agent = Agent( + client=client, + name="teardown-cache-owner", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + run_finished = anyio.Event() + fail_body = anyio.Event() + teardown_finished = anyio.Event() + + async def finish_then_fail() -> None: + try: + async with agent.session(clear_cache_on_success=False) as session: + await session.run("same cached prompt") + run_finished.set() + await fail_body.wait() + raise body_error + except RuntimeError as exc: + assert exc is body_error + finally: + teardown_finished.set() + + async with anyio.create_task_group() as task_group: + task_group.start_soon(finish_then_fail) + await run_finished.wait() + + with pytest.raises(RuntimeError, match="cannot use the same task cache"): + async with agent.session() as duplicate: + await duplicate.run("same cached prompt") + + fail_body.set() + await teardown_finished.wait() + + cached = CacheManager(cache_base_dir=tmp_path).load_state(compute_task_hash("same cached prompt")) + assert cached is not None + assert cached.agent_name == "teardown-cache-owner" + + async with agent.session() as reusable: + finish_params, _, _ = await reusable.run("same cached prompt") + + assert finish_params is not None + assert finish_params.reason == "after teardown" + + +async def test_concurrent_calls_to_same_subagent_are_isolated() -> None: + child = Agent( + client=ConcurrentPromptClient(), + name="worker", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + parent = Agent(client=ScriptedClient([]), name="parent", tools=[child.to_tool()], logger=_quiet_logger()) + results: dict[str, ToolMessage] = {} + + async with parent.session(cache_on_interrupt=False) as session: + + async def call_child(prompt: str) -> None: + results[prompt] = await session.run_tool( + ToolCall(name="worker", arguments=json.dumps({"task": prompt}), tool_call_id=f"call-{prompt}"), + {}, + ) + + async with anyio.create_task_group() as task_group: + task_group.start_soon(call_child, "first prompt") + task_group.start_soon(call_child, "second prompt") + + for prompt, other_prompt in (("first prompt", "second prompt"), ("second prompt", "first prompt")): + assert results[prompt].success + assert prompt in str(results[prompt].content) + assert other_prompt not in str(results[prompt].content) + + +async def test_direct_parent_run_cannot_delegate_or_write_subagent_output_to_cwd( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + child_client = ScriptedClient( + [ + AssistantMessage( + content="write an output", + tool_calls=[ + ToolCall( + name="code_exec", + arguments='{"cmd": "printf CHILD > leaked.txt"}', + tool_call_id="write-child-output", + ) + ], + token_usage=TokenUsage(), + ), + AssistantMessage( + content="child finished", + tool_calls=[ + ToolCall( + name=DEFAULT_FINISH_TOOL_NAME, + arguments='{"reason": "child", "paths": ["leaked.txt"]}', + tool_call_id="finish-child-output", + ) + ], + token_usage=TokenUsage(), + ), + ] + ) + child = Agent( + client=child_client, + name="writer", + max_turns=2, + tools=[LocalCodeExecToolProvider()], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + parent = Agent( + client=ScriptedClient( + [ + AssistantMessage( + content="delegate", + tool_calls=[ToolCall(name="writer", arguments='{"task": "write"}', tool_call_id="delegate-writer")], + token_usage=TokenUsage(), + ), + _finish_response("parent handled error", "finish-parent"), + ] + ), + name="direct-parent", + max_turns=2, + tools=[child.to_tool()], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + + finish_params, history, _ = await parent.run("delegate without a session") + + delegated_result = next( + message for turn in history for message in turn if isinstance(message, ToolMessage) and message.name == "writer" + ) + assert finish_params is not None + assert finish_params.reason == "parent handled error" + assert not delegated_result.success + assert "requires an active parent Agent session" in str(delegated_result.content) + assert "async with parent.session()" in str(delegated_result.content) + assert len(child_client.responses) == 2 + assert not (tmp_path / "leaked.txt").exists() + + +async def test_provider_free_base_run_releases_task_cache_reservation_after_failure() -> None: + run_error = RuntimeError("direct run failed") + agent = Agent( + client=ScriptedClient([run_error, _finish_response("reused", "finish-reused")]), + name="direct-reusable", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + + with pytest.raises(RuntimeError, match="direct run failed") as exc_info: + await agent.run("same direct prompt") + finish_params, _, _ = await agent.run("same direct prompt") + + assert exc_info.value is run_error + assert finish_params is not None + assert finish_params.reason == "reused" + + +async def test_provider_free_base_run_is_independent_inside_an_active_session(tmp_path: Path) -> None: + skills_dir = tmp_path / "skills" + _write_skill(skills_dir, "outer_skill") + input_path = tmp_path / "outer_input.txt" + input_path.write_text("outer data") + + inner_failure = RuntimeError("inner run failed") + inner_client = ScriptedClient( + [ + AssistantMessage( + content="finish independently", + tool_calls=[ + ToolCall( + name=DEFAULT_FINISH_TOOL_NAME, + arguments='{"reason": "inner", "paths": ["missing-from-outer.txt"]}', + tool_call_id="finish-inner", + ) + ], + token_usage=TokenUsage(), + ), + inner_failure, + ] + ) + inner_logger = WarningRecordingLogger() + inner = Agent( + client=inner_client, + name="inner", + max_turns=1, + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=inner_logger, + ) + outer = Agent( + client=ScriptedClient([]), + name="outer", + tools=[LocalCodeExecToolProvider()], + logger=_quiet_logger(), + ) + + async with outer.session( + input_files=input_path, + skills_dir=skills_dir, + cache_on_interrupt=False, + ) as outer_session: + finish_params, _, _ = await inner.run("independent work") + with pytest.raises(RuntimeError, match="inner run failed") as exc_info: + await inner.run("failing independent work") + outer_result = await outer_session.run_tool( + ToolCall(name="code_exec", arguments='{"cmd": "printf OUTER"}', tool_call_id="outer-restored"), + {}, + ) + + assert finish_params is not None + assert finish_params.reason == "inner" + assert exc_info.value is inner_failure + assert outer_result.success + assert "OUTER" in str(outer_result.content) + inner_system_prompt = next( + message for message in inner_client.messages_seen[0] if isinstance(message, SystemMessage) + ) + assert "outer_input.txt" not in str(inner_system_prompt.content) + assert "outer_skill" not in str(inner_system_prompt.content) + assert all("no output_dir is set" not in warning for warning in inner_logger.warnings_seen) + + +async def test_nested_root_is_independent_while_true_subagent_inherits_skills(tmp_path: Path) -> None: + skills_dir = tmp_path / "skills" + _write_skill(skills_dir, "shared_skill") + child_client = ScriptedClient([_finish_response("child", "finish-child")]) + root_client = ScriptedClient([_finish_response("root", "finish-root")]) + child = Agent( + client=child_client, + name="skill_child", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + parent = Agent(client=ScriptedClient([]), name="parent", tools=[child.to_tool()], logger=_quiet_logger()) + independent_root = Agent( + client=root_client, + name="independent", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + + async with parent.session(skills_dir=skills_dir, cache_on_interrupt=False) as parent_session: + async with independent_root.session(cache_on_interrupt=False) as root_session: + await root_session.run("independent work") + child_result = await parent_session.run_tool( + ToolCall(name="skill_child", arguments='{"task": "delegated work"}', tool_call_id="delegate"), + {}, + ) + + assert child_result.success + child_system = [message for message in child_client.messages_seen[0] if isinstance(message, SystemMessage)] + root_system = [message for message in root_client.messages_seen[0] if isinstance(message, SystemMessage)] + assert any("shared_skill" in str(message.content) for message in child_system) + assert all("shared_skill" not in str(message.content) for message in root_system) + + +async def test_session_run_and_run_tool_require_their_exact_active_ambient_context() -> None: + client = ScriptedClient([]) + tool_execution_count = 0 + + def count_execution(_params: BaseModel) -> ToolResult: + nonlocal tool_execution_count + tool_execution_count += 1 + return ToolResult(content="executed") + + agent = Agent( + client=client, + name="guarded", + tools=[Tool(name="count_execution", description="Count executions", executor=count_execution)], + logger=_quiet_logger(), + ) + first = agent.session(cache_on_interrupt=False) + tool_call = ToolCall(name="count_execution", arguments="{}", tool_call_id="count") + + with pytest.raises(RuntimeError, match="own active session context"): + await first.run("before") + with pytest.raises(RuntimeError, match="own active session context"): + await first.run_tool(tool_call, {}) + + async with first: + valid_result = await first.run_tool(tool_call, {}) + async with agent.session(cache_on_interrupt=False): + with pytest.raises(RuntimeError, match="own active session context"): + await first.run("wrong ambient") + with pytest.raises(RuntimeError, match="own active session context"): + await first.run_tool(tool_call, {}) + restored_result = await first.run_tool(tool_call, {}) + + with pytest.raises(RuntimeError, match="own active session context"): + await first.run("after") + with pytest.raises(RuntimeError, match="own active session context"): + await first.run_tool(tool_call, {}) + + assert valid_result.success + assert restored_result.success + assert tool_execution_count == 2 + assert client.messages_seen == [] + + +async def test_failed_session_entry_restores_outer_context_and_releases_custom_resources() -> None: + provider = CustomProvider() + failing = Agent( + client=ScriptedClient([]), + name="failing", + tools=[provider, ViewImageToolProvider()], + logger=_quiet_logger(), + ) + outer = Agent( + client=ScriptedClient([_finish_response("outer", "finish-outer")]), + name="outer", + tools=[], + finish_tool=SIMPLE_FINISH_TOOL, + logger=_quiet_logger(), + ) + + async with outer.session(cache_on_interrupt=False) as outer_session: + with pytest.raises(RuntimeError, match="requires a CodeExecToolProvider"): + async with failing.session(cache_on_interrupt=False): + pass + finish_params, _, _ = await outer_session.run("outer still active") + + assert finish_params is not None + assert finish_params.reason == "outer" + assert (provider.enter_count, provider.exit_count) == (1, 1) + reusable = Agent(client=ScriptedClient([]), name="reusable", tools=[provider], logger=_quiet_logger()) + async with reusable.session(cache_on_interrupt=False): + pass + + +async def test_custom_provider_and_logger_remain_reusable_sequentially() -> None: + provider = CustomProvider() + custom_logger = CustomLogger() + agent = Agent(client=ScriptedClient([]), name="custom", tools=[provider], logger=custom_logger) + + for _ in range(2): + async with agent.session(cache_on_interrupt=False) as session: + result = await session.run_tool( + ToolCall(name="inspect_custom", arguments="{}", tool_call_id="inspect"), + {}, + ) + assert result.success + + assert (provider.enter_count, provider.exit_count) == (2, 2) + assert (custom_logger.enter_count, custom_logger.exit_count) == (2, 2) + + +async def test_cancellation_shields_cleanup_and_holds_custom_reservation_until_it_finishes() -> None: + provider = BlockingCleanupProvider() + agent = Agent(client=ScriptedClient([]), name="cancel-cleanup", tools=[provider], logger=_quiet_logger()) + entered = anyio.Event() + cancelled_session_finished = anyio.Event() + cancel_scope: anyio.CancelScope | None = None + + async def cancel_during_session() -> None: + nonlocal cancel_scope + with anyio.CancelScope() as scope: + cancel_scope = scope + async with agent.session(cache_on_interrupt=False): + entered.set() + await anyio.sleep_forever() + cancelled_session_finished.set() + + async with anyio.create_task_group() as task_group: + task_group.start_soon(cancel_during_session) + await entered.wait() + assert cancel_scope is not None + cancel_scope.cancel() + await provider.cleanup_started.wait() + + overlapping = agent.session(cache_on_interrupt=False) + try: + with pytest.raises(RuntimeError, match="sequential sessions only"): + await overlapping.__aenter__() + finally: + provider.cleanup_allowed.set() + if overlapping._session_state is not None: # noqa: SLF001 + await overlapping.__aexit__(None, None, None) + + await provider.cleanup_finished.wait() + await cancelled_session_finished.wait() + + async with agent.session(cache_on_interrupt=False): + pass + + assert (provider.enter_count, provider.exit_count) == (2, 2) + assert isinstance(provider.first_exit_exception, anyio.get_cancelled_exc_class()) + assert any( + "Tool provider failed during session cleanup" in note for note in provider.first_exit_exception.__notes__ + ) + + +async def test_shared_custom_provider_overlap_across_agents_is_rejected_before_entry() -> None: + provider = CustomProvider() + first = Agent(client=ScriptedClient([]), name="first", tools=[provider], logger=_quiet_logger()) + second = Agent(client=ScriptedClient([]), name="second", tools=[provider], logger=_quiet_logger()) + + async with first.session(cache_on_interrupt=False): + with pytest.raises(RuntimeError, match="sequential sessions only"): + async with second.session(cache_on_interrupt=False): + pass + assert provider.enter_count == 1 + + assert provider.exit_count == 1 + + +async def test_shared_custom_logger_overlap_across_agents_is_rejected_before_entry() -> None: + custom_logger = CustomLogger() + first = Agent(client=ScriptedClient([]), name="first", tools=[], logger=custom_logger) + second = Agent(client=ScriptedClient([]), name="second", tools=[], logger=custom_logger) + + async with first.session(cache_on_interrupt=False): + with pytest.raises(RuntimeError, match="sequential sessions only"): + async with second.session(cache_on_interrupt=False): + pass + assert custom_logger.enter_count == 1 + + assert custom_logger.exit_count == 1 + + +async def test_view_image_uses_each_sessions_active_code_backend() -> None: + configured_exec = LocalCodeExecToolProvider() + agent = Agent( + client=ScriptedClient([]), + name="images", + tools=[configured_exec, ViewImageToolProvider(configured_exec)], + logger=_quiet_logger(), + ) + arrived = {"red": anyio.Event(), "blue": anyio.Event()} + session_execs: dict[str, LocalCodeExecToolProvider] = {} + + async def view_in_session(color: str, other_color: str) -> None: + async with agent.session(cache_on_interrupt=False) as session: + session_tools = session._tools # noqa: SLF001 + session_exec = next(tool for tool in session_tools if isinstance(tool, LocalCodeExecToolProvider)) + session_execs[color] = session_exec + arrived[color].set() + await arrived[other_color].wait() + + image_path = session_exec.temp_dir / "pixel.png" # type: ignore[operator] + Image.new("RGB", (1, 1), color=color).save(image_path) + result = await session.run_tool( + ToolCall(name="view_image", arguments='{"path": "pixel.png"}', tool_call_id=f"view-{color}"), + {}, + ) + assert result.success + + async with anyio.create_task_group() as task_group: + task_group.start_soon(view_in_session, "red", "blue") + task_group.start_soon(view_in_session, "blue", "red") + + assert session_execs["red"] is not session_execs["blue"] + + +async def test_custom_view_image_provider_and_exact_backend_are_reusable_sequentially() -> None: + configured_exec = LocalCodeExecToolProvider() + view_images = CountingViewImageProvider(configured_exec) + agent = Agent( + client=ScriptedClient([]), + name="custom-images", + tools=[configured_exec, view_images], + logger=_quiet_logger(), + ) + + for index, color in enumerate(("red", "blue")): + async with agent.session(cache_on_interrupt=False) as session: + session_tools = session._tools # noqa: SLF001 + session_exec = next(tool for tool in session_tools if isinstance(tool, LocalCodeExecToolProvider)) + assert session_exec is configured_exec + image_path = configured_exec.temp_dir / "pixel.png" # type: ignore[operator] + Image.new("RGB", (1, 1), color=color).save(image_path) + result = await session.run_tool( + ToolCall(name="view_image", arguments='{"path": "pixel.png"}', tool_call_id=f"view-{index}"), + {}, + ) + assert result.success + + assert view_images.enter_count == 2 + assert configured_exec.temp_dir is None + + +async def test_custom_view_image_pair_rejects_shared_backend_overlap_before_entry() -> None: + shared_exec = LocalCodeExecToolProvider() + first_view_images = CountingViewImageProvider(shared_exec) + second_view_images = CountingViewImageProvider(shared_exec) + first = Agent( + client=ScriptedClient([]), + name="first-images", + tools=[shared_exec, first_view_images], + logger=_quiet_logger(), + ) + second = Agent( + client=ScriptedClient([]), + name="second-images", + tools=[shared_exec, second_view_images], + logger=_quiet_logger(), + ) + + async with first.session(cache_on_interrupt=False): + with pytest.raises(RuntimeError, match="sequential sessions only"): + async with second.session(cache_on_interrupt=False): + pass + assert first_view_images.enter_count == 1 + assert second_view_images.enter_count == 0 + + +async def test_shared_subagent_rejects_custom_view_image_backend_pair_before_child_entry() -> None: + child_exec = LocalCodeExecToolProvider() + view_images = CountingViewImageProvider(child_exec) + child = Agent( + client=ScriptedClient([]), + name="custom_shared_child", + tools=[child_exec, view_images], + share_parent_exec_env=True, + logger=_quiet_logger(), + ) + parent = Agent( + client=ScriptedClient([]), + name="custom_shared_parent", + tools=[LocalCodeExecToolProvider(), child.to_tool()], + logger=_quiet_logger(), + ) + + async with parent.session(cache_on_interrupt=False) as session: + child_result = await session.run_tool( + ToolCall(name="custom_shared_child", arguments='{"task": "view it"}', tool_call_id="delegate-custom"), + {}, + ) + parent_result = await session.run_tool( + ToolCall(name="code_exec", arguments='{"cmd": "printf PARENT"}', tool_call_id="parent-after-reject"), + {}, + ) + + assert not child_result.success + assert "cannot use share_parent_exec_env=True" in str(child_result.content) + assert view_images.enter_count == 0 + assert child_exec.temp_dir is None + assert parent_result.success + + +async def test_shared_subagent_view_image_uses_parent_code_backend(tmp_path: Path) -> None: + pixel_path = tmp_path / "pixel.png" + Image.new("RGB", (1, 1), color="red").save(pixel_path) + child_exec = LocalCodeExecToolProvider() + child = Agent( + client=ScriptedClient( + [ + AssistantMessage( + content="view shared image", + tool_calls=[ + ToolCall(name="view_image", arguments='{"path": "pixel.png"}', tool_call_id="view-shared") + ], + token_usage=TokenUsage(), + ), + _finish_response("shared", "finish-shared"), + ] + ), + name="shared_child", + max_turns=2, + tools=[child_exec, ViewImageToolProvider(child_exec)], + finish_tool=SIMPLE_FINISH_TOOL, + share_parent_exec_env=True, + logger=_quiet_logger(), + ) + parent = Agent( + client=ScriptedClient([]), + name="shared_parent", + tools=[LocalCodeExecToolProvider(), child.to_tool()], + logger=_quiet_logger(), + ) + run_metadata: dict[str, list[Any]] = {} + + async with parent.session(input_files=pixel_path, cache_on_interrupt=False) as session: + child_result = await session.run_tool( + ToolCall(name="shared_child", arguments='{"task": "view it"}', tool_call_id="delegate-shared"), + run_metadata, + ) + parent_result = await session.run_tool( + ToolCall(name="code_exec", arguments='{"cmd": "printf PARENT"}', tool_call_id="parent-still-active"), + {}, + ) + + assert child_result.success + assert parent_result.success + assert "PARENT" in str(parent_result.content) + child_metadata = run_metadata["shared_child"][0] + assert isinstance(child_metadata, SubAgentMetadata) + view_results = [ + message + for turn in child_metadata.message_history + for message in turn + if isinstance(message, ToolMessage) and message.name == "view_image" + ] + assert len(view_results) == 1 + assert view_results[0].success + + +async def test_partial_entry_cleanup_failures_do_not_mask_the_entry_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = CleanupFailingProvider() + cleanup_logger = CleanupFailingLogger() + entry_error = RuntimeError("session entry failed") + + def fail_interrupt_installation() -> bool: + raise entry_error + + monkeypatch.setattr("stirrup.core.agent._install_interrupt_handler", fail_interrupt_installation) + agent = Agent( + client=ScriptedClient([]), + name="entry-cleanup", + tools=[provider], + logger=cleanup_logger, + ) + + with pytest.raises(RuntimeError, match="session entry failed") as exc_info: + async with agent.session(): + pass + + assert exc_info.value is entry_error + assert (provider.enter_count, provider.exit_count) == (1, 1) + assert (cleanup_logger.enter_count, cleanup_logger.exit_count) == (1, 1) + assert any("Logger failed during session cleanup" in note for note in entry_error.__notes__) + assert any("Tool provider failed during session cleanup" in note for note in entry_error.__notes__) + + +async def test_body_failure_survives_all_cleanup_failures() -> None: + provider = CleanupFailingProvider() + cleanup_logger = CleanupFailingLogger() + body_error = RuntimeError("session body failed") + agent = Agent( + client=ScriptedClient([]), + name="body-cleanup", + tools=[provider], + logger=cleanup_logger, + ) + + with pytest.raises(RuntimeError, match="session body failed") as exc_info: + async with agent.session(cache_on_interrupt=False): + raise body_error + + assert exc_info.value is body_error + assert (provider.enter_count, provider.exit_count) == (1, 1) + assert (cleanup_logger.enter_count, cleanup_logger.exit_count) == (1, 1) + assert any("Logger failed during session cleanup" in note for note in body_error.__notes__) + assert any("Tool provider failed during session cleanup" in note for note in body_error.__notes__) + + +async def test_cleanup_failure_is_raised_when_the_session_has_no_primary_failure() -> None: + provider = CleanupFailingProvider() + cleanup_logger = CleanupFailingLogger() + agent = Agent( + client=ScriptedClient([]), + name="cleanup-only", + tools=[provider], + logger=cleanup_logger, + ) + + with pytest.raises(RuntimeError, match="logger cleanup failed") as exc_info: + async with agent.session(cache_on_interrupt=False): + pass + + assert exc_info.value is cleanup_logger.exit_error + assert provider.exit_count == 1 + assert any("Tool provider failed during session cleanup" in note for note in exc_info.value.__notes__) + + +async def test_worker_thread_failure_is_cached_without_masking_original_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("stirrup.core.cache.DEFAULT_CACHE_DIR", tmp_path) + original_error = RuntimeError("client failed") + prompt = "cache this worker failure" + + def run_in_worker() -> None: + async def fail_run() -> None: + agent = Agent( + client=ScriptedClient([original_error]), + name="worker-cache", + tools=[], + logger=_quiet_logger(), + ) + async with agent.session() as session: + await session.run(prompt) + + anyio.run(fail_run) + + with pytest.raises(RuntimeError, match="client failed") as exc_info: + await anyio.to_thread.run_sync(run_in_worker) # ty: ignore[unresolved-attribute] + + assert exc_info.value is original_error + cached = CacheManager(cache_base_dir=tmp_path).load_state(compute_task_hash(prompt)) + assert cached is not None + assert cached.agent_name == "worker-cache" + + +async def test_overlapping_root_sessions_restore_signal_handler_after_last_exit() -> None: + agent = Agent(client=ScriptedClient([]), name="signals", tools=[], logger=_quiet_logger()) + first_entered = anyio.Event() + second_entered = anyio.Event() + first_exited = anyio.Event() + original_handler = signal.getsignal(signal.SIGINT) + installed_handler: object | None = None + + async def first_session() -> None: + nonlocal installed_handler + async with agent.session(): + installed_handler = signal.getsignal(signal.SIGINT) + first_entered.set() + await second_entered.wait() + first_exited.set() + + async def second_session() -> None: + await first_entered.wait() + async with agent.session(): + second_entered.set() + await first_exited.wait() + assert signal.getsignal(signal.SIGINT) is installed_handler + + async with anyio.create_task_group() as task_group: + task_group.start_soon(first_session) + task_group.start_soon(second_session) + + assert signal.getsignal(signal.SIGINT) is original_handler diff --git a/uv.lock b/uv.lock index 4af8f53..24408a8 100644 --- a/uv.lock +++ b/uv.lock @@ -5723,7 +5723,7 @@ requires-dist = [ { name = "pillow", specifier = ">=10.4.0" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "python-dotenv", marker = "extra == 'docker'", specifier = ">=1.0.0" }, - { name = "rich", specifier = ">=13.0.0" }, + { name = "rich", specifier = ">=14.1.0" }, { name = "slack-bolt", marker = "extra == 'slack'", specifier = ">=1.20.0" }, { name = "stirrup", extras = ["e2b"], marker = "extra == 'e2b-throttle'" }, { name = "stirrup", extras = ["litellm", "e2b", "e2b-throttle", "docker", "mcp", "browser", "slack"], marker = "extra == 'all'" }, @@ -5739,7 +5739,7 @@ dev = [ { name = "mkdocstrings", specifier = ">=0.30.1" }, { name = "mkdocstrings-python", specifier = ">=1.19.0" }, { name = "pytest", specifier = ">=9.0.0" }, - { name = "rich", specifier = ">=13.0.0" }, + { name = "rich", specifier = ">=14.1.0" }, { name = "ruff", specifier = ">=0.14.4" }, { name = "ty", specifier = ">=0.0.1a32" }, ]