From 296d26760802c06ce75e9d31f17d9cd70839afbf Mon Sep 17 00:00:00 2001 From: "kh.ai" Date: Wed, 28 Jan 2026 19:51:36 -0800 Subject: [PATCH] collab: add fresh_context to send_input --- codex-rs/core/src/agent/control.rs | 5 ++ codex-rs/core/src/codex.rs | 26 ++++++++ codex-rs/core/src/tools/handlers/collab.rs | 65 +++++++++++++++++++ codex-rs/core/src/tools/spec.rs | 9 +++ .../core/templates/agents/orchestrator.md | 1 + codex-rs/protocol/src/protocol.rs | 6 ++ 6 files changed, 112 insertions(+) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 611c0d164c..6cbbed24dd 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -92,6 +92,11 @@ impl AgentControl { result } + pub(crate) async fn reset_conversation(&self, agent_id: ThreadId) -> CodexResult { + let state = self.upgrade()?; + state.send_op(agent_id, Op::ResetConversation).await + } + /// Interrupt the current task for an existing agent thread. pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 70e52eb064..cc6ee01507 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2381,6 +2381,9 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv Op::ThreadRollback { num_turns } => { handlers::thread_rollback(&sess, sub.id.clone(), num_turns).await; } + Op::ResetConversation => { + handlers::reset_conversation(&sess, sub.id.clone()).await; + } Op::RunUserShellCommand { command } => { handlers::run_user_shell_command( &sess, @@ -2893,6 +2896,29 @@ mod handlers { .await; } + pub async fn reset_conversation(sess: &Arc, sub_id: String) { + let has_active_turn = { sess.active_turn.lock().await.is_some() }; + if has_active_turn { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Warning(WarningEvent { + message: "Cannot reset conversation while a turn is in progress.".to_string(), + }), + }) + .await; + return; + } + + let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; + let initial_context = sess.build_initial_context(&turn_context).await; + sess.replace_history(initial_context).await; + { + let mut state = sess.state.lock().await; + state.initial_context_seeded = true; + } + sess.recompute_token_usage(turn_context.as_ref()).await; + } + pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { sess.abort_all_tasks(TurnAbortReason::Interrupted).await; sess.services diff --git a/codex-rs/core/src/tools/handlers/collab.rs b/codex-rs/core/src/tools/handlers/collab.rs index b1666949f6..e2d37fe755 100644 --- a/codex-rs/core/src/tools/handlers/collab.rs +++ b/codex-rs/core/src/tools/handlers/collab.rs @@ -201,6 +201,8 @@ mod send_input { message: String, #[serde(default)] interrupt: bool, + #[serde(default)] + fresh_context: bool, } #[derive(Debug, Serialize)] @@ -222,6 +224,30 @@ mod send_input { "Empty message can't be sent to an agent".to_string(), )); } + if args.interrupt && args.fresh_context { + return Err(FunctionCallError::RespondToModel( + "fresh_context cannot be combined with interrupt; interrupt first, wait, then retry".to_string(), + )); + } + if args.fresh_context { + let status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + if matches!(status, AgentStatus::Running) { + return Err(FunctionCallError::RespondToModel( + "agent must not be running to reset context; wait for completion then retry" + .to_string(), + )); + } + session + .services + .agent_control + .reset_conversation(receiver_thread_id) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err))?; + } if args.interrupt { session .services @@ -898,6 +924,45 @@ mod tests { .expect("shutdown should submit"); } + #[tokio::test] + async fn send_input_resets_before_prompt_when_fresh_context_set() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + session.services.agent_control = manager.agent_control(); + let config = turn.client.config().as_ref().clone(); + let thread = manager.start_thread(config).await.expect("start thread"); + let agent_id = thread.thread_id; + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "send_input", + function_payload(json!({ + "id": agent_id.to_string(), + "message": "hi", + "fresh_context": true + })), + ); + CollabHandler + .handle(invocation) + .await + .expect("send_input should succeed"); + + let ops = manager.captured_ops(); + let ops_for_agent: Vec<&Op> = ops + .iter() + .filter_map(|(id, op)| (*id == agent_id).then_some(op)) + .collect(); + assert_eq!(ops_for_agent.len(), 2); + assert!(matches!(ops_for_agent[0], Op::ResetConversation)); + assert!(matches!(ops_for_agent[1], Op::UserInput { .. })); + + let _ = thread + .thread + .submit(Op::Shutdown {}) + .await + .expect("shutdown should submit"); + } + #[derive(Debug, Deserialize, PartialEq, Eq)] struct WaitResult { status: HashMap, diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index df65a94d8a..e3858e13e9 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -503,6 +503,15 @@ fn create_send_input_tool() -> ToolSpec { ), }, ); + properties.insert( + "fresh_context".to_string(), + JsonSchema::Boolean { + description: Some( + "When true, clear the agent's in-memory context before sending this message." + .to_string(), + ), + }, + ); ToolSpec::Function(ResponsesApiTool { name: "send_input".to_string(), diff --git a/codex-rs/core/templates/agents/orchestrator.md b/codex-rs/core/templates/agents/orchestrator.md index e0976f52ef..481caf841d 100644 --- a/codex-rs/core/templates/agents/orchestrator.md +++ b/codex-rs/core/templates/agents/orchestrator.md @@ -97,6 +97,7 @@ Sub-agents are their to make you go fast and time is a big constraint so leverag - When you ask sub-agent to do the work for you, your only role becomes to coordinate them. Do not perform the actual work while they are working. - When you have plan with multiple step, process them in parallel by spawning one agent per step when this is possible. - Choose the correct agent type. +- When reusing a small agent pool across many independent tasks, prefer `send_input` with `fresh_context=true` so each task runs with a clean context. ## Flow 1. Understand the task. diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index b9a4f2120b..9af0b46dbb 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -281,6 +281,12 @@ pub enum Op { /// responsible for undoing any edits on disk. ThreadRollback { num_turns: u32 }, + /// Reset the in-memory conversation context back to the initial session prefix. + /// + /// This clears prior user/assistant turns without changing the session configuration. + /// Useful for reusing agents as stateless workers across many independent tasks. + ResetConversation, + /// Request a code review from the agent. Review { review_request: ReviewRequest },