diff --git a/buzz-agent-core/Cargo.lock b/buzz-agent-core/Cargo.lock index e88e70b005..eeb3a8f8ba 100644 --- a/buzz-agent-core/Cargo.lock +++ b/buzz-agent-core/Cargo.lock @@ -1872,6 +1872,7 @@ dependencies = [ [[package]] name = "goose" version = "1.44.0" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "agent-client-protocol", "agent-client-protocol-http", @@ -1972,6 +1973,7 @@ dependencies = [ [[package]] name = "goose-acp-macros" version = "1.44.0" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "quote", "syn 2.0.119", @@ -1980,6 +1982,7 @@ dependencies = [ [[package]] name = "goose-download-manager" version = "0.1.0-alpha.5" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "anyhow", "once_cell", @@ -1992,6 +1995,7 @@ dependencies = [ [[package]] name = "goose-provider-types" version = "0.1.0-alpha.5" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "anyhow", "async-stream", @@ -2017,6 +2021,7 @@ dependencies = [ [[package]] name = "goose-providers" version = "0.1.0-alpha.5" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "anyhow", "async-stream", @@ -2040,6 +2045,7 @@ dependencies = [ [[package]] name = "goose-sdk-types" version = "0.1.0-alpha.5" +source = "git+https://github.com/aaif-goose/goose?rev=305849b71709b95b86ed9f11bd3bc939899c0aab#305849b71709b95b86ed9f11bd3bc939899c0aab" dependencies = [ "agent-client-protocol", "agent-client-protocol-schema", diff --git a/buzz-agent-core/Cargo.toml b/buzz-agent-core/Cargo.toml index c7fe9cc408..fdfc62bc92 100644 --- a/buzz-agent-core/Cargo.toml +++ b/buzz-agent-core/Cargo.toml @@ -46,8 +46,8 @@ path = "src/main.rs" # with default-features = false, so a TLS backend only arrives via this feature # (and `rustls-tls` + `native-tls` are mutually exclusive: see the # compile_error! at goose/src/lib.rs:3-4). -goose = { path = "../../goose/crates/goose", default-features = false, features = ["rustls-tls"] } -goose-provider-types = { path = "../../goose/crates/goose-provider-types" } +goose = { git = "https://github.com/aaif-goose/goose", rev = "305849b71709b95b86ed9f11bd3bc939899c0aab", default-features = false, features = ["rustls-tls"] } +goose-provider-types = { git = "https://github.com/aaif-goose/goose", rev = "305849b71709b95b86ed9f11bd3bc939899c0aab" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "time"] } tokio-util = { version = "0.7", features = ["rt"] } diff --git a/buzz-agent-core/src/lib.rs b/buzz-agent-core/src/lib.rs index 4e12339b86..e7eccce5b5 100644 --- a/buzz-agent-core/src/lib.rs +++ b/buzz-agent-core/src/lib.rs @@ -53,7 +53,14 @@ struct Session { /// Set for the duration of a turn; advertised to steer-capable clients. active_run_id: Option, cancel: Option, + /// Model id currently in force for this session (`session/set_model`). + /// Reported back on `usage_update`. model_override: Option, + /// Set by `session/set_model`, consumed by the next `session/prompt`. + /// Applying it at prompt time rather than immediately matches buzz-agent: + /// the override takes effect "from the next prompt" and never mutates a + /// turn already in flight (`buzz-agent/src/lib.rs:494-502`). + pending_model: Option, accumulated_input_tokens: u64, accumulated_output_tokens: u64, } @@ -361,6 +368,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen active_run_id: None, cancel: None, model_override: None, + pending_model: None, accumulated_input_tokens: 0, accumulated_output_tokens: 0, }, @@ -402,7 +410,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire let cancel = CancellationToken::new(); // Single-flight per session, and capture the agent handle. - let (agent, goose_session_id) = { + let (agent, goose_session_id, pending_model) = { let mut sessions = app.sessions.lock().await; let Some(s) = sessions.get_mut(&p.session_id) else { return wire::send( @@ -425,9 +433,30 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire s.busy = true; s.active_run_id = Some(run_id.clone()); s.cancel = Some(cancel.clone()); - (s.agent.clone(), s.goose_session_id.clone()) + // Take the pending override so a `session/set_model` applies exactly + // once, from the next prompt onward (buzz-agent `lib.rs:494-502`). + ( + s.agent.clone(), + s.goose_session_id.clone(), + s.pending_model.take(), + ) }; + // Apply a pending `session/set_model` before the turn starts. + if let Some(model_id) = pending_model { + if let Err(e) = apply_model(&agent, &goose_session_id, &model_id).await { + { + let mut sessions = app.sessions.lock().await; + if let Some(s) = sessions.get_mut(&p.session_id) { + s.busy = false; + s.active_run_id = None; + s.cancel = None; + } + } + return wire::send(wire_tx, wire::err(id, e.json_rpc_code(), &e.to_string())).await; + } + } + // Advertise the run id so steer-capable clients can target this turn. // `_meta` nests INSIDE `update` — see the module docs. wire::send( @@ -542,7 +571,8 @@ async fn set_model(app: &Arc, id: Value, params: Value, wire_tx: &WireSende let mut sessions = app.sessions.lock().await; match sessions.get_mut(&p.session_id) { Some(s) => { - s.model_override = Some(p.model_id); + s.model_override = Some(p.model_id.clone()); + s.pending_model = Some(p.model_id); drop(sessions); wire::send(wire_tx, wire::ok(id, Value::Null)).await; } @@ -645,6 +675,30 @@ async fn cancel_session(app: &Arc, params: Value) { } } +/// Swap the model for an existing session. +/// +/// Rebuilds the provider (base-url/credential resolution lives in goose's +/// registry) and installs it with the new `ModelConfig`. `SharedProvider` is an +/// `Arc>>` precisely so this is hot-swappable +/// (goose `agents/types.rs:11-12`). +async fn apply_model( + agent: &Arc, + session_id: &str, + model_id: &str, +) -> Result<(), AgentError> { + let provider_name = std::env::var("GOOSE_PROVIDER").unwrap_or_else(|_| "openai".to_string()); + let provider = goose::providers::create(&provider_name, Vec::new()) + .await + .map_err(|e| map_provider_error(&e.to_string()))?; + let model_config = goose::model_config::model_config_from_user_config(&provider_name, model_id) + .map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?; + agent + .update_provider(provider, model_config, session_id) + .await + .map_err(|e| map_provider_error(&e.to_string()))?; + Ok(()) +} + /// Preserve buzz-agent's error taxonomy across provider construction, so the /// harness's JSON-RPC code mapping (-32001 auth, -32002 model-not-found) keeps /// its meaning. diff --git a/buzz-agent-core/tests/stdio_turn.rs b/buzz-agent-core/tests/stdio_turn.rs index 1d19f8b17c..cd985d4731 100644 --- a/buzz-agent-core/tests/stdio_turn.rs +++ b/buzz-agent-core/tests/stdio_turn.rs @@ -317,6 +317,66 @@ fn full_turn_over_stdio_carries_the_persona() { // seen strictly BEFORE the session/prompt response was read. } +#[test] +fn set_model_applies_from_the_next_prompt() { + let (base_url, system_rx) = spawn_fake_provider(); + let home = tempfile::tempdir().expect("tempdir"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path()); + + let id = h.request("initialize", json!({ "protocolVersion": 2 })); + let _ = h.await_response(id); + + let id = h.request( + "session/new", + json!({ "cwd": cwd.path().to_str().unwrap(), "mcpServers": [] }), + ); + let (resp, _) = h.await_response(id); + let session_id = resp["result"]["sessionId"].as_str().unwrap().to_string(); + + // Unknown session and empty modelId are both invalid_params. + let id = h.request( + "session/set_model", + json!({ "sessionId": "ses_nope", "modelId": "x" }), + ); + let (resp, _) = h.await_response(id); + assert_eq!(resp["error"]["code"], -32602, "unknown session"); + + let id = h.request( + "session/set_model", + json!({ "sessionId": session_id, "modelId": " " }), + ); + let (resp, _) = h.await_response(id); + assert_eq!(resp["error"]["code"], -32602, "empty modelId"); + + // A real switch is accepted... + let id = h.request( + "session/set_model", + json!({ "sessionId": session_id, "modelId": "fake-model-2" }), + ); + let (resp, _) = h.await_response(id); + assert!(resp["error"].is_null(), "set_model rejected: {resp}"); + + // ...and the next turn still completes (the override is applied at prompt + // time, not mid-flight). + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "hi" }], + }), + ); + let (resp, _) = h.await_response(id); + assert!( + resp["result"]["stopReason"].is_string(), + "turn failed after set_model: {resp}" + ); + + let _ = system_rx + .recv_timeout(Duration::from_secs(30)) + .expect("provider never received a request after set_model"); +} + #[test] fn unknown_method_gets_method_not_found() { let (base_url, _rx) = spawn_fake_provider();