From 1121881f150b965aef4fd28443db7774f0fa334e Mon Sep 17 00:00:00 2001 From: Asaf Varon Date: Wed, 14 Oct 2026 19:29:04 +0300 Subject: [PATCH 1/2] fix: bind HITL resumes to the reviewed tool invocation Checkpoint model responses and completed tool outcomes so approval resumes reuse the reviewed request without repeating earlier side effects. Validate invocation bindings and fail closed on mismatched or legacy continuations. Honor claimed decisions across session-grant and auto-mode changes, preserve execution limits, isolate concurrent resumes, and retain mismatch errors when approval cleanup fails. Document the continuation format and add regression coverage for streaming, nested approvals, provider changes, and replay safety. Validation: make check passed with 1050 tests. Refs #139. --- docs/HUMAN_IN_THE_LOOP.md | 40 +- docs/RUNTIME_LIFECYCLE.md | 23 +- .../0005-checkpointed-hitl-continuation.md | 78 ++ pyproject.toml | 3 +- src/agent_engine/approvals/coordinator.py | 6 +- src/agent_engine/approvals/errors.py | 8 + src/agent_engine/approvals/invocation.py | 1 + .../engine/langgraph/approval_provider.py | 30 +- .../engine/langgraph/checkpointing.py | 16 + src/agent_engine/engine/langgraph/engine.py | 93 ++- .../langgraph/execution/approval_resume.py | 39 + .../engine/langgraph/execution/model_loop.py | 48 +- .../langgraph/nodes/orchestrator_node.py | 8 +- .../engine/langgraph/tools/tool_invoker.py | 28 +- src/agent_engine/runtime/state.py | 3 + tests/approvals/test_coordinator.py | 22 + tests/approvals/test_engine_hitl.py | 3 +- tests/approvals/test_resume_binding.py | 687 ++++++++++++++++++ tests/approvals/test_resume_policy.py | 266 +++++++ tests/engine/test_model_loop_checkpointing.py | 24 + 20 files changed, 1386 insertions(+), 40 deletions(-) create mode 100644 docs/adr/0005-checkpointed-hitl-continuation.md create mode 100644 src/agent_engine/engine/langgraph/execution/approval_resume.py create mode 100644 tests/approvals/test_resume_binding.py create mode 100644 tests/approvals/test_resume_policy.py create mode 100644 tests/engine/test_model_loop_checkpointing.py diff --git a/docs/HUMAN_IN_THE_LOOP.md b/docs/HUMAN_IN_THE_LOOP.md index 356814a7..b6bb2fdf 100644 --- a/docs/HUMAN_IN_THE_LOOP.md +++ b/docs/HUMAN_IN_THE_LOOP.md @@ -205,26 +205,52 @@ 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. + 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..8632f94a 100644 --- a/docs/RUNTIME_LIFECYCLE.md +++ b/docs/RUNTIME_LIFECYCLE.md @@ -77,7 +77,28 @@ 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. +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..c66da4d1 --- /dev/null +++ b/docs/adr/0005-checkpointed-hitl-continuation.md @@ -0,0 +1,78 @@ +# 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. 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..9ddc7b9f 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,36 @@ 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}") + return await app.ainvoke(graph_input, self._thread_config(ctx)) + 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 +656,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 +674,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 +726,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 +747,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 +763,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 +781,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 +790,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 +812,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 +889,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..b2f34cd8 --- /dev/null +++ b/src/agent_engine/engine/langgraph/execution/approval_resume.py @@ -0,0 +1,39 @@ +"""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 + + +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..293e77da --- /dev/null +++ b/tests/approvals/test_resume_binding.py @@ -0,0 +1,687 @@ +"""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_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" From c39c2ddf5c6ee2f69d4507c9eff54cb1b0d91c4b Mon Sep 17 00:00:00 2001 From: Amit Avital Date: Wed, 14 Oct 2026 20:34:30 +0300 Subject: [PATCH 2/2] fix: reject unconsumed approval resumes --- docs/HUMAN_IN_THE_LOOP.md | 4 +- docs/RUNTIME_LIFECYCLE.md | 3 ++ .../0005-checkpointed-hitl-continuation.md | 6 ++- src/agent_engine/engine/langgraph/engine.py | 9 +++- .../langgraph/execution/approval_resume.py | 5 ++ tests/approvals/test_resume_binding.py | 49 +++++++++++++++++++ 6 files changed, 72 insertions(+), 4 deletions(-) diff --git a/docs/HUMAN_IN_THE_LOOP.md b/docs/HUMAN_IN_THE_LOOP.md index b6bb2fdf..a9431be1 100644 --- a/docs/HUMAN_IN_THE_LOOP.md +++ b/docs/HUMAN_IN_THE_LOOP.md @@ -233,7 +233,9 @@ 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. +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. diff --git a/docs/RUNTIME_LIFECYCLE.md b/docs/RUNTIME_LIFECYCLE.md index 8632f94a..c08caf85 100644 --- a/docs/RUNTIME_LIFECYCLE.md +++ b/docs/RUNTIME_LIFECYCLE.md @@ -88,6 +88,9 @@ 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. diff --git a/docs/adr/0005-checkpointed-hitl-continuation.md b/docs/adr/0005-checkpointed-hitl-continuation.md index c66da4d1..785bb5d1 100644 --- a/docs/adr/0005-checkpointed-hitl-continuation.md +++ b/docs/adr/0005-checkpointed-hitl-continuation.md @@ -41,8 +41,10 @@ one invocation. 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. Separate requests - never share this state, and completed steps do not consume a new decision. + 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 diff --git a/src/agent_engine/engine/langgraph/engine.py b/src/agent_engine/engine/langgraph/engine.py index 9ddc7b9f..015492a0 100644 --- a/src/agent_engine/engine/langgraph/engine.py +++ b/src/agent_engine/engine/langgraph/engine.py @@ -436,7 +436,14 @@ async def _invoke_graph( # 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}") - return await app.ainvoke(graph_input, self._thread_config(ctx)) + 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. diff --git a/src/agent_engine/engine/langgraph/execution/approval_resume.py b/src/agent_engine/engine/langgraph/execution/approval_resume.py index b2f34cd8..851041b2 100644 --- a/src/agent_engine/engine/langgraph/execution/approval_resume.py +++ b/src/agent_engine/engine/langgraph/execution/approval_resume.py @@ -33,6 +33,11 @@ def consume(self, approval_id: str, tool_call_id: str) -> None: 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/tests/approvals/test_resume_binding.py b/tests/approvals/test_resume_binding.py index 293e77da..88c11053 100644 --- a/tests/approvals/test_resume_binding.py +++ b/tests/approvals/test_resume_binding.py @@ -508,6 +508,55 @@ def corrupt_resume(payload: dict[str, Any]) -> dict[str, Any]: 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