mirror of
https://github.com/openai/codex.git
synced 2026-05-02 10:26:45 +00:00
965 lines
31 KiB
Rust
965 lines
31 KiB
Rust
use std::collections::HashMap;
|
|
use std::collections::VecDeque;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use crate::client_common::tools::ToolSpec;
|
|
use crate::codex::Session;
|
|
use crate::codex::TurnContext;
|
|
use crate::config::Config;
|
|
use crate::exec_env::create_env;
|
|
use crate::features::Feature;
|
|
use crate::function_tool::FunctionCallError;
|
|
use crate::tools::ToolRouter;
|
|
use crate::tools::code_mode_description::augment_tool_spec_for_code_mode;
|
|
use crate::tools::code_mode_description::code_mode_tool_reference;
|
|
use crate::tools::context::FunctionToolOutput;
|
|
use crate::tools::context::SharedTurnDiffTracker;
|
|
use crate::tools::context::ToolPayload;
|
|
use crate::tools::js_repl::resolve_compatible_node;
|
|
use crate::tools::router::ToolCall;
|
|
use crate::tools::router::ToolCallSource;
|
|
use crate::truncate::TruncationPolicy;
|
|
use crate::truncate::formatted_truncate_text_content_items_with_policy;
|
|
use crate::truncate::truncate_function_output_items_with_policy;
|
|
use crate::unified_exec::resolve_max_tokens;
|
|
use codex_protocol::models::FunctionCallOutputContentItem;
|
|
use serde::Deserialize;
|
|
use serde::Serialize;
|
|
use serde_json::Value as JsonValue;
|
|
use tokio::io::AsyncBufReadExt;
|
|
use tokio::io::AsyncReadExt;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio::io::BufReader;
|
|
use tokio::task::JoinHandle;
|
|
|
|
const CODE_MODE_RUNNER_SOURCE: &str = include_str!("code_mode_runner.cjs");
|
|
const CODE_MODE_BRIDGE_SOURCE: &str = include_str!("code_mode_bridge.js");
|
|
pub(crate) const PUBLIC_TOOL_NAME: &str = "exec";
|
|
pub(crate) const WAIT_TOOL_NAME: &str = "exec_wait";
|
|
pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000;
|
|
|
|
#[derive(Clone)]
|
|
struct ExecContext {
|
|
session: Arc<Session>,
|
|
turn: Arc<TurnContext>,
|
|
tracker: SharedTurnDiffTracker,
|
|
}
|
|
|
|
pub(crate) struct CodeModeProcess {
|
|
child: tokio::process::Child,
|
|
stdin: tokio::process::ChildStdin,
|
|
stdout_lines: tokio::io::Lines<BufReader<tokio::process::ChildStdout>>,
|
|
stderr_task: Option<JoinHandle<String>>,
|
|
pending_messages: HashMap<i32, VecDeque<NodeToHostMessage>>,
|
|
}
|
|
|
|
impl CodeModeProcess {
|
|
fn has_exited(&mut self) -> Result<bool, String> {
|
|
self.child
|
|
.try_wait()
|
|
.map(|status| status.is_some())
|
|
.map_err(|err| format!("failed to inspect {PUBLIC_TOOL_NAME} runner: {err}"))
|
|
}
|
|
|
|
async fn wait_for_exit(&mut self) -> Result<std::process::ExitStatus, String> {
|
|
self.child
|
|
.wait()
|
|
.await
|
|
.map_err(|err| format!("failed to wait for {PUBLIC_TOOL_NAME} runner: {err}"))
|
|
}
|
|
|
|
async fn stderr(&mut self) -> Result<String, String> {
|
|
self.stderr_task
|
|
.take()
|
|
.ok_or_else(|| format!("{PUBLIC_TOOL_NAME} stderr collector missing"))?
|
|
.await
|
|
.map_err(|err| format!("failed to collect {PUBLIC_TOOL_NAME} stderr: {err}"))
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(crate) struct CodeModeYieldedSession {
|
|
pub(crate) session_id: i32,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
|
#[serde(rename_all = "snake_case")]
|
|
enum CodeModeToolKind {
|
|
Function,
|
|
Freeform,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Serialize)]
|
|
struct EnabledTool {
|
|
tool_name: String,
|
|
#[serde(rename = "module")]
|
|
module_path: String,
|
|
namespace: Vec<String>,
|
|
name: String,
|
|
description: String,
|
|
kind: CodeModeToolKind,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
#[serde(tag = "type", rename_all = "snake_case")]
|
|
enum HostToNodeMessage {
|
|
Start {
|
|
session_id: i32,
|
|
enabled_tools: Vec<EnabledTool>,
|
|
stored_values: HashMap<String, JsonValue>,
|
|
source: String,
|
|
},
|
|
Poll {
|
|
session_id: i32,
|
|
yield_time_ms: u64,
|
|
},
|
|
Terminate {
|
|
session_id: i32,
|
|
},
|
|
Response {
|
|
session_id: i32,
|
|
id: String,
|
|
code_mode_result: JsonValue,
|
|
},
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(tag = "type", rename_all = "snake_case")]
|
|
enum NodeToHostMessage {
|
|
ToolCall {
|
|
session_id: i32,
|
|
id: String,
|
|
name: String,
|
|
#[serde(default)]
|
|
input: Option<JsonValue>,
|
|
},
|
|
Yielded {
|
|
session_id: i32,
|
|
content_items: Vec<JsonValue>,
|
|
#[serde(default)]
|
|
max_output_tokens_per_exec_call: Option<usize>,
|
|
},
|
|
Terminated {
|
|
session_id: i32,
|
|
content_items: Vec<JsonValue>,
|
|
#[serde(default)]
|
|
max_output_tokens_per_exec_call: Option<usize>,
|
|
},
|
|
Result {
|
|
session_id: i32,
|
|
content_items: Vec<JsonValue>,
|
|
stored_values: HashMap<String, JsonValue>,
|
|
#[serde(default)]
|
|
error_text: Option<String>,
|
|
#[serde(default)]
|
|
max_output_tokens_per_exec_call: Option<usize>,
|
|
},
|
|
}
|
|
|
|
enum CodeModeSessionAction {
|
|
Start {
|
|
enabled_tools: Vec<EnabledTool>,
|
|
stored_values: HashMap<String, JsonValue>,
|
|
source: String,
|
|
},
|
|
Poll {
|
|
yield_time_ms: u64,
|
|
max_output_tokens: Option<usize>,
|
|
},
|
|
Terminate {
|
|
max_output_tokens: Option<usize>,
|
|
},
|
|
}
|
|
|
|
enum CodeModeSessionProgress {
|
|
Finished(FunctionToolOutput),
|
|
Yielded {
|
|
output: FunctionToolOutput,
|
|
yielded_session: CodeModeYieldedSession,
|
|
},
|
|
}
|
|
|
|
enum CodeModeExecutionStatus {
|
|
Completed,
|
|
Failed,
|
|
Running(i32),
|
|
Terminated,
|
|
}
|
|
|
|
pub(crate) fn instructions(config: &Config) -> Option<String> {
|
|
if !config.features.enabled(Feature::CodeMode) {
|
|
return None;
|
|
}
|
|
|
|
let mut section = String::from("## Exec\n");
|
|
section.push_str(&format!(
|
|
"- Use `{PUBLIC_TOOL_NAME}` for JavaScript execution in a Node-backed `node:vm` context.\n",
|
|
));
|
|
section.push_str(&format!(
|
|
"- `{PUBLIC_TOOL_NAME}` is a freeform/custom tool. Direct `{PUBLIC_TOOL_NAME}` calls must send raw JavaScript tool input. Do not wrap code in JSON, quotes, or markdown code fences.\n",
|
|
));
|
|
section.push_str(&format!(
|
|
"- Direct tool calls remain available while `{PUBLIC_TOOL_NAME}` is enabled.\n",
|
|
));
|
|
section.push_str(&format!(
|
|
"- `{PUBLIC_TOOL_NAME}` uses the same Node runtime resolution as `js_repl`. If needed, point `js_repl_node_path` at the Node binary you want Codex to use.\n",
|
|
));
|
|
section.push_str("- Import nested tools from `tools.js`, for example `import { exec_command } from \"tools.js\"` or `import { ALL_TOOLS } from \"tools.js\"` to inspect the available `{ module, name, description }` entries. Namespaced tools are also available from `tools/<namespace...>.js`; MCP tools use `tools/mcp/<server>.js`, for example `import { append_notebook_logs_chart } from \"tools/mcp/ologs.js\"`. Nested tool calls resolve to their code-mode result values.\n");
|
|
section.push_str(&format!(
|
|
"- Import `{{ output_text, output_image, set_max_output_tokens_per_exec_call, set_yield_time, store, load }}` from `@openai/code_mode` (or `\"openai/code_mode\"`). `output_text(value)` surfaces text back to the model and stringifies non-string objects with `JSON.stringify(...)` when possible. `output_image(imageUrl)` appends an `input_image` content item for `http(s)` or `data:` URLs. `store(key, value)` persists JSON-serializable values across `{PUBLIC_TOOL_NAME}` calls in the current session, and `load(key)` returns a cloned stored value or `undefined`. `set_max_output_tokens_per_exec_call(value)` sets the token budget used to truncate direct `{PUBLIC_TOOL_NAME}` returns; `{WAIT_TOOL_NAME}` uses its own `max_tokens` argument instead and defaults to `10000`. `set_yield_time(value)` asks `{PUBLIC_TOOL_NAME}` to return early if the script is still running after that many milliseconds so `{WAIT_TOOL_NAME}` can resume it later. The returned content starts with a separate `Script completed`, `Script failed`, or `Script running with session ID …` text item that includes wall time. When truncation happens, the final text may include `Total output lines:` and the usual `…N tokens truncated…` marker.\n",
|
|
));
|
|
section.push_str(&format!(
|
|
"- If `{PUBLIC_TOOL_NAME}` returns `Script running with session ID …`, call `{WAIT_TOOL_NAME}` with that `session_id` to keep waiting for more output, completion, or termination.\n",
|
|
));
|
|
section.push_str(
|
|
"- Function tools require JSON object arguments. Freeform tools require raw strings.\n",
|
|
);
|
|
section.push_str("- `add_content(value)` remains available for compatibility. It is synchronous and accepts a content item, an array of content items, or a string. Structured nested-tool results should be converted to text first, for example with `JSON.stringify(...)`.\n");
|
|
section
|
|
.push_str("- Only content passed to `output_text(...)`, `output_image(...)`, or `add_content(value)` is surfaced back to the model.");
|
|
Some(section)
|
|
}
|
|
|
|
pub(crate) async fn execute(
|
|
session: Arc<Session>,
|
|
turn: Arc<TurnContext>,
|
|
tracker: SharedTurnDiffTracker,
|
|
code: String,
|
|
) -> Result<FunctionToolOutput, FunctionCallError> {
|
|
let exec = ExecContext {
|
|
session,
|
|
turn,
|
|
tracker,
|
|
};
|
|
let enabled_tools = build_enabled_tools(&exec).await;
|
|
let stored_values = exec.session.services.code_mode_store.stored_values().await;
|
|
let source = build_source(&code, &enabled_tools).map_err(FunctionCallError::RespondToModel)?;
|
|
let session_id = exec
|
|
.session
|
|
.services
|
|
.code_mode_store
|
|
.allocate_session_id()
|
|
.await;
|
|
let process = ensure_shared_code_mode_process(&exec)
|
|
.await
|
|
.map_err(FunctionCallError::RespondToModel)?;
|
|
let result = {
|
|
let mut process = process.lock().await;
|
|
drive_code_mode_session(
|
|
&exec,
|
|
&mut process,
|
|
session_id,
|
|
CodeModeSessionAction::Start {
|
|
enabled_tools,
|
|
stored_values,
|
|
source,
|
|
},
|
|
)
|
|
.await
|
|
};
|
|
if let Ok(CodeModeSessionProgress::Yielded {
|
|
yielded_session, ..
|
|
}) = &result
|
|
{
|
|
exec.session
|
|
.services
|
|
.code_mode_store
|
|
.store_yielded_session(yielded_session.clone())
|
|
.await;
|
|
}
|
|
match result {
|
|
Ok(CodeModeSessionProgress::Finished(output))
|
|
| Ok(CodeModeSessionProgress::Yielded { output, .. }) => Ok(output),
|
|
Err(error) => Err(FunctionCallError::RespondToModel(error)),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn wait(
|
|
session: Arc<Session>,
|
|
turn: Arc<TurnContext>,
|
|
tracker: SharedTurnDiffTracker,
|
|
session_id: i32,
|
|
yield_time_ms: u64,
|
|
max_output_tokens: Option<usize>,
|
|
terminate: bool,
|
|
) -> Result<FunctionToolOutput, FunctionCallError> {
|
|
let exec = ExecContext {
|
|
session,
|
|
turn,
|
|
tracker,
|
|
};
|
|
let yielded_session = exec
|
|
.session
|
|
.services
|
|
.code_mode_store
|
|
.take_yielded_session(session_id)
|
|
.await
|
|
.ok_or_else(|| {
|
|
FunctionCallError::RespondToModel(format!(
|
|
"{WAIT_TOOL_NAME} session_id {session_id} is not waiting on {WAIT_TOOL_NAME}"
|
|
))
|
|
})?;
|
|
|
|
let process = existing_shared_code_mode_process(&exec).await?;
|
|
let result = {
|
|
let mut process = process.lock().await;
|
|
drive_code_mode_session(
|
|
&exec,
|
|
&mut process,
|
|
yielded_session.session_id,
|
|
if terminate {
|
|
CodeModeSessionAction::Terminate { max_output_tokens }
|
|
} else {
|
|
CodeModeSessionAction::Poll {
|
|
yield_time_ms,
|
|
max_output_tokens,
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
};
|
|
if let Ok(CodeModeSessionProgress::Yielded {
|
|
yielded_session, ..
|
|
}) = &result
|
|
{
|
|
exec.session
|
|
.services
|
|
.code_mode_store
|
|
.store_yielded_session(yielded_session.clone())
|
|
.await;
|
|
}
|
|
|
|
match result {
|
|
Ok(CodeModeSessionProgress::Finished(output))
|
|
| Ok(CodeModeSessionProgress::Yielded { output, .. }) => Ok(output),
|
|
Err(error) => Err(FunctionCallError::RespondToModel(error)),
|
|
}
|
|
}
|
|
|
|
async fn spawn_code_mode_process(exec: &ExecContext) -> Result<CodeModeProcess, String> {
|
|
let node_path = resolve_compatible_node(exec.turn.config.js_repl_node_path.as_deref()).await?;
|
|
let env = create_env(&exec.turn.shell_environment_policy, None);
|
|
let mut cmd = tokio::process::Command::new(&node_path);
|
|
cmd.arg("--experimental-vm-modules");
|
|
cmd.arg("--eval");
|
|
cmd.arg(CODE_MODE_RUNNER_SOURCE);
|
|
cmd.current_dir(&exec.turn.cwd);
|
|
cmd.env_clear();
|
|
cmd.envs(env);
|
|
cmd.stdin(std::process::Stdio::piped())
|
|
.stdout(std::process::Stdio::piped())
|
|
.stderr(std::process::Stdio::piped())
|
|
.kill_on_drop(true);
|
|
|
|
let mut child = cmd
|
|
.spawn()
|
|
.map_err(|err| format!("failed to start {PUBLIC_TOOL_NAME} Node runtime: {err}"))?;
|
|
let stdout = child
|
|
.stdout
|
|
.take()
|
|
.ok_or_else(|| format!("{PUBLIC_TOOL_NAME} runner missing stdout"))?;
|
|
let stderr = child
|
|
.stderr
|
|
.take()
|
|
.ok_or_else(|| format!("{PUBLIC_TOOL_NAME} runner missing stderr"))?;
|
|
let stdin = child
|
|
.stdin
|
|
.take()
|
|
.ok_or_else(|| format!("{PUBLIC_TOOL_NAME} runner missing stdin"))?;
|
|
|
|
let stderr_task = tokio::spawn(async move {
|
|
let mut reader = BufReader::new(stderr);
|
|
let mut buf = Vec::new();
|
|
let _ = reader.read_to_end(&mut buf).await;
|
|
String::from_utf8_lossy(&buf).trim().to_string()
|
|
});
|
|
|
|
Ok(CodeModeProcess {
|
|
child,
|
|
stdin,
|
|
stdout_lines: BufReader::new(stdout).lines(),
|
|
stderr_task: Some(stderr_task),
|
|
pending_messages: HashMap::new(),
|
|
})
|
|
}
|
|
|
|
async fn ensure_shared_code_mode_process(
|
|
exec: &ExecContext,
|
|
) -> Result<Arc<tokio::sync::Mutex<CodeModeProcess>>, String> {
|
|
if let Some(process) = exec.session.services.code_mode_store.process().await {
|
|
let is_running = {
|
|
let mut process_guard = process.lock().await;
|
|
matches!(process_guard.has_exited(), Ok(false))
|
|
};
|
|
if is_running {
|
|
return Ok(process);
|
|
}
|
|
}
|
|
|
|
let process = Arc::new(tokio::sync::Mutex::new(
|
|
spawn_code_mode_process(exec).await?,
|
|
));
|
|
exec.session
|
|
.services
|
|
.code_mode_store
|
|
.store_process(process.clone())
|
|
.await;
|
|
Ok(process)
|
|
}
|
|
|
|
async fn existing_shared_code_mode_process(
|
|
exec: &ExecContext,
|
|
) -> Result<Arc<tokio::sync::Mutex<CodeModeProcess>>, FunctionCallError> {
|
|
let process = exec
|
|
.session
|
|
.services
|
|
.code_mode_store
|
|
.process()
|
|
.await
|
|
.ok_or_else(|| {
|
|
FunctionCallError::RespondToModel(format!(
|
|
"{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}"
|
|
))
|
|
})?;
|
|
let is_running = {
|
|
let mut process_guard = process.lock().await;
|
|
matches!(process_guard.has_exited(), Ok(false))
|
|
};
|
|
if is_running {
|
|
Ok(process)
|
|
} else {
|
|
Err(FunctionCallError::RespondToModel(format!(
|
|
"{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}"
|
|
)))
|
|
}
|
|
}
|
|
|
|
async fn drive_code_mode_session(
|
|
exec: &ExecContext,
|
|
process: &mut CodeModeProcess,
|
|
session_id: i32,
|
|
action: CodeModeSessionAction,
|
|
) -> Result<CodeModeSessionProgress, String> {
|
|
let started_at = std::time::Instant::now();
|
|
let is_terminate = matches!(action, CodeModeSessionAction::Terminate { .. });
|
|
let (message, poll_max_output_tokens) = match action {
|
|
CodeModeSessionAction::Start {
|
|
enabled_tools,
|
|
stored_values,
|
|
source,
|
|
} => (
|
|
HostToNodeMessage::Start {
|
|
session_id,
|
|
enabled_tools,
|
|
stored_values,
|
|
source,
|
|
},
|
|
None,
|
|
),
|
|
CodeModeSessionAction::Poll {
|
|
yield_time_ms,
|
|
max_output_tokens,
|
|
} => (
|
|
HostToNodeMessage::Poll {
|
|
session_id,
|
|
yield_time_ms,
|
|
},
|
|
Some(max_output_tokens),
|
|
),
|
|
CodeModeSessionAction::Terminate { max_output_tokens } => (
|
|
HostToNodeMessage::Terminate { session_id },
|
|
Some(max_output_tokens),
|
|
),
|
|
};
|
|
if let Some(progress) = process_pending_messages(
|
|
exec,
|
|
process,
|
|
session_id,
|
|
poll_max_output_tokens,
|
|
started_at,
|
|
is_terminate,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(progress);
|
|
}
|
|
write_message(&mut process.stdin, &message).await?;
|
|
|
|
if let Some(progress) = process_pending_messages(
|
|
exec,
|
|
process,
|
|
session_id,
|
|
poll_max_output_tokens,
|
|
started_at,
|
|
is_terminate,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(progress);
|
|
}
|
|
|
|
while let Some(line) = process
|
|
.stdout_lines
|
|
.next_line()
|
|
.await
|
|
.map_err(|err| format!("failed to read {PUBLIC_TOOL_NAME} runner stdout: {err}"))?
|
|
{
|
|
if line.trim().is_empty() {
|
|
continue;
|
|
}
|
|
let message: NodeToHostMessage = serde_json::from_str(&line).map_err(|err| {
|
|
format!("invalid {PUBLIC_TOOL_NAME} runner message: {err}; line={line}")
|
|
})?;
|
|
let message_session_id = message_session_id(&message);
|
|
if message_session_id != session_id {
|
|
if let NodeToHostMessage::ToolCall {
|
|
session_id: message_session_id,
|
|
id,
|
|
name,
|
|
input,
|
|
} = message
|
|
{
|
|
let response = HostToNodeMessage::Response {
|
|
session_id: message_session_id,
|
|
id,
|
|
code_mode_result: call_nested_tool(exec.clone(), name, input).await,
|
|
};
|
|
write_message(&mut process.stdin, &response).await?;
|
|
} else {
|
|
process
|
|
.pending_messages
|
|
.entry(message_session_id)
|
|
.or_default()
|
|
.push_back(message);
|
|
}
|
|
continue;
|
|
}
|
|
if let Some(progress) = handle_node_message(
|
|
exec,
|
|
process,
|
|
session_id,
|
|
message,
|
|
poll_max_output_tokens,
|
|
started_at,
|
|
is_terminate,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(progress);
|
|
}
|
|
}
|
|
|
|
let status = process.wait_for_exit().await?;
|
|
let stderr = process.stderr().await?;
|
|
let message = if stderr.is_empty() {
|
|
format!("{PUBLIC_TOOL_NAME} runner exited without returning a result (status {status})")
|
|
} else {
|
|
stderr
|
|
};
|
|
Err(message)
|
|
}
|
|
|
|
async fn process_pending_messages(
|
|
exec: &ExecContext,
|
|
process: &mut CodeModeProcess,
|
|
session_id: i32,
|
|
poll_max_output_tokens: Option<Option<usize>>,
|
|
started_at: std::time::Instant,
|
|
is_terminate: bool,
|
|
) -> Result<Option<CodeModeSessionProgress>, String> {
|
|
loop {
|
|
let Some(message) = process
|
|
.pending_messages
|
|
.get_mut(&session_id)
|
|
.and_then(VecDeque::pop_front)
|
|
else {
|
|
return Ok(None);
|
|
};
|
|
if let Some(progress) = handle_node_message(
|
|
exec,
|
|
process,
|
|
session_id,
|
|
message,
|
|
poll_max_output_tokens,
|
|
started_at,
|
|
is_terminate,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(progress));
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_node_message(
|
|
exec: &ExecContext,
|
|
process: &mut CodeModeProcess,
|
|
session_id: i32,
|
|
message: NodeToHostMessage,
|
|
poll_max_output_tokens: Option<Option<usize>>,
|
|
started_at: std::time::Instant,
|
|
is_terminate: bool,
|
|
) -> Result<Option<CodeModeSessionProgress>, String> {
|
|
match message {
|
|
NodeToHostMessage::ToolCall {
|
|
session_id: message_session_id,
|
|
id,
|
|
name,
|
|
input,
|
|
} => {
|
|
if is_terminate {
|
|
return Ok(None);
|
|
}
|
|
let response = HostToNodeMessage::Response {
|
|
session_id: message_session_id,
|
|
id,
|
|
code_mode_result: call_nested_tool(exec.clone(), name, input).await,
|
|
};
|
|
write_message(&mut process.stdin, &response).await?;
|
|
Ok(None)
|
|
}
|
|
NodeToHostMessage::Yielded {
|
|
content_items,
|
|
max_output_tokens_per_exec_call,
|
|
..
|
|
} => {
|
|
if is_terminate {
|
|
return Ok(None);
|
|
}
|
|
let mut delta_items = output_content_items_from_json_values(content_items)?;
|
|
delta_items = truncate_code_mode_result(
|
|
delta_items,
|
|
poll_max_output_tokens.unwrap_or(max_output_tokens_per_exec_call),
|
|
);
|
|
prepend_script_status(
|
|
&mut delta_items,
|
|
CodeModeExecutionStatus::Running(session_id),
|
|
started_at.elapsed(),
|
|
);
|
|
Ok(Some(CodeModeSessionProgress::Yielded {
|
|
output: FunctionToolOutput::from_content(delta_items, Some(true)),
|
|
yielded_session: CodeModeYieldedSession { session_id },
|
|
}))
|
|
}
|
|
NodeToHostMessage::Terminated {
|
|
content_items,
|
|
max_output_tokens_per_exec_call,
|
|
..
|
|
} => {
|
|
let mut delta_items = output_content_items_from_json_values(content_items)?;
|
|
delta_items = truncate_code_mode_result(
|
|
delta_items,
|
|
poll_max_output_tokens.unwrap_or(max_output_tokens_per_exec_call),
|
|
);
|
|
prepend_script_status(
|
|
&mut delta_items,
|
|
CodeModeExecutionStatus::Terminated,
|
|
started_at.elapsed(),
|
|
);
|
|
Ok(Some(CodeModeSessionProgress::Finished(
|
|
FunctionToolOutput::from_content(delta_items, Some(true)),
|
|
)))
|
|
}
|
|
NodeToHostMessage::Result {
|
|
content_items,
|
|
stored_values,
|
|
error_text,
|
|
max_output_tokens_per_exec_call,
|
|
..
|
|
} => {
|
|
exec.session
|
|
.services
|
|
.code_mode_store
|
|
.replace_stored_values(stored_values)
|
|
.await;
|
|
let mut delta_items = output_content_items_from_json_values(content_items)?;
|
|
let success = error_text.is_none();
|
|
if let Some(error_text) = error_text {
|
|
delta_items.push(FunctionCallOutputContentItem::InputText {
|
|
text: format!("Script error:\n{error_text}"),
|
|
});
|
|
}
|
|
|
|
let mut delta_items = truncate_code_mode_result(
|
|
delta_items,
|
|
poll_max_output_tokens.unwrap_or(max_output_tokens_per_exec_call),
|
|
);
|
|
prepend_script_status(
|
|
&mut delta_items,
|
|
if success {
|
|
CodeModeExecutionStatus::Completed
|
|
} else {
|
|
CodeModeExecutionStatus::Failed
|
|
},
|
|
started_at.elapsed(),
|
|
);
|
|
Ok(Some(CodeModeSessionProgress::Finished(
|
|
FunctionToolOutput::from_content(delta_items, Some(success)),
|
|
)))
|
|
}
|
|
}
|
|
}
|
|
|
|
fn message_session_id(message: &NodeToHostMessage) -> i32 {
|
|
match message {
|
|
NodeToHostMessage::ToolCall { session_id, .. }
|
|
| NodeToHostMessage::Yielded { session_id, .. }
|
|
| NodeToHostMessage::Terminated { session_id, .. }
|
|
| NodeToHostMessage::Result { session_id, .. } => *session_id,
|
|
}
|
|
}
|
|
|
|
async fn write_message(
|
|
stdin: &mut tokio::process::ChildStdin,
|
|
message: &HostToNodeMessage,
|
|
) -> Result<(), String> {
|
|
let line = serde_json::to_string(message)
|
|
.map_err(|err| format!("failed to serialize {PUBLIC_TOOL_NAME} message: {err}"))?;
|
|
stdin
|
|
.write_all(line.as_bytes())
|
|
.await
|
|
.map_err(|err| format!("failed to write {PUBLIC_TOOL_NAME} message: {err}"))?;
|
|
stdin
|
|
.write_all(b"\n")
|
|
.await
|
|
.map_err(|err| format!("failed to write {PUBLIC_TOOL_NAME} message newline: {err}"))?;
|
|
stdin
|
|
.flush()
|
|
.await
|
|
.map_err(|err| format!("failed to flush {PUBLIC_TOOL_NAME} message: {err}"))
|
|
}
|
|
|
|
fn prepend_script_status(
|
|
content_items: &mut Vec<FunctionCallOutputContentItem>,
|
|
status: CodeModeExecutionStatus,
|
|
wall_time: Duration,
|
|
) {
|
|
let wall_time_seconds = ((wall_time.as_secs_f32()) * 10.0).round() / 10.0;
|
|
let header = format!(
|
|
"{}\nWall time {wall_time_seconds:.1} seconds\nOutput:\n",
|
|
match status {
|
|
CodeModeExecutionStatus::Completed => "Script completed".to_string(),
|
|
CodeModeExecutionStatus::Failed => "Script failed".to_string(),
|
|
CodeModeExecutionStatus::Running(session_id) => {
|
|
format!("Script running with session ID {session_id}")
|
|
}
|
|
CodeModeExecutionStatus::Terminated => "Script terminated".to_string(),
|
|
}
|
|
);
|
|
content_items.insert(0, FunctionCallOutputContentItem::InputText { text: header });
|
|
}
|
|
|
|
fn build_source(user_code: &str, enabled_tools: &[EnabledTool]) -> Result<String, String> {
|
|
let enabled_tools_json = serde_json::to_string(enabled_tools)
|
|
.map_err(|err| format!("failed to serialize enabled tools: {err}"))?;
|
|
Ok(CODE_MODE_BRIDGE_SOURCE
|
|
.replace(
|
|
"__CODE_MODE_ENABLED_TOOLS_PLACEHOLDER__",
|
|
&enabled_tools_json,
|
|
)
|
|
.replace("__CODE_MODE_USER_CODE_PLACEHOLDER__", user_code))
|
|
}
|
|
|
|
fn truncate_code_mode_result(
|
|
items: Vec<FunctionCallOutputContentItem>,
|
|
max_output_tokens_per_exec_call: Option<usize>,
|
|
) -> Vec<FunctionCallOutputContentItem> {
|
|
let max_output_tokens = resolve_max_tokens(max_output_tokens_per_exec_call);
|
|
let policy = TruncationPolicy::Tokens(max_output_tokens);
|
|
if items
|
|
.iter()
|
|
.all(|item| matches!(item, FunctionCallOutputContentItem::InputText { .. }))
|
|
{
|
|
let (truncated_items, _) =
|
|
formatted_truncate_text_content_items_with_policy(&items, policy);
|
|
return truncated_items;
|
|
}
|
|
|
|
truncate_function_output_items_with_policy(&items, policy)
|
|
}
|
|
|
|
async fn build_enabled_tools(exec: &ExecContext) -> Vec<EnabledTool> {
|
|
let router = build_nested_router(exec).await;
|
|
let mut out = router
|
|
.specs()
|
|
.into_iter()
|
|
.map(|spec| augment_tool_spec_for_code_mode(spec, true))
|
|
.filter_map(enabled_tool_from_spec)
|
|
.collect::<Vec<_>>();
|
|
out.sort_by(|left, right| left.tool_name.cmp(&right.tool_name));
|
|
out.dedup_by(|left, right| left.tool_name == right.tool_name);
|
|
out
|
|
}
|
|
|
|
fn enabled_tool_from_spec(spec: ToolSpec) -> Option<EnabledTool> {
|
|
let tool_name = spec.name().to_string();
|
|
if tool_name == PUBLIC_TOOL_NAME || tool_name == WAIT_TOOL_NAME {
|
|
return None;
|
|
}
|
|
|
|
let reference = code_mode_tool_reference(&tool_name);
|
|
|
|
let (description, kind) = match spec {
|
|
ToolSpec::Function(tool) => (tool.description, CodeModeToolKind::Function),
|
|
ToolSpec::Freeform(tool) => (tool.description, CodeModeToolKind::Freeform),
|
|
ToolSpec::LocalShell {} | ToolSpec::ImageGeneration { .. } | ToolSpec::WebSearch { .. } => {
|
|
return None;
|
|
}
|
|
};
|
|
|
|
Some(EnabledTool {
|
|
tool_name,
|
|
module_path: reference.module_path,
|
|
namespace: reference.namespace,
|
|
name: reference.tool_key,
|
|
description,
|
|
kind,
|
|
})
|
|
}
|
|
|
|
async fn build_nested_router(exec: &ExecContext) -> ToolRouter {
|
|
let nested_tools_config = exec.turn.tools_config.for_code_mode_nested_tools();
|
|
let mcp_tools = exec
|
|
.session
|
|
.services
|
|
.mcp_connection_manager
|
|
.read()
|
|
.await
|
|
.list_all_tools()
|
|
.await
|
|
.into_iter()
|
|
.map(|(name, tool_info)| (name, tool_info.tool))
|
|
.collect();
|
|
|
|
ToolRouter::from_config(
|
|
&nested_tools_config,
|
|
Some(mcp_tools),
|
|
None,
|
|
exec.turn.dynamic_tools.as_slice(),
|
|
)
|
|
}
|
|
|
|
async fn call_nested_tool(
|
|
exec: ExecContext,
|
|
tool_name: String,
|
|
input: Option<JsonValue>,
|
|
) -> JsonValue {
|
|
if tool_name == PUBLIC_TOOL_NAME {
|
|
return JsonValue::String(format!("{PUBLIC_TOOL_NAME} cannot invoke itself"));
|
|
}
|
|
|
|
let router = build_nested_router(&exec).await;
|
|
|
|
let specs = router.specs();
|
|
let payload = if let Some((server, tool)) = exec.session.parse_mcp_tool_name(&tool_name).await {
|
|
match serialize_function_tool_arguments(&tool_name, input) {
|
|
Ok(raw_arguments) => ToolPayload::Mcp {
|
|
server,
|
|
tool,
|
|
raw_arguments,
|
|
},
|
|
Err(error) => return JsonValue::String(error),
|
|
}
|
|
} else {
|
|
match build_nested_tool_payload(&specs, &tool_name, input) {
|
|
Ok(payload) => payload,
|
|
Err(error) => return JsonValue::String(error),
|
|
}
|
|
};
|
|
|
|
let call = ToolCall {
|
|
tool_name: tool_name.clone(),
|
|
call_id: format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()),
|
|
payload,
|
|
};
|
|
let result = router
|
|
.dispatch_tool_call_with_code_mode_result(
|
|
Arc::clone(&exec.session),
|
|
Arc::clone(&exec.turn),
|
|
Arc::clone(&exec.tracker),
|
|
call,
|
|
ToolCallSource::CodeMode,
|
|
)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(result) => result.code_mode_result(),
|
|
Err(error) => JsonValue::String(error.to_string()),
|
|
}
|
|
}
|
|
|
|
fn tool_kind_for_spec(spec: &ToolSpec) -> CodeModeToolKind {
|
|
if matches!(spec, ToolSpec::Freeform(_)) {
|
|
CodeModeToolKind::Freeform
|
|
} else {
|
|
CodeModeToolKind::Function
|
|
}
|
|
}
|
|
|
|
fn tool_kind_for_name(specs: &[ToolSpec], tool_name: &str) -> Result<CodeModeToolKind, String> {
|
|
specs
|
|
.iter()
|
|
.find(|spec| spec.name() == tool_name)
|
|
.map(tool_kind_for_spec)
|
|
.ok_or_else(|| format!("tool `{tool_name}` is not enabled in {PUBLIC_TOOL_NAME}"))
|
|
}
|
|
|
|
fn build_nested_tool_payload(
|
|
specs: &[ToolSpec],
|
|
tool_name: &str,
|
|
input: Option<JsonValue>,
|
|
) -> Result<ToolPayload, String> {
|
|
let actual_kind = tool_kind_for_name(specs, tool_name)?;
|
|
match actual_kind {
|
|
CodeModeToolKind::Function => build_function_tool_payload(tool_name, input),
|
|
CodeModeToolKind::Freeform => build_freeform_tool_payload(tool_name, input),
|
|
}
|
|
}
|
|
|
|
fn build_function_tool_payload(
|
|
tool_name: &str,
|
|
input: Option<JsonValue>,
|
|
) -> Result<ToolPayload, String> {
|
|
let arguments = serialize_function_tool_arguments(tool_name, input)?;
|
|
Ok(ToolPayload::Function { arguments })
|
|
}
|
|
|
|
fn serialize_function_tool_arguments(
|
|
tool_name: &str,
|
|
input: Option<JsonValue>,
|
|
) -> Result<String, String> {
|
|
match input {
|
|
None => Ok("{}".to_string()),
|
|
Some(JsonValue::Object(map)) => serde_json::to_string(&JsonValue::Object(map))
|
|
.map_err(|err| format!("failed to serialize tool `{tool_name}` arguments: {err}")),
|
|
Some(_) => Err(format!(
|
|
"tool `{tool_name}` expects a JSON object for arguments"
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn build_freeform_tool_payload(
|
|
tool_name: &str,
|
|
input: Option<JsonValue>,
|
|
) -> Result<ToolPayload, String> {
|
|
match input {
|
|
Some(JsonValue::String(input)) => Ok(ToolPayload::Custom { input }),
|
|
_ => Err(format!("tool `{tool_name}` expects a string input")),
|
|
}
|
|
}
|
|
|
|
fn output_content_items_from_json_values(
|
|
content_items: Vec<JsonValue>,
|
|
) -> Result<Vec<FunctionCallOutputContentItem>, String> {
|
|
content_items
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(index, item)| {
|
|
serde_json::from_value(item).map_err(|err| {
|
|
format!("invalid {PUBLIC_TOOL_NAME} content item at index {index}: {err}")
|
|
})
|
|
})
|
|
.collect()
|
|
}
|