core/src/session/turn_tests.rs
67 lines
use super::*;use codex_extension_api::ExtensionData;use codex_extension_api::TurnItemContributor;use codex_protocol::items::AgentMessageContent;use pretty_assertions::assert_eq;use std::sync::Arc;struct RewriteAgentMessageContributor;impl TurnItemContributor for RewriteAgentMessageContributor { fn contribute<'a>( &'a self, _thread_store: &'a ExtensionData, _turn_store: &'a ExtensionData, item: &'a mut TurnItem, ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { Box::pin(async move { if let TurnItem::AgentMessage(agent_message) = item { agent_message.content = vec![AgentMessageContent::Text { text: "plan contributed assistant text".to_string(), }]; } Ok(()) }) }}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(), }], phase: None, metadata: None, }}#[tokio::test]async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() { let (mut session, turn_context) = crate::session::tests::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 mut state = PlanModeStreamState::new(&turn_context.sub_id); let mut last_agent_message = None; let item = assistant_output_text("original assistant text"); let handled = handle_assistant_item_done_in_plan_mode( &session, &turn_context, &turn_store, &item, &mut state, /*previously_active_item*/ None, &mut last_agent_message, ) .await; assert!(handled); assert_eq!( last_agent_message.as_deref(), Some("plan contributed assistant text") );}