agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: adding stream parser (#12666)

Browse Source 
Add a stream parser to extract citations (and others) from a stream.
This support cases where markers are split in differen tokens.

Codex never manage to make this code work so everything was done
manually. Please review correctly and do not touch this part of the code
without a very clear understanding of it
This commit is contained in:
jif-oai
2026-02-25 13:27:58 +00:00
committed by GitHub
Unverified
parent 5a9a5b51b2
commit 5441130e0a
20 changed files with 2070 additions and 371 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex.rs
+253 -97
View File
@@ -42,6 +42,7 @@ use crate::stream_events_utils::HandleOutputCtx;
use crate::stream_events_utils::handle_non_tool_response_item;
use crate::stream_events_utils::handle_output_item_done;
use crate::stream_events_utils::last_assistant_message_from_item;
use crate::stream_events_utils::raw_assistant_output_text_from_item;
use crate::terminal;
use crate::truncate::TruncationPolicy;
use crate::turn_metadata::TurnMetadataState;
@@ -92,6 +93,11 @@ use codex_protocol::request_user_input::RequestUserInputResponse;
use codex_protocol::skill_approval::SkillApprovalResponse;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::OAuthCredentialsStoreMode;
use codex_utils_stream_parser::AssistantTextChunk;
use codex_utils_stream_parser::AssistantTextStreamParser;
use codex_utils_stream_parser::ProposedPlanSegment;
use codex_utils_stream_parser::extract_proposed_plan_text;
use codex_utils_stream_parser::strip_citations;
use futures::future::BoxFuture;
use futures::prelude::*;
use futures::stream::FuturesOrdered;
@@ -174,9 +180,6 @@ use crate::mentions::collect_explicit_app_ids;
use crate::mentions::collect_tool_mentions_from_messages;
use crate::network_policy_decision::execpolicy_network_rule_amendment;
use crate::project_doc::get_user_instructions;
use crate::proposed_plan_parser::ProposedPlanParser;
use crate::proposed_plan_parser::ProposedPlanSegment;
use crate::proposed_plan_parser::extract_proposed_plan_text;
use crate::protocol::AgentMessageContentDeltaEvent;
use crate::protocol::AgentReasoningSectionBreakEvent;
use crate::protocol::ApplyPatchApprovalRequestEvent;
@@ -5618,39 +5621,9 @@ struct ProposedPlanItemState {
completed: bool,
}
/// Per-item plan parsers so we can buffer text while detecting `<proposed_plan>`
/// tags without ever mixing buffered lines across item ids.
struct PlanParsers {
assistant: HashMap<String, ProposedPlanParser>,
}
impl PlanParsers {
fn new() -> Self {
Self {
assistant: HashMap::new(),
}
}
fn assistant_parser_mut(&mut self, item_id: &str) -> &mut ProposedPlanParser {
self.assistant
.entry(item_id.to_string())
.or_insert_with(ProposedPlanParser::new)
}
fn take_assistant_parser(&mut self, item_id: &str) -> Option<ProposedPlanParser> {
self.assistant.remove(item_id)
}
fn drain_assistant_parsers(&mut self) -> Vec<(String, ProposedPlanParser)> {
self.assistant.drain().collect()
}
}
/// Aggregated state used only while streaming a plan-mode response.
/// Includes per-item parsers, deferred agent message bookkeeping, and the plan item lifecycle.
struct PlanModeStreamState {
/// Per-item parsers for assistant streams in plan mode.
plan_parsers: PlanParsers,
/// Agent message items started by the model but deferred until we see non-plan text.
pending_agent_message_items: HashMap<String, TurnItem>,
/// Agent message items whose start notification has been emitted.
@@ -5664,7 +5637,6 @@ struct PlanModeStreamState {
impl PlanModeStreamState {
fn new(turn_id: &str) -> Self {
Self {
plan_parsers: PlanParsers::new(),
pending_agent_message_items: HashMap::new(),
started_agent_message_items: HashSet::new(),
leading_whitespace_by_item: HashMap::new(),
@@ -5673,6 +5645,56 @@ impl PlanModeStreamState {
}
}
#[derive(Debug, Default)]
struct AssistantMessageStreamParsers {
plan_mode: bool,
parsers_by_item: HashMap<String, AssistantTextStreamParser>,
}
type ParsedAssistantTextDelta = AssistantTextChunk;
impl AssistantMessageStreamParsers {
fn new(plan_mode: bool) -> Self {
Self {
plan_mode,
parsers_by_item: HashMap::new(),
}
}
fn parser_mut(&mut self, item_id: &str) -> &mut AssistantTextStreamParser {
let plan_mode = self.plan_mode;
self.parsers_by_item
.entry(item_id.to_string())
.or_insert_with(|| AssistantTextStreamParser::new(plan_mode))
}
fn seed_item_text(&mut self, item_id: &str, text: &str) -> ParsedAssistantTextDelta {
if text.is_empty() {
return ParsedAssistantTextDelta::default();
}
self.parser_mut(item_id).push_str(text)
}
fn parse_delta(&mut self, item_id: &str, delta: &str) -> ParsedAssistantTextDelta {
self.parser_mut(item_id).push_str(delta)
}
fn finish_item(&mut self, item_id: &str) -> ParsedAssistantTextDelta {
let Some(mut parser) = self.parsers_by_item.remove(item_id) else {
return ParsedAssistantTextDelta::default();
};
parser.finish()
}
fn drain_finished(&mut self) -> Vec<(String, ParsedAssistantTextDelta)> {
let parsers_by_item = std::mem::take(&mut self.parsers_by_item);
parsers_by_item
.into_iter()
.map(|(item_id, mut parser)| (item_id, parser.finish()))
.collect()
}
}
impl ProposedPlanItemState {
fn new(turn_id: &str) -> Self {
Self {
@@ -5875,35 +5897,68 @@ async fn handle_plan_segments(
}
}
/// Flush any buffered proposed-plan segments when a specific assistant message ends.
async fn flush_proposed_plan_segments_for_item(
async fn emit_streamed_assistant_text_delta(
sess: &Session,
turn_context: &TurnContext,
state: &mut PlanModeStreamState,
plan_mode_state: Option<&mut PlanModeStreamState>,
item_id: &str,
parsed: ParsedAssistantTextDelta,
) {
let Some(mut parser) = state.plan_parsers.take_assistant_parser(item_id) else {
return;
};
let segments = parser.finish();
if segments.is_empty() {
if parsed.is_empty() {
return;
}
handle_plan_segments(sess, turn_context, state, item_id, segments).await;
if !parsed.citations.is_empty() {
// Citation extraction is intentionally local for now; we strip citations from display text
// but do not yet surface them in protocol events.
let _citations = parsed.citations;
}
if let Some(state) = plan_mode_state {
if !parsed.plan_segments.is_empty() {
handle_plan_segments(sess, turn_context, state, item_id, parsed.plan_segments).await;
}
return;
}
if parsed.visible_text.is_empty() {
return;
}
let event = AgentMessageContentDeltaEvent {
thread_id: sess.conversation_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: item_id.to_string(),
delta: parsed.visible_text,
};
sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event))
.await;
}
/// Flush any remaining assistant plan parsers when the response completes.
async fn flush_proposed_plan_segments_all(
/// Flush buffered assistant text parser state when an assistant message item ends.
async fn flush_assistant_text_segments_for_item(
sess: &Session,
turn_context: &TurnContext,
state: &mut PlanModeStreamState,
plan_mode_state: Option<&mut PlanModeStreamState>,
parsers: &mut AssistantMessageStreamParsers,
item_id: &str,
) {
for (item_id, mut parser) in state.plan_parsers.drain_assistant_parsers() {
let segments = parser.finish();
if segments.is_empty() {
continue;
}
handle_plan_segments(sess, turn_context, state, &item_id, segments).await;
let parsed = parsers.finish_item(item_id);
emit_streamed_assistant_text_delta(sess, turn_context, plan_mode_state, item_id, parsed).await;
}
/// Flush any remaining buffered assistant text parser state at response completion.
async fn flush_assistant_text_segments_all(
sess: &Session,
turn_context: &TurnContext,
mut plan_mode_state: Option<&mut PlanModeStreamState>,
parsers: &mut AssistantMessageStreamParsers,
) {
for (item_id, parsed) in parsers.drain_finished() {
emit_streamed_assistant_text_delta(
sess,
turn_context,
plan_mode_state.as_deref_mut(),
&item_id,
parsed,
)
.await;
}
}
@@ -5924,6 +5979,7 @@ async fn maybe_complete_plan_item_from_message(
}
}
if let Some(plan_text) = extract_proposed_plan_text(&text) {
let (plan_text, _citations) = strip_citations(&plan_text);
if !state.plan_item_state.started {
state.plan_item_state.start(sess, turn_context).await;
}
@@ -6112,6 +6168,7 @@ async fn try_run_sampling_request(
let mut active_item: Option<TurnItem> = None;
let mut should_emit_turn_diff = false;
let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan;
let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode);
let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id));
let receiving_span = trace_span!("receiving_stream");
let outcome: CodexResult<SamplingRequestResult> = loop {
@@ -6151,20 +6208,21 @@ async fn try_run_sampling_request(
ResponseEvent::Created => {}
ResponseEvent::OutputItemDone(item) => {
let previously_active_item = active_item.take();
if let Some(state) = plan_mode_state.as_mut() {
if let Some(previous) = previously_active_item.as_ref() {
let item_id = previous.id();
if matches!(previous, TurnItem::AgentMessage(_)) {
flush_proposed_plan_segments_for_item(
&sess,
&turn_context,
state,
&item_id,
)
.await;
}
}
if handle_assistant_item_done_in_plan_mode(
if let Some(previous) = previously_active_item.as_ref()
&& matches!(previous, TurnItem::AgentMessage(_))
{
let item_id = previous.id();
flush_assistant_text_segments_for_item(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
&item_id,
)
.await;
}
if let Some(state) = plan_mode_state.as_mut()
&& handle_assistant_item_done_in_plan_mode(
&sess,
&turn_context,
&item,
@@ -6173,9 +6231,8 @@ async fn try_run_sampling_request(
&mut last_agent_message,
)
.await
{
continue;
}
{
continue;
}
let mut ctx = HandleOutputCtx {
@@ -6198,6 +6255,28 @@ async fn try_run_sampling_request(
}
ResponseEvent::OutputItemAdded(item) => {
if let Some(turn_item) = handle_non_tool_response_item(&item, plan_mode).await {
let mut turn_item = turn_item;
let mut seeded_parsed: Option<ParsedAssistantTextDelta> = None;
let mut seeded_item_id: Option<String> = None;
if matches!(turn_item, TurnItem::AgentMessage(_))
&& let Some(raw_text) = raw_assistant_output_text_from_item(&item)
{
let item_id = turn_item.id();
let mut seeded =
assistant_message_stream_parsers.seed_item_text(&item_id, &raw_text);
if let TurnItem::AgentMessage(agent_message) = &mut turn_item {
agent_message.content =
vec![codex_protocol::items::AgentMessageContent::Text {
text: if plan_mode {
String::new()
} else {
std::mem::take(&mut seeded.visible_text)
},
}];
}
seeded_parsed = plan_mode.then_some(seeded);
seeded_item_id = Some(item_id);
}
if let Some(state) = plan_mode_state.as_mut()
&& matches!(turn_item, TurnItem::AgentMessage(_))
{
@@ -6208,6 +6287,20 @@ async fn try_run_sampling_request(
} else {
sess.emit_turn_item_started(&turn_context, &turn_item).await;
}
if let (Some(state), Some(item_id), Some(parsed)) = (
plan_mode_state.as_mut(),
seeded_item_id.as_deref(),
seeded_parsed,
) {
emit_streamed_assistant_text_delta(
&sess,
&turn_context,
Some(state),
item_id,
parsed,
)
.await;
}
active_item = Some(turn_item);
}
}
@@ -6237,9 +6330,13 @@ async fn try_run_sampling_request(
token_usage,
can_append: _,
} => {
if let Some(state) = plan_mode_state.as_mut() {
flush_proposed_plan_segments_all(&sess, &turn_context, state).await;
}
flush_assistant_text_segments_all(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
)
.await;
sess.update_token_usage_info(&turn_context, token_usage.as_ref())
.await;
should_emit_turn_diff = true;
@@ -6256,14 +6353,16 @@ async fn try_run_sampling_request(
// UI will show a selection popup from the final ReviewOutput.
if let Some(active) = active_item.as_ref() {
let item_id = active.id();
if let Some(state) = plan_mode_state.as_mut()
&& matches!(active, TurnItem::AgentMessage(_))
{
let segments = state
.plan_parsers
.assistant_parser_mut(&item_id)
.parse(&delta);
handle_plan_segments(&sess, &turn_context, state, &item_id, segments).await;
if matches!(active, TurnItem::AgentMessage(_)) {
let parsed = assistant_message_stream_parsers.parse_delta(&item_id, &delta);
emit_streamed_assistant_text_delta(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&item_id,
parsed,
)
.await;
} else {
let event = AgentMessageContentDeltaEvent {
thread_id: sess.conversation_id.to_string(),
@@ -6329,6 +6428,14 @@ async fn try_run_sampling_request(
}
};
flush_assistant_text_segments_all(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
)
.await;
drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?;
if should_emit_turn_diff {
@@ -6346,23 +6453,10 @@ async fn try_run_sampling_request(
}
pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<String> {
responses.iter().rev().find_map(|item| {
if let ResponseItem::Message { role, content, .. } = item {
if role == "assistant" {
content.iter().rev().find_map(|ci| {
if let ContentItem::OutputText { text } = ci {
Some(text.clone())
} else {
None
}
})
} else {
None
}
} else {
None
}
})
responses
.iter()
.rev()
.find_map(|item| last_assistant_message_from_item(item, false))
}
use crate::memories::prompts::build_memory_tool_developer_instructions;
@@ -6486,6 +6580,68 @@ mod tests {
}
}
#[test]
fn assistant_message_stream_parsers_can_be_seeded_from_output_item_added_text() {
let mut parsers = AssistantMessageStreamParsers::new(false);
let item_id = "msg-1";
let seeded = parsers.seed_item_text(item_id, "hello <oai-mem-citation>doc");
let parsed = parsers.parse_delta(item_id, "1</oai-mem-citation> world");
let tail = parsers.finish_item(item_id);
assert_eq!(seeded.visible_text, "hello ");
assert_eq!(seeded.citations, Vec::<String>::new());
assert_eq!(parsed.visible_text, " world");
assert_eq!(parsed.citations, vec!["doc1".to_string()]);
assert_eq!(tail.visible_text, "");
assert_eq!(tail.citations, Vec::<String>::new());
}
#[test]
fn assistant_message_stream_parsers_seed_buffered_prefix_stays_out_of_finish_tail() {
let mut parsers = AssistantMessageStreamParsers::new(false);
let item_id = "msg-1";
let seeded = parsers.seed_item_text(item_id, "hello <oai-mem-");
let parsed = parsers.parse_delta(item_id, "citation>doc</oai-mem-citation> world");
let tail = parsers.finish_item(item_id);
assert_eq!(seeded.visible_text, "hello ");
assert_eq!(seeded.citations, Vec::<String>::new());
assert_eq!(parsed.visible_text, " world");
assert_eq!(parsed.citations, vec!["doc".to_string()]);
assert_eq!(tail.visible_text, "");
assert_eq!(tail.citations, Vec::<String>::new());
}
#[test]
fn assistant_message_stream_parsers_seed_plan_parser_across_added_and_delta_boundaries() {
let mut parsers = AssistantMessageStreamParsers::new(true);
let item_id = "msg-1";
let seeded = parsers.seed_item_text(item_id, "Intro\n<proposed");
let parsed = parsers.parse_delta(item_id, "_plan>\n- step\n</proposed_plan>\nOutro");
let tail = parsers.finish_item(item_id);
assert_eq!(seeded.visible_text, "Intro\n");
assert_eq!(
seeded.plan_segments,
vec![ProposedPlanSegment::Normal("Intro\n".to_string())]
);
assert_eq!(parsed.visible_text, "Outro");
assert_eq!(
parsed.plan_segments,
vec![
ProposedPlanSegment::ProposedPlanStart,
ProposedPlanSegment::ProposedPlanDelta("- step\n".to_string()),
ProposedPlanSegment::ProposedPlanEnd,
ProposedPlanSegment::Normal("Outro".to_string()),
]
);
assert_eq!(tail.visible_text, "");
assert!(tail.plan_segments.is_empty());
}
fn make_mcp_tool(
server_name: &str,
tool_name: &str,