Files
codex/codex-rs/core/src/client_common.rs
Charlie Marsh beeba1d2fc Share model-visible tool specs across prompts (#37807)
## Why

Building a prompt cloned every model-visible `ToolSpec`, even though the tool
set is immutable for the lifetime of its router.

## What changed

- Store model-visible tool specs as an `Arc<[ToolSpec]>` in `ToolRouter` and
  `Prompt` so prompt construction only clones the shared pointer.
- Keep separately built routers on distinct shared slices so refreshed tool
  sets remain independent.

## Testing

Extend router tests to verify allocation sharing within a router, allocation
separation across rebuilt routers, and the existing deferred-tool filtering.

GitOrigin-RevId: 1e04cd7c4d1b3fa0b494c5c51670232d38ba8ebb
2026-08-10 10:50:49 +00:00

129 lines
4.0 KiB
Rust

pub use codex_api::ResponseEvent;
use codex_protocol::error::Result;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ResponseItem;
use codex_tools::ToolSpec;
use futures::Stream;
use serde_json::Value;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
/// API request payload for a single model turn
#[derive(Debug, Clone)]
pub struct Prompt {
/// Conversation context input items.
pub input: Vec<ResponseItem>,
/// Tools available to the model, including additional tools sourced from
/// external MCP servers.
pub(crate) tools: Arc<[ToolSpec]>,
/// Whether parallel tool calls are permitted for this prompt.
pub(crate) parallel_tool_calls: bool,
pub base_instructions: BaseInstructions,
/// Optional the output schema for the model's response.
pub output_schema: Option<Value>,
/// Whether the Responses API should strictly validate `output_schema`.
pub output_schema_strict: bool,
}
impl Default for Prompt {
fn default() -> Self {
Self {
input: Vec::new(),
tools: Arc::default(),
parallel_tool_calls: false,
base_instructions: BaseInstructions::default(),
output_schema: None,
output_schema_strict: true,
}
}
}
impl Prompt {
pub(crate) fn get_formatted_input_for_request(
&self,
use_responses_lite: bool,
) -> Vec<ResponseItem> {
let mut input = self.input.clone();
if use_responses_lite {
strip_image_details(&mut input);
}
input
}
}
fn strip_image_details(items: &mut [ResponseItem]) {
for item in items {
match item {
ResponseItem::Message { content, .. } => {
for content_item in content {
if let ContentItem::InputImage { detail, .. } = content_item {
*detail = None;
}
}
}
ResponseItem::FunctionCallOutput { output, .. }
| ResponseItem::CustomToolCallOutput { output, .. } => {
if let Some(content) = output.content_items_mut() {
for content_item in content {
if let FunctionCallOutputContentItem::InputImage { detail, .. } =
content_item
{
*detail = None;
}
}
}
}
ResponseItem::AdditionalTools { .. }
| ResponseItem::Reasoning { .. }
| ResponseItem::AgentMessage { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::ToolSearchCall { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::ToolSearchOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::ImageGenerationCall { .. }
| ResponseItem::Compaction { .. }
| ResponseItem::CompactionTrigger { .. }
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other => {}
}
}
}
pub struct ResponseStream {
pub(crate) rx_event: mpsc::Receiver<Result<ResponseEvent>>,
/// Signals the mapper task that the consumer stopped polling before the
/// provider stream reached its own terminal event.
pub(crate) consumer_dropped: CancellationToken,
}
impl Stream for ResponseStream {
type Item = Result<ResponseEvent>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.rx_event.poll_recv(cx)
}
}
impl Drop for ResponseStream {
fn drop(&mut self) {
self.consumer_dropped.cancel();
}
}
#[cfg(test)]
#[path = "client_common_tests.rs"]
mod tests;