diff --git a/docs/HUMAN_IN_THE_LOOP.md b/docs/HUMAN_IN_THE_LOOP.md index 356814a7..a9431be1 100644 --- a/docs/HUMAN_IN_THE_LOOP.md +++ b/docs/HUMAN_IN_THE_LOOP.md @@ -205,26 +205,54 @@ To resume, `Engine.resume(run_id, approval_id, decision)` or the streamed > identity used when the approval was created (for the engine API, via > `caller_session_id` / `X-Session-ID`). Omitting or changing it returns 403. +> **Upgrade note:** Runs suspended before checkpointed model/tool steps existed +> cannot be resumed; they fail and their approval is rejected. Decide or cancel +> pending approvals before upgrading. See +> [`ADR 0005`](adr/0005-checkpointed-hitl-continuation.md). + 1. **Atomically claims** the approval (`PENDING → RESUMING`) — exactly one caller wins (see §7). 2. Moves the run through `RESUMING → RUNNING` and resumes the **same** LangGraph thread from its checkpoint via - `Command(resume={"decision": ...})`. The graph is not restarted, the agent is - not re-selected, and intent/planning are not re-run. -3. The node re-executes, the coordinator interprets the typed decision, and: + `Command(resume={"decision": ..., "approval_id": ..., "tool_call_id": ...})`. + The graph is not restarted, the agent is not re-selected, and intent/planning + are not re-run. +3. The node re-executes, restoring its checkpointed model responses and completed + tool outcomes instead of regenerating them. The gate checks that the resume's + `approval_id` and stable `tool_call_id` match the invocation it is about to + run; a missing or different binding aborts the run as `FAILED` and rejects + the approval. Otherwise the coordinator interprets the typed decision, and: - `ALLOW_ONCE` / `ALLOW_FOR_SESSION` → the original tool call executes and its result flows back to the agent (and `ALLOW_FOR_SESSION` records a session permission); - `DENY` → a structured `[denied]` result is returned and **no** tool/MCP call is made. +The claimed decision takes precedence over auto mode and session grants for that +one pending invocation. A grant added by another run, or an auto-mode change +after rebuilding the engine, cannot bypass a pending denial. A request-local +binding validates the first unfinished tool step before policy evaluation and +is consumed only by the matching approval interrupt. Later calls use the normal +policy; completed steps are restored without consuming this new decision. If a +rebuilt graph removes or blocks the pending invocation before it reaches that +interrupt, the unconsumed binding fails the run and rejects the claimed approval. + An unrecognized or missing resume decision **fails closed to `DENY`**, so an ambiguous resume never triggers a side effect. -Because a LangGraph node re-executes from its start on resume, tool *preparation* -is deterministic and safe to repeat, and all non-idempotent execution happens -**after** the interrupt returns. Duplicate execution on re-entry is prevented by -the idempotency ledger (§7). +A LangGraph node re-executes from its start on resume, but each model response +and each completed leaf-tool outcome is a checkpointed LangGraph task. Replay +returns those saved values, so the model cannot propose a different call than +the one reviewed, and earlier tools, decisions, and session grants are not +repeated. Restored steps must match the node, its input, and the invocation +identity; otherwise the run fails with `CheckpointReplayMismatch`. Runs suspended +before this format existed fail closed on resume. Duplicate execution of a call +that started but did not finish is still prevented by the idempotency ledger +(§7). + +Because the reviewed call must execute exactly as proposed, the checkpoint holds +the model's raw tool arguments and the tool results; only the `ApprovalRecord` +and the approval payload are masked. Treat the checkpoint store as sensitive. Identifiers are kept strictly separate: diff --git a/docs/RUNTIME_LIFECYCLE.md b/docs/RUNTIME_LIFECYCLE.md index ed14e152..c08caf85 100644 --- a/docs/RUNTIME_LIFECYCLE.md +++ b/docs/RUNTIME_LIFECYCLE.md @@ -77,7 +77,31 @@ During a model tool loop, the tool-call message and matching tool result remain structured and ordered for the provider. A run paused for tool approval stores its continuation in the LangGraph checkpoint keyed by `run_id`; session-wide approval grants live in a separate approval repository keyed by -session and tool identity. Normal stream completion does not cancel a suspended +session and tool identity. Model responses and completed leaf-tool outcomes are +checkpointed as individual steps, so resume restores the reviewed tool request +and reuses earlier results instead of regenerating calls or repeating side +effects. This also applies to nested agents and sequential approvals. Resumed +decisions carry the approval id and stable invocation id, which covers the run, +agent path, provider/server, tool name, and raw arguments without the provider's +transient call id. The tool approval gate rejects missing or mismatched bindings +before execution; the displayed and persisted approval arguments stay masked. +The claimed decision also overrides any auto-mode or session-grant change while +the invocation waited. Its binding is request-local and consumed only at the +matching interrupt; later calls return to normal policy evaluation. +Completing a replay with that binding still unconsumed is a mismatch: for +example, a rebuilt graph cannot silently remove the reviewed tool and report its +approval as applied. +Restored steps also validate the node/input or invocation identity, so a changed +route cannot reuse another agent's positional checkpoint. A mismatch aborts the +run, rejects its claimed approval, and removes any newly created pending approval. +The graph state carries an internal continuation version. Resuming a suspended +run created before these checkpointed substeps existed fails closed before model +or tool execution, because replaying that older format could repeat side effects. +Execution limits still account for restored calls before their cached outcomes +are returned; checkpointing does not reset the run's allowed tool budget. +Restored model steps emit no new answer deltas or token usage. The existing tool +execution ledger still protects calls whose step started but did not finish. +Normal stream completion does not cancel a suspended approval. An explicit owner cancellation atomically rejects the pending approval, transitions the run to terminal `CANCELLED`, and prevents later resume. diff --git a/docs/adr/0005-checkpointed-hitl-continuation.md b/docs/adr/0005-checkpointed-hitl-continuation.md new file mode 100644 index 00000000..785bb5d1 --- /dev/null +++ b/docs/adr/0005-checkpointed-hitl-continuation.md @@ -0,0 +1,80 @@ +# ADR 0005 — Checkpoint model and tool steps so resume runs the reviewed call + +- **Status:** Proposed +- **Date:** 2026-10-14 + +## Context + +A run suspended for tool approval resumed by re-executing its LangGraph node +from the start. The node invoked the model again, and the resume command carried +only `{"decision": ...}`. A non-deterministic model could therefore propose a +different tool, agent, or arguments than the human reviewed, and the positional +decision would still apply to it. Earlier tool calls in the same node were +protected from repeated side effects only by the idempotency ledger, and +restored turns reported their token usage a second time. + +## Decision + +The engine replays a suspended node from saved steps and binds each decision to +one invocation. + +- `run_tool_loop` runs every model response as a LangGraph `@task`, and + `ToolInvoker` runs each leaf tool's gate and execution as a `@task`. On resume + LangGraph returns the saved values, so the reviewed tool request is never + regenerated and completed outcomes, denials, and session grants are not + repeated. Execution limits are registered before a step is restored, so the + run's tool budget is unchanged. +- Restored steps validate their identity. A model step carries the node path, + its non-system input, and its turn. A tool step carries the stable + `tool_call_id`. A mismatch, for example after route filtering changes the + replay order, raises `CheckpointReplayMismatch`. +- Resume sends `{"decision", "approval_id", "tool_call_id"}`. The stable id + covers the run, agent path, provider/server, tool name, and raw arguments, but + not the provider's transient call id. The approval gate raises + `ApprovalInvocationMismatch` before execution when the binding is missing or + different. Masked arguments are never used for binding, because distinct + secrets share the same redacted display. +- Either mismatch fails the run, rejects the claimed approval, and cancels any + pending approval created during the failed replay. Orchestrators re-raise both + errors from nested children instead of turning them into tool text. +- Each resume request carries its claimed binding in a context variable. The + first unfinished tool step validates that binding before policy evaluation; + its existing `ToolInvocation.approval_id` marks the claimed decision as + authoritative over auto mode or session grants. The matching interrupt consumes + the binding once, so later calls still follow normal policy. A replay that + completes without consuming it fails closed, covering rebuilt graphs that + remove or block the pending invocation. Separate requests never share this + state, and completed steps do not consume a new decision. +- Approval cleanup is best effort. Repository failures are chained or logged + by type and cannot replace the original replay/consent mismatch. +- Graph state carries `continuation_version`. Resuming a checkpoint without the + current version fails closed before any model or tool step runs, because an + older continuation has no saved steps and replaying it could repeat side + effects. +- Outside a LangGraph runtime (`in_graph_run()` is false) the loop and invoker + run directly, without tasks. + +## Contract changes + +- The runtime dependency floor is `langgraph>=1.0,<2.0`, for functional tasks + inside `StateGraph` nodes and `langgraph.runtime.get_runtime`. +- `GraphState` adds the internal `continuation_version` key. +- The internal resume payload adds `approval_id` and `tool_call_id`. +- New errors: `ApprovalInvocationMismatch` (an `ApprovalError`) and + `CheckpointReplayMismatch` (a `RuntimeError`). + +There is no YAML or HTTP request-shape change. + +## Consequences + +- Runs that are suspended when this version is deployed cannot be resumed; they + fail and their approval is rejected. Decide or cancel pending approvals before + upgrading. +- The checkpoint stores the model's raw tool arguments and the tool results for + every step. It must be treated as a sensitive store. Only the `ApprovalRecord` + and the approval payload are masked. +- Checkpoints grow with each model turn and tool result in a suspended run. +- Restored model steps emit no new answer deltas or token usage, so a resumed run + no longer double-counts tokens. +- The idempotency ledger still protects a call whose step started but did not + finish. diff --git a/pyproject.toml b/pyproject.toml index 7cd5ea94..1507270e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,8 @@ dependencies = [ "uvicorn>=0.30", "httpx>=0.27", "jinja2>=3.1", - "langgraph>=0.2", + # HITL continuation uses functional tasks and the public runtime accessor. + "langgraph>=1.0,<2.0", "langchain>=0.3", "langchain-core>=0.3", "mcp>=1.27,<2", diff --git a/src/agent_engine/approvals/coordinator.py b/src/agent_engine/approvals/coordinator.py index 35235571..3d630613 100644 --- a/src/agent_engine/approvals/coordinator.py +++ b/src/agent_engine/approvals/coordinator.py @@ -4,7 +4,7 @@ together. It contains no tool-specific logic and performs no UI interaction (the injected provider owns that). Its single entry point is :meth:`resolve`: -1. Ask the policy whether approval is required (auto mode → session → ask). +1. Honor a claimed pending approval; otherwise ask policy (auto mode → session → ask). 2. If not required, permit execution immediately (the provider is never called). 3. Otherwise request a typed decision from the provider. 4. Apply the decision: store a session permission only for ``ALLOW_FOR_SESSION`` @@ -76,7 +76,9 @@ def __init__( async def resolve(self, invocation: ToolInvocation, *, auto_mode: bool) -> ApprovalOutcome: query = await self._build_query(invocation, auto_mode=auto_mode) - required = self._policy.requires_approval(query) + # A claimed decision governs this pending invocation even if auto mode + # or a session grant changed while the human was reviewing it. + required = invocation.approval_id is not None or self._policy.requires_approval(query) self._log_evaluated(query, required=required) if not required: diff --git a/src/agent_engine/approvals/errors.py b/src/agent_engine/approvals/errors.py index 8879a157..d2f14fea 100644 --- a/src/agent_engine/approvals/errors.py +++ b/src/agent_engine/approvals/errors.py @@ -46,6 +46,14 @@ def __init__(self, approval_id: str, status: str) -> None: super().__init__(f"approval {approval_id} already processed (status={status})") +class ApprovalInvocationMismatch(ApprovalError): + """A resumed decision does not bind to the invocation at the tool gate.""" + + def __init__(self, approval_id: str) -> None: + self.approval_id = approval_id + super().__init__(f"approval {approval_id} does not match the pending invocation") + + class UnauthorizedApprover(ApprovalError): def __init__(self, approval_id: str) -> None: self.approval_id = approval_id diff --git a/src/agent_engine/approvals/invocation.py b/src/agent_engine/approvals/invocation.py index 6a87e609..150162c4 100644 --- a/src/agent_engine/approvals/invocation.py +++ b/src/agent_engine/approvals/invocation.py @@ -85,6 +85,7 @@ class ToolInvocation: organization_id: str = "" run_id: str | None = None approval_id: str | None = None + """A claimed approval for this invocation; its decision overrides broad grants.""" human_description: str = "" @property diff --git a/src/agent_engine/engine/langgraph/approval_provider.py b/src/agent_engine/engine/langgraph/approval_provider.py index a788c448..623e2f54 100644 --- a/src/agent_engine/engine/langgraph/approval_provider.py +++ b/src/agent_engine/engine/langgraph/approval_provider.py @@ -17,6 +17,9 @@ from agent_engine.approvals.approval_manager import ApprovalManager from agent_engine.approvals.approval_provider import ApprovalProvider, ApprovalRequest from agent_engine.approvals.decision import ApprovalDecision, parse_decision +from agent_engine.approvals.errors import ApprovalInvocationMismatch +from agent_engine.approvals.models import ApprovalStatus +from agent_engine.engine.langgraph.execution.approval_resume import current_approval_resume from agent_engine.logging_config import log from agent_engine.runtime.hooks import current_run_context @@ -96,4 +99,29 @@ async def request_decision(self, request: ApprovalRequest) -> ApprovalDecision: tool_call_id=inv.tool_call_id, ) resume = interrupt(payload) - return parse_decision(resume, default=ApprovalDecision.DENY) + # The stable id covers the run, agent path, provider/server, tool, and + # raw arguments. Masked arguments cannot bind consent: distinct secrets + # have the same redacted display. Provider-generated call ids are absent. + if ( + not isinstance(resume, dict) + or resume.get("approval_id") != record.approval_id + or resume.get("tool_call_id") != inv.tool_call_id + ): + # A changed invocation may have created a new pending row before the + # positional resume was consumed. It must not outlive this failure. + if record.status == ApprovalStatus.PENDING: + try: + await self._approvals.cancel_pending( + run_id=run_id, + approval_id=record.approval_id, + caller_user_id=authorized_user_id, + caller_auth_ref=auth_ref, + ) + except Exception as exc: + raise ApprovalInvocationMismatch(record.approval_id) from exc + raise ApprovalInvocationMismatch(record.approval_id) + decision = parse_decision(resume, default=ApprovalDecision.DENY) + binding = current_approval_resume.get() + if binding is not None: + binding.consume(record.approval_id, inv.tool_call_id) + return decision diff --git a/src/agent_engine/engine/langgraph/checkpointing.py b/src/agent_engine/engine/langgraph/checkpointing.py index 3a18404b..27d8bc68 100644 --- a/src/agent_engine/engine/langgraph/checkpointing.py +++ b/src/agent_engine/engine/langgraph/checkpointing.py @@ -13,6 +13,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.memory import InMemorySaver +from langgraph.runtime import get_runtime from agent_engine.logging_config import log @@ -23,6 +24,21 @@ class CheckpointConfigError(RuntimeError): """A checkpoint connection string was supplied but its backend is unusable.""" +class CheckpointReplayMismatch(RuntimeError): + """A positional checkpoint belongs to a different model/tool invocation.""" + + def __init__(self, node_path: str) -> None: + super().__init__(f"checkpoint replay does not match the invocation at {node_path}") + + +def in_graph_run() -> bool: + """Distinguish a LangGraph runtime from an ordinary LangChain Runnable.""" + try: + return get_runtime() is not None + except RuntimeError: + return False + + @dataclass(frozen=True) class CheckpointerHandle: """A checkpointer plus a flag describing whether it is durable/shared. diff --git a/src/agent_engine/engine/langgraph/engine.py b/src/agent_engine/engine/langgraph/engine.py index 32315f0a..015492a0 100644 --- a/src/agent_engine/engine/langgraph/engine.py +++ b/src/agent_engine/engine/langgraph/engine.py @@ -16,7 +16,7 @@ from agent_engine.approvals.approval_manager import ApprovalManager from agent_engine.approvals.coordinator import ApprovalCoordinator from agent_engine.approvals.decision import ApprovalDecision, parse_decision -from agent_engine.approvals.errors import RunNotFound +from agent_engine.approvals.errors import ApprovalInvocationMismatch, RunNotFound from agent_engine.approvals.in_memory_approval_repository import InMemoryApprovalRepository from agent_engine.approvals.in_memory_session_approval_repository import ( InMemorySessionApprovalRepository, @@ -24,7 +24,7 @@ from agent_engine.approvals.in_memory_tool_execution_repository import ( InMemoryToolExecutionRepository, ) -from agent_engine.approvals.models import ApprovalRecord, RunStatus +from agent_engine.approvals.models import ApprovalRecord, ApprovalStatus, RunStatus from agent_engine.approvals.session_approval_repository import SessionApprovalRepository from agent_engine.approvals.session_approval_store import SessionApprovalStore from agent_engine.approvals.tool_execution_manager import ToolExecutionManager @@ -35,6 +35,11 @@ from agent_engine.engine.langgraph.checkpointing import ( CheckpointerHandle, CheckpointProviderFactory, + CheckpointReplayMismatch, +) +from agent_engine.engine.langgraph.execution.approval_resume import ( + ApprovalResume, + current_approval_resume, ) from agent_engine.engine.langgraph.execution.run_lifecycle import RunLifecycle from agent_engine.engine.langgraph.execution.stream_channel import StreamChannel @@ -83,6 +88,7 @@ from agent_engine.tool_usage.tracker import ToolUsageTracker logger = logging.getLogger(__name__) +_CONTINUATION_VERSION = 1 def _initial_state( @@ -98,6 +104,7 @@ def _initial_state( of their own, so a plain ``run("hi")`` leaves the key absent for filters. """ state: dict[str, Any] = { + "continuation_version": _CONTINUATION_VERSION, "message": message, "history": [ {"role": history_message.role.value, "content": history_message.content} @@ -374,7 +381,13 @@ def _require_built(self, action: str) -> tuple[RunGraph, RunLifecycle]: return self._app, self._lifecycle @contextmanager - def _run_scope(self, ctx: RunContext, *, sinks: StreamSinks | None) -> Iterator[None]: + def _run_scope( + self, + ctx: RunContext, + *, + sinks: StreamSinks | None, + approval_resume: ApprovalResume | None = None, + ) -> Iterator[None]: """Publish the ambient per-run state for the duration of one run. Nodes, tools, and the MCP transport read the run context, the execution @@ -388,6 +401,7 @@ def _run_scope(self, ctx: RunContext, *, sinks: StreamSinks | None) -> Iterator[ that must not disturb the sinks its caller already installed. """ ctx_token = current_run_context.set(ctx) + approval_token = current_approval_resume.set(approval_resume) exec_token = current_execution.set(ExecutionLimiter(self._policy)) stream_token = None if sinks is None else current_streams.set(sinks) try: @@ -403,6 +417,8 @@ def _run_scope(self, ctx: RunContext, *, sinks: StreamSinks | None) -> Iterator[ current_execution.reset(exec_token) with suppress(ValueError): current_run_context.reset(ctx_token) + with suppress(ValueError): + current_approval_resume.reset(approval_token) async def _invoke_graph( self, app: RunGraph, ctx: RunContext, graph_input: Any @@ -412,7 +428,43 @@ async def _invoke_graph( ``graph_input`` is the initial state of a new run, or the ``Command`` that continues a suspended one from its checkpoint. """ - return await app.ainvoke(graph_input, self._thread_config(ctx)) + try: + if isinstance(graph_input, Command): + state = await app.aget_state(self._thread_config(ctx)) + if state.values.get("continuation_version") != _CONTINUATION_VERSION: + # Legacy runs may have already executed auto-approved tools, + # but have no saved substeps. Never replay them to discover + # the mismatch after another side effect has occurred. + raise CheckpointReplayMismatch(f"run {ctx.run_id}") + result = await app.ainvoke(graph_input, self._thread_config(ctx)) + binding = current_approval_resume.get() + if binding is not None: + # A rebuilt graph can remove or block the reviewed tool before + # it reaches the approval gate. Never report that decision as + # applied unless the exact pending invocation consumed it. + binding.ensure_consumed() + return result + except (ApprovalInvocationMismatch, CheckpointReplayMismatch): + # Both invocation paths share this gate. A rejected replay must not + # strand the claimed approval in RESUMING or finalize it as allowed. + resume = graph_input.resume if isinstance(graph_input, Command) else None + approval_id = resume.get("approval_id") if isinstance(resume, dict) else None + if ctx.run_id is not None and isinstance(approval_id, str): + try: + approval = await self._approval_manager.get_approval(ctx.run_id, approval_id) + if approval.status == ApprovalStatus.RESUMING: + await self._approval_manager.finalize(approval_id, approved=False) + except Exception as cleanup_exc: + # Keep the replay failure as the reported cause. + log( + logger, + logging.WARNING, + "claimed approval cleanup failed", + run_id=ctx.run_id, + approval_id=approval_id, + error=type(cleanup_exc).__name__, + ) + raise async def _completed_result( self, @@ -611,7 +663,7 @@ async def resume( """ app, lifecycle = self._require_built("resuming") kind = parse_decision(decision) - ctx = await self._activate_approval_resume( + ctx, approval = await self._activate_approval_resume( lifecycle, run_id=run_id, approval_id=approval_id, @@ -629,9 +681,19 @@ async def resume( decision=kind.value, ) tokens = TokenUsage() - with self._run_scope(ctx, sinks=StreamSinks(token=tokens.add)): + with self._run_scope( + ctx, + sinks=StreamSinks(token=tokens.add), + approval_resume=ApprovalResume(approval.approval_id, approval.tool_call_id), + ): try: - resume_command: Command[Any] = Command(resume={"decision": kind.value}) + resume_command: Command[Any] = Command( + resume={ + "decision": kind.value, + "approval_id": approval.approval_id, + "tool_call_id": approval.tool_call_id, + } + ) result = await self._invoke_graph(app, ctx, resume_command) token_usage = await self._record_token_usage(ctx, tokens) await self._approval_manager.finalize(approval_id, approved=approved) @@ -671,7 +733,7 @@ async def resume_stream( """Resume one approval through the same owned stream used by new runs.""" app, lifecycle = self._require_built("streaming an approval resume") kind = parse_decision(decision) - ctx = await self._activate_approval_resume( + ctx, approval = await self._activate_approval_resume( lifecycle, run_id=run_id, approval_id=approval_id, @@ -692,7 +754,13 @@ async def finalize_approval() -> None: approval_id=approval_id, decision=kind.value, ) - command: Command[Any] = Command(resume={"decision": kind.value}) + command: Command[Any] = Command( + resume={ + "decision": kind.value, + "approval_id": approval.approval_id, + "tool_call_id": approval.tool_call_id, + } + ) execution = cast( AsyncGenerator[RunStreamEvent, None], self._stream_execution( @@ -702,6 +770,7 @@ async def finalize_approval() -> None: command, after_invoke=finalize_approval, started_event=RunStreamEvent(type="resume_started", run_id=run_id), + approval_resume=ApprovalResume(approval.approval_id, approval.tool_call_id), ), ) try: @@ -719,7 +788,7 @@ async def _activate_approval_resume( caller_user_id: str | None, caller_session_id: str | None, access_token: str | None = None, - ) -> RunContext: + ) -> tuple[RunContext, ApprovalRecord]: """Claim and activate one resume without an interruptible ownership gap. The response may be cancelled before its first SSE frame. Shielding this @@ -728,7 +797,7 @@ async def _activate_approval_resume( approval stranded before graph-task ownership begins. """ - async def activate() -> RunContext: + async def activate() -> tuple[RunContext, ApprovalRecord]: approval = await self._approval_manager.claim( run_id=run_id, approval_id=approval_id, @@ -750,14 +819,14 @@ async def activate() -> RunContext: ), ) await lifecycle.activate_resume(ctx) - return ctx + return ctx, approval activation = asyncio.create_task(activate()) try: return await asyncio.shield(activation) except asyncio.CancelledError: try: - ctx = await asyncio.shield(activation) + ctx, _ = await asyncio.shield(activation) except Exception: # The claim lost or activation failed, so there is no active # execution owned by this request to terminally cancel. @@ -827,11 +896,14 @@ async def _stream_execution( *, after_invoke: Callable[[], Awaitable[None]] | None = None, started_event: RunStreamEvent | None = None, + approval_resume: ApprovalResume | None = None, ) -> AsyncIterator[RunStreamEvent]: """Own one active graph leg and cancel it when its consumer leaves.""" channel = StreamChannel() tokens = TokenUsage() - with self._run_scope(ctx, sinks=channel.sinks(token=tokens.add)): + with self._run_scope( + ctx, sinks=channel.sinks(token=tokens.add), approval_resume=approval_resume + ): task = asyncio.create_task( self._stream_graph( app, diff --git a/src/agent_engine/engine/langgraph/execution/approval_resume.py b/src/agent_engine/engine/langgraph/execution/approval_resume.py new file mode 100644 index 00000000..851041b2 --- /dev/null +++ b/src/agent_engine/engine/langgraph/execution/approval_resume.py @@ -0,0 +1,44 @@ +"""A claimed human decision carried only for one resume request.""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass + +from agent_engine.approvals.errors import ApprovalInvocationMismatch + + +@dataclass +class ApprovalResume: + """Bind the first unfinished tool step, then release later calls to policy. + + Completed steps are restored without entering their gate. This request-local + state is shared by nested graph tasks, never stored on the engine or graph. + """ + + approval_id: str + tool_call_id: str + consumed: bool = False + + def for_invocation(self, tool_call_id: str) -> str | None: + if self.consumed: + return None + if tool_call_id != self.tool_call_id: + raise ApprovalInvocationMismatch(self.approval_id) + return self.approval_id + + def consume(self, approval_id: str, tool_call_id: str) -> None: + if not self.consumed: + if approval_id != self.approval_id or tool_call_id != self.tool_call_id: + raise ApprovalInvocationMismatch(self.approval_id) + self.consumed = True + + def ensure_consumed(self) -> None: + """Fail when replay completed without reaching the reviewed invocation.""" + if not self.consumed: + raise ApprovalInvocationMismatch(self.approval_id) + + +current_approval_resume: ContextVar[ApprovalResume | None] = ContextVar( + "current_approval_resume", default=None +) diff --git a/src/agent_engine/engine/langgraph/execution/model_loop.py b/src/agent_engine/engine/langgraph/execution/model_loop.py index 2a6dfa00..7b5b35bb 100644 --- a/src/agent_engine/engine/langgraph/execution/model_loop.py +++ b/src/agent_engine/engine/langgraph/execution/model_loop.py @@ -6,8 +6,10 @@ from collections.abc import Awaitable, Callable from typing import Any -from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.messages import AIMessage, SystemMessage, ToolMessage +from langgraph.func import task +from agent_engine.engine.langgraph.checkpointing import CheckpointReplayMismatch, in_graph_run from agent_engine.engine.langgraph.execution.execution_context import ( ExecutionContextRefresher, ) @@ -19,6 +21,7 @@ ) from agent_engine.runtime.streaming import current_streams from agent_engine.runtime.tool_results import NormalizedToolResult +from agent_engine.tool_usage.models import stable_tool_call_id logger = logging.getLogger(__name__) @@ -55,10 +58,44 @@ async def run_tool_loop( *, refresh_execution_context: ExecutionContextRefresher | None = None, ) -> Any: - """Drive model → tools → model until the model stops calling tools.""" + """Drive model → tools → model, restoring completed steps on HITL replay. + + Model outputs determine the interrupt order and must never be regenerated. + ToolInvoker checkpoints leaf outcomes after registering execution limits. + Orchestrator delegations rebuild their route locally and rely on the child's + checkpointed model/tool steps instead. Restored outputs must match this node + and its original input, even if route filtering changes the replay order. + """ + checkpointed = in_graph_run() + step_key = stable_tool_call_id( + node_path, + [ + {"type": message.type, "content": message.content} + for message in context.messages + if not isinstance(message, SystemMessage) + ], + ) + + @task + async def model_step(turn: int) -> tuple[str, int, Any]: + await _refresh(context, refresh_execution_context) + return step_key, turn, await invoke_model(model, context.messages) + + model_turn = 0 + + async def next_response() -> Any: + nonlocal model_turn + if checkpointed: + restored_key, restored_turn, response = await model_step(model_turn) + if (restored_key, restored_turn) != (step_key, model_turn): + raise CheckpointReplayMismatch(node_path) + model_turn += 1 + return response + await _refresh(context, refresh_execution_context) + return await invoke_model(model, context.messages) + limiter = current_execution.get() - await _refresh(context, refresh_execution_context) - response = await invoke_model(model, context.messages) + response = await next_response() while getattr(response, "tool_calls", None): if limiter is not None: try: @@ -95,8 +132,7 @@ async def run_tool_loop( name=tool_call["name"], ) ) - await _refresh(context, refresh_execution_context) - response = await invoke_model(model, context.messages) + response = await next_response() return response diff --git a/src/agent_engine/engine/langgraph/nodes/orchestrator_node.py b/src/agent_engine/engine/langgraph/nodes/orchestrator_node.py index 62db39cf..f06fe1d0 100644 --- a/src/agent_engine/engine/langgraph/nodes/orchestrator_node.py +++ b/src/agent_engine/engine/langgraph/nodes/orchestrator_node.py @@ -13,7 +13,9 @@ from langgraph.errors import GraphInterrupt from pydantic import BaseModel +from agent_engine.approvals.errors import ApprovalInvocationMismatch from agent_engine.core.spec import OrchestratorSpec +from agent_engine.engine.langgraph.checkpointing import CheckpointReplayMismatch from agent_engine.engine.langgraph.execution.execution_context import ( current_usage_scope, execution_context_refresher, @@ -178,11 +180,11 @@ async def invoke(message: str) -> str: sink_token = current_streams.set(child_sinks) try: result = await entry.node.execute(sub_state) - except GraphInterrupt: + except (GraphInterrupt, ApprovalInvocationMismatch, CheckpointReplayMismatch): # An approval interrupt raised inside a nested child must bubble up # to the LangGraph runtime so the checkpoint is taken — it is - # control flow, not a child failure. Never swallow it, and record - # nothing: the invocation has not reached an outcome yet. + # control flow, not a child failure. Replay/consent mismatches + # also abort the run instead of becoming model-facing tool text. raise except Exception as exc: await self._usage.record_failure(identity, error=str(exc)) diff --git a/src/agent_engine/engine/langgraph/tools/tool_invoker.py b/src/agent_engine/engine/langgraph/tools/tool_invoker.py index ea9fc84a..450a11ef 100644 --- a/src/agent_engine/engine/langgraph/tools/tool_invoker.py +++ b/src/agent_engine/engine/langgraph/tools/tool_invoker.py @@ -19,6 +19,7 @@ from typing import Any from langchain_core.tools import BaseTool +from langgraph.func import task from agent_engine.approvals.coordinator import ApprovalCoordinator from agent_engine.approvals.invocation import ToolInvocation @@ -29,6 +30,8 @@ execution_id_for, ) from agent_engine.core.spec import AgentSpec +from agent_engine.engine.langgraph.checkpointing import CheckpointReplayMismatch, in_graph_run +from agent_engine.engine.langgraph.execution.approval_resume import current_approval_resume from agent_engine.engine.langgraph.tools.agent_tool_binding import AgentToolBinding from agent_engine.engine.langgraph.tools.tool_gate import DenyTool, ExecuteTool, ToolGate from agent_engine.engine.langgraph.tools.tool_result_normalizer import ( @@ -157,7 +160,25 @@ async def invoke(self, tc: dict[str, Any]) -> NormalizedToolResult: if blocked is not None: return NormalizedToolResult.text_only(blocked) - gate = await self._gate_tool_call(call) + if in_graph_run(): + # Count every logical call on replay, but restore completed outcomes + # without repeating approval decisions, grants, hooks, or providers. + @task + async def tool_step() -> tuple[str, dict[str, object]]: + result = await self._gate_and_execute(call) + return call.tool_call_id, result.to_persisted() + + restored_id, result = await tool_step() + if restored_id != call.tool_call_id: + raise CheckpointReplayMismatch(self._node_path) + return NormalizedToolResult.from_persisted(result) + + return await self._gate_and_execute(call) + + async def _gate_and_execute(self, call: _ToolCall) -> NormalizedToolResult: + binding = current_approval_resume.get() + approval_id = binding.for_invocation(call.tool_call_id) if binding is not None else None + gate = await self._gate_tool_call(call, approval_id=approval_id) if isinstance(gate, DenyTool): await self._usage.record_denied(call.identity) return NormalizedToolResult.text_only(gate.message) @@ -419,7 +440,7 @@ def _log_call(self, level: int, event: str, call: _ToolCall, **fields: Any) -> N **fields, ) - async def _gate_tool_call(self, call: _ToolCall) -> ToolGate: + async def _gate_tool_call(self, call: _ToolCall, *, approval_id: str | None = None) -> ToolGate: """Run the call through the approval coordinator. Returns :class:`ExecuteTool` when the tool may run (auto mode, an existing @@ -446,7 +467,6 @@ async def _gate_tool_call(self, call: _ToolCall) -> ToolGate: if run_context else None ) - approval_id_value = run_context.metadata.get("approval_id") if run_context else None invocation = ToolInvocation( tool_call_id=call.tool_call_id, agent_id=self._spec.id, @@ -459,7 +479,7 @@ async def _gate_tool_call(self, call: _ToolCall) -> ToolGate: user_id=user_id or "", organization_id=organization_id or "", run_id=run_context.run_id if run_context else None, - approval_id=(approval_id_value if isinstance(approval_id_value, str) else None), + approval_id=approval_id, human_description=call.description, ) outcome = await self._approval_coordinator.resolve( diff --git a/src/agent_engine/runtime/state.py b/src/agent_engine/runtime/state.py index 734b39de..5e65b9d8 100644 --- a/src/agent_engine/runtime/state.py +++ b/src/agent_engine/runtime/state.py @@ -19,6 +19,9 @@ class GraphState(TypedDict, total=False): ``ToolUsageRepository``, keyed by run, not graph-state propagation. """ + continuation_version: int + """Internal replay format; older suspended runs must fail closed on resume.""" + message: str """The user's incoming message.""" diff --git a/tests/approvals/test_coordinator.py b/tests/approvals/test_coordinator.py index 10ab7d1e..02e2992c 100644 --- a/tests/approvals/test_coordinator.py +++ b/tests/approvals/test_coordinator.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +from dataclasses import replace import pytest @@ -86,6 +87,27 @@ def _inv( ) +@pytest.mark.parametrize("auto_mode", [False, True]) +async def test_claimed_approval_overrides_session_and_auto_bypasses(auto_mode: bool) -> None: + provider = RecordingProvider(ApprovalDecision.DENY) + store = InMemorySessionApprovalStore() + invocation = _inv() + assert invocation.session_key is not None + await store.allow(invocation.session_key) + coordinator = ApprovalCoordinator(provider, session_store=store) + + result = await coordinator.resolve( + replace(invocation, approval_id="claimed"), auto_mode=auto_mode + ) + assert result.execute is False + assert result.decision == ApprovalDecision.DENY + assert len(provider.requests) == 1 + + fresh = await coordinator.resolve(_inv(tool_call_id="fresh"), auto_mode=auto_mode) + assert fresh.execute is True + assert len(provider.requests) == 1 + + # ------------------------------- standard --------------------------------- # diff --git a/tests/approvals/test_engine_hitl.py b/tests/approvals/test_engine_hitl.py index dedb9661..e63d963d 100644 --- a/tests/approvals/test_engine_hitl.py +++ b/tests/approvals/test_engine_hitl.py @@ -336,7 +336,8 @@ async def test_allow_once_resumes_same_run_and_executes_once(tmp_path: Path) -> assert "sent: go" in resumed.answer assert _executions(counter) == 1 # executed exactly once assert resumed.visited == ["writer"] # same run, agent not re-selected as a new route - assert (resumed.input_tokens, resumed.output_tokens) == (6, 3) + # The reviewed model response is restored, so only the post-tool turn adds usage. + assert (resumed.input_tokens, resumed.output_tokens) == (4, 2) assert recovered == resumed diff --git a/tests/approvals/test_resume_binding.py b/tests/approvals/test_resume_binding.py new file mode 100644 index 00000000..88c11053 --- /dev/null +++ b/tests/approvals/test_resume_binding.py @@ -0,0 +1,736 @@ +"""Approval replay must preserve the reviewed action and completed side effects.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import replace +from pathlib import Path +from typing import Any, cast + +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.messages.tool import ToolCall as LCToolCall + +from agent_engine.approvals.approval_manager import ApprovalManager +from agent_engine.approvals.approval_provider import ApprovalRequest +from agent_engine.approvals.errors import ( + ApprovalAlreadyProcessed, + ApprovalError, + ApprovalInvocationMismatch, +) +from agent_engine.approvals.in_memory_approval_repository import InMemoryApprovalRepository +from agent_engine.approvals.invocation import ToolInvocation +from agent_engine.approvals.models import ApprovalStatus +from agent_engine.core.execution import ExecutionPolicy +from agent_engine.core.spec import AgentSpec +from agent_engine.engine.langgraph.approval_provider import InterruptApprovalProvider +from agent_engine.engine.langgraph.engine import LangGraphEngine +from agent_engine.runs.in_memory import InMemoryRunRepository +from agent_engine.runtime.hooks import RunContext, current_run_context +from agent_engine.tool_usage.in_memory import InMemoryToolUsageRepository +from agent_engine.tool_usage.models import ToolInvocationKind, stable_tool_call_id +from tests.approvals.test_engine_hitl import FakeChatModel, _multi_tool_spec +from tests.engine.test_engine_flow import orchestrator, system +from tests.engine.test_execution_limits import ToolThenAnswerModel + + +class DriftingModel(FakeChatModel): + """Change a tool request if the model is asked to regenerate it on resume.""" + + def __init__(self, before: str, after: str, *, change_arguments: bool = False) -> None: + super().__init__() + self.before = before + self.after = after + self.change_arguments = change_arguments + self.proposals = 0 + + def bind_tools(self, tools: list[Any]) -> DriftingModel: + return self + + def _respond(self, messages: list[Any]) -> AIMessage: + if any(isinstance(message, ToolMessage) for message in messages): + return AIMessage(content="done") + self.proposals += 1 + return AIMessage( + content="", + tool_calls=[ + LCToolCall( + name=self.before if self.proposals == 1 else self.after, + args={ + "message": "changed" + if self.change_arguments and self.proposals > 1 + else "reviewed" + }, + id=f"transient-{self.proposals}", + ) + ], + ) + + +class SequentialDriftingModel(DriftingModel): + def _respond(self, messages: list[Any]) -> AIMessage: + completed = sum(isinstance(message, ToolMessage) for message in messages) + if completed == 2: + return AIMessage(content="done") + self.proposals += 1 + name = self.before if completed == 0 else self.after + return AIMessage( + content="", + tool_calls=[ + LCToolCall( + name=name, + args={"message": f"reviewed-{self.proposals}"}, + id=f"transient-{self.proposals}", + ) + ], + ) + + +class SiblingModel(DriftingModel): + def _respond(self, messages: list[Any]) -> AIMessage: + if any(isinstance(message, ToolMessage) for message in messages): + return AIMessage(content="done") + self.proposals += 1 + return AIMessage( + content="", + tool_calls=[ + LCToolCall( + name="create_ticket", + args={"message": f"ticket-{self.proposals}"}, + id=f"ticket-{self.proposals}", + ), + LCToolCall( + name="send_email", args={"message": "reviewed"}, id=f"email-{self.proposals}" + ), + ], + ) + + +def write_tool(base_dir: Path, name: str, output: Path) -> None: + path = base_dir / "plugins" / "tools" / f"{name}.py" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + f"def {name}(message: str) -> str:\n" + f" with open({str(output)!r}, 'a') as output:\n" + f" output.write({name!r} + ':' + message + '\\n')\n" + " return message\n", + encoding="utf-8", + ) + + +async def pending_from_stream(events: AsyncIterator[Any]) -> Any: + collected = [event async for event in events] + assert collected[-1].type == "pending_approval" + return collected[-1] + + +@pytest.mark.parametrize("initial_stream", [False, True]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +@pytest.mark.parametrize("change", ["arguments", "tool", "agent"]) +async def test_resume_executes_only_the_reviewed_invocation( + tmp_path: Path, initial_stream: bool, resumed_stream: bool, change: str +) -> None: + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + write_tool(tmp_path, "archive_email", output) + model = DriftingModel( + "send_email", + "archive_email" if change == "tool" else "send_email", + change_arguments=change == "arguments", + ) + spec = _multi_tool_spec("send_email", "archive_email") + root = DriftingModel("writer", "other") + if change == "agent": + writer = spec.graph + other = replace(writer, node=replace(writer.node, id="other", name="other")) + spec = system(orchestrator("root", [writer, other])) + + builds = 0 + + def factory(*_: Any, **__: Any) -> BaseChatModel: + # GraphBuilder builds the root before its children. + nonlocal builds + builds += 1 + if change == "agent" and builds == 1: + return cast(BaseChatModel, root) + return cast(BaseChatModel, model) + + approvals = InMemoryApprovalRepository() + usage = InMemoryToolUsageRepository() + manager = ApprovalManager(run_repository=InMemoryRunRepository(), approval_repository=approvals) + async with LangGraphEngine( + tmp_path, model_factory=factory, approval_manager=manager, tool_usage_repository=usage + ) as engine: + await engine.build(spec) + if initial_stream: + pending = await pending_from_stream( + engine.stream("review", context=RunContext(run_id="r")) + ) + else: + first = await engine.run("review", context=RunContext(run_id="r")) + pending = first.pending_approval + assert pending is not None + assert pending.tool_name == "send_email" + assert pending.agent_id == "writer" + assert pending.arguments == {"message": "reviewed"} + assert not output.exists() + + if resumed_stream: + events = [ + event + async for event in engine.resume_stream("r", pending.approval_id, "allow once") + ] + assert events[-1].type == "final" + else: + result = await engine.resume("r", pending.approval_id, "allow once") + assert result.status == "completed" + stored = [ + record + for record in await usage.list_for_run("r") + if record.call.kind == ToolInvocationKind.TOOL + ] + approval = await approvals.get(pending.approval_id) + + assert output.read_text().splitlines() == ["send_email:reviewed"] + assert approval is not None + assert len(stored) == 1 + assert stored[0].call.tool_call_id == approval.tool_call_id + assert stored[0].call.agent_id == approval.agent_id == "writer" + assert model.proposals == 1 + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("decision", ["allow once", "allow for session", "deny"]) +async def test_sequential_approvals_do_not_replay_or_rebind_earlier_steps( + tmp_path: Path, nested: bool, decision: str +) -> None: + output = tmp_path / "executed.log" + for name in ("send_email", "archive_email"): + write_tool(tmp_path, name, output) + model = SequentialDriftingModel("send_email", "archive_email") + spec = _multi_tool_spec("send_email", "archive_email") + if nested: + spec = system(orchestrator("root", [orchestrator("inner", [spec.graph])])) + + builds = 0 + + def factory(*_: Any, **__: Any) -> BaseChatModel: + nonlocal builds + builds += 1 + return cast(BaseChatModel, model if builds == (3 if nested else 1) else FakeChatModel()) + + async with LangGraphEngine(tmp_path, model_factory=factory) as engine: + await engine.build(spec) + context = RunContext(run_id="sequential", conversation_id="session", user_id="owner") + first = await pending_from_stream(engine.stream("review", context=context)) + second = await pending_from_stream( + engine.resume_stream( + "sequential", + first.approval_id, + decision, + caller_user_id="owner", + caller_session_id="session", + ) + ) + assert second.tool_name == "archive_email" + assert second.arguments == {"message": "reviewed-2"} + events = [ + event + async for event in engine.resume_stream( + "sequential", + second.approval_id, + "allow once", + caller_user_id="owner", + caller_session_id="session", + ) + ] + assert events[-1].type == "final" + + expected = [] if decision == "deny" else ["send_email:reviewed-1"] + assert output.read_text().splitlines() == [*expected, "archive_email:reviewed-2"] + assert model.proposals == 2 + + +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_completed_session_allowed_sibling_is_not_reexecuted( + tmp_path: Path, resumed_stream: bool +) -> None: + output = tmp_path / "executed.log" + for name in ("create_ticket", "send_email"): + write_tool(tmp_path, name, output) + model = SiblingModel("create_ticket", "send_email") + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **k: cast(BaseChatModel, model) + ) as engine: + await engine.build(_multi_tool_spec("create_ticket", "send_email")) + first = await engine.run( + "review", context=RunContext(run_id="siblings", conversation_id="s", user_id="u") + ) + assert first.pending_approval is not None + second = await engine.resume( + "siblings", + first.pending_approval.approval_id, + "allow for session", + caller_user_id="u", + caller_session_id="s", + ) + assert second.pending_approval is not None + assert second.pending_approval.tool_name == "send_email" + assert output.read_text().splitlines() == ["create_ticket:ticket-1"] + if resumed_stream: + events = [ + event + async for event in engine.resume_stream( + "siblings", + second.pending_approval.approval_id, + "allow once", + caller_user_id="u", + caller_session_id="s", + ) + ] + assert events[-1].type == "final" + else: + result = await engine.resume( + "siblings", + second.pending_approval.approval_id, + "allow once", + caller_user_id="u", + caller_session_id="s", + ) + assert result.status == "completed" + assert output.read_text().splitlines() == ["create_ticket:ticket-1", "send_email:reviewed"] + + +@pytest.mark.parametrize( + "change", ["arguments", "secret", "agent", "provider", "server", "tool", "unbound"] +) +async def test_execution_gate_rejects_a_decision_for_a_different_invocation( + monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + provider = InterruptApprovalProvider(manager) + original = ToolInvocation( + tool_call_id="unused", + agent_id="writer", + tool_name="send_email", + arguments={"message": "reviewed", "token": "first-secret"}, + ) + changes: dict[str, dict[str, Any]] = { + "arguments": {"arguments": {"message": "changed", "token": "first-secret"}}, + "secret": {"arguments": {"message": "reviewed", "token": "changed-secret"}}, + "agent": {"agent_id": "other"}, + "provider": {"provider": "mcp"}, + "server": {"server_id": "other"}, + "tool": {"tool_name": "archive_email"}, + "unbound": {}, + } + + def identified(inv: ToolInvocation) -> ToolInvocation: + return replace( + inv, + tool_call_id=stable_tool_call_id( + "r", inv.agent_id, inv.provider, inv.server_id, inv.tool_name, inv.arguments + ), + ) + + original = identified(original) + changed = identified(replace(original, **changes[change])) + await manager.create_pending( + run_id="r", + thread_id="r", + approval_id="approved", + agent_id="writer", + tool_name="send_email", + tool_call_id=original.tool_call_id, + provider="local", + description="review", + arguments={"message": "reviewed", "token": "***"}, + ) + await manager.claim(run_id="r", approval_id="approved") + payload = ( + {"decision": "allow_once"} + if change == "unbound" + else { + "decision": "allow_once", + "approval_id": "approved", + "tool_call_id": original.tool_call_id, + } + ) + monkeypatch.setattr( + "agent_engine.engine.langgraph.approval_provider.interrupt", lambda _: payload + ) + token = current_run_context.set(RunContext(run_id="r")) + try: + with pytest.raises(ApprovalError, match="invocation"): + await provider.request_decision( + ApprovalRequest(changed, "review", {"message": "reviewed", "token": "***"}) + ) + assert await manager.get_pending("r") is None + finally: + current_run_context.reset(token) + + +@pytest.mark.parametrize("limit", ["max_tool_calls", "max_tool_calls_per_agent"]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_restored_tool_calls_still_count_toward_run_limits( + tmp_path: Path, limit: str, resumed_stream: bool +) -> None: + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + model = ToolThenAnswerModel(3, lambda i: {"message": f"message-{i}"}) + policy = ( + ExecutionPolicy(max_tool_calls=2) + if limit == "max_tool_calls" + else ExecutionPolicy(max_tool_calls_per_agent=2) + ) + spec = replace(_multi_tool_spec("send_email"), execution=policy) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **k: cast(BaseChatModel, model) + ) as engine: + await engine.build(spec) + first = await engine.run("review", context=RunContext(run_id="limited")) + assert first.pending_approval is not None + second = await engine.resume("limited", first.pending_approval.approval_id, "allow once") + assert second.pending_approval is not None + if resumed_stream: + events = [ + event + async for event in engine.resume_stream( + "limited", second.pending_approval.approval_id, "allow once" + ) + ] + assert events[-1].type == "final" + else: + result = await engine.resume( + "limited", second.pending_approval.approval_id, "allow once" + ) + assert result.status == "completed" + assert await engine.get_pending_approval("limited") is None + assert output.read_text().splitlines() == ["send_email:message-0", "send_email:message-1"] + + +async def test_replay_does_not_restore_another_agents_step_when_a_child_is_filtered( + tmp_path: Path, +) -> None: + output = tmp_path / "executed.log" + hidden = tmp_path / "hide_first" + write_tool(tmp_path, "send_email", output) + access = tmp_path / "plugins/access.py" + access.write_text( + "from pathlib import Path\n" + "class AccessResolver:\n" + " def can_access(self, ctx: dict, node_id: str) -> bool:\n" + f" return node_id != 'first' or not Path({str(hidden)!r}).exists()\n", + encoding="utf-8", + ) + writer = _multi_tool_spec("send_email").graph + assert isinstance(writer.node, AgentSpec) + first = replace( + writer, node=replace(writer.node, id="first", name="first", protected=True, auto_mode=True) + ) + spec = system(orchestrator("root", [first, writer])) + + class ParentModel(SiblingModel): + def _respond(self, messages: list[Any]) -> AIMessage: + response = super()._respond(messages) + if response.tool_calls: + response.tool_calls[0]["name"] = "first" + response.tool_calls[1]["name"] = "writer" + return response + + models = [ + ParentModel("first", "writer"), + DriftingModel("send_email", "send_email"), + DriftingModel("send_email", "send_email"), + ] + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **k: cast(BaseChatModel, models.pop(0)) + ) as engine: + await engine.build(spec) + pending = await engine.run("review", context=RunContext(run_id="filtered")) + assert pending.pending_approval is not None + assert output.read_text().splitlines() == ["send_email:reviewed"] + hidden.touch() + with pytest.raises(RuntimeError, match="checkpoint"): + await engine.resume("filtered", pending.pending_approval.approval_id, "allow once") + assert await engine.get_run_status("filtered") == "failed" + assert output.read_text().splitlines() == ["send_email:reviewed"] + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_binding_failure_rejects_the_approval_and_aborts_the_run( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, nested: bool, resumed_stream: bool +) -> None: + from agent_engine.engine.langgraph import approval_provider + + real_interrupt = approval_provider.interrupt + + def corrupt_resume(payload: dict[str, Any]) -> dict[str, Any]: + resume = real_interrupt(payload) + return {**resume, "tool_call_id": "different-invocation"} + + monkeypatch.setattr(approval_provider, "interrupt", corrupt_resume) + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + spec = _multi_tool_spec("send_email") + if nested: + spec = system(orchestrator("root", [spec.graph])) + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as engine: + await engine.build(spec) + first = await engine.run("review", context=RunContext(run_id="bad-binding")) + assert first.pending_approval is not None + approval_id = first.pending_approval.approval_id + with pytest.raises(ApprovalError, match="invocation"): + if resumed_stream: + _ = [ + event + async for event in engine.resume_stream( + "bad-binding", approval_id, "allow once" + ) + ] + else: + await engine.resume("bad-binding", approval_id, "allow once") + assert await engine.get_run_status("bad-binding") == "failed" + assert await engine.get_pending_approval("bad-binding") is None + approval = await manager.get_approval("bad-binding", approval_id) + assert approval.status == ApprovalStatus.REJECTED + assert not output.exists() + + +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_removed_pending_tool_cannot_complete_without_consuming_approval( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, resumed_stream: bool +) -> None: + from agent_engine.engine.langgraph.checkpointing import CheckpointProviderFactory + + handle = CheckpointProviderFactory().create(None) + monkeypatch.setattr(CheckpointProviderFactory, "create", lambda *a, **k: handle) + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + spec = _multi_tool_spec("send_email") + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as first_engine: + await first_engine.build(spec) + first = await first_engine.run("review", context=RunContext(run_id="removed")) + assert first.pending_approval is not None + + assert isinstance(spec.graph.node, AgentSpec) + without_tool = replace(spec, graph=replace(spec.graph, node=replace(spec.graph.node, tools=()))) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as second_engine: + await second_engine.build(without_tool) + approval_id = first.pending_approval.approval_id + with pytest.raises(ApprovalInvocationMismatch): + if resumed_stream: + _ = [ + event + async for event in second_engine.resume_stream( + "removed", approval_id, "allow once" + ) + ] + else: + await second_engine.resume("removed", approval_id, "allow once") + assert await second_engine.get_run_status("removed") == "failed" + record = await manager.get_approval("removed", approval_id) + assert record.status == ApprovalStatus.REJECTED + assert not output.exists() + + +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_legacy_continuation_is_rejected_before_any_replayed_side_effect( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, resumed_stream: bool +) -> None: + from agent_engine.approvals.in_memory_session_approval_repository import ( + InMemorySessionApprovalRepository, + ) + from agent_engine.approvals.invocation import SessionApprovalKey + from agent_engine.engine.langgraph import engine as engine_module + from agent_engine.engine.langgraph.execution import model_loop + from agent_engine.engine.langgraph.tools import tool_invoker + + original_input = engine_module._initial_state + + def legacy_input(*args: Any, **kwargs: Any) -> dict[str, Any]: + state = original_input(*args, **kwargs) + state.pop("continuation_version", None) + return state + + monkeypatch.setattr(engine_module, "_initial_state", legacy_input) + output = tmp_path / "executed.log" + for name in ("create_ticket", "send_email"): + write_tool(tmp_path, name, output) + sessions = InMemorySessionApprovalRepository() + await sessions.allow( + SessionApprovalKey( + "session", + "writer", + "local:local:create_ticket", + system_namespace="hitl", + user_id="owner", + ) + ) + model = SiblingModel("create_ticket", "send_email") + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, model), + approval_manager=manager, + session_approval_repository=sessions, + ) as engine: + await engine.build(_multi_tool_spec("create_ticket", "send_email")) + # Old continuations have neither a version nor checkpointed substeps. + with monkeypatch.context() as legacy: + legacy.setattr(model_loop, "in_graph_run", lambda: False) + legacy.setattr(tool_invoker, "in_graph_run", lambda: False) + first = await engine.run( + "review", + context=RunContext(run_id="legacy", conversation_id="session", user_id="owner"), + ) + assert first.pending_approval is not None + assert first.pending_approval.tool_name == "send_email" + assert output.read_text().splitlines() == ["create_ticket:ticket-1"] + with pytest.raises(RuntimeError, match="checkpoint"): + if resumed_stream: + _ = [ + event + async for event in engine.resume_stream( + "legacy", + first.pending_approval.approval_id, + "allow once", + caller_user_id="owner", + caller_session_id="session", + ) + ] + else: + await engine.resume( + "legacy", + first.pending_approval.approval_id, + "allow once", + caller_user_id="owner", + caller_session_id="session", + ) + assert await engine.get_run_status("legacy") == "failed" + approval = await manager.get_approval("legacy", first.pending_approval.approval_id) + assert approval.status == ApprovalStatus.REJECTED + assert model.proposals == 1 + assert output.read_text().splitlines() == ["create_ticket:ticket-1"] + + +@pytest.mark.parametrize("failure_kind", ["approval", "backend"]) +async def test_gate_cleanup_failure_still_reports_the_binding_mismatch( + monkeypatch: pytest.MonkeyPatch, + failure_kind: str, +) -> None: + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + reviewed = ToolInvocation( + tool_call_id="reviewed-call", agent_id="writer", tool_name="send_email", arguments={} + ) + await manager.create_pending( + run_id="r", + thread_id="r", + approval_id="approved", + agent_id="writer", + tool_name="send_email", + tool_call_id=reviewed.tool_call_id, + provider="local", + description="review", + arguments={}, + ) + await manager.claim(run_id="r", approval_id="approved") + + async def lost_race(**kwargs: Any) -> None: + if failure_kind == "backend": + raise RuntimeError("approval storage unavailable") + raise ApprovalAlreadyProcessed(kwargs["approval_id"], "rejected") + + monkeypatch.setattr(manager, "cancel_pending", lost_race) + monkeypatch.setattr( + "agent_engine.engine.langgraph.approval_provider.interrupt", + lambda _: { + "decision": "allow_once", + "approval_id": "approved", + "tool_call_id": reviewed.tool_call_id, + }, + ) + token = current_run_context.set(RunContext(run_id="r")) + try: + with pytest.raises(ApprovalInvocationMismatch) as raised: + await InterruptApprovalProvider(manager).request_decision( + ApprovalRequest(replace(reviewed, tool_call_id="changed-call"), "review", {}) + ) + finally: + current_run_context.reset(token) + expected = ApprovalAlreadyProcessed if failure_kind == "approval" else RuntimeError + assert isinstance(raised.value.__cause__, expected) + + +@pytest.mark.parametrize("resumed_stream", [False, True]) +@pytest.mark.parametrize("failure_kind", ["approval", "backend"]) +async def test_engine_cleanup_failure_still_reports_the_binding_mismatch( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, resumed_stream: bool, failure_kind: str +) -> None: + from agent_engine.engine.langgraph import approval_provider + + real_interrupt = approval_provider.interrupt + + def corrupt_resume(payload: dict[str, Any]) -> dict[str, Any]: + return {**real_interrupt(payload), "tool_call_id": "different-invocation"} + + monkeypatch.setattr(approval_provider, "interrupt", corrupt_resume) + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + + async def lost_race(approval_id: str, *, approved: bool) -> None: + if failure_kind == "backend": + raise RuntimeError("approval storage unavailable") + raise ApprovalAlreadyProcessed(approval_id, "rejected") + + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as engine: + await engine.build(_multi_tool_spec("send_email")) + first = await engine.run("review", context=RunContext(run_id="cleanup")) + assert first.pending_approval is not None + monkeypatch.setattr(manager, "finalize", lost_race) + with pytest.raises(ApprovalInvocationMismatch): + if resumed_stream: + _ = [ + event + async for event in engine.resume_stream( + "cleanup", first.pending_approval.approval_id, "allow once" + ) + ] + else: + await engine.resume("cleanup", first.pending_approval.approval_id, "allow once") + assert await engine.get_run_status("cleanup") == "failed" + assert not output.exists() diff --git a/tests/approvals/test_resume_policy.py b/tests/approvals/test_resume_policy.py new file mode 100644 index 00000000..6aeafeaa --- /dev/null +++ b/tests/approvals/test_resume_policy.py @@ -0,0 +1,266 @@ +"""Claimed decisions remain authoritative when approval policy changes.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from pathlib import Path +from typing import Any, cast + +import pytest +from langchain_core.language_models import BaseChatModel + +from agent_engine.approvals.approval_manager import ApprovalManager +from agent_engine.approvals.errors import ApprovalInvocationMismatch +from agent_engine.approvals.in_memory_approval_repository import InMemoryApprovalRepository +from agent_engine.approvals.models import ApprovalStatus +from agent_engine.core.spec import AgentSpec +from agent_engine.engine.langgraph.engine import LangGraphEngine +from agent_engine.runs.in_memory import InMemoryRunRepository +from agent_engine.runtime.hooks import RunContext +from tests.approvals.test_engine_hitl import FakeChatModel, _multi_tool_spec +from tests.approvals.test_resume_binding import write_tool +from tests.engine.test_engine_flow import orchestrator, system + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_pending_denial_is_not_bypassed_by_a_new_session_grant( + tmp_path: Path, nested: bool, resumed_stream: bool +) -> None: + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + spec = _multi_tool_spec("send_email") + if nested: + spec = system(orchestrator("root", [spec.graph])) + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as engine: + await engine.build(spec) + first = await engine.run( + "first", context=RunContext(run_id="denied", conversation_id="s", user_id="u") + ) + second = await engine.run( + "second", context=RunContext(run_id="granted", conversation_id="s", user_id="u") + ) + assert first.pending_approval is not None + assert second.pending_approval is not None + await engine.resume( + "granted", + second.pending_approval.approval_id, + "allow for session", + caller_user_id="u", + caller_session_id="s", + ) + assert output.read_text().splitlines() == ["send_email:go"] + if resumed_stream: + events = [ + event + async for event in engine.resume_stream( + "denied", + first.pending_approval.approval_id, + "deny", + caller_user_id="u", + caller_session_id="s", + ) + ] + assert events[-1].type == "final" + else: + result = await engine.resume( + "denied", + first.pending_approval.approval_id, + "deny", + caller_user_id="u", + caller_session_id="s", + ) + assert result.status == "completed" + record = await manager.get_approval("denied", first.pending_approval.approval_id) + assert record.status == ApprovalStatus.REJECTED + assert output.read_text().splitlines() == ["send_email:go"] + # The denial applies to its one pending call, not the other run's grant. + later = await engine.run( + "later", context=RunContext(run_id="later", conversation_id="s", user_id="u") + ) + assert later.status == "completed" + assert output.read_text().splitlines() == ["send_email:go", "send_email:go"] + + +@pytest.mark.parametrize("decision", ["deny", "allow once"]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_pending_decision_survives_an_auto_mode_change_on_rebuild( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, decision: str, resumed_stream: bool +) -> None: + from agent_engine.engine.langgraph.checkpointing import CheckpointProviderFactory + from tests.approvals.test_engine_hitl import ChainedApprovalModel + + handle = CheckpointProviderFactory().create(None) + monkeypatch.setattr(CheckpointProviderFactory, "create", lambda *a, **k: handle) + output = tmp_path / "executed.log" + for name in ("send_email", "archive_email"): + write_tool(tmp_path, name, output) + spec = _multi_tool_spec("send_email", "archive_email") + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, ChainedApprovalModel()), + approval_manager=manager, + ) as first_engine: + await first_engine.build(spec) + pending = await first_engine.run("review", context=RunContext(run_id="rebuilt")) + assert pending.pending_approval is not None + assert not output.exists() + + assert isinstance(spec.graph.node, AgentSpec) + automatic = replace( + spec, graph=replace(spec.graph, node=replace(spec.graph.node, auto_mode=True)) + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, ChainedApprovalModel()), + approval_manager=manager, + ) as second_engine: + await second_engine.build(automatic) + approval_id = pending.pending_approval.approval_id + if resumed_stream: + events = [ + event + async for event in second_engine.resume_stream("rebuilt", approval_id, decision) + ] + assert events[-1].type == "final" + else: + resumed = await second_engine.resume("rebuilt", approval_id, decision) + assert resumed.status == "completed" + record = await manager.get_approval("rebuilt", approval_id) + assert record.status == ( + ApprovalStatus.REJECTED if decision == "deny" else ApprovalStatus.APPROVED + ) + + expected = [] if decision == "deny" else ["send_email:send_email"] + # Subsequent calls still use the new policy once the pending decision is consumed. + assert output.read_text().splitlines() == [*expected, "archive_email:archive_email"] + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_concurrent_resumes_keep_their_claimed_decisions_separate( + tmp_path: Path, nested: bool, resumed_stream: bool +) -> None: + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + spec = _multi_tool_spec("send_email") + if nested: + spec = system(orchestrator("root", [spec.graph])) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()) + ) as engine: + await engine.build(spec) + approvals = [] + for run_id in ("allowed", "denied"): + pending = await engine.run( + "review", context=RunContext(run_id=run_id, conversation_id="s", user_id="u") + ) + assert pending.pending_approval is not None + approvals.append(pending.pending_approval.approval_id) + + async def resume(run_id: str, approval_id: str, decision: str) -> None: + if resumed_stream: + events = [ + event + async for event in engine.resume_stream( + run_id, approval_id, decision, caller_user_id="u", caller_session_id="s" + ) + ] + assert events[-1].type == "final" + else: + result = await engine.resume( + run_id, approval_id, decision, caller_user_id="u", caller_session_id="s" + ) + assert result.status == "completed" + + await asyncio.gather( + resume("allowed", approvals[0], "allow for session"), + resume("denied", approvals[1], "deny"), + ) + assert output.read_text().splitlines() == ["send_email:go"] + + +@pytest.mark.parametrize("resumed_stream", [False, True]) +async def test_rebuilt_provider_cannot_bypass_binding_with_auto_mode( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, resumed_stream: bool +) -> None: + from langchain_core.tools import StructuredTool + + from agent_engine.core.spec import MCPSpec + from agent_engine.engine.langgraph.checkpointing import CheckpointProviderFactory + from agent_engine.engine.langgraph.tools.mcp_connector import MCPConnector + + handle = CheckpointProviderFactory().create(None) + monkeypatch.setattr(CheckpointProviderFactory, "create", lambda *a, **k: handle) + output = tmp_path / "executed.log" + write_tool(tmp_path, "send_email", output) + + def remote_send(message: str) -> str: + output.write_text("mcp:" + message) + return message + + remote = StructuredTool.from_function(remote_send, name="send_email", description="Remote send") + + async def connect(self: MCPConnector, specs: Any) -> dict[str, list[Any]]: + return {"replacement": [remote]} if specs else {} + + monkeypatch.setattr(MCPConnector, "connect", connect) + spec = _multi_tool_spec("send_email") + manager = ApprovalManager( + run_repository=InMemoryRunRepository(), approval_repository=InMemoryApprovalRepository() + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as first_engine: + await first_engine.build(spec) + first = await first_engine.run("review", context=RunContext(run_id="replaced")) + assert first.pending_approval is not None + + assert isinstance(spec.graph.node, AgentSpec) + replaced_spec = replace( + spec, + graph=replace( + spec.graph, + node=replace( + spec.graph.node, + tools=(), + mcps=(MCPSpec(id="replacement", url="http://mcp.invalid/mcp"),), + auto_mode=True, + ), + ), + ) + async with LangGraphEngine( + tmp_path, + model_factory=lambda *a, **k: cast(BaseChatModel, FakeChatModel()), + approval_manager=manager, + ) as second_engine: + await second_engine.build(replaced_spec) + approval_id = first.pending_approval.approval_id + with pytest.raises(ApprovalInvocationMismatch): + if resumed_stream: + _ = [ + event + async for event in second_engine.resume_stream( + "replaced", approval_id, "allow once" + ) + ] + else: + await second_engine.resume("replaced", approval_id, "allow once") + assert not output.exists() + assert await second_engine.get_run_status("replaced") == "failed" + record = await manager.get_approval("replaced", approval_id) + assert record.status == ApprovalStatus.REJECTED diff --git a/tests/engine/test_model_loop_checkpointing.py b/tests/engine/test_model_loop_checkpointing.py new file mode 100644 index 00000000..ad2d68ac --- /dev/null +++ b/tests/engine/test_model_loop_checkpointing.py @@ -0,0 +1,24 @@ +"""Checkpointing is enabled only inside LangGraph, not every Runnable.""" + +from __future__ import annotations + +from typing import Any + +from langchain_core.runnables import RunnableLambda + +from agent_engine.engine.langgraph.execution.model_context import ModelContext +from agent_engine.engine.langgraph.execution.model_loop import run_tool_loop +from tests.approvals.test_engine_hitl import FakeChatModel + + +async def test_model_loop_inside_an_ordinary_runnable_needs_no_graph() -> None: + async def unexpected_tool(call: dict[str, Any]) -> str: + raise AssertionError("model did not request a tool") + + async def run(_: object) -> str: + result = await run_tool_loop( + FakeChatModel(), ModelContext("system", [], "hello"), "writer", unexpected_tool + ) + return str(result.content) + + assert await RunnableLambda(run).ainvoke({}) == "done"