diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 8e290a60d7..e4e9b3b971 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -85,9 +85,6 @@ 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 @@ -209,7 +206,7 @@ async fn build_agent( // Held here rather than pushed into the `Agent` with `update_provider`: // that call writes the provider config to goose's sqlite session row, and // buzz reads the provider and config back through `SessionModel` anyway. - let model = crate::model::SessionModel::new(provider, model_config); + let model = crate::model::SessionModel::new(provider, model_config, model_name); // Persona. This is the seam that does NOT work over plain ACP: goose's own // ACP server never reads `systemPrompt` (rg finds zero hits in @@ -642,7 +639,6 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen busy: false, active_run_id: None, cancel: None, - model_override: None, pending_model: None, hook_extension, prompt, @@ -824,6 +820,11 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire } } + // Snapshot the exact model identity this turn will use after applying any + // pending switch. A later `session/set_model` may queue the following turn, + // but must not relabel usage from this one. + let (_provider, _config, turn_model_id) = model.snapshot().await; + // Advertise the run id so steer-capable clients can target this turn. // `_meta` nests INSIDE `update` — see the module docs. wire::send( @@ -922,25 +923,13 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire if let Some((acc_in, acc_out, acc_cached, acc_written, acc_total, acc_identity)) = accumulated { - let model = { - let sessions = app.sessions.lock().await; - sessions - .get(&p.session_id) - .and_then(|s| s.model_override.clone()) - .or_else(|| app.cfg.model.clone()) - // A session starts fine with only GOOSE_MODEL set - // (build_agent falls back to it), so omitting it here - // emitted `"model": ""` and blanked kind-44200 attribution. - .or_else(|| std::env::var("GOOSE_MODEL").ok()) - .unwrap_or_default() - }; let update = wire::usage_update_payload( Some(acc_in), Some(acc_out), acc_cached.exact_value(), acc_written.exact_value(), acc_total, - &model, + &turn_model_id, acc_identity.as_ref().and_then(|identity| identity.as_ref()), ); wire::send(wire_tx, goose_session_update(&p.session_id, update)).await; @@ -980,7 +969,6 @@ 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.clone()); s.pending_model = Some(p.model_id); drop(sessions); wire::send(wire_tx, wire::ok(id, Value::Null)).await; @@ -1117,7 +1105,9 @@ async fn apply_model(model: &crate::model::SessionModel, model_id: &str) -> Resu .map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?; // Swapped on buzz's own handle. `Agent::update_provider` would do the same // thing plus a write to goose's session row; see [`crate::model`]. - model.set(provider, model_config).await; + model + .set(provider, model_config, model_id.to_string()) + .await; Ok(()) } diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs index 67fe1ae35b..b1524f3371 100644 --- a/crates/buzz-agent/src/loop_drive.rs +++ b/crates/buzz-agent/src/loop_drive.rs @@ -491,8 +491,7 @@ async fn infer( conversation: &Conversation, tokens: &mut super::agent::TurnTokens, ) -> Result, AgentError> { - let provider = ctx.model.provider().await; - let model_config = ctx.model.config().await; + let (provider, model_config, _model_id) = ctx.model.snapshot().await; let system_prompt = ctx.prompt.build(&session.working_dir).await; let tools = ctx.agent.list_tools(ctx.session_id, None).await; diff --git a/crates/buzz-agent/src/model.rs b/crates/buzz-agent/src/model.rs index 0d2dce7336..15792edf6e 100644 --- a/crates/buzz-agent/src/model.rs +++ b/crates/buzz-agent/src/model.rs @@ -13,6 +13,8 @@ //! Two things still need goose's `Agent`: `list_tools` and //! `dispatch_tool_call`. Neither touches the provider (`dispatch_tool_call` //! contains no reference to it) or the session store. +//! The model's exact ACP id is held under the same lock as its provider and +//! config so every inference and its cost attribution describe one state. use std::sync::Arc; @@ -29,33 +31,48 @@ pub struct SessionModel { struct Inner { provider: Arc, config: ModelConfig, + /// Exact ACP model id installed with this provider/config pair. This stays + /// under the same lock so inference and attribution cannot split. + model_id: String, } impl SessionModel { - pub fn new(provider: Arc, config: ModelConfig) -> Self { + pub fn new(provider: Arc, config: ModelConfig, model_id: String) -> Self { Self { - inner: Arc::new(RwLock::new(Inner { provider, config })), + inner: Arc::new(RwLock::new(Inner { + provider, + config, + model_id, + })), } } - /// The provider to stream from. + /// Snapshot the provider and config for one operation, along with the + /// exact model id that owns them. + pub async fn snapshot(&self) -> (Arc, ModelConfig, String) { + let inner = self.inner.read().await; + ( + Arc::clone(&inner.provider), + inner.config.clone(), + inner.model_id.clone(), + ) + } + + /// Provider handle for model discovery. Inference and compaction use + /// [`Self::snapshot`] so their provider and config are one generation. pub async fn provider(&self) -> Arc { Arc::clone(&self.inner.read().await.provider) } - /// The model config, which also carries the context limit compaction needs. - pub async fn config(&self) -> ModelConfig { - self.inner.read().await.config.clone() - } - - /// Swap both, as `session/set_model` does. + /// Swap all three, as `session/set_model` does. /// /// Taken together under one write so a turn can never observe a new - /// provider with the previous model's config. - pub async fn set(&self, provider: Arc, config: ModelConfig) { + /// provider/config while reporting the previous model id (or vice versa). + pub async fn set(&self, provider: Arc, config: ModelConfig, model_id: String) { let mut inner = self.inner.write().await; inner.provider = provider; inner.config = config; + inner.model_id = model_id; } } @@ -92,10 +109,12 @@ mod tests { #[tokio::test] async fn the_configured_provider_and_config_are_what_come_back() { let (provider, config) = model("first"); - let session = SessionModel::new(provider, config); + let session = SessionModel::new(provider, config, "first".to_string()); - assert_eq!(session.provider().await.get_name(), "first"); - assert_eq!(session.config().await.model_name, "first"); + let (provider, config, model_id) = session.snapshot().await; + assert_eq!(provider.get_name(), "first"); + assert_eq!(config.model_name, "first"); + assert_eq!(model_id, "first"); } #[tokio::test] @@ -103,13 +122,17 @@ mod tests { // The pair must move as one: a turn seeing a new provider with the old // config would stream against the wrong context limit. let (provider, config) = model("first"); - let session = SessionModel::new(provider, config); + let session = SessionModel::new(provider, config, "first".to_string()); let (next_provider, next_config) = model("second"); - session.set(next_provider, next_config).await; + session + .set(next_provider, next_config, "second".to_string()) + .await; - assert_eq!(session.provider().await.get_name(), "second"); - assert_eq!(session.config().await.model_name, "second"); + let (provider, config, model_id) = session.snapshot().await; + assert_eq!(provider.get_name(), "second"); + assert_eq!(config.model_name, "second"); + assert_eq!(model_id, "second"); } #[tokio::test] @@ -117,12 +140,17 @@ mod tests { // `SessionModel` is cloned into the turn; a `session/set_model` applied // through one handle has to be visible through the other. let (provider, config) = model("first"); - let session = SessionModel::new(provider, config); + let session = SessionModel::new(provider, config, "first".to_string()); let clone = session.clone(); let (next_provider, next_config) = model("second"); - session.set(next_provider, next_config).await; + session + .set(next_provider, next_config, "second".to_string()) + .await; - assert_eq!(clone.provider().await.get_name(), "second"); + let (provider, config, model_id) = clone.snapshot().await; + assert_eq!(provider.get_name(), "second"); + assert_eq!(config.model_name, "second"); + assert_eq!(model_id, "second"); } } diff --git a/crates/buzz-agent/src/ops.rs b/crates/buzz-agent/src/ops.rs index 5db49085e8..7dfcd1a764 100644 --- a/crates/buzz-agent/src/ops.rs +++ b/crates/buzz-agent/src/ops.rs @@ -461,7 +461,7 @@ impl Operation for BuzzCompactionOperation { conversation: &Conversation, _emit: &Emitter, ) -> Result> { - let provider = self.model.provider().await; + let (provider, model_config, _model_id) = self.model.snapshot().await; if !goose::context_mgmt::check_if_compaction_needed( provider.as_ref(), conversation, @@ -473,7 +473,6 @@ impl Operation for BuzzCompactionOperation { return not_applicable(); } - let model_config = self.model.config().await; let result = goose::context_mgmt::compact_messages( provider.as_ref(), &model_config, diff --git a/crates/buzz-agent/tests/stdio_turn.rs b/crates/buzz-agent/tests/stdio_turn.rs index 8342d7f2fc..dfbe555019 100644 --- a/crates/buzz-agent/tests/stdio_turn.rs +++ b/crates/buzz-agent/tests/stdio_turn.rs @@ -18,7 +18,7 @@ use std::io::{BufRead, BufReader, Write}; use std::net::TcpListener; use std::process::{Child, Command, Stdio}; -use std::sync::mpsc; +use std::sync::{mpsc, Arc, Condvar, Mutex}; use std::time::Duration; use serde_json::{json, Value}; @@ -32,9 +32,43 @@ const FIZZ_PROMPT: &str = "You are Fizz, an energetic maker who turns ideas into /// Returns the captured `system` prompt of the first request over a channel so /// the test can assert the persona survived the trip. fn spawn_fake_provider() -> (String, mpsc::Receiver) { + let (base_url, request_rx, _gate) = spawn_controllable_provider(false); + let (tx, rx) = mpsc::channel(); + std::thread::spawn(move || { + while let Ok(request) = request_rx.recv() { + let _ = tx.send(request.system); + } + }); + (base_url, rx) +} + +#[derive(Debug)] +struct ProviderRequest { + model: String, + system: String, +} + +#[derive(Clone)] +struct ProviderGate(Arc<(Mutex, Condvar)>); + +impl ProviderGate { + fn release(&self) { + let (released, cv) = &*self.0; + *released.lock().expect("provider gate lock") = true; + cv.notify_all(); + } +} + +/// OpenAI-compatible provider whose chat response can be held after the +/// request is captured. This makes in-flight model-switch races deterministic. +fn spawn_controllable_provider( + block_chat_response: bool, +) -> (String, mpsc::Receiver, ProviderGate) { let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); let addr = listener.local_addr().expect("addr"); let (tx, rx) = mpsc::channel(); + let gate = ProviderGate(Arc::new((Mutex::new(!block_chat_response), Condvar::new()))); + let server_gate = gate.clone(); std::thread::spawn(move || { for stream in listener.incoming() { @@ -86,9 +120,13 @@ fn spawn_fake_provider() -> (String, mpsc::Receiver) { continue; } - // Chat completions: capture the system prompt, then answer. + // Chat completions: capture the model and system prompt, then + // optionally hold the response until the test releases it. + let mut request_model = String::new(); + let mut system = String::new(); if let Ok(req) = serde_json::from_slice::(&body) { - let system = req["messages"] + request_model = req["model"].as_str().unwrap_or_default().to_string(); + system = req["messages"] .as_array() .and_then(|msgs| { msgs.iter() @@ -97,8 +135,18 @@ fn spawn_fake_provider() -> (String, mpsc::Receiver) { .map(str::to_owned) }) .unwrap_or_default(); - let _ = tx.send(system); } + let _ = tx.send(ProviderRequest { + model: request_model.clone(), + system, + }); + + let (released, cv) = &*server_gate.0; + let mut released = released.lock().expect("provider gate lock"); + while !*released { + released = cv.wait(released).expect("provider gate wait"); + } + drop(released); // Goose sets `stream: true` on the OpenAI provider // (goose-providers/src/openai.rs:650), so answer with SSE. A @@ -110,7 +158,7 @@ fn spawn_fake_provider() -> (String, mpsc::Receiver) { "id": "chatcmpl-test", "object": "chat.completion.chunk", "created": 1, - "model": "fake-model", + "model": request_model, "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }], }); if !usage.is_null() { @@ -144,7 +192,7 @@ fn spawn_fake_provider() -> (String, mpsc::Receiver) { } }); - (format!("http://{addr}"), rx) + (format!("http://{addr}"), rx, gate) } struct Harness { @@ -198,11 +246,30 @@ impl Harness { /// Read lines until the response with `id` arrives, collecting every /// notification seen along the way (so ordering can be asserted). fn await_response(&mut self, id: i64) -> (Value, Vec) { - let mut notifications = Vec::new(); + self.await_response_collecting(id, Vec::new()) + } + + fn await_response_collecting( + &mut self, + id: i64, + mut notifications: Vec, + ) -> (Value, Vec) { + loop { + let msg = self.read_message(); + if msg.get("id").and_then(Value::as_i64) == Some(id) { + return (msg, notifications); + } + if msg.get("method").is_some() { + notifications.push(msg); + } + } + } + + fn read_message(&mut self) -> Value { loop { let mut line = String::new(); let n = self.stdout.read_line(&mut line).expect("read"); - assert_ne!(n, 0, "agent closed stdout before responding to id={id}"); + assert_ne!(n, 0, "agent closed stdout before the next message"); let Ok(msg) = serde_json::from_str::(&line) else { continue; }; @@ -215,12 +282,7 @@ impl Harness { self.stdin.flush().expect("flush approval"); continue; } - if msg.get("id").and_then(Value::as_i64) == Some(id) { - return (msg, notifications); - } - if msg.get("method").is_some() { - notifications.push(msg); - } + return msg; } } } @@ -431,6 +493,131 @@ fn set_model_applies_from_the_next_prompt() { .expect("provider never received a request after set_model"); } +#[test] +fn in_flight_set_model_does_not_relabel_the_active_turn() { + let (base_url, request_rx, gate) = spawn_controllable_provider(true); + 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(); + + let prompt_id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "hold this turn" }], + }), + ); + let request = request_rx + .recv_timeout(Duration::from_secs(30)) + .expect("provider never received the active turn"); + assert_eq!(request.model, "fake-model"); + + let set_model_id = h.request( + "session/set_model", + json!({ "sessionId": session_id, "modelId": "fake-model-2" }), + ); + let mut before_prompt_response = Vec::new(); + loop { + let msg = h.read_message(); + if msg.get("id").and_then(Value::as_i64) == Some(set_model_id) { + assert!(msg["error"].is_null(), "set_model rejected: {msg}"); + break; + } + if msg.get("method").is_some() { + before_prompt_response.push(msg); + } + } + + gate.release(); + let (resp, notifications) = h.await_response_collecting(prompt_id, before_prompt_response); + assert!( + resp["result"]["stopReason"].is_string(), + "active turn failed after model switch: {resp}" + ); + let usage = notifications + .iter() + .find(|n| { + n["method"] == "_goose/unstable/session/update" + && n["params"]["update"]["sessionUpdate"] == "usage_update" + }) + .expect("active turn emitted no usage update"); + assert_eq!( + usage["params"]["update"]["model"], "fake-model", + "a switch during inference must apply to the next turn, not relabel the old model's usage" + ); + + let next_prompt_id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "next turn" }], + }), + ); + let next = request_rx + .recv_timeout(Duration::from_secs(30)) + .expect("provider never received the next turn"); + assert_eq!(next.model, "fake-model-2"); + let (resp, _) = h.await_response(next_prompt_id); + assert!(resp["result"]["stopReason"].is_string()); +} + +#[test] +fn repeated_model_switches_report_the_model_each_turn_actually_used() { + let (base_url, request_rx, _gate) = spawn_controllable_provider(false); + 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(); + + for model in ["fake-model-2", "fake-model-3"] { + let id = h.request( + "session/set_model", + json!({ "sessionId": session_id, "modelId": model }), + ); + let (resp, _) = h.await_response(id); + assert!(resp["error"].is_null(), "set_model rejected: {resp}"); + + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": format!("use {model}") }], + }), + ); + let request = request_rx + .recv_timeout(Duration::from_secs(30)) + .expect("provider never received switched turn"); + assert_eq!(request.model, model); + let (resp, notifications) = h.await_response(id); + assert!(resp["result"]["stopReason"].is_string()); + let usage = notifications + .iter() + .find(|n| { + n["method"] == "_goose/unstable/session/update" + && n["params"]["update"]["sessionUpdate"] == "usage_update" + }) + .expect("switched turn emitted no usage update"); + assert_eq!(usage["params"]["update"]["model"], model); + } +} + #[test] fn unknown_method_gets_method_not_found() { let (base_url, _rx) = spawn_fake_provider();