PoC: codex-kernel

codex-rs/kernel/src/lib.rs

@@ -0,0 +1,17 @@
+mod prompt;
+mod sampling;
+mod tools;
+
+pub use prompt::Prompt;
+pub use prompt::PromptConfig;
+pub use sampling::PreparedSamplingRequest;
+pub use sampling::SamplingLoopHost;
+pub use sampling::SamplingRequestResult;
+pub use sampling::run_sampling_request_loop;
+pub use tools::CompletedResponseItem;
+pub use tools::KernelToolCall;
+pub use tools::KernelToolExecutor;
+pub use tools::ToolCallError;
+pub use tools::ToolConfig;
+pub use tools::execute_tool_call_with_default_output;
+pub use tools::response_input_to_response_item;

codex-rs/kernel/src/prompt.rs

@@ -0,0 +1,257 @@
+use crate::ToolConfig;
+use codex_protocol::config_types::Personality;
+use codex_protocol::models::BaseInstructions;
+use codex_protocol::models::FunctionCallOutputBody;
+use codex_protocol::models::ResponseItem;
+use codex_tools::ToolSpec;
+use serde::Deserialize;
+use serde_json::Value;
+use std::collections::HashSet;
+
+/// API request payload for a single model turn.
+#[derive(Debug, Clone)]
+pub struct Prompt {
+    pub input: Vec<ResponseItem>,
+    pub tools: Vec<ToolSpec>,
+    pub parallel_tool_calls: bool,
+    pub base_instructions: BaseInstructions,
+    pub personality: Option<Personality>,
+    pub output_schema: Option<Value>,
+    pub output_schema_strict: bool,
+}
+
+/// Retry-stable prompt settings used to rebuild a [`Prompt`] after tool calls or stream retries.
+#[derive(Debug, Clone)]
+pub struct PromptConfig {
+    pub base_instructions: BaseInstructions,
+    pub personality: Option<Personality>,
+    pub output_schema: Option<Value>,
+    pub output_schema_strict: bool,
+}
+
+impl Prompt {
+    pub fn get_formatted_input(&self) -> Vec<ResponseItem> {
+        let mut input = self.input.clone();
+        let is_freeform_apply_patch_tool_present = self.tools.iter().any(|tool| match tool {
+            ToolSpec::Freeform(f) => f.name == "apply_patch",
+            _ => false,
+        });
+        if is_freeform_apply_patch_tool_present {
+            reserialize_shell_outputs(&mut input);
+        }
+
+        input
+    }
+}
+
+impl PromptConfig {
+    pub fn build_prompt(&self, input: Vec<ResponseItem>, tool_config: &ToolConfig) -> Prompt {
+        Prompt {
+            input,
+            tools: tool_config.tools.clone(),
+            parallel_tool_calls: tool_config.parallel_tool_calls,
+            base_instructions: self.base_instructions.clone(),
+            personality: self.personality,
+            output_schema: self.output_schema.clone(),
+            output_schema_strict: self.output_schema_strict,
+        }
+    }
+}
+
+impl Default for Prompt {
+    fn default() -> Self {
+        Self {
+            input: Vec::new(),
+            tools: Vec::new(),
+            parallel_tool_calls: false,
+            base_instructions: BaseInstructions::default(),
+            personality: None,
+            output_schema: None,
+            output_schema_strict: true,
+        }
+    }
+}
+
+impl Default for PromptConfig {
+    fn default() -> Self {
+        Self {
+            base_instructions: BaseInstructions::default(),
+            personality: None,
+            output_schema: None,
+            output_schema_strict: true,
+        }
+    }
+}
+
+fn reserialize_shell_outputs(items: &mut [ResponseItem]) {
+    let mut shell_call_ids: HashSet<String> = HashSet::new();
+
+    items.iter_mut().for_each(|item| match item {
+        ResponseItem::LocalShellCall { call_id, id, .. } => {
+            if let Some(identifier) = call_id.clone().or_else(|| id.clone()) {
+                shell_call_ids.insert(identifier);
+            }
+        }
+        ResponseItem::CustomToolCall {
+            id: _,
+            status: _,
+            call_id,
+            name,
+            input: _,
+        } => {
+            if name == "apply_patch" {
+                shell_call_ids.insert(call_id.clone());
+            }
+        }
+        ResponseItem::FunctionCall { name, call_id, .. }
+            if is_shell_tool_name(name) || name == "apply_patch" =>
+        {
+            shell_call_ids.insert(call_id.clone());
+        }
+        ResponseItem::FunctionCallOutput {
+            call_id, output, ..
+        }
+        | ResponseItem::CustomToolCallOutput {
+            call_id, output, ..
+        } => {
+            if shell_call_ids.remove(call_id)
+                && let Some(structured) = output
+                    .text_content()
+                    .and_then(parse_structured_shell_output)
+            {
+                output.body = FunctionCallOutputBody::Text(structured);
+            }
+        }
+        _ => {}
+    });
+}
+
+fn is_shell_tool_name(name: &str) -> bool {
+    matches!(name, "shell" | "container.exec")
+}
+
+#[derive(Deserialize)]
+struct ExecOutputJson {
+    output: String,
+    metadata: ExecOutputMetadataJson,
+}
+
+#[derive(Deserialize)]
+struct ExecOutputMetadataJson {
+    exit_code: i32,
+    duration_seconds: f32,
+}
+
+fn parse_structured_shell_output(raw: &str) -> Option<String> {
+    let parsed: ExecOutputJson = serde_json::from_str(raw).ok()?;
+    Some(build_structured_output(&parsed))
+}
+
+fn build_structured_output(parsed: &ExecOutputJson) -> String {
+    let mut sections = Vec::new();
+    let exit_code = parsed.metadata.exit_code;
+    sections.push(format!("Exit code: {exit_code}"));
+    let duration_seconds = parsed.metadata.duration_seconds;
+    sections.push(format!("Wall time: {duration_seconds} seconds"));
+
+    let mut output = parsed.output.clone();
+    if let Some((stripped, total_lines)) = strip_total_output_header(&parsed.output) {
+        sections.push(format!("Total output lines: {total_lines}"));
+        output = stripped.to_string();
+    }
+
+    sections.push("Output:".to_string());
+    sections.push(output);
+    sections.join("\n")
+}
+
+fn strip_total_output_header(output: &str) -> Option<(&str, u32)> {
+    let after_prefix = output.strip_prefix("Total output lines: ")?;
+    let (total_segment, remainder) = after_prefix.split_once('\n')?;
+    let total_lines = total_segment.parse::<u32>().ok()?;
+    let remainder = remainder.strip_prefix('\n').unwrap_or(remainder);
+    Some((remainder, total_lines))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Prompt;
+    use codex_protocol::models::FunctionCallOutputPayload;
+    use codex_protocol::models::ResponseItem;
+    use codex_tools::FreeformTool;
+    use codex_tools::FreeformToolFormat;
+    use codex_tools::ToolSpec;
+    use pretty_assertions::assert_eq;
+
+    #[test]
+    fn reserializes_shell_outputs_for_function_and_custom_tool_calls() {
+        let raw_output = r#"{"output":"hello","metadata":{"exit_code":0,"duration_seconds":0.5}}"#;
+        let expected_output = "Exit code: 0\nWall time: 0.5 seconds\nOutput:\nhello";
+        let prompt = Prompt {
+            input: vec![
+                ResponseItem::FunctionCall {
+                    id: None,
+                    name: "shell".to_string(),
+                    namespace: None,
+                    arguments: "{}".to_string(),
+                    call_id: "call-1".to_string(),
+                },
+                ResponseItem::FunctionCallOutput {
+                    call_id: "call-1".to_string(),
+                    output: FunctionCallOutputPayload::from_text(raw_output.to_string()),
+                },
+                ResponseItem::CustomToolCall {
+                    id: None,
+                    status: None,
+                    call_id: "call-2".to_string(),
+                    name: "apply_patch".to_string(),
+                    input: "*** Begin Patch".to_string(),
+                },
+                ResponseItem::CustomToolCallOutput {
+                    call_id: "call-2".to_string(),
+                    name: None,
+                    output: FunctionCallOutputPayload::from_text(raw_output.to_string()),
+                },
+            ],
+            tools: vec![ToolSpec::Freeform(FreeformTool {
+                name: "apply_patch".to_string(),
+                description: "patch".to_string(),
+                format: FreeformToolFormat {
+                    r#type: "grammar".to_string(),
+                    syntax: "lark".to_string(),
+                    definition: "patch".to_string(),
+                },
+            })],
+            ..Prompt::default()
+        };
+
+        assert_eq!(
+            prompt.get_formatted_input(),
+            vec![
+                ResponseItem::FunctionCall {
+                    id: None,
+                    name: "shell".to_string(),
+                    namespace: None,
+                    arguments: "{}".to_string(),
+                    call_id: "call-1".to_string(),
+                },
+                ResponseItem::FunctionCallOutput {
+                    call_id: "call-1".to_string(),
+                    output: FunctionCallOutputPayload::from_text(expected_output.to_string()),
+                },
+                ResponseItem::CustomToolCall {
+                    id: None,
+                    status: None,
+                    call_id: "call-2".to_string(),
+                    name: "apply_patch".to_string(),
+                    input: "*** Begin Patch".to_string(),
+                },
+                ResponseItem::CustomToolCallOutput {
+                    call_id: "call-2".to_string(),
+                    name: None,
+                    output: FunctionCallOutputPayload::from_text(expected_output.to_string()),
+                },
+            ]
+        );
+    }
+}

codex-rs/kernel/src/sampling.rs

@@ -0,0 +1,170 @@
+use crate::KernelToolExecutor;
+use crate::Prompt;
+use crate::PromptConfig;
+use crate::ToolConfig;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result;
+use codex_protocol::models::ResponseItem;
+use codex_protocol::protocol::RateLimitSnapshot;
+use rand::Rng;
+use std::time::Duration;
+use tokio_util::sync::CancellationToken;
+
+const INITIAL_DELAY_MS: u64 = 200;
+const BACKOFF_FACTOR: f64 = 2.0;
+
+/// Host-facing bundle created once before the kernel begins retrying a logical sampling request.
+///
+/// Implementations should place retry-stable state here so retries can rebuild prompts without
+/// repeating expensive setup work such as tool routing or per-turn worker startup.
+pub struct PreparedSamplingRequest<Runtime, Tools> {
+    pub prompt_config: PromptConfig,
+    pub tool_config: ToolConfig,
+    pub runtime: Runtime,
+    pub tools: Tools,
+}
+
+/// Final outcome of a logical sampling request after the model/tool loop reaches a stable state.
+#[derive(Debug)]
+pub struct SamplingRequestResult {
+    pub needs_follow_up: bool,
+    pub last_agent_message: Option<String>,
+}
+
+/// Host adapter for the kernel's retryable sampling loop.
+///
+/// Implementations provide the concrete model transport, tool runtime, history access, and
+/// user-facing retry notifications while allowing `codex-kernel` to stay agnostic of local
+/// filesystem, network, or tool-execution details.
+#[allow(async_fn_in_trait)]
+pub trait SamplingLoopHost {
+    type ClientSession;
+    type Runtime;
+    type Tools: KernelToolExecutor;
+
+    async fn prepare_sampling_request(
+        &self,
+        input: &[ResponseItem],
+        cancellation_token: &CancellationToken,
+    ) -> Result<PreparedSamplingRequest<Self::Runtime, Self::Tools>>;
+
+    async fn history_prompt_input(&self) -> Vec<ResponseItem>;
+
+    async fn run_single_sampling_request(
+        &self,
+        runtime: &Self::Runtime,
+        tools: &Self::Tools,
+        client_session: &mut Self::ClientSession,
+        prompt: &Prompt,
+        cancellation_token: CancellationToken,
+    ) -> Result<SamplingRequestResult>;
+
+    fn stream_max_retries(&self) -> u64;
+
+    fn try_switch_fallback_transport(&self, client_session: &mut Self::ClientSession) -> bool;
+
+    fn should_notify_stream_retry(&self, retries: u64, err: &CodexErr) -> bool;
+
+    async fn handle_context_window_exceeded(&self);
+
+    async fn handle_usage_limit_reached(&self, rate_limits: Option<RateLimitSnapshot>);
+
+    async fn notify_fallback_to_http(&self, err: &CodexErr);
+
+    async fn notify_stream_retry(
+        &self,
+        retries: u64,
+        max_retries: u64,
+        delay: Duration,
+        err: &CodexErr,
+    );
+}
+
+pub async fn run_sampling_request_loop<Host>(
+    host: &Host,
+    client_session: &mut Host::ClientSession,
+    initial_input: Vec<ResponseItem>,
+    cancellation_token: CancellationToken,
+) -> Result<SamplingRequestResult>
+where
+    Host: SamplingLoopHost,
+{
+    let prepared = host
+        .prepare_sampling_request(initial_input.as_slice(), &cancellation_token)
+        .await?;
+    let mut retries = 0;
+    let mut initial_input = Some(initial_input);
+    loop {
+        let prompt_input = if let Some(input) = initial_input.take() {
+            input
+        } else {
+            host.history_prompt_input().await
+        };
+        let prompt = prepared
+            .prompt_config
+            .build_prompt(prompt_input, &prepared.tool_config);
+        let err = match host
+            .run_single_sampling_request(
+                &prepared.runtime,
+                &prepared.tools,
+                client_session,
+                &prompt,
+                cancellation_token.child_token(),
+            )
+            .await
+        {
+            Ok(output) => {
+                return Ok(output);
+            }
+            Err(CodexErr::ContextWindowExceeded) => {
+                host.handle_context_window_exceeded().await;
+                return Err(CodexErr::ContextWindowExceeded);
+            }
+            Err(CodexErr::UsageLimitReached(error)) => {
+                host.handle_usage_limit_reached(
+                    error
+                        .rate_limits
+                        .as_ref()
+                        .map(|snapshot| (**snapshot).clone()),
+                )
+                .await;
+                return Err(CodexErr::UsageLimitReached(error));
+            }
+            Err(err) => err,
+        };
+
+        if !err.is_retryable() {
+            return Err(err);
+        }
+
+        let max_retries = host.stream_max_retries();
+        if retries >= max_retries && host.try_switch_fallback_transport(client_session) {
+            host.notify_fallback_to_http(&err).await;
+            retries = 0;
+            continue;
+        }
+        if retries < max_retries {
+            retries += 1;
+            let delay = match &err {
+                CodexErr::Stream(_, requested_delay) => {
+                    requested_delay.unwrap_or_else(|| backoff(retries))
+                }
+                _ => backoff(retries),
+            };
+            if host.should_notify_stream_retry(retries, &err) {
+                host.notify_stream_retry(retries, max_retries, delay, &err)
+                    .await;
+            }
+            tokio::time::sleep(delay).await;
+        } else {
+            return Err(err);
+        }
+    }
+}
+
+fn backoff(attempt: u64) -> Duration {
+    let exp = BACKOFF_FACTOR.powi(attempt.saturating_sub(1) as i32);
+    let base = (INITIAL_DELAY_MS as f64 * exp) as u64;
+    let jitter = rand::rng().random_range(0.9..1.1);
+    Duration::from_millis((base as f64 * jitter) as u64)
+}

codex-rs/kernel/src/tools.rs

@@ -0,0 +1,226 @@
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result;
+use codex_protocol::models::ResponseInputItem;
+use codex_protocol::models::ResponseItem;
+use codex_tools::ToolSpec;
+use std::fmt;
+use tokio_util::sync::CancellationToken;
+
+/// Retry-stable tool registration for a logical model request.
+///
+/// The host chooses which tools are exposed to the model and whether the model
+/// may issue them in parallel. `codex-kernel` treats this as opaque model-facing
+/// configuration and uses it to rebuild prompts across retries.
+#[derive(Debug, Clone, Default)]
+pub struct ToolConfig {
+    pub tools: Vec<ToolSpec>,
+    pub parallel_tool_calls: bool,
+}
+
+/// Metadata a host-specific tool call must expose so `codex-kernel` can
+/// synthesize a model-visible failure response when execution cannot complete.
+pub trait KernelToolCall: Clone + Send + Sync {
+    fn error_response(&self, message: String) -> ResponseInputItem;
+}
+
+/// Non-fatal tool execution failures are surfaced back to the model as a
+/// complementary output item, while fatal failures abort the turn.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ToolCallError {
+    RespondToModel(String),
+    Fatal(String),
+}
+
+impl fmt::Display for ToolCallError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::RespondToModel(message) => write!(f, "{message}"),
+            Self::Fatal(message) => write!(f, "Fatal error: {message}"),
+        }
+    }
+}
+
+/// Classification of a completed model output item relative to the host's tool
+/// catalog.
+pub enum CompletedResponseItem<Call> {
+    NonTool(ResponseItem),
+    ToolCall {
+        item: ResponseItem,
+        call: Call,
+    },
+    ImmediateResponse {
+        item: ResponseItem,
+        response: ResponseInputItem,
+    },
+}
+
+/// Host-provided tool adapter used by `codex-kernel`'s default model/tool loop.
+///
+/// Implementations own tool-call parsing and execution, while
+/// `codex-kernel` owns the generic behavior of attempting a complementary
+/// tool output item for every tool call the model emits.
+#[allow(async_fn_in_trait)]
+pub trait KernelToolExecutor {
+    type Call: KernelToolCall;
+
+    async fn classify_response_item(
+        &self,
+        item: ResponseItem,
+    ) -> Result<CompletedResponseItem<Self::Call>>;
+
+    async fn execute_tool_call(
+        &self,
+        call: Self::Call,
+        cancellation_token: CancellationToken,
+    ) -> std::result::Result<ResponseInputItem, ToolCallError>;
+}
+
+pub async fn execute_tool_call_with_default_output<Executor>(
+    executor: &Executor,
+    call: Executor::Call,
+    cancellation_token: CancellationToken,
+) -> Result<ResponseInputItem>
+where
+    Executor: KernelToolExecutor,
+{
+    match executor
+        .execute_tool_call(call.clone(), cancellation_token)
+        .await
+    {
+        Ok(response) => Ok(response),
+        Err(ToolCallError::RespondToModel(message)) => Ok(call.error_response(message)),
+        Err(ToolCallError::Fatal(message)) => Err(CodexErr::Fatal(message)),
+    }
+}
+
+pub fn response_input_to_response_item(input: &ResponseInputItem) -> Option<ResponseItem> {
+    match input {
+        ResponseInputItem::FunctionCallOutput { call_id, output } => {
+            Some(ResponseItem::FunctionCallOutput {
+                call_id: call_id.clone(),
+                output: output.clone(),
+            })
+        }
+        ResponseInputItem::CustomToolCallOutput {
+            call_id,
+            name,
+            output,
+        } => Some(ResponseItem::CustomToolCallOutput {
+            call_id: call_id.clone(),
+            name: name.clone(),
+            output: output.clone(),
+        }),
+        ResponseInputItem::McpToolCallOutput { call_id, output } => {
+            let output = output.as_function_call_output_payload();
+            Some(ResponseItem::FunctionCallOutput {
+                call_id: call_id.clone(),
+                output,
+            })
+        }
+        ResponseInputItem::ToolSearchOutput {
+            call_id,
+            status,
+            execution,
+            tools,
+        } => Some(ResponseItem::ToolSearchOutput {
+            call_id: Some(call_id.clone()),
+            status: status.clone(),
+            execution: execution.clone(),
+            tools: tools.clone(),
+        }),
+        _ => None,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::CompletedResponseItem;
+    use super::KernelToolCall;
+    use super::KernelToolExecutor;
+    use super::ToolCallError;
+    use super::execute_tool_call_with_default_output;
+    use codex_protocol::error::Result;
+    use codex_protocol::models::FunctionCallOutputPayload;
+    use codex_protocol::models::ResponseInputItem;
+    use codex_protocol::models::ResponseItem;
+    use pretty_assertions::assert_eq;
+    use tokio_util::sync::CancellationToken;
+
+    #[derive(Clone, Debug, PartialEq, Eq)]
+    struct TestCall {
+        call_id: String,
+    }
+
+    impl KernelToolCall for TestCall {
+        fn error_response(&self, message: String) -> ResponseInputItem {
+            ResponseInputItem::FunctionCallOutput {
+                call_id: self.call_id.clone(),
+                output: FunctionCallOutputPayload::from_text(message),
+            }
+        }
+    }
+
+    struct TestExecutor;
+
+    impl KernelToolExecutor for TestExecutor {
+        type Call = TestCall;
+
+        async fn classify_response_item(
+            &self,
+            item: ResponseItem,
+        ) -> Result<CompletedResponseItem<Self::Call>> {
+            Ok(CompletedResponseItem::NonTool(item))
+        }
+
+        async fn execute_tool_call(
+            &self,
+            call: Self::Call,
+            _cancellation_token: CancellationToken,
+        ) -> std::result::Result<ResponseInputItem, ToolCallError> {
+            match call.call_id.as_str() {
+                "ok" => Ok(ResponseInputItem::FunctionCallOutput {
+                    call_id: call.call_id,
+                    output: FunctionCallOutputPayload::from_text("done".to_string()),
+                }),
+                "retry" => Err(ToolCallError::RespondToModel("try again".to_string())),
+                _ => Err(ToolCallError::Fatal("boom".to_string())),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn execute_tool_call_with_default_output_surfaces_non_fatal_errors() {
+        let response = execute_tool_call_with_default_output(
+            &TestExecutor,
+            TestCall {
+                call_id: "retry".to_string(),
+            },
+            CancellationToken::new(),
+        )
+        .await
+        .expect("tool response");
+
+        assert_eq!(
+            response,
+            ResponseInputItem::FunctionCallOutput {
+                call_id: "retry".to_string(),
+                output: FunctionCallOutputPayload::from_text("try again".to_string()),
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn execute_tool_call_with_default_output_propagates_fatal_errors() {
+        let err = execute_tool_call_with_default_output(
+            &TestExecutor,
+            TestCall {
+                call_id: "fatal".to_string(),
+            },
+            CancellationToken::new(),
+        )
+        .await
+        .expect_err("fatal tool error");
+
+        assert_eq!(err.to_string(), "Fatal error: boom");
+    }
+}