diff --git a/Cargo.lock b/Cargo.lock index caa8c4ace..9f8ed3658 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7863,9 +7863,9 @@ dependencies = [ [[package]] name = "pctx_code_mode" -version = "0.3.0" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d9f5ffc774d0c440e780aaa2593fbf5c5710e5b5c63a6b498116d3f46a09fe0" +checksum = "f55b3476b133078b55bc808aad1ff880428bf3de3127415ec33ceab50ef230d9" dependencies = [ "futures", "pctx_codegen", @@ -7883,9 +7883,9 @@ dependencies = [ [[package]] name = "pctx_codegen" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b05191d9a37d1101f3a78e770e1a2b5521e75ce718f1165fb15ac16025d297b" +checksum = "a56b53c02b41420f6bd3baa7046a5a7509a38114fffa52f737f44c9044eac646" dependencies = [ "biome_formatter", "biome_js_formatter", @@ -7965,9 +7965,9 @@ dependencies = [ [[package]] name = "pctx_registry" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f5a53bf73ba98a5352e788391f47c33af096c696c64e130110b85cc51336e20" +checksum = "251c06ac888c315c7e23d210e4e9ffa88d5dc08225c0d8de96f1263837d72896" dependencies = [ "deno_error", "pctx_config", diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index d96a56019..d9b439cc9 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -217,7 +217,7 @@ aws-lc-rs = { version = "1.17", default-features = false, optional = true } openssl = { version = "0.10.66", default-features = false, optional = true } pem = { version = "3.0.2", default-features = false, features = ["std"], optional = true } -pctx_code_mode = { version = "0.3", default-features = false, optional = true } +pctx_code_mode = { version = "0.4.1", default-features = false, optional = true } # These are needed because temporal_rs 0.1 (a transitive dep via PCTX) enables unstable features on icu_calendar without pinning the dependency version # A fix is available in temporal_rs 0.2 but PCTX has not updated diff --git a/crates/goose/src/agents/platform_extensions/code_execution.rs b/crates/goose/src/agents/platform_extensions/code_execution.rs index 0a9ba0107..11bd3e06b 100644 --- a/crates/goose/src/agents/platform_extensions/code_execution.rs +++ b/crates/goose/src/agents/platform_extensions/code_execution.rs @@ -7,7 +7,7 @@ use async_trait::async_trait; use pctx_code_mode::{ config::ToolDisclosure, descriptions::{tools as tool_descriptions, workflow::get_workflow_description}, - model::{CallbackConfig, ExecuteBashInput, ExecuteInput, GetFunctionDetailsInput}, + model::{CallbackConfig, ExecuteBashInput, ExecuteTypescriptInput, GetFunctionDetailsInput}, registry::{CallbackFn, PctxRegistry}, CodeMode, }; @@ -19,7 +19,6 @@ use schemars::{schema_for, JsonSchema}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::collections::hash_map::DefaultHasher; -use std::collections::{HashMap, HashSet}; use std::future::Future; use std::hash::{Hash, Hasher}; use std::pin::Pin; @@ -30,117 +29,6 @@ use tokio_util::sync::CancellationToken; pub static EXTENSION_NAME: &str = "code_execution"; -fn sanitize_schema_for_code_mode(schema: &mut Value) { - let Some(obj) = schema.as_object_mut() else { - return; - }; - - let Some(defs_key) = ["$defs", "definitions"] - .into_iter() - .find(|key| obj.get(*key).is_some_and(Value::is_object)) - else { - return; - }; - - let names: Vec = obj[defs_key] - .as_object() - .map(|defs| defs.keys().cloned().collect()) - .unwrap_or_default(); - - let mut edges: HashMap> = HashMap::new(); - if let Some(defs) = obj.get(defs_key).and_then(Value::as_object) { - for name in &names { - let mut refs = HashSet::new(); - if let Some(def_value) = defs.get(name) { - collect_ref_targets(def_value, &mut refs); - } - edges.insert(name.clone(), refs); - } - } - - let cuts = find_cycle_edges(&names, &edges); - if cuts.is_empty() { - return; - } - - if let Some(defs) = obj.get_mut(defs_key).and_then(Value::as_object_mut) { - for (from, to) in &cuts { - if let Some(def_value) = defs.get_mut(from) { - neutralize_refs_to(def_value, to); - } - } - } -} - -fn collect_ref_targets(value: &Value, out: &mut HashSet) { - match value { - Value::Object(map) => { - if let Some(Value::String(r)) = map.get("$ref") { - if let Some(name) = r.rsplit('/').next() { - out.insert(name.to_string()); - } - } - map.values().for_each(|v| collect_ref_targets(v, out)); - } - Value::Array(items) => items.iter().for_each(|v| collect_ref_targets(v, out)), - _ => {} - } -} - -fn neutralize_refs_to(value: &mut Value, target: &str) { - let is_target_ref = matches!( - value.as_object().and_then(|map| map.get("$ref")), - Some(Value::String(r)) if r.rsplit('/').next() == Some(target) - ); - if is_target_ref { - *value = json!({}); - return; - } - match value { - Value::Object(map) => map.values_mut().for_each(|v| neutralize_refs_to(v, target)), - Value::Array(items) => items.iter_mut().for_each(|v| neutralize_refs_to(v, target)), - _ => {} - } -} - -fn find_cycle_edges( - names: &[String], - edges: &HashMap>, -) -> Vec<(String, String)> { - enum State { - InProgress, - Done, - } - - fn visit<'a>( - node: &'a str, - edges: &'a HashMap>, - state: &mut HashMap<&'a str, State>, - cuts: &mut Vec<(String, String)>, - ) { - state.insert(node, State::InProgress); - if let Some(targets) = edges.get(node) { - for target in targets { - match state.get(target.as_str()) { - Some(State::InProgress) => cuts.push((node.to_string(), target.clone())), - Some(State::Done) => {} - None => visit(target, edges, state, cuts), - } - } - } - state.insert(node, State::Done); - } - - let mut state: HashMap<&str, State> = HashMap::new(); - let mut cuts = Vec::new(); - for name in names { - if !state.contains_key(name.as_str()) { - visit(name, edges, &mut state, &mut cuts); - } - } - cuts -} - pub struct CodeExecutionClient { info: InitializeResult, context: PlatformExtensionContext, @@ -162,7 +50,7 @@ struct ToolGraphNode { #[derive(Debug, Serialize, Deserialize, JsonSchema)] pub struct ExecuteWithToolGraph { #[serde(flatten)] - input: ExecuteInput, + input: ExecuteTypescriptInput, /// DAG of tool calls showing execution flow. Each node represents a tool call. /// Use depends_on to show data flow (e.g., node 1 uses output from node 0). #[serde(default)] @@ -212,20 +100,12 @@ impl CodeExecutionClient { (tool.name.to_string(), None) }; - let mut input_schema = json!(tool.input_schema); - sanitize_schema_for_code_mode(&mut input_schema); - - let mut output_schema = tool.output_schema.as_ref().map(|s| json!(s)); - if let Some(schema) = output_schema.as_mut() { - sanitize_schema_for_code_mode(schema); - } - cfgs.push(CallbackConfig { name, namespace, description: tool.description.as_ref().map(|d| d.to_string()), - input_schema: Some(input_schema), - output_schema, + input_schema: Some(json!(tool.input_schema)), + output_schema: tool.output_schema.as_ref().map(|s| json!(s)), }) } Some(cfgs) @@ -360,7 +240,10 @@ impl CodeExecutionClient { ) .await?; - Ok(vec![Content::text(output.markdown())]) + Ok(vec![Content::text(format!( + "Exit Code: {}\n\n# STDOUT\n{}\n\n# STDERR\n{}", + output.exit_code, output.stdout, output.stderr + ))]) } /// Handle the execute typescript tool call @@ -776,6 +659,7 @@ impl CodeModeState { #[cfg(test)] mod tests { use super::*; + use pctx_code_mode::model::FunctionId; #[tokio::test] async fn run_in_deno_runtime_times_out_on_hung_execution() { @@ -1009,96 +893,7 @@ mod tests { } #[test] - fn collect_ref_targets_finds_nested_refs() { - let schema = self_referential_any_schema(); - let mut refs = HashSet::new(); - collect_ref_targets(&schema["$defs"]["Any"], &mut refs); - - assert_eq!(refs, HashSet::from(["Any".to_string()])); - } - - #[test] - fn find_cycle_edges_detects_self_loop() { - let mut edges = HashMap::new(); - edges.insert("Any".to_string(), HashSet::from(["Any".to_string()])); - let names = vec!["Any".to_string()]; - - let cuts = find_cycle_edges(&names, &edges); - - assert_eq!(cuts, vec![("Any".to_string(), "Any".to_string())]); - } - - #[test] - fn find_cycle_edges_detects_longer_cycle_without_flagging_acyclic_refs() { - let mut edges = HashMap::new(); - edges.insert("A".to_string(), HashSet::from(["B".to_string()])); - edges.insert("B".to_string(), HashSet::from(["C".to_string()])); - edges.insert("C".to_string(), HashSet::from(["A".to_string()])); - edges.insert("D".to_string(), HashSet::from(["A".to_string()])); - let names = vec![ - "A".to_string(), - "B".to_string(), - "C".to_string(), - "D".to_string(), - ]; - - let cuts = find_cycle_edges(&names, &edges); - - assert_eq!(cuts, vec![("C".to_string(), "A".to_string())]); - } - - #[test] - fn neutralize_refs_to_replaces_matching_refs_only() { - let mut value = json!({ - "anyOf": [ - {"$ref": "#/$defs/Any"}, - {"$ref": "#/$defs/Other"} - ] - }); - - neutralize_refs_to(&mut value, "Any"); - - assert_eq!(value["anyOf"][0], json!({})); - assert_eq!(value["anyOf"][1], json!({"$ref": "#/$defs/Other"})); - } - - #[test] - fn sanitize_schema_for_code_mode_breaks_self_referential_defs() { - let mut schema = self_referential_any_schema(); - - sanitize_schema_for_code_mode(&mut schema); - - let mut refs = HashSet::new(); - collect_ref_targets(&schema["$defs"]["Any"], &mut refs); - assert!( - !refs.contains("Any"), - "cycle should be broken, got: {schema}" - ); - } - - #[test] - fn sanitize_schema_for_code_mode_leaves_acyclic_schemas_untouched() { - let mut schema = json!({ - "type": "object", - "properties": { - "content": {"$ref": "#/$defs/Content"} - }, - "$defs": { - "Content": {"type": "string"} - } - }); - let original = schema.clone(); - - sanitize_schema_for_code_mode(&mut schema); - - assert_eq!(schema, original); - } - - #[test] - fn code_mode_accepts_previously_crashing_self_referential_schema() { - let mut output_schema = self_referential_any_schema(); - sanitize_schema_for_code_mode(&mut output_schema); - + fn code_mode_preserves_types_for_self_referential_schema() { let cfg = CallbackConfig { name: "retain".to_string(), namespace: Some("hindsight".to_string()), @@ -1108,10 +903,29 @@ mod tests { "properties": {"content": {"type": "string"}}, "required": ["content"] })), - output_schema: Some(output_schema), + output_schema: Some(self_referential_any_schema()), }; - let result = CodeMode::default().with_callback(&cfg); - assert!(result.is_ok(), "{:?}", result.err()); + let code_mode = CodeMode::default() + .with_callback(&cfg) + .expect("recursive schemas should be supported"); + let details = code_mode.get_function_details(GetFunctionDetailsInput { + functions: vec![FunctionId { + mod_name: "Hindsight".to_string(), + fn_name: "retain".to_string(), + }], + }); + let function = details + .functions + .first() + .expect("hindsight.retain should have generated details"); + + assert_ne!(function.output_type, "any"); + assert!( + function.types.contains("export type RetainOutputAny =") + && function.types.contains("[key: string]: RetainOutputAny"), + "expected RetainOutputAny to reference itself, got: {}", + function.types + ); } }