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
This commit is contained in:
Charlie Marsh
2026-08-10 10:45:49 +00:00
committed by copyberry
parent c8e6e8555c
commit beeba1d2fc
4 changed files with 16 additions and 8 deletions

View File

@@ -8,6 +8,7 @@ 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;
@@ -21,7 +22,7 @@ pub struct Prompt {
/// Tools available to the model, including additional tools sourced from
/// external MCP servers.
pub(crate) tools: Vec<ToolSpec>,
pub(crate) tools: Arc<[ToolSpec]>,
/// Whether parallel tool calls are permitted for this prompt.
pub(crate) parallel_tool_calls: bool,
@@ -39,7 +40,7 @@ impl Default for Prompt {
fn default() -> Self {
Self {
input: Vec::new(),
tools: Vec::new(),
tools: Arc::default(),
parallel_tool_calls: false,
base_instructions: BaseInstructions::default(),
output_schema: None,

View File

@@ -67,7 +67,7 @@ pub(crate) fn tool_log_payload<'a>(
pub struct ToolRouter {
registry: ToolRegistry,
model_visible_specs: Vec<ToolSpec>,
model_visible_specs: Arc<[ToolSpec]>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -102,12 +102,12 @@ impl ToolRouter {
pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec<ToolSpec>) -> Self {
Self {
registry,
model_visible_specs,
model_visible_specs: model_visible_specs.into(),
}
}
pub(crate) fn model_visible_specs(&self) -> Vec<ToolSpec> {
self.model_visible_specs.clone()
pub(crate) fn model_visible_specs(&self) -> Arc<[ToolSpec]> {
Arc::clone(&self.model_visible_specs)
}
pub(crate) fn deferred_tool_namespaces(&self) -> BTreeMap<String, String> {

View File

@@ -471,9 +471,11 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> {
Vec::new(),
&dynamic_tools,
);
let visible_specs = router.model_visible_specs();
assert!(Arc::ptr_eq(&visible_specs, &router.model_visible_specs()));
assert_eq!(
namespace_function_names(&router.model_visible_specs(), "codex_app"),
namespace_function_names(&visible_specs, "codex_app"),
vec![visible_tool.to_string()]
);
assert_eq!(
@@ -481,6 +483,11 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> {
BTreeMap::from([("codex_app".to_string(), "Codex app tools.".to_string())])
);
let updated_router = test_tool_router(step_context.as_ref(), Vec::new(), Vec::new(), &[]);
let updated_specs = updated_router.model_visible_specs();
assert!(!Arc::ptr_eq(&visible_specs, &updated_specs));
assert!(namespace_function_names(&updated_specs, "codex_app").is_empty());
Ok(())
}

View File

@@ -80,7 +80,7 @@ struct ToolPlanProbe {
impl ToolPlanProbe {
fn from_router(router: ToolRouter) -> Self {
let visible_specs = router.model_visible_specs();
let visible_specs = router.model_visible_specs().to_vec();
let visible_names = visible_specs
.iter()
.map(|spec| spec.name().to_string())