diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..47d0a3c --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,98 @@ +name: CI + +on: + pull_request: + push: + branches: [develop, main] + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + offline: + name: Offline tests (no Docker) + runs-on: ubuntu-latest + defaults: + run: + shell: bash -el {0} + steps: + - uses: actions/checkout@v4 + + - name: Set up conda env (agenticcli) + uses: conda-incubator/setup-miniconda@v3 + with: + miniforge-version: latest + environment-file: environment.yml + activate-environment: agenticcli + + - name: Install langgraph extra (langgraph-backend tests) + run: conda run -n agenticcli pip install -e '.[langgraph]' + + - name: Run offline test suite + run: conda run -n agenticcli python -m pytest -m 'not llm and not docker' -q + + wheel-acceptance: + name: Built-wheel acceptance (console smoke) + runs-on: ubuntu-latest + defaults: + run: + shell: bash -el {0} + steps: + - uses: actions/checkout@v4 + + - name: Set up conda env (agenticcli) + uses: conda-incubator/setup-miniconda@v3 + with: + miniforge-version: latest + environment-file: environment.yml + activate-environment: agenticcli + + # pexpect drives the console over a pty. Installed explicitly via the dev + # extra rather than relied on transitively. + - name: Install dev extra (pexpect) + run: conda run -n agenticcli pip install -e '.[dev]' + + # Its own step rather than part of the offline suite: it builds a wheel + # and installs that wheel's dependencies into a throwaway virtualenv, so + # it is slow and needs the network. The opt-in env var is what the tests + # gate on, so the offline selector stays fast and offline. + # + # Targets the file rather than `-m wheel` alone: a bare marker selection + # still *collects* the whole suite, which would need the langgraph extra + # installed only to collect tests this job does not run. + - name: Run built-wheel acceptance + env: + AGENTIC_WHEEL_ACCEPTANCE: "1" + run: > + conda run -n agenticcli python -m pytest + tests/examples/test_research_demo_wheel.py -m wheel -v + + docker-isolation: + name: Docker isolation tests + runs-on: ubuntu-latest + defaults: + run: + shell: bash -el {0} + steps: + - uses: actions/checkout@v4 + + - name: Set up conda env (agenticcli) + uses: conda-incubator/setup-miniconda@v3 + with: + miniforge-version: latest + environment-file: environment.yml + activate-environment: agenticcli + + - name: Install langgraph extra (for clean full test collection) + run: conda run -n agenticcli pip install -e '.[langgraph]' + + # Keep this tag in sync with settings.sandbox_image's default. Pre-pulling + # avoids first-run pull latency exceeding sandbox_start_timeout. + - name: Pre-pull sandbox image + run: docker pull quay.io/jupyter/scipy-notebook:python-3.12 + + - name: Run docker isolation tests (fail if runtime missing) + env: + SANDBOX_REQUIRE_DOCKER: "1" + run: conda run -n agenticcli python -m pytest -m docker -v diff --git a/CHANGELOG.md b/CHANGELOG.md index 45947a2..b5a9948 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,72 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.6.0] - 2026-08-04 + +Long-running jobs, durable sessions, a Docker sandbox backend, five P0 security +fixes, and a framework-wide correctness pass. **Breaking** — see *Changed* and +*Removed*; the tool-identity, session and turn-result contracts changed. + +### Added +- **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`/`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. the session was deleted, or the run used `session_store='memory'` and the CLI restarted (with the default durable store the session survives a restart) — 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`). + +- **Release acceptance coverage for the shipped CLI.** The composed application is now exercised end to end: a headless suite drives the real `BaseCLIApp.process_input` and `MessageProcessor` against scripted `WorkflowEvent` streams, and a terminal smoke drives the actual console process through startup → `/help` → `/exit` over a pty, requiring no API key, network or LLM. A separate opt-in CI job builds the wheel, installs it into a throwaway virtualenv and runs the same session through the installed `research-demo` console script, checking that the packaged data files ship and that the expected version reaches the wheel name, the distribution metadata and `agentic_cli.__version__`. + +### Changed +- **The persistent permission choice now describes its scope, not a file.** The prompt's third option reads `Allow always for this project` instead of `Allow always (save to project)`, which implied the grant is written into the repository. It never was: interactive grants persist to `~/.{app_name}/project_grants.json`, keyed by the resolved project path. Only the displayed label and the `ALLOW_ALWAYS_CHOICE` constant changed — authorization, grant scope, where grants are stored, rule matching and fail-closed behaviour are unchanged, and the former response wording is still accepted so an answer captured under it keeps meaning "always". +- **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 +- **Legacy JSON session snapshots removed** (`persistence/SessionPersistence`/`SessionSnapshot`, the `_extract_session_data`/`_inject_session_messages` hooks, and the on-exit-save / on-startup-inject path). It saved only on a clean exit (a crash lost the session) and rebuilt tool calls/responses **without their ids** (and dropped thinking), losing fidelity on resume. Superseded by the native durable stores above, which persist continuously with full fidelity. There is no migration of old JSON sessions. +- **`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 multi-agent delegation works again.** Two defects made the built-in agent hand-off unusable, so any app declaring `sub_agents` could not delegate at all. (1) `PermissionPlugin` unwraps `.func` only for *exact* ADK function-tool types, and ADK's auto-injected `TransferToAgentTool` is a `FunctionTool` **subclass** — so the routing tool resolved to no registry identity and was denied as unregistered, even with permissions disabled. Its exact class is now trusted on the same terms as the others (ADK builds it as `FunctionTool(func=transfer_to_agent)` and overrides only `_get_declaration()`, never `run_async`); every other subclass, forged object and same-named tool stays denied, and no name-based authority is restored — the gate remains fail-closed. (2) ADK builds the tool's *model-visible* description from `transfer_to_agent`'s docstring, which through 1.37.0 advises callers to "use TransferToAgentTool instead of this function directly" — guidance meant for Python callers that led models to emit the class name, which ADK rejects. A narrowly scoped before-model plugin corrects that description on the prepared request: only for the exact native tool, only while the misleading text is present, writing `description` alone (name, schema, required fields and the agent-name enum are preserved, and the upstream `__doc__` is never mutated). It is idempotent and becomes a no-op on any ADK release that ships a corrected docstring (fixed upstream in 2.x). +- **`/memory` and `/kb-backfill` no longer error while the workflow is starting.** `app.workflow` raises until the controller reports READY, so both research-demo commands surfaced the generic "Error executing command: Workflow not initialized yet" during background initialization. They now show a warning that initialization is still in progress, matching the built-in commands' behaviour. +- **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 +- **`web_fetch` no longer reaches internal networks (P0-5).** A new `PinnedTransport` resolves every A/AAAA record up front, rejects the request if *any* address is non-global (including NAT64 and 6to4-relay forms), then connects to the validated IP with `sni_hostname` set so TLS still binds to the hostname — closing the DNS-rebinding window between check and connect. Redirects and `robots.txt` go through the same transport, and response bodies are capped while streaming. +- **Sandbox file transfer is symlink-safe (P0-2).** Input staging and output collection use `copy_regular_file_no_follow()` and refuse to follow symlinks; the outputs directory is created `0700`. A planted symlink in the inputs or outputs directory could previously redirect a read or a write outside the sandbox. Note that `O_NOFOLLOW` guards only the final path component, so the parent directory is `lstat`ed as well. +- **`compile_document` host execution is contained (P0-3).** LaTeX builds run in a private temp directory with an exact `TEXMF` env allowlist, an anchored source filename, `RLIMIT_FSIZE`/`RLIMIT_CPU` caps on the child, bounded captured output, and a tail-read of the compiler log. Timeouts are enforced with SIGKILL. +- **`glob`/`grep` are contained to the authorized root (P0-4).** Patterns can no longer escape it, scan and result ceilings are applied before the walk rather than after, directories no longer consume the file budget, and truncation is reported instead of silently returning a partial answer. The ripgrep path bounds its output to a temp file with a capped read. +- **The permission engine fails closed (P1-7).** A missing engine denied nothing; it now denies. +- **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. + ## [0.5.3] - 2026-06-14 ### Added diff --git a/CLAUDE.md b/CLAUDE.md index e66849a..c67f21c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ agentic-cli/ │ ├── __init__.py # Package exports, lazy imports │ ├── config.py # BaseSettings (pydantic-settings) │ ├── settings_mixins.py # Composable settings field groups -│ ├── settings_persistence.py # save_settings() (excludes SECRET_FIELDS) +│ ├── settings_persistence.py # Trust-split save (PROJECT_SETTABLE_KEYS → project, rest → user config; excludes SECRET_FIELDS) │ ├── constants.py # Shared constants, truncate() │ ├── file_utils.py # atomic_write_json / atomic_write_text │ ├── logging.py @@ -157,9 +157,18 @@ Workflow: ### Key Design Patterns - **Tool error handling**: All tools return `{"success": bool, ...}` dicts. Never raise `ToolError`. - **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`. +- **ADK transfer-tool description**: ADK builds `transfer_to_agent`'s model-visible description from its docstring, which through 1.37.0 tells the model to "use TransferToAgentTool instead of this function directly" — advice meant for Python callers that made Gemini emit the class name, which ADK then rejects. `workflow/adk/transfer_tool_description.py` is a `before_model_callback` plugin that rewrites that description on the prepared `LlmRequest`: only when the tool is the **exact** native `TransferToAgentTool` (an application tool sharing the name is untouched) and only while the misleading text is present, so it is idempotent and becomes a **no-op** once the installed ADK ships a corrected docstring (upstream fixed it in 2.x). Only `description` is written — declaration name, parameter schema, required fields and the agent-name enum are preserved, and the upstream function's `__doc__` is never mutated. +- **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). Interactive "Allow always" grants persist to `~/.{app_name}/project_grants.json` keyed by resolved project path — user-owned, never into the project settings file a repo could ship. 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 ADK's own `TransferToAgentTool`), and gates genuine `McpTool` instances (`isinstance`, since ADK creates them on connect) under a synthetic `mcp` capability. `TransferToAgentTool` is a `FunctionTool` **subclass** ADK auto-injects into any agent with `sub_agents`; matching only exact `FunctionTool` denied the built-in routing tool as unregistered, so delegation could not work at all. It is trusted on exactly the same terms as the others — ADK constructs it as `super().__init__(func=transfer_to_agent)` and overrides only `_get_declaration()` (to add the agent-name enum), never `run_async`, so what it invokes is still exactly `self.func`. Listing the **exact class** keeps every other subclass, forged object and same-named tool denied; no name authority is restored. 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 @@ -181,3 +190,26 @@ Available session methods: - **MockVectorStore** and **MockEmbeddingService**: In `knowledge_base/_mocks.py` for testing without ML dependencies - **FAISS tests**: Guard with `pytest.importorskip("faiss")` since FAISS is not installed in dev env - **Integration tests**: `tests/integration/` covers ADK and LangGraph pipeline tests + +### Live LLM tests (real API calls) + +Tests that hit real provider APIs use the existing framework — **don't invent new gating or key handling.** + +- **Marker**: `@pytest.mark.llm`; modules set `pytestmark = [pytest.mark.llm, pytest.mark.skipif(, ...)]`. +- **Key loading is handled by the live-test framework** (`tests/integration/conftest.py`) — keys are not + plain shell env vars, so go through pytest rather than re-deriving it. +- **Run**: `-m llm` (live; needs network — disable the Bash sandbox) or `-m 'not llm'` (offline). + A bare `pytest` run makes real API calls when keys are available. +- **Example**: `tests/integration/test_adk_claude_live.py`. + +### Live docker sandbox tests (real container runtime) + +The `jupyter_docker` backend's isolation boundary is verified against a real daemon. + +- **Marker**: `@pytest.mark.docker`; skipped unless docker/podman is available. The bulk of the + backend is tested offline via the faked `ContainerRuntime` seam (incl. an end-to-end run of the + real `driver.py` as a subprocess) — these live tests only cover what needs a real container. +- **Run**: `-m docker` (needs a container runtime) or `-m 'not llm and not docker'` (offline CI). +- **Fail-loud in CI**: set `SANDBOX_REQUIRE_DOCKER=1` so a missing/broken runtime FAILS instead of + silently skipping (a skip reads as green). CI: `.github/workflows/ci.yml` (offline + docker jobs). +- **Example**: `tests/tools/test_sandbox_docker_live.py`. diff --git a/README.md b/README.md index 00bdbde..30858a1 100644 --- a/README.md +++ b/README.md @@ -6,15 +6,14 @@ A framework for building domain-specific agentic CLI applications powered by LLM Agentic CLI provides the core infrastructure for building interactive CLI applications that leverage LLM agents for complex tasks. It offers: -- **Pluggable Orchestration**: Choose between Google ADK or LangGraph for agent workflows +- **Pluggable Orchestration**: Choose between Google ADK (which runs both Gemini and Anthropic models) or LangGraph for agent workflows - **Rich Terminal UI**: Dual thinking boxes, markdown rendering, and streaming responses via `thinking-prompt` - **Declarative Agents**: Define agents with simple configuration objects - **Native Tool Architecture**: Backend-specific tool factories for ADK and LangGraph, with automatic HITL confirmation for dangerous tools - **Built-in Tools**: Python execution, stateful sandbox, file operations, web search, web fetch, arXiv search - **Knowledge Base**: Semantic + BM25 hybrid search with RRF fusion, per-document markdown sidecars, and agent-authored concept pages - **Semantic Memory**: Embedding-backed memory with lifecycle management, contradiction detection, and forgetting policy -- **Tool Reflection**: Bounded per-tool heuristic memory learned from failures -- **Session Save/Resume**: Persistent conversations across CLI restarts +- **Durable Sessions**: Conversations persist continuously by default (SQLite; PostgreSQL or ephemeral memory optional) and resume across CLI restarts with `--session ` - **Context Window Management**: Native trim detection and token-usage visibility - **Dynamic Model Registry**: Live model discovery from provider APIs - **Type-safe Configuration**: Composable settings mixins (`pydantic-settings`) @@ -57,7 +56,7 @@ Agentic CLI provides the core infrastructure for building interactive CLI applic ## Installation -### Basic Installation (Google ADK) +### Basic Installation (Google ADK — Gemini and Claude) ```bash pip install agentic-cli @@ -186,9 +185,9 @@ Requires: `pip install agentic-cli[langgraph]` |---------|------------|-----------| | Setup complexity | Simple | Moderate | | Cyclical workflows | Limited | Native | -| Multi-provider | Google only | OpenAI, Anthropic, Google (GenAI) | -| State persistence | In-memory | Memory, PostgreSQL, or SQLite | -| Thinking support | Native (Gemini) | Native (Claude & Gemini) | +| Multi-provider | Gemini + Anthropic (Claude runs natively via `DirectAnthropicLlm`) | OpenAI, Anthropic, Google (GenAI) | +| State persistence | SQLite (default), PostgreSQL, or memory | Memory, PostgreSQL, or SQLite | +| Thinking support | Native (Gemini & Claude) | Native (Claude & Gemini) | | Retry handling | Built-in | Built-in with backoff | | Permission gate | PermissionPlugin | wrap_tool_for_permission | | Context trimming | Native | Native | @@ -208,7 +207,7 @@ manager = create_workflow_manager_from_settings(agent_configs=AGENTS, settings=s Settings are organized into composable mixins (`AppSettingsMixin`, `CLISettingsMixin`, `WorkflowSettingsMixin`). `BaseSettings` composes all three. -All settings can be configured via environment variables with the `AGENTIC_` prefix or in a `.env` file: +All settings can be configured via environment variables with the `AGENTIC_` prefix, via constructor arguments, or from layered JSON/dotenv files: ```python from pathlib import Path @@ -217,13 +216,37 @@ from agentic_cli import BaseSettings class MySettings(BaseSettings): model_config = SettingsConfigDict( - env_file=".env", + # Absolute (user-level) paths stay trusted — see the trust note below. + env_file=str(Path.home() / ".my_app" / ".env"), env_prefix="MYAPP_", # Custom prefix ) app_name: str = "my_app" workspace_dir: Path = Path.home() / ".my_app" ``` +### Precedence and the trust boundary + +Sources are consulted highest-precedence first (JSON sources only when the file exists): + +1. **Constructor arguments** — `MySettings(google_api_key=...)` +2. **Environment variables** — `MYAPP_*` / `AGENTIC_*` +3. **Project config** — `./.{app_name}/settings.json` — **untrusted** +4. **User config** — `~/.{app_name}/settings.json` +5. **Dotenv** — whatever `env_file` points at +6. **Field defaults** + +Levels 1, 2 and 4 are trusted; **level 3 is not**, and neither is a +*cwd-relative* `env_file`. A cloned repository can ship both, so both are +filtered down to an explicit allowlist of benign keys (model and behaviour +settings, timeouts, sandbox *resource* limits, `session_store`, log verbosity). +Security-relevant keys set from an untrusted source — executor backend, sandbox +image/mounts/user, OS-sandbox policy, `skills_dirs`, shell sandbox settings, +`workspace_dir`, `raw_llm_logging`, permission rules, and any credential — are +dropped with a logged warning rather than applied. + +**Consequence:** put API keys in real environment variables or an +absolute/user-level file. A key placed in a cwd `.env` is deliberately ignored. + ### Key Settings | Setting | Env Variable | Default | Description | @@ -281,6 +304,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 +338,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 @@ -348,6 +425,45 @@ def search_database(query: str, limit: int = 10) -> dict: **Registration is required for a tool to run.** The permission engine is on by default, and any tool without a capability declaration is denied at call time (fail-closed) on both the ADK and LangGraph backends. You can still pass raw callables into `AgentConfig.tools`, but unless they are registered with `@register_tool` they will be blocked — register every tool with `capabilities=` (or `EXEMPT` for pure, side-effect-free functions). Registration also gives the tool registry metadata and tool-summary formatting. +Identity is the *object the registry issued*, not the name: a callable the +registry never issued is denied however it is named, and is left untouched by +tool assembly. + +**A tool may declare the services it needs** with `requires=` — e.g. +`@register_tool(..., requires="kb_manager")`, or a tuple for several. The +manager builds only the services its agents' tools actually declare. + +**One name means one tool.** `register_tool()` raises on a name that is already +registered. If a tool genuinely needs a per-backend implementation, declare the +shared contract once and register each backend's version against it: + +```python +from agentic_cli.tools import declare_tool, register_tool, ToolCategory +from agentic_cli.workflow.permissions import EXEMPT + +# Declare the shared contract once, under a name your application owns. +declare_tool( + "myapp_scratchpad", + description="Store a short note for the rest of this turn.", + capabilities=EXEMPT, + category=ToolCategory.OTHER, +) + +# Then register one implementation per backend against that contract. +@register_tool( + variant_of="myapp_scratchpad", + capabilities=EXEMPT, + category=ToolCategory.OTHER, +) +def myapp_scratchpad_adk(note: str) -> dict: + """Store a short note for the rest of this turn.""" + return {"success": True, "note": note} +``` + +Variants share one identity and one permission contract while keeping their own +signature and docstring, and a bare name that maps to several variants raises +rather than silently picking one by import order. + ### Capabilities Tool access is gated by the **permission engine** (see the HITL section below). Each registered tool declares what it touches via `capabilities=`: @@ -376,6 +492,20 @@ from agentic_cli.tools.sandbox import sandbox_execute The `/sandbox` CLI command lists and resets sandbox sessions. +#### Docker-isolated sandbox + +`sandbox_execute` can run inside a network-isolated Docker container: + +1. Ensure Docker (or podman) is installed and running, and pull the image once: + `docker pull quay.io/jupyter/scipy-notebook:python-3.12` +2. In settings: `sandbox_execute_enabled = true`, `sandbox_backend = "jupyter_docker"`. + +Each session runs in its own container with `--network none`, a read-only +rootfs, dropped Linux capabilities, and memory/CPU/PID caps. Stage data as +files via `sandbox_data_mounts` (`"/host/path:name"` → `/workspace/data/name`, +read-only). Live network access, package installs, and S3 are not supported in +this mode. + #### Web Search Search the web using pluggable backends (Tavily or Brave): @@ -530,22 +660,17 @@ from agentic_cli.tools import memory_tools - `ForgettingPolicy` with `apply_forgetting()` for bounded retention - Archive filtering and `load_all` with tag/source filters -#### Tool Reflection - -Bounded heuristic memory learned from tool failures: - -```python -from agentic_cli.tools import reflection_tools -# save_reflection(tool_name, error_summary, heuristic) -``` - -Each tool keeps at most N reflections (FIFO eviction). Reflections can be injected into tool descriptions to help agents avoid repeating mistakes. Wired via session-end hook. - #### HITL (Human-in-the-Loop) -Tool calls are gated by the **permission engine** (`workflow/permissions/`). Each tool declares a list of capabilities (e.g. `filesystem.write(path=...)`); the engine evaluates them against rules from four sources (builtin defaults, user `~/.{app_name}/settings.json`, project `./.{app_name}/settings.json`, in-memory session). When no rule matches, the user is prompted with `Allow once / Allow for session / Allow always (save to project) / Deny`. Always-grants persist into the project settings file so the next run picks them up automatically. +Tool calls are gated by the **permission engine** (`workflow/permissions/`). Each tool declares a list of capabilities (e.g. `filesystem.write(path=...)`); the engine evaluates them against rules from four sources (builtin defaults, user `~/.{app_name}/settings.json`, project `./.{app_name}/settings.json`, in-memory session). When no rule matches, the user is prompted with `Allow once` / `Allow for this session` / `Allow always for this project` / `Deny`. -See `docs/superpowers/specs/2026-04-18-permissions-system-design.md` for the full design. +**"Allow always" grants are user-owned, not project-owned.** They persist to +`~/.{app_name}/project_grants.json`, keyed by the *resolved project path*, so +they apply to that checkout on this machine only. They are deliberately not +written into `./.{app_name}/settings.json`: a repository could otherwise ship +pre-approved allow-rules that a clone would silently honour. A clone at a +different path therefore starts with no grants, and a repo-committed grant file +is never loaded. ## CLI Commands @@ -559,10 +684,13 @@ Built-in slash commands available in all apps: | `/exit` | `/quit` | Exit the application | | `/settings` | | Interactive settings editor (with persistence) | | `/sandbox` | `/sb` | List / reset stateful sandbox sessions | +| `/jobs` | | List and manage long-running background jobs (`/jobs all`, `/jobs `, `/jobs cancel `, `/jobs clean`) | +| `/resume` | | Resume the agent with results from finished background jobs | | `/papers` | `/docs` | List knowledge-base documents (filter by source, query, --global) | | `/sessions` | `/sess` | List saved sessions (and delete with `--delete=`) | -Apps can add more. Examples like `research_demo` ship with commands like `/save`, `/resume`, and `/kb-backfill`. +Apps can add more of their own. `research_demo`, for instance, registers +`/memory`, `/files` and `/kb-backfill` on top of the built-ins above. ### Adding Custom Commands @@ -656,7 +784,7 @@ See the `examples/` directory for complete working examples: - `websearch_demo.py` — Web search with multiple backends **Full Applications** -- `research_demo/` — Full-featured research assistant with KB ingest + concept pages, semantic memory, sandbox execution, session save/resume. Installable as a console script. +- `research_demo/` — Full-featured research assistant with KB ingest + concept pages, semantic memory, sandbox execution, and durable sessions. Installed as the `research-demo` console script. Run examples: @@ -673,9 +801,18 @@ pip install agentic-cli[langgraph] python examples/hello_langgraph.py # Research demo (full features) +research-demo # console script, installed with the package +python -m research_demo # equivalent module invocation +python -m research_demo --session my-research # resume a durable session by id + +# From a checkout without installing, run it from the repository root: python -m examples.research_demo ``` +Sessions are durable by default: every run gets a session id, `--session ` +resumes a stored one, and `/sessions` lists or deletes them. The demo adds +`/memory`, `/files` and `/kb-backfill` to the built-in commands. + ## Development ### Running Tests @@ -756,7 +893,6 @@ agentic-cli/ │ │ ├── webfetch_tool.py # web_fetch (orchestrator) │ │ ├── pdf_utils.py # PDF text extraction helpers │ │ ├── memory_tools.py # save/search/update/delete + MemoryStore -│ │ ├── reflection_tools.py # save_reflection + ToolReflectionStore │ │ ├── _core/ # Shared planning/task logic │ │ │ ├── planning.py │ │ │ └── tasks.py @@ -782,8 +918,9 @@ agentic-cli/ │ │ ├── sidecar.py # Per-doc markdown sidecar render/parse │ │ ├── sources.py # SearchSource + ArxivSearchSource │ │ └── _mocks.py # MockEmbeddingService, MockVectorStore, mock BM25 -│ └── persistence/ -│ └── session.py # SessionPersistence (save/resume) +│ └── persistence/ # (sessions are persisted natively by each +│ # orchestrator's store; the legacy JSON +│ # SessionPersistence layer was removed) ├── examples/ │ ├── hello_agent.py # Basic ADK example │ ├── hello_langgraph.py # Basic LangGraph example diff --git a/examples/jobs_demo.py b/examples/jobs_demo.py new file mode 100644 index 0000000..c386f69 --- /dev/null +++ b/examples/jobs_demo.py @@ -0,0 +1,184 @@ +"""Jobs Demo - long-running background jobs with a live status-bar monitor. + +Showcases the job-control substrate end to end: + +1. The agent has ``run_shell_job`` (a typed long-running tool) plus the job + management tools, so it can launch detached shell jobs and check on them. +2. The harness ``JobMonitor`` keeps a live ``jobs: N running``/``✓ done`` + segment in the status bar — independent of the agent loop — so jobs stay + visible while you're idle at the prompt. +3. ``/jobs`` lists/inspects/cancels/cleans jobs. + +Run with: + conda run -n agenticcli python examples/jobs_demo.py + +Things to try once it's up (watch the status bar, bottom of the screen): + > run `sleep 30; echo finished` in the background + > start three jobs: sleep 10, sleep 20, sleep 30 + > what's the status of my jobs? + > /jobs + > /jobs all + > /jobs + +API keys are read from ``~/.research_demo/.env`` (shared with the research +demo) or from standard environment variables (GOOGLE_API_KEY / ANTHROPIC_API_KEY). +""" + +import asyncio +from functools import lru_cache +from pathlib import Path + +from pydantic_settings import SettingsConfigDict +from rich.panel import Panel +from rich.text import Text + +from agentic_cli import BaseCLIApp, BaseSettings +from agentic_cli.cli import AppInfo +from agentic_cli.tools import JOB_MANAGEMENT_TOOLS +from agentic_cli.tools.registry import ToolCategory, register_tool +from agentic_cli.workflow import AgentConfig +from agentic_cli.workflow.permissions import Capability +from agentic_cli.workflow.service_registry import JOB_MANAGER, get_service + + +# ============================================================================= +# Typed long-running tool (demo-only) +# ============================================================================= +# +# This is the "starter" half of the job substrate: a typed long-running tool +# that launches work and returns a job_id immediately. The framework ships the +# generic JobManager + observe-only job_* tools, but NOT a starter like this — +# a starter decides what runs and how, which is an application concern. +# +# SECURITY: this runs the command via the subprocess backend (``sh -c``) and +# does NOT go through the hardened shell tool (``tools/shell/``) — no command +# classifier, no blocked-pattern checks. It is gated only by its default-ASK +# ``longrunning.run_shell_job`` capability. That's acceptable for a local demo +# but is why it is NOT a framework built-in. A production shell-job tool should +# route through the shell security layers (or an OS sandbox) first. + + +@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.", +) +def run_shell_job(command: str, name: str = "", cwd: str = "", tool_context=None) -> dict: + """Start ``command`` as a background job and return its ``job_id``. + + Use the ``job_*`` tools to check status, read logs, fetch the result, or + cancel. The job keeps running across turns and survives a CLI restart. + + Because this is marked ``long_running``, ADK wraps it as a + ``LongRunningFunctionTool`` and injects ``tool_context``; we capture its + ``function_call_id`` so the harness can later deliver the result back to + *this* call and auto-resume the agent (phase-2 push/resume). Session/user + are auto-filled from the active turn by ``JobManager.submit``. + """ + jm = get_service(JOB_MANAGER) + if jm is None: + return {"success": False, "error": "job manager not available"} + rec = jm.submit( + tool="run_shell_job", + backend="subprocess", + spec={"command": command, "cwd": cwd or None}, + name=name or None, + resume_on_complete=True, + call_id=getattr(tool_context, "function_call_id", None), + call_name="run_shell_job", + ) + return {"success": True, "job_id": rec.job_id, "state": rec.state.value} + + +# ============================================================================= +# Settings +# ============================================================================= + + +class Settings(BaseSettings): + """Settings for the Jobs demo. + + Reuses the research demo's ``.env`` (``~/.research_demo/.env``) so API keys + only have to live in one place. API keys may also be set via standard + environment variables (GOOGLE_API_KEY / ANTHROPIC_API_KEY). + """ + + model_config = SettingsConfigDict( + env_file=str(Path.home() / ".research_demo" / ".env"), + extra="ignore", + ) + + def __init__(self, **kwargs): + kwargs.setdefault("app_name", "jobs_demo") + kwargs.setdefault("workspace_dir", Path.home() / ".jobs_demo") + # Cap concurrent jobs low so the queue is easy to observe in the demo. + kwargs.setdefault("max_concurrent_jobs", 2) + # Permissions off for a frictionless demo — every run_shell_job would + # otherwise prompt for approval (it declares a default-ASK capability). + kwargs.setdefault("permissions_enabled", False) + # Auto-resume the agent when a background job finishes (phase-2 demo). + kwargs.setdefault("job_auto_resume", True) + super().__init__(**kwargs) + + +@lru_cache +def get_settings() -> Settings: + return Settings() + + +# ============================================================================= +# Agent & App +# ============================================================================= + +AGENT_PROMPT = """You are a build/ops assistant that runs background jobs. + +When the user asks you to run something that may take a while (builds, tests, +downloads, sleeps), use `run_shell_job` to start it as a detached background +job and immediately report the returned job_id. Do NOT block waiting for it. + +To check on a job, call `job_status(job_id)` — it returns the state, a stdout +tail, and (once finished) the result. Use `job_list` to enumerate jobs and +`job_cancel` to stop one. + +Tell the user they can watch live progress in the status bar at the bottom of +the screen, and inspect jobs anytime with the /jobs command. +""" + + +AGENT_CONFIGS = [ + AgentConfig( + name="ops_assistant", + prompt=AGENT_PROMPT, + tools=[run_shell_job, *JOB_MANAGEMENT_TOOLS], + description="Runs and monitors long-running background shell jobs", + ), +] + + +def _create_app_info() -> AppInfo: + text = Text() + text.append("Jobs Demo\n\n", style="bold cyan") + text.append("Long-running background jobs + live status-bar monitor.\n\n", style="dim") + text.append("Try: ", style="dim") + text.append("run `sleep 30; echo done` in the background\n", style="white") + text.append("Then watch the status bar, or type ", style="dim") + text.append("/jobs", style="white") + return AppInfo( + name="Jobs Demo", + version="0.1.0", + welcome_message=lambda: Panel(text, border_style="cyan"), + echo_thinking=False, + ) + + +if __name__ == "__main__": + app = BaseCLIApp( + app_info=_create_app_info(), + agent_configs=AGENT_CONFIGS, + settings=get_settings(), + ) + asyncio.run(app.run()) diff --git a/examples/research_demo/__main__.py b/examples/research_demo/__main__.py index 8e39462..1c14f07 100644 --- a/examples/research_demo/__main__.py +++ b/examples/research_demo/__main__.py @@ -1,10 +1,12 @@ """Entry point for the Research Demo application. Usage: - research-demo # console script (after pip install -e .) + research-demo # console script (after pip install -e .) python -m research_demo + python -m research_demo --session my-research # resume a durable session """ +import argparse import asyncio from .app import ResearchDemoApp @@ -12,7 +14,16 @@ def main() -> None: """Run the Research Demo application.""" - asyncio.run(ResearchDemoApp().run()) + parser = argparse.ArgumentParser(prog="research-demo") + parser.add_argument( + "--session", + metavar="ID", + default=None, + help="Resume a durable session by id (sessions persist by default; " + "list them with /sessions).", + ) + args = parser.parse_args() + asyncio.run(ResearchDemoApp(session_id=args.session).run()) if __name__ == "__main__": diff --git a/examples/research_demo/agents.py b/examples/research_demo/agents.py index 629bc1b..d15fc41 100644 --- a/examples/research_demo/agents.py +++ b/examples/research_demo/agents.py @@ -1,13 +1,17 @@ """Agent configuration for the Research Demo application. -Multi-agent architecture: +Multi-agent architecture with four agents: - research_coordinator: Root agent that owns workflow state (planning, tasks, HITL) - and delegates academic paper research to the arXiv specialist. + and delegates to the arXiv specialist, data analyst, and report writer. - arxiv_specialist: Leaf agent focused on arXiv paper search, analysis, and ingestion. +- data_analyst: Leaf agent that runs multi-step data analysis in a stateful executor. +- report_writer: Leaf agent that turns analysis artifacts into a compiled LaTeX PDF. Uses framework-provided tools exclusively — no app-specific tools needed. """ +from pathlib import Path + from agentic_cli.workflow import AgentConfig from agentic_cli.tools import ( memory_tools, @@ -25,7 +29,9 @@ diff_compare, grep, glob, + compile_document, ) +from agentic_cli.tools.sandbox import sandbox_execute # --------------------------------------------------------------------------- @@ -102,6 +108,61 @@ """ +# --------------------------------------------------------------------------- +# Data Analyst (leaf agent) +# --------------------------------------------------------------------------- + +DATA_ANALYST_PROMPT = """You are a data-analysis specialist. You run multi-step Python analysis in a stateful, isolated executor (variables and DataFrames persist across calls). + +## Data +- Pre-mounted sample datasets are read-only under `/data/samples/` (e.g. `/data/samples/benchmarks.csv`). Discover them with `os.listdir('/data/samples')`. +- Files handed to you by the coordinator arrive via the tool's `inputs` argument and appear at `inputs/`. Load them by that relative path — never by a host path. + +## Working style +1. Explore first: `df = pd.read_csv('/data/samples/benchmarks.csv'); print(df.info()); print(df.describe())`. +2. Transform/aggregate step by step — the session remembers your DataFrames between calls. +3. Plot with matplotlib (figures are captured automatically). +4. Write FINAL deliverables (cleaned tables, key figures) to `outputs/` — those persist and are shared with other agents. Keep scratch in the working directory. +5. Save a short narrative findings report with `write_file`. + +Report what you found with concrete numbers, and name the files you wrote to `outputs/`. +""" + + +# --------------------------------------------------------------------------- +# Report Writer (leaf agent) +# --------------------------------------------------------------------------- + +def report_writer_prompt() -> str: + """Build the prompt from the *active* settings, so the artifacts/build paths + follow whatever workspace is configured (default ~/.research_demo, a tmp dir + under test, or a user override) instead of a hard-coded path. AgentConfig + calls this (get_prompt()) at agent-assembly time, when settings are set.""" + from agentic_cli.config import get_settings + + settings = get_settings() + artifacts_dir = str( + settings.sandbox_outputs_dir or (Path(settings.workspace_dir) / "artifacts") + ) + reports_dir = str(Path(settings.workspace_dir) / "reports") + return f"""You are a report writer. You turn a completed data analysis into a compiled PDF report using LaTeX. + +## Where things are +- Analysis deliverables (figures as .png, tables as .csv) are in the artifacts directory: {artifacts_dir} +- Author your LaTeX source in the build directory: {reports_dir} +- Deliver the final PDF to: {artifacts_dir}/report.pdf + +## How to work +You have a `report-writer` skill — call `load_skill("report-writer")` for the report structure, the LaTeX template, the compile steps, and the error-recovery guide. In short: +1. `glob("{artifacts_dir}/*.png")` to learn the exact figure filenames. +2. Load the template, fill it in, and `write_file` it to {reports_dir}/report.tex (reference figures by BARE filename). +3. Compile with `compile_document(source_path="{reports_dir}/report.tex", output_pdf="{artifacts_dir}/report.pdf", assets_dir="{artifacts_dir}")`. +4. If it fails, read the returned errors/log_tail, fix the .tex, and recompile (at most 3 tries). + +Report the final PDF path to the user. You only use these tools — you do not run arbitrary code. +""" + + # --------------------------------------------------------------------------- # Research Coordinator (root agent) # --------------------------------------------------------------------------- @@ -143,9 +204,35 @@ Rule of thumb: concept pages > sidecars > chunks. Concept pages and sidecars are synthesis-first; chunks are evidence-first. +## Tool Names + +Call only the tools that appear in your tool declarations, and call each one by +its exact declared name, exactly as written. If a task needs two tools, make two +separate calls. To hand work to another agent, call +`transfer_to_agent(agent_name="")`. + +## When to plan, and when to just do it + +Match the response to the size of the request: + +- **Explicit, bounded request** — one clear operation with its parameters + already given ("search arXiv for two papers on X", "ingest this note", + "read it back"). **Do it now**: execute or delegate directly. Do not write a + plan and do not ask for confirmation; the user already told you exactly what + they want. +- **Open-ended or substantial multi-step research** — a goal rather than an + operation ("research X and write me a report", anything spanning several + tools or agents). **Plan first**: `save_plan(content)` with markdown + checkboxes, show the plan immediately, and wait for confirmation. +- **The user asks for a plan** — always plan, whatever the size. + +Planning is a workflow courtesy, not a safety gate: what you are allowed to do +is enforced by tool permissions, not by whether you planned first. So never use +"I should plan" as a reason to refuse or defer a small, explicit request. + ## Workflow Guidelines -When the user asks you to research something: +When the user asks you to research something open-ended: 1. **Check `kb_search_concepts(topic)`** — reuse any existing synthesis before deriving a new one. 2. Browse the knowledge base with `kb_list` to see what's already ingested. 3. Run `kb_search` only if you need evidence that isn't already summarized in a concept page. @@ -154,6 +241,8 @@ 6. **IMMEDIATELY show the plan** to the user in your response. 7. **WAIT for user confirmation** before executing tasks. 8. For arXiv paper research, **delegate to arxiv_specialist** (it has KB writer access and writes concept pages when 3+ related papers accumulate). +- For multi-step data analysis (datasets, DataFrames, plots), delegate to **data_analyst**. Use `execute_python` only for quick one-off calculations. +- To produce a written/PDF **report** of a completed analysis, delegate to **report_writer** (it compiles a LaTeX report from the analysis artifacts). 9. Execute ONE task at a time, updating the plan after each. 10. Use `web_fetch` to extract information from specific URLs found during research. 11. Use `execute_python` for quick calculations and data validation. @@ -192,13 +281,14 @@ - ALWAYS show the plan after creating it - ALWAYS show progress after completing tasks - Share findings and learnings explicitly in your responses -- Ask for confirmation before starting lengthy work +- Ask for confirmation before starting *lengthy* work — not before a single + explicit operation the user has already spelled out - Be thorough and detailed in your findings and reports """ AGENT_CONFIGS = [ - # Leaf agent: arXiv specialist (must be listed before coordinator) + # Leaf agent: arXiv specialist AgentConfig( name="arxiv_specialist", prompt=ARXIV_SPECIALIST_PROMPT, @@ -216,6 +306,23 @@ ], description="arXiv paper research specialist: search, analyze, save, and catalog academic papers", ), + # Leaf agent: data analyst + AgentConfig( + name="data_analyst", + prompt=DATA_ANALYST_PROMPT, + include_state_tools=False, + tools=[sandbox_execute, read_file, write_file, ask_clarification], + description="Stateful data-analysis specialist: loads datasets and runs multi-step pandas/plotting analysis in an isolated executor.", + ), + # Leaf agent: report writer + AgentConfig( + name="report_writer", + prompt=report_writer_prompt, + include_state_tools=False, + tools=[write_file, read_file, glob, compile_document, ask_clarification], + skills=["report-writer"], + description="Report writer: turns the analysis figures/tables in the artifacts dir into a compiled LaTeX PDF report.", + ), # Root agent: research coordinator (owns workflow state, delegates arXiv work) AgentConfig( name="research_coordinator", @@ -241,7 +348,7 @@ grep, diff_compare, ], - sub_agents=["arxiv_specialist"], + sub_agents=["arxiv_specialist", "data_analyst", "report_writer"], description="Research coordinator with memory, planning, task management, knowledge base, and HITL capabilities", ), ] diff --git a/examples/research_demo/app.py b/examples/research_demo/app.py index b8d54f5..f745dfb 100644 --- a/examples/research_demo/app.py +++ b/examples/research_demo/app.py @@ -60,11 +60,16 @@ class ResearchDemoApp(BaseCLIApp): knowledge base ingestion (incl. PDF full-text extraction) """ - def __init__(self, settings: ResearchDemoSettings | None = None) -> None: + def __init__( + self, + settings: ResearchDemoSettings | None = None, + session_id: str | None = None, + ) -> None: super().__init__( app_info=_create_app_info(), agent_configs=AGENT_CONFIGS, settings=settings or ResearchDemoSettings(), + session_id=session_id, ) def register_commands(self) -> None: diff --git a/examples/research_demo/commands.py b/examples/research_demo/commands.py index 90f56e4..7e74bbe 100644 --- a/examples/research_demo/commands.py +++ b/examples/research_demo/commands.py @@ -10,9 +10,26 @@ from agentic_cli.cli.commands import Command, CommandCategory if TYPE_CHECKING: + from agentic_cli.workflow.base_manager import BaseWorkflowManager + from examples.research_demo.app import ResearchDemoApp +def _ready_workflow(app: "ResearchDemoApp") -> "BaseWorkflowManager | None": + """The workflow manager, or None while it is still coming up. + + ``app.workflow`` raises until the controller reports READY, so a command + that touches it during background initialization would otherwise surface as + the generic "Error executing command" from ``BaseCLIApp._handle_command``. + Built-in commands use exactly this shape (see ``SessionsCommand``); the + caller is expected to warn and return. + """ + try: + return app.workflow + except (RuntimeError, AttributeError): + return None + + class MemoryCommand(Command): """Show persistent memory contents.""" @@ -27,7 +44,15 @@ def __init__(self) -> None: ) async def execute(self, args: str, app: "ResearchDemoApp") -> None: - memory_store = app.workflow.memory_manager if app.workflow else None + workflow = _ready_workflow(app) + if workflow is None: + app.session.add_warning( + "Memory is not available yet — the workflow is still " + "initializing. Try /memory again in a moment." + ) + return + + memory_store = workflow.memory_manager table = Table(title="Persistent Memory", show_header=True) table.add_column("ID", style="dim", width=8) @@ -123,9 +148,12 @@ async def execute(self, args: str, app: "ResearchDemoApp") -> None: from agentic_cli.knowledge_base.manager import BackfillAlreadyRunning from agentic_cli.workflow.service_registry import set_service_registry - workflow = app.workflow + workflow = _ready_workflow(app) if workflow is None: - app.session.add_error("Workflow not initialized") + app.session.add_warning( + "The knowledge base is not available yet — the workflow is " + "still initializing. Try /kb-backfill again in a moment." + ) return project_kb = workflow.kb_manager diff --git a/examples/research_demo/data/benchmarks.csv b/examples/research_demo/data/benchmarks.csv new file mode 100644 index 0000000..516b77b --- /dev/null +++ b/examples/research_demo/data/benchmarks.csv @@ -0,0 +1,6 @@ +model,dataset,accuracy,params_millions,year +resnet50,imagenet,76.1,25.6,2015 +vit_b16,imagenet,77.9,86.6,2020 +convnext_t,imagenet,82.1,28.6,2022 +efficientnet_b0,imagenet,77.1,5.3,2019 +swin_t,imagenet,81.3,28.3,2021 diff --git a/examples/research_demo/settings.py b/examples/research_demo/settings.py index cf86af4..a699be8 100644 --- a/examples/research_demo/settings.py +++ b/examples/research_demo/settings.py @@ -12,12 +12,20 @@ class ResearchDemoSettings(BaseSettings): Demonstrates all P0/P1 features with memory, planning, and HITL. - Settings are loaded from (in order of precedence): - 1. Environment variables (RESEARCH_DEMO_* prefix) - 2. Project config (./settings.json) - 3. User config (~/.research_demo/settings.json) - 4. .env file (~/.research_demo/.env) - 5. Default values + Settings are loaded from (highest precedence first): + + 1. Constructor arguments + 2. Environment variables (``RESEARCH_DEMO_*`` prefix) + 3. Project config ``./.research_demo/settings.json`` — **untrusted**: a + cloned repo can ship one, so only an explicit allowlist of benign keys + is honoured and security-sensitive keys are dropped with a warning + 4. User config ``~/.research_demo/settings.json`` (trusted) + 5. ``~/.research_demo/.env`` — trusted because the path is absolute; a + cwd-relative ``.env`` would be filtered like the project config, so put + API keys here or in real environment variables + 6. Field defaults (including the ones set in ``model_post_init`` below) + + JSON sources are only consulted when the file exists. """ model_config = SettingsConfigDict( @@ -42,3 +50,12 @@ def model_post_init(self, __context): """ if "verbose_thinking" not in self.model_fields_set: object.__setattr__(self, "verbose_thinking", False) + if "sandbox_data_mounts" not in self.model_fields_set: + data_dir = Path(__file__).parent / "data" + object.__setattr__(self, "sandbox_data_mounts", [f"{data_dir}:samples"]) + if "sandbox_outputs_dir" not in self.model_fields_set: + object.__setattr__(self, "sandbox_outputs_dir", str(Path(self.workspace_dir) / "artifacts")) + if "skills_dirs" not in self.model_fields_set: + object.__setattr__( + self, "skills_dirs", [str(Path(__file__).parent / "skills")] + ) diff --git a/examples/research_demo/skills/report-writer/SKILL.md b/examples/research_demo/skills/report-writer/SKILL.md new file mode 100644 index 0000000..c84820c --- /dev/null +++ b/examples/research_demo/skills/report-writer/SKILL.md @@ -0,0 +1,62 @@ +--- +name: report-writer +description: Write and compile a LaTeX analysis report to PDF from figures and tables already produced in the run's artifacts directory. Use when the user wants a written or PDF report of a completed analysis. +--- + +# report-writer + +Produce a compiled PDF analysis report with LaTeX. You author a `.tex` file and +compile it with the `compile_document` tool. You do NOT run arbitrary code. + +## Inputs + +Analysis deliverables (figures as `.png`, tables as `.csv`) live in the run's +**artifacts directory** — the concrete path is in your instructions. Before +writing, list it with `glob` to learn the exact figure filenames: + + glob("/*.png") + +Reference figures by **bare filename** (e.g. `accuracy.png`), never full paths — +`compile_document`'s `assets_dir` makes them resolvable. + +## Report structure + +Write these sections, in order: + +1. **Introduction** — the question and why it matters. +2. **Methods** — the data and how it was analyzed. +3. **Results** — the findings, with figures/tables embedded and referenced. +4. **Discussion** — what the results mean; limitations. +5. **References** — a plain list of sources (no citation engine). + +## Authoring + +1. Load the template: `load_skill_resource("report-writer", "assets/report_template.tex")`. +2. Fill the placeholders (title, author/date, section prose, figure includes, table rows). +3. Escape LaTeX specials in prose: `% & _ # $ { }` → `\% \& \_ \# \$ \{ \}`. +4. `write_file` the filled source to the build directory as `report.tex`. + +## Compile + +Call the tool: + + compile_document( + source_path="/report.tex", + output_pdf="/report.pdf", + assets_dir="", + ) + +- On `success: true`, tell the user the report is at `pdf_path`. Done. +- On `success: false`, read `errors` and `log_tail`, fix the `.tex`, and recompile. + **Retry at most 3 times**, then report the failure with the error if still failing. + +## Common errors → fixes + +- `! Undefined control sequence` — a command from a package you did not + `\usepackage`. Add the package or use a base-LaTeX equivalent. +- `! LaTeX Error: File '' not found` — a figure name is wrong; re-`glob` + the artifacts dir and use the exact filename. +- `! Missing $ inserted` / `! You can't use ...` — an unescaped special + character in prose; escape it. +- `No LaTeX engine on PATH` — the host has no TeX install; tell the user to + install TeX Live or MacTeX. You cannot compile without it. diff --git a/examples/research_demo/skills/report-writer/assets/report_template.tex b/examples/research_demo/skills/report-writer/assets/report_template.tex new file mode 100644 index 0000000..6aa7860 --- /dev/null +++ b/examples/research_demo/skills/report-writer/assets/report_template.tex @@ -0,0 +1,52 @@ +\documentclass[11pt]{article} +\usepackage{graphicx} + +% === Fill in: title / author / date === +\title{TITLE_PLACEHOLDER} +\author{AUTHOR_PLACEHOLDER} +\date{DATE_PLACEHOLDER} + +\begin{document} +\maketitle + +\section{Introduction} +% Fill in: the question and why it matters. +INTRO_PLACEHOLDER + +\section{Methods} +% Fill in: the data and how it was analyzed. +METHODS_PLACEHOLDER + +\section{Results} +% Fill in: findings. Embed figures by BARE filename, e.g.: +% \begin{figure}[h] +% \centering +% \includegraphics[width=0.8\textwidth]{accuracy.png} +% \caption{Caption text.} +% \end{figure} +% +% Simple table (base LaTeX, no extra packages): +% \begin{table}[h] +% \centering +% \begin{tabular}{l r} +% \hline +% Metric & Value \\ +% \hline +% Accuracy & 0.789 \\ +% \hline +% \end{tabular} +% \caption{Caption text.} +% \end{table} +RESULTS_PLACEHOLDER + +\section{Discussion} +% Fill in: interpretation and limitations. +DISCUSSION_PLACEHOLDER + +\section*{References} +% Fill in: a plain list of sources. +\begin{itemize} + \item REFERENCE_PLACEHOLDER +\end{itemize} + +\end{document} diff --git a/pyproject.toml b/pyproject.toml index 36f2c4f..43d82fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "agentic-cli" -version = "0.5.3" +version = "0.6.0" description = "A framework for building domain-specific agentic CLI applications" readme = "README.md" license = "MIT" @@ -22,9 +22,15 @@ classifiers = [ ] dependencies = [ "thinking-prompt>=0.3.0", - "google-adk[genai]>=0.4.0", + "google-adk>=1.34,<2", # >=1.34: Anthropic adaptive thinking (negative budget) + # ADK's Claude path (google.adk.models.anthropic_llm) and our DirectAnthropicLlm + # import `anthropic` at module load; google-adk only pulls it via optional extras, + # so declare it directly (>=0.78 matches google-adk's own floor). + "anthropic>=0.78", + "google-genai>=2.0,<3", # imported directly (google.genai.types); was the adk[genai] extra "pydantic>=2.0.0", "pydantic-settings>=2.0.0", + "PyYAML>=6.0", "structlog>=24.0.0", "rich>=13.0.0", "prompt-toolkit>=3.0.0", @@ -32,8 +38,13 @@ dependencies = [ "numpy>=1.26.0,<2.0", "feedparser>=6.0.0", "pypdf>=4.0.0", + "html2text>=2020.1.16", # eager import in tools/webfetch/converter.py "jupyter_client>=8.0.0", "ipykernel>=6.0.0", + # Durable sessions: ADK DatabaseSessionService (SQLAlchemy async) needs an + # async SQLite driver + greenlet; persistence is on by default. + "aiosqlite>=0.20.0", + "greenlet>=3.0.0", ] [project.optional-dependencies] @@ -41,6 +52,10 @@ dev = [ "pytest>=8.0.0", "pytest-asyncio>=0.24.0", "pytest-cov>=6.0.0", + # The console smoke tests drive the real prompt_toolkit application over a + # pty; it is available transitively today (via ipykernel), which is not a + # guarantee, so declare it. + "pexpect>=4.9.0", ] kb = [ "torch>=2.2.0", @@ -78,6 +93,9 @@ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" markers = [ "llm: tests that require real LLM API calls (deselect with -m 'not llm')", + "docker: tests that require a real container runtime (select with -m docker; set SANDBOX_REQUIRE_DOCKER=1 to fail instead of skip when absent)", + "latex: tests that require a host TeX engine (select with -m latex; set LATEX_REQUIRE=1 to fail instead of skip when absent)", + "wheel: acceptance tests that build and install the wheel (select with -m wheel; excluded from the offline suite because they download dependencies)", ] [tool.ruff] diff --git a/src/agentic_cli/__init__.py b/src/agentic_cli/__init__.py index 43f9d51..19dfdac 100644 --- a/src/agentic_cli/__init__.py +++ b/src/agentic_cli/__init__.py @@ -25,8 +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, @@ -38,7 +42,7 @@ validate_settings, reload_settings, ) -from agentic_cli.settings_persistence import SettingsPersistence +from agentic_cli.settings_persistence import SettingsPersistence, SettingsSaveResult from agentic_cli.workflow.settings import WorkflowSettingsMixin from agentic_cli.settings_mixins import AppSettingsMixin, CLISettingsMixin @@ -73,13 +77,22 @@ 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", "SettingsValidationError", "SettingsPersistence", + "SettingsSaveResult", "get_settings", "set_settings", "set_context_settings", @@ -92,4 +105,4 @@ def __getattr__(name: str): "CLISettingsMixin", ] -__version__ = "0.5.3" +__version__ = "0.6.0" diff --git a/src/agentic_cli/cli/app.py b/src/agentic_cli/cli/app.py index 0f965b9..413ec1f 100644 --- a/src/agentic_cli/cli/app.py +++ b/src/agentic_cli/cli/app.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio from typing import TYPE_CHECKING, Any from prompt_toolkit.completion import Completer, Completion @@ -23,15 +24,22 @@ from agentic_cli.cli.workflow_controller import WorkflowController from agentic_cli.config import BaseSettings from agentic_cli.logging import Loggers, configure_logging +from agentic_cli.settings_persistence import PROJECT_SETTABLE_KEYS if TYPE_CHECKING: - from pathlib import Path + from agentic_cli.settings_persistence import SettingsSaveResult from agentic_cli.workflow import GoogleADKWorkflowManager, EventType, WorkflowEvent from agentic_cli.workflow.base_manager import BaseWorkflowManager from agentic_cli.workflow.config import AgentConfig logger = Loggers.cli() +# Synthetic dialog keys → the settings field their setter actually writes. +# "model" is not a field; update_setting() routes it via set_model() to +# default_model. A synthetic key without an entry here is treated as writing +# a field of the same name. +_UI_KEY_TARGET_FIELDS = {"model": "default_model"} + # === Slash Command Completer === @@ -122,7 +130,13 @@ def __init__( session_id: Optional session ID for save/resume. If provided, the session will be loaded on startup and saved on exit. """ - self._session_id = session_id + # Sessions are durable by default: a fresh id is generated when none is + # given (resumable later via /sessions or --session); an explicit id + # requests a resume. + import uuid + + self._resume_requested = session_id is not None + self._session_id = session_id or uuid.uuid4().hex[:12] # === Configuration === self._app_info = app_info self._settings = settings @@ -144,6 +158,10 @@ def __init__( self._workflow_controller.usage_tracker = self._usage_tracker self._message_processor = MessageProcessor() + # Serializes turns: a user turn and a background-job resume turn must + # never overlap (shared session + thinking boxes). + self._turn_lock = asyncio.Lock() + # === UI: ThinkingPromptSession === completer = SlashCommandCompleter(self.command_registry.get_completions()) @@ -185,6 +203,11 @@ def get_ui_setting_keys(self) -> list[str]: Override to customize which settings appear in the UI. Default: model, thinking_effort + Only project-scoped settings may appear: a key whose target field is + not in PROJECT_SETTABLE_KEYS is excluded by _build_ui_items (with a + warning), because a /settings edit to it would persist in the user + ~/.{app}/settings.json and apply across all projects of the app. + Returns: List of field names that should appear in the settings UI """ @@ -205,6 +228,15 @@ def _build_ui_items(self) -> list[Any]: items: list[tuple[int, Any]] = [] for key in self.get_ui_setting_keys(): + # Project-scope guard: the dialog may only expose settings the + # project file can persist (PROJECT_SETTABLE_KEYS). A user-scoped + # key would save to the user ~/.{app}/settings.json and leak the + # change across all projects of the app. + target_field = _UI_KEY_TARGET_FIELDS.get(key, key) + if target_field not in PROJECT_SETTABLE_KEYS: + logger.warning("user_scoped_setting_excluded_from_ui", key=key) + continue + # Handle special 'model' field with dynamic options if key == "model": available_models = list(self._settings.get_available_models()) @@ -247,16 +279,18 @@ def _build_ui_items(self) -> list[Any]: # Sort by order and return items only return [item for _, item in sorted(items, key=lambda x: x[0])] - async def save_settings(self) -> "Path": - """Save current settings to project config file (./.{app_name}/settings.json). + async def save_settings(self) -> "SettingsSaveResult": + """Save current settings, split by trust. - Uses SettingsPersistence to save non-default settings to the - project-level config file. Secrets (API keys) are never saved. + Allowlisted keys go to the project config + (./.{app_name}/settings.json). User-scoped keys differing from their + defaults go to the user config (~/.{app_name}/settings.json), where + the loader trusts them — the project file may only carry allowlisted + keys since P0-1. Secrets (API keys) are never saved. Returns: - Path to the saved config file + SettingsSaveResult with the written path(s) """ - from pathlib import Path from agentic_cli.settings_persistence import SettingsPersistence persistence = SettingsPersistence(self._settings.app_name) @@ -314,7 +348,9 @@ async def apply_settings(self, changes: dict[str, Any]) -> None: needs_reinit = False new_model = changes.get("model") - reinit_settings = {"model", "thinking_effort"} + # `orchestrator` triggers a reinit too: the controller swaps the backend + # when it changes (otherwise the change only took effect on restart). + reinit_settings = {"model", "thinking_effort", "orchestrator"} for key, value in changes.items(): try: @@ -347,6 +383,8 @@ def _register_builtin_commands(self) -> None: ExitCommand, StatusCommand, SandboxCommand, + JobsCommand, + ResumeCommand, PapersCommand, SessionsCommand, ) @@ -357,6 +395,8 @@ def _register_builtin_commands(self) -> None: self.command_registry.register(ExitCommand()) self.command_registry.register(StatusCommand()) self.command_registry.register(SandboxCommand()) + self.command_registry.register(JobsCommand()) + self.command_registry.register(ResumeCommand()) self.command_registry.register(SettingsCommand()) self.command_registry.register(PapersCommand()) self.command_registry.register(SessionsCommand()) @@ -416,45 +456,127 @@ async def _handle_message(self, message: str) -> None: # Echo user input for regular messages self.session.add_message("user", message) - # Delegate to message processor - await self._message_processor.process( - message=message, - workflow_controller=self._workflow_controller, - ui=self.session, - settings=self._settings, - usage_tracker=self._usage_tracker, - ) + # Delegate to message processor (one turn at a time). + async with self._turn_lock: + await self._message_processor.process( + message=message, + workflow_controller=self._workflow_controller, + ui=self.session, + settings=self._settings, + usage_tracker=self._usage_tracker, + session_id=self._session_id, + ) + + # At the turn boundary, auto-resume any finished background jobs that + # opted in (gated by the job_auto_resume setting). + if getattr(self._settings, "job_auto_resume", False): + await self.resume_finished_jobs() + + 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. 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. - async def _load_session_on_startup(self) -> None: - """Load a saved session after workflow initialization.""" - if not self._session_id: - return + 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 + jm = getattr(self._workflow_controller.workflow, "job_manager", None) + if jm is None: + return 0 + + 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 - if not await self._workflow_controller.ensure_initialized(self.session): - self.session.add_warning("Cannot load session — workflow not initialized.") - return + async def _deliver_resume(self, jm, record) -> bool: + """Run one claimed job's resume turn and close out its transition. - workflow = self._workflow_controller.workflow - loaded = await workflow.load_session(self._session_id) - if loaded: - self.session.add_success(f"Session '{self._session_id}' resumed.") - else: - self.session.add_message("system", f"New session '{self._session_id}'.") + 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). - async def _save_session_on_exit(self) -> None: - """Save the current session on exit.""" - if not self._session_id: - return + Args: + jm: The JobManager holding the claim. + record: The claimed job record. - if not self._workflow_controller.is_ready: + Returns: + True if the result was delivered to the agent. + """ + try: + async with self._turn_lock: + 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, + ) + 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. + + Runs for every startup, not just ``--session``: a fresh auto-generated + id must be adopted too, otherwise turns fall back to the manager's + ``default_session`` and unnamed runs silently share one conversation. + """ + if not await self._workflow_controller.ensure_initialized(self.session): + self.session.add_warning("Cannot adopt session — workflow not initialized.") return workflow = self._workflow_controller.workflow - result = await workflow.save_session(self._session_id) - if result.get("success"): - logger.info("session_saved_on_exit", session_id=self._session_id) - else: - logger.error("session_save_on_exit_failed", error=result.get("error")) + resumed = await workflow.load_session(self._session_id) + if resumed: + self.session.add_success(f"Session '{self._session_id}' resumed.") + elif self._resume_requested: + # An explicit --session that didn't exist yet: tell the user it's new. + self.session.add_message("system", f"New session '{self._session_id}'.") async def _extract_session_facts_on_exit(self) -> None: """Extract key facts from the session into memory on exit (if enabled). @@ -483,9 +605,17 @@ async def run(self) -> None: """Run the main application loop.""" logger.info("repl_starting") + from agentic_cli.cli.job_monitor import JobMonitor + + job_monitor = JobMonitor( + self.session, self._workflow_controller, settings=self._settings + ) + async with self._workflow_controller.background_init(self.session): - if self._session_id: - await self._load_session_on_startup() + # Adopt this run's session id (generated or --session) so every turn + # targets it; a fresh id starts a new durable session rather than + # falling back to the shared 'default_session'. + await self._adopt_session_on_startup() # Register input handler @self.session.on_input @@ -494,15 +624,18 @@ async def handle_input(text: str) -> None: return await self.process_input(text) - # Run the session - user sees prompt immediately! - await self.session.run_async() + # Keep long-running jobs visible in the status bar while the + # session runs, independent of the agent loop. + async with job_monitor.running(): + # 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() - # Save persistent session on exit - if self._session_id: - await self._save_session_on_exit() + # No save-on-exit: durable session stores persist continuously per turn. logger.info("app_ending") self.session.add_message("system", "Goodbye!") diff --git a/src/agentic_cli/cli/builtin_commands.py b/src/agentic_cli/cli/builtin_commands.py index 34fee26..a101bc7 100644 --- a/src/agentic_cli/cli/builtin_commands.py +++ b/src/agentic_cli/cli/builtin_commands.py @@ -168,13 +168,29 @@ async def execute(self, args: str, app: Any) -> None: for s in sessions: manager.reset_session(s["session_id"]) app.session.add_success(f"Reset {len(sessions)} sandbox session(s).") - else: - session_id = rest if rest else "default" - was_active = manager.reset_session(session_id) + elif rest: + was_active = manager.reset_session(rest) if was_active: - app.session.add_success(f"Sandbox session '{session_id}' reset.") + app.session.add_success(f"Sandbox session '{rest}' reset.") + else: + app.session.add_warning(f"Sandbox session '{rest}' was not active.") + else: + # No id: session ids are namespaced per conversation (conv-), + # so a literal "default" rarely matches. Reset the current sandbox + # (the single active session), or ask the user to pick if ambiguous. + sessions = manager.list_sessions() + if not sessions: + app.session.add_message("system", "No active sandbox sessions.") + elif len(sessions) == 1: + sid = sessions[0]["session_id"] + manager.reset_session(sid) + app.session.add_success(f"Sandbox session '{sid}' reset.") else: - app.session.add_warning(f"Sandbox session '{session_id}' was not active.") + ids = ", ".join(s["session_id"] for s in sessions) + app.session.add_warning( + f"Multiple sandbox sessions active ({ids}). " + "Reset one with /sandbox reset , or all with /sandbox reset --all." + ) else: # List sessions sessions = manager.list_sessions() @@ -197,6 +213,126 @@ async def execute(self, args: str, app: Any) -> None: app.session.add_rich(table) +class JobsCommand(Command): + """List and manage long-running jobs.""" + + def __init__(self) -> None: + super().__init__( + name="jobs", + description="List and manage long-running jobs", + usage="/jobs [all | | cancel | clean]", + examples=["/jobs", "/jobs all", "/jobs ", "/jobs cancel ", "/jobs clean"], + category=CommandCategory.WORKFLOW, + ) + + async def execute(self, args: str, app: Any) -> None: + """List running jobs, show one job, cancel, or clean completed jobs.""" + try: + manager = getattr(app.workflow, "job_manager", None) + except (RuntimeError, AttributeError): + manager = None + + if manager is None: + app.session.add_warning( + "Jobs not available. Add a long-running job tool (and the job_* tools) " + "to your agent config to enable them." + ) + return + + parts = self.parse_args(args).positional.split() + + if parts and parts[0] == "cancel": + if len(parts) < 2: + app.session.add_warning("Usage: /jobs cancel ") + return + rec = manager.cancel(parts[1]) + if rec is None: + app.session.add_warning(f"No such job: {parts[1]}") + else: + app.session.add_success(f"Job '{parts[1]}' → {rec.state.value}.") + return + + if parts and parts[0] == "clean": + n = manager.clean() + app.session.add_success(f"Removed {n} completed job(s).") + return + + # A bare id → details + recent logs for that one job. + if parts and parts[0] != "all": + rec = manager.get(parts[0]) + if rec is None: + app.session.add_warning(f"No such job: {parts[0]}") + return + self._render_detail(app, manager, rec) + return + + # Otherwise list: running+queued by default, everything with "all". + jobs = manager.list() if parts and parts[0] == "all" else manager.list(active_only=True) + if not jobs: + scope = "" if (parts and parts[0] == "all") else "running " + app.session.add_message("system", f"No {scope}jobs.") + return + + table = Table(title="Jobs", show_lines=False, padding=(0, 1)) + table.add_column("ID", style="bold cyan", no_wrap=True) + table.add_column("Name", style="white") + table.add_column("Backend", style="dim", no_wrap=True) + table.add_column("State", no_wrap=True) + table.add_column("Elapsed", style="dim", no_wrap=True, justify="right") + table.add_column("Exit", style="dim", no_wrap=True, justify="right") + for r in jobs: + table.add_row( + r.job_id, + r.name, + r.backend, + r.state.value, + f"{r.elapsed_s():.0f}s", + "" if r.exit_code is None else str(r.exit_code), + ) + app.session.add_rich(table) + + def _render_detail(self, app: Any, manager: Any, rec: Any) -> None: + lines = [ + f"id: {rec.job_id}", + f"tool: {rec.tool}", + f"name: {rec.name}", + f"backend: {rec.backend}", + f"state: {rec.state.value}", + f"elapsed: {rec.elapsed_s():.0f}s", + f"exit: {rec.exit_code if rec.exit_code is not None else '-'}", + ] + if rec.error: + lines.append(f"error: {rec.error}") + tail = manager.tail(rec.job_id, 15, "stdout") + body = "\n".join(lines) + if tail: + body += "\n\n-- last stdout --\n" + "\n".join(tail) + app.session.add_rich(Panel(body, title=f"Job {rec.job_id}", expand=False)) + + +class ResumeCommand(Command): + """Resume the agent with finished background-job results.""" + + def __init__(self) -> None: + super().__init__( + name="resume", + description="Resume the agent with results from finished background jobs", + usage="/resume", + examples=["/resume"], + category=CommandCategory.WORKFLOW, + ) + + async def execute(self, args: str, app: Any) -> None: + """Drain any finished resume-flagged jobs through resume turns.""" + resume = getattr(app, "resume_finished_jobs", None) + if resume is None: + app.session.add_warning("Resume is not supported by this app.") + return + n = await resume() + if n == 0: + app.session.add_message("system", "No finished background jobs to resume.") + + class PapersCommand(Command): """List documents in the knowledge base.""" @@ -305,22 +441,37 @@ def __init__(self) -> None: ) async def execute(self, args: str, app: Any) -> None: - """Display saved sessions or delete one.""" - from agentic_cli.persistence.session import SessionPersistence + """Display persisted sessions (from the native store) or delete one.""" + from datetime import datetime parsed = self.parse_args(args) delete_id = parsed.get_option("delete", "", str) or "" - persistence = SessionPersistence(app.settings) + try: + workflow = app.workflow + except (RuntimeError, AttributeError): + workflow = None + if workflow is 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 persistence.delete_session(delete_id): + if await workflow.delete_session(delete_id): app.session.add_success(f"Session '{delete_id}' deleted.") else: app.session.add_error(f"Session '{delete_id}' not found.") return - sessions = persistence.list_sessions() + sessions = await workflow.list_sessions() if not sessions: app.session.add_message("system", "No saved sessions.") return @@ -328,21 +479,19 @@ async def execute(self, args: str, app: Any) -> None: table = Table(title="Saved Sessions", show_lines=False, padding=(0, 1)) table.add_column("Session ID", style="bold cyan") table.add_column("Messages", style="dim", justify="right") - table.add_column("Last Saved", style="dim") - table.add_column("Created", style="dim") + table.add_column("Last Update", style="dim") current_sid = app.session_id for s in sessions: sid = s["session_id"] label = f"* {sid}" if sid == current_sid else sid - saved_at = s.get("saved_at", "")[:19].replace("T", " ") - created_at = s.get("created_at", "")[:10] - table.add_row( - label, - str(s.get("message_count", 0)), - saved_at, - created_at, - ) + last = s.get("last_update") + if isinstance(last, (int, float)): + last_str = datetime.fromtimestamp(last).strftime("%Y-%m-%d %H:%M") + else: + last_str = "" + mc = s.get("message_count") + table.add_row(label, "" if mc is None else str(mc), last_str) app.session.add_rich(table) diff --git a/src/agentic_cli/cli/job_monitor.py b/src/agentic_cli/cli/job_monitor.py new file mode 100644 index 0000000..fa5a9d0 --- /dev/null +++ b/src/agentic_cli/cli/job_monitor.py @@ -0,0 +1,210 @@ +"""Background harness monitor that keeps long-running jobs visible in the UI. + +Milestone 2 of the job-control substrate. The agent loop is **not** involved: +this task runs independently of any LLM turn, so detached jobs advance their +state and stay visible even while the user is idle at the prompt or busy in an +unrelated turn. + +Surface choice — the **status bar only**. ``thinking_prompt`` thinking boxes are +turn-oriented (created and finished within a single turn) and the ``add_*`` +output methods print straight to the console via ``print_formatted_text``, which +would corrupt the live prompt if called from a background coroutine. +``ThinkingPromptSession.set_status`` merely updates the status text and +invalidates the app, which is safe to call from any coroutine. So the monitor +renders job state into a status-bar segment and never prints. + +The monitor owns the *jobs* part of the status line; ``WorkflowController`` +remains the single composer of the full bar (model | tokens | jobs | hints) and +reads the segment the monitor publishes via ``jobs_status_segment``. +""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any, AsyncIterator + +from agentic_cli.logging import Loggers + +if TYPE_CHECKING: + from thinking_prompt import ThinkingPromptSession + + from agentic_cli.cli.workflow_controller import WorkflowController + +logger = Loggers.cli() + +# Glyphs for the transient "just finished" note in the status bar. +_DONE_GLYPH = { + "succeeded": "✓", + "failed": "✗", + "cancelled": "⊘", + "unknown": "?", +} + +_DEFAULT_INTERVAL = 2.0 # seconds between reconcile/render ticks +_DEFAULT_NOTE_TICKS = 5 # how many ticks a finished-job note lingers + + +class JobMonitor: + """Periodically reconcile jobs and render their state to the status bar. + + Args: + ui: The active ``ThinkingPromptSession`` (only ``set_status`` is used, + indirectly, via the controller). + controller: The ``WorkflowController``; provides the lazily-created + ``job_manager`` and composes the full status bar. + interval: Seconds between ticks. + note_ticks: How many ticks a finished-job note stays in the segment. + """ + + def __init__( + self, + ui: "ThinkingPromptSession", + controller: "WorkflowController", + *, + settings: Any = None, + interval: float = _DEFAULT_INTERVAL, + note_ticks: int = _DEFAULT_NOTE_TICKS, + ) -> None: + self._ui = ui + self._controller = controller + self._settings = settings + self._interval = interval + self._note_ticks = note_ticks + self._task: asyncio.Task[None] | None = None + # Last seen state value per job id, to detect terminal transitions. + self._states: dict[str, str] = {} + # Recently-finished notes: ``[label, ticks_remaining]`` entries. + self._notes: list[list[Any]] = [] + # Last segment we published, to avoid redundant status-bar redraws. + self._last_segment: str | None = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def start(self) -> None: + """Start the background poll loop (idempotent).""" + if self._task is None: + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + """Cancel the poll loop and wait for it to unwind.""" + if self._task is not None and not self._task.done(): + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + @asynccontextmanager + async def running(self) -> AsyncIterator["JobMonitor"]: + """Run the monitor for the duration of the ``async with`` block.""" + self.start() + try: + yield self + finally: + await self.stop() + + # ------------------------------------------------------------------ + # Core loop + # ------------------------------------------------------------------ + + def _job_manager(self) -> Any | None: + """Return the JobManager if the workflow is ready and has one.""" + if not self._controller.is_ready: + return None + try: + return self._controller.workflow.job_manager + except (RuntimeError, AttributeError): + return None + + async def _run(self) -> None: + while True: + await asyncio.sleep(self._interval) + try: + self.poll_once() + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 - a transient error must not kill the loop + logger.debug("job_monitor_poll_failed", exc_info=True) + + def poll_once(self) -> str | None: + """Reconcile jobs, publish the segment, and refresh the bar on change. + + Returns the segment string (or ``None``). Exposed for tests so the tick + logic can be exercised without the asyncio loop. + """ + jm = self._job_manager() + if jm is None: + return self._last_segment + # ``list()`` reconciles non-terminal jobs and returns all records. + recs = jm.list() + segment = self._build_segment(recs) + self._controller.jobs_status_segment = segment + if segment != self._last_segment: + self._last_segment = segment + self._controller.update_status_bar(self._ui) + return segment + + # ------------------------------------------------------------------ + # Segment construction + # ------------------------------------------------------------------ + + def _build_segment(self, recs: list[Any]) -> str | None: + """Build the status-bar jobs segment from the current job records.""" + from agentic_cli.tools.jobs.backends import TERMINAL_STATES + + terminal_vals = {s.value for s in TERMINAL_STATES} + running = queued = pending_resume = 0 + seen: set[str] = set() + auto_resume = bool(getattr(self._settings, "job_auto_resume", False)) + + for rec in recs: + seen.add(rec.job_id) + state = rec.state.value + if state == "running": + running += 1 + elif state == "queued": + queued += 1 + if ( + auto_resume + and getattr(rec, "resume_on_complete", False) + and not getattr(rec, "resumed", False) + and state in terminal_vals + ): + pending_resume += 1 + prev = self._states.get(rec.job_id) + # Announce a job that we previously saw active and is now terminal. + if state in terminal_vals and prev is not None and prev not in terminal_vals: + glyph = _DONE_GLYPH.get(state, "•") + self._notes.append([f"{glyph} {self._short(rec.name)}", self._note_ticks]) + self._states[rec.job_id] = state + + # Forget jobs that were cleaned away so bookkeeping doesn't grow. + for job_id in [j for j in self._states if j not in seen]: + del self._states[job_id] + + # Age out transient notes. + for note in self._notes: + note[1] -= 1 + self._notes = [n for n in self._notes if n[1] > 0] + + parts: list[str] = [] + active: list[str] = [] + if running: + active.append(f"{running} running") + if queued: + active.append(f"{queued} queued") + if active: + parts.append("jobs: " + ", ".join(active)) + parts.extend(label for label, _ in self._notes) + if pending_resume: + parts.append(f"↻{pending_resume} to resume") + return " · ".join(parts) if parts else None + + @staticmethod + def _short(name: str, limit: int = 24) -> str: + name = (name or "").strip() or "job" + return name if len(name) <= limit else name[: limit - 1] + "…" diff --git a/src/agentic_cli/cli/message_processor.py b/src/agentic_cli/cli/message_processor.py index a916c30..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 === @@ -144,7 +186,8 @@ async def process( ui: "ThinkingPromptSession", settings: "BaseSettings", usage_tracker: "UsageTracker | None" = None, - ) -> None: + session_id: str | None = None, + ) -> TurnResult: """Process a user message through the workflow. Args: @@ -153,6 +196,13 @@ async def process( ui: UI session for output settings: Application settings usage_tracker: Optional tracker for accumulating LLM token usage + 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): @@ -160,14 +210,141 @@ async def process( "Cannot process message - workflow not initialized. " "Please check your API keys (GOOGLE_API_KEY or ANTHROPIC_API_KEY)." ) - return - - # Import WorkflowEvent here (workflow module is now loaded) - from agentic_cli.workflow import WorkflowEvent + return TurnResult( + TurnStatus.UNAVAILABLE, error="workflow not initialized" + ) bind_context(user_id=settings.default_user) logger.info("handling_message", message_length=len(message)) + def _source(workflow): + return workflow.process( + message=message, + user_id=settings.default_user, + session_id=session_id, + ) + + return await self._run_turn( + _source, workflow_controller, ui, settings, usage_tracker + ) + + async def process_resume( + self, + record, + workflow_controller: "WorkflowController", + ui: "ThinkingPromptSession", + settings: "BaseSettings", + usage_tracker: "UsageTracker | 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). + + Args: + record: The terminal JobRecord to resume from. + workflow_controller: Controller managing workflow lifecycle. + 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 TurnResult( + TurnStatus.UNAVAILABLE, error="workflow not initialized" + ) + + workflow = workflow_controller.workflow + bind_context(user_id=settings.default_user) + icon = "✓" if record.state.value == "succeeded" else "✗" + + # Resume only if the backend supports it AND the originating conversation + # is still available (after a restart the in-memory ADK session is gone). + # Otherwise surface a notice — the result stays reachable by job id. + resumable = hasattr( + workflow, "resume_with_job_result" + ) and await workflow.can_resume(record) + if not resumable: + logger.info( + "job_resume_not_resumable", + job_id=record.job_id, + backend=getattr(workflow, "backend_type", "?"), + ) + ui.add_message( + "system", + f"{icon} Background job '{record.name}' finished " + f"({record.state.value}) while its conversation was unavailable " + f"— fetch the result with /jobs {record.job_id}.", + ) + 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( + "system", + f"↻ Background job '{record.name}' finished ({icon} {record.state.value}) " + "— resuming.", + ) + + def _source(wf): + return wf.resume_with_job_result(record) + + return await self._run_turn( + _source, workflow_controller, ui, settings, usage_tracker + ) + + async def _run_turn( + self, + source_factory, + workflow_controller: "WorkflowController", + ui: "ThinkingPromptSession", + settings: "BaseSettings", + usage_tracker: "UsageTracker | 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, workflow_controller=workflow_controller, @@ -190,107 +367,169 @@ async def process( # (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 workflow.process( - message=message, - user_id=settings.default_user, - ): - 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, @@ -383,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, @@ -409,6 +685,17 @@ 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": + from agentic_cli.cli.sandbox_render import render_sandbox_code + + code = (event.metadata.get("tool_args") or {}).get("code", "") + block = render_sandbox_code(code) + if block is not None: + ui.add_rich(block) async def _handle_tool_result( self, @@ -419,11 +706,22 @@ 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") icon = "+" if success else "x" duration_str = f" ({duration}ms)" if duration else "" + # The stateful executor gets a single combined message: a +/x header with + # the run's output indented under a ╰ marker (mirrors the code block). + if tool_name == "sandbox_execute": + from agentic_cli.cli.sandbox_render import render_sandbox_result + + result = event.metadata.get("result") + result_dict = result if isinstance(result, dict) else {} + state.status_line = f"{icon} {tool_name}{duration_str}" + ui.add_rich(render_sandbox_result(result_dict, success=success)) + return lines = event.content.split("\n") first_line = lines[0] state.status_line = f"{icon} {tool_name}: {first_line}{duration_str}" @@ -585,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/src/agentic_cli/cli/sandbox_render.py b/src/agentic_cli/cli/sandbox_render.py new file mode 100644 index 0000000..7489d75 --- /dev/null +++ b/src/agentic_cli/cli/sandbox_render.py @@ -0,0 +1,124 @@ +"""Rich rendering helpers for the stateful executor (``sandbox_execute``). + +Kept as pure functions returning Rich renderables (or ``None``) so the +layout/truncation logic is unit-testable without a live +``ThinkingPromptSession``. The CLI's ``MessageProcessor`` calls these to show +the code being run (on the tool call) and a single combined result message +(on the result):: + + - Running Python in the stateful executor # on call (green "-") + + + + Stateful Python executor output: # on result (green "+", red "x") + + +Content is rendered at normal brightness; only the ``… (+N more lines)`` +truncation hint is dim. +""" +from __future__ import annotations + +from typing import Any + +from rich.console import RenderableType +from rich.syntax import Syntax +from rich.text import Text + +# Display caps (lines). The code block shows the first N lines of the submitted +# source; the result shows the first M lines of output (or the error). +CODE_MAX_LINES = 25 +OUTPUT_MAX_LINES = 10 + +_HEADER = "Running Python in the stateful executor" +_INDENT = " " # 4 spaces + + +def _append_indented( + out: Text, lines: list[Any], dropped: int, *, body_style: str | None = None +) -> None: + """Append body ``lines`` each indented 4 spaces; plain strings or Rich + ``Text`` (already-highlighted code). A dim ``… (+N more lines)`` hint is + appended when ``dropped`` > 0.""" + for line in lines: + out.append("\n") + out.append(_INDENT) + if isinstance(line, Text): + out.append_text(line) + else: + out.append(line, style=body_style) + if dropped: + out.append("\n") + out.append(f"{_INDENT}… (+{dropped} more lines)", style="dim") + + +def render_sandbox_code(code: str, max_lines: int = CODE_MAX_LINES) -> RenderableType | None: + """Header + syntax-highlighted code for a starting ``sandbox_execute`` call. + + Returns ``None`` when there is nothing to show (empty/blank code). The + source is capped at ``max_lines`` with a dim truncation hint. + """ + if not code or not code.strip(): + return None + lines = code.splitlines() + dropped = max(0, len(lines) - max_lines) + shown_lines = lines[:max_lines] + shown = "\n".join(shown_lines) + # Highlight the whole block once (preserves multi-line context), then split + # into per-line Text so we can indent each line ourselves. + highlighted = Syntax( + shown, "python", theme="ansi_dark", background_color="default" + ).highlight(shown) + code_lines = list(highlighted.split("\n"))[: len(shown_lines)] + + out = Text() + out.append("- ", style="green") # mirrors the "+" shown on completion + out.append(_HEADER) # normal brightness + _append_indented(out, code_lines, dropped) + return out + + +def _select_output(result: dict[str, Any]) -> tuple[str, str | None] | None: + """The (text, style) to show for a finished run, or ``None`` if nothing. + + stdout wins when present (even on failure — stdout is the output); otherwise + a failed run falls back to error then stderr so it is not silent. + """ + stdout = result.get("stdout") or "" + if stdout.strip(): + return stdout, None # normal brightness + if not result.get("success", True): + body = result.get("error") or result.get("stderr") or "" + if body.strip(): + return body, "red" + return None + + +def render_sandbox_result( + result: dict[str, Any], + *, + success: bool = True, + max_lines: int = OUTPUT_MAX_LINES, +) -> RenderableType: + """Single combined result message: a ``+``/``x`` icon and a fixed header + label, with the run's output indented below. + + Header is ``Stateful Python executor output:`` for captured stdout and + ``… error:`` for a failure's error/stderr. With nothing to show the header + stands alone: ``+ Stateful Python executor: no output`` (success) or + ``x Stateful Python executor failed``.""" + icon, icon_style = ("+", "green") if success else ("x", "red") + out = Text() + out.append(f"{icon} ", style=icon_style) + selected = _select_output(result) + if not selected: + out.append( + "Stateful Python executor: no output" if success + else "Stateful Python executor failed" + ) + return out + body, style = selected + label = "error" if style == "red" else "output" # style "red" == error/stderr + out.append(f"Stateful Python executor {label}:") + lines = body.splitlines() + dropped = max(0, len(lines) - max_lines) + _append_indented(out, lines[:max_lines], dropped, body_style=style) + return out diff --git a/src/agentic_cli/cli/settings_command.py b/src/agentic_cli/cli/settings_command.py index 8ff9b44..bd78959 100644 --- a/src/agentic_cli/cli/settings_command.py +++ b/src/agentic_cli/cli/settings_command.py @@ -55,9 +55,14 @@ async def execute(self, args: str, app: "BaseCLIApp") -> None: # Apply settings changes await app.apply_settings(result) - # Save settings to project config file + # Save settings, split by trust (allowlisted → project config, + # user-scoped → user config) try: - path = await app.save_settings() - app.session.add_success(f"Settings saved to {path}") + saved = await app.save_settings() + message = f"Settings saved to {saved.project_path}" + if saved.user_path is not None: + keys = ", ".join(saved.user_scoped_keys) + message += f"; user-scoped ({keys}) saved to {saved.user_path}" + app.session.add_success(message) except Exception as e: app.session.add_warning(f"Settings applied but not saved: {e}") diff --git a/src/agentic_cli/cli/workflow_controller.py b/src/agentic_cli/cli/workflow_controller.py index da5222c..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,23 +155,65 @@ 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 @property 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: @@ -123,106 +228,284 @@ 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: - async def _background_init(self) -> None: + - 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. + + 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 - - def _needs_orchestrator_swap(self, new_model: str | None) -> bool: - """Check if switching to new_model requires a different orchestrator. - - Returns True when the model family changes (e.g. Gemini → Claude or - Claude → Gemini) and the current manager type doesn't match what the - factory would create for the new model. + # 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. + + The backend is chosen purely by ``settings.orchestrator`` and is + model-agnostic (ADK runs Claude natively via ``AnthropicLlm``), so a model + change alone never forces a swap. A swap is needed when the live manager's + ``backend_type`` no longer matches the configured orchestrator — e.g. the + orchestrator setting was changed, leaving a stale manager in place. This + happens regardless of whether a new model was given. + + Compares by ``backend_type`` string rather than importing a backend class, + so it never pulls in the optional ``langgraph`` extra on an ADK-only + install. """ - if new_model is None or self._workflow is None: + if self._workflow is None: return False from agentic_cli.workflow.settings import OrchestratorType - from agentic_cli.workflow.langgraph.manager import LangGraphWorkflowManager orchestrator = getattr(self._settings, "orchestrator", OrchestratorType.ADK) - new_needs_langgraph = orchestrator == OrchestratorType.LANGGRAPH or _is_claude_model(new_model) - - current_is_langgraph = isinstance(self._workflow, LangGraphWorkflowManager) - - return new_needs_langgraph != current_is_langgraph + target_backend = getattr(orchestrator, "value", str(orchestrator)) + return getattr(self._workflow, "backend_type", None) != target_backend async def reinitialize(self, model: str | None = None) -> None: """Reinitialize the workflow with optional new model. - If the new model requires a different orchestrator (e.g. switching from - Gemini on ADK to Claude on LangGraph), the entire workflow manager is - replaced. Otherwise, the existing manager is reinitalized in place. + If the live manager's backend no longer matches ``settings.orchestrator`` + (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 @@ -231,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. @@ -274,6 +647,8 @@ def update_status_bar(self, ui: "ThinkingPromptSession") -> None: token_summary = self.usage_tracker.format_status_bar() if token_summary: parts.append(token_summary) + if self.jobs_status_segment: + parts.append(self.jobs_status_segment) parts.extend(["Ctrl+C: cancel", "/help: commands"]) ui.set_status(" | ".join(parts)) # If still initializing, leave status bar unchanged @@ -308,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/config.py b/src/agentic_cli/config.py index 0adf31b..e913df6 100644 --- a/src/agentic_cli/config.py +++ b/src/agentic_cli/config.py @@ -10,11 +10,17 @@ set_settings(my_settings) settings = get_settings() - 2. Context-based (isolated contexts, multi-tenant): + 2. Context-based (isolated settings lookup, e.g. tests or per-request + overrides): with SettingsContext(my_settings): # Code here sees my_settings via get_settings() settings = get_settings() # Returns my_settings + Note: isolation covers settings *lookup* only. API credentials are + exported to process-global env vars at manager initialization + (provider SDKs read them from the environment), so a SettingsContext + does not isolate credentials between contexts in one process. + Settings Loading Priority (highest to lowest): 1. Environment variables (AGENTIC_* prefix) 2. Project config (./.{app_name}/settings.json) @@ -23,11 +29,13 @@ 5. Default values """ +import re 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 import AliasChoices from pydantic_settings import ( BaseSettings as PydanticBaseSettings, SettingsConfigDict, @@ -37,7 +45,17 @@ from agentic_cli.workflow.settings import WorkflowSettingsMixin from agentic_cli.workflow.models import ModelRegistry from agentic_cli.settings_mixins import AppSettingsMixin, CLISettingsMixin -from agentic_cli.settings_persistence import get_project_config_path, get_user_config_path +# PROJECT_SETTABLE_KEYS lives in settings_persistence so the save-side split +# and the load-side filter below share one definition (and because the reverse +# import would be circular). See its definition for the trust rationale. +from agentic_cli.settings_persistence import ( + PROJECT_SETTABLE_KEYS as _PROJECT_SETTABLE_KEYS, + get_project_config_path, + get_user_config_path, +) +from agentic_cli.logging import Loggers + +logger = Loggers.config() __all__ = [ "BaseSettings", @@ -52,15 +70,100 @@ ] +class _AllowlistFilterSource(PydanticBaseSettingsSource): + """Wrap an untrusted settings source, keeping only allowlisted keys. + + Applied to the project ``settings.json`` and a cwd-relative ``.env``. Any + non-allowlisted key is dropped and logged (one warning per key) so a cloned + repo cannot flip a security boundary. Drops (never raises) a non-allowlisted + key, so a repo cannot flip a boundary by *adding* keys. (Malformed JSON or a + bad-typed allowlisted value is still rejected upstream, as before P0-1 — + this narrows, not removes, that pre-existing surface.) + """ + + def __init__( + self, + settings_cls: Type[PydanticBaseSettings], + inner: PydanticBaseSettingsSource, + label: str, + ) -> None: + super().__init__(settings_cls) + self._inner = inner + self._label = label + + def get_field_value(self, field: Any, field_name: str) -> Tuple[Any, str, bool]: + # Unused: __call__ is overridden to filter the inner source's output. + return None, field_name, False + + def __call__(self) -> dict[str, Any]: + kept: dict[str, Any] = {} + for key, value in self._inner().items(): + if key in _PROJECT_SETTABLE_KEYS: + kept[key] = value + else: + logger.warning( + "untrusted_project_setting_ignored", key=key, source=self._label + ) + 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, + *, + untrusted: bool = False, ) -> PydanticBaseSettingsSource | None: """Create a JSON config source if the file exists. Args: settings_cls: The settings class json_file: Path to JSON config file + untrusted: When True (the project file), keep only allowlisted + (non-security) keys so a cloned workspace cannot flip a security + boundary. Returns: JsonConfigSettingsSource if file exists, None otherwise @@ -70,11 +173,21 @@ def _get_json_config_source( try: from pydantic_settings import JsonConfigSettingsSource - return JsonConfigSettingsSource(settings_cls, json_file=json_file) except ImportError: # Older pydantic-settings without JsonConfigSettingsSource return None + if not untrusted: + return JsonConfigSettingsSource(settings_cls, json_file=json_file) + + # Untrusted (the project file): keep only allowlisted keys so a cloned + # workspace cannot flip a security boundary. + return _AllowlistFilterSource( + settings_cls, + JsonConfigSettingsSource(settings_cls, json_file=json_file), + "project settings.json", + ) + class BaseSettings(WorkflowSettingsMixin, AppSettingsMixin, CLISettingsMixin, PydanticBaseSettings): """Base settings for agentic CLI applications. @@ -104,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. @@ -164,10 +289,12 @@ def settings_customise_sources( app_name = app_name or "agentic_cli" - # Add project-level JSON config (./.app_name/settings.json) + # Add project-level JSON config (./.app_name/settings.json). Untrusted: + # strip security-sensitive keys (a cloned repo can ship this file). project_json = _get_json_config_source( settings_cls, get_project_config_path(app_name), + untrusted=True, ) if project_json: sources.append(project_json) @@ -180,8 +307,33 @@ def settings_customise_sources( if user_json: sources.append(user_json) - # Add dotenv settings - sources.append(dotenv_settings) + # dotenv: a cwd-relative env_file is an untrusted project source (a + # cloned repo can ship ./.env), so filter it like project settings.json. + # We inspect EVERY configured env_file (a str/Path or a list of them): + # if ANY entry is cwd-relative after ~ expansion, the whole dotenv + # source is filtered. DotEnvSettingsSource merges all files into one + # dict, so we cannot filter per-file; over-filtering fails safe. An + # absolute/user-level env_file (including a "~/..." path) and real + # environment variables stay trusted. Consequence: secrets/keys placed + # in a cwd .env are dropped — put them in real env vars or an + # absolute/user-level file. + env_file = settings_cls.model_config.get("env_file") + if env_file is None: + _env_entries: list = [] + elif isinstance(env_file, (list, tuple)): + _env_entries = list(env_file) + else: + _env_entries = [env_file] + _has_cwd_relative = any( + e is not None and not Path(e).expanduser().is_absolute() + for e in _env_entries + ) + if _has_cwd_relative: + sources.append( + _AllowlistFilterSource(settings_cls, dotenv_settings, "cwd .env") + ) + else: + sources.append(dotenv_settings) return tuple(sources) @@ -310,20 +462,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: @@ -331,13 +573,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/file_utils.py b/src/agentic_cli/file_utils.py index 540aa8e..a9212c2 100644 --- a/src/agentic_cli/file_utils.py +++ b/src/agentic_cli/file_utils.py @@ -3,6 +3,8 @@ import fcntl import json import os +import re +import stat import tempfile import time from contextlib import contextmanager @@ -59,6 +61,73 @@ def sanitize_filename(name: str) -> str: return "".join(c if c.isalnum() or c in "-_" else "_" for c in name) +def glob_pattern_escapes_root(pattern: str) -> bool: + """True if a glob pattern would search outside its base directory. + + Rejects absolute patterns and any pattern containing a ``..`` path + component. ``pathlib.Path.glob`` does not normalize ``..``, so a pattern + like ``../*`` escapes the base directory even though the permission engine + only authorized that base. Callers should reject such patterns before + globbing. + """ + if os.path.isabs(pattern): + return True + return ".." in re.split(r"[\\/]+", pattern) + + +def path_is_within(path: Path, root: Path) -> bool: + """True if ``path`` resolves to a location inside ``root``. + + Both operands are fully resolved (following symlinks) before the check, so + a symlink under ``root`` that points outside is correctly rejected. Returns + False on any resolution error (loop, permission) — fail closed. + """ + try: + path.resolve().relative_to(root.resolve()) + return True + except (OSError, ValueError, RuntimeError): + return False + + +def copy_regular_file_no_follow(src: Path, dst: Path) -> None: + """Copy a regular file src -> dst without following symlinks at either end. + + Source: opened O_RDONLY | O_NOFOLLOW | O_NONBLOCK (final component); fstat + must be a regular file (rejects symlink / dir / FIFO / socket / device). + O_NOFOLLOW closes the check->open TOCTOU on the final component; O_NONBLOCK + ensures a FIFO/device source fails fast instead of blocking the open + (regular files ignore O_NONBLOCK for reads). + Dest: written to a private temp file in dst.parent, fsync'd, then + os.replace()'d over dst — a pre-planted symlink at dst is replaced, not + written through. + + Raises OSError (source open incl. ELOOP for a final-component symlink) or + ValueError (source not a regular file). POSIX (macOS/Linux). + """ + fd = os.open(src, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK) + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise ValueError(f"not a regular file: {src}") + dst.parent.mkdir(parents=True, exist_ok=True) + tfd, tmp = tempfile.mkstemp(dir=str(dst.parent), prefix=f".{dst.name}.", suffix=".tmp") + tmp_path = Path(tmp) + try: + with os.fdopen(tfd, "wb") as out: + while True: + chunk = os.read(fd, 1 << 20) + if not chunk: + break + out.write(chunk) + out.flush() + os.fsync(out.fileno()) + os.replace(tmp_path, dst) + except BaseException: + tmp_path.unlink(missing_ok=True) + raise + finally: + os.close(fd) + + def _atomic_write(path: Path, content: str) -> None: """Write content to a file atomically and durably. diff --git a/src/agentic_cli/persistence/__init__.py b/src/agentic_cli/persistence/__init__.py index b184dfe..f91fbb5 100644 --- a/src/agentic_cli/persistence/__init__.py +++ b/src/agentic_cli/persistence/__init__.py @@ -1,11 +1,7 @@ -"""Persistence module for agentic CLI applications.""" +"""Persistence package. -from agentic_cli.persistence.session import ( - SessionPersistence, - SessionSnapshot, -) - -__all__ = [ - "SessionPersistence", - "SessionSnapshot", -] +Session state is now persisted natively by each orchestrator's store (ADK +``DatabaseSessionService`` / LangGraph checkpointer), keyed by session id — see +``BaseWorkflowManager`` session methods. The legacy JSON ``SessionPersistence`` +snapshot layer was removed in favor of those durable, full-fidelity stores. +""" diff --git a/src/agentic_cli/persistence/session.py b/src/agentic_cli/persistence/session.py deleted file mode 100644 index 61c26d3..0000000 --- a/src/agentic_cli/persistence/session.py +++ /dev/null @@ -1,243 +0,0 @@ -"""Session persistence for agentic CLI applications. - -Manages saving and loading of session state including message history. -""" - -import json -from dataclasses import dataclass, field -from datetime import datetime -from pathlib import Path -from typing import TYPE_CHECKING, Any - -from agentic_cli.logging import Loggers -from agentic_cli.file_utils import atomic_write_json, file_lock, sanitize_filename - -if TYPE_CHECKING: - from agentic_cli.config import BaseSettings - -logger = Loggers.persistence() - - -@dataclass -class SessionSnapshot: - """A snapshot of session state that can be persisted. - - Contains message history and optional workflow state. - """ - - session_id: str - created_at: datetime - saved_at: datetime - messages: list[dict] - metadata: dict = field(default_factory=dict) - - def to_dict(self) -> dict: - """Convert to dictionary for serialization.""" - return { - "session_id": self.session_id, - "created_at": self.created_at.isoformat(), - "saved_at": self.saved_at.isoformat(), - "messages": self.messages, - "metadata": self.metadata, - } - - @classmethod - def from_dict(cls, data: dict) -> "SessionSnapshot": - """Create from dictionary.""" - return cls( - session_id=data["session_id"], - created_at=datetime.fromisoformat(data["created_at"]), - saved_at=datetime.fromisoformat(data["saved_at"]), - messages=data["messages"], - metadata=data.get("metadata", {}), - ) - - -class SessionPersistence: - """Manages session state persistence. - - Sessions are saved to: - {workspace_dir}/sessions/ - ├── {session_id}.json - └── ... - """ - - def __init__(self, settings: "BaseSettings"): - """Initialize session persistence. - - Args: - settings: Application settings with workspace_dir - """ - self.sessions_path = settings.sessions_dir - self._ensure_directory_exists() - - def _ensure_directory_exists(self) -> None: - """Create sessions directory if it doesn't exist.""" - self.sessions_path.mkdir(parents=True, exist_ok=True) - - def _get_session_path(self, session_id: str) -> Path: - """Get path for a session file.""" - safe_id = sanitize_filename(session_id) - return self.sessions_path / f"{safe_id}.json" - - def _get_sessions_index_path(self) -> Path: - """Get path to sessions index file.""" - return self.sessions_path / "_sessions_index.json" - - def _load_sessions_index(self) -> dict[str, dict]: - """Load sessions index from disk. - - Returns: - Dict mapping session_id to metadata dict. - """ - index_path = self._get_sessions_index_path() - if not index_path.exists(): - return {} - try: - data = json.loads(index_path.read_text()) - if not isinstance(data, dict): - raise ValueError("sessions index is not a JSON object") - return data.get("sessions", {}) - except (json.JSONDecodeError, ValueError, KeyError, OSError) as exc: - # A corrupt index would otherwise read as empty, and the next save - # would reset it to a single entry — permanently hiding every other - # session whose file still exists. Rebuild from the session files - # instead of silently dropping them. - logger.warning("sessions_index_corrupt_rebuilding", error=str(exc)) - return self._rebuild_sessions_index() - - def _save_sessions_index(self, index: dict[str, dict]) -> None: - """Save sessions index to disk.""" - atomic_write_json(self._get_sessions_index_path(), {"sessions": index}) - - def _update_sessions_index( - self, - session_id: str, - created_at: str, - saved_at: str, - message_count: int, - ) -> None: - """Update sessions index with metadata for a single session.""" - with file_lock(self._get_sessions_index_path()): - sessions_index = self._load_sessions_index() - sessions_index[session_id] = { - "session_id": session_id, - "created_at": created_at, - "saved_at": saved_at, - "message_count": message_count, - } - self._save_sessions_index(sessions_index) - - def _rebuild_sessions_index(self) -> dict[str, dict]: - """Rebuild sessions index from individual session files (migration).""" - index: dict[str, dict] = {} - for session_file in self.sessions_path.glob("*.json"): - if session_file.name.startswith("_"): - continue - try: - with open(session_file, "r") as f: - data = json.load(f) - sid = data["session_id"] - index[sid] = { - "session_id": sid, - "created_at": data["created_at"], - "saved_at": data["saved_at"], - "message_count": len(data["messages"]), - } - except (json.JSONDecodeError, KeyError, TypeError, OSError): - continue - self._save_sessions_index(index) - return index - - def save_snapshot(self, snapshot: SessionSnapshot) -> Path: - """Save a SessionSnapshot to disk. - - Args: - snapshot: The session snapshot to persist. - - Returns: - Path to saved session file. - """ - session_path = self._get_session_path(snapshot.session_id) - atomic_write_json(session_path, snapshot.to_dict()) - - self._update_sessions_index( - session_id=snapshot.session_id, - created_at=snapshot.created_at.isoformat(), - saved_at=snapshot.saved_at.isoformat(), - message_count=len(snapshot.messages), - ) - - logger.info( - "snapshot_saved", - session_id=snapshot.session_id, - message_count=len(snapshot.messages), - path=str(session_path), - ) - return session_path - - def load_session(self, session_id: str) -> SessionSnapshot | None: - """Load a saved session. - - Args: - session_id: Session identifier to load - - Returns: - SessionSnapshot if found, None otherwise - """ - session_path = self._get_session_path(session_id) - if not session_path.exists(): - logger.debug("session_not_found", session_id=session_id) - return None - - with open(session_path, "r") as f: - data = json.load(f) - - snapshot = SessionSnapshot.from_dict(data) - logger.info( - "session_loaded", - session_id=session_id, - message_count=len(snapshot.messages), - ) - return snapshot - - def list_sessions(self) -> list[dict]: - """List all saved sessions. - - Returns: - List of session summaries (id, created_at, message_count) - """ - # Try reading from sessions index first - sessions_index = self._load_sessions_index() - if not sessions_index: - # Fallback: rebuild from files (migration path) - sessions_index = self._rebuild_sessions_index() - - sessions = list(sessions_index.values()) - - # Sort by saved_at descending - sessions.sort(key=lambda x: x["saved_at"], reverse=True) - return sessions - - def delete_session(self, session_id: str) -> bool: - """Delete a saved session. - - Args: - session_id: Session identifier to delete - - Returns: - True if deleted, False if not found - """ - session_path = self._get_session_path(session_id) - try: - session_path.unlink() - except FileNotFoundError: - return False - - # Update sessions index (lock to prevent TOCTOU race) - with file_lock(self._get_sessions_index_path()): - sessions_index = self._load_sessions_index() - sessions_index.pop(session_id, None) - self._save_sessions_index(sessions_index) - - return True diff --git a/src/agentic_cli/settings_persistence.py b/src/agentic_cli/settings_persistence.py index c2a3572..ca1cd0b 100644 --- a/src/agentic_cli/settings_persistence.py +++ b/src/agentic_cli/settings_persistence.py @@ -1,10 +1,17 @@ """Settings persistence utilities. Provides functionality to save settings to JSON files for layered configuration. -Settings are saved to ./.{app_name}/settings.json (project config) by default. + +Saving splits by trust (mirroring the P0-1 load-side allowlist in +``agentic_cli.config``): allowlisted keys go to the project +``./.{app_name}/settings.json``; every other (user-scoped) key goes to the +trusted user ``~/.{app_name}/settings.json``. Without the split, a key the +loader refuses to read from the project file would be written there and +silently dropped on the next start. """ import json +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -18,6 +25,38 @@ "anthropic_api_key", "tavily_api_key", "brave_api_key", + "postgres_uri", # connection string embeds user:password@host +}) + +# Deny-by-default allowlist: the ONLY keys a project ./.{app}/settings.json (or +# a cwd-relative .env) may set. A cloned/untrusted repo must not be able to flip +# a security boundary — executor backend, container image/user, bind mounts, +# outputs dir, OS-sandbox policy, shell backend, raw LLM logging, workspace dir, +# permission rules, or secrets. Every entry below is a benign field that cannot +# select code execution, filesystem/mount scope, container identity/image, +# network policy, secrets, or sensitive logging. Anything not clearly benign — +# and any new field — is excluded automatically. Real environment variables and +# the user ~/.{app}/settings.json remain fully trusted. +# +# Used by BOTH sides of persistence: the load-side filter +# (config._AllowlistFilterSource) and the save-side split +# (SettingsPersistence.save), so writer and reader cannot drift apart. +PROJECT_SETTABLE_KEYS = frozenset({ + # model / behavior + "default_model", "thinking_effort", "orchestrator", + "context_window_trigger_tokens", "context_window_target_tokens", + # retry / request timeouts (not code paths) + "retry_max_attempts", "retry_initial_delay", "retry_backoff_factor", + "anthropic_request_timeout", "python_executor_timeout", "sandbox_timeout", + # sandbox RESOURCE limits (not backend / image / mounts / user / network) + "sandbox_max_sessions", "sandbox_memory_mb", "sandbox_cpus", "sandbox_pids_limit", + # non-exec tool config + "search_backend", + "webfetch_cache_ttl_seconds", "webfetch_max_content_bytes", "webfetch_max_pdf_bytes", + # persistence backend selection (NOT the credential-bearing postgres_uri) + "session_store", + # display / logging verbosity (NOT raw_llm_logging) + "log_level", "log_format", "verbose_thinking", }) # Identity fields set by the application, not the user @@ -37,6 +76,35 @@ def get_user_config_path(app_name: str) -> Path: return Path.home() / f".{app_name}" / "settings.json" +@dataclass(frozen=True) +class SettingsSaveResult: + """Where a settings save landed. + + ``project_path`` always receives the allowlisted keys (or, with an + explicit ``path=``, the full legacy dump). ``user_path`` is set only when + user-scoped (non-allowlisted) keys were written to the user config; + ``user_scoped_keys`` lists the keys added/updated/removed there. + """ + + project_path: Path + user_path: Path | None = None + user_scoped_keys: tuple[str, ...] = () + + +def get_user_project_grants_path(app_name: str) -> Path: + """Path to interactively-granted permission rules + (``~/.{app_name}/project_grants.json``), keyed by resolved project path. + + Lives in USER config — never in the repo — so a cloned workspace carries no + "Allow always" grants: a clone at a different path simply has no entry and + the user re-grants. Structure:: + + { "": {"permissions": {"allow": [...], "deny": [...]}} } + """ + return Path.home() / f".{app_name}" / "project_grants.json" + + + class SettingsPersistence: """Manages loading and saving settings to JSON files. @@ -75,24 +143,37 @@ def save( self, settings: "BaseSettings", path: Path | None = None, - ) -> Path: - """Save settings to JSON config file. - - By default saves to project config (./.{app_name}/settings.json). - Secrets (API keys) and identity fields are never saved. - All other user-configurable settings are saved regardless of - whether they match the schema default, because subclasses may - have different effective defaults. + ) -> SettingsSaveResult: + """Save settings, split by trust to match the load-side allowlist. + + Default save writes two files: + + - Project config (./.{app_name}/settings.json) receives ONLY + allowlisted (``PROJECT_SETTABLE_KEYS``) fields, always all of them — + regardless of whether they match the schema default, because + subclasses may have different effective defaults. The file is fully + rewritten, which also heals stale pre-P0-1 files holding keys the + loader now ignores. + - User config (~/.{app_name}/settings.json) receives the remaining + user-scoped fields, but only those differing from the settings + class default (so code-default changes keep applying for untouched + fields, and the user file stays minimal). A user-scoped field back + at its default is REMOVED from the file. Keys this method does not + manage (hand-stored secrets, unknown/domain keys) are preserved + verbatim; a malformed user file raises rather than being clobbered. + + Secrets (API keys) and identity fields are never written anywhere. + With an explicit ``path=``, the legacy behavior is kept: one full + dump (minus secrets/identity) to that file, no split. Args: settings: Settings instance to save - path: Optional custom path (defaults to project_config_path) + path: Optional custom path (single-file legacy dump) Returns: - Path to the saved config file + SettingsSaveResult with the written path(s) """ - target_path = path or self.project_config_path - target_path.parent.mkdir(parents=True, exist_ok=True) + from agentic_cli.file_utils import atomic_write_text # Get settings as dict, excluding secrets and identity fields data = settings.model_dump( @@ -103,11 +184,77 @@ def save( # Convert Path objects to strings for JSON serialization data = self._serialize_paths(data) - # Write atomically - from agentic_cli.file_utils import atomic_write_text - atomic_write_text(target_path, json.dumps(data, indent=2, default=str)) + if path is not None: + # Explicit target: legacy single-file full dump. + path.parent.mkdir(parents=True, exist_ok=True) + atomic_write_text(path, json.dumps(data, indent=2, default=str)) + return SettingsSaveResult(project_path=path) + + project_data = {k: v for k, v in data.items() if k in PROJECT_SETTABLE_KEYS} + user_updates, user_removals = self._split_user_scoped(settings, data) + + project_path = self.project_config_path + project_path.parent.mkdir(parents=True, exist_ok=True) + atomic_write_text( + project_path, json.dumps(project_data, indent=2, default=str) + ) - return target_path + user_path = self.user_config_path + # Merge-write: never clobber keys we don't manage (e.g. hand-stored + # API keys). A malformed user file raises instead of being replaced. + existing: dict[str, Any] = {} + if user_path.exists(): + existing = json.loads(user_path.read_text()) + if not isinstance(existing, dict): + raise ValueError( + f"User settings file is not a JSON object: {user_path}" + ) + merged = dict(existing) + removed = tuple(k for k in sorted(user_removals) if k in existing) + for key in removed: + del merged[key] + merged.update(user_updates) + + if merged == existing: + return SettingsSaveResult(project_path=project_path) + + user_path.parent.mkdir(parents=True, exist_ok=True) + atomic_write_text(user_path, json.dumps(merged, indent=2, default=str)) + return SettingsSaveResult( + project_path=project_path, + user_path=user_path, + user_scoped_keys=tuple(sorted(user_updates)) + removed, + ) + + def _split_user_scoped( + self, settings: "BaseSettings", data: dict[str, Any] + ) -> tuple[dict[str, Any], set[str]]: + """Partition non-allowlisted dumped fields by deviation from default. + + Returns (updates, removals): ``updates`` maps user-scoped keys whose + live value differs from the settings class default to their dumped + value; ``removals`` holds user-scoped keys back at their default, + whose stale entries should leave the user file. Comparison uses the + instance's own class fields, so subclass default overrides are + respected. + """ + from pydantic_core import PydanticUndefined + + updates: dict[str, Any] = {} + removals: set[str] = set() + fields = type(settings).model_fields + for key, value in data.items(): + if key in PROJECT_SETTABLE_KEYS: + continue + field = fields.get(key) + if field is None: + continue + default = field.get_default(call_default_factory=True) + if default is not PydanticUndefined and getattr(settings, key) == default: + removals.add(key) + else: + updates[key] = value + return updates, removals def load(self, path: Path | None = None) -> dict[str, Any]: """Load settings from JSON config file. diff --git a/src/agentic_cli/tools/__init__.py b/src/agentic_cli/tools/__init__.py index 12125de..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 @@ -50,7 +51,27 @@ fetch_arxiv_paper, ) from agentic_cli.tools.execution_tools import execute_python +from agentic_cli.tools.document import compile_document from agentic_cli.tools.interaction_tools import ask_clarification + +# Long-running job tools (generic observe-only management). Typed long-running +# tools that *start* work (e.g. run_shell_job) are application-provided — see +# examples/jobs_demo.py — because they choose what runs and how. +from agentic_cli.tools.jobs import ( + job_status, + job_result, + job_logs, + job_cancel, + job_list, +) + +# Minimal companion to a typed long-running tool: job_status alone returns +# state + stdout tail + result-when-finished, keeping the agent's tool surface +# small (tool-selection quality drops past ~15-20 tools). Apps that want the +# LLM to also enumerate/cancel jobs can use JOB_MANAGEMENT_TOOLS instead — but +# listing/cancelling is usually a human job via the /jobs command. +JOB_TOOLS = [job_status] +JOB_MANAGEMENT_TOOLS = [job_status, job_result, job_logs, job_cancel, job_list] from agentic_cli.tools.search import web_search from agentic_cli.tools.webfetch_tool import web_fetch from agentic_cli.tools.registry import ( @@ -58,6 +79,7 @@ ToolDefinition, ToolRegistry, get_registry, + declare_tool, register_tool, ) @@ -70,6 +92,7 @@ "ToolDefinition", "ToolRegistry", "get_registry", + "declare_tool", "register_tool", # Executor classes "SafePythonExecutor", @@ -107,7 +130,16 @@ "search_arxiv", "fetch_arxiv_paper", "execute_python", + "compile_document", "ask_clarification", + # Long-running jobs (observe-only; typed starters are app-provided) + "job_status", + "job_result", + "job_logs", + "job_cancel", + "job_list", + "JOB_TOOLS", + "JOB_MANAGEMENT_TOOLS", # Framework tool modules (lazy loaded) "memory_tools", "sandbox_tools", 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/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/document/__init__.py b/src/agentic_cli/tools/document/__init__.py new file mode 100644 index 0000000..6ebd9c8 --- /dev/null +++ b/src/agentic_cli/tools/document/__init__.py @@ -0,0 +1,5 @@ +"""Document-generation tools (LaTeX → PDF).""" + +from agentic_cli.tools.document.compile import compile_document + +__all__ = ["compile_document"] diff --git a/src/agentic_cli/tools/document/compile.py b/src/agentic_cli/tools/document/compile.py new file mode 100644 index 0000000..d2ad630 --- /dev/null +++ b/src/agentic_cli/tools/document/compile.py @@ -0,0 +1,380 @@ +"""Compile a LaTeX document to PDF with a host TeX engine. + +``compile_document`` is a narrow, permission-gated tool: it runs ``latexmk`` +(preferred) or ``pdflatex`` as a guarded subprocess with the two arbitrary-code +vectors disabled — ``latexmk`` with ``-norc`` (so no ``.latexmkrc`` Perl is read +from the build directory or home) and ``pdflatex`` with ``-no-shell-escape`` (no +``\\write18``). It runs on the host (not the container sandbox — an intentional +decoupling); OS-sandbox confinement of the build tree is deferred (spec §9). + +The tool runs with a wall-clock timeout and a private temp build dir, and +returns a structured result. It does not execute arbitrary host code by +default; a ``report_writer``-style agent uses it to turn an authored ``.tex`` +into a PDF. Build intermediates (``*.aux``, ``*.log``) are isolated in a +``tempfile.TemporaryDirectory`` and never written to the source or delivery dir; +only the final PDF is promoted via ``_deliver_no_follow``. + +The subprocess receives only an allowlisted environment (``_ENV_PASSTHROUGH`` + +the ``TEXMF*``/TeX config vars), not the full host environment, so host secrets +aren't handed to the TeX process. Delivery to ``output_pdf`` is a no-follow +atomic write (temp file + ``os.replace``) so an attacker-placed symlink at the +destination can't redirect the write. + +Provisioning is host-based: the engine must be on ``PATH`` (TeX Live / MacTeX). +If neither is found the tool returns a structured error with an install hint. + +The subprocess call and the engine lookup sit behind module-level seams +(``_run``, ``_which``) so the logic is unit-tested offline without a real TeX +install. +""" + +from __future__ import annotations + +import os +import shutil +import signal +import subprocess +import tempfile +import time +from pathlib import Path +from typing import Any + +from agentic_cli.tools.registry import ToolCategory, register_tool +from agentic_cli.workflow.permissions import Capability + +_ENGINES = ("latexmk", "pdflatex") +_LOG_TAIL_LINES = 40 +_LOG_TAIL_BYTES = 64 * 1024 +_MAX_CAPTURE_BYTES = 200_000 +# Cap the size of any single file the TeX child writes (bounds runaway-`.tex` +# disk use, incl. the log/pdf/aux and our redirected stdout/stderr temp files). +_RLIMIT_FSIZE_BYTES = 500 * 1024 * 1024 + +# Only these host env vars reach the TeX process. The tool must not hand the +# whole host environment (API keys, tokens) to a subprocess that — on the +# latexmk path — can execute arbitrary Perl from a .latexmkrc. PATH/HOME are +# needed for the engine binary and kpathsea; TEXINPUTS is set explicitly. +_ENV_PASSTHROUGH = ( + "PATH", "HOME", "TERM", "TMPDIR", "TEMP", "TMP", + "LANG", "LC_ALL", "LC_CTYPE", "SOURCE_DATE_EPOCH", + # latexmk is a Perl program; without its module path it can fail to load. + "PERL5LIB", "PERLLIB", +) + +# TeX's own search/config vars (kpathsea) that don't fall under the TEXMF* +# namespace. TEXINPUTS is deliberately excluded — it is set explicitly below. +_TEX_VARS = ( + "TEXFONTS", "TEXFORMATS", "TEXPOOL", "TEXPSHEADERS", + "TEXCONFIG", "TEXDOCS", "TEXSOURCES", +) + +# The kpathsea TEXMF* configuration variables (exact — a strict secret boundary, +# so a name like TEXMF_SECRET is not passed through). +_TEXMF_VARS = ( + "TEXMFHOME", "TEXMFVAR", "TEXMFCONFIG", "TEXMFCACHE", "TEXMFLOCAL", + "TEXMFDIST", "TEXMFMAIN", "TEXMFSYSVAR", "TEXMFSYSCONFIG", "TEXMFDBS", + "TEXMFCNF", "TEXMFOUTPUT", +) + + +def _build_env(assets_dir: str | None, source_dir: str | None = None) -> dict[str, str]: + """Minimal, allowlisted environment for the TeX subprocess. + + Passes PATH/HOME/locale plus TeX's own configuration variables — an exact + list of ``TEXMF*`` vars and a fixed set of other TeX vars — so a custom + ``TEXMFHOME`` etc. keeps working, but not arbitrary host env (a name like + ``TEXT_API_TOKEN`` starts with "TEX" yet is not a TeX var), and never the + caller's ``TEXINPUTS`` (set explicitly below). + + ``assets_dir`` and ``source_dir`` become TEXINPUTS read roots so figures and + ``\\input`` siblings resolve even though the build runs in a private temp dir. + """ + env = { + k: v + for k, v in os.environ.items() + if k in _ENV_PASSTHROUGH or k in _TEXMF_VARS or k in _TEX_VARS + } + env.setdefault("PATH", os.defpath) + roots = [ + str(Path(r).expanduser().resolve()) + for r in (assets_dir, source_dir) + if r + ] + if roots: + # Trailing empty entry lets kpathsea append its default search path. + env["TEXINPUTS"] = os.pathsep.join(roots) + os.pathsep + return env + + +def _deliver_no_follow(produced: Path, dest: Path) -> None: + """Copy ``produced`` to ``dest`` without following a symlink at ``dest``. + + Writes to a private temp file in dest's directory, then atomically renames + it over dest. ``os.replace`` swaps the destination *name*: if dest is a + symlink the link itself is replaced (not written through), so an + attacker-placed symlink can't redirect the write outside the intended path. + + This protects only the final path component. A symlinked ``dest.parent`` + (or an ancestor) still redirects the write; the permission engine + canonicalizes ``output_pdf`` at check time, but a check→write window + remains. Full parent containment is deferred (spec §9, OS-sandbox). + """ + dest.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + dir=str(dest.parent), prefix=f".{dest.name}.", suffix=".tmp" + ) + tmp_path = Path(tmp_name) + try: + with os.fdopen(fd, "wb") as out, open(produced, "rb") as src: + shutil.copyfileobj(src, out) + out.flush() + os.fsync(out.fileno()) + # Match the produced PDF's mode/mtime (mkstemp is 0600) so the delivered + # file has the readability a consumer expects, as the old copy2 did. + shutil.copystat(produced, tmp_path) + os.replace(tmp_path, dest) + except BaseException: + tmp_path.unlink(missing_ok=True) + raise + + +def _which(name: str) -> str | None: + """Locate an executable on PATH (seam for tests).""" + return shutil.which(name) + + +def _child_rlimits(cpu_seconds: int): + """Build a ``preexec_fn`` that caps the child's file size and CPU time, so a + runaway ``.tex`` can't fill the disk or burn CPU within the wall-clock + timeout. Best-effort (a platform without ``resource`` just skips it).""" + def _apply() -> None: # runs in the forked child, before exec + try: + import resource + resource.setrlimit( + resource.RLIMIT_FSIZE, (_RLIMIT_FSIZE_BYTES, _RLIMIT_FSIZE_BYTES) + ) + resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds)) + except (ValueError, OSError, ImportError): + pass + return _apply + + +def _run( + argv: list[str], *, cwd: str, env: dict[str, str], timeout: float +) -> subprocess.CompletedProcess: + """Run a subprocess in its own process group so a timeout kills the whole + tree (latexmk + its pdflatex grandchild), not just the direct child. + stdout/stderr are captured to temp files and only the last + _MAX_CAPTURE_BYTES of each are retained, bounding host memory. The child + also runs under RLIMIT_FSIZE/RLIMIT_CPU limits. Seam for tests. POSIX + (macOS/Linux), which is what the framework targets.""" + with tempfile.TemporaryFile() as out_f, tempfile.TemporaryFile() as err_f: + proc = subprocess.Popen( + argv, cwd=cwd, env=env, + stdout=out_f, stderr=err_f, + start_new_session=True, + preexec_fn=_child_rlimits(int(timeout) + 30), + ) + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + proc.wait() # reap the killed group + raise + out = _tail_of_file(out_f, _MAX_CAPTURE_BYTES) + err = _tail_of_file(err_f, _MAX_CAPTURE_BYTES) + return subprocess.CompletedProcess(argv, proc.returncode, stdout=out, stderr=err) + + +def _detect_engine(engine: str | None) -> str | None: + """Return the engine to use, or None if unavailable.""" + if engine is not None: + return engine if (engine in _ENGINES and _which(engine)) else None + for candidate in _ENGINES: + if _which(candidate): + return candidate + return None + + +def _build_argv(engine: str, source: str) -> list[str]: + """Compiler argv — never enables shell-escape. + + latexmk runs with ``-norc`` so it won't read ``.latexmkrc`` (arbitrary + Perl) from the build directory or home; pdflatex runs with + ``-no-shell-escape``. Neither path executes arbitrary host code by default. + """ + if engine == "latexmk": + return [ + "latexmk", "-norc", "-pdf", "-interaction=nonstopmode", + "-halt-on-error", source, + ] + return [ + "pdflatex", "-no-shell-escape", "-interaction=nonstopmode", + "-halt-on-error", source, + ] + + +def _safe_source_arg(name: str) -> str: + """Anchor a source filename so a leading ``-`` can't be parsed as an engine + option (arbitrary-exec via ``-pdflatex=CMD`` etc.). ``name`` is a basename + and the subprocess cwd is the source's directory, so ``./`` resolves it.""" + return name if name.startswith("./") else f"./{name}" + + +def _parse_errors(log_text: str) -> list[str]: + """Extract LaTeX error lines (those beginning with '!') from a log.""" + return [ln for ln in log_text.splitlines() if ln.startswith("!")] + + +def _tail_of_file(f, limit: int) -> str: + """Return the last ``limit`` bytes of an open binary temp file, decoded.""" + size = f.seek(0, os.SEEK_END) + f.seek(max(0, size - limit)) + return f.read().decode("utf-8", errors="replace") + + +def _read_log_tail(log_path: Path, fallback: str) -> str: + """Return at most the last _LOG_TAIL_BYTES of the log (decoded), else + ``fallback``. Bounds memory on a runaway compiler log.""" + try: + if not log_path.is_file(): + return fallback + with open(log_path, "rb") as f: + size = f.seek(0, os.SEEK_END) + f.seek(max(0, size - _LOG_TAIL_BYTES)) + return f.read().decode("utf-8", errors="replace") + except OSError: + return fallback + + +@register_tool( + category=ToolCategory.EXECUTION, + capabilities=[ + Capability("document.compile", target_arg="source_path"), + # The tool also reads assets_dir and writes output_pdf when those are + # supplied; scope them explicitly (optional → not exercised when absent). + Capability("filesystem.read", target_arg="assets_dir", optional=True), + Capability("filesystem.write", target_arg="output_pdf", optional=True), + ], + description=( + "Compile a LaTeX source file to PDF using a host TeX engine (latexmk or " + "pdflatex). Runs on the host: latexmk with -norc (no .latexmkrc) and " + "pdflatex with -no-shell-escape, so it does not execute arbitrary host " + "code by default. Returns the PDF path plus any compiler errors. " + "Requires TeX Live/MacTeX on PATH." + ), +) +def compile_document( + source_path: str, + output_pdf: str | None = None, + assets_dir: str | None = None, + engine: str | None = None, + timeout_s: int = 120, +) -> dict[str, Any]: + """Compile a LaTeX file to PDF (guarded subprocess; host TeX engine). + + Args: + source_path: Path to the .tex file to compile. + output_pdf: If set, the produced PDF is copied here (parents created); + build intermediates are isolated in a private temp dir. + assets_dir: Directory prepended to TEXINPUTS so figures/resources resolve + by bare name (e.g. an artifacts dir). + engine: Force an engine ("latexmk"/"pdflatex"); default auto-detects + (latexmk preferred). + timeout_s: Wall-clock timeout; the process is killed on expiry. + + Returns: + dict with success, pdf_path, engine, log_tail, errors, duration_ms, and + (on setup/timeout failure) error. + """ + src = Path(source_path) + if not src.is_file(): + return { + "success": False, "error": f"Source not found: {source_path}", + "pdf_path": None, "engine": None, "log_tail": "", "errors": [], + "duration_ms": 0, + } + + if assets_dir and os.pathsep in assets_dir: + # A path-list separator would turn one authorized filesystem.read target + # into several TEXINPUTS search roots (e.g. "assets:/etc" also reads /etc). + return { + "success": False, + "error": f"assets_dir must be a single path (no {os.pathsep!r}): {assets_dir}", + "pdf_path": None, "engine": None, "log_tail": "", "errors": [], + "duration_ms": 0, + } + + chosen = _detect_engine(engine) + if chosen is None: + looked = engine or "/".join(_ENGINES) + return { + "success": False, + "error": ( + f"No LaTeX engine on PATH (looked for {looked}). " + "Install TeX Live or MacTeX." + ), + "pdf_path": None, "engine": None, "log_tail": "", "errors": [], + "duration_ms": 0, + } + + env = _build_env(assets_dir, source_dir=str(src.parent)) + argv = _build_argv(chosen, _safe_source_arg(src.name)) + start = time.monotonic() + + with tempfile.TemporaryDirectory(prefix="texbuild-", ignore_cleanup_errors=True) as build_dir: + build = Path(build_dir) + try: + shutil.copy2(src, build / src.name) + except OSError as exc: + return { + "success": False, "error": f"Failed to stage source: {exc}", + "pdf_path": None, "engine": chosen, "log_tail": "", "errors": [], + "duration_ms": int((time.monotonic() - start) * 1000), + } + + try: + proc = _run(argv, cwd=str(build), env=env, timeout=float(timeout_s)) + except subprocess.TimeoutExpired: + return { + "success": False, "error": f"Compilation timed out after {timeout_s}s", + "pdf_path": None, "engine": chosen, "log_tail": "", "errors": [], + "duration_ms": int((time.monotonic() - start) * 1000), + } + except OSError as exc: + return { + "success": False, "error": f"Failed to run {chosen}: {exc}", + "pdf_path": None, "engine": chosen, "log_tail": "", "errors": [], + "duration_ms": int((time.monotonic() - start) * 1000), + } + duration_ms = int((time.monotonic() - start) * 1000) + + log_path = build / (src.stem + ".log") + log_text = _read_log_tail(log_path, fallback=proc.stdout or "") + log_tail = "\n".join(log_text.splitlines()[-_LOG_TAIL_LINES:]) + produced = build / (src.stem + ".pdf") + success = proc.returncode == 0 and produced.is_file() + + if not success: + return { + "success": False, "pdf_path": None, "engine": chosen, + "log_tail": log_tail, "errors": _parse_errors(log_text), + "duration_ms": duration_ms, "error": None, + } + + dest = Path(output_pdf) if output_pdf else (src.parent / (src.stem + ".pdf")) + try: + _deliver_no_follow(produced, dest) + except OSError as exc: + return { + "success": False, + "error": f"Failed to deliver PDF to {dest}: {exc}", + "pdf_path": None, "engine": chosen, + "log_tail": log_tail, "errors": [], "duration_ms": duration_ms, + } + + return { + "success": True, "pdf_path": str(dest), "engine": chosen, + "log_tail": log_tail, "errors": [], "duration_ms": duration_ms, + "error": None, + } diff --git a/src/agentic_cli/tools/execution_tools.py b/src/agentic_cli/tools/execution_tools.py index 7d723f1..f4b64ff 100644 --- a/src/agentic_cli/tools/execution_tools.py +++ b/src/agentic_cli/tools/execution_tools.py @@ -54,6 +54,7 @@ def execute_python( enabled=True, writable_paths=getattr(settings, "os_sandbox_writable_paths", []), allow_network=getattr(settings, "os_sandbox_allow_network", False), + strict=getattr(settings, "os_sandbox_strict", False), ) executor = SafePythonExecutor( diff --git a/src/agentic_cli/tools/executor.py b/src/agentic_cli/tools/executor.py index 407c7cb..d7076b7 100644 --- a/src/agentic_cli/tools/executor.py +++ b/src/agentic_cli/tools/executor.py @@ -70,21 +70,27 @@ class SafePythonExecutor: "scipy", "sklearn", "matplotlib", + # sympy.sympify / parse_expr evaluate arbitrary expressions and can be + # coerced into Python code execution — a real OS sandbox must contain it. + "sympy", } # Modules safe to import without OS isolation: pure computation and stdlib # data handling with no file/network/deserialization surface. + # + # NOTE: ``operator`` and ``functools`` are deliberately excluded. Their + # attrgetter/methodcaller (and reduce) take attribute names as *runtime + # strings*, which bypasses the AST underscore-attribute filter and yields + # host RCE (object.__subclasses__ → __init__.__globals__ → os.system). + # Do not re-add them without the OS sandbox as the enforced boundary. CORE_MODULES = { # Math / science - "sympy", "math", "statistics", "cmath", # Collections and utilities "collections", "itertools", - "functools", - "operator", # Data handling "json", "re", @@ -156,12 +162,22 @@ def __init__( self.os_sandbox_policy = os_sandbox_policy # File/network/pickle-capable libraries are only safe behind real OS - # isolation. Without it, restrict imports to the pure-computation core. + # isolation. They are gated on *actual* sandbox availability, not just + # the setting: if the sandbox is requested but no backend exists (so + # execution falls back in-process), these must stay unavailable. sandbox_on = bool(os_sandbox_policy and os_sandbox_policy.enabled) + real_sandbox = sandbox_on and self._real_os_sandbox_available() self.effective_allowed_modules = ( - self.ALLOWED_MODULES if sandbox_on else self.CORE_MODULES + self.ALLOWED_MODULES if real_sandbox else self.CORE_MODULES ) + @staticmethod + def _real_os_sandbox_available() -> bool: + """True if a real OS sandbox backend (not the no-op) is present.""" + from agentic_cli.tools.shell.os_sandbox import get_os_sandbox + + return get_os_sandbox().sandbox_type != "none" + def validate_code(self, code: str) -> tuple[bool, str]: """Validate code for safety. @@ -207,6 +223,17 @@ def validate_code(self, code: str) -> tuple[bool, str]: f"Access to private attribute '{node.attr}' is not allowed", ) + # Reject str.format()/format_map() field names that traverse private + # attributes (e.g. "{0.__class__}"). The attribute path lives inside + # the string literal, so the ast.Attribute check above never sees it. + if isinstance(node, ast.Constant) and isinstance(node.value, str): + if self._format_field_accesses_private(node.value): + return ( + False, + "Format string references a private attribute " + "(potential sandbox escape via str.format)", + ) + # Check for dangerous function calls if isinstance(node, ast.Call): if isinstance(node.func, ast.Name): @@ -215,6 +242,47 @@ def validate_code(self, code: str) -> tuple[bool, str]: return True, "" + @staticmethod + def _format_field_accesses_private(format_string: str) -> bool: + """True if a ``str.format`` field references a ``_``-prefixed attribute. + + e.g. ``"{0.__class__}"`` or ``"{a.__class__.__base__}"``. Such fields + traverse the object graph through the format machinery, bypassing the + AST attribute filter. Item access (``{0[key]}``) is not an + attribute-traversal vector and is left alone. + """ + import string as _string + + try: + parsed = list(_string.Formatter().parse(format_string)) + except (ValueError, IndexError): + # Malformed format string: not our concern here — it will raise at + # runtime inside the subprocess, it cannot escape. + return False + for _literal, field_name, format_spec, _conv in parsed: + if field_name: + for part in field_name.replace("[", ".").split("."): + if part.startswith("_"): + return True + if format_spec and "{" in format_spec: + # Nested replacement field inside the format spec. + if SafePythonExecutor._format_field_accesses_private(format_spec): + return True + return False + + @staticmethod + def _subprocess_env() -> dict[str, str]: + """Parent environment with provider secrets removed. + + The workflow manager exports API keys into ``os.environ``; a code + execution escape must not be able to read them from the child process. + """ + import os as _os + import re as _re + + secret = _re.compile(r"(API_KEY|TOKEN|SECRET|PASSWORD)$", _re.IGNORECASE) + return {k: v for k, v in _os.environ.items() if not secret.search(k)} + def execute( self, code: str, @@ -247,6 +315,18 @@ def execute( return self._execute_in_subprocess(code, context, timeout) + def _sandbox_refused(self, message: str, start_time: float) -> dict[str, Any]: + """Fail-closed result when strict sandboxing is required but unavailable.""" + logger.error("python_executor.os_sandbox_refused", message=message) + elapsed = (time.time() - start_time) * 1000 + return { + "success": False, + "output": "", + "result": None, + "error": message, + "execution_time_ms": round(elapsed, 2), + } + def _execute_in_subprocess( self, code: str, @@ -272,75 +352,73 @@ def _execute_in_subprocess( max_memory_mb=self.max_memory_mb, ) + strict = bool( + self.os_sandbox_policy and getattr(self.os_sandbox_policy, "strict", False) + ) + proc = None try: if self.os_sandbox_policy and self.os_sandbox_policy.enabled: from agentic_cli.tools.shell.os_sandbox import get_os_sandbox sandbox = get_os_sandbox() - # Fail closed when sandboxing was requested but no real - # isolation is available — silently dropping back to a plain - # subprocess would defeat the user's opt-in. if sandbox.sandbox_type == "none": - logger.error( - "python_executor.os_sandbox_required_but_unavailable", + # Requested but no backend. Strict -> refuse; otherwise fall + # back to the in-process executor (already restricted to the + # pure-computation CORE_MODULES) with a warning. + if strict: + return self._sandbox_refused( + "OS sandbox is required (os_sandbox_strict=True) but " + "no supported sandbox tool is available (install " + "sandbox-exec on macOS or bwrap on Linux). Refusing " + "to execute without isolation.", + start_time, + ) + logger.warning( + "python_executor.os_sandbox_unavailable_fallback", sandbox_type=sandbox.sandbox_type, ) - elapsed = (time.time() - start_time) * 1000 - return { - "success": False, - "output": "", - "result": None, - "error": ( - "OS sandbox is required (os_sandbox_enabled=True) " - "but no supported sandbox tool is available on " - "this system (install sandbox-exec on macOS or " - "bwrap on Linux). Refusing to execute without " - "isolation." - ), - "execution_time_ms": round(elapsed, 2), - } - wrap_result = sandbox.wrap_python_command( - [sys.executable, "-c", script], - Path.cwd(), - self.os_sandbox_policy, - ) - if wrap_result.success: - logger.debug( - "python_executor.os_sandbox_wrapped", - sandbox_type=wrap_result.sandbox_type, - ) - proc = subprocess.run( - wrap_result.command, - shell=True, - input=code, - capture_output=True, - text=True, - timeout=timeout, - ) else: - logger.error( - "python_executor.os_sandbox_wrap_failed", - error=wrap_result.error, + wrap_result = sandbox.wrap_python_command( + [sys.executable, "-c", script], + Path.cwd(), + self.os_sandbox_policy, ) - elapsed = (time.time() - start_time) * 1000 - return { - "success": False, - "output": "", - "result": None, - "error": ( + if wrap_result.success: + logger.debug( + "python_executor.os_sandbox_wrapped", + sandbox_type=wrap_result.sandbox_type, + ) + proc = subprocess.run( + wrap_result.command, + shell=True, + input=code, + capture_output=True, + text=True, + timeout=timeout, + env=self._subprocess_env(), + ) + elif strict: + return self._sandbox_refused( f"OS sandbox wrap failed " f"({wrap_result.sandbox_type}): {wrap_result.error}. " - f"Refusing to execute without isolation." - ), - "execution_time_ms": round(elapsed, 2), - } - else: + f"Refusing to execute without isolation.", + start_time, + ) + else: + logger.warning( + "python_executor.os_sandbox_wrap_failed_fallback", + error=wrap_result.error, + ) + + if proc is None: + # No sandbox policy, or a non-strict fallback from above. proc = subprocess.run( [sys.executable, "-c", script], input=code, capture_output=True, text=True, timeout=timeout, + env=self._subprocess_env(), ) except subprocess.TimeoutExpired: elapsed = (time.time() - start_time) * 1000 diff --git a/src/agentic_cli/tools/factories.py b/src/agentic_cli/tools/factories.py index bc232b5..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,18 +374,23 @@ 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] # --------------------------------------------------------------------------- # Sandbox tool # --------------------------------------------------------------------------- -def make_sandbox_tool(sandbox_manager) -> Callable: +def make_sandbox_tool(sandbox_manager, workflow_manager=None) -> Callable: """Create sandbox_execute bound to a SandboxManager. Args: sandbox_manager: SandboxManager instance. + workflow_manager: The owning workflow manager, used to namespace the + default session to the active conversation (so distinct + conversations don't share one kernel/workspace). Returns: sandbox_execute function. @@ -362,6 +400,7 @@ def sandbox_execute( code: str, session_id: str = "default", timeout_seconds: int = 120, + inputs: list[str] | None = None, ) -> dict[str, Any]: """Execute Python code in a stateful sandbox. @@ -369,14 +408,35 @@ def sandbox_execute( code: Python code to execute. session_id: Session identifier for state persistence. timeout_seconds: Maximum execution time in seconds. + inputs: Optional list of host file paths to stage into + inputs/ inside the session before execution. Returns: Dictionary with execution results. """ + # Opt-in gate — the workflow binds THIS tool (base_manager), so the gate + # must live here, not only on the module-level tool. Without it the + # stateful_executor_backend setting is inert in the real path. + from agentic_cli.config import get_settings + from agentic_cli.tools.sandbox.manager import sandbox_disabled_reason + + if getattr(get_settings(), "stateful_executor_backend", "none") == "none": + return {"success": False, "error": sandbox_disabled_reason(get_settings())} + + # Namespace the default session to the active conversation so distinct + # conversations don't share one kernel/workspace. An explicit session_id + # is honored as-is (lets the model keep intentional sub-sessions). + sid = session_id + if sid == "default" and workflow_manager is not None: + active = getattr(workflow_manager, "active_session_id", None) + if active: + sid = f"conv-{active}" + result = sandbox_manager.execute( code=code, - session_id=session_id, + session_id=sid, timeout_seconds=timeout_seconds, + inputs=inputs, ) return { "success": result.success, @@ -389,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] # --------------------------------------------------------------------------- @@ -462,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: @@ -501,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] # --------------------------------------------------------------------------- @@ -564,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/glob_tool.py b/src/agentic_cli/tools/glob_tool.py index 057e03e..784acc0 100644 --- a/src/agentic_cli/tools/glob_tool.py +++ b/src/agentic_cli/tools/glob_tool.py @@ -9,12 +9,15 @@ from pathlib import Path from typing import Any, Literal +from agentic_cli.file_utils import glob_pattern_escapes_root, path_is_within from agentic_cli.tools.registry import ( ToolCategory, register_tool, ) from agentic_cli.workflow.permissions import Capability +_MAX_SCAN = 10_000 # hard ceiling on matches materialized before sorting/limiting + @register_tool( category=ToolCategory.READ, @@ -72,14 +75,38 @@ def glob( "path": str(search_path), } - # Find matching files - matches = list(search_path.glob(pattern)) + # The permission engine authorizes only `path`; a pattern like "../*" or an + # absolute pattern would escape that authorized root, so reject it. + if glob_pattern_escapes_root(pattern): + return { + "success": False, + "error": f"Pattern escapes the search root: {pattern!r}", + "path": str(search_path), + } + + # Find matching files, capping how many we materialize (a pathological + # pattern like "**/*" over a huge tree must not exhaust memory). + matches = [] + scan_truncated = False + for p in search_path.glob(pattern): + matches.append(p) + if len(matches) >= _MAX_SCAN: + scan_truncated = True + break # Filter results filtered = [] for match in matches: - # Skip hidden files if not requested - if not include_hidden and match.name.startswith("."): + # Drop anything resolving outside the authorized root (e.g. a symlink + # inside `path` that points elsewhere) — defense in depth. + if not path_is_within(match, search_path): + continue + + # Skip results with any hidden component, not just a hidden basename — + # a pattern like "**/*" otherwise leaks files under a dot-directory. + if not include_hidden and any( + part.startswith(".") for part in match.relative_to(search_path).parts + ): continue # Skip directories if not requested @@ -97,7 +124,7 @@ def glob( filtered.sort(key=lambda p: p.stat().st_mtime, reverse=True) # Truncate if needed - truncated = len(filtered) > max_results + truncated = scan_truncated or len(filtered) > max_results filtered = filtered[:max_results] # Format output diff --git a/src/agentic_cli/tools/grep_tool.py b/src/agentic_cli/tools/grep_tool.py index 81e93c1..eb48448 100644 --- a/src/agentic_cli/tools/grep_tool.py +++ b/src/agentic_cli/tools/grep_tool.py @@ -5,17 +5,25 @@ """ import functools -import re +import os +import signal import subprocess +import tempfile +import re from pathlib import Path from typing import Any, Literal +from agentic_cli.file_utils import glob_pattern_escapes_root, path_is_within from agentic_cli.tools.registry import ( ToolCategory, register_tool, ) from agentic_cli.workflow.permissions import Capability +_MAX_FILES = 10_000 # cap the number of files the Python fallback scans +_MAX_FILE_BYTES = 5_000_000 # skip files larger than this (avoid reading whole huge files) +_MAX_RG_OUTPUT_BYTES = 10_000_000 # cap ripgrep JSON we read into memory + @register_tool( category=ToolCategory.READ, @@ -70,6 +78,18 @@ def grep( "path": str(search_path), } + # The permission engine authorizes only `path`; a file_pattern like "../*" + # or an absolute pattern would escape that authorized root, so reject it. + if file_pattern and glob_pattern_escapes_root(file_pattern): + return { + "success": False, + "error": f"File pattern escapes the search root: {file_pattern!r}", + "matches": [], + "total_matches": 0, + "files_searched": 0, + "truncated": False, + } + # Try to use ripgrep if available (faster, respects .gitignore) if _ripgrep_available(): return _grep_with_ripgrep( @@ -145,22 +165,35 @@ def _grep_with_ripgrep( cmd.append("--") cmd.append(str(path)) + # Drop RIPGREP_CONFIG_PATH so a host config can't inject flags (e.g. + # --follow, which would make rg traverse symlinks out of the authorized + # root). Containment below is the backstop; this removes the vector. + rg_env = {k: v for k, v in os.environ.items() if k != "RIPGREP_CONFIG_PATH"} + # Capture rg output to a temp file and read back at most _MAX_RG_OUTPUT_BYTES + # so a large tree can't allocate unbounded JSON in host memory. + rg_truncated = False try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=30, - ) - except subprocess.TimeoutExpired: - return { - "success": False, - "error": "Search timed out after 30 seconds", - "matches": [], - "total_matches": 0, - "files_searched": 0, - "truncated": False, - } + with tempfile.TemporaryFile(mode="w+b") as out_f: + proc = subprocess.Popen( + cmd, stdout=out_f, stderr=subprocess.DEVNULL, + env=rg_env, start_new_session=True, + ) + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + proc.wait() + return { + "success": False, + "error": "Search timed out after 30 seconds", + "matches": [], + "total_matches": 0, + "files_searched": 0, + "truncated": False, + } + out_f.seek(0) + raw = out_f.read(_MAX_RG_OUTPUT_BYTES) + rg_truncated = bool(out_f.read(1)) # more output than the cap remained except FileNotFoundError: # Ripgrep not found, fall back to Python return _grep_python( @@ -175,6 +208,8 @@ def _grep_with_ripgrep( output_mode=output_mode, ) + stdout_text = raw.decode("utf-8", errors="replace") + # Parse ripgrep JSON output import json @@ -183,7 +218,7 @@ def _grep_with_ripgrep( file_counts: dict[str, int] = {} total_matches = 0 - for line in result.stdout.strip().split("\n"): + for line in stdout_text.strip().split("\n"): if not line: continue try: @@ -194,6 +229,14 @@ def _grep_with_ripgrep( if data.get("type") == "match": match_data = data.get("data", {}) file_path = match_data.get("path", {}).get("text", "") + # Drop files resolving outside the authorized root (e.g. a symlink + # rg followed) — parity with the Python fallback's containment. rg + # may report a path relative to the search root, so resolve it there. + if file_path: + fp = Path(file_path) + abs_fp = fp if fp.is_absolute() else (path / fp) + if not path_is_within(abs_fp, path): + continue files_searched.add(file_path) file_counts[file_path] = file_counts.get(file_path, 0) + 1 total_matches += 1 @@ -219,7 +262,7 @@ def _grep_with_ripgrep( "matches": matches, "total_matches": total_matches, "files_searched": len(files_searched), - "truncated": len(matches) >= max_results, + "truncated": rg_truncated or len(matches) >= max_results, } @@ -257,25 +300,42 @@ def _grep_python( total_matches = 0 file_counts: dict[str, int] = {} - # Get files to search + # Get files to search, capping how many we materialize. if path.is_file(): - files = [path] + candidates = iter([path]) + elif file_pattern: + candidates = path.rglob(file_pattern) if recursive else path.glob(file_pattern) else: - if file_pattern: - if recursive: - files = list(path.rglob(file_pattern)) - else: - files = list(path.glob(file_pattern)) - else: - if recursive: - files = [f for f in path.rglob("*") if f.is_file()] - else: - files = [f for f in path.iterdir() if f.is_file()] + candidates = path.rglob("*") if recursive else path.iterdir() + + files = [] + scan_truncated = False + for f in candidates: + # Count only files against the budget — directories must not exhaust it + # (otherwise dirs before the files silently drop matches). + if not f.is_file(): + continue + files.append(f) + if len(files) >= _MAX_FILES: + scan_truncated = True + break for file_path in files: if not file_path.is_file(): continue + # Skip files resolving outside the authorized root (e.g. a symlink + # under `path` pointing elsewhere) — defense in depth. + if not path_is_within(file_path, path): + continue + + # Skip files that are too large to read whole (reliability bound). + try: + if file_path.stat().st_size > _MAX_FILE_BYTES: + continue + except OSError: + continue + try: content = file_path.read_text() lines = content.splitlines() @@ -325,5 +385,5 @@ def _grep_python( "matches": matches, "total_matches": total_matches, "files_searched": files_searched, - "truncated": total_matches > max_results, + "truncated": scan_truncated or total_matches > max_results, } diff --git a/src/agentic_cli/tools/jobs/__init__.py b/src/agentic_cli/tools/jobs/__init__.py new file mode 100644 index 0000000..09943c8 --- /dev/null +++ b/src/agentic_cli/tools/jobs/__init__.py @@ -0,0 +1,45 @@ +"""Long-running job substrate (Tier A §3.2). + +`JobManager` + pluggable execution backends behind one interface, plus the +generic observe-only ``job_*`` tools. Typed long-running tools (the ones that +actually start work) are application-provided — see ``run_shell_job`` in +``examples/jobs_demo.py``. See ``docs/plans/2026-06-14-job-control-design.md``. +""" + +from agentic_cli.tools.jobs.backends import ( + InProcessBackend, + JobBackend, + JobState, + SubprocessBackend, + default_backends, +) +from agentic_cli.tools.jobs.manager import ( + JobManager, + JobRecord, + ResumeState, + ResumeStateError, +) +from agentic_cli.tools.jobs.tools import ( + job_cancel, + job_list, + job_logs, + job_result, + job_status, +) + +__all__ = [ + "JobManager", + "JobRecord", + "JobState", + "ResumeState", + "ResumeStateError", + "JobBackend", + "SubprocessBackend", + "InProcessBackend", + "default_backends", + "job_status", + "job_result", + "job_logs", + "job_cancel", + "job_list", +] diff --git a/src/agentic_cli/tools/jobs/backends.py b/src/agentic_cli/tools/jobs/backends.py new file mode 100644 index 0000000..da8cc2b --- /dev/null +++ b/src/agentic_cli/tools/jobs/backends.py @@ -0,0 +1,306 @@ +"""Execution backends for long-running jobs. + +A backend is *where the work runs* (≠ the ADK/LangGraph orchestrator backend). +All backends speak the same :class:`JobBackend` interface to ``JobManager``, so +job lifecycle, persistence, monitoring, and the ``job_*`` tools are uniform +regardless of where a job actually executes. + +Milestone 1 ships two backends: + +- :class:`SubprocessBackend` — detached OS process (``start_new_session=True``), + output streamed to log files, completion recorded via an on-disk ``exit_code`` + sentinel so it survives a CLI restart. +- :class:`InProcessBackend` — a Python callable run on a thread pool. Does *not* + survive a restart and cannot be force-killed (thread); use for lighter work + that should not block the agent turn. +""" + +from __future__ import annotations + +import json +import os +import shlex +import signal +import subprocess +from abc import ABC, abstractmethod +from concurrent.futures import Future, ThreadPoolExecutor +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from agentic_cli.tools.jobs.manager import JobRecord + + +class JobState(str, Enum): + """Lifecycle state of a job.""" + + QUEUED = "queued" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + UNKNOWN = "unknown" + + +TERMINAL_STATES = frozenset( + {JobState.SUCCEEDED, JobState.FAILED, JobState.CANCELLED, JobState.UNKNOWN} +) + + +def _pid_alive(pid: int) -> bool: + """True if a process with ``pid`` exists (signal 0 probe).""" + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True # exists but owned by another user + return True + + +def _read_exit_code(exit_file: Path) -> int | None: + """Read the sentinel exit-code file, or None if absent/unreadable.""" + if not exit_file.exists(): + return None + try: + text = exit_file.read_text().strip() + return int(text) if text else 1 + except (ValueError, OSError): + return 1 + + +def _tail(path: Path, last_n: int) -> list[str]: + """Return the last ``last_n`` lines of a text file (empty if missing).""" + if not path.exists(): + return [] + try: + lines = path.read_text(errors="replace").splitlines() + except OSError: + return [] + return lines[-last_n:] if last_n > 0 else lines + + +class JobBackend(ABC): + """Uniform interface every execution backend implements for ``JobManager``.""" + + 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: + """Begin executing the job. Non-blocking. Sets handle/pid on ``record``.""" + + @abstractmethod + def poll(self, record: "JobRecord", job_dir: Path) -> JobState: + """Return the job's current state (reads sentinel / liveness / remote).""" + + @abstractmethod + def cancel(self, record: "JobRecord", job_dir: Path) -> None: + """Best-effort cancellation of a running job.""" + + def logs( + self, record: "JobRecord", job_dir: Path, last_n: int, stream: str + ) -> list[str]: + """Return the last ``last_n`` lines of stdout/stderr.""" + fname = "stderr.log" if stream == "stderr" else "stdout.log" + return _tail(job_dir / fname, last_n) + + 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.""" + + 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. + self._procs: dict[str, subprocess.Popen] = {} + + def start(self, record: "JobRecord", job_dir: Path) -> None: + command = record.spec.get("command") + if not command: + raise ValueError("subprocess job requires spec['command']") + cwd = record.spec.get("cwd") or None + + inner = ( + command + if isinstance(command, str) + else " ".join(shlex.quote(c) for c in command) + ) + exit_file = job_dir / "exit_code" + # Run the command in a subshell, then record its exit status on disk so + # completion is detectable after a CLI restart (we can't waitpid a + # process we didn't start this session). The subshell ensures a command + # that calls ``exit`` doesn't skip the sentinel write. + wrapped = ( + f"( {inner} )\nrc=$?\nprintf '%s' \"$rc\" > {shlex.quote(str(exit_file))}\n" + ) + + out = open(job_dir / "stdout.log", "wb") + err = open(job_dir / "stderr.log", "wb") + try: + proc = subprocess.Popen( + ["sh", "-c", wrapped], + cwd=cwd, + stdout=out, + stderr=err, + start_new_session=True, + ) + finally: + out.close() + err.close() + self._procs[record.job_id] = proc + record.pid = proc.pid + record.backend_handle = str(proc.pid) + + def poll(self, record: "JobRecord", job_dir: Path) -> JobState: + exit_file = job_dir / "exit_code" + code = _read_exit_code(exit_file) + if code is not None: + record.exit_code = code + self._procs.pop(record.job_id, None) + return JobState.SUCCEEDED if code == 0 else JobState.FAILED + + # If we still hold the live process handle (same session), reap it via + # poll() — this both prevents zombies (which os.kill(pid,0) would report + # as alive) and gives a returncode fallback if the sentinel is missing. + proc = self._procs.get(record.job_id) + if proc is not None: + rc = proc.poll() + if rc is None: + return JobState.RUNNING + code = _read_exit_code(exit_file) # prefer the sentinel + if code is None: + code = rc + record.exit_code = code + self._procs.pop(record.job_id, None) + return JobState.SUCCEEDED if code == 0 else JobState.FAILED + + # No live handle (e.g. after a CLI restart): fall back to PID liveness. + if record.pid and _pid_alive(record.pid): + return JobState.RUNNING + return JobState.UNKNOWN + + def cancel(self, record: "JobRecord", job_dir: Path) -> None: + pid = record.pid + if not pid: + return + try: + os.killpg(os.getpgid(pid), signal.SIGTERM) + except (ProcessLookupError, PermissionError): + pass + self._procs.pop(record.job_id, None) + + +class InProcessBackend(JobBackend): + """Run a Python callable on a thread pool. Not restart-safe; no force-kill.""" + + 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( + max_workers=max_workers, thread_name_prefix="job-inproc" + ) + self._futures: dict[str, Future] = {} + + def start(self, record: "JobRecord", job_dir: Path) -> None: + target = record.spec.get("target") + if not callable(target): + raise ValueError("inprocess job requires a callable spec['target']") + args = record.spec.get("args", ()) + kwargs = record.spec.get("kwargs", {}) + exit_file = job_dir / "exit_code" + result_file = job_dir / "result.json" + + def _run() -> None: + try: + res = target(*args, **kwargs) + try: + result_file.write_text(json.dumps(res, default=str)) + except (TypeError, ValueError, OSError): + pass + exit_file.write_text("0") + except BaseException as exc: # noqa: BLE001 - record any failure + try: + (job_dir / "stderr.log").write_text(repr(exc)) + except OSError: + pass + exit_file.write_text("1") + + self._futures[record.job_id] = self._pool.submit(_run) + record.backend_handle = "thread" + + def poll(self, record: "JobRecord", job_dir: Path) -> JobState: + code = _read_exit_code(job_dir / "exit_code") + if code is not None: + record.exit_code = code + return JobState.SUCCEEDED if code == 0 else JobState.FAILED + fut = self._futures.get(record.job_id) + if fut is None: + return JobState.UNKNOWN # lost across restart + return JobState.RUNNING # done-but-no-sentinel race resolves next poll + + def cancel(self, record: "JobRecord", job_dir: Path) -> None: + fut = self._futures.get(record.job_id) + 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(): + try: + return json.loads(result_file.read_text()) + except (ValueError, OSError): + return None + return None + + +def default_backends() -> dict[str, JobBackend]: + """The backends available out of the box (milestone 1).""" + return {"subprocess": SubprocessBackend(), "inprocess": InProcessBackend()} diff --git a/src/agentic_cli/tools/jobs/manager.py b/src/agentic_cli/tools/jobs/manager.py new file mode 100644 index 0000000..d482236 --- /dev/null +++ b/src/agentic_cli/tools/jobs/manager.py @@ -0,0 +1,1244 @@ +"""JobManager — lifecycle, persistence, and concurrency for long-running jobs. + +`JobManager` is **internal infrastructure, never an LLM-facing tool** (it is to +long-running tools what `SandboxManager` is to `sandbox_execute`). Typed +long-running tools call it from their bodies; the LLM only ever sees the tool. + +Responsibilities: own `JobRecord`s, persist them under a base dir so jobs survive +turns and CLI restarts, reconcile state on read via the execution backends, +enforce a concurrency cap with a queue, and expose a uniform query/manage API. +""" + +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, Iterator + +from agentic_cli.file_utils import atomic_write_json +from agentic_cli.logging import Loggers +from agentic_cli.tools.jobs.backends import ( + TERMINAL_STATES, + JobBackend, + JobState, + default_backends, +) + +if TYPE_CHECKING: + from agentic_cli.config import BaseSettings + +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``. + + ``spec`` may hold live (non-serializable) objects in memory — e.g. an + in-process callable; only JSON-safe entries are written to disk. + """ + + job_id: str + tool: str + backend: str + name: str + state: JobState + spec: dict = field(default_factory=dict) + backend_handle: str | None = None + pid: int | None = None + exit_code: int | None = None + tags: list[str] = field(default_factory=list) + submitted_at: float = field(default_factory=_now) + started_at: float | None = None + finished_at: float | None = None + error: str | None = None + # --- Resume association (phase 2 push/resume; unset = fire-and-forget) --- + session_id: str | None = None # conversation that launched the job + user_id: str | None = None + 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 + 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 + end = self.finished_at or _now() + return round(max(0.0, end - start), 1) + + def to_dict(self) -> dict: + d = asdict(self) + d["state"] = self.state.value + d["spec"] = _json_safe_spec(self.spec) + return d + + @classmethod + def from_dict(cls, d: dict) -> "JobRecord": + d = dict(d) + d["state"] = JobState(d["state"]) + # 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.""" + out = { + "job_id": self.job_id, + "tool": self.tool, + "name": self.name, + "backend": self.backend, + "state": self.state.value, + "elapsed_s": self.elapsed_s(), + "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: + """Keep only JSON-serializable spec entries (drop live callables, etc.).""" + import json + + safe: dict = {} + for k, v in spec.items(): + try: + json.dumps(v) + safe[k] = v + except (TypeError, ValueError): + safe[k] = f"" + return safe + + +class JobManager: + """Lifecycle + persistence + concurrency for long-running jobs.""" + + def __init__( + self, + settings: "BaseSettings | None" = None, + *, + base_dir: Path, + max_concurrent: int = 4, + backends: dict[str, JobBackend] | None = None, + ) -> None: + self._settings = settings + self.base_dir = Path(base_dir) + self.base_dir.mkdir(parents=True, exist_ok=True) + 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() + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def register_backend(self, backend: JobBackend) -> None: + self._backends[backend.name] = backend + + def submit( + self, + *, + tool: str, + backend: str, + spec: dict, + name: str | None = None, + tags: list[str] | None = None, + resume_on_complete: bool = False, + call_id: str | None = None, + call_name: str | None = None, + session_id: str | None = None, + user_id: str | None = None, + ) -> JobRecord: + """Create a job; start it now if under the cap, else queue it. + + When ``resume_on_complete`` is set the job is flagged so the harness can + auto-resume the agent with the result once it finishes (phase 2). The + originating ``session_id``/``user_id`` are auto-filled from the active + workflow turn when not supplied; ``call_id``/``call_name`` (which only + the calling tool knows) identify the pending tool call to answer. + """ + if backend not in self._backends: + raise ValueError( + f"unknown job backend {backend!r}; have {sorted(self._backends)}" + ) + if resume_on_complete and (session_id is None or user_id is None): + session_id, user_id = self._fill_turn_context(session_id, user_id) + with self._lock: + job_id = uuid.uuid4().hex[:12] + rec = JobRecord( + job_id=job_id, + tool=tool, + backend=backend, + name=name or tool, + state=JobState.QUEUED, + spec=dict(spec), + tags=list(tags or []), + resume_on_complete=resume_on_complete, + call_id=call_id, + call_name=call_name or (tool if resume_on_complete else None), + session_id=session_id, + user_id=user_id, + ) + self._records[job_id] = rec + self._job_dir(job_id).mkdir(parents=True, exist_ok=True) + self._persist(rec) + self._maybe_start_queued() + return rec + + @staticmethod + def _fill_turn_context( + session_id: str | None, user_id: str | None + ) -> tuple[str | None, str | None]: + """Best-effort fill session/user from the active workflow turn. + + Reads the ``WORKFLOW`` service from the registry (live during a tool + call). Missing context is non-fatal: the job still runs, but it won't be + resumable, so we warn. + """ + from agentic_cli.workflow.service_registry import WORKFLOW, get_service + + wf = get_service(WORKFLOW) + if wf is not None: + session_id = session_id or getattr(wf, "active_session_id", None) + user_id = user_id or getattr(wf, "active_user_id", None) + if session_id is None or user_id is None: + logger.warning( + "job_resume_context_missing", + have_session=session_id is not None, + have_user=user_id is not None, + ) + 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 + 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( + self, + *, + state: JobState | None = None, + tag: str | None = None, + active_only: bool = False, + ) -> list[JobRecord]: + with self._lock: + self.reconcile() + recs = list(self._records.values()) + if active_only: + recs = [r for r in recs if r.state in (JobState.QUEUED, JobState.RUNNING)] + if state is not None: + recs = [r for r in recs if r.state == state] + if tag is not None: + recs = [r for r in recs if tag in r.tags] + return sorted(recs, key=lambda r: r.submitted_at, reverse=True) + + def tail(self, job_id: str, n: int = 50, stream: str = "stdout") -> list[str]: + with self._lock: + rec = self._records.get(job_id) + if rec is None: + return [] + return self._backends[rec.backend].logs(rec, self._job_dir(job_id), n, stream) + + def result(self, job_id: str) -> Any: + with self._lock: + rec = self._records.get(job_id) + if rec is None: + return None + self._refresh(rec) + if rec.state not in TERMINAL_STATES: + return None + 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 + 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 + + def reconcile(self) -> None: + """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: + 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. + + 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() + 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 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. + """ + with self._lock: + self.reconcile() + recs = [ + r + for r in self._records.values() + 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 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 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 + # ------------------------------------------------------------------ + + def _job_dir(self, job_id: str) -> Path: + return self.base_dir / job_id + + @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. + + 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 status + backend = self._backends.get(rec.backend) + if backend is None: + rec.state = JobState.UNKNOWN + self._persist(rec) + return status + new_state = backend.poll(rec, self._job_dir(rec.job_id)) + 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).""" + if self.running_count() >= self._max_concurrent: + return + queued = sorted( + (r for r in self._records.values() if r.state == JobState.QUEUED), + key=lambda r: r.submitted_at, + ) + for rec in queued: + if self.running_count() >= self._max_concurrent: + break + 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)) + rec.state = JobState.RUNNING + rec.started_at = _now() + except Exception as exc: # noqa: BLE001 - surface launch failures as failed jobs + rec.state = JobState.FAILED + 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, 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.""" + import json + + if not self.base_dir.exists(): + return + with self._lock: + for meta in self.base_dir.glob("*/meta.json"): + try: + rec = JobRecord.from_dict(json.loads(meta.read_text())) + except (ValueError, OSError, TypeError, KeyError): + continue + # In-memory handles (Popen / Future) are gone after a restart, + # 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/src/agentic_cli/tools/jobs/tools.py b/src/agentic_cli/tools/jobs/tools.py new file mode 100644 index 0000000..9d2a84f --- /dev/null +++ b/src/agentic_cli/tools/jobs/tools.py @@ -0,0 +1,143 @@ +"""Generic, observe-only job tools. + +The framework ships only the **observe-only** side of the job substrate here +(capability ``jobs.manage``): tools that read/manage existing jobs but never +*start* one. To keep the agent's tool surface small, ``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_result``/``job_logs``/``job_cancel``/``job_list`` remain available as +opt-in extras (and power the ``/jobs`` command), but are not in the default +bundle. + +The other half — **typed long-running tools** that actually start work and +return a ``job_id`` (declaring ``long_running=True`` and a +``longrunning.`` capability) — is intentionally **not** shipped by +the framework. Such a tool decides *what* runs and *how* (which execution +backend, and any safety checks on its input), so it belongs to the application. +A subprocess-backed ``run_shell_job`` lives in ``examples/jobs_demo.py`` as a +reference: note it runs the command via ``sh -c`` and does **not** go through +the hardened shell tool (``tools/shell/``), so it is a demo, not a built-in. + +``JobManager`` itself is never exposed to the LLM — tools reach it via the +service registry. +""" + +from __future__ import annotations + +from agentic_cli.tools.registry import ToolCategory, register_tool +from agentic_cli.workflow.permissions import Capability +from agentic_cli.workflow.service_registry import JOB_MANAGER, get_service + + +def _manager(): + """Return the JobManager, or an error dict if it isn't available.""" + jm = get_service(JOB_MANAGER) + if jm is None: + return {"success": False, "error": "job manager not available"} + return jm + + +# --------------------------------------------------------------------------- +# Observe-only management tools +# --------------------------------------------------------------------------- + + +@register_tool( + 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. + + Returns state, elapsed time, exit code, and a short stdout tail; once the + job is finished it also includes ``result``. This is the only job tool most + agents need alongside the long-running tool that started the job. + """ + from agentic_cli.tools.jobs.backends import TERMINAL_STATES + + jm = _manager() + if isinstance(jm, dict): + return jm + rec = jm.get(job_id) + if rec is None: + return {"success": False, "error": f"no such job: {job_id}"} + out = rec.summary() + out["success"] = True + out["stdout_tail"] = jm.tail(job_id, 10, "stdout") + if rec.state in TERMINAL_STATES: + out["result"] = jm.result(job_id) + return out + + +@register_tool( + 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.""" + jm = _manager() + if isinstance(jm, dict): + return jm + rec = jm.get(job_id) + if rec is None: + return {"success": False, "error": f"no such job: {job_id}"} + from agentic_cli.tools.jobs.backends import TERMINAL_STATES + + if rec.state not in TERMINAL_STATES: + return {"success": False, "error": f"job {job_id} not finished (state={rec.state.value})"} + return {"success": True, "state": rec.state.value, "result": jm.result(job_id)} + + +@register_tool( + 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``.""" + jm = _manager() + if isinstance(jm, dict): + return jm + if jm.get(job_id) is None: + return {"success": False, "error": f"no such job: {job_id}"} + return {"success": True, "lines": jm.tail(job_id, last_n, stream)} + + +@register_tool( + 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.""" + jm = _manager() + if isinstance(jm, dict): + return jm + rec = jm.cancel(job_id) + if rec is None: + return {"success": False, "error": f"no such job: {job_id}"} + return {"success": True, "state": rec.state.value} + + +@register_tool( + 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.""" + jm = _manager() + if isinstance(jm, dict): + return jm + from agentic_cli.tools.jobs.backends import JobState + + state_filter = JobState(state) if state else None + recs = jm.list(state=state_filter, tag=tag or None) + return {"success": True, "jobs": [r.summary() for r in recs], "count": len(recs)} 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/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/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/registry.py b/src/agentic_cli/tools/registry.py index 9803200..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,21 +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.""" @@ -83,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. @@ -112,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. @@ -119,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, @@ -132,6 +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. @@ -142,28 +331,270 @@ 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) @@ -177,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) @@ -196,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, *, @@ -203,12 +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. + + ``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]: @@ -218,6 +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/sandbox/__init__.py b/src/agentic_cli/tools/sandbox/__init__.py index 994ba14..b81a0d9 100644 --- a/src/agentic_cli/tools/sandbox/__init__.py +++ b/src/agentic_cli/tools/sandbox/__init__.py @@ -6,6 +6,7 @@ from typing import Any +from agentic_cli.config import get_settings from agentic_cli.tools.registry import register_tool, ToolCategory from agentic_cli.workflow.service_registry import require_service, SANDBOX_MANAGER from agentic_cli.workflow.permissions import Capability @@ -13,20 +14,31 @@ @register_tool( category=ToolCategory.EXECUTION, - capabilities=[Capability("python.exec")], + # Distinct from execute_python's ``python.exec`` on purpose: this kernel is + # unsandboxed and stateful, so an "Allow always" for the stateless scratchpad + # must NOT silently authorize it. A deliberate ``python.*`` grant still covers + # both. + capabilities=[Capability("python.exec.stateful"), Capability("filesystem.read", target_arg="inputs")], description=( - "Execute Python code in a stateful sandbox session. " + "Execute Python code in a stateful session. " "State (variables, imports) persists across calls within the same session. " - "The sandbox shares the workspace filesystem — code can read/write files directly. " - "Network access and package installation are blocked — use web_fetch/web_search for HTTP. " - "Use for data analysis, prototyping, and producing work output. " - "Use execute_python instead for quick stateless calculations." + "Isolation depends on stateful_executor_backend: 'docker' runs in a " + "network-isolated container (no network egress; memory/CPU/PID-capped, " + "though disk is not); 'local' runs with host privileges and shared filesystem. " + "Disabled unless explicitly enabled. Use for data analysis, prototyping, " + "and producing work output. Use execute_python for quick stateless calculations. " + "Each `inputs` file is copied to `inputs/` inside the session before " + "the code runs; load it by that relative path (e.g. open('inputs/data.csv')). " + "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, session_id: str = "default", timeout_seconds: int = 120, + inputs: list[str] | None = None, ) -> dict[str, Any]: """Execute Python code in a stateful sandbox. @@ -34,10 +46,22 @@ def sandbox_execute( code: Python code to execute. session_id: Session identifier for state persistence (default: "default"). timeout_seconds: Maximum execution time in seconds. + inputs: Optional list of host file paths to stage into + inputs/ inside the session before execution. Returns: Dictionary with execution results. """ + # Opt-in gate. Gate before touching the service so a disabled deployment + # fails fast. NOTE: the workflow binds the factory tool (tools/factories.py), + # which gates identically — this gate covers the module-level tool used + # outside the workflow. + from agentic_cli.tools.sandbox.manager import sandbox_disabled_reason + + settings = get_settings() + if getattr(settings, "stateful_executor_backend", "none") == "none": + return {"success": False, "error": sandbox_disabled_reason(settings)} + manager = require_service(SANDBOX_MANAGER) if isinstance(manager, dict): return manager @@ -45,6 +69,7 @@ def sandbox_execute( code=code, session_id=session_id, timeout_seconds=timeout_seconds, + inputs=inputs, ) return { "success": result.success, diff --git a/src/agentic_cli/tools/sandbox/backends/base.py b/src/agentic_cli/tools/sandbox/backends/base.py index f9b2dca..6056c06 100644 --- a/src/agentic_cli/tools/sandbox/backends/base.py +++ b/src/agentic_cli/tools/sandbox/backends/base.py @@ -3,12 +3,14 @@ from abc import ABC, abstractmethod from pathlib import Path -from agentic_cli.tools.sandbox.models import ExecutionResult +from agentic_cli.tools.sandbox.models import ExecutionResult, SessionStatus class SandboxBackend(ABC): """Abstract base for sandbox execution backends.""" + backend_name: str = "unknown" + @abstractmethod def execute( self, @@ -16,6 +18,7 @@ def execute( session_id: str, timeout_seconds: int = 120, working_dir: Path | None = None, + inputs: list[str] | None = None, ) -> ExecutionResult: """Execute code in the given session. @@ -44,3 +47,8 @@ def cleanup(self) -> None: def has_session(self, session_id: str) -> bool: """Check if a session exists.""" ... + + def session_status(self, session_id: str) -> SessionStatus: + """Default status derived from has_session(); backends may override.""" + state = "ready" if self.has_session(session_id) else "absent" + return SessionStatus(session_id=session_id, state=state, backend=self.backend_name) diff --git a/src/agentic_cli/tools/sandbox/backends/container_runtime.py b/src/agentic_cli/tools/sandbox/backends/container_runtime.py new file mode 100644 index 0000000..103fb53 --- /dev/null +++ b/src/agentic_cli/tools/sandbox/backends/container_runtime.py @@ -0,0 +1,126 @@ +"""Container runtime seam: assemble and launch `docker run` (test-injectable).""" + +from __future__ import annotations + +import json +import subprocess +from dataclasses import dataclass, field + +from agentic_cli.logging import Loggers + +logger = Loggers.tools() + + +@dataclass +class Mount: + host: str + container: str + read_only: bool = True + + def to_flag(self) -> str: + spec = f"{self.host}:{self.container}" + return spec + ":ro" if self.read_only else spec + + +@dataclass +class ContainerSpec: + image: str + name: str + command: list[str] + network: str = "none" + memory_mb: int = 2048 + cpus: float = 2.0 + pids_limit: int = 256 + user: str = "" + env: dict[str, str] = field(default_factory=dict) + mounts: list[Mount] = field(default_factory=list) + labels: dict[str, str] = field(default_factory=dict) + + +class ContainerHandle: + """Thin wrapper over a running container's process + stdio streams.""" + + def __init__(self, name: str, proc: subprocess.Popen) -> None: + self.name = name + self._proc = proc + self.stdin = proc.stdin + self.stdout = proc.stdout + self.stderr = proc.stderr + + @property + def pid(self) -> int: + return self._proc.pid + + def poll(self) -> int | None: + return self._proc.poll() + + +class DockerContainerRuntime: + """Launch and control containers via the docker/podman CLI.""" + + def __init__(self, exe: str = "docker") -> None: + self._exe = exe + + @staticmethod + def build_run_argv(spec: ContainerSpec, exe: str) -> list[str]: + # Docker's default seccomp profile stays in effect (we never pass + # --privileged or --security-opt seccomp=unconfined), so dangerous + # syscalls remain blocked on top of the dropped capabilities. + argv: list[str] = [ + exe, "run", "--rm", "-i", + "--network", spec.network, + "--read-only", "--tmpfs", "/tmp", + "--cap-drop", "ALL", + "--security-opt", "no-new-privileges", + "--memory", f"{spec.memory_mb}m", + "--memory-swap", f"{spec.memory_mb}m", + "--pids-limit", str(spec.pids_limit), + "--cpus", str(spec.cpus), + "--name", spec.name, + ] + if spec.user: + argv += ["--user", spec.user] + for key, value in spec.env.items(): + argv += ["-e", f"{key}={value}"] + for mount in spec.mounts: + argv += ["-v", mount.to_flag()] + for key, value in spec.labels.items(): + argv += ["--label", f"{key}={value}"] + argv.append(spec.image) + argv += spec.command + return argv + + def start(self, spec: ContainerSpec) -> ContainerHandle: + argv = self.build_run_argv(spec, self._exe) + logger.debug("container_start", name=spec.name, image=spec.image) + proc = subprocess.Popen( + argv, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, # line-buffered + ) + return ContainerHandle(spec.name, proc) + + def kill(self, name: str, signal: str | None = None) -> None: + argv = [self._exe, "kill"] + if signal: + argv += ["--signal", signal] + argv.append(name) + subprocess.run(argv, capture_output=True, check=False) + + def inspect(self, name: str) -> dict: + proc = subprocess.run( + [self._exe, "inspect", name], capture_output=True, text=True, check=False, + ) + if proc.returncode != 0: + return {} + try: + data = json.loads(proc.stdout) + return data[0] if isinstance(data, list) and data else {} + except (json.JSONDecodeError, IndexError): + return {} + + def remove(self, name: str) -> None: + subprocess.run([self._exe, "rm", "-f", name], capture_output=True, check=False) diff --git a/src/agentic_cli/tools/sandbox/backends/detect.py b/src/agentic_cli/tools/sandbox/backends/detect.py new file mode 100644 index 0000000..7eeef1e --- /dev/null +++ b/src/agentic_cli/tools/sandbox/backends/detect.py @@ -0,0 +1,56 @@ +"""Detect an available container runtime (docker, then podman).""" + +from __future__ import annotations + +import shutil +import subprocess +from dataclasses import dataclass +from functools import lru_cache + +from agentic_cli.logging import Loggers + +logger = Loggers.tools() + +_RUNTIMES = ("docker", "podman") + + +@dataclass(frozen=True) +class DockerAvailability: + available: bool + runtime: str + detail: str + + +def _probe_daemon(exe: str) -> bool: + """Return True if ` info` succeeds (daemon reachable).""" + try: + proc = subprocess.run( + [exe, "info"], capture_output=True, timeout=10, check=False, + ) + return proc.returncode == 0 + except (OSError, subprocess.SubprocessError): + return False + + +@lru_cache(maxsize=1) +def detect_docker() -> DockerAvailability: + details: list[str] = [] + for exe in _RUNTIMES: + if shutil.which(exe) is None: + continue + if _probe_daemon(exe): + return DockerAvailability(True, exe, f"{exe} available") + # CLI present but daemon down — keep probing the remaining runtimes + # (the "docker, then podman" design) instead of giving up here. + details.append(f"{exe} CLI found but daemon not reachable") + if details: + return DockerAvailability(False, "", "; ".join(details)) + return DockerAvailability(False, "", "docker/podman not found in PATH") + + +def docker_available() -> bool: + return detect_docker().available + + +def clear_detection_cache() -> None: + detect_docker.cache_clear() diff --git a/src/agentic_cli/tools/sandbox/backends/driver.py b/src/agentic_cli/tools/sandbox/backends/driver.py new file mode 100644 index 0000000..d340af5 --- /dev/null +++ b/src/agentic_cli/tools/sandbox/backends/driver.py @@ -0,0 +1,130 @@ +"""In-container kernel driver. + +Runs INSIDE the sandbox container as `python .../driver.py`. Hosts a jupyter +kernel and speaks NDJSON over stdio: reads {"type":"execute","code","timeout"} +on stdin, writes {"type":"ready"} once and {"type":"result", ...} per request +on stdout. Interrupt is delivered as SIGINT to this process (PID 1); the handler +forwards it to the kernel so a running cell aborts without losing session state. +""" + +from __future__ import annotations + +import json +import os +import secrets +import signal +import sys +from pathlib import Path + +# Dual import: top-level when bind-mounted beside driver.py in-container; +# package path in tests / in-repo. +try: # pragma: no cover - import shim + from kernel_exec import collect_execution, validate_code +except ImportError: # pragma: no cover - import shim + from agentic_cli.tools.sandbox.backends.kernel_exec import collect_execution, validate_code + + +class KernelDriver: + def __init__(self, stdin, stdout, workspace: str | None = None) -> None: + self._stdin = stdin + self._stdout = stdout + self._workspace = Path(workspace or os.environ.get("AGENTIC_SANDBOX_WORKSPACE", "/workspace")) + self._km = None + self._kc = None + # Per-session secret tagging every protocol message. It lives only in + # this (trusted) driver process's memory — never in env/files, and the + # kernel can't ptrace the driver (default docker seccomp blocks it), so + # untrusted kernel code can't learn it. The host rejects any result line + # lacking this token, defeating forged writes to the driver's fd 1 via + # /proc//fd/1 (which the kernel can open, but only write-only). + self._token = secrets.token_hex(16) + + def start(self) -> None: + import subprocess + from jupyter_client import KernelManager + + # Start the kernel in the workspace so relative file writes land in the + # (writable, host-mounted) /workspace rather than the read-only image + # WORKDIR (e.g. /home/jovyan), where they fail or don't persist. The + # kernel inherits the driver's cwd; KernelManager.cwd is not honored by + # all jupyter_client versions, so chdir here is the reliable path. + try: + os.chdir(self._workspace) + except OSError: + pass + km = KernelManager() + # Isolate the kernel's raw std fds (0/1/2) from the driver's. The driver + # uses its own fd 0/1 for the NDJSON protocol with the host; if the + # kernel inherited them, user code could os.write(1, ...) a forged + # {"type":"result"} to desync the session, os.read(0, ...) pending host + # requests, or os.write(2, ...) to spam the host debug log. The kernel + # speaks ZMQ, so it needs none of these; real output/errors flow over + # iopub. + km.start_kernel(stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL) + kc = km.blocking_client() + kc.start_channels() + kc.wait_for_ready(timeout=60) + self._km, self._kc = km, kc + signal.signal(signal.SIGINT, self._on_sigint) + + def _on_sigint(self, signum, frame) -> None: + if self._km is not None: + self._km.interrupt_kernel() + + def handle_request(self, req: dict) -> dict: + code = req.get("code", "") + ok, msg = validate_code(code) + if not ok: + return {"type": "result", "success": False, "stdout": "", "stderr": "", + "result": None, "artifacts": [], "execution_time": 0.0, "error": msg} + msg_id = self._kc.execute(code) + # The HOST owns the timeout/interrupt/kill state machine: it waits the + # per-cell deadline, then sends a cooperative interrupt (SIGINT -> + # _on_sigint -> interrupt_kernel), then hard-kills the container if that + # fails. So the driver must block until the kernel actually goes idle + # (natural completion) or the interrupt aborts the cell. If the driver + # self-timed-out on the same deadline it would return a "timed out" + # result while the cell keeps running in the kernel — the host would + # consume that result WITHOUT interrupting, wedging the session for the + # next request. Passing timeout=None makes collect_execution block. + data = collect_execution(self._kc, msg_id, None, self._workspace) + return {"type": "result", **data} + + def _write(self, obj: dict) -> None: + obj = {**obj, "token": self._token} + self._stdout.write(json.dumps(obj) + "\n") + self._stdout.flush() + + def run(self) -> None: + self.start() + self._write({"type": "ready"}) + for line in self._stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except json.JSONDecodeError: + self._write({"type": "result", "success": False, "error": "invalid request json", + "stdout": "", "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.0}) + continue + self._write(self.handle_request(req)) + + def close(self) -> None: + try: + if self._kc is not None: + self._kc.stop_channels() + except Exception: + pass + try: + if self._km is not None: + self._km.shutdown_kernel(now=True) + except Exception: + pass + + +if __name__ == "__main__": # pragma: no cover - container entrypoint + KernelDriver(sys.stdin, sys.stdout).run() diff --git a/src/agentic_cli/tools/sandbox/backends/jupyter_docker.py b/src/agentic_cli/tools/sandbox/backends/jupyter_docker.py new file mode 100644 index 0000000..958ee77 --- /dev/null +++ b/src/agentic_cli/tools/sandbox/backends/jupyter_docker.py @@ -0,0 +1,396 @@ +"""Docker-backed stateful sandbox backend.""" + +from __future__ import annotations + +import json +import os +import queue +import stat +import threading +import time +from pathlib import Path + +from agentic_cli.logging import Loggers +from agentic_cli.tools.sandbox.models import ExecutionResult, SessionStatus +from agentic_cli.file_utils import copy_regular_file_no_follow, sanitize_filename +from agentic_cli.tools.sandbox.backends.base import SandboxBackend +from agentic_cli.tools.sandbox.backends import kernel_exec +from agentic_cli.tools.sandbox.backends.container_runtime import ( + ContainerSpec, Mount, DockerContainerRuntime, +) +from agentic_cli.tools.sandbox.backends.detect import detect_docker + +logger = Loggers.tools() + +_EOF = object() + + +class SandboxStartError(Exception): + """Raised when a container/kernel fails to become ready.""" + + +class ContainerSession: + """Owns one container: its stdio protocol, reader thread, and status.""" + + def __init__(self, session_id, handle, interrupt, kill, working_dir, + start_timeout=180, interrupt_grace=10.0, backend_name="jupyter_docker"): + self._session_id = session_id + self._handle = handle + self._interrupt = interrupt + self._kill = kill + self._working_dir = working_dir + self._start_timeout = start_timeout + self._interrupt_grace = interrupt_grace + self._backend_name = backend_name + self._queue: queue.Queue = queue.Queue() + self._lock = threading.Lock() + self._status = "starting" + self._token: str | None = None # captured from the driver's ready message + self._reader = threading.Thread(target=self._read_stdout, daemon=True) + self._reader.start() + self._stderr_reader = None + if getattr(handle, "stderr", None) is not None: + self._stderr_reader = threading.Thread(target=self._read_stderr, daemon=True) + self._stderr_reader.start() + + @property + def status(self) -> str: + return self._status + + @property + def container_id(self) -> str: + return getattr(self._handle, "name", "") + + def _read_stdout(self) -> None: + try: + for line in self._handle.stdout: + self._queue.put(line) + finally: + self._queue.put(_EOF) + + def _read_stderr(self) -> None: + try: + for line in self._handle.stderr: + logger.debug("sandbox_stderr", line=line.rstrip()) + except Exception: + pass + + def wait_ready(self) -> None: + try: + line = self._queue.get(timeout=self._start_timeout) + except queue.Empty: + self._kill() + self._status = "dead" + raise SandboxStartError(f"sandbox kernel not ready within {self._start_timeout}s") + if line is _EOF: + self._status = "dead" + raise SandboxStartError("sandbox container exited before becoming ready") + msg = json.loads(line) + if msg.get("type") != "ready": + self._status = "dead" + raise SandboxStartError(f"unexpected startup message: {msg!r}") + # Capture the session token. This ready message is genuine: no user code + # runs before it (the kernel executes nothing until an execute request). + self._token = msg.get("token") + self._status = "ready" + + def execute(self, code: str, timeout: float) -> ExecutionResult: + with self._lock: + if self._handle.poll() is not None: + self._status = "dead" + return ExecutionResult(success=False, + error=f"sandbox container exited unexpectedly ({self._exit_detail()})") + self._status = "busy" + try: + # The host owns the timeout/interrupt/kill state machine; the + # driver blocks until idle, so the request carries only the code. + self._handle.stdin.write(json.dumps({"type": "execute", "code": code}) + "\n") + self._handle.stdin.flush() + + msg = self._read_authenticated(timeout) + if msg is None: # timed out -> cooperative interrupt, then hard kill + self._interrupt() + msg = self._read_authenticated(self._interrupt_grace) + if msg is None: + self._kill() + self._status = "dead" + return ExecutionResult(success=False, + error=f"Execution timed out after {timeout}s; container killed") + if msg is _EOF: + self._status = "dead" + return ExecutionResult(success=False, + error=f"sandbox container exited unexpectedly ({self._exit_detail()})") + + self._status = "ready" + return self._to_result(msg) + except Exception as exc: + # Don't leak the container on an unexpected host-side error + # (e.g. a bad timeout raising in queue.get). + try: + self._kill() + except Exception: + pass + self._status = "dead" + return ExecutionResult(success=False, error=f"sandbox execution failed: {exc}") + + def _read_authenticated(self, timeout: float): + """Return the next AUTHENTICATED protocol message within ``timeout``, + skipping lines that don't carry the session token — those are forged by + user code writing to the driver's fd 1 (e.g. via /proc//fd/1). + Returns the parsed dict, ``None`` on timeout, or ``_EOF``.""" + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + try: + line = self._queue.get(timeout=remaining) + except queue.Empty: + return None + if line is _EOF: + return _EOF + try: + msg = json.loads(line) + except (ValueError, TypeError): + continue # non-JSON garbage — forged + if self._token is not None and msg.get("token") != self._token: + continue # missing/wrong token — forged, skip it + return msg + + def _exit_detail(self) -> str: + """Describe why the container process exited, from the docker-run exit + code (reliable even with --rm, which removes the container before it + could be inspected). 137 = 128+SIGKILL, the cgroup OOM-killer signature.""" + code = self._handle.poll() + if code == 137: + return f"exit {code}; possibly out-of-memory (OOM-killed)" + return f"exit {code}" + + def _to_result(self, data: dict) -> ExecutionResult: + artifacts = [self._translate(p) for p in data.get("artifacts", [])] + return ExecutionResult( + success=data.get("success", False), + stdout=data.get("stdout", ""), + stderr=data.get("stderr", ""), + result=data.get("result"), + artifacts=artifacts, + execution_time=data.get("execution_time", 0.0), + error=data.get("error", ""), + ) + + def _translate(self, container_path: str) -> str: + """Map an in-container /workspace path to the host session dir.""" + prefix = "/workspace" + if self._working_dir is not None and container_path.startswith(prefix): + rel = container_path[len(prefix):].lstrip("/") + return str(Path(self._working_dir) / rel) + return container_path + + def close(self) -> None: + try: + self._kill() + finally: + self._status = "dead" + + +_DRIVER_DIR = "/opt/agentic_sandbox" +# Read-only data mounts land here, a TOP-LEVEL path (sibling of /workspace) — +# NOT nested under the /workspace bind mount. Nesting a second bind mount under +# the session-dir mount makes Docker synthesize the mount-point dir root-owned +# (Linux) / file-sharing-ACL-stamped (macOS), which breaks host-side cleanup. A +# top-level mount point lives in the container's own ephemeral layer instead. +_DATA_DIR = "/data" + + +class JupyterDockerBackend(SandboxBackend): + """Runs each session in its own network-isolated container.""" + + backend_name = "jupyter_docker" + + def __init__(self, settings, runtime=None, detect_fn=None) -> None: + self._settings = settings + self._detect = detect_fn or detect_docker + self._runtime = runtime # lazily created so detect can pick docker/podman + self._sessions: dict[str, ContainerSession] = {} + + def _ensure_runtime(self): + if self._runtime is None: + avail = self._detect() + self._runtime = DockerContainerRuntime(avail.runtime or "docker") + return self._runtime + + def _outputs_dir(self) -> Path: + configured = getattr(self._settings, "sandbox_outputs_dir", "") or "" + base = Path(configured) if configured else Path(self._settings.workspace_dir) / "artifacts" + base.mkdir(parents=True, exist_ok=True) + try: + os.chmod(base, 0o700) # single-user; not world-accessible + except OSError: + pass + return base + + def _collect_outputs(self, working_dir: Path) -> list[str]: + """Copy regular files from /outputs into the shared outputs + dir without following symlinks. Skips symlinks / special files (a kernel + could plant `outputs/x -> /host/secret`) and a symlinked `outputs` dir. + Never raises — best-effort contract so execute() is not disrupted.""" + session_outs = Path(working_dir) / "outputs" + try: + if not stat.S_ISDIR(os.lstat(session_outs).st_mode): + return [] # 'outputs' is a symlink or not a directory + except OSError: + return [] + collected: list[str] = [] + try: + shared = self._outputs_dir() + entries = sorted(session_outs.iterdir()) + except OSError: + logger.warning("sandbox_outputs_unreadable", path=str(session_outs)) + return collected + for src in entries: + dst = shared / src.name + try: + copy_regular_file_no_follow(src, dst) + except (OSError, ValueError): + logger.warning("sandbox_output_skipped", entry=str(src)) + continue + collected.append(str(dst)) + return collected + + def _parse_data_mounts(self) -> list[tuple[str, str]]: + """Parse sandbox_data_mounts into (host_path, sanitized_name) pairs.""" + result = [] + for entry in self._settings.sandbox_data_mounts: + host, _, name = entry.partition(":") + # Sanitize the mount name so a hostile '..'/absolute value can't + # remap the mount outside /data/ (sanitize_filename maps '/' and + # '.' to '_'). + name = sanitize_filename(name or Path(host).name) or "mount" + result.append((host, name)) + return result + + def _build_spec(self, session_id: str, working_dir) -> ContainerSpec: + s = self._settings + here = Path(__file__).parent + # NOTE: /workspace is a writable host bind mount with NO disk quota — + # Docker caps memory/CPU/PIDs but not bind-mount disk. Operators who need + # a hard limit should place workspace_dir on a quota'd filesystem. + mounts = [ + Mount(str(working_dir), "/workspace", read_only=False), + Mount(str(here / "driver.py"), f"{_DRIVER_DIR}/driver.py", read_only=True), + Mount(str(here / "kernel_exec.py"), f"{_DRIVER_DIR}/kernel_exec.py", read_only=True), + ] + for host, name in self._parse_data_mounts(): + mounts.append(Mount(host, f"{_DATA_DIR}/{name}", read_only=True)) + env = { + "HOME": "/tmp", "MPLCONFIGDIR": "/tmp", "IPYTHONDIR": "/tmp", + "JUPYTER_RUNTIME_DIR": "/tmp", "PYTHONDONTWRITEBYTECODE": "1", + "AGENTIC_SANDBOX_WORKSPACE": "/workspace", + } + # Run as the host uid:gid by default so files the non-root kernel writes + # to the /workspace bind mount are owned by the host user — no + # world-writable chmod on the session dir needed. An explicit + # sandbox_container_user overrides. + user = s.sandbox_container_user + if not user and hasattr(os, "getuid"): + user = f"{os.getuid()}:{os.getgid()}" + return ContainerSpec( + image=s.sandbox_image, + name=f"agentic-sbx-{sanitize_filename(session_id)}", + command=["python", f"{_DRIVER_DIR}/driver.py"], + network=s.sandbox_network, + memory_mb=s.sandbox_memory_mb, + cpus=s.sandbox_cpus, + pids_limit=s.sandbox_pids_limit, + user=user, + env=env, + mounts=mounts, + labels={"agentic-sandbox": "1", "agentic-session": session_id}, + ) + + def _start_session(self, session_id: str, working_dir) -> ContainerSession: + # The container runs as the host uid:gid (see _build_spec), so the + # host-owned session dir is writable by the kernel without loosening its + # permissions. No chmod needed. + runtime = self._ensure_runtime() + spec = self._build_spec(session_id, working_dir) + if working_dir is not None: + wd = Path(working_dir) + # Pre-create the outputs/ subdir host-owned so the container (running + # as the host uid) writes finals into a host-owned dir. It is a plain + # subdir of the /workspace mount, not a nested bind mount, so it has + # none of the root-owned/ACL cleanup problems that data mounts had. + # Data mounts need no pre-create: they land at a top-level /data/ + # (see _DATA_DIR) whose mount point lives in the container's own + # layer, never inside the host session dir. + (wd / "outputs").mkdir(parents=True, exist_ok=True) + handle = runtime.start(spec) + name = spec.name + session = ContainerSession( + session_id=session_id, + handle=handle, + interrupt=lambda: runtime.kill(name, "INT"), + kill=lambda: runtime.kill(name), + working_dir=working_dir, + start_timeout=self._settings.sandbox_start_timeout, + backend_name=self.backend_name, + ) + try: + session.wait_ready() + except Exception: + session.close() + raise + self._sessions[session_id] = session + return session + + def execute(self, code, session_id, timeout_seconds=120, working_dir=None, inputs=None) -> ExecutionResult: + avail = self._detect() + if not avail.available: + return ExecutionResult( + success=False, + error=(f"Docker sandbox backend unavailable ({avail.detail}). " + "Refusing to fall back to an unsandboxed kernel."), + ) + ok, msg = kernel_exec.validate_code(code) + if not ok: + return ExecutionResult(success=False, error=msg) + if inputs: + from agentic_cli.tools.sandbox.manager import stage_inputs + try: + stage_inputs(working_dir, inputs) + except ValueError as exc: + return ExecutionResult(success=False, error=f"input staging failed: {exc}") + session = self._sessions.get(session_id) + if session is None or session.status == "dead": + try: + session = self._start_session(session_id, working_dir) + except Exception as exc: + return ExecutionResult(success=False, error=f"Failed to start sandbox: {exc}") + result = session.execute(code, timeout_seconds) + if result.success and working_dir is not None: + extra = self._collect_outputs(Path(working_dir)) + if extra: + result.artifacts = list(result.artifacts) + extra + return result + + def reset_session(self, session_id: str) -> None: + session = self._sessions.pop(session_id, None) + if session is not None: + session.close() + + def cleanup(self) -> None: + for session_id in list(self._sessions): + self.reset_session(session_id) + + def has_session(self, session_id: str) -> bool: + session = self._sessions.get(session_id) + return session is not None and session.status != "dead" + + def session_status(self, session_id: str) -> SessionStatus: + session = self._sessions.get(session_id) + if session is None: + return SessionStatus(session_id=session_id, state="absent", backend=self.backend_name) + return SessionStatus( + session_id=session_id, state=session.status, backend=self.backend_name, + container_id=session.container_id, + ) diff --git a/src/agentic_cli/tools/sandbox/backends/jupyter_local.py b/src/agentic_cli/tools/sandbox/backends/jupyter_local.py index 71d96c7..c834dfd 100644 --- a/src/agentic_cli/tools/sandbox/backends/jupyter_local.py +++ b/src/agentic_cli/tools/sandbox/backends/jupyter_local.py @@ -6,27 +6,18 @@ from __future__ import annotations -import base64 -import re -import time from pathlib import Path -from typing import Any from jupyter_client import KernelManager from jupyter_client.blocking import BlockingKernelClient from agentic_cli.logging import Loggers +from agentic_cli.tools.sandbox.backends import kernel_exec from agentic_cli.tools.sandbox.backends.base import SandboxBackend from agentic_cli.tools.sandbox.models import ExecutionResult logger = Loggers.tools() -# Regex to strip ANSI escape codes from tracebacks -_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") - -# Blocked shell escape patterns -_BLOCKED_MAGICS = frozenset({"pip", "system", "sx"}) - # Initialization code injected into new kernels to block network modules _SANDBOX_INIT_CODE = ''' import sys as _sys @@ -78,22 +69,11 @@ class JupyterLocalBackend(SandboxBackend): pair. State (variables, imports) persists across calls within a session. """ + backend_name = "jupyter_local" + def __init__(self) -> None: self._sessions: dict[str, tuple[KernelManager, BlockingKernelClient]] = {} - @staticmethod - def _validate_code(code: str) -> tuple[bool, str]: - """Pre-scan code for blocked shell escapes and magics.""" - for line in code.splitlines(): - stripped = line.lstrip() - if stripped.startswith("!"): - return False, "Shell commands (!) are not allowed in the sandbox" - if stripped.startswith("%") and not stripped.startswith("%%"): - magic = stripped.lstrip("%").split()[0] if stripped.lstrip("%") else "" - if magic in _BLOCKED_MAGICS: - return False, f"Magic command '%{magic}' is not allowed in the sandbox" - return True, "" - def _get_or_create_session( self, session_id: str, working_dir: Path | None = None, ) -> tuple[KernelManager, BlockingKernelClient]: @@ -102,10 +82,10 @@ def _get_or_create_session( return self._sessions[session_id] km = KernelManager() + start_kwargs: dict = {} if working_dir: - km.cwd = str(working_dir) - - km.start_kernel() + start_kwargs["cwd"] = str(working_dir) + km.start_kernel(**start_kwargs) kc = km.blocking_client() kc.start_channels() kc.wait_for_ready(timeout=30) @@ -133,88 +113,28 @@ def execute( session_id: str, timeout_seconds: int = 120, working_dir: Path | None = None, + inputs: list[str] | None = None, ) -> ExecutionResult: """Execute code in a Jupyter kernel session.""" # Pre-scan for blocked patterns - valid, error = self._validate_code(code) + valid, error = kernel_exec.validate_code(code) if not valid: return ExecutionResult(success=False, error=error) - start = time.monotonic() - _, kc = self._get_or_create_session(session_id, working_dir) - - msg_id = kc.execute(code) - - stdout_parts: list[str] = [] - stderr_parts: list[str] = [] - result_value: str | None = None - artifacts: list[str] = [] - error_text = "" + if working_dir is not None: + (Path(working_dir) / "outputs").mkdir(parents=True, exist_ok=True) - # Collect iopub messages until kernel goes idle - while True: + if inputs: + from agentic_cli.tools.sandbox.manager import stage_inputs try: - msg = kc.get_iopub_msg(timeout=timeout_seconds) - except TimeoutError: - elapsed = time.monotonic() - start - return ExecutionResult( - success=False, - stdout="".join(stdout_parts), - stderr="".join(stderr_parts), - error=f"Execution timed out after {timeout_seconds}s", - execution_time=elapsed, - ) - - # Only process messages from our execution - if msg.get("parent_header", {}).get("msg_id") != msg_id: - continue - - msg_type = msg.get("msg_type", "") - content: dict[str, Any] = msg.get("content", {}) - - if msg_type == "stream": - name = content.get("name", "stdout") - text = content.get("text", "") - if name == "stderr": - stderr_parts.append(text) - else: - stdout_parts.append(text) - - elif msg_type == "execute_result": - data = content.get("data", {}) - result_value = data.get("text/plain", "") - - elif msg_type == "display_data": - data = content.get("data", {}) - if "image/png" in data and working_dir: - # Save image artifact - artifact_dir = working_dir / "artifacts" - artifact_dir.mkdir(parents=True, exist_ok=True) - n = len(artifacts) - path = artifact_dir / f"plot_{n}.png" - img_bytes = base64.b64decode(data["image/png"]) - path.write_bytes(img_bytes) - artifacts.append(str(path)) - logger.debug("artifact_saved", path=str(path)) - - elif msg_type == "error": - traceback_lines = content.get("traceback", []) - raw = "\n".join(traceback_lines) - error_text = _ANSI_RE.sub("", raw) - - elif msg_type == "status" and content.get("execution_state") == "idle": - break + stage_inputs(working_dir, inputs) + except ValueError as exc: + return ExecutionResult(success=False, error=f"input staging failed: {exc}") - elapsed = time.monotonic() - start - return ExecutionResult( - success=not error_text, - stdout="".join(stdout_parts), - stderr="".join(stderr_parts), - result=result_value, - artifacts=artifacts, - execution_time=elapsed, - error=error_text, - ) + _, kc = self._get_or_create_session(session_id, working_dir) + msg_id = kc.execute(code) + data = kernel_exec.collect_execution(kc, msg_id, timeout_seconds, working_dir) + return ExecutionResult(**data) def reset_session(self, session_id: str) -> None: """Restart the kernel for a session.""" diff --git a/src/agentic_cli/tools/sandbox/backends/kernel_exec.py b/src/agentic_cli/tools/sandbox/backends/kernel_exec.py new file mode 100644 index 0000000..bdcde4f --- /dev/null +++ b/src/agentic_cli/tools/sandbox/backends/kernel_exec.py @@ -0,0 +1,110 @@ +"""Shared kernel execution helpers. Dependency-free (stdlib + jupyter_client) +so this module can be bind-mounted next to driver.py and imported in-container.""" + +from __future__ import annotations + +import base64 +import queue +import re +import time +from typing import Any + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +_BLOCKED_MAGICS = frozenset({"pip", "system", "sx"}) + +# Per-stream character cap. A runaway cell (e.g. an interruptible print loop) +# would otherwise buffer output unbounded in-container, ship it all over stdio, +# and bloat host memory / the LLM context. Truncate with a marker instead. +MAX_STREAM_CHARS = 1_000_000 + + +def _append_capped(parts: list[str], current: int, text: str) -> int: + """Append `text` to `parts` up to MAX_STREAM_CHARS. Adds a one-time + truncation marker when the cap is first exceeded, then drops the rest. + Returns the updated character count.""" + if current >= MAX_STREAM_CHARS: + return current + room = MAX_STREAM_CHARS - current + if len(text) <= room: + parts.append(text) + return current + len(text) + parts.append(text[:room]) + parts.append(f"\n...[output truncated at {MAX_STREAM_CHARS} characters]...") + return MAX_STREAM_CHARS + + +def validate_code(code: str) -> tuple[bool, str]: + """Pre-scan code for blocked shell escapes and magics.""" + for line in code.splitlines(): + stripped = line.lstrip() + if stripped.startswith("!"): + return False, "Shell commands (!) are not allowed in the sandbox" + if stripped.startswith("%") and not stripped.startswith("%%"): + body = stripped.lstrip("%") + magic = body.split()[0] if body else "" + if magic in _BLOCKED_MAGICS: + return False, f"Magic command '%{magic}' is not allowed in the sandbox" + return True, "" + + +def collect_execution(kc, msg_id: str, timeout: float, working_dir) -> dict: + """Collect iopub messages for one execution into a result dict.""" + start = time.monotonic() + stdout_parts: list[str] = [] + stderr_parts: list[str] = [] + stdout_len = 0 + stderr_len = 0 + result_value: str | None = None + artifacts: list[str] = [] + error_text = "" + + while True: + try: + msg = kc.get_iopub_msg(timeout=timeout) + except (queue.Empty, TimeoutError): # timeout waiting for kernel output + return { + "success": False, + "stdout": "".join(stdout_parts), + "stderr": "".join(stderr_parts), + "result": result_value, + "artifacts": artifacts, + "execution_time": time.monotonic() - start, + "error": f"Execution timed out after {timeout}s", + } + + if msg.get("parent_header", {}).get("msg_id") != msg_id: + continue + + msg_type = msg.get("msg_type", "") + content: dict[str, Any] = msg.get("content", {}) + + if msg_type == "stream": + text = content.get("text", "") + if content.get("name") == "stderr": + stderr_len = _append_capped(stderr_parts, stderr_len, text) + else: + stdout_len = _append_capped(stdout_parts, stdout_len, text) + elif msg_type == "execute_result": + result_value = content.get("data", {}).get("text/plain", "") + elif msg_type == "display_data": + data = content.get("data", {}) + if "image/png" in data and working_dir is not None: + artifact_dir = working_dir / "artifacts" + artifact_dir.mkdir(parents=True, exist_ok=True) + path = artifact_dir / f"plot_{len(artifacts)}.png" + path.write_bytes(base64.b64decode(data["image/png"])) + artifacts.append(str(path)) + elif msg_type == "error": + error_text = _ANSI_RE.sub("", "\n".join(content.get("traceback", []))) + elif msg_type == "status" and content.get("execution_state") == "idle": + break + + return { + "success": not error_text, + "stdout": "".join(stdout_parts), + "stderr": "".join(stderr_parts), + "result": result_value, + "artifacts": artifacts, + "execution_time": time.monotonic() - start, + "error": error_text, + } diff --git a/src/agentic_cli/tools/sandbox/manager.py b/src/agentic_cli/tools/sandbox/manager.py index 02b1669..5fa83da 100644 --- a/src/agentic_cli/tools/sandbox/manager.py +++ b/src/agentic_cli/tools/sandbox/manager.py @@ -7,12 +7,14 @@ from __future__ import annotations import atexit +import os +import stat from dataclasses import dataclass, field from pathlib import Path from typing import Any, TYPE_CHECKING from agentic_cli.logging import Loggers -from agentic_cli.file_utils import sanitize_filename +from agentic_cli.file_utils import copy_regular_file_no_follow, sanitize_filename from agentic_cli.tools.sandbox.models import ExecutionResult if TYPE_CHECKING: @@ -22,6 +24,50 @@ logger = Loggers.tools() +def stage_inputs(session_dir: Path, inputs: list[str]) -> None: + """Copy each host path into session_dir/inputs/ (the in-sandbox + 'inputs/' contract). Raises ValueError on a missing file or a + basename collision (v1 does not support renaming).""" + if session_dir is None: + raise ValueError("session_dir is required for staging inputs") + if not inputs: + return + inputs_dir = Path(session_dir) / "inputs" + inputs_dir.mkdir(parents=True, exist_ok=True) + if not stat.S_ISDIR(os.lstat(inputs_dir).st_mode): + # 'inputs' is a symlink / not a real dir (kernel-plantable) — O_NOFOLLOW + # and os.replace only guard the final component, not this parent. + raise ValueError("inputs staging directory is not a real directory (symlink?)") + seen: set[str] = set() + for src in inputs: + p = Path(src).expanduser() + name = p.name + if name in seen: + raise ValueError(f"duplicate input basename {name!r}; rename one of the source files") + seen.add(name) + try: + copy_regular_file_no_follow(p, inputs_dir / name) + except FileNotFoundError as exc: + raise ValueError(f"input file not found: {src}") from exc + except (OSError, ValueError) as exc: + raise ValueError(f"input must be a regular file: {src}") from exc + + +def sandbox_disabled_reason(settings) -> str: + """Backend-aware reason string for a disabled sandbox_execute. Shared by + every entry point (module tool + factory tool) so the opt-in gate is + applied consistently and worded accurately for the selected backend.""" + backend = getattr(settings, "stateful_executor_backend", "none") + if backend == "docker": + detail = "The 'docker' backend runs it in a network-isolated, memory/CPU/PID-capped container." + elif backend == "local": + detail = "The 'local' backend runs Python with host privileges and no OS sandbox." + else: + detail = "No stateful executor is configured." + return (f"sandbox_execute is not available. {detail} " + "Set stateful_executor_backend to 'docker' (isolated) or 'local' (unsandboxed) to use it.") + + @dataclass class SandboxSession: """Metadata for an active sandbox session.""" @@ -37,7 +83,7 @@ class SandboxManager: Args: settings: Application settings instance. backend: Optional backend for test injection. If None, created - lazily from settings.sandbox_backend. + lazily from settings.stateful_executor_backend. """ def __init__( @@ -55,15 +101,18 @@ def __init__( def _ensure_backend(self) -> "SandboxBackend": """Lazily create the backend if not injected.""" if self._backend is None: - self._backend = self._create_backend(self._settings.sandbox_backend) + self._backend = self._create_backend(self._settings.stateful_executor_backend) return self._backend def _create_backend(self, backend_name: str) -> "SandboxBackend": """Create a backend by name.""" - if backend_name == "jupyter_local": + if backend_name == "local": from agentic_cli.tools.sandbox.backends.jupyter_local import JupyterLocalBackend return JupyterLocalBackend() - raise ValueError(f"Unknown sandbox backend: {backend_name!r}") + if backend_name == "docker": + from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend + return JupyterDockerBackend(self._settings) + raise ValueError(f"Unknown stateful_executor_backend: {backend_name!r}") def _get_session_dir(self, session_id: str) -> Path: """Get or create the working directory for a session.""" @@ -78,6 +127,7 @@ def execute( code: str, session_id: str = "default", timeout_seconds: int | None = None, + inputs: list[str] | None = None, ) -> ExecutionResult: """Execute code in a sandbox session. @@ -85,10 +135,18 @@ def execute( code: Python code to execute. session_id: Session identifier (default: "default"). timeout_seconds: Execution timeout (uses settings default if None). + inputs: Optional list of host file paths to stage into + inputs/ inside the session before execution. Returns: ExecutionResult with output and metadata. """ + if timeout_seconds is not None and timeout_seconds <= 0: + return ExecutionResult( + success=False, + error=f"timeout_seconds must be a positive number (got {timeout_seconds!r}).", + ) + max_sessions = self._settings.sandbox_max_sessions if session_id not in self._sessions and len(self._sessions) >= max_sessions: return ExecutionResult( @@ -115,8 +173,15 @@ def execute( session_id=session_id, timeout_seconds=timeout_seconds, working_dir=session.working_dir, + inputs=inputs, ) + # A failed start (docker down, startup error) must not leave phantom + # session metadata occupying a max_sessions slot with no backend session. + if not result.success and not backend.has_session(session_id): + self._sessions.pop(session_id, None) + return result + session.execution_count += 1 logger.debug( "sandbox_executed", @@ -144,19 +209,20 @@ def reset_session(self, session_id: str) -> bool: return was_active def list_sessions(self) -> list[dict[str, Any]]: - """List active sandbox sessions. - - Returns: - List of session metadata dicts. - """ - return [ - { + """List active sandbox sessions with live backend status.""" + if not self._sessions: + return [] + backend = self._ensure_backend() + rows: list[dict[str, Any]] = [] + for s in self._sessions.values(): + status = backend.session_status(s.session_id) + rows.append({ "session_id": s.session_id, "working_dir": str(s.working_dir), "execution_count": s.execution_count, - } - for s in self._sessions.values() - ] + "state": status.state, + }) + return rows def cleanup(self) -> None: """Clean up all sessions and the backend.""" diff --git a/src/agentic_cli/tools/sandbox/models.py b/src/agentic_cli/tools/sandbox/models.py index 972d716..adab910 100644 --- a/src/agentic_cli/tools/sandbox/models.py +++ b/src/agentic_cli/tools/sandbox/models.py @@ -14,3 +14,15 @@ class ExecutionResult: artifacts: list[str] = field(default_factory=list) execution_time: float = 0.0 error: str = "" + + +@dataclass +class SessionStatus: + """Live status of a sandbox session.""" + + session_id: str + state: str # "starting" | "ready" | "busy" | "dead" | "absent" + backend: str + execution_count: int = 0 + container_id: str | None = None + detail: str = "" diff --git a/src/agentic_cli/tools/shell/os_sandbox/policy.py b/src/agentic_cli/tools/shell/os_sandbox/policy.py index 0e0d605..59cc9dd 100644 --- a/src/agentic_cli/tools/shell/os_sandbox/policy.py +++ b/src/agentic_cli/tools/shell/os_sandbox/policy.py @@ -75,6 +75,9 @@ class OSSandboxPolicy: Always includes MANDATORY_DENY_WRITE entries. deny_read_paths: Paths to hide entirely from the sandboxed process. allow_network: Whether network access is allowed (Phase 2). + strict: When True, refuse to execute if sandboxing was requested but no + real backend is available (instead of falling back to the restricted + in-process executor). """ enabled: bool = True @@ -84,6 +87,7 @@ class OSSandboxPolicy: ) deny_read_paths: list[str] = field(default_factory=list) allow_network: bool = False + strict: bool = False def resolved_writable_paths(self, working_dir: Path) -> list[Path]: """Resolve all writable paths to absolute, always including working_dir. diff --git a/src/agentic_cli/tools/skills/__init__.py b/src/agentic_cli/tools/skills/__init__.py new file mode 100644 index 0000000..0cbc2c7 --- /dev/null +++ b/src/agentic_cli/tools/skills/__init__.py @@ -0,0 +1,14 @@ +"""Framework skills support, built on ADK's native Agent-Skills system. + +Importing this package registers the ADK skill tool names with the permission +registry (so the ADK PermissionPlugin allows them). +""" + +from agentic_cli.tools.skills.permissions import register_skill_tool_permissions +from agentic_cli.tools.skills.store import SkillStore +from agentic_cli.tools.skills.toolset import make_skill_toolset + +# Ensure skill tool names are permissioned as soon as skills are used. +register_skill_tool_permissions() + +__all__ = ["SkillStore", "make_skill_toolset", "register_skill_tool_permissions"] diff --git a/src/agentic_cli/tools/skills/permissions.py b/src/agentic_cli/tools/skills/permissions.py new file mode 100644 index 0000000..76ef0c2 --- /dev/null +++ b/src/agentic_cli/tools/skills/permissions.py @@ -0,0 +1,45 @@ +"""Register ADK skill tool names with the framework permission registry. + +ADK's ``SkillToolset`` tools aren't decorated with ``@register_tool``, so the +ADK ``PermissionPlugin`` would hard-deny them. We register their names here so +the discovery/read tools are EXEMPT and script execution is permissioned. + +The registered function is a placeholder — the plugin only reads the declared +capabilities, never calls the function. +""" + +from __future__ import annotations + +_SKILL_READ_TOOLS = ("list_skills", "load_skill", "load_skill_resource") + + +def _noop(*args, **kwargs): # pragma: no cover - placeholder, never invoked + return None + + +def register_skill_tool_permissions() -> None: + """Idempotently register skill tool names + capabilities in the registry.""" + from agentic_cli.tools.registry import ToolCategory, get_registry + from agentic_cli.workflow.permissions import EXEMPT + from agentic_cli.workflow.permissions.capabilities import Capability + + reg = get_registry() + + for name in _SKILL_READ_TOOLS: + if name not in reg: + reg.register( + _noop, + name=name, + capabilities=EXEMPT, + category=ToolCategory.KNOWLEDGE, + description=f"ADK skill tool: {name}", + ) + + if "run_skill_script" not in reg: + reg.register( + _noop, + name="run_skill_script", + capabilities=[Capability("skill.script.exec", target_arg="file_path")], + category=ToolCategory.EXECUTION, + description="ADK skill tool: run_skill_script", + ) diff --git a/src/agentic_cli/tools/skills/store.py b/src/agentic_cli/tools/skills/store.py new file mode 100644 index 0000000..54f501d --- /dev/null +++ b/src/agentic_cli/tools/skills/store.py @@ -0,0 +1,51 @@ +"""Load Agent-Skills (SKILL.md folders) into ADK Skill objects. + +Skills follow the Agent Skills specification (same SKILL.md format used by +ADK and Anthropic): a folder with frontmatter (name/description), a markdown +body, and optional ``references/`` / ``assets/`` / ``scripts/`` subfolders. + +This reuses ADK's native loader (``google.adk.skills.load_skill_from_dir``). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + + +class SkillStore: + """Resolve skill references (paths or names) to loaded Skill objects. + + A skill reference is either a path to a skill directory, or a bare name + looked up as a subdirectory of one of ``skill_dirs``. + """ + + def __init__(self, skill_dirs: list[str | Path] | None = None) -> None: + self._search_dirs = [Path(d) for d in (skill_dirs or [])] + + def _load_dir(self, path: Path) -> Any: + from google.adk.skills import load_skill_from_dir + + return load_skill_from_dir(path) + + def _resolve_ref(self, ref: str) -> Path: + path = Path(ref) + if path.is_dir(): + return path + for base in self._search_dirs: + candidate = base / ref + if candidate.is_dir(): + return candidate + searched = [str(d) for d in self._search_dirs] + raise ValueError( + f"Skill {ref!r} not found: not a directory and not under " + f"skills_dirs {searched}." + ) + + def resolve(self, refs: list[str]) -> list[Any]: + """Load the given skill refs, de-duplicated by skill name (last wins).""" + loaded = [self._load_dir(self._resolve_ref(ref)) for ref in refs] + by_name: dict[str, Any] = {} + for skill in loaded: + by_name[skill.name] = skill + return list(by_name.values()) diff --git a/src/agentic_cli/tools/skills/toolset.py b/src/agentic_cli/tools/skills/toolset.py new file mode 100644 index 0000000..31c3b47 --- /dev/null +++ b/src/agentic_cli/tools/skills/toolset.py @@ -0,0 +1,69 @@ +"""Build an ADK ``SkillToolset`` for a set of skills. + +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 + +from typing import Any + + +def make_skill_toolset( + skills: list[Any], + *, + code_executor: Any | None = None, + additional_tools: list[Any] | None = None, +) -> Any: + """Create an ADK SkillToolset; script execution follows the executor. + + Args: + skills: Loaded ADK ``Skill`` objects. + 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. + + Returns: + A configured ``SkillToolset``. + """ + from google.adk.tools.skill_toolset import RunSkillScriptTool, SkillToolset + + toolset = SkillToolset( + skills=skills, + code_executor=code_executor, + additional_tools=additional_tools or [], + ) + if code_executor is None: + 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 new file mode 100644 index 0000000..332e07d --- /dev/null +++ b/src/agentic_cli/tools/tool_resolver.py @@ -0,0 +1,150 @@ +"""Resolve tool references to callables. + +An ``AgentConfig.tools`` entry may be: + +- a callable (function tool) or any already-built tool object -> returned as-is; +- a **bare name** (no dot) -> looked up in the framework ``ToolRegistry`` + (e.g. ``"kb_search"``); +- a **dotted path** -> imported (e.g. ``"my_pkg.tools.my_tool"``). + +This mirrors ADK's own convention (bare name = built-in/registered, dotted = +import). Resolution must run before manager service-detection and tool +assembly, which key on ``tool.__name__``. +""" + +from __future__ import annotations + +import difflib +from importlib import import_module +from typing import Any, Callable + +from agentic_cli.tools.registry import ToolRegistry, get_registry + +_builtins_imported = False + + +def _ensure_builtin_tools_imported() -> None: + """Import ``agentic_cli.tools`` once so framework tools self-register. + + The standard tools register via ``@register_tool`` at import time; importing + the package top-level pulls them all in. Done lazily so callers that only + use callables or dotted paths don't pay the import cost. + """ + global _builtins_imported + if _builtins_imported: + return + import_module("agentic_cli.tools") + _builtins_imported = True + + +def _import_dotted(path: str) -> Any: + """Import an object from a dotted path (``module.sub:obj`` style as ``a.b.c``).""" + module_path, _, obj_name = path.rpartition(".") + if not module_path or not obj_name: + raise ValueError(f"Invalid dotted tool path: {path!r}") + try: + module = import_module(module_path) + except ImportError as exc: + raise ValueError( + f"Cannot import module {module_path!r} for tool {path!r}: {exc}" + ) from exc + try: + return getattr(module, obj_name) + except AttributeError as exc: + raise ValueError( + f"Module {module_path!r} has no attribute {obj_name!r} (tool {path!r})" + ) from exc + + +def _unknown_name_error(name: str, registry: ToolRegistry) -> ValueError: + """Build a helpful error for an unresolved bare tool name.""" + names = [d.name for d in registry.list_tools()] + close = difflib.get_close_matches(name, names, n=3) + hint = f" Did you mean: {', '.join(close)}?" if close else "" + return ValueError( + f"Unknown tool name {name!r} (not in the tool registry).{hint} " + "Use a fully-qualified dotted path for custom tools." + ) + + +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, +) -> Any: + """Resolve a single tool reference to a callable/tool object. + + Args: + ref: A callable, an already-built tool object, a bare registry name, or + a dotted import path. + registry: Registry for bare-name lookups. Defaults to the global + framework registry (auto-populated with built-in tools on miss). + + Returns: + The resolved callable or tool object. + + Raises: + ValueError: If a string ref cannot be resolved. + """ + # Callables and any already-built tool objects pass through unchanged. + if not isinstance(ref, str): + return ref + + name = ref.strip() + if not name: + raise ValueError("Empty tool reference.") + + if "." in name: + return _import_dotted(name) + + # Bare name -> registry lookup. + # 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. + _ensure_builtin_tools_imported() + reg = get_registry() + 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 + + +def resolve_tools( + refs: list[Callable[..., Any] | str] | None, + registry: ToolRegistry | None = None, +) -> list[Any]: + """Resolve a list of tool references (``None`` -> empty list).""" + return [resolve_tool(ref, registry) for ref in (refs or [])] diff --git a/src/agentic_cli/tools/webfetch/fetcher.py b/src/agentic_cli/tools/webfetch/fetcher.py index a5a8a98..bde8ed9 100644 --- a/src/agentic_cli/tools/webfetch/fetcher.py +++ b/src/agentic_cli/tools/webfetch/fetcher.py @@ -2,20 +2,19 @@ from __future__ import annotations -import socket import time from dataclasses import dataclass from urllib.parse import urljoin, urlparse import httpx -from agentic_cli.tools.webfetch.validator import URLValidator +from agentic_cli.tools.webfetch.validator import URLValidator, BlockedAddressError from agentic_cli.tools.webfetch.robots import RobotsTxtChecker +from agentic_cli.tools.webfetch.transport import PinnedTransport @dataclass class RedirectInfo: - """Information about a cross-host redirect.""" from_url: str to_url: str to_host: str @@ -23,7 +22,6 @@ class RedirectInfo: @dataclass class FetchResult: - """Result of a content fetch operation.""" success: bool content: str | bytes | None = None content_type: str | None = None @@ -35,7 +33,6 @@ class FetchResult: @dataclass class CachedResponse: - """Cached HTTP response.""" content: str | bytes content_type: str timestamp: float @@ -43,186 +40,119 @@ class CachedResponse: class ContentFetcher: - """Fetches web content with validation, caching, and redirect handling.""" + """Fetches web content with SSRF-safe pinning, caching, and redirect handling.""" + + MAX_REDIRECTS = 5 def __init__( self, validator: URLValidator, robots_checker: RobotsTxtChecker, + transport: PinnedTransport, cache_ttl_seconds: int = 900, max_content_bytes: int = 102400, max_pdf_bytes: int = 5242880, ) -> None: self._validator = validator self._robots = robots_checker + self._transport = transport self._cache_ttl = cache_ttl_seconds self._max_content_bytes = max_content_bytes self._max_pdf_bytes = max_pdf_bytes self._cache: dict[str, CachedResponse] = {} - MAX_REDIRECTS = 5 - async def fetch(self, url: str, timeout: int = 30) -> FetchResult: """Fetch content from a URL. - Redirects are followed manually so each Location is revalidated by - the SSRF guard before the next request is issued — using - ``follow_redirects=True`` would let httpx contact intermediate hosts - (e.g., 169.254.169.254) before we get a chance to inspect them. - - Args: - url: The URL to fetch. - timeout: Request timeout in seconds. - - Returns: - FetchResult with content or error information. + Redirects are followed manually so each Location is revalidated before + the next request; the PinnedTransport resolves+validates+pins every hop + (and the robots fetch), connecting only to globally-routable IPs. """ - # Check cache first cached = self._get_cached(url) if cached is not None: return FetchResult( - success=True, - content=cached.content, - content_type=cached.content_type, - truncated=cached.truncated, - from_cache=True, + success=True, content=cached.content, content_type=cached.content_type, + truncated=cached.truncated, from_cache=True, ) - # Validate the originally requested URL validation = self._validator.validate(url) if not validation.valid: return FetchResult(success=False, error=validation.error) - # Check robots.txt for the originally requested URL if not await self._robots.can_fetch(url): return FetchResult(success=False, error=f"Blocked by robots.txt for {url}") original_url = url current_url = url - history: list[str] = [] - response: httpx.Response | None = None try: - async with httpx.AsyncClient(follow_redirects=False) as client: + async with httpx.AsyncClient(transport=self._transport, follow_redirects=False) as client: for _ in range(self.MAX_REDIRECTS + 1): - response = await client.get(current_url, timeout=timeout) - - # Post-fetch IP revalidation (mitigates DNS rebinding) - final_host = urlparse(str(response.url)).hostname - if final_host: - try: - post_ip = socket.gethostbyname(final_host) - ip_check = self._validator.validate_ip(post_ip) - if not ip_check.valid: + async with client.stream("GET", current_url, timeout=timeout) as response: + status = response.status_code + headers = response.headers + content_type = headers.get("content-type", "text/html") + is_pdf = "application/pdf" in content_type.lower() + cap = self._max_pdf_bytes if is_pdf else self._max_content_bytes + encoding = response.charset_encoding or "utf-8" + buf = bytearray() + truncated = False + async for chunk in response.aiter_bytes(): + buf.extend(chunk) + if len(buf) > cap: + del buf[cap:] + truncated = True + break + + if status in (301, 302, 303, 307, 308): + location = headers.get("location") + if location: + next_url = urljoin(current_url, location) + next_validation = self._validator.validate(next_url) + if not next_validation.valid: return FetchResult( success=False, - error=f"DNS rebinding detected: {ip_check.error}", + error=f"Redirect to disallowed URL blocked: {next_validation.error}", ) - except socket.gaierror: - pass - - if response.status_code not in (301, 302, 303, 307, 308): - break - - location = response.headers.get("location") - if not location: - break - - next_url = urljoin(str(response.url), location) - next_validation = self._validator.validate(next_url) - if not next_validation.valid: - return FetchResult( - success=False, - error=( - f"Redirect to disallowed URL blocked: " - f"{next_validation.error}" - ), - ) - - # Cross-host redirects must be reported BEFORE issuing - # the next GET — the user's permission grant covers the - # original host only, so contacting another origin is a - # side effect they did not approve. - next_host = urlparse(next_url).netloc - original_host = urlparse(original_url).netloc - if next_host.lower() != original_host.lower(): - return FetchResult( - success=False, - redirect=RedirectInfo( - from_url=original_url, - to_url=next_url, - to_host=next_host, - ), - error=f"Cross-host redirect to {next_host}", - ) - - # Same-host redirect: the new path may itself be - # disallowed by robots.txt even though the original - # path was not. - if not await self._robots.can_fetch(next_url): - return FetchResult( - success=False, - error=f"Blocked by robots.txt for {next_url}", - ) - - history.append(current_url) - current_url = next_url - else: + next_host = urlparse(next_url).netloc + original_host = urlparse(original_url).netloc + if next_host.lower() != original_host.lower(): + return FetchResult( + success=False, + redirect=RedirectInfo(original_url, next_url, next_host), + error=f"Cross-host redirect to {next_host}", + ) + if not await self._robots.can_fetch(next_url): + return FetchResult(success=False, error=f"Blocked by robots.txt for {next_url}") + current_url = next_url + continue + + # Final response (non-redirect, or redirect without Location). + if is_pdf: + content: str | bytes = bytes(buf) + else: + content = buf.decode(encoding, errors="replace") + if truncated: + content += f"\n\n[Content truncated at {cap} bytes]" + self._cache[original_url] = CachedResponse( + content=content, content_type=content_type, + timestamp=time.time(), truncated=truncated, + ) return FetchResult( - success=False, - error=f"Too many redirects (max {self.MAX_REDIRECTS})", + success=True, content=content, content_type=content_type, + truncated=truncated, from_cache=False, ) + else: + return FetchResult(success=False, error=f"Too many redirects (max {self.MAX_REDIRECTS})") + except BlockedAddressError as e: + return FetchResult(success=False, error=f"Blocked (SSRF): {e}") except httpx.TimeoutException: return FetchResult(success=False, error=f"Request timeout after {timeout}s") except httpx.RequestError as e: return FetchResult(success=False, error=f"Request failed: {e}") - assert response is not None # loop runs at least once - - content_type = response.headers.get("content-type", "text/html") - - # Use bytes for PDF, text for everything else - is_pdf = "application/pdf" in content_type.lower() - if is_pdf: - content: str | bytes = response.content - truncated = False - if len(content) > self._max_pdf_bytes: - content = content[:self._max_pdf_bytes] - truncated = True - else: - content = response.text - truncated = False - if len(content) > self._max_content_bytes: - content = content[:self._max_content_bytes] - content += f"\n\n[Content truncated at {self._max_content_bytes} bytes]" - truncated = True - - # Cache the response under the originally requested URL - self._cache[original_url] = CachedResponse( - content=content, - content_type=content_type, - timestamp=time.time(), - truncated=truncated, - ) - - return FetchResult( - success=True, - content=content, - content_type=content_type, - truncated=truncated, - from_cache=False, - ) - def _get_cached(self, url: str) -> CachedResponse | None: - """Get cached response if available and not expired. - - Args: - url: The URL to look up. - - Returns: - CachedResponse if found and valid, None otherwise. - """ if url not in self._cache: return None cached = self._cache[url] @@ -232,5 +162,4 @@ def _get_cached(self, url: str) -> CachedResponse | None: return cached def clear_cache(self) -> None: - """Clear the response cache.""" self._cache.clear() diff --git a/src/agentic_cli/tools/webfetch/robots.py b/src/agentic_cli/tools/webfetch/robots.py index c66c396..6c0862c 100644 --- a/src/agentic_cli/tools/webfetch/robots.py +++ b/src/agentic_cli/tools/webfetch/robots.py @@ -9,67 +9,42 @@ class RobotsTxtChecker: - """Checks robots.txt compliance for URLs. - - Fetches and caches robots.txt per domain, then checks if our - user agent is allowed to access specific paths. - """ + """Checks robots.txt compliance, fetching robots.txt through the shared + SSRF-safe PinnedTransport (so the robots request is resolved-validated-pinned + like every other webfetch request).""" USER_AGENT = "AgenticCLI/1.0" + _MAX_ROBOTS_BYTES = 512 * 1024 - def __init__(self) -> None: - """Initialize the checker with empty cache.""" + def __init__(self, transport: httpx.AsyncBaseTransport | None = None) -> None: + self._transport = transport self._cache: dict[str, RobotFileParser | None] = {} async def can_fetch(self, url: str) -> bool: - """Check if URL can be fetched according to robots.txt. - - Args: - url: The URL to check. - - Returns: - True if allowed (or on error), False if explicitly disallowed. - """ parsed = urlparse(url) domain = f"{parsed.scheme}://{parsed.netloc}" - - # Get or fetch robots.txt for this domain if domain not in self._cache: self._cache[domain] = await self._fetch_robots(domain) - parser = self._cache[domain] if parser is None: - # No robots.txt or error fetching - be permissive - return True - + return True # no robots.txt / fetch error → permissive return parser.can_fetch(self.USER_AGENT, url) async def _fetch_robots(self, domain: str) -> RobotFileParser | None: - """Fetch and parse robots.txt for a domain. - - Args: - domain: The domain (scheme + netloc) to fetch robots.txt for. - - Returns: - Parsed RobotFileParser, or None if not found/error. - """ robots_url = f"{domain}/robots.txt" - try: - async with httpx.AsyncClient() as client: - response = await client.get(robots_url, timeout=10.0) - - if response.status_code != 200: - return None - - parser = RobotFileParser() - parser.parse(response.text.splitlines()) - return parser - + async with httpx.AsyncClient(transport=self._transport) as client: + async with client.stream("GET", robots_url, timeout=10.0) as response: + if response.status_code != 200: + return None + buf = bytearray() + async for chunk in response.aiter_bytes(): + buf.extend(chunk) + if len(buf) > self._MAX_ROBOTS_BYTES: + break + text = buf.decode(response.charset_encoding or "utf-8", errors="replace") + parser = RobotFileParser() + parser.parse(text.splitlines()) + return parser except Exception: - # Network errors, timeouts, etc. - be permissive - return None - - def clear_cache(self) -> None: - """Clear the robots.txt cache.""" - self._cache.clear() + return None # network / SSRF-block / timeout → permissive diff --git a/src/agentic_cli/tools/webfetch/transport.py b/src/agentic_cli/tools/webfetch/transport.py new file mode 100644 index 0000000..90bd7e4 --- /dev/null +++ b/src/agentic_cli/tools/webfetch/transport.py @@ -0,0 +1,52 @@ +"""SSRF-safe HTTP transport: resolve → validate → pin every request.""" + +from __future__ import annotations + +import httpx + +from agentic_cli.tools.webfetch.validator import URLValidator, BlockedAddressError + + +class PinnedTransport(httpx.AsyncBaseTransport): + """Connect every request to a validated pinned IP. + + Wraps an inner transport (default httpx.AsyncHTTPTransport(); a + httpx.MockTransport in tests). Resolves the host (all A+AAAA), requires all + addresses globally routable, and rewrites the request to connect to the + literal validated IP while preserving the Host header and TLS SNI — so cert + verification stays bound to the hostname and the validated address IS the + connected address (no TOCTOU). Sitting at the transport layer, it covers the + initial fetch, every manually-issued redirect hop, and the robots fetch. + """ + + def __init__( + self, + validator: URLValidator, + inner: httpx.AsyncBaseTransport | None = None, + ) -> None: + self._validator = validator + self._inner = inner if inner is not None else httpx.AsyncHTTPTransport() + + def _pin(self, request: httpx.Request) -> None: + host = request.url.host + port = request.url.port or (443 if request.url.scheme == "https" else 80) + pinned_ip = self._validator.resolve_and_validate(host, port) # raises on block + # The client set the Host header from the original URL before this + # transport runs; rewriting url.host to the IP does not touch it. Keep + # Host + TLS SNI bound to the hostname. + request.url = request.url.copy_with(host=pinned_ip) + request.headers.setdefault("Host", host if port in (80, 443) else f"{host}:{port}") + request.extensions["sni_hostname"] = host + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self._pin(request) + return await self._inner.handle_async_request(request) + + async def aclose(self) -> None: + # No-op: this transport is SHARED and long-lived (one per fetcher, used + # by both the fetcher and the robots checker). Each webfetch call opens + # a short-lived httpx.AsyncClient(transport=self) and closes it on exit, + # which would otherwise tear down the shared inner pool under a + # concurrent call. The inner pool persists for the fetcher's lifetime + # (which also enables connection keep-alive across fetches). + return None diff --git a/src/agentic_cli/tools/webfetch/validator.py b/src/agentic_cli/tools/webfetch/validator.py index f78d1ce..840e987 100644 --- a/src/agentic_cli/tools/webfetch/validator.py +++ b/src/agentic_cli/tools/webfetch/validator.py @@ -7,6 +7,14 @@ from dataclasses import dataclass from urllib.parse import urlparse +import httpx + + +class BlockedAddressError(httpx.RequestError): + """A host resolves to a non-global / blocked address (or cannot be + resolved). Subclasses httpx.RequestError so the fetcher's existing + ``except httpx.RequestError`` maps it to a FetchResult error.""" + @dataclass class ValidationResult: @@ -17,137 +25,118 @@ class ValidationResult: resolved_ip: str | None = None +# Ranges ipaddress.is_global marks global but that are SSRF-risky. +_SUPPLEMENTARY_BLOCKED = [ + ipaddress.ip_network("64:ff9b::/96"), # NAT64 well-known prefix (embeds v4) + ipaddress.ip_network("64:ff9b:1::/48"), # NAT64 local-use prefix + ipaddress.ip_network("192.88.99.0/24"), # 6to4 relay anycast (deprecated) +] + + +def _ip_is_safe(ip_obj: ipaddress._BaseAddress) -> bool: + """True only if the address is globally routable and not SSRF-risky.""" + if not ip_obj.is_global: + return False + for net in _SUPPLEMENTARY_BLOCKED: + if ip_obj in net: + return False + mapped = getattr(ip_obj, "ipv4_mapped", None) + if mapped is not None and not _ip_is_safe(mapped): + return False + return True + + +def _as_ip_literal(host: str) -> ipaddress._BaseAddress | None: + """Return the IP if host is an IP literal (brackets stripped), else None.""" + try: + return ipaddress.ip_address(host.strip("[]")) + except ValueError: + return None + + class URLValidator: - """Validates URLs for security (SSRF protection) and policy compliance. + """Validates URLs for SSRF protection and policy compliance. - Checks: - - Allowed schemes (http, https only) - - Private/internal IP addresses blocked - - Configurable domain blocklist with wildcard support + ``validate`` performs non-DNS policy (scheme, blocked domains, IP-literal + safety). Authoritative resolution + pinning happens in + ``resolve_and_validate`` (used by PinnedTransport), which resolves every + A+AAAA and requires all of them to be globally routable. """ ALLOWED_SCHEMES = {"http", "https"} - BLOCKED_NETWORKS = [ - ipaddress.ip_network("127.0.0.0/8"), # Loopback - ipaddress.ip_network("10.0.0.0/8"), # Private A - ipaddress.ip_network("172.16.0.0/12"), # Private B - ipaddress.ip_network("192.168.0.0/16"), # Private C - ipaddress.ip_network("169.254.0.0/16"), # Link-local - ipaddress.ip_network("::1/128"), # IPv6 loopback - ipaddress.ip_network("fc00::/7"), # IPv6 private - ipaddress.ip_network("fe80::/10"), # IPv6 link-local - ] - def __init__(self, blocked_domains: list[str] | None = None) -> None: - """Initialize validator. - - Args: - blocked_domains: List of domains to block. Supports wildcards (*.example.com). - """ self.blocked_domains = blocked_domains or [] def validate(self, url: str) -> ValidationResult: - """Validate a URL for fetching. + """Scheme + blocked-domain + hostname-present + IP-literal safety. - Args: - url: The URL to validate. - - Returns: - ValidationResult with valid=True if OK, or valid=False with error message. + Does NOT resolve DNS — hostname resolution is done (once) by the + transport's resolve_and_validate, which pins the connection. """ - # Parse URL try: parsed = urlparse(url) except Exception as e: return ValidationResult(valid=False, error=f"Malformed URL: {e}") - # Check scheme if parsed.scheme not in self.ALLOWED_SCHEMES: return ValidationResult( valid=False, error=f"Scheme '{parsed.scheme}' not allowed. Use http or https.", ) - # Check hostname exists hostname = parsed.hostname if not hostname: return ValidationResult(valid=False, error="URL must have a hostname") - # Check blocked domains if self._is_domain_blocked(hostname): - return ValidationResult( - valid=False, - error=f"Domain '{hostname}' is blocked by policy", - ) - - # Resolve hostname and check for private IPs - try: - ip_str = socket.gethostbyname(hostname) - ip = ipaddress.ip_address(ip_str) + return ValidationResult(valid=False, error=f"Domain '{hostname}' is blocked by policy") - for network in self.BLOCKED_NETWORKS: - if ip in network: - return ValidationResult( - valid=False, - error=f"Private/internal IP address blocked: {ip_str}", - resolved_ip=ip_str, - ) - - return ValidationResult(valid=True, resolved_ip=ip_str) - - except socket.gaierror as e: + literal = _as_ip_literal(hostname) + if literal is not None and not _ip_is_safe(literal): return ValidationResult( valid=False, - error=f"Could not resolve hostname '{hostname}': {e}", + error=f"Private/internal IP address blocked: {literal}", + resolved_ip=str(literal), ) - def validate_ip(self, ip_str: str) -> ValidationResult: - """Validate an IP address against blocked networks. + return ValidationResult(valid=True) - Args: - ip_str: The IP address string to validate. + def resolve_and_validate(self, host: str, port: int) -> str: + """Resolve every A+AAAA for host and return one validated pinned IP. - Returns: - ValidationResult with valid=True if OK, or valid=False if blocked. + Rejects the whole host (BlockedAddressError) if ANY resolved address is + unsafe, so a split-horizon / rebinding resolver cannot smuggle a private + address alongside a public one. An IP-literal host is validated directly. """ - try: - ip = ipaddress.ip_address(ip_str) - except ValueError as e: - return ValidationResult(valid=False, error=f"Invalid IP address: {e}") - - for network in self.BLOCKED_NETWORKS: - if ip in network: - return ValidationResult( - valid=False, - error=f"Private/internal IP address blocked: {ip_str}", - resolved_ip=ip_str, - ) + literal = _as_ip_literal(host) + if literal is not None: + if not _ip_is_safe(literal): + raise BlockedAddressError(f"blocked non-global address: {host}") + return str(literal) - return ValidationResult(valid=True, resolved_ip=ip_str) + try: + infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) + except socket.gaierror as exc: + raise BlockedAddressError(f"could not resolve {host}: {exc}") from exc + + # sockaddr[0] is the IP; strip any IPv6 zone id ("fe80::1%eth0"). + ips = [ipaddress.ip_address(info[4][0].split("%")[0]) for info in infos] + if not ips: + raise BlockedAddressError(f"no addresses for {host}") + for ip_obj in ips: + if not _ip_is_safe(ip_obj): + raise BlockedAddressError(f"blocked non-global address {ip_obj} for {host}") + return str(ips[0]) def _is_domain_blocked(self, hostname: str) -> bool: - """Check if hostname matches any blocked domain pattern. - - Args: - hostname: The hostname to check. - - Returns: - True if blocked, False otherwise. - """ hostname_lower = hostname.lower() - for pattern in self.blocked_domains: pattern_lower = pattern.lower() - - # Exact match if hostname_lower == pattern_lower: return True - - # Wildcard match (*.example.com matches sub.example.com) if pattern_lower.startswith("*."): suffix = pattern_lower[1:] # .example.com if hostname_lower.endswith(suffix) and hostname_lower != pattern_lower[2:]: return True - return False diff --git a/src/agentic_cli/tools/webfetch_tool.py b/src/agentic_cli/tools/webfetch_tool.py index b37fa07..d50d583 100644 --- a/src/agentic_cli/tools/webfetch_tool.py +++ b/src/agentic_cli/tools/webfetch_tool.py @@ -18,6 +18,7 @@ HTMLToMarkdown, build_summarize_prompt, ) +from agentic_cli.tools.webfetch.transport import PinnedTransport from agentic_cli.workflow.service_registry import require_service, LLM_SUMMARIZER @@ -60,11 +61,13 @@ def get_or_create_fetcher(settings=None) -> ContentFetcher: return _fetcher validator = URLValidator(blocked_domains=settings.webfetch_blocked_domains) - robots_checker = RobotsTxtChecker() + transport = PinnedTransport(validator) + robots_checker = RobotsTxtChecker(transport=transport) _fetcher = ContentFetcher( validator=validator, robots_checker=robots_checker, + transport=transport, cache_ttl_seconds=settings.webfetch_cache_ttl_seconds, max_content_bytes=settings.webfetch_max_content_bytes, max_pdf_bytes=settings.webfetch_max_pdf_bytes, @@ -78,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/__init__.py b/src/agentic_cli/workflow/__init__.py index 0d60607..27f38f4 100644 --- a/src/agentic_cli/workflow/__init__.py +++ b/src/agentic_cli/workflow/__init__.py @@ -12,7 +12,17 @@ # Light imports - always available (fast) from agentic_cli.workflow.events import WorkflowEvent, EventType, UserInputRequest from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.model_settings import ModelSettings, ThinkingSettings +from agentic_cli.workflow.mcp import MCPServerConfig from agentic_cli.workflow.factory import create_workflow_manager_from_settings +from agentic_cli.workflow.agent_loader import ( + load_agents_from_yaml, + create_workflow_manager_from_yaml, +) +from agentic_cli.workflow.adk_config_bridge import ( + load_adk_agent_native, + translate_adk_yaml, +) from agentic_cli.workflow.settings import WorkflowSettingsMixin from agentic_cli.workflow.models import ModelFamily, ModelInfo, ModelRegistry from agentic_cli.workflow.service_registry import get_service, get_service_registry @@ -52,8 +62,15 @@ def __getattr__(name: str): "UserInputRequest", # Factory "create_workflow_manager_from_settings", + "create_workflow_manager_from_yaml", + "load_agents_from_yaml", + "load_adk_agent_native", + "translate_adk_yaml", # Config "AgentConfig", + "ModelSettings", + "ThinkingSettings", + "MCPServerConfig", # Settings mixin "WorkflowSettingsMixin", # Model registry diff --git a/src/agentic_cli/workflow/adk/anthropic_llm.py b/src/agentic_cli/workflow/adk/anthropic_llm.py new file mode 100644 index 0000000..eef27b3 --- /dev/null +++ b/src/agentic_cli/workflow/adk/anthropic_llm.py @@ -0,0 +1,141 @@ +"""Direct-API Anthropic LLM for ADK, with support for extra request params. + +ADK's stock ``AnthropicLlm`` only forwards ``thinking`` to the Anthropic Messages +API — it never sends ``output_config`` (so the discrete effort ladder +``low|medium|high|xhigh|max`` is unreachable), and a bare ``claude-*`` model +string resolves through ``LLMRegistry`` to the *Vertex* ``Claude`` class (which +needs GOOGLE_CLOUD_PROJECT/LOCATION). + +This module provides two things: + +- ``DirectAnthropicLlm`` — a subclass that (a) uses the direct API + (``ANTHROPIC_API_KEY``, inherited from ``AnthropicLlm``) and (b) injects extra + ``messages.create`` kwargs — notably ``output_config={"effort": ...}`` — by + wrapping ADK's ``_anthropic_client`` property. That property is ADK's own + extension seam (the Vertex ``Claude`` class overrides the same property), so + we get ADK's full request-build + streaming aggregation for free. +- ``register_direct_anthropic()`` — points the ADK ``LLMRegistry`` at this class + for ``claude-*`` strings, so plain model strings in agent configs resolve to + the direct API instead of Vertex. + +Note: when a model is resolved *from a string* the registry constructs +``DirectAnthropicLlm(model=...)`` with class defaults (``max_tokens=8192``, +``effort=None``). Per-agent ``max_tokens``/``effort`` must be supplied by +constructing an instance directly (see ``manager._build_model_arg``); the +registry path is the convenience/safety net that fixes the Vertex gotcha. +""" + +from __future__ import annotations + +from functools import cached_property +from typing import Any + +from google.adk.models.anthropic_llm import AnthropicLlm +from google.adk.models.registry import LLMRegistry + +from agentic_cli.logging import Loggers + +logger = Loggers.workflow() + +# Patterns to claim in the registry. The first two mirror ADK's ``Claude`` +# (Vertex) patterns so registering overwrites those entries; the catch-all +# covers newer families (5.x, fable) that ADK's patterns miss. +_CLAUDE_PATTERNS = [r"claude-3-.*", r"claude-.*-4.*", r"claude-.*"] + + +class _ExtraParamMessages: + """Wraps an Anthropic ``messages`` resource, merging fixed extra kwargs into + every ``create`` call. Per-call kwargs win over the fixed extras.""" + + def __init__(self, messages: Any, extra: dict[str, Any]) -> None: + self._messages = messages + self._extra = extra + + def __getattr__(self, name: str) -> Any: + return getattr(self._messages, name) + + async def create(self, **kwargs: Any) -> Any: + merged: dict[str, Any] = {**self._extra, **kwargs} + return await self._messages.create(**merged) + + +class _ExtraParamClient: + """Wraps an AsyncAnthropic client, swapping in an extra-param ``messages``.""" + + def __init__(self, client: Any, extra: dict[str, Any]) -> None: + self._client = client + self._extra = extra + + def __getattr__(self, name: str) -> Any: + return getattr(self._client, name) + + @cached_property + def messages(self) -> _ExtraParamMessages: + return _ExtraParamMessages(self._client.messages, self._extra) + + +class DirectAnthropicLlm(AnthropicLlm): + """Anthropic via the direct API, with support for extra request params. + + Attributes: + effort: Anthropic ``output_config.effort`` (``low|medium|high|xhigh|max``). + Only valid on models that support it (Opus 4.5+/4.6/4.7/4.8, + Sonnet 4.6); Sonnet 4.5 / Haiku 4.5 reject it. Leave ``None`` to omit. + extra_params: Escape hatch merged into every ``messages.create`` call + (e.g. ``{"service_tier": "..."}``). ADK's own per-call kwargs + (model, messages, thinking, ...) always win over these. + request_timeout: Overall client timeout (seconds). Set to a non-default + value so ADK's non-streaming ``messages.create`` skips the SDK's + "streaming required" guard for ``max_tokens`` above ~21k (the guard + only fires when the client uses the default timeout). + max_retries: Client-level retry count (defaults to the anthropic SDK's + when ``None``); wire it from ``settings.retry_max_attempts``. + """ + + effort: str | None = None + extra_params: dict[str, Any] = {} + request_timeout: float | None = None + max_retries: int | None = None + + @staticmethod + def supported_models() -> list[str]: + return list(_CLAUDE_PATTERNS) + + def _extra_create_params(self) -> dict[str, Any]: + """The fixed kwargs to merge into ``messages.create`` for this instance.""" + extra: dict[str, Any] = dict(self.extra_params) + if self.effort: + output_config = dict(extra.get("output_config") or {}) + output_config.setdefault("effort", self.effort) + extra["output_config"] = output_config + return extra + + @cached_property + def _anthropic_client(self): # type: ignore[override] + from anthropic import AsyncAnthropic + + client_kwargs: dict[str, Any] = {} + if self.request_timeout is not None: + client_kwargs["timeout"] = self.request_timeout + if self.max_retries is not None: + client_kwargs["max_retries"] = self.max_retries + client = AsyncAnthropic(**client_kwargs) + extra = self._extra_create_params() + if not extra: + return client + return _ExtraParamClient(client, extra) + + +def register_direct_anthropic() -> None: + """Resolve ``claude-*`` strings to ``DirectAnthropicLlm`` via ``LLMRegistry``. + + Overrides ADK's default (the Vertex ``Claude`` class) for the existing + ``claude-3-.*`` / ``claude-.*-4.*`` patterns and adds a ``claude-.*`` + catch-all for newer families. Idempotent; clears the ``resolve()`` lru_cache + so already-resolved names pick up the change. + """ + LLMRegistry.register(DirectAnthropicLlm) + cache_clear = getattr(LLMRegistry.resolve, "cache_clear", None) + if cache_clear is not None: + cache_clear() + logger.debug("direct_anthropic_registered", patterns=_CLAUDE_PATTERNS) diff --git a/src/agentic_cli/workflow/adk/manager.py b/src/agentic_cli/workflow/adk/manager.py index fe495e8..c95730a 100644 --- a/src/agentic_cli/workflow/adk/manager.py +++ b/src/agentic_cli/workflow/adk/manager.py @@ -10,7 +10,6 @@ from __future__ import annotations -import json import logging from typing import AsyncGenerator, Any, Callable @@ -19,14 +18,16 @@ from google.adk.agents import LlmAgent, Agent from google.adk.planners import BuiltInPlanner from google.adk.sessions import InMemorySessionService, BaseSessionService, Session -from google.adk.events import Event +from google.adk.tools import LongRunningFunctionTool 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 from agentic_cli.workflow.adk.plugins import LLMLoggingPlugin +from agentic_cli.workflow.service_registry import JOB_MANAGER from agentic_cli.config import ( BaseSettings, @@ -41,6 +42,62 @@ # mixed responses. We already handle all part types individually in process_part. logging.getLogger("google_genai.types").setLevel(logging.ERROR) +_RESULT_SUMMARY_LIMIT = 1000 + +# Anthropic/Claude thinking on ADK (native AnthropicLlm, direct API). +# Effort → thinking budget in tokens (mirrors the LangGraph Claude budgets). +_CLAUDE_THINKING_BUDGETS = {"low": 4096, "medium": 10000, "high": 32000} +# Anthropic requires budget_tokens >= 1024 when thinking is enabled. +_ANTHROPIC_MIN_THINKING_BUDGET = 1024 +# Default max_tokens when no ModelSettings.max_tokens is given. +_DEFAULT_ANTHROPIC_MAX_TOKENS = 8192 +# Output headroom reserved above the thinking budget (Anthropic counts thinking +# tokens toward max_tokens and requires max_tokens > budget_tokens). +_ANTHROPIC_OUTPUT_HEADROOM = 8192 + + +def _is_anthropic_model(model: str | None) -> bool: + """True for Claude/Anthropic model ids (string check, no settings needed).""" + return bool(model) and model.startswith("claude") + + +# Claude >= 4.6 deprecates/removes ``budget_tokens`` (400 on 4.7+/Fable); those +# use adaptive thinking + ``output_config.effort``. <= 4.5 keep the legacy +# numeric-budget path (effort also unsupported on Sonnet 4.5 / Haiku 4.5). +_ANTHROPIC_ADAPTIVE_MIN = (4, 6) +# Generic thinking level -> Anthropic effort. ``xhigh``/``max`` are not exposed +# via the generic knob; use a per-agent native override for those. +_GENERIC_TO_EFFORT = {"low": "low", "medium": "medium", "high": "high"} + + +def _claude_version(model: str) -> tuple[int, ...]: + """Numeric version tuple from a Claude id (``claude-opus-4-8`` -> ``(4, 8)``). + + Version components are 1-2 digits; longer numeric segments are date stamps + (``claude-opus-4-20250514``) and must be ignored, else the date would be + read as a minor version and push the model past the adaptive threshold. + """ + return tuple(int(p) for p in model.split("-") if p.isdigit() and len(p) <= 2) + + +def _anthropic_uses_adaptive(model: str | None) -> bool: + """True for Claude models that require adaptive thinking (>= 4.6).""" + return _is_anthropic_model(model) and _claude_version(model) >= _ANTHROPIC_ADAPTIVE_MIN + + +def _summarize_result(result: Any) -> str | None: + """Render a job result as a short string for the resume payload. + + Truncates to keep the resumed turn's context lean; the agent pulls the full + result via the ``job_result`` tool when it needs it. + """ + if result is None: + return None + text = result if isinstance(result, str) else repr(result) + if len(text) > _RESULT_SUMMARY_LIMIT: + return text[:_RESULT_SUMMARY_LIMIT] + f"… (+{len(text) - _RESULT_SUMMARY_LIMIT} chars)" + return text + class GoogleADKWorkflowManager(BaseWorkflowManager): """Config-based workflow manager for agentic applications using Google ADK. @@ -80,6 +137,7 @@ def __init__( app_name: str | None = None, model: str | None = None, on_event: Callable[[WorkflowEvent], WorkflowEvent | None] | None = None, + adk_config_path: str | None = None, ) -> None: """Initialize the workflow manager. @@ -90,6 +148,10 @@ def __init__( app_name: Application name for services (uses settings.app_name if not provided) model: Model override (auto-detected from API keys if not provided) on_event: Optional hook to transform/filter events before yielding + adk_config_path: Optional path to a native ADK ``root_agent.yaml``. + When set, the agent tree is built via ADK ``from_config`` instead + of from ``agent_configs`` (full ADK fidelity; framework service / + state tool injection does not apply). """ super().__init__( agent_configs=agent_configs, @@ -99,8 +161,13 @@ def __init__( on_event=on_event, ) self.session_id = "default_session" + 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 @@ -114,6 +181,12 @@ def __init__( self._llm_logging_plugin: LLMLoggingPlugin | None = None self._task_progress_plugin: "TaskProgressPlugin | None" = None + # Resolve bare ``claude-*`` strings (e.g. from agent YAML / native ADK + # configs) to the direct-API DirectAnthropicLlm instead of Vertex Claude. + from agentic_cli.workflow.adk.anthropic_llm import register_direct_anthropic + + register_direct_anthropic() + logger.debug( "workflow_manager_created", app_name=self.app_name, @@ -141,6 +214,23 @@ async def generate_simple(self, prompt: str, max_tokens: int = 500) -> str: """ await self._ensure_initialized() + # Claude models are not served by the genai client — use the Anthropic + # SDK directly (ANTHROPIC_API_KEY is exported during initialization). + if _is_anthropic_model(self.model): + from anthropic import AsyncAnthropic + + client = AsyncAnthropic() + message = await client.messages.create( + model=self.model, + max_tokens=max_tokens, + messages=[{"role": "user", "content": prompt}], + ) + return "".join( + block.text + for block in message.content + if getattr(block, "type", None) == "text" + ) + from google import genai client = genai.Client() @@ -157,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 @@ -173,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") @@ -185,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", @@ -199,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", @@ -229,57 +350,98 @@ async def reinitialize( sessions_preserved=preserve_sessions, ) - def _get_planner(self) -> BuiltInPlanner | None: - """Get planner with thinking configuration. + def _resolve_model_for_config(self, config: "AgentConfig | None") -> str: + """Return the effective model for an agent (per-config override or default).""" + if config is not None and config.model: + return config.model + return self.model + + def _resolve_thinking( + self, config: "AgentConfig | None" + ) -> ThinkingSettings | None: + """Resolve thinking settings: per-agent override, else global effort. - Gemini 3 models take a discrete ``thinking_level``; Gemini 2.5 models - only understand a numeric ``thinking_budget`` and reject ``thinking_level`` - outright (HTTP 400 "Thinking level is not supported for this model"). - We therefore choose the field that matches the model generation — - sending ``thinking_level`` to a 2.5 model breaks every request. + Returns None when thinking is disabled (no per-agent setting and the + global ``thinking_effort`` is ``"none"``). """ - thinking_effort = self._settings.thinking_effort + ms = config.model_settings if config is not None else None + if ms is not None and ms.thinking is not None: + return ms.thinking + global_effort = self._settings.thinking_effort + if global_effort == "none": + return None + return ThinkingSettings(mode=global_effort) + + def _get_planner( + self, config: "AgentConfig | None" = None + ) -> BuiltInPlanner | None: + """Get planner with thinking configuration for an agent. + + Thinking is resolved per-agent (``config.model_settings.thinking``) with + a fallback to the global ``settings.thinking_effort``. - if thinking_effort == "none": + Per family: Claude >= 4.6 uses adaptive thinking (negative budget) + + effort; Claude <= 4.5 and Gemini 2.5 are numeric-budget-based; Gemini 3 + takes a discrete ``thinking_level`` (Gemini 2.5 rejects it with a 400). + """ + thinking = self._resolve_thinking(config) + if thinking is None or thinking.mode == "none": return None - if not self._settings.supports_thinking_effort(self.model): - logger.debug( - "thinking_not_supported", - model=self.model, - effort=thinking_effort, - ) + model = self._resolve_model_for_config(config) + + if not self._settings.supports_thinking_effort(model): + logger.debug("thinking_not_supported", model=model, mode=thinking.mode) return None - if "gemini-3" in self.model: - thinking_config = self._gemini3_thinking_config(thinking_effort) + if _is_anthropic_model(model): + if _anthropic_uses_adaptive(model): + # Claude >= 4.6 rejects budget_tokens; use adaptive thinking. ADK + # (>= 1.34) maps a negative thinking_budget to {type:"adaptive"}. + # Depth is controlled by output_config.effort, set on the model + # instance (see _anthropic_effort / DirectAnthropicLlm). + thinking_config = types.ThinkingConfig(thinking_budget=-1) + else: + # Legacy Claude (<= 4.5): numeric budget_tokens (>= 1024). The same + # budget sizes max_tokens in _anthropic_max_tokens (Anthropic + # requires max_tokens > budget). + thinking_config = types.ThinkingConfig( + include_thoughts=True, + thinking_budget=self._anthropic_thinking_budget(config), + ) + elif thinking.mode == "budget": + budget = ( + thinking.budget_tokens if thinking.budget_tokens is not None else 12288 + ) + thinking_config = types.ThinkingConfig( + include_thoughts=True, thinking_budget=budget + ) + elif "gemini-3" in model: + thinking_config = self._gemini3_thinking_config(thinking.mode, model) else: - thinking_config = self._gemini25_thinking_config(thinking_effort) - - logger.debug( - "planner_created", - model=self.model, - effort=thinking_effort, - ) + thinking_config = self._gemini25_thinking_config(thinking.mode) + logger.debug("planner_created", model=model, mode=thinking.mode) return BuiltInPlanner(thinking_config=thinking_config) - def _gemini3_thinking_config(self, thinking_effort: str) -> "types.ThinkingConfig": + def _gemini3_thinking_config( + self, effort: str, model: str + ) -> "types.ThinkingConfig": """Build a Gemini 3 thinking config using the discrete ``thinking_level``. Gemini 3 Pro supports only LOW and HIGH (MEDIUM falls back to HIGH); Gemini 3 Flash additionally supports MINIMAL/MEDIUM. """ - is_pro = "pro" in self.model + is_pro = "pro" in model - if thinking_effort == "low": + if effort == "low": level = types.ThinkingLevel.LOW - elif thinking_effort == "medium": + elif effort == "medium": if is_pro: level = types.ThinkingLevel.HIGH logger.debug( "thinking_level_fallback", - model=self.model, + model=model, requested="medium", actual="high", reason="Gemini 3 Pro only supports LOW and HIGH", @@ -291,21 +453,101 @@ def _gemini3_thinking_config(self, thinking_effort: str) -> "types.ThinkingConfi return types.ThinkingConfig(include_thoughts=True, thinking_level=level) - def _gemini25_thinking_config(self, thinking_effort: str) -> "types.ThinkingConfig": + def _gemini25_thinking_config(self, effort: str) -> "types.ThinkingConfig": """Build a Gemini 2.5 thinking config using the numeric ``thinking_budget``. 2.5 models reject ``thinking_level``. Budgets are chosen within the range valid across the 2.5 family (Flash/Pro/Flash-Lite all accept 4096–24576 tokens). """ - budget = {"low": 4096, "medium": 12288, "high": 24576}[thinking_effort] + budget = {"low": 4096, "medium": 12288, "high": 24576}[effort] return types.ThinkingConfig(include_thoughts=True, thinking_budget=budget) - def _get_generate_content_config(self) -> types.GenerateContentConfig: - """Build GenerateContentConfig with retry-enabled HTTP options. + def _anthropic_thinking_budget(self, config: "AgentConfig | None") -> int: + """Resolved Anthropic thinking budget in tokens for an agent. + + Mirrors the LangGraph Claude budgets; ``budget`` mode passes the explicit + value through. Floored at Anthropic's minimum (1024). Returns 0 when + thinking is disabled. The same number sizes the agent's ``max_tokens``. + """ + thinking = self._resolve_thinking(config) + if thinking is None or thinking.mode == "none": + return 0 + if thinking.mode == "budget": + budget = thinking.budget_tokens if thinking.budget_tokens is not None else 12288 + else: + budget = _CLAUDE_THINKING_BUDGETS.get(thinking.mode, 10000) + return max(budget, _ANTHROPIC_MIN_THINKING_BUDGET) + + def _anthropic_effort(self, config: "AgentConfig | None") -> str | None: + """Anthropic ``output_config.effort`` for an agent, or ``None``. + + Only adaptive-capable Claude (>= 4.6) takes effort; older models reject + it and stay on the budget path. Maps the generic ``low/medium/high`` + level; ``none``/``budget`` modes yield no effort (disabled / pure + adaptive — the model picks the depth). + """ + if not _anthropic_uses_adaptive(self._resolve_model_for_config(config)): + return None + thinking = self._resolve_thinking(config) + if thinking is None or thinking.mode in ("none", "budget"): + return None + return _GENERIC_TO_EFFORT.get(thinking.mode) + + def _anthropic_max_tokens(self, config: "AgentConfig | None") -> int: + """``max_tokens`` for an Anthropic agent. + + ADK's ``AnthropicLlm`` reads ``max_tokens`` from the model instance (not + ``GenerateContentConfig``). On the legacy budget path Anthropic counts + thinking tokens toward ``max_tokens`` and requires ``max_tokens > budget``, + so we size it above the budget. Adaptive thinking has no fixed budget, so + we just honour the ``ModelSettings.max_tokens`` floor (or the default). + """ + ms = config.model_settings if config is not None else None + base = ms.max_tokens if (ms is not None and ms.max_tokens) else _DEFAULT_ANTHROPIC_MAX_TOKENS + if _anthropic_uses_adaptive(self._resolve_model_for_config(config)): + return base + budget = self._anthropic_thinking_budget(config) + if budget: + return max(base, budget + _ANTHROPIC_OUTPUT_HEADROOM) + return base + + def _build_model_arg(self, config: "AgentConfig | None"): + """Resolve the ``model`` argument for an ``LlmAgent``. + + Anthropic/Claude models are returned as a direct-API ``DirectAnthropicLlm`` + instance: passing the bare ``claude-*`` string would make ADK's + ``LLMRegistry`` resolve it to the Vertex ``Claude`` class (which needs + GOOGLE_CLOUD_PROJECT/LOCATION). The instance carries the coordinated + ``max_tokens`` and, for adaptive models, ``effort``. Gemini/Gemma pass + through as plain strings for native registry resolution. + """ + model = self._resolve_model_for_config(config) + if not _is_anthropic_model(model): + return model + from agentic_cli.workflow.adk.anthropic_llm import DirectAnthropicLlm + + return DirectAnthropicLlm( + model=model, + max_tokens=self._anthropic_max_tokens(config), + effort=self._anthropic_effort(config), + # An explicit (non-default) timeout keeps large-max_tokens requests + # off the SDK's streaming-required guard; retries honour settings. + request_timeout=getattr(self._settings, "anthropic_request_timeout", 900.0), + max_retries=getattr(self._settings, "retry_max_attempts", None), + ) + + def _get_generate_content_config( + self, config: "AgentConfig | None" = None + ) -> types.GenerateContentConfig: + """Build GenerateContentConfig with retry HTTP options and per-agent params. + + Args: + config: Agent config whose ``model_settings`` (if any) supply + generation params merged on top of the retry options. Returns: - GenerateContentConfig with HttpRetryOptions configured from settings. + GenerateContentConfig with HttpRetryOptions plus any per-agent params. """ http_options = types.HttpOptions( retry_options=types.HttpRetryOptions( @@ -315,7 +557,40 @@ def _get_generate_content_config(self) -> types.GenerateContentConfig: http_status_codes=[500, 502, 503, 504], # Don't auto-retry 429 ) ) - return types.GenerateContentConfig(http_options=http_options) + kwargs: dict[str, Any] = {"http_options": http_options} + ms = config.model_settings if config is not None else None + if ms is not None: + kwargs.update(self._generate_config_kwargs_from_settings(ms)) + return types.GenerateContentConfig(**kwargs) + + def _generate_config_kwargs_from_settings( + self, ms: ModelSettings + ) -> dict[str, Any]: + """Translate neutral ModelSettings into GenerateContentConfig kwargs. + + Maps neutral field names (``max_tokens`` -> ``max_output_tokens``) and + passes ``extra`` through, filtered to valid GenerateContentConfig fields + (unknown keys are logged and dropped). Thinking is handled by the planner. + """ + out: dict[str, Any] = {} + if ms.temperature is not None: + out["temperature"] = ms.temperature + if ms.top_p is not None: + out["top_p"] = ms.top_p + if ms.top_k is not None: + out["top_k"] = ms.top_k + if ms.max_tokens is not None: + out["max_output_tokens"] = ms.max_tokens + if ms.stop_sequences is not None: + out["stop_sequences"] = ms.stop_sequences + if ms.extra: + valid = set(types.GenerateContentConfig.model_fields) + for key, value in ms.extra.items(): + if key in valid: + out[key] = value + else: + logger.warning("model_settings_extra_ignored", key=key) + return out def _get_state_tools(self) -> list: """Return ADK-native state tools using ToolContext.state.""" @@ -324,79 +599,148 @@ def _get_state_tools(self) -> list: ) return [save_plan, get_plan, save_tasks, get_tasks] + def _assemble_agent_tools( + self, config: "AgentConfig", service_map: dict + ) -> list: + """Build an agent's tools: framework tools + state tools + MCP toolsets. + + MCP servers declared on the config are materialized into ADK + ``MCPToolset`` objects (ADK connects lazily) and appended after the + regular tools. Skills are appended as a ``SkillToolset``. The caller is + responsible for ``_wrap_long_running`` over the result. + """ + tools = self._build_tools(config, service_map) + mcp_servers = getattr(config, "mcp_servers", None) or [] + if mcp_servers: + from agentic_cli.workflow.mcp import to_adk_toolset + + for server in mcp_servers: + tools.append(to_adk_toolset(server)) + logger.debug( + "mcp_toolset_attached", agent=config.name, server=server.name + ) + + skill_refs = getattr(config, "skills", None) or [] + if skill_refs: + toolset = self._build_skill_toolset(skill_refs) + if toolset is not None: + tools.append(toolset) + logger.debug("skill_toolset_attached", agent=config.name) + return tools + + def _build_skill_toolset(self, skill_refs: list[str]): + """Resolve skill refs and build an ADK SkillToolset (discovery/read only). + + 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 + + store = SkillStore(getattr(self._settings, "skills_dirs", []) or []) + skills = store.resolve(skill_refs) + if not skills: + return None + return make_skill_toolset(skills) + + def _wrap_long_running(self, tools: list[Callable]) -> list: + """Wrap tools flagged ``long_running`` as ADK ``LongRunningFunctionTool``. + + 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 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 identify_tool + + wrapped: list = [] + for tool in tools: + 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 + and not isinstance(tool, LongRunningFunctionTool) + ): + wrapped.append(LongRunningFunctionTool(func=tool)) + logger.debug("long_running_tool_wrapped", tool=name) + else: + 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] = {} - planner = self._get_planner() - generate_config = self._get_generate_content_config() 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=config.model or self.model, - instruction=config.get_prompt(), - tools=self._build_tools(config, service_map), - description=config.description or None, - planner=planner, - generate_content_config=generate_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=config.model or self.model, - instruction=config.get_prompt(), - tools=self._build_tools(config, service_map), - description=config.description or None, - sub_agents=sub_agent_instances, - planner=planner, - generate_content_config=generate_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 @@ -407,9 +751,19 @@ def _init_plugins(self) -> list: List of BasePlugin instances to pass to Runner(plugins=...). """ from agentic_cli.workflow.adk.task_progress_plugin import TaskProgressPlugin + from agentic_cli.workflow.adk.transfer_tool_description import ( + TransferToolDescriptionPlugin, + ) plugins: list = [PermissionPlugin()] + # ADK's generated description for its own transfer tool tells the model + # to call `TransferToAgentTool` — a name that does not exist as a tool. + # Corrected on the prepared request; a no-op once ADK ships a fixed + # docstring. See transfer_tool_description for why this is safe. + self._transfer_description_plugin = TransferToolDescriptionPlugin() + plugins.append(self._transfer_description_plugin) + # Task progress tracking via ToolContext.state self._task_progress_plugin = TaskProgressPlugin() plugins.append(self._task_progress_plugin) @@ -428,16 +782,51 @@ def _init_plugins(self) -> list: return plugins + def _make_session_service(self) -> BaseSessionService: + """Build the ADK session service from ``session_store`` settings. + + ``memory`` → ephemeral ``InMemorySessionService``; ``sqlite``/``postgres`` + → native ``DatabaseSessionService`` (durable across restarts, full event + fidelity). The sqlite parent dir is created so the engine can open it. + """ + db_url = self._settings.session_db_url() + if db_url is None: + logger.debug("using_in_memory_session_service") + return InMemorySessionService() + + if db_url.startswith("sqlite"): + # sqlite+aiosqlite:/// — ensure the directory exists. + from pathlib import Path + + path = db_url.split(":///", 1)[-1] + if path: + Path(path).parent.mkdir(parents=True, exist_ok=True) + + from google.adk.sessions import DatabaseSessionService + + logger.info("using_database_session_service", store=self._settings.session_store) + return DatabaseSessionService(db_url=db_url) + 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 - self._session_service = InMemorySessionService() - logger.debug("using_in_memory_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: + from agentic_cli.workflow.adk_config_bridge import load_adk_agent_native - # Create agent hierarchy from configs - self._root_agent = self._create_agents() + self._root_agent = load_adk_agent_native(self._adk_config_path) + logger.info("adk_native_config_loaded", path=self._adk_config_path) + else: + self._root_agent = self._create_agents() # Create runner with plugins self._runner = Runner( @@ -454,15 +843,42 @@ 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.""" + """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) # ------------------------------------------------------------------------- @@ -514,6 +930,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 @@ -522,84 +944,74 @@ 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 handling - session = 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)], - ) - - event_count = 0 - - # Build run_config with context window compression if enabled - run_config = None - if self._settings.context_window_enabled: - from google.genai.types import ContextWindowCompressionConfig, SlidingWindow - from google.adk.agents import RunConfig - - run_config = RunConfig( - context_window_compression=ContextWindowCompressionConfig( - trigger_tokens=self._settings.context_window_trigger_tokens, - sliding_window=SlidingWindow( - target_tokens=self._settings.context_window_target_tokens, - ), - ) + # Create message + new_message = types.Content( + role="user", + parts=[types.Part.from_text(text=message)], ) - # Process ADK events directly - retry is handled by HttpRetryOptions - async for adk_event in self._runner.run_async( - session_id=current_session_id, - user_id=user_id, - new_message=new_message, - run_config=run_config, - ): - # Yield LLM events from plugin first (raw capture) - if self._llm_logging_plugin: - for llm_event in self._llm_logging_plugin.drain_events(): - llm_event = self._apply_event_hook(llm_event) - if llm_event: - event_count += 1 - yield llm_event - - # Process ADK event into workflow events - async for workflow_event in self._event_processor.process_event( - adk_event, - current_session_id, + 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(), ): - event_count += 1 - yield workflow_event - - # Drain task progress events buffered by the plugin - if workflow_event.type == EventType.TOOL_RESULT and self._task_progress_plugin: - for progress_event in self._task_progress_plugin.drain_events(): - progress_event = self._apply_event_hook(progress_event) - if progress_event: - event_count += 1 - yield progress_event - - # Final drain — catches progress from the last tool call - if self._task_progress_plugin: - for progress_event in self._task_progress_plugin.drain_events(): - progress_event = self._apply_event_hook(progress_event) - if progress_event: - event_count += 1 - yield progress_event + yield event + + def _build_run_config(self): + """Build a RunConfig with context-window compression if enabled.""" + if not self._settings.context_window_enabled: + return None + from google.genai.types import ContextWindowCompressionConfig, SlidingWindow + from google.adk.agents import RunConfig + + return RunConfig( + context_window_compression=ContextWindowCompressionConfig( + trigger_tokens=self._settings.context_window_trigger_tokens, + sliding_window=SlidingWindow( + target_tokens=self._settings.context_window_target_tokens, + ), + ) + ) - # Drain any remaining LLM events after processing completes + async def _run_and_stream( + self, + *, + session_id: str, + user_id: str, + new_message: "types.Content", + run_config, + ) -> AsyncGenerator[WorkflowEvent, None]: + """Run the ADK runner for one invocation and stream WorkflowEvents. + + Shared by ``process()`` (user message) and ``resume_with_job_result()`` + (function response). The caller must already be inside + ``_workflow_context`` with the session prepared. + """ + event_count = 0 + + # Process ADK events directly - retry is handled by HttpRetryOptions + async for adk_event in self._runner.run_async( + session_id=session_id, + user_id=user_id, + new_message=new_message, + run_config=run_config, + ): + # Yield LLM events from plugin first (raw capture) if self._llm_logging_plugin: for llm_event in self._llm_logging_plugin.drain_events(): llm_event = self._apply_event_hook(llm_event) @@ -607,136 +1019,250 @@ async def process( event_count += 1 yield llm_event - logger.info("message_processed", event_count=event_count) + # Process ADK event into workflow events + async for workflow_event in self._event_processor.process_event( + adk_event, + session_id, + ): + event_count += 1 + yield workflow_event + + # Drain task progress events buffered by the plugin + if workflow_event.type == EventType.TOOL_RESULT and self._task_progress_plugin: + for progress_event in self._task_progress_plugin.drain_events(): + progress_event = self._apply_event_hook(progress_event) + if progress_event: + event_count += 1 + yield progress_event + + # Final drain — catches progress from the last tool call + if self._task_progress_plugin: + for progress_event in self._task_progress_plugin.drain_events(): + progress_event = self._apply_event_hook(progress_event) + if progress_event: + event_count += 1 + yield progress_event + + # Drain any remaining LLM events after processing completes + if self._llm_logging_plugin: + for llm_event in self._llm_logging_plugin.drain_events(): + llm_event = self._apply_event_hook(llm_event) + if llm_event: + event_count += 1 + yield llm_event + + logger.info("message_processed", event_count=event_count) + + async def can_resume(self, record) -> bool: + """True iff the originating ADK session still exists to resume into. + + Requires the resume ids and a session holding the pending call. The + session service is durable by default (``session_store='sqlite'``), so + the session normally survives a CLI restart and the job stays + resumable. This returns False — and the harness surfaces a notice + instead of a dead resume turn — when the session is genuinely gone: it + was deleted, the record lacks its session/user/call ids, or the run used + the explicitly ephemeral ``session_store='memory'`` and the process + restarted. + """ + if not (record.session_id and record.user_id and record.call_id): + return False + if self._session_service is None: + return False + session = await self._session_service.get_session( + app_name=self.app_name, user_id=record.user_id, session_id=record.session_id, + ) + return session is not None + + async def resume_with_job_result( + self, record, result: Any = None + ) -> AsyncGenerator[WorkflowEvent, None]: + """Resume the agent with a finished long-running job's result. + + Delivers the result to the pending long-running tool call as an ADK + ``FunctionResponse`` (matching ``record.call_id``) and re-invokes the + runner, streaming the follow-up turn's events. The originating session + must still contain the pending call (it is persisted); if it's gone or + the record lacks the ids needed to resume, this yields nothing. + + Args: + record: The terminal ``JobRecord`` to resume from. + result: The job's result; fetched from the JobManager if omitted. + """ + session_id = record.session_id + user_id = record.user_id + if not session_id or not user_id or not record.call_id: + logger.warning( + "job_resume_missing_ids", + job_id=getattr(record, "job_id", None), + have_session=bool(session_id), + have_user=bool(user_id), + have_call_id=bool(record.call_id), + ) + return + + bind_context(session_id=session_id, user_id=user_id) + + # 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 + + 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: + """Build the FunctionResponse payload — a summary + pointer, not raw data. + + Keeps the resumed turn's context lean; the agent can pull the full + result/logs via ``job_result``/``job_logs`` using the job_id. + """ + payload: dict[str, Any] = { + "job_id": record.job_id, + "tool": record.tool, + "state": record.state.value, + } + if record.exit_code is not None: + payload["exit_code"] = record.exit_code + if record.error: + payload["error"] = record.error + summary = _summarize_result(result) + if summary is not None: + payload["result_summary"] = summary + payload["hint"] = ( + f"Full result via job_result('{record.job_id}'); " + f"logs via job_logs('{record.job_id}')." + ) + return payload # ------------------------------------------------------------------------- - # Session save/resume hooks + # Sessions (native — DatabaseSessionService persists events continuously) # ------------------------------------------------------------------------- - async def _extract_session_data(self, session_id: str) -> tuple[list[dict], str | None]: - """Extract normalized messages and current agent from ADK session. + async def session_exists(self, session_id: str, *, user_id: str | None = None) -> bool: + """True if the store holds this session with any events. - Returns: - Tuple of (messages list, current agent name or None). + Args: + session_id: Session to look up. + user_id: Owner of the session (defaults to ``settings.default_user``). """ if not self._session_service: - return [], None - + 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, ) - if session is None or not session.events: - return [], None + return session is not None and bool(getattr(session, "events", None)) - messages: list[dict] = [] - for event in session.events: - content = getattr(event, "content", None) - if content is None: - continue + async def list_sessions(self, *, user_id: str | None = None) -> list[dict]: + """List persisted sessions for a user (most recent first). - role = getattr(content, "role", None) - parts = getattr(content, "parts", None) or [] - - for part in parts: - # Text part - if hasattr(part, "text") and part.text: - if role == "user": - messages.append({"role": "user", "content": part.text}) - elif role == "model": - messages.append({"role": "assistant", "content": part.text}) - - # Function call part - elif hasattr(part, "function_call") and part.function_call: - fc = part.function_call - tool_call = { - "id": getattr(fc, "id", fc.name), - "name": fc.name, - "args": dict(fc.args) if fc.args else {}, - } - # Attach to preceding assistant message or create one - if messages and messages[-1]["role"] == "assistant": - messages[-1].setdefault("tool_calls", []).append(tool_call) - else: - messages.append({ - "role": "assistant", - "content": "", - "tool_calls": [tool_call], - }) - - # Function response part - elif hasattr(part, "function_response") and part.function_response: - fr = part.function_response - response_content = json.dumps(fr.response) if isinstance(fr.response, dict) else str(fr.response) - messages.append({ - "role": "tool", - "tool_call_id": getattr(fr, "id", fr.name), - "name": fr.name, - "content": response_content, - }) - - current_agent = getattr(session.events[-1], "author", None) - return messages, current_agent - - async def _inject_session_messages( - self, - session_id: str, - messages: list[dict], - current_agent: str | None = None, - ) -> None: - """Inject normalized messages into the ADK session as real events. - - Uses ``append_event`` so events land in the *stored* session. - ``create_session`` returns a copy of the stored session, so the old - approach of appending to that copy left the stored session empty and - silently lost the restored history on resume. + Args: + user_id: Owner whose sessions to list (defaults to + ``settings.default_user``). """ if not self._session_service: - raise RuntimeError("Session service not initialized") - - # Create a fresh session for the restored conversation - session = await self._session_service.create_session( - app_name=self.app_name, - user_id=self._settings.default_user, - session_id=session_id, + return [] + ref = self.session_ref(user_id=user_id) + resp = await self._session_service.list_sessions( + app_name=ref.app_name, user_id=ref.user_id, ) + sessions = [ + { + "session_id": s.id, + "last_update": getattr(s, "last_update_time", None), + "message_count": len(getattr(s, "events", None) or []), + } + for s in getattr(resp, "sessions", []) + ] + sessions.sort(key=lambda x: x["last_update"] or 0, reverse=True) + return sessions - async def _add(content: types.Content, author: str) -> None: - await self._session_service.append_event( - session, Event(author=author or "user", content=content) - ) + async def delete_session(self, session_id: str, *, user_id: str | None = None) -> bool: + """Delete a persisted session from the store. - for msg in messages: - role = msg["role"] + 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=ref.app_name, + user_id=ref.user_id, + session_id=ref.session_id, + ) + return True - if role == "user": - content = types.Content( - role="user", - parts=[types.Part.from_text(text=msg["content"])], + 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=ref.app_name, + user_id=ref.user_id, + session_id=ref.session_id, + ) + if session is None or not getattr(session, "events", None): + return [] + out: list[dict] = [] + for event in session.events: + content = getattr(event, "content", None) + if content is None: + continue + role = getattr(content, "role", None) + text = " ".join( + p.text for p in (getattr(content, "parts", None) or []) + if getattr(p, "text", None) + ) + if text: + out.append( + {"role": "user" if role == "user" else "assistant", "content": text} ) - await _add(content, "user") - - elif role == "assistant": - parts = [] - if msg.get("content"): - parts.append(types.Part.from_text(text=msg["content"])) - for tc in msg.get("tool_calls", []): - parts.append(types.Part.from_function_call( - name=tc["name"], - args=tc.get("args", {}), - )) - content = types.Content(role="model", parts=parts) - await _add(content, current_agent or "model") - - elif role == "tool": - try: - response = json.loads(msg["content"]) - except (json.JSONDecodeError, TypeError): - response = {"result": msg["content"]} - parts = [types.Part.from_function_response( - name=msg.get("name", "unknown"), - response=response, - )] - content = types.Content(role="user", parts=parts) - await _add(content, current_agent or "user") + return out[-limit:] diff --git a/src/agentic_cli/workflow/adk/permission_plugin.py b/src/agentic_cli/workflow/adk/permission_plugin.py index 6524895..cb02ebd 100644 --- a/src/agentic_cli/workflow/adk/permission_plugin.py +++ b/src/agentic_cli/workflow/adk/permission_plugin.py @@ -1,10 +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. Tool has no capability declaration → deny (author error, loud). -3. Engine absent from service registry → allow (test/dev fallback). -4. Otherwise call engine.check() and return None on allow, error dict on deny. +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. +8. Otherwise call engine.check() and return None on allow, error dict on deny. """ from __future__ import annotations @@ -12,11 +31,16 @@ 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 _CapabilityExempt +from agentic_cli.workflow.permissions.capabilities import Capability, _CapabilityExempt from agentic_cli.workflow.service_registry import PERMISSION_ENGINE, get_service if TYPE_CHECKING: @@ -37,6 +61,105 @@ pass +def _native_transfer_tool_type() -> type | None: + """ADK's exact ``TransferToAgentTool`` class, or None if unavailable.""" + try: + from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool + except ImportError: # pragma: no cover - ADK always ships it today + return None + return TransferToAgentTool + + +# 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``. +# +# ``TransferToAgentTool`` is in the list for the same reason and on the same +# terms. ADK auto-injects it into any agent with ``sub_agents``, and it is a +# ``FunctionTool`` *subclass*, so an exact-type check on ``FunctionTool`` alone +# denied the built-in routing tool as unregistered — delegation could not work +# at all, even with permissions disabled. It is safe to add because ADK +# constructs it as ``super().__init__(func=transfer_to_agent)`` and overrides +# only ``_get_declaration`` (to add the agent-name enum), never ``run_async``: +# what it invokes is still exactly ``self.func``. Listing the exact class keeps +# every other ``FunctionTool`` subclass denied. +_TRUSTED_FUNCTION_TOOL_TYPES = tuple( + t + for t in (FunctionTool, LongRunningFunctionTool, _native_transfer_tool_type()) + if t is not None +) + + +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`` 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: + from google.adk.tools.mcp_tool import McpTool + except Exception: + return False + return isinstance(tool, McpTool) + + +# Synthetic capability for MCP tools; target is the MCP tool name. With no +# matching rule the engine asks the user (default ASK). Allow/deny rules with +# capability ``mcp`` and a tool-name glob target govern MCP access. +_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. + + Fail closed: if permissions are enabled but no engine is wired, deny the + call. In production ``base_manager`` always constructs the engine, so a + missing engine with permissions on is a misconfiguration — allowing would + silently bypass all gating. Only permissions-disabled allows (None). + """ + from agentic_cli.config import get_settings + + if get_settings().permissions_enabled: + logger.warning("permission_engine_missing", tool=tool_name) + return { + "success": False, + "error": "Permission denied: permission engine unavailable", + } + return None + + class PermissionPlugin(BasePlugin): """ADK plugin: gates every tool call through :class:`PermissionEngine`.""" @@ -50,11 +173,30 @@ 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: logger.warning("permission_undeclared", tool=tool.name) return { @@ -64,9 +206,41 @@ async def before_tool_callback( engine = get_service(PERMISSION_ENGINE) if engine is None: - return None # test/dev fallback + 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) + if engine is None: + return _no_engine_result(tool.name) + caps = [Capability("mcp", target_arg=_MCP_TARGET_ARG)] + result = await engine.check(tool.name, caps, {_MCP_TARGET_ARG: tool.name}) if result.allowed: return None return {"success": False, "error": f"Permission denied: {result.reason}"} diff --git a/src/agentic_cli/workflow/adk/plugins.py b/src/agentic_cli/workflow/adk/plugins.py index b3069b8..95c1601 100644 --- a/src/agentic_cli/workflow/adk/plugins.py +++ b/src/agentic_cli/workflow/adk/plugins.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import re import time from collections import deque from datetime import datetime, timezone @@ -24,6 +25,19 @@ logger = Loggers.workflow() +# Best-effort redaction of provider secrets before they hit the log file. +_SECRET_RE = re.compile( + r"(sk-[A-Za-z0-9_\-]{16,}" # Anthropic / OpenAI keys + r"|AIza[A-Za-z0-9_\-]{16,}" # Google API keys + r"|Bearer\s+[A-Za-z0-9._\-]{10,})", # bearer tokens + re.IGNORECASE, +) + + +def _redact_secrets(text: str) -> str: + """Mask obvious API-key/token patterns in a serialized log line.""" + return _SECRET_RE.sub("[REDACTED]", text) + class LLMLoggingPlugin(BasePlugin): """ADK Plugin that captures raw LLM request/response traffic for debugging. @@ -58,10 +72,21 @@ def __init__( self._events: deque[WorkflowEvent] = deque(maxlen=max_events) self._request_timestamps: dict[str, float] = {} - # Initialize log file - log_dir = Path.cwd() / f".{self.app_name}" / "logs" + # Initialize log file under the user's home — NOT the (often + # git-tracked) project dir — since it holds full conversations. + log_dir = Path.home() / f".{self.app_name}" / "logs" log_dir.mkdir(parents=True, exist_ok=True) + try: + log_dir.chmod(0o700) + except OSError: + pass self._log_file: Path = log_dir / "llm_events.jsonl" + # Create (or tighten) with owner-only perms before anything is written. + try: + self._log_file.touch(exist_ok=True) + self._log_file.chmod(0o600) + except OSError: + pass # ------------------------------------------------------------------ # ADK Plugin callbacks @@ -237,8 +262,9 @@ def _write_to_log(self, event: WorkflowEvent) -> None: } try: + line = _redact_secrets(json.dumps(record, default=str)) with open(self._log_file, "a", encoding="utf-8") as f: - f.write(json.dumps(record, default=str) + "\n") + f.write(line + "\n") except OSError: pass diff --git a/src/agentic_cli/workflow/adk/transfer_tool_description.py b/src/agentic_cli/workflow/adk/transfer_tool_description.py new file mode 100644 index 0000000..8aa484c --- /dev/null +++ b/src/agentic_cli/workflow/adk/transfer_tool_description.py @@ -0,0 +1,136 @@ +"""Correct the model-visible description of ADK's ``transfer_to_agent`` tool. + +ADK builds ``TransferToAgentTool``'s declaration from the ``transfer_to_agent`` +function's docstring, which (through 1.37.0, the newest 1.x) contains: + + Note: + For most use cases, you should use TransferToAgentTool instead of this + function directly. + +That paragraph is written for *Python callers*, but it ships to the model as +the tool's description — so the model is told, in the tool it is supposed to +call, to call something else. Gemini 3.1 followed the instruction and emitted +``TransferToAgentTool``, which ADK rejects (``Tool 'TransferToAgentTool' not +found``), and delegation failed outright. + +Upstream fixed the docstring in ADK 2.x, which is a major release outside this +project's ``google-adk>=1.34,<2`` pin. So the description is corrected here, on +the prepared request, under conditions narrow enough that the correction simply +stops applying once the installed ADK no longer needs it: + +- only when the tool object is **exactly** ``TransferToAgentTool`` — an + application tool that happens to be named ``transfer_to_agent`` is left alone; +- only when the misleading sentence is actually present — so it is idempotent, + and an ADK release that fixes the text makes this a no-op; +- only ``description`` is written. The declaration name, parameter schema, + required fields and ADK's enum of valid agent names are untouched, and the + upstream function's ``__doc__`` is never mutated (ADK rebuilds the + declaration per request, so the edit is scoped to one ``LlmRequest``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterator + +from google.adk.plugins.base_plugin import BasePlugin + +from agentic_cli.logging import Loggers + +if TYPE_CHECKING: + from google.adk.agents.callback_context import CallbackContext + from google.adk.models.llm_request import LlmRequest + from google.genai import types + +logger = Loggers.workflow() + +#: The declaration name ADK uses. Also the name the model must call. +TRANSFER_TOOL_NAME = "transfer_to_agent" + +#: The fragment that makes ADK's generated description actively harmful. Its +#: presence is the trigger; its absence means there is nothing to correct. +_MISLEADING_FRAGMENT = "you should use TransferToAgentTool" + +#: Replacement: says plainly what to call, and defers to the parameter schema +#: for *which* agents are valid rather than restating them (ADK's enum is the +#: authority and stays intact). +ROUTING_DESCRIPTION = ( + "Route the current request to another agent. " + f'Call this function by its exact name: {TRANSFER_TOOL_NAME}' + '(agent_name=""). ' + "Choose agent_name from the enum of valid agent names in this tool's " + "parameter schema — those are the only accepted values. " + "Transfer when another agent's description fits the user's request better " + "than your own." +) + + +def _native_transfer_tool_type() -> type | None: + """ADK's exact ``TransferToAgentTool`` class, or None if unavailable.""" + try: + from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool + except ImportError: # pragma: no cover - ADK always ships it today + return None + return TransferToAgentTool + + +def _function_declarations(llm_request: "LlmRequest") -> Iterator["types.FunctionDeclaration"]: + """Every function declaration attached to the prepared request.""" + config = getattr(llm_request, "config", None) + for tool in getattr(config, "tools", None) or []: + for declaration in getattr(tool, "function_declarations", None) or []: + yield declaration + + +def correct_transfer_declaration(llm_request: "LlmRequest") -> bool: + """Rewrite the transfer tool's description on this request, if warranted. + + Returns: + True if a declaration was corrected (useful for tests and logging). + """ + transfer_type = _native_transfer_tool_type() + if transfer_type is None: + return False + + tool = (getattr(llm_request, "tools_dict", None) or {}).get(TRANSFER_TOOL_NAME) + # Exact type, never isinstance: a subclass may declare anything it likes, + # and an application tool sharing the name is not ADK's routing primitive. + if type(tool) is not transfer_type: + return False + + corrected = False + for declaration in _function_declarations(llm_request): + if declaration.name != TRANSFER_TOOL_NAME: + continue + if _MISLEADING_FRAGMENT not in (declaration.description or ""): + continue # already correct (or fixed upstream) — leave it alone + declaration.description = ROUTING_DESCRIPTION + corrected = True + + return corrected + + +class TransferToolDescriptionPlugin(BasePlugin): + """Applies :func:`correct_transfer_declaration` to every model request.""" + + def __init__(self) -> None: + super().__init__(name="transfer_tool_description") + self._corrections = 0 + + @property + def corrections(self) -> int: + """How many requests have been corrected (diagnostics/tests).""" + return self._corrections + + async def before_model_callback( + self, + *, + callback_context: "CallbackContext", + llm_request: "LlmRequest", + ) -> None: + """Fix the transfer description in place; never blocks the request.""" + try: + if correct_transfer_declaration(llm_request): + self._corrections += 1 + except Exception as exc: # noqa: BLE001 - never break a turn over this + logger.warning("transfer_description_correction_failed", error=str(exc)) + return None diff --git a/src/agentic_cli/workflow/adk_config_bridge.py b/src/agentic_cli/workflow/adk_config_bridge.py new file mode 100644 index 0000000..58f58bf --- /dev/null +++ b/src/agentic_cli/workflow/adk_config_bridge.py @@ -0,0 +1,154 @@ +"""Reuse an existing native ADK agent config (``root_agent.yaml``). + +Two strategies: + +- **native** (:func:`load_adk_agent_native`): hand the YAML straight to ADK's + ``from_config`` and run the resulting agent tree as-is. Full ADK fidelity + (planners, callbacks, code executors, ``model_code``/LiteLlm), but framework + service-tool / state-tool auto-injection does not apply, and it only works on + the ADK backend. + +- **translate** (:func:`translate_adk_yaml`): best-effort conversion of the ADK + YAML into the framework's ``AgentConfig`` list, so the normal manager path + (service tools, state tools, permissions, plus the new model_settings/mcp/ + skills features) applies. ADK-only constructs are dropped with a warning. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import yaml + +from agentic_cli.logging import Loggers +from agentic_cli.workflow.config import AgentConfig + +logger = Loggers.workflow() + +# LlmAgentConfig fields with no framework equivalent — dropped during translate. +_DROPPED_FIELDS = ( + "planner", + "code_executor", + "model_code", + "static_instruction", + "input_schema", + "output_schema", + "output_key", + "before_model_callbacks", + "after_model_callbacks", + "before_tool_callbacks", + "after_tool_callbacks", + "before_agent_callbacks", + "after_agent_callbacks", +) + + +def load_adk_agent_native(config_path: str | Path) -> Any: + """Build an agent tree natively via ADK ``from_config`` (full fidelity).""" + from google.adk.agents.config_agent_utils import from_config + + return from_config(str(config_path)) + + +def translate_adk_yaml(config_path: str | Path) -> list[AgentConfig]: + """Translate a native ADK YAML config into framework ``AgentConfig`` objects. + + Best-effort: only ``LlmAgent`` nodes are translated; other agent classes and + ADK-only fields are skipped/dropped with a warning. Sub-agents referenced by + ``config_path`` are translated recursively; ``code`` references are skipped. + """ + path = Path(config_path) + out: list[AgentConfig] = [] + _translate_node(_read_yaml(path), path.parent, out) + return out + + +def _read_yaml(path: Path) -> dict: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError(f"ADK config {path} did not parse to a mapping.") + return data + + +def _translate_node( + data: dict, base_dir: Path, out: list[AgentConfig] +) -> AgentConfig | None: + """Translate one ADK agent node (appending children first, then itself).""" + name = data.get("name") + agent_class = data.get("agent_class") or "LlmAgent" + if agent_class not in ("LlmAgent", "Agent"): + logger.warning( + "adk_translate_unsupported_agent_class", + name=name, + agent_class=agent_class, + ) + return None + + sub_names: list[str] = [] + for ref in data.get("sub_agents") or []: + if not isinstance(ref, dict): + continue + if ref.get("config_path"): + sub_path = base_dir / ref["config_path"] + child = _translate_node(_read_yaml(sub_path), sub_path.parent, out) + if child is not None: + sub_names.append(child.name) + elif ref.get("code"): + logger.warning("adk_translate_code_subagent_skipped", code=ref["code"]) + + tools: list[str] = [] + for tool in data.get("tools") or []: + tool_name = tool.get("name") if isinstance(tool, dict) else tool + if isinstance(tool, dict) and tool.get("args"): + logger.warning("adk_translate_tool_args_dropped", tool=tool_name) + if tool_name: + tools.append(tool_name) + + for field in _DROPPED_FIELDS: + if data.get(field): + logger.warning("adk_translate_field_dropped", name=name, field=field) + + config = AgentConfig( + name=name, + prompt=data.get("instruction", "") or "", + description=data.get("description", "") or "", + model=data.get("model"), + model_settings=_translate_model_settings(data.get("generate_content_config")), + tools=tools, + sub_agents=sub_names, + ) + out.append(config) + return config + + +def _translate_model_settings(gcc: dict | None): + """Map an ADK ``generate_content_config`` dict to ``ModelSettings``.""" + if not gcc: + return None + from agentic_cli.workflow.model_settings import ModelSettings, ThinkingSettings + + kwargs: dict[str, Any] = {} + for src, dst in ( + ("temperature", "temperature"), + ("top_p", "top_p"), + ("top_k", "top_k"), + ("max_output_tokens", "max_tokens"), + ): + if gcc.get(src) is not None: + kwargs[dst] = gcc[src] + if gcc.get("stop_sequences"): + kwargs["stop_sequences"] = gcc["stop_sequences"] + + tc = gcc.get("thinking_config") + if isinstance(tc, dict): + if tc.get("thinking_budget") is not None: + kwargs["thinking"] = ThinkingSettings( + mode="budget", budget_tokens=tc["thinking_budget"] + ) + elif tc.get("thinking_level"): + level = str(tc["thinking_level"]).lower().split(".")[-1] + mode = level if level in ("low", "medium", "high") else "high" + kwargs["thinking"] = ThinkingSettings(mode=mode) + + return ModelSettings(**kwargs) if kwargs else None diff --git a/src/agentic_cli/workflow/agent_loader.py b/src/agentic_cli/workflow/agent_loader.py new file mode 100644 index 0000000..ca965e2 --- /dev/null +++ b/src/agentic_cli/workflow/agent_loader.py @@ -0,0 +1,144 @@ +"""Load agent definitions from a unified YAML config. + +This is the framework's own declarative format (distinct from a native ADK +``root_agent.yaml`` — see ``adk_config_bridge`` for reusing those). It produces +``AgentConfig`` objects, so the result feeds the normal factory/manager path. + +Example YAML:: + + agents: + - name: coordinator + model: gemini-2.5-pro + model_settings: + temperature: 0.2 + thinking: {mode: high} + instruction: | + You are the coordinator... + tools: [kb_search, my_pkg.tools.custom_tool] + sub_agents: [researcher] + - name: researcher + instruction_file: prompts/researcher.md + tools: [web_search] + +A top-level bare list of agents (without the ``agents:`` key) is also accepted. +Tool entries stay as strings here and are resolved to callables when the +workflow manager is constructed (see ``tools.tool_resolver``). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import yaml +from pydantic import AliasChoices, BaseModel, ConfigDict, Field + +from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.mcp import MCPServerConfig +from agentic_cli.workflow.model_settings import ModelSettings + + +class AgentSpec(BaseModel): + """YAML schema for a single agent (validated, then mapped to AgentConfig).""" + + model_config = ConfigDict(protected_namespaces=(), extra="forbid") + + name: str + # Accept either ``prompt:`` or ``instruction:`` for the system prompt. + prompt: str | None = Field( + default=None, + validation_alias=AliasChoices("prompt", "instruction"), + ) + instruction_file: str | None = None + description: str = "" + model: str | None = None + model_settings: ModelSettings | None = None + tools: list[str] = Field(default_factory=list) + mcp_servers: list[MCPServerConfig] = Field(default_factory=list) + skills: list[str] = Field(default_factory=list) + sub_agents: list[str] = Field(default_factory=list) + include_state_tools: bool = True + + +class AgentsFile(BaseModel): + """Top-level YAML schema: a list of agents under ``agents:``.""" + + model_config = ConfigDict(extra="forbid") + + agents: list[AgentSpec] = Field(min_length=1) + + +def _resolve_prompt(spec: AgentSpec, base_dir: Path) -> str: + """Resolve an agent's prompt from inline text or an instruction file.""" + if spec.instruction_file: + path = Path(spec.instruction_file) + if not path.is_absolute(): + path = base_dir / path + return path.read_text(encoding="utf-8") + if spec.prompt is not None: + return spec.prompt + raise ValueError( + f"Agent {spec.name!r}: must set 'prompt'/'instruction' or 'instruction_file'." + ) + + +def load_agents_from_yaml(path: str | Path) -> list[AgentConfig]: + """Load and validate agent configs from a unified YAML file. + + Args: + path: Path to the YAML file. + + Returns: + List of ``AgentConfig`` (tool refs left as strings for later resolution). + + Raises: + FileNotFoundError: If the file does not exist. + ValueError: If the YAML is invalid or a prompt cannot be resolved. + pydantic.ValidationError: If the schema does not validate. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Agent config file not found: {path}") + + data: Any = yaml.safe_load(path.read_text(encoding="utf-8")) + if data is None: + raise ValueError(f"Agent config file is empty: {path}") + # Allow a top-level bare list of agents. + if isinstance(data, list): + data = {"agents": data} + + spec_file = AgentsFile.model_validate(data) + base_dir = path.parent + + configs: list[AgentConfig] = [] + for spec in spec_file.agents: + configs.append( + AgentConfig( + name=spec.name, + prompt=_resolve_prompt(spec, base_dir), + tools=list(spec.tools), + sub_agents=list(spec.sub_agents), + description=spec.description, + model=spec.model, + model_settings=spec.model_settings, + mcp_servers=list(spec.mcp_servers), + skills=list(spec.skills), + include_state_tools=spec.include_state_tools, + ) + ) + return configs + + +def create_workflow_manager_from_yaml( + path: str | Path, + settings: Any, + **kwargs: Any, +): + """Convenience: load agents from YAML and build a workflow manager. + + Equivalent to ``create_workflow_manager_from_settings(load_agents_from_yaml(path), settings)``. + """ + from agentic_cli.workflow.factory import create_workflow_manager_from_settings + + configs = load_agents_from_yaml(path) + return create_workflow_manager_from_settings(configs, settings, **kwargs) diff --git a/src/agentic_cli/workflow/base_manager.py b/src/agentic_cli/workflow/base_manager.py index 2725b19..0815557 100644 --- a/src/agentic_cli/workflow/base_manager.py +++ b/src/agentic_cli/workflow/base_manager.py @@ -13,16 +13,26 @@ 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 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, + JOB_MANAGER, KB_MANAGER, LLM_SUMMARIZER, MEMORY_STORE, @@ -54,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: @@ -89,18 +114,37 @@ 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 + # 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() + # Resolve string / dotted-path tool refs to callables BEFORE manager + # detection and tool assembly (both key on ``tool.__name__``). + self._resolve_config_tool_refs() + # Auto-detect required managers from tools self._required_managers = self._detect_required_managers() @@ -111,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]: @@ -179,6 +252,11 @@ def sandbox_manager(self) -> "SandboxManager | None": """Get the sandbox manager (if required by tools).""" return self._services.get(SANDBOX_MANAGER) + @property + def job_manager(self): + """Get the long-running-job manager (if required by tools).""" + return self._services.get(JOB_MANAGER) + # ------------------------------------------------------------------ # Tool assembly # ------------------------------------------------------------------ @@ -190,23 +268,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. @@ -234,7 +367,7 @@ def _get_service_tool_map(self) -> dict[str, Callable]: if s.get(LLM_SUMMARIZER): tool_map["web_fetch"] = make_webfetch_tool(s[LLM_SUMMARIZER]) if s.get(SANDBOX_MANAGER): - tool_map["sandbox_execute"] = make_sandbox_tool(s[SANDBOX_MANAGER]) + tool_map["sandbox_execute"] = make_sandbox_tool(s[SANDBOX_MANAGER], self) if s.get(ARXIV_SOURCE): for t in make_arxiv_tools(s[ARXIV_SOURCE]): tool_map[t.__name__] = t @@ -247,6 +380,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 @@ -257,59 +394,104 @@ 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"), - } + def _resolve_config_tool_refs(self) -> None: + """Resolve string/dotted-path tool refs in configs to callables. + + Each ``config.tools`` entry may be a callable, a registered tool name, + or a dotted import path. This rewrites every config's ``tools`` list in + place so subsequent service-detection and tool assembly (which key on + ``tool.__name__``) operate on real callables. Callables pass through + unchanged, so the call is idempotent. + """ + from agentic_cli.tools.tool_resolver import resolve_tools + + for config in self._agent_configs: + if config.tools: + 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. - 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. - if "memory_store" in self._required_managers and MEMORY_STORE not in s: + Args: + existing: Service keys already published; those are not rebuilt. + + Returns: + The newly constructed services, keyed by service key. + + 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 @@ -327,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 @@ -350,22 +532,38 @@ 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 "arxiv_source" in self._required_managers and ARXIV_SOURCE 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 + + # User-scoped so long jobs persist across projects and CLI restarts. + jobs_dir = Path.home() / f".{self._settings.app_name}" / "jobs" + s[JOB_MANAGER] = JobManager( + self._settings, + base_dir=jobs_dir, + max_concurrent=getattr(self._settings, "max_concurrent_jobs", 4), + ) + + 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(workdir=Path.cwd(), home=Path.home()) + ctx = PermissionContext( + workdir=Path.cwd(), + home=Path.home(), + app_name=self._settings.app_name, + ) s[PERMISSION_ENGINE] = PermissionEngine( settings=self._settings, workflow=self, ctx=ctx, ) @@ -373,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: @@ -383,13 +611,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). @@ -405,9 +643,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._extract_session_data(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 [] @@ -437,29 +682,131 @@ 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.""" + 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.""" + 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. + + Base default: False (no resume support). Backends that implement + ``resume_with_job_result`` override this to report whether the + originating conversation is still available — e.g. the ADK session that + holds the pending call. Used by the harness to resume vs. surface a + "finished while its conversation was unavailable" notice. + + Sessions are durable by default (``session_store='sqlite'``), so a + conversation normally survives a CLI restart and stays resumable. It is + unavailable when, for example, it was deleted, the record is missing its + session/user/call identifiers, or the run used the explicitly ephemeral + ``session_store='memory'`` and the process restarted. + """ + return False + @contextlib.contextmanager - def _workflow_context(self) -> Iterator[None]: + def _workflow_context( + self, session_id: str | None = None, user_id: str | None = None + ) -> 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. + 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) + turn_token = set_active_turn(self.session_ref(session_id, user_id)) try: yield finally: + reset_active_turn(turn_token) 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 @@ -491,40 +838,157 @@ 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 - from agentic_cli.config import validate_settings - - if validate: - validate_settings(self._settings) + # 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() 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 - # 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. + _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. + + Runs before discovery and service creation so a static configuration + error surfaces immediately and costs nothing. + """ + 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). @@ -603,7 +1067,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. @@ -620,13 +1086,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 @@ -677,131 +1144,132 @@ async def generate_simple(self, prompt: str, max_tokens: int = 500) -> str: # Session save/resume # ------------------------------------------------------------------------- - async def save_session(self, session_id: str | None = None) -> dict: - """Save current session state to disk. - - Extracts messages and current agent from the backend via abstract - hooks, then persists via SessionPersistence.save_snapshot(). + # ------------------------------------------------------------------ + # Sessions (native, durable). The backend store (ADK + # DatabaseSessionService / LangGraph checkpointer) persists conversation + # state continuously, keyed by session_id; there is no separate snapshot. + # ------------------------------------------------------------------ - Args: - session_id: Session ID to save under. Uses backend session_id if None. + @property + def supports_sessions(self) -> bool: + """Whether this backend implements the durable-session hooks. - Returns: - Dict with success status and path. + 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. """ - from datetime import datetime - from agentic_cli.persistence.session import SessionPersistence, SessionSnapshot + return type(self).session_exists is not BaseWorkflowManager.session_exists - sid = session_id or getattr(self, "session_id", "default_session") + def _is_default_user(self, user_id: str | None) -> bool: + """Whether ``user_id`` is (or defaults to) the configured default user. - try: - messages, current_agent = await self._extract_session_data(sid) - - metadata: dict[str, Any] = { - "model": self.model, - "backend_type": self.backend_type, - "app_name": self._app_name, - } - if current_agent: - metadata["current_agent"] = current_agent - - now = datetime.now() - snapshot = SessionSnapshot( - session_id=sid, - created_at=now, - saved_at=now, - messages=messages, - metadata=metadata, - ) + 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. + """ + return user_id is None or user_id == self._settings.default_user - persistence = SessionPersistence(self._settings) - path = persistence.save_snapshot(snapshot) + def session_ref( + self, session_id: str | None = None, user_id: str | None = None + ) -> SessionRef: + """Resolve a full :class:`SessionRef` from partial identity. - logger.info("session_saved", session_id=sid, message_count=len(messages)) - return {"success": True, "session_id": sid, "path": str(path), "message_count": len(messages)} + 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"), + ) - except Exception as exc: - logger.error("session_save_failed", session_id=sid, error=str(exc)) - return {"success": False, "error": str(exc)} + 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. - async def load_session(self, session_id: str) -> bool: - """Load a saved session and inject it into the backend. + Kept for API compatibility / explicit "checkpoint now" intent. Args: - session_id: Session ID to load. + session_id: Session to report (defaults to the manager's current one). + user_id: Owner (defaults to ``settings.default_user``). Returns: - True if session was loaded successfully, False otherwise. + ``{"success": True, "session_id": ..., "user_id": ...}`` — the full + identity, so a caller working on behalf of another user can tell + which conversation was meant. """ - from agentic_cli.persistence.session import SessionPersistence - - persistence = SessionPersistence(self._settings) - snapshot = persistence.load_session(session_id) + ref = self.session_ref(session_id, user_id) + return { + "success": True, + "session_id": ref.session_id, + "user_id": ref.user_id, + } - if snapshot is None: - logger.debug("session_not_found_for_load", session_id=session_id) - return False + 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. - # Warn on backend mismatch (but still load — format is normalized) - saved_backend = snapshot.metadata.get("backend_type") - if saved_backend and saved_backend != self.backend_type: - logger.warning( - "session_backend_mismatch", - saved_backend=saved_backend, - current_backend=self.backend_type, - session_id=session_id, - ) + 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. A backend without durable sessions has nothing to resume, + so it adopts the id and returns False. - current_agent = snapshot.metadata.get("current_agent") - - try: - await self._inject_session_messages(session_id, snapshot.messages, current_agent) - if hasattr(self, "session_id"): - self.session_id = session_id + 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 + if not self.supports_sessions: logger.info( - "session_loaded", - session_id=session_id, - message_count=len(snapshot.messages), + "session_adopted", session_id=session_id, resumed=False, + backend_sessions=False, ) - return True - except Exception as exc: - logger.error("session_load_failed", session_id=session_id, error=str(exc)) 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." + ) - def list_sessions(self) -> list[dict]: - """List all saved sessions. - - Returns: - List of session summary dicts. - """ - from agentic_cli.persistence.session import SessionPersistence - - persistence = SessionPersistence(self._settings) - return persistence.list_sessions() + async def session_exists(self, session_id: str, *, user_id: str | None = None) -> bool: + """Whether the native store already holds this session's state.""" + raise self._no_session_support("session_exists") - @abstractmethod - async def _extract_session_data( - self, session_id: str - ) -> tuple[list[dict], str | None]: - """Extract normalized messages and current agent from backend session. + 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. - Returns: - Tuple of (messages, current_agent_name). - Messages use normalized format: - - {"role": "user", "content": "..."} - - {"role": "assistant", "content": "...", "tool_calls": [...]} - - {"role": "tool", "tool_call_id": "...", "name": "...", "content": "..."} + Used for session-end fact extraction; text-only (no tool-call fidelity). """ - ... + raise self._no_session_support("recent_messages") - @abstractmethod - async def _inject_session_messages( - self, - session_id: str, - messages: list[dict], - current_agent: str | None = None, - ) -> None: - """Inject normalized messages into the backend session.""" - ... + async def list_sessions(self, *, user_id: str | None = None) -> list[dict]: + """List persisted sessions from the native store (most recent first).""" + raise self._no_session_support("list_sessions") + + async def delete_session(self, session_id: str, *, user_id: str | None = None) -> bool: + """Delete a persisted session from the native store.""" + raise self._no_session_support("delete_session") diff --git a/src/agentic_cli/workflow/config.py b/src/agentic_cli/workflow/config.py index 9bed844..b6ebbd1 100644 --- a/src/agentic_cli/workflow/config.py +++ b/src/agentic_cli/workflow/config.py @@ -1,7 +1,13 @@ """Configuration classes for workflow management.""" +import inspect from dataclasses import dataclass, field -from typing import Callable, Any +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 @dataclass @@ -13,24 +19,260 @@ class AgentConfig: Attributes: name: Unique identifier for the agent - prompt: System instruction - either a string or a callable that returns one - tools: List of tool functions the agent can use + 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 + when the workflow manager is constructed. sub_agents: Names of agents that this agent can delegate to description: Short description for routing/logging model: Optional model override (defaults to manager's model) + model_settings: Optional per-agent generation parameters (temperature, + thinking, etc.). Currently consumed by the ADK backend only. + mcp_servers: Optional MCP servers whose tools are exposed to this agent. + Currently consumed by the ADK backend only. + skills: Optional skill references (Agent Skills / SKILL.md folders) — + each a path to a skill directory or a name resolved under + ``settings.skills_dirs``. Currently consumed by the ADK backend only. include_state_tools: Whether to auto-inject plan/task state tools (default True) """ name: str - prompt: str | Callable[[], str] - tools: list[Callable[..., Any]] = field(default_factory=list) + prompt: str | Callable[..., str] + tools: list[Callable[..., Any] | str] = field(default_factory=list) sub_agents: list[str] = field(default_factory=list) description: str = "" model: str | None = None + model_settings: "ModelSettings | None" = None + mcp_servers: "list[MCPServerConfig]" = field(default_factory=list) + 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/src/agentic_cli/workflow/factory.py b/src/agentic_cli/workflow/factory.py index 3ab0a34..edf48f9 100644 --- a/src/agentic_cli/workflow/factory.py +++ b/src/agentic_cli/workflow/factory.py @@ -38,14 +38,16 @@ def create_workflow_manager_from_settings( settings: "BaseSettings", app_name: str | None = None, model: str | None = None, + adk_config_path: str | None = None, + adk_config_mode: str = "native", **kwargs, ) -> "BaseWorkflowManager": """Factory function to create the appropriate workflow manager based on settings. Creates either a GoogleADKWorkflowManager or LangGraphWorkflowManager - based on the settings.orchestrator configuration. Claude models are - automatically routed to LangGraph because ADK's LiteLLM adapter has - critical issues with tool calling, thinking, and streaming. + based purely on the settings.orchestrator configuration. The backend is + model-agnostic: ADK runs Claude natively via the direct-API ``AnthropicLlm`` + (no LiteLLM), so Claude no longer forces a LangGraph swap. Args: agent_configs: List of agent configurations. @@ -67,24 +69,43 @@ def create_workflow_manager_from_settings( """ from agentic_cli.workflow.settings import OrchestratorType + # Reuse an existing native ADK config, if provided. + if adk_config_path: + if adk_config_mode == "translate": + from agentic_cli.workflow.adk_config_bridge import translate_adk_yaml + + agent_configs = translate_adk_yaml(adk_config_path) + # fall through to normal routing with the translated configs + else: # native passthrough — ADK backend only + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + if getattr(settings, "orchestrator", OrchestratorType.ADK) == ( + OrchestratorType.LANGGRAPH + ): + logger.warning( + "adk_native_config_forces_adk", + reason="native ADK config requires the ADK backend", + ) + return GoogleADKWorkflowManager( + agent_configs=agent_configs, + settings=settings, + app_name=app_name, + model=model, + adk_config_path=adk_config_path, + **kwargs, + ) + orchestrator = getattr(settings, "orchestrator", OrchestratorType.ADK) - effective_model = _resolve_effective_model(model, settings) - use_langgraph = orchestrator == OrchestratorType.LANGGRAPH or _is_claude_model( - effective_model - ) - - if use_langgraph and _is_claude_model(effective_model) and orchestrator != OrchestratorType.LANGGRAPH: - logger.info( - "auto_switching_to_langgraph", - model=effective_model, - reason="Claude models require LangGraph orchestrator (ADK LiteLLM adapter has critical issues)", - ) + use_langgraph = orchestrator == OrchestratorType.LANGGRAPH if use_langgraph: try: from agentic_cli.workflow.langgraph import LangGraphWorkflowManager - checkpointer = getattr(settings, "langgraph_checkpointer", "memory") + # Durable sessions: the unified session_store (sqlite by default) + # selects the checkpointer so LangGraph persists like ADK, keyed by + # thread_id == session_id. "memory" stays ephemeral. + checkpointer = getattr(settings, "session_store", "memory") return LangGraphWorkflowManager( agent_configs=agent_configs, settings=settings, diff --git a/src/agentic_cli/workflow/langgraph/graph_builder.py b/src/agentic_cli/workflow/langgraph/graph_builder.py index 2c2fce4..2c97852 100644 --- a/src/agentic_cli/workflow/langgraph/graph_builder.py +++ b/src/agentic_cli/workflow/langgraph/graph_builder.py @@ -220,7 +220,11 @@ def get_llm(self, model: str): kwargs = {"model": model} if thinking and thinking["provider"] == "google": kwargs["include_thoughts"] = thinking.get("include_thoughts", True) - kwargs["thinking_level"] = thinking.get("thinking_level") + # 2.5 → thinking_budget, 3 → thinking_level (see get_thinking_config) + if thinking.get("thinking_budget") is not None: + kwargs["thinking_budget"] = thinking["thinking_budget"] + elif thinking.get("thinking_level") is not None: + kwargs["thinking_level"] = thinking["thinking_level"] return ChatGoogleGenerativeAI(**kwargs) except ImportError: raise ImportError( @@ -275,13 +279,22 @@ def get_thinking_config(self, model: str) -> dict[str, Any] | None: } if model.startswith("gemini-"): - # Gemini 3 Pro only supports "low" and "high", not "medium" - is_gemini_3_pro = "gemini-3" in model and "pro" in model - level = "high" if (effort == "medium" and is_gemini_3_pro) else effort + # Gemini 3 takes a discrete thinking_level; Gemini 2.5 only accepts a + # numeric thinking_budget and rejects thinking_level (400). Mirrors + # the ADK split in GoogleADKWorkflowManager._get_planner. + if "gemini-3" in model: + is_gemini_3_pro = "pro" in model # 3 Pro lacks "medium" + level = "high" if (effort == "medium" and is_gemini_3_pro) else effort + return { + "provider": "google", + "include_thoughts": True, + "thinking_level": level, + } + budget = {"low": 4096, "medium": 12288, "high": 24576}.get(effort, 12288) return { "provider": "google", "include_thoughts": True, - "thinking_level": level, + "thinking_budget": budget, } return None diff --git a/src/agentic_cli/workflow/langgraph/manager.py b/src/agentic_cli/workflow/langgraph/manager.py index bfd42ed..0d719db 100644 --- a/src/agentic_cli/workflow/langgraph/manager.py +++ b/src/agentic_cli/workflow/langgraph/manager.py @@ -27,6 +27,25 @@ logger = Loggers.workflow() +# AgentConfig fields the LangGraph backend does not (yet) consume. +_ADK_ONLY_FIELDS = ("model_settings", "mcp_servers", "skills") + + +def _ignored_adk_only_fields( + agent_configs: list[AgentConfig], +) -> list[tuple[str, list[str]]]: + """Return [(agent_name, [field, ...])] for ADK-only fields set on configs. + + Used to warn that ``model_settings``/``mcp_servers``/``skills`` are ignored + on the LangGraph backend. + """ + affected: list[tuple[str, list[str]]] = [] + for cfg in agent_configs: + present = [f for f in _ADK_ONLY_FIELDS if getattr(cfg, f, None)] + if present: + affected.append((cfg.name, present)) + return affected + class LangGraphWorkflowManager(BaseWorkflowManager): """LangGraph-based workflow manager for agentic applications. @@ -91,6 +110,16 @@ def __init__( on_event=on_event, ) + # These AgentConfig fields are consumed by the ADK backend only; warn + # once if an app set them while running on (or routed to) LangGraph. + ignored = _ignored_adk_only_fields(self._agent_configs) + if ignored: + logger.warning( + "langgraph_ignoring_adk_only_fields", + agents=ignored, + note="model_settings/mcp_servers/skills are consumed by the ADK backend only", + ) + self._checkpointer_type = checkpointer # Graph builder (delegates graph construction + LLM factory) @@ -278,7 +307,7 @@ async def process( logger.info("processing_message_langgraph", message_length=len(message)) # Context setup - with self._workflow_context(): + with self._workflow_context(session_id=current_session_id, user_id=user_id): # Prepare initial state initial_state = { "messages": [{"role": "user", "content": message}], @@ -416,7 +445,7 @@ async def process( yield transformed # Emit task progress after tool results - progress_event = self._build_task_progress(current_session_id) + progress_event = await self._build_task_progress(current_session_id) if progress_event: transformed = self._apply_event_hook(progress_event) if transformed: @@ -424,7 +453,7 @@ async def process( # Final progress check - progress_event = self._build_task_progress(current_session_id) + progress_event = await self._build_task_progress(current_session_id) if progress_event: transformed = self._apply_event_hook(progress_event) if transformed: @@ -432,18 +461,12 @@ async def process( logger.info("message_processed_langgraph") - def _build_task_progress(self, session_id: str) -> "WorkflowEvent | None": - """Build task progress from graph state.""" + async def _build_task_progress(self, session_id: str) -> "WorkflowEvent | None": + """Build task progress from graph state (async — see _get_state_values).""" from agentic_cli.tools._core.tasks import task_progress_data - if not self._compiled_graph: - return None - try: - config = {"configurable": {"thread_id": session_id}} - state = self._compiled_graph.get_state(config) - tasks_data = state.values.get("tasks", []) if state and state.values else [] - except Exception: - return None + values = await self._get_state_values(session_id) + tasks_data = values.get("tasks", []) if values else [] if not tasks_data: return None @@ -523,96 +546,69 @@ async def generate_simple(self, prompt: str, max_tokens: int = 500) -> str: # Session save/resume hooks # ------------------------------------------------------------------------- - def _get_state_values(self, session_id: str) -> dict | None: - """Get state values from LangGraph, or None on failure.""" + async def _get_state_values(self, session_id: str) -> dict | None: + """Get state values from LangGraph, or None on failure. + + Uses the async ``aget_state`` — the persistent checkpointers + (AsyncSqliteSaver/AsyncPostgresSaver) don't implement the sync path. + """ if not self._compiled_graph: return None config = {"configurable": {"thread_id": session_id}} try: - state = self._compiled_graph.get_state(config) + state = await self._compiled_graph.aget_state(config) except Exception: return None if not state or not state.values: return None return state.values - async def _extract_session_data( - self, session_id: str - ) -> tuple[list[dict], str | None]: - """Extract normalized messages and current agent from LangGraph state.""" - values = self._get_state_values(session_id) - if not values: - return [], None + async def session_exists(self, session_id: str) -> bool: + """True if the checkpointer holds state for this thread.""" + return await self._get_state_values(session_id) is not None - current_agent = values.get("current_agent") - raw_messages = values.get("messages", []) - messages: list[dict] = [] + async def list_sessions(self) -> list[dict]: + """List persisted threads (session ids) from the checkpointer.""" + if not self._checkpointer: + return [] + seen: dict[str, Any] = {} + try: + async for cp in self._checkpointer.alist(None): + tid = (getattr(cp, "config", None) or {}).get("configurable", {}).get( + "thread_id" + ) + if tid and tid not in seen: + seen[tid] = cp + except Exception: + logger.debug("langgraph_list_sessions_failed", exc_info=True) + return [] + return [{"session_id": tid, "message_count": None} for tid in seen] + + async def delete_session(self, session_id: str) -> bool: + """Delete a thread's persisted state from the checkpointer.""" + if not self._checkpointer: + return False + try: + await self._checkpointer.adelete_thread(session_id) + return True + except Exception: + logger.debug("langgraph_delete_thread_failed", exc_info=True) + return False - for msg in raw_messages: + async def recent_messages(self, session_id: str, limit: int = 20) -> list[dict]: + """Recent text messages from the thread state (for fact extraction).""" + values = await self._get_state_values(session_id) + if not values: + return [] + out: list[dict] = [] + for msg in values.get("messages", []): msg_type = getattr(msg, "type", "") - + content = getattr(msg, "content", "") if msg_type == "human": - messages.append({"role": "user", "content": msg.content}) - - elif msg_type == "ai": - entry: dict = {"role": "assistant", "content": msg.content or ""} - tool_calls = getattr(msg, "tool_calls", None) - if tool_calls: - entry["tool_calls"] = [ - { - "id": tc.get("id", tc.get("name", "")), - "name": tc["name"], - "args": tc.get("args", {}), - } - for tc in tool_calls - ] - messages.append(entry) - - elif msg_type == "tool": - messages.append({ - "role": "tool", - "tool_call_id": getattr(msg, "tool_call_id", ""), - "name": getattr(msg, "name", "unknown"), - "content": msg.content if isinstance(msg.content, str) else str(msg.content), - }) - - return messages, current_agent - - async def _inject_session_messages( - self, - session_id: str, - messages: list[dict], - current_agent: str | None = None, - ) -> None: - """Inject normalized messages into LangGraph state.""" - if not self._compiled_graph: - raise RuntimeError("LangGraph workflow not initialized") - - from langchain_core.messages import HumanMessage, AIMessage, ToolMessage - - lc_messages = [] - for msg in messages: - role = msg["role"] - if role == "user": - lc_messages.append(HumanMessage(content=msg["content"])) - elif role == "assistant": - kwargs: dict = {"content": msg.get("content", "")} - if msg.get("tool_calls"): - kwargs["tool_calls"] = msg["tool_calls"] - lc_messages.append(AIMessage(**kwargs)) - elif role == "tool": - lc_messages.append(ToolMessage( - content=msg["content"], - tool_call_id=msg.get("tool_call_id", ""), - name=msg.get("name", "unknown"), - )) - - config = {"configurable": {"thread_id": session_id}} - update = {"messages": lc_messages} - if current_agent is not None: - update["current_agent"] = current_agent - - self._compiled_graph.update_state(config, update) + out.append({"role": "user", "content": content}) + elif msg_type == "ai" and content: + out.append({"role": "assistant", "content": content}) + return out[-limit:] # Alias for convenience diff --git a/src/agentic_cli/workflow/langgraph/permission_wrap.py b/src/agentic_cli/workflow/langgraph/permission_wrap.py index 08600fa..d04d6b7 100644 --- a/src/agentic_cli/workflow/langgraph/permission_wrap.py +++ b/src/agentic_cli/workflow/langgraph/permission_wrap.py @@ -4,7 +4,8 @@ Adapter check order matches ADK's PermissionPlugin: 1. EXEMPT tool → returned unwrapped (no engine call ever). 2. Tool has no capability declaration → wrapper returns deny dict at call time. -3. Engine absent from service registry → wrapper runs the original tool (fallback). +3. Engine absent from service registry → fail closed (deny) when permissions + are enabled; run the tool only when permissions are disabled. 4. Otherwise call engine.check(); return on allow, deny dict on deny. """ @@ -47,7 +48,19 @@ async def _guarded(*args: Any, **kwargs: Any) -> Any: "error": "Permission denied: tool has no capability declaration", } engine = get_service(PERMISSION_ENGINE) - if engine is not None: + if engine is None: + # Fail closed: permissions on but no engine wired is a + # misconfiguration (base_manager always builds one) — deny rather + # than run ungated. Only permissions-disabled falls through to run. + from agentic_cli.config import get_settings + + if get_settings().permissions_enabled: + logger.warning("permission_engine_missing", tool=name) + return { + "success": False, + "error": "Permission denied: permission engine unavailable", + } + else: result = await engine.check(name, caps, kwargs) if not result.allowed: return {"success": False, "error": f"Permission denied: {result.reason}"} diff --git a/src/agentic_cli/workflow/mcp.py b/src/agentic_cli/workflow/mcp.py new file mode 100644 index 0000000..fc80c35 --- /dev/null +++ b/src/agentic_cli/workflow/mcp.py @@ -0,0 +1,108 @@ +"""Backend-neutral MCP (Model Context Protocol) server configuration. + +``MCPServerConfig`` describes how to reach an MCP server (stdio subprocess or +remote SSE/HTTP). ``to_adk_toolset`` materializes it into an ADK ``MCPToolset`` +that can be appended to an agent's ``tools`` list; ADK connects lazily on first +use, so construction is synchronous. + +LangGraph materialization (via ``langchain-mcp-adapters``) is deferred — see the +implementation plan. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class MCPServerConfig(BaseModel): + """Connection config for a single MCP server. + + Attributes: + name: Logical name for the server (used for logging / permission targets). + transport: ``stdio`` (local subprocess), ``sse``, or ``http`` + (streamable HTTP). + command/args/env: Used by ``stdio`` — the executable, its arguments, and + extra environment variables. + url/headers: Used by ``sse``/``http`` — the server URL and request headers. + tool_filter: Optional allowlist of MCP tool names to expose. + tool_name_prefix: Optional prefix added to each MCP tool name. + timeout: Optional connection timeout (seconds). + """ + + model_config = ConfigDict(extra="forbid") + + name: str + transport: Literal["stdio", "sse", "http"] = "stdio" + + # stdio + command: str | None = None + args: list[str] = Field(default_factory=list) + env: dict[str, str] = Field(default_factory=dict) + + # sse / http + url: str | None = None + headers: dict[str, str] = Field(default_factory=dict) + + tool_filter: list[str] | None = None + tool_name_prefix: str | None = None + timeout: float | None = None + + @model_validator(mode="after") + def _validate_transport_fields(self) -> "MCPServerConfig": + if self.transport == "stdio": + if not self.command: + raise ValueError("stdio transport requires 'command'") + else: # sse / http + if not self.url: + raise ValueError(f"{self.transport} transport requires 'url'") + return self + + +def build_connection_params(cfg: MCPServerConfig) -> Any: + """Build the ADK connection-params object for an MCP server config. + + Returns one of ``StdioConnectionParams`` / ``SseConnectionParams`` / + ``StreamableHTTPConnectionParams`` depending on transport. + """ + from google.adk.tools.mcp_tool import ( + SseConnectionParams, + StdioConnectionParams, + StreamableHTTPConnectionParams, + ) + + if cfg.transport == "stdio": + from mcp import StdioServerParameters + + server = StdioServerParameters( + command=cfg.command, + args=list(cfg.args), + env=dict(cfg.env) if cfg.env else None, + ) + kwargs: dict[str, Any] = {"server_params": server} + if cfg.timeout is not None: + kwargs["timeout"] = cfg.timeout + return StdioConnectionParams(**kwargs) + + kwargs = {"url": cfg.url, "headers": dict(cfg.headers) if cfg.headers else None} + if cfg.timeout is not None: + kwargs["timeout"] = cfg.timeout + if cfg.transport == "sse": + return SseConnectionParams(**kwargs) + return StreamableHTTPConnectionParams(**kwargs) + + +def to_adk_toolset(cfg: MCPServerConfig) -> Any: + """Materialize an MCPServerConfig into an ADK ``McpToolset``.""" + try: + from google.adk.tools.mcp_tool import McpToolset as _McpToolset + except ImportError: # older ADK only has the (now-deprecated) MCPToolset + from google.adk.tools.mcp_tool import MCPToolset as _McpToolset + + kwargs: dict[str, Any] = {"connection_params": build_connection_params(cfg)} + if cfg.tool_filter is not None: + kwargs["tool_filter"] = list(cfg.tool_filter) + if cfg.tool_name_prefix: + kwargs["tool_name_prefix"] = cfg.tool_name_prefix + return _McpToolset(**kwargs) diff --git a/src/agentic_cli/workflow/model_settings.py b/src/agentic_cli/workflow/model_settings.py new file mode 100644 index 0000000..f62a065 --- /dev/null +++ b/src/agentic_cli/workflow/model_settings.py @@ -0,0 +1,87 @@ +"""Backend-neutral per-agent model parameters. + +``ModelSettings`` carries generation parameters (temperature, top-p/k, max +tokens, stop sequences, thinking/reasoning effort) for a single agent, +independent of the orchestration backend. Each workflow manager translates +these neutral fields into its backend's native config: + +- ADK -> ``google.genai.types.GenerateContentConfig`` + a ``BuiltInPlanner`` + (for the thinking config). +- LangGraph (deferred) -> ``init_chat_model`` keyword arguments. + +The ``extra`` dict is an escape hatch for provider-specific parameters that +have no neutral field; each backend filters it to keys it understands. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + +# Thinking modes. ``none`` disables thinking; ``low``/``medium``/``high`` map to +# the model-family's discrete effort levels; ``budget`` uses an explicit token +# budget (see ``ThinkingSettings.budget_tokens``). +ThinkingMode = Literal["none", "low", "medium", "high", "budget"] + + +class ThinkingSettings(BaseModel): + """Per-agent thinking/reasoning configuration. + + Attributes: + mode: Effort level, or ``"budget"`` to use an explicit token budget. + budget_tokens: Explicit thinking-token budget (only used when + ``mode == "budget"``). + """ + + model_config = ConfigDict(protected_namespaces=()) + + mode: ThinkingMode = "none" + budget_tokens: int | None = Field(default=None, ge=0) + + +class ModelSettings(BaseModel): + """Backend-neutral generation parameters for a single agent. + + All fields are optional; a ``None`` value means "leave the backend/model + default in place". Field names are neutral — backends map them to native + names (e.g. ``max_tokens`` -> ADK ``max_output_tokens``). + + Attributes: + temperature: Sampling temperature. + top_p: Nucleus-sampling probability mass. + top_k: Top-k sampling cutoff. + max_tokens: Maximum output tokens. + stop_sequences: Sequences that stop generation. + thinking: Thinking/reasoning configuration. + extra: Provider-specific parameters passed through to the backend + (filtered per-backend to recognised keys). + """ + + model_config = ConfigDict(protected_namespaces=()) + + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None + max_tokens: int | None = None + stop_sequences: list[str] | None = None + thinking: ThinkingSettings | None = None + extra: dict[str, Any] = Field(default_factory=dict) + + def common_params(self) -> dict[str, Any]: + """Return non-None neutral generation params (excludes thinking/extra). + + Keys use the neutral field names; backends rename as needed. + """ + params: dict[str, Any] = {} + if self.temperature is not None: + params["temperature"] = self.temperature + if self.top_p is not None: + params["top_p"] = self.top_p + if self.top_k is not None: + params["top_k"] = self.top_k + if self.max_tokens is not None: + params["max_tokens"] = self.max_tokens + if self.stop_sequences is not None: + params["stop_sequences"] = self.stop_sequences + return params 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/permissions/capabilities.py b/src/agentic_cli/workflow/permissions/capabilities.py index 0cb1de8..f07aa8e 100644 --- a/src/agentic_cli/workflow/permissions/capabilities.py +++ b/src/agentic_cli/workflow/permissions/capabilities.py @@ -21,6 +21,11 @@ class Capability: name: str # e.g. "filesystem.read" target_arg: str | None = None # arg name holding the target; None → target "*" + optional: bool = False # when the target arg is absent/empty, skip + # this capability (the side effect isn't + # performed) instead of resolving it to a + # spurious target. Only for genuinely + # optional args (e.g. an output path). @dataclass(frozen=True) diff --git a/src/agentic_cli/workflow/permissions/engine.py b/src/agentic_cli/workflow/permissions/engine.py index 998410c..5be4b9c 100644 --- a/src/agentic_cli/workflow/permissions/engine.py +++ b/src/agentic_cli/workflow/permissions/engine.py @@ -32,6 +32,7 @@ from agentic_cli.workflow.permissions.store import ( BUILTIN_RULES, PermissionContext, + load_project_grants, load_rules, ) @@ -108,7 +109,19 @@ def _load_all_rules(self) -> list[Rule]: ) app = self._settings.app_name rules += load_rules(get_user_config_path(app), RuleSource.USER, self._ctx) - rules += load_rules(get_project_config_path(app), RuleSource.PROJECT, self._ctx) + # PROJECT settings.json is untrusted (a cloned repo can ship it): honor + # only deny-rules so a workspace can tighten but never loosen policy. + rules += load_rules( + get_project_config_path(app), + RuleSource.PROJECT, + self._ctx, + allowed_effects=frozenset({Effect.DENY}), + ) + # Interactive "Allow always" grants live in USER config, keyed by the + # resolved project path — trusted (allow+deny). A cloned repo carries + # none (its path won't match), and a repo-shipped permissions.local.json + # is no longer loaded at all. + rules += load_project_grants(app, self._ctx) return rules @property @@ -133,6 +146,11 @@ async def check( resolved = self._resolve(capabilities, args) outcomes = self._evaluate(resolved) + # No capabilities to evaluate (e.g. every cap is optional and its target + # arg was absent) → nothing to gate, allow. + if not outcomes: + return CheckResult(True, "no applicable capabilities") + # DENY wins. deny_hits = [(c, r) for c, r in outcomes if r is not None and r.effect is Effect.DENY] if deny_hits: @@ -145,7 +163,7 @@ async def check( return CheckResult(True, self._fmt_rule_reason(any_r, any_c)) # Ask flow lands in Task 16. - return await self._ask_and_apply(tool_name, resolved, outcomes) + return await self._ask_and_apply(tool_name, resolved, outcomes, args) # ------------------------------------------------------------------ # Helpers @@ -156,9 +174,18 @@ def _resolve( ) -> list[ResolvedCapability]: resolved: list[ResolvedCapability] = [] for cap in capabilities: - raw = "*" if cap.target_arg is None else str(args.get(cap.target_arg, "")) - target = "*" if cap.target_arg is None else get_matcher(cap.name).canonicalize(raw, self._ctx) - resolved.append(ResolvedCapability(cap.name, target)) + if cap.target_arg is None: + resolved.append(ResolvedCapability(cap.name, "*")) + continue + value = args.get(cap.target_arg, "") + if cap.optional and (value is None or value == ""): + # Optional target not supplied → the side effect isn't performed + # this call, so don't resolve (and don't spuriously prompt) it. + continue + matcher = get_matcher(cap.name) + items = value if isinstance(value, (list, tuple)) else [value] + for item in items: + resolved.append(ResolvedCapability(cap.name, matcher.canonicalize(str(item), self._ctx))) return resolved def _evaluate( @@ -193,13 +220,14 @@ async def _ask_and_apply( tool_name: str, resolved: list[ResolvedCapability], outcomes: list[tuple[ResolvedCapability, Rule | None]], + args: dict | None = None, ) -> CheckResult: from agentic_cli.workflow.permissions.prompt import build_request, parse_response from agentic_cli.workflow.permissions.store import append_project_rule unmatched = [cap for cap, r in outcomes if r is None] async with self._ask_lock: - request = build_request(tool_name, resolved) + request = build_request(tool_name, resolved, args) response = await self._workflow.request_user_input(request) scope = parse_response(response) @@ -220,7 +248,7 @@ async def _ask_and_apply( rule = Rule(cap.name, target, Effect.ALLOW, source) self._session_rules.append(rule) if source is RuleSource.PROJECT: - append_project_rule(self._settings.app_name, rule) + append_project_rule(self._settings.app_name, rule, self._ctx.workdir) label = "session" if source is RuleSource.SESSION else "always, saved to project" return CheckResult(True, f"no rule + user allowed ({label})") diff --git a/src/agentic_cli/workflow/permissions/prompt.py b/src/agentic_cli/workflow/permissions/prompt.py index 4bbdbe6..08248d2 100644 --- a/src/agentic_cli/workflow/permissions/prompt.py +++ b/src/agentic_cli/workflow/permissions/prompt.py @@ -4,32 +4,55 @@ import uuid +from agentic_cli.constants import truncate from agentic_cli.workflow.events import InputType, UserInputRequest from agentic_cli.workflow.permissions.capabilities import ResolvedCapability from agentic_cli.workflow.permissions.engine import broaden_target_for_grant from agentic_cli.workflow.permissions.rules import AskScope +# Args that carry executable payloads, shown in the prompt for *.exec grants. +_CODE_ARGS = ("code", "command", "script") +_CODE_PREVIEW_MAX = 1000 + # Strings kept module-level so the UI and parser stay in sync. ALLOW_ONCE_CHOICE = "Allow once" ALLOW_SESSION_CHOICE = "Allow for this session" -ALLOW_ALWAYS_CHOICE = "Allow always (save to project)" +ALLOW_ALWAYS_CHOICE = "Allow always for this project" DENY_CHOICE = "Deny" +# The label this choice used to display. It said "save to project", which +# described the *scope* but implied the grant is written into the repository — +# it is not: grants live in ``~/.{app}/project_grants.json``, keyed by the +# resolved project path (see ``permissions/store.py``). Private: an internal +# parser-compatibility detail, not API. It is still accepted so a response +# captured or queued under the old wording keeps its meaning, but it is never +# offered — callers should use ``ALLOW_ALWAYS_CHOICE``. +_LEGACY_ALLOW_ALWAYS_CHOICE = "Allow always (save to project)" + _CHOICE_TO_SCOPE = { ALLOW_ONCE_CHOICE: AskScope.ONCE, ALLOW_SESSION_CHOICE: AskScope.SESSION, ALLOW_ALWAYS_CHOICE: AskScope.PROJECT, + _LEGACY_ALLOW_ALWAYS_CHOICE: AskScope.PROJECT, DENY_CHOICE: AskScope.DENY, } -def build_request(tool_name: str, capabilities: list[ResolvedCapability]) -> UserInputRequest: +def build_request( + tool_name: str, + capabilities: list[ResolvedCapability], + args: dict | None = None, +) -> UserInputRequest: """Construct a ``UserInputRequest`` (CHOICE) describing the pending grant. The displayed target is the **effective grant scope** — i.e. what will be stored as a rule if the user picks Session or Always. For ``filesystem.*`` that's the parent directory (``/foo/**``) rather than the exact file, so one grant covers every sibling/nested file. + + For code-execution capabilities (``*.exec``) the pending ``code``/``command`` + payload is shown: the capability target is ``*`` (allow-any-code), so the + payload — not the target — is what the user is actually approving. """ lines = [f"Tool `{tool_name}` wants:"] has_broadened_filesystem = False @@ -43,6 +66,11 @@ def build_request(tool_name: str, capabilities: list[ResolvedCapability]) -> Use lines.append("") if has_broadened_filesystem: lines.append("(Grant scope widened to the parent directory.)") + code_preview = _code_preview(capabilities, args) + if code_preview: + lines.append("Code to execute:") + lines.append(code_preview) + lines.append("") lines.append("Allow?") prompt = "\n".join(lines) @@ -61,6 +89,31 @@ def build_request(tool_name: str, capabilities: list[ResolvedCapability]) -> Use ) +def _code_preview( + capabilities: list[ResolvedCapability], args: dict | None +) -> str: + """Truncated preview of the executable payload for ``*.exec`` grants, else ''.""" + if not args: + return "" + # Match the whole exec namespace: "python.exec" AND sub-capabilities like + # "python.exec.stateful" (sandbox_execute). Keying only on the ".exec" + # suffix would skip the stateful kernel — the more dangerous tool. + if not any(cap.name.endswith(".exec") or ".exec." in cap.name + for cap in capabilities): + return "" + for key in _CODE_ARGS: + value = args.get(key) + if isinstance(value, str) and value.strip(): + return truncate(value, _CODE_PREVIEW_MAX) + return "" + + def parse_response(text: str) -> AskScope: - """Parse a choice string into an ``AskScope``. Unknown values deny.""" + """Parse a choice string into an ``AskScope``. Unknown values deny. + + Besides the four labels this module offers, one superseded "allow always" + wording is still accepted, so renaming the displayed choice cannot turn a + user's "always" answer into a denial. That compatibility string is internal + and is never offered as a choice. Anything unrecognised denies. + """ return _CHOICE_TO_SCOPE.get((text or "").strip(), AskScope.DENY) diff --git a/src/agentic_cli/workflow/permissions/store.py b/src/agentic_cli/workflow/permissions/store.py index 2667aa1..8137241 100644 --- a/src/agentic_cli/workflow/permissions/store.py +++ b/src/agentic_cli/workflow/permissions/store.py @@ -12,7 +12,7 @@ from pathlib import Path from agentic_cli.file_utils import atomic_write_text -from agentic_cli.settings_persistence import get_project_config_path +from agentic_cli.settings_persistence import get_user_project_grants_path from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource @@ -23,16 +23,19 @@ class PermissionContext: Attributes: workdir: Absolute current working directory. home: Absolute home directory. + app_name: Application name (drives the ``.{app_name}`` config dir). """ workdir: Path home: Path + app_name: str = "agentic_cli" def substitute(self, s: str) -> str: - """Expand ${workdir} and ${home} in a pattern string.""" + """Expand ${workdir}, ${home} and ${app_name} in a pattern string.""" return ( s.replace("${workdir}", str(self.workdir)) .replace("${home}", str(self.home)) + .replace("${app_name}", self.app_name) ) @@ -57,14 +60,29 @@ def substitute(self, s: str) -> str: Rule("filesystem.write", "${home}/.ssh/**", Effect.DENY, RuleSource.BUILTIN), Rule("filesystem.write", "${home}/.aws/**", Effect.DENY, RuleSource.BUILTIN), Rule("filesystem.write", "${home}/.gnupg/**", Effect.DENY, RuleSource.BUILTIN), + + # The app's own config: writing it would let the agent rewrite permission + # rules / disable the engine (self-escalation). Deny both the project- and + # user-level config dirs. Deny-wins, so this beats any broadened grant. + Rule("filesystem.write", "${workdir}/.${app_name}/**", Effect.DENY, RuleSource.BUILTIN), + Rule("filesystem.write", "${home}/.${app_name}/**", Effect.DENY, RuleSource.BUILTIN), ] -def load_rules(path: Path, source: RuleSource, ctx: PermissionContext) -> list[Rule]: +def load_rules( + path: Path, + source: RuleSource, + ctx: PermissionContext, + allowed_effects: frozenset[Effect] | None = None, +) -> list[Rule]: """Load rules from a settings.json file's ``permissions`` section. Returns an empty list when the file is absent or has no ``permissions`` key. Raises ``ValueError`` if the file is not valid JSON. + + ``allowed_effects`` restricts which effects are honored from this source. + The engine passes ``{Effect.DENY}`` for the untrusted PROJECT file so a + cloned repo can tighten (deny) but never loosen (allow) the policy. """ # Local import to avoid circular dependency: matchers.py imports PermissionContext from here. from agentic_cli.workflow.permissions.matchers import get_matcher # noqa: PLC0415 @@ -79,6 +97,8 @@ def load_rules(path: Path, source: RuleSource, ctx: PermissionContext) -> list[R section = data.get("permissions") or {} rules: list[Rule] = [] for effect_name, effect in (("allow", Effect.ALLOW), ("deny", Effect.DENY)): + if allowed_effects is not None and effect not in allowed_effects: + continue for entry in section.get(effect_name) or []: cap = entry["capability"] target_raw = entry["target"] @@ -87,25 +107,68 @@ def load_rules(path: Path, source: RuleSource, ctx: PermissionContext) -> list[R return rules -def append_project_rule(app_name: str, rule: Rule) -> None: - """Append ``rule`` to the project ``settings.json`` permissions section. +def load_project_grants(app_name: str, ctx: PermissionContext) -> list[Rule]: + """Load interactive 'Allow always' grants for the CURRENT project only. + + Reads ``~/.{app}/project_grants.json`` and returns the allow+deny rules + stored under ``str(ctx.workdir.resolve())`` — trusted (``RuleSource.PROJECT``, + both effects honored), since the user (not a repo) authored them. Returns + ``[]`` when the file or the project's entry is absent. Raises ``ValueError`` + on malformed JSON. + """ + # Local import to avoid a cycle: matchers.py imports PermissionContext here. + from agentic_cli.workflow.permissions.matchers import get_matcher # noqa: PLC0415 + + path = get_user_project_grants_path(app_name) + if not path.exists(): + return [] + try: + data = json.loads(path.read_text()) + except json.JSONDecodeError as exc: + raise ValueError(f"Malformed JSON in {path}: {exc}") from exc + + if not isinstance(data, dict): + raise ValueError(f"Expected a JSON object in {path}, got {type(data).__name__}") + + section = (data.get(str(ctx.workdir.resolve())) or {}).get("permissions") or {} + rules: list[Rule] = [] + for effect_name, effect in (("allow", Effect.ALLOW), ("deny", Effect.DENY)): + for entry in section.get(effect_name) or []: + cap = entry["capability"] + target = get_matcher(cap).canonicalize(entry["target"], ctx) + rules.append(Rule(cap, target, effect, RuleSource.PROJECT)) + return rules + + +def append_project_rule(app_name: str, rule: Rule, project_root: Path) -> None: + """Persist an interactive 'Allow always' grant to the USER-side, path-keyed + grants file (``~/.{app}/project_grants.json``), under ``project_root``'s + resolved path. - Creates the file (and its parent directory) if absent; preserves every - other settings key; dedupes by exact ``(capability, target)``. Atomic - rewrite via ``file_utils.atomic_write_text``. + Kept out of the repo (P0-1) so a cloned workspace carries no grants: a clone + at a different path has no matching entry and the user re-grants. Creates the + file if absent; dedupes by exact ``(capability, target)`` within the + project's section; atomic rewrite via ``atomic_write_text``. Only ``Rule`` instances with ``source == RuleSource.PROJECT`` should be - passed here — this helper doesn't validate (engine enforces the - invariant). + passed here — this helper doesn't validate (engine enforces the invariant). """ - path = get_project_config_path(app_name) + path = get_user_project_grants_path(app_name) try: data = json.loads(path.read_text()) if path.exists() else {} except json.JSONDecodeError as exc: raise ValueError(f"Malformed JSON in {path}: {exc}") from exc + if not isinstance(data, dict): + raise ValueError(f"Expected a JSON object in {path}, got {type(data).__name__}") + + proj_key = str(project_root.resolve()) key = "allow" if rule.effect is Effect.ALLOW else "deny" - section = data.setdefault("permissions", {}).setdefault(key, []) + section = ( + data.setdefault(proj_key, {}) + .setdefault("permissions", {}) + .setdefault(key, []) + ) entry = {"capability": rule.capability, "target": rule.target} if entry not in section: section.append(entry) diff --git a/src/agentic_cli/workflow/retry.py b/src/agentic_cli/workflow/retry.py index 6a115f8..997195b 100644 --- a/src/agentic_cli/workflow/retry.py +++ b/src/agentic_cli/workflow/retry.py @@ -11,8 +11,12 @@ def is_rate_limit_error(error: Exception) -> bool: """Check if an exception is a 429 rate-limit / RESOURCE_EXHAUSTED error.""" + # Gemini/genai errors expose ``.code``; anthropic.RateLimitError (and other + # httpx-based SDKs) expose ``.status_code`` instead. if getattr(error, "code", None) == 429: return True + if getattr(error, "status_code", None) == 429: + return True if "RESOURCE_EXHAUSTED" in str(error): return True return False diff --git a/src/agentic_cli/workflow/service_registry.py b/src/agentic_cli/workflow/service_registry.py index 612922a..698921d 100644 --- a/src/agentic_cli/workflow/service_registry.py +++ b/src/agentic_cli/workflow/service_registry.py @@ -14,6 +14,7 @@ # ---- Well-known registry keys ---- ARXIV_SOURCE = "arxiv_source" +JOB_MANAGER = "job_manager" KB_MANAGER = "kb_manager" LLM_SUMMARIZER = "llm_summarizer" MEMORY_STORE = "memory_store" @@ -22,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/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/src/agentic_cli/workflow/settings.py b/src/agentic_cli/workflow/settings.py index 143008c..b152e0e 100644 --- a/src/agentic_cli/workflow/settings.py +++ b/src/agentic_cli/workflow/settings.py @@ -10,16 +10,26 @@ from enum import Enum from typing import Literal, TYPE_CHECKING -from pydantic import BaseModel, Field +from pydantic import AliasChoices, 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.""" @@ -92,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 @@ -206,12 +228,6 @@ class WorkflowSettingsMixin: description="Workflow orchestrator backend", json_schema_extra={"ui_order": 100}, # Advanced setting ) - langgraph_checkpointer: Literal["memory", "postgres"] | None = Field( - default="memory", - title="LangGraph Checkpointer", - description="LangGraph state persistence type", - json_schema_extra={"ui_order": 101}, - ) # Retry configuration retry_max_attempts: int = Field( @@ -232,6 +248,16 @@ class WorkflowSettingsMixin: description="Multiplier for exponential backoff between retries", json_schema_extra={"ui_order": 112}, ) + anthropic_request_timeout: float = Field( + default=900.0, + title="Anthropic Request Timeout", + description=( + "Overall timeout (seconds) for direct-API Claude requests. A " + "non-default value lets high-thinking (large max_tokens) requests " + "run without the SDK's streaming-required guard." + ), + json_schema_extra={"ui_order": 113}, + ) # Python executor python_executor_timeout: int = Field( @@ -248,11 +274,16 @@ class WorkflowSettingsMixin: ) # Sandbox executor (stateful Jupyter-backed execution) - sandbox_backend: str = Field( - default="jupyter_local", - title="Sandbox Backend", - description="Backend for stateful sandbox execution", - json_schema_extra={"ui_order": 122}, + stateful_executor_backend: Literal["none", "local", "docker"] = Field( + default="none", + title="Stateful Executor Backend", + description=( + "Backend for the stateful sandbox_execute tool. 'none' disables it; " + "'docker' runs in a network-isolated container (recommended); 'local' " + "runs a Jupyter kernel with host privileges (NOT OS-sandboxed). Future: " + "'modal', 'runpod'." + ), + json_schema_extra={"ui_order": 121}, ) sandbox_timeout: int = Field( default=120, @@ -272,25 +303,108 @@ class WorkflowSettingsMixin: description="Additional pip packages to pre-install in sandbox sessions (informational for local backend, drives image build for Docker backend)", json_schema_extra={"ui_order": 125}, ) + sandbox_image: str = Field( + default="quay.io/jupyter/scipy-notebook:python-3.12", + title="Sandbox Image", + description="Container image for the jupyter_docker backend. Must contain ipykernel/jupyter_client. Pin to a digest in production.", + json_schema_extra={"ui_order": 126}, + ) + sandbox_memory_mb: int = Field( + default=2048, + title="Sandbox Memory (MB)", + description="Per-container memory cap for the docker backend; also disables swap.", + json_schema_extra={"ui_order": 127}, + ) + sandbox_cpus: float = Field( + default=2.0, + title="Sandbox CPUs", + description="Per-container CPU cap for the docker backend.", + json_schema_extra={"ui_order": 128}, + ) + sandbox_pids_limit: int = Field( + default=256, + title="Sandbox PID Limit", + description="Per-container process/thread cap for the docker backend.", + json_schema_extra={"ui_order": 129}, + ) + sandbox_network: str = Field( + default="none", + title="Sandbox Network", + description="Docker network mode for the docker backend. v1 supports 'none' only.", + json_schema_extra={"ui_order": 130}, + ) + sandbox_container_user: str = Field( + default="", + title="Sandbox Container User", + description="uid:gid to run the container as; empty uses the image default.", + json_schema_extra={"ui_order": 131}, + ) + sandbox_data_mounts: list[str] = Field( + default_factory=list, + title="Sandbox Data Mounts", + description="Read-only data staged into the container as 'host_path:mount_name' (mounted read-only under /data/).", + json_schema_extra={"ui_order": 132}, + ) + sandbox_start_timeout: int = Field( + default=180, + title="Sandbox Start Timeout", + description="Seconds to wait for container start + image pull + kernel readiness (docker backend).", + json_schema_extra={"ui_order": 133}, + ) + sandbox_outputs_dir: str = Field( + default="", + title="Sandbox Outputs Dir", + description="Shared host dir mounted at /workspace/outputs for FINAL deliverables (default: /artifacts).", + json_schema_extra={"ui_order": 134}, + ) + + @field_validator("sandbox_network") + @classmethod + def _validate_sandbox_network(cls, v: str) -> str: + """Fail closed: the docker backend's no-egress isolation depends on + --network none, so reject any other value rather than silently + weakening it. (v1 supports 'none' only.)""" + if v != "none": + raise ValueError( + f"sandbox_network must be 'none' (got {v!r}). The docker sandbox's " + "network-isolation guarantee depends on it; other modes are not " + "supported in v1." + ) + return v # OS-level sandboxing os_sandbox_enabled: bool = Field( - default=False, + default=True, title="OS Sandbox Enabled", - description="Enable OS-level sandboxing for shell and Python execution (requires sandbox-exec on macOS or bwrap on Linux)", - json_schema_extra={"ui_order": 130}, + description=( + "Wrap Python execution in an OS-level sandbox when a backend is " + "available (sandbox-exec on macOS, bwrap on Linux). When no backend " + "is present, execution falls back to the restricted in-process " + "executor (see os_sandbox_strict) — only pure-computation modules " + "are importable in that case." + ), + json_schema_extra={"ui_order": 134}, + ) + os_sandbox_strict: bool = Field( + default=False, + title="OS Sandbox Strict", + description=( + "Refuse to run code when OS sandboxing is enabled but no backend is " + "available, instead of falling back to the in-process executor." + ), + json_schema_extra={"ui_order": 137}, ) os_sandbox_writable_paths: list[str] = Field( default_factory=list, title="OS Sandbox Writable Paths", description="Additional paths the sandboxed process can write to (working directory is always writable)", - json_schema_extra={"ui_order": 131}, + json_schema_extra={"ui_order": 135}, ) os_sandbox_allow_network: bool = Field( default=False, title="OS Sandbox Allow Network", description="Allow network access from sandboxed processes", - json_schema_extra={"ui_order": 132}, + json_schema_extra={"ui_order": 136}, ) # Permissions @@ -298,53 +412,117 @@ class WorkflowSettingsMixin: default_factory=PermissionsConfig, title="Permissions", description="Declarative allow/deny rules for tool capabilities.", - json_schema_extra={"ui_order": 135}, + json_schema_extra={"ui_order": 138}, ) permissions_enabled: bool = Field( default=True, title="Permissions Enabled", description="Master switch; when False, all tool calls are allowed.", - json_schema_extra={"ui_order": 136}, + json_schema_extra={"ui_order": 139}, + ) + max_concurrent_jobs: int = Field( + default=4, + ge=1, + title="Max Concurrent Jobs", + description="Maximum long-running jobs running at once; excess are queued.", + json_schema_extra={"ui_order": 140}, + ) + job_auto_resume: bool = Field( + default=False, + title="Auto-resume Finished Jobs", + description=( + "When True, a finished long-running job that opted in " + "(resume_on_complete) automatically resumes the agent with its " + "result at the next turn boundary (or via /resume)." + ), + json_schema_extra={"ui_order": 141}, ) + # Session persistence — durable conversations across restarts. + # Drives BOTH backends: ADK uses DatabaseSessionService, LangGraph uses a + # persistent checkpointer (both keyed by session_id). "memory" = ephemeral. + session_store: Literal["memory", "sqlite", "postgres"] = Field( + default="sqlite", + title="Session Store", + description=( + "Where conversations are persisted: sqlite (default, a single file), " + "postgres (shared/multi-instance via Postgres URI), or memory (ephemeral)." + ), + json_schema_extra={"ui_order": 145}, + ) + + # Skills (Agent Skills / SKILL.md folders) + skills_dirs: list[str] = Field( + default_factory=list, + title="Skills Directories", + description="Directories searched for named skills (Agent Skills / SKILL.md folders)", + json_schema_extra={"ui_order": 141}, + ) + # 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( default=None, title="PostgreSQL URI", description="PostgreSQL connection URI for persistent storage", - json_schema_extra={"ui_order": 145}, + json_schema_extra={"ui_order": 146}, ) sqlite_uri: str | None = Field( default=None, title="SQLite URI", description="SQLite connection URI or file path for persistent storage", - json_schema_extra={"ui_order": 146}, + json_schema_extra={"ui_order": 147}, ) store_type: Literal["memory", "postgres"] | None = Field( default="memory", title="Store Type", description="Store type for long-term memory (memory or postgres)", - json_schema_extra={"ui_order": 147}, + json_schema_extra={"ui_order": 148}, ) + def session_db_url(self) -> str | None: + """Async SQLAlchemy URL for the session store, or None when ephemeral. + + Shared by both backends so ADK's DatabaseSessionService and LangGraph's + checkpointer persist to the same place. SQLite is the zero-config + default (``{workspace}/sessions/sessions.db``); Postgres via uri. + """ + store = getattr(self, "session_store", "sqlite") + if store == "memory": + return None + if store == "postgres": + uri = self.postgres_uri + if not uri: + raise ValueError("session_store='postgres' requires postgres_uri") + return uri.replace("postgresql://", "postgresql+asyncpg://", 1) + # sqlite (default) + if self.sqlite_uri: + return self.sqlite_uri.replace("sqlite:///", "sqlite+aiosqlite:///", 1) + path = self.sessions_dir / "sessions.db" + return f"sqlite+aiosqlite:///{path}" + # Shell execution settings (for shell middleware) shell_sandbox_type: Literal["host", "docker"] = Field( default="host", title="Shell Sandbox Type", description="Execution environment for shell commands", - json_schema_extra={"ui_order": 148}, + json_schema_extra={"ui_order": 149}, ) shell_docker_image: str = Field( default="python:3.12-slim", title="Shell Docker Image", description="Docker image to use for sandboxed shell execution", - json_schema_extra={"ui_order": 149}, + json_schema_extra={"ui_order": 150}, ) shell_timeout: int = Field( default=60, title="Shell Timeout", description="Default timeout in seconds for shell commands", - json_schema_extra={"ui_order": 150}, + json_schema_extra={"ui_order": 151}, ) # LLM debugging settings @@ -468,22 +646,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.""" @@ -495,11 +719,25 @@ def set_thinking_effort(self, effort: str) -> None: object.__setattr__(self, "thinking_effort", effort) def export_api_keys_to_env(self) -> None: - """Export API keys to environment variables.""" + """Export configured API keys to provider environment variables. + + Provider SDKs used by the orchestrators (ADK's AnthropicLlm/Gemini, + LangChain clients) read credentials from process env vars. The export + OVERWRITES the env from this settings instance: the key fields bind + only via their env alias (real env vars are the highest-priority + source), so settings and env diverge only when the process env + changed after this instance loaded — e.g. an earlier manager's + export, or a key loaded from a class-specific env_file — and then + this instance's configured value must win. (The previous + set-if-absent guard let the first exporting manager pin credentials + for every later one.) A key unset in settings leaves the environment + untouched. Credentials are process-global; SettingsContext does not + isolate them. + """ import os - if self.google_api_key and not os.environ.get("GOOGLE_API_KEY"): + if self.google_api_key: os.environ["GOOGLE_API_KEY"] = self.google_api_key - if self.anthropic_api_key and not os.environ.get("ANTHROPIC_API_KEY"): + if self.anthropic_api_key: os.environ["ANTHROPIC_API_KEY"] = self.anthropic_api_key diff --git a/src/agentic_cli/workflow/tool_summaries.py b/src/agentic_cli/workflow/tool_summaries.py index ebbf30b..5586c1e 100644 --- a/src/agentic_cli/workflow/tool_summaries.py +++ b/src/agentic_cli/workflow/tool_summaries.py @@ -76,6 +76,16 @@ def _shell_executor(r: dict) -> str: return f"Exit {code} ({dur:.1f}s)" +def _sandbox_execute(r: dict) -> str: + stdout = r.get("stdout", "") + if stdout and stdout.strip(): + return truncate(stdout.strip().splitlines()[0], TOOL_SUMMARY_MAX_LENGTH) + if not r.get("success", True): + err = r.get("error") or r.get("stderr") or "failed" + return truncate(str(err).strip().splitlines()[0], TOOL_SUMMARY_MAX_LENGTH) + return "ok" + + def _save_tasks(r: dict) -> str: return r.get("message", f"{r['count']} tasks saved") @@ -157,6 +167,7 @@ def _list_documents(r: dict) -> str: "edit_file": _edit_file, "execute_python": _execute_python, "shell_executor": _shell_executor, + "sandbox_execute": _sandbox_execute, "save_tasks": _save_tasks, "get_tasks": _get_tasks, "get_plan": _get_plan, diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/cli/test_job_monitor.py b/tests/cli/test_job_monitor.py new file mode 100644 index 0000000..2c61f61 --- /dev/null +++ b/tests/cli/test_job_monitor.py @@ -0,0 +1,217 @@ +"""Tests for the harness JobMonitor (milestone 2: jobs UI monitor). + +The monitor is the background task that keeps long-running jobs visible in the +status bar. These tests drive it against a real ``JobManager`` running real +subprocess jobs, plus a couple of focused unit tests for the segment logic. +""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from agentic_cli.cli.job_monitor import JobMonitor +from agentic_cli.tools.jobs import JobManager, JobState + + +# --- Lightweight fakes --------------------------------------------------------- + + +class FakeUI: + """Captures every status string set on the (fake) session.""" + + def __init__(self) -> None: + self.statuses: list[str] = [] + + def set_status(self, text) -> None: + self.statuses.append(str(text)) + + +class FakeController: + """Minimal stand-in for WorkflowController as the monitor sees it. + + Composes the status bar the same way the real controller does + (model | jobs | hints) so the integration assertions are meaningful. + """ + + def __init__(self, ui: FakeUI, job_manager) -> None: + self.is_ready = True + self.workflow = SimpleNamespace(model="test-model", job_manager=job_manager) + self.usage_tracker = None + self.jobs_status_segment: str | None = None + self._ui = ui + self.status_updates = 0 + + def update_status_bar(self, ui) -> None: + self.status_updates += 1 + parts = [self.workflow.model] + if self.jobs_status_segment: + parts.append(self.jobs_status_segment) + parts.extend(["Ctrl+C: cancel", "/help: commands"]) + ui.set_status(" | ".join(parts)) + + +@pytest.fixture +def jm(tmp_path: Path) -> JobManager: + return JobManager(base_dir=tmp_path / "jobs", max_concurrent=2) + + +async def _await_until(predicate, timeout: float = 5.0, interval: float = 0.02) -> bool: + """Async poll that yields to the loop so background tasks can run.""" + end = time.time() + timeout + while time.time() < end: + if predicate(): + return True + await asyncio.sleep(interval) + return predicate() + + +# --- Segment logic (no asyncio) ------------------------------------------------ + + +class TestSegmentLogic: + def test_no_jobs_is_none(self, jm: JobManager): + ui = FakeUI() + mon = JobMonitor(ui, FakeController(ui, jm)) + assert mon._build_segment([]) is None + + def test_running_and_queued_counts(self, jm: JobManager): + ui = FakeUI() + mon = JobMonitor(ui, FakeController(ui, jm)) + for _ in range(3): + jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "sleep 1"}) + jm.reconcile() # 2 start (cap=2), 1 stays queued + seg = mon._build_segment(jm.list()) + assert seg is not None + assert "2 running" in seg + assert "1 queued" in seg + + def test_finished_job_note_then_ages_out(self, jm: JobManager): + ui = FakeUI() + mon = JobMonitor(ui, FakeController(ui, jm), note_ticks=2) + rec = jm.submit( + tool="run_shell_job", backend="subprocess", + spec={"command": "echo hi"}, name="greet", + ) + # First observe it active so the terminal transition is detectable. + rec.state = JobState.RUNNING + mon._build_segment([rec]) + # Now mark terminal and rebuild → a ✓ note appears. + rec.state = JobState.SUCCEEDED + seg = mon._build_segment([rec]) + assert "✓ greet" in seg + # The note lingers for note_ticks builds, then disappears. + mon._build_segment([rec]) # tick 2 + seg_gone = mon._build_segment([rec]) # aged out + assert seg_gone is None + + def test_no_note_for_jobs_already_terminal_on_first_sight(self, jm: JobManager): + ui = FakeUI() + mon = JobMonitor(ui, FakeController(ui, jm)) + rec = jm.submit(tool="t", backend="subprocess", spec={"command": "echo x"}, name="old") + rec.state = JobState.SUCCEEDED # first time we ever see it: already done + assert mon._build_segment([rec]) is None + + def test_short_truncates_long_names(self): + assert JobMonitor._short("x" * 40, limit=10) == "x" * 9 + "…" + assert JobMonitor._short("", ) == "job" + + def test_pending_resume_cue_when_enabled(self, jm: JobManager): + from agentic_cli.tools.jobs import JobRecord + + ui = FakeUI() + mon = JobMonitor( + ui, FakeController(ui, jm), settings=SimpleNamespace(job_auto_resume=True) + ) + rec = JobRecord( + job_id="j1", tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + ) + seg = mon._build_segment([rec]) + assert seg is not None and "↻1 to resume" in seg + + def test_no_resume_cue_when_disabled(self, jm: JobManager): + from agentic_cli.tools.jobs import JobRecord + + ui = FakeUI() + mon = JobMonitor( + ui, FakeController(ui, jm), settings=SimpleNamespace(job_auto_resume=False) + ) + rec = JobRecord( + job_id="j1", tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + ) + # First-sight terminal job → no ✓ note, and the cue is gated off → None. + assert mon._build_segment([rec]) is None + + def test_no_resume_cue_for_already_resumed(self, jm: JobManager): + from agentic_cli.tools.jobs import JobRecord + + ui = FakeUI() + mon = JobMonitor( + ui, FakeController(ui, jm), settings=SimpleNamespace(job_auto_resume=True) + ) + rec = JobRecord( + job_id="j1", tool="run_shell_job", backend="subprocess", name="build", + state=JobState.SUCCEEDED, resume_on_complete=True, call_id="c1", + resume_state="delivered", + ) + assert mon._build_segment([rec]) is None + + +# --- poll_once publishes + refreshes ------------------------------------------- + + +class TestPollOnce: + def test_publishes_segment_and_refreshes_on_change(self, jm: JobManager): + ui = FakeUI() + ctrl = FakeController(ui, jm) + mon = JobMonitor(ui, ctrl) + + jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "sleep 2"}) + seg = mon.poll_once() + assert seg is not None and "running" in seg + assert ctrl.jobs_status_segment == seg + assert ui.statuses and "running" in ui.statuses[-1] + + def test_no_redraw_when_segment_unchanged(self, jm: JobManager): + ui = FakeUI() + ctrl = FakeController(ui, jm) + mon = JobMonitor(ui, ctrl) + # No jobs → segment None on every tick → no status updates at all. + mon.poll_once() + mon.poll_once() + assert ctrl.status_updates == 0 + + def test_no_manager_is_safe(self, tmp_path: Path): + ui = FakeUI() + ctrl = FakeController(ui, None) # workflow has no job_manager + mon = JobMonitor(ui, ctrl) + assert mon.poll_once() is None + assert ctrl.status_updates == 0 + + +# --- Full background loop against real subprocess jobs ------------------------- + + +class TestBackgroundLoop: + async def test_live_job_appears_then_completes(self, jm: JobManager): + ui = FakeUI() + ctrl = FakeController(ui, jm) + mon = JobMonitor(ui, ctrl, interval=0.05, note_ticks=3) + + async with mon.running(): + jm.submit( + tool="run_shell_job", backend="subprocess", + spec={"command": "sleep 0.3"}, name="sleeper", + ) + # The running job shows up in the status bar. + saw_running = await _await_until(lambda: any("running" in s for s in ui.statuses)) + assert saw_running, ui.statuses + # And once it finishes, a ✓ note for it appears. + saw_done = await _await_until(lambda: any("✓ sleeper" in s for s in ui.statuses)) + assert saw_done, ui.statuses diff --git a/tests/cli/test_process_resume.py b/tests/cli/test_process_resume.py new file mode 100644 index 0000000..6fdb527 --- /dev/null +++ b/tests/cli/test_process_resume.py @@ -0,0 +1,96 @@ +"""MessageProcessor.process_resume: resume vs. graceful notice (milestone 4 / A). + +A finished job resumes only when the backend supports it AND the originating +conversation is still available (workflow.can_resume). Otherwise a one-line +notice is posted and no turn runs — the "finished while its conversation was +unavailable" case (e.g. after a CLI restart, the in-memory ADK session is gone). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from agentic_cli.cli.message_processor import MessageProcessor +from agentic_cli.workflow.events import EventType, WorkflowEvent + +from tests.event_replay import RecordingSession + + +class _ResumableWorkflow: + backend_type = "fake" + + def __init__(self, can: bool = True) -> None: + self._can = can + self.resumed: list[str] = [] + + async def can_resume(self, record) -> bool: + return self._can + + def set_input_callback(self, cb) -> None: # used by _run_turn + pass + + def clear_input_callback(self) -> None: + pass + + async def resume_with_job_result(self, record, result=None): + self.resumed.append(record.job_id) + yield WorkflowEvent(type=EventType.TEXT, content="resumed reply") + + +class _NoResumeWorkflow: + """A backend without resume support (e.g. LangGraph today).""" + + backend_type = "langgraph" + + async def can_resume(self, record) -> bool: + return True # irrelevant: no resume_with_job_result attribute + + +class _Ctrl: + def __init__(self, workflow) -> None: + self.workflow = workflow + + async def ensure_initialized(self, ui=None) -> bool: + return True + + def update_status_bar(self, ui) -> None: + pass + + +def _record(job_id="j1", name="build", state="succeeded"): + return SimpleNamespace(job_id=job_id, name=name, state=SimpleNamespace(value=state)) + + +_SETTINGS = SimpleNamespace(default_user="u", verbose_thinking=False) + + +async def test_resumes_when_resumable(): + wf = _ResumableWorkflow(can=True) + ui = RecordingSession() + await MessageProcessor().process_resume( + record=_record(), workflow_controller=_Ctrl(wf), ui=ui, settings=_SETTINGS + ) + assert wf.resumed == ["j1"] + assert "resumed reply" in ui.responses() + assert any("resuming" in c[2] for c in ui.of("message")) + + +async def test_notifies_when_conversation_unavailable(): + wf = _ResumableWorkflow(can=False) + ui = RecordingSession() + await MessageProcessor().process_resume( + record=_record(), workflow_controller=_Ctrl(wf), ui=ui, settings=_SETTINGS + ) + assert wf.resumed == [] # never resumed + assert not ui.responses() # no model turn + assert any("/jobs j1" in c[2] for c in ui.of("message")) + + +async def test_notifies_when_backend_has_no_resume_support(): + ui = RecordingSession() + await MessageProcessor().process_resume( + record=_record(), workflow_controller=_Ctrl(_NoResumeWorkflow()), ui=ui, + settings=_SETTINGS, + ) + assert not ui.responses() + assert any("/jobs j1" in c[2] for c in ui.of("message")) diff --git a/tests/cli/test_resume_coordinator.py b/tests/cli/test_resume_coordinator.py new file mode 100644 index 0000000..b366c0a --- /dev/null +++ b/tests/cli/test_resume_coordinator.py @@ -0,0 +1,314 @@ +"""Coordinator: BaseCLIApp.resume_finished_jobs drains awaiting jobs (milestone 3). + +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 + +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.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.claimed] + + 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, 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 + ): + self.resumed.append(record.job_id) + return self._result + + +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(result) + app.session = object() + app._settings = SimpleNamespace(job_auto_resume=True) + app._usage_tracker = None + return app, jm + + +async def test_resumes_each_awaiting_job_once(): + recs = [SimpleNamespace(job_id="a"), SimpleNamespace(job_id="b")] + app, jm = _app(recs) + n = await app.resume_finished_jobs() + assert n == 2 + assert app._message_processor.resumed == ["a", "b"] + assert jm.completed == [("a", True, None), ("b", True, None)] + + +async def test_claims_before_processing_and_records_after(): + order: list = [] + app, jm = _app([SimpleNamespace(job_id="a")]) + + 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 == [("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(): + app, _ = _app([], has_jm=False) + assert await app.resume_finished_jobs() == 0 + + +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/cli/test_sandbox_render.py b/tests/cli/test_sandbox_render.py new file mode 100644 index 0000000..2863ea8 --- /dev/null +++ b/tests/cli/test_sandbox_render.py @@ -0,0 +1,142 @@ +"""Tests for sandbox_render — Rich rendering helpers for sandbox_execute. + +Pure functions returning Rich renderables (or None). Content is asserted by +rendering to plain text via a captured Console; styling (brightness/color) is +asserted by inspecting the Text spans. +""" +from __future__ import annotations + +from rich.console import Console +from rich.text import Text + +from agentic_cli.cli.sandbox_render import ( + CODE_MAX_LINES, + OUTPUT_MAX_LINES, + render_sandbox_code, + render_sandbox_result, +) + + +def _to_text(renderable) -> str: + if renderable is None: + return "" + console = Console(width=200, no_color=True) + with console.capture() as cap: + console.print(renderable) + return cap.get() + + +def _styles_at(text: Text, index: int) -> list[str]: + """String forms of every span style covering a character index.""" + return [str(s.style) for s in text.spans if s.start <= index < s.end] + + +class TestRenderCode: + def test_none_for_empty_or_blank_code(self): + assert render_sandbox_code("") is None + assert render_sandbox_code(" \n \n") is None + + def test_header_is_first_line_with_green_dash(self): + block = render_sandbox_code("print('hi')") + assert block.plain.splitlines()[0] == "- Running Python in the stateful executor" + assert any("green" in s for s in _styles_at(block, 0)), block.spans + + def test_short_block_has_no_dim_styling(self): + block = render_sandbox_code("import os\nprint(os.getcwd())") + assert not any("dim" in str(sp.style).lower() for sp in block.spans), block.spans + + def test_code_lines_indented_four_no_marker(self): + out = _to_text(render_sandbox_code("a = 1\nb = 2")) + assert "╰" not in out + lines = out.splitlines() + assert lines[1].startswith(" ") and lines[1].strip() == "a = 1" + assert lines[1].index("a") == 4 + assert lines[2].startswith(" ") and lines[2].strip() == "b = 2" + + def test_short_code_no_more_lines_marker(self): + out = _to_text(render_sandbox_code("import os\nprint(os.getcwd())")) + assert "import os" in out and "print(os.getcwd())" in out + assert "more lines" not in out + + def test_long_code_truncated_to_cap_with_marker(self): + code = "\n".join(f"line_{i} = {i}" for i in range(40)) + out = _to_text(render_sandbox_code(code, max_lines=25)) + assert "line_0 = 0" in out + assert "line_24 = 24" in out + assert "line_25 = 25" not in out + assert "+15 more lines" in out + assert "╰" not in out + + def test_default_cap_is_25(self): + assert CODE_MAX_LINES == 25 + + +class TestRenderResult: + def test_success_fixed_header_then_indented_output(self): + block = render_sandbox_result( + {"success": True, "stdout": "line one\nline two\n"}, success=True + ) + lines = block.plain.splitlines() + assert lines[0] == "+ Stateful Python executor output:" + assert any("green" in s for s in _styles_at(block, 0)), block.spans # "+" green + assert lines[1] == " line one" + assert lines[2] == " line two" + assert "╰" not in block.plain + + def test_no_stdout_says_no_output(self): + out = _to_text(render_sandbox_result({"success": True, "stdout": ""}, success=True)) + assert out.strip() == "+ Stateful Python executor: no output" + + def test_single_line_output_indented_under_header(self): + out = _to_text(render_sandbox_result({"success": True, "stdout": "only line\n"}, success=True)) + assert out.splitlines() == ["+ Stateful Python executor output:", " only line"] + + def test_failure_uses_red_x_and_error_label(self): + block = render_sandbox_result( + {"success": False, "error": "NameError: x", "stdout": ""}, success=False + ) + lines = block.plain.splitlines() + assert lines[0] == "x Stateful Python executor error:" + assert lines[1] == " NameError: x" + assert any("red" in s for s in _styles_at(block, 0)), block.spans + + def test_failure_falls_back_to_stderr(self): + out = _to_text( + render_sandbox_result( + {"success": False, "error": "", "stderr": "Trace boom", "stdout": ""}, + success=False, + ) + ) + assert out.splitlines()[:2] == ["x Stateful Python executor error:", " Trace boom"] + + def test_stdout_on_failure_uses_output_label(self): + out = _to_text( + render_sandbox_result( + {"success": False, "stdout": "partial\n", "error": "later"}, success=False + ) + ) + # body is stdout -> "output" label even though the run failed (red x) + assert out.splitlines()[:2] == ["x Stateful Python executor output:", " partial"] + + def test_failure_no_body_says_failed(self): + out = _to_text( + render_sandbox_result( + {"success": False, "stdout": "", "error": "", "stderr": ""}, success=False + ) + ) + assert out.strip() == "x Stateful Python executor failed" + + def test_output_truncated_to_cap_with_marker(self): + stdout = "\n".join(f"out{i}" for i in range(30)) + out = _to_text( + render_sandbox_result({"success": True, "stdout": stdout}, success=True, max_lines=10) + ) + assert out.splitlines()[0] == "+ Stateful Python executor output:" + assert " out0" in out + assert " out9" in out # 10 output lines shown + assert "out10" not in out + assert "+20 more lines" in out + assert "╰" not in out + + def test_default_cap_is_10(self): + assert OUTPUT_MAX_LINES == 10 diff --git a/tests/cli/test_sessions_command.py b/tests/cli/test_sessions_command.py new file mode 100644 index 0000000..a88f4b3 --- /dev/null +++ b/tests/cli/test_sessions_command.py @@ -0,0 +1,93 @@ +"""/sessions command backed by the native session store (durable sessions).""" + +from __future__ import annotations + +import pytest + +from agentic_cli.cli.builtin_commands import SessionsCommand + +from tests.event_replay import RecordingSession + + +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, *, user_id: str | None = None) -> list[dict]: + return self._sessions + + 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 + self.session = RecordingSession() + self.session_id = session_id + + @property + def workflow(self): + if self._wf is None: + raise RuntimeError("workflow not initialized") + return self._wf + + +async def test_lists_sessions(): + app = _App(_Workflow([ + {"session_id": "a", "last_update": None, "message_count": 3}, + {"session_id": "cur", "last_update": 1_700_000_000.0, "message_count": 1}, + ])) + await SessionsCommand().execute("", app) + assert app.session.of("rich") # rendered a table + + +async def test_empty_message(): + app = _App(_Workflow([])) + await SessionsCommand().execute("", app) + assert any("No saved sessions" in c[2] for c in app.session.of("message")) + + +async def test_delete_existing(): + wf = _Workflow([{"session_id": "a", "last_update": None, "message_count": 1}]) + app = _App(wf) + await SessionsCommand().execute("--delete=a", app) + assert wf.deleted == ["a"] + assert app.session.of("success") + + +async def test_delete_missing_errors(): + app = _App(_Workflow([])) + await SessionsCommand().execute("--delete=zzz", app) + assert app.session.errors() + + +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/cli/test_settings_ui_scope.py b/tests/cli/test_settings_ui_scope.py new file mode 100644 index 0000000..60bfb92 --- /dev/null +++ b/tests/cli/test_settings_ui_scope.py @@ -0,0 +1,93 @@ +"""Guard: the /settings dialog exposes only project-scoped (allowlisted) settings. + +User-scoped keys (not in PROJECT_SETTABLE_KEYS) must never render in the +dialog, regardless of what a domain app returns from get_ui_setting_keys() — +otherwise a /settings edit would persist to the user ~/.{app}/settings.json +and apply across all projects. +""" + +from __future__ import annotations + +import pytest +import structlog + +from agentic_cli.cli.app import BaseCLIApp + +EXCLUDED_EVENT = "user_scoped_setting_excluded_from_ui" + + +def _make_app(settings, keys: list[str]) -> BaseCLIApp: + class _App(BaseCLIApp): + def get_ui_setting_keys(self) -> list[str]: + return keys + + app = _App.__new__(_App) + app._settings = settings + return app + + +@pytest.fixture +def isolated_settings(tmp_path, monkeypatch): + """Hermetic BaseSettings: temp cwd/HOME so no real config files load.""" + home = tmp_path / "home" + home.mkdir() + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(home)) + from agentic_cli.config import BaseSettings + + # A key so get_available_models() is non-empty and "model" can render + return BaseSettings(google_api_key="test-key") + + +def _item_keys(items) -> list[str]: + return [item.key for item in items] + + +class TestSettingsUiScopeGuard: + def test_user_scoped_key_is_excluded_with_warning(self, isolated_settings): + """A non-allowlisted field never renders; the exclusion is logged.""" + app = _make_app(isolated_settings, ["thinking_effort", "raw_llm_logging"]) + + with structlog.testing.capture_logs() as logs: + items = app._build_ui_items() + + assert _item_keys(items) == ["thinking_effort"] + assert any( + e.get("event") == EXCLUDED_EVENT and e.get("key") == "raw_llm_logging" + for e in logs + ) + + def test_allowlisted_keys_render_without_warning(self, isolated_settings): + app = _make_app(isolated_settings, ["thinking_effort", "verbose_thinking"]) + + with structlog.testing.capture_logs() as logs: + items = app._build_ui_items() + + assert set(_item_keys(items)) == {"thinking_effort", "verbose_thinking"} + assert not any(e.get("event") == EXCLUDED_EVENT for e in logs) + + def test_model_synthetic_key_passes_guard(self, isolated_settings): + """"model" is not a field but writes default_model (allowlisted).""" + app = _make_app(isolated_settings, ["model"]) + + with structlog.testing.capture_logs() as logs: + items = app._build_ui_items() + + assert _item_keys(items) == ["model"] + assert not any(e.get("event") == EXCLUDED_EVENT for e in logs) + + def test_dangling_nonexistent_key_is_excluded_with_warning( + self, isolated_settings + ): + """A key for a removed field (e.g. airesearcher's log_activity) now + warns instead of being silently skipped.""" + app = _make_app(isolated_settings, ["log_activity"]) + + with structlog.testing.capture_logs() as logs: + items = app._build_ui_items() + + assert items == [] + assert any( + e.get("event") == EXCLUDED_EVENT and e.get("key") == "log_activity" + for e in logs + ) 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/demo_isolation.py b/tests/demo_isolation.py new file mode 100644 index 0000000..cb52556 --- /dev/null +++ b/tests/demo_isolation.py @@ -0,0 +1,158 @@ +"""Construct fully-isolated ``ResearchDemoSettings`` for tests. + +Plumbing shared by the headless application tests and the live scenario tests; +it defines no test classes, so pytest does not collect it. + +``ResearchDemoSettings`` reads from three developer-owned places, and they are +*not* isolated by the same mechanism: + +=========================== ========================== ==================== +source resolved isolated by +=========================== ========================== ==================== +``./.research_demo/…json`` ``Path.cwd()`` per call ``monkeypatch.chdir`` +``~/.research_demo/…json`` ``Path.home()`` per call ``monkeypatch.setenv`` +``~/.research_demo/.env`` ``Path.home()`` **at an explicit + class-definition time** ``_env_file=`` +=========================== ========================== ==================== + +The third is the trap: ``model_config["env_file"]`` is evaluated when the +module is imported — which, under pytest, is during collection, long before any +fixture runs. Redirecting ``HOME`` afterwards does nothing to it, so a test that +only patched ``HOME`` was still reading the developer's real dotenv (and any +API key in it). Passing ``_env_file`` explicitly is the only thing that moves +it, so every fixture here does. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Any + +# The real home, captured at import (collection time) before any fixture has +# had a chance to redirect HOME. This is what isolation is asserted *against*. +REAL_HOME = Path(os.path.expanduser("~")).resolve() + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from examples.research_demo.settings import ResearchDemoSettings # noqa: E402 + + +def isolated_env_file(home: Path) -> Path: + """Where an isolated run's dotenv lives (it need not exist).""" + return home / ".research_demo" / ".env" + + +def make_isolated_settings( + *, home: Path, workspace: Path, **kwargs: Any +) -> ResearchDemoSettings: + """``ResearchDemoSettings`` that cannot read the developer's configuration. + + The caller is still responsible for redirecting ``HOME`` and ``cwd`` + (``monkeypatch``) — those cover the two JSON sources. This adds the third: + an explicit ``_env_file`` under the temp home, overriding the class's + import-time default. + """ + kwargs.setdefault("_env_file", str(isolated_env_file(home))) + return ResearchDemoSettings(workspace_dir=workspace, **kwargs) + + +def effective_env_file(**kwargs: Any) -> Path | None: + """The dotenv path ``ResearchDemoSettings`` would *actually* read. + + Captured from the ``dotenv_settings`` source pydantic-settings hands to + ``settings_customise_sources`` — the real value in effect, not the class + default and not what the caller hoped it passed. + + Returns: + The resolved path, or None when the dotenv source is disabled. + """ + captured: list[Any] = [] + + class _Probe(ResearchDemoSettings): + @classmethod + def settings_customise_sources( # type: ignore[override] + cls, + settings_cls, + init_settings, + env_settings, + dotenv_settings, + file_secret_settings, + ): + captured.append(getattr(dotenv_settings, "env_file", None)) + return super().settings_customise_sources( + settings_cls, + init_settings, + env_settings, + dotenv_settings, + file_secret_settings, + ) + + _Probe(**kwargs) + assert captured, "settings_customise_sources was never called" + value = captured[0] + if value is None: + return None + if isinstance(value, (list, tuple)): + assert len(value) == 1, f"expected a single env_file, got {value}" + value = value[0] + return Path(value).resolve() + + +def assert_dotenv_isolated(env_file: Path | None, home: Path) -> None: + """The effective dotenv must be inside the temp home, or disabled. + + Asserting "not under the real home" alone would pass for a path that is + simply somewhere else on disk, so this pins the positive form too. + """ + if env_file is None: + return # dotenv disabled outright — isolated by construction + home = home.resolve() + assert env_file.is_relative_to(home), ( + f"the effective dotenv is {env_file}, which is not under the test's " + f"temporary home {home}" + ) + assert not env_file.is_relative_to(REAL_HOME), ( + f"the effective dotenv resolves under the developer's real home: {env_file}" + ) + + +def suppress_global_logging_config(monkeypatch) -> None: + """Stop app construction from reconfiguring logging process-wide. + + ``BaseCLIApp.__init__`` calls ``configure_logging()``, which reconfigures + structlog globally with ``cache_logger_on_first_use=True``. That cannot be + undone afterwards: by the time the constructor returns, the module-level + loggers have already bound and *cached* a logger, so restoring the previous + configuration leaves those caches in place and every later test using + ``structlog.testing.capture_logs()`` silently captures nothing. + + So it is prevented rather than reverted. These tests assert on command + routing and rendering, never on log output, so the application's logging + configuration is not part of what they cover. + """ + monkeypatch.setattr( + "agentic_cli.cli.app.configure_logging", lambda *a, **k: None + ) + + +def assert_no_global_settings_leak() -> None: + """No test may leave a global settings singleton behind. + + ``set_settings()`` writes a process-wide singleton with no teardown, so one + test's settings would silently answer ``get_settings()`` for every later + test that runs outside a manager's ``SettingsContext``. + """ + from agentic_cli import config as _config + + assert _config._settings_instance is None, ( + "a test left a global settings singleton behind " + f"({_config._settings_instance!r}); use the manager's SettingsContext " + "instead of set_settings()" + ) + assert _config._settings_context.get() is None, ( + "a test left settings in the context variable" + ) diff --git a/tests/event_replay.py b/tests/event_replay.py index 5cadf7c..b92a560 100644 --- a/tests/event_replay.py +++ b/tests/event_replay.py @@ -131,6 +131,7 @@ class ReplayWorkflow: def __init__(self, events: list[WorkflowEvent]) -> None: self._events = list(events) self.input_callback = None + self.received_session_id: str | None = None def set_input_callback(self, callback: Any) -> None: self.input_callback = callback @@ -139,6 +140,7 @@ def clear_input_callback(self) -> None: self.input_callback = None async def process(self, message: str, user_id: str, session_id: str | None = None): + self.received_session_id = session_id for event in self._events: yield event diff --git a/tests/examples/__init__.py b/tests/examples/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/examples/console_smoke.py b/tests/examples/console_smoke.py new file mode 100644 index 0000000..233c380 --- /dev/null +++ b/tests/examples/console_smoke.py @@ -0,0 +1,274 @@ +"""Shared driver for the research-demo console smoke tests. + +Plumbing only — no test classes, so pytest does not collect it. Used by both +``test_research_demo_pty.py`` (source checkout, ``python -m research_demo``) +and ``test_research_demo_wheel.py`` (built wheel, the installed +``research-demo`` console script), so the two exercise *the same* interaction +against different installations and different entry points. + +The demo is a full-screen prompt_toolkit application, so it needs a real +terminal: it is driven over a pty with ``pexpect`` and its output is +ANSI-stripped before matching. Matching is on short, stable fragments rather +than a snapshot of raw terminal bytes — the screen is redrawn repeatedly and +line-wrapped to the terminal width, so raw output is not stable. +""" + +from __future__ import annotations + +import io +import json +import os +import re +import subprocess +import sys +import tomllib +from dataclasses import dataclass, field +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def project_version() -> str: + """The version declared in the repository's ``pyproject.toml``. + + Derived, never hardcoded: a release-specific literal would pin the + acceptance suite to one release, going red on every branch that predates a + version bump and needing an edit at each bump. Reading the declared version + instead means these assertions validate *whatever* is being released. + """ + with (_REPO_ROOT / "pyproject.toml").open("rb") as fh: + return tomllib.load(fh)["project"]["version"] + +# CSI/OSC escape sequences, charset selects, and bare CRs, so matching sees +# the text a human reads rather than the bytes that drew it. +_ANSI = re.compile( + r"\x1b\[[0-9;?]*[a-zA-Z]" # CSI + r"|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)" # OSC + r"|\x1b[()][A-Za-z0-9]" # charset select + r"|\x1b[=>]" # keypad mode + r"|\r" +) + +#: The *only* parent variables allowed through to the child. Everything else — +#: provider credentials, cloud authentication, ``RESEARCH_DEMO_*``/``AGENTIC_*`` +#: settings overrides, custom config paths, ``PYTHONPATH`` — is dropped, so the +#: smoke measures the shipped defaults rather than whatever the developer (or +#: CI) happens to export. These are what it takes to *launch* an interpreter and +#: a terminal, nothing about the application. +ALLOWED_PARENT_VARS = ( + "PATH", # find the interpreter/console script and anything it execs + "LANG", # terminal text encoding + "LC_ALL", + "LC_CTYPE", + "TMPDIR", # POSIX temp location (macOS gives every user its own) + "SYSTEMROOT", # Windows needs these to start a process at all + "SystemRoot", + "COMSPEC", +) + +#: Set explicitly on every child, never inherited. +EXPLICIT_CHILD_VARS = ( + "HOME", + "TERM", + "NO_COLOR", + "PYTHONUNBUFFERED", + "PYTHONDONTWRITEBYTECODE", +) + +#: Terminal size. Wide enough that the help table's cells do not wrap, which is +#: what makes short-fragment matching reliable. +TERM_SIZE = (40, 200) + +#: Every wait is bounded. Startup measured ~1s on a warm checkout; 60s is +#: headroom for a cold import on CI, not an invitation to hang. +STARTUP_TIMEOUT = 60 +STEP_TIMEOUT = 30 + + +@dataclass +class SmokeResult: + """What the console session did, for the caller to assert on.""" + + transcript: str + exit_status: int | None + signal_status: int | None + argv: list[str] + steps: list[str] = field(default_factory=list) + + def saw(self, fragment: str) -> bool: + return fragment in self.transcript + + +@dataclass +class ChildImports: + """Where a child process resolved the packages under test.""" + + demo: Path + framework: Path + version: str + executable: Path + + +def child_env(home: Path) -> dict[str, str]: + """A minimal, allowlisted environment for the child process. + + Built up from nothing rather than filtered down from ``os.environ``: a + deny-list only removes the hazards someone thought of, and the ones that + matter here (``RESEARCH_DEMO_*``, ``AGENTIC_*``, ``GOOGLE_APPLICATION_ + CREDENTIALS``, ``VERTEX_*``) are open-ended. + + ``HOME`` is redirected so the demo's user config, user KB and + ``~/.research_demo/.env`` all resolve inside the temp dir; ``PYTHONPATH`` is + absent so the child cannot pick up a checkout the caller did not intend. + """ + env = {k: os.environ[k] for k in ALLOWED_PARENT_VARS if k in os.environ} + env.update( + HOME=str(home), + TERM="xterm-256color", + NO_COLOR="1", + PYTHONUNBUFFERED="1", + PYTHONDONTWRITEBYTECODE="1", + ) + return env + + +def probe_child_imports(python: str, cwd: Path, home: Path) -> ChildImports: + """Ask a child process — same interpreter, cwd and env — what it imports. + + The parent pytest process is *not* evidence: it has the editable install on + ``sys.path`` and whatever the developer's environment supplies. Only the + child can say where the child resolves ``research_demo``. + """ + code = ( + "import json, sys, research_demo, agentic_cli; " + "print(json.dumps({" + "'demo': research_demo.__file__, " + "'framework': agentic_cli.__file__, " + "'version': agentic_cli.__version__, " + "'executable': sys.executable}))" + ) + result = subprocess.run( + [python, "-c", code], + cwd=str(cwd), + env=child_env(home), + capture_output=True, + text=True, + timeout=STEP_TIMEOUT, + ) + if result.returncode != 0: + raise AssertionError( + f"import probe failed ({result.returncode}):\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + payload = json.loads(result.stdout.strip().splitlines()[-1]) + return ChildImports( + demo=Path(payload["demo"]).resolve(), + framework=Path(payload["framework"]).resolve(), + version=payload["version"], + executable=Path(payload["executable"]).resolve(), + ) + + +def run_console_smoke(argv: list[str], cwd: Path, home: Path) -> SmokeResult: + """Drive a console command through startup → /help → /exit. + + Every wait is an ``expect`` on observable output, never a sleep, so the + test is deterministic rather than timing-dependent. + + Args: + argv: The command to launch — ``[python, "-m", "research_demo"]`` for + the source checkout, or ``[".../bin/research-demo"]`` for the + installed console script. + cwd: Working directory for the child — keeps ``./.research_demo/`` + (project KB, permission workdir) out of the repo. + home: Value for ``HOME``. + + Returns: + The full ANSI-stripped transcript plus the process's exit status. + """ + import pexpect + + log = io.StringIO() + child = pexpect.spawn( + argv[0], + list(argv[1:]), + cwd=str(cwd), + env=child_env(home), + encoding="utf-8", + codec_errors="replace", + timeout=STARTUP_TIMEOUT, + dimensions=TERM_SIZE, + ) + child.logfile_read = log + steps: list[str] = [] + try: + # 1. Startup: the welcome panel is drawn and the prompt appears. + child.expect(r">>>", timeout=STARTUP_TIMEOUT) + steps.append("prompt") + + # 2. /help renders the command table. + child.sendline("/help") + child.expect(r"Exit the application", timeout=STEP_TIMEOUT) + steps.append("help") + + # 3. The prompt comes back, so the app is still interactive. + child.expect(r">>>", timeout=STEP_TIMEOUT) + steps.append("prompt-after-help") + + # 4. /exit shuts down cleanly. + child.sendline("/exit") + child.expect(r"Goodbye!", timeout=STEP_TIMEOUT) + steps.append("goodbye") + child.expect(pexpect.EOF, timeout=STEP_TIMEOUT) + steps.append("eof") + finally: + child.close() + + return SmokeResult( + transcript=_ANSI.sub("", log.getvalue()), + exit_status=child.exitstatus, + signal_status=child.signalstatus, + argv=list(argv), + steps=steps, + ) + + +def assert_clean_session(result: SmokeResult) -> None: + """Shared assertions: the demo started, responded, and exited cleanly.""" + assert result.steps == [ + "prompt", + "help", + "prompt-after-help", + "goodbye", + "eof", + ], f"console session did not complete: reached {result.steps} via {result.argv}" + + assert result.saw("Exit the application"), "the /help table was not rendered" + assert result.saw("Goodbye!"), "the app did not say goodbye on /exit" + + for marker in ("Traceback (most recent call last)", "Unhandled exception"): + assert marker not in result.transcript, ( + f"{marker!r} in console output:\n{_tail(result.transcript)}" + ) + + assert result.signal_status is None, ( + f"the console process died on signal {result.signal_status}" + ) + assert result.exit_status == 0, ( + f"console exited with status {result.exit_status}:\n{_tail(result.transcript)}" + ) + + +def _tail(text: str, lines: int = 40) -> str: + return "\n".join(text.splitlines()[-lines:]) + + +def platform_skip_reason() -> str | None: + """Why this platform cannot run a pty smoke, or None if it can.""" + if sys.platform.startswith("win"): + return ( + "the console smoke needs a POSIX pty: Python's `pty` module and " + "`pexpect.spawn` are POSIX-only, and prompt_toolkit's full-screen " + "app cannot be driven through a pipe" + ) + return None diff --git a/tests/examples/test_report_writer_skill.py b/tests/examples/test_report_writer_skill.py new file mode 100644 index 0000000..359f44e --- /dev/null +++ b/tests/examples/test_report_writer_skill.py @@ -0,0 +1,33 @@ +"""Resolution + asset tests for the demo report-writer skill (no LLM).""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.tools.skills import SkillStore, make_skill_toolset + +_SKILLS_DIR = Path(__file__).resolve().parents[2] / "examples" / "research_demo" / "skills" + + +def test_skill_resolves(): + skills = SkillStore([str(_SKILLS_DIR)]).resolve(["report-writer"]) + assert len(skills) == 1 + assert skills[0].name == "report-writer" + assert skills[0].description # non-empty when-to-use + + +def test_template_asset_present(): + tpl = _SKILLS_DIR / "report-writer" / "assets" / "report_template.tex" + assert tpl.is_file() + body = tpl.read_text() + assert "\\documentclass" in body and "\\includegraphics" in body + + +def test_toolset_excludes_scripts(): + skills = SkillStore([str(_SKILLS_DIR)]).resolve(["report-writer"]) + names = {t.name for t in make_skill_toolset(skills)._tools} + assert "run_skill_script" not in names + assert {"list_skills", "load_skill", "load_skill_resource"} <= names diff --git a/tests/examples/test_research_demo_agents.py b/tests/examples/test_research_demo_agents.py new file mode 100644 index 0000000..abaa677 --- /dev/null +++ b/tests/examples/test_research_demo_agents.py @@ -0,0 +1,37 @@ +"""Config-level tests for the research_demo example agents. + +Verifies that AGENT_CONFIGS is wired correctly without making any LLM calls. +""" + + +def test_data_analyst_wired(): + from research_demo.agents import AGENT_CONFIGS + + names = {a.name for a in AGENT_CONFIGS} + assert "data_analyst" in names + analyst = next(a for a in AGENT_CONFIGS if a.name == "data_analyst") + tool_names = {t.__name__ for t in analyst.tools} + assert "sandbox_execute" in tool_names + coord = next(a for a in AGENT_CONFIGS if a.name == "research_coordinator") + assert "data_analyst" in coord.sub_agents + + +def test_report_writer_wired(): + from research_demo.agents import AGENT_CONFIGS + + names = {a.name for a in AGENT_CONFIGS} + assert "report_writer" in names + rw = next(a for a in AGENT_CONFIGS if a.name == "report_writer") + tool_names = {t.__name__ for t in rw.tools} + assert "compile_document" in tool_names + assert "sandbox_execute" not in tool_names # no arbitrary code exec + assert rw.skills == ["report-writer"] + coord = next(a for a in AGENT_CONFIGS if a.name == "research_coordinator") + assert "report_writer" in coord.sub_agents + + +def test_report_writer_skills_dir_configured(): + from research_demo.settings import ResearchDemoSettings + + s = ResearchDemoSettings() + assert any(str(p).rstrip("/").endswith("skills") for p in s.skills_dirs) diff --git a/tests/examples/test_research_demo_app.py b/tests/examples/test_research_demo_app.py new file mode 100644 index 0000000..5124432 --- /dev/null +++ b/tests/examples/test_research_demo_app.py @@ -0,0 +1,313 @@ +"""Headless acceptance tests for the *composed* research-demo application. + +The workflow, the renderer and the individual commands all have component +tests. What none of them covers is the wiring: a real ``ResearchDemoApp`` — its +own command registry, ``BaseCLIApp.process_input``, the real +``MessageProcessor`` — reacting to real input. + +So these build the actual app and drive ``process_input()``. The only things +substituted are the two nondeterministic ones: + +- the **workflow manager**, replaced by a scripted ``WorkflowEvent`` stream + (no LLM, no network) — ``WorkflowEvent`` is the UI-independent boundary the + framework already defines, so scripting it is not a new seam; +- the **UI session**, replaced by ``RecordingSession`` (the same stand-in + ``test_message_processor_render.py`` uses) so assertions are on semantic + render calls rather than terminal bytes. + +Assertions are deliberately about *shape* — which kind of call happened, in +what order, was it a warning or an error — never about generated prose. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from agentic_cli.workflow.events import WorkflowEvent +from tests.demo_isolation import ( + assert_dotenv_isolated, + assert_no_global_settings_leak, + effective_env_file, + isolated_env_file, + make_isolated_settings, + suppress_global_logging_config, +) +from tests.event_replay import RecordingSession, ReplayController + +# Make the `examples` namespace package importable (no __init__.py). +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from examples.research_demo.app import ResearchDemoApp # noqa: E402 + + +class ScriptedWorkflow: + """Workflow-manager stand-in driving one scripted turn per call. + + Extends what ``ReplayWorkflow`` offers with the two things these tests + need: several distinct turns in sequence, and a turn that fails — so + "the app is still usable afterwards" can actually be asserted. + """ + + def __init__(self, turns: list[list[WorkflowEvent] | Exception]) -> None: + self._turns = list(turns) + self.messages: list[str] = [] + self.session_ids: list[str | None] = [] + self.input_callback = None + + def set_input_callback(self, callback) -> None: # noqa: ANN001 + self.input_callback = callback + + def clear_input_callback(self) -> None: + self.input_callback = None + + async def process(self, message: str, user_id: str, session_id: str | None = None): + self.messages.append(message) + self.session_ids.append(session_id) + turn = self._turns.pop(0) if self._turns else [] + if isinstance(turn, Exception): + raise turn + for event in turn: + yield event + + +@pytest.fixture +def demo_app(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ResearchDemoApp: + """A real ResearchDemoApp on an isolated HOME/cwd, with a recording UI. + + HOME and cwd are both redirected: the demo derives its user config, user + KB and ``.env`` from ``~/.research_demo`` and its project KB and permission + workdir from ``./.research_demo``, and neither may touch the developer's + real config or the repo. + + The dotenv needs the third, explicit step (``make_isolated_settings``): + ``model_config["env_file"]`` was frozen to the developer's real + ``~/.research_demo/.env`` when this module was imported, so redirecting + HOME cannot move it. + + The workflow controller is left as the **real** one, uninitialized — that + is what makes ``app.workflow`` raise, which is exactly the state the + readiness tests are about. Tests that need a turn install a scripted + controller themselves. + """ + home = tmp_path / "home" + project = tmp_path / "project" + workspace = tmp_path / "ws" + for d in (home, project, workspace): + d.mkdir(parents=True) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.chdir(project) + + settings = make_isolated_settings( + home=home, + workspace=workspace, + permissions_enabled=False, + ) + # Building a real app would otherwise reconfigure structlog globally and + # break every later test that captures logs (see the helper's docstring). + suppress_global_logging_config(monkeypatch) + + app = ResearchDemoApp(settings=settings) + app.session = RecordingSession() + app._test_home = home # for the isolation assertions below + return app + + +def _script(app: ResearchDemoApp, *turns) -> ScriptedWorkflow: + """Install a scripted workflow behind the app's controller seam.""" + workflow = ScriptedWorkflow(list(turns)) + app._workflow_controller = ReplayController(workflow) + return workflow + + +class TestFixtureIsolation: + """The fixture must not be able to read the developer's configuration. + + Worth its own tests because the failure is silent: a headless test that + quietly picked up a real API key from ``~/.research_demo/.env`` would still + pass, while proving something different from what it claims. + """ + + def test_effective_dotenv_is_under_the_temporary_home(self, demo_app) -> None: + home = demo_app._test_home + env_file = effective_env_file( + workspace_dir=home / "ws", + _env_file=str(isolated_env_file(home)), + ) + assert_dotenv_isolated(env_file, home) + + def test_the_class_default_would_not_have_been_isolated(self) -> None: + """Pin the hazard this fixture works around. + + If the class ever stops resolving its dotenv at import time, this test + fails and the ``_env_file`` plumbing can be simplified away. + """ + from tests.demo_isolation import REAL_HOME, ResearchDemoSettings + + class_default = Path(str(ResearchDemoSettings.model_config["env_file"])) + assert class_default.is_relative_to(REAL_HOME), ( + "ResearchDemoSettings no longer freezes its env_file under the real " + f"home ({class_default}) — re-check whether _env_file is still needed" + ) + + def test_settings_do_not_leak_globally(self, demo_app) -> None: + """Building the app must not install a process-wide settings singleton.""" + assert_no_global_settings_leak() + + +class TestCommandsBeforeReadiness: + """Demo commands must degrade to a warning while the workflow comes up. + + ``app.workflow`` raises until the controller reports READY. A command that + touched it during background init was caught by + ``BaseCLIApp._handle_command`` and surfaced as the generic + "Error executing command: Workflow not initialized yet" — which reads like + a bug rather than "not yet". + """ + + async def test_memory_warns_and_does_not_error(self, demo_app): + with pytest.raises(RuntimeError): + _ = demo_app.workflow # precondition: genuinely not ready + + await demo_app.process_input("/memory") + + assert demo_app.session.errors() == [], ( + f"/memory errored before readiness: {demo_app.session.errors()}" + ) + assert len(demo_app.session.warnings()) == 1 + assert "initializing" in demo_app.session.warnings()[0].lower() + + async def test_kb_backfill_warns_and_does_not_error(self, demo_app): + await demo_app.process_input("/kb-backfill") + + assert demo_app.session.errors() == [], ( + f"/kb-backfill errored before readiness: {demo_app.session.errors()}" + ) + assert len(demo_app.session.warnings()) == 1 + assert "initializing" in demo_app.session.warnings()[0].lower() + + async def test_the_command_is_still_registered_and_echoed(self, demo_app): + """A warning must come from the command, not from it being unknown.""" + await demo_app.process_input("/memory") + + echoed = [c for c in demo_app.session.calls if c[0] == "message"] + assert ("message", "user", "/memory") in echoed + assert not any("Unknown command" in e for e in demo_app.session.errors()) + + +class TestMessageRouting: + """A plain message goes through MessageProcessor and renders its events.""" + + async def test_scripted_turn_is_rendered_in_order(self, demo_app): + workflow = _script( + demo_app, + [ + WorkflowEvent.thinking("Considering the question."), + WorkflowEvent.tool_call("kb_search", {"query": "decoding"}), + WorkflowEvent.tool_result( + "kb_search", {"success": True}, success=True, duration_ms=5 + ), + WorkflowEvent.text("Here is what I found."), + ], + ) + + await demo_app.process_input("what do you know about decoding?") + + assert workflow.messages == ["what do you know about decoding?"] + + kinds = demo_app.session.kinds() + # The user's message is echoed before anything is rendered for the turn. + assert kinds[0] == "message" + assert demo_app.session.calls[0][1] == "user" + # The scripted text reached the UI as a response, after the echo. + assert "Here is what I found." in demo_app.session.responses() + assert kinds.index("response") > 0 + # A tool call opened the events box before the response was written. + assert "start_thinking" in kinds + assert kinds.index("start_thinking") < kinds.index("response") + assert demo_app.session.errors() == [] + + async def test_turn_receives_the_application_session_id(self, demo_app): + workflow = _script(demo_app, [WorkflowEvent.text("ok")]) + + await demo_app.process_input("hello") + + assert workflow.session_ids == [demo_app.session_id] + assert demo_app.session_id, "the app must carry a durable session id" + + async def test_blank_input_is_ignored(self, demo_app): + workflow = _script(demo_app, [WorkflowEvent.text("should not run")]) + + await demo_app.process_input(" ") + + assert workflow.messages == [] + assert demo_app.session.calls == [] + + +class TestUnknownCommand: + async def test_unknown_command_errors_with_a_hint(self, demo_app): + await demo_app.process_input("/definitely-not-a-command") + + errors = demo_app.session.errors() + assert len(errors) == 1 + assert "definitely-not-a-command" in errors[0] + # And the user is pointed somewhere useful. + assert any( + "/help" in c[2] for c in demo_app.session.calls if c[0] == "message" + ) + + async def test_unknown_command_does_not_reach_the_workflow(self, demo_app): + workflow = _script(demo_app, [WorkflowEvent.text("should not run")]) + + await demo_app.process_input("/nope") + + assert workflow.messages == [] + + +class TestRecoveryAfterFailure: + """A failed turn must not wedge the app: the next turn still works.""" + + async def test_failed_turn_is_reported_then_the_next_turn_succeeds(self, demo_app): + workflow = _script( + demo_app, + RuntimeError("backend exploded"), + [WorkflowEvent.text("second turn is fine")], + ) + + await demo_app.process_input("first") + assert demo_app.session.errors(), "a failing turn reported nothing" + first_errors = len(demo_app.session.errors()) + + await demo_app.process_input("second") + + assert workflow.messages == ["first", "second"] + assert "second turn is fine" in demo_app.session.responses() + assert len(demo_app.session.errors()) == first_errors, ( + "the recovery turn produced a new error" + ) + + async def test_a_command_error_does_not_wedge_the_next_turn(self, demo_app): + """An exception inside a command is contained by _handle_command.""" + workflow = _script(demo_app, [WorkflowEvent.text("still working")]) + + boom = demo_app.command_registry.get("memory") + + async def _raise(args, app): # noqa: ANN001 + raise RuntimeError("command exploded") + + monkey = boom.execute + boom.execute = _raise + try: + await demo_app.process_input("/memory") + finally: + boom.execute = monkey + + assert any("command exploded" in e for e in demo_app.session.errors()) + + await demo_app.process_input("carry on") + assert workflow.messages == ["carry on"] + assert "still working" in demo_app.session.responses() diff --git a/tests/examples/test_research_demo_pty.py b/tests/examples/test_research_demo_pty.py new file mode 100644 index 0000000..0eeee06 --- /dev/null +++ b/tests/examples/test_research_demo_pty.py @@ -0,0 +1,219 @@ +"""Real-terminal smoke test for the research-demo console process. + +Everything else in the suite drives Python objects — ``process_input()``, +``MessageProcessor``, the workflow manager. None of it proves the thing a user +actually runs starts up, draws a prompt, answers a command and exits. + +This launches the **actual console process** over a pty and interacts with it: +startup → prompt → ``/help`` → prompt → ``/exit``. It needs no API key, no +network, no Docker and no LLM — the workflow's background initialization is +allowed to fail (there are no credentials), which is itself part of what is +being asserted: a demo with no keys must still come up, answer ``/help`` and +exit 0 rather than crash. + +This module covers the **source checkout** (via ``python -m research_demo`` on +the editable install). ``test_research_demo_wheel.py`` runs the identical +session against the installed console script from a built wheel. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from tests.examples.console_smoke import ( + ALLOWED_PARENT_VARS, + EXPLICIT_CHILD_VARS, + assert_clean_session, + child_env, + platform_skip_reason, + probe_child_imports, + project_version, + run_console_smoke, +) + +_SKIP = platform_skip_reason() +pytestmark = pytest.mark.skipif(_SKIP is not None, reason=_SKIP or "") + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +@pytest.fixture +def sandbox(tmp_path: Path) -> tuple[Path, Path]: + """An isolated (home, cwd) pair for one console run.""" + home = tmp_path / "home" + cwd = tmp_path / "cwd" + home.mkdir() + cwd.mkdir() + return home, cwd + + +def _argv() -> list[str]: + """How the source checkout is launched.""" + return [sys.executable, "-m", "research_demo"] + + +class TestChildImportOrigin: + """Pin *which* installation the spawned child resolves. + + The parent pytest process is not evidence: it has the editable install on + ``sys.path`` already. Only a child launched exactly the way the smoke + launches one — same interpreter, cwd, HOME and stripped environment — can + say where the console under test comes from. Without this, a wheel-only + failure would look like a checkout failure, and vice versa. + """ + + def test_child_resolves_the_checkout(self, sandbox) -> None: + home, cwd = sandbox + + imports = probe_child_imports(sys.executable, cwd=cwd, home=home) + + assert imports.demo.is_relative_to(_REPO_ROOT), ( + f"the child resolved research_demo to {imports.demo}, " + "which is not the checkout this module is meant to cover" + ) + assert imports.framework.is_relative_to(_REPO_ROOT), ( + f"the child resolved agentic_cli to {imports.framework}" + ) + assert imports.version == project_version() + + def test_the_probe_uses_the_same_interpreter_as_the_smoke(self, sandbox) -> None: + home, cwd = sandbox + imports = probe_child_imports(sys.executable, cwd=cwd, home=home) + assert imports.executable == Path(sys.executable).resolve() + assert _argv()[0] == sys.executable + + +class TestChildEnvironmentIsHermetic: + """Nothing from the developer's shell may change what the smoke measures.""" + + def test_only_allowlisted_variables_reach_the_child(self, tmp_path) -> None: + env = child_env(tmp_path) + unexpected = set(env) - set(ALLOWED_PARENT_VARS) - set(EXPLICIT_CHILD_VARS) + assert unexpected == set(), f"unexpected variables in the child env: {unexpected}" + + def test_explicit_variables_are_set(self, tmp_path) -> None: + env = child_env(tmp_path) + assert env["HOME"] == str(tmp_path) + for name in EXPLICIT_CHILD_VARS: + assert name in env, f"{name} was not set on the child" + + @pytest.mark.parametrize( + "name, value", + [ + # Provider credentials + ("GOOGLE_API_KEY", "sk-should-not-leak"), + ("ANTHROPIC_API_KEY", "sk-should-not-leak"), + ("TAVILY_API_KEY", "sk-should-not-leak"), + # Cloud authentication + ("GOOGLE_APPLICATION_CREDENTIALS", "/tmp/creds.json"), + ("GOOGLE_CLOUD_PROJECT", "some-project"), + ("GOOGLE_GENAI_USE_VERTEXAI", "true"), + ("VERTEX_LOCATION", "us-central1"), + # Application settings overrides + ("RESEARCH_DEMO_DEFAULT_MODEL", "some-other-model"), + ("RESEARCH_DEMO_WORKSPACE_DIR", "/tmp/elsewhere"), + ("RESEARCH_DEMO_PERMISSIONS_ENABLED", "false"), + ("AGENTIC_ORCHESTRATOR", "langgraph"), + ("AGENTIC_CLI_LOG_LEVEL", "debug"), + # Import path + ("PYTHONPATH", "/somewhere/else"), + ("PYTHONSTARTUP", "/tmp/startup.py"), + ], + ) + def test_hazardous_variables_never_reach_the_child( + self, tmp_path, monkeypatch, name, value + ) -> None: + """Set it in the parent; it must still be absent from the child env.""" + monkeypatch.setenv(name, value) + + env = child_env(tmp_path) + + assert name not in env, ( + f"{name} leaked into the console child and could change the smoke" + ) + + def test_a_leaked_setting_would_actually_change_the_demo( + self, tmp_path, monkeypatch + ) -> None: + """The pin above is only worth having if such a variable would matter. + + Proves ``RESEARCH_DEMO_DEFAULT_MODEL`` really is load-bearing, so the + allowlist is protecting against something real rather than asserting a + vacuous property. + """ + # `tests.demo_isolation` already puts the repo root on sys.path at + # import; scoped here too so this test does not depend on that and does + # not leave an entry behind if it ever becomes the first importer. + monkeypatch.syspath_prepend(str(_REPO_ROOT)) + from examples.research_demo.settings import ResearchDemoSettings + + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("RESEARCH_DEMO_DEFAULT_MODEL", "sentinel-model") + + settings = ResearchDemoSettings( + workspace_dir=tmp_path / "ws", _env_file=None + ) + assert settings.default_model == "sentinel-model" + + +class TestConsoleSmoke: + """The shipped console entry point, driven as a user would drive it.""" + + def test_startup_help_and_exit(self, sandbox) -> None: + home, cwd = sandbox + + result = run_console_smoke(_argv(), cwd=cwd, home=home) + + assert_clean_session(result) + + def test_a_help_and_exit_session_is_side_effect_free(self, sandbox) -> None: + """Starting up, asking for /help and exiting must write nothing at all. + + The demo resolves its project KB and permission workdir from ``cwd`` + and its user config, user KB and session store from ``HOME``. A session + that never runs a turn should touch none of them — so this asserts both + that the repository is untouched (a leaked path would litter the + checkout) *and* that the temp HOME/cwd are still empty, which is the + stronger statement and the one that would catch a stray write landing + somewhere the repo check does not look. + """ + home, cwd = sandbox + + before = {p.name for p in _REPO_ROOT.iterdir()} + + result = run_console_smoke(_argv(), cwd=cwd, home=home) + assert_clean_session(result) + + assert {p.name for p in _REPO_ROOT.iterdir()} == before, ( + "the console run created entries in the repository root" + ) + assert list(home.rglob("*")) == [], ( + f"the run wrote into HOME: {[str(p) for p in home.rglob('*')][:10]}" + ) + assert list(cwd.rglob("*")) == [], ( + f"the run wrote into cwd: {[str(p) for p in cwd.rglob('*')][:10]}" + ) + + def test_runs_without_any_credentials_in_the_environment( + self, sandbox, monkeypatch + ) -> None: + """Even with keys exported in the parent, the child runs without them. + + ``monkeypatch.setenv`` rather than writing ``os.environ`` directly: a + developer running this with a real ``GOOGLE_API_KEY`` exported would + otherwise have it deleted by the cleanup, because ``os.environ.pop`` + cannot restore a value it never saved. + """ + home, cwd = sandbox + monkeypatch.setenv("GOOGLE_API_KEY", "sk-not-a-real-key") + + env = child_env(home) + assert "GOOGLE_API_KEY" not in env + + result = run_console_smoke(_argv(), cwd=cwd, home=home) + + assert_clean_session(result) diff --git a/tests/examples/test_research_demo_wheel.py b/tests/examples/test_research_demo_wheel.py new file mode 100644 index 0000000..cb40484 --- /dev/null +++ b/tests/examples/test_research_demo_wheel.py @@ -0,0 +1,270 @@ +"""Acceptance smoke against the **built wheel**, not the checkout. + +Every other test in the suite imports the source tree. That cannot catch a +packaging defect: a module or data file left out of +``[tool.hatch.build.targets.wheel]`` is invisible until someone installs the +artifact. This builds the wheel, installs it into an isolated virtualenv, and +runs the same console session ``test_research_demo_pty.py`` runs — through the +**installed ``research-demo`` console script**, which is what a user actually +types, not ``python -m``. + +Building a wheel and installing its dependencies takes minutes and needs the +network, which is exactly what the offline suite is not. So this is opt-in on +two axes: it carries the ``wheel`` marker *and* requires +``AGENTIC_WHEEL_ACCEPTANCE=1``. That keeps the offline selector +(``-m 'not llm and not docker'``) fast and network-free without redefining it, +and CI runs the acceptance as its own step: + + AGENTIC_WHEEL_ACCEPTANCE=1 conda run -n agenticcli python -m pytest \ + tests/examples/test_research_demo_wheel.py -m wheel -v +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import venv +from pathlib import Path + +import pytest + +from tests.examples.console_smoke import ( + assert_clean_session, + platform_skip_reason, + probe_child_imports, + project_version, + run_console_smoke, +) + +_PLATFORM_SKIP = platform_skip_reason() +_OPT_IN = os.environ.get("AGENTIC_WHEEL_ACCEPTANCE") == "1" + +pytestmark = [ + pytest.mark.wheel, + pytest.mark.skipif(_PLATFORM_SKIP is not None, reason=_PLATFORM_SKIP or ""), + pytest.mark.skipif( + not _OPT_IN, + reason=( + "wheel acceptance builds a wheel and installs its dependencies " + "(slow + needs network); set AGENTIC_WHEEL_ACCEPTANCE=1 to run it" + ), + ), +] + +_REPO_ROOT = Path(__file__).resolve().parents[2] + +#: Non-Python files the demo cannot work without. They live under +#: ``examples/research_demo`` and are only in the wheel because hatchling +#: packages that directory — exactly the kind of thing that silently vanishes. +REQUIRED_PACKAGE_DATA = ( + Path("data") / "benchmarks.csv", + Path("skills") / "report-writer" / "SKILL.md", + Path("skills") / "report-writer" / "assets" / "report_template.tex", +) + + +def _run(cmd: list[str], **kwargs) -> subprocess.CompletedProcess: + result = subprocess.run(cmd, capture_output=True, text=True, **kwargs) + if result.returncode != 0: + raise AssertionError( + f"command failed ({result.returncode}): {' '.join(cmd[:4])}…\n" + f"stdout:\n{result.stdout[-2000:]}\nstderr:\n{result.stderr[-2000:]}" + ) + return result + + +@pytest.fixture(scope="module") +def installed_wheel(tmp_path_factory: pytest.TempPathFactory) -> dict: + """Build the wheel and install it into a throwaway virtualenv. + + Module-scoped: building and installing once is the expensive part, and the + tests below only read from the result. + """ + root = tmp_path_factory.mktemp("wheel-acceptance") + dist = root / "dist" + env_dir = root / "venv" + + _run( + [sys.executable, "-m", "pip", "wheel", "--no-deps", "-w", str(dist), + str(_REPO_ROOT)], + cwd=str(root), # outside the repository + ) + wheels = sorted(dist.glob("agentic_cli-*.whl")) + assert len(wheels) == 1, f"expected exactly one wheel, got {wheels}" + + venv.EnvBuilder(with_pip=True, symlinks=True).create(env_dir) + python = env_dir / "bin" / "python" + console = env_dir / "bin" / "research-demo" + assert python.exists(), f"virtualenv has no interpreter at {python}" + + _run([str(python), "-m", "pip", "install", "--quiet", str(wheels[0])], + cwd=str(root)) + + return { + "python": python, + "console": console, + "wheel": wheels[0], + "root": root, + "env_dir": env_dir, + } + + +@pytest.fixture +def sandbox(tmp_path: Path) -> tuple[Path, Path]: + home = tmp_path / "home" + cwd = tmp_path / "cwd" + home.mkdir() + cwd.mkdir() + return home, cwd + + +class TestWheelIsSelfContained: + def test_child_resolves_from_the_isolated_install(self, installed_wheel, sandbox): + """A child launched as the smoke launches one must not see the checkout. + + Probed in a subprocess with the same interpreter, cwd, HOME and + stripped environment the console gets — the parent pytest process has + the editable install on ``sys.path``, so its own imports prove nothing + about the child. + """ + home, cwd = sandbox + + imports = probe_child_imports( + str(installed_wheel["python"]), cwd=cwd, home=home + ) + + env_dir = installed_wheel["env_dir"].resolve() + for label, path in (("demo", imports.demo), ("framework", imports.framework)): + assert path.is_relative_to(env_dir), ( + f"{label} resolved to {path}, outside the isolated venv" + ) + assert not path.is_relative_to(_REPO_ROOT), ( + f"{label} resolved into the repository checkout: {path}" + ) + assert imports.version == project_version() + + def test_package_data_is_present(self, installed_wheel, sandbox): + home, cwd = sandbox + imports = probe_child_imports( + str(installed_wheel["python"]), cwd=cwd, home=home + ) + pkg_dir = imports.demo.parent + + missing = [str(rel) for rel in REQUIRED_PACKAGE_DATA + if not (pkg_dir / rel).is_file()] + assert missing == [], ( + f"the installed wheel is missing package data: {missing}\n" + f"(looked under {pkg_dir})" + ) + + def test_console_script_is_installed_and_executable(self, installed_wheel): + console = installed_wheel["console"] + assert console.exists(), ( + "the `research-demo` console script was not installed by the wheel" + ) + assert os.access(console, os.X_OK), f"{console} is not executable" + + def test_console_script_runs_the_venv_interpreter(self, installed_wheel): + """Its shebang must point into the venv, not at the checkout's python. + + This is what makes running the script equivalent to running the venv's + interpreter — and therefore what lets the import probe above stand for + the console process. + + The shebang path is compared **unresolved**: ``venv/bin/python`` is a + symlink to the base interpreter (``EnvBuilder(symlinks=True)``), so + resolving it lands on the conda binary and says nothing. What makes an + interpreter a venv interpreter is its ``sys.prefix``, which is asserted + separately below. + """ + env_dir = installed_wheel["env_dir"] + shebang = installed_wheel["console"].read_text(errors="replace").splitlines()[0] + assert shebang.startswith("#!"), f"no shebang in the console script: {shebang!r}" + + interpreter = Path(shebang[2:].strip().split()[0]) + assert interpreter.is_relative_to(env_dir), ( + f"console script runs {interpreter}, outside the venv {env_dir}" + ) + + # The property that actually confers isolation: this interpreter's + # site-packages is the venv's, not the base environment's. + prefix = subprocess.run( + [str(interpreter), "-c", "import sys; print(sys.prefix)"], + capture_output=True, text=True, check=True, + ).stdout.strip() + assert Path(prefix) == env_dir, ( + f"the console script's interpreter has sys.prefix={prefix}, " + f"expected the venv at {env_dir}" + ) + + +class TestWheelConsoleSmoke: + def test_installed_console_script_startup_help_and_exit( + self, installed_wheel, sandbox + ): + """The user-facing path: run `research-demo`, not `python -m`.""" + home, cwd = sandbox + + result = run_console_smoke( + [str(installed_wheel["console"])], cwd=cwd, home=home + ) + + assert_clean_session(result) + assert result.argv == [str(installed_wheel["console"])], ( + "the smoke did not execute the installed console script" + ) + + def test_module_invocation_also_works(self, installed_wheel, sandbox): + """`python -m research_demo` from the installed wheel, for parity.""" + home, cwd = sandbox + + result = run_console_smoke( + [str(installed_wheel["python"]), "-m", "research_demo"], + cwd=cwd, + home=home, + ) + + assert_clean_session(result) + + +class TestArtifactCarriesTheDeclaredVersion: + """The built and installed artifact must match ``pyproject.toml``. + + All three surfaces are checked because they can disagree: the wheel + filename comes from the build, the distribution metadata from the install, + and ``__version__`` is a separate literal in the package that a bump can + forget. + """ + + def test_wheel_filename_carries_the_declared_version(self, installed_wheel): + expected = project_version() + assert installed_wheel["wheel"].name.startswith(f"agentic_cli-{expected}-"), ( + f"wheel {installed_wheel['wheel'].name} does not carry the declared " + f"version {expected}" + ) + + def test_installed_metadata_and_module_version_agree(self, installed_wheel): + """Distribution metadata and ``agentic_cli.__version__``, from the venv.""" + expected = project_version() + probe = ( + "import json, importlib.metadata as md, agentic_cli; " + "print(json.dumps({" + "'distribution': md.version('agentic-cli'), " + "'module': agentic_cli.__version__}))" + ) + result = _run( + [str(installed_wheel["python"]), "-c", probe], + cwd=str(installed_wheel["root"]), + ) + found = json.loads(result.stdout.strip().splitlines()[-1]) + + assert found["distribution"] == expected, ( + f"installed distribution metadata says {found['distribution']}, " + f"pyproject.toml declares {expected}" + ) + assert found["module"] == expected, ( + f"agentic_cli.__version__ is {found['module']}, " + f"pyproject.toml declares {expected}" + ) diff --git a/tests/integration/test_adk_claude_live.py b/tests/integration/test_adk_claude_live.py new file mode 100644 index 0000000..7997547 --- /dev/null +++ b/tests/integration/test_adk_claude_live.py @@ -0,0 +1,224 @@ +"""LIVE Claude-on-ADK tests — real Anthropic API calls (opt-in, costs money). + +Marked ``@pytest.mark.llm`` (skipped by default in a plain run; the integration +``conftest`` loads real keys from ``~/.research_demo/.env`` or +``$AGENTIC_TEST_ENV_FILE``). Run explicitly with ``-m llm``: + + conda run -n agenticcli python -m pytest tests/integration/test_adk_claude_live.py -v -m llm + # optional overrides: + # LIVE_CLAUDE_MODEL=claude-sonnet-4-6 default; supports budget thinking + effort + # LIVE_CLAUDE_47PLUS_MODEL=claude-opus-4-8 only if you have access; enables the 400 test + +They exercise ``DirectAnthropicLlm.generate_content_async`` directly (no session +/ runner) so they target the new code with minimal scaffolding. Requests use a +tiny ``max_tokens`` to stay cheap. Assertions are structural (got a response / +got a 400), not exact text. + +Deferred (can't be tested on this branch yet): +- Adaptive thinking via ADK — needs the google-adk >= 1.34 bump (1.33 raises on a + negative thinking_budget before any request is sent). +- Effort routed through the manager's planner — the manager doesn't set effort yet. +""" + +from __future__ import annotations + +import os + +import pytest + +pytest.importorskip("google.adk") +pytest.importorskip("anthropic") + +from google.genai import types # noqa: E402 +from google.adk.models.llm_request import LlmRequest # noqa: E402 + +from agentic_cli.workflow.adk.anthropic_llm import DirectAnthropicLlm # noqa: E402 + + +pytestmark = [ + pytest.mark.llm, + pytest.mark.skipif( + not os.environ.get("ANTHROPIC_API_KEY"), + reason="No Anthropic API key (set ANTHROPIC_API_KEY or ~/.research_demo/.env).", + ), +] + +MODEL = os.getenv("LIVE_CLAUDE_MODEL", "claude-sonnet-4-6") +MODEL_47PLUS = os.getenv("LIVE_CLAUDE_47PLUS_MODEL") # e.g. claude-opus-4-8 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _request(text: str, *, thinking_budget: int | None = None, tools=None) -> LlmRequest: + config = types.GenerateContentConfig( + system_instruction="You are a helpful, concise assistant.", + ) + if thinking_budget is not None: + config.thinking_config = types.ThinkingConfig(thinking_budget=thinking_budget) + if tools is not None: + config.tools = tools + return LlmRequest( + model=MODEL, + contents=[types.Content(role="user", parts=[types.Part(text=text)])], + config=config, + ) + + +async def _collect(llm: DirectAnthropicLlm, request: LlmRequest) -> list: + responses = [] + async for resp in llm.generate_content_async(request, stream=False): + responses.append(resp) + return responses + + +def _all_parts(responses) -> list: + parts = [] + for r in responses: + content = getattr(r, "content", None) + parts.extend(getattr(content, "parts", None) or []) + return parts + + +def _text(responses) -> str: + return "".join(p.text for p in _all_parts(responses) if getattr(p, "text", None)) + + +# --------------------------------------------------------------------------- +# Basic turn + tool call (no thinking) +# --------------------------------------------------------------------------- + + +class TestLiveBasic: + async def test_basic_turn(self): + llm = DirectAnthropicLlm(model=MODEL, max_tokens=64) + responses = await _collect(llm, _request("Reply with exactly one word: pong")) + assert _text(responses).strip(), "expected non-empty text response" + + async def test_tool_call(self): + llm = DirectAnthropicLlm(model=MODEL, max_tokens=256) + tool = types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="get_weather", + description="Get the current weather for a city.", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={"city": types.Schema(type=types.Type.STRING)}, + required=["city"], + ), + ) + ] + ) + req = _request("What's the weather in Paris? Use the tool.", tools=[tool]) + responses = await _collect(llm, req) + calls = [ + p.function_call + for p in _all_parts(responses) + if getattr(p, "function_call", None) + ] + assert calls, "expected a function_call part" + assert calls[0].name == "get_weather" + + +# --------------------------------------------------------------------------- +# Budget thinking (works on <= 4.6; this is the deprecated-but-functional path) +# --------------------------------------------------------------------------- + + +class TestLiveBudgetThinking: + async def test_budget_thinking_succeeds(self): + # max_tokens must exceed the thinking budget (Anthropic requirement). + llm = DirectAnthropicLlm(model=MODEL, max_tokens=4096) + responses = await _collect( + llm, _request("What is 17 * 24? Think it through.", thinking_budget=2048) + ) + # Should not error; should produce some output. A thought part may or may + # not be surfaced depending on the model/display defaults. + assert _text(responses).strip() or any( + getattr(p, "thought", None) for p in _all_parts(responses) + ) + + +# --------------------------------------------------------------------------- +# output_config.effort accepted on the wire (the key new-code assertion) +# --------------------------------------------------------------------------- + + +class TestLiveEffort: + async def test_effort_low_accepted(self): + llm = DirectAnthropicLlm(model=MODEL, max_tokens=256, effort="low") + responses = await _collect(llm, _request("Name one primary color.")) + assert _text(responses).strip(), "effort=low request should return text" + + async def test_effort_high_accepted(self): + llm = DirectAnthropicLlm(model=MODEL, max_tokens=512, effort="high") + responses = await _collect(llm, _request("Name one primary color.")) + assert _text(responses).strip(), "effort=high request should return text" + + +# --------------------------------------------------------------------------- +# Adaptive thinking + effort (ADK >= 1.34; negative budget -> {type:"adaptive"}) +# --------------------------------------------------------------------------- + + +class TestLiveAdaptiveThinking: + async def test_adaptive_plus_effort(self): + llm = DirectAnthropicLlm(model=MODEL, max_tokens=2048, effort="high") + responses = await _collect( + llm, _request("What is 12 * 12? Think briefly.", thinking_budget=-1) + ) + assert _text(responses).strip() or any( + getattr(p, "thought", None) for p in _all_parts(responses) + ) + + +@pytest.mark.skipif( + not MODEL_47PLUS, + reason="Set LIVE_CLAUDE_47PLUS_MODEL (e.g. claude-opus-4-8) to run the 4.7+ adaptive test.", +) +class TestLiveAdaptiveOn47Plus: + async def test_adaptive_succeeds_where_budget_400s(self): + # Same model that 400s on budget_tokens succeeds with adaptive + effort. + llm = DirectAnthropicLlm(model=MODEL_47PLUS, max_tokens=2048, effort="high") + req = LlmRequest( + model=MODEL_47PLUS, + contents=[types.Content(role="user", parts=[types.Part(text="Name one color.")])], + config=types.GenerateContentConfig( + system_instruction="You are a helpful, concise assistant.", + thinking_config=types.ThinkingConfig(thinking_budget=-1), + ), + ) + responses = await _collect(llm, req) + assert _text(responses).strip() or any( + getattr(p, "thought", None) for p in _all_parts(responses) + ) + + +# --------------------------------------------------------------------------- +# Confirm the limitation: budget thinking 400s on Opus 4.7+/Fable +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not MODEL_47PLUS, + reason="Set LIVE_CLAUDE_47PLUS_MODEL (e.g. claude-opus-4-8) to run the 400-confirmation test.", +) +class TestLiveBudgetRejectedOn47Plus: + async def test_budget_thinking_400s(self): + import anthropic + + llm = DirectAnthropicLlm(model=MODEL_47PLUS, max_tokens=4096) + req = LlmRequest( + model=MODEL_47PLUS, + contents=[types.Content(role="user", parts=[types.Part(text="hi")])], + config=types.GenerateContentConfig( + system_instruction="You are a helpful, concise assistant.", + thinking_config=types.ThinkingConfig(thinking_budget=2048), + ), + ) + with pytest.raises(anthropic.BadRequestError): + async for _ in llm.generate_content_async(req, stream=False): + pass diff --git a/tests/integration/test_adk_integration.py b/tests/integration/test_adk_integration.py index ea76bc3..e3224f4 100644 --- a/tests/integration/test_adk_integration.py +++ b/tests/integration/test_adk_integration.py @@ -523,154 +523,89 @@ 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() + async def test_rate_limited_turn_is_not_replayed(self): + from agentic_cli.cli.message_processor import MessageProcessor, TurnStatus - workflow_controller = MagicMock() - workflow_controller.ensure_initialized = AsyncMock(return_value=True) - - 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 - - mock_workflow = MagicMock() - mock_workflow.process = mock_process - workflow_controller.workflow = mock_workflow + workflow_controller, ui, settings, _calls = self._harness() - 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 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 +616,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 +635,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/integration/test_live_durable_sessions.py b/tests/integration/test_live_durable_sessions.py new file mode 100644 index 0000000..748202d --- /dev/null +++ b/tests/integration/test_live_durable_sessions.py @@ -0,0 +1,109 @@ +"""Live end-to-end: durable sessions survive a restart (fresh manager, same db). + +Manager 1 tells the model a codeword (persisted to sqlite via ADK's +DatabaseSessionService); a brand-new Manager 2 over the same store resumes the +same session_id and the model recalls it — proving cross-restart durability with +full fidelity, no custom snapshot. + +@pytest.mark.llm, ADK-only (needs GOOGLE_API_KEY); skipped by default: + + conda run -n agenticcli python -m pytest tests/integration/test_live_durable_sessions.py -v -m llm +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from agentic_cli.config import BaseSettings, set_settings +from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.events import EventType +from agentic_cli.workflow.factory import create_workflow_manager_from_settings + +pytestmark = [ + pytest.mark.llm, + pytest.mark.skipif( + not os.environ.get("GOOGLE_API_KEY"), reason="durable sessions test needs GOOGLE_API_KEY" + ), +] + +_AGENT = AgentConfig( + name="assistant", + prompt="You are a concise assistant. When asked to remember something, do so.", + tools=[], + description="plain assistant", +) + + +async def _run(manager, message: str, session_id: str, user_id: str) -> str: + out: list[str] = [] + async for ev in manager.process(message=message, user_id=user_id, session_id=session_id): + if ev.type == EventType.TEXT: + out.append(ev.content) + return " ".join(out) + + +async def test_session_survives_a_fresh_manager(tmp_path: Path, monkeypatch): + monkeypatch.chdir(tmp_path) + settings = BaseSettings( + workspace_dir=tmp_path, session_store="sqlite", permissions_enabled=False + ) + set_settings(settings) + sid = "durable-codeword" + # The CLI always runs turns as settings.default_user; session queries use the + # same, so mirror that here. + user = settings.default_user + + # --- "Process 1": store a codeword, then tear the manager down. --- + m1 = create_workflow_manager_from_settings(agent_configs=[_AGENT], settings=settings) + await _run(m1, "Remember this codeword exactly: BANANA77. Just acknowledge.", sid, user) + await m1.cleanup() + + # --- "Process 2": a brand-new manager over the same sqlite store. --- + m2 = create_workflow_manager_from_settings(agent_configs=[_AGENT], settings=settings) + await m2.initialize_services() # the app resumes only after init (ensure_initialized) + try: + # The session is already durable in the store (no inject needed). + assert await m2.session_exists(sid) is True + answer = await _run(m2, "What codeword did I ask you to remember?", sid, user) + assert "BANANA77" in answer, f"model did not recall across restart: {answer!r}" + finally: + await m2.cleanup() + + +async def test_langgraph_session_survives_a_fresh_manager(tmp_path: Path, monkeypatch): + """Same durability check on the LangGraph backend (persistent checkpointer). + + Forces gemini onto LangGraph (orchestrator=langgraph) — which also exercises + the gemini-2.5 thinking_budget fix on the LangGraph LLM path. + """ + pytest.importorskip("langgraph") + pytest.importorskip("langchain_google_genai") + pytest.importorskip("langgraph.checkpoint.sqlite.aio") + + monkeypatch.chdir(tmp_path) + settings = BaseSettings( + workspace_dir=tmp_path, + session_store="sqlite", + permissions_enabled=False, + orchestrator="langgraph", # force LangGraph for the gemini default + ) + set_settings(settings) + sid = "lg-durable-codeword" + user = settings.default_user + + m1 = create_workflow_manager_from_settings(agent_configs=[_AGENT], settings=settings) + assert m1.backend_type == "langgraph" + await _run(m1, "Remember this codeword exactly: BANANA77. Just acknowledge.", sid, user) + await m1.cleanup() + + m2 = create_workflow_manager_from_settings(agent_configs=[_AGENT], settings=settings) + await m2.initialize_services() + try: + assert await m2.session_exists(sid) is True + answer = await _run(m2, "What codeword did I ask you to remember?", sid, user) + assert "BANANA77" in answer, f"langgraph did not recall across restart: {answer!r}" + finally: + await m2.cleanup() diff --git a/tests/integration/test_live_job_resume.py b/tests/integration/test_live_job_resume.py new file mode 100644 index 0000000..caf462c --- /dev/null +++ b/tests/integration/test_live_job_resume.py @@ -0,0 +1,243 @@ +"""Live end-to-end validation of ADK push/resume (phase-2 milestone 2). + +Drives the *real* resume mechanic against a live Gemini model in a single +process: turn 1 calls a long-running tool (which starts a job and leaves the +call pending); we then hand the finished result back via +``resume_with_job_result`` and confirm the model actually reacts to it. + +This is the test that answers the open "pending-response shape" question — it is +parametrized on whether the tool's initial return is a ``{"status":"pending"}`` +dict or ``None`` — and proves ADK accepts a second ``FunctionResponse`` for an +already-pending long-running call. + +Marked ``@pytest.mark.llm`` (skipped by default; costs a few real calls). ADK +only, so it requires GOOGLE_API_KEY: + + conda run -n agenticcli python -m pytest tests/integration/test_live_job_resume.py -v -m llm +""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from agentic_cli.config import BaseSettings, set_settings +from agentic_cli.tools.registry import ToolCategory, register_tool +from agentic_cli.workflow.base_manager import BaseWorkflowManager +from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.events import EventType +from agentic_cli.workflow.factory import create_workflow_manager_from_settings +from agentic_cli.workflow.permissions import Capability +from agentic_cli.workflow.service_registry import JOB_MANAGER, get_service + +from tests.integration.helpers import find_events, find_tool_calls + +_has_google = bool(os.environ.get("GOOGLE_API_KEY")) + +pytestmark = [ + pytest.mark.llm, + pytest.mark.skipif(not _has_google, reason="ADK resume needs GOOGLE_API_KEY"), +] + +# Toggled per-parametrization to test both initial-return shapes. +_SHAPE = {"return_none": False} + +_TOKEN = "RESUME_TOKEN_8731" + + +@register_tool( + category=ToolCategory.EXECUTION, + capabilities=[Capability("longrunning.live_resume_probe")], + long_running=True, + description="Start a background shell job; returns a job_id immediately.", +) +def live_resume_probe(command: str = "", tool_context=None): + """Submit a subprocess job flagged for resume, capturing the call id.""" + jm = get_service(JOB_MANAGER) + if jm is None: + return {"success": False, "error": "job manager not available"} + jm.submit( + tool="live_resume_probe", + backend="subprocess", + spec={"command": command or f"echo {_TOKEN}"}, + resume_on_complete=True, + call_id=getattr(tool_context, "function_call_id", None), + call_name="live_resume_probe", + ) + # The shape under test: an explicit pending dict vs. None. + return None if _SHAPE["return_none"] else {"status": "pending"} + + +_AGENT = AgentConfig( + name="job_runner", + prompt=( + "You run background jobs. When asked, call live_resume_probe with a shell " + f"command that prints {_TOKEN} (e.g. 'echo {_TOKEN}'), then tell the user " + "you started it. Never ask for confirmation. You have no other tools." + ), + tools=[live_resume_probe], + description="Starts a long-running background job.", +) + + +@pytest.fixture +def live_settings(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """Isolated settings + a job tool mapped to the JobManager service.""" + app_name = "agentic_cli_live_resume_test" + workspace = tmp_path / "ws" + workspace.mkdir(parents=True) + monkeypatch.chdir(workspace) + + # Make the probe tool require the JobManager (detection is by tool name). + monkeypatch.setattr( + BaseWorkflowManager, + "_TOOL_SERVICE_MAP", + {**BaseWorkflowManager._TOOL_SERVICE_MAP, "live_resume_probe": "job_manager"}, + ) + + settings = BaseSettings( + app_name=app_name, workspace_dir=workspace, permissions_enabled=False + ) + set_settings(settings) + try: + yield settings + finally: + shutil.rmtree(Path.home() / f".{app_name}", ignore_errors=True) + + +def _wait_terminal(jm, job_id: str, timeout: float = 10.0): + from agentic_cli.tools.jobs import JobState + + terminal = {JobState.SUCCEEDED, JobState.FAILED, JobState.CANCELLED, JobState.UNKNOWN} + end = time.time() + timeout + while time.time() < end: + rec = jm.get(job_id) + if rec is not None and rec.state in terminal: + return rec + time.sleep(0.1) + return jm.get(job_id) + + +@pytest.mark.parametrize("return_none", [False, True], ids=["pending_dict", "none"]) +async def test_long_running_tool_then_resume(live_settings, return_none: bool): + _SHAPE["return_none"] = return_none + + manager = create_workflow_manager_from_settings( + agent_configs=[_AGENT], settings=live_settings + ) + assert hasattr(manager, "resume_with_job_result"), "expected the ADK manager" + + async def _auto_input(request) -> str: # noqa: ANN001 + return "yes, proceed" + + manager.set_input_callback(_auto_input) + try: + # --- Turn 1: model calls the long-running tool, job starts pending. --- + session_id = "live-resume-1" + turn1 = [] + async for ev in manager.process( + message="Start a background job now.", user_id="tester", session_id=session_id + ): + turn1.append(ev) + + assert find_tool_calls(turn1, "live_resume_probe"), ( + "model did not call the long-running tool; calls: " + f"{[c.metadata.get('tool_name') for c in find_tool_calls(turn1)]}" + ) + + jm = manager.job_manager + assert jm is not None, "JobManager was not created" + awaiting = [r for r in jm.list() if r.resume_on_complete] + assert len(awaiting) == 1, f"expected one resume-flagged job, got {awaiting}" + record = awaiting[0] + # The crux of association: the tool captured the ADK function_call id. + assert record.call_id, "call_id was not captured from tool_context" + assert record.session_id == session_id and record.user_id == "tester" + + _wait_terminal(jm, record.job_id) + record = jm.get(record.job_id) + + # --- Turn 2: deliver the result back to the pending call and resume. --- + turn2 = [] + async for ev in manager.resume_with_job_result(record): + turn2.append(ev) + + # If ADK rejected a 2nd FunctionResponse for the pending call, run_async + # would have raised above. Reaching here with a model reply == success. + text_events = find_events(turn2, EventType.TEXT) + assert text_events, ( + f"resume produced no model text (shape={'none' if return_none else 'pending_dict'}); " + f"events: {[(e.type, e.content[:40]) for e in turn2]}" + ) + finally: + manager.clear_input_callback() + await manager.cleanup() + + +async def test_full_resume_loop_via_coordinator(live_settings): + """End-to-end: the harness coordinator drains a finished job into a resume turn. + + Exercises the whole milestone-3 path live: awaiting_resume → mark_resumed → + MessageProcessor.process_resume → _run_turn → resume_with_job_result → + model reply, rendered into a RecordingSession. + """ + from agentic_cli.cli.app import BaseCLIApp + from agentic_cli.cli.message_processor import MessageProcessor + from tests.event_replay import RecordingSession + + _SHAPE["return_none"] = False # informative pending dict (validated shape) + + manager = create_workflow_manager_from_settings( + agent_configs=[_AGENT], settings=live_settings + ) + + async def _auto_input(request) -> str: # noqa: ANN001 + return "yes, proceed" + + manager.set_input_callback(_auto_input) + try: + async for _ in manager.process( + message="Start a background job now.", user_id="tester", session_id="loop-1" + ): + pass + + jm = manager.job_manager + flagged = [r for r in jm.list() if r.resume_on_complete] + assert len(flagged) == 1, f"expected one resume-flagged job, got {flagged}" + _wait_terminal(jm, flagged[0].job_id) + + # Drive the real harness coordinator with a recording session. + session = RecordingSession() # is_thinking=True → cancel watcher won't fire + app = BaseCLIApp.__new__(BaseCLIApp) + app._workflow_controller = SimpleNamespace( + is_ready=True, + workflow=manager, + ensure_initialized=lambda ui=None: _true(), + update_status_bar=lambda ui: None, + ) + app._turn_lock = asyncio.Lock() + app._message_processor = MessageProcessor() + app.session = session + app._settings = live_settings + app._usage_tracker = None + + n = await app.resume_finished_jobs() + + assert n == 1 + assert jm.get(flagged[0].job_id).resumed is True + assert session.responses(), ( + f"coordinator resume produced no model response; recorded: {session.kinds()}" + ) + finally: + manager.clear_input_callback() + await manager.cleanup() + + +async def _true() -> bool: + return True diff --git a/tests/integration/test_permission_adk.py b/tests/integration/test_permission_adk.py index 5d36ade..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,14 +122,46 @@ 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, ) assert result == {"success": False, "error": "Permission denied: rule: builtin/deny"} @pytest.mark.asyncio - async def test_engine_absent_allows(self, monkeypatch): + async def test_engine_absent_denies_when_permissions_enabled(self, monkeypatch): + """Fail closed: permissions on but no engine wired -> deny, not allow.""" + from agentic_cli.tools.registry import get_registry + from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin + + monkeypatch.setattr( + "agentic_cli.workflow.adk.permission_plugin.get_service", + lambda k: None, + ) + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=True), + ) + reg = get_registry() + + @reg.register( + name="reader_y_deny", + capabilities=[Capability("filesystem.read", target_arg="path")], + ) + def reader_y_deny(path: str): + return {} + + plugin = PermissionPlugin() + result = await plugin.before_tool_callback( + tool=_adk_tool(reader_y_deny), + tool_args={"path": "/tmp/x"}, + tool_context=None, + ) + assert result is not None and result["success"] is False + + @pytest.mark.asyncio + async def test_engine_absent_allows_when_permissions_disabled(self, monkeypatch): + """Permissions off is the only case a missing engine allows.""" from agentic_cli.tools.registry import get_registry from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin @@ -123,19 +169,23 @@ async def test_engine_absent_allows(self, monkeypatch): "agentic_cli.workflow.adk.permission_plugin.get_service", lambda k: None, ) + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=False), + ) reg = get_registry() @reg.register( - name="reader_y", + name="reader_y_allow", capabilities=[Capability("filesystem.read", target_arg="path")], ) - def reader_y(path: str): + def reader_y_allow(path: str): return {} plugin = PermissionPlugin() result = await plugin.before_tool_callback( - tool=SimpleNamespace(name="reader_y"), + tool=_adk_tool(reader_y_allow), tool_args={"path": "/tmp/x"}, tool_context=None, ) - assert result is None # fallback allow + assert result is None diff --git a/tests/integration/test_permission_langgraph.py b/tests/integration/test_permission_langgraph.py index a9c4e45..4b6765e 100644 --- a/tests/integration/test_permission_langgraph.py +++ b/tests/integration/test_permission_langgraph.py @@ -88,7 +88,10 @@ def write_lg(path: str): assert result == {"success": False, "error": "Permission denied: rule: builtin/deny"} @pytest.mark.asyncio - async def test_engine_absent_allows(self, monkeypatch): + async def test_engine_absent_denies_when_permissions_enabled(self, monkeypatch): + """Fail closed: permissions on but no engine wired -> deny, not run.""" + from types import SimpleNamespace + from agentic_cli.tools.registry import get_registry from agentic_cli.workflow.langgraph.permission_wrap import wrap_tool_for_permission @@ -96,15 +99,48 @@ async def test_engine_absent_allows(self, monkeypatch): "agentic_cli.workflow.langgraph.permission_wrap.get_service", lambda k: None, ) + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=True), + ) + reg = get_registry() + + @reg.register( + name="read_lg_deny", + capabilities=[Capability("filesystem.read", target_arg="path")], + ) + def read_lg_deny(path: str): + return {"ok": True} + + wrapped = wrap_tool_for_permission(read_lg_deny) + result = await wrapped(path="/x") + assert result["success"] is False + + @pytest.mark.asyncio + async def test_engine_absent_allows_when_permissions_disabled(self, monkeypatch): + """Permissions off is the only case a missing engine runs the tool.""" + from types import SimpleNamespace + + from agentic_cli.tools.registry import get_registry + from agentic_cli.workflow.langgraph.permission_wrap import wrap_tool_for_permission + + monkeypatch.setattr( + "agentic_cli.workflow.langgraph.permission_wrap.get_service", + lambda k: None, + ) + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=False), + ) reg = get_registry() @reg.register( - name="read_lg2", + name="read_lg_allow", capabilities=[Capability("filesystem.read", target_arg="path")], ) - def read_lg2(path: str): + def read_lg_allow(path: str): return {"ok": True} - wrapped = wrap_tool_for_permission(read_lg2) + wrapped = wrap_tool_for_permission(read_lg_allow) result = await wrapped(path="/x") assert result == {"ok": True} diff --git a/tests/integration/test_research_scenarios.py b/tests/integration/test_research_scenarios.py index 525e334..b913645 100644 --- a/tests/integration/test_research_scenarios.py +++ b/tests/integration/test_research_scenarios.py @@ -12,14 +12,17 @@ conda run -n agenticcli python -m pytest tests/integration/test_research_scenarios.py -v -m llm -Requires GOOGLE_API_KEY (Gemini/ADK, the demo default) or ANTHROPIC_API_KEY. -For Anthropic, also set AGENTIC_SCENARIO_MODEL=claude-... so the factory routes -to LangGraph. Set AGENTIC_RECORD_EVENTS= to dump each run's event stream -to JSON (replayable by the deterministic render tests). +Requires GOOGLE_API_KEY (Gemini, the demo default) or ANTHROPIC_API_KEY. These +scenarios are pinned to the **ADK** orchestrator, which is what the demo ships +with; ADK runs Claude natively via ``AnthropicLlm``, so setting +``AGENTIC_SCENARIO_MODEL=claude-...`` just changes the model, not the backend. +Set AGENTIC_RECORD_EVENTS= to dump each run's event stream to JSON +(replayable by the deterministic render tests). """ from __future__ import annotations +import contextlib import json import os import sys @@ -27,10 +30,18 @@ import pytest -from agentic_cli.config import BaseSettings, set_settings +from agentic_cli.workflow.adk.transfer_tool_description import TRANSFER_TOOL_NAME from agentic_cli.workflow.events import EventType from agentic_cli.workflow.factory import create_workflow_manager_from_settings - +from agentic_cli.workflow.settings import OrchestratorType + +from tests.demo_isolation import ( + assert_dotenv_isolated, + assert_no_global_settings_leak, + effective_env_file, + isolated_env_file, + make_isolated_settings, +) from tests.event_replay import events_to_dicts from tests.integration.helpers import ( find_events, @@ -43,6 +54,7 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) from examples.research_demo.agents import AGENT_CONFIGS # noqa: E402 +from examples.research_demo.settings import ResearchDemoSettings # noqa: E402 _has_any_key = bool( @@ -56,22 +68,73 @@ @pytest.fixture -def research_settings(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> BaseSettings: - """Settings for live runs: real keys, temp workspace, permission gate off. +def research_settings( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> ResearchDemoSettings: + """The demo's own settings, fully isolated: temp HOME, temp cwd, no gate. + + Uses ``ResearchDemoSettings`` rather than a bare ``BaseSettings`` so the + scenarios exercise what the demo actually runs with — its ``skills_dirs``, + its sandbox data mount and artifacts dir, and its ``research_demo`` + app_name (which is what ``./.{app_name}/`` and ``~/.{app_name}/`` are + derived from). + + Isolation is on both axes, because the demo reads from both: + + - ``HOME`` is redirected, so the user config and user KB resolve into + tmp_path and a developer's real config cannot change what these assert. + - ``cwd`` is redirected, so the project KB and permission workdir + (``./.research_demo/...``) are written under tmp_path, not into the repo. + - the dotenv is passed explicitly (``make_isolated_settings``), because + ``model_config["env_file"]`` was frozen to the real ``~/.research_demo/ + .env`` at import and ``HOME`` cannot move it. + + Provider credentials still arrive the way the live-test framework supplies + them (real environment / the integration conftest) — that is the point of a + live test. What must *not* happen is ``ResearchDemoSettings`` re-reading the + developer's dotenv or JSON config and changing the model, orchestrator or + workspace out from under the scenario. + + No ``set_settings()``: the manager scopes settings to itself for the whole + turn (``BaseWorkflowManager._workflow_context`` → ``set_context_settings``, + and ADK agent construction under ``SettingsContext``), so the global + singleton is unnecessary — and it has no teardown, so setting it would leak + this fixture's settings into every later test in the session. The permission gate is disabled so tool calls run headlessly without a prompt UI — these tests exercise agent behavior, not the permission UX. - - The project knowledge base and the permission workdir are resolved relative - to ``cwd`` (``./.{app_name}/...``), so we chdir into the temp workspace to - keep those writes out of the repo. """ + home = tmp_path / "home" workspace = tmp_path / "research_ws" - workspace.mkdir(parents=True) - monkeypatch.chdir(workspace) - settings = BaseSettings(workspace_dir=workspace, permissions_enabled=False) - set_settings(settings) - return settings + project = tmp_path / "project" + for d in (home, workspace, project): + d.mkdir(parents=True) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.chdir(project) + + settings = make_isolated_settings( + home=home, + workspace=workspace, + permissions_enabled=False, + # The demo ships ADK; pin it so a stray orchestrator setting in the + # environment cannot silently move these scenarios to another backend. + orchestrator=OrchestratorType.ADK, + ) + + # Prove the isolation rather than assume it: a live run that silently read + # the real dotenv could pick up a different model or workspace. + assert_dotenv_isolated( + effective_env_file( + workspace_dir=workspace, + _env_file=str(isolated_env_file(home)), + ), + home, + ) + assert settings.orchestrator is OrchestratorType.ADK + + yield settings + + assert_no_global_settings_leak() def _maybe_record(events: list, name: str) -> None: @@ -85,14 +148,26 @@ def _maybe_record(events: list, name: str) -> None: ) -async def run_agent( - settings: BaseSettings, - message: str, +@contextlib.asynccontextmanager +async def conversation( + settings: ResearchDemoSettings, *, - session_id: str | None = None, - record_name: str | None = None, -) -> list: - """Run one turn through the real demo agents and return the event stream.""" + user_id: str = "tester", + session_id: str = "scenario-session", +): + """One manager and one session, held open across several turns. + + The coordinator's policy is size-based: an explicit, bounded request ("search + arXiv for two papers", "ingest this note") is executed or delegated + immediately, while an open-ended goal — or an explicit request for a plan — + is planned, shown and held for confirmation. The scenarios therefore ask for + one or the other deliberately, and each asserts only its own contract. + + Multiple turns matter where state carries across them: the KB scenario + ingests in one turn and reads back in the next, which only works because the + manager is initialized and cleaned up exactly once and every turn reuses the + same ``user_id`` and ``session_id``. + """ manager = create_workflow_manager_from_settings( agent_configs=AGENT_CONFIGS, settings=settings, @@ -104,23 +179,59 @@ async def _auto_input(request) -> str: # noqa: ANN001 return "yes, proceed" manager.set_input_callback(_auto_input) - try: + turns: list[list] = [] + + async def say(message: str) -> list: + """Send one user message; return that turn's events. + + A workflow exception propagates: these are behavioural contracts, and a + turn that died is a failure, not something to inspect around. Only + ``Exception`` is involved anywhere here — cancellation, Ctrl+C and + SystemExit must keep unwinding an interrupted live run. + """ events: list = [] async for event in manager.process( - message=message, user_id="tester", session_id=session_id + message=message, user_id=user_id, session_id=session_id ): events.append(event) + turns.append(events) + return events + + say.turns = turns # type: ignore[attr-defined] + say.all_events = lambda: [e for t in turns for e in t] # type: ignore[attr-defined] + + try: + yield say finally: manager.clear_input_callback() await manager.cleanup() + +async def run_agent( + settings: ResearchDemoSettings, + message: str, + *, + session_id: str | None = None, + record_name: str | None = None, +) -> list: + """Run one turn through the real demo agents and return the event stream.""" + async with conversation( + settings, session_id=session_id or "scenario-session" + ) as say: + events = await say(message) + if record_name: _maybe_record(events, record_name) return events class TestPlanningScenario: - """Planning: the coordinator should save and show a plan.""" + """Planning, on its own: an explicit plan request produces a saved plan. + + Deliberately does **not** also test delegation. Planning, delegation and KB + behaviour are three independent properties, and bundling them meant one + stochastic policy choice failed all three at once. + """ async def test_creates_and_saves_a_plan(self, research_settings): events = await run_agent( @@ -139,23 +250,87 @@ async def test_creates_and_saves_a_plan(self, research_settings): assert find_events(events, EventType.TEXT), "no text response produced" +def _tool_names(events: list) -> list[str]: + return [c.metadata.get("tool_name") for c in find_tool_calls(events)] + + +#: A bounded, fully-specified arXiv request. Per the coordinator prompt this is +#: an "explicit, bounded operation": delegate and run it, no plan, no +#: confirmation. Shared so the delegation diagnostic and the behavioural +#: scenario exercise the same path. +BOUNDED_ARXIV_REQUEST = ( + "Search arXiv for 2 recent papers on 'speculative decoding' and list their " + "titles. This is a small, explicit request — do it now via the " + "arxiv_specialist; no plan and no confirmation needed." +) + + +#: The class name ADK's own tool description used to recommend. Asserted +#: against here (never put in the agent's prompt, where it would prime the +#: model toward the very mistake being guarded). +_TRANSFER_CLASS_NAME = "TransferToAgentTool" + + +def _assert_no_class_name_transfer(events: list) -> None: + """Delegation must go through the function, never the class name.""" + called = _tool_names(events) + assert _TRANSFER_CLASS_NAME not in called, ( + f"the model called the transfer tool by its class name: {called}" + ) + offending = [ + str(e.content) + for e in find_events(events, EventType.ERROR) + if _TRANSFER_CLASS_NAME in str(e.content) + ] + assert not offending, f"transfer-by-class-name error surfaced: {offending}" + + +def _assert_no_unknown_tool_errors(events: list) -> None: + """Every tool the model called must be one that actually exists.""" + unknown = [ + str(e.content) + for e in find_events(events, EventType.ERROR) + if "not found" in str(e.content) + ] + assert not unknown, ( + f"the model called a tool that does not exist: {unknown}; " + f"tools actually called: {_tool_names(events)}" + ) + + class TestArxivScenario: - """arXiv: the specialist should search arXiv and get results back.""" + """Delegation: a bounded request reaches the specialist and runs. - async def test_searches_arxiv(self, research_settings): - events = await run_agent( - research_settings, - "Search arXiv for 2 recent papers on 'speculative decoding' using " - "search_arxiv and list their titles. Do not ask for confirmation.", - record_name="arxiv_search", - ) + One turn, because the coordinator's prompt says an explicit bounded + operation should just be done. No plan is required or asserted here — that + is :class:`TestPlanningScenario`'s job. + + This is also the delegation regression: the transfer defect is asserted + inside the behavioural run rather than as a separate paid diagnostic, so + the class-name check can never be satisfied vacuously by a run in which no + delegation was attempted — the same test requires the correct + ``transfer_to_agent`` call. + """ - calls = find_tool_calls(events, "search_arxiv") - assert calls, "agent did not call search_arxiv" + async def test_bounded_request_delegates_and_searches(self, research_settings): + async with conversation(research_settings) as say: + events = await say(BOUNDED_ARXIV_REQUEST) + _maybe_record(events, "arxiv_search") + # Delegation happened, by the correct function name. + assert find_tool_calls(events, TRANSFER_TOOL_NAME), ( + f"the coordinator never delegated; tools called: {_tool_names(events)}" + ) + _assert_no_class_name_transfer(events) + _assert_no_unknown_tool_errors(events) + + # And the specialist actually did the work. + assert find_tool_calls(events, "search_arxiv"), ( + f"the specialist never searched arXiv; tools called: {_tool_names(events)}" + ) results = find_tool_results(events, "search_arxiv") assert any(r.metadata.get("success", True) for r in results), ( - "search_arxiv never returned a successful result: " + f"search_arxiv never returned a successful result: " f"{[r.content for r in results]}" ) @@ -172,35 +347,57 @@ def _require_embeddings(self): pytest.importorskip("faiss") pytest.importorskip("sentence_transformers") - async def test_ingest_text_then_search(self, research_settings): + async def test_bounded_ingest_then_bounded_readback(self, research_settings): + """Two bounded operations in one persistent conversation. + + Each turn is explicit and small, so per the coordinator's prompt both + should just run — no plan-policy assertions here. The two turns share a + manager and a session so the read-back sees what the ingest wrote. + + The ingest turn asks for the *outcome* and points at the specialist, + rather than naming a writer tool: the coordinator holds only the KB + readers (``kb_search``/``kb_read``/``kb_list``/``kb_search_concepts``), + and ``kb_ingest_*`` belongs to ``arxiv_specialist``. Telling the + coordinator to call a tool it does not have makes a correct refusal + look like a failure. + """ note = ( "Speculative decoding uses a small draft model to propose tokens that a " "larger target model verifies in parallel, cutting latency without " "changing the output distribution." ) - events = await run_agent( - research_settings, - "Ingest the following note into the knowledge base (use the " - "arxiv_specialist, which has KB write access), then search the knowledge " - "base for 'speculative decoding' to confirm it is retrievable. Do not ask " - f"for confirmation.\n\nNote: {note}", - record_name="kb_ingest_search", - ) + async with conversation(research_settings) as say: + ingest_turn = await say( + "Store this exact note in the knowledge base now. The " + "arxiv_specialist holds the knowledge-base write tools, so hand " + "it over. This is a small explicit request — no plan and no " + f"confirmation needed.\n\nNote: {note}" + ) + readback_turn = await say( + "Now search the knowledge base for 'speculative decoding' and " + "show me what you find. Again, just do it." + ) + _maybe_record(say.all_events(), "kb_ingest_search") ingest_results = [ r - for r in find_events(events, EventType.TOOL_RESULT) + for r in find_events(ingest_turn, EventType.TOOL_RESULT) if str(r.metadata.get("tool_name", "")).startswith("kb_ingest") ] assert ingest_results, ( - "no kb_ingest_* tool result observed; tool calls were: " - f"{[c.metadata.get('tool_name') for c in find_tool_calls(events)]}" + f"no kb_ingest_* tool result; turn-1 tools: {_tool_names(ingest_turn)}" ) assert any(r.metadata.get("success", True) for r in ingest_results), ( - "kb_ingest never succeeded: " - f"{[r.content for r in ingest_results]}" + f"kb_ingest never succeeded: {[r.content for r in ingest_results]}" ) - # And the KB was queried afterward. - assert find_tool_calls(events, "kb_search") or find_tool_calls( - events, "kb_list" - ), "agent did not query the KB after ingesting" + + assert find_tool_calls(readback_turn, "kb_search") or find_tool_calls( + readback_turn, "kb_list" + ), f"the KB was never queried back; turn-2 tools: {_tool_names(readback_turn)}" + + # Both turns must have used real tool names. Asserted here rather than + # as a separate paid run: a spliced name (two real tools joined into + # one) shows up as an unknown-tool error on exactly this path. + both_turns = ingest_turn + readback_turn + _assert_no_unknown_tool_errors(both_turns) + _assert_no_class_name_transfer(both_turns) diff --git a/tests/permissions/test_engine.py b/tests/permissions/test_engine.py index 98fa4b4..09efa6b 100644 --- a/tests/permissions/test_engine.py +++ b/tests/permissions/test_engine.py @@ -7,6 +7,7 @@ from agentic_cli.workflow.permissions.capabilities import Capability from agentic_cli.workflow.permissions.engine import PermissionEngine +from agentic_cli.workflow.permissions.prompt import ALLOW_ALWAYS_CHOICE from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource from agentic_cli.workflow.permissions.store import PermissionContext @@ -190,11 +191,12 @@ async def test_user_allow_session_installs_session_rule(self, ctx, tmp_path): w.request_user_input.assert_not_called() @pytest.mark.asyncio - async def test_user_allow_always_writes_project_file(self, ctx, tmp_path, monkeypatch): + async def test_user_allow_always_writes_user_side_project_grants(self, ctx, tmp_path, monkeypatch): import json monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) w = _stub_workflow() - w.request_user_input = AsyncMock(return_value="Allow always (save to project)") + w.request_user_input = AsyncMock(return_value=ALLOW_ALWAYS_CHOICE) engine = PermissionEngine(settings=_stub_settings(), workflow=w, ctx=ctx) result = await engine.check( @@ -204,8 +206,12 @@ async def test_user_allow_always_writes_project_file(self, ctx, tmp_path, monkey ) assert result.allowed is True - data = json.loads((tmp_path / ".agentic/settings.json").read_text()) - allow = data["permissions"]["allow"] + # Interactive grants persist to the USER-side, path-keyed grants file + # (keyed by the resolved project cwd), never into the repo. + grants = json.loads( + (tmp_path / "home" / ".agentic" / "project_grants.json").read_text() + ) + allow = grants[str(ctx.workdir.resolve())]["permissions"]["allow"] assert len(allow) == 1 assert allow[0]["capability"] == "http.read" assert "example.com" in allow[0]["target"] @@ -281,6 +287,41 @@ async def fake_input(request): assert ask_peak == 1 # never two asks in flight simultaneously +class TestResolveListTargetArg: + def test_list_target_arg_resolves_per_item(self, ctx): + from agentic_cli.workflow.permissions.capabilities import ResolvedCapability + eng = PermissionEngine(settings=_stub_settings(), workflow=_stub_workflow(), ctx=ctx) + resolved = eng._resolve( + [Capability("filesystem.read", target_arg="inputs")], + {"inputs": ["/data/a.csv", "/data/b.csv"]}, + ) + assert all(isinstance(rc, ResolvedCapability) for rc in resolved) + targets = sorted(rc.target for rc in resolved) + assert len(resolved) == 2 + assert any(t.endswith("a.csv") for t in targets) + assert any(t.endswith("b.csv") for t in targets) + + def test_scalar_target_arg_still_yields_single_resolved(self, ctx): + """No regression: a scalar value must still produce exactly one ResolvedCapability.""" + eng = PermissionEngine(settings=_stub_settings(), workflow=_stub_workflow(), ctx=ctx) + resolved = eng._resolve( + [Capability("filesystem.read", target_arg="path")], + {"path": "/data/x.csv"}, + ) + assert len(resolved) == 1 + assert resolved[0].target.endswith("x.csv") + + def test_none_target_arg_still_yields_star(self, ctx): + """No regression: target_arg=None must still produce a single ResolvedCapability with target='*'.""" + eng = PermissionEngine(settings=_stub_settings(), workflow=_stub_workflow(), ctx=ctx) + resolved = eng._resolve( + [Capability("python.exec.stateful")], + {"code": "x = 1"}, + ) + assert len(resolved) == 1 + assert resolved[0].target == "*" + + class TestTargetlessAllowAlwaysRegression: """Regression: after 'Allow always' on a targetless capability (target_arg=None), subsequent calls must not re-prompt. @@ -294,8 +335,9 @@ class TestTargetlessAllowAlwaysRegression: @pytest.mark.asyncio async def test_http_read_allow_always_matches_next_call(self, ctx, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) w = _stub_workflow() - w.request_user_input = AsyncMock(return_value="Allow always (save to project)") + w.request_user_input = AsyncMock(return_value=ALLOW_ALWAYS_CHOICE) engine = PermissionEngine(settings=_stub_settings(), workflow=w, ctx=ctx) # First call: no rule → ask → allow always (saves session + project rule) @@ -333,11 +375,12 @@ async def test_filesystem_grant_broadens_to_parent_directory( """When user picks 'Allow always' for filesystem.write to /foo/bar.txt, the rule covers /foo/** — subsequent writes to /foo/baz.txt must not re-prompt.""" monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) outside = tmp_path / "out" outside.mkdir() w = _stub_workflow() - w.request_user_input = AsyncMock(return_value="Allow always (save to project)") + w.request_user_input = AsyncMock(return_value=ALLOW_ALWAYS_CHOICE) engine = PermissionEngine(settings=_stub_settings(), workflow=w, ctx=ctx) # First write → prompts → allow always. @@ -412,24 +455,98 @@ async def test_memory_and_kb_allowed_by_builtin(self, ctx, tmp_path): @pytest.mark.asyncio async def test_reloaded_wildcard_rule_still_matches(self, ctx, tmp_path, monkeypatch): - """After the project JSON is reloaded (simulating next process run), - a rule stored with target='*' must still match a targetless capability.""" + """A rule granted with target='*' must be stored with the wildcard + preserved (not mangled by matchers) so it can be reloaded correctly. + + Verifies both the WRITE side (grants file persisted with '*' intact) and + the RELOAD side (a fresh engine reloads the persisted grant and allows + the same capability without prompting). + """ + import json monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) # Round 1: grant "allow always" so the rule is persisted. w1 = _stub_workflow() - w1.request_user_input = AsyncMock(return_value="Allow always (save to project)") + w1.request_user_input = AsyncMock(return_value=ALLOW_ALWAYS_CHOICE) engine1 = PermissionEngine(settings=_stub_settings(), workflow=w1, ctx=ctx) - await engine1.check("web_search", [Capability("http.read")], {"query": "x"}) + result1 = await engine1.check("web_search", [Capability("http.read")], {"query": "x"}) + assert result1.allowed is True + + # The grants file must preserve the wildcard target ('*') so it + # survives serialisation and can be reloaded without corruption. + grants = json.loads( + (tmp_path / "home" / ".agentic" / "project_grants.json").read_text() + ) + allow = grants[str(ctx.workdir.resolve())]["permissions"]["allow"] + assert len(allow) == 1 + assert allow[0]["capability"] == "http.read" + assert allow[0]["target"] == "*" # wildcard preserved, not mangled - # Round 2: fresh engine reads rules from disk. - w2 = _stub_workflow() - w2.request_user_input = AsyncMock(return_value="Deny") # would deny if re-asked + # Round 2: a FRESH engine (same settings + ctx → same resolved project + # key) must reload the persisted wildcard grant from project_grants.json + # and allow the same capability WITHOUT prompting. + w2 = _stub_workflow() # default response "Deny" — a prompt here fails the test engine2 = PermissionEngine(settings=_stub_settings(), workflow=w2, ctx=ctx) - result = await engine2.check( - "web_search", - [Capability("http.read")], - {"query": "y"}, + result2 = await engine2.check("web_search", [Capability("http.read")], {"query": "y"}) + assert result2.allowed is True + w2.request_user_input.assert_not_called() + + +class TestOptionalCapability: + """A capability marked optional is only exercised when its target arg is + supplied — so a tool with an optional output/asset path doesn't prompt for a + write/read it isn't performing this call.""" + + def _engine(self, ctx): + return PermissionEngine( + settings=_stub_settings(), workflow=_stub_workflow(), ctx=ctx, + ) + + def test_optional_cap_skipped_when_arg_absent(self, ctx): + engine = self._engine(ctx) + resolved = engine._resolve( + [Capability("filesystem.write", target_arg="output_pdf", optional=True)], + {}, # output_pdf not supplied + ) + assert resolved == [] + + def test_optional_cap_skipped_when_arg_empty(self, ctx): + engine = self._engine(ctx) + resolved = engine._resolve( + [Capability("filesystem.write", target_arg="output_pdf", optional=True)], + {"output_pdf": None}, + ) + assert resolved == [] + + def test_optional_cap_resolved_when_arg_present(self, ctx, tmp_path): + engine = self._engine(ctx) + target = str(tmp_path / "out.pdf") + resolved = engine._resolve( + [Capability("filesystem.write", target_arg="output_pdf", optional=True)], + {"output_pdf": target}, + ) + assert len(resolved) == 1 + assert resolved[0].name == "filesystem.write" + + def test_required_cap_still_resolves_when_arg_absent(self, ctx): + """Non-optional (default) behavior is unchanged: an absent target still + resolves (to be evaluated/asked), never silently skipped.""" + engine = self._engine(ctx) + resolved = engine._resolve( + [Capability("filesystem.write", target_arg="path")], # optional=False + {}, + ) + assert len(resolved) == 1 + + @pytest.mark.asyncio + async def test_check_allows_when_all_optional_caps_absent(self, ctx): + """All-optional caps with absent args resolve to [] — check() must not + crash (IndexError on outcomes[0]) and should allow (nothing to gate).""" + engine = self._engine(ctx) + result = await engine.check( + "some_tool", + [Capability("filesystem.write", target_arg="out", optional=True)], + {}, ) assert result.allowed is True - w2.request_user_input.assert_not_called() diff --git a/tests/permissions/test_grant_provenance.py b/tests/permissions/test_grant_provenance.py new file mode 100644 index 0000000..87129c1 --- /dev/null +++ b/tests/permissions/test_grant_provenance.py @@ -0,0 +1,135 @@ +"""P0-1 grant-provenance tests. + +Interactive 'Allow always' grants live in the USER-side +~/.{app}/project_grants.json keyed by resolved project path — so a cloned repo +(a different path) carries no grants, and a repo-shipped permissions.local.json +is no longer trusted. +""" + +import json +from pathlib import Path + +import pytest + +from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource +from agentic_cli.workflow.permissions.store import ( + PermissionContext, + append_project_rule, + load_project_grants, +) + + +def _ctx(workdir: Path) -> PermissionContext: + return PermissionContext(workdir=workdir, home=Path("/fake/home"), app_name="agentic") + + +class TestLoadProjectGrants: + def test_missing_file_returns_empty(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path / "home")) + assert load_project_grants("agentic", _ctx(tmp_path / "proj")) == [] + + def test_roundtrip_current_project(self, tmp_path, monkeypatch): + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.setenv("HOME", str(home)) + append_project_rule( + "agentic", + Rule("http.read", "https://ok.test/**", Effect.ALLOW, RuleSource.PROJECT), + proj, + ) + rules = load_project_grants("agentic", _ctx(proj)) + assert len(rules) == 1 + assert rules[0].effect is Effect.ALLOW + assert rules[0].source is RuleSource.PROJECT + + def test_grant_under_different_path_not_loaded(self, tmp_path, monkeypatch): + """Simulated clone: a grant recorded for project A is invisible from B.""" + home = tmp_path / "home"; a = tmp_path / "a"; b = tmp_path / "b" + a.mkdir(); b.mkdir() + monkeypatch.setenv("HOME", str(home)) + append_project_rule("agentic", Rule("http.read", "*", Effect.ALLOW, RuleSource.PROJECT), a) + assert load_project_grants("agentic", _ctx(b)) == [] + + def test_malformed_json_raises(self, tmp_path, monkeypatch): + from agentic_cli.settings_persistence import get_user_project_grants_path + monkeypatch.setenv("HOME", str(tmp_path / "home")) + p = get_user_project_grants_path("agentic") + p.parent.mkdir(parents=True) + p.write_text("{not json") + with pytest.raises(ValueError): + load_project_grants("agentic", _ctx(tmp_path / "proj")) + + def test_non_dict_grants_file_raises(self, tmp_path, monkeypatch): + """A project_grants.json containing a JSON array (not an object) must raise ValueError.""" + from agentic_cli.settings_persistence import get_user_project_grants_path + monkeypatch.setenv("HOME", str(tmp_path / "home")) + p = get_user_project_grants_path("agentic") + p.parent.mkdir(parents=True) + p.write_text("[]") + with pytest.raises(ValueError, match="Expected a JSON object"): + load_project_grants("agentic", _ctx(tmp_path / "proj")) + + +class TestLegacyLocalFileNotTrusted: + def test_repo_permissions_local_json_is_ignored(self, tmp_path, monkeypatch): + """A repo-shipped ./.{app}/permissions.local.json allow-rule is no longer + loaded as trusted (P0-1 drops the repo-local trusted load).""" + from unittest.mock import AsyncMock, MagicMock + from agentic_cli.workflow.permissions.engine import PermissionEngine + + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.chdir(proj) + monkeypatch.setenv("HOME", str(home)) + (proj / ".agentic").mkdir() + (proj / ".agentic" / "permissions.local.json").write_text(json.dumps({ + "permissions": {"allow": [{"capability": "http.read", "target": "*"}]} + })) + s = MagicMock(); s.permissions_enabled = True; s.app_name = "agentic" + w = MagicMock(); w.request_user_input = AsyncMock(return_value="Deny") + engine = PermissionEngine(settings=s, workflow=w, ctx=_ctx(proj)) + project_allows = [ + r for r in engine.rules + if r.source is RuleSource.PROJECT and r.effect is Effect.ALLOW + ] + assert project_allows == [] + + +class TestChainBlocked: + @pytest.mark.asyncio + async def test_project_cannot_preauthorize_via_settings_or_local_file( + self, tmp_path, monkeypatch + ): + """The reproduced chain: a cloned repo ships settings.json (sensitive + keys) + a permissions.local.json allow-rule. Neither pre-authorizes: + sensitive settings are dropped AND the repo allow-rule is untrusted, so a + gated tool call still reaches the approval prompt (here: user denies).""" + from unittest.mock import AsyncMock, MagicMock + from agentic_cli.config import BaseSettings + from agentic_cli.workflow.permissions.capabilities import Capability + from agentic_cli.workflow.permissions.engine import PermissionEngine + + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.chdir(proj) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.delenv("AGENTIC_STATEFUL_EXECUTOR_BACKEND", raising=False) + (proj / ".agentic_cli").mkdir() + (proj / ".agentic_cli" / "settings.json").write_text(json.dumps({ + "stateful_executor_backend": "local", + "raw_llm_logging": True, + })) + (proj / ".agentic_cli" / "permissions.local.json").write_text(json.dumps({ + "permissions": {"allow": [{"capability": "http.read", "target": "*"}]} + })) + settings = BaseSettings() # app_name default "agentic_cli" + assert settings.stateful_executor_backend == "none" # sensitive dropped + assert settings.raw_llm_logging is False + + w = MagicMock(); w.request_user_input = AsyncMock(return_value="Deny") + ctx = PermissionContext(workdir=proj, home=home, app_name="agentic_cli") + engine = PermissionEngine(settings=settings, workflow=w, ctx=ctx) + result = await engine.check( + "web_fetch", + [Capability("http.read", target_arg="url")], + {"url": "https://evil.test/x"}, + ) + assert result.allowed is False # repo allow-rule NOT trusted → prompt → denied + w.request_user_input.assert_awaited_once() diff --git a/tests/permissions/test_matchers.py b/tests/permissions/test_matchers.py index 0291023..aac967f 100644 --- a/tests/permissions/test_matchers.py +++ b/tests/permissions/test_matchers.py @@ -33,6 +33,31 @@ def test_matches_glob_star(self, ctx): assert not m.matches("foo*", "barbaz") +class TestCapMatches: + """Capability-name glob: exact / ``ns.*`` prefix / ``*``. This is what keeps + a grant for one tool from silently covering another (e.g. python.exec vs + python.exec.stateful).""" + + def test_exact_and_star(self): + from agentic_cli.workflow.permissions.matchers import _cap_matches + + assert _cap_matches("python.exec", "python.exec") is True + assert _cap_matches("*", "python.exec") is True + + def test_suffix_capability_not_covered_by_exact_rule(self): + from agentic_cli.workflow.permissions.matchers import _cap_matches + + # A grant of the exact 'python.exec' must NOT cover 'python.exec.stateful'. + assert _cap_matches("python.exec", "python.exec.stateful") is False + + def test_namespace_glob_covers_children(self): + from agentic_cli.workflow.permissions.matchers import _cap_matches + + assert _cap_matches("python.*", "python.exec") is True + assert _cap_matches("python.*", "python.exec.stateful") is True + assert _cap_matches("python.*", "shell.exec") is False + + class TestMatcherProtocol: def test_string_glob_matcher_satisfies_protocol(self): assert isinstance(StringGlobMatcher(), Matcher) diff --git a/tests/permissions/test_prompt.py b/tests/permissions/test_prompt.py index 98409e0..b000731 100644 --- a/tests/permissions/test_prompt.py +++ b/tests/permissions/test_prompt.py @@ -12,6 +12,12 @@ build_request, parse_response, ) + +#: The wording this choice used to display. Spelled out here rather than +#: imported: it is an internal parser-compatibility detail, and a test that +#: imported it would only prove the constant equals itself. What matters is that +#: this exact string, as a real user response, still means PROJECT. +SUPERSEDED_ALWAYS_RESPONSE = "Allow always (save to project)" from agentic_cli.workflow.permissions.rules import AskScope @@ -71,6 +77,48 @@ def test_request_id_has_perm_prefix(self): req = build_request("x", []) assert req.request_id.startswith("perm-") + def test_shows_code_preview_for_exec_capability(self): + """A python.exec grant is effectively allow-any-code, so the prompt must + surface the actual code being run, not just the capability.""" + req = build_request( + "execute_python", + [ResolvedCapability("python.exec", "*")], + args={"code": "import os\nos.system('rm -rf ~')"}, + ) + assert "os.system('rm -rf ~')" in req.prompt + + def test_shows_code_preview_for_stateful_exec_capability(self): + """sandbox_execute carries python.exec.stateful (a sub-namespace of + python.exec) — it is MORE dangerous (persistent kernel), so its code must + also be shown at approval, not just python.exec's.""" + req = build_request( + "sandbox_execute", + [ResolvedCapability("python.exec.stateful", "*")], + args={"code": "import os\nos.system('rm -rf ~')"}, + ) + assert "os.system('rm -rf ~')" in req.prompt + + def test_long_code_preview_is_truncated(self): + req = build_request( + "execute_python", + [ResolvedCapability("python.exec", "*")], + args={"code": "x = 1\n" * 500}, + ) + assert len(req.prompt) < 2000 + + def test_no_code_preview_for_non_exec_capability(self): + """A 'code'-shaped arg on a non-exec tool must not trigger a preview.""" + req = build_request( + "write_file", + [ResolvedCapability("filesystem.write", "/foo")], + args={"code": "not actually executed"}, + ) + assert "not actually executed" not in req.prompt + + def test_backward_compatible_without_args(self): + req = build_request("write_file", [ResolvedCapability("filesystem.write", "/foo")]) + assert "write_file" in req.prompt + class TestParseResponse: @pytest.mark.parametrize("text, scope", [ @@ -82,6 +130,27 @@ class TestParseResponse: def test_round_trip(self, text, scope): assert parse_response(text) is scope + def test_displayed_always_choice_names_the_scope_not_a_file(self): + """The label describes scope; it must not imply a repo write. + + Grants go to ``~/.{app}/project_grants.json``, never into the project + directory, so a label saying "save to project" mis-described where the + rule lands. + """ + assert ALLOW_ALWAYS_CHOICE == "Allow always for this project" + assert "save to project" not in ALLOW_ALWAYS_CHOICE + + def test_superseded_always_response_still_parses_as_project(self): + """Renaming the choice must not turn an old "always" into a denial.""" + assert SUPERSEDED_ALWAYS_RESPONSE != ALLOW_ALWAYS_CHOICE + assert parse_response(SUPERSEDED_ALWAYS_RESPONSE) is AskScope.PROJECT + + def test_superseded_wording_is_never_offered(self): + """Accepted for compatibility, but never displayed.""" + request = build_request("t", [ResolvedCapability("http.read", "https://x.test")]) + assert SUPERSEDED_ALWAYS_RESPONSE not in request.choices + assert ALLOW_ALWAYS_CHOICE in request.choices + def test_unknown_defaults_to_deny(self): assert parse_response("whatever") is AskScope.DENY diff --git a/tests/permissions/test_store.py b/tests/permissions/test_store.py index 5f52707..67d0774 100644 --- a/tests/permissions/test_store.py +++ b/tests/permissions/test_store.py @@ -116,64 +116,82 @@ def test_malformed_json_raises(self, tmp_path: Path): class TestAppendProjectRule: - def test_creates_file_when_absent(self, tmp_path, monkeypatch): + """Interactive 'Allow always' grants persist to the USER-side, path-keyed + ~/.{app}/project_grants.json (P0-1) — never into the repo, so a clone + carries no grants.""" + + def _grants_path(self, home: Path, app: str = "agentic") -> Path: + return home / f".{app}" / "project_grants.json" + + def test_creates_user_grants_file_keyed_by_project(self, tmp_path, monkeypatch): + import json from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource from agentic_cli.workflow.permissions.store import append_project_rule - monkeypatch.chdir(tmp_path) + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.setenv("HOME", str(home)) rule = Rule("filesystem.write", "/abs/foo", Effect.ALLOW, RuleSource.PROJECT) - append_project_rule("agentic", rule) + append_project_rule("agentic", rule, proj) - import json - data = json.loads((tmp_path / ".agentic/settings.json").read_text()) - assert data["permissions"]["allow"] == [ + data = json.loads(self._grants_path(home).read_text()) + assert data[str(proj.resolve())]["permissions"]["allow"] == [ {"capability": "filesystem.write", "target": "/abs/foo"} ] - def test_preserves_other_settings_keys(self, tmp_path, monkeypatch): - import json + def test_writes_nothing_into_project_dir(self, tmp_path, monkeypatch): from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource from agentic_cli.workflow.permissions.store import append_project_rule - monkeypatch.chdir(tmp_path) - (tmp_path / ".agentic").mkdir() - (tmp_path / ".agentic/settings.json").write_text(json.dumps({ - "default_model": "claude-sonnet-4", - "thinking_effort": "medium", - })) + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.setenv("HOME", str(home)) + append_project_rule( + "agentic", + Rule("filesystem.write", "/abs/foo", Effect.ALLOW, RuleSource.PROJECT), + proj, + ) + assert not (proj / ".agentic").exists() # nothing dropped inside the repo + + def test_deduplicates_identical_rules(self, tmp_path, monkeypatch): + import json + from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource + from agentic_cli.workflow.permissions.store import append_project_rule + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.setenv("HOME", str(home)) rule = Rule("filesystem.write", "/abs/foo", Effect.ALLOW, RuleSource.PROJECT) - append_project_rule("agentic", rule) + append_project_rule("agentic", rule, proj) + append_project_rule("agentic", rule, proj) - data = json.loads((tmp_path / ".agentic/settings.json").read_text()) - assert data["default_model"] == "claude-sonnet-4" - assert data["thinking_effort"] == "medium" - assert data["permissions"]["allow"][0]["capability"] == "filesystem.write" + data = json.loads(self._grants_path(home).read_text()) + assert len(data[str(proj.resolve())]["permissions"]["allow"]) == 1 - def test_deduplicates_identical_rules(self, tmp_path, monkeypatch): + def test_two_projects_kept_separate(self, tmp_path, monkeypatch): import json from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource from agentic_cli.workflow.permissions.store import append_project_rule - monkeypatch.chdir(tmp_path) - rule = Rule("filesystem.write", "/abs/foo", Effect.ALLOW, RuleSource.PROJECT) - append_project_rule("agentic", rule) - append_project_rule("agentic", rule) + home = tmp_path / "home"; a = tmp_path / "a"; b = tmp_path / "b" + a.mkdir(); b.mkdir() + monkeypatch.setenv("HOME", str(home)) + append_project_rule("agentic", Rule("http.read", "*", Effect.ALLOW, RuleSource.PROJECT), a) + append_project_rule("agentic", Rule("filesystem.write", "/x", Effect.ALLOW, RuleSource.PROJECT), b) - data = json.loads((tmp_path / ".agentic/settings.json").read_text()) - assert len(data["permissions"]["allow"]) == 1 + data = json.loads(self._grants_path(home).read_text()) + assert set(data.keys()) == {str(a.resolve()), str(b.resolve())} def test_writes_deny_section_for_deny_effect(self, tmp_path, monkeypatch): import json from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource from agentic_cli.workflow.permissions.store import append_project_rule - monkeypatch.chdir(tmp_path) - rule = Rule("filesystem.write", "/etc/foo", Effect.DENY, RuleSource.PROJECT) - append_project_rule("agentic", rule) - - data = json.loads((tmp_path / ".agentic/settings.json").read_text()) - assert data["permissions"]["deny"] == [ + home = tmp_path / "home"; proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.setenv("HOME", str(home)) + append_project_rule( + "agentic", + Rule("filesystem.write", "/etc/foo", Effect.DENY, RuleSource.PROJECT), + proj, + ) + data = json.loads(self._grants_path(home).read_text()) + assert data[str(proj.resolve())]["permissions"]["deny"] == [ {"capability": "filesystem.write", "target": "/etc/foo"} ] - assert "allow" not in data["permissions"] or data["permissions"]["allow"] == [] diff --git a/tests/permissions/test_trust_model.py b/tests/permissions/test_trust_model.py new file mode 100644 index 0000000..907ee4c --- /dev/null +++ b/tests/permissions/test_trust_model.py @@ -0,0 +1,180 @@ +"""Trust-model regression tests (P0-2, P0-3). + +The project ``./.{app}/settings.json`` is untrusted — a cloned repo can ship it. +It must not be able to (a) disable the permission engine, or (b) grant +allow-rules; and the agent must not be able to write the app's own config to +self-escalate. See docs/reviews/2026-07-03-*. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agentic_cli.workflow.permissions.capabilities import Capability +from agentic_cli.workflow.permissions.engine import PermissionEngine +from agentic_cli.workflow.permissions.rules import Effect, Rule, RuleSource +from agentic_cli.workflow.permissions.store import BUILTIN_RULES, PermissionContext + + +def _stub_settings(*, enabled: bool = True, app_name: str = "myapp") -> MagicMock: + s = MagicMock() + s.permissions_enabled = enabled + s.app_name = app_name + return s + + +def _stub_workflow(response: str = "Deny") -> MagicMock: + w = MagicMock() + w.request_user_input = AsyncMock(return_value=response) + return w + + +def _engine(workdir: Path, home: Path, app: str = "myapp", response: str = "Deny") -> PermissionEngine: + ctx = PermissionContext(workdir=workdir, home=home, app_name=app) + return PermissionEngine(_stub_settings(app_name=app), _stub_workflow(response), ctx) + + +# --------------------------------------------------------------------------- +# P0-2: project settings can tighten (deny) but never loosen (allow / disable) +# --------------------------------------------------------------------------- + +class TestProjectRuleTrust: + def _dirs(self, tmp_path, monkeypatch): + proj, home = tmp_path / "proj", tmp_path / "home" + proj.mkdir() + home.mkdir() + monkeypatch.chdir(proj) + monkeypatch.setenv("HOME", str(home)) + return proj, home + + def test_project_allow_rules_are_ignored(self, tmp_path, monkeypatch): + proj, home = self._dirs(tmp_path, monkeypatch) + (proj / ".myapp").mkdir() + (proj / ".myapp/settings.json").write_text(json.dumps({ + "permissions": {"allow": [{"capability": "http.read", "target": "*"}]}, + })) + engine = _engine(proj, home) + project_allows = [ + r for r in engine.rules + if r.source is RuleSource.PROJECT and r.effect is Effect.ALLOW + ] + assert project_allows == [] + + def test_project_deny_rules_are_honored(self, tmp_path, monkeypatch): + proj, home = self._dirs(tmp_path, monkeypatch) + (proj / ".myapp").mkdir() + (proj / ".myapp/settings.json").write_text(json.dumps({ + "permissions": {"deny": [{"capability": "http.read", "target": "https://evil.test/**"}]}, + })) + engine = _engine(proj, home) + project_denies = [ + r for r in engine.rules + if r.source is RuleSource.PROJECT and r.effect is Effect.DENY + ] + assert len(project_denies) == 1 + + def test_user_allow_rules_still_honored(self, tmp_path, monkeypatch): + """User config is trusted — it CAN grant allow-rules.""" + proj, home = self._dirs(tmp_path, monkeypatch) + (home / ".myapp").mkdir() + (home / ".myapp/settings.json").write_text(json.dumps({ + "permissions": {"allow": [{"capability": "http.read", "target": "*"}]}, + })) + engine = _engine(proj, home) + user_allows = [ + r for r in engine.rules + if r.source is RuleSource.USER and r.effect is Effect.ALLOW + ] + assert len(user_allows) == 1 + + +class TestProjectCannotDisablePermissions: + def test_project_settings_json_cannot_set_permissions_enabled( + self, tmp_path, monkeypatch + ): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_PERMISSIONS_ENABLED", raising=False) + (tmp_path / ".agentic_cli").mkdir() + (tmp_path / ".agentic_cli/settings.json").write_text( + json.dumps({"permissions_enabled": False}) + ) + from agentic_cli.config import BaseSettings + + settings = BaseSettings() + assert settings.permissions_enabled is True + + +# --------------------------------------------------------------------------- +# P0-3: the agent cannot write the app's own config to self-escalate +# --------------------------------------------------------------------------- + +class TestConfigWriteDeny: + def test_context_substitutes_app_name(self, tmp_path): + ctx = PermissionContext(workdir=tmp_path, home=Path("/home"), app_name="myapp") + assert ctx.substitute("${workdir}/.${app_name}/**") == f"{tmp_path}/.myapp/**" + + def test_builtin_rules_deny_app_config_writes(self): + config_denies = [ + r for r in BUILTIN_RULES + if r.effect is Effect.DENY + and r.capability == "filesystem.write" + and "${app_name}" in r.target + ] + # Both ${workdir} and ${home} config dirs must be covered. + assert any("${workdir}" in r.target for r in config_denies) + assert any("${home}" in r.target for r in config_denies) + + @pytest.mark.asyncio + async def test_write_to_project_config_denied(self, tmp_path): + engine = _engine(tmp_path, tmp_path / "home", response="Allow always") + target = str(tmp_path / ".myapp" / "settings.json") + result = await engine.check( + "write_file", + [Capability("filesystem.write", target_arg="path")], + {"path": target}, + ) + assert result.allowed is False + + @pytest.mark.asyncio + async def test_write_to_user_config_denied(self, tmp_path): + home = tmp_path / "home" + engine = _engine(tmp_path, home, response="Allow always") + target = str(home / ".myapp" / "settings.json") + result = await engine.check( + "write_file", + [Capability("filesystem.write", target_arg="path")], + {"path": target}, + ) + assert result.allowed is False + + @pytest.mark.asyncio + async def test_config_denied_even_with_broadened_workdir_grant(self, tmp_path): + """A prior 'Allow always' that broadened to /** must not reach + the config path — deny-wins.""" + engine = _engine(tmp_path, tmp_path / "home") + engine._session_rules.append( + Rule("filesystem.write", f"{tmp_path}/**", Effect.ALLOW, RuleSource.SESSION) + ) + target = str(tmp_path / ".myapp" / "settings.json") + result = await engine.check( + "write_file", + [Capability("filesystem.write", target_arg="path")], + {"path": target}, + ) + assert result.allowed is False + + @pytest.mark.asyncio + async def test_ordinary_workdir_write_not_overblocked(self, tmp_path): + engine = _engine(tmp_path, tmp_path / "home", response="Allow once") + target = str(tmp_path / "output.txt") + result = await engine.check( + "write_file", + [Capability("filesystem.write", target_arg="path")], + {"path": target}, + ) + assert result.allowed is True diff --git a/tests/test_adk_plugins.py b/tests/test_adk_plugins.py index 8530e83..ced4be8 100644 --- a/tests/test_adk_plugins.py +++ b/tests/test_adk_plugins.py @@ -80,8 +80,9 @@ class TestLLMLoggingPlugin: @pytest.fixture def plugin(self, tmp_path, monkeypatch): - """Create a plugin with log directory in tmp_path.""" + """Create a plugin with log directory under a temp HOME.""" monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path)) return LLMLoggingPlugin( model_name="test-model", app_name="test_app", @@ -92,11 +93,39 @@ def test_plugin_name(self, plugin): assert plugin.name == "llm_logging" def test_log_file_created(self, plugin, tmp_path): - """Log directory and file path should be set up.""" + """Log directory and file path should be set up (under HOME).""" expected_dir = tmp_path / ".test_app" / "logs" assert expected_dir.exists() assert plugin.get_log_file_path() == expected_dir / "llm_events.jsonl" + def test_log_dir_under_home_not_cwd(self, tmp_path, monkeypatch): + """Conversation logs must live under the user's home, not the (often + git-tracked) project directory (P0-6).""" + home, proj = tmp_path / "home", tmp_path / "proj" + home.mkdir() + proj.mkdir() + monkeypatch.chdir(proj) + monkeypatch.setenv("HOME", str(home)) + plugin = LLMLoggingPlugin(app_name="test_app") + log_path = str(plugin.get_log_file_path()) + assert str(home) in log_path + assert str(proj) not in log_path + + def test_log_file_has_restrictive_permissions(self, plugin): + """Full conversations may contain secrets — the file must be 0600.""" + import stat + + mode = stat.S_IMODE(plugin.get_log_file_path().stat().st_mode) + assert mode == 0o600 + + def test_secrets_redacted_in_written_log(self, plugin): + from agentic_cli.workflow.adk.plugins import _redact_secrets + + line = "authorization: Bearer sk-ant-api03-ABCDEFGH1234567890ZYXW" + redacted = _redact_secrets(line) + assert "sk-ant-api03-ABCDEFGH1234567890ZYXW" not in redacted + assert "REDACTED" in redacted + async def test_before_model_captures_request(self, plugin): """before_model_callback should capture request data and buffer an event.""" ctx = _make_mock_callback_context(invocation_id="inv-001") diff --git a/tests/test_config.py b/tests/test_config.py index 054733e..a1338ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -248,6 +248,41 @@ def test_export_api_keys_to_env(self, temp_workspace: Path): else: os.environ.pop("GOOGLE_API_KEY", None) + def test_export_overwrites_stale_env_key(self, temp_workspace: Path, monkeypatch): + """Export re-asserts this instance's key over a mutated env var. + + The key fields bind only via their env alias, so settings and env + diverge when the process env changed after this instance loaded — + exactly what an earlier manager's export does in a two-settings + process. The previous set-if-absent guard then silently kept the + first manager's credentials for every later one. + """ + monkeypatch.setenv("ANTHROPIC_API_KEY", "current-key") + settings = BaseSettings(workspace_dir=temp_workspace) + assert settings.anthropic_api_key == "current-key" + + # Another settings instance exported its key in the meantime + monkeypatch.setenv("ANTHROPIC_API_KEY", "stale-key") + + settings.export_api_keys_to_env() + + assert os.environ["ANTHROPIC_API_KEY"] == "current-key" + + def test_export_leaves_env_when_setting_unset( + self, temp_workspace: Path, monkeypatch, tmp_path: Path + ): + """A key unset in settings never clobbers an externally-set env var.""" + monkeypatch.setenv("HOME", str(tmp_path)) # no user config leakage + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + settings = BaseSettings(workspace_dir=temp_workspace) + assert settings.google_api_key is None + + monkeypatch.setenv("GOOGLE_API_KEY", "external-key") + + settings.export_api_keys_to_env() + + assert os.environ["GOOGLE_API_KEY"] == "external-key" + class TestSettingsContext: """Tests for context-based settings management.""" diff --git a/tests/test_config_trust.py b/tests/test_config_trust.py new file mode 100644 index 0000000..f7664d2 --- /dev/null +++ b/tests/test_config_trust.py @@ -0,0 +1,250 @@ +"""P0-1 config trust-model tests. + +A project ``./.{app}/settings.json`` (and a cwd-relative ``.env``) may set only +non-security allowlisted keys; sensitive keys are dropped with a warning. Real +environment variables and user ``~/.{app}/settings.json`` stay trusted. +""" + +import json +from pathlib import Path + + +def _write_project_settings(root: Path, app: str, data: dict) -> None: + d = root / f".{app}" + d.mkdir(parents=True, exist_ok=True) + (d / "settings.json").write_text(json.dumps(data)) + + +class TestProjectSettingsAllowlist: + def test_allowlisted_key_is_applied(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_DEFAULT_MODEL", raising=False) + _write_project_settings(tmp_path, "agentic_cli", {"default_model": "claude-sonnet-4-6"}) + from agentic_cli.config import BaseSettings + assert BaseSettings().default_model == "claude-sonnet-4-6" + + def test_sensitive_keys_are_dropped(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for var in ("AGENTIC_STATEFUL_EXECUTOR_BACKEND", "AGENTIC_RAW_LLM_LOGGING"): + monkeypatch.delenv(var, raising=False) + _write_project_settings(tmp_path, "agentic_cli", { + "stateful_executor_backend": "local", + "raw_llm_logging": True, + "workspace_dir": "/tmp/evil", + "sandbox_data_mounts": ["/etc:etc"], + "default_model": "kept", + }) + from agentic_cli.config import BaseSettings + s = BaseSettings() + # sensitive → dropped (defaults preserved) + assert s.stateful_executor_backend == "none" + assert s.raw_llm_logging is False + assert s.sandbox_data_mounts == [] + assert str(s.workspace_dir) != "/tmp/evil" + # benign → applied + assert s.default_model == "kept" + + def test_user_config_sensitive_key_is_trusted(self, tmp_path, monkeypatch): + proj = tmp_path / "proj"; proj.mkdir() + monkeypatch.chdir(proj) # cwd has no project settings.json + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_RAW_LLM_LOGGING", raising=False) + ud = tmp_path / "home" / ".agentic_cli" + ud.mkdir(parents=True) + (ud / "settings.json").write_text(json.dumps({"raw_llm_logging": True})) + from agentic_cli.config import BaseSettings + assert BaseSettings().raw_llm_logging is True + + def test_real_env_var_sensitive_key_is_trusted(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.setenv("AGENTIC_STATEFUL_EXECUTOR_BACKEND", "local") + from agentic_cli.config import BaseSettings + assert BaseSettings().stateful_executor_backend == "local" + + def test_dropped_key_logs_warning(self, tmp_path, monkeypatch): + import structlog + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + _write_project_settings(tmp_path, "agentic_cli", {"raw_llm_logging": True}) + from agentic_cli.config import BaseSettings + with structlog.testing.capture_logs() as logs: + BaseSettings() + assert any( + e.get("event") == "untrusted_project_setting_ignored" + and e.get("key") == "raw_llm_logging" + for e in logs + ) + + +class TestCwdDotenvFiltering: + def _subclass_with_env_file(self, env_file): + from agentic_cli.config import BaseSettings + from pydantic_settings import SettingsConfigDict + + class _DomainSettings(BaseSettings): + model_config = SettingsConfigDict( + env_prefix="AGENTIC_", + env_file=env_file, + env_file_encoding="utf-8", + env_nested_delimiter="__", + extra="ignore", + ) + + return _DomainSettings + + def test_cwd_relative_env_drops_sensitive_key(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for v in ("AGENTIC_RAW_LLM_LOGGING", "AGENTIC_DEFAULT_MODEL"): + monkeypatch.delenv(v, raising=False) + (tmp_path / ".env").write_text( + "AGENTIC_RAW_LLM_LOGGING=true\nAGENTIC_DEFAULT_MODEL=envmodel\n" + ) + s = self._subclass_with_env_file(".env")() + assert s.raw_llm_logging is False # sensitive dropped + assert s.default_model == "envmodel" # benign kept + + def test_absolute_env_file_is_trusted(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_RAW_LLM_LOGGING", raising=False) + abs_env = tmp_path / "user.env" + abs_env.write_text("AGENTIC_RAW_LLM_LOGGING=true\n") + s = self._subclass_with_env_file(str(abs_env))() + assert s.raw_llm_logging is True # absolute env_file trusted + + def test_list_env_file_is_trusted(self, tmp_path, monkeypatch): + # A list/tuple env_file must stay trusted (unfiltered) — only a single + # cwd-relative env_file is filtered. + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_RAW_LLM_LOGGING", raising=False) + abs_env = tmp_path / "user.env" + abs_env.write_text("AGENTIC_RAW_LLM_LOGGING=true\n") + s = self._subclass_with_env_file([str(abs_env)])() + assert s.raw_llm_logging is True # list env_file trusted + + def test_tilde_env_file_is_trusted(self, tmp_path, monkeypatch): + # A "~"-prefixed env_file is a user-level path (pydantic expands ~ when + # reading it), so it must stay TRUSTED (unfiltered). + fake_home = tmp_path / "home" + fake_home.mkdir() + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(fake_home)) + monkeypatch.delenv("AGENTIC_RAW_LLM_LOGGING", raising=False) + (fake_home / "user.env").write_text("AGENTIC_RAW_LLM_LOGGING=true\n") + s = self._subclass_with_env_file("~/user.env")() + assert s.raw_llm_logging is True # tilde (user-level) env_file trusted + + def test_list_env_file_with_cwd_relative_entry_is_filtered(self, tmp_path, monkeypatch): + # If ANY entry in a list env_file is cwd-relative, the whole dotenv + # source is filtered (fail-safe over-filter) — a repo could otherwise + # ship ./.env alongside a trusted absolute file and stay unfiltered. + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("AGENTIC_RAW_LLM_LOGGING", raising=False) + abs_env = tmp_path / "abs.env" + abs_env.write_text("") # trusted absolute entry (empty) + (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 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 # --------------------------------------------------------------------------- diff --git a/tests/test_executor_security.py b/tests/test_executor_security.py index 54d7bb1..58d7745 100644 --- a/tests/test_executor_security.py +++ b/tests/test_executor_security.py @@ -6,9 +6,37 @@ in-process restrictions, so they must be unavailable unless os_sandbox is on. """ +import shlex from types import SimpleNamespace +import pytest + from agentic_cli.tools.executor import SafePythonExecutor +from agentic_cli.tools.shell.os_sandbox.policy import OSSandboxPolicy + + +class _FakeSandbox: + """Stand-in OS sandbox. ``sandbox_type='none'`` = unavailable; otherwise the + wrap just runs the plain command (no real isolation — we only test wiring).""" + + def __init__(self, sandbox_type: str = "seatbelt") -> None: + self._type = sandbox_type + + @property + def sandbox_type(self) -> str: + return self._type + + def wrap_python_command(self, argv, working_dir, policy): + from agentic_cli.tools.shell.os_sandbox.base import OSSandboxResult + + return OSSandboxResult(command=shlex.join(argv), sandbox_type=self._type) + + +def _patch_sandbox(monkeypatch, sandbox_type: str) -> None: + monkeypatch.setattr( + "agentic_cli.tools.shell.os_sandbox.get_os_sandbox", + lambda *a, **k: _FakeSandbox(sandbox_type), + ) class TestModuleGatingWithoutSandbox: @@ -53,17 +81,17 @@ def test_core_module_still_works(self): class TestModuleGatingWithSandbox: - """OS sandbox enabled -> the heavy libraries are permitted again.""" + """OS sandbox enabled AND available -> the heavy libraries are permitted.""" - def test_effective_set_is_full_when_sandbox_enabled(self): - policy = SimpleNamespace(enabled=True) - executor = SafePythonExecutor(os_sandbox_policy=policy) + def test_effective_set_is_full_when_sandbox_available(self, monkeypatch): + _patch_sandbox(monkeypatch, "seatbelt") + executor = SafePythonExecutor(os_sandbox_policy=OSSandboxPolicy(enabled=True)) assert executor.effective_allowed_modules == executor.ALLOWED_MODULES assert "numpy" in executor.effective_allowed_modules - def test_numpy_import_allowed_when_sandbox_enabled(self): - policy = SimpleNamespace(enabled=True) - executor = SafePythonExecutor(os_sandbox_policy=policy) + def test_numpy_import_allowed_when_sandbox_available(self, monkeypatch): + _patch_sandbox(monkeypatch, "seatbelt") + executor = SafePythonExecutor(os_sandbox_policy=OSSandboxPolicy(enabled=True)) ok, _ = executor.validate_code("import numpy") assert ok is True @@ -71,3 +99,56 @@ def test_disabled_policy_is_treated_as_no_sandbox(self): policy = SimpleNamespace(enabled=False) executor = SafePythonExecutor(os_sandbox_policy=policy) assert executor.effective_allowed_modules == executor.CORE_MODULES + + +class TestOsSandboxAvailabilityAware: + """Enabling the sandbox but having no real backend must NOT re-enable the + heavy modules — availability, not the flag, gates SANDBOXED_MODULES (5b).""" + + def test_enabled_but_unavailable_is_core_only(self, monkeypatch): + _patch_sandbox(monkeypatch, "none") + ex = SafePythonExecutor(os_sandbox_policy=OSSandboxPolicy(enabled=True)) + assert ex.effective_allowed_modules == ex.CORE_MODULES + assert "numpy" not in ex.effective_allowed_modules + + def test_unavailable_non_strict_falls_back_and_runs(self, monkeypatch): + """Default (non-strict): no backend -> hardened in-process, not refusal.""" + _patch_sandbox(monkeypatch, "none") + ex = SafePythonExecutor( + os_sandbox_policy=OSSandboxPolicy(enabled=True, strict=False) + ) + # numpy still rejected (CORE only) ... + assert ex.execute("import numpy")["success"] is False + # ... but ordinary compute runs rather than being refused. + result = ex.execute("import math\nprint(math.sqrt(16))") + assert result["success"] is True, result + assert "4.0" in result["output"] + + def test_unavailable_strict_refuses(self, monkeypatch): + _patch_sandbox(monkeypatch, "none") + ex = SafePythonExecutor( + os_sandbox_policy=OSSandboxPolicy(enabled=True, strict=True) + ) + result = ex.execute("1 + 1") + assert result["success"] is False + assert "sandbox" in result["error"].lower() + + def test_available_runs_under_sandbox(self, monkeypatch): + _patch_sandbox(monkeypatch, "seatbelt") + ex = SafePythonExecutor(os_sandbox_policy=OSSandboxPolicy(enabled=True)) + result = ex.execute("import math\nprint(math.sqrt(16))") + assert result["success"] is True, result + + +class TestOsSandboxSettingsDefaults: + def test_os_sandbox_enabled_defaults_true(self): + from agentic_cli.config import BaseSettings + + settings = BaseSettings() + assert settings.os_sandbox_enabled is True + + def test_os_sandbox_strict_defaults_false(self): + from agentic_cli.config import BaseSettings + + settings = BaseSettings() + assert settings.os_sandbox_strict is False diff --git a/tests/test_file_utils_nofollow.py b/tests/test_file_utils_nofollow.py new file mode 100644 index 0000000..2d3b977 --- /dev/null +++ b/tests/test_file_utils_nofollow.py @@ -0,0 +1,50 @@ +import os + +import pytest + +from agentic_cli.file_utils import copy_regular_file_no_follow + + +def test_copies_regular_file(tmp_path): + src = tmp_path / "a.txt"; src.write_bytes(b"hello") + dst = tmp_path / "out" / "a.txt" + copy_regular_file_no_follow(src, dst) + assert dst.read_bytes() == b"hello" + assert not dst.is_symlink() + + +def test_rejects_symlink_source(tmp_path): + real = tmp_path / "real.txt"; real.write_bytes(b"secret") + link = tmp_path / "link.txt"; link.symlink_to(real) + with pytest.raises(OSError): # O_NOFOLLOW -> ELOOP on the final component + copy_regular_file_no_follow(link, tmp_path / "out.txt") + + +def test_rejects_fifo_source_without_hanging(tmp_path): + import threading + fifo = tmp_path / "pipe"; os.mkfifo(fifo) + result = {} + def run(): + try: + copy_regular_file_no_follow(fifo, tmp_path / "out.txt") + except (OSError, ValueError): + result["rejected"] = True + t = threading.Thread(target=run, daemon=True); t.start(); t.join(5) + assert not t.is_alive(), "copy hung on a FIFO source (O_NONBLOCK missing?)" + assert result.get("rejected") + + +def test_does_not_follow_dest_symlink(tmp_path): + src = tmp_path / "src.txt"; src.write_bytes(b"NEW") + outside = tmp_path / "outside.txt"; outside.write_bytes(b"ORIGINAL") + deliver = tmp_path / "deliver"; deliver.mkdir() + dst = deliver / "x.txt"; dst.symlink_to(outside) + copy_regular_file_no_follow(src, dst) + assert outside.read_bytes() == b"ORIGINAL" # symlink target untouched + assert dst.read_bytes() == b"NEW" # dst replaced by a real file + assert not dst.is_symlink() + + +def test_missing_source_raises_filenotfound(tmp_path): + with pytest.raises(FileNotFoundError): + copy_regular_file_no_follow(tmp_path / "nope.txt", tmp_path / "out.txt") diff --git a/tests/test_grep_tool_security.py b/tests/test_grep_tool_security.py index 75fff5a..df08b30 100644 --- a/tests/test_grep_tool_security.py +++ b/tests/test_grep_tool_security.py @@ -6,24 +6,32 @@ """ from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import patch from agentic_cli.tools.grep_tool import grep +class _FakeProc: + pid = 4321 + returncode = 0 + + def wait(self, timeout=None): + return 0 + + def _run_grep_capturing_argv(tmp_path: Path, pattern: str, rg_stdout: str = ""): """Call grep() forcing the ripgrep path and capture the argv built.""" captured = {} - def fake_run(cmd, *args, **kwargs): + def fake_popen(cmd, *args, **kwargs): captured["cmd"] = cmd - result = MagicMock() - result.stdout = rg_stdout - result.returncode = 0 - return result + out = kwargs.get("stdout") + if out is not None and rg_stdout: + out.write(rg_stdout.encode()) + return _FakeProc() with patch("agentic_cli.tools.grep_tool._ripgrep_available", return_value=True), \ - patch("agentic_cli.tools.grep_tool.subprocess.run", side_effect=fake_run): + patch("agentic_cli.tools.grep_tool.subprocess.Popen", side_effect=fake_popen): out = grep(pattern=pattern, path=str(tmp_path)) return captured["cmd"], out 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 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_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_langgraph.py b/tests/test_langgraph.py index b8172e3..77aa024 100644 --- a/tests/test_langgraph.py +++ b/tests/test_langgraph.py @@ -112,7 +112,7 @@ def settings(self): return BaseSettings( google_api_key="test-key", orchestrator="langgraph", - langgraph_checkpointer="memory", + session_store="memory", ) @pytest.fixture @@ -285,13 +285,13 @@ def test_settings_langgraph_orchestrator(self): ) assert settings.orchestrator == "langgraph" - def test_settings_langgraph_checkpointer(self): - """Test LangGraph checkpointer setting.""" + def test_settings_session_store(self): + """The unified session_store drives LangGraph persistence.""" settings = BaseSettings( google_api_key="test-key", - langgraph_checkpointer="postgres", + session_store="postgres", ) - assert settings.langgraph_checkpointer == "postgres" + assert settings.session_store == "postgres" class TestLangGraphThinkingConfig: @@ -391,7 +391,9 @@ def test_thinking_config_google_model(self, agent_configs): assert config is not None assert config["provider"] == "google" assert config["include_thoughts"] is True - assert config["thinking_level"] == "high" + # 2.5 uses a numeric budget, NOT thinking_level (which it rejects, 400). + assert config["thinking_budget"] == 24576 + assert "thinking_level" not in config def test_thinking_config_google_medium(self, agent_configs): """Test _get_thinking_config with medium effort for Google.""" diff --git a/tests/test_message_processor_render.py b/tests/test_message_processor_render.py index 30343c6..e5db9ef 100644 --- a/tests/test_message_processor_render.py +++ b/tests/test_message_processor_render.py @@ -78,6 +78,23 @@ def _index_of(ui: RecordingSession, predicate) -> int: return -1 +class TestSessionThreading: + async def test_session_id_threaded_to_workflow(self): + """The app's session id must reach workflow.process so unnamed runs + don't collapse into the manager's fallback 'default_session' (A-1).""" + mp = MessageProcessor() + wf = ReplayWorkflow([WorkflowEvent.text("ok")]) + ctrl = ReplayController(wf) + await mp.process( + message="hi", + workflow_controller=ctrl, + ui=RecordingSession(), + settings=_settings(), + session_id="sess-abc123", + ) + assert wf.received_session_id == "sess-abc123" + + class TestRenderBasics: async def test_text_response_rendered_as_markdown(self): ui, _, _ = await _render([WorkflowEvent.text("hello **world**")]) @@ -194,3 +211,72 @@ async def test_events_survive_serialization(self): assert ui_direct.kinds() == ui_replayed.kinds() assert ui_direct.responses() == ui_replayed.responses() + + +class TestSandboxExecuteRendering: + """sandbox_execute renders a code block on the call and a single combined + result message (header + output) on the result.""" + + async def test_code_block_and_combined_result_rendered(self): + events = [ + WorkflowEvent.tool_call("sandbox_execute", {"code": "x = 41\nprint(x + 1)"}), + WorkflowEvent.tool_result( + "sandbox_execute", + {"success": True, "stdout": "RESULT-42\n"}, + success=True, + duration_ms=12, + ), + ] + ui, _, _ = await _render(events) + rich = ui.rich() + # code block carries the header and the code itself + assert any("Running Python in the stateful executor" in r for r in rich), rich + assert any("x = 41" in r for r in rich), rich + # combined result: fixed header label + output indented below + assert any("Stateful Python executor output:" in r for r in rich), rich + assert any("RESULT-42" in r for r in rich), rich + # exactly two messages: code block (on call) + combined result (on result) + assert len(rich) == 2, rich + + async def test_result_is_single_message_without_output_when_stdout_empty(self): + events = [ + WorkflowEvent.tool_call("sandbox_execute", {"code": "y = 5"}), + WorkflowEvent.tool_result( + "sandbox_execute", + {"success": True, "stdout": ""}, + success=True, + duration_ms=5, + ), + ] + ui, _, _ = await _render(events) + rich = ui.rich() + assert len(rich) == 2, rich + # with no output the result is a single "no output" line + result_msg = [r for r in rich if r.lstrip().startswith("+")][0] + assert result_msg.strip() == "+ Stateful Python executor: no output", result_msg + + async def test_failure_surfaces_error(self): + events = [ + WorkflowEvent.tool_call("sandbox_execute", {"code": "boom()"}), + WorkflowEvent.tool_result( + "sandbox_execute", + {"success": False, "error": "NameError: boom", "stdout": ""}, + success=False, + ), + ] + ui, _, _ = await _render(events) + rich = ui.rich() + assert any("NameError: boom" in r for r in rich), rich + + async def test_non_sandbox_tool_call_emits_no_code_block(self): + """The code-block behavior is scoped to sandbox_execute; a normal tool + call still renders exactly one rich call (its result summary).""" + events = [ + WorkflowEvent.tool_call("read_file", {"path": "/tmp/x"}), + WorkflowEvent.tool_result( + "read_file", {"success": True, "content": "hi", "path": "/tmp/x", "size": 2}, + success=True, + ), + ] + ui, _, _ = await _render(events) + assert len(ui.rich()) == 1, ui.rich() 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/test_packaging.py b/tests/test_packaging.py new file mode 100644 index 0000000..e076ac2 --- /dev/null +++ b/tests/test_packaging.py @@ -0,0 +1,99 @@ +"""Packaging consistency regression pins. + +A clean ``pip install`` uses only pyproject.toml (not environment.yml), so an +eagerly-imported third-party module missing from ``[project.dependencies]`` +breaks on first import. This pins the specific modules that regressed; it is +not an exhaustive import-vs-dependency audit. +""" +from __future__ import annotations + +import tomllib +from pathlib import Path + +_PYPROJECT = Path(__file__).resolve().parent.parent / "pyproject.toml" + + +def _declared_dependencies() -> list[str]: + data = tomllib.loads(_PYPROJECT.read_text()) + return list(data.get("project", {}).get("dependencies", [])) + + +def _dep_names() -> set[str]: + """Normalized distribution names from the dependency specifiers.""" + names = set() + for spec in _declared_dependencies(): + # Strip version/extras/markers: name is the leading run of allowed chars. + name = spec.split(";")[0].strip() + for sep in ("[", ">", "<", "=", "!", "~", " "): + name = name.split(sep)[0] + names.add(name.strip().lower()) + return names + + +def test_html2text_is_declared(): + """converter.py imports html2text at module load (tools/webfetch/converter.py).""" + assert "html2text" in _dep_names(), ( + "html2text is imported eagerly by tools/webfetch/converter.py but is not " + "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/test_persistence.py b/tests/test_persistence.py deleted file mode 100644 index 3731df9..0000000 --- a/tests/test_persistence.py +++ /dev/null @@ -1,249 +0,0 @@ -"""Tests for session persistence.""" - -import json -from datetime import datetime -from pathlib import Path -from unittest.mock import MagicMock - -import pytest - -from agentic_cli.persistence.session import ( - SessionPersistence, - SessionSnapshot, -) -from tests.conftest import MockContext - - -class TestSessionSnapshot: - """Tests for SessionSnapshot dataclass.""" - - def test_snapshot_to_dict(self): - """Test snapshot serialization.""" - snapshot = SessionSnapshot( - session_id="test-session", - created_at=datetime(2024, 1, 1, 10, 0, 0), - saved_at=datetime(2024, 1, 1, 12, 0, 0), - messages=[ - {"content": "Hello", "message_type": "user"}, - {"content": "Hi there", "message_type": "assistant"}, - ], - metadata={"model": "gemini-2.5-pro"}, - ) - - data = snapshot.to_dict() - - assert data["session_id"] == "test-session" - assert data["created_at"] == "2024-01-01T10:00:00" - assert data["saved_at"] == "2024-01-01T12:00:00" - assert len(data["messages"]) == 2 - assert data["metadata"]["model"] == "gemini-2.5-pro" - - def test_snapshot_from_dict(self): - """Test snapshot deserialization.""" - data = { - "session_id": "test-session", - "created_at": "2024-01-01T10:00:00", - "saved_at": "2024-01-01T12:00:00", - "messages": [{"content": "Test", "message_type": "user"}], - "metadata": {"key": "value"}, - } - - snapshot = SessionSnapshot.from_dict(data) - - assert snapshot.session_id == "test-session" - assert snapshot.created_at == datetime(2024, 1, 1, 10, 0, 0) - assert snapshot.saved_at == datetime(2024, 1, 1, 12, 0, 0) - assert len(snapshot.messages) == 1 - assert snapshot.metadata == {"key": "value"} - - def test_snapshot_from_dict_no_metadata(self): - """Test snapshot deserialization without metadata.""" - data = { - "session_id": "test", - "created_at": "2024-01-01T00:00:00", - "saved_at": "2024-01-01T00:00:00", - "messages": [], - } - - snapshot = SessionSnapshot.from_dict(data) - - assert snapshot.metadata == {} - - def test_snapshot_roundtrip(self): - """Test serialization roundtrip.""" - original = SessionSnapshot( - session_id="roundtrip-test", - created_at=datetime(2024, 6, 15, 14, 30, 0), - saved_at=datetime(2024, 6, 15, 15, 0, 0), - messages=[{"content": "Test message", "message_type": "user"}], - metadata={"test": True}, - ) - - data = original.to_dict() - restored = SessionSnapshot.from_dict(data) - - assert restored.session_id == original.session_id - assert restored.created_at == original.created_at - assert restored.saved_at == original.saved_at - assert restored.messages == original.messages - assert restored.metadata == original.metadata - - -class TestSessionPersistence: - """Tests for SessionPersistence class.""" - - def test_init_creates_directory(self, mock_context: MockContext): - """Test that initialization creates sessions directory.""" - persistence = SessionPersistence(mock_context.settings) - - assert persistence.sessions_path.exists() - assert persistence.sessions_path == mock_context.workspace_dir / "sessions" - - def test_session_id_sanitization(self, mock_context: MockContext): - """Test that session IDs are sanitized for filesystem.""" - persistence = SessionPersistence(mock_context.settings) - - # Test various problematic session IDs - test_cases = [ - ("normal-session", "normal-session.json"), - ("session with spaces", "session_with_spaces.json"), - ("session/with/slashes", "session_with_slashes.json"), - ("session:with:colons", "session_with_colons.json"), - ] - - for session_id, expected_filename in test_cases: - path = persistence._get_session_path(session_id) - assert path.name == expected_filename - - def test_list_sessions_empty(self, mock_context: MockContext): - """Test listing sessions when none exist.""" - persistence = SessionPersistence(mock_context.settings) - - sessions = persistence.list_sessions() - - assert sessions == [] - - def test_delete_nonexistent_session(self, mock_context: MockContext): - """Test deleting a session that doesn't exist.""" - persistence = SessionPersistence(mock_context.settings) - - result = persistence.delete_session("nonexistent") - - assert result is False - - def test_load_nonexistent_session(self, mock_context: MockContext): - """Test loading a session that doesn't exist.""" - persistence = SessionPersistence(mock_context.settings) - - snapshot = persistence.load_session("nonexistent") - - assert snapshot is None - - def test_save_and_load_manually(self, mock_context: MockContext): - """Test saving and loading session data manually.""" - persistence = SessionPersistence(mock_context.settings) - - # Create a session file manually - session_id = "manual-test" - session_path = persistence._get_session_path(session_id) - - snapshot_data = { - "session_id": session_id, - "created_at": "2024-01-01T00:00:00", - "saved_at": "2024-01-01T01:00:00", - "messages": [{"content": "Hello", "message_type": "user"}], - "metadata": {"test": True}, - } - - with open(session_path, "w") as f: - json.dump(snapshot_data, f) - - # Load it back - loaded = persistence.load_session(session_id) - - assert loaded is not None - assert loaded.session_id == session_id - assert len(loaded.messages) == 1 - - def test_list_sessions(self, mock_context: MockContext): - """Test listing multiple sessions.""" - persistence = SessionPersistence(mock_context.settings) - - # Create multiple session files - for i in range(3): - session_id = f"session-{i}" - session_path = persistence._get_session_path(session_id) - - snapshot_data = { - "session_id": session_id, - "created_at": f"2024-01-0{i+1}T00:00:00", - "saved_at": f"2024-01-0{i+1}T01:00:00", - "messages": [{"content": f"Message {i}"}], - "metadata": {}, - } - - with open(session_path, "w") as f: - json.dump(snapshot_data, f) - - sessions = persistence.list_sessions() - - assert len(sessions) == 3 - # Should be sorted by saved_at descending - assert sessions[0]["session_id"] == "session-2" - assert sessions[1]["session_id"] == "session-1" - assert sessions[2]["session_id"] == "session-0" - - def test_delete_session(self, mock_context: MockContext): - """Test deleting a session.""" - persistence = SessionPersistence(mock_context.settings) - - # Create a session file - session_id = "to-delete" - session_path = persistence._get_session_path(session_id) - - with open(session_path, "w") as f: - json.dump( - { - "session_id": session_id, - "created_at": "2024-01-01T00:00:00", - "saved_at": "2024-01-01T00:00:00", - "messages": [], - }, - f, - ) - - assert session_path.exists() - - # Delete it - result = persistence.delete_session(session_id) - - assert result is True - assert not session_path.exists() - - def test_list_sessions_skips_invalid(self, mock_context: MockContext): - """Test that list_sessions skips invalid JSON files.""" - persistence = SessionPersistence(mock_context.settings) - - # Create a valid session - valid_path = persistence._get_session_path("valid") - with open(valid_path, "w") as f: - json.dump( - { - "session_id": "valid", - "created_at": "2024-01-01T00:00:00", - "saved_at": "2024-01-01T00:00:00", - "messages": [], - }, - f, - ) - - # Create an invalid session file - invalid_path = persistence.sessions_path / "invalid.json" - with open(invalid_path, "w") as f: - f.write("not valid json{") - - sessions = persistence.list_sessions() - - # Should only return the valid session - assert len(sessions) == 1 - assert sessions[0]["session_id"] == "valid" diff --git a/tests/test_retry.py b/tests/test_retry.py new file mode 100644 index 0000000..36bcfe8 --- /dev/null +++ b/tests/test_retry.py @@ -0,0 +1,32 @@ +"""Tests for rate-limit detection helpers (workflow/retry.py).""" + +from __future__ import annotations + +from agentic_cli.workflow.retry import is_rate_limit_error + + +class _Err(Exception): + def __init__(self, msg: str = "", *, status_code=None, code=None) -> None: + super().__init__(msg) + if status_code is not None: + self.status_code = status_code + if code is not None: + self.code = code + + +def test_gemini_code_429_detected(): + assert is_rate_limit_error(_Err(code=429)) is True + + +def test_resource_exhausted_string_detected(): + assert is_rate_limit_error(_Err("RESOURCE_EXHAUSTED: quota")) is True + + +def test_anthropic_status_code_429_detected(): + """anthropic.RateLimitError carries .status_code, not .code.""" + assert is_rate_limit_error(_Err(status_code=429)) is True + + +def test_non_rate_limit_error_not_detected(): + assert is_rate_limit_error(_Err("boom", status_code=500)) is False + assert is_rate_limit_error(_Err("plain error")) is False diff --git a/tests/test_session_index_recovery.py b/tests/test_session_index_recovery.py deleted file mode 100644 index 26064c0..0000000 --- a/tests/test_session_index_recovery.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Tests that a corrupt sessions index recovers by rebuilding from files -rather than silently hiding existing sessions.""" - -from datetime import datetime - -from agentic_cli.config import BaseSettings -from agentic_cli.persistence.session import SessionPersistence, SessionSnapshot - - -def _persistence(tmp_path) -> SessionPersistence: - return SessionPersistence(BaseSettings(workspace_dir=tmp_path)) - - -def _snapshot(session_id: str) -> SessionSnapshot: - now = datetime.now() - return SessionSnapshot( - session_id=session_id, - created_at=now, - saved_at=now, - messages=[{"role": "user", "content": "hi"}], - ) - - -def test_corrupt_index_rebuilds_from_files(tmp_path): - p = _persistence(tmp_path) - for sid in ("a", "b", "c"): - p.save_snapshot(_snapshot(sid)) - - p._get_sessions_index_path().write_text("{ not valid json ", encoding="utf-8") - - listed = {s["session_id"] for s in p.list_sessions()} - assert listed == {"a", "b", "c"} - - -def test_non_dict_index_rebuilds_from_files(tmp_path): - p = _persistence(tmp_path) - for sid in ("a", "b"): - p.save_snapshot(_snapshot(sid)) - - # A JSON array (not an object) used to raise AttributeError on data.get(...). - p._get_sessions_index_path().write_text('["not", "a", "dict"]', encoding="utf-8") - - listed = {s["session_id"] for s in p.list_sessions()} - assert listed == {"a", "b"} - - -def test_save_after_corrupt_index_does_not_hide_existing(tmp_path): - """Regression: corrupt index, save a new session, list must still show all. - - Previously a corrupt index read as empty, so the next save reset it to a - single entry and every prior session vanished from the listing. - """ - p = _persistence(tmp_path) - p.save_snapshot(_snapshot("a")) - p.save_snapshot(_snapshot("b")) - - p._get_sessions_index_path().write_text("totally broken", encoding="utf-8") - - p.save_snapshot(_snapshot("c")) # triggers index reload (corrupt) -> rebuild - - listed = {s["session_id"] for s in p.list_sessions()} - assert listed == {"a", "b", "c"}, f"existing sessions were hidden: {listed}" diff --git a/tests/test_session_save_resume.py b/tests/test_session_save_resume.py deleted file mode 100644 index 41e35dc..0000000 --- a/tests/test_session_save_resume.py +++ /dev/null @@ -1,488 +0,0 @@ -"""Tests for session save/resume feature. - -Covers: -- SessionPersistence.save_snapshot() -- BaseWorkflowManager.save_session() / load_session() / list_sessions() -- Round-trip: save then load produces same messages -- Cross-backend load logs warning but succeeds -- SessionsCommand displays table -""" - -import json -from datetime import datetime -from typing import AsyncGenerator -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from agentic_cli.config import BaseSettings -from agentic_cli.persistence.session import SessionPersistence, SessionSnapshot -from agentic_cli.workflow.base_manager import BaseWorkflowManager -from agentic_cli.workflow.config import AgentConfig -from agentic_cli.workflow.events import WorkflowEvent - -from tests.conftest import MockContext - - -# --------------------------------------------------------------------------- -# Concrete test subclass of BaseWorkflowManager -# --------------------------------------------------------------------------- - -class _TestWorkflowManager(BaseWorkflowManager): - """Minimal concrete subclass for testing base class session methods.""" - - def __init__(self, settings: BaseSettings, **kwargs): - super().__init__(agent_configs=[], settings=settings, **kwargs) - # In-memory message store for testing - self._stored_messages: list[dict] = [] - self._stored_agent: str | None = None - self.session_id = "default_session" - - @property - def backend_type(self) -> str: - return "test" - - async def _do_initialize(self) -> None: - pass - - def _get_state_tools(self) -> list: - return [] - - async def process( - self, message: str, user_id: str, session_id: str | None = None - ) -> AsyncGenerator[WorkflowEvent, None]: - if False: - yield # type: ignore[misc] - - async def reinitialize( - self, model: str | None = None, preserve_sessions: bool = True - ) -> None: - pass - - async def cleanup(self) -> None: - pass - - async def _extract_session_data(self, session_id: str) -> tuple[list[dict], str | None]: - return self._stored_messages, self._stored_agent - - async def _inject_session_messages( - self, - session_id: str, - messages: list[dict], - current_agent: str | None = None, - ) -> None: - self._stored_messages = messages - self._stored_agent = current_agent - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -SAMPLE_MESSAGES = [ - {"role": "user", "content": "Research quantum computing"}, - { - "role": "assistant", - "content": "I'll search for that.", - "tool_calls": [ - {"id": "tc_1", "name": "web_search", "args": {"query": "quantum computing"}}, - ], - }, - { - "role": "tool", - "tool_call_id": "tc_1", - "name": "web_search", - "content": '{"success": true, "results": []}', - }, - {"role": "assistant", "content": "Based on my research..."}, -] - - -# --------------------------------------------------------------------------- -# Tests: SessionPersistence.save_snapshot() -# --------------------------------------------------------------------------- - - -class TestSaveSnapshot: - """Tests for the new save_snapshot() method.""" - - def test_save_snapshot_creates_file(self, mock_context: MockContext): - persistence = SessionPersistence(mock_context.settings) - snapshot = SessionSnapshot( - session_id="snap-test", - created_at=datetime(2024, 1, 1), - saved_at=datetime(2024, 1, 1, 1, 0), - messages=SAMPLE_MESSAGES, - metadata={"model": "test-model", "backend_type": "test"}, - ) - - path = persistence.save_snapshot(snapshot) - - assert path.exists() - data = json.loads(path.read_text()) - assert data["session_id"] == "snap-test" - assert len(data["messages"]) == 4 - assert data["metadata"]["backend_type"] == "test" - - def test_save_snapshot_updates_index(self, mock_context: MockContext): - persistence = SessionPersistence(mock_context.settings) - snapshot = SessionSnapshot( - session_id="indexed-test", - created_at=datetime(2024, 1, 1), - saved_at=datetime(2024, 1, 1, 1, 0), - messages=[{"role": "user", "content": "hello"}], - ) - - persistence.save_snapshot(snapshot) - - sessions = persistence.list_sessions() - assert len(sessions) == 1 - assert sessions[0]["session_id"] == "indexed-test" - assert sessions[0]["message_count"] == 1 - - def test_save_snapshot_overwrites_existing(self, mock_context: MockContext): - persistence = SessionPersistence(mock_context.settings) - - # Save first version - snap1 = SessionSnapshot( - session_id="overwrite-test", - created_at=datetime(2024, 1, 1), - saved_at=datetime(2024, 1, 1, 1, 0), - messages=[{"role": "user", "content": "v1"}], - ) - persistence.save_snapshot(snap1) - - # Save second version - snap2 = SessionSnapshot( - session_id="overwrite-test", - created_at=datetime(2024, 1, 1), - saved_at=datetime(2024, 1, 1, 2, 0), - messages=[ - {"role": "user", "content": "v1"}, - {"role": "assistant", "content": "v2"}, - ], - ) - persistence.save_snapshot(snap2) - - # Should have latest data - loaded = persistence.load_session("overwrite-test") - assert loaded is not None - assert len(loaded.messages) == 2 - - def test_save_snapshot_load_roundtrip(self, mock_context: MockContext): - persistence = SessionPersistence(mock_context.settings) - now = datetime.now() - snapshot = SessionSnapshot( - session_id="roundtrip", - created_at=now, - saved_at=now, - messages=SAMPLE_MESSAGES, - metadata={"model": "gemini-2.5-flash", "backend_type": "adk"}, - ) - - persistence.save_snapshot(snapshot) - loaded = persistence.load_session("roundtrip") - - assert loaded is not None - assert loaded.session_id == "roundtrip" - assert loaded.messages == SAMPLE_MESSAGES - assert loaded.metadata["model"] == "gemini-2.5-flash" - - -# --------------------------------------------------------------------------- -# Tests: BaseWorkflowManager session methods -# --------------------------------------------------------------------------- - - -class TestWorkflowManagerSession: - """Tests for save_session / load_session / list_sessions on BaseWorkflowManager.""" - - async def test_save_session_persists_to_disk(self, mock_context: MockContext): - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr._stored_messages = SAMPLE_MESSAGES - mgr._stored_agent = "researcher" - mgr._model = "test-model" - mgr._model_resolved = True - - result = await mgr.save_session("my-session") - - assert result["success"] is True - assert result["session_id"] == "my-session" - assert result["message_count"] == 4 - - # Verify file on disk - persistence = SessionPersistence(mock_context.settings) - loaded = persistence.load_session("my-session") - assert loaded is not None - assert loaded.metadata["current_agent"] == "researcher" - assert loaded.metadata["backend_type"] == "test" - - async def test_save_session_uses_default_session_id(self, mock_context: MockContext): - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr.session_id = "custom-default" - mgr._stored_messages = [{"role": "user", "content": "hi"}] - mgr._model = "test-model" - mgr._model_resolved = True - - result = await mgr.save_session() - - assert result["success"] is True - assert result["session_id"] == "custom-default" - - async def test_load_session_injects_messages(self, mock_context: MockContext): - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr._model = "test-model" - mgr._model_resolved = True - - # First save - mgr._stored_messages = SAMPLE_MESSAGES - mgr._stored_agent = "researcher" - await mgr.save_session("inject-test") - - # Clear and load into fresh manager - mgr2 = _TestWorkflowManager(settings=mock_context.settings) - mgr2._model = "test-model" - mgr2._model_resolved = True - loaded = await mgr2.load_session("inject-test") - - assert loaded is True - assert mgr2._stored_messages == SAMPLE_MESSAGES - assert mgr2._stored_agent == "researcher" - assert mgr2.session_id == "inject-test" - - async def test_load_nonexistent_session_returns_false(self, mock_context: MockContext): - mgr = _TestWorkflowManager(settings=mock_context.settings) - - loaded = await mgr.load_session("does-not-exist") - - assert loaded is False - - async def test_cross_backend_load_succeeds_with_warning(self, mock_context: MockContext): - """Loading a session saved by a different backend should work but log a warning.""" - # Save with "adk" backend_type in metadata - persistence = SessionPersistence(mock_context.settings) - snapshot = SessionSnapshot( - session_id="cross-backend", - created_at=datetime.now(), - saved_at=datetime.now(), - messages=SAMPLE_MESSAGES, - metadata={"backend_type": "adk", "model": "gemini-2.5-flash"}, - ) - persistence.save_snapshot(snapshot) - - # Load into "test" backend manager - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr._model = "test-model" - mgr._model_resolved = True - - with patch("agentic_cli.workflow.base_manager.logger") as mock_logger: - loaded = await mgr.load_session("cross-backend") - - assert loaded is True - assert mgr._stored_messages == SAMPLE_MESSAGES - # Should have logged a warning about backend mismatch - mock_logger.warning.assert_called_once() - call_args = mock_logger.warning.call_args - assert call_args[0][0] == "session_backend_mismatch" - - async def test_list_sessions(self, mock_context: MockContext): - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr._model = "test-model" - mgr._model_resolved = True - - # Save two sessions - mgr._stored_messages = [{"role": "user", "content": "a"}] - await mgr.save_session("session-a") - - mgr._stored_messages = [ - {"role": "user", "content": "b"}, - {"role": "assistant", "content": "reply"}, - ] - await mgr.save_session("session-b") - - sessions = mgr.list_sessions() - - assert len(sessions) == 2 - ids = {s["session_id"] for s in sessions} - assert ids == {"session-a", "session-b"} - - async def test_save_session_round_trip_preserves_tool_calls(self, mock_context: MockContext): - """Full round-trip: messages with tool_calls survive save/load.""" - mgr = _TestWorkflowManager(settings=mock_context.settings) - mgr._model = "test-model" - mgr._model_resolved = True - mgr._stored_messages = SAMPLE_MESSAGES - - await mgr.save_session("tool-roundtrip") - - mgr2 = _TestWorkflowManager(settings=mock_context.settings) - mgr2._model = "test-model" - mgr2._model_resolved = True - await mgr2.load_session("tool-roundtrip") - - assert mgr2._stored_messages == SAMPLE_MESSAGES - # Verify tool_calls are preserved - assistant_msg = mgr2._stored_messages[1] - assert "tool_calls" in assistant_msg - assert assistant_msg["tool_calls"][0]["name"] == "web_search" - - # Verify tool result message - tool_msg = mgr2._stored_messages[2] - assert tool_msg["role"] == "tool" - assert tool_msg["tool_call_id"] == "tc_1" - - -# --------------------------------------------------------------------------- -# Tests: SessionsCommand -# --------------------------------------------------------------------------- - - -class TestSessionsCommand: - """Tests for the /sessions command.""" - - async def test_sessions_command_empty(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import SessionsCommand - - cmd = SessionsCommand() - app = MagicMock() - app.settings = mock_context.settings - app.session = MagicMock() - app._session_id = None - - await cmd.execute("", app) - - app.session.add_message.assert_called_once_with("system", "No saved sessions.") - - async def test_sessions_command_lists_sessions(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import SessionsCommand - - # Save a session first - persistence = SessionPersistence(mock_context.settings) - snapshot = SessionSnapshot( - session_id="listed-session", - created_at=datetime(2024, 6, 1), - saved_at=datetime(2024, 6, 1, 12, 0), - messages=[{"role": "user", "content": "hi"}], - ) - persistence.save_snapshot(snapshot) - - cmd = SessionsCommand() - app = MagicMock() - app.settings = mock_context.settings - app.session = MagicMock() - app._session_id = "listed-session" - - await cmd.execute("", app) - - # Should have called add_rich with a table - app.session.add_rich.assert_called_once() - - async def test_sessions_command_delete(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import SessionsCommand - - # Save a session to delete - persistence = SessionPersistence(mock_context.settings) - snapshot = SessionSnapshot( - session_id="to-delete", - created_at=datetime(2024, 6, 1), - saved_at=datetime(2024, 6, 1, 12, 0), - messages=[{"role": "user", "content": "bye"}], - ) - persistence.save_snapshot(snapshot) - - cmd = SessionsCommand() - app = MagicMock() - app.settings = mock_context.settings - app.session = MagicMock() - - await cmd.execute("--delete=to-delete", app) - - app.session.add_success.assert_called_once() - # Verify it's actually gone - assert persistence.load_session("to-delete") is None - - async def test_sessions_command_delete_nonexistent(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import SessionsCommand - - cmd = SessionsCommand() - app = MagicMock() - app.settings = mock_context.settings - app.session = MagicMock() - - await cmd.execute("--delete=nope", app) - - app.session.add_error.assert_called_once() - - -# --------------------------------------------------------------------------- -# Tests: StatusCommand session info -# --------------------------------------------------------------------------- - - -class TestStatusCommandSessionInfo: - """Tests that /status shows session info.""" - - async def test_status_shows_persistent_session(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import StatusCommand - - cmd = StatusCommand() - app = MagicMock() - app._session_id = "my-research" - app.usage_tracker = MagicMock(invocation_count=0) - # Make workflow access raise to skip workflow section - type(app).workflow = property(lambda self: (_ for _ in ()).throw(RuntimeError)) - - await cmd.execute("", app) - - # The table should have been added via add_rich - app.session.add_rich.assert_called_once() - - async def test_status_shows_ephemeral_session(self, mock_context: MockContext): - from agentic_cli.cli.builtin_commands import StatusCommand - - cmd = StatusCommand() - app = MagicMock() - app._session_id = None - app.usage_tracker = MagicMock(invocation_count=0) - type(app).workflow = property(lambda self: (_ for _ in ()).throw(RuntimeError)) - - await cmd.execute("", app) - - app.session.add_rich.assert_called_once() - - -# --------------------------------------------------------------------------- -# Tests: BaseCLIApp session_id parameter -# --------------------------------------------------------------------------- - - -class TestBaseCLIAppSessionId: - """Tests for session_id parameter on BaseCLIApp.""" - - def test_session_id_defaults_to_none(self): - """session_id should default to None when not provided.""" - from agentic_cli.cli.app import BaseCLIApp - - # Just verify the parameter exists in the signature - import inspect - sig = inspect.signature(BaseCLIApp.__init__) - param = sig.parameters.get("session_id") - assert param is not None - assert param.default is None - - def test_session_id_property(self, mock_context: MockContext): - """session_id property should be gettable and settable.""" - from agentic_cli.cli.app import BaseCLIApp - from thinking_prompt import AppInfo - - with patch("agentic_cli.cli.app.ThinkingPromptSession"): - app = BaseCLIApp( - app_info=AppInfo(name="test", version="0.1.0"), - agent_configs=[], - settings=mock_context.settings, - session_id="my-sess", - ) - - assert app.session_id == "my-sess" - app.session_id = "other" - assert app.session_id == "other" diff --git a/tests/test_settings_persistence.py b/tests/test_settings_persistence.py index 41c30f8..c05a5b5 100644 --- a/tests/test_settings_persistence.py +++ b/tests/test_settings_persistence.py @@ -13,17 +13,18 @@ class TestSecretFields: """Tests for SECRET_FIELDS completeness (C1).""" def test_secret_fields_excludes_all_api_keys(self): - """All API key fields in BaseSettings are in SECRET_FIELDS.""" + """Every credential-bearing field in BaseSettings is in SECRET_FIELDS.""" expected = { "google_api_key", "anthropic_api_key", "tavily_api_key", "brave_api_key", + "postgres_uri", # embeds user:password@host } assert SECRET_FIELDS == expected def test_save_excludes_secrets(self, tmp_path): - """API key values are never written to the JSON file.""" + """Secret values are never written to the JSON file.""" from agentic_cli.config import BaseSettings settings = BaseSettings( @@ -31,6 +32,7 @@ def test_save_excludes_secrets(self, tmp_path): anthropic_api_key="secret-anthropic", tavily_api_key="secret-tavily", brave_api_key="secret-brave", + postgres_uri="postgresql://user:pass@host/db", search_backend="tavily", # non-secret ) @@ -86,6 +88,161 @@ def test_save_includes_default_values(self, tmp_path): assert "thinking_effort" in data +class TestTrustSplitSave: + """P0-1 follow-up: save() must agree with the allowlist-filtered loader. + + Default save splits by trust: allowlisted keys → project settings.json, + non-allowlisted (user-scoped) keys → user ~/.{app}/settings.json, so a + /settings change to a security-relevant key survives a restart instead of + being silently dropped by _AllowlistFilterSource. + """ + + @pytest.fixture + def isolated_fs(self, tmp_path, monkeypatch): + """Isolate cwd + HOME so save/load hit temp files only.""" + home = tmp_path / "home" + home.mkdir() + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(home)) + for var in ( + "AGENTIC_RAW_LLM_LOGGING", + "AGENTIC_STATEFUL_EXECUTOR_BACKEND", + "AGENTIC_DEFAULT_MODEL", + ): + monkeypatch.delenv(var, raising=False) + return tmp_path + + def _persistence(self): + return SettingsPersistence(app_name="agentic_cli") + + def test_default_save_keeps_project_file_allowlisted(self, isolated_fs): + """Project file receives only allowlisted keys (and heals stale ones).""" + from agentic_cli.config import BaseSettings + from agentic_cli.settings_persistence import PROJECT_SETTABLE_KEYS + + # Pre-seed a stale kitchen-sink project file (pre-P0-1 artifact) + proj_dir = isolated_fs / ".agentic_cli" + proj_dir.mkdir() + (proj_dir / "settings.json").write_text( + json.dumps({"raw_llm_logging": True, "default_model": "stale-model"}) + ) + + settings = BaseSettings( + raw_llm_logging=True, default_model="claude-sonnet-4-6" + ) + self._persistence().save(settings) + + data = json.loads((proj_dir / "settings.json").read_text()) + assert set(data) <= PROJECT_SETTABLE_KEYS + assert data["default_model"] == "claude-sonnet-4-6" + assert "raw_llm_logging" not in data + + def test_default_save_writes_user_scoped_keys_to_user_config(self, isolated_fs): + """Non-allowlisted keys that differ from defaults land in user config.""" + from agentic_cli.config import BaseSettings + + settings = BaseSettings( + raw_llm_logging=True, stateful_executor_backend="local" + ) + self._persistence().save(settings) + + user_file = isolated_fs / "home" / ".agentic_cli" / "settings.json" + assert user_file.exists() + data = json.loads(user_file.read_text()) + assert data["raw_llm_logging"] is True + assert data["stateful_executor_backend"] == "local" + + def test_default_save_at_defaults_does_not_create_user_config(self, isolated_fs): + """All-default settings write nothing user-scoped.""" + from agentic_cli.config import BaseSettings + + self._persistence().save(BaseSettings()) + + assert not (isolated_fs / "home" / ".agentic_cli" / "settings.json").exists() + + def test_user_config_merge_preserves_unmanaged_keys(self, isolated_fs): + """Hand-stored secrets and unknown keys in user config survive a save.""" + from agentic_cli.config import BaseSettings + + user_dir = isolated_fs / "home" / ".agentic_cli" + user_dir.mkdir(parents=True) + (user_dir / "settings.json").write_text( + json.dumps({"anthropic_api_key": "sk-stored", "domain_custom_key": 7}) + ) + + settings = BaseSettings(raw_llm_logging=True) + self._persistence().save(settings) + + data = json.loads((user_dir / "settings.json").read_text()) + assert data["anthropic_api_key"] == "sk-stored" + assert data["domain_custom_key"] == 7 + assert data["raw_llm_logging"] is True + + def test_revert_to_default_removes_user_config_key(self, isolated_fs): + """Reverting a user-scoped key to its default removes it from user config.""" + from agentic_cli.config import BaseSettings + + user_dir = isolated_fs / "home" / ".agentic_cli" + user_dir.mkdir(parents=True) + (user_dir / "settings.json").write_text( + json.dumps({"raw_llm_logging": True, "domain_custom_key": 7}) + ) + + settings = BaseSettings() # loads raw_llm_logging=True from user config + settings.update_setting("raw_llm_logging", False) # user reverts + self._persistence().save(settings) + + data = json.loads((user_dir / "settings.json").read_text()) + assert "raw_llm_logging" not in data + assert data["domain_custom_key"] == 7 # unmanaged key untouched + + def test_round_trip_user_scoped_change_persists(self, isolated_fs): + """THE regression: a /settings change to a non-allowlisted key must + survive save + fresh load, without tripping the untrusted-key filter.""" + import structlog + + from agentic_cli.config import BaseSettings + + settings = BaseSettings() + settings.update_setting("stateful_executor_backend", "local") + self._persistence().save(settings) + + with structlog.testing.capture_logs() as logs: + reloaded = BaseSettings() + + assert reloaded.stateful_executor_backend == "local" + assert not any( + e.get("event") == "untrusted_project_setting_ignored" for e in logs + ) + + def test_explicit_path_keeps_legacy_full_dump(self, isolated_fs): + """save(path=...) still writes the full single-file dump.""" + from agentic_cli.config import BaseSettings + + out = isolated_fs / "export" / "settings.json" + result = self._persistence().save( + BaseSettings(raw_llm_logging=True), path=out + ) + + data = json.loads(out.read_text()) + assert data["raw_llm_logging"] is True # non-allowlisted key retained + assert result.project_path == out + assert result.user_path is None + + def test_save_result_reports_split(self, isolated_fs): + """Default save reports both paths and the user-scoped keys.""" + from agentic_cli.config import BaseSettings + + result = self._persistence().save(BaseSettings(raw_llm_logging=True)) + + assert result.project_path == isolated_fs / ".agentic_cli" / "settings.json" + assert ( + result.user_path + == isolated_fs / "home" / ".agentic_cli" / "settings.json" + ) + assert "raw_llm_logging" in result.user_scoped_keys + + class TestAtomicWrite: """Tests for atomic settings write (C2).""" diff --git a/tests/test_token_caching.py b/tests/test_token_caching.py index ebf1ae0..25e157d 100644 --- a/tests/test_token_caching.py +++ b/tests/test_token_caching.py @@ -87,7 +87,7 @@ async def mock_ainvoke(messages, **kwargs): import asyncio state = {"messages": [{"role": "user", "content": "hello"}]} - asyncio.get_event_loop().run_until_complete(node_fn(state)) + asyncio.run(node_fn(state)) # Check the SystemMessage sys_msgs = [m for m in captured_messages if isinstance(m, SystemMessage)] @@ -135,7 +135,7 @@ async def mock_ainvoke(messages, **kwargs): import asyncio state = {"messages": [{"role": "user", "content": "hello"}]} - asyncio.get_event_loop().run_until_complete( + asyncio.run( manager._builder._create_agent_node(agent_config, manager.model)(state) ) @@ -179,7 +179,7 @@ async def mock_ainvoke(messages, **kwargs): import asyncio state = {"messages": [{"role": "user", "content": "hello"}]} - asyncio.get_event_loop().run_until_complete( + asyncio.run( manager._builder._create_agent_node(agent_config, manager.model)(state) ) diff --git a/tests/test_tool_summaries.py b/tests/test_tool_summaries.py index 73117d4..d990147 100644 --- a/tests/test_tool_summaries.py +++ b/tests/test_tool_summaries.py @@ -450,3 +450,20 @@ def test_string_result_unchanged(self): success=True, ) assert event.content == "some string result" + + +class TestSandboxExecute: + def test_first_stdout_line(self): + result = {"success": True, "stdout": "shape (5, 5)\nmore\n"} + assert format_tool_summary("sandbox_execute", result) == "shape (5, 5)" + + def test_ok_when_no_stdout(self): + assert format_tool_summary("sandbox_execute", {"success": True, "stdout": ""}) == "ok" + + def test_error_on_failure(self): + result = {"success": False, "error": "NameError: boom", "stdout": ""} + assert format_tool_summary("sandbox_execute", result) == "NameError: boom" + + def test_failure_falls_back_to_stderr(self): + result = {"success": False, "error": "", "stderr": "Traceback here", "stdout": ""} + assert format_tool_summary("sandbox_execute", result) == "Traceback here" diff --git a/tests/test_tools.py b/tests/test_tools.py index a71ee00..0bfa2bb 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1482,7 +1482,7 @@ def test_dangerous_tool_executes_directly(self): from tests.conftest import MockContext from tests.tools.test_sandbox import MockSandboxBackend - with MockContext() as ctx: + with MockContext(stateful_executor_backend="local") as ctx: backend = MockSandboxBackend( ExecutionResult(success=True, stdout="ok\n", result="1") ) diff --git a/tests/test_webfetch.py b/tests/test_webfetch.py index a90267c..f1e4259 100644 --- a/tests/test_webfetch.py +++ b/tests/test_webfetch.py @@ -1,12 +1,36 @@ """Tests for webfetch tool.""" import ipaddress +import socket +import httpx import pytest +from unittest.mock import AsyncMock, patch from agentic_cli.config import BaseSettings +def _pinned_fetcher(monkeypatch, handler, *, resolves_to="93.184.216.34", **kw): + """Build a ContentFetcher whose transport is a PinnedTransport over a + MockTransport(handler); getaddrinfo is stubbed so pinning succeeds.""" + from agentic_cli.tools.webfetch.validator import URLValidator + from agentic_cli.tools.webfetch.transport import PinnedTransport + from agentic_cli.tools.webfetch.robots import RobotsTxtChecker + from agentic_cli.tools.webfetch.fetcher import ContentFetcher + + def _gai(host, port, *a, **k): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (resolves_to, port))] + monkeypatch.setattr(socket, "getaddrinfo", _gai) + + validator = URLValidator() + transport = PinnedTransport(validator, inner=httpx.MockTransport(handler)) + robots = RobotsTxtChecker() # robots.can_fetch is patched in these tests, so + # its transport is irrelevant here (wired in Task 4) + return ContentFetcher( + validator=validator, robots_checker=robots, transport=transport, **kw + ) + + class TestWebFetchSettings: """Tests for webfetch settings fields.""" @@ -61,11 +85,10 @@ def test_invalid_scheme_file(self, validator): result = validator.validate("file:///etc/passwd") assert result.valid is False - def test_localhost_blocked(self, validator): - """Test localhost is blocked.""" - result = validator.validate("http://localhost/api") + def test_localhost_ip_literal_blocked(self, validator): + """A loopback IP literal is blocked by validate() without DNS.""" + result = validator.validate("http://127.0.0.1/api") assert result.valid is False - assert "private" in result.error.lower() or "blocked" in result.error.lower() def test_127_0_0_1_blocked(self, validator): """Test 127.0.0.1 is blocked.""" @@ -120,370 +143,183 @@ def test_malformed_url(self, validator): assert result.valid is False -class TestValidateIP: - """Tests for URLValidator.validate_ip() (C3: extracted IP check).""" - - @pytest.fixture - def validator(self): - from agentic_cli.tools.webfetch.validator import URLValidator - return URLValidator(blocked_domains=[]) - - def test_validate_ip_blocks_private(self, validator): - """Private IPs are blocked.""" - for ip in ("127.0.0.1", "10.0.0.1", "192.168.1.1", "172.16.0.1"): - result = validator.validate_ip(ip) - assert result.valid is False, f"{ip} should be blocked" - assert "blocked" in result.error.lower() - - def test_validate_ip_allows_public(self, validator): - """Public IPs pass validation.""" - result = validator.validate_ip("8.8.8.8") - assert result.valid is True - assert result.resolved_ip == "8.8.8.8" - - def test_validate_ip_invalid_string(self, validator): - """Invalid IP string returns error.""" - result = validator.validate_ip("not-an-ip") - assert result.valid is False - assert "invalid" in result.error.lower() - - -class TestPostFetchRevalidation: - """Tests for DNS rebinding protection in ContentFetcher (C3).""" - - @pytest.mark.asyncio - async def test_post_fetch_revalidation_blocks_rebind(self): - """If DNS resolves to private IP post-fetch, the response is rejected.""" - from agentic_cli.tools.webfetch.fetcher import ContentFetcher - from agentic_cli.tools.webfetch.validator import URLValidator, ValidationResult - from agentic_cli.tools.webfetch.robots import RobotsTxtChecker - from unittest.mock import AsyncMock, patch, MagicMock - - fetcher = ContentFetcher( - validator=URLValidator(blocked_domains=[]), - robots_checker=RobotsTxtChecker(), - ) - - # Mock validator.validate to allow pre-fetch (public IP) - with patch.object(fetcher._validator, "validate") as mock_validate: - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = "evil content" - mock_response.headers = {"content-type": "text/html"} - mock_response.history = [] - mock_response.url = "https://example.com/page" - mock_get.return_value = mock_response - - # Post-fetch DNS resolves to private IP (rebinding attack) - with patch("agentic_cli.tools.webfetch.fetcher.socket.gethostbyname") as mock_dns: - mock_dns.return_value = "127.0.0.1" - result = await fetcher.fetch("https://example.com/page") - - assert result.success is False - assert "rebinding" in result.error.lower() - - -from unittest.mock import AsyncMock, patch - - class TestRobotsTxtChecker: - """Tests for robots.txt compliance.""" - - @pytest.fixture - def checker(self): + def _checker(self, monkeypatch, handler, resolves_to="93.184.216.34"): + from agentic_cli.tools.webfetch.validator import URLValidator + from agentic_cli.tools.webfetch.transport import PinnedTransport from agentic_cli.tools.webfetch.robots import RobotsTxtChecker - return RobotsTxtChecker() + def _gai(host, port, *a, **k): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (resolves_to, port))] + monkeypatch.setattr(socket, "getaddrinfo", _gai) + transport = PinnedTransport(URLValidator(), inner=httpx.MockTransport(handler)) + return RobotsTxtChecker(transport=transport) @pytest.mark.asyncio - async def test_allowed_when_no_robots_txt(self, checker): - """Test URL is allowed when robots.txt doesn't exist.""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.return_value.status_code = 404 - result = await checker.can_fetch("https://example.com/page") - assert result is True + async def test_allowed_when_no_robots_txt(self, monkeypatch): + checker = self._checker(monkeypatch, lambda req: httpx.Response(404)) + assert await checker.can_fetch("https://example.com/page") is True @pytest.mark.asyncio - async def test_allowed_when_robots_txt_error(self, checker): - """Test URL is allowed when robots.txt fetch fails.""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.side_effect = Exception("Network error") - result = await checker.can_fetch("https://example.com/page") - assert result is True # Permissive on error + async def test_blocked_by_robots_txt(self, monkeypatch): + robots = "User-agent: *\nDisallow: /private/\n" + checker = self._checker(monkeypatch, lambda req: httpx.Response(200, text=robots)) + assert await checker.can_fetch("https://example.com/private/secret") is False @pytest.mark.asyncio - async def test_blocked_by_robots_txt(self, checker): - """Test URL is blocked when robots.txt disallows it.""" - robots_content = """ -User-agent: * -Disallow: /private/ -""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = robots_content - mock_get.return_value = mock_response - result = await checker.can_fetch("https://example.com/private/secret") - assert result is False + async def test_allowed_by_robots_txt(self, monkeypatch): + robots = "User-agent: *\nDisallow: /private/\n" + checker = self._checker(monkeypatch, lambda req: httpx.Response(200, text=robots)) + assert await checker.can_fetch("https://example.com/public/page") is True @pytest.mark.asyncio - async def test_allowed_by_robots_txt(self, checker): - """Test URL is allowed when robots.txt permits it.""" - robots_content = """ -User-agent: * -Disallow: /private/ -""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = robots_content - mock_get.return_value = mock_response - result = await checker.can_fetch("https://example.com/public/page") - assert result is True - - @pytest.mark.asyncio - async def test_robots_txt_cached(self, checker): - """Test robots.txt is cached per domain.""" - robots_content = "User-agent: *\nAllow: /" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = robots_content - mock_get.return_value = mock_response - - await checker.can_fetch("https://example.com/page1") - await checker.can_fetch("https://example.com/page2") - - # Should only fetch robots.txt once - assert mock_get.call_count == 1 - - -import time -import httpx + async def test_robots_txt_cached(self, monkeypatch): + calls = {"n": 0} + def handler(req): + calls["n"] += 1 + return httpx.Response(200, text="User-agent: *\nAllow: /\n") + checker = self._checker(monkeypatch, handler) + await checker.can_fetch("https://example.com/page1") + await checker.can_fetch("https://example.com/page2") + assert calls["n"] == 1 # one robots.txt fetch per domain class TestContentFetcher: - """Tests for content fetching with caching and redirect handling.""" - - @pytest.fixture - def fetcher(self): - from agentic_cli.tools.webfetch.fetcher import ContentFetcher - from agentic_cli.tools.webfetch.validator import URLValidator - from agentic_cli.tools.webfetch.robots import RobotsTxtChecker - return ContentFetcher( - validator=URLValidator(blocked_domains=[]), - robots_checker=RobotsTxtChecker(), - cache_ttl_seconds=900, - max_content_bytes=102400, - ) - @pytest.mark.asyncio - async def test_fetch_success(self, fetcher): - """Test successful fetch.""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = "Content" - mock_response.headers = {"content-type": "text/html"} - mock_response.history = [] - mock_response.url = "https://example.com/page" - mock_get.return_value = mock_response - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - result = await fetcher.fetch("https://example.com/page") - + async def test_fetch_success(self, monkeypatch): + fetcher = _pinned_fetcher( + monkeypatch, + lambda req: httpx.Response(200, text="Content", + headers={"content-type": "text/html"}), + ) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/page") assert result.success is True assert "Content" in result.content @pytest.mark.asyncio - async def test_fetch_blocked_by_validator(self, fetcher): - """Test fetch blocked by URL validator.""" + async def test_fetch_pins_to_validated_ip(self, monkeypatch): + seen = {} + def handler(req): + seen["host"] = req.url.host + seen["host_header"] = req.headers.get("Host") + return httpx.Response(200, text="ok", headers={"content-type": "text/html"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + await fetcher.fetch("https://example.com/page") + assert seen["host"] == "93.184.216.34" # connected to pinned IP + assert seen["host_header"] == "example.com" # Host preserved + + @pytest.mark.asyncio + async def test_fetch_blocked_by_validator_ip_literal(self, monkeypatch): + fetcher = _pinned_fetcher(monkeypatch, lambda req: httpx.Response(200)) result = await fetcher.fetch("http://127.0.0.1/internal") assert result.success is False - assert "blocked" in result.error.lower() or "private" in result.error.lower() + assert result.error @pytest.mark.asyncio - async def test_fetch_blocked_by_robots(self, fetcher): - """Test fetch blocked by robots.txt.""" - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = False - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - result = await fetcher.fetch("https://example.com/private/page") - + async def test_fetch_blocked_when_host_resolves_private(self, monkeypatch): + fetcher = _pinned_fetcher(monkeypatch, lambda req: httpx.Response(200), + resolves_to="10.0.0.5") + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://intranet.test/secret") assert result.success is False - assert "robots" in result.error.lower() + assert "block" in (result.error or "").lower() @pytest.mark.asyncio - async def test_fetch_cross_host_redirect_blocks_before_second_request(self, fetcher): - """A 3xx redirect to a different host must be reported as a cross-host - redirect WITHOUT issuing the next GET. The user's permission grant - covers the original host only, so contacting another origin is an - unapproved side effect.""" - redirect_response = AsyncMock() - redirect_response.status_code = 302 - redirect_response.headers = {"location": "https://other.com/page"} - redirect_response.url = httpx.URL("https://example.com/page") - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = redirect_response - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - result = await fetcher.fetch("https://example.com/page") - + async def test_fetch_blocked_by_robots(self, monkeypatch): + fetcher = _pinned_fetcher(monkeypatch, lambda req: httpx.Response(200, text="x")) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=False): + result = await fetcher.fetch("https://example.com/private/page") assert result.success is False - assert result.redirect is not None - assert result.redirect.to_host == "other.com" - # Only the original GET should have been issued. - assert mock_get.call_count == 1 + assert "robots" in result.error.lower() @pytest.mark.asyncio - async def test_fetch_same_host_redirect_blocked_by_robots(self, fetcher): - """A same-host redirect whose target is disallowed by robots.txt - must not be fetched, even though the original URL was allowed.""" - redirect_response = AsyncMock() - redirect_response.status_code = 302 - redirect_response.headers = {"location": "/private/page"} - redirect_response.url = httpx.URL("https://example.com/public") - - robots_calls: list[str] = [] - + async def test_same_host_redirect_blocked_by_robots(self, monkeypatch): + def handler(req): + if req.url.path == "/public": + return httpx.Response(302, headers={"location": "/private/page"}) + return httpx.Response(200, text="x", headers={"content-type": "text/html"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + robots_calls = [] async def fake_can_fetch(u): robots_calls.append(u) return "/private" not in u + with patch.object(fetcher._robots, "can_fetch", side_effect=fake_can_fetch): + result = await fetcher.fetch("https://example.com/public") + assert result.success is False + assert "robots" in result.error.lower() + assert any("/private" in u for u in robots_calls) - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = redirect_response - with patch.object(fetcher._robots, "can_fetch", side_effect=fake_can_fetch): - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - result = await fetcher.fetch("https://example.com/public") - + @pytest.mark.asyncio + async def test_cross_host_redirect_blocks_before_second_request(self, monkeypatch): + calls = {"n": 0} + def handler(req): + calls["n"] += 1 + return httpx.Response(302, headers={"location": "https://other.com/page"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/page") assert result.success is False - assert "robots" in (result.error or "").lower() - # Only the original GET should have fired; redirected URL was blocked - # before its request. - assert mock_get.call_count == 1 - # robots.txt was consulted for both the original and the redirect target. - assert any("/public" in u for u in robots_calls) - assert any("/private/page" in u for u in robots_calls) + assert result.redirect is not None and result.redirect.to_host == "other.com" + assert calls["n"] == 1 # next GET never issued @pytest.mark.asyncio - async def test_fetch_redirect_to_internal_ip_blocked_before_request(self, fetcher): - """A redirect Location pointing at a private IP is blocked without issuing the next request.""" - redirect_response = AsyncMock() - redirect_response.status_code = 302 - redirect_response.headers = {"location": "http://169.254.169.254/latest/meta-data/"} - redirect_response.url = httpx.URL("https://example.com/page") - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = redirect_response - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - # First validate() call (original URL) passes; the second - # call (redirect target) is the real one — we let the real - # validator catch the private IP. - from agentic_cli.tools.webfetch.validator import ValidationResult, URLValidator - real_validator = URLValidator() - with patch.object( - fetcher._validator, "validate", - side_effect=[ - ValidationResult(valid=True, resolved_ip="1.2.3.4"), - real_validator.validate("http://169.254.169.254/latest/meta-data/"), - ], - ): - result = await fetcher.fetch("https://example.com/page") + async def test_same_host_redirect_followed(self, monkeypatch): + def handler(req): + if req.url.path == "/public": + return httpx.Response(302, headers={"location": "/inner"}) + return httpx.Response(200, text="inner page", headers={"content-type": "text/html"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/public") + assert result.success is True + assert "inner page" in result.content + @pytest.mark.asyncio + async def test_redirect_to_internal_ip_blocked(self, monkeypatch): + def handler(req): + return httpx.Response(302, headers={"location": "http://169.254.169.254/latest/"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/page") assert result.success is False assert "redirect" in (result.error or "").lower() - # Only the first GET (the original URL) should have been issued. - assert mock_get.call_count == 1 @pytest.mark.asyncio - async def test_fetch_too_many_redirects(self, fetcher): - """A redirect loop longer than MAX_REDIRECTS is rejected.""" - from agentic_cli.tools.webfetch.fetcher import ContentFetcher - - def make_redirect(i: int): - r = AsyncMock() - r.status_code = 302 - r.headers = {"location": f"https://example.com/hop{i + 1}"} - r.url = httpx.URL(f"https://example.com/hop{i}") - return r - - responses = [make_redirect(i) for i in range(ContentFetcher.MAX_REDIRECTS + 1)] - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_get.side_effect = responses - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - result = await fetcher.fetch("https://example.com/hop0") - + async def test_too_many_redirects(self, monkeypatch): + def handler(req): + n = int(req.url.path.rsplit("hop", 1)[-1]) + return httpx.Response(302, headers={"location": f"https://example.com/hop{n + 1}"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/hop0") assert result.success is False assert "redirect" in (result.error or "").lower() @pytest.mark.asyncio - async def test_fetch_caching(self, fetcher): - """Test responses are cached.""" - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = "Cached content" - mock_response.headers = {"content-type": "text/html"} - mock_response.history = [] - mock_response.url = "https://example.com/page" - mock_get.return_value = mock_response - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - - result1 = await fetcher.fetch("https://example.com/page") - result2 = await fetcher.fetch("https://example.com/page") - - assert result1.from_cache is False - assert result2.from_cache is True - assert mock_get.call_count == 1 + async def test_caching(self, monkeypatch): + calls = {"n": 0} + def handler(req): + calls["n"] += 1 + return httpx.Response(200, text="Cached content", headers={"content-type": "text/html"}) + fetcher = _pinned_fetcher(monkeypatch, handler) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + r1 = await fetcher.fetch("https://example.com/page") + r2 = await fetcher.fetch("https://example.com/page") + assert r1.from_cache is False and r2.from_cache is True + assert calls["n"] == 1 @pytest.mark.asyncio - async def test_fetch_content_truncation(self, fetcher): - """Test content is truncated when too large.""" - large_content = "x" * 200000 - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.text = large_content - mock_response.headers = {"content-type": "text/plain"} - mock_response.history = [] - mock_response.url = "https://example.com/large" - mock_get.return_value = mock_response - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - with patch.object(fetcher._validator, "validate") as mock_validate: - from agentic_cli.tools.webfetch.validator import ValidationResult - mock_validate.return_value = ValidationResult(valid=True, resolved_ip="1.2.3.4") - result = await fetcher.fetch("https://example.com/large") - - assert result.success is True - assert result.truncated is True + async def test_content_truncation_streamed(self, monkeypatch): + big = "x" * 200000 + fetcher = _pinned_fetcher( + monkeypatch, + lambda req: httpx.Response(200, text=big, headers={"content-type": "text/plain"}), + max_content_bytes=102400, + ) + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://example.com/large") + assert result.success is True and result.truncated is True assert len(result.content) <= fetcher._max_content_bytes + 100 @@ -608,69 +444,29 @@ def test_convert_pdf_no_text(self, converter): assert "no extractable text" in result.lower() or "Page" in result @pytest.mark.asyncio - async def test_fetcher_pdf_uses_bytes(self): - """Test fetcher uses response.content (bytes) for PDFs.""" - from agentic_cli.tools.webfetch.fetcher import ContentFetcher - from agentic_cli.tools.webfetch.validator import URLValidator, ValidationResult - from agentic_cli.tools.webfetch.robots import RobotsTxtChecker - - fetcher = ContentFetcher( - validator=URLValidator(blocked_domains=[]), - robots_checker=RobotsTxtChecker(), + async def test_fetcher_pdf_uses_bytes(self, monkeypatch): + pdf = b"%PDF-1.4 fake pdf content" + fetcher = _pinned_fetcher( + monkeypatch, + lambda req: httpx.Response(200, content=pdf, headers={"content-type": "application/pdf"}), ) - - pdf_bytes = b"%PDF-1.4 fake pdf content" - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.content = pdf_bytes - mock_response.text = "garbled text" - mock_response.headers = {"content-type": "application/pdf"} - mock_response.history = [] - mock_response.url = "https://arxiv.org/pdf/2301.00001" - mock_get.return_value = mock_response - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - result = await fetcher.fetch("https://arxiv.org/pdf/2301.00001") - + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://arxiv.org/pdf/2301.00001") assert result.success is True - assert result.content == pdf_bytes # Should be bytes, not garbled text - assert isinstance(result.content, bytes) + assert result.content == pdf and isinstance(result.content, bytes) @pytest.mark.asyncio - async def test_fetcher_pdf_byte_limit(self): - """Test large PDF is truncated at max_pdf_bytes.""" - from agentic_cli.tools.webfetch.fetcher import ContentFetcher - from agentic_cli.tools.webfetch.validator import URLValidator, ValidationResult - from agentic_cli.tools.webfetch.robots import RobotsTxtChecker - - max_pdf = 1000 - fetcher = ContentFetcher( - validator=URLValidator(blocked_domains=[]), - robots_checker=RobotsTxtChecker(), - max_pdf_bytes=max_pdf, + async def test_fetcher_pdf_byte_limit(self, monkeypatch): + big = b"x" * 5000 + fetcher = _pinned_fetcher( + monkeypatch, + lambda req: httpx.Response(200, content=big, headers={"content-type": "application/pdf"}), + max_pdf_bytes=1000, ) - - large_pdf = b"x" * 5000 - - with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.content = large_pdf - mock_response.headers = {"content-type": "application/pdf"} - mock_response.history = [] - mock_response.url = "https://arxiv.org/pdf/2301.00001" - mock_get.return_value = mock_response - - with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock) as mock_robots: - mock_robots.return_value = True - result = await fetcher.fetch("https://arxiv.org/pdf/2301.00001") - - assert result.success is True - assert result.truncated is True - assert len(result.content) == max_pdf + with patch.object(fetcher._robots, "can_fetch", new_callable=AsyncMock, return_value=True): + result = await fetcher.fetch("https://arxiv.org/pdf/2301.00001") + assert result.success is True and result.truncated is True + assert len(result.content) == 1000 class TestWebFetchPDFSetting: @@ -842,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: @@ -1029,3 +824,23 @@ async def _inject_session_messages(self, session_id: str, messages: list[dict], # The manager itself should be the summarizer (has summarize() method) assert manager.llm_summarizer is manager assert hasattr(manager.llm_summarizer, "summarize") + + +class TestFactoryWiring: + def test_get_or_create_fetcher_shares_one_pinned_transport(self): + import agentic_cli.tools.webfetch_tool as wt + from agentic_cli.tools.webfetch.transport import PinnedTransport + from agentic_cli.config import BaseSettings + + orig_fetcher = wt._fetcher + orig_snapshot = wt._fetcher_settings_snapshot + try: + wt._fetcher = None + wt._fetcher_settings_snapshot = None + fetcher = wt.get_or_create_fetcher(BaseSettings()) + assert isinstance(fetcher._transport, PinnedTransport) + assert fetcher._robots._transport is fetcher._transport + assert fetcher._transport._validator is fetcher._validator + finally: + wt._fetcher = orig_fetcher + wt._fetcher_settings_snapshot = orig_snapshot 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 6030c48..c2423b2 100644 --- a/tests/test_workflow_controller.py +++ b/tests/test_workflow_controller.py @@ -1,9 +1,13 @@ """Tests for workflow controller factory and orchestrator routing. -Verifies that Claude models are automatically routed to the LangGraph -orchestrator, while Gemini models use whichever orchestrator is configured. +The backend is chosen purely by ``settings.orchestrator`` and is model-agnostic: +ADK runs Claude natively via the direct-API ``AnthropicLlm`` (no LiteLLM), so +Claude is no longer auto-routed to LangGraph. A model switch alone never forces an +orchestrator swap; only an orchestrator-setting change (leaving a stale manager) +does. """ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -33,7 +37,7 @@ def _make_settings(orchestrator=OrchestratorType.ADK, default_model=None, **extr settings.orchestrator = orchestrator settings.default_model = default_model settings.app_name = "test-app" - settings.langgraph_checkpointer = "memory" + settings.session_store = "memory" for k, v in extra.items(): setattr(settings, k, v) return settings @@ -63,7 +67,7 @@ def _FakeLangGraphWorkflow(model="claude-sonnet-4-5"): return wf -# --- Unit tests for helpers --- +# --- Unit tests for helpers (predicates retained for callers/back-compat) --- class TestIsClaudeModel: @@ -94,11 +98,11 @@ def test_returns_none_when_nothing_set(self): assert _resolve_effective_model(None, settings) is None -# --- Factory routing tests --- +# --- Factory routing tests (backend = orchestrator setting only) --- class TestCreateWorkflowManagerRouting: - """Test that the factory routes Claude models to LangGraph.""" + """The factory routes purely on ``settings.orchestrator`` (model-agnostic).""" def test_gemini_model_with_adk_returns_adk(self, agent_configs): """Gemini model + ADK orchestrator → ADK manager.""" @@ -112,28 +116,40 @@ def test_gemini_model_with_adk_returns_adk(self, agent_configs): mock_adk_cls.assert_called_once() assert result is mock_adk_cls.return_value - @patch("agentic_cli.workflow.langgraph.LangGraphWorkflowManager") - def test_claude_model_with_adk_returns_langgraph( - self, mock_lg_cls, agent_configs - ): - """Claude model + ADK orchestrator → auto-switches to LangGraph.""" + def test_claude_model_with_adk_returns_adk(self, agent_configs): + """Claude model + ADK orchestrator → ADK manager (native AnthropicLlm).""" settings = _make_settings(orchestrator=OrchestratorType.ADK) - result = create_workflow_manager_from_settings( - agent_configs, settings, model="claude-sonnet-4-5" + with patch( + "agentic_cli.workflow.adk.manager.GoogleADKWorkflowManager" + ) as mock_adk_cls: + result = create_workflow_manager_from_settings( + agent_configs, settings, model="claude-sonnet-4-5" + ) + mock_adk_cls.assert_called_once() + assert result is mock_adk_cls.return_value + + def test_claude_model_in_settings_returns_adk(self, agent_configs): + """Claude model in settings.default_model + ADK orchestrator → ADK manager.""" + settings = _make_settings( + orchestrator=OrchestratorType.ADK, + default_model="claude-opus-4", ) - mock_lg_cls.assert_called_once() - assert result is mock_lg_cls.return_value + with patch( + "agentic_cli.workflow.adk.manager.GoogleADKWorkflowManager" + ) as mock_adk_cls: + result = create_workflow_manager_from_settings(agent_configs, settings) + mock_adk_cls.assert_called_once() + assert result is mock_adk_cls.return_value @patch("agentic_cli.workflow.langgraph.LangGraphWorkflowManager") - def test_claude_model_in_settings_returns_langgraph( + def test_claude_model_with_langgraph_returns_langgraph( self, mock_lg_cls, agent_configs ): - """Claude model in settings.default_model → auto-switches to LangGraph.""" - settings = _make_settings( - orchestrator=OrchestratorType.ADK, - default_model="claude-opus-4", + """Claude still runs on LangGraph when that orchestrator is chosen.""" + settings = _make_settings(orchestrator=OrchestratorType.LANGGRAPH) + result = create_workflow_manager_from_settings( + agent_configs, settings, model="claude-sonnet-4-5" ) - result = create_workflow_manager_from_settings(agent_configs, settings) mock_lg_cls.assert_called_once() assert result is mock_lg_cls.return_value @@ -164,19 +180,19 @@ def test_no_model_with_adk_returns_adk(self, agent_configs): class TestWorkflowControllerOrchestratorSwap: - """Test runtime model switch triggers orchestrator swap when needed.""" + """A swap happens only when the manager type no longer matches the setting.""" def _make_controller(self, orchestrator=OrchestratorType.ADK): configs = [AgentConfig(name="test", prompt="Test")] settings = _make_settings(orchestrator=orchestrator) return WorkflowController(configs, settings) - def test_needs_swap_gemini_to_claude(self): - """ADK manager + Claude model → needs swap.""" + def test_no_swap_gemini_to_claude_on_adk(self): + """ADK manager + Claude model → no swap (Claude runs on ADK natively).""" controller = self._make_controller() controller._workflow = _FakeADKWorkflow() - assert controller._needs_orchestrator_swap("claude-sonnet-4-5") is True + assert controller._needs_orchestrator_swap("claude-sonnet-4-5") is False def test_no_swap_gemini_to_gemini(self): """ADK manager + Gemini model → no swap needed.""" @@ -185,27 +201,27 @@ def test_no_swap_gemini_to_gemini(self): assert controller._needs_orchestrator_swap("gemini-2.5-pro") is False - def test_no_swap_claude_to_claude(self): - """LangGraph manager (auto) + Claude model → no swap needed.""" + def test_swap_stale_langgraph_manager_on_adk(self): + """ADK orchestrator + a stale LangGraph manager → swap back to ADK.""" controller = self._make_controller() controller._workflow = _FakeLangGraphWorkflow("claude-sonnet-4-5") - assert controller._needs_orchestrator_swap("claude-opus-4") is False + assert controller._needs_orchestrator_swap("claude-opus-4") is True - def test_needs_swap_claude_to_gemini(self): - """LangGraph manager (auto) + Gemini model → needs swap back to ADK.""" - controller = self._make_controller() - controller._workflow = _FakeLangGraphWorkflow("claude-sonnet-4-5") - - assert controller._needs_orchestrator_swap("gemini-2.5-pro") is True - - def test_no_swap_when_langgraph_orchestrator_and_gemini(self): - """LangGraph orchestrator + Gemini model → no swap (user chose LangGraph).""" + def test_no_swap_when_langgraph_orchestrator(self): + """LangGraph orchestrator + LangGraph manager → no swap (user chose it).""" controller = self._make_controller(orchestrator=OrchestratorType.LANGGRAPH) controller._workflow = _FakeLangGraphWorkflow("gemini-2.5-pro") assert controller._needs_orchestrator_swap("gemini-2.5-flash") is False + def test_swap_stale_adk_manager_on_langgraph(self): + """LangGraph orchestrator + a stale ADK manager → swap to LangGraph.""" + controller = self._make_controller(orchestrator=OrchestratorType.LANGGRAPH) + controller._workflow = _FakeADKWorkflow("gemini-2.5-pro") + + assert controller._needs_orchestrator_swap("gemini-2.5-flash") is True + def test_no_swap_when_model_is_none(self): """No model specified → no swap.""" controller = self._make_controller() @@ -218,23 +234,39 @@ def test_no_swap_when_workflow_is_none(self): controller = self._make_controller() assert controller._needs_orchestrator_swap("claude-sonnet-4-5") is False - @patch("agentic_cli.workflow.langgraph.LangGraphWorkflowManager") - async def test_reinitialize_swaps_orchestrator_for_claude(self, mock_lg_cls): - """Switching from Gemini (ADK) to Claude triggers full manager replacement.""" - controller = self._make_controller() - old_workflow = _FakeADKWorkflow("gemini-2.5-pro") - controller._workflow = old_workflow + def test_swap_when_model_none_but_backend_stale(self): + """A-3: changing the orchestrator alone (model unchanged) must still + swap a stale manager — model=None must not short-circuit the check.""" + controller = self._make_controller(orchestrator=OrchestratorType.ADK) + controller._workflow = _FakeLangGraphWorkflow("gemini-2.5-pro") + assert controller._needs_orchestrator_swap(None) is True + + def test_swap_check_does_not_import_langgraph(self, monkeypatch): + """A-2: the swap check must route on backend_type, not by importing the + LangGraph manager — else changing the model on an ADK-only install + (no `langgraph` extra) raises ImportError.""" + import sys + + monkeypatch.setitem( + sys.modules, "agentic_cli.workflow.langgraph.manager", None + ) + controller = self._make_controller(orchestrator=OrchestratorType.ADK) + controller._workflow = _FakeADKWorkflow() + # Must not raise ImportError: + assert controller._needs_orchestrator_swap("claude-sonnet-4-5") is False - new_workflow = AsyncMock() - mock_lg_cls.return_value = new_workflow + async def test_reinitialize_claude_on_adk_reinits_in_place(self): + """ADK manager + Claude model → no swap; reinitialize in place.""" + controller = self._make_controller() + workflow = _FakeADKWorkflow("gemini-2.5-pro") + controller._workflow = workflow await controller.reinitialize(model="claude-sonnet-4-5") - # Old workflow's reinitialize should NOT have been called - old_workflow.reinitialize.assert_not_called() - # New workflow should have been created and initialized - new_workflow.initialize_services.assert_awaited_once() - assert controller._workflow is new_workflow + workflow.reinitialize.assert_awaited_once_with( + model="claude-sonnet-4-5", preserve_sessions=True + ) + assert controller._workflow is workflow async def test_reinitialize_same_family_calls_existing_reinitialize(self): """Switching Gemini → Gemini calls reinitialize on existing manager.""" @@ -266,8 +298,8 @@ async def test_reinitialize_raises_when_not_initialized(self): with pytest.raises(RuntimeError, match="Cannot reinitialize"): await controller.reinitialize(model="claude-sonnet-4-5") - async def test_swap_claude_to_gemini(self): - """Switching from Claude (LangGraph auto) back to Gemini (ADK) triggers swap.""" + async def test_reinitialize_migrates_stale_langgraph_to_adk(self): + """A stale LangGraph manager under ADK orchestrator is replaced on reinit.""" controller = self._make_controller() old_workflow = _FakeLangGraphWorkflow("claude-sonnet-4-5") controller._workflow = old_workflow @@ -282,3 +314,1084 @@ async def test_swap_claude_to_gemini(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_document_compile.py b/tests/tools/test_document_compile.py new file mode 100644 index 0000000..cc1446d --- /dev/null +++ b/tests/tools/test_document_compile.py @@ -0,0 +1,480 @@ +"""Offline tests for compile_document — subprocess and engine lookup faked.""" +from __future__ import annotations + +import os +import signal +import subprocess +from pathlib import Path + +from agentic_cli.tools.document import compile as mod +from agentic_cli.tools.document import compile_document + + +def _fake_engine(monkeypatch, name="latexmk"): + monkeypatch.setattr(mod, "_which", lambda n: f"/usr/bin/{n}" if n == name else None) + + +def test_no_engine_returns_structured_error(monkeypatch, tmp_path): + monkeypatch.setattr(mod, "_which", lambda n: None) + tex = tmp_path / "r.tex"; tex.write_text("x") + r = compile_document(str(tex)) + assert r["success"] is False + assert "No LaTeX engine" in r["error"] + assert r["engine"] is None + + +def test_missing_source_returns_error(tmp_path): + r = compile_document(str(tmp_path / "nope.tex")) + assert r["success"] is False and "not found" in r["error"] + + +def test_success_delivers_pdf_and_isolates_intermediates(monkeypatch, tmp_path): + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex" + tex.write_text("\\documentclass{article}\\begin{document}hi\\end{document}") + + def fake_run(argv, *, cwd, env, timeout): + # fake_run receives the private build dir as cwd; write artifacts there + (Path(cwd) / "r.pdf").write_bytes(b"%PDF-1.5 fake") + (Path(cwd) / "r.log").write_text("output written on r.pdf") + (Path(cwd) / "r.aux").write_text("\\relax") + return subprocess.CompletedProcess(argv, 0, stdout="ok", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + out = tmp_path / "deliver" / "report.pdf" + r = compile_document(str(tex), output_pdf=str(out), assets_dir=str(tmp_path / "assets")) + assert r["success"] is True + assert r["pdf_path"] == str(out) + assert out.is_file() # PDF promoted to delivery dir + assert not (tmp_path / "r.aux").exists() # intermediates NOT in the source dir + assert not (out.parent / "r.aux").exists() # nor beside the delivered PDF + + +def test_failure_parses_errors(monkeypatch, tmp_path): + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("bad") + + def fake_run(argv, *, cwd, env, timeout): + (Path(cwd) / "r.log").write_text("! Undefined control sequence.\nl.5 \\badcmd\n") + return subprocess.CompletedProcess(argv, 1, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex)) + assert r["success"] is False + assert any("Undefined control sequence" in e for e in r["errors"]) + assert r["pdf_path"] is None + + +def test_argv_never_enables_shell_escape(monkeypatch, tmp_path): + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["argv"] = argv; captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex), assets_dir="/tmp/assets") + assert "-shell-escape" not in captured["argv"] + assert "-halt-on-error" in captured["argv"] + assert captured["argv"][0] == "latexmk" + assert "/tmp/assets" in captured["env"]["TEXINPUTS"] + + +def test_pdflatex_uses_no_shell_escape_flag(monkeypatch, tmp_path): + _fake_engine(monkeypatch, name="pdflatex") + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["argv"] = argv + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex)) + assert captured["argv"][0] == "pdflatex" + assert "-no-shell-escape" in captured["argv"] + + +def test_timeout_returns_structured_failure(monkeypatch, tmp_path): + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + + def fake_run(argv, *, cwd, env, timeout): + raise subprocess.TimeoutExpired(argv, timeout) + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex), timeout_s=1) + assert r["success"] is False and "timed out" in r["error"] + + +# --- Fix wave 1 tests --- + +def test_run_raises_file_not_found_returns_structured_error(monkeypatch, tmp_path): + """Finding 1: _run raising FileNotFoundError must not propagate; must return failure dict.""" + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + + def fake_run(argv, *, cwd, env, timeout): + raise FileNotFoundError("latexmk: not found") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex)) + assert r["success"] is False + assert r["error"] is not None and len(r["error"]) > 0 + assert r["engine"] == "latexmk" + assert r["pdf_path"] is None + assert "duration_ms" in r + + +def test_pdf_copy_oserror_returns_structured_error(monkeypatch, tmp_path): + """Finding 2: OSError during PDF delivery must not propagate; must return failure dict.""" + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + out = tmp_path / "deliver" / "report.pdf" + + def fake_run(argv, *, cwd, env, timeout): + (Path(cwd) / "r.pdf").write_bytes(b"%PDF-1.5 fake") + (Path(cwd) / "r.log").write_text("output written on r.pdf") + return subprocess.CompletedProcess(argv, 0, stdout="ok", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + # Delivery is a no-follow atomic write; force the atomic rename to fail. + monkeypatch.setattr(mod.os, "replace", lambda src, dst: (_ for _ in ()).throw(OSError("disk full"))) + + r = compile_document(str(tex), output_pdf=str(out)) + assert r["success"] is False + assert "deliver" in r["error"].lower() or str(out) in r["error"] + assert r["pdf_path"] is None # build dir (with the PDF) is cleaned up on return + assert "duration_ms" in r + + +def test_forced_unsupported_engine_rejected(monkeypatch, tmp_path): + """Finding 3: forcing engine='xelatex' (not in _ENGINES) must be rejected → No LaTeX engine error.""" + monkeypatch.setattr(mod, "_which", lambda n: f"/usr/bin/{n}" if n == "xelatex" else None) + tex = tmp_path / "r.tex"; tex.write_text("x") + r = compile_document(str(tex), engine="xelatex") + assert r["success"] is False + assert "No LaTeX engine" in r["error"] + + +# --- Fix wave — final review (I1/I2/M1/M2/M3) --- + +def test_run_group_kill_on_timeout(monkeypatch, tmp_path): + """I1 regression: _run starts a new session and kills the process group on + timeout (now via proc.wait, not communicate).""" + import subprocess as _subprocess + + popen_kwargs: dict = {} + killpg_calls: list = [] + + class FakePopen: + pid = 42 + + def __init__(self, argv, **kwargs): + popen_kwargs.update(kwargs) + + def wait(self, timeout=None): + if timeout is not None: + raise _subprocess.TimeoutExpired([], timeout) + return 0 + + monkeypatch.setattr(mod.subprocess, "Popen", FakePopen) + monkeypatch.setattr(mod.os, "getpgid", lambda pid: pid) + monkeypatch.setattr(mod.os, "killpg", lambda pgid, sig: killpg_calls.append((pgid, sig))) + + try: + mod._run(["latexmk"], cwd=str(tmp_path), env={}, timeout=1.0) + except _subprocess.TimeoutExpired: + pass # expected + + assert popen_kwargs.get("start_new_session") is True + assert len(killpg_calls) >= 1 + assert killpg_calls[0][1] == signal.SIGKILL + + +def test_run_caps_captured_output(tmp_path): + import os as _os + import sys as _sys + + argv = [_sys.executable, "-c", "print('x' * 1_000_000)"] + r = mod._run(argv, cwd=str(tmp_path), env={"PATH": _os.environ.get("PATH", "")}, timeout=30) + assert r.returncode == 0 + assert len(r.stdout) <= mod._MAX_CAPTURE_BYTES + 1 # bounded (byte cap + trailing newline) + assert r.stdout.rstrip().endswith("x") # tail retained + + +# --- P0-3 hardening: env allowlist, no-follow delivery, capability scope --- + +def test_env_is_allowlisted_not_full_environ(monkeypatch, tmp_path): + """Env inheritance leak: the TeX process must not receive host secrets; + PATH is preserved and assets_dir still reaches TEXINPUTS.""" + _fake_engine(monkeypatch) + monkeypatch.setenv("MY_SECRET_TOKEN", "sk-must-not-leak") + monkeypatch.setenv("PATH", "/custom/bin") + monkeypatch.setenv("PERL5LIB", "/opt/perl/lib") # latexmk (Perl) needs this + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex), assets_dir="/tmp/assets") + assert "MY_SECRET_TOKEN" not in captured["env"] + assert captured["env"]["PATH"] == "/custom/bin" + assert captured["env"]["PERL5LIB"] == "/opt/perl/lib" + assert "/tmp/assets" in captured["env"]["TEXINPUTS"] + + +def test_delivery_does_not_follow_output_symlink(monkeypatch, tmp_path): + """output_pdf may be an attacker-placed symlink; delivery must not write + through it to the link target.""" + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + outside = tmp_path / "outside.pdf"; outside.write_bytes(b"ORIGINAL") + deliver = tmp_path / "deliver"; deliver.mkdir() + link = deliver / "report.pdf"; link.symlink_to(outside) + + def fake_run(argv, *, cwd, env, timeout): + (Path(cwd) / "r.pdf").write_bytes(b"%PDF-NEW") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex), output_pdf=str(link)) + assert r["success"] is True + assert outside.read_bytes() == b"ORIGINAL" # link target NOT overwritten + assert link.read_bytes() == b"%PDF-NEW" # PDF delivered to the path + assert not link.is_symlink() # symlink replaced by a real file + + +def test_capabilities_scope_assets_and_output(): + """document.compile alone under-authorizes: reading assets_dir and writing + output_pdf need explicit (optional) filesystem capabilities.""" + from agentic_cli.tools.registry import get_registry + + defn = get_registry().get("compile_document") + caps = {(c.name, c.target_arg, c.optional) for c in defn.capabilities} + assert ("document.compile", "source_path", False) in caps + assert ("filesystem.read", "assets_dir", True) in caps + assert ("filesystem.write", "output_pdf", True) in caps + + +def test_env_passes_tex_config_but_not_texinputs(monkeypatch, tmp_path): + """TeX's own TEX* config vars pass through (so a custom TEXMFHOME works), + but a caller-inherited TEXINPUTS is dropped in favor of our controlled one.""" + _fake_engine(monkeypatch) + monkeypatch.setenv("TEXMFHOME", "/home/user/texmf") + monkeypatch.setenv("TEXINPUTS", "/evil/inputs") + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex), assets_dir="/tmp/assets") + assert captured["env"].get("TEXMFHOME") == "/home/user/texmf" + assert "/evil/inputs" not in captured["env"]["TEXINPUTS"] + assert "/tmp/assets" in captured["env"]["TEXINPUTS"] + + +def test_delivery_preserves_readable_mode(monkeypatch, tmp_path): + """Delivered PDF keeps the produced file's mode (copystat), not the + mkstemp default 0600 — a report is not a secret and consumers expect it + readable.""" + import os as _os + + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + out = tmp_path / "deliver" / "report.pdf" + + def fake_run(argv, *, cwd, env, timeout): + p = Path(cwd) / "r.pdf" + p.write_bytes(b"%PDF") + _os.chmod(p, 0o644) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex), output_pdf=str(out)) + assert r["success"] is True + assert out.stat().st_mode & 0o044 # group/other readable, not 0600 + + +# --- P0-3 re-review: env scope, assets_dir separator, latexmk -norc --- + +def test_env_excludes_nontex_vars_starting_with_tex(monkeypatch, tmp_path): + """`k.startswith('TEX')` is too broad — TEXT_API_TOKEN etc. must NOT leak; + only real TeX vars (TEXMF*/known) pass.""" + _fake_engine(monkeypatch) + monkeypatch.setenv("TEXT_API_TOKEN", "sk-must-not-leak") + monkeypatch.setenv("TEXMFHOME", "/home/u/texmf") + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex)) + assert "TEXT_API_TOKEN" not in captured["env"] + assert captured["env"].get("TEXMFHOME") == "/home/u/texmf" + + +def test_env_texmf_uses_exact_allowlist_not_prefix(monkeypatch, tmp_path): + """A real TEXMF var passes; a TEXMF-prefixed non-var (potential secret) does not.""" + _fake_engine(monkeypatch) + monkeypatch.setenv("TEXMFHOME", "/home/u/texmf") + monkeypatch.setenv("TEXMF_SECRET", "leak") + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex)) + assert captured["env"].get("TEXMFHOME") == "/home/u/texmf" + assert "TEXMF_SECRET" not in captured["env"] + + +def test_assets_dir_with_path_separator_rejected(monkeypatch, tmp_path): + """assets_dir authorized as one filesystem path must not smuggle extra + TEXINPUTS roots via os.pathsep (e.g. 'assets:/etc').""" + import os as _os + + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + r = compile_document(str(tex), assets_dir=f"assets{_os.pathsep}/etc") + assert r["success"] is False + assert "assets_dir" in r["error"].lower() + + +def test_latexmk_uses_norc(monkeypatch, tmp_path): + """latexmk must run with -norc so a build-dir/home .latexmkrc (arbitrary + Perl) is not executed.""" + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["argv"] = argv + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document(str(tex)) + assert "-norc" in captured["argv"] + + +def test_option_like_source_name_not_treated_as_flag(monkeypatch, tmp_path): + """A source basename starting with '-' must be anchored (./) so the engine + parses it as a file, not an option — otherwise '-pdflatex=CMD.tex' et al. + execute arbitrary host commands (which -norc does NOT prevent).""" + _fake_engine(monkeypatch) + tex = tmp_path / "-pdflatex=evil.tex" + tex.write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["argv"] = argv + (Path(cwd) / (tex.stem + ".pdf")).write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex)) + assert "-pdflatex=evil.tex" not in captured["argv"] # never a bare option-like token + assert "./-pdflatex=evil.tex" in captured["argv"] # anchored as a path + assert r["success"] is True + + +def test_default_delivery_to_source_dir_without_intermediates(monkeypatch, tmp_path): + """No output_pdf → PDF lands at /.pdf, but the build's + intermediates never touch the source dir.""" + _fake_engine(monkeypatch) + tex = tmp_path / "r.tex"; tex.write_text("x") + + def fake_run(argv, *, cwd, env, timeout): + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + (Path(cwd) / "r.aux").write_text("aux") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + r = compile_document(str(tex)) # no output_pdf + assert r["success"] is True + assert r["pdf_path"] == str(tmp_path / "r.pdf") + assert (tmp_path / "r.pdf").is_file() # delivered to source dir + assert not (tmp_path / "r.aux").exists() # intermediate isolated in temp + + +# --- P0-3 hardening: tail-read .log to bound memory on runaway compiler logs --- + + +def test_read_log_tail_bounds_large_log(tmp_path): + log = tmp_path / "big.log" + log.write_text("START\n" + ("x" * 200_000) + "\n! Real error.\nEND\n") + out = mod._read_log_tail(log, fallback="FB") + assert "END" in out and "! Real error." in out # tail retained + assert "START" not in out # head dropped + assert len(out) <= mod._LOG_TAIL_BYTES + 3 # bounded + + +def test_read_log_tail_missing_returns_fallback(tmp_path): + assert mod._read_log_tail(tmp_path / "nope.log", fallback="FB") == "FB" + + +def test_read_log_tail_oserror_returns_fallback(monkeypatch, tmp_path): + """An existing-but-unreadable log (open raises OSError) returns fallback, not raises.""" + log = tmp_path / "x.log"; log.write_text("data") + + def boom(*a, **k): + raise PermissionError("nope") + + monkeypatch.setattr("builtins.open", boom) + assert mod._read_log_tail(log, fallback="FB") == "FB" + + +def test_texinputs_roots_are_absolute_for_relative_source(monkeypatch, tmp_path): + """Relative source_path/assets_dir must resolve to ABSOLUTE TEXINPUTS roots + (the build runs in a temp dir, so relative roots would resolve there).""" + _fake_engine(monkeypatch) + monkeypatch.chdir(tmp_path) + (tmp_path / "assets").mkdir() + (tmp_path / "r.tex").write_text("x") + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + captured["env"] = env + (Path(cwd) / "r.pdf").write_bytes(b"%PDF") + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(mod, "_run", fake_run) + compile_document("r.tex", assets_dir="assets") # relative paths + entries = [e for e in captured["env"]["TEXINPUTS"].split(os.pathsep) if e] + assert entries + assert all(os.path.isabs(e) for e in entries) + assert str((tmp_path / "assets").resolve()) in entries + + +def test_run_limits_child_file_size(monkeypatch, tmp_path): + """A runaway child writing beyond RLIMIT_FSIZE is killed (SIGXFSZ), not + allowed to fill the disk.""" + import os as _os + import sys as _sys + + monkeypatch.setattr(mod, "_RLIMIT_FSIZE_BYTES", 4096) + argv = [_sys.executable, "-c", "open('big.bin','wb').write(b'x' * 1_000_000)"] + r = mod._run(argv, cwd=str(tmp_path), env={"PATH": _os.environ.get("PATH", "")}, timeout=30) + assert r.returncode != 0 # killed by the file-size limit + assert (tmp_path / "big.bin").stat().st_size <= 4096 * 8 # capped, not 1MB diff --git a/tests/tools/test_document_compile_latex.py b/tests/tools/test_document_compile_latex.py new file mode 100644 index 0000000..5e72181 --- /dev/null +++ b/tests/tools/test_document_compile_latex.py @@ -0,0 +1,37 @@ +"""Real-engine compile — skipped unless a TeX engine is installed. + +Set LATEX_REQUIRE=1 to fail (not skip) when no engine is present. +""" +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +import pytest + +from agentic_cli.tools import compile_document + +pytestmark = pytest.mark.latex + +_HAS_ENGINE = shutil.which("latexmk") or shutil.which("pdflatex") +if not _HAS_ENGINE and os.environ.get("LATEX_REQUIRE") != "1": + pytest.skip("no LaTeX engine (latexmk/pdflatex) on PATH", allow_module_level=True) + + +def test_compiles_minimal_document(tmp_path): + build = tmp_path / "build"; build.mkdir() + tex = build / "r.tex" + tex.write_text( + "\\documentclass{article}\n\\begin{document}\n" + "Hello \\textbf{world}.\n\\end{document}\n" + ) + out = tmp_path / "deliver" / "report.pdf" + r = compile_document(str(tex), output_pdf=str(out)) + assert r["success"] is True, r + assert Path(r["pdf_path"]).is_file() and Path(r["pdf_path"]).stat().st_size > 0 + assert out.is_file() + # Isolation contract: the build runs in a private temp dir (cleaned up), so + # intermediates never land in the source dir or beside the delivered PDF. + assert not (build / "r.log").exists() # source dir stays clean + assert not (out.parent / "r.log").exists() # not beside delivered PDF diff --git a/tests/tools/test_executor_security.py b/tests/tools/test_executor_security.py new file mode 100644 index 0000000..b23076f --- /dev/null +++ b/tests/tools/test_executor_security.py @@ -0,0 +1,157 @@ +"""Security regression tests for SafePythonExecutor (P0-1). + +These encode sandbox-escape vectors that were reproducible against the +in-process AST filter. Each test asserts the escape is *blocked* — i.e. the +executor returns ``{"success": False, ...}`` rather than executing host code +or leaking the object graph. See docs/reviews/2026-07-03-* for the exploits. + +The AST/name filter is defense-in-depth, not a hard boundary (the OS sandbox +is), but the *known* escape primitives must not be reachable from the default +(no-OS-sandbox) configuration. +""" + +from __future__ import annotations + +import subprocess + +import pytest + +from agentic_cli.tools.executor import SafePythonExecutor + + +@pytest.fixture +def executor() -> SafePythonExecutor: + """Default executor: no OS sandbox, CORE_MODULES only.""" + return SafePythonExecutor() + + +# --------------------------------------------------------------------------- +# operator/functools string-getattr escape (confirmed RCE vector) +# --------------------------------------------------------------------------- + +def test_operator_attrgetter_escape_is_blocked(executor: SafePythonExecutor) -> None: + """operator.attrgetter/methodcaller take attr names as runtime strings, + bypassing the AST underscore filter. Reaching object.__subclasses__ and a + module's __globals__ gives host RCE. Must be blocked.""" + code = ( + "import operator\n" + "cls_of = operator.attrgetter('__class__')\n" + "base_of = operator.attrgetter('__base__')\n" + "get_subs = operator.methodcaller('__subclasses__')\n" + "get_globals = operator.attrgetter('__init__.__globals__')\n" + "obj_cls = base_of(cls_of(()))\n" + "hit = 'NO'\n" + "for c in get_subs(obj_cls):\n" + " try:\n" + " gl = get_globals(c)\n" + " if 'os' in gl:\n" + " gl['os'].system('echo pwned')\n" + " hit = 'YES'\n" + " break\n" + " except Exception:\n" + " pass\n" + "print(hit)\n" + ) + result = executor.execute(code) + assert result["success"] is False, f"escape was NOT blocked: {result}" + + +def test_import_operator_is_rejected(executor: SafePythonExecutor) -> None: + result = executor.execute("import operator") + assert result["success"] is False + assert "operator" in result["error"] + + +def test_import_functools_is_rejected(executor: SafePythonExecutor) -> None: + result = executor.execute("import functools") + assert result["success"] is False + assert "functools" in result["error"] + + +def test_from_operator_import_is_rejected(executor: SafePythonExecutor) -> None: + result = executor.execute("from operator import attrgetter") + assert result["success"] is False + + +# --------------------------------------------------------------------------- +# str.format object-graph traversal (info-leak vector, same root cause) +# --------------------------------------------------------------------------- + +def test_str_format_dunder_traversal_is_blocked( + executor: SafePythonExecutor, +) -> None: + """str.format field names can reference dunder attributes (the path lives + inside the format string, not as an AST attribute node), reaching the class + hierarchy and __globals__. Must be blocked.""" + code = "print('{0.__class__.__base__.__subclasses__}'.format(()))" + result = executor.execute(code) + assert result["success"] is False, f"format escape not blocked: {result}" + + +def test_str_format_map_dunder_traversal_is_blocked( + executor: SafePythonExecutor, +) -> None: + code = "print('{a.__class__}'.format_map({'a': ()}))" + result = executor.execute(code) + assert result["success"] is False + + +def test_plain_str_format_still_allowed(executor: SafePythonExecutor) -> None: + """The fix must not break ordinary format usage.""" + result = executor.execute("print('{} {name}'.format(1, name='x'))") + assert result["success"] is True, result + assert "1 x" in result["output"] + + +# --------------------------------------------------------------------------- +# sympy eval surface moved behind the OS sandbox +# --------------------------------------------------------------------------- + +def test_sympy_import_rejected_without_os_sandbox( + executor: SafePythonExecutor, +) -> None: + """sympy.sympify/parse_expr can evaluate arbitrary expressions; sympy must + require the OS sandbox like the other I/O-capable modules.""" + result = executor.execute("import sympy") + assert result["success"] is False + + +# --------------------------------------------------------------------------- +# API keys must not leak into the execution subprocess environment +# --------------------------------------------------------------------------- + +def test_api_keys_scrubbed_from_subprocess_env( + executor: SafePythonExecutor, monkeypatch: pytest.MonkeyPatch +) -> None: + """Any code that reaches the host must not find provider API keys in env.""" + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-secret") + monkeypatch.setenv("GOOGLE_API_KEY", "goog-secret") + monkeypatch.setenv("SOME_TOKEN", "tok") + monkeypatch.setenv("PATH", "/usr/bin") + + captured: dict = {} + real_run = subprocess.run + + def _capture(*args, **kwargs): + captured["env"] = kwargs.get("env") + return real_run(*args, **kwargs) + + monkeypatch.setattr(subprocess, "run", _capture) + executor.execute("print(1)") + + env = captured["env"] + assert env is not None, "executor must pass an explicit scrubbed env" + assert "ANTHROPIC_API_KEY" not in env + assert "GOOGLE_API_KEY" not in env + assert "SOME_TOKEN" not in env + assert env.get("PATH") == "/usr/bin", "non-secret env must be preserved" + + +# --------------------------------------------------------------------------- +# sanity: legitimate compute still works +# --------------------------------------------------------------------------- + +def test_ordinary_computation_still_works(executor: SafePythonExecutor) -> None: + result = executor.execute("import math\nprint(math.sqrt(16))") + assert result["success"] is True, result + assert "4.0" in result["output"] diff --git a/tests/tools/test_glob_grep_containment.py b/tests/tools/test_glob_grep_containment.py new file mode 100644 index 0000000..ab09e3f --- /dev/null +++ b/tests/tools/test_glob_grep_containment.py @@ -0,0 +1,253 @@ +"""P0-4: glob/grep patterns must not escape the permission-authorized root. + +The permission engine authorizes only the ``path`` argument. A ``pattern`` / +``file_pattern`` of ``../*`` or an absolute path would let the tool read +outside the granted directory. These tests pin the containment behavior. +""" +from __future__ import annotations + +import json +import subprocess + +import agentic_cli.tools.glob_tool as glob_mod +import agentic_cli.tools.grep_tool as grep_mod +from agentic_cli.tools.glob_tool import glob +from agentic_cli.tools.grep_tool import grep + + +class _FakeProc: + """Stand-in for a ripgrep subprocess (grep now uses Popen + a temp file).""" + pid = 4321 + returncode = 0 + + def wait(self, timeout=None): + return 0 + + +def test_glob_rejects_parent_escape(tmp_path): + root = tmp_path / "root" + root.mkdir() + (root / "inside.txt").write_text("x") + (tmp_path / "secret.txt").write_text("SECRET") + + r = glob(pattern="../*", path=str(root)) + + assert r["success"] is False + assert "secret" not in str(r).lower() + + +def test_glob_rejects_absolute_pattern(tmp_path): + root = tmp_path / "root" + root.mkdir() + r = glob(pattern="/etc/*", path=str(root)) + assert r["success"] is False + + +def test_glob_normal_pattern_still_works(tmp_path): + root = tmp_path / "root" + root.mkdir() + (root / "a.py").write_text("x") + (root / "b.txt").write_text("y") + r = glob(pattern="*.py", path=str(root)) + assert r["success"] is True + assert r["files"] == ["a.py"] + + +def test_glob_recursive_pattern_still_works(tmp_path): + root = tmp_path / "root" + (root / "sub").mkdir(parents=True) + (root / "sub" / "deep.py").write_text("x") + r = glob(pattern="**/*.py", path=str(root)) + assert r["success"] is True + assert "sub/deep.py" in r["files"] + + +def test_glob_skips_symlink_escaping_root(tmp_path): + root = tmp_path / "root" + root.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (outside / "passwd").write_text("SECRET") + (root / "link").symlink_to(outside, target_is_directory=True) + + r = glob(pattern="link/*", path=str(root)) + + # Pattern itself is legal, but the symlinked result resolves outside root. + assert r["success"] is True + assert all("passwd" not in str(f) for f in r["files"]) + + +def test_grep_rejects_parent_escape_file_pattern(tmp_path, monkeypatch): + root = tmp_path / "root" + root.mkdir() + (root / "a.txt").write_text("needle") + (tmp_path / "secret.txt").write_text("needle SECRET") + # Force the Python fallback — that's the path that follows ../ literally. + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: False) + + r = grep(pattern="needle", path=str(root), file_pattern="../*") + + assert r["success"] is False + + +def test_grep_python_skips_symlink_escaping_root(tmp_path, monkeypatch): + root = tmp_path / "root" + root.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.txt").write_text("needle SECRET") + (root / "link.txt").symlink_to(outside / "secret.txt") + (root / "real.txt").write_text("needle here") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: False) + + r = grep(pattern="needle", path=str(root)) + + assert r["success"] is True + files = {m["file"] for m in r["matches"]} + # link.txt resolves to outside/secret.txt (outside root) → must be skipped, + # even though its own path name doesn't contain "secret". + assert not any("link.txt" in f for f in files) + assert any("real.txt" in f for f in files) + + +def test_grep_ripgrep_filters_outside_root(tmp_path, monkeypatch): + """The ripgrep path (used when rg is installed) must apply the same + containment as the Python fallback — a followed symlink that rg reports + with an outside-root path must be dropped.""" + root = tmp_path / "root" + root.mkdir() + (root / "real.txt").write_text("needle here") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: True) + + outside = tmp_path / "outside" / "secret.txt" + inside = root / "real.txt" + + def fake_popen(cmd, **kwargs): + out = kwargs["stdout"] + lines = "\n".join([ + json.dumps({"type": "match", "data": { + "path": {"text": str(outside)}, "line_number": 1, + "lines": {"text": "needle SECRET\n"}}}), + json.dumps({"type": "match", "data": { + "path": {"text": str(inside)}, "line_number": 1, + "lines": {"text": "needle here\n"}}}), + ]) + "\n" + out.write(lines.encode()) + return _FakeProc() + + monkeypatch.setattr(grep_mod.subprocess, "Popen", fake_popen) + r = grep(pattern="needle", path=str(root)) + files = {m["file"] for m in r["matches"]} + assert any("real.txt" in f for f in files) + assert not any("secret.txt" in f for f in files) + + +def test_grep_ripgrep_scrubs_config_path_env(tmp_path, monkeypatch): + """A host RIPGREP_CONFIG_PATH could inject --follow (defeating containment); + it must not be inherited by the rg subprocess.""" + root = tmp_path / "root" + root.mkdir() + monkeypatch.setenv("RIPGREP_CONFIG_PATH", "/home/user/.rgrc") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: True) + captured = {} + + def fake_popen(cmd, **kwargs): + captured["env"] = kwargs.get("env") + return _FakeProc() + + monkeypatch.setattr(grep_mod.subprocess, "Popen", fake_popen) + grep(pattern="x", path=str(root)) + assert captured["env"] is not None + assert "RIPGREP_CONFIG_PATH" not in captured["env"] + + +def test_grep_ripgrep_output_bounded(tmp_path, monkeypatch): + """rg output beyond _MAX_RG_OUTPUT_BYTES is not read into memory whole; the + result is flagged truncated.""" + root = tmp_path / "root" + root.mkdir() + inside = root / "a.txt" + inside.write_text("needle") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: True) + monkeypatch.setattr(grep_mod, "_MAX_RG_OUTPUT_BYTES", 200) + + def fake_popen(cmd, **kwargs): + out = kwargs["stdout"] + line = (json.dumps({"type": "match", "data": { + "path": {"text": str(inside)}, "line_number": 1, + "lines": {"text": "needle\n"}}}) + "\n").encode() + for _ in range(50): # well over the 200-byte cap + out.write(line) + return _FakeProc() + + monkeypatch.setattr(grep_mod.subprocess, "Popen", fake_popen) + r = grep(pattern="needle", path=str(root)) + assert r["truncated"] is True + + +def test_glob_excludes_hidden_ancestor(tmp_path): + """include_hidden=False must drop results with a hidden ANCESTOR, not just + a hidden basename (e.g. .hidden/secret.txt via **/*).""" + root = tmp_path / "root" + (root / ".hidden").mkdir(parents=True) + (root / ".hidden" / "secret.txt").write_text("s") + (root / "visible.txt").write_text("v") + r = glob(pattern="**/*", path=str(root), include_hidden=False) + assert r["success"] is True + assert all(".hidden" not in f for f in r["files"]) + assert any("visible.txt" in f for f in r["files"]) + + +def test_glob_caps_scanned_matches(tmp_path, monkeypatch): + root = tmp_path / "root" + root.mkdir() + for i in range(6): + (root / f"f{i}.txt").write_text("x") + monkeypatch.setattr(glob_mod, "_MAX_SCAN", 3) + r = glob(pattern="*", path=str(root), max_results=100) + assert r["success"] is True + assert len(r["files"]) == 3 # exactly the ceiling (6 files, all pass filters) + assert r["truncated"] is True + + +def test_grep_python_skips_oversized_files(tmp_path, monkeypatch): + root = tmp_path / "root" + root.mkdir() + (root / "small.txt").write_text("needle here") + (root / "big.txt").write_text("needle " + ("x" * 1000)) + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: False) + monkeypatch.setattr(grep_mod, "_MAX_FILE_BYTES", 100) + r = grep(pattern="needle", path=str(root)) + files = {m["file"] for m in r["matches"]} + assert any("small.txt" in f for f in files) + assert not any("big.txt" in f for f in files) # oversized file skipped + + +def test_grep_python_reports_truncated_when_file_cap_hit(tmp_path, monkeypatch): + """Hitting the _MAX_FILES scan ceiling must be reported as truncated, not + silently dropped (else a partial result claims to be complete).""" + root = tmp_path / "root" + root.mkdir() + for i in range(3): + (root / f"f{i}.txt").write_text("needle") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: False) + monkeypatch.setattr(grep_mod, "_MAX_FILES", 2) # fewer than the 3 files + r = grep(pattern="needle", path=str(root)) + assert r["truncated"] is True + + +def test_grep_python_directories_do_not_consume_file_budget(tmp_path, monkeypatch): + """Directories must not count against _MAX_FILES — otherwise dirs before the + files exhaust the budget and matches are silently missed.""" + root = tmp_path / "root" + root.mkdir() + for i in range(20): + (root / f"dir{i}").mkdir() + (root / "a.txt").write_text("needle") + (root / "b.txt").write_text("needle") + monkeypatch.setattr(grep_mod, "_ripgrep_available", lambda: False) + monkeypatch.setattr(grep_mod, "_MAX_FILES", 2) # exactly the 2 real files + r = grep(pattern="needle", path=str(root)) + files = {m["file"] for m in r["matches"]} + assert any("a.txt" in f for f in files) + assert any("b.txt" in f for f in files) diff --git a/tests/tools/test_jobs.py b/tests/tools/test_jobs.py new file mode 100644 index 0000000..b859676 --- /dev/null +++ b/tests/tools/test_jobs.py @@ -0,0 +1,1830 @@ +"""Tests for the long-running job substrate (JobManager + backends + tools).""" + +from __future__ import annotations + +import json +import os +import subprocess +import time +from pathlib import Path + +import pytest + +from agentic_cli.tools.jobs import JobManager, JobRecord, JobState +from agentic_cli.tools.jobs.backends import JobBackend, default_backends + + +def _wait(jm: JobManager, job_id: str, timeout: float = 5.0) -> JobRecord: + """Poll a job until it reaches a terminal state (or timeout).""" + end = time.time() + timeout + terminal = {JobState.SUCCEEDED, JobState.FAILED, JobState.CANCELLED, JobState.UNKNOWN} + while time.time() < end: + rec = jm.get(job_id) + assert rec is not None + if rec.state in terminal: + return rec + time.sleep(0.05) + 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) + + +class TestSubprocessBackend: + def test_success(self, jm: JobManager): + rec = jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "echo hi"}) + rec = _wait(jm, rec.job_id) + assert rec.state is JobState.SUCCEEDED + assert rec.exit_code == 0 + assert "hi" in "\n".join(jm.tail(rec.job_id, 5, "stdout")) + + def test_failure_via_subshell_sentinel(self, jm: JobManager): + # `exit 3` must not skip the sentinel write (subshell isolation). + rec = jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "exit 3"}) + rec = _wait(jm, rec.job_id) + assert rec.state is JobState.FAILED + assert rec.exit_code == 3 + + def test_cwd_is_respected(self, jm: JobManager, tmp_path: Path): + workdir = tmp_path / "work" + workdir.mkdir() + rec = jm.submit( + tool="run_shell_job", + backend="subprocess", + spec={"command": "pwd", "cwd": str(workdir)}, + ) + rec = _wait(jm, rec.job_id) + out = "\n".join(jm.tail(rec.job_id, 5, "stdout")) + assert str(workdir) in out + + def test_cancel(self, jm: JobManager): + rec = jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "sleep 5"}) + time.sleep(0.2) + cancelled = jm.cancel(rec.job_id) + assert cancelled is not None and cancelled.state is JobState.CANCELLED + + +class TestConcurrency: + def test_cap_and_queue(self, jm: JobManager): + ids = [ + jm.submit(tool="run_shell_job", backend="subprocess", spec={"command": "sleep 0.4"}).job_id + for _ in range(3) + ] + jm.reconcile() + states = [jm.get(i).state for i in ids] # type: ignore[union-attr] + assert states.count(JobState.RUNNING) == 2 + assert states.count(JobState.QUEUED) == 1 + for i in ids: + _wait(jm, i) + jm.reconcile() + assert all(jm.get(i).state is JobState.SUCCEEDED for i in ids) # type: ignore[union-attr] + + +class TestInProcessBackend: + def test_returns_result(self, jm: JobManager): + rec = jm.submit( + tool="calc", backend="inprocess", + spec={"target": lambda a, b: a + b, "args": (2, 3)}, + ) + rec = _wait(jm, rec.job_id) + assert rec.state is JobState.SUCCEEDED + assert jm.result(rec.job_id) == 5 + + def test_exception_marks_failed(self, jm: JobManager): + def boom(): + raise RuntimeError("nope") + + rec = jm.submit(tool="calc", backend="inprocess", spec={"target": boom}) + rec = _wait(jm, rec.job_id) + assert rec.state is JobState.FAILED + + +class TestPersistenceAndReconcile: + def test_reload_terminal_job(self, tmp_path: Path): + base = tmp_path / "jobs" + jm = JobManager(base_dir=base, max_concurrent=2) + rec = _wait(jm, jm.submit(tool="t", backend="subprocess", spec={"command": "echo x"}).job_id) + assert rec.state is JobState.SUCCEEDED + # Fresh manager over the same dir reloads the record. + jm2 = JobManager(base_dir=base) + reloaded = jm2.get(rec.job_id) + assert reloaded is not None and reloaded.state is JobState.SUCCEEDED + + def test_vanished_process_reconciles_to_unknown(self, tmp_path: Path): + base = tmp_path / "jobs" + base.mkdir(parents=True) + # A reliably-dead PID: start a process and reap it. + p = subprocess.Popen(["true"]) + p.wait() + dead_pid = p.pid + + job_id = "deadbeef0001" + (base / job_id).mkdir() + rec = JobRecord( + job_id=job_id, tool="t", backend="subprocess", name="t", + state=JobState.RUNNING, spec={"command": "sleep 999"}, pid=dead_pid, + ) + (base / job_id / "meta.json").write_text(json.dumps(rec.to_dict())) + + # New manager: no live Popen handle, no sentinel, pid dead → UNKNOWN. + jm = JobManager(base_dir=base) + reloaded = jm.get(job_id) + assert reloaded is not None and reloaded.state is JobState.UNKNOWN + + def test_clean_removes_terminal(self, jm: JobManager): + rec = _wait(jm, jm.submit(tool="t", backend="subprocess", spec={"command": "echo x"}).job_id) + assert (jm.base_dir / rec.job_id).exists() + removed = jm.clean() + assert removed >= 1 + assert not (jm.base_dir / rec.job_id).exists() + assert jm.get(rec.job_id) is None + + +class TestManagerGuards: + def test_unknown_backend_raises(self, jm: JobManager): + with pytest.raises(ValueError): + jm.submit(tool="t", backend="does-not-exist", spec={}) + + def test_default_backends_present(self): + b = default_backends() + assert set(b) == {"subprocess", "inprocess"} + assert b["subprocess"].survives_restart is True + assert b["inprocess"].survives_restart is False + + +class TestRegistryAndTools: + def test_long_running_flag(self): + # The framework no longer ships a long-running *starter* tool (those are + # app-provided), so register a throwaway one to prove the flag threads + # through @register_tool. Observe-only tools default to long_running=False. + from agentic_cli.tools.registry import ToolCategory, get_registry, register_tool + from agentic_cli.workflow.permissions import EXEMPT + + @register_tool( + category=ToolCategory.OTHER, + capabilities=EXEMPT, + long_running=True, + description="probe long-running tool for tests", + ) + def _jobs_test_long_running_probe() -> dict: + return {"success": True} + + reg = get_registry() + assert reg.get("_jobs_test_long_running_probe").long_running is True + assert reg.get("job_status").long_running is False + + def test_tools_via_service_registry(self, tmp_path: Path): + # Exercise the observe-only tools against a job submitted directly via + # the JobManager (the starter tool is app-provided, not imported here). + from agentic_cli.tools.jobs import job_list, job_status + from agentic_cli.workflow.service_registry import JOB_MANAGER, set_service_registry + + jm = JobManager(base_dir=tmp_path / "jobs", max_concurrent=2) + token = set_service_registry({JOB_MANAGER: jm}) + try: + rec = jm.submit( + tool="run_shell_job", backend="subprocess", + spec={"command": "echo hi"}, name="greet", + ) + job_id = rec.job_id + + # Poll the public tool until terminal. + end = time.time() + 5 + while time.time() < end: + st = job_status(job_id) + if st["state"] in ("succeeded", "failed"): + break + time.sleep(0.05) + assert st["success"] is True + assert st["state"] == "succeeded" + # job_status returns the result once finished, so the agent needs no + # separate job_result call (minimal tool surface). + assert "result" in st + assert st["result"]["exit_code"] == 0 + + listing = job_list() + assert listing["success"] is True + assert any(j["job_id"] == job_id for j in listing["jobs"]) + finally: + token.var.reset(token) + + def test_minimal_bundle_is_just_job_status(self): + from agentic_cli.tools import JOB_MANAGEMENT_TOOLS, JOB_TOOLS, job_status + + assert JOB_TOOLS == [job_status] + assert len(JOB_MANAGEMENT_TOOLS) == 5 + + def test_tools_error_without_manager(self): + from agentic_cli.tools.jobs import job_status + from agentic_cli.workflow.service_registry import clear_service_registry + + token = clear_service_registry() + try: + res = job_status("whatever") + assert res["success"] is False + assert "not available" in res["error"] + finally: + token.var.reset(token) + + +class TestResumeMetadata: + """Phase-2 association layer: resume metadata on jobs (no behavior change).""" + + def test_defaults_are_fire_and_forget(self, jm: JobManager): + rec = jm.submit(tool="t", backend="subprocess", spec={"command": "echo x"}) + assert rec.resume_on_complete is False + assert rec.session_id is None and rec.user_id is None + assert rec.call_id is None and rec.resumed is False + _wait(jm, rec.job_id) + # A non-resume job never shows up as awaiting resume. + assert jm.awaiting_resume() == [] + + def test_submit_stores_explicit_metadata_and_persists(self, tmp_path: Path): + base = tmp_path / "jobs" + jm = JobManager(base_dir=base, max_concurrent=2) + rec = jm.submit( + tool="run_shell_job", backend="subprocess", spec={"command": "echo x"}, + resume_on_complete=True, call_id="fc-123", call_name="run_shell_job", + session_id="sess-1", user_id="user-1", + ) + assert rec.resume_on_complete is True + assert rec.call_id == "fc-123" + assert rec.call_name == "run_shell_job" + assert rec.session_id == "sess-1" and rec.user_id == "user-1" + # Survives a reload (persisted in meta.json). + jm2 = JobManager(base_dir=base) + reloaded = jm2.get(rec.job_id) + assert reloaded is not None + assert reloaded.resume_on_complete is True + assert reloaded.call_id == "fc-123" + assert reloaded.session_id == "sess-1" + + def test_autofills_session_user_from_workflow_service(self, jm: JobManager): + from types import SimpleNamespace + + from agentic_cli.workflow.service_registry import WORKFLOW, set_service_registry + + fake_wf = SimpleNamespace(active_session_id="sess-A", active_user_id="user-A") + token = set_service_registry({WORKFLOW: fake_wf}) + try: + rec = jm.submit( + tool="run_shell_job", backend="subprocess", spec={"command": "echo x"}, + resume_on_complete=True, call_id="fc-1", + ) + finally: + token.var.reset(token) + assert rec.session_id == "sess-A" + assert rec.user_id == "user-A" + # call_name defaults to the tool name when resume is requested. + assert rec.call_name == "run_shell_job" + + def test_missing_context_is_nonfatal(self, jm: JobManager): + from agentic_cli.workflow.service_registry import clear_service_registry + + token = clear_service_registry() + try: + # No workflow service → can't fill session/user, but the job still runs. + rec = jm.submit( + tool="t", backend="subprocess", spec={"command": "echo x"}, + resume_on_complete=True, call_id="fc-1", + ) + finally: + token.var.reset(token) + assert rec.resume_on_complete is True + assert rec.session_id is None and rec.user_id is None + + def test_awaiting_resume_and_mark_resumed(self, tmp_path: Path): + base = tmp_path / "jobs" + jm = JobManager(base_dir=base, max_concurrent=2) + rec = jm.submit( + tool="run_shell_job", backend="subprocess", spec={"command": "echo x"}, + resume_on_complete=True, call_id="fc-1", session_id="s", user_id="u", + ) + # Also submit a fire-and-forget job that must never appear. + other = jm.submit(tool="t", backend="subprocess", spec={"command": "echo y"}) + _wait(jm, rec.job_id) + _wait(jm, other.job_id) + + awaiting = jm.awaiting_resume() + assert [r.job_id for r in awaiting] == [rec.job_id] + + # Marking it resumed (durably) drops it from the list. + jm.mark_resumed(rec.job_id) + assert jm.awaiting_resume() == [] + 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/tools/test_os_sandbox.py b/tests/tools/test_os_sandbox.py index a97301d..e8f1123 100644 --- a/tests/tools/test_os_sandbox.py +++ b/tests/tools/test_os_sandbox.py @@ -529,8 +529,9 @@ def test_enabled_policy_wraps_python_command(self): # subprocess.run was called with shell=True (wrapped command) assert mock_run.call_args.kwargs.get("shell") is True - def test_enabled_policy_wrap_failure_fails_closed(self): - policy = OSSandboxPolicy(enabled=True) + def test_enabled_policy_wrap_failure_strict_fails_closed(self): + # strict=True: a wrap failure must refuse rather than run unwrapped. + policy = OSSandboxPolicy(enabled=True, strict=True) executor = SafePythonExecutor(os_sandbox_policy=policy) mock_os_sandbox = MagicMock() @@ -550,13 +551,13 @@ def test_enabled_policy_wrap_failure_fails_closed(self): ): result = executor.execute("1 + 1") - # Wrap failure must fail closed rather than running unwrapped. assert result["success"] is False assert "sandbox" in (result.get("error") or "").lower() mock_run.assert_not_called() - def test_enabled_policy_no_real_sandbox_fails_closed(self): - policy = OSSandboxPolicy(enabled=True) + def test_enabled_policy_no_real_sandbox_strict_fails_closed(self): + # strict=True: no backend must refuse. + policy = OSSandboxPolicy(enabled=True, strict=True) executor = SafePythonExecutor(os_sandbox_policy=policy) mock_os_sandbox = MagicMock() @@ -574,3 +575,35 @@ def test_enabled_policy_no_real_sandbox_fails_closed(self): assert "sandbox" in (result.get("error") or "").lower() mock_os_sandbox.wrap_python_command.assert_not_called() mock_run.assert_not_called() + + def test_enabled_policy_no_real_sandbox_non_strict_falls_back(self): + # Default (non-strict): no backend falls back to the in-process executor + # (which is restricted to CORE_MODULES) rather than refusing. + policy = OSSandboxPolicy(enabled=True, strict=False) + executor = SafePythonExecutor(os_sandbox_policy=policy) + + mock_os_sandbox = MagicMock() + mock_os_sandbox.sandbox_type = "none" + + mock_proc = MagicMock() + mock_proc.stdout = ( + "\n__AGENTIC_EXECUTOR_RESULT_SENTINEL__\n" + '{"success": true, "output": "", "result": "2", ' + '"error": "", "execution_time_ms": 0}' + ) + mock_proc.stderr = "" + mock_proc.returncode = 0 + + with patch( + "agentic_cli.tools.executor.subprocess.run", return_value=mock_proc + ) as mock_run, patch( + "agentic_cli.tools.shell.os_sandbox.get_os_sandbox", + return_value=mock_os_sandbox, + ): + result = executor.execute("1 + 1") + + assert result["success"] is True + mock_os_sandbox.wrap_python_command.assert_not_called() + # Fell back to a plain (unwrapped) subprocess. + mock_run.assert_called_once() + assert mock_run.call_args.kwargs.get("shell") is not True diff --git a/tests/tools/test_registry_identity.py b/tests/tools/test_registry_identity.py new file mode 100644 index 0000000..2c42b2c --- /dev/null +++ b/tests/tools/test_registry_identity.py @@ -0,0 +1,1403 @@ +"""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} + + 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/tools/test_sandbox.py b/tests/tools/test_sandbox.py index 9df5630..44f0121 100644 --- a/tests/tools/test_sandbox.py +++ b/tests/tools/test_sandbox.py @@ -14,7 +14,7 @@ import pytest -from agentic_cli.tools.sandbox.models import ExecutionResult +from agentic_cli.tools.sandbox.models import ExecutionResult, SessionStatus from agentic_cli.tools.sandbox.backends.base import SandboxBackend from agentic_cli.tools.sandbox.manager import SandboxManager, SandboxSession from tests.conftest import MockContext @@ -27,19 +27,22 @@ class MockSandboxBackend(SandboxBackend): """Test backend that returns configurable results.""" + backend_name = "mock" + def __init__(self, result: ExecutionResult | None = None) -> None: self._result = result or ExecutionResult(success=True, stdout="ok\n", result="42") self._sessions: set[str] = set() self.execute_calls: list[dict] = [] self.reset_calls: list[str] = [] - def execute(self, code, session_id, timeout_seconds=120, working_dir=None): + def execute(self, code, session_id, timeout_seconds=120, working_dir=None, inputs=None): self._sessions.add(session_id) self.execute_calls.append({ "code": code, "session_id": session_id, "timeout_seconds": timeout_seconds, "working_dir": working_dir, + "inputs": inputs, }) return self._result @@ -54,6 +57,34 @@ def has_session(self, session_id): return session_id in self._sessions +# --------------------------------------------------------------------------- +# SessionStatus +# --------------------------------------------------------------------------- + +class TestSessionStatus: + def test_default_backend_status_ready_when_session_exists(self, tmp_path): + backend = MockSandboxBackend() + backend._sessions.add("s1") + st = backend.session_status("s1") + assert isinstance(st, SessionStatus) + assert st.state == "ready" + assert st.backend == "mock" + assert st.session_id == "s1" + + def test_default_backend_status_absent_when_missing(self): + backend = MockSandboxBackend() + assert backend.session_status("nope").state == "absent" + + def test_list_sessions_includes_state(self): + with MockContext() as ctx: + backend = MockSandboxBackend() + mgr = SandboxManager(ctx.settings, backend=backend) + mgr.execute("x=1", session_id="s1") + rows = mgr.list_sessions() + assert rows[0]["session_id"] == "s1" + assert rows[0]["state"] == "ready" + + # --------------------------------------------------------------------------- # ExecutionResult # --------------------------------------------------------------------------- @@ -131,6 +162,36 @@ def test_max_sessions_enforced(self, tmp_path): assert len(mgr.list_sessions()) == 2 mgr.cleanup() + def test_failed_start_does_not_consume_session_slots(self, tmp_path): + """A failed start (e.g. Docker unavailable) must not leave phantom + session metadata that fills sandbox_max_sessions with no real session.""" + class _FailingBackend(SandboxBackend): + backend_name = "failing" + def execute(self, code, session_id, timeout_seconds=120, working_dir=None, inputs=None): + return ExecutionResult(success=False, error="Docker sandbox backend unavailable") + def reset_session(self, session_id): pass + def cleanup(self): pass + def has_session(self, session_id): return False + + with MockContext(sandbox_max_sessions=2) as ctx: + mgr = SandboxManager(ctx.settings, backend=_FailingBackend()) + for i in range(5): + r = mgr.execute("x = 1", session_id=f"s{i}") + assert r.success is False + assert "Maximum sessions" not in r.error # never blocked by phantom slots + assert mgr.list_sessions() == [] # no slots consumed + mgr.cleanup() + + def test_rejects_non_positive_timeout(self, tmp_path): + """A non-positive timeout must be rejected up front, not fall through to + the backend where it would raise and orphan a container.""" + with MockContext() as ctx: + mgr = SandboxManager(ctx.settings, backend=MockSandboxBackend()) + r = mgr.execute("x = 1", timeout_seconds=-5) + assert r.success is False + assert "timeout" in r.error.lower() + mgr.cleanup() + def test_reset_active_session(self, tmp_path): with MockContext() as ctx: backend = MockSandboxBackend() @@ -202,7 +263,7 @@ def test_working_dir_created(self, tmp_path): class TestSandboxTools: def test_sandbox_execute_success(self, tmp_path): - with MockContext() as ctx: + with MockContext(stateful_executor_backend="local") as ctx: from agentic_cli.workflow.service_registry import set_service_registry backend = MockSandboxBackend( @@ -223,7 +284,7 @@ def test_sandbox_execute_success(self, tmp_path): mgr.cleanup() def test_sandbox_execute_no_manager(self, tmp_path): - with MockContext(): + with MockContext(stateful_executor_backend="local"): from agentic_cli.workflow.service_registry import set_service_registry token = set_service_registry({}) @@ -235,12 +296,129 @@ def test_sandbox_execute_no_manager(self, tmp_path): finally: token.var.reset(token) + def test_sandbox_execute_disabled_by_default(self, tmp_path): + """The Jupyter kernel is unsandboxed host RCE — it must be opt-in.""" + with MockContext(): + from agentic_cli.tools.sandbox import sandbox_execute + result = sandbox_execute("print('hi')") + assert result["success"] is False + assert "stateful" in result["error"].lower() or "backend" in result["error"].lower() + + def test_factory_tool_respects_enabled_flag(self, tmp_path): + """CRITICAL regression: the workflow uses the factory-bound tool + (base_manager wires make_sandbox_tool), which must honor the + stateful_executor_backend opt-in — not just the module-level tool.""" + from agentic_cli.tools.factories import make_sandbox_tool + + with MockContext(stateful_executor_backend="none") as ctx: + mgr = SandboxManager(ctx.settings, backend=MockSandboxBackend()) + tool = make_sandbox_tool(mgr) + r = tool(code="x = 1") + assert r["success"] is False, "factory tool executed despite disabled flag" + assert "stateful" in r["error"].lower() or "backend" in r["error"].lower() + mgr.cleanup() + + with MockContext(stateful_executor_backend="local") as ctx: + mgr = SandboxManager(ctx.settings, backend=MockSandboxBackend()) + tool = make_sandbox_tool(mgr) + r = tool(code="x = 1") + assert r["success"] is True + mgr.cleanup() + + def test_factory_tool_namespaces_default_session_to_conversation(self, tmp_path): + """HIGH regression: with session_id='default', distinct conversations + must NOT share one kernel/workspace — the default is namespaced to the + active conversation. An explicit session_id is still honored verbatim.""" + from agentic_cli.tools.factories import make_sandbox_tool + + class _WF: + active_session_id = "conv-abc" + + with MockContext(stateful_executor_backend="local") as ctx: + backend = MockSandboxBackend() + mgr = SandboxManager(ctx.settings, backend=backend) + tool = make_sandbox_tool(mgr, _WF()) + + tool(code="x = 1") # default -> namespaced to the conversation + assert backend.execute_calls[-1]["session_id"] != "default" + assert "conv-abc" in backend.execute_calls[-1]["session_id"] + + tool(code="y = 1", session_id="explicit") # explicit id untouched + assert backend.execute_calls[-1]["session_id"] == "explicit" + mgr.cleanup() + + def test_disabled_message_is_backend_aware(self, tmp_path): + """The disabled-tool message must reflect the selected backend.""" + from agentic_cli.tools.sandbox import sandbox_execute + with MockContext(stateful_executor_backend="none"): + err = sandbox_execute("print('hi')")["error"].lower() + assert "stateful_executor_backend" in err + + def test_description_makes_no_false_network_claim(self): + """The tool does NOT block network — the description must not claim it + does, since a false safety claim misleads both the model and the user.""" + from agentic_cli.tools.registry import get_registry + + definition = get_registry().get("sandbox_execute") + assert definition is not None + desc = definition.description.lower() + assert "blocked" not in desc + assert "host" in desc # honest: runs with host privileges + + def test_capability_distinct_from_execute_python(self): + """sandbox_execute must NOT share python.exec with execute_python — else + an 'Allow always' for the safe stateless tool silently authorizes the + unsandboxed stateful kernel.""" + from agentic_cli.tools.registry import get_registry + from agentic_cli.tools.execution_tools import execute_python # noqa: F401 + from agentic_cli.workflow.permissions.matchers import _cap_matches + + reg = get_registry() + sandbox_cap_names = [c.name for c in reg.get("sandbox_execute").capabilities] + exec_caps = [c.name for c in reg.get("execute_python").capabilities] + + assert exec_caps == ["python.exec"] + assert "python.exec.stateful" in sandbox_cap_names + # An execute_python grant (rule 'python.exec') must not cover it. + assert _cap_matches("python.exec", "python.exec.stateful") is False + # A deliberate broad 'python.*' grant still covers both. + assert _cap_matches("python.*", "python.exec.stateful") is True + + def test_inputs_declares_filesystem_read(self): + """sandbox_execute must declare filesystem.read for its inputs arg, + so each staged file path is permission-checked identically to read_file.""" + from agentic_cli.tools.registry import get_registry + caps = {(c.name, c.target_arg) for c in get_registry().get("sandbox_execute").capabilities} + assert ("python.exec.stateful", None) in caps + assert ("filesystem.read", "inputs") in caps + # --------------------------------------------------------------------------- # SandboxCommand # --------------------------------------------------------------------------- +class TestBackendSelection: + def test_create_jupyter_docker_backend(self): + from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend + with MockContext(stateful_executor_backend="docker") as ctx: + mgr = SandboxManager(ctx.settings) + backend = mgr._create_backend("docker") + assert isinstance(backend, JupyterDockerBackend) + + def test_unknown_backend_raises(self): + with MockContext() as ctx: + mgr = SandboxManager(ctx.settings) + with pytest.raises(ValueError): + mgr._create_backend("nope") + + def test_description_mentions_backend_dependent_isolation(self): + from agentic_cli.tools.registry import get_registry + definition = get_registry().get("sandbox_execute") + desc = definition.description.lower() + assert "jupyter_docker" in desc or "backend" in desc + + class TestSandboxCommand: @pytest.fixture() def mock_app(self, tmp_path): @@ -311,6 +489,35 @@ async def test_reset_default_session(self, mock_app): await cmd.execute("reset", mock_app) mock_app.session.add_success.assert_called_once() assert len(mgr.list_sessions()) == 0 + + @pytest.mark.asyncio + async def test_reset_no_arg_resets_single_namespaced_session(self, mock_app): + """No-arg reset targets the current sandbox even when it is namespaced + (conv-), not a literal 'default' — the regression from namespacing.""" + with MockContext() as ctx: + from agentic_cli.cli.builtin_commands import SandboxCommand + + mgr = self._make_manager(ctx, sessions=["conv-abc"]) + mock_app.workflow.sandbox_manager = mgr + cmd = SandboxCommand() + + await cmd.execute("reset", mock_app) + mock_app.session.add_success.assert_called_once() + assert len(mgr.list_sessions()) == 0 + mgr.cleanup() + + @pytest.mark.asyncio + async def test_reset_no_arg_multiple_sessions_asks_to_specify(self, mock_app): + with MockContext() as ctx: + from agentic_cli.cli.builtin_commands import SandboxCommand + + mgr = self._make_manager(ctx, sessions=["conv-a", "conv-b"]) + mock_app.workflow.sandbox_manager = mgr + cmd = SandboxCommand() + + await cmd.execute("reset", mock_app) + mock_app.session.add_warning.assert_called_once() # guidance, not a wrong reset + assert len(mgr.list_sessions()) == 2 # nothing reset mgr.cleanup() @pytest.mark.asyncio @@ -369,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.""" @@ -504,6 +711,34 @@ def test_safe_math_still_works(self, tmp_path): backend.cleanup() +class TestStageInputs: + def test_copies_to_inputs_subdir(self, tmp_path): + from agentic_cli.tools.sandbox.manager import stage_inputs + src = tmp_path / "sales.csv"; src.write_text("a,b\n1,2\n") + sess = tmp_path / "sess"; sess.mkdir() + stage_inputs(sess, [str(src)]) + assert (sess / "inputs" / "sales.csv").read_text() == "a,b\n1,2\n" + + def test_missing_file_raises(self, tmp_path): + from agentic_cli.tools.sandbox.manager import stage_inputs + sess = tmp_path / "sess"; sess.mkdir() + with pytest.raises(ValueError): + stage_inputs(sess, [str(tmp_path / "nope.csv")]) + + def test_none_session_dir_raises_value_error(self, tmp_path): + from agentic_cli.tools.sandbox.manager import stage_inputs + with pytest.raises(ValueError, match="session_dir is required"): + stage_inputs(None, ["/x"]) + + def test_basename_collision_raises(self, tmp_path): + from agentic_cli.tools.sandbox.manager import stage_inputs + (tmp_path / "a").mkdir(); (tmp_path / "b").mkdir() + (tmp_path / "a" / "x.csv").write_text("1"); (tmp_path / "b" / "x.csv").write_text("2") + sess = tmp_path / "sess"; sess.mkdir() + with pytest.raises(ValueError): + stage_inputs(sess, [str(tmp_path / "a" / "x.csv"), str(tmp_path / "b" / "x.csv")]) + + class TestJupyterLocalBackend: """Integration tests for JupyterLocalBackend.""" @@ -577,3 +812,19 @@ def test_has_session(self, tmp_path): assert backend.has_session("test") is False finally: backend.cleanup() + + def test_outputs_dir_pre_created(self, tmp_path): + """outputs/ must be pre-created so open('outputs/x','w') works.""" + from agentic_cli.tools.sandbox.backends.jupyter_local import JupyterLocalBackend + + backend = JupyterLocalBackend() + try: + result = backend.execute( + "open('outputs/t.txt','w').write('hi'); print('ok')", + session_id="test", + working_dir=tmp_path, + ) + assert result.success is True + assert (tmp_path / "outputs" / "t.txt").exists() + finally: + backend.cleanup() diff --git a/tests/tools/test_sandbox_container_runtime.py b/tests/tools/test_sandbox_container_runtime.py new file mode 100644 index 0000000..643ca80 --- /dev/null +++ b/tests/tools/test_sandbox_container_runtime.py @@ -0,0 +1,82 @@ +"""Tests for the docker-CLI container runtime seam.""" + +from agentic_cli.tools.sandbox.backends.container_runtime import ( + ContainerSpec, Mount, DockerContainerRuntime, +) + + +def _spec(**kw): + base = dict( + image="img:tag", + name="agentic-sbx-s1", + command=["python", "/opt/agentic_sandbox/driver.py"], + network="none", + memory_mb=2048, + cpus=2.0, + pids_limit=256, + user="1000:1000", + env={"HOME": "/tmp"}, + mounts=[Mount("/host/ws", "/workspace", read_only=False), + Mount("/host/driver.py", "/opt/agentic_sandbox/driver.py")], + labels={"agentic-sandbox": "1"}, + ) + base.update(kw) + return ContainerSpec(**base) + + +def test_argv_carries_isolation_flags(): + argv = DockerContainerRuntime.build_run_argv(_spec(), "docker") + joined = " ".join(argv) + assert argv[:3] == ["docker", "run", "--rm"] + assert "--network none" in joined + assert "--read-only" in argv + assert "--tmpfs" in argv and "/tmp" in argv + assert "ALL" in argv and "--cap-drop" in argv + assert "no-new-privileges" in joined + assert "--memory" in argv and "2048m" in argv + assert "--memory-swap" in argv # swap disabled == memory + assert "--pids-limit" in argv and "256" in argv + assert "--cpus" in argv and "2.0" in argv + assert "--user" in argv and "1000:1000" in argv + assert "-e" in argv and "HOME=/tmp" in argv + assert "/host/ws:/workspace" in argv + assert "/host/driver.py:/opt/agentic_sandbox/driver.py:ro" in argv + assert "--label" in argv and "agentic-sandbox=1" in argv + # image + command are last + assert argv[-3:] == ["img:tag", "python", "/opt/agentic_sandbox/driver.py"] + + +def test_memory_swap_equals_memory(): + argv = DockerContainerRuntime.build_run_argv(_spec(memory_mb=512), "docker") + i = argv.index("--memory-swap") + assert argv[i + 1] == "512m" + + +def test_empty_user_omits_flag(): + argv = DockerContainerRuntime.build_run_argv(_spec(user=""), "docker") + assert "--user" not in argv + + +def test_start_invokes_popen_with_argv(monkeypatch): + seen = {} + + class FakePopen: + def __init__(self, argv, **kw): + seen["argv"] = argv + seen["kw"] = kw + self.stdin = object() + self.stdout = object() + self.stderr = object() + self.pid = 4321 + def poll(self): + return None + + monkeypatch.setattr( + "agentic_cli.tools.sandbox.backends.container_runtime.subprocess.Popen", + FakePopen, + ) + rt = DockerContainerRuntime("docker") + handle = rt.start(_spec()) + assert seen["argv"][0] == "docker" + assert handle.name == "agentic-sbx-s1" + assert handle.poll() is None diff --git a/tests/tools/test_sandbox_container_session.py b/tests/tools/test_sandbox_container_session.py new file mode 100644 index 0000000..fab8506 --- /dev/null +++ b/tests/tools/test_sandbox_container_session.py @@ -0,0 +1,193 @@ +"""Tests for ContainerSession using an in-memory fake container handle.""" + +import io +import json +import queue +import threading + +import pytest + +from agentic_cli.tools.sandbox.backends.jupyter_docker import ( + ContainerSession, SandboxStartError, +) + + +class FakeStdout: + """Blocking line iterator fed by the test; yields until closed (EOF).""" + def __init__(self): + self._q: queue.Queue = queue.Queue() + self._closed = False + def feed(self, line: str): + self._q.put(line) + def eof(self): + self._q.put(None) + def __iter__(self): + return self + def __next__(self): + item = self._q.get() + if item is None: + raise StopIteration + return item + + +class FakeHandle: + def __init__(self): + self.name = "agentic-sbx-test" + self.stdin = io.StringIO() + self.stdout = FakeStdout() + self.stderr = FakeStdout() + self._dead = False + def poll(self): + return 137 if self._dead else None + + +def _session(handle, **kw): + calls = {"interrupt": 0, "kill": 0} + def interrupt(): + calls["interrupt"] += 1 + def kill(): + calls["kill"] += 1 + handle._dead = True + handle.stdout.eof() + working_dir = kw.pop("working_dir", None) + s = ContainerSession( + session_id="test", handle=handle, interrupt=interrupt, kill=kill, + working_dir=working_dir, start_timeout=2, interrupt_grace=1, **kw, + ) + return s, calls + + +def test_wait_ready_consumes_ready_line(): + h = FakeHandle() + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + assert s.status == "ready" + + +def test_wait_ready_times_out_and_kills(): + h = FakeHandle() + s, calls = _session(h) + with pytest.raises(SandboxStartError): + s.wait_ready() # nothing fed -> start_timeout + assert calls["kill"] == 1 + + +def test_execute_returns_result(): + h = FakeHandle() + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + h.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "42\n", + "stderr": "", "result": "42", "artifacts": [], + "execution_time": 0.1, "error": ""}) + "\n") + result = s.execute("print(42)", timeout=5) + assert result.success is True + assert result.stdout == "42\n" + assert s.status == "ready" + + +def test_execute_timeout_interrupts_then_kills(): + h = FakeHandle() + s, calls = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + result = s.execute("while True: pass", timeout=1) # no result fed + assert result.success is False + assert "timed out" in result.error.lower() + assert calls["interrupt"] == 1 + assert calls["kill"] == 1 + assert s.status == "dead" + + +def test_execute_translates_workspace_artifact_paths(tmp_path): + h = FakeHandle() + s, _ = _session(h, working_dir=tmp_path) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + h.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "", "stderr": "", + "result": None, "artifacts": ["/workspace/artifacts/plot_0.png"], + "execution_time": 0.1, "error": ""}) + "\n") + result = s.execute("plot()", timeout=5) + assert result.artifacts == [str(tmp_path / "artifacts" / "plot_0.png")] + + +# Fix 1: stderr drain thread +def test_stderr_drain_thread_exists_and_drains(): + h = FakeHandle() + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + # feed stderr lines, then close + h.stderr.feed("kernel warning 1\n") + h.stderr.feed("kernel warning 2\n") + h.stderr.feed("kernel warning 3\n") + h.stderr.eof() + # a _stderr_reader thread must exist + assert hasattr(s, "_stderr_reader") and s._stderr_reader is not None + # it must finish after EOF (bounded wait) + s._stderr_reader.join(timeout=2) + assert not s._stderr_reader.is_alive(), "stderr drain thread should stop after EOF" + + +# Fix 2b: dead-container early poll() check +def test_execute_on_dead_container_returns_error(): + h = FakeHandle() + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + assert s.status == "ready" + # mark container dead before execute + h._dead = True + h.stdout.eof() + result = s.execute("print(1)", timeout=5) + assert result.success is False + assert "exited unexpectedly" in result.error.lower() + assert s.status == "dead" + + +def test_execute_skips_forged_untokened_result(): + """A forged result (no/wrong token) injected into the driver's stdout — e.g. + by user code writing to /proc//fd/1 — must be skipped; only the + driver's authenticated result is accepted.""" + h = FakeHandle() + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready", "token": "SECRET"}) + "\n") + s.wait_ready() + # forged first (no token), then the real tokened result + h.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "FORGED\n", + "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.0, "error": ""}) + "\n") + h.stdout.feed(json.dumps({"type": "result", "token": "SECRET", "success": True, + "stdout": "REAL\n", "stderr": "", "result": None, + "artifacts": [], "execution_time": 0.0, "error": ""}) + "\n") + result = s.execute("print(1)", timeout=5) + assert result.success is True + assert result.stdout == "REAL\n" # forged line skipped + + +def test_execute_kills_container_on_unexpected_error(): + """An unexpected host-side error mid-execute (e.g. a bad timeout raising in + queue.get) must kill the container, not orphan it.""" + h = FakeHandle() + s, calls = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + result = s.execute("print(1)", timeout=-5) # negative -> queue.get raises ValueError + assert result.success is False + assert calls["kill"] == 1 # container killed, not leaked + assert s.status == "dead" + + +def test_execute_reports_likely_oom_on_137_exit(): + """Exit 137 (128+SIGKILL) is the cgroup OOM-killer signature; the error + should hint at OOM and include the code rather than a bare 'exited'.""" + h = FakeHandle() # FakeHandle.poll() returns 137 when dead + s, _ = _session(h) + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + s.wait_ready() + h._dead = True + h.stdout.eof() + result = s.execute("x = bytearray(10**10)", timeout=5) + assert result.success is False + assert "137" in result.error and "memory" in result.error.lower() diff --git a/tests/tools/test_sandbox_detect.py b/tests/tools/test_sandbox_detect.py new file mode 100644 index 0000000..0c25bad --- /dev/null +++ b/tests/tools/test_sandbox_detect.py @@ -0,0 +1,65 @@ +"""Tests for Docker availability detection.""" + +import subprocess + +import pytest + +from agentic_cli.tools.sandbox.backends import detect + + +@pytest.fixture(autouse=True) +def _clear_cache(): + detect.clear_detection_cache() + yield + detect.clear_detection_cache() + + +def test_available_when_cli_and_daemon_ok(monkeypatch): + monkeypatch.setattr(detect.shutil, "which", lambda exe: "/usr/bin/docker" if exe == "docker" else None) + monkeypatch.setattr(detect, "_probe_daemon", lambda exe: True) + a = detect.detect_docker() + assert a.available is True + assert a.runtime == "docker" + + +def test_unavailable_when_no_cli(monkeypatch): + monkeypatch.setattr(detect.shutil, "which", lambda exe: None) + a = detect.detect_docker() + assert a.available is False + assert a.runtime == "" + assert "not found" in a.detail.lower() + + +def test_unavailable_when_daemon_down(monkeypatch): + monkeypatch.setattr(detect.shutil, "which", lambda exe: "/usr/bin/docker" if exe == "docker" else None) + monkeypatch.setattr(detect, "_probe_daemon", lambda exe: False) + a = detect.detect_docker() + assert a.available is False + assert "daemon" in a.detail.lower() + + +def test_falls_through_to_podman_when_docker_daemon_down(monkeypatch): + """docker CLI present but daemon down must NOT stop the search — podman is + still probed (the 'docker, then podman' design).""" + monkeypatch.setattr(detect.shutil, "which", lambda exe: f"/usr/bin/{exe}") + monkeypatch.setattr(detect, "_probe_daemon", lambda exe: exe == "podman") + a = detect.detect_docker() + assert a.available is True + assert a.runtime == "podman" + + +def test_combined_detail_when_all_runtimes_down(monkeypatch): + monkeypatch.setattr(detect.shutil, "which", lambda exe: f"/usr/bin/{exe}") + monkeypatch.setattr(detect, "_probe_daemon", lambda exe: False) + a = detect.detect_docker() + assert a.available is False + assert "docker" in a.detail and "podman" in a.detail + + +def test_result_is_cached(monkeypatch): + calls = [] + monkeypatch.setattr(detect.shutil, "which", lambda exe: "/usr/bin/docker" if exe == "docker" else None) + monkeypatch.setattr(detect, "_probe_daemon", lambda exe: calls.append(exe) or True) + detect.detect_docker() + detect.detect_docker() + assert len(calls) == 1 # cached diff --git a/tests/tools/test_sandbox_docker_backend.py b/tests/tools/test_sandbox_docker_backend.py new file mode 100644 index 0000000..9fba081 --- /dev/null +++ b/tests/tools/test_sandbox_docker_backend.py @@ -0,0 +1,318 @@ +"""Tests for JupyterDockerBackend with a fake runtime + detect (no docker).""" + +import json + +import pytest + +from agentic_cli.tools.sandbox.backends.detect import DockerAvailability +from agentic_cli.tools.sandbox.backends.container_runtime import ContainerSpec +from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend +from tests.conftest import MockContext +from tests.tools.test_sandbox_container_session import FakeHandle + + +class FakeRuntime: + def __init__(self): + self.started: list[ContainerSpec] = [] + self.killed: list[tuple[str, str | None]] = [] + self._handles: list[FakeHandle] = [] + def start(self, spec: ContainerSpec): + self.started.append(spec) + h = FakeHandle() + h.name = spec.name + h.stdout.feed(json.dumps({"type": "ready"}) + "\n") + self._handles.append(h) + return h + def kill(self, name, signal=None): + self.killed.append((name, signal)) + for h in self._handles: + if h.name == name: + h._dead = True + h.stdout.eof() + def inspect(self, name): + return {} + def remove(self, name): + pass + + +def _backend(available=True): + ctx = MockContext(stateful_executor_backend="docker").__enter__() + rt = FakeRuntime() + detect_fn = lambda: DockerAvailability(available, "docker" if available else "", "test") + backend = JupyterDockerBackend(ctx.settings, runtime=rt, detect_fn=detect_fn) + return backend, rt, ctx + + +def test_fail_closed_when_docker_unavailable(tmp_path): + backend, rt, ctx = _backend(available=False) + try: + result = backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=tmp_path) + assert result.success is False + assert "unavailable" in result.error.lower() + assert rt.started == [] # never tried to start a container + finally: + ctx.__exit__(None, None, None) + + +def test_execute_starts_container_with_isolation_spec(tmp_path): + backend, rt, ctx = _backend() + try: + # feed a result after start: the session reads ready, then result + orig_start = rt.start + def start(spec): + handle = orig_start(spec) + handle.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "1\n", + "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.1, "error": ""}) + "\n") + return handle + rt.start = start + result = backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=tmp_path) + assert result.success is True + spec = rt.started[0] + assert spec.network == "none" + assert spec.image # from settings + assert any(m.container == "/workspace" and not m.read_only for m in spec.mounts) + assert any(m.container.endswith("/driver.py") and m.read_only for m in spec.mounts) + assert any(m.container.endswith("/kernel_exec.py") and m.read_only for m in spec.mounts) + assert spec.labels.get("agentic-session") == "s1" + assert spec.labels.get("agentic-sandbox") == "1" + assert spec.env.get("AGENTIC_SANDBOX_WORKSPACE") == "/workspace" + finally: + ctx.__exit__(None, None, None) + + +def test_reset_session_kills_container(tmp_path): + backend, rt, ctx = _backend() + try: + # start a session via execute (feed ready + one result) + orig = rt.start + def start_with_result(spec): + handle = orig(spec) + handle.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "", "stderr": "", + "result": None, "artifacts": [], "execution_time": 0.0, "error": ""}) + "\n") + return handle + rt.start = start_with_result + backend.execute("x=1", "s1", timeout_seconds=5, working_dir=tmp_path) + assert backend.has_session("s1") + backend.reset_session("s1") + assert not backend.has_session("s1") + assert any(name == "agentic-sbx-s1" for name, _ in rt.killed) + finally: + ctx.__exit__(None, None, None) + + +def test_bad_startup_message_kills_container(tmp_path): + backend, rt, ctx = _backend() + try: + def start_bad(spec): + handle = FakeHandle() + handle.name = spec.name + rt._handles.append(handle) + rt.started.append(spec) + handle.stdout.feed(json.dumps({"type": "boom"}) + "\n") # not "ready" + return handle + rt.start = start_bad + result = backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=tmp_path) + assert result.success is False + assert any(name == "agentic-sbx-s1" for name, _ in rt.killed) # cleaned up, not orphaned + assert not backend.has_session("s1") + finally: + ctx.__exit__(None, None, None) + + +def test_session_status_reports_backend(tmp_path): + backend, rt, ctx = _backend() + try: + st = backend.session_status("absent") + assert st.backend == "jupyter_docker" + assert st.state == "absent" + finally: + ctx.__exit__(None, None, None) + + +# Fix 2a: OSError from runtime.start must not propagate +def test_start_oserror_returns_error_dict(tmp_path): + backend, rt, ctx = _backend() + try: + def raise_oserror(spec): + raise OSError("docker gone") + rt.start = raise_oserror + result = backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=tmp_path) + assert result.success is False + assert "docker gone" in result.error or "failed" in result.error.lower() + finally: + ctx.__exit__(None, None, None) + + +def _feed_result(rt): + """Make the fake runtime feed one successful result after 'ready'.""" + orig = rt.start + def start_with_result(spec): + h = orig(spec) + h.stdout.feed(json.dumps({"type": "result", "success": True, "stdout": "", "stderr": "", + "result": None, "artifacts": [], "execution_time": 0.0, "error": ""}) + "\n") + return h + rt.start = start_with_result + + +# The non-root container must be able to write the /workspace bind mount. Rather +# than make the host session dir world-writable (chmod 0777 — a multi-user-host +# exposure), run the container AS the host uid so files it writes are host-owned. +@pytest.mark.skipif(not hasattr(__import__("os"), "getuid"), reason="POSIX uid only") +def test_container_runs_as_host_uid_and_dir_not_world_writable(tmp_path): + import os + import stat + backend, rt, ctx = _backend() + try: + _feed_result(rt) + wd = tmp_path / "sess" + wd.mkdir(mode=0o700) + backend.execute("x = 1", "s1", timeout_seconds=5, working_dir=wd) + assert rt.started[0].user == f"{os.getuid()}:{os.getgid()}" + # session dir perms untouched — NOT made world-writable + assert stat.S_IMODE(os.stat(wd).st_mode) == 0o700 + finally: + ctx.__exit__(None, None, None) + + +def test_explicit_container_user_overrides_host_uid(tmp_path): + ctx = MockContext(stateful_executor_backend="docker", + sandbox_container_user="1234:5678").__enter__() + rt = FakeRuntime() + backend = JupyterDockerBackend( + ctx.settings, runtime=rt, + detect_fn=lambda: DockerAvailability(True, "docker", "test"), + ) + try: + _feed_result(rt) + backend.execute("x = 1", "s1", timeout_seconds=5, working_dir=tmp_path) + assert rt.started[0].user == "1234:5678" + finally: + ctx.__exit__(None, None, None) + + +def test_execute_stages_inputs_into_session_inputs_dir(tmp_path): + backend, rt, ctx = _backend() + try: + _feed_result(rt) + src = tmp_path / "sales.csv"; src.write_text("x\n1\n") + wd = tmp_path / "sess"; wd.mkdir() + backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=wd, inputs=[str(src)]) + assert (wd / "inputs" / "sales.csv").read_text() == "x\n1\n" + finally: + ctx.__exit__(None, None, None) + + +def test_build_spec_has_no_nested_outputs_mount(tmp_path): + """outputs/ must NOT be a separate bind mount nested inside /workspace. + That pattern causes macOS Docker Desktop ACL/xattr issues that make the + session dir un-removable at teardown.""" + backend, rt, ctx = _backend() + try: + _feed_result(rt) + backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=tmp_path) + spec = rt.started[0] + assert not any(m.container == "/workspace/outputs" for m in spec.mounts), ( + "outputs/ must not be a nested bind mount inside /workspace" + ) + finally: + ctx.__exit__(None, None, None) + + +def test_outputs_mountpoint_pre_created_as_host_user(tmp_path): + """Docker must not create /workspace/outputs as root. + The backend must pre-create /outputs before runtime.start() so + the mount-point directory is owned by the host user (not root), which allows + pytest teardown to remove it and keeps the session dir clean.""" + backend, rt, ctx = _backend() + try: + _feed_result(rt) + wd = tmp_path / "sess" + wd.mkdir() + backend.execute("x = 1", "s1", timeout_seconds=5, working_dir=wd) + assert (wd / "outputs").exists(), "outputs/ mount-point must exist after execute" + assert (wd / "outputs").is_dir(), "outputs/ must be a directory, not a file" + finally: + ctx.__exit__(None, None, None) + + +def test_execute_copies_outputs_to_shared_dir(tmp_path): + """Files written to /outputs/ by the kernel must be copied to + the shared outputs dir after a successful execute, and the shared path (not + the session path) must appear in ExecutionResult.artifacts.""" + backend, rt, ctx = _backend() + try: + _feed_result(rt) + wd = tmp_path / "sess" + wd.mkdir() + # Simulate what the kernel would write inside the container + (wd / "outputs").mkdir() + (wd / "outputs" / "result.csv").write_text("a,b\n1,2\n") + + result = backend.execute("print(1)", "s1", timeout_seconds=5, working_dir=wd) + + shared_dir = backend._outputs_dir() + shared_file = shared_dir / "result.csv" + assert shared_file.exists(), "file must have been copied to shared outputs dir" + assert shared_file.read_text() == "a,b\n1,2\n" + assert str(shared_file) in result.artifacts, ( + "shared path must appear in ExecutionResult.artifacts" + ) + finally: + ctx.__exit__(None, None, None) + + +def test_data_mount_is_not_nested_under_workspace(tmp_path): + """Data is mounted at a top-level /data/ (a sibling of /workspace), not + nested under the /workspace bind mount. Nesting a second bind mount under the + session-dir mount makes Docker synthesize the mount-point dir root-owned + (Linux) / ACL-stamped (macOS), which breaks host-side cleanup; a top-level + mount point lives in the container's own layer and leaves the session dir + untouched, so nothing is pre-created on the host.""" + some_dir = tmp_path / "mydata" + some_dir.mkdir() + ctx = MockContext(stateful_executor_backend="docker", + sandbox_data_mounts=[f"{some_dir}:samples"]).__enter__() + rt = FakeRuntime() + backend = JupyterDockerBackend( + ctx.settings, runtime=rt, + detect_fn=lambda: DockerAvailability(True, "docker", "test"), + ) + try: + _feed_result(rt) + wd = tmp_path / "sess" + wd.mkdir() + backend.execute("x = 1", "s1", timeout_seconds=5, working_dir=wd) + data = [m for m in rt.started[0].mounts if m.container == "/data/samples"] + assert data, "data must be mounted at /data/samples (sibling of /workspace)" + assert not any(m.container.startswith("/workspace/data") for m in rt.started[0].mounts), ( + "data must NOT be nested under the /workspace bind mount" + ) + assert not (wd / "data").exists(), ( + "no host-side data mount point should be created under the session dir" + ) + finally: + ctx.__exit__(None, None, None) + + +def test_data_mount_name_cannot_escape_data_dir(tmp_path): + """A hostile data-mount name (traversal) must not remap the mount point + outside /data/ inside the container.""" + import posixpath + ctx = MockContext(stateful_executor_backend="docker", + sandbox_data_mounts=[f"{tmp_path}:../../etc"]).__enter__() + rt = FakeRuntime() + backend = JupyterDockerBackend( + ctx.settings, runtime=rt, + detect_fn=lambda: DockerAvailability(True, "docker", "test"), + ) + try: + _feed_result(rt) + backend.execute("x = 1", "s1", timeout_seconds=5, working_dir=tmp_path) + data = [m for m in rt.started[0].mounts if m.container.startswith("/data/")] + assert data, "expected a data mount under /data/" + for m in data: + assert ".." not in m.container + assert posixpath.normpath(m.container).startswith("/data/") + finally: + ctx.__exit__(None, None, None) diff --git a/tests/tools/test_sandbox_docker_integration.py b/tests/tools/test_sandbox_docker_integration.py new file mode 100644 index 0000000..89110d8 --- /dev/null +++ b/tests/tools/test_sandbox_docker_integration.py @@ -0,0 +1,151 @@ +"""End-to-end: real driver subprocess (no docker) via a local-process runtime.""" + +from __future__ import annotations + +import os +import signal as _signal +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("jupyter_client") + +from agentic_cli.tools.sandbox.backends import driver as _driver_mod +from agentic_cli.tools.sandbox.backends.detect import DockerAvailability +from agentic_cli.tools.sandbox.backends.container_runtime import ContainerHandle +from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend +from tests.conftest import MockContext + +DRIVER = Path(_driver_mod.__file__) + + +class LocalDriverRuntime: + """Runs `python driver.py` locally, standing in for `docker run`.""" + + def __init__(self) -> None: + self._procs: dict[str, subprocess.Popen] = {} + + def start(self, spec) -> ContainerHandle: + # Find the host workspace mount (container path == "/workspace") + ws = next(m.host for m in spec.mounts if m.container == "/workspace") + # Pass the full environment so the local kernel can find Jupyter/IPython paths, + # but override the workspace var to point at our tmp dir. + env = {**os.environ, "AGENTIC_SANDBOX_WORKSPACE": ws} + proc = subprocess.Popen( + [sys.executable, str(DRIVER)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + cwd=ws, + env=env, + ) + self._procs[spec.name] = proc + return ContainerHandle(spec.name, proc) + + def kill(self, name: str, signal: str | None = None) -> None: + proc = self._procs.get(name) + if proc and proc.poll() is None: + if signal == "INT": # cooperative interrupt, like `docker kill --signal=INT` + proc.send_signal(_signal.SIGINT) + else: + proc.kill() + + def inspect(self, name: str) -> dict: + return {} + + def remove(self, name: str) -> None: + pass + + +@pytest.fixture +def backend(tmp_path): + with MockContext(stateful_executor_backend="docker", sandbox_start_timeout=60) as ctx: + b = JupyterDockerBackend( + ctx.settings, + runtime=LocalDriverRuntime(), + detect_fn=lambda: DockerAvailability(True, "docker", "local"), + ) + try: + yield b + finally: + b.cleanup() + + +def test_stateful_execution_across_calls(backend, tmp_path): + r1 = backend.execute("data = [1, 2, 3]", "s1", timeout_seconds=60, working_dir=tmp_path) + assert r1.success is True, r1.error + r2 = backend.execute("print(sum(data))", "s1", timeout_seconds=60, working_dir=tmp_path) + assert r2.success is True, r2.error + assert "6" in r2.stdout + backend.cleanup() + + +def test_error_surfaces(backend, tmp_path): + r = backend.execute("1/0", "s1", timeout_seconds=60, working_dir=tmp_path) + assert r.success is False + assert "ZeroDivisionError" in r.error + backend.cleanup() + + +def test_user_code_cannot_forge_protocol_via_fd1(backend, tmp_path): + """Raw writes to fd 1 by user code must NOT reach the host NDJSON protocol + channel: the kernel's stdout is isolated from the driver's. Otherwise code + could forge a {"type":"result"} and desync the session.""" + import json + forged = json.dumps({"type": "result", "success": True, "stdout": "FORGED\n", + "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.0, "error": ""}) + code = ("import os\n" + "os.write(1, (" + repr(forged) + " + '\\n').encode())\n" + "print('legit')\n") + r1 = backend.execute(code, "s1", timeout_seconds=30, working_dir=tmp_path) + r2 = backend.execute("print('second')", "s1", timeout_seconds=30, working_dir=tmp_path) + # The security invariant: cell1 gets its OWN real result (not the forged + # dict), and cell2 is NOT desynced. (Whether the forged bytes are discarded + # or surface as inert TEXT in cell1's stdout is kernel-dependent and + # harmless — what matters is they never become a trusted protocol message.) + assert r1.success is True, r1.error + assert "legit" in r1.stdout + assert r2.success is True, r2.error + assert r2.stdout == "second\n" # clean — no forged bytes / leftover bled in + backend.cleanup() + + +def test_inputs_are_loadable_by_relative_path(backend, tmp_path): + """A staged input is available at inputs/ and loadable relatively.""" + src = tmp_path.parent / "iris_like.csv"; src.write_text("a,b\n1,2\n3,4\n") + r = backend.execute("print(open('inputs/iris_like.csv').read().strip())", + "s1", timeout_seconds=30, working_dir=tmp_path, inputs=[str(src)]) + assert r.success is True, r.error + assert "1,2" in r.stdout + backend.cleanup() + + +def test_interrupt_preserves_session_state(backend, tmp_path): + """A runaway cell is aborted by the host's cooperative interrupt, but the + session (kernel + prior state) survives and the next request runs cleanly. + + Regression: the driver used to self-time-out on the host's per-cell deadline + and return a 'timed out' result while the cell kept running in the kernel. + The host then consumed that result WITHOUT interrupting, leaving the kernel + busy — so the next request wedged (empty/desynced output). The host must be + the sole timeout authority. + """ + r0 = backend.execute("kept = 123", "s1", timeout_seconds=60, working_dir=tmp_path) + assert r0.success is True, r0.error + + # Runs far longer than the per-cell timeout -> host sends a cooperative + # interrupt (kill --signal=INT) rather than killing the container. + r1 = backend.execute("import time\nfor _ in range(120):\n time.sleep(1)", + "s1", timeout_seconds=3, working_dir=tmp_path) + assert r1.success is False # aborted + + # Same session still alive with prior state intact. + r2 = backend.execute("print(kept)", "s1", timeout_seconds=60, working_dir=tmp_path) + assert r2.success is True, r2.error + assert "123" in r2.stdout + backend.cleanup() diff --git a/tests/tools/test_sandbox_docker_live.py b/tests/tools/test_sandbox_docker_live.py new file mode 100644 index 0000000..e6ae459 --- /dev/null +++ b/tests/tools/test_sandbox_docker_live.py @@ -0,0 +1,337 @@ +"""Live docker isolation tests — assert the real container boundary. + +These require a real container runtime (docker/podman). Selection: + * ``-m docker`` run only these + * ``-m 'not docker'`` exclude them (used by the offline CI job) + +Gating: normally each test is skipped when no runtime is present. In CI set +``SANDBOX_REQUIRE_DOCKER=1`` so a missing/broken runtime is a hard FAILURE +rather than a silent skip (a skipped test reads as green and would hide a +regression). Mirrors the ``-m llm`` live-test convention. +""" + +import os +import subprocess + +import pytest + +from agentic_cli.tools.sandbox.backends.detect import detect_docker, docker_available +from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend +from agentic_cli.tools.sandbox.manager import SandboxManager +from tests.conftest import MockContext + +# Every test in this module is a docker test (selectable via -m docker). +pytestmark = pytest.mark.docker + +_DOCKER = docker_available() +_requires_docker = pytest.mark.skipif(not _DOCKER, reason="no container runtime available") + + +def test_docker_runtime_present_when_required(): + """Fail loudly (not skip) when CI declares a runtime is required but none is + available — so a broken daemon or failed image pull can't pass as green.""" + if os.environ.get("SANDBOX_REQUIRE_DOCKER") == "1": + assert _DOCKER, "SANDBOX_REQUIRE_DOCKER=1 but no container runtime is available" + elif not _DOCKER: + pytest.skip("no container runtime available") + + +@pytest.fixture +def backend(tmp_path): + with MockContext(stateful_executor_backend="docker") as ctx: + b = JupyterDockerBackend(ctx.settings) + try: + yield b + finally: + b.cleanup() + + +# -------------------------------------------------------------------------- +# Boundary: network, filesystem +# -------------------------------------------------------------------------- + +@_requires_docker +def test_network_is_blocked(backend, tmp_path): + code = ("import socket\n" + "try:\n" + " socket.create_connection(('1.1.1.1', 53), timeout=3); print('OPEN')\n" + "except OSError:\n" + " print('BLOCKED')\n") + r = backend.execute(code, "net", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert "BLOCKED" in r.stdout + + +@_requires_docker +def test_rootfs_is_read_only(backend, tmp_path): + code = ("try:\n" + " open('/etc/passwd', 'a').write('x'); print('WRITABLE')\n" + "except OSError:\n" + " print('READONLY')\n") + r = backend.execute(code, "ro", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert "READONLY" in r.stdout + + +@_requires_docker +def test_workspace_is_writable(backend, tmp_path): + r = backend.execute("open('/workspace/out.txt','w').write('ok'); print('WROTE')", + "ws", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert "WROTE" in r.stdout + assert (tmp_path / "out.txt").read_text() == "ok" + + +@_requires_docker +def test_fd1_write_cannot_forge_protocol(backend, tmp_path): + """Raw writes to fd 1 by user code must not become a trusted host protocol + message: cell1 keeps its own result and cell2 is not desynced.""" + import json + forged = json.dumps({"type": "result", "success": True, "stdout": "FORGED\n", + "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.0, "error": ""}) + code = ("import os\n" + "os.write(1, (" + repr(forged) + " + '\\n').encode())\n" + "print('legit')\n") + r1 = backend.execute(code, "fd1", timeout_seconds=60, working_dir=tmp_path) + r2 = backend.execute("print('second')", "fd1", timeout_seconds=60, working_dir=tmp_path) + assert r1.success is True, r1.error + assert "legit" in r1.stdout + assert r2.success is True, r2.error + assert r2.stdout == "second\n" # not desynced by cell1's raw fd-1 writes + + +@_requires_docker +def test_proc_fd_write_cannot_forge_protocol(backend, tmp_path): + """User code writing forged NDJSON to the driver's fd 1 via /proc//fd/1 + (reachable since the kernel shares the driver's pid/user namespace) must be + rejected as unauthenticated — the session is not forged or desynced.""" + import json + forged = json.dumps({"type": "result", "success": True, "stdout": "PROC-FORGED\n", + "stderr": "", "result": None, "artifacts": [], + "execution_time": 0.0, "error": ""}) + code = ("import os\n" + "for t in (f'/proc/{os.getppid()}/fd/1', '/proc/1/fd/1'):\n" + " try:\n" + " fd = os.open(t, os.O_WRONLY)\n" + " os.write(fd, (" + repr(forged) + " + '\\n').encode())\n" + " os.close(fd)\n" + " except OSError:\n" + " pass\n" + "print('legit')\n") + r1 = backend.execute(code, "proc", timeout_seconds=60, working_dir=tmp_path) + r2 = backend.execute("print('second')", "proc", timeout_seconds=60, working_dir=tmp_path) + assert r1.success is True, r1.error + assert "legit" in r1.stdout and "PROC-FORGED" not in r1.stdout + assert r2.success is True, r2.error + assert r2.stdout == "second\n" # not desynced by the forged /proc writes + + +@_requires_docker +def test_relative_writes_land_in_workspace(backend, tmp_path): + """The kernel starts in /workspace, so a RELATIVE file write persists to the + host mount instead of failing on the read-only image WORKDIR.""" + r = backend.execute("open('out_rel.txt', 'w').write('persisted'); print('ok')", + "cwd", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert (tmp_path / "out_rel.txt").read_text() == "persisted" + + +@_requires_docker +def test_host_path_outside_workspace_not_accessible(backend, tmp_path): + """A host file that is NOT bind-mounted must be invisible: its host path + does not exist inside the container's mount namespace.""" + secret = tmp_path.parent / "host_only_secret.txt" + secret.write_text("top-secret") + code = ("try:\n" + f" print('LEAKED:' + open({str(secret)!r}).read())\n" + "except OSError:\n" + " print('NO_HOST_ACCESS')\n") + r = backend.execute(code, "hostpath", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert "NO_HOST_ACCESS" in r.stdout + + +# -------------------------------------------------------------------------- +# Cross-session isolation +# -------------------------------------------------------------------------- + +@_requires_docker +def test_sessions_are_isolated(backend, tmp_path): + """Distinct session_ids get distinct containers: neither in-memory state + nor workspace files leak between them.""" + dir_a = tmp_path / "a" + dir_b = tmp_path / "b" + dir_a.mkdir() + dir_b.mkdir() + + ra = backend.execute("secret = 'alpha'\nopen('/workspace/a.txt', 'w').write('alpha')", + "sessA", timeout_seconds=60, working_dir=dir_a) + assert ra.success is True, ra.error + + rb = backend.execute( + "import os\n" + "print('VAR_LEAK' if 'secret' in dir() else 'NO_VAR')\n" + "print('FILE_LEAK' if os.path.exists('/workspace/a.txt') else 'NO_FILE')\n", + "sessB", timeout_seconds=60, working_dir=dir_b) + assert rb.success is True, rb.error + assert "NO_VAR" in rb.stdout + assert "NO_FILE" in rb.stdout + + +# -------------------------------------------------------------------------- +# Stateful analysis: multi-call persistence +# -------------------------------------------------------------------------- + +@_requires_docker +def test_analyst_flow_stage_analyze_output(tmp_path): + """Stage an input, run multi-step stateful analysis, write a figure to + outputs/, and confirm it appears on the shared host dir.""" + from agentic_cli.tools.sandbox.manager import SandboxManager + outdir = tmp_path / "artifacts" + src = tmp_path / "rows.csv"; src.write_text("g,v\na,1\na,3\nb,10\n") + with MockContext(stateful_executor_backend="docker", + sandbox_outputs_dir=str(outdir)) as ctx: + mgr = SandboxManager(ctx.settings) + try: + r1 = mgr.execute("import pandas as pd\ndf = pd.read_csv('inputs/rows.csv')\nprint(df.groupby('g').v.mean().to_dict())", + session_id="an", inputs=[str(src)]) + assert r1.success is True, r1.error + assert "'a': 2.0" in r1.stdout or '"a": 2.0' in r1.stdout + r2 = mgr.execute("df.groupby('g').v.mean().to_csv('outputs/means.csv'); print('wrote')", + session_id="an") # same session: df persists + assert r2.success is True, r2.error + assert (outdir / "means.csv").exists() + finally: + mgr.cleanup() + + +@_requires_docker +def test_data_mount_readable_and_session_cleans_up(tmp_path): + """A read-only data mount is readable at /data/ AND leaves the session + dir removable afterward. + + The old design nested the mount at /workspace/data/; Docker then + synthesized that mount-point dir root-owned (Linux) / ACL-stamped (macOS), + so a host-side rmtree of the session dir failed. Mounting at a top-level + /data/ keeps the mount point in the container's own ephemeral layer, + so the host session dir stays clean. This asserts both halves on a real + daemon — the exact scenario the fake-runtime tests cannot reach.""" + import shutil + datadir = tmp_path / "host_data" + datadir.mkdir() + (datadir / "sample.csv").write_text("a,b\n1,2\n") + session_dir = tmp_path / "sess" + session_dir.mkdir() + with MockContext(stateful_executor_backend="docker", + sandbox_data_mounts=[f"{datadir}:samples"]) as ctx: + b = JupyterDockerBackend(ctx.settings) + try: + r = b.execute("print(open('/data/samples/sample.csv').read().strip())", + "dm", timeout_seconds=60, working_dir=session_dir) + assert r.success is True, r.error + assert "a,b" in r.stdout and "1,2" in r.stdout + assert not (session_dir / "data").exists(), ( + "no data mount point should be created under the session dir" + ) + finally: + b.cleanup() + # The exact failure the nested mount caused on macOS: the host could not + # remove the session dir because the synthesized mount point was + # un-rmdir-able. With a top-level mount there is nothing to trip on. + shutil.rmtree(session_dir) + assert not session_dir.exists() + + +# -------------------------------------------------------------------------- +# Cooperative interrupt: a runaway cell is aborted but the session survives +# -------------------------------------------------------------------------- + +@_requires_docker +def test_interrupt_preserves_session_state(backend, tmp_path): + r0 = backend.execute("kept = 123", "intr", timeout_seconds=60, working_dir=tmp_path) + assert r0.success is True, r0.error + + # Runs far longer than the per-cell timeout -> host sends a cooperative + # interrupt (docker kill --signal=INT) instead of killing the container. + r1 = backend.execute("import time\nfor _ in range(120):\n time.sleep(1)", + "intr", timeout_seconds=3, working_dir=tmp_path) + assert r1.success is False # aborted + + # Same session still alive with prior state intact. + r2 = backend.execute("print(kept)", "intr", timeout_seconds=60, working_dir=tmp_path) + assert r2.success is True, r2.error + assert "123" in r2.stdout + + +# -------------------------------------------------------------------------- +# Resource caps (may need first-run threshold tuning per runner cgroup config) +# -------------------------------------------------------------------------- + +@_requires_docker +def test_memory_cap_oom_kills(tmp_path): + """A single allocation far past --memory (swap disabled) is OOM-killed; + the backend surfaces failure rather than a clean success.""" + with MockContext(stateful_executor_backend="docker", sandbox_memory_mb=256) as ctx: + b = JupyterDockerBackend(ctx.settings) + try: + r = b.execute("x = bytearray(1024 * 1024 * 1024) # 1 GiB vs 256 MiB cap", + "oom", timeout_seconds=45, working_dir=tmp_path) + assert r.success is False + finally: + b.cleanup() + + +@_requires_docker +def test_pids_limit_caps_thread_bomb(tmp_path): + """--pids-limit bounds the number of tasks; a thread bomb hits it.""" + with MockContext(stateful_executor_backend="docker", sandbox_pids_limit=128) as ctx: + b = JupyterDockerBackend(ctx.settings) + try: + code = ("import threading, time\n" + "started = 0\n" + "try:\n" + " for _ in range(500):\n" + " threading.Thread(target=lambda: time.sleep(30)).start()\n" + " started += 1\n" + " print('NO_LIMIT', started)\n" + "except RuntimeError:\n" + " print('PIDS_CAPPED', started)\n") + r = b.execute(code, "pids", timeout_seconds=45, working_dir=tmp_path) + assert "PIDS_CAPPED" in r.stdout + finally: + b.cleanup() + + +# -------------------------------------------------------------------------- +# Lifecycle: no leaked containers after cleanup +# -------------------------------------------------------------------------- + +@_requires_docker +def test_outputs_dir_persists_to_shared_host_dir(tmp_path): + outdir = tmp_path / "shared_out" + with MockContext(stateful_executor_backend="docker", sandbox_outputs_dir=str(outdir)) as ctx: + b = JupyterDockerBackend(ctx.settings) + try: + r = b.execute("open('outputs/final.txt','w').write('done'); print('ok')", + "out", timeout_seconds=60, working_dir=tmp_path) + assert r.success is True, r.error + assert (outdir / "final.txt").read_text() == "done" + assert any(a.endswith("final.txt") for a in r.artifacts) + finally: + b.cleanup() + + +@_requires_docker +def test_no_orphaned_containers_after_cleanup(tmp_path): + runtime = detect_docker().runtime or "docker" + with MockContext(stateful_executor_backend="docker") as ctx: + b = JupyterDockerBackend(ctx.settings) + b.execute("x = 1", "orphan1", timeout_seconds=60, working_dir=tmp_path) + b.execute("y = 2", "orphan2", timeout_seconds=60, working_dir=tmp_path) + b.cleanup() + out = subprocess.run( + [runtime, "ps", "-aq", "--filter", "label=agentic-sandbox=1"], + capture_output=True, text=True, check=False, + ) + assert out.stdout.strip() == "", f"orphaned containers remain: {out.stdout!r}" diff --git a/tests/tools/test_sandbox_driver.py b/tests/tools/test_sandbox_driver.py new file mode 100644 index 0000000..eff3bd4 --- /dev/null +++ b/tests/tools/test_sandbox_driver.py @@ -0,0 +1,70 @@ +"""Tests for the in-container kernel driver (real kernel, no docker).""" + +import io +import json +import signal +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("jupyter_client") + +from agentic_cli.tools.sandbox.backends.driver import KernelDriver + + +@pytest.fixture +def driver(tmp_path): + d = KernelDriver(stdin=io.StringIO(), stdout=io.StringIO(), workspace=str(tmp_path)) + d.start() + yield d + d.close() + + +def test_state_persists_across_requests(driver): + driver.handle_request({"type": "execute", "code": "x = 41", "timeout": 30}) + r = driver.handle_request({"type": "execute", "code": "print(x + 1)", "timeout": 30}) + assert r["success"] is True + assert "42" in r["stdout"] + + +def test_error_is_captured(driver): + r = driver.handle_request({"type": "execute", "code": "1/0", "timeout": 30}) + assert r["success"] is False + assert "ZeroDivisionError" in r["error"] + + +def test_run_emits_ready_then_result(tmp_path): + stdin = io.StringIO(json.dumps({"type": "execute", "code": "print('hi')", "timeout": 30}) + "\n") + stdout = io.StringIO() + d = KernelDriver(stdin=stdin, stdout=stdout, workspace=str(tmp_path)) + try: + d.run() # returns at stdin EOF + lines = [json.loads(l) for l in stdout.getvalue().splitlines() if l.strip()] + assert lines[0]["type"] == "ready" + assert lines[1]["type"] == "result" + assert "hi" in lines[1]["stdout"] + finally: + d.close() + + +def test_messages_carry_consistent_session_token(tmp_path): + """Every driver message is tagged with a per-session token so the host can + reject forged lines; ready and result share the same non-empty token.""" + stdin = io.StringIO(json.dumps({"type": "execute", "code": "print('hi')", "timeout": 30}) + "\n") + stdout = io.StringIO() + d = KernelDriver(stdin=stdin, stdout=stdout, workspace=str(tmp_path)) + try: + d.run() + lines = [json.loads(l) for l in stdout.getvalue().splitlines() if l.strip()] + assert lines[0]["type"] == "ready" and lines[0].get("token") + assert lines[1]["type"] == "result" + assert lines[1].get("token") == lines[0]["token"] + finally: + d.close() + + +def test_sigint_handler_interrupts_kernel(tmp_path): + d = KernelDriver(stdin=io.StringIO(), stdout=io.StringIO(), workspace=str(tmp_path)) + d._km = MagicMock() + d._on_sigint(signal.SIGINT, None) + d._km.interrupt_kernel.assert_called_once() diff --git a/tests/tools/test_sandbox_kernel_exec.py b/tests/tools/test_sandbox_kernel_exec.py new file mode 100644 index 0000000..1f9ea0a --- /dev/null +++ b/tests/tools/test_sandbox_kernel_exec.py @@ -0,0 +1,57 @@ +"""Tests for the shared kernel execution helper (no docker needed).""" + +import pytest + +pytest.importorskip("jupyter_client") + +from jupyter_client import KernelManager + +from agentic_cli.tools.sandbox.backends import kernel_exec + + +@pytest.fixture +def kernel(): + km = KernelManager() + km.start_kernel() + kc = km.blocking_client() + kc.start_channels() + kc.wait_for_ready(timeout=30) + yield kc + kc.stop_channels() + km.shutdown_kernel(now=True) + + +def test_validate_rejects_shell_bang(): + ok, msg = kernel_exec.validate_code("!rm -rf /") + assert ok is False + assert "shell" in msg.lower() + + +def test_validate_allows_plain_code(): + assert kernel_exec.validate_code("x = 1\nprint(x)") == (True, "") + + +def test_collect_captures_stdout(kernel, tmp_path): + msg_id = kernel.execute("print('hello')") + out = kernel_exec.collect_execution(kernel, msg_id, timeout=30, working_dir=tmp_path) + assert out["success"] is True + assert "hello" in out["stdout"] + assert out["error"] == "" + + +def test_collect_captures_error(kernel, tmp_path): + msg_id = kernel.execute("raise ValueError('boom')") + out = kernel_exec.collect_execution(kernel, msg_id, timeout=30, working_dir=tmp_path) + assert out["success"] is False + assert "boom" in out["error"] + + +def test_collect_caps_large_stdout(kernel, tmp_path): + """Runaway output is bounded (host memory / LLM context) with a truncation + marker, rather than accumulated unbounded.""" + cap = kernel_exec.MAX_STREAM_CHARS + msg_id = kernel.execute(f"print('A' * {cap * 3})") + out = kernel_exec.collect_execution(kernel, msg_id, timeout=30, working_dir=tmp_path) + assert out["success"] is True + assert len(out["stdout"]) <= cap + 200 # cap plus the short marker + assert "truncated" in out["stdout"].lower() diff --git a/tests/tools/test_sandbox_settings.py b/tests/tools/test_sandbox_settings.py new file mode 100644 index 0000000..c055e34 --- /dev/null +++ b/tests/tools/test_sandbox_settings.py @@ -0,0 +1,38 @@ +"""Tests for docker sandbox settings defaults.""" + +import pytest +from pydantic import ValidationError + +from agentic_cli.config import BaseSettings + + +def test_docker_sandbox_defaults(): + s = BaseSettings() + assert s.sandbox_image.startswith("quay.io/jupyter/") or "jupyter" in s.sandbox_image + assert s.sandbox_memory_mb == 2048 + assert s.sandbox_cpus == 2.0 + assert s.sandbox_pids_limit == 256 + assert s.sandbox_network == "none" + assert s.sandbox_container_user == "" + assert s.sandbox_data_mounts == [] + assert s.sandbox_start_timeout == 180 + # unified backend config (replaces the old two-field enable+backend pattern) + assert s.stateful_executor_backend == "none" + + +def test_stateful_executor_backend_default_and_values(): + from pydantic import ValidationError + assert BaseSettings().stateful_executor_backend == "none" + assert BaseSettings(stateful_executor_backend="docker").stateful_executor_backend == "docker" + assert BaseSettings(stateful_executor_backend="local").stateful_executor_backend == "local" + with pytest.raises(ValidationError): + BaseSettings(stateful_executor_backend="bogus") + + +def test_sandbox_network_must_be_none(): + """The docker backend's no-egress guarantee depends on --network none, so a + non-'none' value is rejected rather than silently weakening isolation.""" + assert BaseSettings(sandbox_network="none").sandbox_network == "none" + for bad in ("host", "bridge", "my-net"): + with pytest.raises(ValidationError): + BaseSettings(sandbox_network=bad) diff --git a/tests/tools/test_sandbox_transfer.py b/tests/tools/test_sandbox_transfer.py new file mode 100644 index 0000000..b3975da --- /dev/null +++ b/tests/tools/test_sandbox_transfer.py @@ -0,0 +1,107 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from agentic_cli.tools.sandbox.manager import stage_inputs +from agentic_cli.tools.sandbox.backends.jupyter_docker import JupyterDockerBackend + + +def _backend(shared: Path): + return JupyterDockerBackend( + settings=SimpleNamespace(sandbox_outputs_dir=str(shared), workspace_dir=str(shared)) + ) + + +def test_stage_inputs_copies_regular_file(tmp_path): + session = tmp_path / "session"; session.mkdir() + src = tmp_path / "data.csv"; src.write_bytes(b"col1,col2") + stage_inputs(session, [str(src)]) + staged = session / "inputs" / "data.csv" + assert staged.read_bytes() == b"col1,col2" and not staged.is_symlink() + + +def test_stage_inputs_rejects_symlink_source(tmp_path): + session = tmp_path / "session"; session.mkdir() + secret = tmp_path / "secret.txt"; secret.write_bytes(b"TOP SECRET") + link = tmp_path / "data.csv"; link.symlink_to(secret) + with pytest.raises(ValueError): + stage_inputs(session, [str(link)]) + + +def test_stage_inputs_does_not_follow_dest_symlink(tmp_path): + # Simulate a kernel that pre-planted inputs/data.csv -> a host file. + session = tmp_path / "session"; session.mkdir() + inputs_dir = session / "inputs"; inputs_dir.mkdir() + outside = tmp_path / "host_secret"; outside.write_bytes(b"ORIGINAL") + (inputs_dir / "data.csv").symlink_to(outside) + src = tmp_path / "data.csv"; src.write_bytes(b"INPUT") + stage_inputs(session, [str(src)]) + assert outside.read_bytes() == b"ORIGINAL" # not written through + assert (inputs_dir / "data.csv").read_bytes() == b"INPUT" # replaced by a real file + assert not (inputs_dir / "data.csv").is_symlink() + + +def test_stage_inputs_duplicate_basename_errors(tmp_path): + session = tmp_path / "session"; session.mkdir() + a = tmp_path / "a" / "x.csv"; a.parent.mkdir(); a.write_text("1") + b = tmp_path / "b" / "x.csv"; b.parent.mkdir(); b.write_text("2") + with pytest.raises(ValueError): + stage_inputs(session, [str(a), str(b)]) + + +def test_collect_outputs_copies_regular_files(tmp_path): + shared = tmp_path / "shared" + wd = tmp_path / "wd"; (wd / "outputs").mkdir(parents=True) + (wd / "outputs" / "plot.png").write_bytes(b"PNG") + got = _backend(shared)._collect_outputs(wd) + assert (shared / "plot.png").read_bytes() == b"PNG" + assert any("plot.png" in g for g in got) + + +def test_collect_outputs_skips_symlink_to_host_secret(tmp_path): + shared = tmp_path / "shared" + wd = tmp_path / "wd"; (wd / "outputs").mkdir(parents=True) + secret = tmp_path / "aws_credentials"; secret.write_bytes(b"AKIA-SECRET") + (wd / "outputs" / "result.txt").symlink_to(secret) # kernel exfil attempt + got = _backend(shared)._collect_outputs(wd) + assert got == [] + assert not (shared / "result.txt").exists() # secret not copied out + + +def test_collect_outputs_skips_when_outputs_is_symlink(tmp_path): + shared = tmp_path / "shared" + wd = tmp_path / "wd"; wd.mkdir() + elsewhere = tmp_path / "elsewhere"; elsewhere.mkdir() + (elsewhere / "x.txt").write_text("data") + (wd / "outputs").symlink_to(elsewhere) # 'outputs' itself a symlink + assert _backend(shared)._collect_outputs(wd) == [] + + +def test_outputs_dir_is_0700(tmp_path): + shared = tmp_path / "shared" + base = _backend(shared)._outputs_dir() + assert (base.stat().st_mode & 0o777) == 0o700 + + +def test_stage_inputs_rejects_symlinked_inputs_dir(tmp_path): + # Kernel plants session/inputs as a symlink to a host dir -> must be rejected, + # not written through (O_NOFOLLOW guards only the final component). + session = tmp_path / "session"; session.mkdir() + outside = tmp_path / "host_dir"; outside.mkdir() + (session / "inputs").symlink_to(outside, target_is_directory=True) + src = tmp_path / "data.csv"; src.write_bytes(b"INPUT") + with pytest.raises(ValueError): + stage_inputs(session, [str(src)]) + assert list(outside.iterdir()) == [] # nothing written into the host dir + + +def test_collect_outputs_never_raises_on_outputs_dir_error(tmp_path, monkeypatch): + shared = tmp_path / "shared" + wd = tmp_path / "wd"; (wd / "outputs").mkdir(parents=True) + (wd / "outputs" / "a.txt").write_text("x") + b = _backend(shared) + def boom(): + raise OSError("boom") + monkeypatch.setattr(b, "_outputs_dir", boom) + assert b._collect_outputs(wd) == [] # best-effort: never raises 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/tools/test_webfetch_ssrf.py b/tests/tools/test_webfetch_ssrf.py new file mode 100644 index 0000000..99d1dcd --- /dev/null +++ b/tests/tools/test_webfetch_ssrf.py @@ -0,0 +1,183 @@ +"""P0-5 SSRF tests: resolve-validate-pin. Fully offline — socket.getaddrinfo is +monkeypatched and httpx.MockTransport is the PinnedTransport inner.""" + +import socket + +import httpx +import pytest + +from agentic_cli.tools.webfetch.transport import PinnedTransport +from agentic_cli.tools.webfetch.validator import ( + URLValidator, + BlockedAddressError, + _ip_is_safe, + _as_ip_literal, +) + + +def stub_getaddrinfo(*ips): + """Return a socket.getaddrinfo replacement yielding the given IP strings.""" + def _stub(host, port, *args, **kwargs): + out = [] + for ip in ips: + is6 = ":" in ip + fam = socket.AF_INET6 if is6 else socket.AF_INET + sa = (ip, port, 0, 0) if is6 else (ip, port) + out.append((fam, socket.SOCK_STREAM, 6, "", sa)) + return out + return _stub + + +import ipaddress + + +class TestIpIsSafe: + def test_public_v4_and_v6_safe(self): + assert _ip_is_safe(ipaddress.ip_address("8.8.8.8")) + assert _ip_is_safe(ipaddress.ip_address("2606:4700:4700::1111")) + + @pytest.mark.parametrize("ip", [ + "127.0.0.1", "10.0.0.1", "192.168.1.1", "172.16.0.1", "169.254.169.254", + "100.64.0.1", "192.0.2.1", "198.18.0.1", "240.0.0.1", "0.0.0.1", + "::1", "fc00::1", "fe80::1", + "::ffff:10.0.0.1", # v4-mapped private + "64:ff9b::a00:1", # NAT64 embedding 10.0.0.1 (is_global=True) + "192.88.99.1", # 6to4 relay anycast (deprecated) + ]) + def test_unsafe_addresses(self, ip): + assert not _ip_is_safe(ipaddress.ip_address(ip)) + + +class TestResolveAndValidate: + def test_public_host_returns_ip(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("93.184.216.34")) + assert URLValidator().resolve_and_validate("example.com", 443) == "93.184.216.34" + + def test_private_host_rejected(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("127.0.0.1")) + with pytest.raises(BlockedAddressError): + URLValidator().resolve_and_validate("localhost", 80) + + def test_mixed_public_and_private_rejected(self, monkeypatch): + # split-horizon: a public + a private answer → reject the whole host + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("93.184.216.34", "10.0.0.1")) + with pytest.raises(BlockedAddressError): + URLValidator().resolve_and_validate("rebind.test", 443) + + def test_ipv6_ula_rejected(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("fc00::1")) + with pytest.raises(BlockedAddressError): + URLValidator().resolve_and_validate("v6.test", 443) + + def test_unresolvable_host_raises_blocked(self, monkeypatch): + def boom(*a, **k): + raise socket.gaierror("nope") + monkeypatch.setattr(socket, "getaddrinfo", boom) + with pytest.raises(BlockedAddressError): + URLValidator().resolve_and_validate("nx.test", 443) + + def test_ip_literal_public_ok_private_blocked(self): + assert URLValidator().resolve_and_validate("8.8.8.8", 443) == "8.8.8.8" + with pytest.raises(BlockedAddressError): + URLValidator().resolve_and_validate("169.254.169.254", 80) + + +class TestValidateNoDNS: + def test_ip_literal_private_blocked_without_dns(self): + # validate() must reject a private IP-literal with no DNS call + r = URLValidator().validate("http://127.0.0.1/x") + assert r.valid is False + + def test_public_hostname_passes_policy(self, monkeypatch): + # validate() must NOT resolve DNS — make resolution explode and assert + # a normal hostname still passes policy. + def _boom(*a, **k): + raise AssertionError("validate() must not call DNS") + monkeypatch.setattr(socket, "getaddrinfo", _boom) + monkeypatch.setattr(socket, "gethostbyname", _boom) + assert URLValidator().validate("https://example.com/x").valid is True + + +def mock_pinned_transport(handler, validator=None): + """PinnedTransport whose inner is an httpx.MockTransport(handler).""" + from agentic_cli.tools.webfetch.transport import PinnedTransport + return PinnedTransport(validator or URLValidator(), inner=httpx.MockTransport(handler)) + + +class TestPinnedTransport: + @pytest.mark.asyncio + async def test_pins_to_validated_ip_preserving_host_and_sni(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("93.184.216.34")) + seen = {} + + def handler(request): + seen["url_host"] = request.url.host + seen["host_header"] = request.headers.get("Host") + seen["sni"] = request.extensions.get("sni_hostname") + return httpx.Response(200, text="ok") + + transport = mock_pinned_transport(handler) + async with httpx.AsyncClient(transport=transport) as client: + resp = await client.get("https://example.com/page") + + assert resp.status_code == 200 + assert seen["url_host"] == "93.184.216.34" # connected to the pinned IP + assert seen["host_header"] == "example.com" # Host preserved + assert seen["sni"] == "example.com" # TLS SNI bound to hostname + + @pytest.mark.asyncio + async def test_blocked_host_raises_and_never_calls_inner(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("169.254.169.254")) + called = {"inner": False} + + def handler(request): + called["inner"] = True + return httpx.Response(200) + + transport = mock_pinned_transport(handler) + async with httpx.AsyncClient(transport=transport) as client: + with pytest.raises(BlockedAddressError): + await client.get("http://metadata.test/latest") + assert called["inner"] is False # never connected + + +class TestRobotsThroughTransport: + @pytest.mark.asyncio + async def test_robots_fetch_for_private_host_is_refused_at_transport(self, monkeypatch): + """A host resolving to a private IP: the robots fetch is refused by the + transport (no connection), and can_fetch stays permissive.""" + from agentic_cli.tools.webfetch.robots import RobotsTxtChecker + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("169.254.169.254")) + inner_called = {"n": 0} + def handler(req): + inner_called["n"] += 1 + return httpx.Response(200, text="User-agent: *\nDisallow: /\n") + checker = RobotsTxtChecker(transport=mock_pinned_transport(handler)) + allowed = await checker.can_fetch("http://metadata.test/x") + assert inner_called["n"] == 0 # never connected to the private IP + assert allowed is True # permissive on the (blocked) fetch error + + +class _SpyInner(httpx.MockTransport): + def __init__(self, handler): + super().__init__(handler) + self.closed = False + + async def aclose(self): + self.closed = True + await super().aclose() + + +class TestSharedTransportLifecycle: + @pytest.mark.asyncio + async def test_client_close_does_not_close_shared_inner(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", stub_getaddrinfo("93.184.216.34")) + inner = _SpyInner(lambda req: httpx.Response(200, text="ok")) + transport = PinnedTransport(URLValidator(), inner=inner) + async with httpx.AsyncClient(transport=transport) as client: + await client.get("https://example.com/x") + assert inner.closed is False # shared inner NOT closed by client exit + # ...and the transport is still usable for a subsequent client + async with httpx.AsyncClient(transport=transport) as client2: + r = await client2.get("https://example.com/y") + assert r.status_code == 200 diff --git a/tests/workflow/test_active_turn_context.py b/tests/workflow/test_active_turn_context.py new file mode 100644 index 0000000..f101877 --- /dev/null +++ b/tests/workflow/test_active_turn_context.py @@ -0,0 +1,108 @@ +"""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()`` 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 + +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", default_user="default_user") + mgr._app_name = "test" + mgr._services = {} + mgr.session_id = "default_session" + return mgr + + +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(): + mgr = _bare_manager() + 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 + + +def test_context_clears_on_exception(): + mgr = _bare_manager() + with pytest.raises(RuntimeError): + with mgr._workflow_context(session_id="s", user_id="u"): + assert mgr.active_session_id == "s" + 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_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_adk_claude.py b/tests/workflow/test_adk_claude.py new file mode 100644 index 0000000..8ca5735 --- /dev/null +++ b/tests/workflow/test_adk_claude.py @@ -0,0 +1,178 @@ +"""Native Claude (direct-API ``AnthropicLlm``) support on the ADK backend. + +Covers the thinking-budget mapping (Claude is budget-based, like Gemini 2.5), +``max_tokens`` coordination (Anthropic requires ``max_tokens > thinking budget`` +and reads it from the model instance), and model-arg construction that avoids the +Vertex ``Claude`` class the registry would otherwise resolve a ``claude-*`` string +to. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from google.adk.models.anthropic_llm import AnthropicLlm # noqa: E402 + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.model_settings import ( # noqa: E402 + ModelSettings, + ThinkingSettings, +) + +CLAUDE = "claude-sonnet-4-5" + + +def _manager(mock_context, model: str) -> GoogleADKWorkflowManager: + return GoogleADKWorkflowManager( + agent_configs=[], settings=mock_context.settings, model=model + ) + + +def _cfg(model_settings=None, model=None) -> AgentConfig: + return AgentConfig(name="a", prompt="p", model=model, model_settings=model_settings) + + +# --------------------------------------------------------------------------- +# Thinking → Anthropic budget mapping (via the planner) +# --------------------------------------------------------------------------- + + +class TestClaudePlannerThinking: + @pytest.mark.parametrize( + "effort,budget", [("low", 4096), ("medium", 10000), ("high", 32000)] + ) + def test_effort_maps_to_budget(self, mock_context, effort, budget): + mgr = _manager(mock_context, CLAUDE) + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode=effort))) + ) + assert planner is not None + assert planner.thinking_config.thinking_budget == budget + assert planner.thinking_config.thinking_level is None # never level-based + + def test_budget_mode_explicit(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="budget", budget_tokens=20000))) + ) + assert planner.thinking_config.thinking_budget == 20000 + + def test_budget_mode_floored_at_anthropic_minimum(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="budget", budget_tokens=200))) + ) + assert planner.thinking_config.thinking_budget == 1024 # floored + + def test_global_effort_fallback(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + mgr._settings.set_thinking_effort("high") + planner = mgr._get_planner(_cfg()) + assert planner.thinking_config.thinking_budget == 32000 + + def test_per_agent_none_disables(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + mgr._settings.set_thinking_effort("high") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="none"))) + ) + assert planner is None + + +# --------------------------------------------------------------------------- +# Model-arg construction (AnthropicLlm instance, not Vertex Claude string) +# --------------------------------------------------------------------------- + + +class TestClaudeModelArg: + def test_claude_returns_anthropic_llm_instance(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + arg = mgr._build_model_arg(_cfg()) + assert isinstance(arg, AnthropicLlm) + assert arg.model == CLAUDE + + def test_gemini_passes_through_as_string(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + assert mgr._build_model_arg(_cfg()) == "gemini-2.5-flash" + + def test_per_agent_model_override_to_claude(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + arg = mgr._build_model_arg(_cfg(model=CLAUDE)) + assert isinstance(arg, AnthropicLlm) + assert arg.model == CLAUDE + + +# --------------------------------------------------------------------------- +# max_tokens coordination (must exceed thinking budget) +# --------------------------------------------------------------------------- + + +class TestClaudeMaxTokens: + def test_default_without_thinking(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + mgr._settings.set_thinking_effort("none") + assert mgr._build_model_arg(_cfg()).max_tokens == 8192 + + def test_exceeds_high_thinking_budget(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + mgr._settings.set_thinking_effort("high") # budget 32000 + arg = mgr._build_model_arg(_cfg()) + assert arg.max_tokens > 32000 + + def test_honors_explicit_floor_without_thinking(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + mgr._settings.set_thinking_effort("none") + arg = mgr._build_model_arg(_cfg(ModelSettings(max_tokens=20000))) + assert arg.max_tokens == 20000 + + def test_grows_above_explicit_floor_when_budget_larger(self, mock_context): + mgr = _manager(mock_context, CLAUDE) + arg = mgr._build_model_arg( + _cfg(ModelSettings(max_tokens=10000, thinking=ThinkingSettings(mode="high"))) + ) + assert arg.max_tokens > 32000 # budget+headroom beats the 10k floor + + +# --------------------------------------------------------------------------- +# generate_simple routes Claude through the Anthropic SDK (not the genai client) +# --------------------------------------------------------------------------- + + +class TestClaudeGenerateSimple: + async def test_routes_to_anthropic_sdk(self, mock_context, monkeypatch): + mgr = _manager(mock_context, CLAUDE) + + async def _noop(): + return None + + monkeypatch.setattr(mgr, "_ensure_initialized", _noop) + + captured: dict = {} + + class _Block: + type = "text" + text = "hello" + + class _Msg: + content = [_Block()] + + class _Client: + class messages: + @staticmethod + async def create(**kwargs): + captured.update(kwargs) + return _Msg() + + import anthropic + + monkeypatch.setattr(anthropic, "AsyncAnthropic", lambda: _Client()) + + out = await mgr.generate_simple("hi", max_tokens=123) + + assert out == "hello" + assert captured["model"] == CLAUDE + assert captured["max_tokens"] == 123 + assert captured["messages"] == [{"role": "user", "content": "hi"}] diff --git a/tests/workflow/test_adk_claude_adaptive.py b/tests/workflow/test_adk_claude_adaptive.py new file mode 100644 index 0000000..5adb91f --- /dev/null +++ b/tests/workflow/test_adk_claude_adaptive.py @@ -0,0 +1,155 @@ +"""Adaptive-thinking + effort policy for Claude on ADK. + +Claude >= 4.6 rejects ``budget_tokens`` (400 on 4.7+/Fable), so the manager +switches those to adaptive thinking (negative budget) + ``output_config.effort``. +Claude <= 4.5 keeps the legacy numeric-budget path with no effort. + +These assert the config the manager produces (planner budget, model effort/ +max_tokens) — they don't make API calls, so they're version-independent (the +negative-budget -> adaptive mapping itself is ADK >= 1.34, exercised live). +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk.manager import ( # noqa: E402 + GoogleADKWorkflowManager, + _anthropic_uses_adaptive, +) +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.model_settings import ( # noqa: E402 + ModelSettings, + ThinkingSettings, +) + +ADAPTIVE = "claude-opus-4-8" # >= 4.6 +LEGACY = "claude-sonnet-4-5" # <= 4.5 + + +def _manager(mock_context, model: str) -> GoogleADKWorkflowManager: + return GoogleADKWorkflowManager( + agent_configs=[], settings=mock_context.settings, model=model + ) + + +def _cfg(model_settings=None, model=None) -> AgentConfig: + return AgentConfig(name="a", prompt="p", model=model, model_settings=model_settings) + + +class TestAdaptiveClassification: + @pytest.mark.parametrize( + "model,adaptive", + [ + ("claude-opus-4-8", True), + ("claude-opus-4-7", True), + ("claude-opus-4-6", True), + ("claude-sonnet-4-6", True), + ("claude-fable-5", True), + ("claude-opus-4-5", False), + ("claude-sonnet-4-5", False), + ("claude-haiku-4-5", False), + ("claude-sonnet-4", False), + ("claude-3-5-sonnet", False), + ("gemini-2.5-pro", False), + # Date-suffixed ids must parse to the version, not the date: a + # dated 4.0 id is < 4.6 (budget path), not adaptive. + ("claude-opus-4-20250514", False), + ("claude-sonnet-4-20250514", False), + ("claude-opus-4-1-20250805", False), + ("claude-haiku-4-5-20251001", False), + ], + ) + def test_uses_adaptive(self, model, adaptive): + assert _anthropic_uses_adaptive(model) is adaptive + + +class TestAdaptivePlanner: + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_uses_negative_budget(self, mock_context, effort): + mgr = _manager(mock_context, ADAPTIVE) + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode=effort))) + ) + assert planner is not None + assert planner.thinking_config.thinking_budget == -1 # adaptive sentinel + assert planner.thinking_config.thinking_level is None + + def test_legacy_uses_positive_budget(self, mock_context): + mgr = _manager(mock_context, LEGACY) + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="high"))) + ) + assert planner.thinking_config.thinking_budget == 32000 + + def test_none_disables_on_adaptive(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + mgr._settings.set_thinking_effort("none") + assert mgr._get_planner(_cfg()) is None + + +class TestAdaptiveEffort: + @pytest.mark.parametrize( + "mode,effort", [("low", "low"), ("medium", "medium"), ("high", "high")] + ) + def test_effort_mapped_for_adaptive(self, mock_context, mode, effort): + mgr = _manager(mock_context, ADAPTIVE) + arg = mgr._build_model_arg( + _cfg(ModelSettings(thinking=ThinkingSettings(mode=mode))) + ) + assert arg.effort == effort + + def test_global_effort_fallback_for_adaptive(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + mgr._settings.set_thinking_effort("medium") + assert mgr._build_model_arg(_cfg()).effort == "medium" + + def test_no_effort_for_legacy(self, mock_context): + mgr = _manager(mock_context, LEGACY) + arg = mgr._build_model_arg( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="high"))) + ) + assert arg.effort is None + + def test_no_effort_for_budget_mode(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + arg = mgr._build_model_arg( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="budget", budget_tokens=5000))) + ) + assert arg.effort is None + + def test_no_effort_when_thinking_disabled(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + mgr._settings.set_thinking_effort("none") + assert mgr._build_model_arg(_cfg()).effort is None + + +class TestAdaptiveMaxTokens: + def test_adaptive_does_not_inflate_max_tokens(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + mgr._settings.set_thinking_effort("high") # would be 32000 budget on legacy + assert mgr._build_model_arg(_cfg()).max_tokens == 8192 # base default + + def test_adaptive_honors_explicit_max_tokens(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + mgr._settings.set_thinking_effort("high") + arg = mgr._build_model_arg(_cfg(ModelSettings(max_tokens=20000))) + assert arg.max_tokens == 20000 + + +class TestClaudeRequestOptions: + def test_legacy_high_effort_sets_request_timeout(self, mock_context): + """Legacy high effort inflates max_tokens past the SDK's non-streaming + ceiling (~21k), so a request_timeout must be set to avoid the guard.""" + mgr = _manager(mock_context, LEGACY) + mgr._settings.set_thinking_effort("high") + arg = mgr._build_model_arg(_cfg()) + assert arg.max_tokens > 21333 + assert arg.request_timeout is not None + + def test_max_retries_from_settings(self, mock_context): + mgr = _manager(mock_context, ADAPTIVE) + arg = mgr._build_model_arg(_cfg()) + assert arg.max_retries == mock_context.settings.retry_max_attempts diff --git a/tests/workflow/test_adk_config_bridge.py b/tests/workflow/test_adk_config_bridge.py new file mode 100644 index 0000000..491fd7e --- /dev/null +++ b/tests/workflow/test_adk_config_bridge.py @@ -0,0 +1,149 @@ +"""Tests for reusing an existing native ADK config (Phase 6).""" + +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.workflow.adk_config_bridge import ( # noqa: E402 + load_adk_agent_native, + translate_adk_yaml, +) +from agentic_cli.workflow.factory import ( # noqa: E402 + create_workflow_manager_from_settings, +) + + +def _write(path: Path, text: str) -> Path: + path.write_text(textwrap.dedent(text), encoding="utf-8") + return path + + +def _native_yaml(tmp_path: Path) -> Path: + return _write( + tmp_path / "root_agent.yaml", + """ + agent_class: LlmAgent + name: root + model: gemini-2.5-flash + instruction: You are root. + """, + ) + + +def _tree_yaml(tmp_path: Path) -> Path: + _write( + tmp_path / "child.yaml", + """ + name: child + model: gemini-2.5-flash + instruction: Do work. + tools: + - name: web_search + """, + ) + return _write( + tmp_path / "root.yaml", + """ + name: coordinator + model: gemini-2.5-pro + description: Root + instruction: Coordinate. + generate_content_config: + temperature: 0.3 + max_output_tokens: 2048 + thinking_config: + thinking_budget: 8000 + tools: + - name: read_file + - name: kb_search + args: + - {name: x, value: 1} + planner: + thinking_config: {thinking_budget: 100} + sub_agents: + - config_path: child.yaml + """, + ) + + +# --------------------------------------------------------------------------- +# Native passthrough +# --------------------------------------------------------------------------- + + +class TestNativeLoad: + def test_loads_agent_tree(self, tmp_path): + agent = load_adk_agent_native(_native_yaml(tmp_path)) + assert agent.name == "root" + + +# --------------------------------------------------------------------------- +# Translate +# --------------------------------------------------------------------------- + + +class TestTranslate: + def test_translates_tree(self, tmp_path): + configs = {c.name: c for c in translate_adk_yaml(_tree_yaml(tmp_path))} + assert set(configs) == {"coordinator", "child"} + + coord = configs["coordinator"] + assert coord.model == "gemini-2.5-pro" + assert coord.description == "Root" + assert coord.get_prompt() == "Coordinate." + assert coord.sub_agents == ["child"] + # tools keep names; args-bearing tool keeps its name (args dropped). + assert coord.tools == ["read_file", "kb_search"] + # generate_content_config -> model_settings + assert coord.model_settings.temperature == 0.3 + assert coord.model_settings.max_tokens == 2048 + assert coord.model_settings.thinking.mode == "budget" + assert coord.model_settings.thinking.budget_tokens == 8000 + + assert configs["child"].tools == ["web_search"] + + def test_no_generate_config_means_no_model_settings(self, tmp_path): + path = _write( + tmp_path / "a.yaml", + """ + name: solo + model: gemini-2.5-flash + instruction: hi + """, + ) + cfg = translate_adk_yaml(path)[0] + assert cfg.model_settings is None + + +# --------------------------------------------------------------------------- +# Factory integration +# --------------------------------------------------------------------------- + + +class TestFactoryIntegration: + def test_native_mode_returns_adk_manager(self, tmp_path, mock_context): + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + mgr = create_workflow_manager_from_settings( + [], mock_context.settings, + adk_config_path=str(_native_yaml(tmp_path)), + adk_config_mode="native", + ) + assert isinstance(mgr, GoogleADKWorkflowManager) + assert mgr._adk_config_path == str(_native_yaml(tmp_path)) + + def test_translate_mode_uses_translated_configs(self, tmp_path, mock_context): + mgr = create_workflow_manager_from_settings( + [], mock_context.settings, + adk_config_path=str(_tree_yaml(tmp_path)), + adk_config_mode="translate", + ) + assert {c.name for c in mgr.agent_configs} == {"coordinator", "child"} + # Translated string tool refs were resolved at manager construction. + coord = next(c for c in mgr.agent_configs if c.name == "coordinator") + assert all(callable(t) for t in coord.tools) diff --git a/tests/workflow/test_adk_direct_anthropic.py b/tests/workflow/test_adk_direct_anthropic.py new file mode 100644 index 0000000..6c5000e --- /dev/null +++ b/tests/workflow/test_adk_direct_anthropic.py @@ -0,0 +1,209 @@ +"""Tests for DirectAnthropicLlm (Seam B) and claude-string registration. + +Covers: effort → output_config injection, the extra-param client wrapper +(merge + per-call precedence + delegation), and that registering the subclass +makes ``LLMRegistry`` resolve ``claude-*`` strings to it instead of Vertex Claude. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") +pytest.importorskip("anthropic") + +from google.adk.models.anthropic_llm import AnthropicLlm, Claude # noqa: E402 +from google.adk.models.registry import LLMRegistry # noqa: E402 +import google.adk.models.registry as _registry_mod # noqa: E402 + +from agentic_cli.workflow.adk.anthropic_llm import ( # noqa: E402 + DirectAnthropicLlm, + _ExtraParamClient, + register_direct_anthropic, +) + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class _FakeMessages: + def __init__(self) -> None: + self.calls: list[dict] = [] + + async def create(self, **kwargs): + self.calls.append(kwargs) + return "RESULT" + + +class _FakeClient: + def __init__(self) -> None: + self.messages = _FakeMessages() + self.base_url = "https://example.test" # for delegation test + + +# --------------------------------------------------------------------------- +# effort → output_config +# --------------------------------------------------------------------------- + + +class TestExtraCreateParams: + def test_effort_maps_to_output_config(self): + llm = DirectAnthropicLlm(model="claude-opus-4-8", effort="high") + assert llm._extra_create_params() == {"output_config": {"effort": "high"}} + + def test_no_effort_no_extra(self): + llm = DirectAnthropicLlm(model="claude-opus-4-8") + assert llm._extra_create_params() == {} + + def test_extra_params_plus_effort_merge(self): + llm = DirectAnthropicLlm( + model="claude-opus-4-8", effort="max", + extra_params={"service_tier": "auto"}, + ) + params = llm._extra_create_params() + assert params["service_tier"] == "auto" + assert params["output_config"] == {"effort": "max"} + + def test_explicit_output_config_effort_preserved(self): + llm = DirectAnthropicLlm( + model="claude-opus-4-8", effort="low", + extra_params={"output_config": {"effort": "max"}}, + ) + # setdefault: an explicit output_config.effort wins over the field + assert llm._extra_create_params()["output_config"] == {"effort": "max"} + + +# --------------------------------------------------------------------------- +# Client wrapper: merge, per-call precedence, delegation +# --------------------------------------------------------------------------- + + +class TestExtraParamClient: + async def test_merges_extra_into_create(self): + fake = _FakeClient() + wrapped = _ExtraParamClient(fake, {"output_config": {"effort": "high"}}) + out = await wrapped.messages.create(model="claude-opus-4-8", messages=[]) + assert out == "RESULT" + call = fake.messages.calls[0] + assert call["output_config"] == {"effort": "high"} + assert call["model"] == "claude-opus-4-8" + assert call["messages"] == [] + + async def test_per_call_kwarg_wins_over_extra(self): + fake = _FakeClient() + wrapped = _ExtraParamClient(fake, {"service_tier": "auto"}) + await wrapped.messages.create(model="m", service_tier="batch") + assert fake.messages.calls[0]["service_tier"] == "batch" + + def test_delegates_unknown_attributes(self): + fake = _FakeClient() + wrapped = _ExtraParamClient(fake, {}) + assert wrapped.base_url == "https://example.test" + + +# --------------------------------------------------------------------------- +# _anthropic_client wiring (no real network / API key) +# --------------------------------------------------------------------------- + + +class TestAnthropicClientWiring: + def test_no_extra_returns_plain_client(self, monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + llm = DirectAnthropicLlm(model="claude-opus-4-8") # no effort + from anthropic import AsyncAnthropic + + assert isinstance(llm._anthropic_client, AsyncAnthropic) + + def test_effort_returns_wrapped_client(self, monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + llm = DirectAnthropicLlm(model="claude-opus-4-8", effort="high") + assert isinstance(llm._anthropic_client, _ExtraParamClient) + + +class TestClientRequestOptions: + def test_client_receives_timeout_and_retries(self, monkeypatch): + import anthropic + + captured: dict = {} + + class _FakeAsyncAnthropic: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(anthropic, "AsyncAnthropic", _FakeAsyncAnthropic) + llm = DirectAnthropicLlm( + model="claude-opus-4-8", request_timeout=900.0, max_retries=5 + ) + _ = llm._anthropic_client + assert captured["timeout"] == 900.0 + assert captured["max_retries"] == 5 + + def test_large_max_tokens_skips_streaming_guard(self, monkeypatch): + """A non-default client timeout is exactly what makes messages.create + skip the 'streaming required' guard for max_tokens > ~21k (B-1).""" + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + from anthropic._constants import DEFAULT_TIMEOUT + + llm = DirectAnthropicLlm( + model="claude-sonnet-4-5", max_tokens=40000, request_timeout=900.0 + ) + assert llm._anthropic_client.timeout != DEFAULT_TIMEOUT + + def test_defaults_leave_client_unchanged(self, monkeypatch): + """No options set → plain client with the SDK default timeout.""" + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + from anthropic._constants import DEFAULT_TIMEOUT + + llm = DirectAnthropicLlm(model="claude-opus-4-8") + assert llm._anthropic_client.timeout == DEFAULT_TIMEOUT + + +# --------------------------------------------------------------------------- +# Registry: claude strings → DirectAnthropicLlm +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _isolated_registry(): + """Snapshot and restore the global LLM registry + resolve cache.""" + snapshot = dict(_registry_mod._llm_registry_dict) + LLMRegistry.resolve.cache_clear() + try: + yield + finally: + _registry_mod._llm_registry_dict.clear() + _registry_mod._llm_registry_dict.update(snapshot) + LLMRegistry.resolve.cache_clear() + + +class TestRegistration: + def test_register_overrides_vertex_claude(self, _isolated_registry): + # Self-contained: force ADK's default (Vertex Claude), confirm it, then + # confirm our registration overrides it for the same string. + LLMRegistry.register(Claude) + LLMRegistry.resolve.cache_clear() + assert LLMRegistry.resolve("claude-sonnet-4-5") is Claude + + register_direct_anthropic() + assert LLMRegistry.resolve("claude-sonnet-4-5") is DirectAnthropicLlm + + def test_register_resolves_3x_and_4x(self, _isolated_registry): + register_direct_anthropic() + assert LLMRegistry.resolve("claude-opus-4-8") is DirectAnthropicLlm + assert LLMRegistry.resolve("claude-3-5-haiku") is DirectAnthropicLlm + + def test_catch_all_covers_newer_families(self, _isolated_registry): + register_direct_anthropic() + # claude-fable-5 isn't matched by ADK's claude-.*-4.* patterns at all. + assert LLMRegistry.resolve("claude-fable-5") is DirectAnthropicLlm + + def test_new_llm_constructs_instance_with_defaults(self, _isolated_registry): + register_direct_anthropic() + inst = LLMRegistry.new_llm("claude-opus-4-8") + assert isinstance(inst, DirectAnthropicLlm) + assert isinstance(inst, AnthropicLlm) # still an AnthropicLlm + assert inst.model == "claude-opus-4-8" + assert inst.max_tokens == 8192 # class default — per-agent override needs an instance + assert inst.effort is None diff --git a/tests/workflow/test_adk_job_resume.py b/tests/workflow/test_adk_job_resume.py new file mode 100644 index 0000000..dbda6a4 --- /dev/null +++ b/tests/workflow/test_adk_job_resume.py @@ -0,0 +1,199 @@ +"""ADK idiomatic push/resume execution (phase-2 milestone 2). + +Covers wrapping ``long_running`` tools as ``LongRunningFunctionTool`` and +``resume_with_job_result`` building the right ``FunctionResponse`` and streaming +the follow-up turn — exercised against a fake runner (no live model). +""" + +from __future__ import annotations + +import asyncio + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("google.adk") + +from google.adk.tools import LongRunningFunctionTool # noqa: E402 + +from agentic_cli.tools.jobs import JobState # noqa: E402 +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.events import EventType, WorkflowEvent # noqa: E402 + + +# --- _wrap_long_running -------------------------------------------------------- + + +def test_wrap_long_running_wraps_only_flagged_tools(): + from agentic_cli.tools.registry import ToolCategory, get_registry, register_tool + from agentic_cli.workflow.permissions import EXEMPT + + @register_tool( + category=ToolCategory.EXECUTION, capabilities=EXEMPT, + long_running=True, description="probe long-running", + ) + def _resume_probe_lr(x: str = "") -> dict: + return {"success": True} + + @register_tool( + category=ToolCategory.OTHER, capabilities=EXEMPT, description="probe normal", + ) + def _resume_probe_normal(x: str = "") -> dict: + return {"success": True} + + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + wrapped = mgr._wrap_long_running([_resume_probe_lr, _resume_probe_normal]) + + assert isinstance(wrapped[0], LongRunningFunctionTool) + assert wrapped[1] is _resume_probe_normal + # Idempotent: an already-wrapped tool is not double-wrapped. + again = mgr._wrap_long_running(wrapped) + assert isinstance(again[0], LongRunningFunctionTool) + assert again[1] is _resume_probe_normal + assert get_registry().get("_resume_probe_lr").long_running is True + + +# --- resume_with_job_result ---------------------------------------------------- + + +class _FakeRunner: + def __init__(self) -> None: + self.calls: list = [] + + async def run_async(self, *, session_id, user_id, new_message, run_config): + self.calls.append( + SimpleNamespace(session_id=session_id, user_id=user_id, new_message=new_message) + ) + for n in ("ev-1", "ev-2"): + yield n + + +def _resume_manager(runner: _FakeRunner, *, session_exists: bool = True) -> GoogleADKWorkflowManager: + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = SimpleNamespace( + app_name="test", context_window_enabled=False, default_user="default_user" + ) + mgr._app_name = "test" + mgr._services = {} + # 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 + + async def _process_event(adk_event, session_id): + yield WorkflowEvent(type=EventType.TEXT, content=str(adk_event)) + + mgr._event_processor = SimpleNamespace(model=None, process_event=_process_event) + + async def _ensure_initialized(): + return None + + mgr._ensure_initialized = _ensure_initialized + + class _SessionService: + async def get_session(self, *, app_name, user_id, session_id): + return object() if session_exists else None + + mgr._session_service = _SessionService() + return mgr + + +def _record(**over): + base = dict( + job_id="j1", tool="run_shell_job", call_id="fc-1", call_name="run_shell_job", + session_id="s1", user_id="u1", state=JobState.SUCCEEDED, exit_code=0, error=None, + ) + base.update(over) + return SimpleNamespace(**base) + + +async def test_resume_sends_function_response_and_streams(): + runner = _FakeRunner() + mgr = _resume_manager(runner) + events = [ev async for ev in mgr.resume_with_job_result(_record(), result="build output")] + + # One WorkflowEvent per adk event from the fake runner. + assert [e.content for e in events] == ["ev-1", "ev-2"] + assert len(runner.calls) == 1 + + part = runner.calls[0].new_message.parts[0] + fr = part.function_response + assert fr.id == "fc-1" + assert fr.name == "run_shell_job" + assert fr.response["state"] == "succeeded" + assert fr.response["exit_code"] == 0 + assert fr.response["result_summary"] == "build output" + assert "job_result('j1')" in fr.response["hint"] + # Active turn was cleared after the resume. + assert mgr.active_session_id is None + + +async def test_resume_fetches_result_from_job_manager_when_omitted(): + runner = _FakeRunner() + mgr = _resume_manager(runner) + mgr._services["job_manager"] = SimpleNamespace(result=lambda jid: f"fetched:{jid}") + + [ev async for ev in mgr.resume_with_job_result(_record())] + fr = runner.calls[0].new_message.parts[0].function_response + assert fr.response["result_summary"] == "fetched:j1" + + +async def test_resume_missing_call_id_yields_nothing(): + runner = _FakeRunner() + mgr = _resume_manager(runner) + events = [ev async for ev in mgr.resume_with_job_result(_record(call_id=None), result="x")] + assert events == [] + assert runner.calls == [] + + +async def test_resume_missing_session_yields_nothing(): + runner = _FakeRunner() + mgr = _resume_manager(runner, session_exists=False) + events = [ev async for ev in mgr.resume_with_job_result(_record(), result="x")] + assert events == [] + assert runner.calls == [] + + +# --- can_resume ---------------------------------------------------------------- + + +class _SessionSvc: + def __init__(self, present: bool) -> None: + self._present = present + + async def get_session(self, *, app_name, user_id, session_id): + return object() if self._present else None + + +def _can_resume_manager(present: bool) -> GoogleADKWorkflowManager: + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._app_name = "test" + mgr._session_service = _SessionSvc(present) + return mgr + + +async def test_can_resume_true_when_session_present(): + assert await _can_resume_manager(True).can_resume(_record()) is True + + +async def test_can_resume_false_when_session_missing(): + # Sessions are durable by default and survive a restart; this is the case + # where the conversation is genuinely gone — deleted, or an explicitly + # ephemeral session_store="memory" run whose process restarted. + assert await _can_resume_manager(False).can_resume(_record()) is False + + +async def test_can_resume_false_when_ids_missing(): + mgr = _can_resume_manager(True) + assert await mgr.can_resume(_record(call_id=None)) is False + assert await mgr.can_resume(_record(session_id=None)) is False diff --git a/tests/workflow/test_adk_mcp_permissions.py b/tests/workflow/test_adk_mcp_permissions.py new file mode 100644 index 0000000..f40250f --- /dev/null +++ b/tests/workflow/test_adk_mcp_permissions.py @@ -0,0 +1,137 @@ +"""MCP tools are gated through the permission engine, not hard-denied (Phase 4).""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.config import BaseSettings # noqa: E402 +from agentic_cli.workflow.adk.permission_plugin import PermissionPlugin # noqa: E402 +from agentic_cli.workflow.permissions import PermissionEngine # noqa: E402 +from agentic_cli.workflow.permissions.prompt import ALLOW_ONCE_CHOICE # noqa: E402 +from agentic_cli.workflow.permissions.rules import ( # noqa: E402 + Effect, + Rule, + RuleSource, +) +from agentic_cli.workflow.permissions.store import PermissionContext # noqa: E402 +from agentic_cli.workflow.service_registry import ( # noqa: E402 + PERMISSION_ENGINE, + set_service_registry, +) + + +_McpToolBase = pytest.importorskip("google.adk.tools.mcp_tool").McpTool + + +class _MCPTool(_McpToolBase): + """A real ``McpTool`` instance without the MCP session plumbing. + + 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: + def __init__(self, name: str): + self.name = name + + +class _StubWorkflow: + def __init__(self, response: str): + 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): + token = set_service_registry({PERMISSION_ENGINE: engine}) + try: + return await PermissionPlugin().before_tool_callback( + tool=tool, tool_args={}, tool_context=None + ) + finally: + token.var.reset(token) + + +class TestMCPPermissions: + async def test_allow_rule_allows(self, tmp_path): + eng = _engine( + tmp_path, + rules=[Rule("mcp", "*", Effect.ALLOW, RuleSource.SESSION)], + ) + assert await _check(eng, _MCPTool("notion_search")) is None + + async def test_deny_rule_denies(self, tmp_path): + eng = _engine( + tmp_path, + rules=[Rule("mcp", "*", Effect.DENY, RuleSource.SESSION)], + ) + res = await _check(eng, _MCPTool("notion_search")) + assert res is not None and res["success"] is False + + async def test_no_rule_asks_user_and_allows(self, tmp_path): + eng = _engine(tmp_path, response=ALLOW_ONCE_CHOICE) + assert await _check(eng, _MCPTool("notion_search")) is None + + async def test_no_rule_asks_user_and_denies(self, tmp_path): + eng = _engine(tmp_path, response="deny") # unknown choice -> DENY + res = await _check(eng, _MCPTool("notion_search")) + assert res is not None and res["success"] is False + + async def test_engine_absent_denies_when_permissions_enabled(self, tmp_path, monkeypatch): + # No engine but permissions enabled -> fail closed (deny), not allow. + from types import SimpleNamespace + + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=True), + ) + token = set_service_registry({}) + try: + res = await PermissionPlugin().before_tool_callback( + tool=_MCPTool("notion_search"), tool_args={}, tool_context=None + ) + finally: + token.var.reset(token) + assert res is not None and res["success"] is False + + async def test_engine_absent_allows_when_permissions_disabled(self, tmp_path, monkeypatch): + # Permissions off is the only case a missing engine allows. + from types import SimpleNamespace + + monkeypatch.setattr( + "agentic_cli.config.get_settings", + lambda: SimpleNamespace(permissions_enabled=False), + ) + token = set_service_registry({}) + try: + res = await PermissionPlugin().before_tool_callback( + tool=_MCPTool("notion_search"), tool_args={}, tool_context=None + ) + finally: + token.var.reset(token) + assert res is None + + async def test_non_mcp_unknown_tool_still_denied(self, tmp_path): + eng = _engine(tmp_path) + res = await _check(eng, _PlainTool("random_tool")) + assert res is not None and res["success"] is False diff --git a/tests/workflow/test_adk_session_resume.py b/tests/workflow/test_adk_session_resume.py deleted file mode 100644 index d0709e5..0000000 --- a/tests/workflow/test_adk_session_resume.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Tests that ADK session resume actually persists restored messages. - -The bug: _inject_session_messages appended events to the copy returned by -create_session, leaving the stored session empty — so resume silently restored -nothing. The fix uses append_event with real Event objects. -""" - -import pytest - -pytest.importorskip("google.adk") - -from google.adk.sessions import InMemorySessionService - -from agentic_cli.config import BaseSettings -from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager - - -def _manager(tmp_path) -> tuple[GoogleADKWorkflowManager, BaseSettings]: - settings = BaseSettings(workspace_dir=tmp_path) - mgr = GoogleADKWorkflowManager( - agent_configs=[], settings=settings, model="gemini-2.0-flash" - ) - mgr._session_service = InMemorySessionService() - return mgr, settings - - -async def test_inject_persists_to_stored_session(tmp_path): - mgr, settings = _manager(tmp_path) - messages = [ - {"role": "user", "content": "hello there"}, - {"role": "assistant", "content": "hi"}, - {"role": "assistant", "content": "", "tool_calls": [ - {"id": "c1", "name": "search", "args": {"q": "x"}} - ]}, - {"role": "tool", "tool_call_id": "c1", "name": "search", "content": '{"result": 42}'}, - ] - - await mgr._inject_session_messages("s1", messages, current_agent="root") - - stored = await mgr._session_service.get_session( - app_name=mgr.app_name, user_id=settings.default_user, session_id="s1" - ) - assert stored is not None - # Was 0 before the fix (events appended to the discarded copy). - assert len(stored.events) == 4 - - -async def test_inject_extract_roundtrip(tmp_path): - mgr, _ = _manager(tmp_path) - messages = [ - {"role": "user", "content": "remember the alpha value"}, - {"role": "assistant", "content": "noted"}, - ] - - await mgr._inject_session_messages("s2", messages, current_agent="root") - got, _agent = await mgr._extract_session_data("s2") - - assert [m["role"] for m in got] == ["user", "assistant"] - assert "remember the alpha value" in got[0]["content"] - assert got[1]["content"] == "noted" diff --git a/tests/workflow/test_adk_transfer_tool.py b/tests/workflow/test_adk_transfer_tool.py new file mode 100644 index 0000000..6b76754 --- /dev/null +++ b/tests/workflow/test_adk_transfer_tool.py @@ -0,0 +1,403 @@ +"""ADK's built-in agent-transfer tool: permission identity and description. + +Two verified defects blocked multi-agent delegation, and both are covered here. + +1. **Permission identity.** ADK auto-injects an exact + ``TransferToAgentTool`` into any agent with ``sub_agents``. It is a + ``FunctionTool`` *subclass*, and the plugin only unwrapped exact + ``FunctionTool``/``LongRunningFunctionTool``, so the routing primitive + resolved to ``_UNVERIFIED`` and was denied as unregistered — even with + permissions disabled. Delegation could not work at all. + +2. **Model-visible description.** ADK builds the declaration from the + ``transfer_to_agent`` docstring, which advises callers to "use + TransferToAgentTool instead of this function directly". The model followed + that advice and emitted ``TransferToAgentTool``, which ADK rejects. + +The identity fix must not become a hole: adding a type to the trusted list is +exactly the kind of change that can quietly re-admit forged tools, so the +negative cases are asserted alongside the positive one. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +pytest.importorskip("google.adk") + +from google.adk.tools import FunctionTool, LongRunningFunctionTool # noqa: E402 +from google.adk.tools.transfer_to_agent_tool import ( # noqa: E402 + TransferToAgentTool, + transfer_to_agent, +) +from google.genai import types # noqa: E402 + +from agentic_cli.workflow.adk.permission_plugin import ( # noqa: E402 + _UNVERIFIED, + PermissionPlugin, +) +from agentic_cli.workflow.permissions import EXEMPT # noqa: E402 +from agentic_cli.workflow.adk.transfer_tool_description import ( # noqa: E402 + ROUTING_DESCRIPTION, + TRANSFER_TOOL_NAME, + TransferToolDescriptionPlugin, + correct_transfer_declaration, +) + +AGENT_NAMES = ["arxiv_specialist", "data_analyst", "report_writer"] + + +def _tool() -> TransferToAgentTool: + return TransferToAgentTool(agent_names=list(AGENT_NAMES)) + + +def _request(tool: Any, declaration: types.FunctionDeclaration | None = None): + """A minimal LlmRequest carrying one tool and its declaration.""" + from google.adk.models.llm_request import LlmRequest + + if declaration is None: + declaration = tool._get_declaration() + return LlmRequest( + config=types.GenerateContentConfig( + tools=[types.Tool(function_declarations=[declaration])] + ), + tools_dict={declaration.name: tool}, + ) + + +#: A description in the shape ADK 1.x generates. The correction-path tests use +#: *this*, not whatever the installed ADK happens to produce: the production +#: workaround is deliberately a no-op once upstream fixes its docstring (2.x +#: already has), so tests that required the installed declaration to be broken +#: would start failing the day the pinned range picks up the fix. What must keep +#: working is the correction itself, given input that needs correcting. +MISLEADING_DESCRIPTION = ( + "Transfer the question to another agent.\n\n" + "This tool hands off control to another agent when it's more suitable to\n" + "answer the user's question according to the agent's description.\n\n" + "Note:\n" + " For most use cases, you should use TransferToAgentTool instead of this\n" + " function directly. TransferToAgentTool provides additional enum " + "constraints\n" + " that prevent LLMs from hallucinating invalid agent names.\n\n" + "Args:\n" + " agent_name: the agent name to transfer to." +) + +SAFE_DESCRIPTION = "Transfer the query to another agent." + + +def _declaration_with(description: str) -> types.FunctionDeclaration: + """The native tool's real declaration, with the description swapped. + + Name, parameters, required fields and ADK's agent-name enum all come from + the genuine tool, so schema assertions stay meaningful. + """ + declaration = _tool()._get_declaration() + declaration.description = description + return declaration + + +def _misleading_request(): + """An exact native tool whose declaration needs correcting.""" + return _request(_tool(), _declaration_with(MISLEADING_DESCRIPTION)) + + +def _description_is_safe(description: str) -> bool: + """No advice to call the class, under either ADK's wording or ours.""" + return ( + "TransferToAgentTool" not in description + and "instead of this function" not in description + ) + + +# --------------------------------------------------------------------------- +# 1. Permission identity +# --------------------------------------------------------------------------- + + +class TestTransferToolResolvesToItsRegisteredDefinition: + def test_exact_transfer_tool_resolves_to_the_exempt_definition(self): + from agentic_cli.tools.registry import identify_tool + + defn = PermissionPlugin._resolve_definition(_tool()) + + assert defn is not _UNVERIFIED, ( + "ADK's built-in transfer tool was denied as unregistered" + ) + assert defn is identify_tool(transfer_to_agent), ( + "resolved to something other than the registered transfer_to_agent" + ) + assert defn.name == TRANSFER_TOOL_NAME + assert defn.capabilities is EXEMPT, ( + "the routing primitive must stay EXEMPT, not acquire capabilities" + ) + + async def test_before_tool_callback_allows_it(self): + result = await PermissionPlugin().before_tool_callback( + tool=_tool(), + tool_args={"agent_name": "arxiv_specialist"}, + tool_context=None, + ) + assert result is None, f"the transfer tool was gated: {result}" + + +class TestIdentityGuaranteesSurvive: + """Widening the trusted-type list must not admit anything else.""" + + async def _denied(self, tool: Any) -> None: + assert PermissionPlugin._resolve_definition(tool) is _UNVERIFIED + result = await PermissionPlugin().before_tool_callback( + tool=tool, tool_args={}, tool_context=None + ) + assert result is not None and result.get("success") is False + assert "not registered" in result.get("error", "") + + async def test_arbitrary_function_tool_subclass_is_denied(self): + class SneakyTool(FunctionTool): + """A subclass may override run_async and do anything.""" + + await self._denied(SneakyTool(func=transfer_to_agent)) + + async def test_transfer_tool_subclass_is_denied(self): + class SubclassedTransfer(TransferToAgentTool): + pass + + await self._denied(SubclassedTransfer(agent_names=list(AGENT_NAMES))) + + async def test_forged_object_named_transfer_to_agent_is_denied(self): + forged = SimpleNamespace(name=TRANSFER_TOOL_NAME) + await self._denied(forged) + + async def test_forged_object_named_after_the_class_is_denied(self): + class TransferToAgentTool: # shadows the real name deliberately + name = TRANSFER_TOOL_NAME + + await self._denied(TransferToAgentTool()) + + async def test_forged_object_carrying_the_real_func_is_denied(self): + """A copied ``.func`` is an attribute, not evidence of behaviour.""" + forged = SimpleNamespace(name=TRANSFER_TOOL_NAME, func=transfer_to_agent) + await self._denied(forged) + + async def test_long_running_function_tool_still_resolves(self): + """The pre-existing trusted types keep working.""" + from agentic_cli.tools.registry import identify_tool + + defn = PermissionPlugin._resolve_definition( + LongRunningFunctionTool(func=transfer_to_agent) + ) + assert defn is identify_tool(transfer_to_agent) + + +# --------------------------------------------------------------------------- +# 2. Model-visible description +# --------------------------------------------------------------------------- + + +class TestTransferDeclarationIsCorrected: + """The correction path, driven by a synthetic misleading declaration. + + These never depend on the installed ADK still being broken — only on the + correction doing the right thing when handed input that needs it. + """ + + def test_exactly_one_transfer_declaration_with_the_right_name(self): + request = _misleading_request() + correct_transfer_declaration(request) + + decls = [ + d + for tool in request.config.tools + for d in tool.function_declarations + if d.name == TRANSFER_TOOL_NAME + ] + assert len(decls) == 1 + assert decls[0].name == TRANSFER_TOOL_NAME + + def test_description_no_longer_advises_calling_the_class(self): + request = _misleading_request() + assert correct_transfer_declaration(request) is True + + description = request.config.tools[0].function_declarations[0].description + assert _description_is_safe(description) + + def test_description_states_the_callable_name(self): + request = _misleading_request() + correct_transfer_declaration(request) + + description = request.config.tools[0].function_declarations[0].description + assert f'{TRANSFER_TOOL_NAME}(agent_name="' in description + + def test_agent_name_enum_and_schema_are_preserved(self): + before = _tool()._get_declaration().parameters.model_dump(exclude_none=True) + + request = _misleading_request() + correct_transfer_declaration(request) + after = request.config.tools[0].function_declarations[0].parameters + + assert after.properties["agent_name"].enum == AGENT_NAMES + assert after.required == ["agent_name"] + assert after.model_dump(exclude_none=True) == before + + def test_correction_is_idempotent(self): + request = _misleading_request() + assert correct_transfer_declaration(request) is True + first = request.config.tools[0].function_declarations[0].description + + # A second pass has nothing left to do. + assert correct_transfer_declaration(request) is False + assert request.config.tools[0].function_declarations[0].description == first + + def test_already_correct_description_is_left_alone(self): + """An ADK release that ships a sane docstring needs no correction.""" + request = _request(_tool(), _declaration_with(SAFE_DESCRIPTION)) + + assert correct_transfer_declaration(request) is False + assert ( + request.config.tools[0].function_declarations[0].description + == SAFE_DESCRIPTION + ) + + def test_the_upstream_function_doc_is_never_mutated(self): + before = transfer_to_agent.__doc__ + correct_transfer_declaration(_misleading_request()) + assert transfer_to_agent.__doc__ == before + + +class TestAgainstTheInstalledAdk: + """One integration check that works whichever docstring ADK ships. + + Either the installed declaration is misleading and the correction fixes it, + or it is already safe and no correction is needed. Both are acceptable; a + declaration that is still misleading *after* the plugin has run is not. + """ + + def test_installed_declaration_ends_up_safe(self): + request = _request(_tool()) + initial = request.config.tools[0].function_declarations[0].description or "" + + corrected = correct_transfer_declaration(request) + final = request.config.tools[0].function_declarations[0].description or "" + + if _description_is_safe(initial): + assert corrected is False, ( + "an already-safe ADK description was rewritten anyway" + ) + else: + assert corrected is True, ( + "the installed ADK description advises calling the class and " + "was not corrected" + ) + + assert _description_is_safe(final), ( + f"the model would still be told to call the class: {final!r}" + ) + # Whichever branch ran, the contract the model needs is intact. + declaration = request.config.tools[0].function_declarations[0] + assert declaration.name == TRANSFER_TOOL_NAME + assert declaration.parameters.properties["agent_name"].enum == AGENT_NAMES + assert declaration.parameters.required == ["agent_name"] + + +class TestOnlyTheNativeToolIsRewritten: + def test_application_tool_with_the_same_name_is_untouched(self): + """A same-named application tool is not ADK's routing primitive.""" + + def transfer_to_agent(agent_name: str) -> dict: # noqa: A001 - deliberate + """You should use TransferToAgentTool instead of this function.""" + return {"success": True} + + app_tool = FunctionTool(func=transfer_to_agent) + declaration = app_tool._get_declaration() + original = declaration.description + + request = _request(app_tool, declaration) + assert correct_transfer_declaration(request) is False + assert request.config.tools[0].function_declarations[0].description == original + + def test_impostor_base_tool_named_transfer_to_agent_is_untouched(self): + """A real BaseTool that merely takes the name is not ADK's tool. + + ``LlmRequest.tools_dict`` is typed to ``BaseTool``, so the realistic + impostor is a genuine tool object with the right name — which is the + case the exact-type check has to reject. + """ + from google.adk.tools import BaseTool + + class ImpostorTool(BaseTool): + def __init__(self) -> None: + super().__init__(name=TRANSFER_TOOL_NAME, description="impostor") + + declaration = types.FunctionDeclaration( + name=TRANSFER_TOOL_NAME, + description="you should use TransferToAgentTool instead", + ) + + request = _request(ImpostorTool(), declaration) + assert correct_transfer_declaration(request) is False + assert ( + request.config.tools[0].function_declarations[0].description + == "you should use TransferToAgentTool instead" + ) + + def test_subclass_of_the_native_tool_is_untouched(self): + class SubclassedTransfer(TransferToAgentTool): + pass + + tool = SubclassedTransfer(agent_names=list(AGENT_NAMES)) + request = _request(tool) + assert correct_transfer_declaration(request) is False + + +class TestPluginWiring: + async def test_plugin_corrects_the_request_and_counts_it(self): + """Deterministic: driven by a synthetic misleading declaration.""" + plugin = TransferToolDescriptionPlugin() + request = _misleading_request() + + await plugin.before_model_callback( + callback_context=SimpleNamespace(invocation_id="i1"), + llm_request=request, + ) + + assert plugin.corrections == 1 + description = request.config.tools[0].function_declarations[0].description + assert description == ROUTING_DESCRIPTION + + async def test_plugin_does_not_count_an_already_safe_request(self): + plugin = TransferToolDescriptionPlugin() + request = _request(_tool(), _declaration_with(SAFE_DESCRIPTION)) + + await plugin.before_model_callback( + callback_context=SimpleNamespace(invocation_id="i1"), + llm_request=request, + ) + + assert plugin.corrections == 0 + + async def test_plugin_never_blocks_a_request(self): + """A malformed request must not fail the turn.""" + plugin = TransferToolDescriptionPlugin() + broken = SimpleNamespace(tools_dict={TRANSFER_TOOL_NAME: _tool()}, config=None) + + result = await plugin.before_model_callback( + callback_context=SimpleNamespace(invocation_id="i1"), llm_request=broken + ) + assert result is None + + def test_manager_registers_the_plugin(self): + """The correction must actually be wired into the runner.""" + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + manager = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + manager._settings = SimpleNamespace( + raw_llm_logging=False, app_name="t", verbose_thinking=False + ) + + plugins = manager._init_plugins() + + assert any(isinstance(p, TransferToolDescriptionPlugin) for p in plugins) 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 diff --git a/tests/workflow/test_agent_loader.py b/tests/workflow/test_agent_loader.py new file mode 100644 index 0000000..bc5c283 --- /dev/null +++ b/tests/workflow/test_agent_loader.py @@ -0,0 +1,151 @@ +"""Tests for the unified YAML agent loader (Phase 3).""" + +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from agentic_cli.workflow.agent_loader import ( + create_workflow_manager_from_yaml, + load_agents_from_yaml, +) +from agentic_cli.workflow.base_manager import BaseWorkflowManager + + +def _write(tmp_path: Path, text: str, name: str = "agents.yaml") -> Path: + path = tmp_path / name + path.write_text(textwrap.dedent(text), encoding="utf-8") + return path + + +class TestLoadAgentsFromYaml: + def test_basic_load(self, tmp_path): + path = _write( + tmp_path, + """ + agents: + - name: coordinator + model: gemini-2.5-pro + description: Routes work + instruction: You are the coordinator. + tools: [kb_search, my_pkg.tools.custom] + sub_agents: [researcher] + - name: researcher + prompt: Do research. + tools: [web_search] + """, + ) + configs = load_agents_from_yaml(path) + assert [c.name for c in configs] == ["coordinator", "researcher"] + coord = configs[0] + assert coord.model == "gemini-2.5-pro" + assert coord.description == "Routes work" + assert coord.get_prompt() == "You are the coordinator." + # Tool refs stay strings (resolved later by the manager). + assert coord.tools == ["kb_search", "my_pkg.tools.custom"] + assert coord.sub_agents == ["researcher"] + + def test_model_settings_parsed(self, tmp_path): + path = _write( + tmp_path, + """ + agents: + - name: a + instruction: hi + model_settings: + temperature: 0.2 + max_tokens: 1000 + thinking: + mode: high + """, + ) + cfg = load_agents_from_yaml(path)[0] + assert cfg.model_settings is not None + assert cfg.model_settings.temperature == 0.2 + assert cfg.model_settings.max_tokens == 1000 + assert cfg.model_settings.thinking.mode == "high" + + def test_instruction_file_relative_to_yaml(self, tmp_path): + (tmp_path / "prompts").mkdir() + (tmp_path / "prompts" / "r.md").write_text("File prompt body", encoding="utf-8") + path = _write( + tmp_path, + """ + agents: + - name: a + instruction_file: prompts/r.md + """, + ) + cfg = load_agents_from_yaml(path)[0] + assert cfg.get_prompt() == "File prompt body" + + def test_bare_top_level_list(self, tmp_path): + path = _write( + tmp_path, + """ + - name: a + instruction: hi + """, + ) + configs = load_agents_from_yaml(path) + assert len(configs) == 1 and configs[0].name == "a" + + def test_defaults(self, tmp_path): + path = _write(tmp_path, "agents:\n - name: a\n instruction: hi\n") + cfg = load_agents_from_yaml(path)[0] + assert cfg.tools == [] + assert cfg.sub_agents == [] + assert cfg.description == "" + assert cfg.model is None + assert cfg.model_settings is None + assert cfg.include_state_tools is True + + def test_missing_prompt_raises(self, tmp_path): + path = _write(tmp_path, "agents:\n - name: a\n") + with pytest.raises(ValueError, match="must set 'prompt'"): + load_agents_from_yaml(path) + + def test_unknown_key_rejected(self, tmp_path): + path = _write( + tmp_path, + "agents:\n - name: a\n instruction: hi\n bogus: 1\n", + ) + with pytest.raises(ValidationError): + load_agents_from_yaml(path) + + def test_empty_agents_rejected(self, tmp_path): + path = _write(tmp_path, "agents: []\n") + with pytest.raises(ValidationError): + load_agents_from_yaml(path) + + def test_empty_file_raises(self, tmp_path): + path = _write(tmp_path, "") + with pytest.raises(ValueError, match="empty"): + load_agents_from_yaml(path) + + def test_missing_file_raises(self, tmp_path): + with pytest.raises(FileNotFoundError): + load_agents_from_yaml(tmp_path / "nope.yaml") + + +class TestCreateManagerFromYaml: + def test_builds_manager_and_resolves_tools(self, tmp_path, mock_context): + path = _write( + tmp_path, + """ + agents: + - name: a + instruction: hi + model: gemini-2.5-flash + tools: [read_file] + """, + ) + mgr = create_workflow_manager_from_yaml(path, mock_context.settings) + assert isinstance(mgr, BaseWorkflowManager) + assert [c.name for c in mgr.agent_configs] == ["a"] + # Tool string resolved to a callable at manager construction. + assert callable(mgr.agent_configs[0].tools[0]) + assert getattr(mgr.agent_configs[0].tools[0], "__name__", None) == "read_file" 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_langgraph_ignored_fields.py b/tests/workflow/test_langgraph_ignored_fields.py new file mode 100644 index 0000000..0da7593 --- /dev/null +++ b/tests/workflow/test_langgraph_ignored_fields.py @@ -0,0 +1,67 @@ +"""LangGraph warns that ADK-only AgentConfig fields are ignored.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("langgraph") + +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.langgraph import manager as mgr_mod # noqa: E402 +from agentic_cli.workflow.langgraph.manager import ( # noqa: E402 + LangGraphWorkflowManager, + _ignored_adk_only_fields, +) +from agentic_cli.workflow.mcp import MCPServerConfig # noqa: E402 +from agentic_cli.workflow.model_settings import ModelSettings # noqa: E402 + + +class TestIgnoredFieldsHelper: + def test_detects_model_settings(self): + cfg = AgentConfig(name="a", prompt="p", model_settings=ModelSettings(temperature=0.5)) + assert _ignored_adk_only_fields([cfg]) == [("a", ["model_settings"])] + + def test_detects_multiple_fields_in_canonical_order(self): + cfg = AgentConfig( + name="a", prompt="p", + skills=["x"], + mcp_servers=[MCPServerConfig(name="s", command="echo")], + ) + # Order follows _ADK_ONLY_FIELDS: model_settings, mcp_servers, skills. + assert _ignored_adk_only_fields([cfg]) == [("a", ["mcp_servers", "skills"])] + + def test_only_affected_agents_listed(self): + a = AgentConfig(name="a", prompt="p", skills=["x"]) + b = AgentConfig(name="b", prompt="p") + assert _ignored_adk_only_fields([a, b]) == [("a", ["skills"])] + + def test_empty_when_no_adk_only_fields(self): + assert _ignored_adk_only_fields([AgentConfig(name="a", prompt="p")]) == [] + + +class TestManagerWarns: + def test_warns_on_construction(self, mock_context, monkeypatch): + calls = [] + monkeypatch.setattr( + mgr_mod.logger, "warning", lambda *a, **k: calls.append((a, k)) + ) + cfg = AgentConfig(name="a", prompt="p", model_settings=ModelSettings(temperature=0.5)) + LangGraphWorkflowManager( + agent_configs=[cfg], settings=mock_context.settings, model="gemini-2.5-flash" + ) + assert any( + a and a[0] == "langgraph_ignoring_adk_only_fields" for a, _ in calls + ) + + def test_no_warning_without_adk_only_fields(self, mock_context, monkeypatch): + calls = [] + monkeypatch.setattr( + mgr_mod.logger, "warning", lambda *a, **k: calls.append((a, k)) + ) + cfg = AgentConfig(name="a", prompt="p") + LangGraphWorkflowManager( + agent_configs=[cfg], settings=mock_context.settings, model="gemini-2.5-flash" + ) + assert not any( + a and a[0] == "langgraph_ignoring_adk_only_fields" for a, _ in calls + ) 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_mcp_config.py b/tests/workflow/test_mcp_config.py new file mode 100644 index 0000000..98f5467 --- /dev/null +++ b/tests/workflow/test_mcp_config.py @@ -0,0 +1,132 @@ +"""Tests for MCP server config + ADK materialization (Phase 4).""" + +from __future__ import annotations + +import textwrap + +import pytest + +pytest.importorskip("google.adk") + +from pydantic import ValidationError # noqa: E402 + +from agentic_cli.workflow.agent_loader import load_agents_from_yaml # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.mcp import ( # noqa: E402 + MCPServerConfig, + build_connection_params, + to_adk_toolset, +) + + +class TestMCPServerConfigValidation: + def test_stdio_requires_command(self): + with pytest.raises(ValidationError): + MCPServerConfig(name="x", transport="stdio") + + def test_sse_requires_url(self): + with pytest.raises(ValidationError): + MCPServerConfig(name="x", transport="sse") + + def test_http_requires_url(self): + with pytest.raises(ValidationError): + MCPServerConfig(name="x", transport="http") + + def test_unknown_field_rejected(self): + with pytest.raises(ValidationError): + MCPServerConfig(name="x", command="npx", bogus=1) + + def test_stdio_valid(self): + cfg = MCPServerConfig(name="x", command="npx") + assert cfg.transport == "stdio" + + +class TestBuildConnectionParams: + def test_stdio(self): + from google.adk.tools.mcp_tool import StdioConnectionParams + + cfg = MCPServerConfig( + name="notion", transport="stdio", command="npx", + args=["-y", "srv"], env={"A": "B"}, + ) + conn = build_connection_params(cfg) + assert isinstance(conn, StdioConnectionParams) + assert conn.server_params.command == "npx" + assert conn.server_params.args == ["-y", "srv"] + assert conn.server_params.env == {"A": "B"} + + def test_sse(self): + from google.adk.tools.mcp_tool import SseConnectionParams + + cfg = MCPServerConfig( + name="s", transport="sse", url="https://x/sse", headers={"H": "V"} + ) + conn = build_connection_params(cfg) + assert isinstance(conn, SseConnectionParams) + assert conn.url == "https://x/sse" + assert conn.headers == {"H": "V"} + + def test_http(self): + from google.adk.tools.mcp_tool import StreamableHTTPConnectionParams + + cfg = MCPServerConfig(name="s", transport="http", url="https://x/mcp") + conn = build_connection_params(cfg) + assert isinstance(conn, StreamableHTTPConnectionParams) + assert conn.url == "https://x/mcp" + + +class TestToAdkToolset: + def test_returns_mcptoolset(self): + from google.adk.tools.mcp_tool import McpToolset + + ts = to_adk_toolset(MCPServerConfig(name="s", command="echo")) + assert isinstance(ts, McpToolset) + + def test_with_tool_filter(self): + from google.adk.tools.mcp_tool import McpToolset + + ts = to_adk_toolset( + MCPServerConfig(name="s", command="echo", tool_filter=["a", "b"]) + ) + assert isinstance(ts, McpToolset) + + +class TestAgentLoaderMCP: + def test_yaml_mcp_servers_parsed(self, tmp_path): + path = tmp_path / "a.yaml" + path.write_text( + textwrap.dedent( + """ + agents: + - name: a + instruction: hi + mcp_servers: + - name: notion + transport: stdio + command: npx + args: ["-y", "notion-mcp"] + """ + ), + encoding="utf-8", + ) + cfg = load_agents_from_yaml(path)[0] + assert len(cfg.mcp_servers) == 1 + assert cfg.mcp_servers[0].name == "notion" + assert cfg.mcp_servers[0].command == "npx" + + +class TestADKManagerAttachesMCP: + def test_agent_gets_mcp_toolset(self, mock_context): + from google.adk.tools.mcp_tool import McpToolset + + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + cfg = AgentConfig( + name="a", prompt="p", tools=[], include_state_tools=False, + mcp_servers=[MCPServerConfig(name="s", command="echo")], + ) + mgr = GoogleADKWorkflowManager( + agent_configs=[cfg], settings=mock_context.settings, model="gemini-2.5-flash" + ) + tools = mgr._assemble_agent_tools(cfg, mgr._get_service_tool_map()) + assert any(isinstance(t, McpToolset) for t in tools) diff --git a/tests/workflow/test_memory_wiring.py b/tests/workflow/test_memory_wiring.py index 3d79d75..c4fd93e 100644 --- a/tests/workflow/test_memory_wiring.py +++ b/tests/workflow/test_memory_wiring.py @@ -42,13 +42,8 @@ async def reinitialize(self, model: str | None = None, preserve_sessions: bool = async def cleanup(self) -> None: pass - async def _extract_session_data(self, session_id: str) -> tuple[list[dict], str | None]: - return getattr(self, "_fake_messages", []), None - - async def _inject_session_messages( - self, session_id: str, messages: list[dict], current_agent: str | None = None - ) -> None: - pass + async def recent_messages(self, session_id: str, limit: int = 20) -> list[dict]: + return getattr(self, "_fake_messages", []) def _get_state_tools(self) -> list: return [] diff --git a/tests/workflow/test_model_settings.py b/tests/workflow/test_model_settings.py new file mode 100644 index 0000000..1e8fbf4 --- /dev/null +++ b/tests/workflow/test_model_settings.py @@ -0,0 +1,180 @@ +"""Tests for per-agent ModelSettings (Phase 1) and its ADK translation. + +Covers: +- the backend-neutral ``ModelSettings`` model, +- ADK ``_get_generate_content_config`` (neutral params -> GenerateContentConfig), +- ADK ``_get_planner`` per-agent thinking resolution (override + global fallback). +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("google.adk") + +from google.genai import types # noqa: E402 + +from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager # noqa: E402 +from agentic_cli.workflow.config import AgentConfig # noqa: E402 +from agentic_cli.workflow.model_settings import ( # noqa: E402 + ModelSettings, + ThinkingSettings, +) + + +def _manager(mock_context, model: str) -> GoogleADKWorkflowManager: + """A manager with a pinned model (no API call / init needed).""" + return GoogleADKWorkflowManager( + agent_configs=[], settings=mock_context.settings, model=model + ) + + +def _cfg(model_settings=None, model=None) -> AgentConfig: + return AgentConfig( + name="a", prompt="p", model=model, model_settings=model_settings + ) + + +# --------------------------------------------------------------------------- +# ModelSettings model +# --------------------------------------------------------------------------- + + +class TestModelSettingsModel: + def test_defaults_all_none(self): + ms = ModelSettings() + assert ms.common_params() == {} + assert ms.thinking is None + assert ms.extra == {} + + def test_common_params_only_non_none(self): + ms = ModelSettings(temperature=0.5, max_tokens=100) + assert ms.common_params() == {"temperature": 0.5, "max_tokens": 100} + + def test_thinking_defaults(self): + ts = ThinkingSettings() + assert ts.mode == "none" + assert ts.budget_tokens is None + + +# --------------------------------------------------------------------------- +# ADK GenerateContentConfig translation +# --------------------------------------------------------------------------- + + +class TestGenerateContentConfig: + def test_no_settings_only_http_options(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + cfg = mgr._get_generate_content_config(None) + assert cfg.http_options is not None + assert cfg.temperature is None + + def test_config_without_model_settings(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + cfg = mgr._get_generate_content_config(_cfg()) + assert cfg.http_options is not None + assert cfg.temperature is None + + def test_neutral_params_mapped(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + ms = ModelSettings( + temperature=0.3, top_p=0.9, top_k=40, max_tokens=1000, + stop_sequences=["END"], + ) + cfg = mgr._get_generate_content_config(_cfg(ms)) + assert cfg.temperature == 0.3 + assert cfg.top_p == 0.9 + assert cfg.top_k == 40 + assert cfg.max_output_tokens == 1000 # neutral max_tokens -> ADK name + assert cfg.stop_sequences == ["END"] + assert cfg.http_options is not None # retry options preserved + + def test_extra_valid_key_passthrough_invalid_dropped(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + # Pick a real GenerateContentConfig field that isn't one of our neutral ones. + neutral = { + "temperature", "top_p", "top_k", "max_output_tokens", + "stop_sequences", "http_options", + } + valid_key = next( + k for k in types.GenerateContentConfig.model_fields if k not in neutral + ) + ms = ModelSettings( + temperature=0.5, extra={valid_key: 7, "definitely_not_a_field": 1} + ) + out = mgr._generate_config_kwargs_from_settings(ms) + assert out["temperature"] == 0.5 + assert out[valid_key] == 7 + assert "definitely_not_a_field" not in out + + +# --------------------------------------------------------------------------- +# ADK planner / thinking resolution +# --------------------------------------------------------------------------- + + +class TestPlannerThinking: + def test_per_agent_overrides_global_disabled(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + mgr._settings.set_thinking_effort("none") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="high"))) + ) + assert planner is not None + assert planner.thinking_config.thinking_budget == 24576 # 2.5 high + + def test_per_agent_none_disables_even_if_global_high(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + mgr._settings.set_thinking_effort("high") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="none"))) + ) + assert planner is None + + def test_global_fallback_when_no_model_settings(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + mgr._settings.set_thinking_effort("low") + planner = mgr._get_planner(_cfg()) + assert planner is not None + assert planner.thinking_config.thinking_budget == 4096 # 2.5 low + + def test_global_none_returns_no_planner(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + mgr._settings.set_thinking_effort("none") + assert mgr._get_planner(_cfg()) is None + + def test_budget_mode_uses_explicit_budget(self, mock_context): + mgr = _manager(mock_context, "gemini-2.5-flash") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="budget", budget_tokens=5000))) + ) + assert planner.thinking_config.thinking_budget == 5000 + assert planner.thinking_config.thinking_level is None + + def test_gemini3_uses_thinking_level(self, mock_context): + mgr = _manager(mock_context, "gemini-3-pro-preview") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="high"))) + ) + assert planner.thinking_config.thinking_level == types.ThinkingLevel.HIGH + assert planner.thinking_config.thinking_budget is None + + def test_per_agent_model_override_selects_family(self, mock_context): + # Manager default is 2.5 (budget path); agent overrides to gemini-3 + # flash, which must take the thinking_level path. + mgr = _manager(mock_context, "gemini-2.5-flash") + planner = mgr._get_planner( + _cfg( + ModelSettings(thinking=ThinkingSettings(mode="medium")), + model="gemini-3-flash-preview", + ) + ) + assert planner.thinking_config.thinking_level == types.ThinkingLevel.MEDIUM + assert planner.thinking_config.thinking_budget is None + + def test_unsupported_model_returns_no_planner(self, mock_context): + mgr = _manager(mock_context, "gpt-4o") + planner = mgr._get_planner( + _cfg(ModelSettings(thinking=ThinkingSettings(mode="high"))) + ) + assert planner is None diff --git a/tests/workflow/test_model_validation.py b/tests/workflow/test_model_validation.py new file mode 100644 index 0000000..a6b1c15 --- /dev/null +++ b/tests/workflow/test_model_validation.py @@ -0,0 +1,653 @@ +"""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 _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. + + ``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_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: + 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] 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)) 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_session_store.py b/tests/workflow/test_session_store.py new file mode 100644 index 0000000..62a1410 --- /dev/null +++ b/tests/workflow/test_session_store.py @@ -0,0 +1,416 @@ +"""Resumable sessions: session_db_url resolution + ADK session-service selection. + +Milestone 1 of the durable-sessions feature: the unified `session_store` setting +resolves to one async SQLAlchemy URL (shared by both backends), and the ADK +manager builds a DatabaseSessionService (durable) vs InMemory (ephemeral). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from agentic_cli.config import BaseSettings + + +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() + assert url is not None + assert url.startswith("sqlite+aiosqlite:///") + assert url.endswith("sessions/sessions.db") + + def test_memory_is_none(self, tmp_path: Path): + assert _settings(tmp_path, session_store="memory").session_db_url() is None + + def test_postgres_normalized_to_async(self, tmp_path: Path): + url = _settings( + tmp_path, session_store="postgres", postgres_uri="postgresql://u@h/db" + ).session_db_url() + assert url == "postgresql+asyncpg://u@h/db" + + def test_postgres_requires_uri(self, tmp_path: Path): + with pytest.raises(ValueError): + _settings(tmp_path, session_store="postgres").session_db_url() + + def test_explicit_sqlite_uri_normalized(self, tmp_path: Path): + url = _settings( + tmp_path, session_store="sqlite", sqlite_uri="sqlite:////tmp/custom.db" + ).session_db_url() + assert url == "sqlite+aiosqlite:////tmp/custom.db" + + +class TestAdkSessionServiceSelection: + @pytest.fixture(autouse=True) + async def _require_adk(self, _closing_session_services): + pytest.importorskip("google.adk") + + def _manager(self, settings): + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + mgr = GoogleADKWorkflowManager.__new__(GoogleADKWorkflowManager) + mgr._settings = settings + return mgr + + 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) + + 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")) + svc = mgr._make_session_service() + assert isinstance(svc, DatabaseSessionService) + assert (tmp_path / "sessions").is_dir() + + +class TestAdkNativeSessions: + """Native session query/manage against a real sqlite DatabaseSessionService.""" + + @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(self, mgr, settings, 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=settings.default_user, 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_exists_list_recent_delete(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + await self._seed(mgr, settings, "sess-x", "remember the alpha value") + + assert await mgr.session_exists("sess-x") is True + assert await mgr.session_exists("missing") is False + + listed = await mgr.list_sessions() + assert any(s["session_id"] == "sess-x" for s in listed) + + recent = await mgr.recent_messages("sess-x") + assert recent and recent[-1]["content"] == "remember the alpha value" + + assert await mgr.delete_session("sess-x") is True + assert await mgr.session_exists("sess-x") is False + + async def test_load_session_reports_resume(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + await self._seed(mgr, settings, "sess-y", "hi") + # Existing session → resumed=True and id adopted. + assert await mgr.load_session("sess-y") is True + assert mgr.session_id == "sess-y" + # Unknown session → new (False) but still adopted. + assert await mgr.load_session("brand-new") is False + assert mgr.session_id == "brand-new" + + async def test_persists_across_fresh_manager(self, tmp_path: Path): + mgr, settings = self._manager(tmp_path) + await self._seed(mgr, settings, "sess-z", "durable") + # 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", + } diff --git a/tests/workflow/test_skills.py b/tests/workflow/test_skills.py new file mode 100644 index 0000000..ccab741 --- /dev/null +++ b/tests/workflow/test_skills.py @@ -0,0 +1,172 @@ +"""Tests for skills support (Phase 5) — ADK SkillToolset wiring, scripts gated.""" + +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest + +pytest.importorskip("google.adk") + +from agentic_cli.tools.skills import ( # noqa: E402 + SkillStore, + make_skill_toolset, + register_skill_tool_permissions, +) +from agentic_cli.workflow.config import AgentConfig # noqa: E402 + + +def _make_skill(parent: Path, name: str = "pdf-tools", desc: str = "Work with PDFs.") -> Path: + """Create a minimal valid skill directory (dir name must match the name).""" + skill_dir = parent / name + (skill_dir / "scripts").mkdir(parents=True) + (skill_dir / "SKILL.md").write_text( + textwrap.dedent( + f"""\ + --- + name: {name} + description: {desc} + --- + # {name} + Follow these instructions. + """ + ), + encoding="utf-8", + ) + (skill_dir / "scripts" / "run.py").write_text("print('hi')\n", encoding="utf-8") + return skill_dir + + +# --------------------------------------------------------------------------- +# SkillStore +# --------------------------------------------------------------------------- + + +class TestSkillStore: + def test_resolve_by_path(self, tmp_path): + skill_dir = _make_skill(tmp_path) + skills = SkillStore().resolve([str(skill_dir)]) + assert [s.name for s in skills] == ["pdf-tools"] + + def test_resolve_by_name_via_dirs(self, tmp_path): + _make_skill(tmp_path) + skills = SkillStore([str(tmp_path)]).resolve(["pdf-tools"]) + assert skills[0].name == "pdf-tools" + + def test_resolve_dedupes_by_name(self, tmp_path): + skill_dir = _make_skill(tmp_path) + skills = SkillStore([str(tmp_path)]).resolve([str(skill_dir), "pdf-tools"]) + assert len(skills) == 1 + + def test_resolve_unknown_raises(self, tmp_path): + with pytest.raises(ValueError, match="not found"): + SkillStore([str(tmp_path)]).resolve(["nope"]) + + +# --------------------------------------------------------------------------- +# make_skill_toolset (script gating) +# --------------------------------------------------------------------------- + + +class TestMakeSkillToolset: + def test_scripts_disabled_by_default(self, tmp_path): + skills = SkillStore().resolve([str(_make_skill(tmp_path))]) + ts = make_skill_toolset(skills) + names = {t.name for t in ts._tools} + assert "run_skill_script" not in names + assert {"list_skills", "load_skill", "load_skill_resource"} <= names + + 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, 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 +# --------------------------------------------------------------------------- + + +class TestSkillPermissions: + def test_skill_tools_registered(self): + from agentic_cli.tools.registry import get_registry + from agentic_cli.workflow.permissions.capabilities import _CapabilityExempt + + register_skill_tool_permissions() + reg = get_registry() + for name in ("list_skills", "load_skill", "load_skill_resource"): + assert isinstance(reg.get(name).capabilities, _CapabilityExempt) + # run_skill_script is permissioned (non-exempt capability list). + run = reg.get("run_skill_script") + assert run is not None and not isinstance( + run.capabilities, _CapabilityExempt + ) + + +# --------------------------------------------------------------------------- +# ADK manager + YAML loader integration +# --------------------------------------------------------------------------- + + +class TestSkillsIntegration: + def test_manager_attaches_skill_toolset(self, tmp_path, mock_context): + from google.adk.tools.skill_toolset import SkillToolset + + from agentic_cli.workflow.adk.manager import GoogleADKWorkflowManager + + skill_dir = _make_skill(tmp_path) + cfg = AgentConfig( + name="a", prompt="p", tools=[], include_state_tools=False, + skills=[str(skill_dir)], + ) + mgr = GoogleADKWorkflowManager( + agent_configs=[cfg], settings=mock_context.settings, model="gemini-2.5-flash" + ) + tools = mgr._assemble_agent_tools(cfg, mgr._get_service_tool_map()) + skill_toolsets = [t for t in tools if isinstance(t, SkillToolset)] + assert len(skill_toolsets) == 1 + # Scripts disabled by default -> run_skill_script not exposed. + assert "run_skill_script" not in {t.name for t in skill_toolsets[0]._tools} + + def test_yaml_loader_parses_skills(self, tmp_path): + from agentic_cli.workflow.agent_loader import load_agents_from_yaml + + path = tmp_path / "agents.yaml" + path.write_text( + textwrap.dedent( + """ + agents: + - name: a + instruction: hi + skills: [pdf-tools, ./skills/sql] + """ + ), + encoding="utf-8", + ) + cfg = load_agents_from_yaml(path)[0] + assert cfg.skills == ["pdf-tools", "./skills/sql"] diff --git a/tests/workflow/test_tool_resolver.py b/tests/workflow/test_tool_resolver.py new file mode 100644 index 0000000..a247521 --- /dev/null +++ b/tests/workflow/test_tool_resolver.py @@ -0,0 +1,161 @@ +"""Tests for tool reference resolution (Phase 2). + +Covers ``resolve_tool``/``resolve_tools`` (callable / bare-name / dotted-path) +and the base-manager integration that resolves config tool refs before +service detection. +""" + +from __future__ import annotations + +import math +from typing import AsyncGenerator + +import pytest + +from agentic_cli.tools.registry import ToolCategory, ToolRegistry +from agentic_cli.tools.tool_resolver import resolve_tool, resolve_tools +from agentic_cli.workflow.base_manager import BaseWorkflowManager +from agentic_cli.workflow.config import AgentConfig +from agentic_cli.workflow.events import WorkflowEvent +from agentic_cli.workflow.permissions import EXEMPT + + +# --------------------------------------------------------------------------- +# resolve_tool / resolve_tools +# --------------------------------------------------------------------------- + + +class TestResolveTool: + def test_callable_passes_through(self): + def my_tool(): + pass + + assert resolve_tool(my_tool) is my_tool + + def test_non_string_object_passes_through(self): + # e.g. an already-built toolset object + sentinel = object() + assert resolve_tool(sentinel) is sentinel + + def test_dotted_path_imports(self): + assert resolve_tool("math.sqrt") is math.sqrt + + def test_dotted_path_bad_module_raises(self): + with pytest.raises(ValueError, match="Cannot import module"): + resolve_tool("nonexistent_module_xyz.foo") + + def test_dotted_path_bad_attr_raises(self): + with pytest.raises(ValueError, match="has no attribute"): + resolve_tool("math.nonexistent_attr_xyz") + + def test_bare_name_via_explicit_registry(self): + reg = ToolRegistry() + + def custom(): + pass + + reg.register(custom, capabilities=EXEMPT, category=ToolCategory.OTHER) + assert resolve_tool("custom", registry=reg) is custom + + def test_unknown_bare_name_raises_with_hint(self): + reg = ToolRegistry() + + def kb_search(): + pass + + reg.register(kb_search, capabilities=EXEMPT, category=ToolCategory.OTHER) + with pytest.raises(ValueError, match="Did you mean: kb_search"): + resolve_tool("kb_searh", registry=reg) # typo + + def test_empty_ref_raises(self): + with pytest.raises(ValueError, match="Empty tool reference"): + resolve_tool(" ") + + def test_bare_name_resolves_builtin_via_default_registry(self): + # Triggers lazy import of agentic_cli.tools to populate the registry. + tool = resolve_tool("read_file") + assert callable(tool) + assert getattr(tool, "__name__", None) == "read_file" + + +class TestResolveTools: + def test_none_returns_empty(self): + assert resolve_tools(None) == [] + + def test_mixed_list(self): + def cb(): + pass + + out = resolve_tools([cb, "math.sqrt"]) + assert out == [cb, math.sqrt] + + +# --------------------------------------------------------------------------- +# Base-manager integration +# --------------------------------------------------------------------------- + + +class _MiniManager(BaseWorkflowManager): + """Minimal concrete manager to exercise base-class resolution/detection.""" + + @property + def backend_type(self) -> str: + return "test" + + async def _do_initialize(self) -> None: + pass + + async def process( + self, message: str, user_id: str, session_id: str | None = None + ) -> AsyncGenerator[WorkflowEvent, None]: + if False: + yield # type: ignore[misc] + + async def reinitialize(self, model=None, preserve_sessions=True) -> None: + pass + + async def cleanup(self) -> None: + pass + + async def _extract_session_data(self, session_id: str): + return [], None + + async def _inject_session_messages( + self, session_id: str, messages, current_agent=None + ) -> None: + pass + + def _get_state_tools(self) -> list: + return [] + + +class TestBaseManagerIntegration: + def test_string_tool_resolved_and_service_detected(self, mock_context): + # "kb_search" is a registered service tool -> resolution must make it a + # callable AND detection must flag the kb_manager service. + cfg = AgentConfig( + name="a", prompt="p", tools=["kb_search"], include_state_tools=False + ) + mgr = _MiniManager(agent_configs=[cfg], settings=mock_context.settings) + + assert callable(cfg.tools[0]) + assert getattr(cfg.tools[0], "__name__", None) == "kb_search" + assert "kb_manager" in mgr.required_managers + + def test_callable_tools_unaffected(self, mock_context): + def my_tool(): + pass + + cfg = AgentConfig( + name="a", prompt="p", tools=[my_tool], include_state_tools=False + ) + mgr = _MiniManager(agent_configs=[cfg], settings=mock_context.settings) + assert cfg.tools == [my_tool] + assert mgr.required_managers == set() + + def test_dotted_path_tool_resolved(self, mock_context): + cfg = AgentConfig( + name="a", prompt="p", tools=["math.sqrt"], include_state_tools=False + ) + _MiniManager(agent_configs=[cfg], settings=mock_context.settings) + assert cfg.tools[0] is math.sqrt 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)