From 854585472fc306d506c9e639c554a19dd58a53c4 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Sun, 2 Aug 2026 23:42:34 -0400 Subject: [PATCH 01/11] refactor(tools)!: registry-owned tool identity and declared variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Permission gating and tool assembly resolved tools by ``__name__``, so an unregistered callable named after a genuine tool inherited its capabilities — including EXEMPT — and was handed its services and its long-running contract. Identity is now a per-``ToolRegistry`` ``id(obj) -> (weakref, definition)`` map confirmed with ``is``: no name path, no equality, and a recycled address inherits nothing. Everything the framework issues is bound at its construction site (registered callables, factory service-bound variants, renamed wrappers, the native ADK skill tools). Anything unbound is denied, and left exactly as the application supplied it during assembly. Keeping the map on the instance is also what lets a short-lived registry and its closures be collected. A name may now mean only one thing: ``register()`` raises on any duplicate. Sharing is declared, never inferred — ``declare_tool(name, ...)`` states a contract with no backend-neutral implementation, and each backend registers its own with ``register_tool(..., variant_of=name)``. The ADK and LangGraph save_plan/get_plan/save_tasks/get_tasks are declared once in ``tools/_core/state_tools.py``; previously they contested the name and import order decided the winner. ``canonical_for()`` keeps assembly from ever yielding a declaration's absent ``func``, and ``replace=True`` retires the old definition's identities so they resolve to nothing. ``service_registry.KNOWN_SERVICE_KEYS`` lands here rather than with the tool declarations that use it: ``registry._validate_requires`` imports it at module scope, so the mechanism and its vocabulary cannot be separated. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/tools/__init__.py | 3 + src/agentic_cli/tools/_core/state_tools.py | 51 ++ src/agentic_cli/tools/adk/state_tools.py | 20 +- src/agentic_cli/tools/factories.py | 85 ++- .../tools/langgraph/state_tools.py | 20 +- src/agentic_cli/tools/registry.py | 493 +++++++++++++++++- src/agentic_cli/tools/skills/toolset.py | 21 + src/agentic_cli/tools/tool_resolver.py | 35 +- src/agentic_cli/workflow/adk/manager.py | 19 +- .../workflow/adk/permission_plugin.py | 139 ++++- src/agentic_cli/workflow/base_manager.py | 69 ++- src/agentic_cli/workflow/service_registry.py | 27 + tests/integration/test_permission_adk.py | 34 +- tests/test_packaging.py | 60 +++ tests/tools/test_state_tool_aliases.py | 144 +++++ tests/workflow/test_adk_mcp_permissions.py | 16 +- .../workflow/test_permission_tool_identity.py | 362 +++++++++++++ 17 files changed, 1513 insertions(+), 85 deletions(-) create mode 100644 src/agentic_cli/tools/_core/state_tools.py create mode 100644 tests/tools/test_state_tool_aliases.py create mode 100644 tests/workflow/test_permission_tool_identity.py diff --git a/src/agentic_cli/tools/__init__.py b/src/agentic_cli/tools/__init__.py index b666ae9..063b967 100644 --- a/src/agentic_cli/tools/__init__.py +++ b/src/agentic_cli/tools/__init__.py @@ -7,6 +7,7 @@ - ToolDefinition: Metadata-rich tool definitions - ToolRegistry: Registry for tool management and discovery - register_tool: Decorator for easy tool registration + - declare_tool: Declare a tool implemented only by backend-native variants Framework Tools: - memory_tools: Working and long-term memory tools @@ -78,6 +79,7 @@ ToolDefinition, ToolRegistry, get_registry, + declare_tool, register_tool, ) @@ -90,6 +92,7 @@ "ToolDefinition", "ToolRegistry", "get_registry", + "declare_tool", "register_tool", # Executor classes "SafePythonExecutor", diff --git a/src/agentic_cli/tools/_core/state_tools.py b/src/agentic_cli/tools/_core/state_tools.py new file mode 100644 index 0000000..5e0cf0d --- /dev/null +++ b/src/agentic_cli/tools/_core/state_tools.py @@ -0,0 +1,51 @@ +"""Backend-neutral declarations for the plan/task state tools. + +These four tools have no neutral implementation: reading and writing plan and +task state is inherently backend-native (ADK's ``ToolContext.state``, +LangGraph's graph state and ``Command`` updates), so the signatures and the +model-visible schemas differ per backend. + +What *is* neutral is the contract — the name, the description and the +permission declaration — so it is declared here, once, and each backend +registers its implementation as a variant of it +(``register_tool(..., variant_of="save_plan")``). Without this, the two backend +modules contested the same registry names and whichever imported first decided +what a bare ``"save_plan"`` in an ``AgentConfig`` resolved to. + +Importing either backend's state tools imports this module first, so the +declarations always exist before a variant registers against them. +""" + +from __future__ import annotations + +from agentic_cli.tools.registry import ToolCategory, declare_tool +from agentic_cli.workflow.permissions import EXEMPT + +# Plan/task state is the agent's own scratch space — no external side effects, +# so nothing to gate. The declaration is what both backends share. +_STATE_TOOLS = ( + ( + "save_plan", + "Save or update the execution plan as markdown with checkboxes.", + ), + ("get_plan", "Retrieve the current execution plan."), + ( + "save_tasks", + "Write the complete task list. This replaces the existing list.", + ), + ("get_tasks", "Retrieve the current task list, optionally filtered."), +) + + +def declare_state_tools() -> None: + """Declare the state tools' shared contract. Idempotent.""" + for name, description in _STATE_TOOLS: + declare_tool( + name, + description=description, + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + +declare_state_tools() diff --git a/src/agentic_cli/tools/adk/state_tools.py b/src/agentic_cli/tools/adk/state_tools.py index 00f3b42..dcabd68 100644 --- a/src/agentic_cli/tools/adk/state_tools.py +++ b/src/agentic_cli/tools/adk/state_tools.py @@ -13,11 +13,17 @@ from agentic_cli.tools._core.planning import summarize_checkboxes from agentic_cli.tools._core.tasks import validate_tasks, normalize_tasks, filter_tasks +from agentic_cli.tools._core.state_tools import declare_state_tools from agentic_cli.tools.registry import ToolCategory, register_tool from agentic_cli.workflow.permissions import EXEMPT +# The shared contract must exist before these variants register against it. +declare_state_tools() -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) + +@register_tool( + variant_of="save_plan", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def save_plan(content: str, tool_context: ToolContext) -> dict[str, Any]: """Save or update the execution plan as markdown with checkboxes. @@ -36,7 +42,9 @@ def save_plan(content: str, tool_context: ToolContext) -> dict[str, Any]: return {"success": True, "message": message} -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="get_plan", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def get_plan(tool_context: ToolContext) -> dict[str, Any]: """Retrieve the current execution plan. @@ -51,7 +59,9 @@ def get_plan(tool_context: ToolContext) -> dict[str, Any]: return {"success": True, "content": plan} -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="save_tasks", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def save_tasks( tasks: list[dict[str, Any]], tool_context: ToolContext ) -> dict[str, Any]: @@ -91,7 +101,9 @@ def save_tasks( } -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="get_tasks", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def get_tasks( status: str = "", priority: str = "", diff --git a/src/agentic_cli/tools/factories.py b/src/agentic_cli/tools/factories.py index 3b8ff3f..c836766 100644 --- a/src/agentic_cli/tools/factories.py +++ b/src/agentic_cli/tools/factories.py @@ -18,6 +18,34 @@ from typing import Any, Callable +def _issued(*pairs: "tuple[Callable, Callable]") -> list[Callable]: + """Stamp factory-built closures with the identity of the tool they re-bind. + + Each pair is ``(variant, registered)`` — the closure this factory built, + and the module-level callable it is a service-bound version of. Identity + comes from that *exact* callable, never from a name lookup: if an + application has deliberately taken the name over + (``register_tool(..., replace=True)``), the framework's original is retired + and the variant is simply left unbound, so tool assembly keeps the + application's tool instead of quietly running a different implementation. + + Args: + *pairs: ``(variant, registered)`` pairs, in the order to return. + + Returns: + The variants, in order. + """ + from agentic_cli.tools.registry import bind_tool_identity, identify_tool + + variants: list[Callable] = [] + for variant, registered in pairs: + definition = identify_tool(registered) + if definition is not None: + bind_tool_identity(variant, definition) + variants.append(variant) + return variants + + # --------------------------------------------------------------------------- # Memory tools # --------------------------------------------------------------------------- @@ -87,7 +115,12 @@ def delete_memory( delete_memory.__name__ = "delete_memory" delete_memory.__doc__ = _orig_delete.__doc__ - return [save_memory, search_memory, update_memory, delete_memory] + return _issued( + (save_memory, _orig_save), + (search_memory, _orig_search), + (update_memory, _orig_update), + (delete_memory, _orig_delete), + ) # --------------------------------------------------------------------------- @@ -251,16 +284,16 @@ async def kb_search_concepts( kb_search_concepts.__name__ = "kb_search_concepts" kb_search_concepts.__doc__ = _orig_search_concepts.__doc__ - return [ - kb_search, - kb_ingest_text, - kb_ingest_file, - kb_ingest_url, - kb_read, - kb_list, - kb_write_concept, - kb_search_concepts, - ] + return _issued( + (kb_search, _orig_search), + (kb_ingest_text, _orig_ingest_text), + (kb_ingest_file, _orig_ingest_file), + (kb_ingest_url, _orig_ingest_url), + (kb_read, _orig_read), + (kb_list, _orig_list), + (kb_write_concept, _orig_write_concept), + (kb_search_concepts, _orig_search_concepts), + ) # --------------------------------------------------------------------------- @@ -341,7 +374,9 @@ async def web_fetch(url: str, prompt: str, timeout: int = 30) -> dict[str, Any]: } web_fetch.__name__ = "web_fetch" - return web_fetch + from agentic_cli.tools.webfetch_tool import web_fetch as _orig_web_fetch + + return _issued((web_fetch, _orig_web_fetch))[0] # --------------------------------------------------------------------------- @@ -414,7 +449,9 @@ def sandbox_execute( } sandbox_execute.__name__ = "sandbox_execute" - return sandbox_execute + from agentic_cli.tools.sandbox import sandbox_execute as _orig_sandbox_execute + + return _issued((sandbox_execute, _orig_sandbox_execute))[0] # --------------------------------------------------------------------------- @@ -487,7 +524,15 @@ async def fetch_arxiv_paper(arxiv_id: str) -> dict[str, Any]: search_arxiv.__name__ = "search_arxiv" fetch_arxiv_paper.__name__ = "fetch_arxiv_paper" - return [search_arxiv, fetch_arxiv_paper] + from agentic_cli.tools.arxiv_tools import ( + fetch_arxiv_paper as _orig_fetch_paper, + search_arxiv as _orig_search_arxiv, + ) + + return _issued( + (search_arxiv, _orig_search_arxiv), + (fetch_arxiv_paper, _orig_fetch_paper), + ) def make_ingest_arxiv_tool(arxiv_source, kb_manager) -> Callable: @@ -526,7 +571,11 @@ async def ingest_arxiv_paper( ) ingest_arxiv_paper.__name__ = "ingest_arxiv_paper" - return ingest_arxiv_paper + from agentic_cli.tools.arxiv_tools import ( + ingest_arxiv_paper as _orig_ingest_paper, + ) + + return _issued((ingest_arxiv_paper, _orig_ingest_paper))[0] # --------------------------------------------------------------------------- @@ -589,4 +638,8 @@ async def ask_clarification( } ask_clarification.__name__ = "ask_clarification" - return [ask_clarification] + from agentic_cli.tools.interaction_tools import ( + ask_clarification as _orig_ask_clarification, + ) + + return _issued((ask_clarification, _orig_ask_clarification)) diff --git a/src/agentic_cli/tools/langgraph/state_tools.py b/src/agentic_cli/tools/langgraph/state_tools.py index 9c049f6..336f7c3 100644 --- a/src/agentic_cli/tools/langgraph/state_tools.py +++ b/src/agentic_cli/tools/langgraph/state_tools.py @@ -17,11 +17,17 @@ from agentic_cli.tools._core.planning import summarize_checkboxes from agentic_cli.tools._core.tasks import validate_tasks, normalize_tasks, filter_tasks +from agentic_cli.tools._core.state_tools import declare_state_tools from agentic_cli.tools.registry import ToolCategory, register_tool from agentic_cli.workflow.permissions import EXEMPT +# The shared contract must exist before these variants register against it. +declare_state_tools() -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) + +@register_tool( + variant_of="save_plan", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def save_plan( content: str, tool_call_id: Annotated[str, InjectedToolCallId], @@ -43,7 +49,9 @@ def save_plan( }) -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="get_plan", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def get_plan( state: Annotated[dict, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId], @@ -62,7 +70,9 @@ def get_plan( }) -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="save_tasks", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def save_tasks( tasks: list[dict[str, Any]], tool_call_id: Annotated[str, InjectedToolCallId], @@ -108,7 +118,9 @@ def save_tasks( }) -@register_tool(capabilities=EXEMPT, category=ToolCategory.PLANNING) +@register_tool( + variant_of="get_tasks", capabilities=EXEMPT, category=ToolCategory.PLANNING +) def get_tasks( status: str = "", priority: str = "", diff --git a/src/agentic_cli/tools/registry.py b/src/agentic_cli/tools/registry.py index 22189cb..7b0b723 100644 --- a/src/agentic_cli/tools/registry.py +++ b/src/agentic_cli/tools/registry.py @@ -8,7 +8,9 @@ from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable +import functools import inspect +import weakref from agentic_cli.workflow.permissions.capabilities import ( Capability, @@ -16,6 +18,10 @@ EXEMPT, ) from agentic_cli.workflow.permissions.capabilities import _CapabilityExempt +from agentic_cli.workflow.service_registry import ( + KNOWN_SERVICE_KEYS, + SERVICE_KEY_HINTS, +) class ToolCategory(Enum): @@ -61,22 +67,38 @@ class ToolCategory(Enum): class ToolDefinition: """Metadata-rich tool definition. + ``name`` is the tool's canonical identity: it is what the model calls, what + permission rules match, and what service detection keys on. ``func`` is + guaranteed to expose it as ``__name__`` (see :meth:`ToolRegistry.register`) + so backends that derive a tool name from the callable agree with it. + Attributes: name: Tool name (defaults to function name) description: Human-readable description category: Tool category for organization (READ, WRITE, NETWORK, etc.) capabilities: Capability declarations for the permission engine + requires: Service keys the tool needs at runtime (see + ``service_registry.KNOWN_SERVICE_KEYS``); managers create exactly + these, lazily. is_async: Whether the tool is async - func: The actual tool function + func: The backend-neutral implementation, or **None** for a tool that + only exists as backend-native variants (see ``variants``). A bare + name that resolves to such a tool is an error, not a guess. + variants: Backend-native implementations declared with + ``register(..., variant_of=name)``. They share this tool's identity + and permission metadata, and may differ in signature and docstring + — that is what makes them native. """ name: str description: str - func: Callable[..., Any] + func: Callable[..., Any] | None capabilities: CapabilitiesSpec category: ToolCategory = ToolCategory.OTHER + requires: tuple[str, ...] = () is_async: bool = False long_running: bool = False # tool starts a background job; see tools/jobs/ + variants: tuple[Callable[..., Any], ...] = () def __post_init__(self): """Infer is_async from function.""" @@ -84,6 +106,38 @@ def __post_init__(self): self.is_async = True +def bind_tool_identity(obj: Any, definition: "ToolDefinition") -> None: + """Record, in the default registry, that ``obj`` *is* ``definition``'s tool. + + Called by the registry on registration, and by the framework whenever it + hands a backend something other than the registered callable for the same + tool: a service-bound factory variant, the canonical-name wrapper, or a + backend-native tool object the framework itself constructed (ADK skill + tools). Binding is the only way to acquire capabilities — an object the + framework never issued stays unbound and is gated as unregistered. + + Bindings made by some *other* :class:`ToolRegistry` are deliberately + invisible here: the framework trusts the registry it owns, not one an + application happens to construct. + """ + _default_registry.bind_identity(obj, definition) + + +def identify_tool(obj: Any) -> "ToolDefinition | None": + """Resolve what ``obj`` is, per the default registry, by object identity. + + Strictly this object: no name lookup, no equality, no attribute traversal. + Callers that legitimately need to look *inside* a backend wrapper must + unwrap it themselves, and only for wrapper types they trust (see + ``workflow/adk/permission_plugin.py``). + + Returns: + The bound ``ToolDefinition``, or None when this object was never issued + by the default registry — callers must treat None as "not a tool". + """ + return _default_registry.identify(obj) + + def _validate_capabilities(caps: Any, tool_name: str) -> CapabilitiesSpec: """Validate and return a capabilities value for a tool registration. @@ -113,6 +167,114 @@ def _validate_capabilities(caps: Any, tool_name: str) -> CapabilitiesSpec: ) +def _validate_requires(requires: Any, tool_name: str) -> tuple[str, ...]: + """Validate declared service keys against the constructible service keys. + + A key the manager cannot construct would be a silent no-op: nothing would + be created and the tool would fail at call time with a missing service. The + declarable set is therefore exactly what + ``_ensure_managers_initialized`` knows how to build. + """ + if requires is None: + return () + if isinstance(requires, str): + requires = (requires,) + if not isinstance(requires, (list, tuple, set, frozenset)): + raise TypeError( + f"Tool {tool_name!r}: requires must be a string or a sequence of " + f"service keys, got {type(requires)!r}." + ) + keys = tuple(dict.fromkeys(requires)) # de-duplicate, keep order + for key in keys: + if not isinstance(key, str) or not key.strip(): + raise ValueError( + f"Tool {tool_name!r}: every requires entry must be a non-empty " + f"service-key string, got {key!r}." + ) + unknown = [k for k in keys if k not in KNOWN_SERVICE_KEYS] + if unknown: + hints = " ".join( + SERVICE_KEY_HINTS[k] for k in sorted(unknown) if k in SERVICE_KEY_HINTS + ) + raise ValueError( + f"Tool {tool_name!r}: unknown required service(s): " + f"{', '.join(sorted(unknown))}. Known services: " + f"{', '.join(sorted(KNOWN_SERVICE_KEYS))}." + + (f" {hints}" if hints else "") + ) + return keys + + +def _variant_sort_key(variant: Callable[..., Any]) -> tuple[str, str]: + """Order variants by where they are defined, not by import order.""" + target = getattr(variant, "__wrapped__", variant) + return ( + getattr(target, "__module__", "") or "", + getattr(target, "__qualname__", "") or "", + ) + + +def _ordered_variants( + variants: "tuple[Callable[..., Any], ...]", +) -> "tuple[Callable[..., Any], ...]": + """Deterministic variant order: module-qualified, import-order independent.""" + return tuple(sorted(variants, key=_variant_sort_key)) + + +def _same_declaration( + existing: "ToolDefinition", + capabilities: CapabilitiesSpec, + requires: tuple[str, ...], + category: "ToolCategory", + long_running: bool, +) -> bool: + """Whether a re-registration describes the *same tool*, differently implemented. + + Everything the framework acts on — what it may do, what it needs, whether + it is long-running — must match. Only the callable may differ, which is the + legitimate case of one tool with two backend-native implementations. + """ + return ( + existing.capabilities == capabilities + and existing.requires == requires + and existing.category == category + and existing.long_running == long_running + ) + + +def _with_canonical_name(func: Callable[..., Any], name: str) -> Callable[..., Any]: + """Return a callable whose ``__name__`` is the registered tool name. + + Backends derive the model-visible tool name from ``func.__name__`` (ADK + does), while permission lookup, service detection and tool assembly key on + the registry name. When ``register_tool(name=...)`` renames a tool those two + diverge, and the permission engine then looks up a name that isn't + registered. Wrapping keeps a single identity. + + The wrapper preserves the signature (via ``__wrapped__``), docstring, + annotations and async-ness, so schema generation and the + ``{"success": bool}`` return contract are unaffected. + """ + if getattr(func, "__name__", None) == name: + return func + + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + return await func(*args, **kwargs) + + else: + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + wrapper.__name__ = name + wrapper.__qualname__ = name + return wrapper + + class ToolRegistry: """Registry for managing and discovering tools. @@ -120,10 +282,32 @@ class ToolRegistry: - Tool registration with metadata - Tool lookup by name or category - Tool list generation for agents + - **Identity**: which exact objects *are* the tools it issued + + Identity is owned per registry, not globally. Permission gating and tool + assembly ask the framework's default registry (``get_registry()``), so a + tool registered into some other ``ToolRegistry`` is not one of *its* tools: + it is left as-is during assembly and denied at permission time. Keeping the + map on the instance is also what lets a short-lived registry — with its + definitions and their closures — be garbage collected. """ def __init__(self): self._tools: dict[str, ToolDefinition] = {} + # id(obj) -> (weak ref to obj, the definition it implements). + # + # Not keyed by the object itself: a ``WeakKeyDictionary`` resolves by + # ``hash``/``__eq__``, which any object can define to collide with a + # registered callable. Keyed by ``id`` instead, with every hit + # confirmed by ``is`` against the weak reference, so a recycled address + # cannot inherit a dead object's capabilities. The weak reference also + # drops the entry when the object dies, so per-manager service-bound + # variants do not accumulate. + self._identity: dict[int, tuple[weakref.ref, ToolDefinition]] = {} + # Callables that *were* one of this registry's tools until a + # replace=True took the name over. Weak, and distinct from "never + # registered" — see _retire_identities. + self._retired: "weakref.WeakSet[Any]" = weakref.WeakSet() def register( self, @@ -133,7 +317,10 @@ def register( description: str | None = None, category: ToolCategory = ToolCategory.OTHER, capabilities: CapabilitiesSpec, + requires: str | tuple[str, ...] | list[str] | None = None, long_running: bool = False, + replace: bool = False, + variant_of: str | None = None, ) -> Callable[..., Any]: """Register a tool function. @@ -145,31 +332,269 @@ def my_tool(query: str) -> dict: Or called directly: registry.register(my_tool, category=ToolCategory.READ, capabilities=EXEMPT) + ``requires=`` declares the service keys the tool needs (e.g. + ``requires="kb_manager"``); workflow managers read that metadata to + create exactly those services, lazily. Unknown keys raise. + ``long_running=True`` marks a tool that starts a background job (it should return a ``job_id`` and delegate to ``JobManager``); see ``tools/jobs/``. + + A name means **one tool**. Registering it again raises ``ValueError`` + rather than taking it over: the two definitions would disagree about + what the tool is, and consumers keyed on the name (permission rules, + the service-tool map) would follow whichever they happened to ask. + Matching capabilities are *not* grounds for sharing a name either — + they say nothing about what the model sees or what the callable does. + + Sharing is declared, never inferred: + + - ``variant_of=""`` registers a **backend-native variant** of an + already-declared tool (see :func:`declare_tool`). It binds to that + tool's identity and permission metadata, and may differ in signature + and docstring, which is the whole point. Its declaration + (capabilities/requires/category/long_running) must match exactly. + - ``replace=True`` takes the name over deliberately; the previous + definition's identities are retired, so its callables resolve to + nothing and are denied rather than inheriting the replacement's + capabilities. + + Returns: + The callable to bind at the definition site. When ``name`` renames + the tool this is a thin wrapper carrying the canonical + ``__name__``, so every consumer sees one identity. + + Raises: + ValueError: If the name is already registered and ``replace`` is + False (or the capabilities/requires declarations are invalid). """ def decorator(f: Callable[..., Any]) -> Callable[..., Any]: - tool_name = name or f.__name__ + tool_name = variant_of or name or f.__name__ tool_desc = description or (f.__doc__ or "").split("\n")[0].strip() validated_caps = _validate_capabilities(capabilities, tool_name) + validated_requires = _validate_requires(requires, tool_name) + canonical = _with_canonical_name(f, tool_name) + + existing = self._tools.get(tool_name) + + if variant_of is not None: + if existing is None: + raise ValueError( + f"Tool {variant_of!r} is not declared, so " + f"{f.__qualname__!r} cannot be a variant of it. Declare " + "the tool first with declare_tool()." + ) + if not _same_declaration( + existing, + validated_caps, + validated_requires, + category, + long_running, + ): + raise ValueError( + f"Variant of {variant_of!r} declares a different " + "declaration than the tool it implements: capabilities, " + "requires, category and long_running must match exactly " + "(a variant shares the tool's permission contract)." + ) + existing.variants = _ordered_variants( + existing.variants + (canonical,) + ) + self.bind_identity(canonical, existing) + if canonical is not f: + self.bind_identity(f, existing) + return canonical + + if existing is not None: + if not replace: + raise ValueError( + f"Tool {tool_name!r} is already registered (by " + f"{getattr(existing.func, '__qualname__', existing.func)!r}). " + "Pick a different name, pass variant_of= for a " + "backend-native implementation of the same tool, or " + "replace=True to take it over deliberately." + ) + self._retire_identities(existing) definition = ToolDefinition( name=tool_name, description=tool_desc, - func=f, + func=canonical, capabilities=validated_caps, category=category, + requires=validated_requires, long_running=long_running, ) self._tools[tool_name] = definition - return f + # Bind identity for every callable that legitimately *is* this tool: + # the canonical (possibly renamed) wrapper handed to backends, and + # the original, which a caller of ``register(func, name=...)`` may + # keep using. + self.bind_identity(canonical, definition) + if canonical is not f: + self.bind_identity(f, definition) + return canonical if func is not None: return decorator(func) return decorator + def bind_identity(self, obj: Any, definition: "ToolDefinition") -> None: + """Record that ``obj`` *is* the tool ``definition`` describes. + + Objects that cannot be weak-referenced (rare; some C callables) are + skipped rather than bound by id alone, since a recycled id would + otherwise hand a later object someone else's capabilities. They resolve + to no definition, which fails closed. + """ + key = id(obj) + identity = self._identity + + def _drop(ref: "weakref.ref") -> None: + # Only clear our own entry: by the time this runs the id may + # already have been re-used and re-bound by a live object. + entry = identity.get(key) + if entry is not None and entry[0] is ref: + del identity[key] + + try: + ref = weakref.ref(obj, _drop) + except TypeError: # not weak-referenceable + return + identity[key] = (ref, definition) + + def declare( + self, + name: str, + *, + description: str, + capabilities: CapabilitiesSpec, + category: ToolCategory = ToolCategory.OTHER, + requires: str | tuple[str, ...] | list[str] | None = None, + long_running: bool = False, + ) -> "ToolDefinition": + """Declare a tool that exists only as backend-native variants. + + The declaration owns the name and the permission contract; each backend + then registers its own implementation with + ``register(..., variant_of=name)``. Neither backend can win the name by + importing first, and a bare-name reference resolves deterministically — + to nothing, because there is no backend-neutral implementation to give. + + Idempotent: re-declaring the same contract returns the existing + definition (module import order must not matter), while a conflicting + re-declaration raises. + + Returns: + The declared ``ToolDefinition`` (``func`` is None). + + Raises: + ValueError: If the name already has a different declaration. + """ + validated_caps = _validate_capabilities(capabilities, name) + validated_requires = _validate_requires(requires, name) + + existing = self._tools.get(name) + if existing is not None: + if ( + existing.func is not None + or existing.description != description + or not _same_declaration( + existing, validated_caps, validated_requires, category, long_running + ) + ): + raise ValueError( + f"Tool {name!r} is already registered with a different " + "declaration; declare_tool() cannot take it over." + ) + return existing + + definition = ToolDefinition( + name=name, + description=description, + func=None, + capabilities=validated_caps, + category=category, + requires=validated_requires, + long_running=long_running, + ) + self._tools[name] = definition + return definition + + def _retire_identities(self, definition: "ToolDefinition") -> None: + """Unbind every object that used to *be* ``definition``'s tool. + + A replaced tool must not keep its capabilities through a callable the + application still holds: those bindings now describe a tool this + registry no longer has. Retired objects resolve to nothing, which is + denied at permission time and left alone during assembly. + + They are also remembered (weakly) as *retired*, which is different from + "never registered": a backend's state-tool variant that has been + replaced must not be auto-injected, while an application's own + unregistered callable of the same name is nobody's business but its + author's. + """ + stale = [ + key for key, (_, bound) in self._identity.items() if bound is definition + ] + for key in stale: + ref, _ = self._identity.pop(key) + obj = ref() + if obj is not None: + try: + self._retired.add(obj) + except TypeError: # pragma: no cover - unhashable + pass + for variant in definition.variants: + try: + self._retired.add(variant) + except TypeError: # pragma: no cover - unhashable + pass + + def is_retired(self, obj: Any) -> bool: + """Whether ``obj`` implemented a tool this registry has since replaced.""" + try: + return obj in self._retired + except TypeError: # pragma: no cover - unhashable + return False + + def canonical_for(self, obj: Any) -> Any: + """The canonical-named callable for a tool object this registry issued. + + ``register(func, name=...)`` and ``register(func, variant_of=...)`` + both hand back a wrapper carrying the registered name while the caller + may keep the original. Assembly must give the backend the wrapper, so + the model-visible name is the tool's identity — and for a declared tool + the canonical form is *that variant's* wrapper, since the declaration + itself has no implementation to substitute. + + Returns ``obj`` unchanged when it is already canonical, or when this + registry did not issue it. + """ + definition = self.identify(obj) + if definition is None: + return obj + if getattr(obj, "__name__", None) == definition.name: + return obj + for candidate in (definition.func, *definition.variants): + if candidate is not None and getattr(candidate, "__wrapped__", None) is obj: + return candidate + return definition.func if definition.func is not None else obj + + def identify(self, obj: Any) -> "ToolDefinition | None": + """The definition ``obj`` was bound to *in this registry*, or None.""" + if obj is None: + return None + entry = self._identity.get(id(obj)) + if entry is None: + return None + ref, definition = entry + if ref() is not obj: # id recycled after the bound object died + return None + return definition + def get(self, name: str) -> ToolDefinition | None: """Get a tool definition by name.""" return self._tools.get(name) @@ -183,8 +608,12 @@ def list_by_category(self, category: ToolCategory) -> list[ToolDefinition]: return [t for t in self._tools.values() if t.category == category] def get_functions(self) -> list[Callable[..., Any]]: - """Get all tool functions (for passing to agents).""" - return [t.func for t in self._tools.values()] + """Get all tool functions (for passing to agents). + + Declared-only tools are skipped: they have no backend-neutral + implementation to hand out (see :meth:`declare`). + """ + return [t.func for t in self._tools.values() if t.func is not None] def __len__(self) -> int: return len(self._tools) @@ -202,6 +631,33 @@ def get_registry() -> ToolRegistry: return _default_registry +def declare_tool( + name: str, + *, + description: str, + capabilities: CapabilitiesSpec, + category: ToolCategory = ToolCategory.OTHER, + requires: str | tuple[str, ...] | list[str] | None = None, + long_running: bool = False, + registry: "ToolRegistry | None" = None, +) -> "ToolDefinition": + """Declare a tool implemented only by backend-native variants. + + See :meth:`ToolRegistry.declare`. Defaults to the framework registry. + """ + # `registry or ...` would fall through for an *empty* registry: the class + # defines __len__, so a registry with no tools yet is falsy. + target = _default_registry if registry is None else registry + return target.declare( + name, + description=description, + capabilities=capabilities, + category=category, + requires=requires, + long_running=long_running, + ) + + def register_tool( func: Callable[..., Any] | None = None, *, @@ -209,14 +665,30 @@ def register_tool( description: str | None = None, category: ToolCategory = ToolCategory.OTHER, capabilities: CapabilitiesSpec, + requires: str | tuple[str, ...] | list[str] | None = None, long_running: bool = False, + replace: bool = False, + variant_of: str | None = None, ) -> Callable[..., Any]: """Register a tool with the default registry. ``capabilities`` is a required keyword argument. Pass ``EXEMPT`` to opt out of the permission engine, or a list of ``Capability`` instances to declare - the resources this tool accesses. ``long_running=True`` marks a tool that - starts a background job (see ``tools/jobs/``). + the resources this tool accesses. + + ``requires`` declares which **framework-provided** services the tool needs + (``"kb_manager"``, ``"memory_store"``, … — the full set is + ``service_registry.KNOWN_SERVICE_KEYS``), so managers create exactly those, + lazily. A downstream tool may request any of them without editing the + framework; anything else raises, because nothing would construct it. There + is no mechanism for registering new service *types*. + + ``long_running=True`` marks a tool that starts a background job (see + ``tools/jobs/``). + + A name may be registered once; a collision raises ``ValueError`` unless + ``variant_of=`` (a backend-native implementation of a declared tool) or + ``replace=True`` is passed deliberately (see ``ToolRegistry.register``). """ def _outer(f: Callable[..., Any]) -> Callable[..., Any]: @@ -226,7 +698,10 @@ def _outer(f: Callable[..., Any]) -> Callable[..., Any]: description=description, category=category, capabilities=capabilities, + requires=requires, long_running=long_running, + replace=replace, + variant_of=variant_of, ) if func is not None: diff --git a/src/agentic_cli/tools/skills/toolset.py b/src/agentic_cli/tools/skills/toolset.py index 8ce9865..be5e906 100644 --- a/src/agentic_cli/tools/skills/toolset.py +++ b/src/agentic_cli/tools/skills/toolset.py @@ -42,4 +42,25 @@ def make_skill_toolset( toolset._tools = [ t for t in toolset._tools if not isinstance(t, RunSkillScriptTool) ] + _bind_skill_tool_identities(toolset) return toolset + + +def _bind_skill_tool_identities(toolset: Any) -> None: + """Give each skill tool object the registry identity it implements. + + ADK's skill tools wrap no callable, so the permission plugin cannot verify + them the way it verifies a function tool. Binding happens *here*, where the + framework itself constructs them and their concrete types are known — + rather than letting the plugin resolve them by ``tool.name``, which any + application could pick to impersonate an EXEMPT tool. + """ + from agentic_cli.tools.registry import bind_tool_identity, get_registry + + # The names/capabilities themselves are registered when the ``skills`` + # package is imported, which importing this module guarantees. + registry = get_registry() + for tool in getattr(toolset, "_tools", []) or []: + definition = registry.get(getattr(tool, "name", "")) + if definition is not None: + bind_tool_identity(tool, definition) diff --git a/src/agentic_cli/tools/tool_resolver.py b/src/agentic_cli/tools/tool_resolver.py index c94e93a..332e07d 100644 --- a/src/agentic_cli/tools/tool_resolver.py +++ b/src/agentic_cli/tools/tool_resolver.py @@ -67,6 +67,34 @@ def _unknown_name_error(name: str, registry: ToolRegistry) -> ValueError: ) +def _qualified(variant: Any) -> str: + """Module-qualified name of a variant — two backends' tools share a name.""" + target = getattr(variant, "__wrapped__", variant) + module = getattr(target, "__module__", "") or "" + qualname = getattr(target, "__qualname__", None) or repr(target) + return f"{module}.{qualname}" if module else qualname + + +def _ambiguous_name_error(defn: Any) -> ValueError: + """Build the error for a bare name with no backend-neutral implementation. + + Some tools exist only as backend-native variants (the plan/task state tools + read and write ADK's ``ToolContext.state`` or LangGraph's graph state). A + bare name cannot choose between them — and choosing by import order, which + is what happened before they were declared, could hand an ADK agent a + LangGraph tool. + """ + variants = ", ".join( + _qualified(v) for v in getattr(defn, "variants", ()) + ) + return ValueError( + f"Tool name {defn.name!r} is ambiguous: it has no backend-neutral " + f"implementation, only backend-native variants ({variants or 'none yet'}). " + "Let the workflow manager inject it (AgentConfig.include_state_tools=True), " + "or reference the backend's implementation by dotted path." + ) + + def resolve_tool( ref: Callable[..., Any] | str | Any, registry: ToolRegistry | None = None, @@ -97,7 +125,10 @@ def resolve_tool( return _import_dotted(name) # Bare name -> registry lookup. - reg = registry or get_registry() + # Not ``registry or get_registry()``: ToolRegistry defines __len__, so a + # caller's empty registry is falsy and would be silently replaced by the + # global one — resolving names it never registered. + reg = get_registry() if registry is None else registry defn = reg.get(name) if defn is None and registry is None: # Default registry may not have imported the built-ins yet. @@ -106,6 +137,8 @@ def resolve_tool( defn = reg.get(name) if defn is None: raise _unknown_name_error(name, reg) + if defn.func is None: + raise _ambiguous_name_error(defn) return defn.func diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index a31eb8d..2c95e80 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -615,18 +615,21 @@ def _wrap_long_running(self, tools: list[Callable]) -> list: A long-running tool returns a ``job_id`` immediately and the model is instructed not to re-call it while pending; the eventual result is delivered later as a ``FunctionResponse`` (see ``resume_with_job_result``). - Detection is by registered tool name; non-long-running tools and any - already-wrapped tools (including toolset objects) pass through unchanged. - Permission gating is unaffected — ADK gates via ``PermissionPlugin`` (by - name), not by wrapping the callable. + Detection is by registry identity — the exact object the default + registry issued, never a matching name — so a plain callable an + application names after a long-running tool keeps its ordinary + call-and-return contract instead of being told to leave a job pending + that nothing will ever complete. Non-long-running tools and any + already-wrapped tools (including toolset objects) pass through + unchanged. Permission gating is unaffected: ADK gates via + ``PermissionPlugin``, which resolves the same identity. """ - from agentic_cli.tools.registry import get_registry + from agentic_cli.tools.registry import identify_tool - reg = get_registry() wrapped: list = [] for tool in tools: - name = getattr(tool, "__name__", "") - defn = reg.get(name) if name else None + defn = identify_tool(tool) + name = defn.name if defn is not None else getattr(tool, "__name__", "") if ( defn is not None and defn.long_running diff --git a/src/agentic_cli/workflow/adk/permission_plugin.py b/src/agentic_cli/workflow/adk/permission_plugin.py index 1088e10..2b4935e 100644 --- a/src/agentic_cli/workflow/adk/permission_plugin.py +++ b/src/agentic_cli/workflow/adk/permission_plugin.py @@ -1,13 +1,29 @@ """ADK plugin that gates tool calls via PermissionEngine. +Tool identity is *object* identity, resolved through the registry's identity +binding. It is never ``tool.name`` (ADK derives that from the callable, so an +unregistered function named after a registered tool would inherit its +capabilities — an EXEMPT one would be waved straight through), never the tool's +class name, and never equality. + Adapter check order (mirrors LangGraph wrapper for consistency): -1. EXEMPT tool → allow, no engine call. -2. Unregistered MCP toolset tool → gate through the engine under a synthetic - ``mcp`` capability (no rule → ASK). -3. Tool has no capability declaration → deny (author error, loud). -4. Engine absent from service registry → fail closed (deny) when permissions +1. Resolve the definition bound to this exact object. Framework-issued tools — + registered callables, service-bound factory variants, and the native ADK + tool objects the framework constructs (skill tools) — are bound at their + construction site. +2. Failing that, unwrap ``.func`` only for ADK's own function-tool types + (:data:`_TRUSTED_FUNCTION_TOOL_TYPES`), whose contract is to invoke exactly + that callable, and resolve the callable by identity. +3. EXEMPT tool → allow, no engine call. +4. A genuine ``McpTool`` instance (isinstance, not class name) → gate under a + synthetic ``mcp`` capability (no rule → ASK); its tools are created inside + ADK when the server connects, so they cannot be bound in advance. +5. Anything still unresolved → deny. That includes an unregistered callable and + any tool object the framework did not issue, whatever it calls itself. +6. No capability declaration → deny (author error, loud). +7. Engine absent from service registry → fail closed (deny) when permissions are enabled; allow only when permissions are disabled. -5. Otherwise call engine.check() and return None on allow, error dict on deny. +8. Otherwise call engine.check() and return None on allow, error dict on deny. """ from __future__ import annotations @@ -15,9 +31,14 @@ from typing import Any, TYPE_CHECKING from google.adk.plugins.base_plugin import BasePlugin +from google.adk.tools import FunctionTool, LongRunningFunctionTool from agentic_cli.logging import Loggers -from agentic_cli.tools.registry import ToolCategory, get_registry, register_tool +from agentic_cli.tools.registry import ( + ToolCategory, + identify_tool, + register_tool, +) from agentic_cli.workflow.permissions import EXEMPT from agentic_cli.workflow.permissions.capabilities import Capability, _CapabilityExempt from agentic_cli.workflow.service_registry import PERMISSION_ENGINE, get_service @@ -40,18 +61,37 @@ pass +# ADK's own function-tool types: their documented contract is to call exactly +# ``self.func``, so the callable's identity is the tool's identity. Matched by +# exact type — a subclass may override ``run_async`` and run something else +# while still advertising a genuine ``func``. +_TRUSTED_FUNCTION_TOOL_TYPES = (FunctionTool, LongRunningFunctionTool) + + +def _trusted_wrapped_callable(tool: "BaseTool") -> Any | None: + """The callable an ADK function tool will actually invoke, or None. + + Only exact trusted types are unwrapped: ``.func`` on anything else is just + an attribute, and an attribute is not evidence of what the tool does. + """ + if type(tool) not in _TRUSTED_FUNCTION_TOOL_TYPES: + return None + return getattr(tool, "func", None) + + def _is_mcp_tool(tool: "BaseTool") -> bool: - """True if ``tool`` is an ADK MCP toolset tool (not in our registry).""" + """True if ``tool`` really is an ADK MCP toolset tool. + + ``isinstance`` against the class ADK ships — never the class *name*, which + any application can choose. ``McpTool`` is the base class and ``MCPTool`` a + deprecated subclass, so one check covers both. If the MCP extra is not + installed nothing can be an MCP tool, and the caller denies. + """ try: - # McpTool is the base class; MCPTool is a deprecated subclass, so an - # isinstance check against McpTool catches both. from google.adk.tools.mcp_tool import McpTool - - if isinstance(tool, McpTool): - return True except Exception: - pass - return type(tool).__name__ in {"MCPTool", "McpTool"} + return False + return isinstance(tool, McpTool) # Synthetic capability for MCP tools; target is the MCP tool name. With no @@ -60,6 +100,24 @@ def _is_mcp_tool(tool: "BaseTool") -> bool: _MCP_TARGET_ARG = "__mcp_target__" +class _Sentinel: + """Resolution outcome that is not a ``ToolDefinition``.""" + + __slots__ = ("_label",) + + def __init__(self, label: str) -> None: + self._label = label + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"<{self._label}>" + + +# The tool carries no registry identity: deny. +_UNVERIFIED = _Sentinel("unverified-tool") +# A genuine MCP tool: gate under the synthetic ``mcp`` capability. +_MCP = _Sentinel("mcp-tool") + + def _no_engine_result(tool_name: str) -> dict | None: """Return value when the permission engine is absent from the registry. @@ -92,17 +150,31 @@ async def before_tool_callback( tool_args: dict[str, Any], tool_context: "ToolContext | None", ) -> dict | None: - defn = get_registry().get(tool.name) - caps = defn.capabilities if defn else None + defn = self._resolve_definition(tool) + + if defn is _UNVERIFIED: + # Nothing the framework issued. Its name (or its class name) may + # match a registered tool; that proves nothing. + logger.warning("permission_unregistered_tool", tool=tool.name) + return { + "success": False, + "error": ( + "Permission denied: tool is not registered " + "(register it with @register_tool to declare capabilities)" + ), + } + + if defn is _MCP: + # MCP tools are created inside ADK when the server connects, so + # they carry no binding; gate them under a synthetic capability. + return await self._check_mcp(tool) + + caps = defn.capabilities if isinstance(caps, _CapabilityExempt): return None if not caps: - # MCP toolset tools aren't registered; gate them through the engine - # under a synthetic 'mcp' capability (no rule → ASK). - if _is_mcp_tool(tool): - return await self._check_mcp(tool) logger.warning("permission_undeclared", tool=tool.name) return { "success": False, @@ -111,13 +183,34 @@ async def before_tool_callback( engine = get_service(PERMISSION_ENGINE) if engine is None: - return _no_engine_result(tool.name) + return _no_engine_result(defn.name) - result = await engine.check(tool.name, caps, tool_args) + result = await engine.check(defn.name, caps, tool_args) if result.allowed: return None return {"success": False, "error": f"Permission denied: {result.reason}"} + @staticmethod + def _resolve_definition(tool: "BaseTool"): + """Resolve what this exact tool object is authorised to do. + + Returns the bound ``ToolDefinition``; ``_MCP`` for a genuine MCP tool + (gated under a synthetic capability); or ``_UNVERIFIED`` for anything + the framework did not issue, which the caller denies. There is no + name-based path: a name is chosen by whoever built the tool. + """ + defn = identify_tool(tool) + if defn is not None: + return defn + + func = _trusted_wrapped_callable(tool) + if func is not None: + return identify_tool(func) or _UNVERIFIED + + if _is_mcp_tool(tool): + return _MCP + return _UNVERIFIED + async def _check_mcp(self, tool: "BaseTool") -> dict | None: """Gate an MCP tool through the engine under a synthetic capability.""" engine = get_service(PERMISSION_ENGINE) diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index fb088c1..5851768 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -206,23 +206,78 @@ def _build_tools( Replaces service tools with closure-bound factory versions and auto-injects backend-specific state tools when requested. + + Entries are matched by *registry identity* — the exact object the + default registry issued — never by ``__name__``. A plain callable an + application happens to name ``kb_search`` is not the framework's tool: + substituting the service-bound variant for it would silently run + different code, and it is denied at permission time anyway. It is + therefore passed through untouched, as is a tool registered into some + other ``ToolRegistry``. + + Conversely ``register(func, name=...)`` leaves the caller holding a + callable whose ``__name__`` is the private implementation name; that + one *is* bound, so it resolves to its service-bound variant, or is + replaced by the canonical callable so the model sees the registered + name. """ + from agentic_cli.tools.registry import get_registry, identify_tool + if service_map is None: service_map = self._get_service_tool_map() + registry = get_registry() result = [] for tool in config.tools or []: - name = getattr(tool, "__name__", "") - if name in service_map: - result.append(service_map[name]) - else: + definition = identify_tool(tool) + if definition is None: result.append(tool) + continue + variant = service_map.get(definition.name) + if variant is not None and identify_tool(variant) is definition: + # The variant must *be* this tool, not merely share its name: + # an application that took the name over (register(..., + # replace=True)) would otherwise have the framework's + # implementation run in place of its own. + result.append(variant) + continue + # A renamed tool's (or a declared variant's) original callable: + # hand the backend the canonical one so the model-visible name is + # the tool's identity. Never None — see ``canonical_for``. + result.append(registry.canonical_for(tool)) if config.include_state_tools: - result.extend(self._get_state_tools()) + result.extend(self._injectable_state_tools(result)) return result + def _injectable_state_tools(self, assembled: list) -> list[Callable]: + """State tools worth auto-injecting, given what the agent already has. + + A state tool the application has taken over (``replace=True``) leaves + the backend's variant *retired* — no identity, no capabilities, and the + replacement is what the agent should call. Injecting it anyway would + hand the model two tools with the same name, one of them denied. + + An unregistered state tool that collides with nothing is left alone: + a backend may legitimately supply its own. + """ + from agentic_cli.tools.registry import get_registry + + registry = get_registry() + present = {getattr(tool, "__name__", "") for tool in assembled} + injectable = [] + for tool in self._get_state_tools(): + name = getattr(tool, "__name__", "") + if name in present: + logger.debug("state_tool_already_present", tool=name) + continue + if registry.is_retired(tool): + logger.debug("state_tool_retired", tool=name) + continue + injectable.append(tool) + return injectable + def _get_service_tool_map(self) -> dict[str, Callable]: """Create service tools via factories, returning name→function map. @@ -263,6 +318,10 @@ def _get_service_tool_map(self) -> dict[str, Callable]: for t in make_interaction_tools(self): tool_map[t.__name__] = t + # The factories bind each closure to the definition of the exact + # module-level tool it re-binds (see ``factories._issued``), so a + # variant carries identity only while the framework still owns that + # tool — never merely because the names match. return tool_map @abstractmethod diff --git a/src/agentic_cli/workflow/service_registry.py b/src/agentic_cli/workflow/service_registry.py index c0a7ced..698921d 100644 --- a/src/agentic_cli/workflow/service_registry.py +++ b/src/agentic_cli/workflow/service_registry.py @@ -23,6 +23,33 @@ USER_KB_MANAGER = "user_kb_manager" WORKFLOW = "workflow" +# Service keys a tool may declare via ``register_tool(requires=...)``. +# +# Every entry must be something a manager can actually construct on demand +# (see ``BaseWorkflowManager._ensure_managers_initialized``). Keys that are +# always present (PERMISSION_ENGINE, WORKFLOW) say nothing when declared, and +# USER_KB_MANAGER is not independently constructible — it is created together +# with KB_MANAGER — so declaring it would validate and then provide nothing. +KNOWN_SERVICE_KEYS = frozenset({ + ARXIV_SOURCE, + JOB_MANAGER, + KB_MANAGER, + LLM_SUMMARIZER, + MEMORY_STORE, + SANDBOX_MANAGER, +}) + +# Extra guidance for keys that look declarable but are not. +SERVICE_KEY_HINTS = { + USER_KB_MANAGER: ( + f"{USER_KB_MANAGER!r} is created together with {KB_MANAGER!r}; " + f"declare requires={KB_MANAGER!r} to get both the project- and " + "user-scoped knowledge bases." + ), + PERMISSION_ENGINE: f"{PERMISSION_ENGINE!r} is always available.", + WORKFLOW: f"{WORKFLOW!r} is always available.", +} + # ---- ContextVar and accessors ---- diff --git a/tests/integration/test_permission_adk.py b/tests/integration/test_permission_adk.py index 02ffb8e..3be8a4b 100644 --- a/tests/integration/test_permission_adk.py +++ b/tests/integration/test_permission_adk.py @@ -8,6 +8,18 @@ from agentic_cli.workflow.service_registry import PERMISSION_ENGINE +def _adk_tool(func): + """Wrap a registered callable the way ADK does before dispatching it. + + The plugin resolves capabilities by *identity*, never by ``tool.name``, so + a name-only stand-in is (correctly) denied as unregistered — see + ``tests/workflow/test_permission_tool_identity.py``. + """ + from google.adk.tools import FunctionTool + + return FunctionTool(func=func) + + @pytest.fixture def stub_engine(): from agentic_cli.workflow.permissions.rules import CheckResult @@ -34,7 +46,7 @@ def exempt_x(): plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="exempt_x"), tool_args={}, tool_context=None, + tool=_adk_tool(exempt_x), tool_args={}, tool_context=None, ) assert result is None stub_engine.check.assert_not_called() @@ -47,16 +59,18 @@ async def test_missing_declaration_denies(self, monkeypatch, stub_engine): "agentic_cli.workflow.adk.permission_plugin.get_service", lambda k: stub_engine if k == PERMISSION_ENGINE else None, ) + def never_registered(): + """Never passed through @register_tool.""" + return {} + plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="never_registered"), + tool=_adk_tool(never_registered), tool_args={}, tool_context=None, ) - assert result == { - "success": False, - "error": "Permission denied: tool has no capability declaration", - } + assert result is not None and result["success"] is False + assert "not registered" in result["error"] @pytest.mark.asyncio async def test_allow_calls_engine_and_passes(self, monkeypatch, stub_engine): @@ -78,7 +92,7 @@ def reader_x(path: str): plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="reader_x"), + tool=_adk_tool(reader_x), tool_args={"path": "/tmp/x"}, tool_context=None, ) @@ -108,7 +122,7 @@ def writer_x(path: str): plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="writer_x"), + tool=_adk_tool(writer_x), tool_args={"path": "/etc/x"}, tool_context=None, ) @@ -139,7 +153,7 @@ def reader_y_deny(path: str): plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="reader_y_deny"), + tool=_adk_tool(reader_y_deny), tool_args={"path": "/tmp/x"}, tool_context=None, ) @@ -170,7 +184,7 @@ def reader_y_allow(path: str): plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="reader_y_allow"), + tool=_adk_tool(reader_y_allow), tool_args={"path": "/tmp/x"}, tool_context=None, ) diff --git a/tests/test_packaging.py b/tests/test_packaging.py index 5ee3b5d..e076ac2 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -37,3 +37,63 @@ def test_html2text_is_declared(): "declared in pyproject.toml [project.dependencies] — a clean pip install " "breaks when the webfetch tools are imported." ) + + +# --- Optional-capability isolation ------------------------------------------- +# +# The heavyweight capability stacks live behind extras (``kb``: torch / +# sentence-transformers / faiss / bm25; ``langgraph``). Importing the base +# package — or the tool package, which the tool resolver imports to discover +# built-ins — must not require any of them, and must not drag them in. + +_OPTIONAL_ROOTS = { + "torch", + "sentence_transformers", + "faiss", + "bm25s", + "rank_bm25", + "langgraph", + "langchain", + "langchain_core", +} + + +def _import_with_optionals_blocked(module_name: str) -> set[str]: + """Import ``module_name`` in a subprocess with the extras blocked. + + Returns: + The set of optional roots that ended up in ``sys.modules`` anyway. + """ + import json + import subprocess + import sys + + script = f""" +import builtins, importlib, json, sys +blocked = {sorted(_OPTIONAL_ROOTS)!r} +_real = builtins.__import__ +def _guard(name, *a, **kw): + if name.split(".")[0] in blocked: + raise ImportError("blocked optional dependency: " + name) + return _real(name, *a, **kw) +builtins.__import__ = _guard +importlib.import_module({module_name!r}) +builtins.__import__ = _real +print(json.dumps([m for m in blocked if m in sys.modules])) +""" + proc = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True + ) + assert proc.returncode == 0, ( + f"importing {module_name} requires an optional dependency:\n{proc.stderr}" + ) + return set(json.loads(proc.stdout.strip().splitlines()[-1])) + + +def test_base_package_imports_without_optional_extras(): + assert _import_with_optionals_blocked("agentic_cli") == set() + + +def test_tools_package_imports_without_optional_extras(): + """Tool auto-discovery must not pull the kb/langgraph stacks.""" + assert _import_with_optionals_blocked("agentic_cli.tools") == set() diff --git a/tests/tools/test_state_tool_aliases.py b/tests/tools/test_state_tool_aliases.py new file mode 100644 index 0000000..2583272 --- /dev/null +++ b/tests/tools/test_state_tool_aliases.py @@ -0,0 +1,144 @@ +"""Backend-native tool variants are explicit, and bare names are deterministic. + +The ADK and LangGraph state tools (``save_plan``/``get_plan``/``save_tasks``/ +``get_tasks``) are the same *tool* with two native implementations: different +signatures (``ToolContext`` vs ``InjectedState``/``Command``) and different +docstrings, hence different model-visible schemas. + +They used to contest the registry name, and whichever module imported second +won — silently, because their permission metadata happened to match. So +``ToolDefinition.func`` (what a bare ``"save_plan"`` in an ``AgentConfig`` +resolves to) depended on import order, and could hand an ADK agent a LangGraph +tool. They now register as declared *variants* of a backend-neutral contract, +and a bare name that has no neutral implementation fails as ambiguous instead +of guessing. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +import textwrap + +import pytest + + +def _probe(*imports: str) -> dict: + """Import the given modules, in order, in a fresh interpreter.""" + script = textwrap.dedent( + """ + import json + {imports} + from agentic_cli.tools.registry import get_registry, identify_tool + from agentic_cli.tools.tool_resolver import resolve_tool + + registry = get_registry() + definition = registry.get("save_plan") + out = {{ + "declared": definition is not None, + "has_neutral_implementation": bool( + definition is not None and definition.func is not None + ), + "variants": len(definition.variants) if definition else 0, + "capabilities_exempt": definition is not None + and not isinstance(definition.capabilities, list), + }} + try: + resolve_tool("save_plan") + out["bare_name"] = "resolved" + except ValueError as exc: + out["bare_name"] = "ambiguous" if "ambiguous" in str(exc) else "error" + + bound = [] + for module_name, attr in ( + ("agentic_cli.tools.adk.state_tools", "save_plan"), + ("agentic_cli.tools.langgraph.state_tools", "save_plan"), + ): + import importlib + try: + module = importlib.import_module(module_name) + except ImportError: + continue + bound.append(identify_tool(getattr(module, attr)) is definition) + out["all_variants_bound"] = bool(bound) and all(bound) + print(json.dumps(out)) + """ + ).format(imports="\n".join(imports)) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.strip().splitlines()[-1]) + + +ADK_FIRST = ( + "import agentic_cli.tools.adk.state_tools", + "import agentic_cli.tools.langgraph.state_tools", +) +LANGGRAPH_FIRST = ( + "import agentic_cli.tools.langgraph.state_tools", + "import agentic_cli.tools.adk.state_tools", +) + + +@pytest.fixture(scope="module", autouse=True) +def _require_langgraph(): + pytest.importorskip("langgraph") + + +class TestImportOrderIsIrrelevant: + """The registry must look the same whichever backend module loads first.""" + + def test_both_orders_agree(self): + assert _probe(*ADK_FIRST) == _probe(*LANGGRAPH_FIRST) + + def test_both_orders_register_two_variants(self): + for order in (ADK_FIRST, LANGGRAPH_FIRST): + out = _probe(*order) + assert out["declared"] is True + assert out["variants"] == 2, out + assert out["all_variants_bound"] is True + + def test_bare_name_is_ambiguous_in_both_orders(self): + for order in (ADK_FIRST, LANGGRAPH_FIRST): + out = _probe(*order) + assert out["has_neutral_implementation"] is False + assert out["bare_name"] == "ambiguous", out + + def test_a_single_backend_still_registers_its_variant(self): + out = _probe("import agentic_cli.tools.adk.state_tools") + assert out["declared"] is True + assert out["variants"] == 1 + assert out["bare_name"] == "ambiguous" + + +class TestVariantsShareOneContract: + """In-process: both variants gate under the same declared capabilities.""" + + def test_both_variants_resolve_to_the_declared_tool(self): + from agentic_cli.tools.adk import state_tools as adk_state_tools + from agentic_cli.tools.langgraph import state_tools as lg_state_tools + from agentic_cli.tools.registry import get_registry, identify_tool + + definition = get_registry().get("save_plan") + assert definition is not None + assert identify_tool(adk_state_tools.save_plan) is definition + assert identify_tool(lg_state_tools.save_plan) is definition + + def test_the_declaration_has_no_neutral_implementation(self): + from agentic_cli.tools.registry import get_registry + + definition = get_registry().get("save_plan") + assert definition.func is None + assert len(definition.variants) == 2 + + def test_bare_name_resolution_names_the_variants(self): + from agentic_cli.tools.tool_resolver import resolve_tool + + with pytest.raises(ValueError, match="ambiguous") as exc: + resolve_tool("save_plan") + + message = str(exc.value) + assert "save_plan" in message + assert "include_state_tools" in message diff --git a/tests/workflow/test_adk_mcp_permissions.py b/tests/workflow/test_adk_mcp_permissions.py index 7fd8e04..f40250f 100644 --- a/tests/workflow/test_adk_mcp_permissions.py +++ b/tests/workflow/test_adk_mcp_permissions.py @@ -22,14 +22,20 @@ ) -class _MCPTool: - """Stand-in for an ADK MCP tool (detected by class name).""" +_McpToolBase = pytest.importorskip("google.adk.tools.mcp_tool").McpTool - def __init__(self, name: str): - self.name = name +class _MCPTool(_McpToolBase): + """A real ``McpTool`` instance without the MCP session plumbing. -_MCPTool.__name__ = "MCPTool" + It must genuinely *be* one: detection is ``isinstance`` against the class + ADK ships, never the class name (which any application can choose — see + ``tests/workflow/test_permission_tool_identity.py``). + """ + + def __init__(self, name: str): + object.__setattr__(self, "name", name) + object.__setattr__(self, "description", f"MCP tool {name}") class _PlainTool: diff --git a/tests/workflow/test_permission_tool_identity.py b/tests/workflow/test_permission_tool_identity.py new file mode 100644 index 0000000..8b2e1a7 --- /dev/null +++ b/tests/workflow/test_permission_tool_identity.py @@ -0,0 +1,362 @@ +"""Permission gating binds to registry identity, not to a tool's name. + +``before_tool_callback`` resolved capabilities with +``get_registry().get(tool.name)``. ADK derives that name from the callable, so +an *unregistered* function whose ``__name__`` collides with a registered tool +inherited that tool's capabilities — a raw callable named ``ask_clarification`` +was allowed outright, because the genuine registered tool is EXEMPT. + +Identity now comes from a registry-owned binding attached when a callable is +registered (or when the framework produces a service-bound/renamed variant of +one), so name equality alone proves nothing. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from google.adk.tools import FunctionTool, LongRunningFunctionTool # noqa: E402 + +from agentic_cli.config import BaseSettings # noqa: E402 +from agentic_cli.tools.registry import ( # noqa: E402 + ToolCategory, + ToolRegistry, + get_registry, +) +from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin # noqa: E402 +from agentic_cli.workflow.permissions import EXEMPT, PermissionEngine # noqa: E402 +from agentic_cli.workflow.permissions.capabilities import Capability # noqa: E402 +from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource # noqa: E402 +from agentic_cli.workflow.permissions.store import PermissionContext # noqa: E402 +from agentic_cli.workflow.service_registry import ( # noqa: E402 + PERMISSION_ENGINE, + set_service_registry, +) + + +class _StubWorkflow: + """Answers a HITL permission prompt with a fixed choice.""" + + def __init__(self, response: str = "deny"): + self._response = response + + async def request_user_input(self, request): + return self._response + + +def _engine(tmp_path, response="deny", rules=None) -> PermissionEngine: + settings = BaseSettings(google_api_key="test") + ctx = PermissionContext(workdir=tmp_path, home=tmp_path) + eng = PermissionEngine(settings=settings, workflow=_StubWorkflow(response), ctx=ctx) + if rules: + eng._session_rules.extend(rules) + return eng + + +async def _check(engine, tool, tool_args=None): + token = set_service_registry({PERMISSION_ENGINE: engine}) + try: + return await PermissionPlugin().before_tool_callback( + tool=tool, tool_args=tool_args or {}, tool_context=None + ) + finally: + token.var.reset(token) + + +def _denied(result) -> bool: + return isinstance(result, dict) and result.get("success") is False + + +class TestImpostorCallables: + """A raw callable must not inherit a registered tool's capabilities.""" + + async def test_impostor_named_after_an_exempt_tool_is_denied(self, tmp_path): + # The genuine ask_clarification is registered EXEMPT. + from agentic_cli.tools.interaction_tools import ask_clarification # noqa: F401 + + assert get_registry().get("ask_clarification") is not None + + def ask_clarification(question: str) -> dict: # noqa: F811 - deliberate collision + """An unregistered impostor with a colliding name.""" + return {"success": True} + + tool = FunctionTool(func=ask_clarification) + assert tool.name == "ask_clarification" + + result = await _check(_engine(tmp_path), tool) + assert _denied(result), "an unregistered callable inherited EXEMPT status" + + async def test_impostor_named_after_a_permissioned_tool_is_denied(self, tmp_path): + from agentic_cli.tools.file_read import read_file # noqa: F401 + + def read_file(path: str) -> dict: # noqa: F811 - deliberate collision + """Impostor.""" + return {"success": True} + + allow_all = Rule( + capability="fs.read", target="**", effect=Effect.ALLOW, + source=RuleSource.SESSION, + ) + result = await _check( + _engine(tmp_path, rules=[allow_all]), + FunctionTool(func=read_file), + {"path": str(tmp_path / "x")}, + ) + assert _denied(result), "an unregistered callable used a registered rule" + + async def test_genuine_tool_is_still_allowed(self, tmp_path): + from agentic_cli.tools.interaction_tools import ask_clarification + + result = await _check(_engine(tmp_path), FunctionTool(func=ask_clarification)) + assert result is None, "the genuine EXEMPT tool must pass" + + +class TestRegisteredIdentity: + """Registration binds identity for both the decorator and direct forms. + + These register into the *default* registry: identity is registry-owned, and + the plugin trusts only the framework's own (see + ``tests/tools/test_registry_identity.py::TestIdentityIsPerRegistry``). + """ + + async def test_renamed_tool_uses_its_declared_capabilities(self, tmp_path): + registry = get_registry() + + @registry.register( + name="public_read", + capabilities=[Capability("fs.read", target_arg="path")], + category=ToolCategory.READ, + ) + def _internal_impl(path: str) -> dict: + """Read something.""" + return {"success": True} + + tool = FunctionTool(func=_internal_impl) + assert tool.name == "public_read" + + allow = Rule( + capability="fs.read", target="**", effect=Effect.ALLOW, + source=RuleSource.SESSION, + ) + allowed = await _check( + _engine(tmp_path, rules=[allow]), tool, {"path": str(tmp_path / "f")} + ) + assert allowed is None + + denied = await _check( + _engine(tmp_path, response="deny"), tool, {"path": str(tmp_path / "f")} + ) + assert _denied(denied), "the declared capability was not evaluated" + + async def test_direct_register_call_binds_both_callables(self, tmp_path): + """``registry.register(func, name=...)`` — original and returned wrapper.""" + registry = get_registry() + + def _impl(path: str) -> dict: + """Impl.""" + return {"success": True} + + returned = registry.register( + _impl, name="direct_public", capabilities=EXEMPT + ) + + for callable_ in (returned, _impl): + result = await _check(_engine(tmp_path), FunctionTool(func=callable_)) + assert result is None, f"{callable_!r} was not recognised as registered" + + async def test_long_running_wrapper_keeps_identity(self, tmp_path): + from agentic_cli.tools.interaction_tools import ask_clarification + + tool = LongRunningFunctionTool(func=ask_clarification) + assert await _check(_engine(tmp_path), tool) is None + + +class TestServiceBoundTools: + """Factory-produced (closure-bound) tools are framework-issued variants.""" + + async def test_factory_bound_tool_is_recognised(self, tmp_path): + from unittest.mock import MagicMock + + from agentic_cli.tools.factories import make_memory_tools + + tools = {t.__name__: t for t in make_memory_tools(MagicMock())} + save_memory = tools["save_memory"] + assert save_memory is not get_registry().get("save_memory").func + + allow = Rule( + capability="memory.write", target="**", effect=Effect.ALLOW, + source=RuleSource.SESSION, + ) + result = await _check( + _engine(tmp_path, rules=[allow]), + FunctionTool(func=save_memory), + {"content": "hi"}, + ) + assert result is None, "the service-bound variant was not recognised" + + async def test_factory_bound_tool_still_obeys_deny(self, tmp_path): + """Its declared capability is really evaluated — a deny rule bites.""" + from unittest.mock import MagicMock + + from agentic_cli.tools.factories import make_memory_tools + + save_memory = {t.__name__: t for t in make_memory_tools(MagicMock())}["save_memory"] + deny = Rule( + capability="memory.write", target="*", effect=Effect.DENY, + source=RuleSource.SESSION, + ) + result = await _check( + _engine(tmp_path, rules=[deny]), + FunctionTool(func=save_memory), + {"content": "hi"}, + ) + assert _denied(result) + + +class TestNonCallableTools: + """Backend-native tool objects are bound by the framework, never by name.""" + + async def test_skill_tools_are_resolved(self, tmp_path): + import pathlib + import tempfile + + from agentic_cli.tools.skills import SkillStore, make_skill_toolset + + skill_dir = pathlib.Path(tempfile.mkdtemp()) / "demo-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text( + "---\nname: demo-skill\ndescription: d\n---\nBody\n" + ) + toolset = make_skill_toolset(SkillStore().resolve([str(skill_dir)])) + list_skills = next(t for t in toolset._tools if t.name == "list_skills") + assert not hasattr(list_skills, "func") + + assert await _check(_engine(tmp_path), list_skills) is None + + async def test_unknown_native_tool_object_is_denied(self, tmp_path): + class _Native: + name = "totally_unknown" + + assert _denied(await _check(_engine(tmp_path), _Native())) + + +class TestNameIsNotAuthority: + """A tool's *name* must never grant it another tool's capabilities.""" + + async def test_custom_basetool_named_after_an_exempt_tool_is_denied(self, tmp_path): + """The forgery the name fallback allowed: subclass BaseTool, pick a name.""" + from google.adk.tools import BaseTool + + from agentic_cli.tools.interaction_tools import ask_clarification # noqa: F401 + + assert get_registry().get("ask_clarification") is not None + + class _Impostor(BaseTool): + async def run_async(self, *, args, tool_context): # pragma: no cover + return {"success": True} + + tool = _Impostor(name="ask_clarification", description="impostor") + assert tool.name == "ask_clarification" + + result = await _check(_engine(tmp_path), tool) + assert _denied(result), "a BaseTool inherited EXEMPT status from its name" + + async def test_custom_basetool_named_after_a_skill_tool_is_denied(self, tmp_path): + """Registered-by-name skill tools must not be impersonable either.""" + from google.adk.tools import BaseTool + + from agentic_cli.tools.skills import register_skill_tool_permissions + + register_skill_tool_permissions() + assert get_registry().get("list_skills") is not None + + class _Impostor(BaseTool): + async def run_async(self, *, args, tool_context): # pragma: no cover + return {"success": True} + + result = await _check( + _engine(tmp_path), _Impostor(name="list_skills", description="impostor") + ) + assert _denied(result), "a BaseTool inherited a skill tool's EXEMPT status" + + async def test_class_named_like_an_mcp_tool_is_denied(self, tmp_path): + """MCP detection must not key on a forgeable class name. + + The engine is given a blanket ``mcp`` ALLOW rule, so reaching the MCP + path at all means the impostor is waved through. + """ + + class McpTool: # not google.adk.tools.mcp_tool.McpTool + def __init__(self, name: str): + self.name = name + + allow_mcp = Rule( + capability="mcp", target="**", effect=Effect.ALLOW, + source=RuleSource.SESSION, + ) + result = await _check( + _engine(tmp_path, rules=[allow_mcp]), McpTool("remote_op") + ) + assert _denied(result), "a class *named* McpTool got the MCP capability path" + + async def test_real_mcp_tool_is_gated_by_the_engine(self, tmp_path): + """A genuine ADK McpTool instance still reaches the synthetic 'mcp' rule.""" + McpTool = pytest.importorskip("google.adk.tools.mcp_tool").McpTool + + class _RealEnough(McpTool): + def __init__(self): # bypass the MCP session plumbing + object.__setattr__(self, "name", "remote_op") + object.__setattr__(self, "description", "remote op") + + result = await _check(_engine(tmp_path, response="deny"), _RealEnough()) + assert _denied(result) + + +class TestForgedIdentity: + """Identity is object identity — equality and attributes cannot forge it.""" + + async def test_equality_colliding_callable_is_denied(self, tmp_path): + """An object that compares equal to a registered tool is not that tool.""" + from agentic_cli.tools.interaction_tools import ask_clarification + + genuine = get_registry().get("ask_clarification").func + + class _Collider: + """Hashes and compares equal to the genuine registered callable.""" + + __name__ = "ask_clarification" + + def __hash__(self): + return hash(genuine) + + def __eq__(self, other): + return other is genuine + + def __call__(self, question: str) -> dict: # pragma: no cover + return {"success": True} + + collider = _Collider() + assert collider == genuine and hash(collider) == hash(genuine) + assert collider is not ask_clarification + + result = await _check(_engine(tmp_path), FunctionTool(func=collider)) + assert _denied(result), "an equality-colliding object forged tool identity" + + async def test_untrusted_wrapper_exposing_a_genuine_func_is_denied(self, tmp_path): + """``.func`` is only trusted on ADK's own function-tool types.""" + from google.adk.tools import BaseTool + + from agentic_cli.tools.interaction_tools import ask_clarification + + class _Wrapper(BaseTool): + """Advertises the genuine callable but runs whatever it likes.""" + + async def run_async(self, *, args, tool_context): # pragma: no cover + return {"success": True} + + tool = _Wrapper(name="totally_other", description="w") + object.__setattr__(tool, "func", ask_clarification) + + assert _denied(await _check(_engine(tmp_path), tool)) From 0d2fc14ebe99b533ec7c8679fc609195d5fa650f Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:39 -0400 Subject: [PATCH 02/11] feat(tools): tools declare the services they require MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``BaseWorkflowManager._TOOL_SERVICE_MAP`` was a central name→service table, so a downstream tool that needed a framework service could only get one by editing the framework — and it matched on ``__name__``, which meant an application function named ``kb_search`` had a knowledge base built for it even though the permission engine would deny the call. Tools now declare their own needs with ``@register_tool(..., requires=...)``, validated at registration against ``service_registry.KNOWN_SERVICE_KEYS`` so only genuinely constructible services can be declared (``user_kb_manager`` is created together with ``kb_manager``, and the error says so). Detection reads that metadata off the registry by identity, so ``register(func, name=...)``'s original callable still declares its services while an unregistered lookalike declares nothing. There is still no mechanism for registering new service *types*. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- examples/jobs_demo.py | 3 + src/agentic_cli/tools/arxiv_tools.py | 3 + src/agentic_cli/tools/jobs/tools.py | 5 + src/agentic_cli/tools/knowledge_tools.py | 8 + src/agentic_cli/tools/memory_tools.py | 4 + src/agentic_cli/tools/sandbox/__init__.py | 1 + src/agentic_cli/tools/webfetch_tool.py | 1 + src/agentic_cli/workflow/base_manager.py | 61 +- tests/test_knowledge_tools.py | 9 +- tests/test_webfetch.py | 9 +- tests/tools/test_registry_identity.py | 1390 +++++++++++++++++++++ tests/tools/test_sandbox.py | 10 +- 12 files changed, 1447 insertions(+), 57 deletions(-) create mode 100644 tests/tools/test_registry_identity.py diff --git a/examples/jobs_demo.py b/examples/jobs_demo.py index 775a076..c386f69 100644 --- a/examples/jobs_demo.py +++ b/examples/jobs_demo.py @@ -61,6 +61,9 @@ @register_tool( category=ToolCategory.EXECUTION, capabilities=[Capability("longrunning.run_shell_job")], + # Declares its own service need, so the manager creates the JobManager for + # an app that ships only this starter (no framework job_* tools required). + requires="job_manager", long_running=True, description="Run a shell command as a detached background job; returns a job_id immediately.", ) diff --git a/src/agentic_cli/tools/arxiv_tools.py b/src/agentic_cli/tools/arxiv_tools.py index 2b195e6..bfbe910 100644 --- a/src/agentic_cli/tools/arxiv_tools.py +++ b/src/agentic_cli/tools/arxiv_tools.py @@ -125,6 +125,7 @@ async def _fetch_arxiv_paper_with_source(source, arxiv_id: str) -> dict[str, Any capabilities=[Capability("http.read")], description="Search arXiv for academic papers by query, category, or date range. Use this to find research papers on a topic.", + requires="arxiv_source", ) def search_arxiv( query: str, @@ -174,6 +175,7 @@ def search_arxiv( capabilities=[Capability("http.read")], description="Fetch metadata for a specific arXiv paper by ID or URL. Returns title, authors, abstract, categories, and PDF URL.", + requires="arxiv_source", ) async def fetch_arxiv_paper( arxiv_id: str, @@ -301,6 +303,7 @@ async def _ingest_arxiv_paper_with_services( capabilities=[Capability("http.read"), Capability("kb.write")], description="Download an arXiv paper's PDF, extract text, and ingest it into the knowledge base. Use this to add a specific arXiv paper to long-term storage so it can be searched later.", + requires=("arxiv_source", "kb_manager"), ) async def ingest_arxiv_paper( arxiv_id: str, diff --git a/src/agentic_cli/tools/jobs/tools.py b/src/agentic_cli/tools/jobs/tools.py index 310f11c..9d2a84f 100644 --- a/src/agentic_cli/tools/jobs/tools.py +++ b/src/agentic_cli/tools/jobs/tools.py @@ -46,6 +46,7 @@ def _manager(): category=ToolCategory.EXECUTION, capabilities=[Capability("jobs.manage")], description="Check a background job: state, exit code, a stdout tail, and the result once finished.", + requires="job_manager", ) def job_status(job_id: str) -> dict: """One-stop check for a background job. @@ -74,6 +75,7 @@ def job_status(job_id: str) -> dict: category=ToolCategory.EXECUTION, capabilities=[Capability("jobs.manage")], description="Get the result of a finished background job.", + requires="job_manager", ) def job_result(job_id: str) -> dict: """Return the job's result, or an error if it isn't finished yet.""" @@ -94,6 +96,7 @@ def job_result(job_id: str) -> dict: category=ToolCategory.EXECUTION, capabilities=[Capability("jobs.manage")], description="Read recent log lines (stdout/stderr) of a background job.", + requires="job_manager", ) def job_logs(job_id: str, last_n: int = 50, stream: str = "stdout") -> dict: """Return the last ``last_n`` lines of the job's ``stdout`` or ``stderr``.""" @@ -109,6 +112,7 @@ def job_logs(job_id: str, last_n: int = 50, stream: str = "stdout") -> dict: category=ToolCategory.EXECUTION, capabilities=[Capability("jobs.manage")], description="Cancel a running background job.", + requires="job_manager", ) def job_cancel(job_id: str) -> dict: """Best-effort cancel a running job.""" @@ -125,6 +129,7 @@ def job_cancel(job_id: str) -> dict: category=ToolCategory.EXECUTION, capabilities=[Capability("jobs.manage")], description="List background jobs, optionally filtered by state or tag.", + requires="job_manager", ) def job_list(state: str = "", tag: str = "") -> dict: """List jobs (most recent first), optionally filtered by state/tag.""" diff --git a/src/agentic_cli/tools/knowledge_tools.py b/src/agentic_cli/tools/knowledge_tools.py index 58ca36b..3ea923f 100644 --- a/src/agentic_cli/tools/knowledge_tools.py +++ b/src/agentic_cli/tools/knowledge_tools.py @@ -680,6 +680,7 @@ def _find_document_in_kbs(doc_id_or_title: str) -> tuple: category=ToolCategory.KNOWLEDGE, capabilities=[Capability("kb.read")], description="Search the local knowledge base for relevant documents using semantic similarity. Use this when you need to find previously ingested papers, notes, or documents.", + requires="kb_manager", ) def kb_search( query: str, @@ -711,6 +712,7 @@ def kb_search( "you already have in memory; use kb_ingest_file for local files and " "kb_ingest_url for remote URLs." ), + requires="kb_manager", ) async def kb_ingest_text( content: str, @@ -758,6 +760,7 @@ async def kb_ingest_text( "extracted automatically. Triggers a filesystem.read permission " "check for the supplied path." ), + requires="kb_manager", ) async def kb_ingest_file( path: str, @@ -807,6 +810,7 @@ async def kb_ingest_file( "papers, prefer ingest_arxiv_paper. Triggers an http.read " "permission check for the supplied URL." ), + requires="kb_manager", ) async def kb_ingest_url( url: str, @@ -851,6 +855,7 @@ async def kb_ingest_url( "sidecar (summary, key claims, entities) by default. Pass full=True " "to get the raw extracted text up to max_chars." ), + requires="kb_manager", ) async def kb_read( doc_id_or_title: str, @@ -875,6 +880,7 @@ async def kb_read( category=ToolCategory.KNOWLEDGE, capabilities=[Capability("kb.read")], description="List documents in the knowledge base with summaries. Filter by query or source type. Returns summaries, not full content.", + requires="kb_manager", ) def kb_list( query: str = "", @@ -907,6 +913,7 @@ def kb_list( "agent-writable, grep-searchable, and human-readable. `sources` " "must cite at least one valid document ID from the KB." ), + requires="kb_manager", ) async def kb_write_concept( title: str, @@ -944,6 +951,7 @@ async def kb_write_concept( "Case-insensitive substring match; title hits rank above body " "hits. Use when asking 'what does the KB know about X?'." ), + requires="kb_manager", ) async def kb_search_concepts( query: str, diff --git a/src/agentic_cli/tools/memory_tools.py b/src/agentic_cli/tools/memory_tools.py index 2f6c37f..b8eb1d5 100644 --- a/src/agentic_cli/tools/memory_tools.py +++ b/src/agentic_cli/tools/memory_tools.py @@ -518,6 +518,7 @@ def _delete_memory_with_store( category=ToolCategory.MEMORY, capabilities=[Capability("memory.write")], description="Save information to persistent memory that survives across sessions. Use this to remember user preferences, important facts, or learnings for future conversations.", + requires="memory_store", ) def save_memory( content: str, @@ -541,6 +542,7 @@ def save_memory( category=ToolCategory.MEMORY, capabilities=[Capability("memory.read")], description="Search persistent memory by keyword/substring. Use this to recall previously saved facts, preferences, or learnings.", + requires="memory_store", ) def search_memory( query: str, @@ -564,6 +566,7 @@ def search_memory( category=ToolCategory.MEMORY, capabilities=[Capability("memory.write")], description="Update an existing memory item", + requires="memory_store", ) def update_memory( item_id: str, @@ -587,6 +590,7 @@ def update_memory( category=ToolCategory.MEMORY, capabilities=[Capability("memory.write")], description="Delete a memory item", + requires="memory_store", ) def delete_memory( item_id: str, diff --git a/src/agentic_cli/tools/sandbox/__init__.py b/src/agentic_cli/tools/sandbox/__init__.py index 1509e65..b81a0d9 100644 --- a/src/agentic_cli/tools/sandbox/__init__.py +++ b/src/agentic_cli/tools/sandbox/__init__.py @@ -32,6 +32,7 @@ "Write scratch/intermediate files to the working directory; write FINAL deliverables " "(figures, tables) to `outputs/` — those persist and are shared with other agents." ), + requires="sandbox_manager", ) def sandbox_execute( code: str, diff --git a/src/agentic_cli/tools/webfetch_tool.py b/src/agentic_cli/tools/webfetch_tool.py index 8369425..d50d583 100644 --- a/src/agentic_cli/tools/webfetch_tool.py +++ b/src/agentic_cli/tools/webfetch_tool.py @@ -81,6 +81,7 @@ def get_or_create_fetcher(settings=None) -> ContentFetcher: category=ToolCategory.NETWORK, capabilities=[Capability("http.read", target_arg="url")], description="Fetch a web page, convert it to markdown, and summarize it using an LLM based on your prompt. Use this to extract specific information from a URL (e.g., documentation, articles).", + requires="llm_summarizer", ) async def web_fetch(url: str, prompt: str, timeout: int = 30) -> dict[str, Any]: """Fetch web content and summarize it using an LLM. diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index 5851768..ea2b67f 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -332,39 +332,6 @@ def _get_state_tools(self) -> list[Callable]: """ ... - # Mapping from tool function name to the service(s) it requires. - # Value may be a single service key or a tuple of keys for tools - # that compose multiple services. - _TOOL_SERVICE_MAP: dict[str, str | tuple[str, ...]] = { - "save_memory": "memory_store", - "search_memory": "memory_store", - "update_memory": "memory_store", - "delete_memory": "memory_store", - "kb_search": "kb_manager", - "kb_ingest_text": "kb_manager", - "kb_ingest_file": "kb_manager", - "kb_ingest_url": "kb_manager", - "kb_read": "kb_manager", - "kb_list": "kb_manager", - "kb_write_concept": "kb_manager", - "kb_search_concepts": "kb_manager", - "web_fetch": "llm_summarizer", - "sandbox_execute": "sandbox_manager", - "search_arxiv": "arxiv_source", - "fetch_arxiv_paper": "arxiv_source", - "ingest_arxiv_paper": ("arxiv_source", "kb_manager"), - # Long-running jobs: the observe-only tools need the JobManager service. - # ``run_shell_job`` is an application-provided typed starter (see - # examples/jobs_demo.py), not a framework tool — its name is mapped here - # by convention so an app can add it without also adding observe tools. - "run_shell_job": "job_manager", - "job_status": "job_manager", - "job_result": "job_manager", - "job_logs": "job_manager", - "job_cancel": "job_manager", - "job_list": "job_manager", - } - def _resolve_config_tool_refs(self) -> None: """Resolve string/dotted-path tool refs in configs to callables. @@ -381,26 +348,34 @@ def _resolve_config_tool_refs(self) -> None: config.tools = resolve_tools(config.tools) def _detect_required_managers(self) -> set[str]: - """Detect which services are needed by scanning tool names. + """Detect which services the configured tools declared they need. + + Each tool declares its own dependencies via + ``register_tool(..., requires=...)``, so an extension can ship a + service-backed tool without editing the framework. Tools that are not + registered (plain callables) declare nothing and need nothing. + + Resolution is by registry identity only: a renamed tool's original + callable still declares its services, but a callable the default + registry never issued declares nothing, however it is named. Building a + knowledge base because an application named a function ``kb_search`` + would be work done for a tool that is denied at permission time. Returns: Set of required service keys (e.g. ``{"kb_manager", "memory_store"}``). """ + from agentic_cli.tools.registry import identify_tool + required: set[str] = set() for config in self._agent_configs: for tool in config.tools or []: - name = getattr(tool, "__name__", "") - service = self._TOOL_SERVICE_MAP.get(name) - if service is None: - continue - if isinstance(service, tuple): - required.update(service) - else: - required.add(service) + definition = identify_tool(tool) + if definition is not None: + required.update(definition.requires) return required def _ensure_managers_initialized(self) -> None: - """Create managers based on detected requirements. + """Create and publish the services detected from tool metadata. Called during initialize_services() to lazily create only the managers that are actually needed by the configured tools. diff --git a/tests/test_knowledge_tools.py b/tests/test_knowledge_tools.py index 2bc0ce9..39a768b 100644 --- a/tests/test_knowledge_tools.py +++ b/tests/test_knowledge_tools.py @@ -938,10 +938,11 @@ def test_service_roundtrip(self): finally: token.var.reset(token) - def test_kb_manager_detected_via_tool_service_map(self): - """Verify kb tools are detected via _TOOL_SERVICE_MAP (not @requires).""" - from agentic_cli.workflow.base_manager import BaseWorkflowManager - assert "kb_search" in BaseWorkflowManager._TOOL_SERVICE_MAP + def test_kb_tools_declare_the_kb_manager_service(self): + """KB tools carry their own service requirement in the registry.""" + from agentic_cli.tools.registry import get_registry + + assert get_registry().get("kb_search").requires == ("kb_manager",) def test_base_manager_has_kb_manager_slot(self): from agentic_cli.workflow.base_manager import BaseWorkflowManager diff --git a/tests/test_webfetch.py b/tests/test_webfetch.py index d9d0f6b..f1e4259 100644 --- a/tests/test_webfetch.py +++ b/tests/test_webfetch.py @@ -638,12 +638,11 @@ async def test_web_fetch_no_summarizer(self): finally: token.var.reset(token) - def test_web_fetch_detected_via_tool_service_map(self): - """Test web_fetch is detected via _TOOL_SERVICE_MAP.""" - from agentic_cli.workflow.base_manager import BaseWorkflowManager + def test_web_fetch_declares_the_llm_summarizer_service(self): + """web_fetch carries its own service requirement in the registry.""" + from agentic_cli.tools.registry import get_registry - assert "web_fetch" in BaseWorkflowManager._TOOL_SERVICE_MAP - assert BaseWorkflowManager._TOOL_SERVICE_MAP["web_fetch"] == "llm_summarizer" + assert get_registry().get("web_fetch").requires == ("llm_summarizer",) class TestWorkflowManagerIntegration: diff --git a/tests/tools/test_registry_identity.py b/tests/tools/test_registry_identity.py new file mode 100644 index 0000000..5fb20f9 --- /dev/null +++ b/tests/tools/test_registry_identity.py @@ -0,0 +1,1390 @@ +"""Canonical tool identity and registry-declared service requirements. + +Two defects: + +1. ``register_tool(name="public_name")`` stored the public name but handed the + original callable to the backend, which derives the model-visible tool name + from ``func.__name__``. Permission lookup (keyed on the backend's name) then + missed the registry entry entirely. +2. Service detection read a central ``_TOOL_SERVICE_MAP`` keyed by tool name, + so an extension could not ship a service-backed tool without editing the + framework. Tools now declare ``requires=``. +""" + +from __future__ import annotations + +import asyncio +import inspect +from typing import Callable + +import pytest + +from agentic_cli.tools.registry import ToolCategory, ToolRegistry +from agentic_cli.workflow.permissions import EXEMPT +from agentic_cli.workflow.permissions.capabilities import Capability + + +class TestCanonicalName: + def test_renamed_tool_exposes_the_registered_name(self): + registry = ToolRegistry() + + @registry.register(name="public_name", capabilities=EXEMPT) + def _internal_impl(query: str) -> dict: + """Do a thing.""" + return {"success": True, "query": query} + + defn = registry.get("public_name") + assert defn is not None + # What the backend will call the tool == what the registry knows. + assert defn.func.__name__ == "public_name" + assert _internal_impl.__name__ == "public_name" + + def test_renamed_tool_still_runs_and_keeps_its_contract(self): + registry = ToolRegistry() + + @registry.register(name="public_name", capabilities=EXEMPT) + def _internal_impl(query: str, limit: int = 5) -> dict: + """Do a thing.""" + return {"success": True, "query": query, "limit": limit} + + result = registry.get("public_name").func("hello") + assert result == {"success": True, "query": "hello", "limit": 5} + + def test_signature_and_docstring_survive_the_rename(self): + registry = ToolRegistry() + + @registry.register(name="public_name", capabilities=EXEMPT) + def _internal_impl(query: str, limit: int = 5) -> dict: + """First line of docs.""" + return {"success": True} + + func = registry.get("public_name").func + assert list(inspect.signature(func).parameters) == ["query", "limit"] + assert func.__doc__.startswith("First line of docs.") + assert registry.get("public_name").description == "First line of docs." + assert inspect.signature(func).parameters["limit"].default == 5 + + def test_async_tool_stays_async(self): + registry = ToolRegistry() + + @registry.register(name="public_async", capabilities=EXEMPT) + async def _internal_async(x: int) -> dict: + """Async thing.""" + return {"success": True, "x": x} + + func = registry.get("public_async").func + assert inspect.iscoroutinefunction(func) + assert registry.get("public_async").is_async is True + assert asyncio.run(func(3)) == {"success": True, "x": 3} + + def test_unrenamed_tool_is_not_wrapped(self): + """No rename, no wrapper — the registered callable is the original.""" + registry = ToolRegistry() + + def plain_tool() -> dict: + """Plain.""" + return {"success": True} + + returned = registry.register(plain_tool, capabilities=EXEMPT) + assert returned is plain_tool + assert registry.get("plain_tool").func is plain_tool + + def test_permission_lookup_finds_the_renamed_tool(self): + """The plugin looks the tool up by the backend-visible name.""" + registry = ToolRegistry() + + @registry.register( + name="public_name", + capabilities=[Capability("fs.read", target_arg="path")], + ) + def _internal_impl(path: str) -> dict: + """Read.""" + return {"success": True} + + backend_visible_name = registry.get("public_name").func.__name__ + defn = registry.get(backend_visible_name) + assert defn is not None, "permission lookup would fail closed on a real tool" + assert defn.capabilities[0].name == "fs.read" + + +class TestIdentityIsObjectIdentity: + """The identity map is keyed by ``is``, not by hash/eq or by name. + + Identity is owned *per registry*, so these use ``registry.identify()``; + the module-level ``identify_tool()`` answers for the default registry only + (see :class:`TestIdentityIsPerRegistry`). + """ + + def test_registered_callable_resolves(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + assert registry.identify(a_tool) is registry.get("a_tool") + + def test_equality_colliding_object_does_not_resolve(self): + """``WeakKeyDictionary`` semantics let this forge identity; ``is`` does not.""" + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + genuine = registry.get("a_tool").func + + class _Collider: + __name__ = "a_tool" + + def __hash__(self): + return hash(genuine) + + def __eq__(self, other): + return other is genuine + + def __call__(self): # pragma: no cover + return {"success": True} + + collider = _Collider() + assert collider == genuine and hash(collider) == hash(genuine) + assert registry.identify(collider) is None + + def test_same_name_different_object_does_not_resolve(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + def a_tool_impostor() -> dict: # noqa: D401 + """Impostor.""" + return {"success": True} + + a_tool_impostor.__name__ = "a_tool" + assert registry.identify(a_tool_impostor) is None + + def test_binding_does_not_keep_the_object_alive(self): + """Entries are weak, so per-manager service-bound tools are collectable.""" + import gc + import weakref + + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + definition = registry.get("a_tool") + + def _variant() -> dict: + """Variant.""" + return {"success": True} + + registry.bind_identity(_variant, definition) + assert registry.identify(_variant) is definition + + ref = weakref.ref(_variant) + del _variant + gc.collect() + assert ref() is None + + def test_non_weakrefable_object_is_not_bound(self): + """An unbindable object fails closed rather than being keyed by id alone. + + Binding it by id would hand its capabilities to whatever object next + lands on that address. + """ + import weakref + + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + class _NoWeakref: + __slots__ = () # no __weakref__ slot + + def __call__(self): # pragma: no cover + return {"success": True} + + obj = _NoWeakref() + with pytest.raises(TypeError): + weakref.ref(obj) + + registry.bind_identity(obj, registry.get("a_tool")) + assert registry.identify(obj) is None + + +class TestDuplicateNamePolicy: + """Registering a name twice is an error, not a silent takeover. + + The second registration replaced ``_tools[name]`` and left the first + definition's identity bindings in place, so the same tool name resolved to + two different definitions depending on which callable you asked about — and + the framework's service map, keyed by name, would hand an agent the + *framework's* implementation for a name an application had taken over. + """ + + def test_conflicting_registration_raises(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """First.""" + return {"success": True} + + with pytest.raises(ValueError, match="already registered"): + + @registry.register( + name="a_tool", + capabilities=[Capability("fs.read", target_arg="path")], + ) + def _second(path: str) -> dict: + """Second, and it wants more.""" + return {"success": True} + + def test_conflicting_requires_raises(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """First.""" + return {"success": True} + + with pytest.raises(ValueError, match="already registered"): + + @registry.register( + name="a_tool", capabilities=EXEMPT, requires="kb_manager" + ) + def _second() -> dict: + """Second.""" + return {"success": True} + + def test_identical_metadata_is_not_enough_to_alias(self): + """Same capabilities, different docstring and signature — still a clash. + + Metadata equality says nothing about what the model sees or what the + callable does, so it cannot be grounds for silently sharing a name. + Sharing must be declared (see :class:`TestDeclaredVariants`). + """ + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT, category=ToolCategory.PLANNING) + def a_tool(content: str) -> dict: + """Save the thing.""" + return {"success": True} + + with pytest.raises(ValueError, match="already registered"): + + @registry.register( + name="a_tool", capabilities=EXEMPT, category=ToolCategory.PLANNING + ) + def _other_backend(content: str, extra: int = 0) -> dict: + """Save the thing, differently, with another argument.""" + return {"success": True} + + +class TestDeclaredVariants: + """Backend-native implementations share metadata only when they say so.""" + + def _declared(self): + from agentic_cli.tools.registry import declare_tool + + registry = ToolRegistry() + declare_tool( + "a_tool", + description="Do the thing.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + registry=registry, + ) + return registry + + def test_a_declaration_has_no_implementation(self): + registry = self._declared() + definition = registry.get("a_tool") + + assert definition.func is None + assert definition.variants == () + + def test_variants_bind_to_the_declaration(self): + registry = self._declared() + definition = registry.get("a_tool") + + @registry.register( + variant_of="a_tool", capabilities=EXEMPT, category=ToolCategory.PLANNING + ) + def a_tool(content: str, ctx: object) -> dict: + """Backend A.""" + return {"success": True} + + @registry.register( + variant_of="a_tool", capabilities=EXEMPT, category=ToolCategory.PLANNING + ) + def _backend_b(content: str, state: dict) -> dict: + """Backend B, other signature entirely.""" + return {"success": True} + + assert registry.identify(a_tool) is definition + assert registry.identify(_backend_b) is definition + assert len(registry.get("a_tool").variants) == 2 + assert registry.get("a_tool").func is None + + def test_a_variant_may_not_change_the_capabilities(self): + registry = self._declared() + + with pytest.raises(ValueError, match="different declaration"): + + @registry.register( + variant_of="a_tool", + capabilities=[Capability("fs.read", target_arg="path")], + category=ToolCategory.PLANNING, + ) + def _greedy(path: str) -> dict: + """Backend that wants more.""" + return {"success": True} + + def test_a_variant_may_not_change_requires(self): + registry = self._declared() + + with pytest.raises(ValueError, match="different declaration"): + + @registry.register( + variant_of="a_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + requires="kb_manager", + ) + def _needy() -> dict: + """Backend that wants a service.""" + return {"success": True} + + def test_variant_of_an_unknown_tool_raises(self): + registry = ToolRegistry() + + with pytest.raises(ValueError, match="not declared"): + + @registry.register(variant_of="nope", capabilities=EXEMPT) + def _orphan() -> dict: + """Nothing to be a variant of.""" + return {"success": True} + + def test_a_declaration_cannot_be_registered_over(self): + registry = self._declared() + + with pytest.raises(ValueError, match="already registered"): + + @registry.register(name="a_tool", capabilities=EXEMPT) + def _takeover() -> dict: + """Takeover.""" + return {"success": True} + + def test_colliding_with_a_builtin_service_tool_raises(self): + from agentic_cli.tools.knowledge_tools import kb_search # noqa: F401 + from agentic_cli.tools.registry import get_registry, register_tool + + assert get_registry().get("kb_search") is not None + + with pytest.raises(ValueError, match="kb_search"): + + @register_tool(name="kb_search", capabilities=EXEMPT) + def _my_kb_search(query: str) -> dict: + """An application's own search.""" + return {"success": True, "mine": True} + + def test_the_builtin_survives_a_rejected_collision(self): + """The registry is unchanged, so the framework's tool still runs.""" + from agentic_cli.tools.knowledge_tools import kb_search + from agentic_cli.tools.registry import get_registry, identify_tool, register_tool + + before = get_registry().get("kb_search") + + with pytest.raises(ValueError): + + @register_tool(name="kb_search", capabilities=EXEMPT) + def _my_kb_search(query: str) -> dict: + """Impostor.""" + return {"success": True, "mine": True} + + assert get_registry().get("kb_search") is before + assert identify_tool(kb_search) is before + + def test_replace_retires_the_old_identities(self): + from agentic_cli.tools.registry import identify_tool # noqa: F401 + + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """First.""" + return {"success": True} + + first = a_tool + + @registry.register(name="a_tool", capabilities=EXEMPT, replace=True) + def _second() -> dict: + """Second.""" + return {"success": True} + + assert registry.identify(first) is None, "a retired tool kept its capabilities" + assert registry.identify(_second) is registry.get("a_tool") + + def test_replace_keeps_the_new_definition_reachable(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """First.""" + return {"success": True} + + @registry.register( + name="a_tool", + capabilities=[Capability("fs.read", target_arg="path")], + replace=True, + ) + def _second(path: str) -> dict: + """Second.""" + return {"success": True} + + definition = registry.get("a_tool") + assert definition.capabilities[0].name == "fs.read" + + +class TestServiceSubstitutionIsExact: + """A service variant replaces the *definition it implements*, not a name. + + ``_build_tools`` looked up ``service_map[definition.name]``. If an + application deliberately replaced a built-in service tool, the framework + still substituted its own closure — so a different implementation ran than + the one the agent was configured with. + """ + + @staticmethod + def _manager(tools): + from unittest.mock import MagicMock + + from agentic_cli.workflow.config import AgentConfig + + settings = MagicMock() + settings.app_name = "test-app" + config = AgentConfig( + name="a", prompt="p", tools=list(tools), include_state_tools=False + ) + return _stub_manager_cls()(agent_configs=[config], settings=settings), config + + def test_replaced_builtin_runs_the_replacement(self): + """The proof: the assembled tool is the application's implementation.""" + from agentic_cli.tools.registry import ToolRegistry as _Registry # noqa: F401 + from agentic_cli.tools.registry import get_registry, register_tool + + from agentic_cli.tools.knowledge_tools import kb_search as _builtin # noqa: F401 + + registry = get_registry() + original = registry.get("kb_search") + try: + + @register_tool(name="kb_search", capabilities=EXEMPT, replace=True) + def _app_kb_search(query: str) -> dict: + """The application's own search.""" + return {"success": True, "implementation": "application"} + + mgr, config = self._manager([_app_kb_search]) + service_map = self._framework_kb_variant() + + built = mgr._build_tools(config, service_map=service_map) + + assert len(built) == 1 + assert built[0]("q")["implementation"] == "application", ( + "the framework's service variant replaced the application's tool" + ) + finally: + # Restore the registry for the rest of the session. + registry._tools["kb_search"] = original + registry.bind_identity(original.func, original) + + @staticmethod + def _framework_kb_variant() -> dict: + from unittest.mock import MagicMock + + from agentic_cli.tools.factories import make_kb_tools + + return {t.__name__: t for t in make_kb_tools(MagicMock())} + + def test_genuine_builtin_is_still_substituted(self): + from agentic_cli.tools.knowledge_tools import kb_search + + mgr, config = self._manager([kb_search]) + service_map = self._framework_kb_variant() + + built = mgr._build_tools(config, service_map=service_map) + + assert built == [service_map["kb_search"]] + + def test_unbound_variant_is_not_substituted(self): + """A map entry that is not the registered tool proves nothing.""" + from agentic_cli.tools.knowledge_tools import kb_search + + mgr, config = self._manager([kb_search]) + impostor = lambda query: {"success": True} # noqa: E731 + impostor.__name__ = "kb_search" + + built = mgr._build_tools(config, service_map={"kb_search": impostor}) + + assert built == [kb_search] + + +class TestVariantContractIsComplete: + """The declared-variant mechanism, exercised the way callers use it.""" + + def _declared(self, description: str = "Do the thing."): + from agentic_cli.tools.registry import declare_tool + + registry = ToolRegistry() + declare_tool( + "a_tool", + description=description, + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + registry=registry, + ) + return registry + + def test_direct_register_returns_the_canonical_callable(self): + registry = self._declared() + + def _backend_impl(content: str) -> dict: + """Backend implementation with its own name.""" + return {"success": True} + + returned = registry.register( + _backend_impl, + variant_of="a_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + assert returned is not None + assert callable(returned) + assert returned.__name__ == "a_tool" + + def test_module_level_register_tool_returns_the_canonical_callable(self): + from agentic_cli.tools.registry import declare_tool, register_tool + + declare_tool( + "_variant_probe_tool", + description="Probe.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + def _probe_impl(content: str) -> dict: + """Probe implementation.""" + return {"success": True} + + returned = register_tool( + _probe_impl, + variant_of="_variant_probe_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + assert returned is not None and returned.__name__ == "_variant_probe_tool" + + def test_assembly_never_yields_none_for_a_variant(self): + """A config listing the variant's *original* callable must still work.""" + from unittest.mock import MagicMock + + from agentic_cli.tools.registry import declare_tool, register_tool + from agentic_cli.workflow.config import AgentConfig + + declare_tool( + "_assembly_probe_tool", + description="Probe.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + def _assembly_impl(content: str) -> dict: + """Backend implementation under a private name.""" + return {"success": True} + + register_tool( + _assembly_impl, + variant_of="_assembly_probe_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + settings = MagicMock() + settings.app_name = "test-app" + config = AgentConfig( + name="a", prompt="p", tools=[_assembly_impl], include_state_tools=False + ) + mgr = _stub_manager_cls()(agent_configs=[config], settings=settings) + + built = mgr._build_tools(config, service_map={}) + + assert built and built[0] is not None, "assembly produced a None tool" + assert built[0].__name__ == "_assembly_probe_tool" + + def test_redeclaring_with_a_different_description_raises(self): + from agentic_cli.tools.registry import declare_tool + + registry = self._declared() + + with pytest.raises(ValueError, match="different declaration"): + declare_tool( + "a_tool", + description="Something else entirely.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + registry=registry, + ) + + def test_redeclaring_the_same_contract_is_idempotent(self): + from agentic_cli.tools.registry import declare_tool + + registry = self._declared() + first = registry.get("a_tool") + + again = declare_tool( + "a_tool", + description="Do the thing.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + registry=registry, + ) + assert again is first + + def test_variant_order_is_deterministic(self): + """Ordered by module and qualified name, not by registration order.""" + registry = self._declared() + + def _zeta(content: str) -> dict: + """Z.""" + return {"success": True} + + def _alpha(content: str) -> dict: + """A.""" + return {"success": True} + + for impl in (_zeta, _alpha): + registry.register( + impl, + variant_of="a_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + keys = [ + (v.__module__, getattr(v, "__wrapped__", v).__qualname__) + for v in registry.get("a_tool").variants + ] + assert keys == sorted(keys) + + def test_ambiguity_error_distinguishes_variants_by_module(self): + pytest.importorskip("langgraph") + + from agentic_cli.tools.adk import state_tools as _adk # noqa: F401 + from agentic_cli.tools.langgraph import state_tools as _lg # noqa: F401 + from agentic_cli.tools.tool_resolver import resolve_tool + + with pytest.raises(ValueError, match="ambiguous") as exc: + resolve_tool("save_plan") + + message = str(exc.value) + assert "agentic_cli.tools.adk.state_tools" in message + assert "agentic_cli.tools.langgraph.state_tools" in message + + +class TestReplacedStateToolIsNotInjected: + """``replace=True`` on a state tool must retire it from auto-injection.""" + + def test_retired_variants_are_not_auto_injected(self): + pytest.importorskip("google.adk") + from unittest.mock import MagicMock + + from agentic_cli.tools.registry import get_registry, register_tool + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + from agentic_cli.workflow.config import AgentConfig + + from agentic_cli.tools.adk import state_tools as _adk # noqa: F401 + + registry = get_registry() + original = registry.get("save_plan") + assert original is not None + try: + + @register_tool(name="save_plan", capabilities=EXEMPT, replace=True) + def _app_save_plan(content: str) -> dict: + """The application's own plan tool.""" + return {"success": True, "implementation": "application"} + + settings = MagicMock() + settings.app_name = "test-app" + config = AgentConfig( + name="a", prompt="p", tools=[_app_save_plan], include_state_tools=True + ) + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = settings + mgr._services = {} + + built = mgr._build_tools(config, service_map={}) + + names = [getattr(t, "__name__", "") for t in built] + assert names.count("save_plan") == 1, ( + f"a retired variant was injected alongside the replacement: {names}" + ) + assert built[names.index("save_plan")]("x")["implementation"] == ( + "application" + ) + finally: + registry._tools["save_plan"] = original + for variant in original.variants: + registry.bind_identity(variant, original) + + +class TestIdentityIsPerRegistry: + """Each registry owns its bindings; the framework trusts only its own.""" + + def test_default_registry_does_not_see_another_registrys_binding(self): + from agentic_cli.tools.registry import identify_tool + + other = ToolRegistry() + + @other.register(capabilities=EXEMPT) + def foreign() -> dict: + """Registered elsewhere.""" + return {"success": True} + + assert other.identify(foreign) is other.get("foreign") + assert identify_tool(foreign) is None + + def test_default_registry_sees_its_own_binding(self): + from agentic_cli.tools.registry import ( + get_registry, + identify_tool, + register_tool, + ) + + @register_tool(capabilities=EXEMPT) + def _default_registry_probe_tool() -> dict: + """Registered in the framework's registry.""" + return {"success": True} + + assert identify_tool(_default_registry_probe_tool) is get_registry().get( + "_default_registry_probe_tool" + ) + + +class TestIdentityDoesNotPinDeadObjects: + """A binding must not keep an otherwise-dead registry graph alive. + + The identity map used to be a module-level dict holding the + ``ToolDefinition`` *strongly*, and a definition holds its callable — so the + weak reference to that callable could never fire. Every local + ``ToolRegistry`` (one per test, per short-lived tool set) leaked its + definitions and closures for the life of the process. + """ + + def test_local_registry_graph_is_collectable(self): + import gc + import weakref + + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def a_tool() -> dict: + """A tool.""" + return {"success": True} + + definition = registry.get("a_tool") + refs = { + "registry": weakref.ref(registry), + "definition": weakref.ref(definition), + "callable": weakref.ref(definition.func), + } + + del registry, definition, a_tool + gc.collect() + + alive = [name for name, ref in refs.items() if ref() is not None] + assert alive == [], f"identity binding pinned {alive}" + + def test_renamed_wrapper_graph_is_collectable(self): + import gc + import weakref + + registry = ToolRegistry() + + def _impl() -> dict: + """Impl.""" + return {"success": True} + + canonical = registry.register(_impl, name="public", capabilities=EXEMPT) + definition = registry.get("public") + refs = { + "registry": weakref.ref(registry), + "definition": weakref.ref(definition), + "canonical": weakref.ref(canonical), + "original": weakref.ref(_impl), + } + + del registry, definition, canonical, _impl + gc.collect() + + alive = [name for name, ref in refs.items() if ref() is not None] + assert alive == [], f"identity binding pinned {alive}" + + def test_service_bound_variant_is_collectable(self): + """Per-manager factory closures must not accumulate.""" + import gc + import weakref + + from agentic_cli.tools.registry import bind_tool_identity, get_registry + + definition = get_registry().get("read_file") + + def variant() -> dict: + """A service-bound variant of a long-lived registered tool.""" + return {"success": True} + + bind_tool_identity(variant, definition) + ref = weakref.ref(variant) + del variant + gc.collect() + assert ref() is None + + +class TestRequiresMetadata: + def test_requires_is_recorded(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT, requires="kb_manager") + def kb_thing() -> dict: + """KB thing.""" + return {"success": True} + + assert registry.get("kb_thing").requires == ("kb_manager",) + + def test_multiple_services_are_recorded_in_order(self): + registry = ToolRegistry() + + @registry.register( + capabilities=EXEMPT, requires=("arxiv_source", "kb_manager") + ) + def composite() -> dict: + """Composite.""" + return {"success": True} + + assert registry.get("composite").requires == ("arxiv_source", "kb_manager") + + def test_no_requires_defaults_to_empty(self): + registry = ToolRegistry() + + @registry.register(capabilities=EXEMPT) + def plain() -> dict: + """Plain.""" + return {"success": True} + + assert registry.get("plain").requires == () + + def test_unknown_service_key_raises(self): + registry = ToolRegistry() + + with pytest.raises(ValueError, match="unknown required service"): + + @registry.register(capabilities=EXEMPT, requires="not_a_service") + def bad() -> dict: + """Bad.""" + return {"success": True} + + def test_wrong_type_raises(self): + registry = ToolRegistry() + + with pytest.raises(TypeError, match="requires must be"): + + @registry.register(capabilities=EXEMPT, requires=object()) + def bad() -> dict: + """Bad.""" + return {"success": True} + + def test_builtin_service_tools_declare_their_services(self): + """The shipped service-backed tools carry their own metadata.""" + from agentic_cli.tools.arxiv_tools import ingest_arxiv_paper # noqa: F401 + from agentic_cli.tools.knowledge_tools import kb_search # noqa: F401 + from agentic_cli.tools.memory_tools import save_memory # noqa: F401 + from agentic_cli.tools.registry import get_registry + + reg = get_registry() + assert reg.get("kb_search").requires == ("kb_manager",) + assert reg.get("save_memory").requires == ("memory_store",) + assert reg.get("ingest_arxiv_paper").requires == ( + "arxiv_source", + "kb_manager", + ) + assert reg.get("read_file").requires == () + + +class TestManagerDetectionUsesRegistry: + """Managers derive required services from tool metadata, not a name map.""" + + def test_detection_reads_declared_requires(self): + from unittest.mock import MagicMock + + from agentic_cli.workflow.base_manager import BaseWorkflowManager + from agentic_cli.workflow.config import AgentConfig + + class _Manager(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + from agentic_cli.tools.knowledge_tools import kb_search + from agentic_cli.tools.memory_tools import save_memory + + settings = MagicMock() + settings.app_name = "test-app" + mgr = _Manager( + agent_configs=[AgentConfig(name="a", prompt="p", tools=[kb_search, save_memory])], + settings=settings, + ) + assert mgr.required_managers == {"kb_manager", "memory_store"} + + def test_extension_tool_needs_no_framework_edit(self): + """A tool defined outside the framework still gets its service created.""" + from unittest.mock import MagicMock + + from agentic_cli.tools.registry import register_tool + from agentic_cli.workflow.base_manager import BaseWorkflowManager + from agentic_cli.workflow.config import AgentConfig + + @register_tool(capabilities=EXEMPT, requires="sandbox_manager") + def _extension_tool() -> dict: + """An app-provided tool the framework has never heard of.""" + return {"success": True} + + class _Manager(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + settings = MagicMock() + settings.app_name = "test-app" + mgr = _Manager( + agent_configs=[AgentConfig(name="a", prompt="p", tools=[_extension_tool])], + settings=settings, + ) + assert mgr.required_managers == {"sandbox_manager"} + + def test_central_tool_service_map_is_gone(self): + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + assert not hasattr(BaseWorkflowManager, "_TOOL_SERVICE_MAP") + + +def _stub_manager_cls(): + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + class _Manager(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + return _Manager + + +class TestAssemblyRequiresBoundIdentity: + """Assembly must key on the bound definition, never on a matching name. + + ``lookup_definition()`` fell back to ``registry.get(tool.__name__)``, so a + plain callable an application happened to name ``kb_search`` was handed the + framework's service-bound variant, had its services constructed, and was + wrapped as long-running — all for a function the registry never issued and + that the permission engine then (correctly) denied. + """ + + @staticmethod + def _manager(tools): + from unittest.mock import MagicMock + + from agentic_cli.workflow.config import AgentConfig + + settings = MagicMock() + settings.app_name = "test-app" + config = AgentConfig( + name="a", prompt="p", tools=list(tools), include_state_tools=False + ) + return _stub_manager_cls()(agent_configs=[config], settings=settings), config + + def test_raw_same_name_callable_declares_no_services(self): + def kb_search(query: str) -> dict: + """An application's own function that happens to share a name.""" + return {"success": True} + + mgr, _ = self._manager([kb_search]) + assert mgr.required_managers == set() + + def test_raw_same_name_callable_is_not_substituted(self): + def kb_search(query: str) -> dict: + """Impostor.""" + return {"success": True} + + mgr, config = self._manager([kb_search]) + service_variant = lambda query: {"success": True} # noqa: E731 + built = mgr._build_tools(config, service_map={"kb_search": service_variant}) + + assert built == [kb_search], "an unregistered callable was replaced" + + def test_raw_same_name_callable_is_not_wrapped_long_running(self): + pytest.importorskip("google.adk") + from unittest.mock import MagicMock + + from google.adk.tools import LongRunningFunctionTool + + from agentic_cli.tools.registry import get_registry, register_tool + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + if get_registry().get("_lr_reference_tool") is None: + + @register_tool(capabilities=EXEMPT, long_running=True) + def _lr_reference_tool() -> dict: + """A genuine long-running tool.""" + return {"success": True} + + def _lr_reference_tool() -> dict: # noqa: F811 - deliberate collision + """Impostor.""" + return {"success": True} + + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = MagicMock() + wrapped = mgr._wrap_long_running([_lr_reference_tool]) + + assert wrapped == [_lr_reference_tool] + assert not isinstance(wrapped[0], LongRunningFunctionTool) + + async def test_raw_same_name_callable_is_denied_at_permission_time(self, tmp_path): + pytest.importorskip("google.adk") + from google.adk.tools import FunctionTool + + from agentic_cli.config import BaseSettings + from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin + from agentic_cli.workflow.permissions import PermissionEngine + from agentic_cli.workflow.permissions.store import PermissionContext + from agentic_cli.workflow.service_registry import ( + PERMISSION_ENGINE, + set_service_registry, + ) + + def kb_search(query: str) -> dict: + """Impostor.""" + return {"success": True} + + class _Stub: + async def request_user_input(self, request): + return "deny" + + engine = PermissionEngine( + settings=BaseSettings(google_api_key="test"), + workflow=_Stub(), + ctx=PermissionContext(workdir=tmp_path, home=tmp_path), + ) + token = set_service_registry({PERMISSION_ENGINE: engine}) + try: + result = await PermissionPlugin().before_tool_callback( + tool=FunctionTool(func=kb_search), tool_args={}, tool_context=None + ) + finally: + token.var.reset(token) + assert isinstance(result, dict) and result["success"] is False + + def test_tool_from_another_registry_is_left_alone(self): + """A private registry is not the framework's registry.""" + other = ToolRegistry() + + @other.register( + capabilities=EXEMPT, requires="sandbox_manager", long_running=True + ) + def foreign_tool() -> dict: + """Registered somewhere else entirely.""" + return {"success": True} + + mgr, config = self._manager([foreign_tool]) + assert mgr.required_managers == set() + assert mgr._build_tools(config, service_map={}) == [foreign_tool] + + async def test_tool_from_another_registry_is_denied(self, tmp_path): + pytest.importorskip("google.adk") + from google.adk.tools import FunctionTool + + from agentic_cli.config import BaseSettings + from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin + from agentic_cli.workflow.permissions import PermissionEngine + from agentic_cli.workflow.permissions.store import PermissionContext + from agentic_cli.workflow.service_registry import ( + PERMISSION_ENGINE, + set_service_registry, + ) + + other = ToolRegistry() + + @other.register(capabilities=EXEMPT) + def foreign_exempt_tool() -> dict: + """EXEMPT — but only according to a registry nobody trusts.""" + return {"success": True} + + class _Stub: + async def request_user_input(self, request): + return "deny" + + engine = PermissionEngine( + settings=BaseSettings(google_api_key="test"), + workflow=_Stub(), + ctx=PermissionContext(workdir=tmp_path, home=tmp_path), + ) + token = set_service_registry({PERMISSION_ENGINE: engine}) + try: + result = await PermissionPlugin().before_tool_callback( + tool=FunctionTool(func=foreign_exempt_tool), + tool_args={}, + tool_context=None, + ) + finally: + token.var.reset(token) + assert isinstance(result, dict) and result["success"] is False, ( + "a private registry's EXEMPT declaration was honoured" + ) + + def test_string_registry_reference_still_resolves(self): + """A config may name a tool; that is a registry lookup, not a guess.""" + from unittest.mock import MagicMock + + from agentic_cli.tools.factories import make_kb_tools + from agentic_cli.tools.knowledge_tools import kb_search # noqa: F401 + + mgr, config = self._manager(["kb_search"]) + assert mgr.required_managers == {"kb_manager"} + + service_map = {t.__name__: t for t in make_kb_tools(MagicMock())} + built = mgr._build_tools(config, service_map=service_map) + assert built == [service_map["kb_search"]] + + +class TestDirectRegisterOriginalCallable: + """``register(func, name=..., ...)`` — the caller may keep using ``func``. + + Assembly used to key on ``func.__name__``, which for a renamed tool is the + private implementation name. The tool's declared services were then never + created, ``long_running`` never applied, and the model saw the private name. + """ + + # The original callable, registered once for the whole module: a name may + # only be registered once (see TestDuplicateNamePolicy). + _original: "Callable | None" = None + + @classmethod + def _register_renamed(cls): + from agentic_cli.tools.registry import get_registry, register_tool + + if cls._original is None: + + def _private_impl(query: str) -> dict: + """A renamed, service-backed, long-running tool.""" + return {"success": True} + + register_tool( + _private_impl, + name="renamed_public_tool", + capabilities=EXEMPT, + requires="sandbox_manager", + long_running=True, + ) + cls._original = _private_impl + + assert get_registry().get("renamed_public_tool") is not None + assert cls._original.__name__ == "_private_impl" + return cls._original + + def test_declared_services_are_detected(self): + from unittest.mock import MagicMock + + from agentic_cli.workflow.config import AgentConfig + + original = self._register_renamed() + settings = MagicMock() + settings.app_name = "test-app" + mgr = _stub_manager_cls()( + agent_configs=[AgentConfig(name="a", prompt="p", tools=[original])], + settings=settings, + ) + assert mgr.required_managers == {"sandbox_manager"} + + def test_assembled_tool_carries_the_registered_name(self): + from unittest.mock import MagicMock + + from agentic_cli.workflow.config import AgentConfig + + original = self._register_renamed() + settings = MagicMock() + settings.app_name = "test-app" + config = AgentConfig( + name="a", prompt="p", tools=[original], include_state_tools=False + ) + mgr = _stub_manager_cls()(agent_configs=[config], settings=settings) + + built = mgr._build_tools(config, service_map={}) + assert [getattr(t, "__name__", "") for t in built] == ["renamed_public_tool"] + + def test_long_running_wrapping_uses_identity(self): + pytest.importorskip("google.adk") + from unittest.mock import MagicMock + + from google.adk.tools import LongRunningFunctionTool + + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + original = self._register_renamed() + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = MagicMock() + + wrapped = mgr._wrap_long_running([original]) + assert isinstance(wrapped[0], LongRunningFunctionTool) + + +class TestRequiresAreConstructible: + """Only services a manager can actually build may be declared.""" + + def test_user_kb_manager_is_not_declarable(self): + """It is created with kb_manager, never on its own — declaring it + validated and then provided nothing.""" + registry = ToolRegistry() + + with pytest.raises(ValueError, match="user_kb_manager") as exc: + + @registry.register(capabilities=EXEMPT, requires="user_kb_manager") + def _tool() -> dict: + """Tool.""" + return {"success": True} + + assert "kb_manager" in str(exc.value) + + def test_kb_manager_creates_both_scopes(self): + """The documented replacement really provides the user-scoped KB too.""" + from unittest.mock import MagicMock + + from agentic_cli.workflow.service_registry import ( + KB_MANAGER, + USER_KB_MANAGER, + ) + + from tests.conftest import MockContext + + with MockContext(google_api_key="k", knowledge_base_use_mock=True) as ctx: + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + class _Manager(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + mgr = _Manager(agent_configs=[], settings=ctx.settings) + mgr._required_managers = {"kb_manager"} + mgr._ensure_managers_initialized() + + assert mgr.services.get(KB_MANAGER) is not None + assert mgr.services.get(USER_KB_MANAGER) is not None + + def test_always_present_services_are_not_declarable(self): + registry = ToolRegistry() + for key in ("permission_engine", "workflow"): + with pytest.raises(ValueError, match="always available"): + + @registry.register(capabilities=EXEMPT, requires=key) + def _tool() -> dict: + """Tool.""" + return {"success": True} + + def test_empty_string_requires_is_rejected(self): + registry = ToolRegistry() + with pytest.raises(ValueError, match="non-empty"): + + @registry.register(capabilities=EXEMPT, requires="") + def _tool() -> dict: + """Tool.""" + return {"success": True} + + def test_non_string_element_is_rejected(self): + registry = ToolRegistry() + with pytest.raises(ValueError, match="non-empty"): + + @registry.register(capabilities=EXEMPT, requires=("kb_manager", None)) + def _tool() -> dict: + """Tool.""" + return {"success": True} diff --git a/tests/tools/test_sandbox.py b/tests/tools/test_sandbox.py index 132778a..44f0121 100644 --- a/tests/tools/test_sandbox.py +++ b/tests/tools/test_sandbox.py @@ -576,12 +576,12 @@ async def test_no_workflow(self, mock_app): # --------------------------------------------------------------------------- class TestManagerAutoDetection: - def test_sandbox_detected_via_tool_service_map(self): - """Verify sandbox_execute is detected via _TOOL_SERVICE_MAP.""" - from agentic_cli.workflow.base_manager import BaseWorkflowManager + def test_sandbox_execute_declares_the_sandbox_service(self): + """sandbox_execute carries its own service requirement in the registry.""" + from agentic_cli.tools.registry import get_registry + from agentic_cli.tools.sandbox import sandbox_execute # noqa: F401 - assert "sandbox_execute" in BaseWorkflowManager._TOOL_SERVICE_MAP - assert BaseWorkflowManager._TOOL_SERVICE_MAP["sandbox_execute"] == "sandbox_manager" + assert get_registry().get("sandbox_execute").requires == ("sandbox_manager",) def test_base_manager_detects_sandbox(self, tmp_path): """BaseWorkflowManager picks up sandbox_manager from tool configs.""" From 9c624e1a4bb70474898f7ad105269e682b4a61b1 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:40 -0400 Subject: [PATCH 03/11] feat(workflow)!: validate agent graphs before allocating anything MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multiple roots were silently accepted and all but one tree was unreachable — the runner starts from a single root, so those agents could never run. Duplicate names, dangling ``sub_agents`` references, self-references, delegation cycles and a child claimed by two parents were likewise only discovered as a half-built hierarchy, after model discovery and service creation had already paid for themselves. ``validate_agent_graph()`` now returns a validated ``AgentGraph`` (config map, dependency-ordered build order, root) and raises ``AgentGraphError`` naming the offending agents — before any allocation or network call, since a bad graph is a static configuration error. Agents are built in dependency order, so declaration order no longer changes the hierarchy. Prompt factories are resolved under the manager's settings rather than the global singleton: a factory may take no arguments (including all-defaulted ones) or exactly one settings argument. Other signatures, ``async def`` factories and non-string results are rejected by name instead of producing a coroutine object as an agent's instruction. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/workflow/adk/manager.py | 128 ++++---- src/agentic_cli/workflow/base_manager.py | 12 + src/agentic_cli/workflow/config.py | 237 +++++++++++++- tests/workflow/test_adk_agent_construction.py | 67 ++++ tests/workflow/test_agent_graph.py | 290 ++++++++++++++++++ 5 files changed, 664 insertions(+), 70 deletions(-) create mode 100644 tests/workflow/test_adk_agent_construction.py create mode 100644 tests/workflow/test_agent_graph.py diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index 2c95e80..52eff00 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -22,7 +22,7 @@ from agentic_cli.workflow.base_manager import BaseWorkflowManager from agentic_cli.workflow.events import WorkflowEvent, EventType -from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.config import AgentConfig, validate_agent_graph from agentic_cli.workflow.model_settings import ModelSettings, ThinkingSettings from agentic_cli.workflow.adk.event_processor import ADKEventProcessor from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin @@ -641,81 +641,70 @@ def _wrap_long_running(self, tools: list[Callable]) -> list: wrapped.append(tool) return wrapped + def _build_agent( + self, + config: "AgentConfig", + service_map: dict, + sub_agents: list[Agent] | None = None, + ) -> Agent: + """Construct one ``LlmAgent`` from a config (with resolved sub-agents).""" + return LlmAgent( + name=config.name, + model=self._build_model_arg(config), + instruction=config.get_prompt(self._settings), + tools=self._wrap_long_running( + self._assemble_agent_tools(config, service_map) + ), + description=config.description, + sub_agents=sub_agents or [], + planner=self._get_planner(config), + generate_content_config=self._get_generate_content_config(config), + ) + def _create_agents(self) -> Agent: - """Create agent hierarchy from configs. + """Create the agent hierarchy from configs, in dependency order. + + The graph is validated first (``validate_agent_graph``) so duplicate + names, dangling ``sub_agents`` references, self-references, cycles and + shapes ADK cannot represent fail with the offending agent named, + instead of silently producing a coordinator with missing children. + + Construction then follows a topological order, so an agent is always + built after the agents it delegates to — the previous leaves-then- + coordinators split silently dropped a sub-agent that was itself a + coordinator. + + Prompt factories are called under this manager's settings context (see + ``AgentConfig.get_prompt``), so a prompt that reads settings sees the + manager's instance rather than the global singleton. Returns: - Root agent (the first agent with sub_agents, or first agent if none have sub_agents) + Root agent (the first agent with sub_agents, or the first config). + + Raises: + AgentGraphError: If the configured graph is invalid. """ - # Build a map of agent configs by name - config_map = {config.name: config for config in self._agent_configs} + graph = validate_agent_graph(self._agent_configs, backend=self.backend_type) - # Build agents (non-coordinators first, then coordinators) agent_map: dict[str, Agent] = {} service_map = self._get_service_tool_map() - # First pass: create agents without sub_agents (leaf agents) - for config in self._agent_configs: - if not config.sub_agents: - agent_map[config.name] = LlmAgent( - name=config.name, - model=self._build_model_arg(config), - instruction=config.get_prompt(), - tools=self._wrap_long_running( - self._assemble_agent_tools(config, service_map) - ), - description=config.description or None, - planner=self._get_planner(config), - generate_content_config=self._get_generate_content_config(config), - ) - logger.debug("agent_created", name=config.name, type="leaf") - - # Second pass: create agents with sub_agents (coordinators) - for config in self._agent_configs: - if config.sub_agents: - sub_agent_instances = [] - for sub_name in config.sub_agents: - if sub_name in agent_map: - sub_agent_instances.append(agent_map[sub_name]) - else: - logger.warning( - "sub_agent_not_found", - coordinator=config.name, - sub_agent=sub_name, - ) - - agent_map[config.name] = LlmAgent( - name=config.name, - model=self._build_model_arg(config), - instruction=config.get_prompt(), - tools=self._wrap_long_running( - self._assemble_agent_tools(config, service_map) - ), - description=config.description or None, - sub_agents=sub_agent_instances, - planner=self._get_planner(config), - generate_content_config=self._get_generate_content_config(config), - ) + from agentic_cli.config import SettingsContext + + # Prompt factories run inside the manager's settings context. + with SettingsContext(self._settings): + for name in graph.build_order: + config = graph.config_map[name] + sub_agents = [agent_map[sub] for sub in config.sub_agents] + agent_map[name] = self._build_agent(config, service_map, sub_agents) logger.debug( "agent_created", - name=config.name, - type="coordinator", - sub_agents=[a.name for a in sub_agent_instances], + name=name, + type="coordinator" if sub_agents else "leaf", + sub_agents=[a.name for a in sub_agents], ) - # Find root agent (first with sub_agents, or first in list) - root_agent = None - for config in self._agent_configs: - if config.sub_agents: - root_agent = agent_map[config.name] - break - - if root_agent is None and self._agent_configs: - root_agent = agent_map[self._agent_configs[0].name] - - if root_agent is None: - raise RuntimeError("No agents configured") - + root_agent = agent_map[graph.root_name] logger.info("agents_created", root=root_agent.name, total=len(agent_map)) return root_agent @@ -804,6 +793,17 @@ async def _do_initialize(self) -> None: required_managers=list(self._required_managers), ) + def _validate_agent_graph(self) -> None: + """Reject an unbuildable agent graph before anything is allocated. + + Runs at the top of initialization — ahead of model discovery, service + construction and the session service — so a static configuration error + costs no network call and leaves nothing to roll back. + """ + if self._adk_config_path: + return # native ADK config: ADK owns the topology + validate_agent_graph(self._agent_configs, backend=self.backend_type) + async def _ensure_initialized(self) -> None: """Ensure services are initialized before processing.""" if not self._initialized: diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index ea2b67f..4ca5007 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -620,6 +620,10 @@ async def initialize_services(self, validate: bool = True) -> None: if self._initialized: return + # Validate the declared agent graph before anything is allocated or + # any network call is made: a bad graph is a static configuration + # error and should not cost a model listing or an embedding model. + self._validate_agent_graph() from agentic_cli.config import validate_settings if validate: @@ -646,6 +650,14 @@ async def initialize_services(self, validate: bool = True) -> None: await self._do_initialize() self._initialized = True + def _validate_agent_graph(self) -> None: + """Validate the declared agent graph. Backends may narrow this. + + Runs before discovery and service creation so a static configuration + error surfaces immediately and costs nothing. + """ + return None + @abstractmethod async def _do_initialize(self) -> None: """Backend-specific initialization (create agents/graph). diff --git a/src/agentic_cli/workflow/config.py b/src/agentic_cli/workflow/config.py index 84d1792..b6ebbd1 100644 --- a/src/agentic_cli/workflow/config.py +++ b/src/agentic_cli/workflow/config.py @@ -1,9 +1,11 @@ """Configuration classes for workflow management.""" +import inspect from dataclasses import dataclass, field from typing import Callable, Any, TYPE_CHECKING if TYPE_CHECKING: + from agentic_cli.config import BaseSettings from agentic_cli.workflow.model_settings import ModelSettings from agentic_cli.workflow.mcp import MCPServerConfig @@ -17,7 +19,10 @@ class AgentConfig: Attributes: name: Unique identifier for the agent - prompt: System instruction - either a string or a callable that returns one + prompt: System instruction — either a string or a callable returning one. + The callable may take no arguments, or a single ``settings`` + argument, which receives the manager's settings instance (see + ``get_prompt``). tools: Tools the agent can use. Each entry is a callable, a registered tool name (e.g. "kb_search"), or a dotted import path (e.g. "my_pkg.tools.my_tool"). String refs are resolved to callables @@ -36,7 +41,7 @@ class AgentConfig: """ name: str - prompt: str | Callable[[], str] + prompt: str | Callable[..., str] tools: list[Callable[..., Any] | str] = field(default_factory=list) sub_agents: list[str] = field(default_factory=list) description: str = "" @@ -46,8 +51,228 @@ class AgentConfig: skills: list[str] = field(default_factory=list) include_state_tools: bool = True - def get_prompt(self) -> str: - """Get the prompt string, calling the getter if needed.""" - if callable(self.prompt): + def get_prompt(self, settings: "BaseSettings | None" = None) -> str: + """Get the prompt string, calling the factory if the prompt is callable. + + Supported factory shapes, in the order they are tried: + + 1. **Callable with no arguments** — including one whose parameters all + have defaults (``lambda prefix="x": ...``). Called as-is; it should + read ``get_settings()``, which the manager binds to its own instance + while building agents. + 2. **Callable taking exactly one settings argument** — passed the + manager's settings explicitly. + + Anything else (two required parameters, a required parameter that is + not settings) is rejected: guessing would either drop the caller's + intent or pass settings into an unrelated slot. + + Args: + settings: The manager's settings, when available. + + Returns: + The resolved system instruction. + + Raises: + AgentGraphError: If the factory's signature is unsupported, it is + async, or it does not return a string. The message names the + agent. + """ + if not callable(self.prompt): + return self.prompt + + if inspect.iscoroutinefunction(self.prompt): + raise AgentGraphError( + f"Agent {self.name!r}: async prompt factories are not supported " + "— the instruction is resolved synchronously while agents are " + "built. Use a plain function." + ) + + result = self._call_prompt_factory(settings) + if inspect.isawaitable(result): + raise AgentGraphError( + f"Agent {self.name!r}: the prompt factory returned an awaitable; " + "it must return a string." + ) + if not isinstance(result, str): + raise AgentGraphError( + f"Agent {self.name!r}: the prompt factory returned " + f"{type(result).__name__}, not a string." + ) + return result + + def _call_prompt_factory(self, settings: "BaseSettings | None"): + """Invoke the factory with the arity it actually supports.""" + try: + signature = inspect.signature(self.prompt) + except (TypeError, ValueError): # builtins / C callables + return self.prompt() + + # Preferred and backward-compatible: if it binds with no arguments + # (including all-defaulted parameters), call it with none. + try: + signature.bind() + except TypeError: + pass + else: return self.prompt() - return self.prompt + + if settings is None: + raise AgentGraphError( + f"Agent {self.name!r}: the prompt factory requires an argument " + f"{signature}, but this caller resolved the prompt without " + "settings. Use a zero-argument factory here." + ) + + try: + signature.bind(settings) + except TypeError: + raise AgentGraphError( + f"Agent {self.name!r}: unsupported prompt factory signature " + f"{signature}. A prompt factory must take no arguments, or " + "exactly one argument that receives the settings instance." + ) from None + return self.prompt(settings) + + +class AgentGraphError(ValueError): + """A configured agent graph cannot be built. + + Raised before any backend object is constructed, so the message names the + offending agents rather than surfacing as a partially-built hierarchy. + """ + + +@dataclass(frozen=True) +class AgentGraph: + """A validated agent graph ready to construct, in dependency order. + + Attributes: + config_map: Agent name → config. + build_order: Names ordered so every agent follows its sub-agents. + root_name: The agent the runner starts from. + """ + + config_map: dict[str, AgentConfig] + build_order: tuple[str, ...] + root_name: str + + +def validate_agent_graph( + configs: list[AgentConfig], backend: str = "adk" +) -> AgentGraph: + """Validate an agent graph and return it in dependency (topological) order. + + Checks, in order, so the first failure is the most fundamental: + + 1. at least one agent; + 2. no duplicate agent names; + 3. every ``sub_agents`` entry resolves to a configured agent; + 4. no agent lists itself as a sub-agent; + 5. no delegation cycle; + 6. no agent is a sub-agent of two parents — ADK agents hold a single + ``parent_agent``, so a shared child is a tree the backend cannot build. + + Args: + configs: The declared agents. + backend: Backend name, used only in error messages. + + Returns: + The validated graph plus a build order in which every agent comes after + the agents it delegates to. + + Raises: + AgentGraphError: With the offending agent name(s) in the message. + """ + if not configs: + raise AgentGraphError("No agents configured: at least one AgentConfig is required.") + + config_map: dict[str, AgentConfig] = {} + duplicates: list[str] = [] + for config in configs: + if config.name in config_map: + duplicates.append(config.name) + config_map[config.name] = config + if duplicates: + raise AgentGraphError( + f"Duplicate agent name(s): {', '.join(sorted(set(duplicates)))}. " + "Agent names must be unique." + ) + + missing = [ + f"{config.name} -> {sub}" + for config in configs + for sub in config.sub_agents + if sub not in config_map + ] + if missing: + raise AgentGraphError( + f"Unknown sub_agents reference(s): {', '.join(missing)}. " + f"Known agents: {', '.join(sorted(config_map))}." + ) + + self_refs = [c.name for c in configs if c.name in c.sub_agents] + if self_refs: + raise AgentGraphError( + f"Agent(s) list themselves as sub_agents: {', '.join(sorted(self_refs))}." + ) + + parents: dict[str, str] = {} + shared: list[str] = [] + for config in configs: + for sub in config.sub_agents: + if sub in parents: + shared.append(f"{sub} (of {parents[sub]} and {config.name})") + else: + parents[sub] = config.name + if shared: + raise AgentGraphError( + f"The {backend} backend requires a tree: agent(s) with more than one " + f"parent: {', '.join(sorted(shared))}." + ) + + order = _topological_order(config_map) + # The root is the agent nobody delegates to. Selecting "first config with + # sub_agents" made the root depend on declaration order (listing a + # sub-coordinator before its parent promoted the child to root) and + # silently accepted a forest: only one root is ever run, so every other + # tree — and every agent under it — was unreachable. + roots = [c.name for c in configs if c.name not in parents] + if len(roots) > 1: + raise AgentGraphError( + f"Agent graph has {len(roots)} roots: {', '.join(sorted(roots))}. " + "Exactly one agent may be unreferenced — the runner starts from a " + "single root, so agents under any other root are unreachable. Add " + "the extra root(s) to a coordinator's sub_agents." + ) + return AgentGraph(config_map=config_map, build_order=order, root_name=roots[0]) + + +def _topological_order(config_map: dict[str, AgentConfig]) -> tuple[str, ...]: + """Order agent names so each follows its sub-agents. + + Raises: + AgentGraphError: If a delegation cycle is found (names the cycle). + """ + order: list[str] = [] + done: set[str] = set() + visiting: list[str] = [] + + def _visit(name: str) -> None: + if name in done: + return + if name in visiting: + cycle = visiting[visiting.index(name):] + [name] + raise AgentGraphError( + f"Delegation cycle in sub_agents: {' -> '.join(cycle)}." + ) + visiting.append(name) + for sub in config_map[name].sub_agents: + _visit(sub) + visiting.pop() + done.add(name) + order.append(name) + + for name in config_map: + _visit(name) + return tuple(order) diff --git a/tests/workflow/test_adk_agent_construction.py b/tests/workflow/test_adk_agent_construction.py new file mode 100644 index 0000000..c3e8e95 --- /dev/null +++ b/tests/workflow/test_adk_agent_construction.py @@ -0,0 +1,67 @@ +"""Real ADK agent construction from declarative ``AgentConfig``s. + +These tests build actual ``LlmAgent`` objects through +``GoogleADKWorkflowManager._create_agents()`` rather than pre-seeding manager +internals, so a config the framework documents must really construct. + +Regression: ``description`` defaulted to ``""`` on ``AgentConfig`` but was +converted to ``None`` on the way into ``LlmAgent``. ADK types the field as +``str``, so the documented minimal config (README quick-start: name + prompt + +tools) raised ``ValidationError`` during initialization. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from tests.conftest import MockContext # noqa: E402 + + +def _manager(configs: list[AgentConfig], settings) -> GoogleADKWorkflowManager: + return GoogleADKWorkflowManager(agent_configs=configs, settings=settings) + + +def test_minimal_config_builds_a_real_adk_agent(): + """The documented minimal config (no description) must construct.""" + with MockContext(google_api_key="test-key") as ctx: + mgr = _manager( + [AgentConfig(name="assistant", prompt="You are helpful.")], + ctx.settings, + ) + root = mgr._create_agents() + + assert root.name == "assistant" + # ADK types description as ``str``; the empty default must survive as "". + assert root.description == "" + + +def test_explicit_description_is_preserved(): + with MockContext(google_api_key="test-key") as ctx: + mgr = _manager( + [AgentConfig(name="a", prompt="p", description="Does a thing")], + ctx.settings, + ) + root = mgr._create_agents() + + assert root.description == "Does a thing" + + +def test_coordinator_with_sub_agents_builds(): + """A coordinator + leaf pair (README example) constructs end to end.""" + with MockContext(google_api_key="test-key") as ctx: + mgr = _manager( + [ + AgentConfig(name="coordinator", prompt="Route.", sub_agents=["worker"]), + AgentConfig(name="worker", prompt="Work."), + ], + ctx.settings, + ) + root = mgr._create_agents() + + assert root.name == "coordinator" + assert [a.name for a in root.sub_agents] == ["worker"] + assert root.description == "" diff --git a/tests/workflow/test_agent_graph.py b/tests/workflow/test_agent_graph.py new file mode 100644 index 0000000..33829e3 --- /dev/null +++ b/tests/workflow/test_agent_graph.py @@ -0,0 +1,290 @@ +"""Agent-graph validation and settings-scoped construction. + +Before construction the graph was only implicitly checked: a missing sub-agent +was logged and dropped, duplicates/cycles/self-references were unchecked, and +the two-pass build (leaves, then coordinators) depended on declaration order — +a coordinator whose child was itself a coordinator lost that child. Callable +prompts were also evaluated outside the manager's settings context, so they +resolved against the global singleton. +""" + +from __future__ import annotations + +import pytest + +from agentic_cli.workflow.config import ( + AgentConfig, + AgentGraphError, + validate_agent_graph, +) + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from tests.conftest import MockContext # noqa: E402 + + +class TestGraphValidation: + def test_empty_graph_is_rejected(self): + with pytest.raises(AgentGraphError, match="No agents configured"): + validate_agent_graph([]) + + def test_duplicate_names_are_rejected(self): + configs = [ + AgentConfig(name="dup", prompt="a"), + AgentConfig(name="dup", prompt="b"), + ] + with pytest.raises(AgentGraphError, match="Duplicate agent name.*dup"): + validate_agent_graph(configs) + + def test_missing_sub_agent_is_rejected_not_dropped(self): + configs = [AgentConfig(name="coord", prompt="p", sub_agents=["ghost"])] + with pytest.raises(AgentGraphError) as exc: + validate_agent_graph(configs) + assert "coord -> ghost" in str(exc.value) + + def test_self_reference_is_rejected(self): + configs = [AgentConfig(name="loop", prompt="p", sub_agents=["loop"])] + with pytest.raises(AgentGraphError, match="themselves as sub_agents.*loop"): + validate_agent_graph(configs) + + def test_cycle_is_rejected_and_named(self): + configs = [ + AgentConfig(name="a", prompt="p", sub_agents=["b"]), + AgentConfig(name="b", prompt="p", sub_agents=["c"]), + AgentConfig(name="c", prompt="p", sub_agents=["a"]), + ] + with pytest.raises(AgentGraphError, match="Delegation cycle") as exc: + validate_agent_graph(configs) + assert "a" in str(exc.value) and "c" in str(exc.value) + + def test_shared_child_is_rejected_as_unsupported_topology(self): + configs = [ + AgentConfig(name="p1", prompt="p", sub_agents=["shared"]), + AgentConfig(name="p2", prompt="p", sub_agents=["shared"]), + AgentConfig(name="shared", prompt="p"), + ] + with pytest.raises(AgentGraphError, match="more than one parent"): + validate_agent_graph(configs) + + def test_build_order_puts_children_first(self): + configs = [ + AgentConfig(name="top", prompt="p", sub_agents=["mid"]), + AgentConfig(name="mid", prompt="p", sub_agents=["leaf"]), + AgentConfig(name="leaf", prompt="p"), + ] + graph = validate_agent_graph(configs) + order = list(graph.build_order) + assert order.index("leaf") < order.index("mid") < order.index("top") + assert graph.root_name == "top" + + +class TestConstructionOrderIndependence: + """Declaration order must not change the built hierarchy.""" + + def _nested_configs(self, order: str) -> list[AgentConfig]: + top = AgentConfig(name="top", prompt="p", sub_agents=["mid"]) + mid = AgentConfig(name="mid", prompt="p", sub_agents=["leaf"]) + leaf = AgentConfig(name="leaf", prompt="p") + return {"top-first": [top, mid, leaf], "leaf-first": [leaf, mid, top]}[order] + + @pytest.mark.parametrize("order", ["top-first", "leaf-first"]) + def test_two_level_hierarchy_is_built_completely(self, order: str): + with MockContext(google_api_key="test-key") as ctx: + mgr = GoogleADKWorkflowManager( + agent_configs=self._nested_configs(order), settings=ctx.settings + ) + root = mgr._create_agents() + + assert root.name == "top" + assert [a.name for a in root.sub_agents] == ["mid"] + # The nested coordinator must keep its own child. + assert [a.name for a in root.sub_agents[0].sub_agents] == ["leaf"] + + +class TestPromptSettingsScope: + """Prompt factories resolve against the manager's settings, not the global.""" + + def test_zero_arg_factory_sees_manager_settings(self): + seen: list[str] = [] + + def _prompt() -> str: + from agentic_cli.config import get_settings + + seen.append(get_settings().app_name) + return "instruction" + + with MockContext(google_api_key="test-key", app_name="manager-app") as ctx: + from agentic_cli.config import BaseSettings, set_settings + + # Global singleton points somewhere else entirely. + set_settings(BaseSettings(app_name="global-app")) + mgr = GoogleADKWorkflowManager( + agent_configs=[AgentConfig(name="a", prompt=_prompt)], + settings=ctx.settings, + ) + mgr._create_agents() + + assert seen == ["manager-app"] + + def test_factory_taking_settings_is_passed_them(self): + received: list[object] = [] + + def _prompt(settings) -> str: + received.append(settings) + return f"app={settings.app_name}" + + with MockContext(google_api_key="test-key", app_name="manager-app") as ctx: + mgr = GoogleADKWorkflowManager( + agent_configs=[AgentConfig(name="a", prompt=_prompt)], + settings=ctx.settings, + ) + root = mgr._create_agents() + + assert received == [ctx.settings] + assert root.instruction == "app=manager-app" + + def test_plain_string_prompt_is_unchanged(self): + config = AgentConfig(name="a", prompt="literal") + assert config.get_prompt() == "literal" + assert config.get_prompt(settings=object()) == "literal" + + def test_zero_arg_factory_without_settings_still_works(self): + config = AgentConfig(name="a", prompt=lambda: "made up") + assert config.get_prompt() == "made up" + + +class TestPromptFactorySignatures: + """Only the two documented shapes are accepted; the rest fail by name.""" + + def _config(self, prompt): + return AgentConfig(name="scribe", prompt=prompt) + + def test_zero_argument_factory(self): + assert self._config(lambda: "made up").get_prompt() == "made up" + + def test_factory_with_only_defaulted_parameters_is_called_with_none(self): + """``lambda prefix="default": ...`` binds with zero args — keep doing that.""" + config = self._config(lambda prefix="default": f"{prefix}!") + assert config.get_prompt() == "default!" + # Even when settings are available, the zero-arg form wins. + assert config.get_prompt(settings=object()) == "default!" + + def test_single_required_parameter_receives_settings(self): + sentinel = object() + received: list[object] = [] + + def _prompt(settings): + received.append(settings) + return "ok" + + assert self._config(_prompt).get_prompt(settings=sentinel) == "ok" + assert received == [sentinel] + + def test_single_required_parameter_without_settings_is_an_error(self): + with pytest.raises(AgentGraphError, match="scribe"): + self._config(lambda settings: "x").get_prompt() + + def test_two_required_parameters_are_rejected(self): + with pytest.raises(AgentGraphError, match="unsupported prompt factory"): + self._config(lambda settings, extra: "x").get_prompt(settings=object()) + + def test_var_args_factory_is_called_with_no_arguments(self): + def _prompt(*args): + return f"args={len(args)}" + + assert self._config(_prompt).get_prompt(settings=object()) == "args=0" + + def test_non_string_result_is_rejected(self): + with pytest.raises(AgentGraphError, match="not a string"): + self._config(lambda: 42).get_prompt() + + def test_async_factory_is_rejected(self): + async def _prompt(): + return "nope" + + with pytest.raises(AgentGraphError, match="async prompt factories"): + self._config(_prompt).get_prompt() + + def test_error_names_the_agent(self): + config = AgentConfig(name="researcher", prompt=lambda a, b: "x") + with pytest.raises(AgentGraphError) as exc: + config.get_prompt(settings=object()) + assert "researcher" in str(exc.value) + + def test_plain_string_prompt_is_never_called(self): + assert AgentConfig(name="a", prompt="literal").get_prompt() == "literal" + + +class TestSingleRoot: + """A forest is rejected: only one root ever runs.""" + + def test_two_roots_are_rejected_and_both_named(self): + configs = [ + AgentConfig(name="alpha", prompt="p", sub_agents=["helper"]), + AgentConfig(name="helper", prompt="p"), + AgentConfig(name="beta", prompt="p"), # never reachable + ] + with pytest.raises(AgentGraphError, match="2 roots") as exc: + validate_agent_graph(configs) + message = str(exc.value) + assert "alpha" in message and "beta" in message + + def test_two_standalone_agents_are_rejected(self): + configs = [ + AgentConfig(name="one", prompt="p"), + AgentConfig(name="two", prompt="p"), + ] + with pytest.raises(AgentGraphError, match="roots"): + validate_agent_graph(configs) + + def test_single_agent_is_the_root(self): + graph = validate_agent_graph([AgentConfig(name="solo", prompt="p")]) + assert graph.root_name == "solo" + + def test_single_tree_is_accepted(self): + configs = [ + AgentConfig(name="coord", prompt="p", sub_agents=["a", "b"]), + AgentConfig(name="a", prompt="p"), + AgentConfig(name="b", prompt="p"), + ] + assert validate_agent_graph(configs).root_name == "coord" + + def test_manager_rejects_a_forest_before_it_builds_anything(self): + with MockContext(google_api_key="test-key") as ctx: + mgr = GoogleADKWorkflowManager( + agent_configs=[ + AgentConfig(name="one", prompt="p"), + AgentConfig(name="two", prompt="p"), + ], + settings=ctx.settings, + ) + with pytest.raises(AgentGraphError, match="roots"): + mgr._create_agents() + + +class TestValidationHappensBeforeAllocation: + """A static graph error must not cost discovery or service construction.""" + + async def test_bad_graph_fails_before_discovery_and_services(self): + from unittest.mock import AsyncMock, MagicMock + + with MockContext(google_api_key="test-key") as ctx: + mgr = GoogleADKWorkflowManager( + agent_configs=[ + AgentConfig(name="coord", prompt="p", sub_agents=["ghost"]), + ], + settings=ctx.settings, + ) + refresh = AsyncMock() + mgr._model_registry = MagicMock(refresh=refresh) + created: list[str] = [] + mgr._ensure_managers_initialized = lambda: created.append("services") + mgr._make_session_service = lambda: created.append("session") or object() + + with pytest.raises(AgentGraphError, match="ghost"): + await mgr.initialize_services() + + refresh.assert_not_awaited() + assert created == [], "resources were allocated before validation" + assert mgr.is_initialized is False From 03d2ec60aa1ccf87818a8500b926fd9ad6bb4233 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:41 -0400 Subject: [PATCH 04/11] fix(config): validate every runtime-effective model; per-provider discovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A manager's own model — ``Manager(model=...)``, ``reinitialize(model=...)``, or a ``settings.get_model()`` cached before discovery ran — was never in settings and so never validated: an unusable id reached the provider at first use, and a deprecated alias kept being sent even though ``check_model()`` already knew its replacement (only ``set_model()`` wrote one back). Every effective model now goes through one all-or-nothing pass via the internal ``_validate_settings_with_models()``, and the resolved ids are applied only after all of them validate. ``validate_settings()`` itself is unchanged and still returns None. Validation also runs *after* discovery: the static fallback list would otherwise reject a model that exists but postdates it. Discovery authority is tracked per provider, so a Google outage no longer makes the Claude listing non-authoritative — nor an Anthropic outage make Anthropic's hardcoded fallbacks authoritative. An unknown model is never silently swapped for a near neighbour; it is an error naming what was asked for. Also here, since they are the same listing path: the provider SDK clients are closed after a listing rather than left to GC with their connection pools open, and the blocking Anthropic listing runs on a worker thread so startup does not block the event loop. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/config.py | 127 ++++- src/agentic_cli/workflow/base_manager.py | 44 +- src/agentic_cli/workflow/models.py | 173 ++++++- src/agentic_cli/workflow/settings.py | 86 ++- tests/test_model_registry.py | 20 +- tests/workflow/test_model_validation.py | 634 +++++++++++++++++++++++ 6 files changed, 1027 insertions(+), 57 deletions(-) create mode 100644 tests/workflow/test_model_validation.py diff --git a/src/agentic_cli/config.py b/src/agentic_cli/config.py index 1d5da9f..50cd7fc 100644 --- a/src/agentic_cli/config.py +++ b/src/agentic_cli/config.py @@ -31,7 +31,7 @@ from contextvars import ContextVar, Token from pathlib import Path -from typing import Generator, Any, Tuple, Type +from typing import Callable, Generator, Any, Sequence, Tuple, Type from contextlib import contextmanager from pydantic_settings import ( @@ -405,20 +405,110 @@ class SettingsValidationError(Exception): pass -def validate_settings(settings: BaseSettings) -> None: +def _effective_models( + settings: BaseSettings, agent_configs: Any | None +) -> list[tuple[str, str, Callable[[str], None]]]: + """Every model that will actually be used, with a way to rewrite it. + + The configured ``default_model`` plus each agent's ``model`` override — an + override pointing at a provider with no credential fails just as hard as a + bad default, only later and less legibly. + + Each entry carries an ``apply`` callback that writes a resolved id back to + where the model came from. A deprecated alias resolves to its live + replacement, and that replacement has to reach the runtime: a config value + loaded from settings.json or the environment never passes through + ``set_model()``, so validating it and discarding the result left the dead + id to be sent to the provider. + """ + entries: list[tuple[str, str, Callable[[str], None]]] = [] + + if settings.default_model: + + def _apply_default(resolved: str) -> None: + object.__setattr__(settings, "default_model", resolved) + + entries.append(("default_model", settings.default_model, _apply_default)) + + for config in agent_configs or []: + model = getattr(config, "model", None) + if not model: + continue + + def _apply_override(resolved: str, cfg: Any = config) -> None: + cfg.model = resolved + + entries.append( + (f"agent '{getattr(config, 'name', '?')}'", model, _apply_override) + ) + return entries + + +def validate_settings( + settings: BaseSettings, agent_configs: Any | None = None +) -> None: """Validate settings for runtime use. Performs validation that can only be done at runtime: - API key availability - - Model compatibility + - Model availability and provider credentials, for the default model *and* + every per-agent model override - Path accessibility + Every effective model goes through ``settings.check_model()`` — the same + rules ``set_model()`` applies, so the setter and startup validation can + never disagree. Model availability is judged per provider: a model is + rejected as unknown only when *its own* provider answered the listing. + + Not purely a check: a **deprecated alias is rewritten in place** to the live + model it resolves to, on ``settings.default_model`` and on each + ``AgentConfig.model``. That is the only point at which a value loaded from + settings.json or the environment can be corrected, and the runtime reads + those attributes directly. + + Rewrites are **all-or-nothing**: nothing is written until every model has + validated. Applying them as each model was checked left the configuration + half-rewritten by a call that raised, so a retry validated something the + user never wrote. + Args: settings: Settings to validate + agent_configs: Optional agent configs whose ``model`` overrides are + validated — and, when deprecated, upgraded — alongside + ``default_model``. Raises: SettingsValidationError: If validation fails """ + _validate_settings_with_models(settings, agent_configs) + + +def _validate_settings_with_models( + settings: BaseSettings, + agent_configs: Any | None = None, + extra_models: "Sequence[tuple[str, str]] | None" = None, +) -> dict[str, str]: + """:func:`validate_settings` plus models that do not live in settings. + + Internal. A workflow manager's own model — ``Manager(model=...)``, + ``reinitialize(model=...)``, or one cached before discovery ran — has to be + validated in the *same* all-or-nothing pass as ``default_model`` and the + agent overrides, but it is the manager's state to rewrite, not settings'. + So it is passed in here and its resolution handed back, keeping + ``validate_settings()``'s public contract (returns None) intact. + + Args: + settings: Settings to validate. + agent_configs: Optional agent configs, as for ``validate_settings``. + extra_models: ``(label, model)`` pairs to validate alongside them. + + Returns: + ``{label: resolved_model}`` for every entry in ``extra_models`` (the + input model when it needed no rewrite). + + Raises: + SettingsValidationError: If validation fails. + """ errors = [] if not settings.has_any_api_key: @@ -426,13 +516,30 @@ def validate_settings(settings: BaseSettings) -> None: "No API keys configured. Set GOOGLE_API_KEY or ANTHROPIC_API_KEY." ) - if settings.default_model: - available = settings.get_available_models() - if settings.default_model not in available: - errors.append( - f"Configured model '{settings.default_model}' is not available. " - f"Available models: {', '.join(available) if available else 'none'}" - ) + entries = list(_effective_models(settings, agent_configs)) + resolutions: dict[str, str] = {} + + def _record(label: str) -> Callable[[str], None]: + def _apply(resolved: str) -> None: + resolutions[label] = resolved + + return _apply + + for label, model in extra_models or (): + entries.append((label, model, _record(label))) + + pending: list[tuple[Callable[[str], None], str]] = [] + for label, model, apply_resolved in entries: + try: + resolved = settings.check_model(model, label=label) + except ValueError as exc: + errors.append(str(exc)) + continue + pending.append((apply_resolved, resolved)) if errors: raise SettingsValidationError("\n".join(errors)) + + for apply_resolved, resolved in pending: + apply_resolved(resolved) + return resolutions diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index 4ca5007..d5f029a 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -624,20 +624,20 @@ async def initialize_services(self, validate: bool = True) -> None: # any network call is made: a bad graph is a static configuration # error and should not cost a model listing or an embedding model. self._validate_agent_graph() - from agentic_cli.config import validate_settings - - if validate: - validate_settings(self._settings) self._settings.export_api_keys_to_env() - # Refresh model registry from APIs + # Discover models BEFORE validating them: the static fallback list + # would otherwise reject a model that exists but predates the list. await self._model_registry.refresh( google_api_key=self._settings.google_api_key, anthropic_api_key=self._settings.anthropic_api_key, ) self._settings.set_model_registry(self._model_registry) + if validate: + self._validate_models() + # Create services BEFORE backend init so _build_tools() can # produce factory-bound tools during agent/graph creation. # Offloaded to a worker thread because constructors here may @@ -650,6 +650,40 @@ async def initialize_services(self, validate: bool = True) -> None: await self._do_initialize() self._initialized = True + # Label for this manager's own model in validation errors. + _MODEL_OVERRIDE_LABEL = "workflow model" + + def _validate_models(self) -> None: + """Validate every model this manager will actually use, and normalize it. + + ``settings.default_model`` and the per-agent overrides are validated by + ``validate_settings``. This manager's *own* model is not in settings at + all — it comes from ``Manager(model=...)``, ``reinitialize(model=...)``, + or a ``settings.get_model()`` cached before discovery ran — so it is + passed into the same all-or-nothing pass rather than checked separately: + an unusable id must fail startup, and a deprecated one must be replaced + by the id the runtime then sends. + + Raises: + SettingsValidationError: If any effective model is unusable. + """ + from agentic_cli.config import _validate_settings_with_models + + extras: list[tuple[str, str]] = [] + if self._model_resolved and self._model: + extras.append((self._MODEL_OVERRIDE_LABEL, self._model)) + + resolved = _validate_settings_with_models( + self._settings, self._agent_configs, extras + ) + + replacement = resolved.get(self._MODEL_OVERRIDE_LABEL) + if replacement is not None and replacement != self._model: + logger.info( + "model_override_resolved", requested=self._model, model=replacement + ) + self._model = replacement + def _validate_agent_graph(self) -> None: """Validate the declared agent graph. Backends may narrow this. diff --git a/src/agentic_cli/workflow/models.py b/src/agentic_cli/workflow/models.py index 6f89d1b..5abec18 100644 --- a/src/agentic_cli/workflow/models.py +++ b/src/agentic_cli/workflow/models.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio import logging import re from dataclasses import dataclass, field @@ -16,6 +17,21 @@ logger = logging.getLogger(__name__) +def _close_quietly(client: Any) -> None: + """Release a provider SDK client after a listing, if it can be closed. + + Both provider clients hold an HTTP connection pool and expose ``close()``; + a listing is a one-shot call, so the pool would otherwise linger until GC. + """ + close = getattr(client, "close", None) + if close is None: + return + try: + close() + except Exception as exc: # noqa: BLE001 - closing must not fail a listing + logger.debug("Failed to close provider client: %s", exc) + + class ModelFamily(str, Enum): """Model provider families.""" @@ -24,6 +40,26 @@ class ModelFamily(str, Enum): GPT = "gpt" +class DiscoveryState(str, Enum): + """How much the registry actually knows about one provider's models. + + Tracked per family, because providers fail independently: an Anthropic + outage must not make the Gemini list non-authoritative, and it must not + make the *Anthropic* fallback list authoritative either. + + - ``UNATTEMPTED`` — no key, or no refresh yet. Nothing is known; an + unrecognised model of that family cannot be rejected. + - ``SUCCEEDED`` — the provider answered with models. The list is + authoritative: an unknown model of that family is an error. + - ``DEGRADED`` — the listing failed or came back empty, so the hardcoded + fallbacks stand in. Known-incomplete; never authoritative. + """ + + UNATTEMPTED = "unattempted" + SUCCEEDED = "succeeded" + DEGRADED = "degraded" + + @dataclass class ModelInfo: """Metadata for a single model.""" @@ -75,12 +111,58 @@ def __init__(self) -> None: self._models: dict[str, ModelInfo] = {} self._defaults: dict[ModelFamily, str] = dict(self.FALLBACK_DEFAULTS) self._refreshed = False + # Per-family outcome of the last refresh(); empty until one runs. + self._discovery: dict[ModelFamily, DiscoveryState] = {} + # Families whose listing failed during the in-flight refresh(). + self._degraded_families: set[ModelFamily] = set() @property def is_refreshed(self) -> bool: """Whether the registry has been populated from APIs.""" return self._refreshed + def authority_for(self, family: ModelFamily) -> DiscoveryState: + """Discovery state for one provider family. + + Falls back to the global ``is_refreshed`` flag only when no refresh has + recorded per-family outcomes (a registry whose internals were seeded + directly), so callers always get a definite answer. + """ + recorded = self._discovery.get(family) + if recorded is not None: + return recorded + return DiscoveryState.SUCCEEDED if self._refreshed else DiscoveryState.UNATTEMPTED + + def is_authoritative_for(self, model_id: str) -> bool: + """Whether this registry's list can reject ``model_id`` as unknown. + + True only when *that model's* provider answered the listing — a + different provider's outage is irrelevant. + """ + try: + family = self.get_family(model_id) + except ValueError: + return False + return self.authority_for(family) is DiscoveryState.SUCCEEDED + + @property + def discovery_complete(self) -> bool: + """Whether every attempted provider answered its listing. + + Coarse, kept for callers that want one flag; prefer + :meth:`is_authoritative_for` when judging a specific model. + """ + if not self._refreshed: + return False + attempted = [ + state + for state in self._discovery.values() + if state is not DiscoveryState.UNATTEMPTED + ] + if not attempted: + return True # internals seeded directly; nothing contradicts it + return all(state is DiscoveryState.SUCCEEDED for state in attempted) + # ------------------------------------------------------------------ # Public sync API # ------------------------------------------------------------------ @@ -203,24 +285,33 @@ def supports_thinking(self, model_id: str) -> bool: def resolve_model(self, model_id: str) -> str: """Validate a model ID against the registry. - If the model exists, returns it as-is. If deprecated or missing, - attempts to find the closest match in the same family and tier. + A **deprecated** model is transparently upgraded to the closest live + model in its family and tier, with a warning — that is an alias, and + the user's intent is unambiguous. + + A model that is simply **unknown** is never silently swapped for + something else: if this registry is authoritative for that family + (its provider answered the listing) the id is rejected; otherwise it is + accepted as-is, because a degraded or unattempted listing cannot prove + the model does not exist. Args: model_id: Model identifier to resolve. Returns: - Resolved model ID (may differ from input if deprecated). + The resolved model ID — the input, or the replacement for a + deprecated alias. Raises: - ValueError: If no suitable model can be found. + ValueError: If the model's provider cannot be determined, or the + provider's authoritative listing does not contain it. """ # Exact match if model_id in self._models: info = self._models[model_id] if not info.deprecated: return model_id - # Deprecated — find replacement + # Deprecated alias — upgrade, loudly. replacement = self._find_closest_match(model_id, info.family) if replacement: logger.warning( @@ -229,12 +320,10 @@ def resolve_model(self, model_id: str) -> str: replacement, ) return replacement + raise ValueError( + f"Model '{model_id}' is deprecated and no replacement is available." + ) - # Not in registry — if not refreshed, accept anything - if not self._refreshed: - return model_id - - # Try to find closest match try: family = self.get_family(model_id) except ValueError: @@ -242,14 +331,9 @@ def resolve_model(self, model_id: str) -> str: f"Model '{model_id}' is not available and its family cannot be determined." ) - replacement = self._find_closest_match(model_id, family) - if replacement: - logger.warning( - "Model '%s' is not available, using '%s' instead", - model_id, - replacement, - ) - return replacement + if not self.is_authoritative_for(model_id): + # Discovery was never attempted, or this provider's listing failed. + return model_id available = self.get_available_models(family) raise ValueError( @@ -327,14 +411,21 @@ async def refresh( anthropic_api_key: Anthropic API key for listing Claude models. """ models: dict[str, ModelInfo] = {} + self._degraded_families = set() + self._discovery = { + ModelFamily.GEMINI: DiscoveryState.UNATTEMPTED, + ModelFamily.CLAUDE: DiscoveryState.UNATTEMPTED, + } if google_api_key: google_models = await self._fetch_google_models(google_api_key) + self._record_discovery(ModelFamily.GEMINI, google_models) for m in google_models: models[m.id] = m if anthropic_api_key: anthropic_models = await self._fetch_anthropic_models(anthropic_api_key) + self._record_discovery(ModelFamily.CLAUDE, anthropic_models) for m in anthropic_models: models[m.id] = m @@ -355,6 +446,20 @@ async def refresh( "No models fetched from APIs, using fallback lists" ) + def _record_discovery( + self, family: ModelFamily, fetched: list[ModelInfo] + ) -> None: + """Record whether a provider's listing can be trusted. + + A failed listing (the fetcher substituted fallbacks) and an empty one + are treated the same: nothing was learned, so the family is DEGRADED. + """ + if family in self._degraded_families or not fetched: + self._discovery[family] = DiscoveryState.DEGRADED + logger.warning("Model discovery degraded for %s", family.value) + else: + self._discovery[family] = DiscoveryState.SUCCEEDED + # Patterns to exclude from Google model listings _GOOGLE_EXCLUDE_PATTERNS = re.compile( r"-\d{3}$" # point releases: -001, -002 @@ -370,7 +475,16 @@ async def _fetch_google_models(self, api_key: str) -> list[ModelInfo]: Filters out non-text models (image/video/audio/embedding), open-source models (Gemma), specialized previews, and duplicate aliases. + + The listing itself is a blocking SDK call, so it runs on a worker + thread — this coroutine is awaited during startup on the CLI's event + loop, which must stay responsive. """ + return await asyncio.to_thread(self._fetch_google_models_sync, api_key) + + def _fetch_google_models_sync(self, api_key: str) -> list[ModelInfo]: + """Blocking Google model listing (see :meth:`_fetch_google_models`).""" + client = None try: from google import genai @@ -412,14 +526,27 @@ async def _fetch_google_models(self, api_key: str) -> list[ModelInfo]: except Exception as exc: logger.warning("Failed to fetch Google models: %s", exc) - # Return fallback models + # Fallbacks are incomplete; mark this family degraded so callers + # don't treat the list as authoritative. + self._degraded_families.add(ModelFamily.GEMINI) return [ ModelInfo(id=mid, family=ModelFamily.GEMINI, supports_thinking="2.5" in mid or "3" in mid) for mid in self.FALLBACK_GOOGLE ] + finally: + _close_quietly(client) async def _fetch_anthropic_models(self, api_key: str) -> list[ModelInfo]: - """Fetch models from Anthropic API.""" + """Fetch models from the Anthropic API. + + Runs the blocking SDK listing on a worker thread so the event loop + stays responsive during startup. + """ + return await asyncio.to_thread(self._fetch_anthropic_models_sync, api_key) + + def _fetch_anthropic_models_sync(self, api_key: str) -> list[ModelInfo]: + """Blocking Anthropic model listing (see :meth:`_fetch_anthropic_models`).""" + client = None try: import anthropic @@ -453,11 +580,15 @@ async def _fetch_anthropic_models(self, api_key: str) -> list[ModelInfo]: except Exception as exc: logger.warning("Failed to fetch Anthropic models: %s", exc) - # Return fallback models + # Fallbacks are incomplete; mark this family degraded so callers + # don't treat the list as authoritative. + self._degraded_families.add(ModelFamily.CLAUDE) return [ ModelInfo(id=mid, family=ModelFamily.CLAUDE, supports_thinking=True) for mid in self.FALLBACK_ANTHROPIC ] + finally: + _close_quietly(client) @staticmethod def _normalize_anthropic_id(model_id: str) -> str: diff --git a/src/agentic_cli/workflow/settings.py b/src/agentic_cli/workflow/settings.py index 852970d..9811b17 100644 --- a/src/agentic_cli/workflow/settings.py +++ b/src/agentic_cli/workflow/settings.py @@ -12,14 +12,24 @@ from pydantic import BaseModel, Field, field_validator +from agentic_cli.logging import Loggers from agentic_cli.workflow.models import ModelFamily, ModelRegistry if TYPE_CHECKING: pass +logger = Loggers.config() + # Thinking effort levels (module-level constant for backward compatibility) THINKING_EFFORT_LEVELS = ModelRegistry.THINKING_EFFORT_LEVELS +# Which environment variable supplies each provider's credential (for error +# messages — the value is never echoed). +_PROVIDER_ENV_VAR = { + ModelFamily.GEMINI: "GOOGLE_API_KEY", + ModelFamily.CLAUDE: "ANTHROPIC_API_KEY", +} + class PermissionRuleConfig(BaseModel): """A single permission rule — serialised to settings.json as JSON.""" @@ -625,22 +635,68 @@ def supports_thinking_effort(self, model: str | None = None) -> bool: registry = self._get_registry() return registry.supports_thinking(model) - def set_model(self, model: str) -> None: - """Set the default model.""" + def check_model(self, model: str, *, label: str = "model") -> str: + """Validate one model against credentials and discovery authority. + + The single rule set shared by ``set_model()`` and ``validate_settings()`` + so a model the setter accepts can never be rejected at startup (or the + reverse): + + 1. the provider must be derivable from the id; + 2. that provider's credential must be configured; + 3. a deprecated alias resolves to its replacement (warned); + 4. an unknown model is rejected only when that provider's listing is + authoritative — a degraded/unattempted listing cannot disprove it, + and is logged instead. + + Args: + model: Model identifier to check. + label: What is being checked, for the error message. + + Returns: + The resolved model id (differs only for a deprecated alias). + + Raises: + ValueError: With an actionable message; never includes a credential. + """ registry = self._get_registry() - if registry.is_refreshed: - # Validate and possibly resolve deprecated models - resolved = registry.resolve_model(model) - object.__setattr__(self, "default_model", resolved) - else: - # Pre-refresh: validate against fallback list - available = self.get_available_models() - if model not in available: - raise ValueError( - f"Model '{model}' is not available. " - f"Available models: {', '.join(available)}" - ) - object.__setattr__(self, "default_model", model) + try: + family = registry.get_family(model) + except ValueError: + raise ValueError( + f"Model '{model}' ({label}) is not available: its provider " + "cannot be determined from the model id." + ) from None + + if not self._has_credential_for(family): + env_var = _PROVIDER_ENV_VAR.get(family, "the provider API key") + raise ValueError( + f"Model '{model}' ({label}) is not available: it needs a " + f"{family.value} credential. Set {env_var}." + ) + + resolved = registry.resolve_model(model) # raises when authoritative + if resolved == model and model not in self.get_available_models(): + # Not authoritative (else resolve_model would have raised), so the + # static list simply lags reality. + logger.warning("model_not_in_static_list", model=model, source=label) + return resolved + + def _has_credential_for(self, family: ModelFamily) -> bool: + """Whether the credential a model family needs is configured.""" + if family is ModelFamily.GEMINI: + return self.has_google_key + if family is ModelFamily.CLAUDE: + return self.has_anthropic_key + return False + + def set_model(self, model: str) -> None: + """Set the default model, validating it exactly as startup would. + + Raises: + ValueError: If the model is unusable (see :meth:`check_model`). + """ + object.__setattr__(self, "default_model", self.check_model(model)) def set_thinking_effort(self, effort: str) -> None: """Set the thinking effort level.""" diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 39bdbd6..1ffbdb1 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -142,9 +142,17 @@ def test_set_default(self): class TestResolveModel: """Tests for model resolution.""" - def test_resolve_before_refresh_accepts_any(self): + def test_resolve_before_refresh_accepts_a_known_family(self): + """Without discovery the registry cannot disprove a model id.""" reg = ModelRegistry() - assert reg.resolve_model("any-model") == "any-model" + assert reg.resolve_model("gemini-9.9-experimental") == "gemini-9.9-experimental" + assert reg.resolve_model("claude-opus-9") == "claude-opus-9" + + def test_resolve_rejects_an_id_with_no_determinable_provider(self): + """Well-formedness is not a discovery question — always an error.""" + reg = ModelRegistry() + with pytest.raises(ValueError, match="cannot be determined"): + reg.resolve_model("any-model") def test_resolve_exact_match(self): reg = ModelRegistry() @@ -167,16 +175,16 @@ def test_resolve_deprecated_finds_replacement(self): resolved = reg.resolve_model("gemini-2.0-pro") assert resolved == "gemini-2.5-pro" - def test_resolve_missing_finds_closest(self): + def test_resolve_missing_is_rejected_not_substituted(self): + """An explicitly chosen unknown model must never become another one.""" reg = ModelRegistry() reg._models["gemini-2.5-flash"] = ModelInfo( id="gemini-2.5-flash", family=ModelFamily.GEMINI ) reg._refreshed = True - # Missing pro model → falls back to flash (only available) - resolved = reg.resolve_model("gemini-3-pro-preview") - assert resolved == "gemini-2.5-flash" + with pytest.raises(ValueError, match="not available"): + reg.resolve_model("gemini-3-pro-preview") def test_resolve_missing_unknown_family_raises(self): reg = ModelRegistry() diff --git a/tests/workflow/test_model_validation.py b/tests/workflow/test_model_validation.py new file mode 100644 index 0000000..5e4ad4d --- /dev/null +++ b/tests/workflow/test_model_validation.py @@ -0,0 +1,634 @@ +"""Model/provider validation: ordered, complete, and off the event loop. + +Three defects are covered: + +1. ``validate_settings`` ran *before* the registry refresh, so a model that + exists but predates the static fallback list was rejected at startup. +2. Per-agent ``AgentConfig.model`` overrides were never checked, so an + override pointing at a provider with no credential failed mid-run. +3. ``ModelRegistry._fetch_*_models`` are async but called blocking provider + SDKs, stalling the CLI's event loop during startup. +""" + +from __future__ import annotations + +import asyncio +import threading +from unittest.mock import MagicMock + +import pytest + +from agentic_cli.config import SettingsValidationError, validate_settings +from agentic_cli.workflow.base_manager import BaseWorkflowManager +from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.models import ModelFamily, ModelInfo, ModelRegistry +from tests.conftest import MockContext + +# A well-formed Gemini id that is deliberately absent from FALLBACK_GOOGLE. +DISCOVERED = "gemini-9.9-flash-preview" + + +class _TestManager(BaseWorkflowManager): + """Minimal concrete manager; backend init is a no-op.""" + + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + +def _discovering_registry(*models: ModelInfo) -> ModelRegistry: + """A registry whose provider fetches return fixed models.""" + registry = ModelRegistry() + + async def _google(api_key): + return [m for m in models if m.family is ModelFamily.GEMINI] + + async def _anthropic(api_key): + return [m for m in models if m.family is ModelFamily.CLAUDE] + + registry._fetch_google_models = _google + registry._fetch_anthropic_models = _anthropic + return registry + + +class TestDiscoveryBeforeValidation: + async def test_dynamically_discovered_model_is_accepted(self): + """A model absent from the static list but present in the API listing.""" + with MockContext(google_api_key="k", default_model=DISCOVERED) as ctx: + assert DISCOVERED not in ModelRegistry.FALLBACK_GOOGLE + mgr = _TestManager(agent_configs=[], settings=ctx.settings) + mgr._model_registry = _discovering_registry( + ModelInfo(id=DISCOVERED, family=ModelFamily.GEMINI) + ) + + await mgr.initialize_services() + + assert mgr.is_initialized + + async def test_refreshed_registry_still_rejects_unknown_model(self): + """Once discovery succeeded, its list is authoritative.""" + with MockContext(google_api_key="k", default_model="gemini-not-real") as ctx: + mgr = _TestManager(agent_configs=[], settings=ctx.settings) + mgr._model_registry = _discovering_registry( + ModelInfo(id=DISCOVERED, family=ModelFamily.GEMINI) + ) + + with pytest.raises(SettingsValidationError, match="not available"): + await mgr.initialize_services() + + +class TestPerAgentOverrideValidation: + def test_override_without_provider_credential_fails(self): + """A Claude override with only a Google key must fail loudly.""" + with MockContext(google_api_key="k") as ctx: + configs = [ + AgentConfig(name="root", prompt="p"), + AgentConfig(name="claude_helper", prompt="p", model="claude-sonnet-4-6"), + ] + with pytest.raises(SettingsValidationError) as exc: + validate_settings(ctx.settings, agent_configs=configs) + + message = str(exc.value) + assert "claude_helper" in message + assert "ANTHROPIC_API_KEY" in message + # Errors must name the missing credential, never its value. + assert "k" != message and "api_key=" not in message + + def test_override_with_credential_passes(self): + with MockContext(google_api_key="k", anthropic_api_key="a") as ctx: + configs = [ + AgentConfig(name="root", prompt="p"), + AgentConfig(name="helper", prompt="p", model="claude-sonnet-4-6"), + ] + validate_settings(ctx.settings, agent_configs=configs) + + def test_unknown_provider_in_override_fails(self): + with MockContext(google_api_key="k") as ctx: + configs = [AgentConfig(name="odd", prompt="p", model="mystery-model-1")] + with pytest.raises(SettingsValidationError, match="provider cannot be determined"): + validate_settings(ctx.settings, agent_configs=configs) + + +class TestDiscoveryUnavailable: + """Offline behaviour stays deterministic: no discovery, no false rejections.""" + + def test_unrecognised_but_well_formed_model_is_not_rejected(self): + with MockContext(google_api_key="k", default_model=DISCOVERED) as ctx: + assert ctx.settings._get_registry().is_refreshed is False + validate_settings(ctx.settings) # warns, does not raise + + def test_missing_credential_still_fails_offline(self): + with MockContext(google_api_key="k", default_model="claude-sonnet-4-6") as ctx: + with pytest.raises(SettingsValidationError, match="ANTHROPIC_API_KEY"): + validate_settings(ctx.settings) + + async def test_failed_discovery_does_not_make_fallbacks_authoritative( + self, monkeypatch + ): + """A provider outage must not turn the stale fallback list into truth.""" + def _boom(api_key=None): + raise RuntimeError("provider down") + + monkeypatch.setattr("google.genai.Client", _boom) + + with MockContext(google_api_key="k", default_model=DISCOVERED) as ctx: + registry = ModelRegistry() + mgr = _TestManager(agent_configs=[], settings=ctx.settings) + mgr._model_registry = registry + + await mgr.initialize_services() + + assert mgr.is_initialized + # Fallbacks were substituted, so the list is not authoritative. + assert registry.is_refreshed is True + assert registry.discovery_complete is False + + +class TestBlockingFetchOffEventLoop: + """Provider listings are blocking SDK calls; they must not run on the loop.""" + + async def test_google_listing_runs_on_worker_thread(self, monkeypatch): + calling_thread: list[threading.Thread] = [] + + class _FakeModels: + def list(self): + calling_thread.append(threading.current_thread()) + return [] + + class _FakeClient: + def __init__(self, api_key=None): + self.models = _FakeModels() + + monkeypatch.setattr("google.genai.Client", _FakeClient) + + registry = ModelRegistry() + await registry._fetch_google_models("key") + + assert calling_thread, "the SDK listing was never called" + assert calling_thread[0] is not threading.main_thread() + + async def test_anthropic_listing_runs_on_worker_thread(self, monkeypatch): + calling_thread: list[threading.Thread] = [] + + class _FakeModels: + def list(self, limit=None): + calling_thread.append(threading.current_thread()) + return MagicMock(data=[]) + + class _FakeAnthropic: + def __init__(self, api_key=None): + self.models = _FakeModels() + + monkeypatch.setattr("anthropic.Anthropic", _FakeAnthropic) + + registry = ModelRegistry() + await registry._fetch_anthropic_models("key") + + assert calling_thread, "the SDK listing was never called" + assert calling_thread[0] is not threading.main_thread() + + async def test_event_loop_stays_responsive_during_refresh(self, monkeypatch): + """A slow provider listing must not stall other coroutines.""" + release = threading.Event() + ticks = 0 + + class _FakeModels: + def list(self): + release.wait(timeout=5) + return [] + + class _FakeClient: + def __init__(self, api_key=None): + self.models = _FakeModels() + + monkeypatch.setattr("google.genai.Client", _FakeClient) + + registry = ModelRegistry() + fetch = asyncio.create_task(registry._fetch_google_models("key")) + for _ in range(3): + await asyncio.sleep(0.01) + ticks += 1 + release.set() + await fetch + + assert ticks == 3 # the loop kept running while the SDK call blocked + + +class TestPerProviderDiscoveryAuthority: + """Providers fail independently; authority is tracked per family.""" + + def _registry(self, *, google_ok: bool, anthropic_ok: bool) -> ModelRegistry: + registry = ModelRegistry() + + async def _google(api_key): + if google_ok: + return [ModelInfo(id=DISCOVERED, family=ModelFamily.GEMINI)] + registry._degraded_families.add(ModelFamily.GEMINI) + return [ + ModelInfo(id=m, family=ModelFamily.GEMINI) + for m in ModelRegistry.FALLBACK_GOOGLE + ] + + async def _anthropic(api_key): + if anthropic_ok: + return [ModelInfo(id="claude-real-9", family=ModelFamily.CLAUDE)] + registry._degraded_families.add(ModelFamily.CLAUDE) + return [ + ModelInfo(id=m, family=ModelFamily.CLAUDE) + for m in ModelRegistry.FALLBACK_ANTHROPIC + ] + + registry._fetch_google_models = _google + registry._fetch_anthropic_models = _anthropic + return registry + + async def test_google_success_anthropic_failure_states(self): + from agentic_cli.workflow.models import DiscoveryState + + registry = self._registry(google_ok=True, anthropic_ok=False) + await registry.refresh(google_api_key="g", anthropic_api_key="a") + + assert registry.authority_for(ModelFamily.GEMINI) is DiscoveryState.SUCCEEDED + assert registry.authority_for(ModelFamily.CLAUDE) is DiscoveryState.DEGRADED + assert registry.is_authoritative_for(DISCOVERED) is True + assert registry.is_authoritative_for("claude-sonnet-4-6") is False + + async def test_google_success_rejects_unknown_gemini_despite_anthropic_outage(self): + registry = self._registry(google_ok=True, anthropic_ok=False) + await registry.refresh(google_api_key="g", anthropic_api_key="a") + + with pytest.raises(ValueError, match="not available"): + registry.resolve_model("gemini-does-not-exist") + + async def test_anthropic_outage_does_not_make_its_fallbacks_authoritative(self): + registry = self._registry(google_ok=True, anthropic_ok=False) + await registry.refresh(google_api_key="g", anthropic_api_key="a") + + # Not in the (fallback) Claude list, but the listing failed: accept it. + assert registry.resolve_model("claude-brand-new-1") == "claude-brand-new-1" + + async def test_unattempted_provider_is_never_authoritative(self): + from agentic_cli.workflow.models import DiscoveryState + + registry = self._registry(google_ok=True, anthropic_ok=True) + await registry.refresh(google_api_key="g") # no Anthropic key + + assert registry.authority_for(ModelFamily.CLAUDE) is DiscoveryState.UNATTEMPTED + assert registry.resolve_model("claude-anything-1") == "claude-anything-1" + + async def test_empty_listing_is_treated_as_degraded(self): + from agentic_cli.workflow.models import DiscoveryState + + registry = ModelRegistry() + + async def _empty(api_key): + return [] + + registry._fetch_google_models = _empty + await registry.refresh(google_api_key="g") + + assert registry.authority_for(ModelFamily.GEMINI) is DiscoveryState.DEGRADED + assert registry.resolve_model("gemini-whatever") == "gemini-whatever" + + +class TestSetterValidatorConsistency: + """set_model() and validate_settings() must never disagree.""" + + def _settings_with(self, ctx, registry): + ctx.settings.set_model_registry(registry) + return ctx.settings + + async def test_setter_accepts_what_the_validator_accepts(self): + with MockContext(google_api_key="k") as ctx: + registry = _discovering_registry( + ModelInfo(id=DISCOVERED, family=ModelFamily.GEMINI) + ) + await registry.refresh(google_api_key="k") + settings = self._settings_with(ctx, registry) + + settings.set_model(DISCOVERED) + assert settings.default_model == DISCOVERED + validate_settings(settings) # must not raise + + async def test_setter_rejects_what_the_validator_rejects(self): + with MockContext(google_api_key="k") as ctx: + registry = _discovering_registry( + ModelInfo(id=DISCOVERED, family=ModelFamily.GEMINI) + ) + await registry.refresh(google_api_key="k") + settings = self._settings_with(ctx, registry) + + with pytest.raises(ValueError, match="not available"): + settings.set_model("gemini-nope") + + object.__setattr__(settings, "default_model", "gemini-nope") + with pytest.raises(SettingsValidationError, match="not available"): + validate_settings(settings) + + def test_setter_rejects_a_model_without_its_credential(self): + with MockContext(google_api_key="k") as ctx: + with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): + ctx.settings.set_model("claude-sonnet-4-6") + + def test_setter_accepts_an_unknown_model_when_not_authoritative(self): + """Offline, the static list cannot disprove a well-formed model id.""" + with MockContext(google_api_key="k") as ctx: + ctx.settings.set_model(DISCOVERED) + assert ctx.settings.default_model == DISCOVERED + validate_settings(ctx.settings) + + async def test_deprecated_alias_is_upgraded_by_both(self): + with MockContext(google_api_key="k") as ctx: + registry = ModelRegistry() + registry._models = { + "gemini-old": ModelInfo( + id="gemini-old", family=ModelFamily.GEMINI, deprecated=True + ), + "gemini-2.5-pro": ModelInfo( + id="gemini-2.5-pro", family=ModelFamily.GEMINI + ), + } + registry._refreshed = True + settings = self._settings_with(ctx, registry) + + settings.set_model("gemini-old") + assert settings.default_model == "gemini-2.5-pro" + validate_settings(settings) + + +class TestDeprecatedAliasesApplyAtRuntime: + """A deprecated alias must be *replaced*, not merely warned about. + + ``check_model()`` returns the live replacement, but only ``set_model()`` + was writing it back. A ``default_model`` loaded from settings.json/env, and + every ``AgentConfig.model`` override, kept the dead id and was then sent to + the provider. None of these tests call ``set_model()``. + """ + + @staticmethod + def _deprecating_registry() -> ModelRegistry: + registry = ModelRegistry() + registry._models = { + "gemini-old": ModelInfo( + id="gemini-old", family=ModelFamily.GEMINI, deprecated=True + ), + "gemini-2.5-pro": ModelInfo(id="gemini-2.5-pro", family=ModelFamily.GEMINI), + } + registry._refreshed = True + return registry + + def test_configured_default_is_replaced(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._deprecating_registry()) + # As if loaded from settings.json / GOOGLE_MODEL, not via set_model(). + object.__setattr__(settings, "default_model", "gemini-old") + + validate_settings(settings) + + assert settings.default_model == "gemini-2.5-pro" + assert settings.get_model() == "gemini-2.5-pro" + + def test_agent_override_is_replaced(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._deprecating_registry()) + config = AgentConfig(name="a", prompt="p", model="gemini-old") + + validate_settings(settings, agent_configs=[config]) + + assert config.model == "gemini-2.5-pro" + + def test_live_model_is_left_alone(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._deprecating_registry()) + object.__setattr__(settings, "default_model", "gemini-2.5-pro") + config = AgentConfig(name="a", prompt="p", model="gemini-2.5-pro") + + validate_settings(settings, agent_configs=[config]) + + assert settings.default_model == "gemini-2.5-pro" + assert config.model == "gemini-2.5-pro" + + async def test_manager_runs_on_the_replacement(self): + """End to end: what the backend actually uses after initialization.""" + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + object.__setattr__(settings, "default_model", "gemini-old") + config = AgentConfig(name="a", prompt="p", model="gemini-old") + + manager = _TestManager(agent_configs=[config], settings=settings) + manager._model_registry = self._deprecating_registry() + # refresh() is a no-op for a directly-seeded registry. + manager._model_registry.refresh = lambda **kw: asyncio.sleep(0) + + await manager.initialize_services() + + assert manager.model == "gemini-2.5-pro" + assert config.model == "gemini-2.5-pro" + + +class TestManagerModelIsValidated: + """Every model the *runtime* will actually send must be validated. + + ``validate_settings`` covered ``settings.default_model`` and the per-agent + overrides, but a manager's own model — ``GoogleADKWorkflowManager(model=...)``, + ``reinitialize(model=...)``, or one cached from an earlier + ``settings.get_model()`` — bypassed it entirely: an unusable id reached the + provider, and a deprecated one was never swapped for its replacement. + """ + + @staticmethod + def _deprecating_registry() -> ModelRegistry: + registry = ModelRegistry() + registry._models = { + "gemini-old": ModelInfo( + id="gemini-old", family=ModelFamily.GEMINI, deprecated=True + ), + "gemini-2.5-pro": ModelInfo(id="gemini-2.5-pro", family=ModelFamily.GEMINI), + } + registry._refreshed = True + return registry + + def _manager(self, settings, cls=_TestManager, **kwargs): + manager = cls(agent_configs=[], settings=settings, **kwargs) + manager._model_registry = self._deprecating_registry() + manager._model_registry.refresh = lambda **kw: asyncio.sleep(0) + return manager + + async def test_explicit_constructor_model_is_normalized(self): + with MockContext(google_api_key="k") as ctx: + manager = self._manager(ctx.settings, model="gemini-old") + await manager.initialize_services() + assert manager.model == "gemini-2.5-pro" + + async def test_cached_model_is_normalized(self): + """A model resolved before discovery may since have been deprecated.""" + with MockContext(google_api_key="k") as ctx: + manager = self._manager(ctx.settings) + manager._model = "gemini-old" + manager._model_resolved = True + + await manager.initialize_services() + + assert manager.model == "gemini-2.5-pro" + + + async def test_explicit_model_without_a_credential_fails(self): + with MockContext(google_api_key="k") as ctx: + manager = self._manager(ctx.settings, model="claude-sonnet-4-6") + with pytest.raises(SettingsValidationError, match="ANTHROPIC_API_KEY"): + await manager.initialize_services() + + async def test_unknown_explicit_model_is_rejected(self): + with MockContext(google_api_key="k") as ctx: + manager = self._manager(ctx.settings, model="gemini-nope") + with pytest.raises(SettingsValidationError, match="not available"): + await manager.initialize_services() + + +class TestValidateSettingsReturnContract: + """``validate_settings()`` is a checker: it returns None, always. + + Extra-model resolutions are an internal need of the workflow manager and + must not change what the public function hands back. + """ + + def test_returns_none(self): + with MockContext(google_api_key="k") as ctx: + assert validate_settings(ctx.settings) is None + + def test_returns_none_with_agent_configs(self): + with MockContext(google_api_key="k") as ctx: + config = AgentConfig(name="a", prompt="p", model="gemini-2.5-flash") + assert validate_settings(ctx.settings, agent_configs=[config]) is None + + def test_public_signature_takes_no_extra_models(self): + import inspect + + params = inspect.signature(validate_settings).parameters + assert list(params) == ["settings", "agent_configs"] + + +class TestRewritesAreAllOrNothing: + """A deprecated-alias rewrite must not land when validation later fails. + + Rewrites were applied as each model was checked, so a bad agent override + left ``settings.default_model`` already mutated by a validation that raised + — the next attempt then validated a different configuration than the user + wrote. + """ + + @staticmethod + def _registry() -> ModelRegistry: + registry = ModelRegistry() + registry._models = { + "gemini-old": ModelInfo( + id="gemini-old", family=ModelFamily.GEMINI, deprecated=True + ), + "gemini-2.5-pro": ModelInfo(id="gemini-2.5-pro", family=ModelFamily.GEMINI), + } + registry._refreshed = True + return registry + + def test_default_model_is_untouched_when_an_override_fails(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._registry()) + object.__setattr__(settings, "default_model", "gemini-old") + bad = AgentConfig(name="a", prompt="p", model="claude-sonnet-4-6") + + with pytest.raises(SettingsValidationError): + validate_settings(settings, agent_configs=[bad]) + + assert settings.default_model == "gemini-old", ( + "a rewrite was applied by a validation that failed" + ) + + def test_agent_override_is_untouched_when_the_default_fails(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._registry()) + object.__setattr__(settings, "default_model", "claude-sonnet-4-6") + good = AgentConfig(name="a", prompt="p", model="gemini-old") + + with pytest.raises(SettingsValidationError): + validate_settings(settings, agent_configs=[good]) + + assert good.model == "gemini-old" + + def test_all_rewrites_land_when_everything_validates(self): + with MockContext(google_api_key="k") as ctx: + settings = ctx.settings + settings.set_model_registry(self._registry()) + object.__setattr__(settings, "default_model", "gemini-old") + config = AgentConfig(name="a", prompt="p", model="gemini-old") + + validate_settings(settings, agent_configs=[config]) + + assert settings.default_model == "gemini-2.5-pro" + assert config.model == "gemini-2.5-pro" + + +class TestProviderClientRelease: + """Listing clients are one-shot; their connection pools must be released.""" + + async def test_google_client_is_closed(self, monkeypatch): + closed: list[bool] = [] + + class _FakeClient: + def __init__(self, api_key=None): + self.models = MagicMock(list=lambda: []) + + def close(self): + closed.append(True) + + monkeypatch.setattr("google.genai.Client", _FakeClient) + await ModelRegistry()._fetch_google_models("k") + assert closed == [True] + + async def test_anthropic_client_is_closed(self, monkeypatch): + closed: list[bool] = [] + + class _FakeAnthropic: + def __init__(self, api_key=None): + self.models = MagicMock(list=lambda limit=None: MagicMock(data=[])) + + def close(self): + closed.append(True) + + monkeypatch.setattr("anthropic.Anthropic", _FakeAnthropic) + await ModelRegistry()._fetch_anthropic_models("k") + assert closed == [True] + + async def test_client_is_closed_even_when_the_listing_fails(self, monkeypatch): + closed: list[bool] = [] + + class _FakeClient: + def __init__(self, api_key=None): + self.models = MagicMock( + list=MagicMock(side_effect=RuntimeError("provider down")) + ) + + def close(self): + closed.append(True) + + monkeypatch.setattr("google.genai.Client", _FakeClient) + await ModelRegistry()._fetch_google_models("k") + assert closed == [True] From 81ac3ace2a21153cd69322b9c536450fba3547f8 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:42 -0400 Subject: [PATCH 05/11] fix(config): credential constructor arguments actually bind MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``BaseSettings(google_api_key="…")`` silently bound nothing. Each credential field declared only its environment-variable ``validation_alias``, so with ``populate_by_name`` off the field name was not an accepted input at all and ``extra="ignore"`` swallowed the kwarg — the setting kept its default and the caller got an unauthenticated client with no error. Every credential now accepts both names via ``AliasChoices``, with the env name first so a real environment variable still wins within a source. Because the field name is now accepted, a misspelled credential kwarg would be dropped just as quietly, so constructor kwargs that *look* like credentials but match no field are rejected by name. Credential values are also kept out of ``repr()``. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/config.py | 57 ++++++++++++++++ src/agentic_cli/workflow/settings.py | 24 +++++-- tests/test_config_trust.py | 97 ++++++++++++++++++++++++++++ 3 files changed, 172 insertions(+), 6 deletions(-) diff --git a/src/agentic_cli/config.py b/src/agentic_cli/config.py index 50cd7fc..e913df6 100644 --- a/src/agentic_cli/config.py +++ b/src/agentic_cli/config.py @@ -29,11 +29,13 @@ 5. Default values """ +import re from contextvars import ContextVar, Token from pathlib import Path from typing import Callable, Generator, Any, Sequence, Tuple, Type from contextlib import contextmanager +from pydantic import AliasChoices from pydantic_settings import ( BaseSettings as PydanticBaseSettings, SettingsConfigDict, @@ -105,6 +107,49 @@ def __call__(self) -> dict[str, Any]: return kept +# Constructor kwargs matching this shape are credentials; an unrecognised one +# must fail loudly instead of being swallowed by ``extra="ignore"``. +_CREDENTIAL_KEY_RE = re.compile(r"(?i)(api_?key|secret|token|password|credential)") + + +def _accepted_input_names(settings_cls: Type[PydanticBaseSettings]) -> set[str]: + """Every name the model accepts for a field: its own name and any aliases.""" + names: set[str] = set() + for field_name, field in settings_cls.model_fields.items(): + names.add(field_name) + alias = field.validation_alias + if isinstance(alias, str): + names.add(alias) + elif isinstance(alias, AliasChoices): + names.update(c for c in alias.choices if isinstance(c, str)) + if isinstance(field.alias, str): + names.add(field.alias) + return names + + +def _reject_unknown_credential_kwargs( + settings_cls: Type[PydanticBaseSettings], values: dict[str, Any] +) -> None: + """Raise on a credential-shaped kwarg the model would silently drop. + + Raises: + ValueError: If a kwarg looks like a credential but matches no field or + alias. The message names the key only — never its value. + """ + accepted = _accepted_input_names(settings_cls) + # Leading underscore = pydantic-settings' own kwargs (_env_file, + # _secrets_dir, …), not settings fields. + unknown = [k for k in values if not k.startswith("_") and k not in accepted] + bad = [k for k in unknown if _CREDENTIAL_KEY_RE.search(k)] + if not bad: + return + known = sorted(n for n in _accepted_input_names(settings_cls) if _CREDENTIAL_KEY_RE.search(n)) + raise ValueError( + f"Unknown credential setting(s): {', '.join(sorted(bad))}. " + f"{settings_cls.__name__} accepts: {', '.join(known)}." + ) + + def _get_json_config_source( settings_cls: Type[PydanticBaseSettings], json_file: Path, @@ -172,6 +217,18 @@ class BaseSettings(WorkflowSettingsMixin, AppSettingsMixin, CLISettingsMixin, Py extra="ignore", ) + def __init__(self, **values: Any) -> None: + """Construct settings, rejecting credential kwargs that would be dropped. + + ``extra="ignore"`` (needed so config files may carry keys a given app + does not define) means a mistyped constructor argument vanishes without + a word. That is tolerable for an ordinary setting and dangerous for a + credential — the app then runs unauthenticated, or silently on a + different key. Only credential-shaped unknown kwargs raise. + """ + _reject_unknown_credential_kwargs(type(self), values) + super().__init__(**values) + def update_setting(self, key: str, value: Any) -> None: """Update a single setting, using dedicated setters where required. diff --git a/src/agentic_cli/workflow/settings.py b/src/agentic_cli/workflow/settings.py index 9811b17..6ba6db5 100644 --- a/src/agentic_cli/workflow/settings.py +++ b/src/agentic_cli/workflow/settings.py @@ -10,7 +10,7 @@ from enum import Enum from typing import Literal, TYPE_CHECKING -from pydantic import BaseModel, Field, field_validator +from pydantic import AliasChoices, BaseModel, Field, field_validator from agentic_cli.logging import Loggers from agentic_cli.workflow.models import ModelFamily, ModelRegistry @@ -102,26 +102,38 @@ class WorkflowSettingsMixin: json_schema_extra={"ui_order": 27}, ) - # API Keys (common across all domains, never saved to JSON) + # API Keys (common across all domains, never saved to JSON). + # + # Each accepts BOTH the provider's environment variable name and its Python + # field name (``AliasChoices``): the bare env alias made + # ``BaseSettings(google_api_key=...)`` bind nothing at all — the value was + # dropped by ``extra="ignore"`` and the field kept its default. The env name + # is listed first, so a real environment variable still wins within a source. + # Values are kept out of ``repr()`` and out of every persisted file (see + # ``settings_persistence.SECRET_FIELDS``). google_api_key: str | None = Field( default=None, description="Google API key for Gemini models", - validation_alias="GOOGLE_API_KEY", + validation_alias=AliasChoices("GOOGLE_API_KEY", "google_api_key"), + repr=False, ) anthropic_api_key: str | None = Field( default=None, description="Anthropic API key for Claude models", - validation_alias="ANTHROPIC_API_KEY", + validation_alias=AliasChoices("ANTHROPIC_API_KEY", "anthropic_api_key"), + repr=False, ) tavily_api_key: str | None = Field( default=None, description="Tavily API key for web search", - validation_alias="TAVILY_API_KEY", + validation_alias=AliasChoices("TAVILY_API_KEY", "tavily_api_key"), + repr=False, ) brave_api_key: str | None = Field( default=None, description="Brave Search API key for web search", - validation_alias="BRAVE_API_KEY", + validation_alias=AliasChoices("BRAVE_API_KEY", "brave_api_key"), + repr=False, ) # Web search configuration diff --git a/tests/test_config_trust.py b/tests/test_config_trust.py index 9d92c8d..f7664d2 100644 --- a/tests/test_config_trust.py +++ b/tests/test_config_trust.py @@ -151,3 +151,100 @@ def test_list_env_file_with_cwd_relative_entry_is_filtered(self, tmp_path, monke (tmp_path / ".env").write_text("AGENTIC_RAW_LLM_LOGGING=true\n") # cwd-relative s = self._subclass_with_env_file([str(abs_env), ".env"])() assert s.raw_llm_logging is False # cwd-relative entry → whole source filtered + + +class TestCredentialInputSurface: + """Credential fields accept their field name *and* the provider env name. + + The bare ``validation_alias`` bound only the env var, so a programmatic + ``BaseSettings(google_api_key=...)`` was silently dropped by + ``extra="ignore"``. Widening the alias must not widen the P0-1 trust + boundary: an untrusted project file still cannot inject a key. + """ + + def _clean_env(self, monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for var in ("GOOGLE_API_KEY", "ANTHROPIC_API_KEY", "TAVILY_API_KEY", "BRAVE_API_KEY"): + monkeypatch.delenv(var, raising=False) + + def test_constructor_field_name_is_retained(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + from agentic_cli.config import BaseSettings + + settings = BaseSettings( + google_api_key="ctor-google", anthropic_api_key="ctor-anthropic" + ) + assert settings.google_api_key == "ctor-google" + assert settings.anthropic_api_key == "ctor-anthropic" + assert settings.has_any_api_key is True + + def test_environment_variable_still_binds(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + monkeypatch.setenv("ANTHROPIC_API_KEY", "env-key") + from agentic_cli.config import BaseSettings + + assert BaseSettings().anthropic_api_key == "env-key" + + def test_constructor_beats_environment(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + monkeypatch.setenv("ANTHROPIC_API_KEY", "env-key") + from agentic_cli.config import BaseSettings + + assert BaseSettings(anthropic_api_key="ctor-key").anthropic_api_key == "ctor-key" + + def test_secrets_stay_out_of_repr(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + from agentic_cli.config import BaseSettings + + settings = BaseSettings(google_api_key="super-secret") + assert "super-secret" not in repr(settings) + assert "super-secret" not in str(settings) + + def test_misspelled_credential_kwarg_raises(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + import pytest + + from agentic_cli.config import BaseSettings + + with pytest.raises(ValueError, match="Unknown credential setting"): + BaseSettings(anthropic_apikey="typo") + + def test_pydantic_settings_own_kwargs_are_not_mistaken_for_credentials( + self, tmp_path, monkeypatch + ): + """``_secrets_dir`` matches the credential shape but is a library kwarg.""" + self._clean_env(monkeypatch, tmp_path) + from agentic_cli.config import BaseSettings + + secrets_dir = tmp_path / "secrets" + secrets_dir.mkdir() + BaseSettings(_secrets_dir=str(secrets_dir)) # must not raise + + def test_unknown_non_credential_kwarg_still_ignored(self, tmp_path, monkeypatch): + """Only credential-shaped keys are strict; config files stay permissive.""" + self._clean_env(monkeypatch, tmp_path) + from agentic_cli.config import BaseSettings + + BaseSettings(some_future_option=True) # must not raise + + def test_project_settings_json_still_cannot_inject_a_key(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + _write_project_settings( + tmp_path, + "agentic_cli", + {"google_api_key": "from-untrusted-repo", "GOOGLE_API_KEY": "also-untrusted"}, + ) + from agentic_cli.config import BaseSettings + + assert BaseSettings().google_api_key is None + + def test_cwd_dotenv_still_cannot_inject_a_key(self, tmp_path, monkeypatch): + self._clean_env(monkeypatch, tmp_path) + (tmp_path / ".env").write_text("GOOGLE_API_KEY=from-untrusted-repo\n") + from agentic_cli.config import BaseSettings + + class _DomainSettings(BaseSettings): + model_config = {**BaseSettings.model_config, "env_file": ".env"} + + assert _DomainSettings().google_api_key is None From 68f8724481c46ed9fe01c9c25079582454ccb56b Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:44 -0400 Subject: [PATCH 06/11] feat(workflow)!: SessionRef and user-scoped session APIs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A conversation's identity is the ``(app_name, user_id, session_id)`` triple, but the session hooks took only a session id and implicitly used the manager's default user — so listing, deleting, or reading the history of another user's session silently answered about the wrong conversation, or came back empty. ``session_exists``/``list_sessions``/``delete_session``/``recent_messages``/ ``load_session``/``save_session`` now take an optional ``user_id``, defaulting to ``settings.default_user`` only when the caller omits it, and ``on_session_end(session=SessionRef(...))`` reads the conversation it is given. Base-class helpers still call the backend hooks *without* ``user_id`` when it is the default, so a downstream override that never added the parameter keeps working. The in-flight ``(user, session)`` is a ContextVar set with a token by ``_workflow_context()``, so concurrent turns on one manager stay isolated and nesting restores the outer turn. A backend with no durable store now leaves ``supports_sessions`` False and the base hooks raise ``NotImplementedError`` rather than answering ``False``/``[]`` — an empty list read as "you have no saved sessions" when the truth was "this backend keeps none", which ``/sessions`` now says outright. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/cli/builtin_commands.py | 9 + src/agentic_cli/workflow/adk/manager.py | 63 +++-- src/agentic_cli/workflow/base_manager.py | 176 ++++++++++--- src/agentic_cli/workflow/sessions.py | 61 +++++ tests/cli/test_sessions_command.py | 25 +- tests/workflow/test_active_turn_context.py | 68 ++++- tests/workflow/test_session_store.py | 285 ++++++++++++++++++++- 7 files changed, 627 insertions(+), 60 deletions(-) create mode 100644 src/agentic_cli/workflow/sessions.py diff --git a/src/agentic_cli/cli/builtin_commands.py b/src/agentic_cli/cli/builtin_commands.py index 766a2d4..a101bc7 100644 --- a/src/agentic_cli/cli/builtin_commands.py +++ b/src/agentic_cli/cli/builtin_commands.py @@ -455,6 +455,15 @@ async def execute(self, args: str, app: Any) -> None: app.session.add_warning("Sessions not available yet — workflow is initializing.") return + # An empty list would read as "you have no saved sessions"; say plainly + # that this backend keeps none. + if not getattr(workflow, "supports_sessions", False): + app.session.add_warning( + f"The {getattr(workflow, 'backend_type', 'current')} backend does not " + "persist sessions, so there are none to list or delete." + ) + return + if delete_id: if await workflow.delete_session(delete_id): app.session.add_success(f"Session '{delete_id}' deleted.") diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index 52eff00..8ca5138 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -1095,23 +1095,35 @@ def _job_result_payload(record, result: Any) -> dict: # Sessions (native — DatabaseSessionService persists events continuously) # ------------------------------------------------------------------------- - async def session_exists(self, session_id: str) -> bool: - """True if the store holds this session with any events.""" + async def session_exists(self, session_id: str, *, user_id: str | None = None) -> bool: + """True if the store holds this session with any events. + + Args: + session_id: Session to look up. + user_id: Owner of the session (defaults to ``settings.default_user``). + """ if not self._session_service: return False + ref = self.session_ref(session_id, user_id) session = await self._session_service.get_session( - app_name=self.app_name, - user_id=self._settings.default_user, - session_id=session_id, + app_name=ref.app_name, + user_id=ref.user_id, + session_id=ref.session_id, ) return session is not None and bool(getattr(session, "events", None)) - async def list_sessions(self) -> list[dict]: - """List persisted sessions for the current user (most recent first).""" + async def list_sessions(self, *, user_id: str | None = None) -> list[dict]: + """List persisted sessions for a user (most recent first). + + Args: + user_id: Owner whose sessions to list (defaults to + ``settings.default_user``). + """ if not self._session_service: return [] + ref = self.session_ref(user_id=user_id) resp = await self._session_service.list_sessions( - app_name=self.app_name, user_id=self._settings.default_user, + app_name=ref.app_name, user_id=ref.user_id, ) sessions = [ { @@ -1124,25 +1136,40 @@ async def list_sessions(self) -> list[dict]: sessions.sort(key=lambda x: x["last_update"] or 0, reverse=True) return sessions - async def delete_session(self, session_id: str) -> bool: - """Delete a persisted session from the store.""" + async def delete_session(self, session_id: str, *, user_id: str | None = None) -> bool: + """Delete a persisted session from the store. + + Args: + session_id: Session to delete. + user_id: Owner of the session (defaults to ``settings.default_user``). + """ if not self._session_service: return False + ref = self.session_ref(session_id, user_id) await self._session_service.delete_session( - app_name=self.app_name, - user_id=self._settings.default_user, - session_id=session_id, + app_name=ref.app_name, + user_id=ref.user_id, + session_id=ref.session_id, ) return True - async def recent_messages(self, session_id: str, limit: int = 20) -> list[dict]: - """Recent text messages from the stored session (for fact extraction).""" + async def recent_messages( + self, session_id: str, limit: int = 20, *, user_id: str | None = None + ) -> list[dict]: + """Recent text messages from the stored session (for fact extraction). + + Args: + session_id: Session to read. + limit: Maximum number of messages returned (most recent last). + user_id: Owner of the session (defaults to ``settings.default_user``). + """ if not self._session_service: return [] + ref = self.session_ref(session_id, user_id) session = await self._session_service.get_session( - app_name=self.app_name, - user_id=self._settings.default_user, - session_id=session_id, + app_name=ref.app_name, + user_id=ref.user_id, + session_id=ref.session_id, ) if session is None or not getattr(session, "events", None): return [] diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index d5f029a..7469f44 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -20,6 +20,12 @@ from agentic_cli.workflow.events import WorkflowEvent, UserInputRequest from agentic_cli.workflow.config import AgentConfig from agentic_cli.workflow.models import ModelRegistry +from agentic_cli.workflow.sessions import ( + SessionRef, + get_active_turn, + reset_active_turn, + set_active_turn, +) from agentic_cli.workflow.service_registry import ( set_service_registry, ARXIV_SOURCE, @@ -474,13 +480,23 @@ async def summarize(self, content: str, prompt: str) -> str: """ return await self.generate_simple(prompt, max_tokens=12000) - async def on_session_end(self, messages: list[dict] | None = None) -> list[str]: + async def on_session_end( + self, + messages: list[dict] | None = None, + *, + session: "SessionRef | None" = None, + ) -> list[str]: """Hook called when a session ends. Optionally extracts facts. Override in downstream apps for custom session-end behavior. Args: messages: Recent messages from the session (optional). + session: Which conversation to read when ``messages`` is omitted. + Defaults to the turn still in context, else this manager's + current session under ``settings.default_user`` — so a session + belonging to another user is read as *that* user rather than + silently coming back empty. Returns: List of extracted facts (empty if disabled or no messages). @@ -496,9 +512,16 @@ async def on_session_end(self, messages: list[dict] | None = None) -> list[str]: # backend session (same source/sid save_session uses) so the CLI can # invoke this with no arguments on exit. if messages is None: - sid = getattr(self, "session_id", "default_session") + ref = session or get_active_turn() or self.session_ref() try: - messages = await self.recent_messages(sid) + if self._is_default_user(ref.user_id): + # Compatible call for backends that predate the user_id + # parameter (see _user_scoped_kwargs). + messages = await self.recent_messages(ref.session_id) + else: + messages = await self.recent_messages( + ref.session_id, user_id=ref.user_id + ) except Exception: logger.debug("session_fact_extraction_extract_failed", exc_info=True) return [] @@ -528,15 +551,27 @@ async def on_session_end(self, messages: list[dict] | None = None) -> list[str]: store.store(fact, tags=["auto-extracted", "session"]) return facts + @property + def active_turn(self) -> SessionRef | None: + """Identity of the turn running in this context, or None when idle. + + Context-local, not manager-local: concurrent ``process()`` calls on one + manager (possible for framework consumers — the CLI serializes turns) + each see their own value. + """ + return get_active_turn() + @property def active_session_id(self) -> str | None: """Session id of the in-flight ``process()`` call, or None when idle.""" - return self._active_session_id + ref = get_active_turn() + return ref.session_id if ref else None @property def active_user_id(self) -> str | None: """User id of the in-flight ``process()`` call, or None when idle.""" - return self._active_user_id + ref = get_active_turn() + return ref.user_id if ref else None async def can_resume(self, record) -> bool: """Whether a finished job can be resumed into its conversation now. @@ -556,22 +591,24 @@ def _workflow_context( ) -> Iterator[None]: """Context manager that exposes the service registry to tools. - Sets a single ContextVar (the service registry) so tools can - call ``get_service(key)`` during execution, and records the active - session/user for the duration of the turn so JobManager can associate - a launched job with the conversation that started it. + Sets ContextVars (settings, the service registry, and the active turn) + so tools can call ``get_service(key)`` during execution and the + JobManager can associate a launched job with the conversation that + started it. + + All three are restored from tokens on exit, so a nested context + restores the outer turn rather than clearing it, and concurrent turns + on one manager never see each other's identity. """ from agentic_cli.config import set_context_settings settings_token = set_context_settings(self._settings) registry_token = set_service_registry(self._services) - self._active_session_id = session_id - self._active_user_id = user_id + turn_token = set_active_turn(self.session_ref(session_id, user_id)) try: yield finally: - self._active_session_id = None - self._active_user_id = None + reset_active_turn(turn_token) registry_token.var.reset(registry_token) settings_token.var.reset(settings_token) @@ -850,47 +887,126 @@ async def generate_simple(self, prompt: str, max_tokens: int = 500) -> str: # state continuously, keyed by session_id; there is no separate snapshot. # ------------------------------------------------------------------ - async def save_session(self, session_id: str | None = None) -> dict: - """No-op flush — durable stores persist as the turn runs. + @property + def supports_sessions(self) -> bool: + """Whether this backend implements the durable-session hooks. + + Derived from the subclass actually overriding ``session_exists``, so a + backend opts in by implementing the capability rather than by setting a + flag that can drift from the code. + """ + return type(self).session_exists is not BaseWorkflowManager.session_exists + + def _is_default_user(self, user_id: str | None) -> bool: + """Whether ``user_id`` is (or defaults to) the configured default user. - Kept for API compatibility / explicit "checkpoint now" intent. Returns - the session id that is (already) persisted. + Base-class helpers call user-scoped backend hooks *without* the + ``user_id`` keyword in that case, so a backend that predates the + parameter (LangGraph, and any downstream override) keeps working; an + explicit non-default user is always passed through, so it can never be + silently serviced as the default user. """ - sid = session_id or getattr(self, "session_id", "default_session") - return {"success": True, "session_id": sid} + return user_id is None or user_id == self._settings.default_user - async def load_session(self, session_id: str) -> bool: + def session_ref( + self, session_id: str | None = None, user_id: str | None = None + ) -> SessionRef: + """Resolve a full :class:`SessionRef` from partial identity. + + Unsupplied parts default to this manager's app name, the configured + ``default_user`` and the manager's current session id. Defaults apply + only to what the caller omitted: an explicit ``user_id`` is never + replaced by the default user. + """ + return SessionRef( + app_name=self.app_name, + user_id=user_id or self._settings.default_user, + session_id=session_id or getattr(self, "session_id", "default_session"), + ) + + async def save_session( + self, session_id: str | None = None, *, user_id: str | None = None + ) -> dict: + """No-op flush — durable stores persist as the turn runs. + + Kept for API compatibility / explicit "checkpoint now" intent. + + Args: + session_id: Session to report (defaults to the manager's current one). + user_id: Owner (defaults to ``settings.default_user``). + + Returns: + ``{"success": True, "session_id": ..., "user_id": ...}`` — the full + identity, so a caller working on behalf of another user can tell + which conversation was meant. + """ + ref = self.session_ref(session_id, user_id) + return { + "success": True, + "session_id": ref.session_id, + "user_id": ref.user_id, + } + + async def load_session(self, session_id: str, *, user_id: str | None = None) -> bool: """Adopt ``session_id`` for resume; the native store already holds it. Returns True if that session already has content (i.e. a real resume), False if it's new — but the id is adopted either way so the next turn - continues it. + continues it. A backend without durable sessions has nothing to resume, + so it adopts the id and returns False. + + Args: + session_id: Session to adopt. + user_id: Owner to look the session up as (defaults to + ``settings.default_user``). """ if hasattr(self, "session_id"): self.session_id = session_id - exists = await self.session_exists(session_id) + if not self.supports_sessions: + logger.info( + "session_adopted", session_id=session_id, resumed=False, + backend_sessions=False, + ) + return False + if self._is_default_user(user_id): + exists = await self.session_exists(session_id) + else: + exists = await self.session_exists(session_id, user_id=user_id) logger.info("session_adopted", session_id=session_id, resumed=exists) return exists # ---- Backend hooks (override in ADK / LangGraph managers) ---- + # + # Each takes an optional ``user_id`` so a session created for one user + # stays reachable through the public API. Backends that cannot persist + # sessions must not answer with a misleading "no" — the base raises. + + def _no_session_support(self, operation: str) -> NotImplementedError: + """Error for a session operation the backend does not implement.""" + return NotImplementedError( + f"{type(self).__name__} does not implement durable sessions " + f"({operation}). Check ``supports_sessions`` before calling." + ) - async def session_exists(self, session_id: str) -> bool: + async def session_exists(self, session_id: str, *, user_id: str | None = None) -> bool: """Whether the native store already holds this session's state.""" - return False + raise self._no_session_support("session_exists") - async def recent_messages(self, session_id: str, limit: int = 20) -> list[dict]: + async def recent_messages( + self, session_id: str, limit: int = 20, *, user_id: str | None = None + ) -> list[dict]: """Recent ``{role, content}`` text messages from the native session. Used for session-end fact extraction; text-only (no tool-call fidelity). """ - return [] + raise self._no_session_support("recent_messages") - async def list_sessions(self) -> list[dict]: + async def list_sessions(self, *, user_id: str | None = None) -> list[dict]: """List persisted sessions from the native store (most recent first).""" - return [] + raise self._no_session_support("list_sessions") - async def delete_session(self, session_id: str) -> bool: + async def delete_session(self, session_id: str, *, user_id: str | None = None) -> bool: """Delete a persisted session from the native store.""" - return False + raise self._no_session_support("delete_session") diff --git a/src/agentic_cli/workflow/sessions.py b/src/agentic_cli/workflow/sessions.py new file mode 100644 index 0000000..dcb5530 --- /dev/null +++ b/src/agentic_cli/workflow/sessions.py @@ -0,0 +1,61 @@ +"""Backend-neutral conversation identity. + +Every durable session is addressed by the triple ``(app_name, user_id, +session_id)`` — the ADK session services key on exactly that, and any +replacement backend has to carry the same information. ``SessionRef`` makes +that identity explicit so a session created for one user cannot be looked up, +listed, or deleted as another user's by accident. +""" + +from __future__ import annotations + +from contextvars import ContextVar, Token +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class SessionRef: + """Identity of one conversation in a session store. + + Attributes: + app_name: Namespace of the owning application. + user_id: Owner of the conversation. + session_id: Conversation identifier, unique within (app_name, user_id). + """ + + app_name: str + user_id: str + session_id: str + + def __str__(self) -> str: # pragma: no cover - trivial + return f"{self.app_name}/{self.user_id}/{self.session_id}" + + +# The conversation whose turn is currently executing. A ContextVar, not a +# manager attribute: one manager instance may drive several turns at once when +# it is embedded in a server (the CLI serializes turns, framework consumers do +# not), and concurrent tasks must never observe each other's identity. Each +# asyncio task gets its own copy of the context, so isolation is automatic. +_active_turn: ContextVar["SessionRef | None"] = ContextVar( + "agentic_cli_active_turn", default=None +) + + +def set_active_turn(ref: "SessionRef | None") -> Token: + """Mark ``ref`` as the turn running in this context. + + Returns: + Token for :func:`reset_active_turn` — reset restores the *previous* + value, so nested turns do not erase the outer one. + """ + return _active_turn.set(ref) + + +def reset_active_turn(token: Token) -> None: + """Restore the active turn recorded before the matching :func:`set_active_turn`.""" + _active_turn.reset(token) + + +def get_active_turn() -> "SessionRef | None": + """The turn executing in this context, or None when idle.""" + return _active_turn.get() diff --git a/tests/cli/test_sessions_command.py b/tests/cli/test_sessions_command.py index 40124ec..a88f4b3 100644 --- a/tests/cli/test_sessions_command.py +++ b/tests/cli/test_sessions_command.py @@ -10,20 +10,32 @@ class _Workflow: + supports_sessions = True + backend_type = "adk" + def __init__(self, sessions: list[dict]) -> None: self._sessions = sessions self.deleted: list[str] = [] - async def list_sessions(self) -> list[dict]: + async def list_sessions(self, *, user_id: str | None = None) -> list[dict]: return self._sessions - async def delete_session(self, session_id: str) -> bool: + async def delete_session( + self, session_id: str, *, user_id: str | None = None + ) -> bool: if any(s["session_id"] == session_id for s in self._sessions): self.deleted.append(session_id) return True return False +class _SessionlessWorkflow: + """A backend with no durable session store.""" + + supports_sessions = False + backend_type = "custom" + + class _App: def __init__(self, workflow, session_id: str = "cur") -> None: self._wf = workflow @@ -70,3 +82,12 @@ async def test_not_ready_warns(): app = _App(None) await SessionsCommand().execute("", app) assert app.session.warnings() + + +async def test_backend_without_sessions_says_so(): + """An empty list would read as 'no saved sessions' — be explicit instead.""" + app = _App(_SessionlessWorkflow()) + await SessionsCommand().execute("", app) + warnings = app.session.warnings() + assert warnings and any("does not" in str(w) for w in warnings) + assert not app.session.of("rich") diff --git a/tests/workflow/test_active_turn_context.py b/tests/workflow/test_active_turn_context.py index 2574fa5..f101877 100644 --- a/tests/workflow/test_active_turn_context.py +++ b/tests/workflow/test_active_turn_context.py @@ -1,13 +1,18 @@ -"""The manager exposes the active session/user during a turn (phase-2 association). +"""The active turn is context-local, not manager-local. ``JobManager`` reads ``active_session_id``/``active_user_id`` off the WORKFLOW service to associate a resume-on-complete job with the conversation that -launched it. ``_workflow_context()`` sets these for the turn and clears them on -exit (even on error). +launched it. ``_workflow_context()`` publishes them for the duration of a turn. + +They used to live in manager instance attributes, which cross-wired two turns +running on one manager (possible for framework consumers — only the CLI +serializes turns) and made a nested context erase the outer turn on exit. They +are now a ``ContextVar`` restored from a token. """ from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest @@ -15,15 +20,16 @@ pytest.importorskip("google.adk") from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.sessions import SessionRef, get_active_turn # noqa: E402 def _bare_manager() -> GoogleADKWorkflowManager: """A manager instance without running __init__ (concrete subclass of base).""" mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) - mgr._settings = SimpleNamespace(app_name="test") + mgr._settings = SimpleNamespace(app_name="test", default_user="default_user") + mgr._app_name = "test" mgr._services = {} - mgr._active_session_id = None - mgr._active_user_id = None + mgr.session_id = "default_session" return mgr @@ -31,6 +37,7 @@ def test_idle_active_ids_are_none(): mgr = _bare_manager() assert mgr.active_session_id is None assert mgr.active_user_id is None + assert mgr.active_turn is None def test_context_sets_and_clears_active_ids(): @@ -38,6 +45,7 @@ def test_context_sets_and_clears_active_ids(): with mgr._workflow_context(session_id="sess-1", user_id="user-1"): assert mgr.active_session_id == "sess-1" assert mgr.active_user_id == "user-1" + assert mgr.active_turn == SessionRef("test", "user-1", "sess-1") assert mgr.active_session_id is None assert mgr.active_user_id is None @@ -50,3 +58,51 @@ def test_context_clears_on_exception(): raise RuntimeError("boom") assert mgr.active_session_id is None assert mgr.active_user_id is None + + +def test_nested_context_restores_outer_turn(): + """The inner context must restore the outer turn, not clear it.""" + mgr = _bare_manager() + with mgr._workflow_context(session_id="outer", user_id="alice"): + with mgr._workflow_context(session_id="inner", user_id="bob"): + assert mgr.active_session_id == "inner" + assert mgr.active_user_id == "bob" + assert mgr.active_session_id == "outer" + assert mgr.active_user_id == "alice" + assert mgr.active_turn is None + + +async def test_concurrent_turns_do_not_cross_wire(): + """Two turns on one manager must each see their own identity throughout.""" + mgr = _bare_manager() + observed: dict[str, list[tuple[str | None, str | None]]] = {"a": [], "b": []} + both_inside = asyncio.Barrier(2) + + async def _turn(tag: str, session_id: str, user_id: str) -> None: + with mgr._workflow_context(session_id=session_id, user_id=user_id): + observed[tag].append((mgr.active_session_id, mgr.active_user_id)) + # Force overlap: neither task leaves its context until both entered. + await both_inside.wait() + observed[tag].append((mgr.active_session_id, mgr.active_user_id)) + + await asyncio.gather( + _turn("a", "sess-a", "alice"), + _turn("b", "sess-b", "bob"), + ) + + assert observed["a"] == [("sess-a", "alice"), ("sess-a", "alice")] + assert observed["b"] == [("sess-b", "bob"), ("sess-b", "bob")] + assert get_active_turn() is None + + +async def test_turn_identity_visible_to_spawned_tasks(): + """Tools run in tasks spawned inside the turn; they inherit its context.""" + mgr = _bare_manager() + + async def _tool() -> tuple[str | None, str | None]: + return mgr.active_session_id, mgr.active_user_id + + with mgr._workflow_context(session_id="sess-1", user_id="alice"): + result = await asyncio.create_task(_tool()) + + assert result == ("sess-1", "alice") diff --git a/tests/workflow/test_session_store.py b/tests/workflow/test_session_store.py index 5995329..62a1410 100644 --- a/tests/workflow/test_session_store.py +++ b/tests/workflow/test_session_store.py @@ -18,6 +18,44 @@ def _settings(tmp_path: Path, **over) -> BaseSettings: return BaseSettings(workspace_dir=tmp_path, **over) +@pytest.fixture +async def _closing_session_services(monkeypatch): + """Close every session service the test creates. + + ``DatabaseSessionService`` owns a SQLAlchemy async engine whose connection + worker thread outlives a service that is merely dropped. When the engine is + eventually finalized that thread raises, and pytest reports it as a + ``PytestUnhandledThreadExceptionWarning`` against whichever *unrelated* + test happens to be running at the time. Closing them here keeps the failure + attributable — and matches the ownership contract the manager itself obeys + (see ``BaseWorkflowManager._aclose_owned``). + """ + import inspect + + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + created: list[object] = [] + original = GoogleADKWorkflowManager._make_session_service + + def _tracked(self): + service = original(self) + created.append(service) + return service + + monkeypatch.setattr( + GoogleADKWorkflowManager, "_make_session_service", _tracked + ) + yield created + + for service in created: + close = getattr(service, "aclose", None) or getattr(service, "close", None) + if close is None: + continue + result = close() + if inspect.isawaitable(result): + await result + + class TestSessionDbUrl: def test_sqlite_default(self, tmp_path: Path): url = _settings(tmp_path, session_store="sqlite").session_db_url() @@ -47,7 +85,7 @@ def test_explicit_sqlite_uri_normalized(self, tmp_path: Path): class TestAdkSessionServiceSelection: @pytest.fixture(autouse=True) - def _require_adk(self): + async def _require_adk(self, _closing_session_services): pytest.importorskip("google.adk") def _manager(self, settings): @@ -57,13 +95,13 @@ def _manager(self, settings): mgr._settings = settings return mgr - def test_memory_uses_in_memory_service(self, tmp_path: Path): + async def test_memory_uses_in_memory_service(self, tmp_path: Path): from google.adk.sessions import InMemorySessionService mgr = self._manager(_settings(tmp_path, session_store="memory")) assert isinstance(mgr._make_session_service(), InMemorySessionService) - def test_sqlite_uses_database_service_and_creates_dir(self, tmp_path: Path): + async def test_sqlite_uses_database_service_and_creates_dir(self, tmp_path: Path): from google.adk.sessions import DatabaseSessionService mgr = self._manager(_settings(tmp_path, session_store="sqlite")) @@ -76,7 +114,7 @@ class TestAdkNativeSessions: """Native session query/manage against a real sqlite DatabaseSessionService.""" @pytest.fixture(autouse=True) - def _require_adk(self): + async def _require_adk(self, _closing_session_services): pytest.importorskip("google.adk") def _manager(self, tmp_path: Path): @@ -137,3 +175,242 @@ async def test_persists_across_fresh_manager(self, tmp_path: Path): # A second manager over the same sqlite file sees the session. mgr2, _ = self._manager(tmp_path) assert await mgr2.session_exists("sess-z") is True + + +class TestAdkSessionUserScope: + """Session APIs accept an explicit user_id, defaulting to settings.default_user. + + process() always accepted arbitrary user_id (and job-resume threads + record.user_id), but the query/manage APIs hard-coded default_user — + sessions created for another user were invisible to them. + """ + + @pytest.fixture(autouse=True) + async def _require_adk(self, _closing_session_services): + pytest.importorskip("google.adk") + + def _manager(self, tmp_path: Path): + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + settings = _settings(tmp_path, session_store="sqlite") + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = settings + mgr._app_name = "test_app" + mgr.session_id = "default_session" + mgr._session_service = mgr._make_session_service() + return mgr, settings + + async def _seed_as(self, mgr, user_id: str, sid: str, text: str): + from google.adk.events import Event + from google.genai import types + + s = await mgr._session_service.create_session( + app_name=mgr.app_name, user_id=user_id, session_id=sid + ) + await mgr._session_service.append_event( + session=s, + event=Event( + author="user", + content=types.Content( + role="user", parts=[types.Part.from_text(text=text)] + ), + ), + ) + + async def test_session_apis_scope_to_given_user(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + await self._seed_as(mgr, "alice", "sess-a", "alice message") + await self._seed_as(mgr, settings.default_user, "sess-d", "default message") + + # Default scope: unchanged behavior, sees only default_user's sessions + assert await mgr.session_exists("sess-d") is True + assert await mgr.session_exists("sess-a") is False + + # Explicit user scope reaches alice's session through every API + assert await mgr.session_exists("sess-a", user_id="alice") is True + + listed = await mgr.list_sessions(user_id="alice") + assert [s["session_id"] for s in listed] == ["sess-a"] + + recent = await mgr.recent_messages("sess-a", user_id="alice") + assert recent and recent[-1]["content"] == "alice message" + + assert await mgr.delete_session("sess-a", user_id="alice") is True + assert await mgr.session_exists("sess-a", user_id="alice") is False + + # Default user's session untouched by alice-scoped operations + assert await mgr.session_exists("sess-d") is True + + async def test_session_ref_resolves_partial_identity(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + + ref = mgr.session_ref() + assert (ref.app_name, ref.user_id, ref.session_id) == ( + "test_app", settings.default_user, "default_session", + ) + + explicit = mgr.session_ref("sess-x", "alice") + assert (explicit.user_id, explicit.session_id) == ("alice", "sess-x") + + async def test_load_session_honours_explicit_user(self, tmp_path: Path): + """A session adopted for another user must be seen as a real resume.""" + mgr, _ = self._manager(tmp_path) + await self._seed_as(mgr, "alice", "sess-a", "alice message") + + assert await mgr.load_session("sess-a") is False # default user: not theirs + assert await mgr.load_session("sess-a", user_id="alice") is True + assert mgr.session_id == "sess-a" + + async def test_adk_reports_session_support(self, tmp_path: Path): + mgr, _ = self._manager(tmp_path) + assert mgr.supports_sessions is True + + +class TestSessionsUnsupportedFailExplicitly: + """A backend without durable sessions must not answer with a bare False/[].""" + + def _manager(self): + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + class _NoSessions(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + return None + + from unittest.mock import MagicMock + + settings = MagicMock() + settings.app_name = "test-app" + settings.default_user = "default_user" + return _NoSessions(agent_configs=[], settings=settings) + + def test_supports_sessions_is_false(self): + assert self._manager().supports_sessions is False + + async def test_hooks_raise_not_implemented(self): + mgr = self._manager() + for call in ( + mgr.session_exists("s"), + mgr.list_sessions(), + mgr.delete_session("s"), + mgr.recent_messages("s"), + ): + with pytest.raises(NotImplementedError, match="durable sessions"): + await call + + async def test_load_session_adopts_without_raising(self): + """Adoption still works — there is simply nothing to resume.""" + mgr = self._manager() + assert await mgr.load_session("sess-1") is False + + +class TestSessionEndScope: + """Fact extraction must read the session it actually belongs to.""" + + @pytest.fixture(autouse=True) + def _require_adk(self): + pytest.importorskip("google.adk") + + def _manager(self, tmp_path: Path): + from unittest.mock import MagicMock + + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + settings = _settings(tmp_path, session_store="sqlite") + object.__setattr__(settings, "auto_extract_session_facts", True) + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = settings + mgr._app_name = "test_app" + mgr.session_id = "sess-current" + mgr._services = {"memory_store": MagicMock()} + mgr._session_service = None + return mgr, settings + + async def test_reads_the_default_user_session_by_default(self, tmp_path: Path): + """The default user is passed positionally, so pre-``user_id`` + backend overrides keep working.""" + mgr, settings = self._manager(tmp_path) + seen: list[tuple[str, str | None]] = [] + + async def _recent(session_id, limit=20, *, user_id=None): + seen.append((session_id, user_id)) + return [] + + mgr.recent_messages = _recent + await mgr.on_session_end() + + assert seen == [("sess-current", None)] + + async def test_legacy_override_without_user_id_still_works(self, tmp_path: Path): + """A backend that predates the parameter must not break (LangGraph).""" + mgr, _ = self._manager(tmp_path) + seen: list[str] = [] + + async def _legacy_recent(session_id, limit=20): + seen.append(session_id) + return [] + + mgr.recent_messages = _legacy_recent + await mgr.on_session_end() + + assert seen == ["sess-current"] + + async def test_explicit_session_ref_is_honoured(self, tmp_path: Path): + from agentic_cli.workflow.sessions import SessionRef + + mgr, _ = self._manager(tmp_path) + seen: list[tuple[str, str | None]] = [] + + async def _recent(session_id, limit=20, *, user_id=None): + seen.append((session_id, user_id)) + return [] + + mgr.recent_messages = _recent + await mgr.on_session_end( + session=SessionRef(app_name="test_app", user_id="alice", session_id="sess-a") + ) + + assert seen == [("sess-a", "alice")], "another user's session was not read" + + async def test_active_turn_identity_is_used_when_available(self, tmp_path: Path): + mgr, _ = self._manager(tmp_path) + seen: list[tuple[str, str | None]] = [] + + async def _recent(session_id, limit=20, *, user_id=None): + seen.append((session_id, user_id)) + return [] + + mgr.recent_messages = _recent + with mgr._workflow_context(session_id="sess-live", user_id="bob"): + await mgr.on_session_end() + + assert seen == [("sess-live", "bob")] + + async def test_save_session_reports_full_identity(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + + assert await mgr.save_session() == { + "success": True, + "session_id": "sess-current", + "user_id": settings.default_user, + } + assert await mgr.save_session("s2", user_id="alice") == { + "success": True, + "session_id": "s2", + "user_id": "alice", + } From 2c89da468f0d35bc89ff7c64be8afcdf945fa252 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:45 -0400 Subject: [PATCH 07/11] fix(workflow): serialize turns, make init transactional, own resources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A manager ran turns concurrently against one backend, and lifecycle mutation could land in the middle of one. ``process()``/``resume_with_job_result()`` now enter through ``_turn_admission()`` (turn lock) while ``initialize_services``/``reinitialize``/``cleanup`` hold the lifecycle lock *and* the turn lock. Lock order is lifecycle → turn, so a turn initializes before taking the turn lock — which is what keeps the two from deadlocking, and is also why a queued cleanup could run in between and leave the turn holding admission to a released backend (a ``None`` runner, surfacing as an AttributeError deep inside ADK). Admission therefore re-checks ``_backend_ready()`` while holding the turn lock and reinitializes once, or fails cleanly. Initialization is transactional. Services are built on a worker thread into a *local* dict and published only while the attempt still owns init: writing straight into the manager meant a cancelled attempt was followed, moments later, by that uncancellable thread publishing into a manager that had already been cleaned up. A cancelled attempt now releases whatever the thread went on to build, and a constructor that raises releases its predecessors — nothing was published, so nobody else could ever have closed them. ``cleanup()`` is idempotent and awaits an async ``close()`` on owned resources. A failed in-place reinitialization *keeps* the manager: it rolled itself back to uninitialized, but still owns the session service its own ``reinitialize(preserve_sessions=True)`` restored, and releasing it threw away the conversation for a failure the user could correct and retry. The HITL input callback becomes a per-manager ContextVar for the same reason: it is one manager with possibly two consumers, and a plain attribute let the second one capture a running turn's prompt and let either one unregister the other's callback. ``WorkflowController`` gains a derived ``WorkflowState`` that can never drift, serializes every lifecycle transition on its own lock, and never publishes after ``close()``. Construction in the init executor is shielded and tracked by a single-shot claim: cancelling the await used to cancel the asyncio future, after which asyncio silently discarded the manager the (uncancellable) thread returned. Shutdown runs in a task the controller owns and callers join under a shield, so a cancelled caller cannot abandon a teardown half-done — including the cleanup of an abandoned construction, which a later ``close()`` joins. ``_init_error`` is cleared on every success, so state, ``ensure_initialized()`` and the status bar can no longer disagree about a recovered controller. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/cli/workflow_controller.py | 476 +++++++- src/agentic_cli/workflow/adk/manager.py | 250 ++-- src/agentic_cli/workflow/base_manager.py | 342 +++++- tests/integration/test_adk_integration.py | 13 +- tests/test_input_callback.py | 69 +- tests/test_workflow.py | 8 +- tests/test_workflow_controller.py | 1082 +++++++++++++++++ tests/tools/test_registry_identity.py | 13 + tests/workflow/test_base_manager_init_lock.py | 78 ++ tests/workflow/test_lifecycle_races.py | 287 +++++ tests/workflow/test_model_validation.py | 19 + tests/workflow/test_resource_ownership.py | 457 +++++++ tests/workflow/test_turn_serialization.py | 316 +++++ 13 files changed, 3195 insertions(+), 215 deletions(-) create mode 100644 tests/workflow/test_base_manager_init_lock.py create mode 100644 tests/workflow/test_lifecycle_races.py create mode 100644 tests/workflow/test_resource_ownership.py create mode 100644 tests/workflow/test_turn_serialization.py diff --git a/src/agentic_cli/cli/workflow_controller.py b/src/agentic_cli/cli/workflow_controller.py index 5d27e2f..7500d9d 100644 --- a/src/agentic_cli/cli/workflow_controller.py +++ b/src/agentic_cli/cli/workflow_controller.py @@ -14,6 +14,7 @@ import asyncio from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager +from enum import Enum from typing import TYPE_CHECKING, AsyncIterator from agentic_cli.logging import Loggers @@ -35,14 +36,76 @@ logger = Loggers.cli() +class WorkflowState(str, Enum): + """Lifecycle state of a :class:`WorkflowController`. + + Derived from the controller's internals rather than stored, so the reported + state can never drift from reality: + + - ``UNINITIALIZED`` — no init attempted (or none since the last failure was + cleared). + - ``INITIALIZING`` — a background init task is in flight. + - ``READY`` — a manager finished ``initialize_services()`` and is published. + - ``FAILED`` — the last init attempt raised; ``init_error`` holds it. + ``start_background_init()`` clears it and retries. + - ``CLOSED`` — ``close()`` ran; the controller is terminal. + """ + + UNINITIALIZED = "uninitialized" + INITIALIZING = "initializing" + READY = "ready" + FAILED = "failed" + CLOSED = "closed" + + +class _Construction: + """A manager being built in the init executor, settled exactly once. + + The worker thread cannot be cancelled, so whoever stops waiting for it must + still take responsibility for what it eventually returns. ``claim()`` is the + single-shot token that decides who does: shutdown if it gets there first, + otherwise the future's done-callback. + """ + + __slots__ = ("future", "_settled") + + def __init__(self, future: "asyncio.Future") -> None: + self.future = future + self._settled = False + + def claim(self) -> bool: + """Take responsibility for the result. True for exactly one caller.""" + if self._settled: + return False + self._settled = True + return True + + def release(self) -> None: + """Give the claim back, for a claimer that could not finish. + + A settler cancelled between claiming and releasing the manager would + otherwise strand it: the claim is single-shot, so nobody else — not + even the future's own callback — could take over. + """ + self._settled = False + + class WorkflowController: """Manages the workflow manager lifecycle. Encapsulates: - - Background initialization in ThreadPoolExecutor - - Readiness checking (blocking and non-blocking) - - Reinitialization when model/settings change - - Cleanup of pending init tasks + - Background initialization in ThreadPoolExecutor, single-flight: repeated + ``start_background_init()`` calls join the in-flight attempt instead of + building a second manager + - Readiness checking (blocking and non-blocking); the manager is published + only after ``initialize_services()`` succeeds, so ``is_ready`` implies a + fully-initialized manager and never masks ``init_error`` + - Reinitialization when model/settings change (an orchestrator swap + initializes the replacement first, swaps atomically, then cleans up the + old manager) + - ``close()``: idempotent shutdown — cancels pending init, shuts down the + init executor, and cleans up the live manager; invoked on app exit via + ``background_init()`` Example: controller = WorkflowController( @@ -92,6 +155,24 @@ def _create_workflow() -> "BaseWorkflowManager": self._workflow: "BaseWorkflowManager | None" = None self._init_task: asyncio.Task[None] | None = None self._init_error: Exception | None = None + self._closed = False + # Serializes every lifecycle transition (start_background_init, + # reinitialize/swap, close). Each of them read-modify-writes + # ``_workflow`` across awaits, so two in flight could publish two + # managers — leaking one — or publish after close(). The background + # init *task* never takes this lock (close() awaits that task while + # holding it), it checks ``_closed`` before publishing instead. + self._lifecycle_lock = asyncio.Lock() + # A manager currently being constructed in the init executor, and every + # in-flight cleanup of one the controller ended up owning — whether it + # finished building after we stopped waiting, or its release outlived + # the cancelled caller that asked for it. Both exist because neither a + # worker thread nor an unreachable manager's cleanup can be abandoned. + self._construction: "_Construction | None" = None + self._orphan_cleanups: set[asyncio.Task] = set() + # Single-flight shutdown, owned by the controller so a cancelled caller + # cannot abandon it half-done (see close()). + self._close_task: asyncio.Task[None] | None = None self.usage_tracker: "UsageTracker | None" = None # Status-bar jobs segment, published by JobMonitor; None when idle. self.jobs_status_segment: str | None = None @@ -100,17 +181,39 @@ def _create_workflow() -> "BaseWorkflowManager": def workflow(self) -> "BaseWorkflowManager": """Get the workflow manager. + Only ever a ``READY`` one. A manager can be published but unusable — + a failed in-place reinitialization leaves it uninitialized, and it is + deliberately retained so a retry can reuse its (possibly in-memory) + session store — and handing that out would look like success. + Raises: - RuntimeError: If workflow is not yet initialized + RuntimeError: If no fully initialized workflow is available. """ - if self._workflow is None: + if self._workflow is None or self.state is not WorkflowState.READY: raise RuntimeError("Workflow not initialized yet") return self._workflow + @property + def state(self) -> WorkflowState: + """Current lifecycle state (derived, never stored — cannot drift).""" + if self._closed: + return WorkflowState.CLOSED + if self._workflow is not None: + # A published manager that failed an in-place reinitialization is + # no longer usable, whatever the controller last recorded. + if getattr(self._workflow, "is_initialized", True): + return WorkflowState.READY + return WorkflowState.FAILED + if self._init_task is not None and not self._init_task.done(): + return WorkflowState.INITIALIZING + if self._init_error is not None: + return WorkflowState.FAILED + return WorkflowState.UNINITIALIZED + @property def is_ready(self) -> bool: - """Check if workflow is initialized and ready.""" - return self._workflow is not None + """True only when a fully initialized manager is published.""" + return self.state is WorkflowState.READY @property def init_error(self) -> Exception | None: @@ -125,79 +228,243 @@ def model(self) -> str | None: return self._workflow.model async def start_background_init(self) -> None: - """Start background initialization of workflow manager. + """Start background initialization of the workflow manager. Creates an async task that: 1. Creates workflow manager in ThreadPoolExecutor 2. Calls initialize_services() to preload LLM, build graph, etc. - This is non-blocking - the task runs in the background. - """ - self._init_task = asyncio.create_task(self._background_init()) + Non-blocking, single-flight and retryable: + + - already ``READY`` → no-op; + - already ``INITIALIZING`` → no-op (the in-flight attempt is joined by + ``ensure_initialized()``), so two callers can never build two managers; + - ``FAILED`` → the recorded error is cleared and a fresh attempt starts. + + A manager that is published but uninitialized (a failed in-place + reinitialization) is *revived* rather than replaced: it still owns the + session service that reinitialization preserved, and with + ``session_store='memory'`` building a replacement would silently throw + the conversation away. Only if reviving it fails is it released. - async def _background_init(self) -> None: + Raises: + RuntimeError: If the controller has been closed. + """ + async with self._lifecycle_lock: + if self._closed: + raise RuntimeError("WorkflowController is closed") + if self.state is WorkflowState.READY: + return + if self._init_task is not None and not self._init_task.done(): + return + revive, self._workflow = self._workflow, None + self._init_error = None + self._init_task = asyncio.create_task(self._background_init(revive)) + + async def _background_init( + self, revive: "BaseWorkflowManager | None" = None + ) -> None: """Initialize workflow manager in background. Creates the workflow manager and calls initialize_services() to preload LLM, build graph, and set up checkpointing. This avoids lag on the first user message. + + The manager is published to ``self._workflow`` only after + initialize_services() succeeds, so ``is_ready`` / ``ensure_initialized()`` + never report a partially-initialized or failed manager as ready — and + only if the controller has not been closed in the meantime, so shutdown + never leaves a live backend behind. A manager that fails (or is + cancelled) mid-init is cleaned up rather than leaked. + + Args: + revive: An existing, uninitialized manager to re-initialize instead + of building a new one (see ``start_background_init``). """ loop = asyncio.get_running_loop() def _create_workflow() -> "BaseWorkflowManager": return self._create_fn() + manager: "BaseWorkflowManager | None" = revive try: - logger.debug("background_init_starting") + logger.debug("background_init_starting", reviving=revive is not None) # Step 1: Create workflow manager (sync, in thread pool) - self._workflow = await loop.run_in_executor( - self._init_executor, _create_workflow - ) + if manager is None: + manager = await self._construct_manager(loop, _create_workflow) # Step 2: Initialize services (async - builds graph, loads LLM, etc.) - await self._workflow.initialize_services() - - logger.info("background_init_complete", model=self._workflow.model) - + await manager.initialize_services() + + if self._closed: + # close() ran while we were initializing; it has already taken + # its snapshot of _workflow, so publishing now would strand + # this manager. Release it instead. + await self._cleanup_manager(manager) + logger.debug("background_init_discarded_after_close") + return + + self._workflow = manager + # A recorded failure must not outlive its recovery: state, + # ensure_initialized() and the status bar all read this field. + self._init_error = None + logger.info("background_init_complete", model=manager.model) + + except asyncio.CancelledError: + if manager is not None: + await self._cleanup_manager(manager) + raise except Exception as e: self._init_error = e + if manager is not None: + await self._cleanup_manager(manager) logger.debug("background_init_failed", error=str(e)) + async def _construct_manager(self, loop, create_fn) -> "BaseWorkflowManager": + """Build a manager in the init executor, keeping ownership of the result. + + The await is **shielded**: cancelling it would cancel the asyncio future + too, and asyncio then discards whatever the (uncancellable) worker + thread returns — a fully constructed manager, unreachable and never + cleaned up. Shielded, the future survives, so shutdown can settle it or + its done-callback can release it. + """ + construction = _Construction(loop.run_in_executor(self._init_executor, create_fn)) + self._construction = construction + try: + manager = await asyncio.shield(construction.future) + except BaseException: + self._abandon_construction(construction) + raise + construction.claim() # the result is ours; nobody else may release it + if self._construction is construction: + self._construction = None + return manager + + def _spawn_cleanup(self, manager: "BaseWorkflowManager") -> "asyncio.Task | None": + """Release a manager in a task the **controller** owns. + + Cleanup is not the caller's to abandon. By the time it starts, the + manager is already unreachable — nothing else holds a reference — so a + cleanup cancelled halfway leaks its backend (session service, sandbox, + job manager) for the life of the process, with no one left to retry. + Running it in a tracked task means a cancelled caller only stops + *waiting*, and a later ``close()`` can join what it left running. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: # pragma: no cover - loop already gone + logger.warning("orphan_manager_not_cleaned") + return None + task = loop.create_task(self._cleanup_manager(manager)) + self._orphan_cleanups.add(task) + task.add_done_callback(self._orphan_cleanups.discard) + return task + + def _abandon_construction(self, construction: "_Construction") -> None: + """Arrange for a manager we no longer want to be released, once. + + Shutdown settles ``_construction`` deterministically; this callback is + the fallback for a cancellation that is never followed by ``close()``. + Whichever runs first claims the result, so it is cleaned exactly once + and never published. + """ + + def _on_done(future: "asyncio.Future") -> None: + if not construction.claim(): + return # shutdown got there first + if future.cancelled() or future.exception() is not None: + return + self._spawn_cleanup(future.result()) + + construction.future.add_done_callback(_on_done) + + async def _settle_construction(self) -> None: + """Release anything the init executor is still building. + + Cancellation-safe *through completion*: ``cancel_init()`` is public and + may be awaited directly, so a caller can be cancelled at either of the + two waits here. + + - Still waiting on the worker thread: the claim is **handed back** and + the fallback callback re-armed, so the manager is released exactly + once by whichever of that callback or a later ``close()`` gets there + first. + - Already releasing the manager: the cleanup belongs to the controller + and keeps running; it stays tracked in ``_orphan_cleanups``, which is + what a later ``close()`` joins. Cancelling here once consumed the + claim *and* aborted the cleanup, leaving a half-released manager that + nothing could finish. + """ + construction, self._construction = self._construction, None + if construction is not None and construction.claim(): + try: + manager = await asyncio.shield(construction.future) + except asyncio.CancelledError: + construction.release() + self._construction = construction + self._abandon_construction(construction) + raise + except Exception as exc: # noqa: BLE001 - shutdown must not fail + logger.debug("construction_failed_during_shutdown", error=str(exc)) + else: + cleanup = self._spawn_cleanup(manager) + if cleanup is not None: + await asyncio.shield(cleanup) + if self._orphan_cleanups: + # Shielded for the same reason: these are the controller's tasks, + # and gather() would otherwise cancel them along with its awaiter. + await asyncio.shield( + asyncio.gather(*list(self._orphan_cleanups), return_exceptions=True) + ) + async def ensure_initialized( self, ui: "ThinkingPromptSession | None" = None, ) -> bool: """Wait for background initialization to complete. + Reports readiness truthfully: it awaits the in-flight attempt (if any) + and then answers from the resulting state, so a failed initialization + is never reported as ready. + + Recovers from ``FAILED``: a fresh attempt is started (settings may have + been corrected since), so a failed reinitialization does not wedge the + session until restart. + Args: ui: Optional UI session for showing "waiting" feedback Returns: - True if initialization succeeded, False otherwise + True if a fully initialized manager is available, False otherwise """ - if self._workflow is not None: - return True - - if self._init_task is None: + if self._closed: return False - if not self._init_task.done(): + if self.state is WorkflowState.FAILED: + await self.start_background_init() + + if self._init_task is not None and not self._init_task.done(): # Show user we're waiting for initialization if ui is not None: ctx = ui.start_thinking(lambda: "Waiting for initialization...", content_format="ansi") try: - await self._init_task + await asyncio.shield(self._init_task) + except asyncio.CancelledError: + if self._init_task.cancelled(): + return False + raise finally: if ui is not None: ctx.finish(add_to_history=False) - if self._init_error: - if ui is not None: - ui.add_error(f"Initialization failed: {self._init_error}") - return False - - return self._workflow is not None + # Readiness is the state, not the presence of a past error — the two + # agreed only as long as every success remembered to clear the error. + ready = self.state is WorkflowState.READY + if not ready and self._init_error is not None and ui is not None: + ui.add_error(f"Initialization failed: {self._init_error}") + return ready def _needs_orchestrator_swap(self, new_model: str | None = None) -> bool: """Check if the current manager still matches the orchestrator setting. @@ -229,6 +496,17 @@ async def reinitialize(self, model: str | None = None) -> None: (e.g. the orchestrator setting was changed), the entire workflow manager is replaced. Otherwise, the existing manager is reinitialized in place. + Either outcome is well-defined: on success a fully initialized manager + is published; on failure the controller enters ``FAILED`` with + ``init_error`` set, so nothing can observe a READY controller wrapping + an uninitialized manager (``workflow`` refuses to hand it out). + Recovery is a fresh ``start_background_init()`` + (``ensure_initialized()`` triggers one), which revives that same + manager so its preserved sessions survive. + + Serialized against every other lifecycle transition, so two concurrent + swaps cannot both publish and leak one of the replacements. + Args: model: Optional new model to use @@ -236,35 +514,125 @@ async def reinitialize(self, model: str | None = None) -> None: RuntimeError: If workflow is not initialized Exception: If reinitialization fails """ - if self._workflow is None: - raise RuntimeError("Cannot reinitialize - workflow not initialized") + async with self._lifecycle_lock: + if self._workflow is None: + raise RuntimeError("Cannot reinitialize - workflow not initialized") - if self._needs_orchestrator_swap(model): - logger.info( - "orchestrator_swap", - old_model=self._workflow.model, - new_model=model, - ) - self._workflow = create_workflow_manager_from_settings( - agent_configs=self._agent_configs, - settings=self._settings, - app_name=self._app_name, - model=model, - ) - await self._workflow.initialize_services() - else: + if self._needs_orchestrator_swap(model): + await self._swap_orchestrator(model) + else: + await self._reinitialize_in_place(model) + + async def _swap_orchestrator(self, model: str | None) -> None: + """Replace the manager with one for the configured orchestrator. + + Initializes the replacement fully before swapping, so a failed init + leaves the working manager in place; whichever manager ends up unused + is cleaned up. The caller holds ``_lifecycle_lock``. + """ + logger.info( + "orchestrator_swap", old_model=self._workflow.model, new_model=model + ) + new_workflow = create_workflow_manager_from_settings( + agent_configs=self._agent_configs, + settings=self._settings, + app_name=self._app_name, + model=model, + ) + try: + await new_workflow.initialize_services() + except Exception: + await self._cleanup_manager(new_workflow) + raise + if self._closed: + await self._cleanup_manager(new_workflow) + raise RuntimeError("WorkflowController is closed") + old_workflow, self._workflow = self._workflow, new_workflow + self._init_error = None + await self._cleanup_manager(old_workflow) + + async def _reinitialize_in_place(self, model: str | None) -> None: + """Reinitialize the live manager. The caller holds ``_lifecycle_lock``. + + On failure the manager is **kept**, not released: it rolled itself back + to uninitialized (so ``state`` is FAILED and ``workflow`` refuses to + hand it out), but it still owns the session service its own + ``reinitialize(preserve_sessions=True)`` restored. Releasing it here + closed that service — with ``session_store='memory'`` the conversation + went with it, for a failure the user could correct and retry. + """ + try: await self._workflow.reinitialize(model=model, preserve_sessions=True) + except Exception as e: + self._init_error = e + logger.warning("reinitialize_failed", error=str(e)) + raise + # Success: drop any error recorded by an earlier attempt, so state, + # ensure_initialized() and the status bar cannot disagree. + self._init_error = None async def cancel_init(self) -> None: - """Cancel pending initialization task if still running.""" + """Cancel pending initialization task and shut the init executor down. + + Terminal for the executor: after this, ``start_background_init()`` can + no longer schedule work, so use it as part of shutdown (see ``close()``) + rather than to abort one attempt. + + Waits for any manager still under construction in the executor and + releases it — the worker thread is not cancellable, and a manager it + returns after we stop waiting would otherwise be unreachable. + """ if self._init_task and not self._init_task.done(): self._init_task.cancel() try: await self._init_task except asyncio.CancelledError: pass + await self._settle_construction() self._init_executor.shutdown(wait=False) + @staticmethod + async def _cleanup_manager(manager: "BaseWorkflowManager") -> None: + """Best-effort manager cleanup; a failing cleanup is logged, not raised.""" + try: + await manager.cleanup() + except Exception as e: + logger.warning("workflow_manager_cleanup_failed", error=str(e)) + + async def close(self) -> None: + """Release the controller: cancel pending init, clean up the manager. + + Idempotent and terminal — repeated calls join the same shutdown and the + controller stays ``CLOSED``. Invoked from application shutdown (the + ``background_init()`` context manager exit). + + The teardown runs in a task the **controller** owns, and callers join it + under a shield: whoever asked for the shutdown may be cancelled (Ctrl+C + during exit, a cancelled task group) without abandoning a manager + half-cleaned or a construction still running in the executor. A later + ``close()`` therefore waits for that work rather than returning because + the flag is already set. + """ + # Set synchronously, before any await: nothing may publish from here on. + self._closed = True + if self._close_task is None: + self._close_task = asyncio.create_task(self._close_once()) + await asyncio.shield(self._close_task) + + async def _close_once(self) -> None: + """The actual teardown. Runs exactly once; owned by the controller. + + Takes the lifecycle lock, so it waits for an in-flight reinitialization + or swap instead of racing it, and nothing can publish afterwards: a + background init that finishes later sees ``_closed`` and releases its + manager rather than installing it. + """ + async with self._lifecycle_lock: + await self.cancel_init() + manager, self._workflow = self._workflow, None + if manager is not None: + await self._cleanup_manager(manager) + def update_status_bar(self, ui: "ThinkingPromptSession") -> None: """Update UI status bar with current workflow status. @@ -315,12 +683,12 @@ async def wait_and_update() -> None: try: yield finally: - # Cancel init task if still running - await self.cancel_init() - # Also cancel status update task if still running + # Cancel the status-update task first so it isn't woken by the + # init cancellation below, then release everything we own. if not update_task.done(): update_task.cancel() try: await update_task except asyncio.CancelledError: pass + await self.close() diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index 8ca5138..625bddf 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -164,6 +164,10 @@ def __init__( self._adk_config_path = adk_config_path self._session_service: BaseSessionService | None = None + # True while reinitialize(preserve_sessions=True) is in flight: the + # live session service is being carried across and must survive both a + # successful rebuild and a rollback. + self._session_service_pinned = False self._root_agent: Agent | None = None self._runner: Runner | None = None @@ -243,15 +247,38 @@ async def generate_simple(self, prompt: str, max_tokens: int = 500) -> str: async def cleanup(self) -> None: """Clean up workflow manager resources. - Releases resources and resets state. Call this before - shutting down or when reinitializing with new settings. + Releases resources and resets state. Call this before shutting down or + when reinitializing with new settings. Idempotent: references are + detached before being closed, so a repeat call is a no-op. Takes the + lifecycle and turn locks, so it cannot tear the runner down while a + turn is streaming. + """ + async with self._lifecycle_lock: + async with self._turn_lock: + await self._release_resources() + + async def _release_resources(self, keep_session_service: bool = False) -> None: + """Release owned resources; never raises. + + The session service is *closed*, not merely dropped — the durable + ``DatabaseSessionService`` holds a SQLAlchemy engine whose connection + pool leaks otherwise. ``reinitialize(preserve_sessions=True)`` keeps it + instead, and then reuses that same instance rather than building a + replacement that would be discarded unclosed. + + Args: + keep_session_service: Retain (do not close) the session service. """ logger.debug("cleaning_up_workflow_manager") - # Clear runner and agents + # Detach first so a concurrent/repeat cleanup can't double-close. A + # pinned service is being carried across a reinitialize, so a rollback + # inside that window must not close it either. + session_service = None + if not (keep_session_service or self._session_service_pinned): + session_service, self._session_service = self._session_service, None self._runner = None self._root_agent = None - self._session_service = None self._initialized = False # Clear LLM logging plugin @@ -259,7 +286,9 @@ async def cleanup(self) -> None: self._llm_logging_plugin.clear() self._llm_logging_plugin = None - # Clean up managers (sandbox, etc.) + await self._aclose_owned(session_service, "session_service") + + # Clean up managers (sandbox, jobs, etc.) self._cleanup_managers() logger.info("workflow_manager_cleaned_up") @@ -271,13 +300,24 @@ async def reinitialize( ) -> None: """Reinitialize the workflow manager with new configuration. - Use this method when settings change (e.g., model switch) to - properly recreate agents and runners with the new configuration. + Use this method when settings change (e.g. a model switch) to recreate + agents and runners with the new configuration. + + Transactional: on failure the manager ends up *uninitialized* ( + ``is_initialized`` False) with nothing left allocated, rather than + holding a half-built runner that would answer as if it were ready. The + caller (``WorkflowController``) turns that into a FAILED state. A + preserved session service survives both outcomes, so a retry can still + continue the same conversations. Args: model: Optional new model to use. If None, re-resolves from settings. - preserve_sessions: If True, keeps existing session data (default). - If False, creates fresh session service. + preserve_sessions: If True, keeps the existing session service + (default) and reuses it for the new runner. If False, the old + one is closed and a fresh one created. + + Raises: + Exception: Whatever initialization raised, after rollback. """ logger.info( "reinitializing_workflow_manager", @@ -285,29 +325,24 @@ async def reinitialize( preserve_sessions=preserve_sessions, ) - # Store session service if preserving - old_session_service = self._session_service if preserve_sessions else None - - # Clean up current state - await self.cleanup() - - # Update model - self._reset_model(model) - - # Reinitialize services - await self.initialize_services() - - # Restore session service if preserving - if old_session_service is not None and preserve_sessions: - self._session_service = old_session_service - # Update runner with preserved session service - if self._runner and self._root_agent: - self._runner = Runner( - app_name=self.app_name, - agent=self._root_agent, - session_service=self._session_service, - plugins=self._init_plugins(), - ) + async with self._lifecycle_lock: + async with self._turn_lock: + preserved = self._session_service if preserve_sessions else None + await self._release_resources(keep_session_service=preserve_sessions) + self._reset_model(model) + self._session_service_pinned = preserve_sessions + try: + # _do_initialize reuses self._session_service when set, so + # no replacement service is built for a preserved one. + await self._initialize_locked() + except BaseException: + # _initialize_locked already rolled back what it created; + # restore the preserved service so a retry can use it. + if preserved is not None: + self._session_service = preserved + raise + finally: + self._session_service_pinned = False logger.info( "workflow_manager_reinitialized", @@ -765,9 +800,13 @@ async def _do_initialize(self) -> None: """ADK-specific initialization: session service, agents, runner.""" logger.info("initializing_services", app_name=self.app_name) - # Create session service (durable DatabaseSessionService by default; - # InMemory only when session_store='memory'). - self._session_service = self._make_session_service() + # Create the session service (durable DatabaseSessionService by + # default; InMemory only when session_store='memory') — unless one is + # already held, which is how reinitialize(preserve_sessions=True) + # carries live conversations across without building a replacement + # that would then be discarded unclosed. + if self._session_service is None: + self._session_service = self._make_session_service() # Create agent hierarchy — natively from an ADK config, or from configs. if self._adk_config_path: @@ -805,14 +844,30 @@ def _validate_agent_graph(self) -> None: validate_agent_graph(self._agent_configs, backend=self.backend_type) async def _ensure_initialized(self) -> None: - """Ensure services are initialized before processing.""" + """Ensure services are initialized before processing. + + Called *before* the turn lock is taken, so a turn never waits on the + lifecycle lock while holding the turn lock (that ordering is what keeps + cleanup/reinitialize from deadlocking against a running turn). Because + of that ordering a cleanup can still land in between, which is what + ``_turn_admission`` re-checks with ``_backend_ready``. + """ if not self._initialized: await self.initialize_services() - if not self._runner or not self._session_service or not self._root_agent: + if not self._backend_ready(): raise RuntimeError( "Workflow Manager failed to initialize. Check API keys and configuration." ) + def _backend_ready(self) -> bool: + """True when the runner, session service and agent tree are all live.""" + return bool( + self._initialized + and self._runner is not None + and self._session_service is not None + and self._root_agent is not None + ) + # ------------------------------------------------------------------------- # Session handling (inlined from SessionHandler) # ------------------------------------------------------------------------- @@ -864,6 +919,12 @@ async def process( This method sets up a settings context so that all tools called during processing will use this manager's settings instance. + Holds the manager's turn lock for the whole stream: the ADK plugins' + event buffers are manager-scoped, so overlapping turns would drain each + other's events. Overlapping callers queue; the lock is released on + cancellation. Admission also re-verifies that the backend is still live + (a cleanup can land between initialization and the turn lock). + Args: message: User message user_id: User identifier @@ -872,34 +933,33 @@ async def process( Yields: WorkflowEvent objects representing workflow output """ - await self._ensure_initialized() - current_session_id = session_id or self.session_id bind_context(session_id=current_session_id, user_id=user_id) logger.info("processing_message", message_length=len(message)) - # Sync event processor model (may have been lazily resolved) - self._event_processor.model = self.model + async with self._turn_admission(): + # Sync event processor model (may have been lazily resolved) + self._event_processor.model = self.model - # Context setup - with self._workflow_context(session_id=current_session_id, user_id=user_id): - # Session handling - await self._get_or_create_session(user_id, current_session_id) + # Context setup + with self._workflow_context(session_id=current_session_id, user_id=user_id): + # Session handling + await self._get_or_create_session(user_id, current_session_id) - # Create message - new_message = types.Content( - role="user", - parts=[types.Part.from_text(text=message)], - ) + # Create message + new_message = types.Content( + role="user", + parts=[types.Part.from_text(text=message)], + ) - async for event in self._run_and_stream( - session_id=current_session_id, - user_id=user_id, - new_message=new_message, - run_config=self._build_run_config(), - ): - yield event + async for event in self._run_and_stream( + session_id=current_session_id, + user_id=user_id, + new_message=new_message, + run_config=self._build_run_config(), + ): + yield event def _build_run_config(self): """Build a RunConfig with context-window compression if enabled.""" @@ -1013,8 +1073,6 @@ async def resume_with_job_result( record: The terminal ``JobRecord`` to resume from. result: The job's result; fetched from the JobManager if omitted. """ - await self._ensure_initialized() - session_id = record.session_id user_id = record.user_id if not session_id or not user_id or not record.call_id: @@ -1028,43 +1086,47 @@ async def resume_with_job_result( return bind_context(session_id=session_id, user_id=user_id) - self._event_processor.model = self.model - with self._workflow_context(session_id=session_id, user_id=user_id): - # The pending call lives in the existing session; don't create a new - # empty one (that would have no call to answer). - session = await self._session_service.get_session( - app_name=self.app_name, user_id=user_id, session_id=session_id, - ) - if session is None: - logger.warning("job_resume_session_missing", job_id=record.job_id, - session_id=session_id) - return - - if result is None: - jm = self._services.get(JOB_MANAGER) - if jm is not None: - result = jm.result(record.job_id) - - function_response = types.FunctionResponse( - id=record.call_id, - name=record.call_name or record.tool, - response=self._job_result_payload(record, result), - ) - new_message = types.Content( - role="user", - parts=[types.Part(function_response=function_response)], - ) + # Same admission as process(): a resume is a turn, must not interleave + # with a user turn, and must not run against a released backend. + async with self._turn_admission(): + self._event_processor.model = self.model - logger.info("job_resume_started", job_id=record.job_id, - call_id=record.call_id, state=record.state.value) - async for event in self._run_and_stream( - session_id=session_id, - user_id=user_id, - new_message=new_message, - run_config=self._build_run_config(), - ): - yield event + with self._workflow_context(session_id=session_id, user_id=user_id): + # The pending call lives in the existing session; don't create a + # new empty one (that would have no call to answer). + session = await self._session_service.get_session( + app_name=self.app_name, user_id=user_id, session_id=session_id, + ) + if session is None: + logger.warning("job_resume_session_missing", job_id=record.job_id, + session_id=session_id) + return + + if result is None: + jm = self._services.get(JOB_MANAGER) + if jm is not None: + result = jm.result(record.job_id) + + function_response = types.FunctionResponse( + id=record.call_id, + name=record.call_name or record.tool, + response=self._job_result_payload(record, result), + ) + new_message = types.Content( + role="user", + parts=[types.Part(function_response=function_response)], + ) + + logger.info("job_resume_started", job_id=record.job_id, + call_id=record.call_id, state=record.state.value) + async for event in self._run_and_stream( + session_id=session_id, + user_id=user_id, + new_message=new_message, + run_config=self._build_run_config(), + ): + yield event @staticmethod def _job_result_payload(record, result: Any) -> dict: diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index 7469f44..344542e 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -13,8 +13,11 @@ from __future__ import annotations +import asyncio import contextlib +import inspect from abc import ABC, abstractmethod +from contextvars import ContextVar, Token from typing import Any, AsyncGenerator, Awaitable, Callable, Iterator, TYPE_CHECKING from agentic_cli.workflow.events import WorkflowEvent, UserInputRequest @@ -61,6 +64,21 @@ class BaseWorkflowManager(ABC): - Event streaming - User input request/response flow + Concurrency contract: + A manager runs **one turn at a time**. ``process()`` and + ``resume_with_job_result()`` serialize on an internal turn lock, so + overlapping callers queue rather than interleave — the backend's + per-invocation event buffers are manager-scoped, and interleaving would + let one invocation drain another's events. Run separate managers for + genuine parallelism. (The HITL input callback is context-local, so it + does not depend on that serialization; see ``set_input_callback``.) + + Lifecycle mutation (``initialize_services``/``reinitialize``/ + ``cleanup``) additionally takes the turn lock, so the backend is never + torn down under a running generator. Both locks are released on + cancellation. The active session/user identity remains a ContextVar, + so it stays correct for nested and task-spawned work. + Example: class CustomWorkflowManager(BaseWorkflowManager): async def initialize_services(self) -> None: @@ -96,20 +114,29 @@ def __init__( self._settings = settings or get_settings() self._app_name = app_name or self._settings.app_name self._initialized = False + # --- Concurrency contract (see the class docstring) --- + # _lifecycle_lock serializes initialize/reinitialize/cleanup: background + # init and a first user message may call them concurrently, and the + # _initialized guard alone is check-then-act. + # _turn_lock serializes turns, and is taken by lifecycle mutation so it + # cannot tear the backend down under a running generator. + # Lock order is lifecycle → turn: a turn releases the lifecycle lock + # (inside _ensure_initialized) *before* taking the turn lock, so the + # two can never deadlock. + self._lifecycle_lock = asyncio.Lock() + self._turn_lock = asyncio.Lock() self._on_event = on_event # Model resolution (lazy) self._model: str | None = model self._model_resolved: bool = model is not None - # User input handling (callback-only) - self._user_input_callback: Callable[[UserInputRequest], Awaitable[str]] | None = None - - # Active turn — set per process() call via _workflow_context(); read by - # JobManager to associate a long-running job with the session/user that - # launched it (phase 2 push/resume). None when no turn is in flight. - self._active_session_id: str | None = None - self._active_user_id: str | None = None + # User input handling (callback-only), context-local — see + # set_input_callback(). One ContextVar per manager; there are only ever + # a handful of managers in a process. + self._user_input_callback: ContextVar[ + Callable[[UserInputRequest], Awaitable[str]] | None + ] = ContextVar(f"agentic_cli_input_callback_{id(self):x}", default=None) # Model registry self._model_registry = ModelRegistry() @@ -128,13 +155,42 @@ def __init__( def set_input_callback( self, callback: Callable[[UserInputRequest], Awaitable[str]] - ) -> None: - """Register a callback for handling user input requests from tools.""" - self._user_input_callback = callback + ) -> "Token | None": + """Register a callback for handling user input requests from tools. + + **Context-local**, not manager-global. The turn lock serialises + ``process()``, but callbacks are installed *before* it: with a single + manager attribute, a second consumer that installed its callback while + a turn was already running would answer that turn's prompt, and the + first consumer's ``clear_input_callback()`` would then unregister the + second's. A task started after this call inherits the value (the + context is copied at ``create_task`` time), which is exactly the + consumer → turn relationship. + + Returns: + The ContextVar token, for an exact ``clear_input_callback(token)``. + Callers may ignore it. + """ + return self._user_input_callback.set(callback) - def clear_input_callback(self) -> None: - """Remove the registered user input callback.""" - self._user_input_callback = None + def clear_input_callback(self, token: "Token | None" = None) -> None: + """Remove the user input callback *this context* registered. + + Args: + token: The token ``set_input_callback()`` returned. Passing it + restores whatever was installed before; without it the value is + cleared for this context only. Either way another consumer's + callback is untouched. + """ + if token is not None: + try: + self._user_input_callback.reset(token) + return + except ValueError: + # Token from a different context (the caller crossed tasks); + # fall through and clear this context's value instead. + pass + self._user_input_callback.set(None) @property def agent_configs(self) -> list[AgentConfig]: @@ -383,14 +439,59 @@ def _detect_required_managers(self) -> set[str]: def _ensure_managers_initialized(self) -> None: """Create and publish the services detected from tool metadata. - Called during initialize_services() to lazily create only the - managers that are actually needed by the configured tools. - Populates ``self._services`` which is exposed to tools via - the service registry ContextVar. + Synchronous convenience wrapper around + :meth:`_build_services`/:meth:`_publish_services`; initialization uses + the transactional path in :meth:`_construct_services` instead. """ - s = self._services + self._publish_services(self._build_services(frozenset(self._services))) + + def _publish_services(self, built: dict[str, Any]) -> None: + """Adopt constructed services without displacing anything already live.""" + for key, service in built.items(): + self._services.setdefault(key, service) + + def _build_services( + self, existing: frozenset[str] = frozenset() + ) -> dict[str, Any]: + """Construct the required services into a **fresh** dict. + + Pure construction: it never touches ``self._services``. Constructors + here can load heavy dependencies (the sentence-transformers model inside + ``EmbeddingService``), so this runs on a worker thread — and cancelling + the coroutine that awaits it does not stop that thread. Writing results + straight into the manager therefore published services *after* a + rolled-back or cleaned-up initialization, leaking whatever the thread + had built. The caller publishes, and only while it still owns the + attempt (see :meth:`_construct_services`). + + Transactional in itself: if a later constructor raises, everything this + call already built is released before the error propagates. Nothing has + been published at that point, so nobody else could ever close it — an + abandoned SandboxManager or JobManager would keep its pool alive for + the life of the process. + + Args: + existing: Service keys already published; those are not rebuilt. + + Returns: + The newly constructed services, keyed by service key. - if "memory_store" in self._required_managers and MEMORY_STORE not in s: + Raises: + Exception: Whatever a service constructor raised, after rollback. + """ + s: dict[str, Any] = {} + try: + self._build_services_into(s, existing) + except BaseException: + self._close_services(s) + raise + return s + + def _build_services_into( + self, s: dict[str, Any], existing: frozenset[str] + ) -> None: + """Construct the required services into ``s``. See :meth:`_build_services`.""" + if "memory_store" in self._required_managers and MEMORY_STORE not in existing: from agentic_cli.tools.memory_tools import MemoryStore embedding_service = None @@ -408,7 +509,7 @@ def _ensure_managers_initialized(self) -> None: s[MEMORY_STORE] = MemoryStore(self._settings, embedding_service=embedding_service) - if "kb_manager" in self._required_managers and KB_MANAGER not in s: + if "kb_manager" in self._required_managers and KB_MANAGER not in existing: from pathlib import Path from agentic_cli.knowledge_base import KnowledgeBaseManager @@ -431,14 +532,14 @@ def _ensure_managers_initialized(self) -> None: else: s[USER_KB_MANAGER] = s[KB_MANAGER] - if "llm_summarizer" in self._required_managers and LLM_SUMMARIZER not in s: + if "llm_summarizer" in self._required_managers and LLM_SUMMARIZER not in existing: s[LLM_SUMMARIZER] = self - if "sandbox_manager" in self._required_managers and SANDBOX_MANAGER not in s: + if "sandbox_manager" in self._required_managers and SANDBOX_MANAGER not in existing: from agentic_cli.tools.sandbox.manager import SandboxManager s[SANDBOX_MANAGER] = SandboxManager(self._settings) - if "job_manager" in self._required_managers and JOB_MANAGER not in s: + if "job_manager" in self._required_managers and JOB_MANAGER not in existing: from pathlib import Path from agentic_cli.tools.jobs import JobManager @@ -450,12 +551,12 @@ def _ensure_managers_initialized(self) -> None: max_concurrent=getattr(self._settings, "max_concurrent_jobs", 4), ) - if "arxiv_source" in self._required_managers and ARXIV_SOURCE not in s: + if "arxiv_source" in self._required_managers and ARXIV_SOURCE not in existing: from agentic_cli.tools.arxiv_source import ArxivSearchSource s[ARXIV_SOURCE] = ArxivSearchSource() # Always construct the PermissionEngine (all agents may need it) - if PERMISSION_ENGINE not in s: + if PERMISSION_ENGINE not in existing: from pathlib import Path from agentic_cli.workflow.permissions import PermissionContext, PermissionEngine ctx = PermissionContext( @@ -470,6 +571,36 @@ def _ensure_managers_initialized(self) -> None: # Always ensure workflow reference is available s[WORKFLOW] = self + async def _construct_services(self) -> None: + """Build services off the event loop and publish them transactionally. + + The build runs on a worker thread that cancellation cannot interrupt, + so the result is published only if this attempt is still the one that + owns initialization. If the await is cancelled, whatever the thread + goes on to build is *released* rather than published — otherwise a + rolled-back initialization would leave a live sandbox or job manager + behind that nothing would ever close. + """ + build = asyncio.ensure_future( + asyncio.to_thread(self._build_services, frozenset(self._services)) + ) + try: + built = await asyncio.shield(build) + except BaseException: + build.add_done_callback(self._discard_built_services) + raise + self._publish_services(built) + + def _discard_built_services(self, build: "asyncio.Future[dict[str, Any]]") -> None: + """Release services constructed for an attempt that no longer owns init.""" + if build.cancelled() or build.exception() is not None: + return + built = build.result() + if not built: + return + logger.warning("services_discarded_after_rollback", services=sorted(built)) + self._close_services(built) + async def summarize(self, content: str, prompt: str) -> str: """Summarize content using the configured LLM. Args: @@ -612,12 +743,65 @@ def _workflow_context( registry_token.var.reset(registry_token) settings_token.var.reset(settings_token) + async def _aclose_owned(self, resource: Any, label: str) -> None: + """Close one resource this manager owns; awaits an async close. + + The close contract is duck-typed on ``aclose()``/``close()`` (ADK's + ``DatabaseSessionService`` exposes an async ``close()``; the in-memory + one exposes none) and must be idempotent: callers null out their + reference first, so a second cleanup passes ``None`` and does nothing. + Never raises — a failing close must not block shutdown. + + Args: + resource: The owned resource, or None. + label: Name used in the failure log. + """ + if resource is None: + return + closer = getattr(resource, "aclose", None) or getattr(resource, "close", None) + if closer is None: + return + try: + result = closer() + if inspect.isawaitable(result): + await result + except Exception as exc: # noqa: BLE001 - shutdown must not fail + logger.warning("resource_close_failed", resource=label, error=str(exc)) + + # Services that own OS resources, and the sync method that releases them. + _SYNC_SERVICE_CLOSERS = ( + (SANDBOX_MANAGER, "cleanup"), + (JOB_MANAGER, "close"), + ) + + @classmethod + def _close_services(cls, services: dict[str, Any]) -> None: + """Release the owned resources in a service mapping. Never raises. + + Each closer is isolated — one raising must not leave the rest open. + Used both for the live registry and for services a rolled-back + initialization constructed but never published. + """ + for key, method in cls._SYNC_SERVICE_CLOSERS: + service = services.get(key) + if service is None: + continue + try: + getattr(service, method)() + except Exception as exc: # noqa: BLE001 - shutdown must not fail + logger.warning("resource_close_failed", resource=key, error=str(exc)) + def _cleanup_managers(self) -> None: - """Clean up all manager resources (call from subclass cleanup).""" - sandbox = self._services.get(SANDBOX_MANAGER) - if sandbox is not None: - sandbox.cleanup() - self._services = {} + """Release the synchronous resources this manager owns. + + Only services this manager created are released (see + ``_build_services``). Idempotent: the registry is emptied, so a second + call finds nothing. The registry is cleared regardless of failures. + """ + try: + self._close_services(self._services) + finally: + self._services = {} @property @abstractmethod @@ -649,11 +833,22 @@ def model(self) -> str: async def initialize_services(self, validate: bool = True) -> None: """Initialize backend services asynchronously. + + Concurrency-safe, idempotent and transactional: concurrent callers + serialize on the lifecycle lock, late arrivals see ``_initialized`` and + return, and a failed attempt releases whatever it had allocated instead + of leaving the manager half-built. + Args: validate: If True, validate settings before initialization. Raises: SettingsValidationError: If settings validation fails. """ + async with self._lifecycle_lock: + await self._initialize_locked(validate=validate) + + async def _initialize_locked(self, validate: bool = True) -> None: + """Initialization body. The caller must hold ``_lifecycle_lock``.""" if self._initialized: return @@ -677,14 +872,16 @@ async def initialize_services(self, validate: bool = True) -> None: # Create services BEFORE backend init so _build_tools() can # produce factory-bound tools during agent/graph creation. - # Offloaded to a worker thread because constructors here may - # load heavy dependencies (e.g. the sentence-transformers model - # inside EmbeddingService) that would otherwise block the event - # loop — which keeps the prompt unresponsive at startup. - import asyncio as _asyncio - - await _asyncio.to_thread(self._ensure_managers_initialized) - await self._do_initialize() + # Construction is offloaded to a worker thread (heavy constructors) + # and published transactionally — see _construct_services. + try: + await self._construct_services() + await self._do_initialize() + except BaseException: + # Roll back: services (and any backend resource the partial + # _do_initialize created) must not outlive the failed attempt. + await self._release_resources() + raise self._initialized = True # Label for this manager's own model in validation errors. @@ -729,6 +926,64 @@ def _validate_agent_graph(self) -> None: """ return None + async def _ensure_initialized(self) -> None: + """Initialize on demand. Backends override to add readiness checks.""" + if not self._initialized: + await self.initialize_services() + + def _backend_ready(self) -> bool: + """Whether the backend resources a turn needs are live right now. + + Backends override to check their own handles (ADK: runner, session + service, root agent). Used by :meth:`_turn_admission` to detect a + cleanup that landed between a turn's initialization and its admission. + """ + return self._initialized + + @contextlib.asynccontextmanager + async def _turn_admission(self) -> "AsyncGenerator[None, None]": + """Hold the turn lock with a *live* backend behind it. + + Initialization takes the lifecycle lock, so a turn must initialize + **before** taking the turn lock — that ordering is what stops + cleanup (lifecycle → turn) from deadlocking against a running turn. + It also leaves a window: a cleanup already queued on the turn lock runs + first and releases everything the turn just initialized, and the turn + was then admitted to a torn-down backend (a ``None`` runner, surfacing + as an ``AttributeError`` deep inside ADK). + + Readiness is therefore re-checked *while holding the turn lock*. If the + backend was released underneath, the lock is dropped and initialization + retried once — anything worse fails cleanly rather than running against + released resources. + + Raises: + RuntimeError: If the backend cannot be made ready. + """ + for attempt in (1, 2): + await self._ensure_initialized() + await self._turn_lock.acquire() + if self._backend_ready(): + try: + yield + finally: + self._turn_lock.release() + return + self._turn_lock.release() + logger.info("turn_admission_retry", attempt=attempt) + raise RuntimeError( + f"{type(self).__name__} was released while this turn waited for " + "admission and could not be reinitialized. Retry the request." + ) + + async def _release_resources(self) -> None: + """Release everything this manager owns. Idempotent, never raises. + + Backends override to add their own resources (ADK closes the session + service); the base releases the service registry. + """ + self._cleanup_managers() + @abstractmethod async def _do_initialize(self) -> None: """Backend-specific initialization (create agents/graph). @@ -807,7 +1062,9 @@ async def request_user_input(self, request: UserInputRequest) -> str: Called by tools that need user interaction. Requires ``set_input_callback()`` to be set by the consumer (e.g. - MessageProcessor) before any tool invokes this method. + MessageProcessor) before any tool invokes this method. The callback is + resolved from the *current context*, so a tool always reaches the + consumer that started its turn. Args: request: The user input request. @@ -824,13 +1081,14 @@ async def request_user_input(self, request: UserInputRequest) -> str: tool_name=request.tool_name, ) - if self._user_input_callback is None: + callback = self._user_input_callback.get() + if callback is None: raise RuntimeError( "No user input callback registered. " "Call set_input_callback() before invoking tools that require user input." ) - return await self._user_input_callback(request) + return await callback(request) # Async context manager support diff --git a/tests/integration/test_adk_integration.py b/tests/integration/test_adk_integration.py index ea76bc3..1c4ccd2 100644 --- a/tests/integration/test_adk_integration.py +++ b/tests/integration/test_adk_integration.py @@ -667,10 +667,15 @@ async def mock_process(**kwargs): class TestUserInputCallback: - """Tests for the direct _user_input_callback path in request_user_input.""" + """Tests for the registered-callback path in request_user_input. + + The callback lives in a per-manager ContextVar (it is context-local, so two + consumers cannot capture each other's prompts), hence set_input_callback() + rather than an attribute assignment. + """ async def test_callback_invoked_when_set(self, mock_settings, simple_agent_config): - """When _user_input_callback is set, request_user_input calls it directly.""" + """With a callback registered, request_user_input calls it directly.""" from agentic_cli.workflow.events import UserInputRequest, InputType manager = _create_manager(mock_settings, simple_agent_config) @@ -681,7 +686,7 @@ async def fake_callback(request: UserInputRequest) -> str: captured_requests.append(request) return "user answer" - manager._user_input_callback = fake_callback + manager.set_input_callback(fake_callback) request = UserInputRequest( request_id="req-1", @@ -700,7 +705,7 @@ async def test_request_user_input_without_callback_raises(self, mock_settings, s from agentic_cli.workflow.events import UserInputRequest, InputType manager = _create_manager(mock_settings, simple_agent_config) - assert manager._user_input_callback is None + assert manager._user_input_callback.get() is None request = UserInputRequest( request_id="req-2", diff --git a/tests/test_input_callback.py b/tests/test_input_callback.py index 38b4fc2..a397758 100644 --- a/tests/test_input_callback.py +++ b/tests/test_input_callback.py @@ -1,26 +1,61 @@ -"""Tests for input callback public API on BaseWorkflowManager.""" +"""Tests for input callback public API on BaseWorkflowManager. -from unittest.mock import AsyncMock, MagicMock, patch +The callback is held in a per-manager ContextVar rather than a plain attribute +(see ``set_input_callback``), so these assert the observable behaviour — +whether ``request_user_input`` reaches the callback — rather than the storage. +""" + +from contextvars import ContextVar +from unittest.mock import AsyncMock, patch + +import pytest + +from agentic_cli.workflow.events import UserInputRequest + + +def _manager(): + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + with patch.object(BaseWorkflowManager, "__abstractmethods__", set()): + manager = BaseWorkflowManager.__new__(BaseWorkflowManager) + manager._user_input_callback = ContextVar("test_input_callback", default=None) + return manager + + +def _request() -> UserInputRequest: + return UserInputRequest(request_id="r", tool_name="t", prompt="p") class TestInputCallbackAPI: - def test_set_input_callback(self): - from agentic_cli.workflow.base_manager import BaseWorkflowManager + async def test_set_input_callback(self): + manager = _manager() + callback = AsyncMock(return_value="answer") + + manager.set_input_callback(callback) + + assert await manager.request_user_input(_request()) == "answer" + callback.assert_awaited_once() + + async def test_clear_input_callback(self): + manager = _manager() + manager.set_input_callback(AsyncMock(return_value="answer")) + + manager.clear_input_callback() - with patch.object(BaseWorkflowManager, '__abstractmethods__', set()): - manager = BaseWorkflowManager.__new__(BaseWorkflowManager) - manager._user_input_callback = None + with pytest.raises(RuntimeError, match="No user input callback"): + await manager.request_user_input(_request()) - callback = AsyncMock() - manager.set_input_callback(callback) - assert manager._user_input_callback is callback + async def test_clear_with_token_restores_the_previous_callback(self): + manager = _manager() + manager.set_input_callback(AsyncMock(return_value="outer")) + token = manager.set_input_callback(AsyncMock(return_value="inner")) - def test_clear_input_callback(self): - from agentic_cli.workflow.base_manager import BaseWorkflowManager + assert await manager.request_user_input(_request()) == "inner" - with patch.object(BaseWorkflowManager, '__abstractmethods__', set()): - manager = BaseWorkflowManager.__new__(BaseWorkflowManager) - manager._user_input_callback = AsyncMock() + manager.clear_input_callback(token) + assert await manager.request_user_input(_request()) == "outer" - manager.clear_input_callback() - assert manager._user_input_callback is None + async def test_no_callback_raises(self): + manager = _manager() + with pytest.raises(RuntimeError, match="No user input callback"): + await manager.request_user_input(_request()) diff --git a/tests/test_workflow.py b/tests/test_workflow.py index 72b75db..c462a6d 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -463,7 +463,7 @@ async def test_reinitialize_with_new_model(self, mock_settings, agent_configs): from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager with patch.object( - GoogleADKWorkflowManager, "initialize_services", new_callable=AsyncMock + GoogleADKWorkflowManager, "_initialize_locked", new_callable=AsyncMock ): manager = GoogleADKWorkflowManager( agent_configs=agent_configs, @@ -486,7 +486,7 @@ async def test_reinitialize_resolves_model_from_settings( mock_settings.get_model.return_value = "resolved-model" with patch.object( - GoogleADKWorkflowManager, "initialize_services", new_callable=AsyncMock + GoogleADKWorkflowManager, "_initialize_locked", new_callable=AsyncMock ): manager = GoogleADKWorkflowManager( agent_configs=agent_configs, @@ -508,7 +508,7 @@ async def test_reinitialize_preserves_sessions_by_default( from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager with patch.object( - GoogleADKWorkflowManager, "initialize_services", new_callable=AsyncMock + GoogleADKWorkflowManager, "_initialize_locked", new_callable=AsyncMock ): manager = GoogleADKWorkflowManager( agent_configs=agent_configs, @@ -534,7 +534,7 @@ async def test_reinitialize_can_discard_sessions( from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager with patch.object( - GoogleADKWorkflowManager, "initialize_services", new_callable=AsyncMock + GoogleADKWorkflowManager, "_initialize_locked", new_callable=AsyncMock ): manager = GoogleADKWorkflowManager( agent_configs=agent_configs, diff --git a/tests/test_workflow_controller.py b/tests/test_workflow_controller.py index 9005ed4..c2423b2 100644 --- a/tests/test_workflow_controller.py +++ b/tests/test_workflow_controller.py @@ -7,6 +7,7 @@ does. """ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -313,3 +314,1084 @@ async def test_reinitialize_migrates_stale_langgraph_to_adk(self): old_workflow.reinitialize.assert_not_called() new_workflow.initialize_services.assert_awaited_once() assert controller._workflow is new_workflow +# --- Lifecycle: readiness, atomic swap, close() --- + + +def _make_lifecycle_controller(orchestrator=OrchestratorType.ADK): + configs = [AgentConfig(name="test", prompt="Test")] + return WorkflowController(configs, _make_settings(orchestrator=orchestrator)) + + +def _blocked_init_workflow(): + """Fake manager whose initialize_services blocks until released.""" + import asyncio + + wf = _FakeADKWorkflow() + started, release = asyncio.Event(), asyncio.Event() + + async def _slow_init(): + started.set() + await release.wait() + + wf.initialize_services = _slow_init + return wf, started, release + + +class TestControllerLifecycle: + """is_ready / ensure_initialized must reflect completed service init.""" + + async def test_not_ready_until_services_initialized(self): + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + controller._create_fn = lambda: wf + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + + assert controller.is_ready is False + with pytest.raises(RuntimeError): + controller.workflow + + release.set() + await controller._init_task + assert controller.is_ready is True + assert controller.workflow is wf + + async def test_failed_service_init_is_not_ready_and_cleans_up(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.initialize_services = AsyncMock(side_effect=RuntimeError("boom")) + controller._create_fn = lambda: wf + + await controller.start_background_init() + await controller._init_task + + assert controller.is_ready is False + assert isinstance(controller.init_error, RuntimeError) + wf.cleanup.assert_awaited_once() + + async def test_ensure_initialized_retries_a_failed_attempt(self): + """A FAILED controller retries, so a corrected setting can recover it.""" + controller = _make_lifecycle_controller() + failing = _FakeADKWorkflow() + failing.initialize_services = AsyncMock(side_effect=RuntimeError("boom")) + good = _FakeADKWorkflow() + managers = [failing, good] + controller._create_fn = lambda: managers.pop(0) + + await controller.start_background_init() + await controller._init_task + assert controller.is_ready is False + + assert await controller.ensure_initialized() is True + assert controller.workflow is good + assert controller.init_error is None + + async def test_ensure_initialized_waits_for_inflight_services(self): + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + controller._create_fn = lambda: wf + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + + ensure_task = asyncio.create_task(controller.ensure_initialized()) + await asyncio.sleep(0.05) + assert not ensure_task.done() + + release.set() + assert await ensure_task is True + + +class TestOrchestratorSwapLifecycle: + """Swap must initialize the replacement first, then replace atomically.""" + + def _controller_needing_swap(self): + # Live ADK manager while settings demand LangGraph → swap required + controller = _make_lifecycle_controller( + orchestrator=OrchestratorType.LANGGRAPH + ) + old = _FakeADKWorkflow() + controller._workflow = old + return controller, old + + async def test_swap_failure_keeps_old_manager(self): + controller, old = self._controller_needing_swap() + new = _FakeLangGraphWorkflow() + new.initialize_services = AsyncMock(side_effect=RuntimeError("init failed")) + + with patch( + "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings", + return_value=new, + ): + with pytest.raises(RuntimeError, match="init failed"): + await controller.reinitialize() + + assert controller._workflow is old + old.cleanup.assert_not_awaited() + new.cleanup.assert_awaited_once() + + async def test_swap_success_replaces_then_cleans_old(self): + controller, old = self._controller_needing_swap() + new = _FakeLangGraphWorkflow() + + with patch( + "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings", + return_value=new, + ): + await controller.reinitialize() + + assert controller._workflow is new + old.cleanup.assert_awaited_once() + + +class TestControllerClose: + """close() releases the manager and is safe to call repeatedly.""" + + async def test_close_cleans_manager_and_is_idempotent(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + controller._workflow = wf + + await controller.close() + wf.cleanup.assert_awaited_once() + assert controller.is_ready is False + + await controller.close() + wf.cleanup.assert_awaited_once() # still once — idempotent + + async def test_background_init_cm_closes_manager_on_exit(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + controller._create_fn = lambda: wf + ui = MagicMock() + + async with controller.background_init(ui): + assert await controller.ensure_initialized() is True + + wf.cleanup.assert_awaited_once() + + +class TestControllerLifecycleSerialization: + """init / reinitialize / swap / close must not interleave. + + They all read-modify-write ``_workflow`` across awaits, so two of them in + flight could publish two managers (leaking one), or publish one *after* + ``close()`` had already run — leaving a live backend behind at shutdown. + """ + + async def test_nothing_is_published_after_close(self): + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + wf.is_initialized = True + controller._create_fn = lambda: wf + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + + closing = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + release.set() + await asyncio.wait_for(closing, timeout=5) + + assert controller._workflow is None, "a manager was published after close()" + wf.cleanup.assert_awaited() + assert controller.state.value == "closed" + + async def test_close_waits_for_an_in_flight_reinitialize(self): + import asyncio + + controller = _make_lifecycle_controller() + old = _FakeADKWorkflow() + old.is_initialized = True + controller._workflow = old + + entered, release = asyncio.Event(), asyncio.Event() + + async def _slow_reinit(model=None, preserve_sessions=True): + entered.set() + await release.wait() + + old.reinitialize = _slow_reinit + + reinit = asyncio.create_task(controller.reinitialize()) + await asyncio.wait_for(entered.wait(), timeout=5) + + closing = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + assert not closing.done(), "close() ran through a live reinitialization" + + release.set() + await asyncio.wait_for(reinit, timeout=5) + await asyncio.wait_for(closing, timeout=5) + + assert controller._workflow is None + old.cleanup.assert_awaited() + + async def test_concurrent_swaps_clean_every_losing_candidate(self): + import asyncio + + controller = _make_lifecycle_controller( + orchestrator=OrchestratorType.LANGGRAPH + ) + old = _FakeADKWorkflow() + old.is_initialized = True + controller._workflow = old + + built = [] + + def _build(**kwargs): + new = _FakeLangGraphWorkflow() + new.is_initialized = True + built.append(new) + return new + + with patch( + "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings", + side_effect=_build, + ): + await asyncio.gather( + controller.reinitialize(), controller.reinitialize() + ) + + assert controller._workflow in built + survivors = [m for m in built if m is controller._workflow] + losers = [m for m in built if m is not controller._workflow] + assert len(survivors) == 1 + for loser in losers: + loser.cleanup.assert_awaited(), "a losing swap candidate leaked" + old.cleanup.assert_awaited() + + async def test_init_started_during_reinitialize_does_not_race(self): + import asyncio + + controller = _make_lifecycle_controller() + old = _FakeADKWorkflow() + old.is_initialized = True + controller._workflow = old + + entered, release = asyncio.Event(), asyncio.Event() + + async def _slow_reinit(model=None, preserve_sessions=True): + entered.set() + await release.wait() + + old.reinitialize = _slow_reinit + + reinit = asyncio.create_task(controller.reinitialize()) + await asyncio.wait_for(entered.wait(), timeout=5) + + starting = asyncio.create_task(controller.start_background_init()) + await asyncio.sleep(0.05) + assert not starting.done(), "an init was scheduled mid-reinitialization" + + release.set() + await asyncio.wait_for(reinit, timeout=5) + await asyncio.wait_for(starting, timeout=5) + + assert controller._workflow is old # already READY → no new manager + await controller.close() + + +class TestConstructionInTheExecutorIsTracked: + """A manager built by the init thread must never be silently dropped. + + ``run_in_executor``'s future was awaited unshielded, so cancelling the init + task cancelled the *asyncio* future while the worker thread carried on. + When the thread returned, asyncio discarded the result — a fully + constructed manager (with whatever it had already opened) that nothing + could reach, let alone clean up. + """ + + def _blocking_create(self): + """A ``_create_fn`` that blocks in the worker thread until released.""" + import threading + + entered = threading.Event() + release = threading.Event() + made: list = [] + + def _create(): + entered.set() + release.wait(timeout=5) + wf = _FakeADKWorkflow() + wf.is_initialized = True + made.append(wf) + return wf + + return _create, entered, release, made + + async def test_cancelled_construction_is_cleaned_exactly_once(self): + import asyncio + + controller = _make_lifecycle_controller() + create, entered, release, made = self._blocking_create() + controller._create_fn = create + + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + controller._init_task.cancel() + with pytest.raises(asyncio.CancelledError): + await controller._init_task + + release.set() + await controller.close() + + assert made, "the worker thread never produced a manager" + assert controller._workflow is None, "a cancelled build was published" + made[0].cleanup.assert_awaited_once() + + async def test_close_during_construction_cleans_the_manager(self): + import asyncio + + controller = _make_lifecycle_controller() + create, entered, release, made = self._blocking_create() + controller._create_fn = create + + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + closing = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + release.set() + await asyncio.wait_for(closing, timeout=5) + + assert made + assert controller._workflow is None + made[0].cleanup.assert_awaited_once() + + async def test_manager_built_after_close_is_never_published(self): + import asyncio + + from agentic_cli.cli.workflow_controller import WorkflowState + + controller = _make_lifecycle_controller() + create, entered, release, made = self._blocking_create() + controller._create_fn = create + + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + closing = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + release.set() + await asyncio.wait_for(closing, timeout=5) + + assert controller.state is WorkflowState.CLOSED + with pytest.raises(RuntimeError): + controller.workflow + + +class TestCloseIsCancellationSafe: + """Shutdown must finish even if whoever asked for it goes away. + + ``close()`` did its teardown inline, so a cancelled caller abandoned it + mid-way: a manager still being built in the executor was never released, a + manager already published was left half-cleaned, and a later ``close()`` + returned immediately because ``_closed`` was already True — reporting a + shutdown that never happened. + """ + + async def test_cancelled_close_still_releases_a_pending_construction(self): + import asyncio + import threading + + controller = _make_lifecycle_controller() + entered, release = threading.Event(), threading.Event() + made: list = [] + + def _create(): + entered.set() + release.wait(timeout=5) + wf = _FakeADKWorkflow() + wf.is_initialized = True + made.append(wf) + return wf + + controller._create_fn = _create + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + closing = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + + release.set() + await controller.close() # joins the teardown the cancelled caller started + + assert made, "the worker thread never produced a manager" + assert controller._workflow is None + made[0].cleanup.assert_awaited_once() + + async def test_cancelled_close_still_cleans_a_published_manager(self): + import asyncio + + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + + entered, release = asyncio.Event(), asyncio.Event() + finished: list[str] = [] + + async def _slow_cleanup(): + entered.set() + await release.wait() + finished.append("cleanup") + + wf.cleanup = AsyncMock(side_effect=_slow_cleanup) + + closing = asyncio.create_task(controller.close()) + await asyncio.wait_for(entered.wait(), timeout=5) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + + release.set() + await controller.close() + + wf.cleanup.assert_awaited_once() + assert finished == ["cleanup"], "cleanup was abandoned mid-way" + assert controller._workflow is None + + async def test_a_later_close_joins_the_first(self): + import asyncio + + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + + entered, release = asyncio.Event(), asyncio.Event() + finished: list[str] = [] + + async def _slow_cleanup(): + entered.set() + await release.wait() + finished.append("cleanup") + + wf.cleanup = AsyncMock(side_effect=_slow_cleanup) + + first = asyncio.create_task(controller.close()) + await asyncio.wait_for(entered.wait(), timeout=5) + + second = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + assert not second.done(), "a second close() reported a shutdown still running" + + release.set() + await asyncio.wait_for(asyncio.gather(first, second), timeout=5) + wf.cleanup.assert_awaited_once() + assert finished == ["cleanup"] + + async def test_cancelling_the_join_does_not_stop_the_teardown(self): + """The teardown is the controller's, not the caller's.""" + import asyncio + + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + + entered, release = asyncio.Event(), asyncio.Event() + finished: list[str] = [] + + async def _slow_cleanup(): + entered.set() + await release.wait() + finished.append("cleanup") + + wf.cleanup = AsyncMock(side_effect=_slow_cleanup) + + first = asyncio.create_task(controller.close()) + await asyncio.wait_for(entered.wait(), timeout=5) + second = asyncio.create_task(controller.close()) + await asyncio.sleep(0.05) + + first.cancel() + second.cancel() + for task in (first, second): + with pytest.raises(asyncio.CancelledError): + await task + + release.set() + await controller.close() + + wf.cleanup.assert_awaited_once() + assert finished == ["cleanup"], "a cancelled caller abandoned the teardown" + assert controller.state.value == "closed" + + +class TestCancelInitIsCancellationSafe: + """``cancel_init()`` is public and may be awaited directly. + + Cancelling it mid-settle consumed the construction claim and then never + released the manager: the claim is single-shot, so the fallback callback + could not take over either, and a later ``close()`` found nothing to settle. + """ + + def _blocking_create(self): + import threading + + entered, release = threading.Event(), threading.Event() + made: list = [] + + def _create(): + entered.set() + release.wait(timeout=5) + wf = _FakeADKWorkflow() + wf.is_initialized = True + made.append(wf) + return wf + + return _create, entered, release, made + + async def test_cancelled_cancel_init_does_not_strand_the_manager(self): + import asyncio + + controller = _make_lifecycle_controller() + create, entered, release, made = self._blocking_create() + controller._create_fn = create + + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + cancelling = asyncio.create_task(controller.cancel_init()) + await asyncio.sleep(0.05) + cancelling.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelling + + release.set() + await controller.close() + + assert made, "the worker thread never produced a manager" + assert controller._workflow is None, "a stranded manager was published" + made[0].cleanup.assert_awaited_once() + + async def test_cancelled_cancel_init_without_close_still_cleans(self): + """The fallback callback must still own the result.""" + import asyncio + + controller = _make_lifecycle_controller() + create, entered, release, made = self._blocking_create() + controller._create_fn = create + + await controller.start_background_init() + await asyncio.to_thread(entered.wait, 5) + + cancelling = asyncio.create_task(controller.cancel_init()) + await asyncio.sleep(0.05) + cancelling.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelling + + release.set() + for _ in range(100): + if made and made[0].cleanup.await_count: + break + await asyncio.sleep(0.02) + + assert made + made[0].cleanup.assert_awaited_once() + assert controller._workflow is None + + +class TestConstructionCleanupSurvivesCancellation: + """The *cleanup* of an abandoned construction is the controller's too. + + ``_settle_construction()`` was cancellation-safe only while awaiting the + worker thread. Once it had the manager it awaited ``manager.cleanup()`` + inline, so a caller cancelled at that point cancelled the cleanup itself — + with the construction claim already consumed and ``_construction`` cleared, + neither the fallback callback nor a later ``close()`` could finish it. + """ + + def _controller_with_slow_cleanup(self): + """A controller whose next-built manager blocks inside ``cleanup()``.""" + import asyncio + import threading + + controller = _make_lifecycle_controller() + made: list = [] + finished: list[str] = [] + entered, release = threading.Event(), threading.Event() + cleanup_entered, cleanup_release = asyncio.Event(), asyncio.Event() + + async def _slow_cleanup() -> None: + cleanup_entered.set() + await cleanup_release.wait() + finished.append("cleanup") + + def _create(): + entered.set() + release.wait(timeout=5) + wf = _FakeADKWorkflow() + wf.is_initialized = True + wf.cleanup = AsyncMock(side_effect=_slow_cleanup) + made.append(wf) + return wf + + controller._create_fn = _create + return SimpleNamespace( + controller=controller, + made=made, + finished=finished, + entered=entered, + release=release, + cleanup_entered=cleanup_entered, + cleanup_release=cleanup_release, + ) + + async def _cancel_inside_cleanup(self, h): + """Drive ``cancel_init()`` until cleanup has entered, then cancel it.""" + import asyncio + + await h.controller.start_background_init() + await asyncio.to_thread(h.entered.wait, 5) + + cancelling = asyncio.create_task(h.controller.cancel_init()) + await asyncio.sleep(0.05) # reach the await on the worker thread + h.release.set() + await asyncio.wait_for(h.cleanup_entered.wait(), timeout=5) + + cancelling.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelling + return cancelling + + async def test_close_joins_a_cleanup_the_cancelled_caller_left_running(self): + import asyncio + + h = self._controller_with_slow_cleanup() + await self._cancel_inside_cleanup(h) + + closing = asyncio.create_task(h.controller.close()) + await asyncio.sleep(0.05) + assert not closing.done(), "close() did not join the pending cleanup" + + h.cleanup_release.set() + await asyncio.wait_for(closing, timeout=5) + + assert h.finished == ["cleanup"], "the cancelled caller aborted the cleanup" + assert h.made[0].cleanup.await_count == 1 + assert h.controller._workflow is None + + async def test_cleanup_completes_without_a_later_close(self): + import asyncio + + h = self._controller_with_slow_cleanup() + await self._cancel_inside_cleanup(h) + + h.cleanup_release.set() + for _ in range(200): + if h.finished: + break + await asyncio.sleep(0.02) + + assert h.finished == ["cleanup"], "nobody finished the abandoned cleanup" + assert h.made[0].cleanup.await_count == 1 + + await h.controller.close() + assert h.made[0].cleanup.await_count == 1, "the manager was cleaned twice" + + +class TestInitErrorIsClearedOnSuccess: + """A recorded failure must not outlive the recovery. + + ``_init_error`` was only cleared at the *start* of a background init, so + after a failed reinitialization that then succeeded the controller was + ``READY`` while still holding the old exception — and every consumer keyed + off a different field: ``ensure_initialized()`` returned False (it checked + the error), ``state``/``workflow`` said ready, and the status bar showed + "Init failed - check API keys" over a working session. + """ + + def _recovered_controller(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + controller._init_error = RuntimeError("an earlier failure") + return controller, wf + + async def test_successful_reinitialize_clears_the_error(self): + controller, _ = self._recovered_controller() + + await controller.reinitialize(model="gemini-2.5-flash") + + assert controller.init_error is None + + async def test_successful_swap_clears_the_error(self): + controller = _make_lifecycle_controller( + orchestrator=OrchestratorType.LANGGRAPH + ) + old = _FakeADKWorkflow() + old.is_initialized = True + controller._workflow = old + controller._init_error = RuntimeError("an earlier failure") + + new = _FakeLangGraphWorkflow() + new.is_initialized = True + with patch( + "agentic_cli.cli.workflow_controller.create_workflow_manager_from_settings", + return_value=new, + ): + await controller.reinitialize() + + assert controller.init_error is None + + async def test_successful_background_init_clears_the_error(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._create_fn = lambda: wf + controller._init_error = RuntimeError("an earlier failure") + + assert await controller.ensure_initialized() is True + assert controller.init_error is None + + async def test_everything_agrees_after_recovery(self): + from agentic_cli.cli.workflow_controller import WorkflowState + + controller, wf = self._recovered_controller() + await controller.reinitialize() + + assert controller.state is WorkflowState.READY + assert controller.is_ready is True + assert controller.workflow is wf + assert await controller.ensure_initialized() is True + + ui = MagicMock() + controller.update_status_bar(ui) + status = ui.set_status.call_args[0][0] + assert "Init failed" not in status + assert wf.model in status + + +class TestFailedReinitPreservesSessions: + """A failed in-place reinitialization must not discard the conversation. + + ``GoogleADKWorkflowManager.reinitialize(preserve_sessions=True)`` restores + the session service it was carrying when initialization fails. The + controller then cleaned the manager up anyway, which closed that service — + with ``session_store='memory'`` the whole conversation went with it, for a + failure the user could fix (a bad model id) and retry. + """ + + def _controller_with_failing_reinit(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + + async def _fail(model=None, preserve_sessions=True): + wf.is_initialized = False # the manager rolled itself back + raise RuntimeError("reinit boom") + + wf.reinitialize = _fail + return controller, wf + + async def test_failed_reinit_keeps_the_manager_alive(self): + from agentic_cli.cli.workflow_controller import WorkflowState + + controller, wf = self._controller_with_failing_reinit() + + with pytest.raises(RuntimeError, match="reinit boom"): + await controller.reinitialize() + + assert controller.state is WorkflowState.FAILED + assert controller.is_ready is False + wf.cleanup.assert_not_awaited(), "the preserved session service was closed" + + async def test_retry_revives_the_same_manager(self): + controller, wf = self._controller_with_failing_reinit() + + async def _init_ok(): + wf.is_initialized = True + + wf.initialize_services = _init_ok + replacement = _FakeADKWorkflow() + controller._create_fn = lambda: replacement + + with pytest.raises(RuntimeError): + await controller.reinitialize() + + assert await controller.ensure_initialized() is True + assert controller.workflow is wf, "sessions were dropped for a fresh manager" + replacement.cleanup.assert_not_awaited() + + async def test_unrevivable_manager_is_released_and_replaced(self): + controller, wf = self._controller_with_failing_reinit() + wf.initialize_services = AsyncMock(side_effect=RuntimeError("still broken")) + + replacement = _FakeADKWorkflow() + replacement.is_initialized = True + controller._create_fn = lambda: replacement + + with pytest.raises(RuntimeError): + await controller.reinitialize() + + assert await controller.ensure_initialized() is False + wf.cleanup.assert_awaited_once() + + assert await controller.ensure_initialized() is True + assert controller.workflow is replacement + + async def test_failed_manager_is_never_handed_out(self): + controller, wf = self._controller_with_failing_reinit() + + with pytest.raises(RuntimeError, match="reinit boom"): + await controller.reinitialize() + + with pytest.raises(RuntimeError, match="not initialized"): + controller.workflow + + +class TestControllerStateMachine: + """Explicit lifecycle states, derived from the controller's internals.""" + + async def test_state_progression_uninitialized_to_ready(self): + import asyncio + + from agentic_cli.cli.workflow_controller import WorkflowState + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + controller._create_fn = lambda: wf + + assert controller.state is WorkflowState.UNINITIALIZED + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + assert controller.state is WorkflowState.INITIALIZING + + release.set() + await controller._init_task + assert controller.state is WorkflowState.READY + + await controller.close() + assert controller.state is WorkflowState.CLOSED + + async def test_state_failed_after_init_error(self): + from agentic_cli.cli.workflow_controller import WorkflowState + + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.initialize_services = AsyncMock(side_effect=RuntimeError("boom")) + controller._create_fn = lambda: wf + + await controller.start_background_init() + await controller._init_task + + assert controller.state is WorkflowState.FAILED + + async def test_closed_controller_is_not_ready_and_refuses_init(self): + controller = _make_lifecycle_controller() + await controller.close() + + assert controller.is_ready is False + assert await controller.ensure_initialized() is False + with pytest.raises(RuntimeError, match="closed"): + await controller.start_background_init() + + +class TestControllerSingleFlightInit: + """Two callers must never build two managers for one controller.""" + + async def test_concurrent_start_background_init_creates_one_manager(self): + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + created = [] + + def _create(): + created.append(wf) + return wf + + controller._create_fn = _create + + await asyncio.gather(*(controller.start_background_init() for _ in range(5))) + await asyncio.wait_for(started.wait(), timeout=5) + release.set() + await controller._init_task + + assert len(created) == 1 + assert controller.workflow is wf + + async def test_concurrent_ensure_initialized_callers_all_see_ready(self): + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + controller._create_fn = lambda: wf + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + + waiters = [ + asyncio.create_task(controller.ensure_initialized()) for _ in range(4) + ] + await asyncio.sleep(0.05) + assert not any(w.done() for w in waiters) + + release.set() + assert await asyncio.gather(*waiters) == [True] * 4 + + async def test_cancelled_caller_does_not_kill_shared_init(self): + """One caller giving up must not cancel the shared init for everyone.""" + import asyncio + + controller = _make_lifecycle_controller() + wf, started, release = _blocked_init_workflow() + controller._create_fn = lambda: wf + + await controller.start_background_init() + await asyncio.wait_for(started.wait(), timeout=5) + + giving_up = asyncio.create_task(controller.ensure_initialized()) + await asyncio.sleep(0.05) + giving_up.cancel() + with pytest.raises(asyncio.CancelledError): + await giving_up + + release.set() + assert await controller.ensure_initialized() is True + assert controller.workflow is wf + + +class TestControllerInitRetry: + """After a failure, a fresh start_background_init() retries cleanly.""" + + async def test_retry_after_failure_succeeds_and_clears_error(self): + controller = _make_lifecycle_controller() + failing = _FakeADKWorkflow() + failing.initialize_services = AsyncMock(side_effect=RuntimeError("boom")) + good = _FakeADKWorkflow() + managers = [failing, good] + controller._create_fn = lambda: managers.pop(0) + + await controller.start_background_init() + await controller._init_task + assert controller.init_error is not None + failing.cleanup.assert_awaited_once() + + await controller.start_background_init() + await controller._init_task + + assert controller.init_error is None + assert controller.is_ready is True + assert controller.workflow is good + + async def test_start_is_noop_once_ready(self): + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + created = [] + + def _create(): + created.append(wf) + return wf + + controller._create_fn = _create + + await controller.start_background_init() + await controller._init_task + await controller.start_background_init() + + assert len(created) == 1 + + +class TestReinitializeTransaction: + """A failed in-place reinitialization must not leave the controller READY.""" + + def _ready_controller(self): + from agentic_cli.cli.workflow_controller import WorkflowState + + controller = _make_lifecycle_controller() + wf = _FakeADKWorkflow() + wf.is_initialized = True + controller._workflow = wf + assert controller.state is WorkflowState.READY + return controller, wf + + async def test_failed_inplace_reinit_enters_failed_and_withholds(self): + """FAILED, and the manager is not handed out — but it is kept. + + See ``TestFailedReinitPreservesSessions``: releasing it here closed + the session service the manager had just preserved. + """ + from agentic_cli.cli.workflow_controller import WorkflowState + + controller, wf = self._ready_controller() + + async def _fail(model=None, preserve_sessions=True): + wf.is_initialized = False # the manager rolled itself back + raise RuntimeError("reinit boom") + + wf.reinitialize = _fail + + with pytest.raises(RuntimeError, match="reinit boom"): + await controller.reinitialize(model="gemini-2.5-flash") + + assert controller.state is WorkflowState.FAILED + assert controller.is_ready is False + assert isinstance(controller.init_error, RuntimeError) + with pytest.raises(RuntimeError): + controller.workflow + wf.cleanup.assert_not_awaited() + + async def test_uninitialized_manager_is_never_reported_ready(self): + """State is derived from the manager, not from 'we published one'.""" + from agentic_cli.cli.workflow_controller import WorkflowState + + controller, wf = self._ready_controller() + wf.is_initialized = False + + assert controller.state is WorkflowState.FAILED + assert controller.is_ready is False + + async def test_successful_reinit_stays_ready(self): + from agentic_cli.cli.workflow_controller import WorkflowState + + controller, wf = self._ready_controller() + await controller.reinitialize(model="gemini-2.5-flash") + + assert controller.state is WorkflowState.READY + assert controller.workflow is wf + + async def test_recovery_after_failed_reinit_restores_readiness(self): + controller, wf = self._ready_controller() + + async def _fail(model=None, preserve_sessions=True): + wf.is_initialized = False + raise RuntimeError("reinit boom") + + async def _init_ok(): + wf.is_initialized = True + + wf.reinitialize = _fail + wf.initialize_services = _init_ok + replacement = _FakeADKWorkflow() + replacement.is_initialized = True + controller._create_fn = lambda: replacement + + from agentic_cli.cli.workflow_controller import WorkflowState + + with pytest.raises(RuntimeError): + await controller.reinitialize() + + assert await controller.ensure_initialized() is True + assert controller.state is WorkflowState.READY + # Recovery revives the failed manager (see + # TestFailedReinitPreservesSessions), so no replacement is built. + replacement.initialize_services.assert_not_awaited() diff --git a/tests/tools/test_registry_identity.py b/tests/tools/test_registry_identity.py index 5fb20f9..2c42b2c 100644 --- a/tests/tools/test_registry_identity.py +++ b/tests/tools/test_registry_identity.py @@ -1388,3 +1388,16 @@ def test_non_string_element_is_rejected(self): def _tool() -> dict: """Tool.""" return {"success": True} + + def test_every_declarable_key_is_constructible(self): + """The declarable set must not drift from what the manager can build.""" + import inspect + + from agentic_cli.workflow.base_manager import BaseWorkflowManager + from agentic_cli.workflow.service_registry import KNOWN_SERVICE_KEYS + + source = inspect.getsource(BaseWorkflowManager._build_services_into) + for key in KNOWN_SERVICE_KEYS: + assert f'"{key}" in self._required_managers' in source, ( + f"{key} is declarable but _build_services_into never constructs it" + ) diff --git a/tests/workflow/test_base_manager_init_lock.py b/tests/workflow/test_base_manager_init_lock.py new file mode 100644 index 0000000..7207905 --- /dev/null +++ b/tests/workflow/test_base_manager_init_lock.py @@ -0,0 +1,78 @@ +"""Concurrent initialize_services() must run the init body exactly once. + +The guard at the top of BaseWorkflowManager.initialize_services() was +check-then-act: a user message arriving while background init is mid-flight +(manager._ensure_initialized → initialize_services) raced the background +call and ran the whole body twice (duplicate registry refresh, duplicate +service creation). An asyncio lock serializes them. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from agentic_cli.workflow.base_manager import BaseWorkflowManager + + +class _CountingManager(BaseWorkflowManager): + """Minimal concrete manager counting _do_initialize runs.""" + + def __init__(self, settings): + super().__init__(agent_configs=[], settings=settings) + self.do_init_calls = 0 + + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + self.do_init_calls += 1 + # Yield so a concurrent initialize_services() can interleave + await asyncio.sleep(0.02) + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + pass + + async def cleanup(self): + pass + + +def _settings(): + s = MagicMock() + s.app_name = "test-app" + s.google_api_key = None + s.anthropic_api_key = None + return s + + +def _manager(): + m = _CountingManager(_settings()) + m._model_registry = MagicMock(refresh=AsyncMock()) + m._ensure_managers_initialized = lambda: None + return m + + +async def test_concurrent_initialize_services_runs_once(): + m = _manager() + + await asyncio.gather( + m.initialize_services(validate=False), + m.initialize_services(validate=False), + ) + + assert m.do_init_calls == 1 + assert m.is_initialized + + +async def test_sequential_initialize_services_is_idempotent(): + m = _manager() + + await m.initialize_services(validate=False) + await m.initialize_services(validate=False) + + assert m.do_init_calls == 1 diff --git a/tests/workflow/test_lifecycle_races.py b/tests/workflow/test_lifecycle_races.py new file mode 100644 index 0000000..5a8165f --- /dev/null +++ b/tests/workflow/test_lifecycle_races.py @@ -0,0 +1,287 @@ +"""Lifecycle races between a turn, a cleanup, and a cancelled initialization. + +Two defects: + +1. ``process()`` initializes *before* taking the turn lock (that ordering is + what keeps cleanup from deadlocking against a running turn). A cleanup that + was already queued therefore ran in between, and the turn woke up holding + admission to a manager whose runner had just been released — it then used + ``None`` as a runner deep inside ADK. +2. Service construction runs on a worker thread (``asyncio.to_thread``) and + wrote straight into ``self._services``. Cancelling the awaiting coroutine + does not stop the thread, so a rolled-back initialization was followed, + moments later, by that thread publishing services into a manager that had + already been cleaned up — leaking a sandbox/job manager nobody would close. +""" + +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.base_manager import BaseWorkflowManager # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.events import EventType, WorkflowEvent # noqa: E402 +from agentic_cli.workflow.service_registry import SANDBOX_MANAGER # noqa: E402 +from tests.conftest import MockContext # noqa: E402 + + +# --------------------------------------------------------------------------- +# 1. A turn admitted after a cleanup must not use released resources +# --------------------------------------------------------------------------- + + +class _AdmissionHarness: + """An ADK manager whose stream and initialization the test drives.""" + + def __init__(self, ctx) -> None: + self.manager = GoogleADKWorkflowManager( + agent_configs=[AgentConfig(name="a", prompt="p")], settings=ctx.settings + ) + self.manager._event_processor = SimpleNamespace(model=None) + self.inits = 0 + self.streams: list[str] = [] + self.runner_at_stream: list[object] = [] + + async def _do_initialize() -> None: + self.inits += 1 + self.manager._session_service = SimpleNamespace( + get_session=self._get_session + ) + self.manager._root_agent = SimpleNamespace(name="a") + self.manager._runner = SimpleNamespace(name=f"runner-{self.inits}") + + async def _get_or_create(user_id, session_id): + return SimpleNamespace(id=session_id) + + async def _stream(*, session_id, user_id, new_message, run_config): + self.streams.append(session_id) + self.runner_at_stream.append(self.manager._runner) + yield WorkflowEvent(type=EventType.TEXT, content=session_id) + + self.manager._do_initialize = _do_initialize + self.manager._get_or_create_session = _get_or_create + self.manager._run_and_stream = _stream + self.manager._model_registry = MagicMock(refresh=AsyncMock()) + self.manager._ensure_managers_initialized = lambda: None + self.manager._validate_agent_graph = lambda: None + + @staticmethod + async def _get_session(**kwargs): + return SimpleNamespace(id=kwargs.get("session_id")) + + async def drain(self, session_id: str) -> list[str]: + return [ + e.content + async for e in self.manager.process("hi", "u", session_id=session_id) + ] + + +def _admission_harness(): + ctx = MockContext(google_api_key="test-key") + ctx.__enter__() + return _AdmissionHarness(ctx), ctx + + +class TestTurnAdmissionRechecksReadiness: + async def test_turn_queued_behind_cleanup_reinitializes(self): + """Cleanup lands between the turn's init and its admission.""" + h, ctx = _admission_harness() + try: + await h.manager.initialize_services(validate=False) + assert h.inits == 1 + + # Hold the turn lock so the turn queues, then clean up behind it. + await h.manager._turn_lock.acquire() + turn = asyncio.create_task(h.drain("sess-a")) + await asyncio.sleep(0.05) + assert not turn.done() + + # Release the manager's resources while the turn waits for + # admission (the turn already passed _ensure_initialized). + await h.manager._release_resources() + assert h.manager._runner is None + h.manager._turn_lock.release() + + assert await asyncio.wait_for(turn, timeout=2) == ["sess-a"] + assert h.inits == 2, "the turn ran against the released backend" + assert h.runner_at_stream[-1] is not None + assert h.runner_at_stream[-1].name == "runner-2" + finally: + ctx.__exit__(None, None, None) + + async def test_turn_fails_cleanly_when_the_backend_cannot_be_revived(self): + """No silent AttributeError on a ``None`` runner.""" + h, ctx = _admission_harness() + try: + await h.manager.initialize_services(validate=False) + + async def _do_nothing() -> None: + self_inits = None # noqa: F841 - deliberately leaves it unready + return None + + await h.manager._turn_lock.acquire() + turn = asyncio.create_task(h.drain("sess-a")) + await asyncio.sleep(0.05) + + await h.manager._release_resources() + h.manager._do_initialize = _do_nothing + h.manager._turn_lock.release() + + with pytest.raises(RuntimeError, match="initiali"): + await asyncio.wait_for(turn, timeout=2) + assert h.streams == [], "the turn streamed from a released backend" + finally: + ctx.__exit__(None, None, None) + + async def test_resume_turn_rechecks_too(self): + h, ctx = _admission_harness() + try: + await h.manager.initialize_services(validate=False) + record = SimpleNamespace( + job_id="j1", session_id="sess-r", user_id="u", call_id="c1", + call_name="t", tool="t", state=SimpleNamespace(value="succeeded"), + exit_code=0, error=None, + ) + + await h.manager._turn_lock.acquire() + + async def _drain_resume(): + return [ + e.content + async for e in h.manager.resume_with_job_result(record, "ok") + ] + + resume = asyncio.create_task(_drain_resume()) + await asyncio.sleep(0.05) + + await h.manager._release_resources() + h.manager._turn_lock.release() + + assert await asyncio.wait_for(resume, timeout=2) == ["sess-r"] + assert h.inits == 2 + finally: + ctx.__exit__(None, None, None) + + +# --------------------------------------------------------------------------- +# 2. Worker-thread service construction is transactional +# --------------------------------------------------------------------------- + + +class _SlowServiceManager(BaseWorkflowManager): + """Builds services on a worker thread, slowly, and records the closes.""" + + def __init__(self, settings, gate: threading.Event) -> None: + super().__init__(agent_configs=[], settings=settings) + self._gate = gate + self.built: list[object] = [] + self._required_managers = {"sandbox_manager"} + + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + await self._release_resources() + + def _make_sandbox_manager(self): + # Called on the worker thread; blocks until the test releases it. + self._gate.wait(timeout=5) + service = MagicMock() + service.cleanup = MagicMock() + self.built.append(service) + return service + + +def _slow_manager(gate: threading.Event) -> _SlowServiceManager: + settings = MagicMock() + settings.app_name = "test-app" + settings.google_api_key = None + settings.anthropic_api_key = None + manager = _SlowServiceManager(settings, gate) + manager._model_registry = MagicMock(refresh=AsyncMock()) + return manager + + +class TestWorkerThreadConstructionIsTransactional: + async def test_cancelled_init_does_not_publish_services(self, monkeypatch): + gate = threading.Event() + manager = _slow_manager(gate) + monkeypatch.setattr( + "agentic_cli.tools.sandbox.manager.SandboxManager", + lambda settings: manager._make_sandbox_manager(), + ) + + task = asyncio.create_task(manager.initialize_services(validate=False)) + await asyncio.sleep(0.05) # let the worker thread start and block + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + gate.set() # the thread finishes *after* the rollback + for _ in range(100): + if manager.built: + break + await asyncio.sleep(0.02) + + assert manager.services.get(SANDBOX_MANAGER) is None, ( + "a cancelled initialization published services into a live manager" + ) + assert manager.is_initialized is False + + async def test_services_built_after_cancellation_are_released(self, monkeypatch): + gate = threading.Event() + manager = _slow_manager(gate) + monkeypatch.setattr( + "agentic_cli.tools.sandbox.manager.SandboxManager", + lambda settings: manager._make_sandbox_manager(), + ) + + task = asyncio.create_task(manager.initialize_services(validate=False)) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + gate.set() + for _ in range(100): + if manager.built and manager.built[0].cleanup.called: + break + await asyncio.sleep(0.02) + + assert manager.built, "the worker thread never finished" + manager.built[0].cleanup.assert_called_once() + + async def test_successful_init_publishes_normally(self, monkeypatch): + gate = threading.Event() + gate.set() + manager = _slow_manager(gate) + monkeypatch.setattr( + "agentic_cli.tools.sandbox.manager.SandboxManager", + lambda settings: manager._make_sandbox_manager(), + ) + + await manager.initialize_services(validate=False) + + assert manager.services.get(SANDBOX_MANAGER) is manager.built[0] + assert manager.is_initialized is True diff --git a/tests/workflow/test_model_validation.py b/tests/workflow/test_model_validation.py index 5e4ad4d..a6b1c15 100644 --- a/tests/workflow/test_model_validation.py +++ b/tests/workflow/test_model_validation.py @@ -444,6 +444,17 @@ async def test_manager_runs_on_the_replacement(self): assert config.model == "gemini-2.5-pro" +class _ReinitManager(_TestManager): + """A manager whose reinitialize really re-runs initialization.""" + + async def reinitialize(self, model=None, preserve_sessions=True): + async with self._lifecycle_lock: + async with self._turn_lock: + self._initialized = False + self._reset_model(model) + await self._initialize_locked() + + class TestManagerModelIsValidated: """Every model the *runtime* will actually send must be validated. @@ -489,6 +500,14 @@ async def test_cached_model_is_normalized(self): assert manager.model == "gemini-2.5-pro" + async def test_reinitialize_override_is_normalized(self): + with MockContext(google_api_key="k") as ctx: + manager = self._manager(ctx.settings, cls=_ReinitManager) + await manager.initialize_services() + + await manager.reinitialize(model="gemini-old") + + assert manager.model == "gemini-2.5-pro" async def test_explicit_model_without_a_credential_fails(self): with MockContext(google_api_key="k") as ctx: diff --git a/tests/workflow/test_resource_ownership.py b/tests/workflow/test_resource_ownership.py new file mode 100644 index 0000000..4971822 --- /dev/null +++ b/tests/workflow/test_resource_ownership.py @@ -0,0 +1,457 @@ +"""Owned async resources are closed exactly once, on every shutdown path. + +``cleanup()`` used to drop the session service by assignment. The durable +``DatabaseSessionService`` owns a SQLAlchemy engine with an async ``close()``, +so dropping the reference leaked its connection pool. Cleanup now awaits the +close contract, stays idempotent, and never closes a service it handed over +(``reinitialize(preserve_sessions=True)``). +""" + +from __future__ import annotations + +import asyncio + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 + + +class _AsyncClosable: + """Stand-in for DatabaseSessionService: async close(), counted.""" + + def __init__(self) -> None: + self.closes = 0 + + async def close(self) -> None: + self.closes += 1 + + +class _SyncClosable: + def __init__(self) -> None: + self.closes = 0 + + def close(self) -> None: + self.closes += 1 + + +class _NotClosable: + """Stand-in for InMemorySessionService: nothing to release.""" + + +class _FailingClosable: + async def close(self) -> None: + raise RuntimeError("close blew up") + + +def _manager(session_service) -> GoogleADKWorkflowManager: + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = SimpleNamespace(app_name="test", default_user="u") + mgr._app_name = "test" + mgr._services = {} + mgr._session_service = session_service + mgr._runner = object() + mgr._root_agent = object() + mgr._initialized = True + mgr._llm_logging_plugin = None + mgr._model = "gemini-2.5-flash" + mgr._model_resolved = True + mgr._session_service_pinned = False + mgr._lifecycle_lock = asyncio.Lock() + mgr._turn_lock = asyncio.Lock() + return mgr + + +class TestSessionServiceClose: + async def test_cleanup_awaits_async_close(self): + service = _AsyncClosable() + mgr = _manager(service) + + await mgr.cleanup() + + assert service.closes == 1 + assert mgr._session_service is None + assert mgr.is_initialized is False + + async def test_cleanup_is_idempotent(self): + service = _AsyncClosable() + mgr = _manager(service) + + await mgr.cleanup() + await mgr.cleanup() + + assert service.closes == 1 + + async def test_service_without_close_is_tolerated(self): + mgr = _manager(_NotClosable()) + await mgr.cleanup() # must not raise + assert mgr._session_service is None + + async def test_failing_close_does_not_block_shutdown(self): + mgr = _manager(_FailingClosable()) + await mgr.cleanup() # swallowed and logged + assert mgr._session_service is None + + async def test_sync_close_is_supported(self): + service = _SyncClosable() + mgr = _manager(service) + await mgr.cleanup() + assert service.closes == 1 + + +def _real_manager(monkeypatch=None): + """A manager built through ``__init__`` with the network stubbed out. + + ``_do_initialize`` is replaced by a faithful stand-in: it creates the + session service only when one is not already held (which is exactly how + ``reinitialize(preserve_sessions=True)`` avoids building a replacement). + """ + from unittest.mock import AsyncMock, MagicMock + + from agentic_cli.workflow.config import AgentConfig + from tests.conftest import MockContext + + ctx = MockContext(google_api_key="test-key") + ctx.__enter__() + mgr = GoogleADKWorkflowManager( + agent_configs=[AgentConfig(name="a", prompt="p")], settings=ctx.settings + ) + mgr._model_registry = MagicMock(refresh=AsyncMock(), discovery_complete=False) + mgr._ensure_managers_initialized = lambda: None + + created: list[_AsyncClosable] = [] + + def _make_service(): + service = _AsyncClosable() + created.append(service) + return service + + mgr._make_session_service = _make_service + state = {"boom": False} + + async def _do_init(): + if mgr._session_service is None: + mgr._session_service = mgr._make_session_service() + if state["boom"]: + raise RuntimeError("backend init failed") + mgr._runner = object() + mgr._root_agent = object() + + mgr._do_initialize = _do_init + return mgr, created, state, ctx + + +class TestReinitializeTransaction: + """reinitialize() either fully succeeds or leaves nothing half-built.""" + + async def test_preserved_service_is_reused_not_replaced(self): + mgr, created, _state, ctx = _real_manager() + try: + await mgr.initialize_services() + original = mgr._session_service + assert len(created) == 1 + + await mgr.reinitialize(preserve_sessions=True) + + assert mgr._session_service is original, "the live service was swapped" + assert len(created) == 1, "a replacement service was built and discarded" + assert original.closes == 0, "the preserved service was closed" + assert mgr.is_initialized is True + finally: + ctx.__exit__(None, None, None) + + async def test_discarded_service_is_closed_and_replaced(self): + mgr, created, _state, ctx = _real_manager() + try: + await mgr.initialize_services() + original = mgr._session_service + + await mgr.reinitialize(preserve_sessions=False) + + assert original.closes == 1 + assert mgr._session_service is not original + assert len(created) == 2 + finally: + ctx.__exit__(None, None, None) + + async def test_failed_reinit_keeps_preserved_service_and_uninitializes(self): + mgr, created, state, ctx = _real_manager() + try: + await mgr.initialize_services() + original = mgr._session_service + + state["boom"] = True + with pytest.raises(RuntimeError, match="backend init failed"): + await mgr.reinitialize(preserve_sessions=True) + + assert original.closes == 0, "the preserved service was lost" + assert mgr._session_service is original + assert mgr.is_initialized is False, "a failed reinit must not look ready" + assert len(created) == 1 + finally: + ctx.__exit__(None, None, None) + + async def test_failed_reinit_closes_the_replacement_it_created(self): + """preserve_sessions=False: the new service must not leak on failure.""" + mgr, created, state, ctx = _real_manager() + try: + await mgr.initialize_services() + original = mgr._session_service + + state["boom"] = True + with pytest.raises(RuntimeError): + await mgr.reinitialize(preserve_sessions=False) + + assert original.closes == 1 # discarded on purpose + assert len(created) == 2 + assert created[1].closes == 1, "the replacement service leaked" + assert mgr._session_service is None + assert mgr.is_initialized is False + finally: + ctx.__exit__(None, None, None) + + +class TestDirectInitializationFailure: + """A direct initialize_services() failure rolls its own resources back.""" + + async def test_partial_initialization_is_rolled_back(self): + mgr, created, state, ctx = _real_manager() + try: + state["boom"] = True + with pytest.raises(RuntimeError, match="backend init failed"): + await mgr.initialize_services() + + assert len(created) == 1 + assert created[0].closes == 1, "the session service leaked" + assert mgr._session_service is None + assert mgr.is_initialized is False + assert mgr.services == {} + finally: + ctx.__exit__(None, None, None) + + async def test_manager_can_be_initialized_after_a_failure(self): + mgr, created, state, ctx = _real_manager() + try: + state["boom"] = True + with pytest.raises(RuntimeError): + await mgr.initialize_services() + + state["boom"] = False + await mgr.initialize_services() + + assert mgr.is_initialized is True + assert mgr._session_service is created[-1] + finally: + ctx.__exit__(None, None, None) + + +class TestLifecycleSerialization: + """close() and reinitialize() must not interleave.""" + + async def test_concurrent_cleanup_and_reinitialize(self): + mgr, created, _state, ctx = _real_manager() + try: + await mgr.initialize_services() + + order: list[str] = [] + real_do_init = mgr._do_initialize + + async def _slow_init(): + order.append("reinit-start") + await asyncio.sleep(0.02) + await real_do_init() + order.append("reinit-end") + + mgr._do_initialize = _slow_init + + async def _cleanup(): + await asyncio.sleep(0.005) + order.append("cleanup-start") + await mgr.cleanup() + order.append("cleanup-end") + + await asyncio.gather(mgr.reinitialize(preserve_sessions=True), _cleanup()) + + # cleanup must wait for the whole reinit, never interleave with it + assert order.index("reinit-end") < order.index("cleanup-end") + assert mgr.is_initialized is False # cleanup ran last + finally: + ctx.__exit__(None, None, None) + + +class TestOwnedServicesReleased: + """_cleanup_managers releases only services this manager created.""" + + async def test_job_manager_and_sandbox_are_closed(self): + closed: list[str] = [] + mgr = _manager(_NotClosable()) + mgr._services = { + "job_manager": SimpleNamespace(close=lambda: closed.append("jobs")), + "sandbox_manager": SimpleNamespace(cleanup=lambda: closed.append("sandbox")), + } + + await mgr.cleanup() + + assert sorted(closed) == ["jobs", "sandbox"] + assert mgr._services == {} + + async def test_second_cleanup_finds_nothing_to_close(self): + closed: list[str] = [] + mgr = _manager(_NotClosable()) + mgr._services = {"job_manager": SimpleNamespace(close=lambda: closed.append("jobs"))} + + await mgr.cleanup() + await mgr.cleanup() + + assert closed == ["jobs"] + + +class TestPartialConstructionRollsBack: + """A service constructor that raises must not strand its predecessors. + + ``_build_services`` builds into a local dict and hands it to the caller to + publish. When a *later* constructor raised, that dict was simply dropped — + so an already-built SandboxManager (a container/process pool) or JobManager + (a thread pool) was never published and never closed: nothing could ever + release it. + """ + + @staticmethod + def _manager_needing(*services: str): + from unittest.mock import MagicMock + + from agentic_cli.workflow.base_manager import BaseWorkflowManager + + class _Manager(BaseWorkflowManager): + def _get_state_tools(self): + return [] + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + return None + + async def process(self, message, user_id, session_id=None): + raise NotImplementedError + + async def reinitialize(self, model=None, preserve_sessions=True): + return None + + async def cleanup(self): + await self._release_resources() + + settings = MagicMock() + settings.app_name = "test-app" + settings.max_concurrent_jobs = 2 + mgr = _Manager(agent_configs=[], settings=settings) + mgr._required_managers = set(services) + return mgr + + def test_sandbox_is_released_when_a_later_constructor_raises(self, monkeypatch): + sandbox = SimpleNamespace(cleanup=lambda: closed.append("sandbox")) + closed: list[str] = [] + + monkeypatch.setattr( + "agentic_cli.tools.sandbox.manager.SandboxManager", + lambda settings: sandbox, + ) + monkeypatch.setattr( + "agentic_cli.tools.jobs.JobManager", + _raising_ctor("jobs blew up"), + ) + + mgr = self._manager_needing("sandbox_manager", "job_manager") + with pytest.raises(RuntimeError, match="jobs blew up"): + mgr._build_services() + + assert closed == ["sandbox"], "an already-built service was stranded" + + def test_job_manager_is_released_when_a_later_constructor_raises( + self, monkeypatch + ): + closed: list[str] = [] + jobs = SimpleNamespace(close=lambda: closed.append("jobs")) + + monkeypatch.setattr("agentic_cli.tools.jobs.JobManager", lambda *a, **k: jobs) + monkeypatch.setattr( + "agentic_cli.tools.arxiv_source.ArxivSearchSource", + _raising_ctor("arxiv blew up"), + ) + + mgr = self._manager_needing("job_manager", "arxiv_source") + with pytest.raises(RuntimeError, match="arxiv blew up"): + mgr._build_services() + + assert closed == ["jobs"] + + async def test_initialization_failure_leaves_nothing_published( + self, monkeypatch + ): + closed: list[str] = [] + monkeypatch.setattr( + "agentic_cli.tools.sandbox.manager.SandboxManager", + lambda settings: SimpleNamespace(cleanup=lambda: closed.append("sandbox")), + ) + monkeypatch.setattr( + "agentic_cli.tools.jobs.JobManager", _raising_ctor("jobs blew up") + ) + + mgr = self._manager_needing("sandbox_manager", "job_manager") + mgr._model_registry = SimpleNamespace(refresh=_noop_refresh) + + with pytest.raises(RuntimeError, match="jobs blew up"): + await mgr.initialize_services(validate=False) + + assert closed == ["sandbox"] + assert mgr.services == {} + assert mgr.is_initialized is False + + +def _raising_ctor(message: str): + def _ctor(*args, **kwargs): + raise RuntimeError(message) + + return _ctor + + +async def _noop_refresh(**kwargs): + return None + + +class TestCloserIsolation: + """One resource's close failure must not skip the others.""" + + async def test_failing_sync_closer_does_not_block_the_rest(self): + closed: list[str] = [] + + def _boom(): + raise RuntimeError("sandbox cleanup blew up") + + mgr = _manager(_AsyncClosable()) + mgr._services = { + "sandbox_manager": SimpleNamespace(cleanup=_boom), + "job_manager": SimpleNamespace(close=lambda: closed.append("jobs")), + } + + await mgr.cleanup() + + assert closed == ["jobs"], "a failing closer skipped the next resource" + assert mgr._services == {} + + async def test_failing_sync_closer_does_not_block_the_session_service(self): + service = _AsyncClosable() + mgr = _manager(service) + mgr._services = { + "sandbox_manager": SimpleNamespace( + cleanup=lambda: (_ for _ in ()).throw(RuntimeError("boom")) + ) + } + + await mgr.cleanup() + + assert service.closes == 1 diff --git a/tests/workflow/test_turn_serialization.py b/tests/workflow/test_turn_serialization.py new file mode 100644 index 0000000..9b07b0d --- /dev/null +++ b/tests/workflow/test_turn_serialization.py @@ -0,0 +1,316 @@ +"""One turn at a time per manager, and lifecycle mutation waits for it. + +The active session/user identity is a ContextVar, so it is already per-turn. +Everything *else* a turn touches is manager-scoped: the HITL input callback +(``set_input_callback``) and the ADK plugins' event buffers (drained by +``_run_and_stream``). Two overlapping turns would route a permission/HITL answer +to the wrong request and let one invocation drain the other's events, and a +cleanup could tear the runner down mid-stream. + +``process()``/``resume_with_job_result()`` therefore hold a turn lock, and +``cleanup()``/``reinitialize()`` take it too. Initialization happens *before* +the turn lock so the two lock orders can never deadlock. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.events import ( # noqa: E402 + EventType, + UserInputRequest, + WorkflowEvent, +) +from tests.conftest import MockContext # noqa: E402 + + +class _Harness: + """A manager whose stream is driven by the test, with no real ADK runner.""" + + def __init__(self, ctx) -> None: + self.manager = GoogleADKWorkflowManager( + agent_configs=[AgentConfig(name="a", prompt="p")], settings=ctx.settings + ) + self.manager._initialized = True + self.manager._event_processor = SimpleNamespace(model=None) + # Turn admission re-checks readiness while holding the turn lock, so + # the double presents a complete backend (runner + agent + sessions). + self.manager._session_service = SimpleNamespace() + self.manager._runner = SimpleNamespace(name="runner") + self.manager._root_agent = SimpleNamespace(name="a") + self.events: list[str] = [] + self.release = asyncio.Event() + self.entered = asyncio.Event() + # When set, the stream asks the user a question mid-turn (as a HITL + # tool would) once it is released. + self.prompt_mid_turn = False + + async def _ensure() -> None: + return None + + async def _get_or_create(user_id, session_id): + return SimpleNamespace(id=session_id) + + async def _stream(*, session_id, user_id, new_message, run_config): + self.events.append(f"start:{session_id}") + self.entered.set() + await self.release.wait() + if self.prompt_mid_turn: + answer = await self.manager.request_user_input( + UserInputRequest( + request_id=f"req-{session_id}", + tool_name="ask_clarification", + prompt="which?", + ) + ) + self.events.append(f"answer:{session_id}:{answer}") + yield WorkflowEvent(type=EventType.TEXT, content=session_id) + self.events.append(f"end:{session_id}") + + self.manager._ensure_initialized = _ensure + self.manager._get_or_create_session = _get_or_create + self.manager._run_and_stream = _stream + + async def drain(self, session_id: str) -> list[str]: + return [ + e.content + async for e in self.manager.process("hi", "u", session_id=session_id) + ] + + +def _harness(): + ctx = MockContext(google_api_key="test-key") + ctx.__enter__() + return _Harness(ctx), ctx + + +class TestTurnSerialization: + async def test_second_turn_waits_for_the_first(self): + h, ctx = _harness() + try: + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + second = asyncio.create_task(h.drain("sess-b")) + await asyncio.sleep(0.05) + + assert h.events == ["start:sess-a"], "turns overlapped" + assert not second.done() + + h.release.set() + assert await first == ["sess-a"] + assert await second == ["sess-b"] + assert h.events == [ + "start:sess-a", "end:sess-a", "start:sess-b", "end:sess-b", + ] + finally: + ctx.__exit__(None, None, None) + + async def test_running_turn_keeps_its_own_hitl_callback(self): + """A second consumer's callback must not capture the first turn's prompt. + + The turn lock only serialises ``process()``; callbacks are installed + *before* it, so a manager-global callback attribute let the second + consumer answer the first turn's question (and then the first + consumer's ``clear_input_callback()`` unregistered the second's). + The callback is therefore context-local. + """ + h, ctx = _harness() + try: + h.prompt_mid_turn = True + observed: list[str] = [] + + async def _cb_a(request): + observed.append(f"a:{request.request_id}") + return "from-a" + + async def _cb_b(request): # pragma: no cover - must never run + observed.append(f"b:{request.request_id}") + return "from-b" + + h.manager.set_input_callback(_cb_a) + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + # A second consumer installs its own callback and starts a turn. + h.manager.set_input_callback(_cb_b) + second = asyncio.create_task(h.drain("sess-b")) + await asyncio.sleep(0.05) + assert h.events == ["start:sess-a"], "turns overlapped" + + h.release.set() + await asyncio.gather(first, second) + + assert observed == [ + "a:req-sess-a", + "b:req-sess-b", + ], "a turn's prompt was answered by another consumer's callback" + assert "answer:sess-a:from-a" in h.events + assert "answer:sess-b:from-b" in h.events + finally: + ctx.__exit__(None, None, None) + + async def test_clearing_one_callback_does_not_unregister_another(self): + """One consumer tidying up must not unregister a concurrent consumer.""" + h, ctx = _harness() + try: + ready = asyncio.Event() + + async def _cb(request): + return "answer" + + h.manager.set_input_callback(_cb) + + async def _consumer(): + await ready.wait() + return await h.manager.request_user_input( + UserInputRequest( + request_id="r", tool_name="t", prompt="which?" + ) + ) + + # Created after the install, so it carries this callback. + consumer = asyncio.create_task(_consumer()) + + # A different consumer finishes its turn and clears *its* callback. + h.manager.clear_input_callback() + + ready.set() + assert await asyncio.wait_for(consumer, timeout=2) == "answer" + finally: + ctx.__exit__(None, None, None) + + async def test_resume_turn_shares_the_lock(self): + h, ctx = _harness() + try: + record = SimpleNamespace( + job_id="j1", session_id="sess-r", user_id="u", call_id="c1", + call_name="t", tool="t", state=SimpleNamespace(value="succeeded"), + exit_code=0, error=None, + ) + + async def _get_session(**kwargs): + return SimpleNamespace(id="sess-r") + + h.manager._session_service = SimpleNamespace(get_session=_get_session) + + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + async def _drain_resume(): + return [ + e.content + async for e in h.manager.resume_with_job_result(record, "ok") + ] + + resume = asyncio.create_task(_drain_resume()) + await asyncio.sleep(0.05) + assert h.events == ["start:sess-a"], "a resume ran during a user turn" + + h.release.set() + await first + await resume + finally: + ctx.__exit__(None, None, None) + + +class TestCancellationReleasesTheLock: + async def test_cancelled_turn_frees_the_manager(self): + h, ctx = _harness() + try: + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + + h.release.set() + second = asyncio.create_task(h.drain("sess-b")) + assert await asyncio.wait_for(second, timeout=2) == ["sess-b"] + finally: + ctx.__exit__(None, None, None) + + +class TestLifecycleWaitsForTurns: + async def test_cleanup_does_not_run_during_a_turn(self): + h, ctx = _harness() + try: + order: list[str] = [] + real_release = h.manager._release_resources + + async def _tracked(keep_session_service: bool = False): + order.append("cleanup") + await real_release(keep_session_service) + + h.manager._release_resources = _tracked + + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + cleanup = asyncio.create_task(h.manager.cleanup()) + await asyncio.sleep(0.05) + assert order == [], "cleanup tore the backend down mid-turn" + + h.release.set() + await first + await cleanup + assert order == ["cleanup"] + finally: + ctx.__exit__(None, None, None) + + async def test_reinitialize_does_not_run_during_a_turn(self): + h, ctx = _harness() + try: + order: list[str] = [] + + async def _init(validate: bool = True): + order.append("reinit") + + h.manager._initialize_locked = _init + h.manager._reset_model = lambda model: None + + first = asyncio.create_task(h.drain("sess-a")) + await asyncio.wait_for(h.entered.wait(), timeout=2) + + reinit = asyncio.create_task(h.manager.reinitialize()) + await asyncio.sleep(0.05) + assert order == [], "reinitialize ran under an active turn" + + h.release.set() + await first + await reinit + assert order == ["reinit"] + finally: + ctx.__exit__(None, None, None) + + async def test_turn_after_cleanup_reinitializes_without_deadlock(self): + """Lock order (lifecycle → turn) must not deadlock a turn that inits.""" + h, ctx = _harness() + try: + h.release.set() + await h.manager.cleanup() + + inits: list[int] = [] + + async def _ensure() -> None: + inits.append(1) + # A real _ensure_initialized rebuilds the backend; admission + # verifies that it did. + h.manager._initialized = True + h.manager._session_service = SimpleNamespace() + h.manager._runner = SimpleNamespace(name="runner") + h.manager._root_agent = SimpleNamespace(name="a") + + h.manager._ensure_initialized = _ensure + assert await asyncio.wait_for(h.drain("sess-x"), timeout=2) == ["sess-x"] + assert inits == [1] + finally: + ctx.__exit__(None, None, None) From 61b7950cd414bf2048cce2149e01efd228312268 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:46 -0400 Subject: [PATCH 08/11] fix(cli): make the turn boundary safe; no harness-level replay MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ADK appends the turn's input — the user message, or a resumed ``FunctionResponse`` — to the session while setting up the invocation, before the first event. There is therefore no point at which re-running a turn is side-effect free, and the harness's "Retry in Ns?" dialog was replaying turns whose input had already been persisted. The event source is now invoked exactly once and a surfaced rate limit fails the turn explaining that; transient retries belong to the provider client (ADK ``HttpRetryOptions``, Anthropic ``retry_max_attempts``). ``MessageProcessor.process()`` returns a typed ``TurnResult`` (COMPLETED/CANCELLED/FAILED/UNAVAILABLE) instead of a bare bool, so a caller can tell "the user cancelled" from "the turn failed". ``EventType.ERROR`` had no handler at all and was silently swallowed; it is now rendered as it arrives, with ``recoverable=True`` a warning that leaves the outcome to the stream and anything else failing the turn. Cancelling the caller used to leave the child consumer task driving the workflow while the turn was torn down around it; the consumer is now cancelled and awaited *before* the input callback is cleared, so no tool is left asking a question nobody owns. The HITL dialog's ``finally`` reopened a replacement events box even while unwinding, stranding a thinking box on screen — the box now reopens only on success and is finished exactly once. Session-fact extraction moves inside the ``background_init`` context: it needs the live session store and an LLM call, and leaving the context closes the manager first. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/cli/app.py | 7 +- src/agentic_cli/cli/message_processor.py | 368 +++++++++++++++++----- tests/cli/test_turn_boundary.py | 368 ++++++++++++++++++++++ tests/cli/test_turn_retry_safety.py | 177 +++++++++++ tests/integration/test_adk_integration.py | 132 ++------ tests/test_dual_thinking_boxes.py | 22 +- 6 files changed, 877 insertions(+), 197 deletions(-) create mode 100644 tests/cli/test_turn_boundary.py create mode 100644 tests/cli/test_turn_retry_safety.py diff --git a/src/agentic_cli/cli/app.py b/src/agentic_cli/cli/app.py index 7cf270b..504b3ac 100644 --- a/src/agentic_cli/cli/app.py +++ b/src/agentic_cli/cli/app.py @@ -571,8 +571,11 @@ async def handle_input(text: str) -> None: # Run the session - user sees prompt immediately! await self.session.run_async() - # Extract session facts into memory on exit (if enabled) - await self._extract_session_facts_on_exit() + # Extract session facts into memory (if enabled) while the workflow + # is still alive: leaving this context closes the manager, and fact + # extraction needs the live session store and an LLM call. + await self._extract_session_facts_on_exit() + # No save-on-exit: durable session stores persist continuously per turn. logger.info("app_ending") diff --git a/src/agentic_cli/cli/message_processor.py b/src/agentic_cli/cli/message_processor.py index 2a38a97..0a3a4ec 100644 --- a/src/agentic_cli/cli/message_processor.py +++ b/src/agentic_cli/cli/message_processor.py @@ -9,6 +9,7 @@ import asyncio from contextlib import suppress from dataclasses import dataclass, field +from enum import Enum from typing import TYPE_CHECKING, ClassVar from agentic_cli.logging import Loggers, bind_context @@ -64,6 +65,46 @@ def _esc(text: str) -> str: return rich_to_ansi("\n".join(lines)) +# === Turn results === + + +class TurnStatus(str, Enum): + """Outcome of one processed turn. + + - ``COMPLETED`` — the event stream ran to the end. + - ``CANCELLED`` — the user pressed Ctrl+C. + - ``FAILED`` — the workflow raised; ``TurnResult.error`` says why. + - ``UNAVAILABLE`` — the turn never started (workflow not initialized, or a + job whose originating conversation is gone). + """ + + COMPLETED = "completed" + CANCELLED = "cancelled" + FAILED = "failed" + UNAVAILABLE = "unavailable" + + +@dataclass(frozen=True) +class TurnResult: + """Explicit outcome of ``MessageProcessor.process``/``process_resume``. + + Callers that own durable state (the background-job coordinator) must key + off ``delivered`` rather than "the coroutine returned without raising" — + that is what let a failed resume be recorded as delivered and dropped. + """ + + status: TurnStatus + error: str | None = None + # True once the turn produced output or ran a tool, i.e. a replay of the + # whole turn would duplicate side effects. + partial: bool = False + + @property + def delivered(self) -> bool: + """True only when the turn ran to completion.""" + return self.status is TurnStatus.COMPLETED + + # === Event Processing State === @@ -81,6 +122,13 @@ class _EventProcessingState: in_hitl: bool = False thinking_content: list[str] = field(default_factory=list) response_content: list[str] = field(default_factory=list) + # Set once the turn emits assistant text or runs a tool: the turn is then + # observably partial, which the caller records on a failed/cancelled + # TurnResult. (It is not a retry gate — the harness never replays a turn.) + side_effects_seen: bool = False + # First non-recoverable ERROR event seen. The stream may keep going (the + # backend decides), but the turn's outcome is FAILED, never delivered. + fatal_error: str | None = None # Prevents double-counting when LangGraph emits both CONTEXT_TRIMMED and LLM_USAGE _context_trimmed_this_invocation: bool = False @@ -88,12 +136,6 @@ def get_status(self) -> str: """Return the current status line for the events thinking box.""" return self.status_line - def reset_for_retry(self) -> None: - """Reset state for retry after rate limit.""" - self.status_line = "Processing..." - self.thinking_content.clear() - self.response_content.clear() - # === Message Processor === @@ -145,7 +187,7 @@ async def process( settings: "BaseSettings", usage_tracker: "UsageTracker | None" = None, session_id: str | None = None, - ) -> None: + ) -> TurnResult: """Process a user message through the workflow. Args: @@ -157,6 +199,10 @@ async def process( session_id: Session to run the turn in. Passed explicitly so every turn targets the app's durable session rather than the manager's fallback (which would collapse unnamed runs into one session). + + Returns: + The turn's outcome; ``TurnResult.delivered`` is True only when the + event stream ran to completion. """ # Wait for initialization if needed if not await workflow_controller.ensure_initialized(ui): @@ -164,7 +210,9 @@ async def process( "Cannot process message - workflow not initialized. " "Please check your API keys (GOOGLE_API_KEY or ANTHROPIC_API_KEY)." ) - return + return TurnResult( + TurnStatus.UNAVAILABLE, error="workflow not initialized" + ) bind_context(user_id=settings.default_user) logger.info("handling_message", message_length=len(message)) @@ -176,7 +224,9 @@ def _source(workflow): session_id=session_id, ) - await self._run_turn(_source, workflow_controller, ui, settings, usage_tracker) + return await self._run_turn( + _source, workflow_controller, ui, settings, usage_tracker + ) async def process_resume( self, @@ -185,12 +235,12 @@ async def process_resume( ui: "ThinkingPromptSession", settings: "BaseSettings", usage_tracker: "UsageTracker | None" = None, - ) -> None: + ) -> TurnResult: """Resume the agent with a finished long-running job's result. Streams ``workflow.resume_with_job_result(record)`` through the exact same rendering path as a user turn (events box, tool results, token - accounting, Ctrl+C). A no-op if the backend can't resume. + accounting, Ctrl+C). Args: record: The terminal JobRecord to resume from. @@ -198,9 +248,16 @@ async def process_resume( ui: UI session for output. settings: Application settings. usage_tracker: Optional tracker for accumulating LLM token usage. + + Returns: + The turn's outcome. ``UNAVAILABLE`` when the workflow is not ready + or the originating conversation is gone; the caller must not record + the job as delivered unless ``TurnResult.delivered`` is True. """ if not await workflow_controller.ensure_initialized(ui): - return + return TurnResult( + TurnStatus.UNAVAILABLE, error="workflow not initialized" + ) workflow = workflow_controller.workflow bind_context(user_id=settings.default_user) @@ -224,7 +281,10 @@ async def process_resume( f"({record.state.value}) while its conversation was unavailable " f"— fetch the result with /jobs {record.job_id}.", ) - return + return TurnResult( + TurnStatus.UNAVAILABLE, + error="originating conversation is no longer available", + ) logger.info("resuming_job", job_id=record.job_id, state=record.state.value) ui.add_message( @@ -236,7 +296,9 @@ async def process_resume( def _source(wf): return wf.resume_with_job_result(record) - await self._run_turn(_source, workflow_controller, ui, settings, usage_tracker) + return await self._run_turn( + _source, workflow_controller, ui, settings, usage_tracker + ) async def _run_turn( self, @@ -245,13 +307,43 @@ async def _run_turn( ui: "ThinkingPromptSession", settings: "BaseSettings", usage_tracker: "UsageTracker | None" = None, - ) -> None: + ) -> TurnResult: """Drive one turn from an event-source factory through the UI. Shared by ``process`` (user message) and ``process_resume`` (job result). ``source_factory(workflow)`` returns the WorkflowEvent async generator to consume; everything else (events box, HITL callback, Ctrl+C cancel, rate-limit retry, token accounting) is identical. + + The event source is invoked **exactly once**. The harness never replays + a turn: ADK accepts and persists the input (the user message, or a + resumed ``FunctionResponse``) into the session while setting up the + invocation — before the first event is yielded — so there is no point + at which re-running ``source_factory`` is side-effect free. Even a 429 + on the first model call leaves that input in the session, and a replay + would duplicate the turn and repeat any tool calls it made. + + Retrying belongs at the provider/model-client boundary, which can + prove nothing was accepted: ADK's ``HttpRetryOptions`` (configured in + ``_get_generate_content_config``) retries transient 5xx inside the + client, and the Anthropic client retries per ``retry_max_attempts``. + A rate limit that still surfaces here fails the turn explicitly. + + Cancellation is symmetric: whether the *user* cancels (Ctrl+C) or the + *caller* cancels this coroutine, the consumer task is cancelled and + awaited to completion **before** the HITL callback and turn state are + torn down — otherwise a tool could still be running, and asking for + input, with nothing left to answer it. + + Errors reported as ``EventType.ERROR`` are rendered as they arrive. A + recoverable one is a warning and the stream decides the outcome; a + non-recoverable one makes the turn ``FAILED`` (``delivered`` False) + even if the stream then ends normally, so a caller owning durable + state does not record a failed delivery as delivered. + + Returns: + The turn's outcome as a :class:`TurnResult`. ``partial`` reports + whether the turn had already emitted output or run a tool. """ state = _EventProcessingState( usage_tracker=usage_tracker, @@ -275,104 +367,169 @@ async def _run_turn( # (state.get_status) drives its display. events_ctx: "ThinkingContext | None" = None + def _finish_events_box() -> None: + """Finish the events box if it is open. Idempotent by construction. + + Every path goes through this, so a box is finished exactly once: + reopening one that was already closed (or closing one twice) is + visible to the user as a stray or duplicated panel. + """ + nonlocal events_ctx + if state.thinking_started and events_ctx is not None: + events_ctx.finish(add_to_history=False) + events_ctx = None + state.thinking_started = False + + def _open_events_box() -> None: + nonlocal events_ctx + events_ctx = ui.start_thinking(state.get_status, content_format="ansi") + state.thinking_started = True + # Set up direct callback so HITL tools can prompt the user without # deadlocking the workflow runner. async def _handle_input(request: "UserInputRequest") -> str: - nonlocal events_ctx # Mark the HITL window before tearing down the events box so the # cancel watcher doesn't read "no active boxes" as a Ctrl+C. state.in_hitl = True + _finish_events_box() try: - if state.thinking_started and events_ctx is not None: - events_ctx.finish(add_to_history=False) - state.thinking_started = False - response = await self._prompt_user_input(request, ui) - finally: - events_ctx = ui.start_thinking( - state.get_status, content_format="ansi" - ) - state.thinking_started = True + except BaseException: + # The turn is unwinding (cancelled, or the dialog failed). + # Reopening the events box here would leave a panel on screen + # that nothing downstream will ever finish. state.in_hitl = False + raise + _open_events_box() + state.in_hitl = False return response + # The callback is context-local, so installing and clearing it here + # affects only this turn; no token round-trip is needed (and a manager + # implementing the older no-argument clear stays compatible). workflow.set_input_callback(_handle_input) + result = TurnResult(TurnStatus.FAILED, error="turn did not run") + proc_task: "asyncio.Task[None] | None" = None try: - while True: - try: - events_ctx = ui.start_thinking( - state.get_status, content_format="ansi" + try: + _open_events_box() + + # Consume the event stream in a cancellable task so Ctrl+C + # can abort an in-flight run. thinking_prompt's Ctrl+C + # binding finishes all thinking boxes (so ui.is_thinking + # flips False) but never cancels our coroutine, so we watch + # for that and cancel the task ourselves. + async def _consume() -> None: + async for event in source_factory(workflow): + handler = dispatch.get(event.type) + if handler is not None: + await handler( + self, event, state, ui, settings, workflow + ) + + proc_task = asyncio.create_task(_consume()) + if await self._watch_for_cancel(proc_task, ui, state): + # Ctrl+C already finished every active box; just drop + # our now-dead references so the next turn starts clean + # (finishing them again would double-close). + events_ctx = None + state.thinking_started = False + self._task_box = None + self._last_task_content = None + ui.add_warning("Cancelled.") + workflow_controller.update_status_bar(ui) + logger.info("message_cancelled_by_user") + result = TurnResult( + TurnStatus.CANCELLED, + error="cancelled by user", + partial=state.side_effects_seen, ) - state.thinking_started = True - - # Consume the event stream in a cancellable task so Ctrl+C - # can abort an in-flight run. thinking_prompt's Ctrl+C - # binding finishes all thinking boxes (so ui.is_thinking - # flips False) but never cancels our coroutine, so we watch - # for that and cancel the task ourselves. - async def _consume() -> None: - async for event in source_factory(workflow): - handler = dispatch.get(event.type) - if handler is not None: - await handler( - self, event, state, ui, settings, workflow - ) - - proc_task = asyncio.create_task(_consume()) - if await self._watch_for_cancel(proc_task, ui, state): - # Ctrl+C already finished every active box; just drop - # our now-dead references so the next turn starts clean. - state.thinking_started = False - self._task_box = None - self._last_task_content = None - ui.add_warning("Cancelled.") - workflow_controller.update_status_bar(ui) - logger.info("message_cancelled_by_user") - break - + else: # Finish events box only (don't add status to history) - if state.thinking_started and events_ctx is not None: - events_ctx.finish(add_to_history=False) + _finish_events_box() # Ensure final token counts are reflected in status bar workflow_controller.update_status_bar(ui) - logger.debug("message_handled_successfully") - break # Success — exit retry loop - - except Exception as e: - if state.thinking_started and events_ctx is not None: - events_ctx.finish(add_to_history=False) - state.thinking_started = False - - # Check for 429 rate limit errors — prompt user to wait and retry - from agentic_cli.workflow.retry import ( - is_rate_limit_error, - parse_retry_delay, - ) - - if is_rate_limit_error(e): - delay = parse_retry_delay(e) or 60.0 - retry = await ui.yes_no_dialog( - title="Rate Limited", - text=f"API rate limit reached. Retry in {delay:.0f}s?", + if state.fatal_error is not None: + # The stream ended, but it reported a failure. Never + # "delivered". + logger.info("turn_failed_by_error_event") + result = TurnResult( + TurnStatus.FAILED, + error=state.fatal_error, + partial=state.side_effects_seen, ) - if retry: - ui.add_warning(f"Waiting {delay:.0f}s before retrying...") - await asyncio.sleep(delay) - state.reset_for_retry() - continue # Retry the loop - - # Non-429 or user chose cancel + else: + logger.debug("message_handled_successfully") + result = TurnResult(TurnStatus.COMPLETED) + + except Exception as e: + _finish_events_box() + + from agentic_cli.workflow.retry import is_rate_limit_error + + if is_rate_limit_error(e): + logger.warning("turn_rate_limited", partial=state.side_effects_seen) + ui.add_error( + f"Rate limited: {e}\n" + "The turn was not retried: the backend accepts and " + "persists the input before the first event, so running " + "it again would duplicate this turn (and repeat any " + "tool calls it made). Send the request again once the " + "limit resets." + ) + else: ui.add_error(f"Workflow error: {e}") - break + result = TurnResult( + TurnStatus.FAILED, + error=str(e), + partial=state.side_effects_seen, + ) finally: + # Settle the consumer *first*: it may still be driving the workflow + # (a caller cancelling us does not touch it), and tearing the + # callback down under a live tool would strand a HITL prompt. + await self._settle(proc_task) + # Then close whatever is still open — on a cancelled turn none of + # the paths above ran, and a box left open outlives the turn. + _finish_events_box() workflow.clear_input_callback() # Cache plain-text task content for cold start on next turn # (not get_content() which returns already-richified ANSI) self._last_task_progress = ( self._last_task_content if self._task_box else None ) + return result + + @staticmethod + async def _settle(proc_task: "asyncio.Task[None] | None") -> None: + """Ensure the consumer task is finished before the turn is torn down. + + A no-op on the normal paths (the task is already done). It matters when + *this* coroutine is cancelled while waiting on the task: cancellation + does not propagate into it, so it would keep consuming the workflow + generator after its owner is gone. + + Uses ``asyncio.wait`` rather than awaiting the task, so neither the + task's ``CancelledError`` nor its exception is re-raised out of a + ``finally`` block — the original outcome must survive. + """ + if proc_task is None or proc_task.done(): + MessageProcessor._retrieve_exception(proc_task) + return + proc_task.cancel() + with suppress(asyncio.CancelledError): + await asyncio.wait({proc_task}) + MessageProcessor._retrieve_exception(proc_task) + + @staticmethod + def _retrieve_exception(proc_task: "asyncio.Task[None] | None") -> None: + """Mark a finished task's exception retrieved (no 'never retrieved' log).""" + if proc_task is None or not proc_task.done() or proc_task.cancelled(): + return + with suppress(asyncio.InvalidStateError): + proc_task.exception() async def _watch_for_cancel( self, @@ -465,6 +622,43 @@ async def _handle_text( """Handle TEXT events — stream response to console.""" ui.add_response(event.content, markdown=True) state.response_content.append(event.content) + # Text is already on screen and in the durable session: replaying the + # turn would emit it twice. + state.side_effects_seen = True + + async def _handle_error( + self, + event: "WorkflowEvent", + state: _EventProcessingState, + ui: "ThinkingPromptSession", + settings: "BaseSettings", + workflow: object, + ) -> None: + """Handle ERROR events — render, and fail the turn unless recoverable. + + ``WorkflowEvent.error(..., recoverable=True)`` means the backend + handled it and is carrying on (a retried tool, a degraded feature): it + is surfaced as a warning and the stream still decides the outcome. + Anything else is a failure the caller must see in the ``TurnResult``, + because a background-job coordinator keys durable state off + ``delivered`` — an unrendered, unreported error was recorded as a + successful delivery. + """ + recoverable = bool(event.metadata.get("recoverable", False)) + code = event.metadata.get("error_code") + suffix = f" [{code}]" if code else "" + if recoverable: + ui.add_warning(f"{event.content}{suffix}") + state.status_line = f"! {event.content}" + logger.warning("workflow_error_event", recoverable=True, code=code) + return + + ui.add_error(f"{event.content}{suffix}") + state.status_line = f"x {event.content}" + logger.error("workflow_error_event", recoverable=False, code=code) + if state.fatal_error is None: + # Keep the first: later ones are usually the cascade. + state.fatal_error = event.content async def _handle_thinking( self, @@ -491,6 +685,8 @@ async def _handle_tool_call( """Handle TOOL_CALL events — update status line.""" tool_name = event.metadata.get("tool_name", "unknown") state.status_line = f"Calling: {tool_name}" + # A tool ran; the turn is no longer safe to replay wholesale. + state.side_effects_seen = True # For the stateful executor, show the code being run (syntax-highlighted, # first N lines) so the run is visible, not just a status blip. if tool_name == "sandbox_execute": @@ -510,6 +706,7 @@ async def _handle_tool_result( workflow: object, ) -> None: """Handle TOOL_RESULT events — display result summary.""" + state.side_effects_seen = True tool_name = event.metadata.get("tool_name", "unknown") success = event.metadata.get("success", True) duration = event.metadata.get("duration_ms") @@ -686,6 +883,7 @@ def _get_event_dispatch(cls) -> dict: cls._EVENT_DISPATCH = { EventType.TEXT: cls._handle_text, + EventType.ERROR: cls._handle_error, EventType.THINKING: cls._handle_thinking, EventType.TOOL_CALL: cls._handle_tool_call, EventType.TOOL_RESULT: cls._handle_tool_result, diff --git a/tests/cli/test_turn_boundary.py b/tests/cli/test_turn_boundary.py new file mode 100644 index 0000000..e330cd5 --- /dev/null +++ b/tests/cli/test_turn_boundary.py @@ -0,0 +1,368 @@ +"""The whole turn boundary is safe, not just the happy path. + +Three defects: + +1. Cancelling the caller of ``MessageProcessor`` left the child consumer task + running: ``_run_turn``'s ``finally`` cleared the HITL callback and the turn + state while the workflow generator was still being driven, so a tool could + still be executing with no callback to answer it and no owner to await it. +2. ``EventType.ERROR`` had no handler. A backend that reported a failure as an + event rendered nothing and the turn was still reported ``COMPLETED``, so a + background-job resume recorded a failed delivery as delivered. +3. (See ``tests/workflow/test_turn_serialization.py`` for the HITL callback + being context-local rather than manager-global.) +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from agentic_cli.cli.message_processor import MessageProcessor, TurnStatus +from agentic_cli.workflow.events import EventType, WorkflowEvent + +from tests.event_replay import RecordingSession, RecordingThinkingContext + + +class _UI(RecordingSession): + """RecordingSession plus the dialog surface ``_run_turn`` may reach for.""" + + def __init__(self) -> None: + super().__init__() + self.dialogs = 0 + + async def yes_no_dialog(self, title: str = "", text: str = "") -> bool: + self.dialogs += 1 + return True + + +class _Workflow: + """Tracks whether the HITL callback is currently installed.""" + + def __init__(self) -> None: + self.callback = None + + def set_input_callback(self, cb): + self.callback = cb + return None + + def clear_input_callback(self, token=None) -> None: + self.callback = None + + +def _controller(workflow): + return SimpleNamespace(workflow=workflow, update_status_bar=lambda ui: None) + + +def _settings(): + return SimpleNamespace(verbose_thinking=False, default_user="u") + + +async def _run(processor, ui, source_factory, workflow=None): + workflow = workflow or _Workflow() + return await processor._run_turn( + source_factory, _controller(workflow), ui, _settings(), None + ) + + +class TestCallerCancellation: + """A cancelled caller must not leave the event stream running behind it.""" + + async def test_child_consumer_is_cancelled_and_awaited(self): + started = asyncio.Event() + observed: dict = {"cancelled": False, "closed": False, "callback_at_cancel": "unset"} + workflow = _Workflow() + + def _source(_wf): + async def _gen(): + started.set() + try: + await asyncio.sleep(30) + yield WorkflowEvent(type=EventType.TEXT, content="never") + except asyncio.CancelledError: + observed["cancelled"] = True + observed["callback_at_cancel"] = workflow.callback + raise + finally: + observed["closed"] = True + + return _gen() + + turn = asyncio.create_task( + _run(MessageProcessor(), _UI(), _source, workflow) + ) + await asyncio.wait_for(started.wait(), timeout=2) + + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + + assert observed["cancelled"] is True, "the event stream outlived its caller" + assert observed["closed"] is True, "the consumer was never awaited" + assert observed["callback_at_cancel"] is not None, ( + "the HITL callback was cleared while a tool could still ask for input" + ) + assert workflow.callback is None, "the callback was not cleared afterwards" + + async def test_cancellation_does_not_leave_a_pending_task(self): + """Nothing is left for the loop to garbage-collect mid-flight.""" + started = asyncio.Event() + + def _source(_wf): + async def _gen(): + started.set() + await asyncio.sleep(30) + yield WorkflowEvent(type=EventType.TEXT, content="never") + + return _gen() + + turn = asyncio.create_task(_run(MessageProcessor(), _UI(), _source)) + await asyncio.wait_for(started.wait(), timeout=2) + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + + others = [ + t + for t in asyncio.all_tasks() + if t is not asyncio.current_task() and not t.done() + ] + assert others == [], f"orphaned tasks survived the cancelled turn: {others}" + + +class _CountingContext(RecordingThinkingContext): + """A thinking context that counts how often it was finished.""" + + def __init__(self, session, label: str) -> None: + super().__init__(session, label) + self.finish_count = 0 + + def finish(self, **kwargs) -> None: + self.finish_count += 1 + super().finish(**kwargs) + + +class _HitlUI(_UI): + """Tracks every thinking context and blocks inside the input dialog.""" + + def __init__(self) -> None: + super().__init__() + self.contexts: list[_CountingContext] = [] + self.dialog_open = asyncio.Event() + + def start_thinking(self, *args, **kwargs) -> _CountingContext: + label = kwargs.get("title") or "events" + self.calls.append(("start_thinking", label, {})) + ctx = _CountingContext(self, label) + self.contexts.append(ctx) + return ctx + + async def input_dialog(self, title: str = "", text: str = "", default: str = ""): + self.dialog_open.set() + await asyncio.sleep(30) # the user is still typing when we're cancelled + return "answer" # pragma: no cover + + +class _HitlWorkflow(_Workflow): + """Records the teardown order and drives the installed HITL callback.""" + + def __init__(self, trace: list[str]) -> None: + super().__init__() + self._trace = trace + + def clear_input_callback(self, token=None) -> None: + self._trace.append("callback-cleared") + super().clear_input_callback(token) + + +class TestCancellationDuringHitl: + """Cancelling while a HITL dialog is open must not strand a thinking box. + + The dialog's ``finally`` unconditionally opened a *replacement* events box + — including while the turn was unwinding — so a cancelled HITL turn left a + box on screen that nothing would ever finish. + """ + + def _turn(self): + trace: list[str] = [] + ui = _HitlUI() + workflow = _HitlWorkflow(trace) + + def _source(wf): + async def _gen(): + try: + await wf.callback( + SimpleNamespace( + request_id="r", + tool_name="ask_clarification", + prompt="which?", + input_type=None, + choices=None, + default=None, + ) + ) + yield WorkflowEvent(type=EventType.TEXT, content="never") + finally: + trace.append("consumer-done") + + return _gen() + + return trace, ui, workflow, _source + + async def test_every_thinking_context_is_finished_exactly_once(self): + trace, ui, workflow, source = self._turn() + + turn = asyncio.create_task(_run(MessageProcessor(), ui, source, workflow)) + await asyncio.wait_for(ui.dialog_open.wait(), timeout=2) + + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + + assert ui.contexts, "no thinking context was ever opened" + counts = [ctx.finish_count for ctx in ui.contexts] + assert counts == [1] * len(ui.contexts), ( + f"thinking contexts were not finished exactly once: {counts}" + ) + + async def test_consumer_settles_before_the_callback_is_cleared(self): + trace, ui, workflow, source = self._turn() + + turn = asyncio.create_task(_run(MessageProcessor(), ui, source, workflow)) + await asyncio.wait_for(ui.dialog_open.wait(), timeout=2) + + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + + assert trace.index("consumer-done") < trace.index("callback-cleared") + + async def test_normal_hitl_turn_still_reopens_the_events_box(self): + """The replacement box is right on the *success* path — keep it.""" + ui = _HitlUI() + ui.input_dialog = _answering_dialog + workflow = _HitlWorkflow([]) + + def _source(wf): + async def _gen(): + answer = await wf.callback( + SimpleNamespace( + request_id="r", tool_name="t", prompt="p", + input_type=None, choices=None, default=None, + ) + ) + yield WorkflowEvent(type=EventType.TEXT, content=answer) + + return _gen() + + result = await _run(MessageProcessor(), ui, _source, workflow) + + assert result.status is TurnStatus.COMPLETED + assert "answer" in ui.responses() + assert len(ui.contexts) == 2, "the events box was not reopened after the dialog" + assert [ctx.finish_count for ctx in ui.contexts] == [1, 1] + + async def test_cancel_outside_hitl_finishes_the_events_box(self): + started = asyncio.Event() + ui = _HitlUI() + workflow = _HitlWorkflow([]) + + def _source(_wf): + async def _gen(): + started.set() + await asyncio.sleep(30) + yield WorkflowEvent(type=EventType.TEXT, content="never") + + return _gen() + + turn = asyncio.create_task(_run(MessageProcessor(), ui, _source, workflow)) + await asyncio.wait_for(started.wait(), timeout=2) + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + + assert [ctx.finish_count for ctx in ui.contexts] == [1] + + +async def _answering_dialog(title: str = "", text: str = "", default: str = ""): + return "answer" + + +class _EventSource: + """Yields a fixed list of events, once.""" + + def __init__(self, *events: WorkflowEvent) -> None: + self._events = list(events) + self.invocations = 0 + + def __call__(self, _workflow): + self.invocations += 1 + + async def _gen(): + for event in self._events: + yield event + + return _gen() + + +class TestErrorEvents: + """An ERROR event is rendered, and a fatal one fails the turn.""" + + async def test_non_recoverable_error_is_rendered(self): + ui = _UI() + await _run( + MessageProcessor(), + ui, + _EventSource(WorkflowEvent.error("model refused the request")), + ) + assert any("model refused the request" in e for e in ui.errors()) + + async def test_non_recoverable_error_fails_the_turn(self): + result = await _run( + MessageProcessor(), + _UI(), + _EventSource(WorkflowEvent.error("backend exploded")), + ) + assert result.status is TurnStatus.FAILED + assert result.delivered is False + assert "backend exploded" in (result.error or "") + + async def test_error_after_output_is_reported_partial(self): + result = await _run( + MessageProcessor(), + _UI(), + _EventSource( + WorkflowEvent(type=EventType.TEXT, content="half an answer"), + WorkflowEvent.error("then it died"), + ), + ) + assert result.status is TurnStatus.FAILED + assert result.partial is True + + async def test_recoverable_error_is_rendered_and_the_turn_completes(self): + """A recoverable error is informational: the stream owns the outcome.""" + ui = _UI() + result = await _run( + MessageProcessor(), + ui, + _EventSource( + WorkflowEvent.error("one tool retried", recoverable=True), + WorkflowEvent(type=EventType.TEXT, content="done anyway"), + ), + ) + assert any("one tool retried" in w for w in ui.warnings()) + assert result.status is TurnStatus.COMPLETED + assert result.delivered is True + + async def test_first_fatal_error_is_the_reported_one(self): + result = await _run( + MessageProcessor(), + _UI(), + _EventSource( + WorkflowEvent.error("first failure"), + WorkflowEvent.error("cascade"), + ), + ) + assert result.error == "first failure" diff --git a/tests/cli/test_turn_retry_safety.py b/tests/cli/test_turn_retry_safety.py new file mode 100644 index 0000000..ba5df96 --- /dev/null +++ b/tests/cli/test_turn_retry_safety.py @@ -0,0 +1,177 @@ +"""The harness never replays a turn. + +``_run_turn`` used to re-invoke the event-source factory after a 429, first +unconditionally and then "only before the first visible event". Both are wrong: +ADK's ``Runner`` appends the input to the session while setting up the +invocation — before any event is yielded — so a 429 on the very first model +call has already persisted the user message (or the resumed +``FunctionResponse``). Replaying duplicates the turn. + +Retrying now happens only inside the provider client, which can prove nothing +was accepted (ADK ``HttpRetryOptions`` for transient 5xx). A surfaced rate limit +fails the turn with an explicit message. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from agentic_cli.cli.message_processor import ( + MessageProcessor, + TurnResult, + TurnStatus, +) +from agentic_cli.workflow.events import EventType, WorkflowEvent + +from tests.event_replay import RecordingSession + + +class _RateLimited(Exception): + """Shaped like a provider 429 for ``is_rate_limit_error``.""" + + def __init__(self) -> None: + super().__init__("429 RESOURCE_EXHAUSTED: rate limit exceeded") + + +class _UI(RecordingSession): + """RecordingSession plus the dialog surface ``_run_turn`` may reach for.""" + + def __init__(self) -> None: + super().__init__() + self.dialogs = 0 + + async def yes_no_dialog(self, title: str = "", text: str = "") -> bool: + self.dialogs += 1 + return True + + +def _controller(): + workflow = MagicMock() + workflow.set_input_callback = MagicMock() + workflow.clear_input_callback = MagicMock() + return SimpleNamespace(workflow=workflow, update_status_bar=lambda ui: None) + + +def _settings(): + return SimpleNamespace(verbose_thinking=False, default_user="u") + + +def _tool_call_event() -> WorkflowEvent: + return WorkflowEvent( + type=EventType.TOOL_CALL, + content="calling", + metadata={"tool_name": "write_file", "tool_args": {}}, + ) + + +async def _run(processor, ui, source_factory): + return await processor._run_turn( + source_factory, _controller(), ui, _settings(), None + ) + + +class _DurableSource: + """An event source that persists its input before yielding anything. + + Mirrors ADK: ``Runner.run_async`` appends ``new_message`` to the session + during invocation setup, so the append happens even when the first model + call raises. + """ + + def __init__(self, *, events=(), raises=None) -> None: + self.appended: list[str] = [] + self._events = list(events) + self._raises = raises + + def __call__(self, workflow): + self.appended.append("input") + + async def _gen(): + for event in self._events: + yield event + if self._raises is not None: + raise self._raises + + return _gen() + + +class TestNoReplay: + async def test_rate_limit_before_any_event_does_not_replay(self): + """The 'nothing ran yet' boundary does not exist — the input is stored.""" + source = _DurableSource(raises=_RateLimited()) + ui = _UI() + + result = await _run(MessageProcessor(), ui, source) + + assert source.appended == ["input"], "the turn was replayed" + assert ui.dialogs == 0, "the user was offered a retry that is not safe" + assert result.status is TurnStatus.FAILED + assert result.partial is False + + async def test_rate_limit_after_a_tool_ran_does_not_replay(self): + source = _DurableSource(events=[_tool_call_event()], raises=_RateLimited()) + ui = _UI() + + result = await _run(MessageProcessor(), ui, source) + + assert source.appended == ["input"] + assert result.status is TurnStatus.FAILED + assert result.partial is True + + async def test_resume_source_is_invoked_exactly_once(self): + """A resumed FunctionResponse is persisted too — never re-delivered.""" + source = _DurableSource( + events=[WorkflowEvent(type=EventType.TEXT, content="partial")], + raises=_RateLimited(), + ) + ui = _UI() + + result = await _run(MessageProcessor(), ui, source) + + assert source.appended == ["input"], "the job result was delivered twice" + assert result.delivered is False + + async def test_non_rate_limit_failure_also_runs_once(self): + source = _DurableSource(raises=RuntimeError("boom")) + ui = _UI() + + result = await _run(MessageProcessor(), ui, source) + + assert source.appended == ["input"] + assert result.status is TurnStatus.FAILED + assert "boom" in (result.error or "") + + +class TestRateLimitReporting: + async def test_error_explains_the_turn_was_not_replayed(self): + ui = _UI() + result = await _run(MessageProcessor(), ui, _DurableSource(raises=_RateLimited())) + + errors = " ".join(str(e) for e in ui.errors()) + assert "Rate limited" in errors + assert "not retried" in errors.lower() + assert result.error and "429" in result.error + + async def test_no_retry_dialog_is_ever_shown(self): + ui = _UI() + await _run(MessageProcessor(), ui, _DurableSource(raises=_RateLimited())) + assert ui.dialogs == 0 + + +class TestTurnResultContract: + async def test_success_is_delivered(self): + source = _DurableSource( + events=[WorkflowEvent(type=EventType.TEXT, content="hi")] + ) + result = await _run(MessageProcessor(), _UI(), source) + + assert result == TurnResult(TurnStatus.COMPLETED) + assert result.delivered is True + assert source.appended == ["input"] + + async def test_failure_is_not_delivered(self): + result = await _run( + MessageProcessor(), _UI(), _DurableSource(raises=RuntimeError("boom")) + ) + assert result.delivered is False diff --git a/tests/integration/test_adk_integration.py b/tests/integration/test_adk_integration.py index 1c4ccd2..e3224f4 100644 --- a/tests/integration/test_adk_integration.py +++ b/tests/integration/test_adk_integration.py @@ -523,148 +523,78 @@ async def test_no_auto_clear_with_pending(self): class TestMessageProcessorRateLimit: - """Tests that MessageProcessor handles 429 errors with user prompt.""" + """A surfaced 429 fails the turn; the harness never replays it. - async def test_rate_limit_retry_on_user_accept(self): - """When user accepts retry, processor waits and retries.""" - from agentic_cli.cli.message_processor import MessageProcessor - - processor = MessageProcessor() + ADK appends the user message to the session while setting up the + invocation, so even a 429 raised before the first event has already + persisted the turn's input. Re-invoking the source would duplicate it. + Retrying belongs to the provider client (HttpRetryOptions). + """ - # Mock workflow controller + def _harness(self): workflow_controller = MagicMock() workflow_controller.ensure_initialized = AsyncMock(return_value=True) - # First call raises 429, second call succeeds - call_count = 0 + calls = {"n": 0} async def mock_process(**kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - error = Exception("RESOURCE_EXHAUSTED: retry in 5s") - error.code = 429 - raise error - # Second call: yield a text event - yield WorkflowEvent.text("Success!", "session") + calls["n"] += 1 + error = Exception("RESOURCE_EXHAUSTED: retry in 5s") + error.code = 429 + raise error + yield # pragma: no cover - makes it an async generator mock_workflow = MagicMock() mock_workflow.process = mock_process workflow_controller.workflow = mock_workflow - # Mock UI ui = MagicMock() - ctx_mock = MagicMock() - ui.start_thinking.return_value = ctx_mock + ui.start_thinking.return_value = MagicMock() ui.add_response = MagicMock() ui.add_warning = MagicMock() ui.add_error = MagicMock() ui.add_rich = MagicMock() - ui.yes_no_dialog = AsyncMock(return_value=True) # User accepts retry + ui.yes_no_dialog = AsyncMock(return_value=True) - # Mock settings settings = MagicMock() settings.default_user = "test-user" settings.verbose_thinking = False + return workflow_controller, ui, settings, calls - with patch("agentic_cli.cli.message_processor.asyncio.sleep", new_callable=AsyncMock) as mock_sleep: - await processor.process( - message="test", - workflow_controller=workflow_controller, - ui=ui, - settings=settings, - ) - - # Verify retry happened - assert call_count == 2 - ui.yes_no_dialog.assert_called_once() - mock_sleep.assert_called_once_with(5.0) - ui.add_warning.assert_called_once() - # Success path should have been reached - ui.add_response.assert_called_once_with("Success!", markdown=True) - ui.add_error.assert_not_called() - - async def test_rate_limit_cancel_on_user_decline(self): - """When user declines retry, processor shows error and stops.""" - from agentic_cli.cli.message_processor import MessageProcessor - - processor = MessageProcessor() - - workflow_controller = MagicMock() - workflow_controller.ensure_initialized = AsyncMock(return_value=True) + async def test_rate_limited_turn_is_not_replayed(self): + from agentic_cli.cli.message_processor import MessageProcessor, TurnStatus - async def mock_process(**kwargs): - error = Exception("RESOURCE_EXHAUSTED: retry in 30s") - error.code = 429 - raise error - yield # make it an async generator # noqa: E501 - - mock_workflow = MagicMock() - mock_workflow.process = mock_process - workflow_controller.workflow = mock_workflow - - ui = MagicMock() - ctx_mock = MagicMock() - ui.start_thinking.return_value = ctx_mock - ui.add_error = MagicMock() - ui.add_rich = MagicMock() - ui.yes_no_dialog = AsyncMock(return_value=False) # User declines - - settings = MagicMock() - settings.default_user = "test-user" - settings.verbose_thinking = False + workflow_controller, ui, settings, calls = self._harness() - await processor.process( + result = await MessageProcessor().process( message="test", workflow_controller=workflow_controller, ui=ui, settings=settings, ) - ui.yes_no_dialog.assert_called_once() - ui.add_error.assert_called_once() - assert "Workflow error" in ui.add_error.call_args[0][0] + assert calls["n"] == 1, "the turn was replayed after a rate limit" + ui.yes_no_dialog.assert_not_called() + ui.add_response.assert_not_called() + assert result.status is TurnStatus.FAILED + assert result.delivered is False - async def test_non_rate_limit_error_not_retried(self): - """Non-429 errors are not retried, shown as workflow error.""" + async def test_rate_limit_error_explains_no_replay(self): from agentic_cli.cli.message_processor import MessageProcessor - processor = MessageProcessor() - - workflow_controller = MagicMock() - workflow_controller.ensure_initialized = AsyncMock(return_value=True) - - async def mock_process(**kwargs): - raise RuntimeError("Something broke") - yield # noqa: E501 + workflow_controller, ui, settings, _calls = self._harness() - mock_workflow = MagicMock() - mock_workflow.process = mock_process - workflow_controller.workflow = mock_workflow - - ui = MagicMock() - ctx_mock = MagicMock() - ui.start_thinking.return_value = ctx_mock - ui.add_error = MagicMock() - ui.add_rich = MagicMock() - ui.yes_no_dialog = AsyncMock() - - settings = MagicMock() - settings.default_user = "test-user" - settings.verbose_thinking = False - - await processor.process( + await MessageProcessor().process( message="test", workflow_controller=workflow_controller, ui=ui, settings=settings, ) - # Should NOT prompt user for retry - ui.yes_no_dialog.assert_not_called() ui.add_error.assert_called_once() - assert "Something broke" in ui.add_error.call_args[0][0] - + message = ui.add_error.call_args[0][0] + assert "Rate limited" in message + assert "not retried" in message.lower() class TestUserInputCallback: """Tests for the registered-callback path in request_user_input. diff --git a/tests/test_dual_thinking_boxes.py b/tests/test_dual_thinking_boxes.py index d51d9cf..5daa06e 100644 --- a/tests/test_dual_thinking_boxes.py +++ b/tests/test_dual_thinking_boxes.py @@ -48,16 +48,20 @@ def test_get_status_default(self): state = _EventProcessingState() assert state.get_status() == "Processing..." - def test_reset_for_retry_does_not_clear_task_fields(self): - """reset_for_retry() should not reference task display fields.""" + def test_no_retry_reset_hook_remains(self): + """The turn is never replayed, so there is no retry-reset state hook. + + ``reset_for_retry()`` existed only to re-run a turn after a 429; the + harness no longer does that (ADK has already persisted the input). + """ + assert not hasattr(_EventProcessingState(), "reset_for_retry") + + def test_side_effects_flag_tracks_visible_progress(self): + """It reports whether the turn got far enough to be observably partial.""" state = _EventProcessingState() - state.status_line = "Something" - state.thinking_content.append("thought") - state.response_content.append("response") - state.reset_for_retry() - assert state.status_line == "Processing..." - assert state.thinking_content == [] - assert state.response_content == [] + assert state.side_effects_seen is False + state.side_effects_seen = True + assert state.side_effects_seen is True # --------------------------------------------------------------------------- From a152c1b62ca2ab78b867ebc50065d2b68a333291 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:47 -0400 Subject: [PATCH 09/11] fix(jobs): make the whole persisted job record coherent across processes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The jobs directory is user-scoped, so two CLIs routinely hold the same records — and every mutator acted on its own in-memory snapshot. Delivery has an explicit lifecycle: ``JobRecord.resumed`` (a bool set *before* the turn ran) becomes ``resume_state`` (pending → resuming → delivered|failed) with ``resume_error`` and ``resume_owner``, claimed with ``begin_resume()`` and closed with ``complete_resume()`` on success, failure *and* cancellation. Both transitions happen under a cross-process ``flock`` against the record on disk, so two CLIs can no longer deliver one result into two conversations. A record found ``resuming`` at startup is recovered as failed rather than replayed — the interrupted turn may already have run tools. Execution is claimed the same way: ``exec_owner`` (``::``) is written *before* the backend is started, so a queued job cannot be launched twice, and an interrupted launch is failed rather than replayed. Every other metadata write reloads the durable record first — a plain state write carried this manager's stale resume fields and erased another process's live claim — and terminal transitions are monotonic, so a stale snapshot cannot rewrite a recorded success as CANCELLED. Reading distinguishes *deleted* from *unreadable*: a record another manager cleaned away is forgotten rather than resurrected, while an unparseable one is left untouched and fails closed. Startup recovery honours the same distinction rather than writing a verdict decided on state it failed to read. A foreign job is polled (backends publish outcomes durably) but its ``UNKNOWN`` — "I hold no handle for this" — is ignored, so an observer sees a job finish without marking a healthy one terminal. Whether a foreign job can be *cancelled* is now a declared backend capability rather than inferred from restart-safety, which answers a different question. ``InProcessBackend.close()`` no longer cancels submitted work. A queued job is already durably RUNNING under a live owner, and a cancelled future writes no exit-code sentinel — so no manager could ever resolve it, and it stayed RUNNING forever. If the cross-process lock cannot be taken, claims, launches, recovery, reconcile, clean and cancel all fail closed rather than writing unsynchronized. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/cli/app.py | 77 +- src/agentic_cli/tools/jobs/__init__.py | 9 +- src/agentic_cli/tools/jobs/backends.py | 37 + src/agentic_cli/tools/jobs/manager.py | 942 ++++++++++++++- tests/cli/test_job_monitor.py | 3 +- tests/cli/test_resume_coordinator.py | 257 +++- tests/tools/test_jobs.py | 1515 +++++++++++++++++++++++- tests/workflow/test_adk_job_resume.py | 15 +- 8 files changed, 2784 insertions(+), 71 deletions(-) diff --git a/src/agentic_cli/cli/app.py b/src/agentic_cli/cli/app.py index 504b3ac..413ec1f 100644 --- a/src/agentic_cli/cli/app.py +++ b/src/agentic_cli/cli/app.py @@ -476,9 +476,20 @@ async def resume_finished_jobs(self) -> int: """Resume the agent for each finished, resume-flagged background job. Each resume is a serialized turn (via the turn lock) so it never - overlaps a user turn or another resume. Returns the number resumed. - Used at turn boundaries (auto, gated) and by the /resume command - (explicit, ungated). + overlaps a user turn or another resume. Used at turn boundaries (auto, + gated) and by the /resume command (explicit, ungated). + + Delivery follows the job's resume lifecycle: the job is *claimed* + (pending → resuming) before the turn runs, so a crash cannot silently + re-deliver it, and only recorded delivered once the turn actually + completed — a failed or cancelled resume is recorded as such instead of + being dropped. + + Returns: + The number of jobs a resume turn was run for (unchanged meaning: + "how many were picked up"). Whether each one was delivered is + recorded on the job and reported by ``/jobs``; a failed delivery is + surfaced to the user by the turn itself. """ if not self._workflow_controller.is_ready: return 0 @@ -486,19 +497,67 @@ async def resume_finished_jobs(self) -> int: if jm is None: return 0 - records = jm.awaiting_resume() - for record in records: - # Mark before running so a crash mid-resume can't double-fire. - jm.mark_resumed(record.job_id) + attempted = 0 + for record in jm.awaiting_resume(): + # Claim first: durable, so an interrupted delivery is recoverable + # and two coordinators can't both deliver the same result. + if not jm.begin_resume(record.job_id): + continue + attempted += 1 + await self._deliver_resume(jm, record) + return attempted + + async def _deliver_resume(self, jm, record) -> bool: + """Run one claimed job's resume turn and close out its transition. + + Every exit path closes the claim, so no record can be left ``RESUMING`` + in this process: a normal outcome records delivered/failed, an + exception records failed with the reason, and a cancellation records + failed before re-raising (cancellation still propagates). + + Args: + jm: The JobManager holding the claim. + record: The claimed job record. + + Returns: + True if the result was delivered to the agent. + """ + try: async with self._turn_lock: - await self._message_processor.process_resume( + result = await self._message_processor.process_resume( record=record, workflow_controller=self._workflow_controller, ui=self.session, settings=self._settings, usage_tracker=self._usage_tracker, ) - return len(records) + except asyncio.CancelledError: + jm.complete_resume( + record.job_id, delivered=False, error="resume cancelled" + ) + logger.info("job_resume_cancelled", job_id=record.job_id) + raise + except Exception as exc: # noqa: BLE001 - the claim must always close + jm.complete_resume(record.job_id, delivered=False, error=str(exc)) + logger.warning( + "job_resume_raised", job_id=record.job_id, error=str(exc) + ) + self.session.add_error( + f"Background job '{record.name}' could not be resumed: {exc}" + ) + return False + + jm.complete_resume( + record.job_id, delivered=result.delivered, error=result.error + ) + if not result.delivered: + logger.info( + "job_resume_not_delivered", + job_id=record.job_id, + status=result.status.value, + error=result.error, + ) + return result.delivered async def _adopt_session_on_startup(self) -> None: """Adopt this run's session id so the manager targets it from turn one. diff --git a/src/agentic_cli/tools/jobs/__init__.py b/src/agentic_cli/tools/jobs/__init__.py index 651dd91..09943c8 100644 --- a/src/agentic_cli/tools/jobs/__init__.py +++ b/src/agentic_cli/tools/jobs/__init__.py @@ -13,7 +13,12 @@ SubprocessBackend, default_backends, ) -from agentic_cli.tools.jobs.manager import JobManager, JobRecord +from agentic_cli.tools.jobs.manager import ( + JobManager, + JobRecord, + ResumeState, + ResumeStateError, +) from agentic_cli.tools.jobs.tools import ( job_cancel, job_list, @@ -26,6 +31,8 @@ "JobManager", "JobRecord", "JobState", + "ResumeState", + "ResumeStateError", "JobBackend", "SubprocessBackend", "InProcessBackend", diff --git a/src/agentic_cli/tools/jobs/backends.py b/src/agentic_cli/tools/jobs/backends.py index 41a383a..da8cc2b 100644 --- a/src/agentic_cli/tools/jobs/backends.py +++ b/src/agentic_cli/tools/jobs/backends.py @@ -87,6 +87,13 @@ class JobBackend(ABC): name: str = "base" survives_restart: bool = False streams_logs: bool = False + # Whether a *different* manager can cancel a job this backend started. + # True only when the job is addressed through a durable handle (a pid, a + # remote id) rather than an object in the starter's memory — otherwise + # "cancelled" would be recorded while the job kept running elsewhere. + # Distinct from ``survives_restart``: a backend can publish a readable + # outcome without being remotely controllable. + cancels_foreign_jobs: bool = False @abstractmethod def start(self, record: "JobRecord", job_dir: Path) -> None: @@ -111,6 +118,14 @@ def result(self, record: "JobRecord", job_dir: Path) -> Any: """Return the job's result (backend-specific).""" return {"exit_code": record.exit_code, "stdout_tail": self.logs(record, job_dir, 20, "stdout")} + def close(self) -> None: + """Release resources the backend owns. Idempotent; default no-op. + + Backends that own OS resources (thread pools, connections) override + this; running work is *not* cancelled — see each backend's docstring. + """ + return None + class SubprocessBackend(JobBackend): """Detached subprocess; restart-safe via an on-disk ``exit_code`` sentinel.""" @@ -118,6 +133,7 @@ class SubprocessBackend(JobBackend): name = "subprocess" survives_restart = True streams_logs = True + cancels_foreign_jobs = True # addressed by pid def __init__(self) -> None: # job_id -> Popen, kept so cancel() can signal the process group. @@ -205,6 +221,7 @@ class InProcessBackend(JobBackend): name = "inprocess" survives_restart = False streams_logs = False + cancels_foreign_jobs = False # the Future lives in the starting manager def __init__(self, max_workers: int = 8) -> None: self._pool = ThreadPoolExecutor( @@ -254,6 +271,26 @@ def cancel(self, record: "JobRecord", job_dir: Path) -> None: if fut is not None: fut.cancel() # only succeeds if not yet started; running threads continue + def close(self) -> None: + """Stop accepting new work; let everything already submitted finish. + + Idempotent, and does not block. Submitted work is deliberately **not** + cancelled, per ``JobManager.close()``'s contract — here that is a + correctness requirement, not just a courtesy. A job's outcome is + published only by ``_run`` writing the ``exit_code`` sentinel, and its + record is already durably RUNNING under a live owner by the time it is + queued. Dropping the future left a record no manager could ever + resolve: this one no longer holds the future, and any other reads + UNKNOWN from a live foreign owner and (correctly) declines to believe + it — so the job stayed RUNNING forever and was never deliverable. + + A thread already running a job cannot be interrupted anyway; queued + jobs now share that fate. The cost is bounded by the queue: the pool's + threads are joined at interpreter exit, so a long backlog delays + process exit rather than being silently discarded. + """ + self._pool.shutdown(wait=False) + def result(self, record: "JobRecord", job_dir: Path) -> Any: result_file = job_dir / "result.json" if result_file.exists(): diff --git a/src/agentic_cli/tools/jobs/manager.py b/src/agentic_cli/tools/jobs/manager.py index 530ed21..d482236 100644 --- a/src/agentic_cli/tools/jobs/manager.py +++ b/src/agentic_cli/tools/jobs/manager.py @@ -11,12 +11,16 @@ from __future__ import annotations +import contextlib +import os +import socket import threading import time import uuid from dataclasses import asdict, dataclass, field +from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Iterator from agentic_cli.file_utils import atomic_write_json from agentic_cli.logging import Loggers @@ -32,11 +36,181 @@ logger = Loggers.workflow() +# Persisted values of the terminal states, for comparing against raw JSON. +_TERMINAL_VALUES = frozenset(state.value for state in TERMINAL_STATES) + def _now() -> float: return time.time() +def _owner_token() -> str: + """Identity of the process holding a resume claim (``:``).""" + return f"{socket.gethostname()}:{os.getpid()}" + + +def _claim_owner_is_live(owner: str | None) -> bool: + """Whether the process that holds a claim is still running. + + Accepts both the two-part resume token (``:``) and the + three-part execution token (``::``). + + Startup recovery must fail only claims whose owner is *gone*: another CLI + process delivering a result right now would otherwise have its claim + yanked, and would then write ``DELIVERED`` over a record this process had + already marked failed. + + Unknown is treated as "live" wherever we genuinely cannot tell (a claim + from another host), because the failure mode of leaving a stale claim is a + result you can still read with ``/jobs ``, while the failure mode of + recovering a live one is a duplicated turn. A claim with no owner at all + predates this bookkeeping and cannot belong to a running process. + """ + if not owner: + return False + parts = owner.split(":") + if len(parts) < 2: + return False + host, pid_text = parts[0], parts[1] + if host and host != socket.gethostname(): + return True # another machine — not ours to judge + try: + pid = int(pid_text) + except ValueError: + return False + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except OSError: + return True # exists but not signalable (different user) — still alive + return True + + +class DurableRead(str, Enum): + """Outcome of reading a job's persisted metadata. + + ``MISSING`` and ``UNREADABLE`` are deliberately separate: a record another + manager cleaned away must be forgotten, while one that cannot be parsed + must be left untouched (acting on it, or rewriting it, would destroy state + we failed to read). + """ + + PRESENT = "present" + MISSING = "missing" + UNREADABLE = "unreadable" + + +class ClaimLockUnavailable(RuntimeError): + """The cross-process resume lock could not be established. + + Raised rather than swallowed: without the lock, two processes can both + decide a job is theirs to deliver. Callers that could *cause* a duplicate + delivery (claiming, startup recovery) must treat this as "not mine" and do + nothing; callers that only *prevent* one (recording a completed delivery) + proceed best-effort. + """ + + +@contextlib.contextmanager +def _file_lock(path: Path) -> Iterator[None]: + """Exclusive lock over ``path``, held across processes. + + Resume claims are the one piece of job state two CLI processes contend + for, so every read-modify-write of the resume fields happens inside this. + + Raises: + ClaimLockUnavailable: If no real lock could be taken — the lock file + could not be created, or the platform offers no lock primitive. + Silently continuing would leave the caller believing the section + was serialized when it was not. + """ + try: + fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600) + except OSError as exc: + raise ClaimLockUnavailable(f"cannot open {path}: {exc}") from exc + try: + try: + _acquire_lock(fd) + except OSError as exc: # any OS-level refusal is "no lock" + raise ClaimLockUnavailable(f"cannot lock {path}: {exc}") from exc + try: + yield + finally: + _release_lock(fd) + finally: + os.close(fd) + + +def _acquire_lock(fd: int) -> None: + """Take an exclusive lock on ``fd``, or raise ``ClaimLockUnavailable``.""" + try: + import fcntl + except ImportError: # pragma: no cover - non-POSIX + _acquire_lock_windows(fd) + return + try: + fcntl.flock(fd, fcntl.LOCK_EX) + except OSError as exc: + raise ClaimLockUnavailable(f"flock failed: {exc}") from exc + + +def _acquire_lock_windows(fd: int) -> None: # pragma: no cover - non-POSIX + try: + import msvcrt + + msvcrt.locking(fd, msvcrt.LK_LOCK, 1) + except Exception as exc: # noqa: BLE001 + raise ClaimLockUnavailable(f"no lock primitive: {exc}") from exc + + +def _release_lock(fd: int) -> None: + try: + import fcntl + + fcntl.flock(fd, fcntl.LOCK_UN) + except ImportError: # pragma: no cover - non-POSIX + try: + import msvcrt + + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + except Exception: # noqa: BLE001 + pass + except OSError as exc: # noqa: BLE001 - unlock failure is not actionable + logger.debug("job_claim_unlock_failed", error=str(exc)) + + +class ResumeStateError(RuntimeError): + """An invalid resume-lifecycle transition was requested. + + Signals a coordinator bug (completing a delivery nobody claimed), not a + user-facing condition — the valid "not mine to deliver" answer is + ``begin_resume()`` returning False. + """ + + +class ResumeState(str, Enum): + """Delivery lifecycle of a finished job's result back into the agent turn. + + ``PENDING`` → ``RESUMING`` → ``DELIVERED`` | ``FAILED``. + + The middle state exists so a crash can be told apart from a completed + delivery: marking a job delivered *before* running the turn loses the + result when the turn never ran, and marking it *after* re-delivers it if the + process dies mid-turn. A record found in ``RESUMING`` at startup is + therefore recovered as ``FAILED`` (its result is still readable via + ``/jobs ``) rather than replayed, because the interrupted turn may + already have executed tools. + """ + + PENDING = "pending" + RESUMING = "resuming" + DELIVERED = "delivered" + FAILED = "failed" + + @dataclass class JobRecord: """One job's metadata. Persisted as ``//meta.json``. @@ -65,7 +239,22 @@ class JobRecord: resume_on_complete: bool = False # wake the agent with the result when terminal call_id: str | None = None # ADK function_call_id / LangGraph tool_call_id call_name: str | None = None # function name to answer on resume - resumed: bool = False # guard against double-resume + resume_state: str = ResumeState.PENDING.value # see ResumeState + resume_error: str | None = None # why delivery failed, when it did + resume_owner: str | None = None # ":" holding a RESUMING claim + # ":" that launched the job and owns its execution. Set when the + # launch is claimed, so a second CLI sharing the jobs directory neither + # relaunches a queued job nor writes off a running one whose handle lives + # in the owner's memory. + exec_owner: str | None = None + + @property + def resumed(self) -> bool: + """True once delivery reached a terminal state (kept for compatibility).""" + return self.resume_state in ( + ResumeState.DELIVERED.value, + ResumeState.FAILED.value, + ) def elapsed_s(self) -> float: start = self.started_at or self.submitted_at @@ -82,11 +271,20 @@ def to_dict(self) -> dict: def from_dict(cls, d: dict) -> "JobRecord": d = dict(d) d["state"] = JobState(d["state"]) - return cls(**{k: d.get(k) for k in cls.__dataclass_fields__}) # type: ignore[attr-defined] + # Records written before the resume lifecycle carried a bool ``resumed``. + legacy_resumed = d.pop("resumed", None) + if "resume_state" not in d and legacy_resumed is not None: + d["resume_state"] = ( + ResumeState.DELIVERED.value if legacy_resumed else ResumeState.PENDING.value + ) + fields = cls.__dataclass_fields__ # type: ignore[attr-defined] + # Only pass keys the record actually declares, so a missing optional + # key falls back to its default instead of becoming None. + return cls(**{k: v for k, v in d.items() if k in fields}) def summary(self) -> dict: """Compact, JSON-safe view for tools / UI.""" - return { + out = { "job_id": self.job_id, "tool": self.tool, "name": self.name, @@ -96,6 +294,11 @@ def summary(self) -> dict: "exit_code": self.exit_code, "tags": self.tags, } + if self.resume_on_complete: + out["resume_state"] = self.resume_state + if self.resume_error: + out["resume_error"] = self.resume_error + return out def _json_safe_spec(spec: dict) -> dict: @@ -129,7 +332,20 @@ def __init__( self._max_concurrent = max(1, int(max_concurrent)) self._backends = backends or default_backends() self._lock = threading.RLock() + # Depth of the cross-process claim transaction this manager holds. + # POSIX ``flock`` is per file-description, so a second ``open`` in the + # same process would block against our own lock — the transaction is + # made re-entrant here instead (always entered under ``_lock``). + self._claim_depth = 0 + # Execution ownership is per *manager*, not per process: the backend + # handle for a running job lives in this instance, so a sibling manager + # in the same process is as unable to poll it as another CLI would be. + self._instance_id = uuid.uuid4().hex[:8] self._records: dict[str, JobRecord] = {} + # Job ids this manager has seen on disk. A record that was persisted + # and is now missing was *deleted* by another manager; one that was + # never persisted is simply new. Writing must tell them apart. + self._persisted: set[str] = set() self._load_existing() # ------------------------------------------------------------------ @@ -214,12 +430,21 @@ def _fill_turn_context( return session_id, user_id def get(self, job_id: str) -> JobRecord | None: + """The job, refreshed from durable state, or None if it is gone.""" with self._lock: rec = self._records.get(job_id) if rec is None: return None - self._refresh(rec) - self._maybe_start_queued() + try: + with self._claim_transaction(): + if self._refresh(rec) is DurableRead.MISSING: + self._forget(job_id) + return None + self._maybe_start_queued() + except ClaimLockUnavailable as exc: + logger.warning( + "job_refresh_skipped_unlocked", job_id=job_id, error=str(exc) + ) return rec def list( @@ -258,46 +483,127 @@ def result(self, job_id: str) -> Any: return self._backends[rec.backend].result(rec, self._job_dir(job_id)) def cancel(self, job_id: str) -> JobRecord | None: + """Cancel a job, if this manager can actually cancel it. + + Reloads the durable record first: a job that has already finished stays + finished (terminal transitions are monotonic — a stale view must not + rewrite a recorded success as CANCELLED). A job another live manager is + executing is only cancelled when the backend can reach it from here + (``cancels_foreign_jobs``); otherwise the record is returned unchanged + rather than reported cancelled while it keeps running. + + Returns: + The record — check ``state`` for what actually happened — or None + if the job is unknown. + """ with self._lock: rec = self._records.get(job_id) if rec is None: return None - if rec.state in TERMINAL_STATES: + try: + with self._claim_transaction(): + status = self._reload_durable(rec) + if status is DurableRead.MISSING: + self._forget(job_id) + return None + if status is DurableRead.UNREADABLE: + logger.warning("job_cancel_unreadable", job_id=job_id) + return rec + if rec.state in TERMINAL_STATES: + return rec + if self._execution_is_foreign(rec) and not self._can_cancel_foreign( + rec + ): + logger.warning( + "job_cancel_not_reachable", + job_id=job_id, + owner=rec.exec_owner, + ) + return rec + self._backends[rec.backend].cancel(rec, self._job_dir(job_id)) + rec.state = JobState.CANCELLED + rec.finished_at = _now() + self._persist(rec, owns_shared_state=True) + self._maybe_start_queued() + return rec + except ClaimLockUnavailable as exc: + logger.warning( + "job_cancel_skipped_unlocked", job_id=job_id, error=str(exc) + ) return rec - self._backends[rec.backend].cancel(rec, self._job_dir(job_id)) - rec.state = JobState.CANCELLED - rec.finished_at = _now() - self._persist(rec) - self._maybe_start_queued() - return rec def reconcile(self) -> None: - """Refresh non-terminal jobs from their backends, then promote queued.""" + """Refresh non-terminal jobs from durable state and their backends. + + Runs inside the cross-process transaction: reloading a record, polling + for it and writing the result is a read-modify-write of state other + managers share. Skipped entirely if the lock is unavailable — stale + reads are safe, unsynchronised writes are not. + """ with self._lock: - for rec in self._records.values(): - if rec.state not in TERMINAL_STATES: - self._refresh(rec) - self._maybe_start_queued() + try: + with self._claim_transaction(): + gone = [ + rec.job_id + for rec in list(self._records.values()) + if self._refresh(rec) is DurableRead.MISSING + ] + for job_id in gone: + self._forget(job_id) + self._maybe_start_queued() + except ClaimLockUnavailable as exc: + logger.warning("job_reconcile_skipped_unlocked", error=str(exc)) def clean(self) -> int: - """Remove terminal jobs (records + dirs). Returns the count removed.""" + """Remove terminal jobs (records + dirs). Returns the count removed. + + Each candidate is re-read under the cross-process lock first: deletion + is irreversible, and a stale terminal snapshot would take a job another + manager is still running with it. + """ import shutil with self._lock: self.reconcile() - removed = [r for r in self._records.values() if r.state in TERMINAL_STATES] - for rec in removed: - self._records.pop(rec.job_id, None) - shutil.rmtree(self._job_dir(rec.job_id), ignore_errors=True) - return len(removed) + try: + with self._claim_transaction(): + removed: list[JobRecord] = [] + for rec in list(self._records.values()): + status = self._reload_durable(rec) + if status is DurableRead.MISSING: + self._forget(rec.job_id) + continue + if status is DurableRead.UNREADABLE: + continue # never delete state we could not read + if rec.state not in TERMINAL_STATES: + continue + if rec.resume_state == ResumeState.RESUMING.value and ( + _claim_owner_is_live(rec.resume_owner) + ): + # A delivery is in flight; removing the record now + # would strip the result out from under it. + logger.debug( + "job_clean_skipped_delivering", job_id=rec.job_id + ) + continue + removed.append(rec) + self._records.pop(rec.job_id, None) + self._persisted.discard(rec.job_id) + shutil.rmtree(self._job_dir(rec.job_id), ignore_errors=True) + return len(removed) + except ClaimLockUnavailable as exc: + logger.warning("job_clean_skipped_unlocked", error=str(exc)) + return 0 def running_count(self) -> int: with self._lock: return sum(1 for r in self._records.values() if r.state == JobState.RUNNING) def awaiting_resume(self) -> list[JobRecord]: - """Terminal jobs flagged for resume that haven't been resumed yet. + """Terminal jobs flagged for resume whose result is still undelivered. + Only ``PENDING`` records are returned: one already ``RESUMING`` is being + delivered right now, and ``DELIVERED``/``FAILED`` are terminal. Reconciles first so freshly-finished jobs are included; sorted oldest-finished-first so the coordinator drains in completion order. """ @@ -306,17 +612,222 @@ def awaiting_resume(self) -> list[JobRecord]: recs = [ r for r in self._records.values() - if r.resume_on_complete and not r.resumed and r.state in TERMINAL_STATES + if r.resume_on_complete + and r.resume_state == ResumeState.PENDING.value + and r.state in TERMINAL_STATES ] return sorted(recs, key=lambda r: r.finished_at or r.submitted_at) - def mark_resumed(self, job_id: str) -> None: - """Mark a job as resumed (durably) so it is never resumed twice.""" + def begin_resume(self, job_id: str) -> bool: + """Claim a job for delivery: ``PENDING`` → ``RESUMING``. + + The claim is atomic and durable **across processes**: two CLIs sharing + one jobs directory (the default — it is user-scoped, not project-scoped) + would otherwise both see ``PENDING`` in their own memory and deliver the + same result into two conversations. The transition therefore happens + under an inter-process lock, re-reading the record from disk so the + decision is made on the shared state rather than this manager's cache. + + Only a job that is genuinely deliverable can be claimed: it must exist, + be terminal, be flagged ``resume_on_complete``, and still be + ``PENDING``. + + If the lock cannot be taken the claim **fails closed** (returns False): + an undelivered result is still readable with ``/jobs ``, whereas an + unsynchronised claim can deliver the same result into two conversations. + + Returns: + True if this caller now owns delivery; False otherwise. A False + result means "not yours to deliver" — never an error. + """ with self._lock: rec = self._records.get(job_id) - if rec is not None and not rec.resumed: - rec.resumed = True - self._persist(rec) + if rec is None: + return False + if not rec.resume_on_complete: + return False + try: + with self._claim_transaction(): + # Terminality is judged on the durable record, not on this + # manager's snapshot: a job that finished elsewhere is + # deliverable even if we still have it as RUNNING. + status = self._reload_durable(rec) + if status is DurableRead.MISSING: + self._forget(job_id) + return False + if status is DurableRead.UNREADABLE: + logger.warning("job_resume_claim_unreadable", job_id=job_id) + return False + if rec.state not in TERMINAL_STATES: + return False + if rec.resume_state != ResumeState.PENDING.value: + return False + rec.resume_state = ResumeState.RESUMING.value + rec.resume_owner = _owner_token() + rec.resume_error = None + self._persist(rec, owns_shared_state=True) + except ClaimLockUnavailable as exc: + logger.warning( + "job_resume_claim_unlocked", job_id=job_id, error=str(exc) + ) + return False + return True + + def complete_resume( + self, job_id: str, *, delivered: bool, error: str | None = None + ) -> None: + """Close out a claimed delivery: ``RESUMING`` → ``DELIVERED``/``FAILED``. + + Only the record *this process claimed* transitions. Anything else is a + coordinator bug — closing out a job that was never claimed, or one + another process is delivering, would rewrite state someone else owns — + so it raises rather than silently overwriting. + + Args: + job_id: The claimed job. + delivered: True only when the resume turn actually ran to completion. + error: Why delivery failed, recorded for ``/jobs``. + + If the lock cannot be taken the durable record is **left as it is** — + still ``RESUMING``, owned by this process. That is the failed-safe + outcome: an unlocked rewrite could regress a delivery another process + had just recorded, whereas a record stuck in ``RESUMING`` is recovered + as failed once this process is gone, and is never re-delivered. The + in-memory record still transitions, so this manager's own ``/jobs`` + view is accurate. + + Raises: + ResumeStateError: If the job is unknown, was not claimed, or is + claimed by someone else. + """ + with self._lock: + rec = self._records.get(job_id) + if rec is None: + raise ResumeStateError(f"Unknown job {job_id!r}: nothing to complete.") + try: + with self._claim_transaction(): + self._complete_resume_locked(rec, delivered, error) + except ClaimLockUnavailable as exc: + logger.warning( + "job_resume_complete_unpersisted", + job_id=job_id, + delivered=delivered, + error=str(exc), + ) + rec.resume_state = ( + ResumeState.DELIVERED.value + if delivered + else ResumeState.FAILED.value + ) + rec.resume_error = None if delivered else error + rec.resume_owner = None + + def _complete_resume_locked( + self, rec: JobRecord, delivered: bool, error: str | None + ) -> None: + """Body of :meth:`complete_resume`; the caller holds the transaction.""" + status = self._reload_durable(rec) + if status is not DurableRead.PRESENT: + # The record was cleaned away (or cannot be read) while the turn + # ran. Recording the outcome would recreate a deleted job; there is + # nothing left to deliver to. + logger.warning( + "job_resume_complete_record_gone", + job_id=rec.job_id, + status=status.value, + ) + if status is DurableRead.MISSING: + self._forget(rec.job_id) + return + if rec.resume_state != ResumeState.RESUMING.value: + raise ResumeStateError( + f"Job {rec.job_id!r} is {rec.resume_state!r}, not " + f"{ResumeState.RESUMING.value!r}: complete_resume() must " + "follow a successful begin_resume()." + ) + owner = _owner_token() + if rec.resume_owner not in (None, owner): + raise ResumeStateError( + f"Job {rec.job_id!r} is being delivered by {rec.resume_owner!r}, " + f"not {owner!r}: only the claiming process may complete it." + ) + rec.resume_state = ( + ResumeState.DELIVERED.value if delivered else ResumeState.FAILED.value + ) + rec.resume_error = None if delivered else error + rec.resume_owner = None + self._persist(rec, owns_shared_state=True) + + def mark_resumed(self, job_id: str) -> None: + """Mark a job's result delivered (durably) so it is never resumed twice. + + Deprecated compatibility shim for the pre-lifecycle API: it claims and + completes in one step. Prefer ``begin_resume()``/``complete_resume()``, + which survives a crash between claiming and delivering. A job that + cannot be claimed is left alone. + """ + if self.begin_resume(job_id): + self.complete_resume(job_id, delivered=True) + + def _recover_interrupted_resumes(self) -> None: + """Fail records left mid-delivery by a *dead* owner (called on load). + + The interrupted turn may already have executed tools, so it is never + replayed automatically; the result stays readable via ``/jobs ``. + + A claim whose owner is still running belongs to another live CLI + delivering it right now (the jobs directory is shared per user) and is + left alone — taking it over would duplicate the turn and race that + process's ``complete_resume``. + + The decision is made on state re-read **inside** the lock: the records + were loaded from disk before it was taken, and in that window another + process may have claimed a job or finished delivering one. Acting on + the pre-lock snapshot rewrote a completed delivery as failed. A record + that window deleted, or left unreadable, is skipped outright — see + :meth:`_recovery_target`. + + If the lock is unavailable nothing is recovered — leaving a stale + ``RESUMING`` record costs a result you can still read with + ``/jobs ``, while a wrong recovery corrupts another process's + state. + """ + try: + with self._claim_transaction(): + for rec in list(self._records.values()): + if not self._recovery_target(rec): + continue + if rec.resume_state != ResumeState.RESUMING.value: + continue + if _claim_owner_is_live(rec.resume_owner): + logger.debug( + "job_resume_claim_active", + job_id=rec.job_id, + owner=rec.resume_owner, + ) + continue + rec.resume_state = ResumeState.FAILED.value + rec.resume_error = "resume interrupted before the turn completed" + rec.resume_owner = None + self._persist(rec, owns_shared_state=True) + logger.warning("job_resume_interrupted", job_id=rec.job_id) + except ClaimLockUnavailable as exc: + logger.warning("job_resume_recovery_skipped", error=str(exc)) + + def close(self) -> None: + """Release resources owned by the manager's backends. + + Idempotent. Jobs themselves are not cancelled: a detached subprocess is + meant to outlive the CLI, and its state is recovered from disk on the + next start. Only in-process resources (thread pools) are released. + """ + for backend in self._backends.values(): + try: + backend.close() + except Exception as exc: # noqa: BLE001 - shutdown must not fail + logger.warning( + "job_backend_close_failed", backend=backend.name, error=str(exc) + ) # ------------------------------------------------------------------ # Internals @@ -325,24 +836,287 @@ def mark_resumed(self, job_id: str) -> None: def _job_dir(self, job_id: str) -> Path: return self.base_dir / job_id - def _persist(self, rec: JobRecord) -> None: + @property + def _claim_lock_path(self) -> Path: + """One lock for the whole jobs directory; claims are rare and brief.""" + return self.base_dir / ".resume.lock" + + @contextlib.contextmanager + def _claim_transaction(self) -> Iterator[None]: + """Serialize a read-modify-write of the resume fields, cross-process. + + Re-entrant within this manager (see ``_claim_depth``) so a claim path + can persist without deadlocking against its own ``flock``. + + Raises: + ClaimLockUnavailable: propagated from :func:`_file_lock`. + """ + with self._lock: + if self._claim_depth: + self._claim_depth += 1 + try: + yield + finally: + self._claim_depth -= 1 + return + with _file_lock(self._claim_lock_path): + self._claim_depth = 1 + try: + yield + finally: + self._claim_depth = 0 + + def _persist(self, rec: JobRecord, *, owns_shared_state: bool = False) -> None: + """Write a record's metadata, never clobbering a newer delivery state. + + ``meta.json`` holds both this process's job bookkeeping *and* the + cross-process resume fields. A plain state write (a poll that finds the + job finished, a cancel, a launch) carries whatever resume fields this + manager last read, which may be older than what another process has + since written — persisting them verbatim erased that process's claim, + and the job could then be claimed and delivered twice. + + So unless the caller owns the resume fields (the claim transitions, + which set them under the lock), they are re-read from disk first and + adopted into ``rec``. + + Args: + rec: The record to write. + owns_shared_state: True only for a caller holding the claim + transaction that just set these fields. + """ + if owns_shared_state: + self._write(rec) + return + try: + with self._claim_transaction(): + status, data = self._read_persisted(rec.job_id) + if status is DurableRead.UNREADABLE: + logger.warning("job_persist_skipped_unreadable", job_id=rec.job_id) + return + if status is DurableRead.MISSING and rec.job_id in self._persisted: + # Another manager cleaned this job away; writing it back + # would resurrect a deleted record. + self._forget(rec.job_id) + return + if data is not None: + durable = data.get("state") + if durable in _TERMINAL_VALUES and durable != rec.state.value: + # Terminal transitions are monotonic: some manager has + # already recorded how this job ended, and a snapshot + # taken before that cannot say otherwise. Adopt it. + self._reload_durable(rec) + logger.info( + "job_terminal_state_preserved", + job_id=rec.job_id, + state=durable, + ) + return + self._adopt_shared_fields(rec, data) + self._write(rec) + except ClaimLockUnavailable as exc: + # No lock means no safe read-modify-write: re-reading and rewriting + # is exactly the race the lock prevents, so an interleaved delivery + # would be regressed. Skip the write entirely — the in-memory + # change is retried by the next reconcile — unless there is nothing + # on disk yet to clobber, in which case creating the record is safe. + meta = self._job_dir(rec.job_id) / "meta.json" + if meta.exists() or rec.job_id in self._persisted: + logger.warning( + "job_persist_skipped_unlocked", job_id=rec.job_id, error=str(exc) + ) + return + logger.warning( + "job_persist_created_unlocked", job_id=rec.job_id, error=str(exc) + ) + self._write(rec) + + def _read_persisted(self, job_id: str) -> "tuple[DurableRead, dict | None]": + """Read a record from disk, distinguishing *gone* from *unreadable*. + + The two demand opposite responses: a record another manager cleaned + away must be forgotten (recreating it would resurrect a deleted job), + while one we simply cannot parse must be left exactly as it is. + """ + import json + + meta = self._job_dir(job_id) / "meta.json" + try: + text = meta.read_text() + except FileNotFoundError: + return DurableRead.MISSING, None + except OSError as exc: + logger.warning("job_meta_unreadable", job_id=job_id, error=str(exc)) + return DurableRead.UNREADABLE, None + try: + data = json.loads(text) + except ValueError as exc: + logger.warning("job_meta_corrupt", job_id=job_id, error=str(exc)) + return DurableRead.UNREADABLE, None + if not isinstance(data, dict): + logger.warning("job_meta_corrupt", job_id=job_id, error="not an object") + return DurableRead.UNREADABLE, None + return DurableRead.PRESENT, data + + def _forget(self, job_id: str) -> None: + """Drop a record another manager deleted. Never recreate it.""" + self._persisted.discard(job_id) + if self._records.pop(job_id, None) is not None: + logger.info("job_record_deleted_elsewhere", job_id=job_id) + + # The lifecycle a record's *executor* owns and persists. Every manager + # reads these back before acting: an in-memory copy is only ever a snapshot + # of what some manager last wrote. + _DURABLE_LIFECYCLE_FIELDS = ( + "exit_code", + "error", + "started_at", + "finished_at", + "pid", + "backend_handle", + ) + + def _adopt_shared_fields(self, rec: JobRecord, data: dict) -> None: + """Adopt the ownership/delivery fields from a persisted snapshot.""" + if "exec_owner" in data: + rec.exec_owner = data["exec_owner"] + state = data.get("resume_state") + if state is None: + # A record written before the lifecycle existed. + legacy = data.get("resumed") + if legacy is None: + return + state = ( + ResumeState.DELIVERED.value if legacy else ResumeState.PENDING.value + ) + rec.resume_state = state + rec.resume_error = data.get("resume_error") + rec.resume_owner = data.get("resume_owner") + + def _write(self, rec: JobRecord) -> None: + """Write the record and remember that it now exists on disk.""" atomic_write_json(self._job_dir(rec.job_id) / "meta.json", rec.to_dict()) + self._persisted.add(rec.job_id) + + def _reload_durable(self, rec: JobRecord) -> DurableRead: + """Reload the whole persisted lifecycle into ``rec``. + + A manager's in-memory record is a snapshot; the file is the shared + truth. Every mutator reloads before deciding, so a stale view cannot + overwrite an outcome another manager already recorded, an observer + learns that a foreign job finished, and a *terminal* record still picks + up the delivery lifecycle that moves after it. + + Returns: + What the read found. ``MISSING`` means another manager deleted the + record and the caller must forget it rather than write it back; + ``UNREADABLE`` means the caller must not act on it at all. + """ + status, data = self._read_persisted(rec.job_id) + if status is not DurableRead.PRESENT: + return status + raw_state = data.get("state") + if raw_state is not None: + try: + rec.state = JobState(raw_state) + except ValueError: # pragma: no cover - unknown state on disk + pass + for name in self._DURABLE_LIFECYCLE_FIELDS: + if name in data: + setattr(rec, name, data[name]) + self._adopt_shared_fields(rec, data) + return status + + def _recovery_target(self, rec: JobRecord) -> bool: + """Whether startup recovery may judge — and rewrite — this record. + + The durable read is the *decision*, not a side effect. Recovery writes + with ``owns_shared_state=True``, which by design does not re-check the + file, so both non-PRESENT answers have to be handled here: + + - ``MISSING`` — another manager cleaned the job away between our load + and now. Forget it; writing a verdict would resurrect a deleted job. + - ``UNREADABLE`` — we have no basis for a verdict at all, and rewriting + would destroy exactly the state we failed to read. Leave it alone. + + Returns: + True only if the record was reloaded and may be acted on. + """ + status = self._reload_durable(rec) + if status is DurableRead.PRESENT: + return True + if status is DurableRead.MISSING: + self._forget(rec.job_id) + else: + logger.warning("job_recovery_skipped_unreadable", job_id=rec.job_id) + return False + + def _can_cancel_foreign(self, rec: JobRecord) -> bool: + """Whether this manager could actually cancel a job it did not start. + + A declared backend capability (``cancels_foreign_jobs``), not an + inference: publishing a readable outcome and being remotely + controllable are different properties, and conflating them would mark a + record CANCELLED while the job kept running in the owning process. + """ + backend = self._backends.get(rec.backend) + return bool(backend is not None and backend.cancels_foreign_jobs) + + def _exec_owner_token(self) -> str: + """This manager's execution identity (``::``).""" + return f"{_owner_token()}:{self._instance_id}" + + def _execution_is_foreign(self, rec: JobRecord) -> bool: + """Whether a live manager *other than this one* owns the execution.""" + owner = rec.exec_owner + return ( + bool(owner) + and owner != self._exec_owner_token() + and _claim_owner_is_live(owner) + ) + + def _refresh(self, rec: JobRecord) -> DurableRead: + """Bring a job up to date: durable state first, then the backend. + + The durable record comes first for *every* record, terminal ones + included — a finished job's delivery lifecycle keeps moving, and that + is state another manager owns. + + A foreign job is still polled: backends publish their outcome durably + (the ``exit_code`` sentinel), so the observer can read it. What it must + not do is believe ``UNKNOWN`` — that answer means "I hold no handle for + this", which is true of every job another manager started, and taking + it at face value would mark a healthy job terminal and hand its + "result" to a resume. - def _refresh(self, rec: JobRecord) -> None: - """Poll the backend for a non-terminal job and persist any change.""" + Returns: + The durable read status, so callers can forget a deleted record. + """ + status = self._reload_durable(rec) + if status is not DurableRead.PRESENT: + return status if rec.state in TERMINAL_STATES or rec.state == JobState.QUEUED: - return + return status backend = self._backends.get(rec.backend) if backend is None: rec.state = JobState.UNKNOWN self._persist(rec) - return + return status new_state = backend.poll(rec, self._job_dir(rec.job_id)) - if new_state != rec.state: - rec.state = new_state - if new_state in TERMINAL_STATES and rec.finished_at is None: - rec.finished_at = _now() - self._persist(rec) + if new_state == rec.state: + return status + if new_state is JobState.UNKNOWN and self._execution_is_foreign(rec): + logger.debug( + "job_poll_unknown_foreign", + job_id=rec.job_id, + owner=rec.exec_owner, + ) + return status + rec.state = new_state + if new_state in TERMINAL_STATES and rec.finished_at is None: + rec.finished_at = _now() + self._persist(rec) + return status def _maybe_start_queued(self) -> None: """Start queued jobs up to the concurrency cap (caller holds the lock).""" @@ -358,6 +1132,54 @@ def _maybe_start_queued(self) -> None: self._start(rec) def _start(self, rec: JobRecord) -> None: + """Launch a queued job — at most once across every process. + + The jobs directory is user-scoped, so two CLIs can hold the same queued + record. Execution is therefore claimed the same way delivery is: under + the cross-process lock, against the record *on disk*, recording the + owner durably **before** the backend is touched. If the lock is + unavailable the launch is skipped — a late job beats two of them. + """ + try: + with self._claim_transaction(): + if not self._claim_execution(rec): + return + self._launch(rec) + except ClaimLockUnavailable as exc: + logger.warning( + "job_launch_skipped_unlocked", job_id=rec.job_id, error=str(exc) + ) + + def _claim_execution(self, rec: JobRecord) -> bool: + """Take ownership of a queued job's launch. Caller holds the transaction.""" + status, data = self._read_persisted(rec.job_id) + if status is DurableRead.UNREADABLE: + logger.warning("job_launch_skipped_unreadable", job_id=rec.job_id) + return False + if status is DurableRead.MISSING and rec.job_id in self._persisted: + self._forget(rec.job_id) + return False + if data is not None: + try: + rec.state = JobState(data.get("state", rec.state)) + except ValueError: # pragma: no cover - unknown state on disk + pass + rec.exec_owner = data.get("exec_owner") + if rec.state != JobState.QUEUED: + return False + if rec.exec_owner is not None: + # Already claimed: either it is running elsewhere, or its launcher + # died and startup recovery will fail it. Never launch it twice. + logger.debug( + "job_launch_owned_elsewhere", job_id=rec.job_id, owner=rec.exec_owner + ) + return False + rec.exec_owner = self._exec_owner_token() + self._persist(rec, owns_shared_state=True) + return True + + def _launch(self, rec: JobRecord) -> None: + """Start the backend for a claimed job. Caller holds the transaction.""" backend = self._backends[rec.backend] try: backend.start(rec, self._job_dir(rec.job_id)) @@ -368,7 +1190,35 @@ def _start(self, rec: JobRecord) -> None: rec.error = f"launch failed: {exc}" rec.finished_at = _now() logger.warning("job_launch_failed", job_id=rec.job_id, error=str(exc)) - self._persist(rec) + self._persist(rec, owns_shared_state=True) + + def _recover_interrupted_launches(self) -> None: + """Fail queued jobs whose launcher died mid-claim (called on load). + + The claim is written before the backend is touched, so a record left + QUEUED with a dead owner may or may not have started something. It is + never relaunched — a duplicate side effect is worse than a job that + must be resubmitted — and is failed with that stated plainly. + + Only records the durable read still vouches for are judged; see + :meth:`_recovery_target`. + """ + try: + with self._claim_transaction(): + for rec in list(self._records.values()): + if not self._recovery_target(rec): + continue + if rec.state != JobState.QUEUED or rec.exec_owner is None: + continue + if _claim_owner_is_live(rec.exec_owner): + continue + rec.state = JobState.FAILED + rec.error = "launch interrupted before the job started" + rec.finished_at = _now() + self._persist(rec, owns_shared_state=True) + logger.warning("job_launch_interrupted", job_id=rec.job_id) + except ClaimLockUnavailable as exc: + logger.warning("job_launch_recovery_skipped", error=str(exc)) def _load_existing(self) -> None: """Load persisted job records on startup and reconcile their state.""" @@ -383,6 +1233,12 @@ def _load_existing(self) -> None: except (ValueError, OSError, TypeError, KeyError): continue # In-memory handles (Popen / Future) are gone after a restart, - # so non-restart-safe running jobs become UNKNOWN. + # so non-restart-safe running jobs become UNKNOWN — unless a + # live process still owns their execution (see _refresh). self._records[rec.job_id] = rec + self._persisted.add(rec.job_id) + # Recover before reconciling: reconcile() starts queued jobs, and + # an interrupted launch must never be one of them. + self._recover_interrupted_launches() + self._recover_interrupted_resumes() self.reconcile() diff --git a/tests/cli/test_job_monitor.py b/tests/cli/test_job_monitor.py index 988df54..2c61f61 100644 --- a/tests/cli/test_job_monitor.py +++ b/tests/cli/test_job_monitor.py @@ -157,7 +157,8 @@ def test_no_resume_cue_for_already_resumed(self, jm: JobManager): ) rec = JobRecord( job_id="j1", tool="run_shell_job", backend="subprocess", name="build", - state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", resumed=True, + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + resume_state="delivered", ) assert mon._build_segment([rec]) is None diff --git a/tests/cli/test_resume_coordinator.py b/tests/cli/test_resume_coordinator.py index 1e1efe1..b366c0a 100644 --- a/tests/cli/test_resume_coordinator.py +++ b/tests/cli/test_resume_coordinator.py @@ -3,6 +3,10 @@ Each finished, resume-flagged job becomes one serialized resume turn. Tested on a bare app (no real ThinkingPromptSession) with fake controller / job manager / message processor. + +Delivery uses the job's resume lifecycle: claim (pending → resuming) *before* +the turn, record delivered/failed *after* it. Marking delivery up-front — the +previous behaviour — silently dropped every resume whose turn then failed. """ from __future__ import annotations @@ -10,37 +14,53 @@ import asyncio from types import SimpleNamespace +import pytest + from agentic_cli.cli.app import BaseCLIApp +from agentic_cli.cli.message_processor import TurnResult, TurnStatus class _FakeJM: + """Minimal stand-in implementing the resume lifecycle contract.""" + def __init__(self, records: list) -> None: self._records = records - self.marked: list[str] = [] + self.claimed: list[str] = [] + self.completed: list[tuple[str, bool, str | None]] = [] def awaiting_resume(self) -> list: - return [r for r in self._records if r.job_id not in self.marked] + return [r for r in self._records if r.job_id not in self.claimed] - def mark_resumed(self, job_id: str) -> None: - self.marked.append(job_id) + def begin_resume(self, job_id: str) -> bool: + if job_id in self.claimed: + return False + self.claimed.append(job_id) + return True + + def complete_resume(self, job_id: str, *, delivered: bool, error=None) -> None: + self.completed.append((job_id, delivered, error)) class _FakeMessageProcessor: - def __init__(self) -> None: + def __init__(self, result: TurnResult | None = None) -> None: self.resumed: list[str] = [] + self._result = result or TurnResult(TurnStatus.COMPLETED) - async def process_resume(self, *, record, workflow_controller, ui, settings, usage_tracker): + async def process_resume( + self, *, record, workflow_controller, ui, settings, usage_tracker + ): self.resumed.append(record.job_id) + return self._result -def _app(records: list, *, ready: bool = True, has_jm: bool = True): +def _app(records: list, *, ready: bool = True, has_jm: bool = True, result=None): app = BaseCLIApp.__new__(BaseCLIApp) jm = _FakeJM(records) if has_jm else None app._workflow_controller = SimpleNamespace( is_ready=ready, workflow=SimpleNamespace(job_manager=jm) ) app._turn_lock = asyncio.Lock() - app._message_processor = _FakeMessageProcessor() + app._message_processor = _FakeMessageProcessor(result) app.session = object() app._settings = SimpleNamespace(job_auto_resume=True) app._usage_tracker = None @@ -53,23 +73,62 @@ async def test_resumes_each_awaiting_job_once(): n = await app.resume_finished_jobs() assert n == 2 assert app._message_processor.resumed == ["a", "b"] - assert jm.marked == ["a", "b"] + assert jm.completed == [("a", True, None), ("b", True, None)] -async def test_marks_resumed_before_processing(): +async def test_claims_before_processing_and_records_after(): order: list = [] app, jm = _app([SimpleNamespace(job_id="a")]) - real_mark = jm.mark_resumed - jm.mark_resumed = lambda jid: (order.append(("mark", jid)), real_mark(jid))[1] + real_begin = jm.begin_resume + jm.begin_resume = lambda jid: (order.append(("claim", jid)), real_begin(jid))[1] + real_complete = jm.complete_resume + jm.complete_resume = lambda jid, **kw: ( + order.append(("done", jid, kw["delivered"])), + real_complete(jid, **kw), + )[1] async def _proc(*, record, **kw): order.append(("proc", record.job_id)) + return TurnResult(TurnStatus.COMPLETED) app._message_processor.process_resume = _proc await app.resume_finished_jobs() - assert order == [("mark", "a"), ("proc", "a")] + assert order == [("claim", "a"), ("proc", "a"), ("done", "a", True)] + + +async def test_failed_resume_is_recorded_as_failed_not_delivered(): + """A resume whose turn failed must not be recorded as delivered.""" + app, jm = _app( + [SimpleNamespace(job_id="a")], + result=TurnResult(TurnStatus.FAILED, error="workflow blew up"), + ) + + n = await app.resume_finished_jobs() + + assert n == 1 # picked up... + assert jm.completed == [("a", False, "workflow blew up")] # ...but not delivered + + +async def test_unavailable_conversation_is_recorded_as_failed(): + app, jm = _app( + [SimpleNamespace(job_id="a")], + result=TurnResult(TurnStatus.UNAVAILABLE, error="conversation gone"), + ) + + assert await app.resume_finished_jobs() == 1 + assert jm.completed == [("a", False, "conversation gone")] + + +async def test_unclaimable_job_is_skipped(): + """A job another coordinator already claimed must not be delivered twice.""" + app, jm = _app([SimpleNamespace(job_id="a")]) + jm.begin_resume = lambda jid: False + + assert await app.resume_finished_jobs() == 0 + assert app._message_processor.resumed == [] + assert jm.completed == [] async def test_no_manager_returns_zero(): @@ -81,3 +140,175 @@ async def test_not_ready_returns_zero(): app, _ = _app([SimpleNamespace(job_id="a")], ready=False) assert await app.resume_finished_jobs() == 0 assert app._message_processor.resumed == [] + + +class TestShutdownOrder: + """Fact extraction must run before the controller closes the manager. + + ``background_init``'s exit now calls ``controller.close()``, which cleans up + and drops the manager. Extraction placed after that block would find + ``is_ready`` False and silently do nothing. + """ + + def test_extraction_runs_inside_the_controller_context(self): + import ast + import inspect + import textwrap + + from agentic_cli.cli.app import BaseCLIApp + + tree = ast.parse(textwrap.dedent(inspect.getsource(BaseCLIApp.run))) + + def _is_background_init(node: ast.AsyncWith) -> bool: + return any( + isinstance(item.context_expr, ast.Call) + and isinstance(item.context_expr.func, ast.Attribute) + and item.context_expr.func.attr == "background_init" + for item in node.items + ) + + blocks = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.AsyncWith) and _is_background_init(n) + ] + assert len(blocks) == 1, "run() no longer has a single background_init block" + + def _calls_extraction(node: ast.AST) -> bool: + return any( + isinstance(n, ast.Attribute) + and n.attr == "_extract_session_facts_on_exit" + for n in ast.walk(node) + ) + + assert _calls_extraction(blocks[0]), ( + "_extract_session_facts_on_exit() must be called inside the " + "background_init block — the controller closes the workflow on exit" + ) + outside = [n for n in tree.body[0].body if n is not blocks[0]] + assert not any(_calls_extraction(n) for n in outside), ( + "_extract_session_facts_on_exit() is also called after cleanup" + ) + + async def test_extraction_is_skipped_once_the_controller_closed(self): + from agentic_cli.cli.app import BaseCLIApp + + calls: list[str] = [] + app = BaseCLIApp.__new__(BaseCLIApp) + app._settings = SimpleNamespace(auto_extract_session_facts=True) + app._workflow_controller = SimpleNamespace( + is_ready=False, + workflow=SimpleNamespace( + on_session_end=lambda: calls.append("extract") + ), + ) + + await app._extract_session_facts_on_exit() + assert calls == [] + + +class _StrictJM(_FakeJM): + """Enforces the real lifecycle: complete only after a claim, once.""" + + def __init__(self, records: list) -> None: + super().__init__(records) + self.open_claims: set[str] = set() + + def begin_resume(self, job_id: str) -> bool: + if not super().begin_resume(job_id): + return False + self.open_claims.add(job_id) + return True + + def complete_resume(self, job_id: str, *, delivered: bool, error=None) -> None: + from agentic_cli.tools.jobs.manager import ResumeStateError + + if job_id not in self.open_claims: + raise ResumeStateError(f"{job_id} was not claimed") + self.open_claims.discard(job_id) + super().complete_resume(job_id, delivered=delivered, error=error) + + +def _strict_app(records: list, processor=None): + app = BaseCLIApp.__new__(BaseCLIApp) + jm = _StrictJM(records) + app._workflow_controller = SimpleNamespace( + is_ready=True, workflow=SimpleNamespace(job_manager=jm) + ) + app._turn_lock = asyncio.Lock() + app._message_processor = processor or _FakeMessageProcessor() + app.session = SimpleNamespace(add_error=lambda msg: None) + app._settings = SimpleNamespace(job_auto_resume=True) + app._usage_tracker = None + return app, jm + + +class TestClaimIsAlwaysClosed: + """No record may be left RESUMING once the coordinator returns.""" + + async def test_processor_exception_closes_the_claim(self): + class _Raising: + async def process_resume(self, **kwargs): + raise RuntimeError("processor exploded") + + app, jm = _strict_app([SimpleNamespace(job_id="a", name="build")], _Raising()) + + assert await app.resume_finished_jobs() == 1 + assert jm.open_claims == set(), "the job is stuck RESUMING" + assert jm.completed == [("a", False, "processor exploded")] + + async def test_can_resume_failure_closes_the_claim(self): + """``can_resume()`` raising inside process_resume is still a closed claim.""" + + class _Raising: + async def process_resume(self, **kwargs): + raise ConnectionError("session store unreachable") + + app, jm = _strict_app([SimpleNamespace(job_id="a", name="build")], _Raising()) + + await app.resume_finished_jobs() + assert jm.open_claims == set() + assert jm.completed[0][1] is False + + async def test_cancellation_closes_the_claim_and_propagates(self): + started = asyncio.Event() + + class _Hanging: + async def process_resume(self, **kwargs): + started.set() + await asyncio.Event().wait() + + app, jm = _strict_app([SimpleNamespace(job_id="a", name="build")], _Hanging()) + + task = asyncio.create_task(app.resume_finished_jobs()) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert jm.open_claims == set(), "cancellation left the job RESUMING" + assert jm.completed == [("a", False, "resume cancelled")] + + async def test_normal_delivery_closes_the_claim(self): + app, jm = _strict_app([SimpleNamespace(job_id="a", name="build")]) + assert await app.resume_finished_jobs() == 1 + assert jm.open_claims == set() + assert jm.completed == [("a", True, None)] + + async def test_duplicate_coordinators_deliver_once(self): + record = SimpleNamespace(job_id="a", name="build") + app, jm = _strict_app([record]) + second_app = BaseCLIApp.__new__(BaseCLIApp) + second_app._workflow_controller = app._workflow_controller + second_app._turn_lock = asyncio.Lock() + second_app._message_processor = _FakeMessageProcessor() + second_app.session = app.session + second_app._settings = app._settings + second_app._usage_tracker = None + + counts = await asyncio.gather( + app.resume_finished_jobs(), second_app.resume_finished_jobs() + ) + + assert sorted(counts) == [0, 1], "both coordinators claimed the same job" + assert len(jm.completed) == 1 diff --git a/tests/tools/test_jobs.py b/tests/tools/test_jobs.py index 577276d..b859676 100644 --- a/tests/tools/test_jobs.py +++ b/tests/tools/test_jobs.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os import subprocess import time from pathlib import Path @@ -10,7 +11,7 @@ import pytest from agentic_cli.tools.jobs import JobManager, JobRecord, JobState -from agentic_cli.tools.jobs.backends import default_backends +from agentic_cli.tools.jobs.backends import JobBackend, default_backends def _wait(jm: JobManager, job_id: str, timeout: float = 5.0) -> JobRecord: @@ -26,6 +27,21 @@ def _wait(jm: JobManager, job_id: str, timeout: float = 5.0) -> JobRecord: return jm.get(job_id) # type: ignore[return-value] +def _this_host() -> str: + import socket + + return socket.gethostname() + + +def _dead_pid() -> int: + """A pid that is certainly not running: a child we started and reaped.""" + import sys + + proc = subprocess.Popen([sys.executable, "-c", "pass"]) + proc.wait() + return proc.pid + + @pytest.fixture def jm(tmp_path: Path) -> JobManager: return JobManager(base_dir=tmp_path / "jobs", max_concurrent=2) @@ -315,3 +331,1500 @@ def test_awaiting_resume_and_mark_resumed(self, tmp_path: Path): jm_reloaded = JobManager(base_dir=base) assert jm_reloaded.get(rec.job_id).resumed is True # type: ignore[union-attr] assert jm_reloaded.awaiting_resume() == [] + + +class TestResumeLifecycle: + """pending → resuming → delivered/failed, with crash recovery.""" + + def _jm(self, tmp_path): + from agentic_cli.tools.jobs import JobManager + + return JobManager(base_dir=tmp_path / "jobs") + + def _finished_job(self, jm, job_id="j1"): + from agentic_cli.tools.jobs import JobRecord + from agentic_cli.tools.jobs.backends import JobState + + rec = JobRecord( + job_id=job_id, tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + session_id="s", user_id="u", finished_at=1.0, + ) + jm._records[rec.job_id] = rec + jm._persist(rec) + return rec + + def test_new_job_is_pending(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + rec = self._finished_job(jm) + assert rec.resume_state == ResumeState.PENDING.value + assert [r.job_id for r in jm.awaiting_resume()] == ["j1"] + + def test_claim_is_exclusive_and_hides_from_awaiting(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + rec = self._finished_job(jm) + + assert jm.begin_resume("j1") is True + assert jm.begin_resume("j1") is False # already claimed + assert rec.resume_state == ResumeState.RESUMING.value + assert jm.awaiting_resume() == [] + + def test_claim_of_unknown_job_is_false(self, tmp_path): + assert self._jm(tmp_path).begin_resume("nope") is False + + def test_complete_records_delivered_and_failed(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + rec = self._finished_job(jm) + jm.begin_resume("j1") + jm.complete_resume("j1", delivered=True) + assert rec.resume_state == ResumeState.DELIVERED.value + assert rec.resumed is True + + rec2 = self._finished_job(jm, "j2") + jm.begin_resume("j2") + jm.complete_resume("j2", delivered=False, error="turn failed") + assert rec2.resume_state == ResumeState.FAILED.value + assert rec2.resume_error == "turn failed" + assert jm.awaiting_resume() == [] # terminal either way + + def test_interrupted_resume_is_recovered_as_failed_not_replayed(self, tmp_path): + """A crash between claim and delivery must not silently re-deliver.""" + import json + + from agentic_cli.tools.jobs import JobManager + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + self._finished_job(jm) + jm.begin_resume("j1") # process dies here + + # The claim records *who* holds it. Rewrite it to a process that is + # genuinely gone, which is what "the CLI crashed" looks like on disk. + meta = tmp_path / "jobs" / "j1" / "meta.json" + data = json.loads(meta.read_text()) + data["resume_owner"] = f"{_this_host()}:{_dead_pid()}" + meta.write_text(json.dumps(data)) + + reloaded = JobManager(base_dir=tmp_path / "jobs") + rec = reloaded._records["j1"] + assert rec.resume_state == ResumeState.FAILED.value + assert "interrupted" in (rec.resume_error or "") + assert reloaded.awaiting_resume() == [] # no automatic replay + + def test_ownerless_interrupted_resume_is_recovered(self, tmp_path): + """Records written before claims were owned still recover.""" + import json + + from agentic_cli.tools.jobs import JobManager + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + self._finished_job(jm) + jm.begin_resume("j1") + + meta = tmp_path / "jobs" / "j1" / "meta.json" + data = json.loads(meta.read_text()) + data.pop("resume_owner", None) + meta.write_text(json.dumps(data)) + + reloaded = JobManager(base_dir=tmp_path / "jobs") + assert reloaded._records["j1"].resume_state == ResumeState.FAILED.value + + def test_legacy_resumed_flag_migrates(self, tmp_path): + """Records written before the lifecycle carried a bool ``resumed``.""" + import json + + from agentic_cli.tools.jobs import JobManager + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + self._finished_job(jm) + meta = tmp_path / "jobs" / "j1" / "meta.json" + data = json.loads(meta.read_text()) + data.pop("resume_state", None) + data["resumed"] = True + meta.write_text(json.dumps(data)) + + reloaded = JobManager(base_dir=tmp_path / "jobs") + assert reloaded._records["j1"].resume_state == ResumeState.DELIVERED.value + assert reloaded.awaiting_resume() == [] + + +class TestCrossProcessResumeClaims: + """Two JobManagers over one jobs directory = two CLI processes. + + The claim used to consult only in-memory records, so both would see + ``PENDING`` and both would deliver the same job result into their own + conversation. And each manager's startup recovery flipped the *other's* + live ``RESUMING`` claim to FAILED behind its back. + """ + + @staticmethod + def _seed_finished_job(base: Path, job_id: str = "j1") -> None: + """Persist a terminal, resume-flagged job without holding a manager.""" + seeder = JobManager(base_dir=base) + rec = JobRecord( + job_id=job_id, tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + session_id="s", user_id="u", finished_at=1.0, + ) + seeder._records[rec.job_id] = rec + seeder._job_dir(rec.job_id).mkdir(parents=True, exist_ok=True) + seeder._persist(rec) + seeder.close() + + def test_exactly_one_of_two_managers_claims(self, tmp_path: Path): + base = tmp_path / "jobs" + self._seed_finished_job(base) + + first = JobManager(base_dir=base) + second = JobManager(base_dir=base) + try: + assert [r.job_id for r in first.awaiting_resume()] == ["j1"] + assert [r.job_id for r in second.awaiting_resume()] == ["j1"] + + claims = [first.begin_resume("j1"), second.begin_resume("j1")] + assert claims.count(True) == 1, ( + f"both managers claimed the same job result: {claims}" + ) + finally: + first.close() + second.close() + + def test_loser_sees_the_claim_after_reloading(self, tmp_path: Path): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + self._seed_finished_job(base) + + first = JobManager(base_dir=base) + second = JobManager(base_dir=base) + try: + assert first.begin_resume("j1") is True + assert second.begin_resume("j1") is False + assert second._records["j1"].resume_state == ResumeState.RESUMING.value + assert second.awaiting_resume() == [] + finally: + first.close() + second.close() + + def test_startup_does_not_fail_a_live_claim(self, tmp_path: Path): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + self._seed_finished_job(base) + + holder = JobManager(base_dir=base) + try: + assert holder.begin_resume("j1") is True + + # A second process starts while the first is mid-delivery. + newcomer = JobManager(base_dir=base) + try: + assert ( + newcomer._records["j1"].resume_state + == ResumeState.RESUMING.value + ), "a live claim was recovered as failed" + assert holder._records["j1"].resume_state == ResumeState.RESUMING.value + + holder.complete_resume("j1", delivered=True) + assert ( + holder._records["j1"].resume_state == ResumeState.DELIVERED.value + ) + finally: + newcomer.close() + finally: + holder.close() + + def test_completing_another_processes_claim_raises(self, tmp_path: Path): + """Only the claiming process closes a claim out.""" + from agentic_cli.tools.jobs.manager import ResumeState, ResumeStateError + + base = tmp_path / "jobs" + self._seed_finished_job(base) + + mine = JobManager(base_dir=base) + try: + assert mine.begin_resume("j1") is True + + # Another (live) process took the claim between our claim and our + # completion — on disk, that is what it looks like. + meta = base / "j1" / "meta.json" + data = json.loads(meta.read_text()) + data["resume_owner"] = f"{_this_host()}:{os.getpid() + 1}" + meta.write_text(json.dumps(data)) + + with pytest.raises(ResumeStateError, match="being delivered by"): + mine.complete_resume("j1", delivered=True) + + # And the other process's claim is intact. + assert ( + json.loads(meta.read_text())["resume_state"] + == ResumeState.RESUMING.value + ) + finally: + mine.close() + + +class _GatedBackend(JobBackend): + """Jobs stay RUNNING until the shared gate is opened, then SUCCEEDED. + + Two managers share one gate, so the test controls exactly when each of + them *observes* the finish (and therefore when each writes metadata). + """ + + name = "gated" + survives_restart = True + + def __init__(self, gate: dict) -> None: + self._gate = gate + + def start(self, record, job_dir): # pragma: no cover - never started here + record.state = JobState.RUNNING + + def poll(self, record, job_dir): + return JobState.SUCCEEDED if self._gate.get("done") else JobState.RUNNING + + def cancel(self, record, job_dir): # pragma: no cover + return None + + +class TestResumeStateIsAtomicAcrossProcesses: + """A normal metadata write must not clobber another process's claim. + + Only ``begin_resume``/``complete_resume`` were lock-and-reload aware. + Every *other* persist (``_refresh`` after a poll, ``cancel``, ``submit``) + wrote the whole record from memory — including a stale ``resume_state`` — + so a second CLI merely *observing* a job erased the first one's claim and + could then claim it too, delivering the same result into two conversations. + """ + + @staticmethod + def _manager(base: Path, gate: dict) -> JobManager: + # The backend must be known before _load_existing reconciles, or the + # record is written off as UNKNOWN on the way in. + return JobManager( + base_dir=base, + backends={**default_backends(), "gated": _GatedBackend(gate)}, + ) + + @classmethod + def _seed_running_job(cls, base: Path, gate: dict, job_id: str = "j1") -> None: + seeder = cls._manager(base, gate) + rec = JobRecord( + job_id=job_id, tool="run_shell_job", backend="gated", name="build", + state=JobState.RUNNING, resume_on_complete=True, call_id="c1", + session_id="s", user_id="u", + ) + seeder._records[rec.job_id] = rec + seeder._job_dir(rec.job_id).mkdir(parents=True, exist_ok=True) + seeder._persist(rec) + seeder.close() + + def test_observer_persist_does_not_erase_a_claim(self, tmp_path: Path): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + gate: dict = {} + self._seed_running_job(base, gate) + + first = self._manager(base, gate) + second = self._manager(base, gate) # both still see it RUNNING + try: + gate["done"] = True + + # A observes the finish, then claims it. + assert first.get("j1").state is JobState.SUCCEEDED + assert first.begin_resume("j1") is True + + # B observes the same finish — a plain metadata write, from a + # record whose resume fields predate A's claim. + assert second.get("j1").state is JobState.SUCCEEDED + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.RESUMING.value, ( + "an observer's persist erased the other process's claim" + ) + assert second.begin_resume("j1") is False, "the job was claimed twice" + finally: + first.close() + second.close() + + def test_observer_persist_does_not_erase_a_completed_delivery( + self, tmp_path: Path + ): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + gate: dict = {} + self._seed_running_job(base, gate) + + first = self._manager(base, gate) + second = self._manager(base, gate) + try: + gate["done"] = True + first.get("j1") + assert first.begin_resume("j1") is True + first.complete_resume("j1", delivered=True) + + second.get("j1") # observer write + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.DELIVERED.value + assert second.begin_resume("j1") is False + finally: + first.close() + second.close() + + def test_stale_recovery_cannot_overwrite_delivered(self, tmp_path: Path): + """Startup recovery must re-read under the lock before deciding.""" + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + gate: dict = {} + self._seed_running_job(base, gate) + + holder = self._manager(base, gate) + stale = self._manager(base, gate) + try: + gate["done"] = True + holder.get("j1") + assert holder.begin_resume("j1") is True + + # ``stale`` loaded the record before the claim; give it the view a + # crashed-owner claim would have, then complete the real delivery. + rec = stale._records["j1"] + rec.resume_state = ResumeState.RESUMING.value + rec.resume_owner = f"{_this_host()}:{_dead_pid()}" + holder.complete_resume("j1", delivered=True) + + stale._recover_interrupted_resumes() + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.DELIVERED.value, ( + "stale recovery overwrote a completed delivery" + ) + assert stale._records["j1"].resume_state == ResumeState.DELIVERED.value + finally: + holder.close() + stale.close() + + def test_unavailable_lock_fails_the_claim_closed(self, tmp_path: Path, monkeypatch): + """No lock, no claim — a duplicate delivery is worse than a late one.""" + from agentic_cli.tools.jobs import manager as jobs_manager + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + def _no_lock(fd): + raise OSError("flock unavailable") + + monkeypatch.setattr(jobs_manager, "_acquire_lock", _no_lock) + + assert jm.begin_resume("j1") is False + assert jm._records["j1"].resume_state == ResumeState.PENDING.value + finally: + jm.close() + + def test_unavailable_lock_blocks_startup_recovery( + self, tmp_path: Path, monkeypatch + ): + from agentic_cli.tools.jobs import manager as jobs_manager + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + claimer = JobManager(base_dir=base) + try: + assert claimer.begin_resume("j1") is True + meta = base / "j1" / "meta.json" + data = json.loads(meta.read_text()) + data["resume_owner"] = f"{_this_host()}:{_dead_pid()}" + meta.write_text(json.dumps(data)) + + def _no_lock(fd): + raise OSError("flock unavailable") + + monkeypatch.setattr(jobs_manager, "_acquire_lock", _no_lock) + + reloaded = JobManager(base_dir=base) + try: + assert ( + reloaded._records["j1"].resume_state + == ResumeState.RESUMING.value + ), "recovery rewrote state it could not lock" + finally: + reloaded.close() + finally: + claimer.close() + + +class _InProcessLikeBackend(JobBackend): + """Only the process that started a job can see it running. + + Mirrors the in-process backend: the handle lives in memory, so another + process polling the same record can only answer UNKNOWN. + """ + + name = "inproc" + survives_restart = False + + def __init__(self) -> None: + self._handles: set[str] = set() + + def start(self, record, job_dir): + self._handles.add(record.job_id) + + def poll(self, record, job_dir): + return JobState.RUNNING if record.job_id in self._handles else JobState.UNKNOWN + + def cancel(self, record, job_dir): # pragma: no cover + self._handles.discard(record.job_id) + + +class _CountingBackend(JobBackend): + """Records every launch, so a double start is visible.""" + + name = "counting" + survives_restart = True + + def __init__(self, log: list[str]) -> None: + self._log = log + + def start(self, record, job_dir): + self._log.append(record.job_id) + + def poll(self, record, job_dir): + return JobState.RUNNING + + def cancel(self, record, job_dir): # pragma: no cover + return None + + +class TestExecutionOwnership: + """Who is *running* a job is shared state, just like who is delivering it. + + Two CLIs share the jobs directory. Without a recorded execution owner, a + second one polled a job whose handle lives in the first one's memory, got + UNKNOWN, and marked a perfectly healthy job terminal — making it resumable + and its result deliverable. And two managers that both saw a job QUEUED + both launched it. + """ + + @staticmethod + def _manager(base: Path, backend: JobBackend) -> JobManager: + return JobManager( + base_dir=base, backends={**default_backends(), backend.name: backend} + ) + + def test_second_manager_leaves_a_live_in_process_job_alone(self, tmp_path: Path): + base = tmp_path / "jobs" + first = self._manager(base, _InProcessLikeBackend()) + try: + rec = first.submit( + tool="t", backend="inproc", spec={}, resume_on_complete=True, + session_id="s", user_id="u", + ) + assert rec.state is JobState.RUNNING + + second = self._manager(base, _InProcessLikeBackend()) + try: + mine = second._records[rec.job_id] + assert mine.state is JobState.RUNNING, ( + "another process's running job was written off as UNKNOWN" + ) + + second.reconcile() + assert second._records[rec.job_id].state is JobState.RUNNING + assert second.awaiting_resume() == [], ( + "a live job became deliverable to a second process" + ) + + on_disk = json.loads((base / rec.job_id / "meta.json").read_text()) + assert on_disk["state"] == JobState.RUNNING.value + finally: + second.close() + finally: + first.close() + + def test_a_dead_owners_job_is_still_reconciled(self, tmp_path: Path): + """Ownership defers to a *live* process only.""" + base = tmp_path / "jobs" + first = self._manager(base, _InProcessLikeBackend()) + try: + rec = first.submit(tool="t", backend="inproc", spec={}) + meta = base / rec.job_id / "meta.json" + data = json.loads(meta.read_text()) + data["exec_owner"] = f"{_this_host()}:{_dead_pid()}" + meta.write_text(json.dumps(data)) + + second = self._manager(base, _InProcessLikeBackend()) + try: + assert second._records[rec.job_id].state is JobState.UNKNOWN + finally: + second.close() + finally: + first.close() + + def test_two_managers_launch_a_queued_job_once(self, tmp_path: Path): + base = tmp_path / "jobs" + starts: list[str] = [] + first = self._manager(base, _CountingBackend(starts)) + second = self._manager(base, _CountingBackend(starts)) + try: + queued = JobRecord( + job_id="q1", tool="t", backend="counting", name="n", + state=JobState.QUEUED, + ) + first._records["q1"] = queued + first._job_dir("q1").mkdir(parents=True, exist_ok=True) + first._persist(queued) + # The second process loaded its own copy while it was still queued. + second._records["q1"] = JobRecord.from_dict(queued.to_dict()) + + first._maybe_start_queued() + second._maybe_start_queued() + + assert starts == ["q1"], f"the job was launched {len(starts)} times" + finally: + first.close() + second.close() + + def test_launch_claim_is_recorded_durably(self, tmp_path: Path): + base = tmp_path / "jobs" + starts: list[str] = [] + jm = self._manager(base, _CountingBackend(starts)) + try: + rec = jm.submit(tool="t", backend="counting", spec={}) + on_disk = json.loads((base / rec.job_id / "meta.json").read_text()) + # :: — the handle lives in a specific + # manager, so a sibling in the same process is a different owner. + assert on_disk["exec_owner"].startswith(f"{_this_host()}:{os.getpid()}:") + finally: + jm.close() + + def test_interrupted_launch_is_failed_not_replayed(self, tmp_path: Path): + base = tmp_path / "jobs" + starts: list[str] = [] + seeder = self._manager(base, _CountingBackend(starts)) + queued = JobRecord( + job_id="q1", tool="t", backend="counting", name="n", + state=JobState.QUEUED, exec_owner=f"{_this_host()}:{_dead_pid()}", + ) + seeder._records["q1"] = queued + seeder._job_dir("q1").mkdir(parents=True, exist_ok=True) + seeder._persist(queued, owns_shared_state=True) + seeder.close() + + reloaded = self._manager(base, _CountingBackend(starts)) + try: + rec = reloaded._records["q1"] + assert rec.state is JobState.FAILED + assert "interrupted" in (rec.error or "") + assert starts == [], "an interrupted launch was replayed" + finally: + reloaded.close() + + def test_unlocked_launch_is_skipped(self, tmp_path: Path, monkeypatch): + """Fail closed: a late launch beats a double launch.""" + from agentic_cli.tools.jobs import manager as jobs_manager + + base = tmp_path / "jobs" + starts: list[str] = [] + jm = self._manager(base, _CountingBackend(starts)) + try: + queued = JobRecord( + job_id="q1", tool="t", backend="counting", name="n", + state=JobState.QUEUED, + ) + jm._records["q1"] = queued + jm._job_dir("q1").mkdir(parents=True, exist_ok=True) + jm._persist(queued) + + def _boom(fd): + raise OSError("flock unavailable") + + monkeypatch.setattr(jobs_manager, "_acquire_lock", _boom) + jm._maybe_start_queued() + + assert starts == [] + assert jm._records["q1"].state is JobState.QUEUED + finally: + jm.close() + + +class _GatedInProcessBackend(JobBackend): + """In-process semantics: only the starter can poll, gated completion.""" + + name = "gated_inproc" + survives_restart = False + + def __init__(self, gate: dict) -> None: + self._gate = gate + self._handles: set[str] = set() + + def start(self, record, job_dir): + self._handles.add(record.job_id) + + def poll(self, record, job_dir): + if record.job_id not in self._handles: + return JobState.UNKNOWN # no handle here — cannot judge + return JobState.SUCCEEDED if self._gate.get("done") else JobState.RUNNING + + def cancel(self, record, job_dir): + self._cancelled = True + self._handles.discard(record.job_id) + + +class _SentinelBackend(JobBackend): + """Restart-safe *and* remotely cancellable (a durable handle, like a pid).""" + + name = "sentinel" + survives_restart = True + cancels_foreign_jobs = True + + def __init__(self, gate: dict) -> None: + self._gate = gate + self.cancelled: list[str] = [] + + def start(self, record, job_dir): + return None + + def poll(self, record, job_dir): + return JobState.SUCCEEDED if self._gate.get("done") else JobState.RUNNING + + def cancel(self, record, job_dir): + self.cancelled.append(record.job_id) + + +class TestJobRecordLifecycleCoherence: + """The whole persisted record is shared state, not just its owner fields. + + Execution ownership stopped a second manager from mangling a running job, + but it also stopped it from ever learning the job had *finished*: the + observer skipped the record entirely, so a completed job stayed RUNNING in + its view forever. And every other mutator (cancel, clean, recovery) still + acted on whatever it had in memory, so a stale copy could overwrite a + terminal outcome another manager had already recorded. + """ + + @staticmethod + def _manager(base: Path, backend: JobBackend) -> JobManager: + return JobManager( + base_dir=base, backends={**default_backends(), backend.name: backend} + ) + + def test_observer_sees_the_owners_terminal_state(self, tmp_path: Path): + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _GatedInProcessBackend(gate)) + try: + rec = owner.submit(tool="t", backend="gated_inproc", spec={}) + assert rec.state is JobState.RUNNING + + observer = self._manager(base, _GatedInProcessBackend(gate)) + try: + assert observer._records[rec.job_id].state is JobState.RUNNING + + gate["done"] = True + owner.reconcile() # the owner records the outcome durably + assert owner._records[rec.job_id].state is JobState.SUCCEEDED + + observer.reconcile() + assert observer._records[rec.job_id].state is JobState.SUCCEEDED, ( + "a completed job stayed RUNNING for the observer" + ) + finally: + observer.close() + finally: + owner.close() + + def test_observer_reads_a_restart_safe_sentinel_itself(self, tmp_path: Path): + """A shared sentinel needs no owner to report it.""" + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _SentinelBackend(gate)) + try: + rec = owner.submit(tool="t", backend="sentinel", spec={}) + observer = self._manager(base, _SentinelBackend(gate)) + try: + gate["done"] = True + observer.reconcile() + assert observer._records[rec.job_id].state is JobState.SUCCEEDED + finally: + observer.close() + finally: + owner.close() + + def test_stale_cancel_does_not_overwrite_a_completed_job(self, tmp_path: Path): + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _SentinelBackend(gate)) + observer_backend = _SentinelBackend(gate) + try: + rec = owner.submit(tool="t", backend="sentinel", spec={}) + observer = self._manager(base, observer_backend) + try: + gate["done"] = True + owner.reconcile() + assert owner._records[rec.job_id].state is JobState.SUCCEEDED + + # The observer still believes it is running. + assert observer._records[rec.job_id].state is JobState.RUNNING + cancelled = observer.cancel(rec.job_id) + + assert cancelled.state is JobState.SUCCEEDED, ( + "a stale cancel overwrote a completed job" + ) + on_disk = json.loads((base / rec.job_id / "meta.json").read_text()) + assert on_disk["state"] == JobState.SUCCEEDED.value + assert observer_backend.cancelled == [] + finally: + observer.close() + finally: + owner.close() + + def test_foreign_in_process_job_is_not_reported_cancelled(self, tmp_path: Path): + """This backend instance holds no handle — it cannot cancel anything.""" + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _GatedInProcessBackend(gate)) + try: + rec = owner.submit(tool="t", backend="gated_inproc", spec={}) + observer = self._manager(base, _GatedInProcessBackend(gate)) + try: + result = observer.cancel(rec.job_id) + + assert result.state is JobState.RUNNING, ( + "a job this manager cannot cancel was reported cancelled" + ) + on_disk = json.loads((base / rec.job_id / "meta.json").read_text()) + assert on_disk["state"] == JobState.RUNNING.value + assert owner._records[rec.job_id].state is JobState.RUNNING + finally: + observer.close() + finally: + owner.close() + + def test_a_restart_safe_foreign_job_can_still_be_cancelled(self, tmp_path: Path): + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _SentinelBackend(gate)) + observer_backend = _SentinelBackend(gate) + try: + rec = owner.submit(tool="t", backend="sentinel", spec={}) + observer = self._manager(base, observer_backend) + try: + result = observer.cancel(rec.job_id) + assert result.state is JobState.CANCELLED + assert observer_backend.cancelled == [rec.job_id] + finally: + observer.close() + finally: + owner.close() + + def test_terminal_state_is_monotonic_on_write(self, tmp_path: Path): + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + rec = jm._records["j1"] + rec.state = JobState.CANCELLED # a stale view, about to be written + + jm._persist(rec) + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["state"] == JobState.SUCCEEDED.value + assert rec.state is JobState.SUCCEEDED, "the stale view was not corrected" + finally: + jm.close() + + def test_clean_reloads_before_removing(self, tmp_path: Path): + """A stale terminal view must not delete a job that is still running.""" + base = tmp_path / "jobs" + gate: dict = {} + owner = self._manager(base, _GatedInProcessBackend(gate)) + try: + rec = owner.submit(tool="t", backend="gated_inproc", spec={}) + observer = self._manager(base, _GatedInProcessBackend(gate)) + try: + observer._records[rec.job_id].state = JobState.UNKNOWN # stale + + assert observer.clean() == 0, "a live job was cleaned away" + assert (base / rec.job_id / "meta.json").exists() + finally: + observer.close() + finally: + owner.close() + + def test_launch_recovery_respects_a_durable_outcome(self, tmp_path: Path): + """A queued+dead-owner view must not overwrite a recorded success.""" + base = tmp_path / "jobs" + gate: dict = {"done": True} + owner = self._manager(base, _SentinelBackend(gate)) + try: + rec = owner.submit(tool="t", backend="sentinel", spec={}) + owner.reconcile() + assert owner._records[rec.job_id].state is JobState.SUCCEEDED + + stale = self._manager(base, _SentinelBackend(gate)) + try: + mine = stale._records[rec.job_id] + mine.state = JobState.QUEUED + mine.exec_owner = f"{_this_host()}:{_dead_pid()}" + + stale._recover_interrupted_launches() + + assert mine.state is JobState.SUCCEEDED + on_disk = json.loads((base / rec.job_id / "meta.json").read_text()) + assert on_disk["state"] == JobState.SUCCEEDED.value + finally: + stale.close() + finally: + owner.close() + + def test_unlocked_reconcile_is_skipped(self, tmp_path: Path, monkeypatch): + from agentic_cli.tools.jobs import manager as jobs_manager + + base = tmp_path / "jobs" + gate: dict = {} + jm = self._manager(base, _SentinelBackend(gate)) + try: + rec = jm.submit(tool="t", backend="sentinel", spec={}) + gate["done"] = True + + def _boom(fd): + raise OSError("flock unavailable") + + monkeypatch.setattr(jobs_manager, "_acquire_lock", _boom) + jm.reconcile() + + assert jm._records[rec.job_id].state is JobState.RUNNING + finally: + jm.close() + + +def _wait_for_gate(path: str, timeout: float = 5.0) -> str: + """Block until ``path`` appears. Module level so the spec stays copyable.""" + import os + + deadline = time.time() + timeout + while time.time() < deadline: + if os.path.exists(path): + return "done" + time.sleep(0.02) + return "timeout" + + +class TestObserverSeesRealBackendCompletion: + """The shipped in-process backend, observed from a second manager. + + ``InProcessBackend.poll`` reads the on-disk ``exit_code`` sentinel *before* + consulting its in-memory future, so a second manager can read the outcome — + it just answers UNKNOWN while the sentinel is absent. Skipping the poll + entirely (because the record is foreign-owned) meant the observer never saw + the job finish; polling and taking UNKNOWN at face value would have marked + a live job terminal. Poll, and ignore only the "I cannot tell" answer. + """ + + @staticmethod + def _manager(base: Path) -> JobManager: + return JobManager(base_dir=base) + + def test_observer_sees_the_sentinel_after_the_owner_stops_polling( + self, tmp_path: Path + ): + base = tmp_path / "jobs" + gate_file = tmp_path / "gate" + owner = self._manager(base) + observer = None + try: + rec = owner.submit( + tool="t", + backend="inprocess", + spec={"target": _wait_for_gate, "kwargs": {"path": str(gate_file)}}, + resume_on_complete=True, + session_id="s", + user_id="u", + ) + assert rec.state is JobState.RUNNING + + observer = self._manager(base) + observer.reconcile() + assert observer._records[rec.job_id].state is JobState.RUNNING, ( + "an UNKNOWN poll from a foreign backend marked a live job terminal" + ) + + gate_file.write_text("go") + owner.close() # the owner stops polling; its process stays alive + + deadline = time.time() + 5 + while time.time() < deadline: + observer.reconcile() + if observer._records[rec.job_id].state is JobState.SUCCEEDED: + break + time.sleep(0.05) + + assert observer._records[rec.job_id].state is JobState.SUCCEEDED, ( + "the observer never saw the durable sentinel" + ) + assert [r.job_id for r in observer.awaiting_resume()] == [rec.job_id] + finally: + if observer is not None: + observer.close() + owner.close() + + def test_foreign_cancellation_is_a_backend_capability(self, tmp_path: Path): + """Not inferred from restart-safety — the two are different questions.""" + from agentic_cli.tools.jobs.backends import ( + InProcessBackend, + SubprocessBackend, + ) + + assert InProcessBackend.cancels_foreign_jobs is False + assert SubprocessBackend.cancels_foreign_jobs is True + + +class TestTerminalRecordsAreReloaded: + """A terminal record still has shared state: its delivery lifecycle.""" + + @staticmethod + def _pair(base: Path): + TestCrossProcessResumeClaims._seed_finished_job(base) + return JobManager(base_dir=base), JobManager(base_dir=base) + + def test_get_list_and_awaiting_reflect_a_completed_delivery( + self, tmp_path: Path + ): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + owner, observer = self._pair(base) + try: + assert [r.job_id for r in observer.awaiting_resume()] == ["j1"] + + assert owner.begin_resume("j1") is True + owner.complete_resume("j1", delivered=True) + + assert observer.get("j1").resume_state == ResumeState.DELIVERED.value + listed = {r.job_id: r for r in observer.list()} + assert listed["j1"].resume_state == ResumeState.DELIVERED.value + assert observer.awaiting_resume() == [] + finally: + owner.close() + observer.close() + + def test_begin_resume_reloads_before_judging_terminality(self, tmp_path: Path): + base = tmp_path / "jobs" + owner, observer = self._pair(base) + try: + # The observer's snapshot predates the job finishing. + observer._records["j1"].state = JobState.RUNNING + + assert observer.begin_resume("j1") is True, ( + "a stale non-terminal snapshot refused a deliverable job" + ) + finally: + owner.close() + observer.close() + + +class TestDeletedAndUnreadableRecords: + """A record another manager removed must stay removed.""" + + @staticmethod + def _pair(base: Path): + TestCrossProcessResumeClaims._seed_finished_job(base) + return JobManager(base_dir=base), JobManager(base_dir=base) + + def _assert_gone(self, base: Path) -> None: + assert not (base / "j1").exists(), "a deleted job record was recreated" + + def test_get_does_not_resurrect_a_cleaned_job(self, tmp_path: Path): + base = tmp_path / "jobs" + cleaner, stale = self._pair(base) + try: + assert cleaner.clean() == 1 + assert stale.get("j1") is None + assert "j1" not in stale._records + self._assert_gone(base) + finally: + cleaner.close() + stale.close() + + def test_cancel_does_not_resurrect_a_cleaned_job(self, tmp_path: Path): + base = tmp_path / "jobs" + cleaner, stale = self._pair(base) + try: + stale._records["j1"].state = JobState.RUNNING # stale, looks cancellable + assert cleaner.clean() == 1 + + assert stale.cancel("j1") is None + self._assert_gone(base) + finally: + cleaner.close() + stale.close() + + def test_begin_resume_does_not_resurrect_a_cleaned_job(self, tmp_path: Path): + base = tmp_path / "jobs" + cleaner, stale = self._pair(base) + try: + assert cleaner.clean() == 1 + assert stale.begin_resume("j1") is False + self._assert_gone(base) + finally: + cleaner.close() + stale.close() + + def test_complete_resume_does_not_resurrect_a_cleaned_job(self, tmp_path: Path): + import shutil + + base = tmp_path / "jobs" + cleaner, stale = self._pair(base) + try: + assert stale.begin_resume("j1") is True + # An active delivery is protected from clean() (see + # test_clean_leaves_an_active_delivery_alone), so model the record + # being removed out from under us directly. + shutil.rmtree(base / "j1") + + stale.complete_resume("j1", delivered=True) # must not raise + self._assert_gone(base) + finally: + cleaner.close() + stale.close() + + def test_unreadable_state_fails_closed(self, tmp_path: Path): + """Corrupt metadata is not permission to act — or to overwrite.""" + base = tmp_path / "jobs" + owner, observer = self._pair(base) + try: + meta = base / "j1" / "meta.json" + meta.write_text("{ this is not json") + + assert observer.begin_resume("j1") is False + assert "j1" in observer._records, "a corrupt record was forgotten" + assert meta.read_text() == "{ this is not json", ( + "a corrupt record was overwritten" + ) + finally: + owner.close() + observer.close() + + def test_clean_leaves_an_active_delivery_alone(self, tmp_path: Path): + base = tmp_path / "jobs" + deliverer, cleaner = self._pair(base) + try: + assert deliverer.begin_resume("j1") is True + + assert cleaner.clean() == 0, "a job being delivered was cleaned away" + assert (base / "j1" / "meta.json").exists() + finally: + deliverer.close() + cleaner.close() + + +class TestNoUnsafeUnlockedWrites: + """Without the lock there is no safe read-modify-write of shared fields. + + The fallback path read ``meta.json`` and then rewrote it whole — which is + precisely the race the lock exists to prevent. A delivery completed between + that read and that write was silently regressed from DELIVERED, and the + job could then be claimed and delivered a second time. + """ + + @staticmethod + def _no_lock(monkeypatch): + from agentic_cli.tools.jobs import manager as jobs_manager + + def _boom(fd): + raise OSError("flock unavailable") + + monkeypatch.setattr(jobs_manager, "_acquire_lock", _boom) + + @staticmethod + def _write_state(base: Path, job_id: str, **fields) -> None: + """Stand in for another process's write, straight to disk.""" + meta = base / job_id / "meta.json" + data = json.loads(meta.read_text()) + data.update(fields) + meta.write_text(json.dumps(data)) + + def test_unlocked_write_leaves_an_existing_record_untouched( + self, tmp_path: Path, monkeypatch + ): + """No lock, no rewrite. + + The read-then-rewrite fallback preserved the shared fields *only* when + nothing changed between its read and its write — which is the race the + lock exists to prevent, so it was never a safe fallback. The property + under test is therefore the absence of the write itself: an unlocked + persist must not touch a record another process may be inside. + """ + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + assert jm.begin_resume("j1") is True + rec = jm._records["j1"] + meta = base / "j1" / "meta.json" + before = meta.read_text() + + self._no_lock(monkeypatch) + rec.state = JobState.CANCELLED # an ordinary change we would persist + jm._persist(rec) + + assert meta.read_text() == before, ( + "an unlocked persist rewrote a record it could not lock" + ) + finally: + jm.close() + + def test_write_after_a_concurrent_complete_does_not_regress_it( + self, tmp_path: Path, monkeypatch + ): + """read → (another process completes) → write, with no lock held.""" + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + assert jm.begin_resume("j1") is True + rec = jm._records["j1"] + + self._no_lock(monkeypatch) + # Between our read and our write, the delivery completes elsewhere. + self._write_state( + base, + "j1", + resume_state=ResumeState.DELIVERED.value, + resume_owner=None, + ) + + jm._persist(rec) # an ordinary metadata write + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.DELIVERED.value, ( + "an unlocked write regressed a completed delivery" + ) + finally: + jm.close() + + def test_unlocked_write_creates_a_record_that_has_none( + self, tmp_path: Path, monkeypatch + ): + """Nothing on disk means nothing to clobber — the write must happen.""" + base = tmp_path / "jobs" + jm = JobManager(base_dir=base) + try: + self._no_lock(monkeypatch) + rec = JobRecord( + job_id="new", tool="t", backend="subprocess", name="n", + state=JobState.QUEUED, + ) + jm._job_dir("new").mkdir(parents=True, exist_ok=True) + jm._persist(rec) + + assert (base / "new" / "meta.json").exists() + finally: + jm.close() + + def test_unlocked_complete_resume_does_not_regress_delivered( + self, tmp_path: Path, monkeypatch + ): + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + assert jm.begin_resume("j1") is True + + self._no_lock(monkeypatch) + self._write_state( + base, + "j1", + resume_state=ResumeState.DELIVERED.value, + resume_owner=None, + ) + + jm.complete_resume("j1", delivered=False, error="turn failed") + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.DELIVERED.value + finally: + jm.close() + + def test_unlocked_complete_resume_leaves_the_record_claimed( + self, tmp_path: Path, monkeypatch + ): + """Failed-safe: still RESUMING on disk, so it is never re-delivered.""" + from agentic_cli.tools.jobs.manager import ResumeState + + base = tmp_path / "jobs" + TestCrossProcessResumeClaims._seed_finished_job(base) + jm = JobManager(base_dir=base) + try: + assert jm.begin_resume("j1") is True + self._no_lock(monkeypatch) + + jm.complete_resume("j1", delivered=True) # must not raise + + on_disk = json.loads((base / "j1" / "meta.json").read_text()) + assert on_disk["resume_state"] == ResumeState.RESUMING.value + assert jm.awaiting_resume() == [], "the job became deliverable again" + finally: + jm.close() + + +class TestBackendClose: + """Owned executors must not outlive the manager.""" + + def test_inprocess_backend_shuts_pool_down(self): + from agentic_cli.tools.jobs.backends import InProcessBackend + + backend = InProcessBackend(max_workers=2) + backend.close() + assert backend._pool._shutdown is True + + def test_close_is_idempotent(self): + from agentic_cli.tools.jobs.backends import InProcessBackend + + backend = InProcessBackend(max_workers=1) + backend.close() + backend.close() # must not raise + + def test_subprocess_backend_close_is_noop(self): + from agentic_cli.tools.jobs.backends import SubprocessBackend + + SubprocessBackend().close() + + def test_manager_close_closes_every_backend(self, tmp_path): + from agentic_cli.tools.jobs import JobManager + + jm = JobManager(base_dir=tmp_path / "jobs") + jm.close() + jm.close() # idempotent + assert jm._backends["inprocess"]._pool._shutdown is True + + +class TestResumeTransitionGuards: + """Only a deliverable job can be claimed; only a claim can be completed.""" + + def _jm(self, tmp_path): + from agentic_cli.tools.jobs import JobManager + + return JobManager(base_dir=tmp_path / "jobs") + + def _record(self, jm, **over): + from agentic_cli.tools.jobs import JobRecord + from agentic_cli.tools.jobs.backends import JobState + + fields = dict( + job_id="j1", tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + session_id="s", user_id="u", finished_at=1.0, + ) + fields.update(over) + rec = JobRecord(**fields) + jm._records[rec.job_id] = rec + jm._persist(rec) + return rec + + def test_running_job_cannot_be_claimed(self, tmp_path): + from agentic_cli.tools.jobs.backends import JobState + + jm = self._jm(tmp_path) + self._record(jm, state=JobState.RUNNING) + assert jm.begin_resume("j1") is False + + def test_job_without_resume_flag_cannot_be_claimed(self, tmp_path): + jm = self._jm(tmp_path) + self._record(jm, resume_on_complete=False) + assert jm.begin_resume("j1") is False + + def test_completing_an_unclaimed_job_raises(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeStateError + + jm = self._jm(tmp_path) + rec = self._record(jm) + with pytest.raises(ResumeStateError, match="pending"): + jm.complete_resume("j1", delivered=True) + assert rec.resume_state == "pending" # state untouched + + def test_completing_twice_raises(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeStateError + + jm = self._jm(tmp_path) + self._record(jm) + jm.begin_resume("j1") + jm.complete_resume("j1", delivered=True) + with pytest.raises(ResumeStateError): + jm.complete_resume("j1", delivered=False, error="late") + + def test_completing_an_unknown_job_raises(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeStateError + + with pytest.raises(ResumeStateError, match="Unknown job"): + self._jm(tmp_path).complete_resume("nope", delivered=True) + + def test_mark_resumed_shim_claims_then_completes(self, tmp_path): + from agentic_cli.tools.jobs.manager import ResumeState + + jm = self._jm(tmp_path) + rec = self._record(jm) + jm.mark_resumed("j1") + assert rec.resume_state == ResumeState.DELIVERED.value + jm.mark_resumed("j1") # already terminal: no-op, no raise + assert rec.resume_state == ResumeState.DELIVERED.value + + +class TestStartupRecoveryHonorsDurableReads: + """Recovery reloads durable state — and must act on what the read *said*. + + ``_recover_interrupted_launches`` / ``_recover_interrupted_resumes`` called + ``_reload_durable()`` and threw the answer away. A record another manager + deleted between the initial load and recovery was therefore rewritten from + the stale in-memory copy (``owns_shared_state=True`` writes unconditionally), + resurrecting a job that had been cleaned away; a record that could not be + parsed was overwritten with a recovery verdict decided on state we had + failed to read. + """ + + @staticmethod + def _seed(base: Path, **fields) -> JobManager: + """A manager holding one loaded record, matching what is on disk.""" + jm = JobManager(base_dir=base) + rec = JobRecord( + job_id="j1", tool="run_shell_job", backend="subprocess", name="build", + session_id="s", user_id="u", **fields, + ) + jm._records["j1"] = rec + jm._job_dir("j1").mkdir(parents=True, exist_ok=True) + jm._persist(rec) + return jm + + @staticmethod + def _interrupted_launch(base: Path) -> JobManager: + return TestStartupRecoveryHonorsDurableReads._seed( + base, + state=JobState.QUEUED, + exec_owner=f"{_this_host()}:{_dead_pid()}:abcd1234", + ) + + @staticmethod + def _interrupted_resume(base: Path) -> JobManager: + from agentic_cli.tools.jobs.manager import ResumeState + + jm = TestStartupRecoveryHonorsDurableReads._seed( + base, state=JobState.SUCCEEDED, resume_on_complete=True, finished_at=1.0 + ) + rec = jm._records["j1"] + rec.resume_state = ResumeState.RESUMING.value + rec.resume_owner = f"{_this_host()}:{_dead_pid()}" + jm._persist(rec, owns_shared_state=True) + return jm + + @pytest.mark.parametrize("kind", ["launches", "resumes"]) + def test_deleted_record_is_not_resurrected(self, tmp_path: Path, kind: str): + base = tmp_path / "jobs" + jm = ( + self._interrupted_launch(base) + if kind == "launches" + else self._interrupted_resume(base) + ) + try: + meta = base / "j1" / "meta.json" + assert meta.exists() + meta.unlink() # another manager cleaned the job away + + getattr(jm, f"_recover_interrupted_{kind}")() + + assert not meta.exists(), "recovery recreated a deleted record" + assert "j1" not in jm._records, "a deleted record was not forgotten" + finally: + jm.close() + + @pytest.mark.parametrize("kind", ["launches", "resumes"]) + def test_unreadable_record_is_left_untouched(self, tmp_path: Path, kind: str): + base = tmp_path / "jobs" + jm = ( + self._interrupted_launch(base) + if kind == "launches" + else self._interrupted_resume(base) + ) + try: + meta = base / "j1" / "meta.json" + meta.write_text("{not json") + before = meta.read_bytes() + + getattr(jm, f"_recover_interrupted_{kind}")() + + assert meta.read_bytes() == before, "recovery overwrote unreadable state" + finally: + jm.close() + + +def _touch_marker(marker: str) -> str: + Path(marker).write_text("ran") + return "ran" + + +class TestInProcessShutdownDoesNotStrandQueuedJobs: + """``close()`` promises jobs are not cancelled — the pool broke that promise. + + ``shutdown(cancel_futures=True)`` dropped work that was already submitted + and already recorded RUNNING. Nothing then wrote the ``exit_code`` sentinel, + so another manager polling the record got UNKNOWN, saw a live foreign + execution owner, correctly refused to believe the UNKNOWN — and left the + job RUNNING forever, undeliverable. + """ + + def test_queued_job_still_reaches_a_terminal_state(self, tmp_path: Path): + from agentic_cli.tools.jobs.backends import InProcessBackend + + base = tmp_path / "jobs" + gate, marker = tmp_path / "gate", tmp_path / "marker" + + # One worker, two jobs: the second is submitted (and RUNNING) but its + # future is still queued behind the first. + owner = JobManager( + base_dir=base, + max_concurrent=4, + backends={"inprocess": InProcessBackend(max_workers=1)}, + ) + try: + blocking = owner.submit( + tool="wait", backend="inprocess", + spec={"target": _wait_for_gate, "kwargs": {"path": str(gate), "timeout": 10.0}}, + ) + queued = owner.submit( + tool="mark", backend="inprocess", + spec={"target": _touch_marker, "args": (str(marker),)}, + resume_on_complete=True, call_id="c1", session_id="s", user_id="u", + ) + assert owner._records[blocking.job_id].state is JobState.RUNNING + assert owner._records[queued.job_id].state is JobState.RUNNING + finally: + owner.close() + + gate.write_text("go") # let the running job finish and the queue drain + + observer = JobManager(base_dir=base) + try: + rec = _wait(observer, queued.job_id, timeout=10.0) + assert rec.state is JobState.SUCCEEDED, ( + f"a submitted job was stranded as {rec.state.value}" + ) + assert marker.exists(), "the queued job never ran" + assert observer.begin_resume(queued.job_id) is True + finally: + observer.close() diff --git a/tests/workflow/test_adk_job_resume.py b/tests/workflow/test_adk_job_resume.py index 21f19a8..7dcc327 100644 --- a/tests/workflow/test_adk_job_resume.py +++ b/tests/workflow/test_adk_job_resume.py @@ -7,6 +7,8 @@ from __future__ import annotations +import asyncio + from types import SimpleNamespace import pytest @@ -69,15 +71,22 @@ async def run_async(self, *, session_id, user_id, new_message, run_config): def _resume_manager(runner: _FakeRunner, *, session_exists: bool = True) -> GoogleADKWorkflowManager: mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) - mgr._settings = SimpleNamespace(app_name="test", context_window_enabled=False) + mgr._settings = SimpleNamespace( + app_name="test", context_window_enabled=False, default_user="default_user" + ) mgr._app_name = "test" mgr._services = {} - mgr._active_session_id = None - mgr._active_user_id = None + # Turns serialize on the manager's turn lock (see the concurrency contract). + mgr._turn_lock = asyncio.Lock() + mgr._lifecycle_lock = asyncio.Lock() mgr._model = "gemini-2.5-flash" mgr._model_resolved = True mgr._on_event = None + # Turn admission re-checks that the backend is live while holding the turn + # lock, so the double has to present a complete one. + mgr._initialized = True mgr._runner = runner + mgr._root_agent = SimpleNamespace(name="root") mgr._llm_logging_plugin = None mgr._task_progress_plugin = None From a0f66dbdd5cdddceffc7e452f6c3d9166b6baeb7 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:48 -0400 Subject: [PATCH 10/11] feat(skills)!: expose run_skill_script only when a code executor is supplied MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``skill_scripts_enabled`` advertised ADK's ``run_skill_script`` to the model while no supported manager path wires a code executor, so every call it provoked answered ``NO_CODE_EXECUTOR``. A switch that cannot make the thing work is worse than no switch: the setting is removed and the tool is exposed exactly when ``make_skill_toolset`` is given an executor — the thing that actually makes it work. In practice scripts stay off; the parameter is there for a caller that owns one. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- src/agentic_cli/tools/skills/toolset.py | 23 +++++++++++--------- src/agentic_cli/workflow/adk/manager.py | 13 ++++++------ src/agentic_cli/workflow/settings.py | 11 +++++----- tests/workflow/test_skills.py | 28 +++++++++++++++++++++++-- 4 files changed, 51 insertions(+), 24 deletions(-) diff --git a/src/agentic_cli/tools/skills/toolset.py b/src/agentic_cli/tools/skills/toolset.py index be5e906..31c3b47 100644 --- a/src/agentic_cli/tools/skills/toolset.py +++ b/src/agentic_cli/tools/skills/toolset.py @@ -1,9 +1,13 @@ """Build an ADK ``SkillToolset`` for a set of skills. -Wraps ADK's native toolset. When script execution is disabled (the default), -the ``run_skill_script`` tool is removed so it isn't advertised to the model; -the discovery/read tools (``list_skills``/``load_skill``/``load_skill_resource``) -and the L1 metadata prompt injection still work. +Wraps ADK's native toolset. ``run_skill_script`` is exposed only when a code +executor is actually supplied — the tool cannot run a script without one, and +advertising it regardless produced a guaranteed ``NO_CODE_EXECUTOR`` failure. +The discovery/read tools (``list_skills``/``load_skill``/``load_skill_resource``) +and the L1 metadata prompt injection always work. + +No supported manager path wires an executor today, so in practice scripts stay +off; the parameter exists for a caller that owns one. """ from __future__ import annotations @@ -14,17 +18,16 @@ def make_skill_toolset( skills: list[Any], *, - scripts_enabled: bool = False, code_executor: Any | None = None, additional_tools: list[Any] | None = None, ) -> Any: - """Create an ADK SkillToolset, optionally excluding script execution. + """Create an ADK SkillToolset; script execution follows the executor. Args: skills: Loaded ADK ``Skill`` objects. - scripts_enabled: If False (default), ``run_skill_script`` is removed. - code_executor: ADK code executor for script execution (only meaningful - when ``scripts_enabled`` is True). + code_executor: ADK code executor for script execution. When None + (the default), ``run_skill_script`` is removed from the toolset + rather than offered and then failing. additional_tools: Tools surfaced when a skill with ``adk_additional_tools`` frontmatter is activated. @@ -38,7 +41,7 @@ def make_skill_toolset( code_executor=code_executor, additional_tools=additional_tools or [], ) - if not scripts_enabled: + if code_executor is None: toolset._tools = [ t for t in toolset._tools if not isinstance(t, RunSkillScriptTool) ] diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index 625bddf..f84f105 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -629,11 +629,13 @@ def _assemble_agent_tools( return tools def _build_skill_toolset(self, skill_refs: list[str]): - """Resolve skill refs and build an ADK SkillToolset (scripts gated). + """Resolve skill refs and build an ADK SkillToolset (discovery/read only). - Script execution is disabled unless ``settings.skill_scripts_enabled`` - is True (and a code executor is wired — a future enhancement), so by - default only discovery/read tools are exposed. + This path supplies no code executor, so ``run_skill_script`` is not + exposed: only the discovery/read tools and the L1 metadata injection. + (There used to be a ``skill_scripts_enabled`` setting that advertised + the tool anyway; with no executor it could only ever answer + ``NO_CODE_EXECUTOR``.) """ from agentic_cli.tools.skills import SkillStore, make_skill_toolset @@ -641,8 +643,7 @@ def _build_skill_toolset(self, skill_refs: list[str]): skills = store.resolve(skill_refs) if not skills: return None - scripts_enabled = getattr(self._settings, "skill_scripts_enabled", False) - return make_skill_toolset(skills, scripts_enabled=scripts_enabled) + return make_skill_toolset(skills) def _wrap_long_running(self, tools: list[Callable]) -> list: """Wrap tools flagged ``long_running`` as ADK ``LongRunningFunctionTool``. diff --git a/src/agentic_cli/workflow/settings.py b/src/agentic_cli/workflow/settings.py index 6ba6db5..b152e0e 100644 --- a/src/agentic_cli/workflow/settings.py +++ b/src/agentic_cli/workflow/settings.py @@ -458,12 +458,11 @@ def _validate_sandbox_network(cls, v: str) -> str: description="Directories searched for named skills (Agent Skills / SKILL.md folders)", json_schema_extra={"ui_order": 141}, ) - skill_scripts_enabled: bool = Field( - default=False, - title="Skill Scripts Enabled", - description="Allow executing scripts bundled with skills (requires a code executor; disabled by default)", - json_schema_extra={"ui_order": 142}, - ) + # NOTE: ``skill_scripts_enabled`` was removed. Turning it on exposed ADK's + # ``run_skill_script`` while the supported manager path supplies no code + # executor, so every call answered ``NO_CODE_EXECUTOR``. Script execution is + # now enabled by passing a code executor to ``make_skill_toolset`` — the + # thing that actually makes it work — instead of by a switch that cannot. # Persistence settings (LangGraph) postgres_uri: str | None = Field( diff --git a/tests/workflow/test_skills.py b/tests/workflow/test_skills.py index de5e060..ccab741 100644 --- a/tests/workflow/test_skills.py +++ b/tests/workflow/test_skills.py @@ -77,11 +77,35 @@ def test_scripts_disabled_by_default(self, tmp_path): assert "run_skill_script" not in names assert {"list_skills", "load_skill", "load_skill_resource"} <= names - def test_scripts_enabled_includes_run_tool(self, tmp_path): + def test_run_tool_appears_only_with_a_code_executor(self, tmp_path): + """The tool is exposed by the thing that makes it work, not by a flag. + + A flag could enable it without an executor, and every call then failed + with NO_CODE_EXECUTOR. + """ skills = SkillStore().resolve([str(_make_skill(tmp_path))]) - ts = make_skill_toolset(skills, scripts_enabled=True) + ts = make_skill_toolset(skills, code_executor=object()) assert "run_skill_script" in {t.name for t in ts._tools} + def test_manager_path_never_exposes_the_script_tool(self, tmp_path): + """The supported manager path supplies no executor, so scripts stay off.""" + from types import SimpleNamespace + + import pytest + + pytest.importorskip("google.adk") + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = SimpleNamespace(skills_dirs=[]) + toolset = mgr._build_skill_toolset([str(_make_skill(tmp_path))]) + assert "run_skill_script" not in {t.name for t in toolset._tools} + + def test_removed_setting_is_gone(self): + from agentic_cli.config import BaseSettings + + assert "skill_scripts_enabled" not in BaseSettings.model_fields + # --------------------------------------------------------------------------- # Permission registration From 5a73940e7ba2bdfb7c2faabd85e279634f63c1d5 Mon Sep 17 00:00:00 2001 From: Andrey Shiryaev <7615137+shoom1@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:13:49 -0400 Subject: [PATCH 11/11] docs: record the review-pass contracts and publish the import surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CHANGELOG/README/CLAUDE.md for the preceding ten commits, plus the top-level exports the new contracts are addressed by — ``SessionRef``, ``TurnResult``, ``TurnStatus``, ``WorkflowState``, ``AgentGraphError`` — and a test that pins the import surface so a package reshuffle cannot quietly drop one. Claude-Session: https://claude.ai/code/session_01BpktWb9vKN5exdT9MbLQYh --- CHANGELOG.md | 43 ++++++++++++- CLAUDE.md | 10 ++- README.md | 60 +++++++++++++++++- src/agentic_cli/__init__.py | 11 +++- tests/test_import_surface.py | 114 +++++++++++++++++++++++++++++++++++ 5 files changed, 231 insertions(+), 7 deletions(-) create mode 100644 tests/test_import_surface.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 19501b6..68e11cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,10 +11,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Durable sessions by default** (resumable conversations across restarts): conversation state is now persisted continuously by each orchestrator's **native** store, keyed by session id — ADK via `DatabaseSessionService` (SQLAlchemy async), LangGraph via a persistent checkpointer (`thread_id == session_id`). A single `session_store` setting (`sqlite` default, `postgres`, or `memory` for ephemeral) drives both backends; `BaseSettings.session_db_url()` resolves the shared async URL (`sqlite+aiosqlite:///{workspace}/sessions/sessions.db` by default). Persistence is **on by default** — each run gets a fresh durable session id, and **`--session ` resumes** a stored one (creating it if new). Durability is per-event/per-step (crash-safe mid-turn), full-fidelity (real events incl. function-call ids — which also makes long-running resume survivable across restarts), and needs no save-on-exit. `/sessions` lists/deletes from the native store (`list_sessions`/`delete_session` on the manager); new `BaseWorkflowManager.session_exists`/`recent_messages`. New deps `aiosqlite` + `greenlet` (SQLAlchemy async). Validated live end-to-end across a fresh manager over the same sqlite (`tests/integration/test_live_durable_sessions.py`). Built on ADK 1.33 — no ADK 2.0 upgrade needed (2.0's breaking Workflow-Runtime rewrite adds no session primitive `DatabaseSessionService` doesn't already provide). - **Long-running job substrate** (`tools/jobs/`, Tier A milestone 1): typed long-running tools start detached work via an internal `JobManager` over pluggable execution backends behind one `JobBackend` interface (ships **subprocess** + **in-process**). The LLM only ever sees the tool — `JobManager` is internal infrastructure, never an LLM-facing tool, and there is no generic `job_submit`. Includes restart-safe completion (on-disk `exit_code` sentinel; subprocesses run detached with `start_new_session`), persistence under `~/.{app_name}/jobs/`, a concurrency cap + queue (`max_concurrent_jobs`, default 4), observe-only management tools (capability `jobs.manage`) — `job_status` is the recommended companion to a long-running tool (it returns state, a stdout tail, and the result once finished, so most agents need only it; `JOB_TOOLS == [job_status]`), with `job_result`/`job_logs`/`job_cancel`/`job_list` as opt-in extras (`JOB_MANAGEMENT_TOOLS`) that also power `/jobs` — the `@register_tool(long_running=True)` flag, and a `/jobs` command (`/jobs`, `/jobs all`, `/jobs `, `/jobs cancel `, `/jobs clean`). The framework ships only the generic substrate + observe-only tools: a typed long-running **starter** tool (the half that actually launches work and declares `long_running=True` + a `longrunning.` capability) is application-provided, since it decides what runs and how. A subprocess-backed `run_shell_job` ships as a reference in `examples/jobs_demo.py` — it runs `sh -c` directly and does **not** go through the hardened shell tool (`tools/shell/`), so it is intentionally a demo, not a built-in. Auto-ingest-on-completion (push/resume) is deferred to a later milestone. - **Harness Jobs UI monitor** (`cli/job_monitor.py`, Tier A milestone 2): a background `JobMonitor` task — started for the lifetime of the CLI session, independent of the agent loop — periodically reconciles the `JobManager` (so detached jobs advance state with no LLM turn) and renders a live jobs segment into the status bar (`jobs: 2 running, 1 queued`), with a transient `✓`/`✗`/`⊘` note when a job finishes. The status bar is the only background-safe UI surface (`thinking_prompt` boxes are turn-oriented and `add_*` prints directly, which would corrupt the live prompt; `set_status` only invalidates the app); `WorkflowController` stays the single composer of the bar and reads the segment the monitor publishes. New `examples/jobs_demo.py` exercises it interactively. -- **Long-running job push/resume auto-ingest — ADK** (Tier A phase 2): when a long-running job that opted in (`resume_on_complete`) finishes, the agent is automatically resumed with its result — no polling. On ADK the result is delivered to the pending call as a `FunctionResponse` (`GoogleADKWorkflowManager.resume_with_job_result`); long-running tools are wrapped as `LongRunningFunctionTool` so the model leaves the call pending. The harness coordinator (`BaseCLIApp.resume_finished_jobs`) drains finished jobs into **serialized resume turns at turn boundaries** — one turn at a time via a turn lock, never overlapping a user turn or the live prompt — rendered through the same UI path as a user turn (`MessageProcessor.process_resume`, sharing `_run_turn` with `process`). Gated by the opt-in `job_auto_resume` setting (default off); `/resume` triggers it on demand; the status bar shows `↻N to resume`. The resume association (`session_id`/`user_id`/`call_id`/`call_name`/`resumed`) is tracked on the `JobRecord` and auto-filled from the active turn (`JobManager.submit(resume_on_complete=True)` reads the active session/user; `awaiting_resume`/`mark_resumed` are the coordinator's query/commit API). The coordinator and association layer are backend-agnostic; only `resume_with_job_result` is ADK-specific so far (LangGraph resume is not yet wired). A resume runs only when the originating conversation is still available (`BaseWorkflowManager.can_resume`, default False; ADK checks the session holds the pending call); when it isn't — e.g. after a CLI restart, since ADK's default session is in-memory — the harness posts a "finished while its conversation was unavailable — fetch with `/jobs `" notice instead of firing a dead resume turn (the persisted `↻N to resume` status-bar cue surfaces it across restarts; the result stays reachable by id). Validated live end-to-end (`tests/integration/test_live_job_resume.py`). +- **Long-running job push/resume auto-ingest — ADK** (Tier A phase 2): when a long-running job that opted in (`resume_on_complete`) finishes, the agent is automatically resumed with its result — no polling. On ADK the result is delivered to the pending call as a `FunctionResponse` (`GoogleADKWorkflowManager.resume_with_job_result`); long-running tools are wrapped as `LongRunningFunctionTool` so the model leaves the call pending. The harness coordinator (`BaseCLIApp.resume_finished_jobs`) drains finished jobs into **serialized resume turns at turn boundaries** — one turn at a time via a turn lock, never overlapping a user turn or the live prompt — rendered through the same UI path as a user turn (`MessageProcessor.process_resume`, sharing `_run_turn` with `process`). Gated by the opt-in `job_auto_resume` setting (default off); `/resume` triggers it on demand; the status bar shows `↻N to resume`. The resume association (`session_id`/`user_id`/`call_id`/`call_name`/`resume_state`) is tracked on the `JobRecord` and auto-filled from the active turn (`JobManager.submit(resume_on_complete=True)` reads the active session/user; `awaiting_resume`/`begin_resume`/`complete_resume` are the coordinator's query/claim/commit API). The coordinator and association layer are backend-agnostic; only `resume_with_job_result` is ADK-specific so far (LangGraph resume is not yet wired). A resume runs only when the originating conversation is still available (`BaseWorkflowManager.can_resume`, default False; ADK checks the session holds the pending call); when it isn't — e.g. after a CLI restart, since ADK's default session is in-memory — the harness posts a "finished while its conversation was unavailable — fetch with `/jobs `" notice instead of firing a dead resume turn (the persisted `↻N to resume` status-bar cue surfaces it across restarts; the result stays reachable by id). Validated live end-to-end (`tests/integration/test_live_job_resume.py`). + +### Changed +- **Background-job resume has an explicit delivery lifecycle.** `JobRecord.resumed` (a bool set *before* the turn ran) is replaced by `resume_state`: `pending → resuming → delivered | failed`, plus `resume_error` and `resume_owner`. The coordinator claims a job with `JobManager.begin_resume()` (only a terminal, `resume_on_complete`, pending record can be claimed) and closes it with `complete_resume(delivered=...)` — on success, failure, *and* cancellation, so no job is left `resuming`. A record found `resuming` at startup is recovered as `failed` (its result stays readable via `/jobs `) rather than replayed, since the interrupted turn may already have run tools. `mark_resumed()` remains as a deprecated claim-and-complete shim; `JobRecord.resumed` is now a read-only property meaning "delivery reached a terminal state". Pre-existing records with the old boolean migrate on load. `ResumeState`/`ResumeStateError` are exported from `agentic_cli.tools.jobs`. +- **Rate-limited turns are no longer replayed by the CLI harness.** ADK appends the turn's input (the user message, or a resumed `FunctionResponse`) to the session while setting up the invocation — before the first event — so there is no point at which re-running the turn is side-effect free. `MessageProcessor` now invokes the event source exactly once and returns a failed `TurnResult` explaining that the turn was not retried; transient retries stay inside the provider client (ADK `HttpRetryOptions`, Anthropic `retry_max_attempts`). The "Retry in Ns?" dialog is gone. +- **Session APIs are user-scoped.** `session_exists`/`list_sessions`/`delete_session`/`recent_messages`/`load_session`/`save_session` take an optional `user_id` (defaulting to `settings.default_user` only when the caller omits it), and `on_session_end(session=SessionRef(...))` reads the conversation it is given. Identity is the `SessionRef(app_name, user_id, session_id)` triple, exported from `agentic_cli`. Backends without a durable store now leave `supports_sessions` False and the base hooks raise `NotImplementedError` instead of answering `False`/`[]`; `/sessions` says so explicitly. +- **Agent graphs are validated before anything is allocated.** Duplicate names, dangling `sub_agents` references, self-references, delegation cycles, a child with two parents, and **more than one root** now raise `AgentGraphError` (exported from `agentic_cli`) naming the offending agents — before model discovery, service creation, or the session service. Multiple roots were previously accepted and all but one tree was silently unreachable. Agents are built in dependency order, so declaration order no longer changes the hierarchy. Prompt factories are evaluated under the manager's settings and may take zero arguments or exactly one settings argument; other signatures, async factories, and non-string results are rejected by name. +- **Tools declare the services they need.** `@register_tool(..., requires="kb_manager")` replaces the central `BaseWorkflowManager._TOOL_SERVICE_MAP`, so a downstream tool can request any framework-provided service (`service_registry.KNOWN_SERVICE_KEYS`) without editing the framework; unknown or non-constructible keys raise at registration (`user_kb_manager` is not declarable — `kb_manager` creates both scopes). There is no mechanism for registering new service *types*. +- **Model discovery authority is tracked per provider.** A Google outage no longer makes the Claude listing non-authoritative, and an Anthropic outage no longer makes Anthropic's hardcoded fallbacks authoritative. `set_model()` and `validate_settings()` share one rule set (`BaseSettings.check_model()`), so the setter cannot accept a model startup would reject. An unknown model is never silently swapped for another; deprecated aliases still resolve, with a warning. A workflow manager's own model override joins the same all-or-nothing validation pass through the internal `_validate_settings_with_models()`; `validate_settings()` itself is unchanged and still returns `None`. +- **A tool name may only mean one thing, and sharing one is declared.** `ToolRegistry.register()` / `register_tool()` now raise `ValueError` for *any* already-registered name — matching capabilities are not grounds for silently aliasing callables with different docstrings or model-visible schemas. Two new opt-ins replace the guesswork: `declare_tool(name, ...)` declares a tool with **no backend-neutral implementation** (`ToolDefinition.func is None`, plus a new `variants` tuple), and `register_tool(..., variant_of=name)` registers a backend-native implementation of it, sharing the identity and permission contract while keeping its own signature. The ADK and LangGraph `save_plan`/`get_plan`/`save_tasks`/`get_tasks` are now declared once in `tools/_core/state_tools.py` and registered as variants: previously they contested the name and whichever module imported last silently won, so a bare `"save_plan"` in an `AgentConfig` resolved by import order and could hand an ADK agent a LangGraph tool. Such a bare name now raises "ambiguous" instead. `replace=True` still retires the old definition's identity bindings, and service substitution still requires the variant to be the *same* `ToolDefinition`. +- **An empty `ToolRegistry` is no longer silently replaced by the global one.** `ToolRegistry` defines `__len__`, so `registry or get_registry()` in `resolve_tool()` fell through whenever the caller's registry had no tools yet — resolving names it had never registered. +- **`declare_tool()` is exported from `agentic_cli.tools`** alongside `register_tool`, with `register_tool(..., variant_of=...)` for backend-native implementations. Minimal usage: declare the contract once (`declare_tool("save_plan", description=..., capabilities=EXEMPT)`), then register each backend's implementation against it. Variants are ordered by defining module rather than import order, ambiguity errors name them module-qualified, re-declaring with a different description raises, and a `replace=True` excludes the retired backend variants from `include_state_tools` injection so the model never sees a duplicate name. `ToolRegistry.canonical_for()` is new — assembly uses it to substitute a variant's canonical callable instead of a declaration's absent `func`. +- **Backends declare whether a *foreign* job can be cancelled** (`JobBackend.cancels_foreign_jobs`, default False; True for `SubprocessBackend`). Previously inferred from `survives_restart`, which answers a different question: a backend can publish a readable outcome without being remotely controllable. +- **The whole persisted job record is shared state.** Execution ownership stopped a second CLI from mangling a running job but also stopped it from ever learning the job had *finished* — the observer skipped the record, so a completed job stayed RUNNING in its view. Every mutator (`reconcile`, `cancel`, `clean`, both recovery paths) now reloads the durable record under the cross-process lock before acting, an observer picks up the outcome its owner persisted (or reads the sentinel itself when the backend is restart-safe), and **terminal transitions are monotonic**: a stale snapshot can no longer rewrite a recorded success as CANCELLED, and `clean()` cannot delete a job another manager is still running. A job whose backend cannot reach it from here (an in-process handle in another manager) is no longer *reported* cancelled — `cancel()` returns the record unchanged. Lock-unavailable behaviour stays fail-closed: `reconcile`/`clean`/`cancel` skip rather than write. A foreign job is now **polled** (backends publish their outcome durably, so an observer can read it) and only its `UNKNOWN` answer — "I hold no handle for this" — is ignored, so an observer sees a job finish even while its owner is alive but no longer polling. Terminal records are reloaded too, so `get()`/`list()`/`awaiting_resume()` reflect a delivery another manager completed, and `begin_resume()` judges terminality on the durable record rather than its own snapshot. Reading distinguishes **deleted** from **unreadable**: a record another manager cleaned away is forgotten rather than resurrected by a stale `get`/`cancel`/`begin_resume`/`complete_resume`, while an unparseable one is left untouched and fails closed. `clean()` will not delete a job whose delivery is in flight. +- **Job execution ownership is recorded.** `JobRecord.exec_owner` (`::`) is claimed under the cross-process lock *before* a backend is started, so two CLIs sharing the (user-scoped) jobs directory can no longer both launch the same QUEUED job, and a second CLI no longer polls a job whose in-process handle lives in another manager — which answered UNKNOWN and made a healthy job terminal and deliverable. A launch interrupted mid-claim is failed on the next start (`launch interrupted before the job started`), never replayed. +- **Tool identity is owned per `ToolRegistry`.** `bind_tool_identity()`/`identify_tool()` are unchanged as module-level helpers for the framework's default registry, and `ToolRegistry.bind_identity()`/`identify()` are new. A tool registered into an application's own `ToolRegistry` is no longer visible to the framework's identity checks: it passes through tool assembly untouched and is denied by the permission engine. Keeping the map on the instance is also what lets a short-lived registry (with its definitions and their closures) be garbage collected — the previous module-level map held every definition strongly, forever. + +### Removed +- **`skill_scripts_enabled` setting removed.** Turning it on exposed ADK's `run_skill_script` while the supported manager path supplies no code executor, so every call answered `NO_CODE_EXECUTOR`. Script execution is now enabled by passing a `code_executor` to `make_skill_toolset()` — the thing that actually makes it work — and `make_skill_toolset(scripts_enabled=...)` is gone. Skill discovery/read tools and the L1 metadata injection are unaffected. + +### Fixed +- **ADK permission gating no longer trusts a tool's name (P0).** `PermissionPlugin` resolved capabilities via `get_registry().get(tool.name)`, and ADK derives that name from the callable — so an unregistered function named `ask_clarification` inherited the genuine tool's EXEMPT status and ran ungated. Capabilities are now resolved through a registry-owned identity binding, and **only** that: the map is keyed by `id()` and every hit is confirmed with `is` against a weak reference, so a forged `__eq__`/`__hash__` cannot impersonate a registered callable; there is no name fallback, so a custom `BaseTool` named after an EXEMPT tool is denied; `.func` is unwrapped only for the exact ADK types whose contract is to call it (`FunctionTool`, `LongRunningFunctionTool`), so a wrapper merely *exposing* a genuine callable is denied; and MCP detection is `isinstance(tool, McpTool)` rather than a class-name match. Native ADK tool objects the framework builds (skill tools) are bound explicitly at construction. Renamed tools, long-running wrappers, service-bound factory variants, skill tools, and real MCP tools are unaffected. +- **Tool assembly uses registered identity, not `__name__`.** `register(func, name=..., requires=..., long_running=True)` leaves the caller holding a callable whose `__name__` is the private implementation name; an `AgentConfig` listing it never got its declared services created, never got the `LongRunningFunctionTool` wrapper, never picked up its service-bound variant, and exposed the private name to the model. All four now resolve through the registry, and **only** by identity: an application's own callable that happens to share a registered tool's name is no longer given that tool's services, substituted for its service-bound variant, or wrapped as long-running — it stays itself, and is denied at permission time. String tool references and renamed registered callables are unaffected. +- **Job resume metadata is atomic across processes.** Only the claim transitions took the cross-process lock; every *other* metadata write (a poll that found the job finished, a cancel, a launch) rewrote the whole record from memory, including stale resume fields — so a second CLI merely observing a job erased the first one's claim and could then deliver the same result. A plain write now re-reads and preserves the on-disk resume fields under the same lock, startup recovery re-reads inside the lock before deciding (it was acting on a pre-lock snapshot and could rewrite a completed delivery as failed), and a lock that cannot be taken fails the claim and skips recovery instead of proceeding unsynchronised. **Without the lock nothing shared is written at all**: an unlocked read-then-rewrite is the very race the lock prevents, so a persist that cannot lock is skipped (only the *creation* of a record that does not exist yet is safe), and `complete_resume()` leaves the durable record `RESUMING` — recovered as failed later, never re-delivered — rather than risk regressing a `DELIVERED` another process just wrote. +- **Every runtime-effective model is validated and normalized.** A manager's own model — `Manager(model=...)`, `reinitialize(model=...)`, or one cached from a `settings.get_model()` that ran before discovery — bypassed validation entirely: an unusable id reached the provider, and a deprecated one was never replaced. It now joins `default_model` and the per-agent overrides in one pass, and rewrites are applied only after *all* of them validate (they used to be written as each model was checked, leaving the configuration half-rewritten by a call that raised). +- **A failed service constructor releases its predecessors.** `_build_services()` builds into a local dict; when a later constructor raised, that dict was dropped with an already-built SandboxManager or JobManager inside it — never published, so nothing could ever close its pool. +- **Shutdown survives a cancelled caller.** `WorkflowController.close()` did its teardown inline, so cancelling whoever asked for it (Ctrl+C during exit, a cancelled task group) abandoned a manager half-cleaned or a construction still running in the executor — and a later `close()` returned immediately because `_closed` was already set, reporting a shutdown that never happened. The teardown now runs in a single-flight task the controller owns; callers join it under a shield, so a cancelled caller cannot stop it and a later `close()` waits for it. `cancel_init()` is public and may be awaited directly: cancelling it mid-settle used to consume the construction claim and strand the manager (the claim is single-shot, so even the fallback callback could not take over); the claim is now handed back and the callback re-armed. +- **A manager built in the init executor is never orphaned.** The `run_in_executor` future was awaited unshielded, so cancelling initialization cancelled the asyncio future while the (uncancellable) worker thread carried on — and asyncio then discarded the manager it returned. The await is now shielded and the construction tracked: shutdown waits for it and releases it, with a done-callback as the fallback, so it is cleaned exactly once and never published. +- **A recovered controller no longer reports the old failure.** `WorkflowController._init_error` was cleared only at the start of a background init, so after a failed reinitialization that then succeeded the controller was `READY` while `ensure_initialized()` returned False and the status bar read "Init failed - check API keys". Every successful init/reinitialize/swap now clears it, and readiness is derived from the state rather than from the presence of a past error. +- **Cancelling during a HITL prompt no longer strands a thinking box.** The dialog's `finally` unconditionally opened a replacement events box, including while the turn was unwinding, leaving a panel nothing would finish; the box is now reopened only when the prompt actually returns, and every path finishes it exactly once. +- **Background-job resume claims are safe across processes.** `begin_resume()`/`complete_resume()` consulted only in-memory records, so two CLIs sharing the (user-scoped) jobs directory both saw `pending` and delivered the same result into two conversations. The transition now happens under an inter-process file lock, re-reading the record from disk. A claim records its owner (`:`), and startup recovery fails only claims whose owner is *gone* — another live process mid-delivery is left alone instead of having its claim yanked. Only the claiming process may `complete_resume()`. +- **Deprecated model aliases are actually replaced at runtime.** `check_model()` returned the live replacement but only `set_model()` wrote it back, so a `default_model` loaded from `settings.json`/the environment — and every `AgentConfig.model` override — kept the dead id and sent it to the provider. `validate_settings()` now rewrites both in place. +- **The turn boundary is safe end to end.** Cancelling the caller of `MessageProcessor` left the child consumer task driving the workflow while the turn's callback and state were torn down; it is now cancelled and awaited first. The HITL input callback is **context-local** (`set_input_callback()` stores into a per-manager `ContextVar` and returns a token), so a second consumer cannot capture a running turn's prompt and one consumer's `clear_input_callback()` cannot unregister another's. `EventType.ERROR` is now handled: rendered as it arrives, and a non-recoverable one makes the turn `FAILED`/`delivered=False` (previously it was silently dropped and the turn reported `COMPLETED`, so a failed job resume was recorded as delivered). `recoverable=True` is surfaced as a warning and leaves the outcome to the stream. +- **Lifecycle races closed.** A turn admitted after a queued cleanup ran against released resources (a `None` runner); admission now re-checks readiness while holding the turn lock and reinitializes once, or fails cleanly. Service construction runs on a worker thread that cancellation cannot stop, and used to publish into the manager after a rollback; it now builds into a local dict, publishes only while the attempt still owns initialization, and releases anything the thread finished building after a cancellation. `WorkflowController` serializes init/reinitialize/swap/close, so concurrent swaps cannot leak a manager and nothing can publish after `close()`; `controller.workflow` refuses to hand out a manager that is not `READY`. +- **A failed reinitialization no longer discards the conversation.** The controller cleaned up the manager whose own `reinitialize(preserve_sessions=True)` had just restored its session service — closing it, which with `session_store='memory'` took the whole conversation with it. The manager is now retained (state `FAILED`, not handed out) and the next initialization revives it; only if reviving fails is it released and replaced. +- **The documented minimal `AgentConfig` could not construct an ADK agent.** `description` defaulted to `""` but was converted to `None`, which ADK (typed `str`) rejects — so the README quick-start raised `ValidationError`. +- **Workflow lifecycle is transactional.** `initialize_services()` rolls back what a failed attempt allocated; `reinitialize(preserve_sessions=True)` reuses the live session service instead of building a replacement it then discarded unclosed, and preserves it across a failure; a failed in-place reinitialization leaves the controller `FAILED` rather than READY over an uninitialized manager (`WorkflowController.state`, exported as `WorkflowState`). `cleanup()` is idempotent, awaits async `close()` on owned resources, and isolates each closer so one failure cannot skip the rest. `ensure_initialized()` retries a failed attempt. +- **A manager serializes its turns.** `process()`/`resume_with_job_result()` hold a turn lock (released on cancellation), and lifecycle mutation takes it too, so overlapping consumers cannot drain each other's plugin event buffers and cleanup cannot tear the backend down mid-stream. Per-turn session/user identity remains a `ContextVar`, correct under nesting and in spawned tasks. +- **Credentials accept constructor arguments.** `BaseSettings(google_api_key=...)` bound nothing (the field had only an env alias, and `extra="ignore"` swallowed the kwarg); the fields now accept both forms, stay out of `repr()`, and a misspelled credential kwarg raises instead of vanishing. +- **Model validation runs after discovery and covers per-agent overrides**, so a model that exists but predates the static fallback list is no longer rejected at startup, and an `AgentConfig.model` pointing at a provider with no credential fails at startup instead of mid-run. Provider listings run off the event loop and their clients are closed. ### Security -- **Project config can no longer flip security boundaries (P0-1).** A cloned/untrusted repo's `./.{app}/settings.json` (and a cwd-relative `.env`) is now restricted to an explicit deny-by-default allowlist of benign keys (`_PROJECT_SETTABLE_KEYS`) — model/behavior, retry & request timeouts, sandbox *resource* limits, non-exec tool config, `session_store`, and display/logging verbosity. Security-sensitive fields set by a project file — `stateful_executor_backend`, `sandbox_image`/`sandbox_container_user`/`sandbox_data_mounts`/`sandbox_outputs_dir`, the `os_sandbox_*` policy, `skill_scripts_enabled`, `skills_dirs`, `shell_sandbox_type`/`shell_docker_image`, `raw_llm_logging`, `workspace_dir`, permission rules, and secrets — are dropped with a logged warning rather than rejected (the filter drops non-allowlisted keys instead of raising). Real environment variables and the user `~/.{app}/settings.json` remain fully trusted. Previously only `permissions_enabled` was stripped, so a repo could select the host executor, bind mounts, and container image. **Consequence:** put secrets/keys in real environment variables or a user-level file, not in a cwd `.env`. +- **Project config can no longer flip security boundaries (P0-1).** A cloned/untrusted repo's `./.{app}/settings.json` (and a cwd-relative `.env`) is now restricted to an explicit deny-by-default allowlist of benign keys (`_PROJECT_SETTABLE_KEYS`) — model/behavior, retry & request timeouts, sandbox *resource* limits, non-exec tool config, `session_store`, and display/logging verbosity. Security-sensitive fields set by a project file — `stateful_executor_backend`, `sandbox_image`/`sandbox_container_user`/`sandbox_data_mounts`/`sandbox_outputs_dir`, the `os_sandbox_*` policy, `skills_dirs`, `shell_sandbox_type`/`shell_docker_image`, `raw_llm_logging`, `workspace_dir`, permission rules, and secrets — are dropped with a logged warning rather than rejected (the filter drops non-allowlisted keys instead of raising). Real environment variables and the user `~/.{app}/settings.json` remain fully trusted. Previously only `permissions_enabled` was stripped, so a repo could select the host executor, bind mounts, and container image. **Consequence:** put secrets/keys in real environment variables or a user-level file, not in a cwd `.env`. - **Interactive "Allow always" grants moved out of the repo (P0-1).** Persistent permission grants now live in `~/.{app}/project_grants.json`, keyed by the resolved project path, instead of `./.{app}/permissions.local.json` (which a repo could force-track and ship as trusted allow-rules). A clone at a different path carries no grants (re-grant on first use); a repo-shipped `permissions.local.json` is no longer loaded. **No migration** — existing local grant files are ignored; re-grant when prompted. ### Removed diff --git a/CLAUDE.md b/CLAUDE.md index fbc8268..3525f54 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -159,7 +159,15 @@ Workflow: - **Tool registration**: Use `@register_tool(category=..., capabilities=..., description=...)` decorator. `capabilities=` is required — pass `EXEMPT` for tools that need no permission check or a list of `Capability(name, target_arg=...)` tuples the engine matches against rules. Tools are auto-discovered via the global `ToolRegistry`. - **Permissions**: `workflow/permissions/` holds a framework-independent engine that evaluates declared capabilities against rules from four sources (builtin, user `~/.{app_name}/settings.json`, project `./.{app_name}/settings.json`, in-memory session). ADK + LangGraph gate tool calls via `workflow/adk/permission_plugin.py::PermissionPlugin` and `workflow/langgraph/permission_wrap.py::wrap_tool_for_permission`. - **Service registry**: Tools access services and shared state via `get_service(key)` from `workflow.service_registry`. A single ContextVar holds a `dict[str, Any]` set by the workflow manager during processing. Complex services (KBManager, SandboxManager, MemoryStore) are lazily created; simple state (plan string, task list) lives directly in the registry dict. -- **Manager detection**: `BaseWorkflowManager._detect_required_managers()` scans each agent's tool names against the `_TOOL_SERVICE_MAP` (name → service key, in `base_manager.py`); `_ensure_managers_initialized()` then lazily instantiates only the services actually needed (KBManager, SandboxManager, MemoryStore, …). Adding a new service-backed tool means adding its name → service entry to `_TOOL_SERVICE_MAP`. (There is no `@requires` decorator.) +- **Manager detection**: tools declare their own service needs — `@register_tool(..., requires="kb_manager")` (or a tuple for several) — and the key is validated against `service_registry.KNOWN_SERVICE_KEYS` at registration (only *constructible* services are declarable; `user_kb_manager` is created together with `kb_manager`). `BaseWorkflowManager._detect_required_managers()` reads that metadata off the registry for each agent's tools (by registry identity, so `register(func, name=...)`'s original callable still declares its services); `_build_services()` then lazily constructs only the services actually needed, into a local dict that is released in full if a later constructor raises (nothing is published, so nothing else could close it). A downstream tool may request any framework-provided service without editing the framework; there is no mechanism for registering new service *types*, and no central name→service map. +- **Canonical tool names + permission identity**: `ToolDefinition.name` is the single identity. `register_tool(name="public_name")` wraps the callable so `func.__name__` is the registered name (backends derive the model-visible name from the callable). Permission gating **and tool assembly** resolve through a **registry-owned identity binding**: each `ToolRegistry` owns a `id(obj) → (weakref, definition)` map (`registry.bind_identity()`/`identify()`; the module-level `tools.registry.bind_tool_identity()`/`identify_tool()` answer for `get_registry()`). Every hit is confirmed with `is` against the weak reference, so neither a name, a class name, nor a forged `__eq__`/`__hash__` can stand in for it, and a recycled address inherits nothing. Identity is **per registry** — a tool registered into an application's own `ToolRegistry` is not one of the framework's, so it stays untouched during assembly and is denied at permission time — which is also what lets a short-lived registry, its definitions and their closures be garbage collected. Everything the framework issues is bound at its construction site: registered callables, factory service-bound variants, renamed wrappers, and the native ADK tool objects the framework builds (skill tools, in `tools/skills/toolset.py`). `workflow/adk/permission_plugin.py` unwraps `.func` only for the exact ADK types whose contract is to call it (`FunctionTool`, `LongRunningFunctionTool`), and gates genuine `McpTool` instances (`isinstance`, since ADK creates them on connect) under a synthetic `mcp` capability. Anything unbound is denied — and, in assembly (service detection, service-tool substitution, canonicalization, long-running wrapping), left exactly as the application supplied it: a plain callable named `kb_search` is not the framework's tool and must not be given its services, its service-bound variant, or its long-running contract. Substitution additionally requires the service variant to *be* that same definition (`identify_tool(variant) is definition`); factories bind each closure to the exact module-level tool it re-binds, so a tool an application has taken over keeps its own implementation. +- **One name, one tool**: `ToolRegistry.register()` raises on any name that is already registered — matching capabilities are *not* grounds for sharing one, since they say nothing about the docstring or the model-visible schema. Sharing is declared, never inferred: `declare_tool(name, ...)` (exported from `agentic_cli.tools`) declares a tool that has **no backend-neutral implementation** (`ToolDefinition.func is None`), and each backend registers its own with `register_tool(..., variant_of=name)` — same identity and permission contract, its own signature and docstring. Re-declaring the same contract is idempotent; changing its description or capabilities raises. `definition.variants` is ordered by defining module, so it never depends on import order, and assembly substitutes a variant's canonical callable via `registry.canonical_for()` (never `None`). A `replace=True` retires the previous definition's callables, and retired backend variants are excluded from `include_state_tools` injection, so the model never sees two tools with one name. That is how the ADK and LangGraph `save_plan`/`get_plan`/`save_tasks`/`get_tasks` coexist (declared in `tools/_core/state_tools.py`); previously they contested the name and import order decided the winner. A bare-name reference to a declared-only tool raises "ambiguous" rather than guessing a backend. `replace=True` takes a name over deliberately and **retires the old definition's identities**, so its callables resolve to nothing (denied, and left alone by assembly) rather than inheriting the replacement's capabilities. +- **Turn/lifecycle concurrency**: a manager runs one turn at a time — `process()`/`resume_with_job_result()` enter through `_turn_admission()` (which holds `_turn_lock`), and `initialize_services`/`reinitialize`/`cleanup` hold `_lifecycle_lock` **and** `_turn_lock`. Lock order is lifecycle → turn; a turn initializes *before* taking the turn lock, which is what keeps the two from deadlocking — and because of that a cleanup can land in between, so admission re-checks `_backend_ready()` while holding the turn lock and reinitializes once (or fails cleanly) rather than running against released resources. Initialization is transactional: services are built on a worker thread into a *local* dict and published only while the attempt still owns init (a cancelled attempt releases what the thread went on to build), and a failed attempt rolls back. A failed in-place reinitialization leaves the manager uninitialized; the controller reports `FAILED` and refuses to hand it out, but **keeps** it so a retry can revive it with its preserved (possibly in-memory) sessions intact. `WorkflowController` serializes init/reinitialize/swap/close on its own lifecycle lock, and a background init that finishes after `close()` releases its manager instead of publishing it. +- **HITL callback is context-local**: `set_input_callback()` stores into a per-manager `ContextVar`, so a second consumer installing its callback cannot capture a running turn's prompt, and one consumer's `clear_input_callback()` cannot unregister another's. `MessageProcessor` cancels and awaits its consumer task **before** clearing the callback, so no tool is left asking a question nobody owns. +- **No harness-level turn replay**: ADK persists a turn's input during invocation setup, so the CLI never re-invokes an event source. Retries belong to the provider client (`HttpRetryOptions`). `MessageProcessor` returns a typed `TurnResult`. +- **Session identity**: durable sessions are addressed by `SessionRef(app_name, user_id, session_id)` (`workflow/sessions.py`). Every session hook (`session_exists`/`list_sessions`/`delete_session`/`recent_messages`/`load_session`) takes an optional `user_id`, defaulting to `settings.default_user` only when the caller omits it. Backends without a session store leave `supports_sessions` False, and the base hooks raise `NotImplementedError` rather than answering with a misleading `False`/`[]`. +- **Active turn**: the in-flight `(user, session)` is a `ContextVar` (`workflow/sessions.py::get_active_turn`), set with a token by `_workflow_context()` — concurrent turns on one manager stay isolated and nesting restores the outer turn. +- **Resource ownership**: `cleanup()` is idempotent and awaits an async `close()` on owned resources (`BaseWorkflowManager._aclose_owned`); `WorkflowController.close()` is the single shutdown path (cancel init → shut executor → clean manager). - **Atomic writes**: Use `atomic_write_json`/`atomic_write_text` from `file_utils.py` for file persistence. ### Console Output diff --git a/README.md b/README.md index 5358101..80364aa 100644 --- a/README.md +++ b/README.md @@ -281,6 +281,11 @@ dynamic_agent = AgentConfig( tools=[tool_a, tool_b], ) +# A prompt factory may also take the manager's settings explicitly. Either way +# it is evaluated under the manager's settings, not the global singleton. +def get_scoped_prompt(settings): + return f"You are {settings.app_name}." + # Coordinator with sub-agents coordinator = AgentConfig( name="coordinator", @@ -310,13 +315,62 @@ configs = [coordinator, researcher, analyst] | Field | Type | Description | |-------|------|-------------| -| `name` | str | Unique identifier | -| `prompt` | str \| Callable | System instruction | +| `name` | str | Unique identifier (must be unique across the config list) | +| `prompt` | str \| Callable | System instruction; a callable may take no arguments or a single `settings` argument | | `tools` | list[Callable] | Available tool functions | | `sub_agents` | list[str] | Names of agents this one can delegate to | -| `description` | str | Short description for routing | +| `description` | str | Short description for routing (defaults to `""`) | | `model` | str \| None | Model override (defaults to manager's model) | +The agent graph is validated before anything is allocated — ahead of model +discovery and service creation. Duplicate names, unknown `sub_agents` +references, self-references, delegation cycles, a sub-agent shared by two +parents, and **more than one root** raise `AgentGraphError` naming the +offending agents. Exactly one agent may be unreferenced: the runner starts from +a single root, so agents under any other root would never run. Agents are built +in dependency order, so declaration order does not matter. A per-agent `model` +override is validated against the configured provider credentials at startup, +alongside `default_model`. + +A callable `prompt` may take **no arguments** (including all-defaulted ones) or +**exactly one argument**, which receives the manager's settings. Other +signatures, `async def` factories, and non-string results raise +`AgentGraphError` naming the agent. + +### Turn and lifecycle contracts + +A workflow manager runs **one turn at a time**: `process()` and +`resume_with_job_result()` serialize on a turn lock, and +`initialize_services()`/`reinitialize()`/`cleanup()` take it too, so the backend +is never torn down mid-stream. Admission also re-verifies the backend is live +after acquiring the lock, so a turn queued behind a cleanup reinitializes (or +fails cleanly) instead of running against released resources. Run separate +managers for genuine parallelism. + +`MessageProcessor.process()` returns a `TurnResult` (`TurnStatus.COMPLETED` / +`CANCELLED` / `FAILED` / `UNAVAILABLE`); a turn is never replayed by the +harness, so a surfaced rate limit fails it explicitly. An `EventType.ERROR` +event is rendered as it arrives — `recoverable=True` is a warning and the +stream still decides the outcome, anything else makes the turn `FAILED` +(`delivered` False). Cancelling the caller cancels and awaits the event +consumer before the turn is torn down. + +The HITL input callback is **context-local**: `set_input_callback()` binds it +for the calling context (and any task started from it), so two consumers of one +manager cannot capture each other's prompts. + +`WorkflowController` exposes a derived `WorkflowState` +(`uninitialized`/`initializing`/`ready`/`failed`/`closed`) that never reports +`ready` over an uninitialized manager, and `controller.workflow` raises unless +the state is `ready`. Lifecycle transitions (init, reinitialize, orchestrator +swap, close) are serialized, so nothing is published after `close()`. A failed +in-place reinitialization keeps the manager — it still owns the session service +it preserved — and the next initialization revives it rather than discarding +the conversation. + +`SessionRef(app_name, user_id, session_id)` is the conversation identity used by +every session API. All of these are importable from `agentic_cli`. + ## Tools ### Creating Tools diff --git a/src/agentic_cli/__init__.py b/src/agentic_cli/__init__.py index 558523b..ef9bd15 100644 --- a/src/agentic_cli/__init__.py +++ b/src/agentic_cli/__init__.py @@ -25,9 +25,12 @@ from agentic_cli.cli.app import BaseCLIApp from agentic_cli.workflow.factory import create_workflow_manager_from_settings from agentic_cli.cli.commands import Command, CommandRegistry -from agentic_cli.workflow.config import AgentConfig +from agentic_cli.cli.message_processor import TurnResult, TurnStatus +from agentic_cli.cli.workflow_controller import WorkflowState +from agentic_cli.workflow.config import AgentConfig, AgentGraphError from agentic_cli.workflow.model_settings import ModelSettings, ThinkingSettings from agentic_cli.workflow.events import WorkflowEvent, EventType +from agentic_cli.workflow.sessions import SessionRef from agentic_cli.config import ( BaseSettings, SettingsContext, @@ -74,10 +77,16 @@ def __getattr__(name: str): "GoogleADKWorkflowManager", # lazy (Google ADK) "LangGraphWorkflowManager", # lazy (requires langgraph extra) "AgentConfig", + "AgentGraphError", "ModelSettings", "ThinkingSettings", "WorkflowEvent", "EventType", + # Lifecycle / turn contracts + "SessionRef", + "TurnResult", + "TurnStatus", + "WorkflowState", # Settings "BaseSettings", "SettingsContext", diff --git a/tests/test_import_surface.py b/tests/test_import_surface.py new file mode 100644 index 0000000..b74cd16 --- /dev/null +++ b/tests/test_import_surface.py @@ -0,0 +1,114 @@ +"""The framework-facing contracts are importable from their natural packages. + +These types are what an embedding application programs against — a turn's +outcome, a controller's state, a conversation's identity, a job's delivery +state, and the graph error it must handle at startup. Each was reachable only +from a private module path. +""" + +from __future__ import annotations + +import pytest + + +class TestTopLevelExports: + @pytest.mark.parametrize( + "name", + ["SessionRef", "AgentGraphError", "TurnResult", "TurnStatus", "WorkflowState"], + ) + def test_exported_from_package_root(self, name: str): + import agentic_cli + + assert hasattr(agentic_cli, name), f"agentic_cli.{name} is not importable" + assert name in agentic_cli.__all__, f"{name} is missing from __all__" + + def test_root_exports_are_the_defining_objects(self): + """No shadow copies: the export is the class the framework uses.""" + import agentic_cli + from agentic_cli.cli.message_processor import TurnResult, TurnStatus + from agentic_cli.cli.workflow_controller import WorkflowState + from agentic_cli.workflow.config import AgentGraphError + from agentic_cli.workflow.sessions import SessionRef + + assert agentic_cli.SessionRef is SessionRef + assert agentic_cli.AgentGraphError is AgentGraphError + assert agentic_cli.TurnResult is TurnResult + assert agentic_cli.TurnStatus is TurnStatus + assert agentic_cli.WorkflowState is WorkflowState + + +class TestToolsExports: + """``declare_tool`` is how an application declares a tool it implements + per-backend; it sits next to ``register_tool`` in the tools package.""" + + def test_declare_tool_is_exported(self): + from agentic_cli import tools + + assert hasattr(tools, "declare_tool") + assert "declare_tool" in tools.__all__ + + def test_declare_tool_is_the_defining_object(self): + from agentic_cli import tools + from agentic_cli.tools.registry import declare_tool + + assert tools.declare_tool is declare_tool + + def test_minimal_usage(self): + """Declare a contract, register a backend variant against it.""" + from agentic_cli.tools import ToolCategory, declare_tool, register_tool + from agentic_cli.tools.registry import ToolRegistry + from agentic_cli.workflow.permissions import EXEMPT + + registry = ToolRegistry() + declare_tool( + "doc_probe_tool", + description="A tool each backend implements natively.", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + registry=registry, + ) + + def _native(content: str) -> dict: + """Backend-native implementation.""" + return {"success": True} + + returned = registry.register( + _native, + variant_of="doc_probe_tool", + capabilities=EXEMPT, + category=ToolCategory.PLANNING, + ) + + assert returned.__name__ == "doc_probe_tool" + assert registry.identify(_native) is registry.get("doc_probe_tool") + assert register_tool is not None # exported alongside + + +class TestJobsExports: + def test_resume_lifecycle_types_are_exported(self): + from agentic_cli.tools import jobs + + assert hasattr(jobs, "ResumeState") + assert hasattr(jobs, "ResumeStateError") + assert "ResumeState" in jobs.__all__ + assert "ResumeStateError" in jobs.__all__ + + def test_resume_state_values(self): + from agentic_cli.tools.jobs import ResumeState + + assert [s.value for s in ResumeState] == [ + "pending", + "resuming", + "delivered", + "failed", + ] + + +class TestSessionRefShape: + def test_is_a_frozen_triple(self): + from agentic_cli import SessionRef + + ref = SessionRef(app_name="app", user_id="u", session_id="s") + assert (ref.app_name, ref.user_id, ref.session_id) == ("app", "u", "s") + with pytest.raises(Exception): + ref.user_id = "other" # frozen