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

Reduce Goose embedding surface in buzz-agent

Browse Source 
Replace the monolithic Agent extension manager with a bounded direct RMCP registry while preserving lifecycle hooks, permission gating, skill loading, process-tree cleanup, and tool-result limits. Keep turn state in Buzz memory instead of a Goose sqlite session.

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Signed-off-by: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz>
This commit is contained in:
Jimmy
2026-09-01 15:57:34 +10:00
parent d0744b222b
commit ef482e886d
9 changed files with 1395 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
crates/buzz-agent/src/hooks.rs
View File

@@ -34,9 +34,7 @@
use std::sync::Arc;
use goose::agents::Agent;
use goose::session::session_manager::Session;
use rmcp::model::CallToolRequestParams;
use crate::mcp::McpRegistry;
/// Default cap on consecutive vetoes, mirroring goose's own
/// `stop_hook_block_cap` (`agent.rs:108-119`). A tool that always objects must
@@ -69,12 +67,8 @@ fn qualified(extension: &str, tool: &str) -> String {
/// no such extension, tool error, malformed result — is `None`. A broken hook
/// must never trap a turn; that is also goose's rule for its own hooks
/// ("a misbehaving hook MUST NOT block", `hooks/mod.rs:361`).
pub async fn stop_objection(
agent: &Arc<Agent>,
session: &Session,
extension: &str,
) -> Option<String> {
let text = call_hook(agent, session, extension, "_Stop").await?;
pub async fn stop_objection(mcp: &Arc<McpRegistry>, extension: &str) -> Option<String> {
let text = call_hook(mcp, extension, "_Stop").await?;
let trimmed = text.trim();
// Empty response is the documented "no objection" signal
// (`buzz-dev-mcp/src/todo.rs:100`).
@@ -85,65 +79,40 @@ pub async fn stop_objection(
///
/// buzz-agent called this so the todo list survives truncation; goose emits
/// `AgentEvent::HistoryReplaced` when it compacts, which is our trigger.
pub async fn post_compact_state(
agent: &Arc<Agent>,
session: &Session,
extension: &str,
) -> Option<String> {
let text = call_hook(agent, session, extension, "_PostCompact").await?;
pub async fn post_compact_state(mcp: &Arc<McpRegistry>, extension: &str) -> Option<String> {
let text = call_hook(mcp, extension, "_PostCompact").await?;
let trimmed = text.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
}
/// Dispatch a zero-argument hook tool and flatten its result to text.
async fn call_hook(
agent: &Arc<Agent>,
session: &Session,
extension: &str,
tool: &str,
) -> Option<String> {
let params = CallToolRequestParams::new(qualified(extension, tool))
.with_arguments(serde_json::Map::new());
let (_id, result) = agent
.dispatch_tool_call(
params,
format!("buzz_hook_{tool}"),
// Deliberately not cancellable: this runs at a turn boundary and
// must resolve either way.
None,
session,
async fn call_hook(mcp: &Arc<McpRegistry>, extension: &str, tool: &str) -> Option<String> {
let cancel = tokio_util::sync::CancellationToken::new();
let out = match mcp
.call_rmcp(
&qualified(extension, tool),
&format!("buzz_hook_{tool}"),
&serde_json::Value::Object(serde_json::Map::new()),
&cancel,
)
.await;
let call = match result {
Ok(c) => c,
Err(e) => {
tracing::debug!(tool, error = %e, "hook tool unavailable");
.await
{
Ok(out) => out,
Err(error) => {
tracing::debug!(tool, error = ?error, "hook tool unavailable");
return None;
}
};
let out = match call.result.await {
Ok(o) => o,
Err(e) => {
tracing::debug!(tool, error = ?e, "hook tool failed");
return None;
}
};
if out.is_error.unwrap_or(false) {
tracing::debug!(tool, "hook tool returned is_error");
return None;
}
let mut text = String::new();
for content in out.content.iter() {
if let Some(raw) = content.as_text() {
text.push_str(&raw.text);
}
}
Some(text)
Some(
out.content
.iter()
.filter_map(|content| content.as_text().map(|text| text.text.as_str()))
.collect(),
)
}
#[cfg(test)]

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

@@ -24,12 +24,12 @@ pub mod agent;
pub mod config;
pub mod hooks;
pub mod loop_drive;
pub mod mcp;
pub mod model;
pub mod ops;
pub mod permission;
pub mod prompt;
pub mod provider;
pub mod session_store;
pub mod skills;
pub mod steer;
pub mod tools;
@@ -50,10 +50,6 @@ use serde_json::{json, Value};
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use goose::agents::{Agent, AgentConfig, ExtensionConfig, GoosePlatform};
use goose::config::PermissionManager;
use goose::session::session_manager::SessionType;
use config::{Config, MAX_PROMPT_BYTES, MAX_SYSTEM_PROMPT_BYTES, PROTOCOL_VERSION};
use types::McpServerStdio;
use wire::{
@@ -64,7 +60,7 @@ use wire::{
/// One live ACP session, wrapping a Goose `Agent`.
struct Session {
agent: Arc<Agent>,
mcp: Arc<crate::mcp::McpRegistry>,
/// Session id, shared with goose so tool dispatch and provider request
/// attribution agree on a name. No database sits behind it: the
/// conversation lives in `history` and in the turn's own
@@ -142,7 +138,7 @@ async fn build_agent(
mcp_servers: &[McpServerStdio],
) -> Result<
(
Arc<Agent>,
Arc<crate::mcp::McpRegistry>,
String,
crate::model::SessionModel,
Option<String>,
@@ -150,226 +146,38 @@ async fn build_agent(
),
AgentError,
> {
let session_manager = crate::session_store::session_manager();
let permission_manager = PermissionManager::instance();
let agent = Agent::with_config(AgentConfig::new(
session_manager.clone(),
permission_manager,
None,
cfg.goose_mode,
// Session naming is a Goose UX affordance; the harness names sessions.
true,
GoosePlatform::GooseCli,
));
// One session row, and it is genuinely required -- I tried removing it.
//
// `Agent::add_extension_inner` (goose agent.rs:1465) resolves the
// extension's working directory by loading the session from the store, and
// fails setup outright if the row is absent: "Failed to get session ...
// Session not found". Every MCP server buzz mounts goes through that path,
// so without the row an agent has no tools at all.
//
// Nothing else needs it. The conversation stays in memory
// (`crate::turn_state`) and the provider config stays on buzz's own handle
// (`crate::model`), so a conversation of any length adds zero further
// rows. Removing this last one means an upstream change: goose would have
// to take the working directory as an argument rather than reading it back
// out of the session.
let session = session_manager
.create_session(
std::path::PathBuf::from(cwd),
"buzz-agent".to_string(),
SessionType::Acp,
cfg.goose_mode,
)
.await
.map_err(|e| AgentError::Llm(format!("session create: {e}")))?;
let session_id = session.id.clone();
// Provider/model names come from the `GOOSE_*` variables `Config::from_env`
// projected out of the `BUZZ_AGENT_*` ones; goose's registry owns base-url
// and credential resolution from there.
let provider_name = std::env::var("GOOSE_PROVIDER").unwrap_or_else(|_| "openai".to_string());
let model_name = cfg
.model
.clone()
.or_else(|| std::env::var("GOOSE_MODEL").ok())
.ok_or_else(|| AgentError::Llm("no model configured".into()))?;
let provider = build_provider(&provider_name).await?;
let model_config =
goose::model_config::model_config_from_user_config(&provider_name, &model_name)
.map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?;
// 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.
.map_err(|error| AgentError::LlmModelNotFound(error.to_string()))?;
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
// goose/crates/goose/src), and both PRs that would have wired it —
// buzz#1290 and goose#9971 — are closed unmerged. Driving the library
// directly is what makes Fizz (and `[Base]`) actually arrive.
// buzz-agent owns the loop, so it owns the prompt: goose's own
// `PromptManager` is only reachable from `Agent::reply`, which we no longer
// call. See `crate::prompt`.
let prompt = crate::prompt::SessionPrompt::new(cfg.goose_mode);
if let Some(sp) = system_prompt.or(cfg.system_prompt.as_deref()) {
if !sp.trim().is_empty() {
prompt.set_override(sp.to_string()).await;
if let Some(system_prompt) = system_prompt.or(cfg.system_prompt.as_deref()) {
if !system_prompt.trim().is_empty() {
prompt.set_override(system_prompt.to_string()).await;
}
}
for server in mcp_servers {
let envs: HashMap<String, String> = server
.env
.iter()
.map(|e| (e.name.clone(), e.value.clone()))
.collect();
let ext = ExtensionConfig::Stdio {
name: server.name.clone(),
description: String::new(),
cmd: server.command.clone(),
args: server.args.clone(),
envs: goose::agents::extension::Envs::new(envs),
env_keys: Vec::new(),
timeout: Some(300),
cwd: Some(cwd.to_string()),
bundled: None,
available_tools: Vec::new(),
};
if let Err(e) = agent.add_extension(ext, &session_id).await {
// Match buzz-agent: a failed MCP server is a hard session error,
// not a silently tool-less agent.
return Err(AgentError::Mcp(format!("{}: {e}", server.name)));
}
let mcp = Arc::new(crate::mcp::McpRegistry::spawn_all(cfg, mcp_servers, cwd).await?);
let skill_index = crate::skills::skill_index(mcp.skills());
if !skill_index.is_empty() {
prompt.add_extra("skills", skill_index).await;
}
// AGENTS.md hints are goose's job, not ours.
//
// `PromptManager::build_system_prompt` already calls `.with_hints()`
// (`prompt_manager.rs:255`), which runs the same git-root-anchored walk
// buzz-agent used to run itself. Adding our own copy as a prompt extra put
// every AGENTS.md into the system prompt *twice* -- verified against a
// scratch repo before this change, once under our `## Project Hints` and
// again under goose's `### Project Hints`.
//
// goose's loader is also a superset: `.goosehints` as well as `AGENTS.md`,
// the filename list configurable via `CONTEXT_FILE_NAMES`, gitignore-aware
// filtering, and `@file` reference expansion bounded at the git root.
// Skills are goose's too, for the same reason as hints.
//
// buzz-agent grew its own discovery, index and `load_skill` tool
// (`hints.rs` + `builtin.rs` + `builtin_client.rs`, ~1.2k lines) before
// goose had one. goose has since shipped `goose::skills` and registers a
// `SkillsClient` as a platform extension, so ours was a strictly smaller
// duplicate: same three project directories, but goose also reads the
// global ones, plugin-installed skills and its own builtins, and supports
// `args` templating we never had.
//
// Registered by handing over a *client* rather than by name, and that is
// the whole point: goose's platform factory
// (`platform_extensions/mod.rs:219-221`) calls plain `SkillsClient::new`,
// which leaves goose's own bundled skills switched on. Those are
// `goose-doc-guide` and `web-search` (`skills/builtins/`) — neither is a
// Buzz skill, and `web-search` advertises a capability Buzz does not
// provide, so every Buzz agent would offer the model a tool-shaped promise
// to shell out to `uvx ddgs`. `with_builtin_skills(false)` is only
// reachable off the constructor, so buzz builds the client itself and
// registers it with `add_client`.
//
// What that route costs: `add_client` does not consult
// `PLATFORM_EXTENSIONS`, so nothing here inherits the table's
// `unprefixed_tools: true`. It does not have to — `is_unprefixed_extension`
// (`extension_manager.rs:392-400`) keys off the *config*, and this passes
// the same `ExtensionConfig::Platform { name: "skills" }` the factory route
// does, so the lookup still hits the table entry and the model still sees a
// bare `load_skill`. Verified by driving both routes: same tool name, same
// filesystem skills, builtins gone.
//
// The prompt index is filtered to match. A skill listed in the index but
// absent from the client is a dead `load_skill` reference.
let skills: Vec<_> = goose::skills::discover_skills(Some(std::path::Path::new(cwd)))
.into_iter()
.filter(|s| s.source_type != goose::custom_requests::SourceType::BuiltinSkill)
.collect();
if !skills.is_empty() {
let ext = ExtensionConfig::Platform {
name: goose::skills::EXTENSION_NAME.to_string(),
description: "Skills".to_string(),
display_name: Some("Skills".to_string()),
bundled: Some(true),
available_tools: Vec::new(),
};
// The factory receives a context carrying the session, because
// `SkillsClient::new` reads `session.working_dir` for discovery and
// falls back to the *process* cwd without it — which for a
// desktop-spawned agent is not the nest.
let mut ctx = agent.extension_manager.get_context().clone();
ctx.extension_manager = Some(Arc::downgrade(&agent.extension_manager));
ctx.session = Some(Arc::new(session.clone()));
match goose::skills::SkillsClient::new(ctx) {
Ok(client) => {
let client = client.with_builtin_skills(false);
let info = goose::agents::mcp_client::McpClientTrait::get_info(&client).cloned();
agent
.extension_manager
.add_client(
goose::skills::EXTENSION_NAME.to_string(),
ext,
Arc::new(client),
info,
)
.await;
prompt
.add_extra("skills", crate::skills::skill_index(&skills))
.await;
}
Err(e) => {
// Unlike an MCP server, a missing skills extension is not
// fatal: the agent still has every other tool. Losing it
// silently would be worse than losing it loudly.
tracing::warn!(error = %e, "skills extension unavailable");
}
}
}
let agent = Arc::new(agent);
// Find the extension carrying `_Stop`, by asking rather than assuming the
// name — `buzz-acp` derives it from the MCP binary's file stem
// (`buzz-acp/src/lib.rs:4145-4149`), so it is not a fixed string.
//
// KNOWN DEVIATION: buzz-agent hid `_`-prefixed tools from the model
// (`agent.rs:328-336`) while still calling them itself. Goose's
// `available_tools` allowlist gates advertising *and* dispatch through the
// same cache (`extension_manager.rs:1421`, `:1698`), so hiding them would
// also make them undispatchable — which would break the veto. They stay
// visible; `hook_tool_guidance()` tells the model to leave them alone.
let mut hook_extension = None;
for tool in agent.list_tools(&session_id, None).await {
if tool.name.ends_with("___Stop") {
hook_extension = tool
.name
.strip_suffix("___Stop")
.map(|prefix| prefix.to_string());
break;
}
}
if let Some(ext) = &hook_extension {
let hook_extension = mcp.hook_extension("_Stop");
if let Some(extension) = &hook_extension {
prompt
.add_extra("buzz_hook_tools", hook_tool_guidance())
.await;
tracing::info!(extension = %ext, "lifecycle hooks available");
tracing::info!(extension, "lifecycle hooks available");
}
Ok((agent, session_id, model, hook_extension, prompt))
Ok((mcp, uuid_like(), model, hook_extension, prompt))
}
/// Keep the model's hands off the lifecycle hooks.
@@ -597,7 +405,7 @@ async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSen
.await;
}
let (agent, goose_session_id, model, hook_extension, prompt) =
let (mcp, goose_session_id, model, hook_extension, prompt) =
match build_agent(&app.cfg, &p.cwd, p.system_prompt.as_deref(), &p.mcp_servers).await {
Ok(v) => v,
Err(e) => {
@@ -631,7 +439,7 @@ async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSen
sessions.insert(
sid.clone(),
Session {
agent,
mcp,
goose_session_id,
history: Vec::new(),
working_dir: std::path::PathBuf::from(&p.cwd),
@@ -758,7 +566,7 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
// Single-flight per session, and capture the agent handle.
let (
agent,
mcp,
goose_session_id,
pending_model,
hook_extension,
@@ -793,7 +601,7 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
// 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.mcp.clone(),
s.goose_session_id.clone(),
s.pending_model.take(),
s.hook_extension.clone(),
@@ -839,7 +647,7 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
let (result, tokens, conversation) = loop_drive::run_turn(
loop_drive::TurnContext {
agent: &agent,
mcp: &mcp,
session_id: &goose_session_id,
wire_tx,
cancel: &cancel,

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

@@ -46,8 +46,8 @@ use std::sync::Arc;
use futures::StreamExt;
use tokio_util::sync::CancellationToken;
use goose::agents::Agent;
use goose::session::session_manager::Session;
use crate::mcp::McpRegistry;
use crate::turn_state::TurnSession;
use goose_provider_types::conversation::message::{Message, MessageContent, ToolRequest};
use goose_provider_types::conversation::Conversation;
@@ -63,7 +63,7 @@ use goose_agent::operation::{ConversationEffect, Emitter};
/// goose's own `GooseEffect`: the extra variants there (recipes, extension
/// data, recorded usage) exist for operations buzz does not run, so naming the
/// narrower type keeps unreachable cases out of `apply_effects` entirely.
type BuzzMachine<'a> = StateMachine<'a, Session, ConversationEffect>;
type BuzzMachine<'a> = StateMachine<'a, TurnSession, ConversationEffect>;
use crate::wire::WireSender;
@@ -103,7 +103,7 @@ enum Round {
/// Everything a round needs that does not change between rounds.
pub struct TurnContext<'a> {
pub agent: &'a Arc<Agent>,
pub mcp: &'a Arc<McpRegistry>,
pub session_id: &'a str,
pub wire_tx: &'a WireSender,
pub cancel: &'a CancellationToken,
@@ -158,11 +158,8 @@ pub async fn run_turn(
// The turn's conversation lives here, not in a database. See
// `crate::turn_state` for why goose never needed one.
let model_config = ctx
.agent
.model_config_for_session(ctx.session_id)
.await
.ok();
let (_provider, model_config, _model_id) = ctx.model.snapshot().await;
let model_config = Some(model_config);
let mut state = crate::turn_state::TurnState::new(
ctx.session_id.to_string(),
ctx.working_dir.clone(),
@@ -183,9 +180,8 @@ pub async fn run_turn(
let outcome = crate::ops::Outcome::new();
let reply_guard_tools = if ctx.require_reply {
Some(
ctx.agent
.list_tools(ctx.session_id, None)
.await
ctx.mcp
.rmcp_tools()
.into_iter()
.map(|tool| tool.name.to_string())
.collect(),
@@ -198,14 +194,14 @@ pub async fn run_turn(
outcome.clone(),
ctx.cancel.clone(),
ctx.hook_extension
.map(|extension| (Arc::clone(ctx.agent), extension.to_string())),
.map(|extension| (Arc::clone(ctx.mcp), extension.to_string())),
reply_guard_tools,
);
let start_machine = crate::ops::round_start(
ctx.steers.clone(),
crate::ops::BuzzCompactionOperation::new(
ctx.model.clone(),
Arc::clone(ctx.agent),
Arc::clone(ctx.mcp),
ctx.hook_extension.map(str::to_string),
ctx.session_id.to_string(),
),
@@ -321,7 +317,7 @@ enum Gate {
/// Called at the top of every round, and again when a round wants to end --
/// that second call is what gives `_Stop` its veto, since its operation only
/// applies to a conversation whose last message ends the turn.
async fn gate(machine: &BuzzMachine<'_>, session: &Session, emitter: &Emitter) -> Gate {
async fn gate(machine: &BuzzMachine<'_>, session: &TurnSession, emitter: &Emitter) -> Gate {
match machine.step(session, emitter).await {
Ok(Some(result)) if result.yield_to_client => Gate::Ended,
Ok(Some(result)) => Gate::Applied(result.effects),
@@ -462,8 +458,8 @@ async fn round(
}
let results = crate::tools::execute(
ctx.agent,
state.session(),
ctx.mcp,
&state.session().id,
ctx.wire_tx,
ctx.cancel,
ctx.permissions,
@@ -487,14 +483,14 @@ struct Inference {
/// arrive and accumulating usage.
async fn infer(
ctx: &TurnContext<'_>,
session: &Session,
session: &TurnSession,
conversation: &Conversation,
tokens: &mut super::agent::TurnTokens,
) -> Result<Option<Inference>, AgentError> {
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;
let tools = ctx.mcp.rmcp_tools();
let mut stream = provider
.stream(
@@ -825,7 +821,7 @@ mod tests {
"summary"
);
assert_eq!(
state.session().usage.total_tokens,
state.session().total_tokens,
None,
"a stale total would re-trigger compaction immediately"
);
@@ -834,7 +830,7 @@ mod tests {
// current occupancy, not a delta to add to the discarded conversation.
state.set_total_tokens(Some(12_000));
assert_eq!(
state.session().usage.total_tokens,
state.session().total_tokens,
Some(12_000),
"the post-compaction round must not accumulate the old 190k total"
);

1155
crates/buzz-agent/src/mcp.rs Normal file
View File

File diff suppressed because it is too large Load Diff

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

@@ -20,17 +20,17 @@
use std::sync::{Arc, Mutex};
use crate::mcp::McpRegistry;
use crate::turn_state::TurnSession;
use anyhow::Result;
use async_trait::async_trait;
use goose::agents::Agent;
use goose::conversation::Conversation;
use goose::session::Session;
use goose_agent::machine::{StateMachine, Step};
use goose_agent::operation::{
applied, assistant_turn_count, messages_since_kickoff, not_applicable, yielded,
ConversationEffect, Emitter, Operation, OperationResult,
};
use goose_provider_types::conversation::message::{Message, ToolRequest};
use goose_provider_types::conversation::Conversation;
use crate::types::StopReason;
@@ -97,14 +97,14 @@ impl BuzzMaxRoundsOperation {
}
#[async_trait]
impl Operation<Session, ConversationEffect> for BuzzMaxRoundsOperation {
impl Operation<TurnSession, ConversationEffect> for BuzzMaxRoundsOperation {
fn name(&self) -> &'static str {
"buzz_max_rounds"
}
async fn run(
&self,
_session: &Session,
_session: &TurnSession,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
@@ -181,7 +181,7 @@ fn ends_turn(messages: &[Message]) -> bool {
/// spans exactly that same turn. Three objections still end the turn; a new
/// prompt still starts from zero.
pub struct BuzzStopVetoOperation {
agent: Arc<Agent>,
mcp: Arc<McpRegistry>,
extension: String,
block_cap: u32,
}
@@ -189,9 +189,9 @@ pub struct BuzzStopVetoOperation {
impl BuzzStopVetoOperation {
/// `extension` is the MCP extension serving `_Stop`; without one there is
/// no hook to ask, and the operation never applies.
pub fn new(agent: Arc<Agent>, extension: String, block_cap: u32) -> Self {
pub fn new(mcp: Arc<McpRegistry>, extension: String, block_cap: u32) -> Self {
Self {
agent,
mcp,
extension,
block_cap,
}
@@ -203,14 +203,14 @@ impl BuzzStopVetoOperation {
const OBJECTED: &str = "objected";
#[async_trait]
impl Operation<Session, ConversationEffect> for BuzzStopVetoOperation {
impl Operation<TurnSession, ConversationEffect> for BuzzStopVetoOperation {
fn name(&self) -> &'static str {
"buzz_stop_veto"
}
async fn run(
&self,
session: &Session,
_session: &TurnSession,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
@@ -233,9 +233,7 @@ impl Operation<Session, ConversationEffect> for BuzzStopVetoOperation {
// A broken or absent hook must never trap a turn -- `stop_objection`
// maps every failure to `None`, which is "no objection".
let Some(objection) =
crate::hooks::stop_objection(&self.agent, session, &self.extension).await
else {
let Some(objection) = crate::hooks::stop_objection(&self.mcp, &self.extension).await else {
return not_applicable();
};
@@ -318,14 +316,14 @@ impl BuzzReplyGuardOperation {
}
#[async_trait]
impl Operation<Session, ConversationEffect> for BuzzReplyGuardOperation {
impl Operation<TurnSession, ConversationEffect> for BuzzReplyGuardOperation {
fn name(&self) -> &'static str {
"buzz_reply_guard"
}
async fn run(
&self,
_session: &Session,
_session: &TurnSession,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
@@ -392,14 +390,14 @@ impl BuzzSteerOperation {
}
#[async_trait]
impl Operation<Session, ConversationEffect> for BuzzSteerOperation {
impl Operation<TurnSession, ConversationEffect> for BuzzSteerOperation {
fn name(&self) -> &'static str {
"buzz_steer"
}
async fn run(
&self,
_session: &Session,
_session: &TurnSession,
_conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
@@ -428,7 +426,7 @@ impl Operation<Session, ConversationEffect> for BuzzSteerOperation {
/// mid-work and announce its own housekeeping to the humans watching.
pub struct BuzzCompactionOperation {
model: crate::model::SessionModel,
agent: Arc<Agent>,
mcp: Arc<McpRegistry>,
hook_extension: Option<String>,
session_id: String,
}
@@ -436,13 +434,13 @@ pub struct BuzzCompactionOperation {
impl BuzzCompactionOperation {
pub fn new(
model: crate::model::SessionModel,
agent: Arc<Agent>,
mcp: Arc<McpRegistry>,
hook_extension: Option<String>,
session_id: String,
) -> Self {
Self {
model,
agent,
mcp,
hook_extension,
session_id,
}
@@ -450,23 +448,37 @@ impl BuzzCompactionOperation {
}
#[async_trait]
impl Operation<Session, ConversationEffect> for BuzzCompactionOperation {
impl Operation<TurnSession, ConversationEffect> for BuzzCompactionOperation {
fn name(&self) -> &'static str {
"buzz_compaction"
}
async fn run(
&self,
session: &Session,
session: &TurnSession,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<ConversationEffect>> {
let (provider, model_config, _model_id) = self.model.snapshot().await;
// Goose's legacy threshold facade still takes its application Session.
// Keep that coupling at this single adapter until the granular context
// crate exposes thresholding over model config + token occupancy.
let legacy_session = goose::session::Session {
id: session.id.clone(),
working_dir: session.working_dir.clone(),
model_config: session.model_config.clone(),
usage: goose_provider_types::conversation::token_usage::Usage {
total_tokens: session.total_tokens,
..Default::default()
},
conversation: Some(conversation.clone()),
..Default::default()
};
if !goose::context_mgmt::check_if_compaction_needed(
provider.as_ref(),
conversation,
None,
session,
&legacy_session,
)
.await?
{
@@ -491,9 +503,7 @@ impl Operation<Session, ConversationEffect> for BuzzCompactionOperation {
let mut effects = vec![ConversationEffect::ReplaceConversation(result.conversation)];
if let Some(extension) = &self.hook_extension {
if let Some(reported) =
crate::hooks::post_compact_state(&self.agent, session, extension).await
{
if let Some(reported) = crate::hooks::post_compact_state(&self.mcp, extension).await {
// `[PostCompact]` prefix preserved from the inline version:
// it is how the model tells re-injected state apart from a
// user turn, and dropping it would change what it reads.
@@ -518,18 +528,18 @@ pub fn round_gate(
max_rounds: u32,
outcome: Outcome,
cancel: tokio_util::sync::CancellationToken,
stop_veto: Option<(Arc<Agent>, String)>,
stop_veto: Option<(Arc<McpRegistry>, String)>,
reply_guard_tools: Option<Vec<String>>,
) -> StateMachine<'static, Session, ConversationEffect> {
) -> StateMachine<'static, TurnSession, ConversationEffect> {
// Order is precedence. The round budget is checked first: once it is
// spent the turn ends, and asking `_Stop` to veto a turn we are ending
// anyway would dispatch a tool call whose answer cannot be honoured.
let mut steps: Vec<Step<Session, ConversationEffect>> = vec![Step::Operation(Arc::new(
let mut steps: Vec<Step<TurnSession, ConversationEffect>> = vec![Step::Operation(Arc::new(
BuzzMaxRoundsOperation::new(max_rounds, outcome),
))];
if let Some((agent, extension)) = stop_veto {
if let Some((mcp, extension)) = stop_veto {
steps.push(Step::Operation(Arc::new(BuzzStopVetoOperation::new(
agent,
mcp,
extension,
crate::hooks::stop_block_cap(),
))));
@@ -557,7 +567,7 @@ pub fn round_start(
steers: crate::steer::SteerQueue,
compaction: BuzzCompactionOperation,
cancel: tokio_util::sync::CancellationToken,
) -> StateMachine<'static, Session, ConversationEffect> {
) -> StateMachine<'static, TurnSession, ConversationEffect> {
// Steer before compaction: a steer that arrives just as the window fills
// should be part of what gets summarised, not appended to a conversation
// that was compacted a moment earlier without it.
@@ -578,8 +588,8 @@ mod tests {
/// A bare agent with no extensions: enough to dispatch a tool call that
/// will fail to resolve, which is the "hook unavailable" path.
async fn agent() -> Arc<Agent> {
Arc::new(Agent::new())
async fn agent() -> Arc<McpRegistry> {
Arc::new(McpRegistry::empty())
}
fn emitter() -> Emitter {
@@ -621,7 +631,7 @@ mod tests {
let outcome = Outcome::new();
let op = BuzzMaxRoundsOperation::new(3, outcome.clone());
let result = op
.run(&Session::default(), &conversation(2), &emitter())
.run(&TurnSession::default(), &conversation(2), &emitter())
.await
.unwrap();
assert!(matches!(result, OperationResult::NotApplicable));
@@ -633,7 +643,7 @@ mod tests {
let outcome = Outcome::new();
let op = BuzzMaxRoundsOperation::new(3, outcome.clone());
let result = op
.run(&Session::default(), &conversation(3), &emitter())
.run(&TurnSession::default(), &conversation(3), &emitter())
.await
.unwrap();
match result {
@@ -662,7 +672,7 @@ mod tests {
let op = BuzzStopVetoOperation::new(agent().await, "buzz-dev-mcp".to_string(), 3);
let result = op
.run(
&Session::default(),
&TurnSession::default(),
&tool_calling_conversation(),
&emitter(),
)
@@ -677,7 +687,7 @@ mod tests {
// A broken hook must not be able to hold a turn open.
let op = BuzzStopVetoOperation::new(agent().await, "no-such-extension".to_string(), 3);
let result = op
.run(&Session::default(), &conversation(1), &emitter())
.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap();
assert!(matches!(result, OperationResult::NotApplicable));
@@ -713,7 +723,7 @@ mod tests {
// At the cap the turn is allowed to end.
let result = op
.run(&Session::default(), &conversation, &emitter())
.run(&TurnSession::default(), &conversation, &emitter())
.await
.unwrap();
assert!(matches!(result, OperationResult::NotApplicable));
@@ -742,7 +752,7 @@ mod tests {
async fn the_reply_guard_reminds_when_nothing_was_published() {
let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
let result = op
.run(&Session::default(), &conversation(1), &emitter())
.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap();
match result {
@@ -764,7 +774,7 @@ mod tests {
] {
let result = op
.run(
&Session::default(),
&TurnSession::default(),
&turn_with(shell_request(command)),
&emitter(),
)
@@ -782,7 +792,7 @@ mod tests {
let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
let result = op
.run(
&Session::default(),
&TurnSession::default(),
&turn_with(shell_request("ls -la")),
&emitter(),
)
@@ -796,7 +806,7 @@ mod tests {
let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["real__shell".to_string()]);
let result = op
.run(
&Session::default(),
&TurnSession::default(),
&turn_with(shell_request("buzz messages send --channel x --content hi")),
&emitter(),
)
@@ -822,7 +832,7 @@ mod tests {
}
let result = op
.run(
&Session::default(),
&TurnSession::default(),
&Conversation::new_unvalidated(messages),
&emitter(),
)
@@ -840,7 +850,7 @@ mod tests {
let op = BuzzSteerOperation::new(steers.clone());
match op
.run(&Session::default(), &conversation(1), &emitter())
.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap()
{
@@ -854,7 +864,7 @@ mod tests {
// Drained, not peeked: a steer delivered twice would repeat the
// instruction to the model on the next round.
assert!(matches!(
op.run(&Session::default(), &conversation(1), &emitter())
op.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap(),
OperationResult::NotApplicable
@@ -868,7 +878,7 @@ mod tests {
// always-applying operation would spin it.
let op = BuzzSteerOperation::new(crate::steer::SteerQueue::new());
assert!(matches!(
op.run(&Session::default(), &conversation(1), &emitter())
op.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap(),
OperationResult::NotApplicable
@@ -881,7 +891,7 @@ mod tests {
// `_Stop` off. An operator with a misbehaving hook depends on it.
let op = BuzzStopVetoOperation::new(agent().await, "buzz-dev-mcp".to_string(), 0);
let result = op
.run(&Session::default(), &conversation(1), &emitter())
.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap();
assert!(matches!(result, OperationResult::NotApplicable));
@@ -903,7 +913,7 @@ mod tests {
// actually saying it.
let op = BuzzMaxRoundsOperation::new(1, Outcome::new());
let result = op
.run(&Session::default(), &conversation(1), &emitter())
.run(&TurnSession::default(), &conversation(1), &emitter())
.await
.unwrap();
match result {
@@ -943,7 +953,7 @@ mod tests {
let steered = Conversation::new_unvalidated(messages);
let result = op
.run(&Session::default(), &steered, &emitter())
.run(&TurnSession::default(), &steered, &emitter())
.await
.unwrap();
assert!(

107
crates/buzz-agent/src/session_store.rs
View File

@@ -1,107 +0,0 @@
//! Where buzz-agent's conversations are persisted.
//!
//! goose's [`SessionManager`] is the conversation store the loop reads and
//! writes each round. Its `instance()` singleton resolves to
//! `Paths::data_dir()` — for a normal install, the *user's own* goose data
//! directory. A buzz-agent using it would write its turns into the same
//! `sessions.db` as the person's `goose` CLI history: not a correctness bug
//! (SQLite handles concurrent writers), but wrong on two counts. Agent
//! conversations get interleaved with a human's personal sessions, and
//! buzz-agent grows an ever-larger file it does not own and cannot prune.
//!
//! So buzz-agent picks its own directory instead. `BUZZ_AGENT_SESSION_DIR`
//! (absolute) is the explicit knob; otherwise we derive a per-user default
//! under the OS data directory. Only if neither is available do we fall back
//! to goose's singleton, because a degraded shared store still beats refusing
//! to start.
//!
//! Note this is deliberately *not* `GOOSE_PATH_ROOT`: that would relocate
//! goose's config and keyring lookups too, and buzz-agent wants goose's
//! provider configuration to keep resolving exactly where goose puts it.
use std::path::PathBuf;
use std::sync::Arc;
use goose::session::session_manager::SessionManager;
/// Absolute path to a directory buzz-agent should keep its session store in.
pub const SESSION_DIR_ENV: &str = "BUZZ_AGENT_SESSION_DIR";
/// Resolve the session-store directory, or `None` to use goose's default.
///
/// Split out from [`session_manager`] so the precedence is testable without
/// touching the filesystem or building a SQLite pool.
fn resolve_dir(env_value: Option<&str>, data_dir: Option<PathBuf>) -> Option<PathBuf> {
// An explicit relative path is a misconfiguration we should not silently
// resolve against an arbitrary CWD -- the agent's CWD is the user's
// workspace, and a stray sessions.db there would be surprising.
if let Some(path) = env_value.map(PathBuf::from).filter(|p| p.is_absolute()) {
return Some(path);
}
data_dir.map(|dir| dir.join("buzz-agent"))
}
/// The session store this process should use.
///
/// Resolved once and shared: `SessionManager::new` builds its own SQLite pool,
/// so calling it per session or per round would open a new pool each time
/// against the same file. Every call site must get this one.
///
/// Falls back to goose's shared singleton only when no private directory can
/// be determined at all.
pub fn session_manager() -> Arc<SessionManager> {
static STORE: std::sync::OnceLock<Arc<SessionManager>> = std::sync::OnceLock::new();
STORE
.get_or_init(|| {
Arc::new(
match resolve_dir(
std::env::var(SESSION_DIR_ENV).ok().as_deref(),
dirs::data_dir(),
) {
Some(dir) => {
tracing::debug!(dir = %dir.display(), "session store");
SessionManager::new(dir)
}
None => {
tracing::warn!(
"no data dir available; falling back to goose's shared session store"
);
SessionManager::instance()
}
},
)
})
.clone()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn an_absolute_env_path_wins() {
let dir = resolve_dir(Some("/tmp/buzz-sessions"), Some(PathBuf::from("/data")));
assert_eq!(dir, Some(PathBuf::from("/tmp/buzz-sessions")));
}
#[test]
fn a_relative_env_path_is_ignored() {
// Resolving this against the agent's CWD would drop a sessions.db in
// whatever workspace the user happened to open.
let dir = resolve_dir(Some("relative/path"), Some(PathBuf::from("/data")));
assert_eq!(dir, Some(PathBuf::from("/data/buzz-agent")));
}
#[test]
fn the_default_is_private_to_buzz_agent() {
// The point of the whole module: never the bare goose data dir, which
// is where the user's own CLI history lives.
let dir = resolve_dir(None, Some(PathBuf::from("/data")));
assert_eq!(dir, Some(PathBuf::from("/data/buzz-agent")));
}
#[test]
fn without_a_data_dir_we_defer_to_goose() {
assert_eq!(resolve_dir(None, None), None);
}
}

46
crates/buzz-agent/src/skills.rs
View File

@@ -14,6 +14,52 @@
use goose::custom_requests::{SourceEntry, SourceType};
/// Provider-native definition for the in-process skill loader.
pub fn load_skill_tool() -> rmcp::model::Tool {
let schema = serde_json::json!({
"type": "object",
"required": ["name"],
"properties": {
"name": { "type": "string", "description": "Skill name, or skill-name/path for a supporting file." },
"args": { "type": "string", "description": "Optional arguments for a parameterized skill." }
}
});
rmcp::model::Tool::new(
"load_skill",
"Load a skill's full content into your context so you can follow its instructions.",
schema.as_object().cloned().unwrap_or_default(),
)
}
/// Execute the in-process skill loader using Goose's public skill rendering.
pub fn load_skill(
skills: &[SourceEntry],
arguments: &serde_json::Value,
) -> rmcp::model::CallToolResult {
let name = arguments
.get("name")
.and_then(serde_json::Value::as_str)
.unwrap_or("");
let args = arguments.get("args").and_then(serde_json::Value::as_str);
if let Some(skill) = skills.iter().find(|skill| skill.name == name) {
return match goose::skills::loaded_skill_context_with_args(skill, args) {
Ok(rendered) => {
rmcp::model::CallToolResult::success(vec![rmcp::model::ContentBlock::text(
rendered,
)])
}
Err(error) => {
rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
"Failed to parse skill arguments: {error}"
))])
}
};
}
rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
"Skill '{name}' not found."
))])
}
/// Render the index of available skills for the system prompt.
///
/// Empty for an empty slate, so the caller can skip the prompt extra entirely

39
crates/buzz-agent/src/tools.rs
View File

@@ -23,8 +23,7 @@ use futures::stream::{FuturesUnordered, StreamExt};
use rmcp::model::{CallToolResult, ContentBlock, ErrorData};
use tokio_util::sync::CancellationToken;
use goose::agents::Agent;
use goose::session::session_manager::Session;
use crate::mcp::McpRegistry;
use goose_provider_types::conversation::message::{ToolRequest, ToolResult};
use crate::permission::{AskSubject, PermissionBroker, PermissionDecision};
@@ -41,8 +40,8 @@ pub type Outcome = (String, ToolResult<CallToolResult>);
/// produced them.
#[allow(clippy::too_many_arguments)]
pub async fn execute(
agent: &Arc<Agent>,
session: &Session,
mcp: &Arc<McpRegistry>,
session_id: &str,
wire_tx: &WireSender,
cancel: &CancellationToken,
permissions: &Arc<PermissionBroker>,
@@ -53,8 +52,8 @@ pub async fn execute(
let mut pending = FuturesUnordered::new();
for (index, request) in requests.iter().enumerate() {
pending.push(run_one(
agent,
session,
mcp,
session_id,
wire_tx,
cancel,
permissions,
@@ -73,7 +72,7 @@ pub async fn execute(
for outcome in collected.into_iter().flatten() {
let (id, result) = outcome;
let result = reflect_on_failure(result, reflections);
emit_terminal(wire_tx, session, &id, &result).await;
emit_terminal(wire_tx, session_id, &id, &result).await;
results.push((id, result));
}
results
@@ -100,8 +99,8 @@ pub fn cancelled_results(requests: &[ToolRequest]) -> Vec<Outcome> {
#[allow(clippy::too_many_arguments)]
async fn run_one(
agent: &Arc<Agent>,
session: &Session,
mcp: &Arc<McpRegistry>,
session_id: &str,
wire_tx: &WireSender,
cancel: &CancellationToken,
permissions: &Arc<PermissionBroker>,
@@ -117,7 +116,7 @@ async fn run_one(
Err(e) => return (index, (id, Err(e.clone()))),
};
crate::agent::emit_tool_call(wire_tx, &session.id, &id, &call).await;
crate::agent::emit_tool_call(wire_tx, session_id, &id, &call).await;
// Ask the client to authorize this call. Every non-authorizing outcome
// fails closed: the model sees an ordinary tool failure and the turn
@@ -131,7 +130,7 @@ async fn run_one(
.request_permission(
wire_tx,
protocol_version,
&session.id,
session_id,
AskSubject {
tool_call_id: &id,
tool_name: &call.name,
@@ -174,15 +173,15 @@ async fn run_one(
);
}
let (_request_id, dispatched) = agent
.dispatch_tool_call(call, id.clone(), Some(cancel.clone()), session)
let arguments = call
.arguments
.clone()
.map(serde_json::Value::Object)
.unwrap_or(serde_json::Value::Null);
let result = mcp
.call_rmcp(call.name.as_ref(), &id, &arguments, cancel)
.await;
let result = match dispatched {
Ok(call_result) => call_result.result.await,
Err(e) => Err(e),
};
(index, (id, result))
}
@@ -232,7 +231,7 @@ fn reflect_on_failure(
async fn emit_terminal(
wire_tx: &WireSender,
session: &Session,
session_id: &str,
id: &str,
result: &ToolResult<CallToolResult>,
) {
@@ -240,7 +239,7 @@ async fn emit_terminal(
Ok(out) => out.is_error.unwrap_or(false),
Err(_) => true,
};
crate::agent::emit_tool_call_update(wire_tx, &session.id, id, failed).await;
crate::agent::emit_tool_call_update(wire_tx, session_id, id, failed).await;
}
#[cfg(test)]

153
crates/buzz-agent/src/turn_state.rs
View File

@@ -1,95 +1,79 @@
//! The conversation for one turn, held in memory.
//! The conversation and small amount of runtime metadata for one turn.
//!
//! buzz-agent used to round-trip every message through goose's
//! `SessionManager` — a sqlite database — once per round. That was never
//! required: none of the goose APIs buzz calls reads the store.
//!
//! * `Provider::stream` takes `&[Message]`, not a session.
//! * `Agent::dispatch_tool_call` takes `&Session` but reads only
//! `working_dir` (for the subdirectory-hint tracker) and `id` (for the
//! `PreToolUse` hook context).
//! * `context_mgmt::{check_if_compaction_needed, compact_messages}` take a
//! `&Session` / `session_id`, but the id is only carried into a tokio
//! task-local that becomes an `agent-session-id` HTTP header for provider
//! request attribution (`goose/src/session_context.rs`), and the session is
//! read for `usage.total_tokens` and `model_config` — both plain fields.
//!
//! So the turn owns a [`Session`] value directly. goose still gets the
//! `&Session` its signatures ask for; it just isn't backed by a database.
//!
//! # Why this matters beyond dependency hygiene
//!
//! A `SessionManager` per agent process meant every buzz agent grew a sqlite
//! file it never read back. Worse, before the store was isolated they all
//! shared the user's own goose CLI history. An agent turn is ephemeral — the
//! durable record of a Buzz conversation is the relay, not a local database.
//! Buzz persists conversation history on the relay. The Goose GDK loop and
//! provider APIs only need an in-memory [`Conversation`], so carrying Goose's
//! application-level `Session` (and its sqlite-oriented fields) here couples
//! the embedder to state it does not use.
use goose::conversation::Conversation;
use goose::session::session_manager::Session;
use goose_provider_types::conversation::message::Message;
use goose_provider_types::conversation::Conversation;
/// One turn's conversation and the session metadata goose needs alongside it.
/// Runtime metadata passed through Goose's generic state machine.
#[derive(Default)]
pub struct TurnSession {
/// ACP/Goose attribution id for this turn.
pub id: String,
/// Working directory used by prompt discovery and tool hints.
pub working_dir: std::path::PathBuf,
/// Model configuration used by the temporary legacy compaction adapter.
pub model_config: Option<goose_providers::model::ModelConfig>,
/// Latest provider-reported occupancy for the current conversation.
pub total_tokens: Option<i32>,
conversation: Conversation,
}
impl goose_agent::machine::MachineSession for TurnSession {
fn id(&self) -> &str {
&self.id
}
fn conversation(&self) -> Option<&Conversation> {
Some(&self.conversation)
}
}
/// One turn's in-memory conversation and runtime metadata.
pub struct TurnState {
session: Session,
session: TurnSession,
}
impl TurnState {
/// Start a turn from the session's identity and working directory.
///
/// `model_config` is carried because `check_if_compaction_needed` reads it
/// from the session to resolve the context limit; without it goose falls
/// back to a default limit and compaction fires at the wrong point.
pub fn new(
id: String,
working_dir: std::path::PathBuf,
model_config: Option<goose_providers::model::ModelConfig>,
) -> Self {
let mut session = Session {
id,
working_dir,
model_config,
..Session::default()
};
session.conversation = Some(Conversation::new_unvalidated(Vec::new()));
Self { session }
Self {
session: TurnSession {
id,
working_dir,
model_config,
total_tokens: None,
conversation: Conversation::new_unvalidated(Vec::new()),
},
}
}
/// The session value goose's APIs take.
pub fn session(&self) -> &Session {
pub fn session(&self) -> &TurnSession {
&self.session
}
/// This turn's conversation.
pub fn conversation(&self) -> Conversation {
self.session.conversation.clone().unwrap_or_default()
self.session.conversation.clone()
}
/// Append a message to the turn.
pub fn push(&mut self, message: Message) {
let mut messages = self
.session
.conversation
.take()
.map(|conversation| conversation.messages().to_vec())
.unwrap_or_default();
let mut messages = self.session.conversation.messages().to_vec();
messages.push(message);
self.session.message_count = messages.len();
self.session.conversation = Some(Conversation::new_unvalidated(messages));
self.session.conversation = Conversation::new_unvalidated(messages);
}
/// Replace the conversation wholesale, as compaction does.
pub fn replace(&mut self, conversation: Conversation) {
self.session.message_count = conversation.messages().len();
self.session.conversation = Some(conversation);
self.session.conversation = conversation;
}
/// Record the running token total.
///
/// `check_if_compaction_needed` prefers `session.usage.total_tokens` over
/// re-counting the conversation, so keeping this current is what makes
/// compaction fire on the provider's numbers rather than an estimate.
pub fn set_total_tokens(&mut self, total: Option<i32>) {
self.session.usage.total_tokens = total;
self.session.total_tokens = total;
}
}
@@ -102,66 +86,37 @@ mod tests {
}
#[test]
fn a_new_turn_starts_empty_but_has_a_conversation() {
// `Some(empty)` rather than `None`: goose's compaction check bails on
// a session with no conversation at all.
let state = state();
assert!(state.session().conversation.is_some());
assert!(state.conversation().messages().is_empty());
fn a_new_turn_starts_empty() {
assert!(state().conversation().messages().is_empty());
}
#[test]
fn pushed_messages_are_visible_to_goose() {
fn pushed_messages_are_visible() {
let mut state = state();
state.push(Message::user().with_text("hello"));
state.push(Message::assistant().with_text("hi"));
let conversation = state.conversation();
assert_eq!(conversation.messages().len(), 2);
// goose reads the conversation through the session, so the two views
// must not drift.
assert_eq!(
state
.session()
.conversation
.as_ref()
.map(|c| c.messages().len()),
Some(2)
);
assert_eq!(state.session().message_count, 2);
assert_eq!(state.conversation().messages().len(), 2);
}
#[test]
fn replace_swaps_the_whole_conversation() {
let mut state = state();
state.push(Message::user().with_text("one"));
state.push(Message::user().with_text("two"));
state.replace(Conversation::new_unvalidated(vec![
Message::user().with_text("summary")
]));
assert_eq!(state.conversation().messages().len(), 1);
assert_eq!(state.session().message_count, 1);
}
#[test]
fn total_tokens_reach_the_session_goose_reads() {
let mut state = state();
fn runtime_metadata_is_preserved() {
let mut state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work"), None);
state.set_total_tokens(Some(1234));
assert_eq!(state.session().usage.total_tokens, Some(1234));
}
#[test]
fn the_working_dir_and_id_goose_reads_are_preserved() {
// These two fields are the entirety of what `dispatch_tool_call` uses
// the session for; getting them wrong breaks subdirectory hints and
// the PreToolUse hook context.
let state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work"), None);
assert_eq!(state.session().id, "abc");
assert_eq!(
state.session().working_dir,
std::path::PathBuf::from("/work")
);
assert_eq!(state.session().total_tokens, Some(1234));
}
}