Handle partial answers consistently across agent workflows (#51249)

## Why

Partial answers contain answer fragments while a turn is still in progress. Treating them as final answers defers mailbox delivery and can count empty fragments as empty final responses; filtering them from forked histories and guardian transcripts loses answer content.

## What changed

- Allow mailbox preemption after `partial_answer` messages and keep mailbox delivery open.
- Preserve partial answers in forked agent histories and include them as protected assistant entries in guardian transcripts.
- Exclude empty partial answers from empty-final-response accounting for goals.

## Testing

Add regression coverage for partial-answer phases in live events, sampling continuations, follow-up requests, and forked child requests, plus mailbox preemption, goal accounting, and guardian transcript inclusion.

GitOrigin-RevId: 208c8933e211b42b6c2d0bec6952527b8c9c120e
This commit is contained in:
rb-oai
2026-10-06 03:27:04 +00:00
committed by copyberry
parent 8b6bb1c77d
commit 989c01a41a
16 changed files with 408 additions and 22 deletions
+4 -1
View File
@@ -89,7 +89,10 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_context_baselines: bool
RolloutItem::ResponseItem(envelope) => match &envelope.item {
ResponseItem::Message { role, phase, .. } => match role.as_str() {
"system" | "developer" | "user" => true,
"assistant" => *phase == Some(MessagePhase::FinalAnswer),
"assistant" => matches!(
phase,
Some(MessagePhase::PartialAnswer | MessagePhase::FinalAnswer)
),
_ => false,
},
ResponseItem::FunctionCallOutput { call_id: None, .. }
+18
View File
@@ -2449,6 +2449,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() {
),
},
assistant_message("parent commentary", Some(MessagePhase::Commentary)),
assistant_message("parent answer fragment", Some(MessagePhase::PartialAnswer)),
assistant_message("parent final answer", Some(MessagePhase::FinalAnswer)),
standalone_output,
assistant_message("parent unknown phase", /*phase*/ None),
@@ -2529,6 +2530,22 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() {
};
assert!(assistant_position(&parent_thread.session.clone_history().await).is_some());
assert_eq!(assistant_position(&history), None);
let expected_partial_answer = parent_thread
.session
.clone_history()
.await
.raw_items()
.find(|item| {
matches!(
item,
ResponseItem::Message {
phase: Some(MessagePhase::PartialAnswer),
..
}
)
})
.cloned()
.expect("parent answer fragment should be recorded");
let expected_final_answer = parent_thread
.session
.clone_history()
@@ -2587,6 +2604,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() {
let expected_history = [
expected_parent_seed,
expected_developer_message,
expected_partial_answer,
expected_final_answer,
expected_standalone_output,
ContextualUserFragment::into(MultiAgentRoleInstructions::Configured(
+5 -3
View File
@@ -2750,11 +2750,13 @@ async fn try_run_sampling_request(
cancellation_token: cancellation_token.child_token(),
};
// TODO: Allow mailbox preemption after PartialAnswer and keep delivery open
// in the completed-item deferral predicates.
let preempt_for_mailbox_mail = match &item {
ResponseItem::Message { role, phase, .. } => {
role == "assistant" && matches!(phase, Some(MessagePhase::Commentary))
role == "assistant"
&& matches!(
phase,
Some(MessagePhase::Commentary | MessagePhase::PartialAnswer)
)
}
ResponseItem::Reasoning { .. } => true,
ResponseItem::AgentMessage { .. } => false,
+2 -2
View File
@@ -291,7 +291,7 @@ pub(crate) async fn finalize_non_tool_response_item(
Some(combined)
};
let defers_mailbox_delivery_to_next_turn =
!matches!(agent_message.phase, Some(MessagePhase::Commentary))
matches!(agent_message.phase, Some(MessagePhase::FinalAnswer) | None)
&& last_agent_message.is_some();
(
agent_message.memory_citation.clone(),
@@ -532,7 +532,7 @@ fn completed_item_defers_mailbox_delivery_to_next_turn(
) -> bool {
match item {
ResponseItem::Message { role, phase, .. } => {
if role != "assistant" || matches!(phase, Some(MessagePhase::Commentary)) {
if role != "assistant" || !matches!(phase, Some(MessagePhase::FinalAnswer) | None) {
return false;
}
// Treat `None` like final-answer text so untagged providers default
@@ -479,14 +479,16 @@ async fn finalized_turn_item_defers_mailbox_for_contributed_visible_text() {
assert!(finalized.facts.defers_mailbox_delivery_to_next_turn);
}
#[test_case::test_case(MessagePhase::Commentary; "commentary")]
#[test_case::test_case(MessagePhase::PartialAnswer; "partial_answer")]
#[tokio::test]
async fn finalized_turn_item_keeps_mailbox_open_for_commentary_text() {
async fn finalized_turn_item_keeps_mailbox_open_for_nonterminal_text(phase: MessagePhase) {
let (mut session, turn_context) = make_session_and_context().await;
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor));
session.services.extensions = Arc::new(builder.build());
let turn_store = ExtensionData::new(turn_context.sub_id.clone());
let item = assistant_output_text_with_phase("still working", Some(MessagePhase::Commentary));
let item = assistant_output_text_with_phase("still working", Some(phase));
let finalized = finalize_non_tool_response_item(
&session,
@@ -545,9 +547,10 @@ fn completed_item_defers_mailbox_delivery_for_unknown_phase_messages() {
));
}
#[test]
fn completed_item_keeps_mailbox_delivery_open_for_commentary_messages() {
let item = assistant_output_text_with_phase("still working", Some(MessagePhase::Commentary));
#[test_case::test_case(MessagePhase::Commentary; "commentary")]
#[test_case::test_case(MessagePhase::PartialAnswer; "partial_answer")]
fn completed_item_keeps_mailbox_delivery_open_for_nonterminal_messages(phase: MessagePhase) {
let item = assistant_output_text_with_phase("still working", Some(phase));
assert!(!completed_item_defers_mailbox_delivery_to_next_turn(
&item, /*plan_mode*/ false,
+3
View File
@@ -104,6 +104,9 @@ mod agent_message_board;
#[path = "scenarios_mailbox_preemption_tests.rs"]
mod mailbox_preemption;
#[path = "scenarios_partial_answers.rs"]
mod partial_answers;
#[path = "scenarios_guardian_extra_policy.rs"]
mod guardian_extra_policy;
@@ -25,13 +25,16 @@ use tokio::sync::oneshot;
enum Boundary {
Reasoning,
Commentary,
PartialAnswer,
}
#[tracing_test::traced_test]
#[test_case(Boundary::Reasoning, false; "reasoning_default")]
#[test_case(Boundary::Commentary, false; "commentary_default")]
#[test_case(Boundary::PartialAnswer, false; "partial_answer_default")]
#[test_case(Boundary::Reasoning, true; "reasoning_deferred")]
#[test_case(Boundary::Commentary, true; "commentary_deferred")]
#[test_case(Boundary::PartialAnswer, true; "partial_answer_deferred")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn mailbox_preemption_preserves_response_when_deferred(
boundary: Boundary,
@@ -44,14 +47,17 @@ async fn mailbox_preemption_preserves_response_when_deferred(
responses::ev_reasoning_item("boundary", &["Preparing the next action"], &[]),
"reasoning",
),
Boundary::Commentary => {
Boundary::Commentary | Boundary::PartialAnswer => {
let phase = match boundary {
Boundary::Commentary => "commentary",
Boundary::PartialAnswer => "partial_answer",
Boundary::Reasoning => unreachable!(),
};
let mut done = responses::ev_assistant_message("boundary", "I will update the plan.");
done["item"]["phase"] = json!("commentary");
(
responses::ev_message_item_added("boundary", ""),
done,
"commentary",
)
done["item"]["phase"] = json!(phase);
let mut added = responses::ev_message_item_added("boundary", "");
added["item"]["phase"] = json!(phase);
(added, done, phase)
}
};
let (streaming, _completions) = start_streaming_sse_server(vec![
@@ -0,0 +1,216 @@
//! Partial answers retain their phase in live events and subsequent model requests.
use super::*;
use codex_protocol::models::MessagePhase;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn partial_answers_preserve_phase_across_sampling_continuation() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let mut added = ev_message_item_added("fragment", "");
added["item"]["phase"] = json!("partial_answer");
let mut partial = ev_assistant_message("fragment", "The first result is ready.");
partial["item"]["phase"] = json!("partial_answer");
let mut continuing = ev_completed("partial-response");
continuing["response"]["end_turn"] = json!(false);
let mut commentary = ev_assistant_message("progress", "Checking the remaining result.");
commentary["item"]["phase"] = json!("commentary");
let mut final_answer = ev_assistant_message("answer", "The second result is ready.");
final_answer["item"]["phase"] = json!("final_answer");
let mut completed = ev_completed("final-response");
completed["response"]["end_turn"] = json!(true);
let mock = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("partial-response"),
added,
ev_output_text_delta("The first result is ready."),
partial,
continuing,
]),
sse(vec![
ev_response_created("final-response"),
commentary,
final_answer,
completed,
]),
sse(vec![
ev_response_created("follow-up-response"),
ev_assistant_message("follow-up", "Both results are still in context."),
ev_completed("follow-up-response"),
]),
],
)
.await;
let test = test_codex()
.with_config(configure_scenario_catalog)
.build_with_auto_env(&server)
.await?;
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![text(
"Give me both results.",
)]))
.await?;
let started = wait_for_event_match(&test.codex, |event| match event {
EventMsg::ItemStarted(event) => match &event.item {
TurnItem::AgentMessage(item) => Some((item.id.clone(), item.phase.clone())),
_ => None,
},
_ => None,
})
.await;
assert_eq!(
started,
("fragment".to_string(), Some(MessagePhase::PartialAnswer))
);
wait_for_event(&test.codex, |event| {
assert!(!matches!(event, EventMsg::Error(_)), "{event:?}");
matches!(event, EventMsg::TurnComplete(_))
})
.await;
test.submit_text_turn("Summarize the results.").await?;
let requests = mock.requests();
assert_eq!(requests.len(), 3);
let phases = |index: usize| {
requests[index].body_json()["input"]
.as_array()
.expect("request input")
.iter()
.filter(|item| item["type"] == "message" && item["role"] == "assistant")
.map(|item| item["phase"].clone())
.collect::<Vec<_>>()
};
assert_eq!(phases(/*index*/ 1), vec![json!("partial_answer")]);
assert_eq!(
phases(/*index*/ 2),
vec![
json!("partial_answer"),
json!("commentary"),
json!("final_answer")
]
);
insta::assert_snapshot!(
"partial_answer_continuation",
context_snapshot::format_request_history_snapshot(
"A nonterminal answer continues sampling; its phase survives an ordinary follow-up.",
&requests,
&ContextSnapshotOptions::default().rewrite_known_segments(),
)
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn partial_answers_survive_forked_child_requests() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let mut partial = ev_assistant_message("fragment", "The first result is ready.");
partial["item"]["phase"] = json!("partial_answer");
let mut commentary = ev_assistant_message("progress", "Checking the remaining result.");
commentary["item"]["phase"] = json!("commentary");
let mut final_answer = ev_assistant_message("answer", "The second result is ready.");
final_answer["item"]["phase"] = json!("final_answer");
let done = sse(vec![
ev_assistant_message("done", "Done."),
ev_completed("done"),
]);
let mock = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("results"),
partial,
commentary,
final_answer,
ev_completed("results"),
]),
sse(vec![
ev_function_call_with_namespace(
"spawn-worker",
"collaboration",
"spawn_agent",
&json!({
"task_name": "worker",
"message": "Check both inherited results.",
"fork_turns": "all",
})
.to_string(),
),
ev_completed("spawn"),
]),
// Parent continuation and child inference may arrive in either order.
done.clone(),
done,
],
)
.await;
let test = test_codex()
.with_config(|config| {
configure_scenario_catalog(config);
config
.features
.enable(Feature::Collab)
.expect("enable agents");
config
.features
.enable(Feature::MultiAgentV2)
.expect("enable agent paths");
})
.build_with_auto_env(&server)
.await?;
test.submit_turn("Give me both results.").await?;
test.submit_turn("Have a worker check both results.")
.await?;
let child_id = test
.thread_manager
.list_thread_ids()
.await
.into_iter()
.find(|id| *id != test.session_configured.thread_id)
.expect("forked child");
let child = test.thread_manager.get_thread(child_id).await?;
wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await;
let child_requests = mock
.requests()
.into_iter()
.filter(|request| {
request.body_json()["client_metadata"]["thread_id"] == child_id.to_string()
})
.collect::<Vec<_>>();
assert_eq!(child_requests.len(), 1);
let body = child_requests[0].body_json();
let assistant_messages = body["input"]
.as_array()
.expect("child request input")
.iter()
.filter(|item| item["type"] == "message" && item["role"] == "assistant")
.map(|item| (item["phase"].clone(), item["content"].clone()))
.collect::<Vec<_>>();
assert_eq!(
assistant_messages,
vec![
(
json!("partial_answer"),
json!([{"type":"output_text", "text":"The first result is ready."}])
),
(
json!("final_answer"),
json!([{"type":"output_text", "text":"The second result is ready."}])
),
]
);
insta::assert_snapshot!(
"partial_answer_fork",
context_snapshot::format_request_history_snapshot(
"A forked worker receives both stable answer fragments, with their phases preserved and parent commentary filtered out.",
&child_requests,
&ContextSnapshotOptions::default().rewrite_known_segments(),
)
);
child.shutdown_and_wait().await?;
Ok(())
}
@@ -0,0 +1,21 @@
---
source: core/tests/suite/scenarios_mailbox_preemption_tests.rs
expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)"
---
Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.
## Window 1
-- request 1 (request) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Plan the next action while the worker finishes.
-- request 2 (request) --
04:message/assistant:
I will update the plan.
05:agent_message
06:agent_message
@@ -0,0 +1,23 @@
---
source: core/tests/suite/scenarios_mailbox_preemption_tests.rs
expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)"
---
Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.
## Window 1
-- request 1 (request) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Plan the next action while the worker finishes.
-- request 2 (request) --
04:message/assistant:
I will update the plan.
05:function_call/update_plan:{"plan":[{"status":"in_progress","step":"Reply to the user"}]}
06:function_call_output:Plan updated
07:agent_message
08:agent_message
@@ -0,0 +1,26 @@
---
source: core/tests/suite/scenarios_partial_answers.rs
expression: "context_snapshot::format_request_history_snapshot(\"A nonterminal answer continues sampling; its phase survives an ordinary follow-up.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)"
---
Scenario: A nonterminal answer continues sampling; its phase survives an ordinary follow-up.
## Window 1
-- request 1 (turn) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Give me both results.
-- request 2 (turn) --
04:message/assistant:
The first result is ready.
-- request 3 (turn) --
05:message/assistant:
Checking the remaining result.
06:message/assistant:
The second result is ready.
07:message/user:
Summarize the results.
@@ -0,0 +1,43 @@
---
source: core/tests/suite/scenarios_partial_answers.rs
expression: "context_snapshot::format_request_history_snapshot(\"A forked worker receives both stable answer fragments, with their phases preserved and parent commentary filtered out.\",\n&child_requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)"
---
Scenario: A forked worker receives both stable answer fragments, with their phases preserved and parent commentary filtered out.
## Window 1
-- request 1 (turn) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Give me both results.
04:message/assistant:
The first result is ready.
05:message/assistant:
The second result is ready.
06:message/user:
Have a worker check both results.
07:message/developer:
You are an agent in a team of agents collaborating to complete a task.
You can spawn sub-agents to handle subtasks, and those sub-agents can spawn their own sub-agents. All ag...nd capable, and have access to the same set of tools. [hash=876a93c54d3f61d6]
You can use `spawn_agent` to create a new agent, `followup_task` to give an existing agent a new task an... `send_message` to pass a message to a running agent. [hash=824c12425c7626dc]
Child agents can also spawn their own sub-agents.
When you provide a response in the final channel, that content is immediately delivered back to your parent agent.
<OMITTED 15 LINES; ~225 TOKENS; hash=fa366829bdb214be>
- All agents use the same current working directory.
- As a result, edits made by one agent are immediately visible to all other agents.
When calling `wait_agent`, prefer longer waits (minutes) to avoid busy polling.
There are 4 available concurrency slots, meaning that up to 4 agents can be active at once, including you.
Full-history forks (`fork_turns` omitted or `"all"`) inherit the parent model and reasoning effort and d...fork_turns` to `"none"` or a positive integer string. [hash=db3545796947884f]
08:message/developer:
<MULTI_AGENT_MODE>
09:agent_message
+1 -1
View File
@@ -172,7 +172,7 @@ impl GoalAccountingState {
});
turn.has_activity |= has_text || message.questions.is_some();
turn.empty_final |=
!has_text && !matches!(message.phase, Some(MessagePhase::Commentary));
!has_text && matches!(message.phase, Some(MessagePhase::FinalAnswer) | None);
}
TurnItem::Reasoning(reasoning) => {
turn.has_activity |= reasoning
+8 -1
View File
@@ -76,6 +76,7 @@ fn empty_continuations_require_three_turns_without_activity_or_goal_changes() {
("goal", 5),
("reset", 6),
("missing final", 6),
("partial answer", 6),
] {
let state = GoalAccountingState::default();
for turn in 1..=blocking_turn {
@@ -87,7 +88,13 @@ fn empty_continuations_require_three_turns_without_activity_or_goal_changes() {
};
state.start_turn(&id, ModeKind::Default, &TokenUsage::default());
state.mark_turn_goal_active(&id, goal_id);
if !(turn == 3 && interruption == "missing final") {
if turn == 3 && interruption == "partial answer" {
let mut partial = empty_final.clone();
if let TurnItem::AgentMessage(message) = &mut partial {
message.phase = Some(MessagePhase::PartialAnswer);
}
state.record_item(&id, &partial);
} else if !(turn == 3 && interruption == "missing final") {
state.record_item(&id, &empty_final);
}
// Admission can finish after output arrives, but before turn-stop evaluation.
+4 -2
View File
@@ -155,8 +155,10 @@ pub fn collect_transcript(
ConversationTranscriptEntryKind::Developer
}
"assistant"
if matches!(phase, None | Some(MessagePhase::FinalAnswer))
&& !InterAgentCommunication::is_message_content(content) =>
if matches!(
phase,
None | Some(MessagePhase::PartialAnswer | MessagePhase::FinalAnswer)
) && !InterAgentCommunication::is_message_content(content) =>
{
ConversationTranscriptEntryKind::ProtectedAssistant
}
@@ -131,6 +131,15 @@ fn registered_transcript_filters_roles_and_preserves_node_repl_tool_attribution(
phase: None,
internal_chat_message_metadata_passthrough: None,
},
ResponseItem::Message {
id: None,
role: "assistant".to_string(),
content: vec![ContentItem::OutputText {
text: "The first result is ready.".to_string(),
}],
phase: Some(MessagePhase::PartialAnswer),
internal_chat_message_metadata_passthrough: None,
},
ResponseItem::Message {
id: None,
role: "assistant".to_string(),
@@ -187,6 +196,10 @@ fn registered_transcript_filters_roles_and_preserves_node_repl_tool_attribution(
"Inspect the workspace."
),
entry(ConversationTranscriptEntryKind::Developer, &approved_action),
entry(
ConversationTranscriptEntryKind::ProtectedAssistant,
"The first result is ready."
),
entry(
ConversationTranscriptEntryKind::ProtectedAssistant,
"Inspection complete."