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
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/Cargo.toml
+1
View File
@@ -50,6 +50,7 @@ codex-utils-pty = { workspace = true }
codex-utils-readiness = { workspace = true }
codex-secrets = { workspace = true }
codex-utils-string = { workspace = true }
codex-utils-stream-parser = { workspace = true }
codex-windows-sandbox = { package = "codex-windows-sandbox", path = "../windows-sandbox-rs" }
csv = { workspace = true }
dirs = { workspace = true }
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,
codex-rs/core/src/lib.rs
-2
View File
@@ -56,13 +56,11 @@ mod message_history;
mod model_provider_info;
pub mod path_utils;
pub mod personality_migration;
mod proposed_plan_parser;
mod sandbox_tags;
pub mod sandboxing;
mod session_prefix;
mod shell_detect;
mod stream_events_utils;
mod tagged_block_parser;
pub mod test_support;
mod text_encoding;
pub mod token_data;
codex-rs/core/src/proposed_plan_parser.rs
-185
View File
@@ -1,185 +0,0 @@
use crate::tagged_block_parser::TagSpec;
use crate::tagged_block_parser::TaggedLineParser;
use crate::tagged_block_parser::TaggedLineSegment;
const OPEN_TAG: &str = "<proposed_plan>";
const CLOSE_TAG: &str = "</proposed_plan>";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PlanTag {
ProposedPlan,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ProposedPlanSegment {
Normal(String),
ProposedPlanStart,
ProposedPlanDelta(String),
ProposedPlanEnd,
}
/// Parser for `<proposed_plan>` blocks emitted in plan mode.
///
/// This is a thin wrapper around the generic line-based tag parser. It maps
/// tag-aware segments into plan-specific segments for downstream consumers.
#[derive(Debug)]
pub(crate) struct ProposedPlanParser {
parser: TaggedLineParser<PlanTag>,
}
impl ProposedPlanParser {
pub(crate) fn new() -> Self {
Self {
parser: TaggedLineParser::new(vec![TagSpec {
open: OPEN_TAG,
close: CLOSE_TAG,
tag: PlanTag::ProposedPlan,
}]),
}
}
pub(crate) fn parse(&mut self, delta: &str) -> Vec<ProposedPlanSegment> {
self.parser
.parse(delta)
.into_iter()
.map(map_plan_segment)
.collect()
}
pub(crate) fn finish(&mut self) -> Vec<ProposedPlanSegment> {
self.parser
.finish()
.into_iter()
.map(map_plan_segment)
.collect()
}
}
fn map_plan_segment(segment: TaggedLineSegment<PlanTag>) -> ProposedPlanSegment {
match segment {
TaggedLineSegment::Normal(text) => ProposedPlanSegment::Normal(text),
TaggedLineSegment::TagStart(PlanTag::ProposedPlan) => {
ProposedPlanSegment::ProposedPlanStart
}
TaggedLineSegment::TagDelta(PlanTag::ProposedPlan, text) => {
ProposedPlanSegment::ProposedPlanDelta(text)
}
TaggedLineSegment::TagEnd(PlanTag::ProposedPlan) => ProposedPlanSegment::ProposedPlanEnd,
}
}
pub(crate) fn strip_proposed_plan_blocks(text: &str) -> String {
let mut parser = ProposedPlanParser::new();
let mut out = String::new();
for segment in parser.parse(text).into_iter().chain(parser.finish()) {
if let ProposedPlanSegment::Normal(delta) = segment {
out.push_str(&delta);
}
}
out
}
pub(crate) fn extract_proposed_plan_text(text: &str) -> Option<String> {
let mut parser = ProposedPlanParser::new();
let mut plan_text = String::new();
let mut saw_plan_block = false;
for segment in parser.parse(text).into_iter().chain(parser.finish()) {
match segment {
ProposedPlanSegment::ProposedPlanStart => {
saw_plan_block = true;
plan_text.clear();
}
ProposedPlanSegment::ProposedPlanDelta(delta) => {
plan_text.push_str(&delta);
}
ProposedPlanSegment::ProposedPlanEnd | ProposedPlanSegment::Normal(_) => {}
}
}
saw_plan_block.then_some(plan_text)
}
#[cfg(test)]
mod tests {
use super::ProposedPlanParser;
use super::ProposedPlanSegment;
use super::strip_proposed_plan_blocks;
use pretty_assertions::assert_eq;
#[test]
fn streams_proposed_plan_segments() {
let mut parser = ProposedPlanParser::new();
let mut segments = Vec::new();
for chunk in [
"Intro text\n<prop",
"osed_plan>\n- step 1\n",
"</proposed_plan>\nOutro",
] {
segments.extend(parser.parse(chunk));
}
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
ProposedPlanSegment::Normal("Intro text\n".to_string()),
ProposedPlanSegment::ProposedPlanStart,
ProposedPlanSegment::ProposedPlanDelta("- step 1\n".to_string()),
ProposedPlanSegment::ProposedPlanEnd,
ProposedPlanSegment::Normal("Outro".to_string()),
]
);
}
#[test]
fn preserves_non_tag_lines() {
let mut parser = ProposedPlanParser::new();
let mut segments = parser.parse(" <proposed_plan> extra\n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![ProposedPlanSegment::Normal(
" <proposed_plan> extra\n".to_string()
)]
);
}
#[test]
fn closes_unterminated_plan_block_on_finish() {
let mut parser = ProposedPlanParser::new();
let mut segments = parser.parse("<proposed_plan>\n- step 1\n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
ProposedPlanSegment::ProposedPlanStart,
ProposedPlanSegment::ProposedPlanDelta("- step 1\n".to_string()),
ProposedPlanSegment::ProposedPlanEnd,
]
);
}
#[test]
fn closes_tag_line_without_trailing_newline() {
let mut parser = ProposedPlanParser::new();
let mut segments = parser.parse("<proposed_plan>\n- step 1\n</proposed_plan>");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
ProposedPlanSegment::ProposedPlanStart,
ProposedPlanSegment::ProposedPlanDelta("- step 1\n".to_string()),
ProposedPlanSegment::ProposedPlanEnd,
]
);
}
#[test]
fn strips_proposed_plan_blocks_from_text() {
let text = "before\n<proposed_plan>\n- step\n</proposed_plan>\nafter";
assert_eq!(strip_proposed_plan_blocks(text), "before\nafter");
}
}
codex-rs/core/src/stream_events_utils.rs
+104 -19
View File
@@ -3,6 +3,7 @@ use std::sync::Arc;
use codex_protocol::config_types::ModeKind;
use codex_protocol::items::TurnItem;
use codex_utils_stream_parser::strip_citations;
use tokio_util::sync::CancellationToken;
use crate::codex::Session;
@@ -11,17 +12,42 @@ use crate::error::CodexErr;
use crate::error::Result;
use crate::function_tool::FunctionCallError;
use crate::parse_turn_item;
use crate::proposed_plan_parser::strip_proposed_plan_blocks;
use crate::tools::parallel::ToolCallRuntime;
use crate::tools::router::ToolRouter;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_utils_stream_parser::strip_proposed_plan_blocks;
use futures::Future;
use tracing::debug;
use tracing::instrument;
fn strip_hidden_assistant_markup(text: &str, plan_mode: bool) -> String {
let (without_citations, _citations) = strip_citations(text);
if plan_mode {
strip_proposed_plan_blocks(&without_citations)
} else {
without_citations
}
}
pub(crate) fn raw_assistant_output_text_from_item(item: &ResponseItem) -> Option<String> {
if let ResponseItem::Message { role, content, .. } = item
&& role == "assistant"
{
let combined = content
.iter()
.filter_map(|ci| match ci {
codex_protocol::models::ContentItem::OutputText { text } => Some(text.as_str()),
_ => None,
})
.collect::<String>();
return Some(combined);
}
None
}
/// Handle a completed output item from the model stream, recording it and
/// queuing any tool execution futures. This records items immediately so
/// history and rollout stay in sync even if the turn is later cancelled.
@@ -169,7 +195,7 @@ pub(crate) async fn handle_non_tool_response_item(
| ResponseItem::Reasoning { .. }
| ResponseItem::WebSearchCall { .. } => {
let mut turn_item = parse_turn_item(item)?;
if plan_mode && let TurnItem::AgentMessage(agent_message) = &mut turn_item {
if let TurnItem::AgentMessage(agent_message) = &mut turn_item {
let combined = agent_message
.content
.iter()
@@ -177,7 +203,7 @@ pub(crate) async fn handle_non_tool_response_item(
codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
})
.collect::<String>();
let stripped = strip_proposed_plan_blocks(&combined);
let stripped = strip_hidden_assistant_markup(&combined, plan_mode);
agent_message.content =
vec![codex_protocol::items::AgentMessageContent::Text { text: stripped }];
}
@@ -195,25 +221,15 @@ pub(crate) fn last_assistant_message_from_item(
item: &ResponseItem,
plan_mode: bool,
) -> Option<String> {
if let ResponseItem::Message { role, content, .. } = item
&& role == "assistant"
{
let combined = content
.iter()
.filter_map(|ci| match ci {
codex_protocol::models::ContentItem::OutputText { text } => Some(text.as_str()),
_ => None,
})
.collect::<String>();
if let Some(combined) = raw_assistant_output_text_from_item(item) {
if combined.is_empty() {
return None;
}
return if plan_mode {
let stripped = strip_proposed_plan_blocks(&combined);
(!stripped.trim().is_empty()).then_some(stripped)
} else {
Some(combined)
};
let stripped = strip_hidden_assistant_markup(&combined, plan_mode);
if stripped.trim().is_empty() {
return None;
}
return Some(stripped);
}
None
}
@@ -248,3 +264,72 @@ pub(crate) fn response_input_to_response_item(input: &ResponseInputItem) -> Opti
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::handle_non_tool_response_item;
use super::last_assistant_message_from_item;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use pretty_assertions::assert_eq;
fn assistant_output_text(text: &str) -> ResponseItem {
ResponseItem::Message {
id: Some("msg-1".to_string()),
role: "assistant".to_string(),
content: vec![ContentItem::OutputText {
text: text.to_string(),
}],
end_turn: Some(true),
phase: None,
}
}
#[tokio::test]
async fn handle_non_tool_response_item_strips_citations_from_assistant_message() {
let item = assistant_output_text("hello<oai-mem-citation>doc1</oai-mem-citation> world");
let turn_item = handle_non_tool_response_item(&item, false)
.await
.expect("assistant message should parse");
let TurnItem::AgentMessage(agent_message) = turn_item else {
panic!("expected agent message");
};
let text = agent_message
.content
.iter()
.map(|entry| match entry {
codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
})
.collect::<String>();
assert_eq!(text, "hello world");
}
#[test]
fn last_assistant_message_from_item_strips_citations_and_plan_blocks() {
let item = assistant_output_text(
"before<oai-mem-citation>doc1</oai-mem-citation>\n<proposed_plan>\n- x\n</proposed_plan>\nafter",
);
let message = last_assistant_message_from_item(&item, true)
.expect("assistant text should remain after stripping");
assert_eq!(message, "before\nafter");
}
#[test]
fn last_assistant_message_from_item_returns_none_for_citation_only_message() {
let item = assistant_output_text("<oai-mem-citation>doc1</oai-mem-citation>");
assert_eq!(last_assistant_message_from_item(&item, false), None);
}
#[test]
fn last_assistant_message_from_item_returns_none_for_plan_only_hidden_message() {
let item = assistant_output_text("<proposed_plan>\n- x\n</proposed_plan>");
assert_eq!(last_assistant_message_from_item(&item, true), None);
}
}
codex-rs/core/src/tagged_block_parser.rs
-314
View File
@@ -1,314 +0,0 @@
//! Line-based tag block parsing for streamed text.
//!
//! The parser buffers each line until it can disprove that the line is a tag,
//! which is required for tags that must appear alone on a line. For example,
//! Proposed Plan output uses `<proposed_plan>` and `</proposed_plan>` tags
//! on their own lines so clients can stream plan content separately.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct TagSpec<T> {
pub(crate) open: &'static str,
pub(crate) close: &'static str,
pub(crate) tag: T,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum TaggedLineSegment<T> {
Normal(String),
TagStart(T),
TagDelta(T, String),
TagEnd(T),
}
/// Stateful line parser that splits input into normal text vs tag blocks.
///
/// How it works:
/// - While reading a line, we buffer characters until the line either finishes
/// (`\n`) or stops matching any tag prefix (after `trim_start`).
/// - If it stops matching a tag prefix, the buffered line is immediately
/// emitted as text and we continue in "plain text" mode until the next
/// newline.
/// - When a full line is available, we compare it to the open/close tags; tag
/// lines emit TagStart/TagEnd, otherwise the line is emitted as text.
/// - `finish()` flushes any buffered line and auto-closes an unterminated tag,
/// which keeps streaming resilient to missing closing tags.
#[derive(Debug, Default)]
pub(crate) struct TaggedLineParser<T>
where
T: Copy + Eq,
{
specs: Vec<TagSpec<T>>,
active_tag: Option<T>,
detect_tag: bool,
line_buffer: String,
}
impl<T> TaggedLineParser<T>
where
T: Copy + Eq,
{
pub(crate) fn new(specs: Vec<TagSpec<T>>) -> Self {
Self {
specs,
active_tag: None,
detect_tag: true,
line_buffer: String::new(),
}
}
/// Parse a streamed delta into line-aware segments.
pub(crate) fn parse(&mut self, delta: &str) -> Vec<TaggedLineSegment<T>> {
let mut segments = Vec::new();
let mut run = String::new();
for ch in delta.chars() {
if self.detect_tag {
if !run.is_empty() {
self.push_text(std::mem::take(&mut run), &mut segments);
}
self.line_buffer.push(ch);
if ch == '\n' {
self.finish_line(&mut segments);
continue;
}
let slug = self.line_buffer.trim_start();
if slug.is_empty() || self.is_tag_prefix(slug) {
continue;
}
// This line cannot be a tag line, so flush it immediately.
let buffered = std::mem::take(&mut self.line_buffer);
self.detect_tag = false;
self.push_text(buffered, &mut segments);
continue;
}
run.push(ch);
if ch == '\n' {
self.push_text(std::mem::take(&mut run), &mut segments);
self.detect_tag = true;
}
}
if !run.is_empty() {
self.push_text(run, &mut segments);
}
segments
}
/// Flush any buffered text and close an unterminated tag block.
pub(crate) fn finish(&mut self) -> Vec<TaggedLineSegment<T>> {
let mut segments = Vec::new();
if !self.line_buffer.is_empty() {
let buffered = std::mem::take(&mut self.line_buffer);
let without_newline = buffered.strip_suffix('\n').unwrap_or(&buffered);
let slug = without_newline.trim_start().trim_end();
if let Some(tag) = self.match_open(slug)
&& self.active_tag.is_none()
{
push_segment(&mut segments, TaggedLineSegment::TagStart(tag));
self.active_tag = Some(tag);
} else if let Some(tag) = self.match_close(slug)
&& self.active_tag == Some(tag)
{
push_segment(&mut segments, TaggedLineSegment::TagEnd(tag));
self.active_tag = None;
} else {
// The buffered line never proved to be a tag line.
self.push_text(buffered, &mut segments);
}
}
if let Some(tag) = self.active_tag.take() {
push_segment(&mut segments, TaggedLineSegment::TagEnd(tag));
}
self.detect_tag = true;
segments
}
fn finish_line(&mut self, segments: &mut Vec<TaggedLineSegment<T>>) {
let line = std::mem::take(&mut self.line_buffer);
let without_newline = line.strip_suffix('\n').unwrap_or(&line);
let slug = without_newline.trim_start().trim_end();
if let Some(tag) = self.match_open(slug)
&& self.active_tag.is_none()
{
push_segment(segments, TaggedLineSegment::TagStart(tag));
self.active_tag = Some(tag);
self.detect_tag = true;
return;
}
if let Some(tag) = self.match_close(slug)
&& self.active_tag == Some(tag)
{
push_segment(segments, TaggedLineSegment::TagEnd(tag));
self.active_tag = None;
self.detect_tag = true;
return;
}
self.detect_tag = true;
self.push_text(line, segments);
}
fn push_text(&self, text: String, segments: &mut Vec<TaggedLineSegment<T>>) {
if let Some(tag) = self.active_tag {
push_segment(segments, TaggedLineSegment::TagDelta(tag, text));
} else {
push_segment(segments, TaggedLineSegment::Normal(text));
}
}
fn is_tag_prefix(&self, slug: &str) -> bool {
let slug = slug.trim_end();
self.specs
.iter()
.any(|spec| spec.open.starts_with(slug) || spec.close.starts_with(slug))
}
fn match_open(&self, slug: &str) -> Option<T> {
self.specs
.iter()
.find(|spec| spec.open == slug)
.map(|spec| spec.tag)
}
fn match_close(&self, slug: &str) -> Option<T> {
self.specs
.iter()
.find(|spec| spec.close == slug)
.map(|spec| spec.tag)
}
}
fn push_segment<T>(segments: &mut Vec<TaggedLineSegment<T>>, segment: TaggedLineSegment<T>)
where
T: Copy + Eq,
{
match segment {
TaggedLineSegment::Normal(delta) => {
if delta.is_empty() {
return;
}
if let Some(TaggedLineSegment::Normal(existing)) = segments.last_mut() {
existing.push_str(&delta);
return;
}
segments.push(TaggedLineSegment::Normal(delta));
}
TaggedLineSegment::TagDelta(tag, delta) => {
if delta.is_empty() {
return;
}
if let Some(TaggedLineSegment::TagDelta(existing_tag, existing)) = segments.last_mut()
&& *existing_tag == tag
{
existing.push_str(&delta);
return;
}
segments.push(TaggedLineSegment::TagDelta(tag, delta));
}
TaggedLineSegment::TagStart(tag) => {
segments.push(TaggedLineSegment::TagStart(tag));
}
TaggedLineSegment::TagEnd(tag) => {
segments.push(TaggedLineSegment::TagEnd(tag));
}
}
}
#[cfg(test)]
mod tests {
use super::TagSpec;
use super::TaggedLineParser;
use super::TaggedLineSegment;
use pretty_assertions::assert_eq;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Tag {
Block,
}
fn parser() -> TaggedLineParser<Tag> {
TaggedLineParser::new(vec![TagSpec {
open: "<tag>",
close: "</tag>",
tag: Tag::Block,
}])
}
#[test]
fn buffers_prefix_until_tag_is_decided() {
let mut parser = parser();
let mut segments = parser.parse("<t");
segments.extend(parser.parse("ag>\nline\n</tag>\n"));
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
TaggedLineSegment::TagStart(Tag::Block),
TaggedLineSegment::TagDelta(Tag::Block, "line\n".to_string()),
TaggedLineSegment::TagEnd(Tag::Block),
]
);
}
#[test]
fn rejects_tag_lines_with_extra_text() {
let mut parser = parser();
let mut segments = parser.parse("<tag> extra\n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![TaggedLineSegment::Normal("<tag> extra\n".to_string())]
);
}
#[test]
fn closes_unterminated_tag_on_finish() {
let mut parser = parser();
let mut segments = parser.parse("<tag>\nline\n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
TaggedLineSegment::TagStart(Tag::Block),
TaggedLineSegment::TagDelta(Tag::Block, "line\n".to_string()),
TaggedLineSegment::TagEnd(Tag::Block),
]
);
}
#[test]
fn accepts_tags_with_trailing_whitespace() {
let mut parser = parser();
let mut segments = parser.parse("<tag> \nline\n</tag> \n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![
TaggedLineSegment::TagStart(Tag::Block),
TaggedLineSegment::TagDelta(Tag::Block, "line\n".to_string()),
TaggedLineSegment::TagEnd(Tag::Block),
]
);
}
#[test]
fn passes_through_plain_text() {
let mut parser = parser();
let mut segments = parser.parse("plain text\n");
segments.extend(parser.finish());
assert_eq!(
segments,
vec![TaggedLineSegment::Normal("plain text\n".to_string())]
);
}
}
codex-rs/core/templates/memories/read_path.md
+41
View File
@@ -52,6 +52,47 @@ When to update memory:
- When user explicitly asks to remember/update memory, revise memory_summary.md
and/or MEMORY.md.
Memory citation requirements:
- If ANY relevant memory files were used: append exactly one
`<oai-mem-citation>` block as the VERY LAST content of the final reply.
Normal responses should include the answer first, then append the
`<oai-mem-citation>` block at the end.
- Use this exact structure for programmatic parsing:
```
<oai-mem-citation>
<citation_entries>
MEMORY.md:234-236|note=[responsesapi citation extraction code pointer]
rollout_summaries/2026-02-17T21-23-02-LN3m-weekly_memory_report_pivot_from_git_history.md:10-12|note=[weekly report format]
</citation_entries>
<rollout_ids>
019c6e27-e55b-73d1-87d8-4e01f1f75043
019c7714-3b77-74d1-9866-e1f484aae2ab
</rollout_ids>
</oai-mem-citation>
```
- `citation_entries` is for rendering:
- one citation entry per line
- format: `<file>:<line_start>-<line_end>|note=[<how memory was used>]`
- use file paths relative to the memory base path (for example, `MEMORY.md`,
`rollout_summaries/...`, `skills/...`)
- only cite files actually used under the memory base path (do not cite
workspace files as memory citations)
- if you used `MEMORY.md` and then a rollout summary/skill file, cite both
- list entries in order of importance (most important first)
- `note` should be short, single-line, and use simple characters only (avoid
unusual symbols, no newlines)
- `rollout_ids` is for us to track what previous rollouts you find useful:
- include one rollout id per line
- rollout ids should look like UUIDs (for example,
`019c6e27-e55b-73d1-87d8-4e01f1f75043`)
- include unique ids only; do not repeat ids
- an empty `<rollout_ids>` section is allowed if no rollout ids are available
- you can find rollout ids in rollout summary files and MEMORY.md
- do not include file paths or notes in this section
- Never include memory citations inside pull-request messages.
- Never cite blank lines; double-check ranges.
========= MEMORY_SUMMARY BEGINS =========
{{ memory_summary }}
========= MEMORY_SUMMARY ENDS =========
codex-rs/core/tests/suite/items.rs
+308
View File
@@ -505,6 +505,314 @@ async fn plan_mode_strips_plan_from_agent_messages() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn plan_mode_streaming_citations_are_stripped_across_added_deltas_and_done()
-> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let TestCodex {
codex,
session_configured,
..
} = test_codex().build(&server).await?;
let added_text = "Intro <oai-mem-";
let deltas = [
"citation>outer-doc</oai-mem-citation>\n<proposed",
"_plan>\n- Step 1<oai-mem-",
"citation>plan-doc</oai-mem-citation>\n- Step 2\n</proposed_plan>\nOu",
"tro",
];
let full_message = format!("{added_text}{}", deltas.concat());
let mut events = vec![
ev_response_created("resp-1"),
ev_message_item_added("msg-1", added_text),
];
for delta in deltas {
events.push(ev_output_text_delta(delta));
}
events.push(ev_assistant_message("msg-1", &full_message));
events.push(ev_completed("resp-1"));
mount_sse_once(&server, sse(events)).await;
let collaboration_mode = CollaborationMode {
mode: ModeKind::Plan,
settings: Settings {
model: session_configured.model.clone(),
reasoning_effort: None,
developer_instructions: None,
},
};
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "please plan with citations".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
cwd: std::env::current_dir()?,
approval_policy: codex_protocol::protocol::AskForApproval::Never,
sandbox_policy: codex_protocol::protocol::SandboxPolicy::DangerFullAccess,
model: session_configured.model.clone(),
effort: None,
summary: codex_protocol::config_types::ReasoningSummary::Auto,
collaboration_mode: Some(collaboration_mode),
personality: None,
})
.await?;
let mut agent_started = None;
let mut agent_started_idx = None;
let mut agent_completed = None;
let mut agent_completed_idx = None;
let mut plan_started = None;
let mut plan_started_idx = None;
let mut plan_completed = None;
let mut plan_completed_idx = None;
let mut agent_deltas = Vec::new();
let mut plan_deltas = Vec::new();
let mut first_agent_delta_idx = None;
let mut last_agent_delta_idx = None;
let mut first_plan_delta_idx = None;
let mut last_plan_delta_idx = None;
let mut idx = 0usize;
let turn_complete_idx = loop {
let ev = wait_for_event(&codex, |_| true).await;
match ev {
EventMsg::ItemStarted(ItemStartedEvent {
item: TurnItem::AgentMessage(item),
..
}) => {
agent_started_idx = Some(idx);
agent_started = Some(item);
}
EventMsg::ItemStarted(ItemStartedEvent {
item: TurnItem::Plan(item),
..
}) => {
plan_started_idx = Some(idx);
plan_started = Some(item);
}
EventMsg::AgentMessageContentDelta(event) => {
if first_agent_delta_idx.is_none() {
first_agent_delta_idx = Some(idx);
}
last_agent_delta_idx = Some(idx);
agent_deltas.push(event.delta);
}
EventMsg::PlanDelta(event) => {
if first_plan_delta_idx.is_none() {
first_plan_delta_idx = Some(idx);
}
last_plan_delta_idx = Some(idx);
plan_deltas.push(event.delta);
}
EventMsg::ItemCompleted(ItemCompletedEvent {
item: TurnItem::AgentMessage(item),
..
}) => {
agent_completed_idx = Some(idx);
agent_completed = Some(item);
}
EventMsg::ItemCompleted(ItemCompletedEvent {
item: TurnItem::Plan(item),
..
}) => {
plan_completed_idx = Some(idx);
plan_completed = Some(item);
}
EventMsg::TurnComplete(_) => {
break idx;
}
_ => {}
}
idx += 1;
};
let agent_started = agent_started.expect("agent item start should be emitted");
let agent_completed = agent_completed.expect("agent item completion should be emitted");
let plan_started = plan_started.expect("plan item start should be emitted");
let plan_completed = plan_completed.expect("plan item completion should be emitted");
assert_eq!(agent_started.id, agent_completed.id);
assert_eq!(plan_started.id, plan_completed.id);
assert_eq!(plan_started.text, "");
let agent_started_text: String = agent_started
.content
.iter()
.map(|entry| match entry {
AgentMessageContent::Text { text } => text.as_str(),
})
.collect();
let agent_completed_text: String = agent_completed
.content
.iter()
.map(|entry| match entry {
AgentMessageContent::Text { text } => text.as_str(),
})
.collect();
let agent_delta_text = agent_deltas.concat();
let plan_delta_text = plan_deltas.concat();
assert_eq!(agent_started_text, "");
assert_eq!(agent_delta_text, "Intro \nOutro");
assert_eq!(agent_completed_text, "Intro \nOutro");
assert_eq!(plan_delta_text, "- Step 1\n- Step 2\n");
assert_eq!(plan_completed.text, "- Step 1\n- Step 2\n");
for text in [
agent_started_text.as_str(),
agent_delta_text.as_str(),
agent_completed_text.as_str(),
plan_delta_text.as_str(),
plan_completed.text.as_str(),
] {
assert!(!text.contains("<oai-mem-citation>"));
assert!(!text.contains("</oai-mem-citation>"));
}
let agent_started_idx = agent_started_idx.expect("agent start index");
let agent_completed_idx = agent_completed_idx.expect("agent completion index");
let plan_started_idx = plan_started_idx.expect("plan start index");
let plan_completed_idx = plan_completed_idx.expect("plan completion index");
let first_agent_delta_idx = first_agent_delta_idx.expect("agent delta index");
let last_agent_delta_idx = last_agent_delta_idx.expect("agent delta index");
let first_plan_delta_idx = first_plan_delta_idx.expect("plan delta index");
let last_plan_delta_idx = last_plan_delta_idx.expect("plan delta index");
assert!(agent_started_idx < first_agent_delta_idx);
assert!(plan_started_idx < first_plan_delta_idx);
assert!(last_agent_delta_idx < agent_completed_idx);
assert!(last_plan_delta_idx < plan_completed_idx);
assert!(agent_completed_idx < turn_complete_idx);
assert!(plan_completed_idx < turn_complete_idx);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn plan_mode_streaming_proposed_plan_tag_split_across_added_and_delta_is_parsed()
-> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let TestCodex {
codex,
session_configured,
..
} = test_codex().build(&server).await?;
let added_text = "Intro\n<proposed";
let deltas = ["_plan>\n- Step 1\n</proposed_plan>\nOutro"];
let full_message = format!("{added_text}{}", deltas.concat());
let mut events = vec![
ev_response_created("resp-1"),
ev_message_item_added("msg-1", added_text),
];
for delta in deltas {
events.push(ev_output_text_delta(delta));
}
events.push(ev_assistant_message("msg-1", &full_message));
events.push(ev_completed("resp-1"));
mount_sse_once(&server, sse(events)).await;
let collaboration_mode = CollaborationMode {
mode: ModeKind::Plan,
settings: Settings {
model: session_configured.model.clone(),
reasoning_effort: None,
developer_instructions: None,
},
};
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "please plan".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
cwd: std::env::current_dir()?,
approval_policy: codex_protocol::protocol::AskForApproval::Never,
sandbox_policy: codex_protocol::protocol::SandboxPolicy::DangerFullAccess,
model: session_configured.model.clone(),
effort: None,
summary: codex_protocol::config_types::ReasoningSummary::Auto,
collaboration_mode: Some(collaboration_mode),
personality: None,
})
.await?;
let mut agent_started = None;
let mut agent_completed = None;
let mut plan_started = None;
let mut plan_completed = None;
let mut agent_deltas = Vec::new();
let mut plan_deltas = Vec::new();
loop {
let ev = wait_for_event(&codex, |_| true).await;
match ev {
EventMsg::ItemStarted(ItemStartedEvent {
item: TurnItem::AgentMessage(item),
..
}) => agent_started = Some(item),
EventMsg::ItemStarted(ItemStartedEvent {
item: TurnItem::Plan(item),
..
}) => plan_started = Some(item),
EventMsg::AgentMessageContentDelta(event) => agent_deltas.push(event.delta),
EventMsg::PlanDelta(event) => plan_deltas.push(event.delta),
EventMsg::ItemCompleted(ItemCompletedEvent {
item: TurnItem::AgentMessage(item),
..
}) => agent_completed = Some(item),
EventMsg::ItemCompleted(ItemCompletedEvent {
item: TurnItem::Plan(item),
..
}) => plan_completed = Some(item),
EventMsg::TurnComplete(_) => break,
_ => {}
}
}
let agent_started = agent_started.expect("agent item start should be emitted");
let agent_completed = agent_completed.expect("agent item completion should be emitted");
let plan_started = plan_started.expect("plan item start should be emitted");
let plan_completed = plan_completed.expect("plan item completion should be emitted");
let agent_started_text: String = agent_started
.content
.iter()
.map(|entry| match entry {
AgentMessageContent::Text { text } => text.as_str(),
})
.collect();
let agent_completed_text: String = agent_completed
.content
.iter()
.map(|entry| match entry {
AgentMessageContent::Text { text } => text.as_str(),
})
.collect();
assert_eq!(agent_started_text, "");
assert_eq!(agent_deltas.concat(), "Intro\nOutro");
assert_eq!(agent_completed_text, "Intro\nOutro");
assert_eq!(plan_started.text, "");
assert_eq!(plan_deltas.concat(), "- Step 1\n");
assert_eq!(plan_completed.text, "- Step 1\n");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn plan_mode_handles_missing_plan_close_tag() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));