fix(core): preserve serialized approval resume ownership (#4613)
This commit is contained in:
Kazuhiro Sera
+177
-1
@@ -182,6 +182,7 @@ def _default_run_state_validation_error(
|
||||
CURRENT_SCHEMA_VERSION = "1.17"
|
||||
_PROGRAMMATIC_TOOL_CALLING_MIN_SCHEMA_VERSION = "1.13"
|
||||
_HOSTED_MCP_APPROVALS_MIN_SCHEMA_VERSION = "1.14"
|
||||
_CURRENT_RESPONSE_OWNERSHIP_MIN_SCHEMA_VERSION = "1.17"
|
||||
# Keep this mapping in chronological order. Every schema bump must add a one-line summary here.
|
||||
SCHEMA_VERSION_SUMMARIES: dict[str, str] = {
|
||||
"1.0": "Initial RunState snapshot format for HITL pause/resume flows.",
|
||||
@@ -213,7 +214,10 @@ SCHEMA_VERSION_SUMMARIES: dict[str, str] = {
|
||||
"Persists Docker network-isolation state and lets an exact call approval decision "
|
||||
"override a sticky decision for the same tool."
|
||||
),
|
||||
"1.17": "Persists Docker container labels across sandbox resume and replacement.",
|
||||
"1.17": (
|
||||
"Persists Docker container labels and current-response generated-item ownership across "
|
||||
"resume flows."
|
||||
),
|
||||
}
|
||||
SUPPORTED_SCHEMA_VERSIONS = frozenset(SCHEMA_VERSION_SUMMARIES)
|
||||
|
||||
@@ -1457,6 +1461,46 @@ class RunState(Generic[TContext, TAgent]):
|
||||
indexes.append(session_index)
|
||||
return indexes
|
||||
|
||||
def _current_response_generated_item_ownership(
|
||||
self,
|
||||
generated_items: Sequence[RunItem],
|
||||
) -> dict[str, Any] | None:
|
||||
"""Record the response range and approval occurrences from live item identities."""
|
||||
from .run_internal.run_steps import NextStepInterruption
|
||||
|
||||
if self._last_processed_response is None:
|
||||
return None
|
||||
if not isinstance(self._current_step, NextStepInterruption):
|
||||
return None
|
||||
|
||||
processed_items = self._last_processed_response.new_items
|
||||
interruptions = self._current_step.interruptions
|
||||
if not processed_items or not interruptions or len(processed_items) > len(generated_items):
|
||||
return None
|
||||
|
||||
candidate_starts = [
|
||||
start
|
||||
for start in range(len(generated_items) - len(processed_items) + 1)
|
||||
if all(
|
||||
generated_items[start + offset] is item
|
||||
for offset, item in enumerate(processed_items)
|
||||
)
|
||||
]
|
||||
if len(candidate_starts) != 1:
|
||||
return None
|
||||
|
||||
start = candidate_starts[0]
|
||||
indexes_by_identity: dict[int, list[int]] = {}
|
||||
for index in range(start + len(processed_items), len(generated_items)):
|
||||
indexes_by_identity.setdefault(id(generated_items[index]), []).append(index)
|
||||
interruption_indexes: list[int] = []
|
||||
for item in interruptions:
|
||||
indexes = indexes_by_identity.pop(id(item), [])
|
||||
if len(indexes) != 1:
|
||||
return None
|
||||
interruption_indexes.append(indexes[0])
|
||||
return {"start": start, "end": len(generated_items), "interruptions": interruption_indexes}
|
||||
|
||||
def _serialize_context_payload(
|
||||
self,
|
||||
*,
|
||||
@@ -1805,6 +1849,14 @@ class RunState(Generic[TContext, TAgent]):
|
||||
"generated_session_item_indexes": self._generated_session_item_indexes(generated_items),
|
||||
}
|
||||
|
||||
current_response_generated_item_ownership = self._current_response_generated_item_ownership(
|
||||
generated_items
|
||||
)
|
||||
if current_response_generated_item_ownership is not None:
|
||||
result["current_response_generated_item_ownership"] = (
|
||||
current_response_generated_item_ownership
|
||||
)
|
||||
|
||||
result["generated_items"] = [
|
||||
self._serialize_item(item, agent_identity_keys_by_id=agent_identity_keys_by_id)
|
||||
for item in generated_items
|
||||
@@ -4250,6 +4302,24 @@ async def _build_run_state_from_json(
|
||||
current_step_data.get("data", {}).get("llm_end_hooks_started", True)
|
||||
),
|
||||
)
|
||||
_restore_current_response_item_identities(
|
||||
state,
|
||||
serialized_generated_items=serialized_generated_items,
|
||||
generated_source_indexes=generated_source_indexes,
|
||||
last_processed_response_data=last_processed_response_data,
|
||||
current_step_data=current_step_data,
|
||||
current_response_generated_item_ownership=(
|
||||
state_json.get("current_response_generated_item_ownership")
|
||||
if (schema_major, schema_minor)
|
||||
>= tuple(
|
||||
int(part)
|
||||
for part in _CURRENT_RESPONSE_OWNERSHIP_MIN_SCHEMA_VERSION.split(
|
||||
".", maxsplit=1
|
||||
)
|
||||
)
|
||||
else None
|
||||
),
|
||||
)
|
||||
if state._current_step.response_accepted:
|
||||
state._clear_generated_items_last_processed_marker()
|
||||
for approval_item in state._current_step.interruptions:
|
||||
@@ -5396,6 +5466,112 @@ def _deserialize_items_with_source_indexes(
|
||||
return items, source_indexes
|
||||
|
||||
|
||||
def _restore_current_response_item_identities(
|
||||
state: RunState[Any],
|
||||
*,
|
||||
serialized_generated_items: Any,
|
||||
generated_source_indexes: Sequence[int],
|
||||
last_processed_response_data: Any,
|
||||
current_step_data: Mapping[str, Any],
|
||||
current_response_generated_item_ownership: Any,
|
||||
) -> None:
|
||||
"""Relink one response from explicit generated-item ownership after deserialization."""
|
||||
from .run_internal.run_steps import NextStepInterruption
|
||||
|
||||
processed_response = state._last_processed_response
|
||||
if processed_response is None:
|
||||
return
|
||||
current_step = state._current_step
|
||||
if not isinstance(current_step, NextStepInterruption):
|
||||
return
|
||||
if not isinstance(serialized_generated_items, list):
|
||||
return
|
||||
if not isinstance(last_processed_response_data, Mapping):
|
||||
return
|
||||
|
||||
serialized_processed_items = last_processed_response_data.get("new_items")
|
||||
if not isinstance(serialized_processed_items, list) or not serialized_processed_items:
|
||||
return
|
||||
if len(processed_response.new_items) != len(serialized_processed_items):
|
||||
return
|
||||
current_step_payload = current_step_data.get("data")
|
||||
if not isinstance(current_step_payload, Mapping):
|
||||
return
|
||||
serialized_interruptions = current_step_payload.get("interruptions")
|
||||
if not isinstance(serialized_interruptions, list) or not serialized_interruptions:
|
||||
return
|
||||
if len(current_step.interruptions) != len(serialized_interruptions):
|
||||
return
|
||||
ownership = current_response_generated_item_ownership
|
||||
if not isinstance(ownership, Mapping):
|
||||
return
|
||||
source_start = ownership.get("start")
|
||||
source_end = ownership.get("end")
|
||||
interruption_indexes = ownership.get("interruptions")
|
||||
if type(source_start) is not int or type(source_end) is not int:
|
||||
return
|
||||
if source_start < 0 or source_end != len(serialized_generated_items):
|
||||
return
|
||||
processed_end = source_start + len(serialized_processed_items)
|
||||
# Handoff filters can clear prior items without resetting the model turn count.
|
||||
if processed_end > source_end:
|
||||
return
|
||||
if not isinstance(interruption_indexes, list):
|
||||
return
|
||||
if len(interruption_indexes) != len(serialized_interruptions) or any(
|
||||
type(index) is not int or index < processed_end or index >= source_end
|
||||
for index in interruption_indexes
|
||||
):
|
||||
return
|
||||
if len(set(interruption_indexes)) != len(interruption_indexes):
|
||||
return
|
||||
source_indexes = [*range(source_start, processed_end), *interruption_indexes]
|
||||
serialized_current_response_items = [*serialized_processed_items, *serialized_interruptions]
|
||||
if any(
|
||||
serialized_generated_items[source_index] != expected_item
|
||||
for source_index, expected_item in zip(
|
||||
source_indexes,
|
||||
serialized_current_response_items,
|
||||
strict=True,
|
||||
)
|
||||
):
|
||||
return
|
||||
|
||||
restored_indexes_by_source: dict[int, list[int]] = {}
|
||||
for restored_index, source_index in enumerate(generated_source_indexes):
|
||||
restored_indexes_by_source.setdefault(source_index, []).append(restored_index)
|
||||
|
||||
restored_current_response_items: list[RunItem] = []
|
||||
for source_index in range(source_start, source_end):
|
||||
restored_indexes = restored_indexes_by_source.get(source_index)
|
||||
if restored_indexes is None or len(restored_indexes) != 1:
|
||||
return
|
||||
restored_current_response_items.append(state._generated_items[restored_indexes[0]])
|
||||
|
||||
processed_item_count = len(serialized_processed_items)
|
||||
restored_processed_items = restored_current_response_items[:processed_item_count]
|
||||
restored_interruptions = [
|
||||
restored_current_response_items[index - source_start] for index in interruption_indexes
|
||||
]
|
||||
if not all(isinstance(item, ToolApprovalItem) for item in restored_interruptions):
|
||||
return
|
||||
|
||||
# The complete current response must be the same terminal suffix in both histories.
|
||||
session_start = len(state._session_items) - len(restored_current_response_items)
|
||||
if session_start < 0 or any(
|
||||
generated_item is not session_item
|
||||
for generated_item, session_item in zip(
|
||||
restored_current_response_items,
|
||||
state._session_items[session_start:],
|
||||
strict=True,
|
||||
)
|
||||
):
|
||||
return
|
||||
|
||||
processed_response.new_items = restored_processed_items
|
||||
current_step.interruptions = cast(list[ToolApprovalItem], restored_interruptions)
|
||||
|
||||
|
||||
def _clone_original_input(original_input: str | list[Any]) -> str | list[Any]:
|
||||
"""Return a deep copy of the original input so later mutations don't leak into saved state."""
|
||||
if isinstance(original_input, str):
|
||||
|
||||
@@ -2572,16 +2572,363 @@ async def test_resumed_approved_tool_final_persists_complete_post_verdict_batch(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["non_streamed", "streamed"])
|
||||
@pytest.mark.parametrize("attach_session", [False, True], ids=["without-session", "with-session"])
|
||||
@pytest.mark.parametrize("mixed_tool_position", [None, "before", "after"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_ambiguous_serialized_approval_state_fails_before_tool_execution(
|
||||
async def test_serialized_later_turn_approval_with_output_guardrail_resumes(
|
||||
mode: str,
|
||||
attach_session: bool,
|
||||
mixed_tool_position: str | None,
|
||||
) -> None:
|
||||
tool_calls = 0
|
||||
tool_calls = {"normal": 0, "approval": 0}
|
||||
|
||||
@function_tool(name_override="normal_tool")
|
||||
def normal_tool() -> str:
|
||||
tool_calls["normal"] += 1
|
||||
return "normal-result"
|
||||
|
||||
@function_tool(name_override="approval_tool", needs_approval=True)
|
||||
def approval_tool() -> str:
|
||||
tool_calls["approval"] += 1
|
||||
return "approved-result"
|
||||
|
||||
def output_guardrail(
|
||||
_context: RunContextWrapper[Any],
|
||||
_agent: Agent[Any],
|
||||
_output: Any,
|
||||
) -> GuardrailFunctionOutput:
|
||||
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False)
|
||||
|
||||
approval_response = [get_function_tool_call("approval_tool", "{}", call_id="call-approved")]
|
||||
if mixed_tool_position is not None:
|
||||
extra_call = get_function_tool_call("normal_tool", "{}", call_id="call-extra")
|
||||
approval_response.insert(0 if mixed_tool_position == "before" else 1, extra_call)
|
||||
expected_normal_calls = 1 if mixed_tool_position is None else 2
|
||||
model = ScriptedModel(
|
||||
[
|
||||
[get_function_tool_call("normal_tool", "{}", call_id="call-normal")],
|
||||
approval_response,
|
||||
[get_text_message("done")],
|
||||
]
|
||||
)
|
||||
agent = Agent(
|
||||
name="test",
|
||||
model=model,
|
||||
tools=[normal_tool, approval_tool],
|
||||
output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)],
|
||||
)
|
||||
session = SimpleListSession() if attach_session else None
|
||||
|
||||
async def run_once(input_value: Any) -> Any:
|
||||
if mode == "non_streamed":
|
||||
return await Runner.run(agent, input_value, session=session)
|
||||
result = Runner.run_streamed(agent, input_value, session=session)
|
||||
await consume_stream(result)
|
||||
return result
|
||||
|
||||
first = await run_once("Use normal_tool, then approval_tool")
|
||||
state = first.to_state()
|
||||
assert state._current_turn == 2
|
||||
assert [item.tool_name for item in first.interruptions] == ["approval_tool"]
|
||||
assert tool_calls == {"normal": expected_normal_calls, "approval": 0}
|
||||
|
||||
serialized = state.to_json()
|
||||
assert serialized["current_response_generated_item_ownership"] == {
|
||||
"start": 2,
|
||||
"end": 4 if mixed_tool_position is None else 6,
|
||||
"interruptions": [
|
||||
3 if mixed_tool_position is None else (5 if mixed_tool_position == "before" else 4)
|
||||
],
|
||||
}
|
||||
|
||||
released_payload = json.loads(json.dumps(serialized))
|
||||
released_payload["$schemaVersion"] = "1.16"
|
||||
released = await RunState.from_string(agent, json.dumps(released_payload))
|
||||
released.approve(released.get_interruptions()[0])
|
||||
with pytest.raises(UserError, match="current response boundary cannot be proven"):
|
||||
await run_once(released)
|
||||
assert tool_calls == {"normal": expected_normal_calls, "approval": 0}
|
||||
|
||||
malformed_payload = json.loads(json.dumps(serialized))
|
||||
malformed_payload["current_response_generated_item_ownership"]["interruptions"][0] = True
|
||||
malformed = await RunState.from_string(agent, json.dumps(malformed_payload))
|
||||
malformed.approve(malformed.get_interruptions()[0])
|
||||
with pytest.raises(UserError, match="current response boundary cannot be proven"):
|
||||
await run_once(malformed)
|
||||
assert tool_calls == {"normal": expected_normal_calls, "approval": 0}
|
||||
|
||||
restored_once = await RunState.from_string(agent, json.dumps(serialized))
|
||||
reserialized = restored_once.to_json()
|
||||
assert (
|
||||
reserialized["current_response_generated_item_ownership"]
|
||||
== serialized["current_response_generated_item_ownership"]
|
||||
)
|
||||
|
||||
restored = await RunState.from_string(agent, json.dumps(reserialized))
|
||||
assert restored._last_processed_response is not None
|
||||
assert any(
|
||||
restored._last_processed_response.new_items[0] is item for item in restored._generated_items
|
||||
)
|
||||
restored.approve(restored.get_interruptions()[0])
|
||||
|
||||
resumed = await run_once(restored)
|
||||
|
||||
assert resumed.final_output == "done"
|
||||
assert tool_calls == {"normal": expected_normal_calls, "approval": 1}
|
||||
if session is not None:
|
||||
saved_items = await session.get_items()
|
||||
saved_tool_items = [
|
||||
item
|
||||
for item in saved_items
|
||||
if isinstance(item, dict)
|
||||
and item.get("type") in {"function_call", "function_call_output"}
|
||||
]
|
||||
expected_tool_items = [
|
||||
("function_call", "call-normal"),
|
||||
("function_call_output", "call-normal"),
|
||||
("function_call", "call-approved"),
|
||||
("function_call_output", "call-approved"),
|
||||
]
|
||||
if mixed_tool_position is not None:
|
||||
expected_tool_items.insert(
|
||||
2 if mixed_tool_position == "before" else 3, ("function_call", "call-extra")
|
||||
)
|
||||
expected_tool_items.insert(4, ("function_call_output", "call-extra"))
|
||||
assert [(item.get("type"), item.get("call_id")) for item in saved_tool_items] == (
|
||||
expected_tool_items
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["non_streamed", "streamed"])
|
||||
@pytest.mark.parametrize(
|
||||
"session_ownership",
|
||||
[
|
||||
"valid",
|
||||
"missing-prefix",
|
||||
"missing",
|
||||
"invalid",
|
||||
"missing-current",
|
||||
"prefix-anchor",
|
||||
"nonterminal-session",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_serialized_mixed_approval_guardrail_preserves_only_accepted_outputs(
|
||||
mode: str,
|
||||
session_ownership: str,
|
||||
) -> None:
|
||||
tool_calls = {"normal": 0, "approval": 0}
|
||||
|
||||
@function_tool
|
||||
def normal_tool() -> str:
|
||||
tool_calls["normal"] += 1
|
||||
return "accepted-result" if tool_calls["normal"] == 1 else "sibling-secret"
|
||||
|
||||
@function_tool(needs_approval=True)
|
||||
def approval_tool() -> str:
|
||||
tool_calls["approval"] += 1
|
||||
return "approved-secret"
|
||||
|
||||
model = ScriptedModel(
|
||||
[
|
||||
[get_function_tool_call("normal_tool", "{}", call_id="call-normal")],
|
||||
[
|
||||
get_function_tool_call("normal_tool", "{}", call_id="call-extra"),
|
||||
get_function_tool_call("approval_tool", "{}", call_id="call-approved"),
|
||||
],
|
||||
]
|
||||
)
|
||||
agent = Agent(
|
||||
name="test",
|
||||
model=model,
|
||||
tools=[normal_tool, approval_tool],
|
||||
tool_use_behavior={"stop_at_tool_names": ["approval_tool"]},
|
||||
output_guardrails=[
|
||||
OutputGuardrail(
|
||||
guardrail_function=lambda *_: GuardrailFunctionOutput(
|
||||
output_info=None, tripwire_triggered=True
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
session = SimpleListSession()
|
||||
|
||||
async def run_once(input_value: Any) -> Any:
|
||||
if mode == "non_streamed":
|
||||
return await Runner.run(agent, input_value, session=session)
|
||||
result = Runner.run_streamed(agent, input_value, session=session)
|
||||
await consume_stream(result)
|
||||
return result
|
||||
|
||||
first = await run_once("Use the tools")
|
||||
payload = first.to_state().to_json()
|
||||
if session_ownership == "missing":
|
||||
payload.pop("generated_session_item_indexes")
|
||||
elif session_ownership == "missing-prefix":
|
||||
payload["generated_session_item_indexes"][0] = None
|
||||
elif session_ownership == "invalid":
|
||||
payload["generated_session_item_indexes"][0] = True
|
||||
elif session_ownership == "missing-current":
|
||||
current_start = payload["current_response_generated_item_ownership"]["start"]
|
||||
payload["generated_session_item_indexes"][current_start] = None
|
||||
elif session_ownership == "prefix-anchor":
|
||||
earlier_copies = json.loads(json.dumps(payload["last_processed_response"]["new_items"]))
|
||||
payload["session_items"][:0] = earlier_copies
|
||||
payload["generated_session_item_indexes"] = [
|
||||
index + len(earlier_copies) if index is not None else None
|
||||
for index in payload["generated_session_item_indexes"]
|
||||
]
|
||||
current_start = payload["current_response_generated_item_ownership"]["start"]
|
||||
for offset in range(len(earlier_copies)):
|
||||
payload["generated_session_item_indexes"][current_start + offset] = offset
|
||||
elif session_ownership == "nonterminal-session":
|
||||
payload["session_items"].append(json.loads(json.dumps(payload["session_items"][0])))
|
||||
restored = await RunState.from_string(agent, json.dumps(payload))
|
||||
restored = await RunState.from_string(agent, restored.to_string())
|
||||
restored.approve(restored.get_interruptions()[0])
|
||||
if session_ownership not in {"valid", "missing-prefix"}:
|
||||
saved_before = json.loads(json.dumps(await session.get_items()))
|
||||
state_session_before = restored.to_json()["session_items"]
|
||||
with pytest.raises(UserError, match="current response boundary cannot be proven"):
|
||||
await run_once(restored)
|
||||
assert tool_calls == {"normal": 2, "approval": 0}
|
||||
assert await session.get_items() == saved_before
|
||||
assert restored.to_json()["session_items"] == state_session_before
|
||||
assert "accepted-result" in json.dumps(state_session_before)
|
||||
return
|
||||
with pytest.raises(OutputGuardrailTripwireTriggered):
|
||||
await run_once(restored)
|
||||
|
||||
assert tool_calls == {"normal": 2, "approval": 1}
|
||||
outputs = [
|
||||
(item["call_id"], item["output"])
|
||||
for item in await session.get_items()
|
||||
if isinstance(item, dict) and item.get("type") == "function_call_output"
|
||||
]
|
||||
assert outputs == [
|
||||
("call-normal", "accepted-result"),
|
||||
("call-extra", run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT),
|
||||
("call-approved", run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT),
|
||||
]
|
||||
serialized = restored.to_string()
|
||||
assert "accepted-result" in serialized
|
||||
assert "sibling-secret" not in serialized
|
||||
assert "approved-secret" not in serialized
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["non_streamed", "streamed"])
|
||||
@pytest.mark.parametrize("attach_session", [False, True], ids=["without-session", "with-session"])
|
||||
@pytest.mark.parametrize("tripwire", [False, True], ids=["accepted", "blocked"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_serialized_filtered_handoff_approval_with_empty_prefix_resumes(
|
||||
mode: str,
|
||||
attach_session: bool,
|
||||
tripwire: bool,
|
||||
) -> None:
|
||||
tool_calls = 0
|
||||
|
||||
@function_tool(needs_approval=True)
|
||||
def approval_tool() -> str:
|
||||
nonlocal tool_calls
|
||||
tool_calls += 1
|
||||
return "approved-secret"
|
||||
|
||||
def clear_generated_history(data: HandoffInputData) -> HandoffInputData:
|
||||
return HandoffInputData(
|
||||
input_history=data.input_history,
|
||||
pre_handoff_items=(),
|
||||
new_items=(),
|
||||
run_context=data.run_context,
|
||||
)
|
||||
|
||||
target = Agent(
|
||||
name="target",
|
||||
model=ScriptedModel(
|
||||
[[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]]
|
||||
),
|
||||
tools=[approval_tool],
|
||||
tool_use_behavior="stop_on_first_tool",
|
||||
output_guardrails=[
|
||||
OutputGuardrail(
|
||||
guardrail_function=lambda *_: GuardrailFunctionOutput(
|
||||
output_info=None, tripwire_triggered=tripwire
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
starting = Agent(
|
||||
name="starting",
|
||||
model=ScriptedModel([[get_handoff_tool_call(target)]]),
|
||||
handoffs=[handoff(target, input_filter=clear_generated_history)],
|
||||
)
|
||||
session = SimpleListSession() if attach_session else None
|
||||
|
||||
async def run_once(input_value: Any) -> Any:
|
||||
if mode == "non_streamed":
|
||||
return await Runner.run(starting, input_value, session=session)
|
||||
result = Runner.run_streamed(starting, input_value, session=session)
|
||||
await consume_stream(result)
|
||||
return result
|
||||
|
||||
first = await run_once("Transfer and run approval_tool")
|
||||
state = first.to_state()
|
||||
assert state._current_turn == 2
|
||||
assert tool_calls == 0
|
||||
ownership = {"start": 0, "end": 2, "interruptions": [1]}
|
||||
assert state.to_json()["current_response_generated_item_ownership"] == ownership
|
||||
restored = await RunState.from_string(starting, state.to_string())
|
||||
assert restored.to_json()["current_response_generated_item_ownership"] == ownership
|
||||
restored = await RunState.from_string(starting, restored.to_string())
|
||||
restored.approve(restored.get_interruptions()[0])
|
||||
|
||||
if tripwire:
|
||||
with pytest.raises(OutputGuardrailTripwireTriggered):
|
||||
await run_once(restored)
|
||||
assert "approved-secret" not in restored.to_string()
|
||||
else:
|
||||
resumed = await run_once(restored)
|
||||
assert resumed.final_output == "approved-secret"
|
||||
assert tool_calls == 1
|
||||
if session is not None:
|
||||
saved_outputs = [
|
||||
item
|
||||
for item in await session.get_items()
|
||||
if isinstance(item, dict) and item.get("type") == "function_call_output"
|
||||
]
|
||||
assert [(item["call_id"], item["output"]) for item in saved_outputs] == [
|
||||
(
|
||||
"call-approved",
|
||||
run_loop._OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT if tripwire else "approved-secret",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["non_streamed", "streamed"])
|
||||
@pytest.mark.parametrize(
|
||||
"corruption",
|
||||
[
|
||||
"missing-prefix",
|
||||
"nonterminal-indexes",
|
||||
"invalid-interruption-index",
|
||||
"processed-mismatch",
|
||||
"interruption-mismatch",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_ambiguous_serialized_approval_state_fails_before_tool_execution(
|
||||
mode: str,
|
||||
corruption: str,
|
||||
) -> None:
|
||||
tool_calls = {"normal": 0, "approval": 0}
|
||||
|
||||
@function_tool(name_override="normal_tool")
|
||||
def normal_tool() -> str:
|
||||
tool_calls["normal"] += 1
|
||||
return "normal-result"
|
||||
|
||||
@function_tool(name_override="approval_tool", needs_approval=True)
|
||||
def approval_tool() -> str:
|
||||
tool_calls["approval"] += 1
|
||||
return "secret-result"
|
||||
|
||||
def output_guardrail(
|
||||
@@ -2592,19 +2939,43 @@ async def test_ambiguous_serialized_approval_state_fails_before_tool_execution(
|
||||
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False)
|
||||
|
||||
model = ScriptedModel(
|
||||
[[get_function_tool_call("approval_tool", "{}", call_id="call-approved")]]
|
||||
[
|
||||
[get_function_tool_call("normal_tool", "{}", call_id="call-normal")],
|
||||
[
|
||||
get_function_tool_call("normal_tool", "{}", call_id="call-extra"),
|
||||
get_function_tool_call("approval_tool", "{}", call_id="call-approved"),
|
||||
],
|
||||
]
|
||||
)
|
||||
agent = Agent(
|
||||
name="test",
|
||||
model=model,
|
||||
tools=[approval_tool],
|
||||
tools=[normal_tool, approval_tool],
|
||||
output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail)],
|
||||
)
|
||||
first = await Runner.run(agent, "Use approval_tool")
|
||||
first = await Runner.run(agent, "Use normal_tool, then approval_tool")
|
||||
state = first.to_state()
|
||||
state._current_turn = 2
|
||||
state._current_turn_persisted_item_count = 1
|
||||
restored = await RunState.from_json(agent, state.to_json())
|
||||
assert state._current_turn == 2
|
||||
assert tool_calls == {"normal": 2, "approval": 0}
|
||||
serialized = state.to_json()
|
||||
ownership = serialized["current_response_generated_item_ownership"]
|
||||
assert ownership == {"start": 2, "end": 6, "interruptions": [5]}
|
||||
|
||||
if corruption == "missing-prefix":
|
||||
serialized["generated_items"] = serialized["generated_items"][2:]
|
||||
elif corruption == "nonterminal-indexes":
|
||||
trailing_item = json.loads(json.dumps(serialized["generated_items"][0]))
|
||||
trailing_item["raw_item"]["call_id"] = "call-trailing"
|
||||
serialized["generated_items"].append(trailing_item)
|
||||
elif corruption == "invalid-interruption-index":
|
||||
ownership["interruptions"] = [3]
|
||||
elif corruption == "processed-mismatch":
|
||||
serialized["generated_items"][ownership["start"]]["raw_item"]["call_id"] = "call-mismatch"
|
||||
elif corruption == "interruption-mismatch":
|
||||
serialized["generated_items"][ownership["interruptions"][0]]["raw_item"]["call_id"] = (
|
||||
"call-mismatch"
|
||||
)
|
||||
restored = await RunState.from_json(agent, serialized)
|
||||
restored.approve(restored.get_interruptions()[0])
|
||||
|
||||
with pytest.raises(UserError, match="current response boundary cannot be proven"):
|
||||
@@ -2614,7 +2985,7 @@ async def test_ambiguous_serialized_approval_state_fails_before_tool_execution(
|
||||
result = Runner.run_streamed(agent, restored, session=None)
|
||||
await consume_stream(result)
|
||||
|
||||
assert tool_calls == 0
|
||||
assert tool_calls == {"normal": 2, "approval": 0}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["non_streamed", "streamed"])
|
||||
|
||||
Reference in New Issue
Block a user