jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-28 21:11:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(buzz-agent): keep turn model attribution coherent

Browse Source 
Store the provider, model config, and exact ACP model id as one atomic
session snapshot. Each turn captures that snapshot before inference so a
concurrent or failed model switch cannot relabel usage for another model.

Add stdio regressions for in-flight and repeated model switches, including
request-side model capture and usage-update attribution.

Signed-off-by: Thinker <75d8a808fa21bb8d1812e080cf471db601c6e8e6a62dcd516d37531e83a7bb77@meshllm.communities.buzz.xyz>
This commit is contained in:
Thinker
2026-09-01 14:58:14 +10:00
parent 6ee2a3e56e
commit d0744b222b
5 changed files with 262 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
crates/buzz-agent/src/lib.rs
View File

@@ -85,9 +85,6 @@ struct Session {
/// Set for the duration of a turn; advertised to steer-capable clients.
active_run_id: Option<String>,
cancel: Option<CancellationToken>,
/// Model id currently in force for this session (`session/set_model`).
/// Reported back on `usage_update`.
model_override: Option<String>,
/// 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<App>, 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<App>, 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<App>, 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<App>, 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(())
}

3
crates/buzz-agent/src/loop_drive.rs
View File

@@ -491,8 +491,7 @@ async fn infer(
conversation: &Conversation,
tokens: &mut super::agent::TurnTokens,
) -> Result<Option<Inference>, 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;

70
crates/buzz-agent/src/model.rs
View File

@@ -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<dyn Provider>,
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<dyn Provider>, config: ModelConfig) -> Self {
pub fn new(provider: Arc<dyn Provider>, 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<dyn Provider>, 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<dyn Provider> {
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<dyn Provider>, config: ModelConfig) {
/// provider/config while reporting the previous model id (or vice versa).
pub async fn set(&self, provider: Arc<dyn Provider>, 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");
}
}

3
crates/buzz-agent/src/ops.rs
View File

@@ -461,7 +461,7 @@ impl Operation<Session, ConversationEffect> for BuzzCompactionOperation {
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
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<Session, ConversationEffect> for BuzzCompactionOperation {
return not_applicable();
}
let model_config = self.model.config().await;
let result = goose::context_mgmt::compact_messages(
provider.as_ref(),
&model_config,

215
crates/buzz-agent/tests/stdio_turn.rs
View File

@@ -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<String>) {
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<bool>, 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<ProviderRequest>, 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<String>) {
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::<Value>(&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<String>) {
.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<String>) {
"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<String>) {
}
});
(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<Value>) {
let mut notifications = Vec::new();
self.await_response_collecting(id, Vec::new())
}
fn await_response_collecting(
&mut self,
id: i64,
mut notifications: Vec<Value>,
) -> (Value, Vec<Value>) {
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::<Value>(&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();