codex-rs/code-mode-protocol/src/host/payload.rs

use std::num::TryFromIntError;
use std::time::Duration;

use codex_protocol::ToolName;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value as JsonValue;

use crate::CellId;
use crate::CodeModeNestedToolCall;
use crate::CodeModeSessionCellExecutionLimits;
use crate::CodeModeToolKind;
use crate::ExecuteRequest;
use crate::FunctionCallOutputContentItem;
use crate::ImageDetail;
use crate::MissingCodeModeHostDuration;
use crate::RuntimeResponse;
use crate::ToolDefinition;
use crate::WaitOutcome;
use crate::WaitRequest;

/// The per-cell execution limits carried by a V1 session-open request.
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct WireSessionCellExecutionLimits {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_yield_time_ms: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_heap_size_bytes: Option<u64>,
}

impl TryFrom<CodeModeSessionCellExecutionLimits> for WireSessionCellExecutionLimits {
    type Error = TryFromIntError;

    fn try_from(value: CodeModeSessionCellExecutionLimits) -> Result<Self, Self::Error> {
        Ok(Self {
            max_yield_time_ms: value.max_yield_time_ms,
            max_heap_size_bytes: value.max_heap_size_bytes.map(u64::try_from).transpose()?,
        })
    }
}

impl TryFrom<WireSessionCellExecutionLimits> for CodeModeSessionCellExecutionLimits {
    type Error = TryFromIntError;

    fn try_from(value: WireSessionCellExecutionLimits) -> Result<Self, Self::Error> {
        Ok(Self {
            max_yield_time_ms: value.max_yield_time_ms,
            max_heap_size_bytes: value.max_heap_size_bytes.map(usize::try_from).transpose()?,
        })
    }
}

/// A cell identifier with a wire representation owned by protocol V1.
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(transparent)]
pub struct WireCellId(String);

impl WireCellId {
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl From<CellId> for WireCellId {
    fn from(value: CellId) -> Self {
        Self(value.as_str().to_string())
    }
}

impl From<&CellId> for WireCellId {
    fn from(value: &CellId) -> Self {
        Self(value.as_str().to_string())
    }
}

impl From<WireCellId> for CellId {
    fn from(value: WireCellId) -> Self {
        Self::new(value.0)
    }
}

/// The V1 wire representation of a tool's stable name.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireToolName {
    pub name: String,
    pub namespace: Option<String>,
}

impl From<ToolName> for WireToolName {
    fn from(value: ToolName) -> Self {
        Self {
            name: value.name,
            namespace: value.namespace,
        }
    }
}

impl From<WireToolName> for ToolName {
    fn from(value: WireToolName) -> Self {
        Self::new(value.namespace, value.name)
    }
}

/// The tool invocation shape supported by protocol V1.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum WireToolKind {
    Function,
    Freeform,
}

impl From<CodeModeToolKind> for WireToolKind {
    fn from(value: CodeModeToolKind) -> Self {
        match value {
            CodeModeToolKind::Function => Self::Function,
            CodeModeToolKind::Freeform => Self::Freeform,
        }
    }
}

impl From<WireToolKind> for CodeModeToolKind {
    fn from(value: WireToolKind) -> Self {
        match value {
            WireToolKind::Function => Self::Function,
            WireToolKind::Freeform => Self::Freeform,
        }
    }
}

/// A V1 tool definition embedded in an execute request.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireToolDefinition {
    pub name: String,
    pub tool_name: WireToolName,
    pub description: String,
    pub kind: WireToolKind,
    pub input_schema: Option<JsonValue>,
    pub output_schema: Option<JsonValue>,
}

impl From<ToolDefinition> for WireToolDefinition {
    fn from(value: ToolDefinition) -> Self {
        Self {
            name: value.name,
            tool_name: value.tool_name.into(),
            description: value.description,
            kind: value.kind.into(),
            input_schema: value.input_schema,
            output_schema: value.output_schema,
        }
    }
}

impl From<WireToolDefinition> for ToolDefinition {
    fn from(value: WireToolDefinition) -> Self {
        Self {
            name: value.name,
            tool_name: value.tool_name.into(),
            description: value.description,
            kind: value.kind.into(),
            input_schema: value.input_schema,
            output_schema: value.output_schema,
        }
    }
}

/// The complete execute request shape supported by protocol V1.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireExecuteRequest {
    pub tool_call_id: String,
    pub enabled_tools: Vec<WireToolDefinition>,
    pub source: String,
    pub yield_time_ms: Option<u64>,
    pub max_output_tokens: Option<i32>,
}

impl TryFrom<ExecuteRequest> for WireExecuteRequest {
    type Error = TryFromIntError;

    fn try_from(value: ExecuteRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            tool_call_id: value.tool_call_id,
            enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(),
            source: value.source,
            yield_time_ms: value.yield_time_ms,
            max_output_tokens: value.max_output_tokens.map(i32::try_from).transpose()?,
        })
    }
}

impl TryFrom<WireExecuteRequest> for ExecuteRequest {
    type Error = TryFromIntError;

    fn try_from(value: WireExecuteRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            tool_call_id: value.tool_call_id,
            enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(),
            source: value.source,
            yield_time_ms: value.yield_time_ms,
            max_output_tokens: value.max_output_tokens.map(usize::try_from).transpose()?,
        })
    }
}

/// The complete wait request shape supported by protocol V1.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireWaitRequest {
    pub cell_id: WireCellId,
    pub yield_time_ms: u64,
}

impl From<WaitRequest> for WireWaitRequest {
    fn from(value: WaitRequest) -> Self {
        Self {
            cell_id: value.cell_id.into(),
            yield_time_ms: value.yield_time_ms,
        }
    }
}

impl From<WireWaitRequest> for WaitRequest {
    fn from(value: WireWaitRequest) -> Self {
        Self {
            cell_id: value.cell_id.into(),
            yield_time_ms: value.yield_time_ms,
        }
    }
}

/// Image detail values accepted in a V1 runtime response.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum WireImageDetail {
    Auto,
    Low,
    High,
    Original,
}

impl From<ImageDetail> for WireImageDetail {
    fn from(value: ImageDetail) -> Self {
        match value {
            ImageDetail::Auto => Self::Auto,
            ImageDetail::Low => Self::Low,
            ImageDetail::High => Self::High,
            ImageDetail::Original => Self::Original,
        }
    }
}

impl From<WireImageDetail> for ImageDetail {
    fn from(value: WireImageDetail) -> Self {
        match value {
            WireImageDetail::Auto => Self::Auto,
            WireImageDetail::Low => Self::Low,
            WireImageDetail::High => Self::High,
            WireImageDetail::Original => Self::Original,
        }
    }
}

/// One output item emitted by a V1 runtime response.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "type", rename_all = "snake_case")]
pub enum WireContentItem {
    InputText {
        text: String,
    },
    InputImage {
        image_url: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        detail: Option<WireImageDetail>,
    },
    InputAudio {
        audio_url: String,
    },
}

impl From<FunctionCallOutputContentItem> for WireContentItem {
    fn from(value: FunctionCallOutputContentItem) -> Self {
        match value {
            FunctionCallOutputContentItem::InputText { text } => Self::InputText { text },
            FunctionCallOutputContentItem::InputImage { image_url, detail } => Self::InputImage {
                image_url,
                detail: detail.map(Into::into),
            },
            FunctionCallOutputContentItem::InputAudio { audio_url } => {
                Self::InputAudio { audio_url }
            }
        }
    }
}

impl From<WireContentItem> for FunctionCallOutputContentItem {
    fn from(value: WireContentItem) -> Self {
        match value {
            WireContentItem::InputText { text } => Self::InputText { text },
            WireContentItem::InputImage { image_url, detail } => Self::InputImage {
                image_url,
                detail: detail.map(Into::into),
            },
            WireContentItem::InputAudio { audio_url } => Self::InputAudio { audio_url },
        }
    }
}

/// Runtime output returned over the V1 host connection.
///
/// Host time is required and covers this request, not the cell lifetime. The
/// app-server and host run at the same version; no negotiation is needed.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub enum WireRuntimeResponse {
    Yielded {
        cell_id: WireCellId,
        content_items: Vec<WireContentItem>,
        code_mode_host_duration_ns: u64,
    },
    Terminated {
        cell_id: WireCellId,
        content_items: Vec<WireContentItem>,
        code_mode_host_duration_ns: u64,
    },
    Result {
        cell_id: WireCellId,
        content_items: Vec<WireContentItem>,
        error_text: Option<String>,
        code_mode_host_duration_ns: u64,
    },
}

impl TryFrom<RuntimeResponse> for WireRuntimeResponse {
    type Error = MissingCodeModeHostDuration;

    /// Preserves the response's timing; the host handler must record it first.
    fn try_from(value: RuntimeResponse) -> Result<Self, Self::Error> {
        Ok(match value {
            RuntimeResponse::Yielded {
                cell_id,
                content_items,
                code_mode_host_duration,
            } => {
                let code_mode_host_duration =
                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
                Self::Yielded {
                    cell_id: cell_id.into(),
                    content_items: content_items.into_iter().map(Into::into).collect(),
                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
                        .unwrap_or(u64::MAX),
                }
            }
            RuntimeResponse::Terminated {
                cell_id,
                content_items,
                code_mode_host_duration,
            } => {
                let code_mode_host_duration =
                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
                Self::Terminated {
                    cell_id: cell_id.into(),
                    content_items: content_items.into_iter().map(Into::into).collect(),
                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
                        .unwrap_or(u64::MAX),
                }
            }
            RuntimeResponse::Result {
                cell_id,
                content_items,
                error_text,
                code_mode_host_duration,
            } => {
                let code_mode_host_duration =
                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
                Self::Result {
                    cell_id: cell_id.into(),
                    content_items: content_items.into_iter().map(Into::into).collect(),
                    error_text,
                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
                        .unwrap_or(u64::MAX),
                }
            }
        })
    }
}

impl From<WireRuntimeResponse> for RuntimeResponse {
    fn from(value: WireRuntimeResponse) -> Self {
        match value {
            WireRuntimeResponse::Yielded {
                cell_id,
                content_items,
                code_mode_host_duration_ns,
            } => Self::Yielded {
                cell_id: cell_id.into(),
                content_items: content_items.into_iter().map(Into::into).collect(),
                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
            },
            WireRuntimeResponse::Terminated {
                cell_id,
                content_items,
                code_mode_host_duration_ns,
            } => Self::Terminated {
                cell_id: cell_id.into(),
                content_items: content_items.into_iter().map(Into::into).collect(),
                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
            },
            WireRuntimeResponse::Result {
                cell_id,
                content_items,
                error_text,
                code_mode_host_duration_ns,
            } => Self::Result {
                cell_id: cell_id.into(),
                content_items: content_items.into_iter().map(Into::into).collect(),
                error_text,
                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
            },
        }
    }
}

/// Whether a waited-for cell remained live in protocol V1.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub enum WireWaitOutcome {
    LiveCell(WireRuntimeResponse),
    MissingCell(WireRuntimeResponse),
}

impl TryFrom<WaitOutcome> for WireWaitOutcome {
    type Error = MissingCodeModeHostDuration;

    fn try_from(value: WaitOutcome) -> Result<Self, Self::Error> {
        Ok(match value {
            WaitOutcome::LiveCell(response) => Self::LiveCell(response.try_into()?),
            WaitOutcome::MissingCell(response) => Self::MissingCell(response.try_into()?),
        })
    }
}

impl From<WireWaitOutcome> for WaitOutcome {
    fn from(value: WireWaitOutcome) -> Self {
        match value {
            WireWaitOutcome::LiveCell(response) => Self::LiveCell(response.into()),
            WireWaitOutcome::MissingCell(response) => Self::MissingCell(response.into()),
        }
    }
}

/// A nested tool invocation sent over the V1 host connection.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireNestedToolCall {
    pub cell_id: WireCellId,
    pub runtime_tool_call_id: String,
    pub tool_name: WireToolName,
    pub tool_kind: WireToolKind,
    pub input: Option<JsonValue>,
}

impl From<CodeModeNestedToolCall> for WireNestedToolCall {
    fn from(value: CodeModeNestedToolCall) -> Self {
        Self {
            cell_id: value.cell_id.into(),
            runtime_tool_call_id: value.runtime_tool_call_id,
            tool_name: value.tool_name.into(),
            tool_kind: value.tool_kind.into(),
            input: value.input,
        }
    }
}

impl From<WireNestedToolCall> for CodeModeNestedToolCall {
    fn from(value: WireNestedToolCall) -> Self {
        Self {
            cell_id: value.cell_id.into(),
            runtime_tool_call_id: value.runtime_tool_call_id,
            tool_name: value.tool_name.into(),
            tool_kind: value.tool_kind.into(),
            input: value.input,
        }
    }
}