| use codex_protocol::ToolName; |
| use serde::Deserialize; |
| use serde::Serialize; |
| use serde_json::Value as JsonValue; |
| use std::collections::BTreeMap; |
|
|
| use crate::PUBLIC_TOOL_NAME; |
| use crate::json_schema_types::render_json_schema_to_typescript; |
|
|
| const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1; |
| const DEFERRED_NESTED_TOOLS_GUIDANCE: &str = r#"Some deferred nested tools may be omitted from this description. They are still available on the global `tools` object and listed in `ALL_TOOLS`. |
| To find one, filter `ALL_TOOLS` by `name` and `description`."#; |
| const LEGACY_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument."#; |
| const UNIFIED_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string } | ImageContent)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`."#; |
| const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/compose tool calls |
| - Evaluates the provided JavaScript code in a fresh V8 isolate as an async module. |
| - All nested tools are available on the global `tools` object, for example `await tools.exec_command(...)`. Tool names are exposed as normalized JavaScript identifiers, for example `await tools.mcp__ologs__get_profile(...)`. |
| - Nested tool methods take either a string or an object as their input argument. |
| - Nested tools return either an object or a string, based on the description. |
| - Runs raw JavaScript -- no Node, no file system, no network access, no console. |
| - Accepts raw JavaScript source text, not JSON, quoted strings, or markdown code fences. |
| - You may optionally start the tool input with a first-line pragma like `// @exec: {"yield_time_ms": 10000, "max_output_tokens": 1000}`. |
| - `yield_time_ms` asks `exec` to yield early if the script is still running. Defaults to 10000 ms. |
| - `max_output_tokens` sets the token budget for direct `exec` results. Defaults to 10000 tokens. |
| - When the JS code is fully evaluated, the isolate's lifetime ends and unawaited promises are silently discarded. |
| |
| - Global helpers: |
| - `exit()`: Immediately ends the current script successfully (like an early return from the top level). |
| - `text(value: string | number | boolean | undefined | null)`: Appends a text item. Non-string values are stringified with `JSON.stringify(...)` when possible. |
| - `image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument. |
| - `audio(audioUrlOrItem: string | { audio_url: string } | AudioContent)`: Appends an audio item. `audio_url` should be a base64-encoded `data:` URL. To forward an MCP tool audio block, pass an individual `AudioContent` block from `result.content`, for example `audio(result.content[0])`. |
| - `generatedImage(result: { image_url: string; output_hint?: string })`: Appends an image-generation result and its optional output hint. HTTP(S) URLs are not supported. |
| - `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session. |
| - `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing. |
| - `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`. |
| - `setTimeout(callback: () => void, delayMs?: number)`: schedules a callback to run later and returns a timeout id. Pending timeouts do not keep `exec` alive by themselves; await an explicit promise if you need to wait for one. |
| - `clearTimeout(timeoutId?: number)`: cancels a timeout created by `setTimeout`. |
| - `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries. |
| - `yield_control()`: yields the accumulated output to the model immediately while the script keeps running."#; |
| const WAIT_DESCRIPTION_TEMPLATE: &str = r#"- Use `wait` only after `exec` returns `Script running with cell ID ...`. |
| - `cell_id` identifies the running `exec` cell to resume. |
| - `yield_time_ms` controls how long to wait for more output before yielding again. Defaults to 10000 ms. |
| - `max_tokens` limits how much new output this wait call returns. Defaults to 10000 tokens. |
| - `terminate: true` stops the running cell; false or omitted waits for output. |
| - `wait` returns only the new output since the last yield, or the final completion or termination result for that cell. |
| - If the cell is still running, `wait` may yield again with the same `cell_id`. |
| - If the cell has already finished, `wait` returns the completed result and closes the cell."#; |
| |
| const MCP_TYPESCRIPT_PREAMBLE: &str = r#"type Role = "user" | "assistant"; |
| type MetaObject = Record<string, unknown>; |
| type Annotations = { |
| audience?: Role[]; |
| priority?: number; |
| lastModified?: string; |
| }; |
| type Icon = { |
| src: string; |
| mimeType?: string; |
| sizes?: string[]; |
| theme?: "light" | "dark"; |
| }; |
| type TextResourceContents = { |
| uri: string; |
| mimeType?: string; |
| _meta?: MetaObject; |
| text: string; |
| }; |
| type BlobResourceContents = { |
| uri: string; |
| mimeType?: string; |
| _meta?: MetaObject; |
| blob: string; |
| }; |
| type TextContent = { |
| type: "text"; |
| text: string; |
| annotations?: Annotations; |
| _meta?: MetaObject; |
| }; |
| type ImageContent = { |
| type: "image"; |
| data: string; |
| mimeType: string; |
| annotations?: Annotations; |
| _meta?: MetaObject; |
| }; |
| type AudioContent = { |
| type: "audio"; |
| data: string; |
| mimeType: string; |
| annotations?: Annotations; |
| _meta?: MetaObject; |
| }; |
| type ResourceLink = { |
| icons?: Icon[]; |
| name: string; |
| title?: string; |
| uri: string; |
| description?: string; |
| mimeType?: string; |
| annotations?: Annotations; |
| size?: number; |
| _meta?: MetaObject; |
| type: "resource_link"; |
| }; |
| type EmbeddedResource = { |
| type: "resource"; |
| resource: TextResourceContents | BlobResourceContents; |
| annotations?: Annotations; |
| _meta?: MetaObject; |
| }; |
| type ContentBlock = |
| | TextContent |
| | ImageContent |
| | AudioContent |
| | ResourceLink |
| | EmbeddedResource; |
| type CallToolResult<TStructured = { [key: string]: unknown }> = { |
| _meta?: MetaObject; |
| content: ContentBlock[]; |
| isError?: boolean; |
| structuredContent?: TStructured; |
| [key: string]: unknown; |
| };"#; |
|
|
| pub const CODE_MODE_PRAGMA_PREFIX: &str = "// @exec:"; |
|
|
| #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] |
| #[serde(rename_all = "snake_case")] |
| pub enum CodeModeToolKind { |
| Function, |
| Freeform, |
| } |
|
|
| #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] |
| pub struct ToolDefinition { |
| pub name: String, |
| pub tool_name: ToolName, |
| pub description: String, |
| pub kind: CodeModeToolKind, |
| pub input_schema: Option<JsonValue>, |
| pub output_schema: Option<JsonValue>, |
| } |
|
|
| #[derive(Clone, Debug, Eq, PartialEq)] |
| pub struct ToolNamespaceDescription { |
| pub name: String, |
| pub description: String, |
| } |
|
|
| #[derive(Debug, Default, Deserialize, PartialEq, Eq)] |
| #[serde(deny_unknown_fields)] |
| struct CodeModeExecPragma { |
| #[serde(default)] |
| yield_time_ms: Option<u64>, |
| #[serde(default)] |
| max_output_tokens: Option<usize>, |
| } |
|
|
| #[derive(Debug, PartialEq, Eq)] |
| pub struct ParsedExecSource { |
| pub code: String, |
| pub yield_time_ms: Option<u64>, |
| pub max_output_tokens: Option<usize>, |
| } |
|
|
| pub fn parse_exec_source(input: &str) -> Result<ParsedExecSource, String> { |
| if input.trim().is_empty() { |
| return Err( |
| "exec expects raw JavaScript source text (non-empty). Provide JS only, optionally with first-line `// @exec: {\"yield_time_ms\": 10000, \"max_output_tokens\": 1000}`.".to_string(), |
| ); |
| } |
|
|
| let mut args = ParsedExecSource { |
| code: input.to_string(), |
| yield_time_ms: None, |
| max_output_tokens: None, |
| }; |
|
|
| let mut lines = input.splitn(2, '\n'); |
| let first_line = lines.next().unwrap_or_default(); |
| let rest = lines.next().unwrap_or_default(); |
| let trimmed = first_line.trim_start(); |
| let Some(pragma) = trimmed.strip_prefix(CODE_MODE_PRAGMA_PREFIX) else { |
| return Ok(args); |
| }; |
|
|
| if rest.trim().is_empty() { |
| return Err( |
| "exec pragma must be followed by JavaScript source on subsequent lines".to_string(), |
| ); |
| } |
|
|
| let directive = pragma.trim(); |
| if directive.is_empty() { |
| return Err( |
| "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`" |
| .to_string(), |
| ); |
| } |
|
|
| let value: serde_json::Value = serde_json::from_str(directive).map_err(|err| { |
| format!( |
| "exec pragma must be valid JSON with supported fields `yield_time_ms` and `max_output_tokens`: {err}" |
| ) |
| })?; |
| let object = value.as_object().ok_or_else(|| { |
| "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`" |
| .to_string() |
| })?; |
| for key in object.keys() { |
| match key.as_str() { |
| "yield_time_ms" | "max_output_tokens" => {} |
| _ => { |
| return Err(format!( |
| "exec pragma only supports `yield_time_ms` and `max_output_tokens`; got `{key}`" |
| )); |
| } |
| } |
| } |
|
|
| let pragma: CodeModeExecPragma = serde_json::from_value(value).map_err(|err| { |
| format!( |
| "exec pragma fields `yield_time_ms` and `max_output_tokens` must be non-negative safe integers: {err}" |
| ) |
| })?; |
| if pragma |
| .yield_time_ms |
| .is_some_and(|yield_time_ms| yield_time_ms > MAX_JS_SAFE_INTEGER) |
| { |
| return Err( |
| "exec pragma field `yield_time_ms` must be a non-negative safe integer".to_string(), |
| ); |
| } |
| if pragma.max_output_tokens.is_some_and(|max_output_tokens| { |
| u64::try_from(max_output_tokens) |
| .map(|max_output_tokens| max_output_tokens > MAX_JS_SAFE_INTEGER) |
| .unwrap_or(true) |
| }) { |
| return Err( |
| "exec pragma field `max_output_tokens` must be a non-negative safe integer".to_string(), |
| ); |
| } |
|
|
| args.code = rest.to_string(); |
| args.yield_time_ms = pragma.yield_time_ms; |
| args.max_output_tokens = pragma.max_output_tokens; |
| Ok(args) |
| } |
|
|
| pub fn is_code_mode_nested_tool(tool_name: &str) -> bool { |
| tool_name != crate::PUBLIC_TOOL_NAME && tool_name != crate::WAIT_TOOL_NAME |
| } |
|
|
| #[derive(Clone, Copy, Debug, Eq, PartialEq)] |
| pub enum ImageDetailVisibility { |
| Visible, |
| Hidden, |
| } |
|
|
| pub fn build_exec_tool_description( |
| enabled_tools: &[ToolDefinition], |
| deferred_tools: &[ToolDefinition], |
| namespace_descriptions: &BTreeMap<String, ToolNamespaceDescription>, |
| default_exec_yield_time_ms: u64, |
| code_mode_only: bool, |
| image_detail_visibility: ImageDetailVisibility, |
| ) -> String { |
| let mut sections = Vec::new(); |
| sections.push(EXEC_DESCRIPTION_TEMPLATE.replace( |
| "Defaults to 10000 ms.", |
| &format!("Defaults to {default_exec_yield_time_ms} ms."), |
| )); |
| if image_detail_visibility == ImageDetailVisibility::Hidden { |
| sections[0] = sections[0].replace( |
| LEGACY_IMAGE_HELPER_DESCRIPTION, |
| UNIFIED_IMAGE_HELPER_DESCRIPTION, |
| ); |
| } |
| if !deferred_tools.is_empty() { |
| sections.push(DEFERRED_NESTED_TOOLS_GUIDANCE.to_string()); |
| } |
| if !code_mode_only { |
| return sections.join("\n\n"); |
| } |
|
|
| let has_mcp_tools = enabled_tools |
| .iter() |
| .chain(deferred_tools) |
| .any(|tool| mcp_structured_content_schema(tool.output_schema.as_ref()).is_some()); |
| if has_mcp_tools { |
| sections.push(format!( |
| "Shared MCP Types:\n```ts\n{MCP_TYPESCRIPT_PREAMBLE}\n```" |
| )); |
| } |
|
|
| if !enabled_tools.is_empty() { |
| let mut current_namespace: Option<&str> = None; |
| let mut nested_tool_sections = Vec::with_capacity(enabled_tools.len()); |
|
|
| for tool in enabled_tools { |
| let name = tool.name.as_str(); |
| let nested_description = render_code_mode_sample_for_definition(tool); |
| let namespace_description = tool |
| .tool_name |
| .namespace |
| .as_ref() |
| .and_then(|namespace| namespace_descriptions.get(namespace)); |
| let next_namespace = namespace_description |
| .map(|namespace_description| namespace_description.name.as_str()); |
| if next_namespace != current_namespace { |
| if let Some(namespace_description) = namespace_description { |
| let namespace_description_text = namespace_description.description.trim(); |
| if !namespace_description_text.is_empty() { |
| nested_tool_sections.push(format!( |
| "## {}\n{namespace_description_text}", |
| namespace_description.name |
| )); |
| } |
| } |
| current_namespace = next_namespace; |
| } |
|
|
| let global_name = normalize_code_mode_identifier(name); |
| let nested_description = nested_description.trim(); |
| if nested_description.is_empty() { |
| nested_tool_sections.push(render_tool_heading(&global_name, name)); |
| } else { |
| nested_tool_sections.push(format!( |
| "{}\n{nested_description}", |
| render_tool_heading(&global_name, name) |
| )); |
| } |
| } |
|
|
| sections.push(nested_tool_sections.join("\n\n")); |
| } |
|
|
| sections.join("\n\n") |
| } |
|
|
| pub fn build_wait_tool_description() -> &'static str { |
| WAIT_DESCRIPTION_TEMPLATE |
| } |
|
|
| pub fn normalize_code_mode_identifier(tool_key: &str) -> String { |
| let mut identifier = String::new(); |
|
|
| for (index, ch) in tool_key.chars().enumerate() { |
| let is_valid = if index == 0 { |
| ch == '_' || ch == '$' || ch.is_ascii_alphabetic() |
| } else { |
| ch == '_' || ch == '$' || ch.is_ascii_alphanumeric() |
| }; |
|
|
| if is_valid { |
| identifier.push(ch); |
| } else { |
| identifier.push('_'); |
| } |
| } |
|
|
| if identifier.is_empty() { |
| "_".to_string() |
| } else { |
| identifier |
| } |
| } |
|
|
| pub fn augment_tool_definition(mut definition: ToolDefinition) -> ToolDefinition { |
| if definition.name != PUBLIC_TOOL_NAME { |
| definition.description = render_code_mode_sample_for_definition(&definition); |
| } |
| definition |
| } |
|
|
| pub fn enabled_tool_metadata(definition: &ToolDefinition) -> EnabledToolMetadata { |
| EnabledToolMetadata { |
| tool_name: definition.tool_name.clone(), |
| global_name: normalize_code_mode_identifier(&definition.name), |
| description: definition.description.clone(), |
| kind: definition.kind, |
| } |
| } |
|
|
| #[derive(Clone, Debug, Eq, PartialEq, Serialize)] |
| pub struct EnabledToolMetadata { |
| pub tool_name: ToolName, |
| pub global_name: String, |
| pub description: String, |
| pub kind: CodeModeToolKind, |
| } |
|
|
| pub fn render_code_mode_sample( |
| description: &str, |
| tool_name: &str, |
| input_name: &str, |
| input_type: String, |
| output_type: String, |
| ) -> String { |
| let declaration = format!( |
| "declare const tools: {{ {} }};", |
| render_code_mode_tool_declaration(tool_name, input_name, &input_type, &output_type) |
| ); |
| format!("{description}\n\nexec tool declaration:\n```ts\n{declaration}\n```") |
| } |
|
|
| fn render_code_mode_sample_for_definition(definition: &ToolDefinition) -> String { |
| let input_name = match definition.kind { |
| CodeModeToolKind::Function => "args", |
| CodeModeToolKind::Freeform => "input", |
| }; |
| let input_type = match definition.kind { |
| CodeModeToolKind::Function => definition |
| .input_schema |
| .as_ref() |
| .map(render_json_schema_to_typescript) |
| .unwrap_or_else(|| "unknown".to_string()), |
| CodeModeToolKind::Freeform => "string".to_string(), |
| }; |
| let output_type = if let Some(structured_content_schema) = |
| mcp_structured_content_schema(definition.output_schema.as_ref()) |
| { |
| let structured_content_type = render_json_schema_to_typescript(structured_content_schema); |
| if structured_content_type == "unknown" { |
| "CallToolResult".to_string() |
| } else { |
| format!("CallToolResult<{structured_content_type}>") |
| } |
| } else { |
| definition |
| .output_schema |
| .as_ref() |
| .map(render_json_schema_to_typescript) |
| .unwrap_or_else(|| "unknown".to_string()) |
| }; |
| render_code_mode_sample( |
| &definition.description, |
| &definition.name, |
| input_name, |
| input_type, |
| output_type, |
| ) |
| } |
|
|
| fn render_code_mode_tool_declaration( |
| tool_name: &str, |
| input_name: &str, |
| input_type: &str, |
| output_type: &str, |
| ) -> String { |
| let tool_name = normalize_code_mode_identifier(tool_name); |
| format!("{tool_name}({input_name}: {input_type}): Promise<{output_type}>;") |
| } |
|
|
| fn render_tool_heading(global_name: &str, raw_name: &str) -> String { |
| if global_name == raw_name { |
| format!("### `{global_name}`") |
| } else { |
| format!("### `{global_name}` (`{raw_name}`)") |
| } |
| } |
|
|
| fn mcp_structured_content_schema(output_schema: Option<&JsonValue>) -> Option<&JsonValue> { |
| let output_schema = output_schema?; |
| let properties = output_schema |
| .get("properties") |
| .and_then(JsonValue::as_object)?; |
| let content_schema = properties.get("content").and_then(JsonValue::as_object)?; |
| if content_schema.get("type").and_then(JsonValue::as_str) != Some("array") { |
| return None; |
| } |
|
|
| if content_schema |
| .get("items") |
| .and_then(JsonValue::as_object) |
| .is_none_or(|items| items.get("type").and_then(JsonValue::as_str) != Some("object")) |
| { |
| return None; |
| } |
|
|
| if properties |
| .get("isError") |
| .and_then(JsonValue::as_object) |
| .is_none_or(|schema| schema.get("type").and_then(JsonValue::as_str) != Some("boolean")) |
| { |
| return None; |
| } |
|
|
| if properties |
| .get("_meta") |
| .and_then(JsonValue::as_object) |
| .is_none_or(|schema| schema.get("type").and_then(JsonValue::as_str) != Some("object")) |
| { |
| return None; |
| } |
|
|
| Some( |
| properties |
| .get("structuredContent") |
| .unwrap_or(&JsonValue::Bool(true)), |
| ) |
| } |
|
|
| #[cfg(test)] |
| mod tests { |
| use super::CodeModeToolKind; |
| use super::ImageDetailVisibility; |
| use super::ParsedExecSource; |
| use super::ToolDefinition; |
| use super::ToolNamespaceDescription; |
| use super::augment_tool_definition; |
| use super::build_exec_tool_description; |
| use super::normalize_code_mode_identifier; |
| use super::parse_exec_source; |
| use codex_protocol::ToolName; |
| use pretty_assertions::assert_eq; |
| use serde_json::Value as JsonValue; |
| use serde_json::json; |
| use std::collections::BTreeMap; |
|
|
| fn mcp_call_tool_result_schema(structured_content_schema: JsonValue) -> JsonValue { |
| json!({ |
| "type": "object", |
| "properties": { |
| "content": { |
| "type": "array", |
| "items": { |
| "type": "object" |
| } |
| }, |
| "structuredContent": structured_content_schema, |
| "isError": { "type": "boolean" }, |
| "_meta": { "type": "object" } |
| }, |
| "required": ["content"], |
| "additionalProperties": false |
| }) |
| } |
|
|
| #[test] |
| fn parse_exec_source_without_pragma() { |
| assert_eq!( |
| parse_exec_source("text('hi')").unwrap(), |
| ParsedExecSource { |
| code: "text('hi')".to_string(), |
| yield_time_ms: None, |
| max_output_tokens: None, |
| } |
| ); |
| } |
|
|
| #[test] |
| fn parse_exec_source_with_pragma() { |
| assert_eq!( |
| parse_exec_source("// @exec: {\"yield_time_ms\": 10}\ntext('hi')").unwrap(), |
| ParsedExecSource { |
| code: "text('hi')".to_string(), |
| yield_time_ms: Some(10), |
| max_output_tokens: None, |
| } |
| ); |
| } |
|
|
| #[test] |
| fn normalize_identifier_rewrites_invalid_characters() { |
| assert_eq!( |
| "mcp__ologs__get_profile", |
| normalize_code_mode_identifier("mcp__ologs__get_profile") |
| ); |
| assert_eq!( |
| "hidden_dynamic_tool", |
| normalize_code_mode_identifier("hidden-dynamic-tool") |
| ); |
| } |
|
|
| #[test] |
| fn augment_tool_definition_appends_typed_declaration() { |
| let definition = ToolDefinition { |
| name: "hidden_dynamic_tool".to_string(), |
| tool_name: ToolName::plain("hidden_dynamic_tool"), |
| description: "Test tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": { "city": { "type": "string" } }, |
| "required": ["city"], |
| "additionalProperties": false |
| })), |
| output_schema: Some(json!({ |
| "type": "object", |
| "properties": { "ok": { "type": "boolean" } }, |
| "required": ["ok"] |
| })), |
| }; |
|
|
| let description = augment_tool_definition(definition).description; |
| assert!(description.contains("declare const tools")); |
| assert!( |
| description.contains( |
| "hidden_dynamic_tool(args: { city: string; }): Promise<{ ok: boolean; }>;" |
| ) |
| ); |
| } |
|
|
| #[test] |
| fn augment_tool_definition_includes_property_descriptions_as_comments() { |
| let definition = ToolDefinition { |
| name: "weather_tool".to_string(), |
| tool_name: ToolName::plain("weather_tool"), |
| description: "Weather tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": { |
| "weather": { |
| "type": "array", |
| "description": "look up weather for a given list of locations", |
| "items": { |
| "type": "object", |
| "properties": { |
| "location": { "type": "string" } |
| }, |
| "required": ["location"] |
| } |
| } |
| }, |
| "required": ["weather"] |
| })), |
| output_schema: Some(json!({ |
| "type": "object", |
| "properties": { |
| "forecast": { |
| "type": "string", |
| "description": "human readable weather forecast" |
| } |
| }, |
| "required": ["forecast"] |
| })), |
| }; |
|
|
| let description = augment_tool_definition(definition).description; |
| assert!(description.contains( |
| r#"weather_tool(args: { |
| // look up weather for a given list of locations |
| weather: Array<{ location: string; }>; |
| }): Promise<{ |
| // human readable weather forecast |
| forecast: string; |
| }>;"# |
| )); |
| } |
|
|
| #[test] |
| fn code_mode_types_structured_content_result_refs() { |
| let definition = ToolDefinition { |
| name: "mcp__sample__search".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "search"), |
| description: "Search".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(mcp_call_tool_result_schema(json!({ |
| "type": "object", |
| "properties": { |
| "results": { |
| "type": "array", |
| "items": { "$ref": "#/definitions/Result~1item~0v1" } |
| } |
| }, |
| "required": ["results"], |
| "additionalProperties": false, |
| "definitions": { |
| "Result/item~v1": { |
| "type": "object", |
| "properties": { |
| "id": { "type": "string" }, |
| "score": { "type": "number" } |
| }, |
| "required": ["id", "score"], |
| "additionalProperties": false |
| } |
| } |
| }))), |
| }; |
|
|
| let description = augment_tool_definition(definition).description; |
| assert!(description.contains( |
| "mcp__sample__search(args: {}): Promise<CallToolResult<{ results: Array<{ id: string; score: number; }>; }>>;" |
| )); |
| } |
|
|
| #[test] |
| fn code_mode_only_description_includes_nested_tools() { |
| let description = build_exec_tool_description( |
| &[ToolDefinition { |
| name: "foo".to_string(), |
| tool_name: ToolName::plain("foo"), |
| description: "bar".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: None, |
| output_schema: None, |
| }], |
| &[], |
| &BTreeMap::new(), |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| true, |
| ImageDetailVisibility::Visible, |
| ); |
| assert!(description.contains( |
| "### `foo` |
| bar" |
| )); |
| assert!(!description.contains("do not attempt to use any other tools directly")); |
| } |
|
|
| #[test] |
| fn exec_description_mentions_timeout_helpers() { |
| let description = build_exec_tool_description( |
| &[], |
| &[], |
| &BTreeMap::new(), |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| false, |
| ImageDetailVisibility::Visible, |
| ); |
| assert!(description.contains("`audio(audioUrlOrItem:")); |
| assert!(description.contains("`setTimeout(callback: () => void, delayMs?: number)`")); |
| assert!(description.contains("`clearTimeout(timeoutId?: number)`")); |
| } |
|
|
| #[test] |
| fn code_mode_only_description_groups_namespace_instructions_once() { |
| let namespace_descriptions = BTreeMap::from([( |
| "mcp__sample__".to_string(), |
| ToolNamespaceDescription { |
| name: "mcp__sample".to_string(), |
| description: "Shared namespace guidance.".to_string(), |
| }, |
| )]); |
| let description = build_exec_tool_description( |
| &[ |
| ToolDefinition { |
| name: "mcp__sample__alpha".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "alpha"), |
| description: "First tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(mcp_call_tool_result_schema(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| }))), |
| }, |
| ToolDefinition { |
| name: "mcp__sample__beta".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "beta"), |
| description: "Second tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(mcp_call_tool_result_schema(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| }))), |
| }, |
| ], |
| &[], |
| &namespace_descriptions, |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| true, |
| ImageDetailVisibility::Visible, |
| ); |
| assert_eq!(description.matches("## mcp__sample").count(), 1); |
| assert!(description.contains("## mcp__sample\nShared namespace guidance.")); |
| assert!(description.contains( |
| "declare const tools: { mcp__sample__alpha(args: {}): Promise<CallToolResult<{}>>; };" |
| )); |
| assert!(description.contains( |
| "declare const tools: { mcp__sample__beta(args: {}): Promise<CallToolResult<{}>>; };" |
| )); |
| } |
|
|
| #[test] |
| fn code_mode_only_description_omits_empty_namespace_sections() { |
| let namespace_descriptions = BTreeMap::from([( |
| "mcp__sample__".to_string(), |
| ToolNamespaceDescription { |
| name: "mcp__sample".to_string(), |
| description: String::new(), |
| }, |
| )]); |
| let description = build_exec_tool_description( |
| &[ToolDefinition { |
| name: "mcp__sample__alpha".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "alpha"), |
| description: "First tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(mcp_call_tool_result_schema(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| }))), |
| }], |
| &[], |
| &namespace_descriptions, |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| true, |
| ImageDetailVisibility::Visible, |
| ); |
|
|
| assert!(!description.contains("## mcp__sample")); |
| assert!(description.contains("### `mcp__sample__alpha`")); |
| } |
|
|
| #[test] |
| fn code_mode_only_description_renders_shared_mcp_types_once() { |
| let first_tool = augment_tool_definition(ToolDefinition { |
| name: "mcp__sample__alpha".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "alpha"), |
| description: "First tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(json!({ |
| "type": "object", |
| "properties": { |
| "content": { |
| "type": "array", |
| "items": { |
| "type": "object" |
| } |
| }, |
| "structuredContent": { |
| "type": "object", |
| "properties": { |
| "echo": { "type": "string" } |
| }, |
| "required": ["echo"], |
| "additionalProperties": false |
| }, |
| "isError": { "type": "boolean" }, |
| "_meta": { "type": "object" } |
| }, |
| "required": ["content"], |
| "additionalProperties": false |
| })), |
| }); |
| let second_tool = augment_tool_definition(ToolDefinition { |
| name: "mcp__sample__beta".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "beta"), |
| description: "Second tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(json!({ |
| "type": "object", |
| "properties": { |
| "content": { |
| "type": "array", |
| "items": { |
| "type": "object" |
| } |
| }, |
| "structuredContent": { |
| "type": "object", |
| "properties": { |
| "count": { "type": "integer" } |
| }, |
| "required": ["count"], |
| "additionalProperties": false |
| }, |
| "isError": { "type": "boolean" }, |
| "_meta": { "type": "object" } |
| }, |
| "required": ["content"], |
| "additionalProperties": false |
| })), |
| }); |
|
|
| let description = build_exec_tool_description( |
| &[ |
| ToolDefinition { |
| name: first_tool.name, |
| tool_name: first_tool.tool_name, |
| description: "First tool".to_string(), |
| kind: first_tool.kind, |
| input_schema: first_tool.input_schema, |
| output_schema: first_tool.output_schema, |
| }, |
| ToolDefinition { |
| name: second_tool.name, |
| tool_name: second_tool.tool_name, |
| description: "Second tool".to_string(), |
| kind: second_tool.kind, |
| input_schema: second_tool.input_schema, |
| output_schema: second_tool.output_schema, |
| }, |
| ], |
| &[], |
| &BTreeMap::new(), |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| true, |
| ImageDetailVisibility::Visible, |
| ); |
|
|
| assert_eq!( |
| description |
| .matches("type CallToolResult<TStructured = { [key: string]: unknown }>") |
| .count(), |
| 1 |
| ); |
| assert_eq!(description.matches("Shared MCP Types:").count(), 1); |
| } |
|
|
| #[test] |
| fn code_mode_only_description_renders_shared_mcp_types_for_deferred_tools() { |
| let deferred_tool = ToolDefinition { |
| name: "mcp__sample__alpha".to_string(), |
| tool_name: ToolName::namespaced("mcp__sample__", "alpha"), |
| description: "Deferred tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: Some(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| })), |
| output_schema: Some(mcp_call_tool_result_schema(json!({ |
| "type": "object", |
| "properties": {}, |
| "additionalProperties": false |
| }))), |
| }; |
|
|
| let description = build_exec_tool_description( |
| &[], |
| &[deferred_tool], |
| &BTreeMap::new(), |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| true, |
| ImageDetailVisibility::Visible, |
| ); |
|
|
| assert!(description.contains("Some deferred nested tools may be omitted")); |
| assert!(description.contains("Shared MCP Types:")); |
| assert!(!description.contains("### `mcp__sample__alpha`")); |
| } |
|
|
| #[test] |
| fn exec_description_mentions_deferred_nested_tools_when_available() { |
| let description = build_exec_tool_description( |
| &[], |
| &[ToolDefinition { |
| name: "deferred_tool".to_string(), |
| tool_name: ToolName::plain("deferred_tool"), |
| description: "Deferred tool".to_string(), |
| kind: CodeModeToolKind::Function, |
| input_schema: None, |
| output_schema: None, |
| }], |
| &BTreeMap::new(), |
| crate::DEFAULT_EXEC_YIELD_TIME_MS, |
| false, |
| ImageDetailVisibility::Visible, |
| ); |
|
|
| assert!(description.contains("Some deferred nested tools may be omitted")); |
| assert!(description.contains("filter `ALL_TOOLS` by `name` and `description`")); |
| assert!(!description.contains("do not print the full `ALL_TOOLS` array")); |
| } |
| } |
|
|