mod config;
mod declarations;
mod description;
mod sandbox;
mod source;
mod store;
use super::{
EXTENSION_FACTORIES, ExtensionFactory, PRIORITY_CODEMODE, command_arg,
util::{notify_text, value_text},
};
use crate::{
core::{
auth,
extensions::{
Extension, ExtensionCommand, ExtensionMode, ExtensionSetting, ExtensionTool,
GrammarSampling, SubcommandDef, ToolExecCtx, ToolExposure, ToolLoadout,
ToolLoadoutChanges, UiNotifyLevel,
},
model_resolver,
model_resolver::ModelEntry,
provider::{
self, AgentMessage, ClassifierContext, ImageContent, ModelConfig, ModelType, Usage,
usage::combine_usage,
},
tools::{ToolError, ToolResult, ToolResultAttachment},
},
extensions::tool_search,
modes::interactive::{app::App, handlers::register_slash_command},
utils::{output_files, truncate},
};
use base64::Engine as _;
use description::CHARS_PER_TOKEN;
use futures_util::future::BoxFuture;
use sandbox::{SandboxRequest, ScriptError, ScriptGlobal, ScriptHost, ScriptOutput, ScriptTool};
use serde_json::{Value, json};
use std::{
collections::{HashMap, HashSet},
sync::{Arc, Mutex as StdMutex},
time::Instant,
};
use tokio::sync::Semaphore;
pub const TOOL_NAME: &str = "codemode";
const EXT: &str = "codemode";
const CMD: &str = "codemode";
const MAX_CONCURRENT_MODEL_CALLS: usize = 4;
const MODEL_CALL_ERROR_PREVIEW_CHARS: usize = 500;
const MODEL_CALL_ARG_PREVIEW_CHARS: usize = 200;
const MODELS_CLASSIFY: &str = "models.classify";
const MODELS_GENERATE_IMAGES: &str = "models.generateImages";
const DEFAULT_MAX_OUTPUT_TOKENS: u64 = 10_000;
const IMAGE_FILE_EXTENSIONS: &[(&str, &str)] = &[
("image/png", ".png"),
("image/jpeg", ".jpg"),
("image/gif", ".gif"),
("image/webp", ".webp"),
];
const SUBCOMMANDS: &[SubcommandDef] = &[
SubcommandDef {
name: "settings",
description: "Open the codemode settings panel (same as `/codemode` with no arguments)",
},
SubcommandDef {
name: "show",
description: "Print the current mode and inlineBudget (config file path included)",
},
];
#[linkme::distributed_slice(EXTENSION_FACTORIES)]
static CODEMODE_FACTORY: ExtensionFactory = ExtensionFactory {
priority: PRIORITY_CODEMODE,
make: || -> Arc<dyn Extension> { Arc::new(Codemode) },
};
pub struct Codemode;
impl Extension for Codemode {
fn name(&self) -> &str {
EXT
}
fn description(&self) -> &str {
"Run JavaScript that orchestrates the other tools in a QuickJS sandbox"
}
fn default_enabled(&self) -> bool {
false
}
fn modes(&self) -> Vec<ExtensionMode> {
vec![ExtensionMode::Dev, ExtensionMode::Creator]
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![tool_definition()]
}
fn hosts_scripts(&self) -> bool {
true
}
fn prepare_loadout(&self, loadout: &ToolLoadout) -> ToolLoadoutChanges {
let cfg = config::load_config();
let deferred: HashSet<String> = loadout
.callable
.iter()
.filter(|t| t.exposure == ToolExposure::Deferred)
.map(|t| t.name.clone())
.collect();
let is_direct = |name: &str| {
loadout
.declared
.iter()
.find(|t| t.name == name)
.is_some_and(|t| t.exposure == ToolExposure::Direct)
};
let changes = description::build_loadout(
TOOL_NAME,
&loadout.declared,
&loadout.callable,
&is_direct,
&deferred,
cfg.mode(),
Some(cfg.inline_budget()),
);
ToolLoadoutChanges {
descriptions: changes.descriptions.into_iter().collect(),
hidden_declarations: changes.hidden_declarations,
}
}
fn execute_tool_async(
&self,
name: String,
args: Value,
ctx: ToolExecCtx,
) -> BoxFuture<'static, std::result::Result<ToolResult, ToolError>> {
Box::pin(async move {
if name != TOOL_NAME {
return Err(ToolError(format!("codemode: unknown tool: {name}")));
}
let Some(code) = args.get("code").and_then(|v| v.as_str()) else {
return Err(ToolError(
"codemode: missing string argument `code`".to_string(),
));
};
run_script(code, ctx).await
})
}
fn on_session_switched(&self, _session_path: Option<&str>, _messages: &[AgentMessage]) {
store::reset();
}
fn settings(&self) -> Vec<ExtensionSetting> {
config::panel_settings()
}
fn apply_setting(&self, key: &str, value: &str) -> std::result::Result<(), String> {
config::apply_panel_choice(key, value).map(|_| ())
}
fn commands(&self) -> Vec<ExtensionCommand> {
vec![ExtensionCommand {
name: CMD.to_string(),
description:
"Open the codemode settings panel; /codemode show prints the current config"
.to_string(),
busy_safe: true, subcommands: SUBCOMMANDS.to_vec(),
}]
}
fn on_registered(&self) {
register_slash_command(EXT, CMD, command_codemode);
}
}
fn command_codemode(app: &mut App, raw: &str) -> bool {
match command_arg(raw).trim() {
"" | "settings" => {
if app.ext_settings.is_open_for(EXT) {
app.ext_settings.close();
} else {
app.ext_settings.open(EXT);
}
app.dirty = true;
}
"show" => {
let cfg = config::load_config();
let lines = [
"codemode config".to_string(),
format!(" mode: {}", cfg.mode_label()),
format!(" inlineBudget: {}", cfg.inline_budget()),
format!(" file: {}", config::config_path().display()),
];
notify_text(&lines.join("\n"), UiNotifyLevel::Info);
}
other => notify_text(
&format!("codemode: unknown subcommand: {other}"),
UiNotifyLevel::Error,
),
}
false
}
fn tool_definition() -> ExtensionTool {
let mut tool = ExtensionTool::simple(
TOOL_NAME,
&description::create_description(&[], &HashSet::new(), None),
json!({
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Raw JavaScript source. Top-level await and return work. May start with a `// @options: {\"max_output_tokens\": 1000}` line."
}
},
"required": ["code"]
}),
"Run JavaScript that calls other tools (chains, loops, Promise.all, filtering large results)",
);
tool.exposure = ToolExposure::ModelOnly; tool.grammar_sampling = Some(GrammarSampling {
openai_lark: Some(source::CODEMODE_SOURCE_GRAMMAR.to_string()),
openai_regex: None,
});
tool.prompt_guidelines = vec![
"Use codemode to batch or chain several tool calls, or to filter large tool output down to what you need, instead of issuing many individual tool calls. Batch independent calls in one codemode call using await Promise.allSettled([...])."
.to_string(),
];
tool
}
#[derive(Default)]
struct ModelGlobalsState {
calls: Vec<Value>,
classify_count: usize,
image_count: usize,
generated_images: usize,
usage: Option<Usage>,
}
#[derive(Clone)]
struct Host {
ctx: ToolExecCtx,
tools: Arc<HashMap<String, ExtensionTool>>,
samples: Arc<HashMap<String, String>>,
model_globals: Arc<StdMutex<ModelGlobalsState>>,
model_call_gate: Arc<Semaphore>,
}
impl Host {
fn script_value(result: &ToolResult, tool: Option<&ExtensionTool>) -> Result<Value, String> {
if let Some(tool) = tool
&& tool.output_schema.is_some()
&& let Some(structured) = &result.structured_content
{
return Ok(structured.clone());
}
if result.is_error {
let err = if result.text.is_empty() {
match tool {
Some(tool) => format!("Tool \"{}\" failed", tool.name),
None => "Tool failed".to_string(),
}
} else {
result.text.clone()
};
return Err(err);
}
Ok(Value::String(result.text.clone()))
}
fn search_tools(&self, args: &[Value]) -> Result<Value, String> {
let Some(query) = args.first().and_then(|v| v.as_str()) else {
return Err("searchTools() expects a query string".to_string());
};
let opts = args.get(1).and_then(|v| v.as_object());
let limit = match opts.and_then(|o| o.get("limit")) {
Some(v) => match v.as_u64() {
Some(n) if n > 0 => n as usize,
_ => return Err("searchTools() limit must be a positive integer".to_string()),
},
None => tool_search::DEFAULT_LIMIT,
};
let namespace = match opts.and_then(|o| o.get("namespace")) {
Some(Value::Null) | None => None,
Some(Value::String(s)) => Some(namespace_key(s)),
Some(_) => return Err("searchTools() namespace must be a string".to_string()),
};
let documents: Vec<_> = self
.tools
.values()
.filter(|t| {
namespace
.as_ref()
.map(|ns| {
t.namespace
.as_ref()
.is_some_and(|own| namespace_key(own) == *ns)
})
.unwrap_or(true)
})
.map(tool_search::tool_document)
.collect();
let matches = tool_search::rank(query, &documents, limit);
Ok(Value::Array(
matches
.into_iter()
.map(|(name, _)| {
json!({
"name": declarations::to_codemode_identifier(&name),
"description": self.samples.get(&name).cloned().unwrap_or_default(),
})
})
.collect(),
))
}
fn describe_tool(&self, args: &[Value]) -> Result<Value, String> {
let Some(name) = args.first().and_then(|v| v.as_str()) else {
return Err("describeTool() expects a tool name".to_string());
};
let found = self
.tools
.values()
.find(|t| t.name == name || declarations::to_codemode_identifier(&t.name) == name);
Ok(match found.and_then(|t| self.samples.get(&t.name)) {
Some(sample) => Value::String(sample.clone()),
None => Value::Null,
})
}
fn describe_namespace(&self, args: &[Value]) -> Result<Value, String> {
let Some(name) = args.first().and_then(|v| v.as_str()) else {
return Err("describe_namespace() expects a namespace name".to_string());
};
Ok(match namespace_tools(&self.tools, &self.samples, name) {
Some((namespace, tools)) => json!({
"namespace": namespace,
"tools": tools
.into_iter()
.map(|(name, description)| json!({ "name": name, "description": description }))
.collect::<Vec<_>>(),
}),
None => Value::Null,
})
}
fn codemode_call_id(&self) -> String {
self.ctx
.parent_tool_call_id
.clone()
.unwrap_or_else(|| TOOL_NAME.to_string())
}
fn models_of_type(&self, args: &[Value], available_only: bool) -> Result<Value, String> {
let kind = model_type_arg(args.first())?;
let provider = provider_arg(args.get(1))?;
Ok(Value::Array(catalog_models(kind, provider, available_only)))
}
fn model_of_type(&self, args: &[Value]) -> Result<Value, String> {
let kind = model_type_arg(args.first())?;
let expected =
|| "models.getModelOfType() expects a type, a provider, and an id".to_string();
let provider = args.get(1).and_then(|v| v.as_str()).ok_or_else(expected)?;
let id = args.get(2).and_then(|v| v.as_str()).ok_or_else(expected)?;
Ok(catalog_model(kind, provider, id))
}
async fn classify(&self, args: &[Value]) -> Result<Value, String> {
let model = resolve_model_call(MODELS_CLASSIFY, ModelType::Classifier, args.first())?;
let context = classifier_context_arg(args.get(1))?;
let index = self.begin_model_call(MODELS_CLASSIFY, &model_ref(&model));
let _permit = self.model_call_gate.acquire().await.ok();
let started = Instant::now();
let result = provider::classify(&model, &context).await;
self.finish_model_call(
index,
started.elapsed(),
&result.stop_reason,
result.error_message.as_deref(),
result.usage.as_ref(),
);
serde_json::to_value(&result).map_err(|e| e.to_string())
}
async fn generate_images(&self, args: &[Value]) -> Result<Value, String> {
let model = resolve_model_call(MODELS_GENERATE_IMAGES, ModelType::Image, args.first())?;
let input = images_context_arg(args.get(1))?;
let index = self.begin_model_call(MODELS_GENERATE_IMAGES, &model_ref(&model));
let _permit = self.model_call_gate.acquire().await.ok();
let started = Instant::now();
let result = provider::generate_images(&model, &input).await;
let images = result
.output
.iter()
.filter(|block| matches!(block, ImageContent::Image { .. }))
.count();
if images > 0 {
self.model_globals.lock().unwrap().generated_images += images;
}
self.finish_model_call(
index,
started.elapsed(),
&result.stop_reason,
result.error_message.as_deref(),
result.usage.as_ref(),
);
serde_json::to_value(&result).map_err(|e| e.to_string())
}
fn begin_model_call(&self, global: &str, ref_text: &str) -> usize {
let mut state = self.model_globals.lock().unwrap();
let count = if global == MODELS_GENERATE_IMAGES {
state.image_count += 1;
state.image_count
} else {
state.classify_count += 1;
state.classify_count
};
state.calls.push(json!({
"id": format!("{}/{global}/{count}", self.codemode_call_id()),
"name": global,
"args": ref_text,
"status": "running",
"durationMs": 0,
}));
state.calls.len() - 1
}
fn finish_model_call(
&self,
index: usize,
elapsed: std::time::Duration,
stop_reason: &str,
error_message: Option<&str>,
usage: Option<&Usage>,
) {
let mut state = self.model_globals.lock().unwrap();
if let Some(record) = state.calls.get_mut(index).and_then(|v| v.as_object_mut()) {
record.insert("durationMs".to_string(), json!(elapsed.as_millis() as u64));
record.insert(
"status".to_string(),
json!(if stop_reason == "stop" { "ok" } else { "error" }),
);
if let Some(message) = error_message {
record.insert(
"error".to_string(),
json!(truncate::truncate_chars_within(
message,
MODEL_CALL_ERROR_PREVIEW_CHARS,
"..."
)),
);
}
if let Some(usage) = usage {
record.insert("cost".to_string(), json!(usage.cost.total));
}
}
if let Some(usage) = usage {
state.usage = Some(match state.usage.take() {
Some(prev) => combine_usage(&prev, usage),
None => usage.clone(),
});
}
}
}
fn model_type_arg(value: Option<&Value>) -> Result<ModelType, String> {
match value.and_then(|v| v.as_str()) {
Some("chat") => Ok(ModelType::Chat),
Some("image") => Ok(ModelType::Image),
Some("classifier") => Ok(ModelType::Classifier),
other => Err(format!(
"Unknown model type {}. Use \"chat\", \"image\", or \"classifier\".",
match other {
Some(value) => format!("\"{value}\""),
None => "undefined".to_string(),
}
)),
}
}
fn provider_arg(value: Option<&Value>) -> Result<Option<String>, String> {
match value {
None | Some(Value::Null) => Ok(None),
Some(Value::String(provider)) => Ok(Some(provider.clone())),
Some(_) => Err("provider must be a string".to_string()),
}
}
fn model_ref(model: &ModelConfig) -> String {
format!("{}/{}", model.provider, model.model_id)
}
fn describe_script_value(value: Option<&Value>) -> String {
match value {
None => "undefined".to_string(),
Some(value) => {
truncate::truncate_chars_within(&value.to_string(), MODEL_CALL_ARG_PREVIEW_CHARS, "...")
}
}
}
fn resolve_model_call(
global: &str,
kind: ModelType,
value: Option<&Value>,
) -> Result<ModelConfig, String> {
let type_name = kind.as_str();
let list_hint = format!(
"List the {type_name} models you can use with models.getAvailableOfType(\"{type_name}\")."
);
let fields = value.and_then(|v| v.as_object());
let provider = fields
.and_then(|f| f.get("provider"))
.and_then(|v| v.as_str());
let id = fields.and_then(|f| f.get("id")).and_then(|v| v.as_str());
let (Some(provider), Some(id)) = (provider, id) else {
let undefined_hint = match value {
None | Some(Value::Null) => {
" models.getModelOfType() returns undefined for an unknown provider or id."
}
_ => "",
};
return Err(format!(
"{global}() expects a model of type \"{type_name}\" as its first argument, got {}.{undefined_hint} {list_hint}",
describe_script_value(value),
));
};
if let Ok(model) = model_resolver::configured_model_of_type(provider, id, kind, None, None) {
return Ok(model);
}
let actual = [ModelType::Chat, ModelType::Image, ModelType::Classifier]
.into_iter()
.find(|other| {
*other != kind && model_resolver::find_model_of_type(provider, id, *other).is_ok()
});
Err(match actual {
Some(other) => format!(
"\"{provider}/{id}\" is a {} model, not a {type_name} model. {list_hint}",
other.as_str()
),
None => format!("Unknown {type_name} model \"{provider}/{id}\". {list_hint}"),
})
}
fn classifier_context_arg(value: Option<&Value>) -> Result<ClassifierContext, String> {
const EXPECTED: &str = r#"{ state: { ... }, images?: [{ type: "image", data: <base64>, mimeType }], questions: { <id>: { type: "choice" | "score" | "bool", instructions, criteria } } }"#;
let fail =
|problem: String| format!("{MODELS_CLASSIFY}() {problem}. Expected context: {EXPECTED}.");
let Some(context) = value.and_then(|v| v.as_object()) else {
return Err(fail(format!(
"expects a context object as its second argument, got {}",
describe_script_value(value)
)));
};
let state = context.get("state");
if !state.is_some_and(Value::is_object) {
return Err(fail(format!(
"context.state must be an object, got {}",
describe_script_value(state)
)));
}
if let Some(images) = context.get("images") {
let Some(blocks) = images.as_array() else {
return Err(fail(format!(
"context.images must be an array, got {}",
describe_script_value(Some(images))
)));
};
for (index, block) in blocks.iter().enumerate() {
let is_image_block = block.get("type").and_then(Value::as_str) == Some("image")
&& block.get("data").and_then(Value::as_str).is_some()
&& block.get("mimeType").and_then(Value::as_str).is_some();
if !is_image_block {
return Err(fail(format!(
"context.images[{index}] must be an image block, got {}",
describe_script_value(Some(block))
)));
}
}
}
match context.get("questions").and_then(Value::as_object) {
Some(questions) if !questions.is_empty() => {}
_ => {
return Err(fail(format!(
"context.questions must map question IDs to questions, got {}",
describe_script_value(context.get("questions"))
)));
}
}
serde_json::from_value(Value::Object(context.clone())).map_err(|err| {
format!("{MODELS_CLASSIFY}() context is malformed: {err}. Expected context: {EXPECTED}.")
})
}
fn images_context_arg(value: Option<&Value>) -> Result<Vec<ImageContent>, String> {
const EXPECTED: &str = r#"{ input: [{ type: "text", text: <prompt> }, ...optional { type: "image", data: <base64>, mimeType } references] }"#;
let fail = |problem: String| {
format!("{MODELS_GENERATE_IMAGES}() {problem}. Expected context: {EXPECTED}.")
};
let Some(context) = value.and_then(|v| v.as_object()) else {
return Err(fail(format!(
"expects a context object as its second argument, got {}",
describe_script_value(value)
)));
};
let input = context.get("input").and_then(|v| v.as_array());
let Some(input) = input.filter(|blocks| !blocks.is_empty()) else {
return Err(fail(format!(
"context.input must be a non-empty array of blocks, got {}",
describe_script_value(context.get("input"))
)));
};
input
.iter()
.enumerate()
.map(|(index, block)| {
serde_json::from_value::<ImageContent>(block.clone()).map_err(|_| {
fail(format!(
"context.input[{index}] must be a text or image block, got {}",
describe_script_value(Some(block))
))
})
})
.collect()
}
fn model_info(entry: &ModelEntry) -> Value {
json!({
"type": entry.model_type.as_str(),
"provider": &entry.provider,
"id": &entry.id,
"name": &entry.name,
"api": &entry.api,
"input": &entry.input,
"contextWindow": entry.context_window,
"maxTokens": entry.max_tokens,
"reasoning": entry.reasoning,
})
}
fn catalog_models(kind: ModelType, provider: Option<String>, available_only: bool) -> Vec<Value> {
let configured: Option<HashSet<String>> =
available_only.then(|| auth::list_configured_providers().into_iter().collect());
let providers: Vec<String> = match provider {
Some(provider) => vec![provider],
None => model_resolver::SUPPORTED_PROVIDERS
.iter()
.map(|p| (*p).to_string())
.collect(),
};
let mut out: Vec<Value> = Vec::new();
for provider in providers {
if let Some(configured) = &configured
&& !configured.contains(&provider)
{
continue;
}
for (id, _) in model_resolver::list_models_of_type(&provider, kind) {
if let Ok(entry) = model_resolver::find_model_of_type(&provider, &id, kind) {
out.push(model_info(&entry));
}
}
}
out
}
fn namespace_key(name: &str) -> String {
let name = name
.strip_prefix("mcp__")
.or_else(|| name.strip_prefix("mcp_"))
.unwrap_or(name);
name.replace('-', "_")
}
fn namespace_tools(
tools: &HashMap<String, ExtensionTool>,
samples: &HashMap<String, String>,
name: &str,
) -> Option<(String, Vec<(String, String)>)> {
let key = namespace_key(name);
let mut namespace: Option<String> = None;
let mut found: Vec<(String, String)> = Vec::new();
for tool in tools.values() {
let Some(own) = tool.namespace.as_deref() else {
continue;
};
if namespace_key(own) != key {
continue;
}
namespace = Some(own.to_string());
found.push((
declarations::to_codemode_identifier(&tool.name),
samples.get(&tool.name).cloned().unwrap_or_default(),
));
}
found.sort();
namespace.map(|namespace| (namespace, found))
}
fn catalog_model(kind: ModelType, provider: &str, id: &str) -> Value {
match model_resolver::find_model_of_type(provider, id, kind) {
Ok(entry) => model_info(&entry),
Err(_) => Value::Null,
}
}
impl ScriptHost for Host {
fn call_tool(
&self,
name: String,
args: Option<Value>,
) -> BoxFuture<'static, std::result::Result<Value, String>> {
let ctx = self.ctx.clone();
let tools = self.tools.clone();
Box::pin(async move {
let tool = match tools.get(&name) {
Some(tool) => tool.clone(),
None => return Err(format!("Unknown tool \"{name}\"")),
};
let args = args.unwrap_or(Value::Null);
match ctx.execute_tool(name, args).await {
Ok(result) => Host::script_value(&result, Some(&tool)),
Err(e) => Err(e.0),
}
})
}
fn call_global(
&self,
name: String,
args: Vec<Value>,
) -> BoxFuture<'static, std::result::Result<Value, String>> {
if name == MODELS_CLASSIFY || name == MODELS_GENERATE_IMAGES {
let host = self.clone();
return Box::pin(async move {
if name == MODELS_GENERATE_IMAGES {
host.generate_images(&args).await
} else {
host.classify(&args).await
}
});
}
let result = match name.as_str() {
"models.getModelsOfType" => self.models_of_type(&args, false),
"models.getAvailableOfType" => self.models_of_type(&args, true),
"models.getModelOfType" => self.model_of_type(&args),
"searchTools" => self.search_tools(&args),
"describeTool" => self.describe_tool(&args),
"describeNamespace" => self.describe_namespace(&args),
other => Err(format!("Unknown global \"{other}\"")),
};
Box::pin(async move { result })
}
}
async fn run_script(code: &str, ctx: ToolExecCtx) -> std::result::Result<ToolResult, ToolError> {
if let Some(entries) = &ctx.session_branch_entries {
store::seed_from_entries(entries);
}
let parsed = source::parse_source(code).map_err(ToolError)?;
let tools: Vec<ExtensionTool> = ctx.script_tools.iter().cloned().collect();
let samples: HashMap<String, String> = tools
.iter()
.map(|t| {
(
t.name.clone(),
declarations::render_tool_sample(
&t.name,
&t.description,
&t.prompt_guidelines,
Some(&t.parameters),
t.output_schema.as_ref(),
),
)
})
.collect();
let by_name: HashMap<String, ExtensionTool> =
tools.iter().map(|t| (t.name.clone(), t.clone())).collect();
let script_tools: Vec<ScriptTool> = tools
.iter()
.map(|t| ScriptTool {
name: t.name.clone(),
js_name: declarations::to_codemode_identifier(&t.name),
description: samples.get(&t.name).cloned().unwrap_or_default(),
})
.collect();
let request = SandboxRequest {
code: parsed.code,
tools: script_tools,
globals: vec![
ScriptGlobal {
name: "searchTools".to_string(),
spread: true,
},
ScriptGlobal {
name: "describeTool".to_string(),
spread: true,
},
ScriptGlobal {
name: "describeNamespace".to_string(),
spread: true,
},
ScriptGlobal {
name: "models.getModelsOfType".to_string(),
spread: true,
},
ScriptGlobal {
name: "models.getAvailableOfType".to_string(),
spread: true,
},
ScriptGlobal {
name: "models.getModelOfType".to_string(),
spread: true,
},
ScriptGlobal {
name: MODELS_CLASSIFY.to_string(),
spread: true,
},
ScriptGlobal {
name: MODELS_GENERATE_IMAGES.to_string(),
spread: true,
},
],
store: store::snapshot(),
timeout_ms: parsed.options.timeout_ms,
cancel: ctx.parent_abort.clone(),
};
let model_globals = Arc::new(StdMutex::new(ModelGlobalsState::default()));
let host = Arc::new(Host {
ctx: ctx.clone(),
tools: Arc::new(by_name),
samples: Arc::new(samples),
model_globals: model_globals.clone(),
model_call_gate: Arc::new(Semaphore::new(MAX_CONCURRENT_MODEL_CALLS)),
});
let started = Instant::now();
let outcome = sandbox::run(request, host).await;
let wall_seconds = started.elapsed().as_secs_f64();
let mut script_output: Vec<ScriptOutput> = outcome.output.clone();
if outcome.ok
&& let Some(value) = &outcome.value
{
script_output.push(ScriptOutput::Text(value_text(value)));
}
let mut items = format_output(&script_output);
if outcome.ok {
store::apply_writes(&outcome.store_writes);
let generated = model_globals.lock().unwrap().generated_images;
let shown = items
.iter()
.filter(|i| matches!(i, FormattedItem::Image { .. }))
.count();
if let Some(note) = unshown_images_note(generated, shown) {
items.push(FormattedItem::Text(note));
}
} else {
items.push(FormattedItem::Text(format!(
"Script error:\n{}",
format_error(
outcome.error.as_ref(),
outcome.error_kind,
&nested_call_summary(&ctx)
)
)));
}
let items = join_adjacent_text(items);
let texts: Vec<String> = items
.iter()
.filter_map(|item| match item {
FormattedItem::Text(text) => Some(text.clone()),
FormattedItem::Image { .. } => None,
})
.collect();
let max_tokens = parsed
.options
.max_output_tokens
.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS);
let (mut body, full_output_path) = truncate_output(&texts, max_tokens);
let attachments = save_images(&items, &mut body);
let header = format!(
"Script {}\nWall time {wall_seconds:.1} seconds\nOutput:\n",
if outcome.ok { "completed" } else { "failed" }
);
let text = format!("{header}{body}");
let (model_calls, model_usage) = {
let mut state = model_globals.lock().unwrap();
(std::mem::take(&mut state.calls), state.usage.take())
};
let mut result = if outcome.ok {
ToolResult::text(text)
} else {
ToolResult::error(text)
};
result.attachments = attachments;
result.usage = model_usage;
let mut details = serde_json::Map::new();
if let Some(path) = full_output_path {
details.insert("fullOutputPath".to_string(), json!(path));
}
if !model_calls.is_empty() {
details.insert("modelCalls".to_string(), Value::Array(model_calls));
}
if !details.is_empty() {
result.details = Some(Value::Object(details));
}
Ok(result)
}
enum FormattedItem {
Text(String),
Image {
data: String,
mime_type: String,
},
}
fn format_output(items: &[ScriptOutput]) -> Vec<FormattedItem> {
let total = items
.iter()
.filter(|item| matches!(item, ScriptOutput::Text(_)))
.count();
let mut out: Vec<FormattedItem> = Vec::new();
let mut console_lines: Vec<&str> = Vec::new();
let mut index = 0usize;
for item in items {
match item {
ScriptOutput::Image { data, mime_type } => out.push(FormattedItem::Image {
data: data.clone(),
mime_type: mime_type.clone(),
}),
ScriptOutput::Console(line) => console_lines.push(line),
ScriptOutput::Text(text) => {
index += 1;
out.push(FormattedItem::Text(if total > 1 {
format!("==> text {index}/{total} <==\n{text}")
} else {
text.clone()
}));
}
}
}
if !console_lines.is_empty() {
out.push(FormattedItem::Text(format!(
"<console_output>\n{}\n</console_output>",
console_lines.join("\n")
)));
}
out
}
fn join_adjacent_text(items: Vec<FormattedItem>) -> Vec<FormattedItem> {
let mut joined: Vec<FormattedItem> = Vec::new();
for item in items {
match (joined.last_mut(), item) {
(Some(FormattedItem::Text(last)), FormattedItem::Text(text)) => {
if !last.is_empty() && !last.ends_with('\n') {
last.push('\n');
}
last.push_str(&text);
}
(_, item) => joined.push(item),
}
}
joined
}
fn save_images(items: &[FormattedItem], text: &mut String) -> Vec<ToolResultAttachment> {
let mut attachments: Vec<ToolResultAttachment> = Vec::new();
let mut labels: HashMap<String, String> = HashMap::new();
for item in items {
let FormattedItem::Image { data, mime_type } = item else {
continue;
};
let label = labels
.entry(data.clone())
.or_insert_with(|| image_output_label(data, mime_type));
text.push('\n');
text.push_str(label);
attachments.push(ToolResultAttachment {
data_base64: data.clone(),
mime_type: mime_type.clone(),
original_size: None,
converted_from: None,
});
}
attachments
}
fn image_output_label(data: &str, mime_type: &str) -> String {
let bytes = match base64::engine::general_purpose::STANDARD.decode(data) {
Ok(bytes) => bytes,
Err(err) => return format!("[Image could not be saved: invalid base64 ({err})]"),
};
let kind = format!("{mime_type}, {}", truncate::format_size(bytes.len()));
let Some((_, extension)) = IMAGE_FILE_EXTENSIONS.iter().find(|(m, _)| *m == mime_type) else {
return format!(
"[Image ({kind}) could not be saved: no file extension for image type {mime_type}]"
);
};
match output_files::write_output_file("codemode", extension, &bytes) {
Ok(path) => format!("[Image saved to {} ({kind})]", path.display()),
Err(err) => format!("[Image ({kind}) could not be saved: {err}]"),
}
}
fn unshown_images_note(generated: usize, shown: usize) -> Option<String> {
(generated > 0 && shown == 0).then(|| {
format!(
"Note: models.generateImages() returned {generated} image{} that the script did not show. Show each image block of result.output with image(block).",
if generated == 1 { "" } else { "s" }
)
})
}
fn nested_call_summary(ctx: &ToolExecCtx) -> Vec<(String, bool)> {
let log = ctx.nested_calls.lock().unwrap();
log.calls
.iter()
.filter_map(|call| {
let name = call.get("toolName").and_then(|v| v.as_str())?.to_string();
let is_error = call
.get("isError")
.and_then(|v| v.as_bool())
.unwrap_or(false);
Some((name, is_error))
})
.collect()
}
fn format_call_summary(calls: &[(String, bool)]) -> String {
if calls.is_empty() {
return "No tool calls were made.".to_string();
}
let listed: Vec<String> = calls
.iter()
.map(|(name, is_error)| format!("{name} ({})", if *is_error { "error" } else { "ok" }))
.collect();
format!(
"Tool calls made before the failure (they are not undone): {}",
listed.join(", ")
)
}
fn format_error(
error: Option<&ScriptError>,
kind: sandbox::ErrorKind,
calls: &[(String, bool)],
) -> String {
let name = error.map(|e| e.name.as_str()).unwrap_or("Error");
let message = error.map(|e| e.message.as_str()).unwrap_or("script failed");
let head = match (kind, error) {
(sandbox::ErrorKind::Script, Some(err)) => err
.stack
.clone()
.unwrap_or_else(|| format!("{name}: {message}")),
(sandbox::ErrorKind::Timeout, _) => format!("Script timed out: {message}"),
(sandbox::ErrorKind::Aborted, _) => format!("Script aborted: {message}"),
(sandbox::ErrorKind::Sandbox, _) => format!("Script sandbox failed: {message}"),
(sandbox::ErrorKind::Script, None) => format!("{name}: {message}"),
};
format!("{head}\n\n{}", format_call_summary(calls))
}
fn truncate_output(texts: &[String], max_tokens: u64) -> (String, Option<String>) {
let combined = texts.join("\n");
let budget = (max_tokens as usize).saturating_mul(CHARS_PER_TOKEN);
if texts.is_empty() || combined.chars().count() <= budget {
return (combined, None);
}
let head_chars = budget / 2;
let tail_chars = budget - head_chars;
let total_chars = combined.chars().count();
let head: String = combined.chars().take(head_chars).collect();
let tail: String = combined
.chars()
.skip(total_chars.saturating_sub(tail_chars))
.collect();
let removed = total_chars.saturating_sub(head_chars + tail_chars);
let lines = combined.lines().count();
let mut text = format!(
"Warning: truncated output (original token count: {})\nTotal output lines: {lines}\n\n{head}…{} tokens truncated…{tail}",
total_chars.div_ceil(CHARS_PER_TOKEN),
removed.div_ceil(CHARS_PER_TOKEN)
);
match spill_output(&combined) {
Ok(path) => {
text.push_str(&format!(
"\n\n[Full output: {path} (read with offset/limit)]"
));
(text, Some(path))
}
Err(err) => {
text.push_str(&format!("\n\n[Could not save the full output: {err}]"));
(text, None)
}
}
}
fn spill_output(text: &str) -> std::result::Result<String, String> {
output_files::write_output_file("codemode", ".txt", text.as_bytes())
.map(|path| path.to_string_lossy().to_string())
.map_err(|e| e.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::extensions::ExtensionUiRequest;
fn setup() -> (
crate::test_support::AgentDirGuard,
std::sync::MutexGuard<'static, ()>,
) {
let guard = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
(crate::test_support::AgentDirGuard::temp(), guard)
}
fn drain_notify_text() -> String {
let mut out = String::new();
while let Some(req) = crate::core::extensions::take_pending_ui() {
match req {
ExtensionUiRequest::Notify { text, .. } => out.push_str(&text),
ExtensionUiRequest::NotifyRich { spans, .. } => {
for s in spans {
out.push_str(&s.text);
}
}
_ => {}
}
}
out
}
#[test]
fn subcommands_match_parser() {
let names: Vec<&str> = SUBCOMMANDS.iter().map(|s| s.name).collect();
assert_eq!(names, vec!["settings", "show"]);
}
#[test]
fn command_toggles_panel_and_shows_config() {
let (_ad, _g) = setup();
crate::core::extensions::register_extension(Codemode);
crate::core::extensions::set_extension_enabled(EXT, true);
let mut st = App::new();
assert!(!command_codemode(&mut st, "codemode"));
assert!(st.ext_settings.is_open_for(EXT), "无参应打开设置面板");
let keys: Vec<&str> = st
.ext_settings
.items
.iter()
.map(|i| i.key.as_str())
.collect();
assert_eq!(keys, vec!["mode", "inlineBudget"], "面板应展示两项配置");
assert!(!command_codemode(&mut st, "codemode"));
assert!(!st.ext_settings.is_open_for(EXT), "再次无参应关闭面板");
assert!(!command_codemode(&mut st, "codemode settings"));
assert!(st.ext_settings.is_open_for(EXT), "settings 应打开设置面板");
assert!(!command_codemode(&mut st, "codemode settings"));
assert!(
!st.ext_settings.is_open_for(EXT),
"再次 settings 应关闭面板"
);
while crate::core::extensions::take_pending_ui().is_some() {}
assert!(!command_codemode(&mut st, "codemode show"));
let text = drain_notify_text();
assert!(text.contains("mode: on"), "got: {text}");
assert!(text.contains("inlineBudget: 3000"), "got: {text}");
assert!(text.contains("codemode.json"), "got: {text}");
assert!(!command_codemode(&mut st, "codemode bogus"));
let text = drain_notify_text();
assert!(text.contains("unknown subcommand"), "got: {text}");
crate::core::extensions::set_extension_enabled(EXT, false);
_ = crate::core::extensions::unregister_extension(EXT);
}
#[test]
fn tool_definition_shape() {
let tool = tool_definition();
assert_eq!(tool.name, TOOL_NAME);
assert_eq!(tool.exposure, ToolExposure::ModelOnly);
assert_eq!(tool.parameters["required"][0], json!("code"));
assert_eq!(
tool.parameters["properties"]["code"]["type"],
json!("string")
);
assert!(tool.description.contains("tools"));
assert_eq!(tool.prompt_guidelines.len(), 1);
let grammar = tool
.grammar_sampling
.as_ref()
.expect("codemode 应声明语法约束采样");
assert!(
grammar
.openai_lark
.as_deref()
.unwrap_or_default()
.contains("options_source"),
"lark 语法应约束首行选项行"
);
assert!(grammar.openai_regex.is_none());
}
#[test]
fn script_seeds_store_from_ctx_branch_entries() {
let (_ad, _g) = setup();
store::reset();
let ctx_with = |entries: Option<Value>| ToolExecCtx {
execute_tool: crate::core::extensions::unavailable_tool_exec(),
parent_tool_call_id: None,
nested_calls: Default::default(),
session_branch_entries: entries.map(|e| std::sync::Arc::new(vec![e])),
script_tools: Default::default(),
cwd: "/tmp/x".to_string(),
make_sub_agent: std::sync::Arc::new(|_| Err("unused".to_string())),
parent_abort: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
agent_id: None,
depth: 0,
parent_model: None,
script_call: false,
};
let rt = tokio::runtime::Runtime::new().unwrap();
let branch_entry = json!({
"customType": store::ENTRY_TYPE,
"data": { "set": { "branch": "from-branch" } },
});
let out = rt
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({ "code": "return load(\"branch\")" }),
ctx_with(Some(branch_entry)),
))
.expect("script runs");
assert!(out.text.contains("from-branch"), "got: {}", out.text);
let out = rt
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({ "code": "return load(\"branch\")" }),
ctx_with(None),
))
.expect("script runs");
assert!(out.text.contains("from-branch"), "got: {}", out.text);
let out = rt
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({ "code": "return String(load(\"branch\"))" }),
ctx_with(Some(json!({ "customType": "user", "data": {} }))),
))
.expect("script runs");
assert!(out.text.contains("undefined"), "got: {}", out.text);
store::reset();
}
#[test]
fn script_generate_images_validates_arguments() {
let (_ad, _g) = setup();
let run = |code: &str| {
let ctx = ToolExecCtx {
execute_tool: crate::core::extensions::unavailable_tool_exec(),
parent_tool_call_id: None,
nested_calls: Default::default(),
session_branch_entries: None,
script_tools: Default::default(),
cwd: "/tmp/x".to_string(),
make_sub_agent: std::sync::Arc::new(|_| Err("unused".to_string())),
parent_abort: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
agent_id: None,
depth: 0,
parent_model: None,
script_call: false,
};
tokio::runtime::Runtime::new()
.unwrap()
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({ "code": code }),
ctx,
))
.expect("script runs")
};
let out = run("return typeof models.generateImages");
assert!(out.text.contains("function"), "got: {}", out.text);
let out = run("return await models.generateImages({}, {})");
assert!(out.text.contains("Script failed"), "got: {}", out.text);
assert!(
out.text
.contains("models.generateImages() expects a model of type \"image\""),
"got: {}",
out.text
);
assert!(
out.text.contains("models.getAvailableOfType(\"image\")"),
"got: {}",
out.text
);
}
#[test]
fn unshown_images_note_only_when_nothing_shown() {
let note = unshown_images_note(1, 0).unwrap();
assert!(note.contains("returned 1 image that"));
assert!(note.contains("with image(block)"));
assert!(
unshown_images_note(3, 0)
.unwrap()
.contains("returned 3 images that")
);
assert!(unshown_images_note(3, 1).is_none());
assert!(unshown_images_note(0, 0).is_none());
}
fn render_text(items: &[ScriptOutput]) -> String {
join_adjacent_text(format_output(items))
.iter()
.filter_map(|item| match item {
FormattedItem::Text(text) => Some(text.clone()),
FormattedItem::Image { .. } => None,
})
.collect::<Vec<_>>()
.join("\n")
}
#[test]
fn format_output_labels_items_and_groups_console() {
let text = render_text(&[
ScriptOutput::Text("first".to_string()),
ScriptOutput::Console("log a".to_string()),
ScriptOutput::Text("second".to_string()),
ScriptOutput::Console("log b".to_string()),
]);
assert!(text.contains("==> text 1/2 <==\nfirst"), "{text}");
assert!(text.contains("==> text 2/2 <==\nsecond"), "{text}");
assert!(
text.contains("<console_output>\nlog a\nlog b\n</console_output>"),
"{text}"
);
assert!(
text.find("<console_output>").unwrap() > text.find("second").unwrap(),
"console 块应排在其它输出之后: {text}"
);
let single = render_text(&[ScriptOutput::Text("only".to_string())]);
assert_eq!(single.trim(), "only");
assert!(!single.contains("==> text"), "{single}");
}
#[test]
fn join_adjacent_text_inserts_a_separator() {
fn first(items: Vec<FormattedItem>) -> String {
match items.into_iter().next() {
Some(FormattedItem::Text(text)) => text,
_ => panic!("应合并成一段文本"),
}
}
assert_eq!(
first(join_adjacent_text(vec![
FormattedItem::Text("a".to_string()),
FormattedItem::Text("b".to_string()),
])),
"a\nb"
);
assert_eq!(
first(join_adjacent_text(vec![
FormattedItem::Text("a\n".to_string()),
FormattedItem::Text("b".to_string()),
])),
"a\nb"
);
let items = join_adjacent_text(vec![
FormattedItem::Text("a".to_string()),
FormattedItem::Image {
data: "AA==".to_string(),
mime_type: "image/png".to_string(),
},
]);
assert_eq!(items.len(), 2);
}
#[test]
fn script_images_are_spilled_with_a_path_label() {
let data = base64::engine::general_purpose::STANDARD.encode(b"fake-png");
let items = format_output(&[
ScriptOutput::Image {
data: data.clone(),
mime_type: "image/png".to_string(),
},
ScriptOutput::Image {
data,
mime_type: "image/png".to_string(),
},
]);
let mut text = String::new();
let attachments = save_images(&items, &mut text);
assert_eq!(attachments.len(), 2, "两次展示都是附件");
let labels: Vec<&str> = text
.lines()
.filter(|line| line.starts_with("[Image saved to "))
.collect();
assert_eq!(labels.len(), 2, "每张图都要留一条提示: {text}");
assert_eq!(labels[0], labels[1], "同一张图的提示应相同");
assert!(
labels[0].ends_with("(image/png, 8B)]"),
"提示形状应与 pi 一致: {}",
labels[0]
);
let path = labels[0]
.trim_start_matches("[Image saved to ")
.split(" (")
.next()
.unwrap();
assert_eq!(std::fs::read(path).unwrap(), b"fake-png");
let _ = std::fs::remove_file(path);
}
#[test]
fn script_images_report_save_failures_in_the_label() {
let items = format_output(&[ScriptOutput::Image {
data: base64::engine::general_purpose::STANDARD.encode(b"x"),
mime_type: "image/bmp".to_string(),
}]);
let mut text = String::new();
let attachments = save_images(&items, &mut text);
assert!(
text.contains("could not be saved: no file extension for image type image/bmp"),
"{text}"
);
assert_eq!(attachments.len(), 1, "附件仍要保留");
}
#[test]
fn script_value_exposes_structured_output() {
let mut tool = ExtensionTool::simple("read", "Read a file", json!({}), "read");
tool.output_schema = Some(json!({ "anyOf": [{ "type": "string" }, { "type": "object" }] }));
let result =
ToolResult::text("Read image file [image/png]").with_structured_content(json!({
"type": "image",
"data": "AAAA",
"mimeType": "image/png",
"note": "Read image file [image/png]",
}));
let value = Host::script_value(&result, Some(&tool)).unwrap();
assert_eq!(value["type"], "image");
assert_eq!(value["data"], "AAAA");
assert_eq!(value["mimeType"], "image/png");
let plain = ExtensionTool::simple("plain", "d", json!({}), "s");
let value = Host::script_value(&result, Some(&plain)).unwrap();
assert_eq!(value, json!("Read image file [image/png]"));
}
#[test]
fn classifier_context_shape_is_validated() {
let ok = json!({
"state": { "message": "hi" },
"images": [{ "type": "image", "data": "AAAA", "mimeType": "image/png" }],
"questions": {
"q": { "type": "choice", "instructions": "pick", "criteria": { "a": "A" } }
}
});
assert!(classifier_context_arg(Some(&ok)).is_ok());
let err = classifier_context_arg(Some(&json!("nope"))).unwrap_err();
assert!(err.contains("expects a context object"), "{err}");
let bad = json!({
"state": "x",
"questions": { "q": { "type": "bool", "instructions": "i", "criteria": { "true": "y", "false": "n" } } }
});
let err = classifier_context_arg(Some(&bad)).unwrap_err();
assert!(err.contains("context.state must be an object"), "{err}");
let bad = json!({
"state": {},
"images": [{ "type": "text", "text": "x" }],
"questions": { "q": { "type": "bool", "instructions": "i", "criteria": { "true": "y", "false": "n" } } }
});
let err = classifier_context_arg(Some(&bad)).unwrap_err();
assert!(
err.contains("context.images[0] must be an image block"),
"{err}"
);
let err =
classifier_context_arg(Some(&json!({ "state": {}, "questions": {} }))).unwrap_err();
assert!(
err.contains("context.questions must map question IDs to questions"),
"{err}"
);
let bad = json!({ "state": {}, "questions": { "q": { "type": "nope" } } });
let err = classifier_context_arg(Some(&bad)).unwrap_err();
assert!(err.contains("context is malformed"), "{err}");
}
#[test]
fn script_describe_namespace_lists_members() {
let (_ad, _g) = setup();
let namespaced = |name: &str, description: &str| {
let mut tool = ExtensionTool::simple(
name,
description,
json!({"type": "object", "properties": {}}),
"",
);
tool.namespace = Some("dev-radius".to_string());
tool
};
let ctx = ToolExecCtx {
execute_tool: crate::core::extensions::unavailable_tool_exec(),
parent_tool_call_id: None,
nested_calls: Default::default(),
session_branch_entries: None,
script_tools: std::sync::Arc::new(vec![
namespaced("mcp__dev-radius__put", "write a thing"),
namespaced("mcp__dev-radius__get", "read a thing"),
]),
cwd: "/tmp/x".to_string(),
make_sub_agent: std::sync::Arc::new(|_| Err("unused".to_string())),
parent_abort: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
agent_id: None,
depth: 0,
parent_model: None,
script_call: false,
};
let rt = tokio::runtime::Runtime::new().unwrap();
let out = rt
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({
"code": "const ns = await describeNamespace('mcp__dev_radius'); \
return ns.namespace + '|' + ns.tools.map((t) => t.name).join(',')"
}),
ctx.clone(),
))
.expect("script runs");
assert!(
out.text
.contains("dev-radius|mcp__dev_radius__get,mcp__dev_radius__put"),
"got: {}",
out.text
);
let out = rt
.block_on(Codemode.execute_tool_async(
TOOL_NAME.to_string(),
json!({ "code": "return String(await describeNamespace('nope'))" }),
ctx,
))
.expect("script runs");
assert!(out.text.contains("null"), "got: {}", out.text);
}
#[test]
fn truncation_passthrough() {
let (text, path) = truncate_output(&["a".to_string(), "b".to_string()], 10);
assert_eq!(text, "a\nb");
assert!(path.is_none());
}
#[test]
fn truncation_keeps_head_and_tail() {
let text = "x".repeat(1000);
let (out, path) = truncate_output(std::slice::from_ref(&text), 10);
assert!(out.starts_with("Warning: truncated output"), "{out}");
assert!(out.contains("tokens truncated"), "{out}");
assert!(out.contains("[Full output:"), "{out}");
let path = path.expect("应落盘");
assert_eq!(std::fs::read_to_string(&path).unwrap(), text);
let _ = std::fs::remove_file(&path);
}
#[test]
fn error_text_includes_stack_and_call_summary() {
let err = sandbox::ScriptError {
name: "TypeError".to_string(),
message: "nope".to_string(),
stack: Some("TypeError: nope\n at x".to_string()),
};
let text = format_error(
Some(&err),
sandbox::ErrorKind::Script,
&[("read".to_string(), false), ("bash".to_string(), true)],
);
assert!(text.starts_with("TypeError: nope"));
assert!(text.contains("read (ok), bash (error)"));
assert_eq!(
format_error(None, sandbox::ErrorKind::Timeout, &[]),
"Script timed out: script failed\n\nNo tool calls were made."
);
assert_eq!(
format_error(None, sandbox::ErrorKind::Script, &[]),
"Error: script failed\n\nNo tool calls were made."
);
}
#[test]
fn value_text_shapes() {
assert_eq!(value_text(&json!("hi")), "hi");
assert_eq!(value_text(&json!({"a": 1})), "{\"a\":1}");
assert_eq!(value_text(&json!(null)), "null");
}
#[test]
fn model_type_and_provider_arguments_are_validated() {
assert_eq!(
model_type_arg(Some(&json!("chat"))).unwrap(),
ModelType::Chat
);
assert_eq!(
model_type_arg(Some(&json!("classifier"))).unwrap(),
ModelType::Classifier
);
assert_eq!(
model_type_arg(Some(&json!("vision"))).unwrap_err(),
"Unknown model type \"vision\". Use \"chat\", \"image\", or \"classifier\"."
);
assert!(model_type_arg(None).unwrap_err().contains("undefined"));
assert_eq!(provider_arg(None).unwrap(), None);
assert_eq!(provider_arg(Some(&json!(null))).unwrap(), None);
assert_eq!(
provider_arg(Some(&json!("openai"))).unwrap().as_deref(),
Some("openai")
);
assert_eq!(
provider_arg(Some(&json!(7))).unwrap_err(),
"provider must be a string"
);
}
#[test]
fn resolve_model_call_uses_only_provider_and_id() {
let _guard = crate::test_support::AgentDirGuard::temp();
let model = resolve_model_call(
MODELS_CLASSIFY,
ModelType::Classifier,
Some(&json!({
"provider": "typesafe",
"id": "jev-latest",
"apiKey": "sk-secret",
"headers": { "authorization": "Bearer sk-secret" },
})),
)
.unwrap();
assert_eq!(model.provider, "typesafe");
assert_eq!(model.model_id, "jev-latest");
assert!(
!model.api_key.contains("sk-secret"),
"脚本自带凭据不得透传：{}",
model.api_key
);
let err = resolve_model_call(
MODELS_CLASSIFY,
ModelType::Classifier,
Some(&json!({ "provider": "typesafe" })),
)
.unwrap_err();
assert!(
err.starts_with("models.classify() expects a model of type \"classifier\"")
&& err.contains("getAvailableOfType(\"classifier\")"),
"{err}"
);
let err = resolve_model_call(
MODELS_CLASSIFY,
ModelType::Classifier,
Some(&json!({ "provider": "typesafe", "id": "nope" })),
)
.unwrap_err();
assert!(
err.starts_with("Unknown classifier model \"typesafe/nope\""),
"{err}"
);
let err = resolve_model_call(
MODELS_GENERATE_IMAGES,
ModelType::Image,
Some(&json!({ "provider": "typesafe", "id": "jev-latest" })),
)
.unwrap_err();
assert!(
err.contains("is a classifier model, not a image model"),
"{err}"
);
let err = resolve_model_call(MODELS_CLASSIFY, ModelType::Classifier, None).unwrap_err();
assert!(err.contains("got undefined"), "{err}");
}
#[test]
fn images_context_requires_non_empty_blocks() {
let input = images_context_arg(Some(&json!({
"input": [
{ "type": "text", "text": "a red fox" },
{ "type": "image", "data": "AA==", "mimeType": "image/png" },
]
})))
.unwrap();
assert_eq!(input.len(), 2);
assert!(matches!(&input[0], ImageContent::Text { text } if text == "a red fox"));
assert!(
matches!(&input[1], ImageContent::Image { data, mime_type } if data == "AA==" && mime_type == "image/png")
);
for bad in [
json!(null),
json!({ "prompt": "a red fox" }),
json!({ "input": [] }),
json!({ "input": [{ "type": "text" }] }),
json!({ "input": ["a red fox"] }),
] {
let err = images_context_arg(Some(&bad)).unwrap_err();
assert!(
err.starts_with("models.generateImages()")
&& err.contains("Expected context: { input: ["),
"{bad}: {err}"
);
}
}
#[test]
fn catalog_lookups_return_model_info() {
let _guard = crate::test_support::AgentDirGuard::temp();
let chat = catalog_models(ModelType::Chat, Some("openrouter".into()), false);
assert!(!chat.is_empty());
assert!(
chat.iter()
.all(|m| m["type"] == "chat" && m["provider"] == "openrouter")
);
assert!(
chat.iter()
.all(|m| m["api"].is_string() && m["input"].is_array() && m["id"].is_string())
);
let classifiers = catalog_models(ModelType::Classifier, Some("openrouter".into()), false);
assert!(!classifiers.is_empty());
assert!(classifiers.iter().all(|m| m["type"] == "classifier"));
let one = catalog_model(ModelType::Classifier, "openrouter", "~typesafe/jev-latest");
assert_eq!(one["type"], "classifier");
assert_eq!(one["provider"], "openrouter");
assert_eq!(
catalog_model(ModelType::Chat, "openrouter", "~typesafe/jev-latest"),
Value::Null
);
assert_eq!(
catalog_model(ModelType::Classifier, "openrouter", "nope"),
Value::Null
);
}
#[test]
fn available_models_require_credentials() {
let _guard = crate::test_support::AgentDirGuard::temp();
let configured: HashSet<String> = auth::list_configured_providers().into_iter().collect();
let available = catalog_models(ModelType::Chat, None, true);
assert!(
available
.iter()
.all(|m| configured.contains(m["provider"].as_str().unwrap())),
"{available:?}"
);
assert!(available.len() <= catalog_models(ModelType::Chat, None, false).len());
}
#[test]
fn namespace_aliases_normalize() {
let key = namespace_key("mcp__dev-radius");
assert_eq!(key, namespace_key("mcp__dev_radius"));
assert_eq!(key, namespace_key("dev-radius"));
assert_eq!(key, namespace_key("dev_radius"));
assert_ne!(key, namespace_key("dev"));
assert_ne!(key, namespace_key("dev_radius2"));
assert_eq!(namespace_key(""), "");
assert_eq!(namespace_key("mcp__"), "");
}
#[test]
fn namespace_tools_lists_sorted_members() {
let tool = |name: &str, namespace: Option<&str>| {
let mut tool = ExtensionTool::simple(name, "d", json!({}), "");
tool.namespace = namespace.map(|n| n.to_string());
tool
};
let tools: HashMap<String, ExtensionTool> = [
tool("mcp__dev-radius__b", Some("dev-radius")),
tool("mcp__dev-radius__a", Some("dev-radius")),
tool("plain_tool", None),
tool("other_ns_tool", Some("other")),
]
.into_iter()
.map(|t| (t.name.clone(), t))
.collect();
let samples: HashMap<String, String> = tools
.keys()
.map(|k| (k.clone(), format!("sample of {k}")))
.collect();
let (namespace, found) = namespace_tools(&tools, &samples, "mcp__dev_radius").unwrap();
assert_eq!(namespace, "dev-radius", "返回工具表里声明的原名");
assert_eq!(
found,
vec![
(
"mcp__dev_radius__a".to_string(),
"sample of mcp__dev-radius__a".to_string()
),
(
"mcp__dev_radius__b".to_string(),
"sample of mcp__dev-radius__b".to_string()
),
],
"按脚本标识符排序，且只含该命名空间的工具"
);
assert!(namespace_tools(&tools, &samples, "dev").is_none());
assert!(namespace_tools(&tools, &samples, "mcp__plain_tool").is_none());
}
}