Skip to content

fix: retry resumed Session writes before model calls #4621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
AAA-413 wants to merge 1 commit into from
Closed

fix: retry resumed Session writes before model calls #4621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion .agents/skills/implementation-final-review/scripts/test_skill_contract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,10 @@ def test_quality_gates_cover_prior_failure_modes(self) -> None:
"contract-surface inventory",
"every consumer, forwarding branch, and adapter",
"Search adjacent contract surfaces even when they are absent from the diff",
"do not add it to the current task manifest, report it as a current-pull-request finding, or let it block clean review",
(
"do not add it to the current task manifest, report it as a "
"current-pull-request finding, or let it block clean review"
),
"await-boundary matrix",
"a newer operation that starts and completes while suspended",
"current active state is insufficient",
Expand Down
9 changes: 9 additions & 0 deletions src/agents/result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,15 @@ def _populate_state_from_result(
state._generated_prompt_cache_key = source_state._generated_prompt_cache_key
state._pending_input = copy.deepcopy(source_state._pending_input)
state._current_step = source_state._current_step
if source_state._current_step is not None:
state._current_agent = source_state._current_agent
state._current_turn_persisted_item_count = (
source_state._current_turn_persisted_item_count
)
if source_state._pending_session_items:
state._pending_session_items = list(source_state._pending_session_items)
state._pending_session_id = source_state._pending_session_id
state._pending_session_store = source_state._pending_session_store
else:
state._generated_prompt_cache_key = getattr(result, "_generated_prompt_cache_key", None)
state._pending_input = copy.deepcopy(getattr(result, "_pending_input_for_state", []))
Expand Down
59 changes: 59 additions & 0 deletions src/agents/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,7 @@
prepare_input_with_session,
reconcile_nested_history_owned_session_item_refs,
resumed_turn_items,
retry_pending_resumed_turn_session_items,
save_result_to_session,
save_resumed_turn_items,
session_items_for_turn,
Expand Down Expand Up @@ -1160,6 +1161,14 @@ def _mark_response_hooks_started() -> None:
),
store=store_setting,
wrapper=context_wrapper,
run_state=(
run_state
if isinstance(
turn_result.next_step,
NextStepRunAgain | NextStepHandoff,
)
else None
),
)
)

Expand Down Expand Up @@ -1293,6 +1302,8 @@ def _mark_response_hooks_started() -> None:
owner_starts=blocked_output_owner_starts,
blocked_message=blocked_message,
)
if run_state is not None:
run_state._current_step = NextStepRunAgain()
list.extend(session_items, retained_items)
try:
await save_final_turn_items_after_guardrails(
Expand All @@ -1308,6 +1319,7 @@ def _mark_response_hooks_started() -> None:
response_id=turn_result.model_response.response_id,
store=store_setting,
wrapper=context_wrapper,
track_pending_write=True,
)
except BaseException as persistence_error:
raise _safe_redacted_persistence_error(
Expand Down Expand Up @@ -1340,6 +1352,7 @@ def _mark_response_hooks_started() -> None:
response_id=turn_result.model_response.response_id,
store=store_setting,
wrapper=context_wrapper,
track_pending_write=True,
)
raise

Expand All @@ -1350,6 +1363,13 @@ def _mark_response_hooks_started() -> None:
current_agent,
run_config,
)
if run_state is not None:
run_state._current_step = NextStepFinalOutput(
turn_result.next_step.output
)
run_state._output_guardrail_results = list(
output_guardrail_results
)
await save_final_turn_items_after_guardrails(
session=session,
run_state=run_state,
Expand All @@ -1359,6 +1379,7 @@ def _mark_response_hooks_started() -> None:
response_id=turn_result.model_response.response_id,
store=store_setting,
wrapper=context_wrapper,
track_pending_write=True,
)
current_step = getattr(run_state, "_current_step", None)
approvals_from_state = approvals_from_step(current_step)
Expand Down Expand Up @@ -1409,6 +1430,44 @@ def _mark_response_hooks_started() -> None:
if run_state._current_step is None:
run_state._current_step = NextStepRunAgain()

await retry_pending_resumed_turn_session_items(
session=session,
run_state=run_state,
response_id=(
run_state._model_responses[-1].response_id
if run_state._model_responses
else None
),
wrapper=context_wrapper,
)

if isinstance(run_state._current_step, NextStepFinalOutput):
pending_final_output = run_state._current_step.output
run_state._current_step = None
result = RunResult(
input=original_input,
new_items=session_items,
raw_responses=model_responses,
final_output=pending_final_output,
_last_agent=current_agent,
input_guardrail_results=input_guardrail_results,
output_guardrail_results=output_guardrail_results,
tool_input_guardrail_results=tool_input_guardrail_results,
tool_output_guardrail_results=tool_output_guardrail_results,
context_wrapper=context_wrapper,
interruptions=[],
_tool_use_tracker_snapshot=_tool_use_tracker_snapshot(),
max_turns=max_turns,
)
result._current_turn = current_turn
result._model_input_items = list(generated_items)
result._replay_from_model_input_items = list(generated_items) != list(
session_items
)
result._trace_state = run_state._trace_state
result._original_input = copy_input_items(original_input)
return _finalize_result(result)

pending_input = run_state.pending_input
if pending_input:
pending_guardrails = current_agent.input_guardrails + (
Expand Down
19 changes: 14 additions & 5 deletions src/agents/run_internal/agent_runner_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -565,23 +565,32 @@ async def save_final_turn_items_after_guardrails(
reasoning_item_id_policy: ReasoningItemIdPolicy | None = None,
store: bool | None = None,
wrapper: RunContextWrapper[Any] | None = None,
track_pending_write: bool = False,
) -> int:
"""Persist deferred final-turn items without skipping a partially persisted resumed turn."""
if not session_persistence_enabled or not items:
return 0
if input_guardrails_triggered(input_guardrail_results):
return 0
if run_state is not None and run_state._current_turn_persisted_item_count > 0:
run_state._current_turn_persisted_item_count = await save_resumed_turn_items(
resumable_state = (
run_state
if track_pending_write
and run_state is not None
and isinstance(run_state._current_step, NextStepRunAgain | NextStepFinalOutput)
else None
)
if resumable_state is not None:
resumable_state._current_turn_persisted_item_count = await save_resumed_turn_items(
session=session,
items=items,
persisted_count=run_state._current_turn_persisted_item_count,
persisted_count=resumable_state._current_turn_persisted_item_count,
response_id=response_id,
reasoning_item_id_policy=run_state._reasoning_item_id_policy,
reasoning_item_id_policy=resumable_state._reasoning_item_id_policy,
store=store,
wrapper=wrapper,
run_state=resumable_state,
)
return run_state._current_turn_persisted_item_count
return resumable_state._current_turn_persisted_item_count
return await save_result_to_session(
session,
[],
Expand Down
37 changes: 37 additions & 0 deletions src/agents/run_internal/run_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@
prepare_input_with_session,
reconcile_nested_history_owned_session_item_refs,
resumed_turn_items,
retry_pending_resumed_turn_session_items,
rewind_session_items,
save_result_to_session,
save_resumed_turn_items,
Expand Down Expand Up @@ -399,6 +400,12 @@ async def _save_resumed_stream_items(
reasoning_item_id_policy=streamed_result._reasoning_item_id_policy,
store=store,
wrapper=streamed_result.context_wrapper,
run_state=(
run_state
if run_state is not None
and isinstance(run_state._current_step, NextStepRunAgain | NextStepFinalOutput)
else None
),
)
if run_state is not None:
run_state._current_turn_persisted_item_count = (
Expand Down Expand Up @@ -501,6 +508,7 @@ async def _finalize_streamed_final_output(
response_id: str | None,
store_setting: bool | None,
on_persisted_after_guardrails: Callable[[bool], None] | None = None,
track_pending_write: bool = False,
) -> None:
output_guardrail_result_start = len(streamed_result.output_guardrail_results)
redacted_persistence_error: BaseException | None = None
Expand Down Expand Up @@ -556,6 +564,8 @@ async def _finalize_streamed_final_output(
owner_starts=owner_starts,
blocked_message=blocked_message,
)
if track_pending_write and streamed_result._state is not None:
streamed_result._state._current_step = NextStepRunAgain()
if retained_items:
try:
await save_items(retained_items, response_id, store_setting)
Expand Down Expand Up @@ -622,6 +632,11 @@ async def _finalize_streamed_final_output(
agent,
run_config,
)
if track_pending_write and streamed_result._state is not None:
streamed_result._state._current_step = NextStepFinalOutput(output)
streamed_result._state._output_guardrail_results = list(
streamed_result.output_guardrail_results
)

# Saved as one ordered batch so the session mirrors the model response. Doing it in two
# halves would both reorder the turn and, because the first save advances the turn's
Expand Down Expand Up @@ -1429,6 +1444,7 @@ async def _save_max_turns_items(
owner_starts=blocked_output_owner_starts,
response_id=turn_result.model_response.response_id,
store_setting=store_setting,
track_pending_write=True,
)
if streamed_result._stored_exception is not None:
break
Expand Down Expand Up @@ -1458,6 +1474,27 @@ async def _save_max_turns_items(
if streamed_result.is_complete:
break

if run_state is not None:
streamed_result._current_turn_persisted_item_count = (
await retry_pending_resumed_turn_session_items(
session=session,
run_state=run_state,
response_id=(
run_state._model_responses[-1].response_id
if run_state._model_responses
else None
),
wrapper=streamed_result.context_wrapper,
)
)

if isinstance(run_state._current_step, NextStepFinalOutput):
streamed_result.final_output = run_state._current_step.output
run_state._current_step = None
streamed_result.is_complete = True
streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
break

if run_state is not None and run_state._pending_input:
if run_state._current_step is None:
run_state._current_step = NextStepRunAgain()
Expand Down
Loading