fix(run-state): preserve agent-tool owner provenance across restore - #5046
namtran1812 wants to merge 2 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 23533b3bc9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if owner_signature is not None: | ||
| entry["agent_tool_owner_signature"] = owner_signature |
There was a problem hiding this comment.
When an interrupted Agent.as_tool() owner has literal instructions, prompt variables, model settings, manifest data, or capability configuration, _agent_identity_signature() returns those values as raw canonical JSON, and this assignment embeds that JSON directly in the persisted RunState. Applications that store or transmit snapshots now disclose private prompts and potentially credential-bearing configuration that was never previously part of the state; persist a non-reversible digest or another explicitly non-sensitive identifier instead.
AGENTS.md reference: AGENTS.md:L16-L17
Useful? React with 👍 / 👎.
| if ( | ||
| not isinstance(current_owner, Agent) | ||
| or _agent_identity_signature(current_owner) != persisted_owner_signature | ||
| ): |
There was a problem hiding this comment.
When a restored graph replaces the inner agent with another same-name agent containing same-name tools, this comparison accepts the replacement even if those tools have different callbacks, schemas, approval policies, or guardrails: _agent_identity_signature() reduces ordinary function tools to type/name/routing fields and omits all of those behavior-affecting values. The nested approved call can therefore execute the replacement implementation and its side effects, reproducing the provenance bypass this patch is intended to close; use provenance that distinguishes behavior-changing owner configurations.
AGENTS.md reference: AGENTS.md:L115-L119
Useful? React with 👍 / 👎.
| nested_state_data = entry.get("agent_run_state") | ||
| owner_signature = entry.get("agent_tool_owner_signature") | ||
| if owner_signature is not None and not isinstance(owner_signature, str): | ||
| raise validation_error_factory( | ||
| "Run state Agent.as_tool() owner provenance has an invalid type.", | ||
| UserError, | ||
| ) |
There was a problem hiding this comment.
When a current-version 1.17 snapshot is corrupted or edited to remove agent_tool_owner_signature, this parser accepts the missing value as None, and _rebind_function_run() subsequently skips the provenance check altogether. Removing one field therefore restores the original replacement vulnerability for any interrupted Agent.as_tool() call; only snapshots carrying an older schema version should receive the legacy fallback, while current-schema nested states should require valid provenance.
AGENTS.md reference: AGENTS.md:L119-L119
Useful? React with 👍 / 👎.
| pending_result = _SerializedAgentToolRunResult( | ||
| nested_state, | ||
| agent_tool_owner_signature=function_action.agent_tool_owner_signature, | ||
| ) |
There was a problem hiding this comment.
When the configured agent-tool owner has been replaced, provenance is merely attached here after _build_run_state_from_json() has already recursively restored the nested state; that restoration calls the replacement inner agent's get_all_tools(), which evaluates public is_enabled callbacks and may perform MCP discovery. Those callbacks or requests can produce side effects during RunState.from_json() before the mismatch is rejected later by Runner.run(), so compare the persisted signature with the bound owner before recursively restoring the nested state.
AGENTS.md reference: AGENTS.md:L118-L119
Useful? React with 👍 / 👎.
| function_tool = getattr(function_run, "function_tool", None) | ||
| owner = getattr(function_tool, "_agent_instance", None) | ||
| owner_signature = _agent_identity_signature(owner) if isinstance(owner, Agent) else None |
There was a problem hiding this comment.
When the original inner agent uses a stateful AgentHooks implementation, normal callbacks before the interruption commonly mutate public fields such as an event list or counter; _agent_identity_signature() recursively includes those fields, so the persisted value represents runtime history rather than stable configuration. Reconstructing the same logical graph in another process with a fresh hook instance then produces a different signature and incorrectly rejects the resume, contrary to the supported equivalent-owner restoration path; derive provenance only from stable declarative identity.
AGENTS.md reference: AGENTS.md:L117-L117
Useful? React with 👍 / 👎.
| function_tool = getattr(function_run, "function_tool", None) | ||
| owner = getattr(function_tool, "_agent_instance", None) | ||
| owner_signature = _agent_identity_signature(owner) if isinstance(owner, Agent) else None |
There was a problem hiding this comment.
When a snapshot created for owner A is restored against replacement owner B and then serialized again before Runner.run() performs reconciliation, this code recomputes the signature from the already rebound function_run.function_tool and writes B's signature instead of retaining the pending result's persisted A signature. Restoring that second snapshot against B now passes the provenance check and can execute the replacement, so reserialization must carry forward pending_run_result.agent_tool_owner_signature rather than laundering it through the currently bound tool.
AGENTS.md reference: AGENTS.md:L115-L115
Useful? React with 👍 / 👎.
1 participant
Summary
Persist
Agent.as_tool()owner provenance for interrupted nested runs so owner replacement can still be detected afterRunStateserialization and restoration.Previously, replacing an interrupted agent tool before
RunState.from_json()could cause the restored action to resolve directly to the replacement wrapper. This loses the original Python object identity before interrupted-turn reconciliation, allowing the existing replacement check to be bypassed when the replacement uses the same agent and tool names.This change:
Agent.as_tool()stateThe new field is included in the current unreleased RunState schema version (
1.17).Testing
tests/test_tool_name_collision_policy.py: 40 passedtests/test_run_state.py: 446 passedtests/test_run_state_compatibility_corpus.py: 117 passed