mirror of
https://github.com/block/buzz.git
synced 2026-10-28 13:01:10 +00:00
Signed-off-by: Wes <wesbillman@users.noreply.github.com> Co-authored-by: Pinky <44b8e82baa6e0e254e0208d68f335c283c94e7b78dd1fa10d5a49d3f13dd0435@sprout-oss.stage.blox.sqprod.co>
5142 lines
208 KiB
Rust
5142 lines
208 KiB
Rust
#![deny(unsafe_code)]
|
||
|
||
mod acp;
|
||
mod config;
|
||
mod engram_fetch;
|
||
mod filter;
|
||
mod observer;
|
||
mod pool;
|
||
mod queue;
|
||
mod relay;
|
||
mod setup_mode;
|
||
mod usage;
|
||
|
||
pub use usage::TurnUsage;
|
||
|
||
use std::collections::{HashMap, HashSet};
|
||
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
|
||
use acp::{AcpClient, EnvVar, McpServer};
|
||
use anyhow::Result;
|
||
use buzz_core::kind::{
|
||
KIND_MEMBER_ADDED_NOTIFICATION, KIND_MEMBER_REMOVED_NOTIFICATION, KIND_STREAM_MESSAGE,
|
||
KIND_STREAM_REMINDER, KIND_WORKFLOW_APPROVAL_REQUESTED,
|
||
};
|
||
use buzz_core::observer::{
|
||
decrypt_observer_payload, encrypt_observer_payload, OBSERVER_FRAME_TELEMETRY,
|
||
OBSERVER_MAX_PLAINTEXT_LEN,
|
||
};
|
||
use clap::Parser;
|
||
use config::{
|
||
AuthAgentArgs, AuthMethodsArgs, AuthenticateArgs, Config, DedupMode, ModelsArgs,
|
||
MultipleEventHandling, RespondTo, SubscribeMode,
|
||
};
|
||
use filter::SubscriptionRule;
|
||
use futures_util::FutureExt;
|
||
use nostr::{PublicKey, ToBech32};
|
||
use pool::{
|
||
AgentPool, ControlSignal, IdleSwitchResult, OwnedAgent, PromptContext, PromptOutcome,
|
||
PromptResult, PromptSource, SessionState, TimeoutKind,
|
||
};
|
||
use queue::{CancelReason, EventQueue, FlushBatch, QueuedEvent, ThreadTags};
|
||
use relay::{HarnessRelay, RelayEventPublisher};
|
||
use tokio::sync::{mpsc, watch};
|
||
use tracing_subscriber::EnvFilter;
|
||
use uuid::Uuid;
|
||
|
||
/// Check if argv[1] matches a subcommand name, before any clap parsing.
|
||
///
|
||
/// This avoids clap rejecting harness flags (like `--private-key`) that aren't
|
||
/// declared on the subcommand's `Parser`. The `models` path has its own
|
||
/// dedicated parser; the default path uses the existing `CliArgs`.
|
||
///
|
||
/// **Constraint**: subcommand must be argv[1] — flags before the subcommand
|
||
/// name (e.g., `buzz-acp --verbose models`) are not supported.
|
||
fn is_subcommand(name: &str) -> bool {
|
||
std::env::args().nth(1).map(|a| a == name).unwrap_or(false)
|
||
}
|
||
|
||
/// Timeout for lightweight helper subcommands (spawn + initialize + model/method probes).
|
||
const MODELS_TIMEOUT: Duration = Duration::from_secs(10);
|
||
|
||
/// Timeout for `buzz-acp authenticate`. Browser-based vendor auth can require
|
||
/// human interaction, so it must not share the short probe timeout.
|
||
const AUTHENTICATE_TIMEOUT: Duration = Duration::from_secs(10 * 60);
|
||
|
||
/// Publish a kind:20001 presence update event via the WebSocket connection.
|
||
///
|
||
/// Ephemeral kinds (20000-29999) are rejected by the HTTP bridge, so presence
|
||
/// updates must be routed through the WS path.
|
||
///
|
||
/// Content is a bare status string (`"online"`, `"away"`, `"offline"`) matching
|
||
/// the desktop client's format. The relay stores this in Redis and synthesizes
|
||
/// it back on presence queries.
|
||
async fn publish_presence(
|
||
publisher: &relay::RelayEventPublisher,
|
||
keys: &nostr::Keys,
|
||
status: &str,
|
||
) -> Result<(), relay::RelayError> {
|
||
use buzz_core::kind::KIND_PRESENCE_UPDATE;
|
||
use nostr::{EventBuilder, Kind};
|
||
|
||
let event = EventBuilder::new(Kind::Custom(KIND_PRESENCE_UPDATE as u16), status)
|
||
.tags([])
|
||
.sign_with_keys(keys)
|
||
.map_err(|e| relay::RelayError::Http(format!("presence sign error: {e}")))?;
|
||
publisher.publish_event(event).await?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Resolve the agent's owner pubkey at startup.
|
||
///
|
||
/// Priority:
|
||
/// 1. `BUZZ_AUTH_TAG` env var — NIP-OA attestation signed by the owner.
|
||
/// Verified against the agent's own pubkey to extract the owner pubkey.
|
||
/// 2. `--agent-owner` CLI flag / `BUZZ_ACP_AGENT_OWNER` env var.
|
||
fn resolve_agent_owner(config: &Config) -> Option<String> {
|
||
// Try BUZZ_AUTH_TAG first (NIP-OA attestation).
|
||
if let Ok(auth_tag) = std::env::var("BUZZ_AUTH_TAG") {
|
||
if !auth_tag.is_empty() {
|
||
let agent_pk = config.keys.public_key();
|
||
match buzz_sdk::nip_oa::verify_auth_tag(&auth_tag, &agent_pk) {
|
||
Ok(owner_pk) => {
|
||
let owner_hex = owner_pk.to_hex().to_ascii_lowercase();
|
||
tracing::info!("owner resolved from BUZZ_AUTH_TAG: {owner_hex}");
|
||
return Some(owner_hex);
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("BUZZ_AUTH_TAG verification failed: {e} — falling back");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Fall back to --agent-owner config.
|
||
config.agent_owner.clone()
|
||
}
|
||
|
||
/// Cache for the agent's owner pubkey.
|
||
///
|
||
/// Owner is now provided via `--agent-owner` config flag (no REST lookup).
|
||
/// Cache for the agent's owner pubkey + sibling lookups.
|
||
///
|
||
/// Siblings are other agents whose NIP-OA auth tag proves the same owner.
|
||
/// Lookup results are cached for the process lifetime (attestations are immutable).
|
||
struct OwnerCache {
|
||
pubkey: Option<String>,
|
||
/// author_hex → is_sibling (true = same owner, false = not)
|
||
siblings: std::sync::Mutex<HashMap<String, bool>>,
|
||
}
|
||
|
||
impl OwnerCache {
|
||
fn new(initial: Option<String>) -> Self {
|
||
Self {
|
||
pubkey: initial,
|
||
siblings: std::sync::Mutex::new(HashMap::new()),
|
||
}
|
||
}
|
||
|
||
/// Return the cached owner pubkey.
|
||
fn get(&self) -> Option<&str> {
|
||
self.pubkey.as_deref()
|
||
}
|
||
|
||
/// Check if author is a known sibling (cached result).
|
||
fn is_known_sibling(&self, author: &str) -> Option<bool> {
|
||
self.siblings.lock().ok()?.get(author).copied()
|
||
}
|
||
|
||
/// Cache a sibling lookup result.
|
||
fn cache_sibling(&self, author: String, is_sibling: bool) {
|
||
if let Ok(mut map) = self.siblings.lock() {
|
||
// Cap at 256 entries to prevent unbounded growth.
|
||
if map.len() >= 256 {
|
||
map.clear();
|
||
}
|
||
map.insert(author, is_sibling);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Check if `author` is the owner OR a sibling (same owner via NIP-OA).
|
||
///
|
||
/// For unknown authors, queries their kind:0 profile to extract the NIP-OA
|
||
/// auth tag and verify the owner matches. Result is cached.
|
||
async fn is_owner_or_sibling(
|
||
author: &str,
|
||
owner_cache: &OwnerCache,
|
||
rest_client: &relay::RestClient,
|
||
) -> bool {
|
||
let my_owner = match owner_cache.get() {
|
||
Some(o) => o,
|
||
None => return false, // no owner configured — fail closed
|
||
};
|
||
|
||
// Direct owner check.
|
||
if author == my_owner {
|
||
return true;
|
||
}
|
||
|
||
// Check sibling cache.
|
||
if let Some(cached) = owner_cache.is_known_sibling(author) {
|
||
return cached;
|
||
}
|
||
|
||
// Query the author's kind:0 profile to check for NIP-OA auth tag.
|
||
let is_sibling = check_sibling_via_profile(author, my_owner, rest_client).await;
|
||
owner_cache.cache_sibling(author.to_string(), is_sibling);
|
||
is_sibling
|
||
}
|
||
|
||
/// Inbound author gate decision: does this author's event fire a turn?
|
||
///
|
||
/// Coarse security policy applied before subscription rules. Both `OwnerOnly`
|
||
/// and `Allowlist` accept the owner and same-owner siblings; `Allowlist`
|
||
/// additionally accepts the explicit external pubkey list.
|
||
async fn author_allowed(
|
||
respond_to: &RespondTo,
|
||
allowlist: &HashSet<String>,
|
||
author: &str,
|
||
owner_cache: &OwnerCache,
|
||
rest_client: &relay::RestClient,
|
||
) -> bool {
|
||
match respond_to {
|
||
RespondTo::Anyone => true,
|
||
RespondTo::Nobody => false,
|
||
RespondTo::OwnerOnly => is_owner_or_sibling(author, owner_cache, rest_client).await,
|
||
RespondTo::Allowlist => {
|
||
allowlist.contains(author)
|
||
|| is_owner_or_sibling(author, owner_cache, rest_client).await
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Query an author's kind:0 profile and check if their NIP-OA auth tag
|
||
/// proves the same owner as us.
|
||
async fn check_sibling_via_profile(
|
||
author: &str,
|
||
expected_owner: &str,
|
||
rest_client: &relay::RestClient,
|
||
) -> bool {
|
||
let filter = nostr::Filter::new()
|
||
.kind(nostr::Kind::Metadata)
|
||
.author(match nostr::PublicKey::from_hex(author) {
|
||
Ok(pk) => pk,
|
||
Err(_) => return false,
|
||
})
|
||
.limit(1);
|
||
|
||
let resp = match tokio::time::timeout(Duration::from_millis(2000), rest_client.query(&[filter]))
|
||
.await
|
||
{
|
||
Ok(Ok(v)) => v,
|
||
_ => return false, // timeout or error — fail closed
|
||
};
|
||
|
||
// Look for an "auth" tag in the profile event.
|
||
let events = match resp.as_array() {
|
||
Some(arr) => arr,
|
||
None => return false,
|
||
};
|
||
let event = match events.first() {
|
||
Some(e) => e,
|
||
None => return false,
|
||
};
|
||
let tags = match event.get("tags").and_then(|t| t.as_array()) {
|
||
Some(t) => t,
|
||
None => return false,
|
||
};
|
||
|
||
// Find ["auth", owner_pk, conditions, sig] and verify the Schnorr signature.
|
||
// Don't trust the relay — verify ourselves.
|
||
let agent_pk = match nostr::PublicKey::from_hex(author) {
|
||
Ok(pk) => pk,
|
||
Err(_) => return false,
|
||
};
|
||
|
||
for tag in tags {
|
||
let parts = match tag.as_array() {
|
||
Some(p) if p.len() >= 4 => p,
|
||
_ => continue,
|
||
};
|
||
if parts[0].as_str() != Some("auth") {
|
||
continue;
|
||
}
|
||
let tag_owner = match parts[1].as_str() {
|
||
Some(o) => o,
|
||
None => continue,
|
||
};
|
||
// Only verify if the owner field matches ours.
|
||
if !tag_owner.eq_ignore_ascii_case(expected_owner) {
|
||
continue;
|
||
}
|
||
// Cryptographically verify the NIP-OA attestation signature.
|
||
let tag_json = serde_json::to_string(tag).unwrap_or_default();
|
||
match buzz_sdk::nip_oa::verify_auth_tag(&tag_json, &agent_pk) {
|
||
Ok(_) => {
|
||
tracing::debug!(author, expected_owner, "sibling verified via NIP-OA");
|
||
return true;
|
||
}
|
||
Err(e) => {
|
||
tracing::debug!(author, "NIP-OA auth tag verification failed: {e}");
|
||
}
|
||
}
|
||
}
|
||
|
||
false
|
||
}
|
||
|
||
fn spawn_relay_observer_publisher(
|
||
observer: observer::ObserverHandle,
|
||
publisher: RelayEventPublisher,
|
||
keys: nostr::Keys,
|
||
agent_pubkey_hex: String,
|
||
owner_pubkey_hex: String,
|
||
owner_pubkey: PublicKey,
|
||
) -> tokio::task::JoinHandle<()> {
|
||
tokio::spawn(async move {
|
||
let mut coalescer = ObserverChunkCoalescer::default();
|
||
for event in observer.snapshot() {
|
||
for event in coalescer.ingest(event) {
|
||
publish_relay_observer_event(
|
||
&publisher,
|
||
&keys,
|
||
&agent_pubkey_hex,
|
||
&owner_pubkey_hex,
|
||
&owner_pubkey,
|
||
event,
|
||
)
|
||
.await;
|
||
}
|
||
}
|
||
|
||
let mut rx = observer.subscribe();
|
||
let mut flush_interval = tokio::time::interval(std::time::Duration::from_millis(500));
|
||
flush_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||
loop {
|
||
tokio::select! {
|
||
result = rx.recv() => {
|
||
match result {
|
||
Ok(event) => {
|
||
for event in coalescer.ingest(event) {
|
||
publish_relay_observer_event(
|
||
&publisher, &keys, &agent_pubkey_hex,
|
||
&owner_pubkey_hex, &owner_pubkey, event,
|
||
).await;
|
||
}
|
||
}
|
||
Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => {
|
||
for event in coalescer.flush() {
|
||
publish_relay_observer_event(
|
||
&publisher, &keys, &agent_pubkey_hex,
|
||
&owner_pubkey_hex, &owner_pubkey, event,
|
||
).await;
|
||
}
|
||
tracing::warn!(dropped = count, "relay observer publisher lagged");
|
||
}
|
||
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
|
||
for event in coalescer.flush() {
|
||
publish_relay_observer_event(
|
||
&publisher, &keys, &agent_pubkey_hex,
|
||
&owner_pubkey_hex, &owner_pubkey, event,
|
||
).await;
|
||
}
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
_ = flush_interval.tick() => {
|
||
// Periodic flush ensures live streaming even during continuous chunk delivery.
|
||
for event in coalescer.flush() {
|
||
publish_relay_observer_event(
|
||
&publisher, &keys, &agent_pubkey_hex,
|
||
&owner_pubkey_hex, &owner_pubkey, event,
|
||
).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
})
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct ObserverChunkCoalescer {
|
||
pending: Vec<PendingObserverChunk>,
|
||
}
|
||
|
||
struct PendingObserverChunk {
|
||
key: ObserverChunkKey,
|
||
event: observer::ObserverEvent,
|
||
text: String,
|
||
}
|
||
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
struct ObserverChunkKey {
|
||
update_type: String,
|
||
message_id: Option<String>,
|
||
channel_id: Option<String>,
|
||
session_id: Option<String>,
|
||
turn_id: Option<String>,
|
||
agent_index: Option<usize>,
|
||
}
|
||
|
||
/// Flush coalesced chunks before they exceed the NIP-44 plaintext limit (65,535 bytes).
|
||
/// Leave headroom for the JSON envelope wrapping the text. This is a SOFT pre-flush
|
||
/// of raw text below the hard cap; `fit_observer_event_to_budget` (the final ceiling,
|
||
/// keyed to `OBSERVER_MAX_PLAINTEXT_LEN` in buzz-core/observer.rs:25) is what actually
|
||
/// guarantees the serialized frame fits. Edit one of these two and review the other.
|
||
const OBSERVER_CHUNK_MAX_TEXT_BYTES: usize = 60_000;
|
||
|
||
impl ObserverChunkCoalescer {
|
||
fn ingest(&mut self, event: observer::ObserverEvent) -> Vec<observer::ObserverEvent> {
|
||
let Some((key, text)) = observer_chunk_key_and_text(&event) else {
|
||
let mut events = self.flush();
|
||
events.push(event);
|
||
return events;
|
||
};
|
||
|
||
if let Some(pending) = self.pending.iter_mut().find(|pending| pending.key == key) {
|
||
// Flush before appending if this would exceed the plaintext size limit.
|
||
if pending.text.len() + text.len() >= OBSERVER_CHUNK_MAX_TEXT_BYTES {
|
||
let events = self.flush();
|
||
// Start a new pending entry with the current chunk.
|
||
self.pending.push(PendingObserverChunk { key, event, text });
|
||
return events;
|
||
}
|
||
pending.text.push_str(&text);
|
||
pending.event.seq = event.seq;
|
||
pending.event.timestamp = event.timestamp;
|
||
return Vec::new();
|
||
}
|
||
|
||
self.pending.push(PendingObserverChunk { key, event, text });
|
||
Vec::new()
|
||
}
|
||
|
||
fn flush(&mut self) -> Vec<observer::ObserverEvent> {
|
||
self.pending
|
||
.drain(..)
|
||
.map(|mut pending| {
|
||
set_observer_chunk_text(&mut pending.event.payload, pending.text);
|
||
pending.event
|
||
})
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
fn observer_chunk_key_and_text(
|
||
event: &observer::ObserverEvent,
|
||
) -> Option<(ObserverChunkKey, String)> {
|
||
let update = event.payload.get("params")?.get("update")?;
|
||
let update_type = update.get("sessionUpdate")?.as_str()?;
|
||
if !matches!(
|
||
update_type,
|
||
"agent_message_chunk" | "user_message_chunk" | "agent_thought_chunk"
|
||
) {
|
||
return None;
|
||
}
|
||
|
||
let text = update.get("content")?.get("text")?.as_str()?.to_string();
|
||
let message_id = update
|
||
.get("messageId")
|
||
.and_then(|value| value.as_str())
|
||
.map(ToOwned::to_owned);
|
||
|
||
Some((
|
||
ObserverChunkKey {
|
||
update_type: update_type.to_string(),
|
||
message_id,
|
||
channel_id: event.channel_id.clone(),
|
||
session_id: event.session_id.clone(),
|
||
turn_id: event.turn_id.clone(),
|
||
agent_index: event.agent_index,
|
||
},
|
||
text,
|
||
))
|
||
}
|
||
|
||
fn set_observer_chunk_text(payload: &mut serde_json::Value, text: String) {
|
||
let Some(content) = payload
|
||
.get_mut("params")
|
||
.and_then(|params| params.get_mut("update"))
|
||
.and_then(|update| update.get_mut("content"))
|
||
else {
|
||
return;
|
||
};
|
||
|
||
if let Some(content_object) = content.as_object_mut() {
|
||
content_object.insert("text".to_string(), serde_json::Value::String(text));
|
||
}
|
||
}
|
||
|
||
/// Bytes of head and tail to retain from an elided string leaf — the value
|
||
/// shown to the renderer at each end. The ONLY tuning knob here: large enough
|
||
/// that a clipped diff/tool-result still shows real content, small enough that
|
||
/// eliding actually shrinks the frame.
|
||
const OBSERVER_LEAF_RETAIN_BYTES: usize = 3_000;
|
||
|
||
/// Trim an oversized observer telemetry frame so its SERIALIZED form fits under
|
||
/// `OBSERVER_MAX_PLAINTEXT_LEN`, instead of dropping the whole frame (silent
|
||
/// telemetry loss). The common case — a frame already under budget — is left
|
||
/// byte-identical.
|
||
///
|
||
/// The cap is measured in SERIALIZED bytes (JSON escaping makes serialized
|
||
/// length differ from raw), so the stop condition is always a full reserialize
|
||
/// of the whole frame: that counts the envelope, the variable `Option<String>`
|
||
/// IDs, and any elision markers exactly. No separate margin constant is needed.
|
||
///
|
||
/// Termination is provable: each iteration elides the largest string leaf that
|
||
/// would STRICTLY shrink the serialized frame, then reserializes. Shrinkability
|
||
/// is re-evaluated against each leaf's CURRENT value, so a leaf already at its
|
||
/// retained floor can never be re-elided — the loop strictly decreases the
|
||
/// serialized length each pass and is bounded by the leaf count. When no leaf
|
||
/// can shrink the frame and it still overflows, the payload is replaced with a
|
||
/// tiny stub, which trivially fits. Monotone decrease, bounded below by the stub.
|
||
///
|
||
/// **Signature choice (`&mut`, double-serialize accepted):** on the common
|
||
/// under-budget path this serializes the frame once to decide it fits, then
|
||
/// `encrypt_observer_payload` serializes it again — one extra `to_string` of an
|
||
/// already-small frame. Reusing that string would mean changing buzz-core's
|
||
/// `encrypt_observer_payload` signature or adding a parallel encrypt path; both
|
||
/// are out of this change's scope (buzz-core stays untouched). The clean `&mut`
|
||
/// signature with one cheap redundant serialize is the deliberate tradeoff.
|
||
fn fit_observer_event_to_budget(event: &mut observer::ObserverEvent) {
|
||
if serialized_len(event) <= OBSERVER_MAX_PLAINTEXT_LEN {
|
||
return;
|
||
}
|
||
|
||
// Raw size of the payload we are about to trim, captured before mutation so
|
||
// the stub's `originalBytes` reports source bytes discarded, not serialized
|
||
// overflow — consistent with the per-leaf marker's raw byte count.
|
||
let original_payload_bytes = serde_json::to_string(&event.payload)
|
||
.map(|s| s.len())
|
||
.unwrap_or(0);
|
||
|
||
// Elide the largest shrinkable leaf, reserialize, repeat. Each successful
|
||
// elision strictly shrinks the serialized frame, and a floored leaf can
|
||
// never be re-elided, so the loop is bounded by the leaf count.
|
||
while let Some(leaf) = largest_shrinkable_leaf(&mut event.payload) {
|
||
elide_leaf(leaf);
|
||
if serialized_len(event) <= OBSERVER_MAX_PLAINTEXT_LEN {
|
||
return;
|
||
}
|
||
}
|
||
|
||
// No leaf can shrink the frame further and it still overflows: replace the
|
||
// whole payload with a stub that is trivially under-cap.
|
||
event.payload = serde_json::json!({
|
||
"elided": format!("{} payload too large", event.kind),
|
||
"originalBytes": original_payload_bytes,
|
||
});
|
||
}
|
||
|
||
fn serialized_len(event: &observer::ObserverEvent) -> usize {
|
||
serde_json::to_string(event).map(|s| s.len()).unwrap_or(0)
|
||
}
|
||
|
||
/// Find the longest string leaf that would STRICTLY shrink if elided, returning
|
||
/// a mutable handle to it. A leaf shrinks only if `head + marker + tail` is
|
||
/// shorter than its current value (the marker-pushback guard); a leaf already at
|
||
/// its retained floor fails this test and is skipped, which is what bounds the
|
||
/// loop. Returns `None` when no leaf can shrink.
|
||
fn largest_shrinkable_leaf(value: &mut serde_json::Value) -> Option<&mut serde_json::Value> {
|
||
// First pass: find the byte length of the best candidate without holding a
|
||
// borrow, then re-descend to return the matching mutable reference. Two
|
||
// immutable-style passes keep the borrow checker happy without unsafe.
|
||
let best_len = max_shrinkable_len(value)?;
|
||
find_leaf_with_len(value, best_len)
|
||
}
|
||
|
||
/// Largest current length among string leaves that can strictly shrink.
|
||
fn max_shrinkable_len(value: &serde_json::Value) -> Option<usize> {
|
||
match value {
|
||
serde_json::Value::String(s) if leaf_shrinks(s) => Some(s.len()),
|
||
serde_json::Value::String(_) => None,
|
||
serde_json::Value::Array(items) => items.iter().filter_map(max_shrinkable_len).max(),
|
||
serde_json::Value::Object(map) => map.values().filter_map(max_shrinkable_len).max(),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
/// Return the first string leaf whose current length equals `target` and that
|
||
/// can strictly shrink. Used after `max_shrinkable_len` to re-acquire a mutable
|
||
/// borrow of the chosen leaf.
|
||
fn find_leaf_with_len(
|
||
value: &mut serde_json::Value,
|
||
target: usize,
|
||
) -> Option<&mut serde_json::Value> {
|
||
match value {
|
||
serde_json::Value::String(s) if s.len() == target && leaf_shrinks(s) => Some(value),
|
||
serde_json::Value::Array(items) => items
|
||
.iter_mut()
|
||
.find_map(|item| find_leaf_with_len(item, target)),
|
||
serde_json::Value::Object(map) => map
|
||
.values_mut()
|
||
.find_map(|item| find_leaf_with_len(item, target)),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
/// True when eliding `s` to head + marker + tail yields a strictly shorter raw
|
||
/// string. The marker width grows with `N` (bytes removed), so a leaf only
|
||
/// marginally larger than the retained ends must NOT be touched.
|
||
fn leaf_shrinks(s: &str) -> bool {
|
||
let (head_end, tail_start) = elision_boundaries(s);
|
||
tail_start > head_end && {
|
||
let removed = tail_start - head_end;
|
||
let marker = elision_marker(removed);
|
||
head_end + marker.len() + (s.len() - tail_start) < s.len()
|
||
}
|
||
}
|
||
|
||
/// Replace the middle of a string leaf with `…[elided N bytes]…`, keeping a head
|
||
/// and tail slice on UTF-8 char boundaries. `N` is RAW bytes removed.
|
||
fn elide_leaf(leaf: &mut serde_json::Value) {
|
||
let serde_json::Value::String(s) = leaf else {
|
||
return;
|
||
};
|
||
let (head_end, tail_start) = elision_boundaries(s);
|
||
let removed = tail_start - head_end;
|
||
let mut elided = String::with_capacity(head_end + 32 + (s.len() - tail_start));
|
||
elided.push_str(&s[..head_end]);
|
||
elided.push_str(&elision_marker(removed));
|
||
elided.push_str(&s[tail_start..]);
|
||
*s = elided;
|
||
}
|
||
|
||
fn elision_marker(removed_bytes: usize) -> String {
|
||
format!("…[elided {removed_bytes} bytes]…")
|
||
}
|
||
|
||
/// Byte offsets bounding the elided middle, snapped to char boundaries so we
|
||
/// never split a multi-byte char. Returns `(head_end, tail_start)` with
|
||
/// `head_end <= tail_start`.
|
||
fn elision_boundaries(s: &str) -> (usize, usize) {
|
||
let head_end = floor_char_boundary(s, OBSERVER_LEAF_RETAIN_BYTES.min(s.len()));
|
||
let tail_start = ceil_char_boundary(s, s.len().saturating_sub(OBSERVER_LEAF_RETAIN_BYTES));
|
||
(head_end, tail_start.max(head_end))
|
||
}
|
||
|
||
fn floor_char_boundary(s: &str, mut i: usize) -> usize {
|
||
while i > 0 && !s.is_char_boundary(i) {
|
||
i -= 1;
|
||
}
|
||
i
|
||
}
|
||
|
||
fn ceil_char_boundary(s: &str, mut i: usize) -> usize {
|
||
while i < s.len() && !s.is_char_boundary(i) {
|
||
i += 1;
|
||
}
|
||
i
|
||
}
|
||
|
||
async fn publish_relay_observer_event(
|
||
publisher: &RelayEventPublisher,
|
||
keys: &nostr::Keys,
|
||
agent_pubkey_hex: &str,
|
||
owner_pubkey_hex: &str,
|
||
owner_pubkey: &PublicKey,
|
||
mut event: observer::ObserverEvent,
|
||
) {
|
||
// Trim oversized frames to fit the plaintext cap rather than letting
|
||
// encrypt_observer_payload reject and drop them whole (silent telemetry loss).
|
||
fit_observer_event_to_budget(&mut event);
|
||
let encrypted = match encrypt_observer_payload(keys, owner_pubkey, &event) {
|
||
Ok(encrypted) => encrypted,
|
||
Err(error) => {
|
||
tracing::warn!("failed to encrypt relay observer event: {error}");
|
||
return;
|
||
}
|
||
};
|
||
let builder = match buzz_sdk::build_agent_observer_frame(
|
||
owner_pubkey_hex,
|
||
agent_pubkey_hex,
|
||
OBSERVER_FRAME_TELEMETRY,
|
||
&encrypted,
|
||
) {
|
||
Ok(builder) => builder,
|
||
Err(error) => {
|
||
tracing::warn!("failed to build relay observer event: {error}");
|
||
return;
|
||
}
|
||
};
|
||
let signed = match builder.sign_with_keys(keys) {
|
||
Ok(event) => event,
|
||
Err(error) => {
|
||
tracing::warn!("failed to sign relay observer event: {error}");
|
||
return;
|
||
}
|
||
};
|
||
if let Err(error) = publisher.publish_event(signed).await {
|
||
tracing::warn!("relay observer event dropped: {error}");
|
||
}
|
||
}
|
||
|
||
/// Maximum age (seconds) for an observer control frame to be considered fresh.
|
||
const OBSERVER_CONTROL_FRESHNESS_SECS: i64 = 300;
|
||
|
||
fn handle_relay_observer_control_event(
|
||
keys: &nostr::Keys,
|
||
event: nostr::Event,
|
||
pool: &mut AgentPool,
|
||
observer: Option<&observer::ObserverHandle>,
|
||
owner_pubkey_hex: &str,
|
||
) {
|
||
// Defense-in-depth: verify signature even though the relay already checked.
|
||
if let Err(e) = buzz_core::verify_event(&event) {
|
||
tracing::warn!(error = %e, "observer control frame failed signature verification");
|
||
return;
|
||
}
|
||
|
||
// Defense-in-depth: verify the sender is the resolved owner.
|
||
if event.pubkey.to_hex() != owner_pubkey_hex {
|
||
tracing::warn!(
|
||
sender = %event.pubkey,
|
||
expected = %owner_pubkey_hex,
|
||
"observer control frame from non-owner — dropping"
|
||
);
|
||
return;
|
||
}
|
||
|
||
// Freshness: reject stale/replayed frames outside ±5 minute window.
|
||
let now = chrono::Utc::now().timestamp();
|
||
let event_ts = event.created_at.as_secs() as i64;
|
||
if (event_ts - now).unsigned_abs() > OBSERVER_CONTROL_FRESHNESS_SECS as u64 {
|
||
tracing::warn!(
|
||
event_ts,
|
||
now,
|
||
"observer control frame outside freshness window — dropping"
|
||
);
|
||
return;
|
||
}
|
||
|
||
let payload = match decrypt_observer_payload::<serde_json::Value>(keys, &event) {
|
||
Ok(payload) => payload,
|
||
Err(error) => {
|
||
tracing::warn!("failed to decrypt observer control frame: {error}");
|
||
return;
|
||
}
|
||
};
|
||
|
||
let command_type = payload.get("type").and_then(|value| value.as_str());
|
||
match command_type {
|
||
Some("cancel_turn") => {
|
||
handle_cancel_turn_control(&payload, pool, observer);
|
||
}
|
||
Some("switch_model") => {
|
||
handle_switch_model_control(&payload, pool, observer);
|
||
}
|
||
_ => {
|
||
tracing::debug!(payload = %payload, "ignoring unknown observer control frame");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Handle a `cancel_turn` control frame: signal the in-flight task to cancel.
|
||
fn handle_cancel_turn_control(
|
||
payload: &serde_json::Value,
|
||
pool: &mut AgentPool,
|
||
observer: Option<&observer::ObserverHandle>,
|
||
) {
|
||
let Some(channel_id) = payload
|
||
.get("channelId")
|
||
.and_then(|value| value.as_str())
|
||
.and_then(|value| value.parse::<Uuid>().ok())
|
||
else {
|
||
tracing::warn!("observer cancel_turn control frame missing valid channelId");
|
||
return;
|
||
};
|
||
|
||
let fired = signal_in_flight_task(pool, channel_id, ControlSignal::Cancel);
|
||
let status = if fired { "sent" } else { "no_active_turn" };
|
||
if let Some(observer) = observer {
|
||
observer.emit(
|
||
"control_result",
|
||
None,
|
||
&observer::ObserverContext {
|
||
channel_id: Some(channel_id.to_string()),
|
||
session_id: None,
|
||
turn_id: None,
|
||
},
|
||
serde_json::json!({
|
||
"type": "cancel_turn",
|
||
"status": status,
|
||
}),
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Handle a `switch_model` control frame (Phase 3a, Option ii).
|
||
///
|
||
/// Busy path: deliver `SwitchModel` over the in-flight task's oneshot — the
|
||
/// task cancels the turn, sets `desired_model`, and requeues the batch so it
|
||
/// re-runs on a fresh session under the new model. A catalog miss surfaces
|
||
/// post-cancel via `create_session_and_apply_model` (the turn restarts on the
|
||
/// unchanged model + an `unsupported_model` result).
|
||
///
|
||
/// Idle path: validate against the cached catalog *before* invalidating
|
||
/// (pre-cancel guard), then set `desired_model` + invalidate. The override
|
||
/// takes visible effect on the agent's next turn.
|
||
fn handle_switch_model_control(
|
||
payload: &serde_json::Value,
|
||
pool: &mut AgentPool,
|
||
observer: Option<&observer::ObserverHandle>,
|
||
) {
|
||
let Some(channel_id) = payload
|
||
.get("channelId")
|
||
.and_then(|value| value.as_str())
|
||
.and_then(|value| value.parse::<Uuid>().ok())
|
||
else {
|
||
tracing::warn!("observer switch_model control frame missing valid channelId");
|
||
return;
|
||
};
|
||
let Some(model_id) = payload.get("modelId").and_then(|value| value.as_str()) else {
|
||
tracing::warn!("observer switch_model control frame missing modelId");
|
||
return;
|
||
};
|
||
|
||
// A turn is in flight for this channel iff a task_map entry exists. The
|
||
// agent is moved out of the pool during a turn, so the control oneshot is
|
||
// the only reachable lever; an idle channel has no such entry.
|
||
let turn_in_flight = pool
|
||
.task_map()
|
||
.values()
|
||
.any(|m| m.channel_id == Some(channel_id));
|
||
|
||
let status = if turn_in_flight {
|
||
// Busy path: deliver over the oneshot. `false` means the oneshot was
|
||
// already consumed this turn (a prior cancel/interrupt) — the turn is
|
||
// already ending, so the switch cannot land on it.
|
||
if signal_in_flight_task(
|
||
pool,
|
||
channel_id,
|
||
ControlSignal::SwitchModel(model_id.to_string()),
|
||
) {
|
||
"sent"
|
||
} else {
|
||
"turn_ending"
|
||
}
|
||
} else {
|
||
// Idle path: validate against the cached catalog before invalidating.
|
||
match pool.switch_idle_agent_model(channel_id, model_id) {
|
||
IdleSwitchResult::Switched => "switched",
|
||
IdleSwitchResult::UnsupportedModel => "unsupported_model",
|
||
IdleSwitchResult::NoIdleAgent => "no_active_turn",
|
||
}
|
||
};
|
||
|
||
if let Some(observer) = observer {
|
||
observer.emit(
|
||
"control_result",
|
||
None,
|
||
&observer::ObserverContext {
|
||
channel_id: Some(channel_id.to_string()),
|
||
session_id: None,
|
||
turn_id: None,
|
||
},
|
||
serde_json::json!({
|
||
"type": "switch_model",
|
||
"status": status,
|
||
"modelId": model_id,
|
||
}),
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Maximum crashes in a 60-second window before a slot's circuit opens.
|
||
const CIRCUIT_BREAKER_THRESHOLD: usize = 3;
|
||
/// Window for circuit-breaker crash counting.
|
||
const CIRCUIT_BREAKER_WINDOW: Duration = Duration::from_secs(60);
|
||
/// Cooldown before a tripped circuit breaker allows a probe respawn.
|
||
const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(300); // 5 minutes
|
||
/// Base backoff delay for respawn (doubles per recent crash, capped at 30s).
|
||
const RESPAWN_BASE_DELAY: Duration = Duration::from_secs(1);
|
||
/// Maximum respawn backoff delay.
|
||
const RESPAWN_MAX_DELAY: Duration = Duration::from_secs(30);
|
||
|
||
/// Per-slot circuit breaker state.
|
||
///
|
||
/// `crash_times` holds timestamps of recent crashes within `CIRCUIT_BREAKER_WINDOW`.
|
||
/// `open_until` is set when the threshold is hit; the circuit stays open until that
|
||
/// instant, then allows one probe respawn (half-open). If the probe crashes, the
|
||
/// circuit re-opens for another `CIRCUIT_BREAKER_COOLDOWN` period.
|
||
///
|
||
/// All state transitions go through methods on this struct — callers never
|
||
/// manipulate `crash_times` or `open_until` directly.
|
||
struct SlotCircuit {
|
||
crash_times: Vec<std::time::Instant>,
|
||
open_until: Option<std::time::Instant>,
|
||
/// True while a background respawn/refill task is in flight for this slot.
|
||
/// Prevents duplicate spawns from maintenance ticks that fire before the
|
||
/// previous spawn_and_init completes.
|
||
respawn_in_flight: bool,
|
||
}
|
||
|
||
/// Result of [`SlotCircuit::record_crash`].
|
||
enum CrashVerdict {
|
||
/// Respawn is allowed after sleeping for this duration (jittered backoff).
|
||
Respawn(Duration),
|
||
/// Circuit is open — do not respawn.
|
||
CircuitOpen,
|
||
/// Circuit was open but cooldown has elapsed — one probe respawn is allowed
|
||
/// (no backoff sleep). If the probe crashes, the next `record_crash` will
|
||
/// immediately re-open the circuit.
|
||
HalfOpenProbe,
|
||
}
|
||
|
||
impl SlotCircuit {
|
||
/// Record a crash and decide whether to respawn.
|
||
///
|
||
/// This is the **single canonical path** for all crash → respawn decisions.
|
||
/// Called by `respawn_agent_into`, `recover_panicked_agent`, and slot refill.
|
||
fn record_crash(&mut self) -> CrashVerdict {
|
||
let now = std::time::Instant::now();
|
||
|
||
// Half-open: cooldown elapsed → allow one probe.
|
||
if let Some(open_until) = self.open_until {
|
||
if now >= open_until {
|
||
// Pre-seed crash_times to threshold-1 so that if the probe
|
||
// itself crashes on the *next* call, the threshold is hit
|
||
// immediately and the circuit re-opens. This implements a
|
||
// "prove stability for one full window" policy.
|
||
self.crash_times.clear();
|
||
for _ in 0..(CIRCUIT_BREAKER_THRESHOLD - 1) {
|
||
self.crash_times.push(now);
|
||
}
|
||
self.open_until = None;
|
||
return CrashVerdict::HalfOpenProbe;
|
||
} else {
|
||
return CrashVerdict::CircuitOpen;
|
||
}
|
||
}
|
||
|
||
// Record this crash and prune old entries.
|
||
self.crash_times.push(now);
|
||
self.crash_times
|
||
.retain(|&t| now.duration_since(t) < CIRCUIT_BREAKER_WINDOW);
|
||
|
||
let recent = self.crash_times.len();
|
||
|
||
if recent >= CIRCUIT_BREAKER_THRESHOLD {
|
||
self.open_until = Some(now + CIRCUIT_BREAKER_COOLDOWN);
|
||
return CrashVerdict::CircuitOpen;
|
||
}
|
||
|
||
// Exponential backoff: 1s * 2^(recent-1), capped at 30s, with ±20% jitter.
|
||
let base = RESPAWN_BASE_DELAY.saturating_mul(1u32 << (recent - 1).min(5));
|
||
let capped = base.min(RESPAWN_MAX_DELAY);
|
||
let jitter = (std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.unwrap_or_default()
|
||
.subsec_nanos() as f64)
|
||
/ 1_000_000_000.0; // 0.0..1.0
|
||
let factor = 0.8 + jitter * 0.4; // 0.8..1.2
|
||
CrashVerdict::Respawn(capped.mul_f64(factor))
|
||
}
|
||
|
||
/// Mark a spawn failure — opens the circuit so the slot isn't retried
|
||
/// on every heartbeat tick. Uses fresh `Instant::now()` so spawn latency
|
||
/// doesn't shorten the effective cooldown.
|
||
fn mark_spawn_failed(&mut self) {
|
||
self.open_until = Some(std::time::Instant::now() + CIRCUIT_BREAKER_COOLDOWN);
|
||
}
|
||
|
||
/// Check if an empty slot can be refilled. Unlike `record_crash`, this
|
||
/// does NOT record a new crash — it only checks whether the circuit
|
||
/// allows a respawn attempt.
|
||
///
|
||
/// Returns `true` if respawn is allowed. For half-open probes, pre-seeds
|
||
/// crash_times so the next crash re-opens immediately. For normal refills
|
||
/// (no circuit was ever opened), crash history is preserved so the breaker
|
||
/// can still trip if the refilled agent crashes quickly.
|
||
fn can_refill(&mut self) -> bool {
|
||
let now = std::time::Instant::now();
|
||
match self.open_until {
|
||
Some(open_until) => {
|
||
if now >= open_until {
|
||
// Half-open probe: pre-seed crash_times.
|
||
self.crash_times.clear();
|
||
for _ in 0..(CIRCUIT_BREAKER_THRESHOLD - 1) {
|
||
self.crash_times.push(now);
|
||
}
|
||
self.open_until = None;
|
||
true
|
||
} else {
|
||
false // cooldown not elapsed
|
||
}
|
||
}
|
||
None => true, // no circuit open — normal refill, preserve crash history
|
||
}
|
||
}
|
||
}
|
||
|
||
/// True if any slot has a respawn task in flight. Used to prevent premature
|
||
/// "all agents dead" exits — a respawning agent may succeed in seconds.
|
||
fn any_respawn_in_flight(crash_history: &[SlotCircuit]) -> bool {
|
||
crash_history.iter().any(|s| s.respawn_in_flight)
|
||
}
|
||
|
||
/// Result of a background respawn task.
|
||
struct RespawnResult {
|
||
index: usize,
|
||
/// Tuple: (initialized client, protocol version, supports_goose_steer).
|
||
/// The third element is always `true` — the supervisor uses
|
||
/// try-and-tolerate for the steer extension.
|
||
result: Result<(AcpClient, u32, String)>,
|
||
}
|
||
|
||
/// Outcome of a non-cancelling steer attempt, forwarded from a per-attempt
|
||
/// watcher task (which awaits the `SteerRequest.ack_tx` oneshot) back to
|
||
/// the main loop's `select!`. The main loop drives queue side-effects from
|
||
/// this — it cannot await the oneshot itself without blocking the relay
|
||
/// stream.
|
||
///
|
||
/// Carries enough identity to operate on the right withheld event in
|
||
/// `EventQueue::withheld_native_steer`: `channel_id` is the routing key,
|
||
/// `event_id` is the hex id of the single event the steer carried.
|
||
struct SteerAckEvent {
|
||
channel_id: Uuid,
|
||
event_id: String,
|
||
/// `Ok` if the read loop sent any of the locked `SteerAck` variants.
|
||
/// `Err` if the oneshot was dropped without a send — should not happen
|
||
/// under the current read-loop drains, but if it ever does the main
|
||
/// loop treats it as `PromptCompletedNeutral` (release withheld, no
|
||
/// fallback signal) to avoid leaking the withheld event.
|
||
ack: std::result::Result<pool::SteerAck, tokio::sync::oneshot::error::RecvError>,
|
||
}
|
||
|
||
/// RAII guard that ensures a `RespawnResult` is sent even if the task panics.
|
||
/// Without this, a panicked respawn task would leave `respawn_in_flight = true`
|
||
/// permanently, silently losing the slot forever.
|
||
struct RespawnGuard {
|
||
index: usize,
|
||
tx: mpsc::Sender<RespawnResult>,
|
||
sent: bool,
|
||
}
|
||
|
||
impl RespawnGuard {
|
||
fn new(index: usize, tx: mpsc::Sender<RespawnResult>) -> Self {
|
||
Self {
|
||
index,
|
||
tx,
|
||
sent: false,
|
||
}
|
||
}
|
||
|
||
/// Send the result and disarm the guard. Uses `try_send` (sync) so there
|
||
/// is no await boundary between marking `sent` and actually enqueueing —
|
||
/// cancellation cannot slip between the two.
|
||
fn send(mut self, result: Result<(AcpClient, u32, String)>) {
|
||
// Invariant: try_send succeeds because the channel capacity equals the
|
||
// slot count, and respawn_in_flight guarantees at most one outstanding
|
||
// result per slot. If this ever fails, the channel sizing or the
|
||
// respawn_in_flight guard has drifted — that's a bug, not a transient.
|
||
match self.tx.try_send(RespawnResult {
|
||
index: self.index,
|
||
result,
|
||
}) {
|
||
Ok(()) => self.sent = true,
|
||
Err(e) => {
|
||
tracing::error!(
|
||
agent = self.index,
|
||
"respawn result channel full or closed: {e}"
|
||
);
|
||
// Drop will fire and send a failure result as fallback.
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Drop for RespawnGuard {
|
||
fn drop(&mut self) {
|
||
if !self.sent {
|
||
tracing::error!(
|
||
agent = self.index,
|
||
"respawn task exited without sending result — sending failure"
|
||
);
|
||
// Best-effort: try_send in Drop (can't await).
|
||
let _ = self.tx.try_send(RespawnResult {
|
||
index: self.index,
|
||
result: Err(anyhow::anyhow!("respawn task panicked or was cancelled")),
|
||
});
|
||
}
|
||
}
|
||
}
|
||
|
||
//
|
||
// Sync env-var propagation must run before the tokio runtime starts so that
|
||
// any child processes inherit the correct environment. This must happen in the
|
||
// sync entry point — `std::env::set_var` is only safe before tokio spawns
|
||
// worker threads (Rust 2024 edition safety requirement).
|
||
|
||
pub fn run() -> Result<()> {
|
||
config::propagate_legacy_env_vars();
|
||
tokio_main()
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn tokio_main() -> Result<()> {
|
||
// Install the ring crypto provider for rustls (required for wss:// connections).
|
||
rustls::crypto::ring::default_provider()
|
||
.install_default()
|
||
.expect("failed to install rustls crypto provider");
|
||
if is_subcommand("models") {
|
||
// Strip the subcommand token so clap doesn't reject it as a positional.
|
||
// Keeps argv[0] (binary name) and passes everything after the subcommand.
|
||
let filtered: Vec<String> = std::env::args()
|
||
.enumerate()
|
||
.filter(|(i, _)| *i != 1)
|
||
.map(|(_, a)| a)
|
||
.collect();
|
||
let args = ModelsArgs::parse_from(&filtered);
|
||
return run_models(args).await;
|
||
}
|
||
|
||
if is_subcommand("auth-methods") {
|
||
let filtered: Vec<String> = std::env::args()
|
||
.enumerate()
|
||
.filter(|(i, _)| *i != 1)
|
||
.map(|(_, a)| a)
|
||
.collect();
|
||
let args = AuthMethodsArgs::parse_from(&filtered);
|
||
return run_auth_methods(args).await;
|
||
}
|
||
|
||
if is_subcommand("authenticate") {
|
||
let filtered: Vec<String> = std::env::args()
|
||
.enumerate()
|
||
.filter(|(i, _)| *i != 1)
|
||
.map(|(_, a)| a)
|
||
.collect();
|
||
let args = AuthenticateArgs::parse_from(&filtered);
|
||
return run_authenticate(args).await;
|
||
}
|
||
|
||
tracing_subscriber::fmt()
|
||
.with_env_filter(
|
||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("buzz_acp=info")),
|
||
)
|
||
.compact()
|
||
.init();
|
||
|
||
let mut config = Config::from_cli().map_err(|e| anyhow::anyhow!("configuration error: {e}"))?;
|
||
|
||
// ── Setup-mode early branch ───────────────────────────────────────────────
|
||
//
|
||
// When the desktop determines an agent is not ready (missing credentials,
|
||
// model, or provider), it spawns buzz-acp with BUZZ_ACP_SETUP_PAYLOAD set.
|
||
// We enter the minimal setup-listener path and never start the agent pool.
|
||
if let Some(payload) = setup_mode::SetupPayload::from_env()
|
||
.map_err(|e| anyhow::anyhow!("setup payload error: {e}"))?
|
||
{
|
||
tracing::info!("buzz-acp: setup payload present, entering setup-listener mode");
|
||
return setup_mode::run_setup_listener(config, payload).await;
|
||
}
|
||
|
||
tracing::info!("buzz-acp starting: {}", config.summary());
|
||
|
||
let observer = config
|
||
.relay_observer
|
||
.then(observer::ObserverHandle::in_process);
|
||
if let Some(handle) = &observer {
|
||
handle.emit(
|
||
"harness_started",
|
||
None,
|
||
&observer::ObserverContext::default(),
|
||
serde_json::json!({
|
||
"relayUrl": config.relay_url,
|
||
"agentCommand": config.agent_command,
|
||
"agentArgs": config.agent_args,
|
||
"parallelism": config.agents,
|
||
"relayObserver": config.relay_observer,
|
||
}),
|
||
);
|
||
}
|
||
|
||
//
|
||
// Finding #10: one agent failing to start must not kill the whole pool.
|
||
// We attempt each spawn under a 60-second timeout; failures are logged and
|
||
// skipped. If ALL agents fail we return an error. A partial pool is valid —
|
||
// the harness continues with reduced capacity and logs a warning.
|
||
let mut agent_slots: Vec<Option<OwnedAgent>> = Vec::with_capacity(config.agents as usize);
|
||
for i in 0..config.agents as usize {
|
||
// Spawn OUTSIDE the timeout so we always own the child for cleanup.
|
||
// This matches the run_models pattern and prevents zombie leaks on
|
||
// init timeout (the cancelled future would drop the AcpClient via
|
||
// Drop which is best-effort only).
|
||
let spawn_result = AcpClient::spawn(
|
||
&config.agent_command,
|
||
&config.agent_args,
|
||
&config.persona_env_vars,
|
||
config.has_generated_codex_config,
|
||
)
|
||
.await;
|
||
match spawn_result {
|
||
Ok(mut acp) => {
|
||
acp.set_observer(observer.clone(), i);
|
||
match tokio::time::timeout(Duration::from_secs(60), acp.initialize()).await {
|
||
Ok(Ok(init_result)) => {
|
||
tracing::info!(agent = i, "agent initialized: {init_result}");
|
||
let protocol_version =
|
||
init_result["protocolVersion"].as_u64().unwrap_or(1) as u32;
|
||
tracing::info!(
|
||
agent = i,
|
||
name = init_result
|
||
.get("agentInfo")
|
||
.or_else(|| init_result.get("serverInfo"))
|
||
.and_then(|info| info.get("name"))
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("unknown"),
|
||
"agent initialized — non-cancelling steer enabled (try-and-tolerate)"
|
||
);
|
||
acp.observe(
|
||
"agent_initialized",
|
||
serde_json::json!({
|
||
"agentIndex": i,
|
||
"initializeResult": init_result,
|
||
}),
|
||
);
|
||
let agent_name = normalized_agent_name(&init_result);
|
||
agent_slots.push(Some(OwnedAgent {
|
||
index: i,
|
||
acp,
|
||
state: SessionState::default(),
|
||
model_capabilities: None,
|
||
desired_model: config.model.clone(),
|
||
model_overridden: false,
|
||
agent_name,
|
||
goose_system_prompt_supported: None,
|
||
protocol_version,
|
||
}));
|
||
}
|
||
Ok(Err(e)) => {
|
||
tracing::error!(agent = i, "agent initialize failed: {e}");
|
||
acp.shutdown().await;
|
||
agent_slots.push(None);
|
||
}
|
||
Err(_) => {
|
||
tracing::error!(agent = i, "agent timed out during init (60s)");
|
||
acp.shutdown().await;
|
||
agent_slots.push(None);
|
||
}
|
||
}
|
||
}
|
||
Err(e) => {
|
||
tracing::error!(agent = i, "agent failed to spawn: {e}");
|
||
agent_slots.push(None);
|
||
}
|
||
}
|
||
}
|
||
let live_count = agent_slots.iter().filter(|s| s.is_some()).count();
|
||
if live_count == 0 {
|
||
return Err(anyhow::anyhow!(
|
||
"all {} agents failed to start — cannot continue",
|
||
config.agents
|
||
));
|
||
}
|
||
if live_count < config.agents as usize {
|
||
tracing::warn!(
|
||
"started {}/{} agents — continuing with reduced pool",
|
||
live_count,
|
||
config.agents
|
||
);
|
||
}
|
||
tracing::info!("agent_pool_ready agents={}", live_count);
|
||
let mut pool = AgentPool::from_slots(agent_slots);
|
||
|
||
//
|
||
// Finding #22: capture a startup watermark BEFORE connecting to the relay.
|
||
// This timestamp is used for membership notification replay (via
|
||
// startup_watermark) and as the initial subscribe_since for channels
|
||
// discovered at startup. The Subscribe handler falls back to
|
||
// subscribe_since when last_seen is None, closing the blind spot
|
||
// between "agents ready" and "first REQ sent".
|
||
let startup_watermark: u64 = std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.unwrap_or_default()
|
||
.as_secs();
|
||
|
||
let pubkey_hex = config.keys.public_key().to_hex();
|
||
|
||
// Parse BUZZ_AUTH_TAG into a nostr::Tag for NIP-OA relay membership delegation.
|
||
let relay_auth_tag: Option<nostr::Tag> = std::env::var("BUZZ_AUTH_TAG")
|
||
.ok()
|
||
.filter(|s| !s.is_empty())
|
||
.and_then(|s| buzz_sdk::nip_oa::parse_auth_tag(&s).ok());
|
||
|
||
let mut relay =
|
||
HarnessRelay::connect(&config.relay_url, &config.keys, &pubkey_hex, relay_auth_tag)
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("relay connect error: {e}"))?;
|
||
|
||
// Finding #22: tell the relay background task the watermark so it can use
|
||
// `since = watermark - 5s` on the first REQ instead of `since=now`.
|
||
// Best-effort: a failure here is non-fatal (we just lose the startup window
|
||
// protection, which is the same as the pre-fix behaviour).
|
||
if let Err(e) = relay.set_startup_watermark(startup_watermark).await {
|
||
tracing::warn!("failed to set startup watermark: {e}");
|
||
}
|
||
|
||
tracing::info!("connected to relay at {}", config.relay_url);
|
||
|
||
relay
|
||
.subscribe_membership_notifications()
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("membership notification subscribe error: {e}"))?;
|
||
tracing::info!("subscribed to membership notifications");
|
||
|
||
let presence_publisher = relay.event_publisher();
|
||
let presence_keys = config.keys.clone();
|
||
|
||
// Priority: BUZZ_AUTH_TAG (NIP-OA attestation) → --agent-owner flag.
|
||
let startup_owner: Option<String> = resolve_agent_owner(&config);
|
||
if let Some(ref owner) = startup_owner {
|
||
tracing::info!("agent owner: {owner}");
|
||
} else {
|
||
tracing::info!("no agent owner configured");
|
||
}
|
||
// Warn if owner-dependent mode but no owner resolved yet.
|
||
if startup_owner.is_none() {
|
||
match &config.respond_to {
|
||
RespondTo::OwnerOnly => {
|
||
tracing::warn!(
|
||
"respond-to=owner-only but no owner is set — all events will be \
|
||
dropped. Set BUZZ_AUTH_TAG or --agent-owner, or use --respond-to=anyone."
|
||
);
|
||
}
|
||
RespondTo::Allowlist => {
|
||
tracing::warn!(
|
||
"respond-to=allowlist but no owner is set — allowlisted pubkeys \
|
||
will still be accepted, but owner-based matching is unavailable \
|
||
until owner is resolved."
|
||
);
|
||
}
|
||
_ => {} // anyone/nobody don't depend on owner
|
||
}
|
||
}
|
||
let owner_cache = OwnerCache::new(startup_owner.clone());
|
||
|
||
let mut relay_observer_control_rx = None;
|
||
let mut relay_observer_publisher_task = None;
|
||
let mut relay_observer_publisher = None;
|
||
if config.relay_observer {
|
||
if let (Some(observer), Some(owner_pubkey_hex)) =
|
||
(observer.clone(), owner_cache.pubkey.clone())
|
||
{
|
||
match PublicKey::from_hex(&owner_pubkey_hex) {
|
||
Ok(owner_pubkey) => {
|
||
relay_observer_publisher = Some((
|
||
observer,
|
||
relay.event_publisher(),
|
||
config.keys.clone(),
|
||
pubkey_hex.clone(),
|
||
owner_pubkey_hex,
|
||
owner_pubkey,
|
||
));
|
||
relay
|
||
.subscribe_observer_controls()
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("observer control subscribe error: {e}"))?;
|
||
relay_observer_control_rx = relay.take_observer_control_rx();
|
||
tracing::info!("relay observer enabled");
|
||
}
|
||
Err(error) => {
|
||
tracing::warn!("relay observer disabled: invalid owner pubkey: {error}");
|
||
}
|
||
}
|
||
} else {
|
||
tracing::warn!(
|
||
"relay observer requested but no agent owner was resolved at startup; \
|
||
observer frames will not be published"
|
||
);
|
||
}
|
||
}
|
||
|
||
let channel_info_map = relay
|
||
.discover_channels()
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("channel discovery error: {e}"))?;
|
||
|
||
tracing::info!("discovered {} channel(s)", channel_info_map.len());
|
||
let channel_ids: Vec<Uuid> = channel_info_map.keys().copied().collect();
|
||
|
||
let rules: Vec<SubscriptionRule> = match config.subscribe_mode {
|
||
SubscribeMode::Mentions => {
|
||
vec![SubscriptionRule {
|
||
name: "mentions".into(),
|
||
channels: filter::ChannelScope::All("all".into()),
|
||
kinds: config.kinds_override.clone().unwrap_or_else(|| {
|
||
vec![
|
||
KIND_STREAM_MESSAGE,
|
||
KIND_WORKFLOW_APPROVAL_REQUESTED,
|
||
KIND_STREAM_REMINDER,
|
||
]
|
||
}),
|
||
require_mention: !config.no_mention_filter,
|
||
filter: None,
|
||
compiled_filter: None,
|
||
consecutive_timeouts: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
|
||
prompt_tag: Some("@mention".into()),
|
||
}]
|
||
}
|
||
SubscribeMode::All => {
|
||
vec![SubscriptionRule {
|
||
name: "all".into(),
|
||
channels: filter::ChannelScope::All("all".into()),
|
||
kinds: config.kinds_override.clone().unwrap_or_default(),
|
||
require_mention: false,
|
||
filter: None,
|
||
compiled_filter: None,
|
||
consecutive_timeouts: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
|
||
prompt_tag: Some("all".into()),
|
||
}]
|
||
}
|
||
SubscribeMode::Config => {
|
||
// load_rules() already warns if the config file has zero rules.
|
||
config::load_rules(&config.config_path)?
|
||
}
|
||
};
|
||
|
||
let channel_filters = config::resolve_channel_filters(&config, &channel_ids, &rules);
|
||
if channel_filters.is_empty() {
|
||
tracing::warn!("no channel subscriptions resolved — agent will sit idle");
|
||
}
|
||
let mut subscribed_channel_ids = HashSet::with_capacity(channel_filters.len());
|
||
for (channel_id, filter) in &channel_filters {
|
||
if let Err(e) = relay.subscribe_channel(*channel_id, filter.clone()).await {
|
||
tracing::warn!("failed to subscribe to channel {channel_id}: {e}");
|
||
} else {
|
||
subscribed_channel_ids.insert(*channel_id);
|
||
tracing::info!("subscribed to channel {channel_id}");
|
||
}
|
||
}
|
||
|
||
if let Some((observer, publisher, keys, agent_pubkey, owner_pubkey, owner)) =
|
||
relay_observer_publisher.take()
|
||
{
|
||
relay_observer_publisher_task = Some(spawn_relay_observer_publisher(
|
||
observer,
|
||
publisher,
|
||
keys,
|
||
agent_pubkey,
|
||
owner_pubkey,
|
||
owner,
|
||
));
|
||
}
|
||
|
||
// Online means the harness can receive work, not merely that its socket is
|
||
// connected. Publishing after channel subscriptions gives desktop callers
|
||
// a durable readiness boundary before they send a startup mention.
|
||
if config.presence_enabled {
|
||
match publish_presence(&presence_publisher, &presence_keys, "online").await {
|
||
Ok(_) => tracing::info!("presence set to online"),
|
||
Err(e) => tracing::warn!("failed to set initial presence: {e}"),
|
||
}
|
||
}
|
||
|
||
let dedup_mode = config.dedup_mode;
|
||
let mut queue =
|
||
EventQueue::new(dedup_mode).with_in_flight_deadline(config.max_turn_duration_secs);
|
||
|
||
let base_prompt_content = config.base_prompt_content.take();
|
||
let ctx = Arc::new(PromptContext {
|
||
mcp_servers: build_mcp_servers(&config),
|
||
initial_message: config.initial_message.clone(),
|
||
idle_timeout: Duration::from_secs(config.idle_timeout_secs),
|
||
max_turn_duration: Duration::from_secs(config.max_turn_duration_secs),
|
||
turn_liveness_interval: Duration::from_secs(config.turn_liveness_secs),
|
||
dedup_mode: config.dedup_mode,
|
||
system_prompt: config.system_prompt.clone(),
|
||
team_instructions: config.team_instructions.clone(),
|
||
base_prompt: if config.no_base_prompt {
|
||
None
|
||
} else if let Some(content) = base_prompt_content {
|
||
Some(Box::leak(content.into_boxed_str()))
|
||
} else {
|
||
Some(include_str!("base_prompt.md"))
|
||
},
|
||
heartbeat_prompt: config.heartbeat_prompt.clone(),
|
||
cwd: std::env::current_dir()
|
||
.unwrap_or_else(|_| std::path::PathBuf::from("/"))
|
||
.to_string_lossy()
|
||
.to_string(),
|
||
rest_client: relay.rest_client(),
|
||
channel_info: channel_info_map,
|
||
context_message_limit: config.context_message_limit,
|
||
max_turns_per_session: config.max_turns_per_session,
|
||
permission_mode: config.permission_mode,
|
||
agent_keys: config.keys.clone(),
|
||
agent_owner_pubkey: startup_owner
|
||
.as_deref()
|
||
.and_then(|hex| nostr::PublicKey::from_hex(hex).ok()),
|
||
memory_enabled: config.memory_enabled,
|
||
harness_name: crate::config::normalize_agent_command_identity(&config.agent_command),
|
||
});
|
||
|
||
if !config.memory_enabled {
|
||
tracing::info!(
|
||
target: "engram::core",
|
||
"NIP-AE core memory injection disabled (re-enable by removing --no-memory / BUZZ_ACP_NO_MEMORY)"
|
||
);
|
||
}
|
||
|
||
let mut heartbeat = if config.heartbeat_interval_secs > 0 {
|
||
let interval = Duration::from_secs(config.heartbeat_interval_secs);
|
||
Some(tokio::time::interval_at(
|
||
tokio::time::Instant::now() + interval,
|
||
interval,
|
||
))
|
||
} else {
|
||
None
|
||
};
|
||
let mut heartbeat_in_flight = false;
|
||
|
||
let mut presence_heartbeat = if config.presence_enabled {
|
||
let interval = Duration::from_secs(60);
|
||
Some(tokio::time::interval_at(
|
||
tokio::time::Instant::now() + interval,
|
||
interval,
|
||
))
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let mut typing_refresh = if config.typing_enabled {
|
||
let interval = Duration::from_secs(3);
|
||
Some(tokio::time::interval_at(
|
||
tokio::time::Instant::now() + interval,
|
||
interval,
|
||
))
|
||
} else {
|
||
None
|
||
};
|
||
let mut typing_channels: HashMap<Uuid, ThreadTags> = HashMap::new();
|
||
let mut presence_task: Option<tokio::task::JoinHandle<()>> = None;
|
||
|
||
// Runs at the TOP of every loop iteration via Instant check — cannot be
|
||
// starved by the biased select. Slot refill spawns background tasks so
|
||
// spawn_and_init never blocks the main loop.
|
||
let maintenance_interval = Duration::from_secs(30);
|
||
let mut last_maintenance = std::time::Instant::now();
|
||
|
||
// Channel for background respawn tasks to return completed agents.
|
||
// Bounded to agent count — at most one respawn per slot in flight.
|
||
let (respawn_tx, mut respawn_rx) = mpsc::channel::<RespawnResult>(config.agents as usize);
|
||
// JoinSet for respawn tasks so shutdown can abort them.
|
||
let mut respawn_tasks: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
|
||
|
||
// Channel for non-cancelling steer ack watchers to forward outcomes back
|
||
// to the main loop. Each `pool.send_steer(...) == Ok(())` spawns a
|
||
// short-lived task that awaits the `SteerRequest.ack_tx` oneshot and
|
||
// forwards a `SteerAckEvent`. Unbounded because:
|
||
// 1. The producer count is bounded by in-flight goose turns
|
||
// (`agents` slots, capacity-1 `steer_tx` each), so the channel
|
||
// cannot legitimately back up under steady state.
|
||
// 2. We must never drop a steer outcome — losing an ack would leak a
|
||
// withheld event in `EventQueue::withheld_native_steer` until
|
||
// `IN_FLIGHT_DEADLINE_SECS` expires.
|
||
let (steer_ack_tx, mut steer_ack_rx) = mpsc::unbounded_channel::<SteerAckEvent>();
|
||
|
||
// ── Step 7: Shutdown signal ───────────────────────────────────────────────
|
||
let (shutdown_tx, mut shutdown_rx) = watch::channel(());
|
||
|
||
let tx = shutdown_tx.clone();
|
||
tokio::spawn(async move {
|
||
tokio::signal::ctrl_c().await.ok();
|
||
let _ = tx.send(());
|
||
});
|
||
|
||
#[cfg(unix)]
|
||
{
|
||
let tx = shutdown_tx.clone();
|
||
tokio::spawn(async move {
|
||
use tokio::signal::unix::{signal, SignalKind};
|
||
let mut sigterm = signal(SignalKind::terminate()).expect("SIGTERM handler");
|
||
sigterm.recv().await;
|
||
let _ = tx.send(());
|
||
});
|
||
}
|
||
|
||
// Track the newest membership notification timestamp per channel.
|
||
// On reconnect the relay replays events newest-first, so the first event
|
||
// per channel is authoritative. Any later event with ts < newest is stale.
|
||
// Exact duplicates (same event ID) are caught by seen_membership_ids.
|
||
//
|
||
// Uses strict `<` (not `<=`) so that legitimate live events at the same
|
||
// second are both processed. The seen_membership_ids set handles exact
|
||
// replays that share the same timestamp.
|
||
let mut membership_newest_ts: HashMap<Uuid, u64> = HashMap::new();
|
||
// Two-generation dedup for membership event replays (bounded, no amnesia).
|
||
// Rotates at 1000 entries instead of clearing the entire set at 2000.
|
||
let mut seen_membership_current: HashSet<String> = HashSet::new();
|
||
let mut seen_membership_previous: HashSet<String> = HashSet::new();
|
||
|
||
// Channels the agent has been removed from. When a checked-out agent is
|
||
// returned to the pool, its sessions for these channels are stripped, and
|
||
// failed/panicked batches for these channels are dropped instead of requeued.
|
||
//
|
||
// Cleared on re-add (KIND_MEMBER_ADDED_NOTIFICATION) so re-joined channels
|
||
// regain session affinity.
|
||
//
|
||
// Known limitation: if a batch is in-flight when the channel is removed AND
|
||
// re-added before the batch returns, the stale batch may be requeued. This
|
||
// is acceptable because: (a) the agent is a member again and has access,
|
||
// (b) the events are from the agent's authorized history, (c) the window
|
||
// is extremely narrow (membership changes are rare, prompt turns are seconds),
|
||
// and (d) fixing this would require per-channel epoch tracking on TaskMeta
|
||
// and PromptResult — significant complexity for a benign edge case. If strict
|
||
// causal invalidation is needed, add a monotonic epoch counter per channel
|
||
// and capture it in TaskMeta at dispatch time.
|
||
let mut removed_channels: HashSet<Uuid> = HashSet::new();
|
||
|
||
//
|
||
// One SlotCircuit per agent slot. crash_times entries are pruned to the last
|
||
// CIRCUIT_BREAKER_WINDOW on each respawn attempt. The Vec is indexed by
|
||
// agent slot index, so it must be sized to the configured pool capacity
|
||
// (not the live count, which may be smaller after partial startup).
|
||
let mut crash_history: Vec<SlotCircuit> = (0..config.agents as usize)
|
||
.map(|_| SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
})
|
||
.collect();
|
||
|
||
//
|
||
// Branches 1 & 2 both need to borrow `pool`, but they access different
|
||
// fields (result_rx vs join_set). We use `rx_and_join_set()` to split the
|
||
// borrow, yielding a typed enum so the outer code can dispatch cleanly.
|
||
enum PoolEvent {
|
||
Result(Box<PromptResult>),
|
||
Panic(tokio::task::JoinError),
|
||
SteerAck(SteerAckEvent),
|
||
}
|
||
|
||
loop {
|
||
if last_maintenance.elapsed() >= maintenance_interval {
|
||
last_maintenance = std::time::Instant::now();
|
||
queue.compact_expired_state();
|
||
|
||
// Slot refill: spawn background tasks for empty slots whose
|
||
// circuit breaker allows it. spawn_and_init runs off the main
|
||
// loop so it never blocks event processing.
|
||
for (idx, slot) in crash_history.iter_mut().enumerate() {
|
||
if pool.slot_alive(idx) || slot.respawn_in_flight {
|
||
continue;
|
||
}
|
||
if !slot.can_refill() {
|
||
continue;
|
||
}
|
||
slot.respawn_in_flight = true;
|
||
tracing::info!(agent = idx, "slot refill: spawning background respawn");
|
||
let cmd = config.agent_command.clone();
|
||
let args = config.agent_args.clone();
|
||
let env = config.persona_env_vars.clone();
|
||
let has_codex = config.has_generated_codex_config;
|
||
let observer = observer.clone();
|
||
let guard = RespawnGuard::new(idx, respawn_tx.clone());
|
||
respawn_tasks.spawn(async move {
|
||
let result = spawn_and_init(&cmd, &args, &env, has_codex, idx, observer).await;
|
||
guard.send(result);
|
||
});
|
||
}
|
||
|
||
// Flush requeued batches whose retry_after has expired. Without
|
||
// this, a batch requeued during crash recovery can sit idle
|
||
// indefinitely on quiet channels — dispatch_pending is only
|
||
// called on relay events or pool results, neither of which
|
||
// arrive when the channel is silent.
|
||
if queue.has_flushable_work() {
|
||
for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) {
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
}
|
||
|
||
let mut respawn_collected = false;
|
||
while let Ok(rr) = respawn_rx.try_recv() {
|
||
crash_history[rr.index].respawn_in_flight = false;
|
||
match rr.result {
|
||
Ok((acp, protocol_version, agent_name)) => {
|
||
let agent = OwnedAgent {
|
||
index: rr.index,
|
||
acp,
|
||
state: SessionState::default(),
|
||
model_capabilities: None,
|
||
desired_model: config.model.clone(),
|
||
model_overridden: false,
|
||
agent_name,
|
||
goose_system_prompt_supported: None,
|
||
protocol_version,
|
||
};
|
||
pool.return_agent(agent);
|
||
tracing::info!(agent = rr.index, "respawn complete");
|
||
respawn_collected = true;
|
||
}
|
||
Err(e) => {
|
||
crash_history[rr.index].mark_spawn_failed();
|
||
tracing::warn!(agent = rr.index, "respawn failed: {e} — circuit re-opened");
|
||
}
|
||
}
|
||
}
|
||
// Flush requeued events that were waiting for a live agent. Without
|
||
// this, batches requeued during crash recovery sit idle until the
|
||
// next relay event arrives — which can be minutes on quiet channels.
|
||
if respawn_collected {
|
||
for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) {
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
|
||
// Borrow result_rx and join_set simultaneously via split-borrow helper.
|
||
let pool_event: Option<PoolEvent> = {
|
||
let (result_rx, join_set) = pool.rx_and_join_set();
|
||
tokio::select! {
|
||
biased;
|
||
// Finding #24: recv() returning None means all senders dropped
|
||
// (pool was torn down). Break cleanly instead of panicking.
|
||
r = result_rx.recv() => match r {
|
||
Some(result) => Some(PoolEvent::Result(Box::new(result))),
|
||
None => {
|
||
tracing::info!("result channel closed — exiting main loop");
|
||
break;
|
||
}
|
||
},
|
||
// Guard: join_next() returns None immediately when JoinSet is
|
||
// empty, which would cause a tight spin. Only poll when there
|
||
// are in-flight tasks.
|
||
Some(Err(e)) = join_set.join_next(), if !join_set.is_empty() => {
|
||
Some(PoolEvent::Panic(e))
|
||
}
|
||
// Goose-native steer ack from a watcher task. Outcomes drive
|
||
// queue side-effects (drop / release withheld event) and
|
||
// optionally the cancel+merge fallback signal. See the
|
||
// `Some(PoolEvent::SteerAck(...))` match arm below for the
|
||
// locked semantics (Eva + Max + Perci).
|
||
Some(ack_event) = steer_ack_rx.recv() => {
|
||
Some(PoolEvent::SteerAck(ack_event))
|
||
}
|
||
control_event = async {
|
||
match relay_observer_control_rx.as_mut() {
|
||
Some(rx) => rx.recv().await,
|
||
None => std::future::pending().await,
|
||
}
|
||
} => {
|
||
let _ = result_rx;
|
||
match control_event {
|
||
Some(event) => {
|
||
if let Some(ref owner_hex) = owner_cache.pubkey {
|
||
handle_relay_observer_control_event(&config.keys, event, &mut pool, observer.as_ref(), owner_hex);
|
||
} else {
|
||
tracing::warn!("observer control frame received but no owner resolved — dropping");
|
||
}
|
||
}
|
||
None => {
|
||
relay_observer_control_rx = None;
|
||
tracing::warn!("relay observer control channel closed");
|
||
}
|
||
}
|
||
None
|
||
}
|
||
// Remaining branches don't touch pool — evaluated when pool is idle.
|
||
buzz_event = relay.next_event() => {
|
||
let _ = result_rx; // end split borrow before relay handling
|
||
match buzz_event {
|
||
Some(buzz_event) => {
|
||
let kind_u32 = buzz_event.event.kind.as_u16() as u32;
|
||
|
||
if kind_u32 == KIND_MEMBER_ADDED_NOTIFICATION
|
||
|| kind_u32 == KIND_MEMBER_REMOVED_NOTIFICATION
|
||
{
|
||
let ch = buzz_event.channel_id;
|
||
let ts = buzz_event.event.created_at.as_secs();
|
||
let eid = buzz_event.event.id.to_hex();
|
||
|
||
// Two-layer membership dedup:
|
||
//
|
||
// 1. Exact duplicate rejection (seen_membership_ids):
|
||
// Catches the same event replayed on reconnect.
|
||
//
|
||
// 2. Timestamp watermark (membership_newest_ts):
|
||
// Uses strict `<` so that older events from reconnect
|
||
// replay are dropped, but legitimate live events at the
|
||
// same second are both processed. This is safe because
|
||
// exact duplicates are already caught by layer 1.
|
||
//
|
||
// Why not `<=`? That would suppress legitimate live
|
||
// add→remove (or remove→add) sequences in the same second,
|
||
// leaving the harness in the wrong membership state.
|
||
// Two-generation dedup: check both sets before inserting.
|
||
if seen_membership_current.contains(&eid)
|
||
|| seen_membership_previous.contains(&eid)
|
||
{
|
||
tracing::debug!(
|
||
channel_id = %ch,
|
||
kind = kind_u32,
|
||
"skipping duplicate membership notification (same event_id)"
|
||
);
|
||
continue;
|
||
}
|
||
seen_membership_current.insert(eid);
|
||
// Rotate at 1000: current → previous, no amnesia window.
|
||
if seen_membership_current.len() >= 1000 {
|
||
seen_membership_previous =
|
||
std::mem::take(&mut seen_membership_current);
|
||
}
|
||
if let Some(&newest) = membership_newest_ts.get(&ch) {
|
||
if ts < newest {
|
||
tracing::debug!(
|
||
channel_id = %ch,
|
||
kind = kind_u32,
|
||
ts,
|
||
newest,
|
||
"skipping stale membership notification (older than newest)"
|
||
);
|
||
continue;
|
||
}
|
||
}
|
||
membership_newest_ts.insert(ch, ts);
|
||
|
||
if kind_u32 == KIND_MEMBER_ADDED_NOTIFICATION {
|
||
// Clear removal tracking so sessions are not
|
||
// stripped for a legitimately re-added channel.
|
||
removed_channels.remove(&ch);
|
||
|
||
if subscribed_channel_ids.contains(&ch) {
|
||
tracing::debug!(channel_id = %ch, "membership notification: channel already subscribed");
|
||
} else if let Some(filter) = config::resolve_dynamic_channel_filter(&config, ch, &rules) {
|
||
tracing::info!(channel_id = %ch, "membership notification: subscribing to new channel");
|
||
if let Err(e) = relay.subscribe_channel_from(ch, filter, Some(ts)).await {
|
||
tracing::warn!("failed to subscribe to new channel {ch}: {e}");
|
||
} else {
|
||
subscribed_channel_ids.insert(ch);
|
||
}
|
||
} else {
|
||
tracing::debug!(channel_id = %ch, "membership notification: no matching rules — skipping");
|
||
}
|
||
} else {
|
||
subscribed_channel_ids.remove(&ch);
|
||
tracing::info!(channel_id = %ch, "membership notification: unsubscribing from channel");
|
||
if let Err(e) = relay.unsubscribe_channel(ch).await {
|
||
tracing::warn!("failed to unsubscribe from channel {ch}: {e}");
|
||
}
|
||
// Drain queued events and invalidate sessions for the
|
||
// removed channel. Events already in-flight will
|
||
// complete normally (the relay may reject actions if
|
||
// the agent lost access).
|
||
let drained_ids = queue.drain_channel(ch);
|
||
let invalidated = pool.invalidate_channel_sessions(ch);
|
||
// Track removed channels so checked-out agents get
|
||
// their sessions stripped when they return to the pool.
|
||
removed_channels.insert(ch);
|
||
typing_channels.remove(&ch);
|
||
// Best-effort: clean up 👀 on drained events.
|
||
// Note: the relay revokes membership before
|
||
// emitting the notification, so this DELETE may
|
||
// 403 on non-open channels. Stale 👀 in that
|
||
// case is a known limitation — fix belongs in
|
||
// the relay (clean up bot reactions on removal).
|
||
if !drained_ids.is_empty() {
|
||
let rc = ctx.rest_client.clone();
|
||
let ids = drained_ids.clone();
|
||
tokio::spawn(async move {
|
||
for eid in &ids {
|
||
pool::reaction_remove(&rc, eid, "👀").await;
|
||
}
|
||
});
|
||
}
|
||
if !drained_ids.is_empty() || invalidated > 0 {
|
||
tracing::info!(
|
||
channel_id = %ch,
|
||
drained = drained_ids.len(),
|
||
invalidated,
|
||
"cleaned up after membership removal"
|
||
);
|
||
}
|
||
}
|
||
continue;
|
||
}
|
||
|
||
if config.ignore_self && buzz_event.event.pubkey.to_hex() == pubkey_hex {
|
||
tracing::debug!(channel_id = %buzz_event.channel_id, "dropping self-authored event");
|
||
continue;
|
||
}
|
||
|
||
// Check: kind:9, content "!shutdown", from owner, mentions THIS agent.
|
||
let is_shutdown = is_owner_control_command(
|
||
&buzz_event.event,
|
||
kind_u32,
|
||
"!shutdown",
|
||
&pubkey_hex,
|
||
);
|
||
if is_shutdown {
|
||
let owner = owner_cache.get();
|
||
if let Some(owner) = owner {
|
||
if buzz_event.event.pubkey.to_hex() == *owner {
|
||
tracing::info!(
|
||
channel_id = %buzz_event.channel_id,
|
||
sender = %buzz_event.event.pubkey.to_hex(),
|
||
"shutdown command from owner — exiting gracefully"
|
||
);
|
||
let _ = shutdown_tx.send(());
|
||
continue;
|
||
}
|
||
}
|
||
// Not from owner — fall through to normal prompt handling.
|
||
// Don't drop it — it's a regular message that happens to
|
||
// contain "!shutdown" from a non-owner.
|
||
}
|
||
|
||
// Mirrors !shutdown: kind:9, content "!cancel", from
|
||
// owner, mentions THIS agent. Must be BEFORE
|
||
// queue.push() — the event content is moved by push.
|
||
//
|
||
// Mode-independent: !cancel fires regardless of
|
||
// --multiple-event-handling. It is explicit user
|
||
// intent, not an automatic policy decision.
|
||
let is_cancel = is_owner_control_command(
|
||
&buzz_event.event,
|
||
kind_u32,
|
||
"!cancel",
|
||
&pubkey_hex,
|
||
);
|
||
if is_cancel {
|
||
if let Some(owner) = owner_cache.get() {
|
||
if buzz_event.event.pubkey.to_hex() == *owner {
|
||
let fired = signal_in_flight_task(
|
||
&mut pool,
|
||
buzz_event.channel_id,
|
||
ControlSignal::Cancel,
|
||
);
|
||
if !fired {
|
||
tracing::warn!(
|
||
channel_id = %buzz_event.channel_id,
|
||
"!cancel received but no in-flight task — no-op"
|
||
);
|
||
}
|
||
continue; // consume event — do NOT push to queue
|
||
}
|
||
}
|
||
// Not from owner — fall through to normal prompt handling.
|
||
}
|
||
|
||
// Mirrors !shutdown / !cancel: kind:9, content
|
||
// "!rotate", from owner, mentions THIS agent.
|
||
//
|
||
// Rotation is explicit owner intent to start the
|
||
// next turn in this channel with a fresh ACP
|
||
// session. It is consumed by the harness and never
|
||
// forwarded to the agent. If a turn is in-flight,
|
||
// cancel it, drop its triggering batch, and
|
||
// invalidate the channel session when the task
|
||
// returns. If idle, invalidate the cached channel
|
||
// session immediately. Queued future events remain
|
||
// queued and will create a fresh session on dispatch.
|
||
let is_rotate = is_owner_control_command(
|
||
&buzz_event.event,
|
||
kind_u32,
|
||
"!rotate",
|
||
&pubkey_hex,
|
||
);
|
||
if is_rotate {
|
||
if let Some(owner) = owner_cache.get() {
|
||
if buzz_event.event.pubkey.to_hex() == *owner {
|
||
let fired = signal_in_flight_task(
|
||
&mut pool,
|
||
buzz_event.channel_id,
|
||
ControlSignal::Rotate,
|
||
);
|
||
if fired {
|
||
tracing::info!(
|
||
channel_id = %buzz_event.channel_id,
|
||
"!rotate received — cancelling in-flight turn and rotating session"
|
||
);
|
||
} else {
|
||
let invalidated = pool.invalidate_channel_sessions(buzz_event.channel_id);
|
||
tracing::info!(
|
||
channel_id = %buzz_event.channel_id,
|
||
invalidated,
|
||
"!rotate received — invalidated idle channel session(s)"
|
||
);
|
||
}
|
||
continue; // consume event — do NOT push to queue
|
||
}
|
||
}
|
||
// Not from owner — fall through to normal prompt handling.
|
||
}
|
||
|
||
// Coarse security policy: drop events from disallowed
|
||
// authors before they reach subscription rules or the
|
||
// agent. Must be AFTER !shutdown (owner can always
|
||
// shut down regardless of gate mode).
|
||
//
|
||
// Both OwnerOnly and Allowlist accept events from
|
||
// "siblings" — pubkeys whose agent_owner_pubkey
|
||
// matches this agent's owner (e.g. other bots
|
||
// launched by the same human). Allowlist adds the
|
||
// explicit pubkey list on top, for external people;
|
||
// it never revokes same-owner team bots.
|
||
{
|
||
let author = buzz_event.event.pubkey.to_hex();
|
||
let allowed = author_allowed(
|
||
&config.respond_to,
|
||
&config.respond_to_allowlist,
|
||
&author,
|
||
&owner_cache,
|
||
&ctx.rest_client,
|
||
)
|
||
.await;
|
||
if !allowed {
|
||
tracing::debug!(
|
||
channel_id = %buzz_event.channel_id,
|
||
author = %buzz_event.event.pubkey.to_hex(),
|
||
mode = %config.respond_to,
|
||
"inbound author gate — dropping event"
|
||
);
|
||
continue;
|
||
}
|
||
}
|
||
|
||
let matched = filter::match_event(&buzz_event.event, buzz_event.channel_id, &rules, &pubkey_hex).await;
|
||
let prompt_tag = match matched {
|
||
Some(m) => m.prompt_tag,
|
||
None => {
|
||
tracing::debug!(channel_id = %buzz_event.channel_id, kind = buzz_event.event.kind.as_u16(), "event matched no rule — dropping");
|
||
continue;
|
||
}
|
||
};
|
||
// Capture author pubkey before queue.push() moves
|
||
// buzz_event.event (needed for mode gate below).
|
||
let author_hex = buzz_event.event.pubkey.to_hex();
|
||
let event_id_hex = buzz_event.event.id.to_hex();
|
||
// Clone for the non-cancelling steer fork, which
|
||
// needs the event to render the steer body. The
|
||
// clone is unconditional because we don't know
|
||
// yet whether the mode gate will demand a steer
|
||
// — checking `multiple_event_handling` here
|
||
// would couple the queueing path to the mode
|
||
// and break the existing invariant that every
|
||
// accepted event goes through `queue.push`
|
||
// first. `nostr::Event::clone` is cheap (Arc-
|
||
// backed payload) so the cost is negligible.
|
||
let event_for_steer = buzz_event.event.clone();
|
||
let prompt_tag_for_steer = prompt_tag.clone();
|
||
let accepted = queue.push(QueuedEvent {
|
||
channel_id: buzz_event.channel_id,
|
||
event: buzz_event.event,
|
||
received_at: std::time::Instant::now(),
|
||
prompt_tag,
|
||
});
|
||
// 👀 — immediate "seen" reaction, only if the event
|
||
// was actually queued (not dropped by DedupMode::Drop).
|
||
// Fire-and-forget: on rare fast-failure paths the
|
||
// guard's cleanup may race with this add, leaving a
|
||
// cosmetic stale 👀. Acceptable — see ReactionGuard docs.
|
||
if accepted {
|
||
let rc = ctx.rest_client.clone();
|
||
let eid = event_id_hex.clone();
|
||
tokio::spawn(async move {
|
||
pool::reaction_add(&rc, &eid, "👀").await;
|
||
});
|
||
}
|
||
// Event is already queued. If mode requires it AND
|
||
// the channel has an in-flight task, fire cancel —
|
||
// OR take the non-cancelling (ACP steer) fork for Steer signals.
|
||
if accepted && queue.is_channel_in_flight(buzz_event.channel_id) {
|
||
// Author eligibility (owner ∪ allowlist ∪ siblings)
|
||
// is already enforced by the inbound author gate
|
||
// above, so the mid-turn signal fires for every
|
||
// event that reaches here.
|
||
let signal = mode_gate_signal(
|
||
config.multiple_event_handling,
|
||
&author_hex,
|
||
owner_cache.get(),
|
||
);
|
||
if let Some(signal) = signal {
|
||
// Try-and-tolerate fork: when the mode
|
||
// wants a Steer, attempt the non-cancelling
|
||
// path first for any agent. On accept,
|
||
// withhold the queued event and spawn an
|
||
// ack watcher; the main loop's
|
||
// `PoolEvent::SteerAck` arm decides
|
||
// success/release/fallback. On reject
|
||
// (including `-32601 method_not_found`
|
||
// from agents that don't implement the
|
||
// extension), fall through to the universal
|
||
// cancel+merge `Steer` signal so the event
|
||
// still reaches the agent.
|
||
let native_attempted = matches!(signal, ControlSignal::Steer)
|
||
&& try_native_steer(
|
||
&mut pool,
|
||
&mut queue,
|
||
buzz_event.channel_id,
|
||
event_for_steer,
|
||
prompt_tag_for_steer,
|
||
&steer_ack_tx,
|
||
);
|
||
if !native_attempted {
|
||
signal_in_flight_task(
|
||
&mut pool,
|
||
buzz_event.channel_id,
|
||
signal,
|
||
);
|
||
}
|
||
}
|
||
}
|
||
for (channel_id, thread_tags) in
|
||
dispatch_pending(&mut pool, &mut queue, &ctx)
|
||
{
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
None => {
|
||
tracing::warn!("relay event stream ended — requesting reconnect");
|
||
if let Err(e) = relay.reconnect().await {
|
||
tracing::error!("relay background task is gone: {e} — exiting");
|
||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
None
|
||
}
|
||
_ = async {
|
||
match heartbeat.as_mut() {
|
||
Some(hb) => hb.tick().await,
|
||
None => std::future::pending().await,
|
||
}
|
||
} => {
|
||
let _ = result_rx;
|
||
if queue.has_flushable_work() {
|
||
tracing::debug!("heartbeat_skipped_events");
|
||
for (channel_id, thread_tags) in
|
||
dispatch_pending(&mut pool, &mut queue, &ctx)
|
||
{
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
} else if pool.any_idle() {
|
||
dispatch_heartbeat(&mut pool, &ctx, &mut heartbeat_in_flight);
|
||
} else {
|
||
tracing::debug!("heartbeat_skipped_busy");
|
||
}
|
||
None
|
||
}
|
||
_ = async {
|
||
match presence_heartbeat.as_mut() {
|
||
Some(t) => t.tick().await,
|
||
None => std::future::pending().await,
|
||
}
|
||
} => {
|
||
let _ = result_rx;
|
||
// Abort previous heartbeat if still in flight (prevents race on shutdown).
|
||
if let Some(h) = presence_task.take() {
|
||
h.abort();
|
||
}
|
||
let pp = presence_publisher.clone();
|
||
let pk = presence_keys.clone();
|
||
presence_task = Some(tokio::spawn(async move {
|
||
if let Err(e) = publish_presence(&pp, &pk, "online").await {
|
||
tracing::warn!("presence heartbeat failed: {e}");
|
||
}
|
||
}));
|
||
None
|
||
}
|
||
_ = async {
|
||
match typing_refresh.as_mut() {
|
||
Some(t) => t.tick().await,
|
||
None => std::future::pending().await,
|
||
}
|
||
} => {
|
||
let _ = result_rx;
|
||
// Use try_publish (non-blocking) for typing indicators —
|
||
// they're ephemeral and must not block the main loop during
|
||
// relay reconnection (#35).
|
||
for (&ch, thread_tags) in &typing_channels {
|
||
if let Ok(event) = relay.build_typing_event(
|
||
ch,
|
||
thread_tags.root_event_id.as_deref(),
|
||
thread_tags.parent_event_id.as_deref(),
|
||
) {
|
||
if let Err(e) = relay.try_publish_event(event) {
|
||
tracing::debug!("typing indicator dropped for {ch}: {e}");
|
||
}
|
||
}
|
||
}
|
||
None
|
||
}
|
||
_ = shutdown_rx.changed() => {
|
||
tracing::info!("shutting down");
|
||
break;
|
||
}
|
||
}
|
||
};
|
||
|
||
match pool_event {
|
||
Some(PoolEvent::Result(result)) => {
|
||
// Stop typing indicator for the completed channel.
|
||
if let PromptSource::Channel(ch) = &result.source {
|
||
typing_channels.remove(ch);
|
||
}
|
||
if handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
*result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
observer.clone(),
|
||
Some(&ctx.rest_client),
|
||
) == LoopAction::Exit
|
||
{
|
||
break;
|
||
}
|
||
if drain_ready_join_results(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut typing_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
observer.clone(),
|
||
) == LoopAction::Exit
|
||
{
|
||
break;
|
||
}
|
||
for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) {
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
Some(PoolEvent::Panic(join_error)) => {
|
||
tracing::error!("agent task panicked: {join_error}");
|
||
recover_panicked_agent(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
join_error,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut typing_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
observer.clone(),
|
||
);
|
||
if pool.live_count() == 0 && !any_respawn_in_flight(&crash_history) {
|
||
tracing::error!("all agents dead — exiting");
|
||
break;
|
||
}
|
||
for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) {
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
Some(PoolEvent::SteerAck(SteerAckEvent {
|
||
channel_id,
|
||
event_id,
|
||
ack,
|
||
})) => {
|
||
// Goose-native steer attempt resolved. Locked semantics
|
||
// (Eva + Max + Perci, unanimous on Option X):
|
||
//
|
||
// Success
|
||
// The agent received the steer via the non-cancelling
|
||
// path. Drop the withheld event so normal dispatch
|
||
// never redelivers it.
|
||
//
|
||
// Err(_) where the write never landed (Transport /
|
||
// ExpectedRunIdMissing):
|
||
// Delivery state of the underlying message is "never
|
||
// attempted on the wire". Release withheld back to the
|
||
// queue front AND issue the cancel+merge fallback so
|
||
// the message still reaches the agent.
|
||
//
|
||
// Err(AgentError { code: -32601, .. })
|
||
// The agent returned method_not_found — it does not
|
||
// implement the steer extension. Release withheld AND
|
||
// fire the cancel+merge fallback so the message still
|
||
// reaches the agent via the universal path.
|
||
//
|
||
// Err(AgentError { code: other, .. })
|
||
// The write landed and the agent returned a JSON-RPC
|
||
// error at the application level (e.g. wrong run id).
|
||
// The agent's turn is still running (or just completed).
|
||
// Release withheld for normal dispatch; do NOT fire the
|
||
// fallback signal — the agent already saw the steer
|
||
// attempt. If the turn is still running, normal dispatch
|
||
// re-delivers when it completes. If the turn already
|
||
// ended, there is nothing to cancel.
|
||
//
|
||
// PromptCompletedNeutral
|
||
// The read loop wrote the steer (or was preparing to)
|
||
// but the prompt completed before the response landed.
|
||
// Delivery state is unknown — but the prompt completing
|
||
// means there is no in-flight turn to signal anymore.
|
||
// Release withheld for normal dispatch; do NOT fire
|
||
// the fallback signal (it would target a turn that
|
||
// just ended; normal dispatch already handles
|
||
// redelivery via the released queue entry).
|
||
//
|
||
// Err(PromptCompleted)
|
||
// `SteerError::PromptCompleted` is returned synchronously
|
||
// by `pool::send_steer` when no task is in flight (handled
|
||
// in `try_native_steer`'s Err branch, which falls through
|
||
// to cancel+merge). It is never routed through the ack
|
||
// channel, so this variant never appears in `SteerAckEvent`.
|
||
//
|
||
// Watcher Err (oneshot dropped)
|
||
// Should not happen — the read loop drains
|
||
// pending_steer on every return path. If it does,
|
||
// treat as PromptCompletedNeutral to avoid leaking
|
||
// the withheld event in `withheld_native_steer`.
|
||
let (release_withheld, drop_withheld, signal_fallback) = match &ack {
|
||
Ok(pool::SteerAck::Success) => (false, true, false),
|
||
// -32601 = method_not_found: agent does not implement the
|
||
// steer extension. Fire cancel+merge so the message still
|
||
// reaches the agent.
|
||
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { code, .. }))
|
||
if *code == -32601 =>
|
||
{
|
||
(true, false, true)
|
||
}
|
||
// AgentError: write landed, agent rejected it at the
|
||
// application level (e.g. wrong run id). Release for
|
||
// normal dispatch; no fallback signal (the turn is still
|
||
// running or just ended — either way there is nothing to
|
||
// cancel).
|
||
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { .. })) => {
|
||
(true, false, false)
|
||
}
|
||
// Transport / ExpectedRunIdMissing: write never landed.
|
||
// Release and fire the cancel+merge fallback so the
|
||
// message still reaches the agent.
|
||
Ok(pool::SteerAck::Err(_)) => (true, false, true),
|
||
Ok(pool::SteerAck::PromptCompletedNeutral) => (true, false, false),
|
||
Err(_recv_err) => (true, false, false),
|
||
};
|
||
tracing::info!(
|
||
channel = %channel_id,
|
||
event_id = %event_id,
|
||
?ack,
|
||
release_withheld,
|
||
drop_withheld,
|
||
signal_fallback,
|
||
"non-cancelling steer ack received"
|
||
);
|
||
if drop_withheld {
|
||
queue.remove_event(channel_id, &event_id);
|
||
}
|
||
if release_withheld {
|
||
queue.release_native_steer(channel_id, &event_id);
|
||
}
|
||
if signal_fallback {
|
||
// Universal cancel+merge fallback. Note: the
|
||
// queued event has already been released to the
|
||
// front of `queues[channel_id]`, so the cancel
|
||
// will pick it up as part of the merged batch and
|
||
// re-prompt the agent.
|
||
signal_in_flight_task(&mut pool, channel_id, ControlSignal::Steer);
|
||
}
|
||
// After releasing a withheld event, give dispatch a chance
|
||
// to re-flush. If the prompt is still in flight, the
|
||
// channel stays `in_flight_channels` and `flush_next`
|
||
// skips it — but a Steer fallback signal sent above will
|
||
// tear down the in-flight task; on its completion the
|
||
// queue drains. We still try here in case the in-flight
|
||
// task has already returned.
|
||
for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) {
|
||
typing_channels.insert(channel_id, thread_tags);
|
||
}
|
||
}
|
||
None => {} // relay/heartbeat/shutdown branches handled inline above
|
||
}
|
||
}
|
||
|
||
tracing::info!("shutdown: waiting for in-flight prompts");
|
||
// 30 s is generous for in-flight prompts to be cancelled; using
|
||
// max_turn_duration here would cause Ctrl+C to hang for up to an hour.
|
||
let grace = Duration::from_secs(30);
|
||
// Best-effort drain of both join_set and result_rx during the grace period.
|
||
// Tasks that finish normally send their OwnedAgent through result_rx — we
|
||
// explicitly shut them down here to reap child processes. If the grace
|
||
// period expires, remaining tasks are aborted and fall back to
|
||
// AcpClient::Drop (start_kill + try_wait — best-effort, not guaranteed).
|
||
let (rx_ref, js_ref) = pool.rx_and_join_set();
|
||
let shutdown_result = tokio::time::timeout(grace, async {
|
||
loop {
|
||
tokio::select! {
|
||
result = js_ref.join_next() => {
|
||
match result {
|
||
Some(Err(e)) => tracing::warn!("task error during shutdown: {e}"),
|
||
Some(Ok(())) => {}
|
||
None => break, // join_set empty
|
||
}
|
||
}
|
||
maybe_result = rx_ref.recv() => {
|
||
if let Some(mut pr) = maybe_result {
|
||
let idx = pr.agent.index;
|
||
pr.agent.acp.shutdown().await;
|
||
tracing::debug!(agent = idx, "reaped checked-out agent on shutdown");
|
||
}
|
||
// If None, channel closed — tasks are done.
|
||
}
|
||
}
|
||
}
|
||
})
|
||
.await;
|
||
if shutdown_result.is_err() {
|
||
tracing::warn!("grace period expired, aborting remaining tasks");
|
||
pool.join_set.shutdown().await;
|
||
}
|
||
// Drain any remaining results that arrived after join_set drained but
|
||
// before tasks were aborted.
|
||
while let Ok(mut pr) = pool.result_rx_try_recv() {
|
||
let idx = pr.agent.index;
|
||
pr.agent.acp.shutdown().await;
|
||
tracing::debug!(agent = idx, "reaped late-arriving agent on shutdown");
|
||
}
|
||
// Explicitly shut down idle agents still sitting in their slots.
|
||
for slot in pool.agents_mut().iter_mut() {
|
||
if let Some(agent) = slot.take() {
|
||
let idx = agent.index;
|
||
let mut acp = agent.acp;
|
||
acp.shutdown().await;
|
||
tracing::debug!(agent = idx, "reaped idle agent on shutdown");
|
||
}
|
||
}
|
||
drop(pool);
|
||
|
||
// Abort any in-flight respawn tasks. They may be sleeping in backoff or
|
||
// running spawn_and_init — either way, we don't want them spawning new
|
||
// children after the main loop has exited. RespawnGuard::Drop sends a
|
||
// failure result for aborted tasks, so respawn_in_flight is cleared.
|
||
respawn_tasks.shutdown().await;
|
||
|
||
// Drain any respawn results that completed before the abort. Explicitly
|
||
// shut down returned agents instead of relying on AcpClient::Drop.
|
||
while let Ok(rr) = respawn_rx.try_recv() {
|
||
if let Ok((mut acp, _, _)) = rr.result {
|
||
acp.shutdown().await;
|
||
tracing::debug!(agent = rr.index, "reaped respawned agent on shutdown");
|
||
}
|
||
}
|
||
|
||
// Cancel any in-flight presence heartbeat before sending offline.
|
||
if let Some(h) = presence_task.take() {
|
||
h.abort();
|
||
}
|
||
|
||
// Best-effort: set presence to offline before exiting.
|
||
if config.presence_enabled {
|
||
match tokio::time::timeout(
|
||
Duration::from_secs(2),
|
||
publish_presence(&presence_publisher, &presence_keys, "offline"),
|
||
)
|
||
.await
|
||
{
|
||
Ok(Ok(_)) => tracing::info!("presence set to offline"),
|
||
Ok(Err(e)) => tracing::warn!("failed to set offline presence: {e}"),
|
||
Err(_) => tracing::warn!("offline presence timed out"),
|
||
}
|
||
}
|
||
|
||
if let Some(handle) = relay_observer_publisher_task.take() {
|
||
handle.abort();
|
||
}
|
||
|
||
// Graceful relay shutdown — sends WebSocket close frame and waits up to 5s
|
||
// for the background task to finish, rather than aborting immediately (#40).
|
||
relay.shutdown().await;
|
||
|
||
tracing::info!("buzz-acp stopped");
|
||
Ok(())
|
||
}
|
||
|
||
#[derive(PartialEq)]
|
||
enum LoopAction {
|
||
Continue,
|
||
Exit,
|
||
}
|
||
|
||
fn event_mentions_agent(event: &nostr::Event, agent_pubkey_hex: &str) -> bool {
|
||
event.tags.iter().any(|t| {
|
||
t.as_slice().first().map(|s| s.as_str()) == Some("p")
|
||
&& t.as_slice().get(1).map(|s| s.as_str()) == Some(agent_pubkey_hex)
|
||
})
|
||
}
|
||
|
||
fn is_owner_control_command(
|
||
event: &nostr::Event,
|
||
kind_u32: u32,
|
||
command: &str,
|
||
agent_pubkey_hex: &str,
|
||
) -> bool {
|
||
kind_u32 == KIND_STREAM_MESSAGE
|
||
&& event.content.trim() == command
|
||
&& event_mentions_agent(event, agent_pubkey_hex)
|
||
}
|
||
|
||
// ── signal_in_flight_task ─────────────────────────────────────────────────────
|
||
|
||
/// Decide which [`ControlSignal`] (if any) to send to an in-flight turn when a
|
||
/// new, already-author-gated event arrives for that channel.
|
||
///
|
||
/// Returns `None` to leave the in-flight turn untouched (the event waits in the
|
||
/// queue and is delivered when the turn completes). Author eligibility — owner
|
||
/// ∪ allowlist ∪ siblings — is enforced upstream by the inbound author gate, so
|
||
/// `Steer`/`Interrupt` apply to every event that reaches this point; only
|
||
/// `OwnerInterrupt` re-checks authorship (owner-only) here.
|
||
///
|
||
/// `owner` is the resolved owner pubkey hex, if known.
|
||
fn mode_gate_signal(
|
||
handling: MultipleEventHandling,
|
||
author_hex: &str,
|
||
owner: Option<&str>,
|
||
) -> Option<ControlSignal> {
|
||
match handling {
|
||
MultipleEventHandling::Queue => None,
|
||
MultipleEventHandling::Steer => Some(ControlSignal::Steer),
|
||
MultipleEventHandling::Interrupt => Some(ControlSignal::Interrupt),
|
||
MultipleEventHandling::OwnerInterrupt => match owner {
|
||
Some(o) if author_hex == o => Some(ControlSignal::Interrupt),
|
||
_ => None,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// Send a control signal to the in-flight task for `channel_id`.
|
||
/// Returns `true` if a signal was sent, `false` if no in-flight task was found.
|
||
fn signal_in_flight_task(
|
||
pool: &mut AgentPool,
|
||
channel_id: uuid::Uuid,
|
||
mode: ControlSignal,
|
||
) -> bool {
|
||
let entry = pool
|
||
.task_map_mut()
|
||
.values_mut()
|
||
.find(|m| m.channel_id == Some(channel_id));
|
||
|
||
if let Some(meta) = entry {
|
||
if let Some(tx) = meta.control_tx.take() {
|
||
tracing::info!(channel = %channel_id, ?mode, "control signal sent to in-flight task");
|
||
let _ = tx.send(mode);
|
||
return true;
|
||
}
|
||
}
|
||
false
|
||
}
|
||
|
||
/// Attempt the non-cancelling (ACP) steer for a freshly-queued event.
|
||
///
|
||
/// Caller invariants:
|
||
/// - `event` has already been pushed into `EventQueue::queues[channel_id]`
|
||
/// via [`EventQueue::push`] — its `event.id` must still be locatable
|
||
/// there so [`EventQueue::mark_native_steer_pending`] can move it to the
|
||
/// side table.
|
||
/// - `multiple_event_handling` resolved to `ControlSignal::Steer`; this
|
||
/// function is the non-cancelling fork of that signal.
|
||
///
|
||
/// Returns `true` if the native attempt was accepted by the read loop
|
||
/// (capacity-1 mpsc `try_send` succeeded, event withheld synchronously,
|
||
/// ack watcher spawned). On `true` the caller MUST NOT issue the
|
||
/// universal cancel+merge `ControlSignal::Steer` fallback — the watcher
|
||
/// will issue it from the ack arm if the native attempt fails.
|
||
///
|
||
/// Returns `false` if `pool.send_steer` failed (no in-flight task,
|
||
/// `steer_tx` already full from a prior in-flight steer, or read loop
|
||
/// torn down). The caller MUST fall through to
|
||
/// `signal_in_flight_task(channel_id, ControlSignal::Steer)` so the
|
||
/// event still reaches the agent via the universal path.
|
||
///
|
||
/// The withheld event is NOT released here on `false` because no withhold
|
||
/// was established: `mark_native_steer_pending` only runs on `Ok(())`.
|
||
fn try_native_steer(
|
||
pool: &mut AgentPool,
|
||
queue: &mut EventQueue,
|
||
channel_id: uuid::Uuid,
|
||
event: nostr::Event,
|
||
prompt_tag: String,
|
||
steer_ack_tx: &mpsc::UnboundedSender<SteerAckEvent>,
|
||
) -> bool {
|
||
// Build the steer body: framing strings come from
|
||
// `queue::native_steer_framing()` (Eva's drift-proof requirement —
|
||
// native and cancel+merge fallback share these so the agent gets the
|
||
// same orientation regardless of transport). The single event block
|
||
// is rendered by `queue::format_event_block`, the same function
|
||
// `queue::format_prompt` uses internally for `[Buzz event: …]`
|
||
// sections, so the rendering also cannot drift.
|
||
//
|
||
// Passing `None` for `channel_info` / `profile_lookup` is intentional:
|
||
// native steer is a *delta* into a live turn — the agent already saw
|
||
// channel context and the actor's profile in the original prompt,
|
||
// duplicating it here would defeat the point of non-cancelling
|
||
// steering (which is to inject only what's new).
|
||
let (header, closing) = queue::native_steer_framing();
|
||
let event_id_hex = event.id.to_hex();
|
||
let be = queue::BatchEvent {
|
||
event,
|
||
prompt_tag: prompt_tag.clone(),
|
||
received_at: std::time::Instant::now(),
|
||
};
|
||
let event_block = queue::format_event_block(channel_id, None, &be, None);
|
||
let body = format!("{header}\n\n[Buzz event: {prompt_tag}]\n{event_block}\n\n{closing}");
|
||
|
||
let (ack_tx, ack_rx) = tokio::sync::oneshot::channel::<pool::SteerAck>();
|
||
let request = pool::SteerRequest {
|
||
prompt_blocks: vec![body],
|
||
ack_tx,
|
||
};
|
||
|
||
match pool.send_steer(channel_id, request) {
|
||
Ok(()) => {
|
||
// Withhold the queued event synchronously BEFORE spawning
|
||
// the watcher: this closes the race where `mark_complete`
|
||
// clears `in_flight_channels` and a stray `flush_next` could
|
||
// re-deliver the event via normal dispatch. See
|
||
// `EventQueue::mark_native_steer_pending` docs at queue.rs:606.
|
||
let withheld = queue.mark_native_steer_pending(channel_id, &event_id_hex);
|
||
if !withheld {
|
||
// Race: the event was already drained out of the queue
|
||
// before we got here (e.g. a concurrent flush picked it
|
||
// up). The steer is on the wire; if it succeeds the
|
||
// agent gets it via the native path AND normal
|
||
// dispatch — duplicate delivery is benign (agent gets
|
||
// the same message twice). Log so this is visible if it
|
||
// ever happens in production.
|
||
tracing::warn!(
|
||
channel = %channel_id,
|
||
event_id = %event_id_hex,
|
||
"native steer accepted by read loop but event was not in queue to withhold \
|
||
— possible duplicate delivery if steer succeeds"
|
||
);
|
||
}
|
||
let ack_tx_clone = steer_ack_tx.clone();
|
||
let event_id_for_watcher = event_id_hex.clone();
|
||
tokio::spawn(async move {
|
||
let ack = ack_rx.await;
|
||
let _ = ack_tx_clone.send(SteerAckEvent {
|
||
channel_id,
|
||
event_id: event_id_for_watcher,
|
||
ack,
|
||
});
|
||
});
|
||
true
|
||
}
|
||
Err(e) => {
|
||
tracing::info!(
|
||
channel = %channel_id,
|
||
error = ?e,
|
||
"non-cancelling steer not accepted — falling back to cancel+merge"
|
||
);
|
||
false
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── dispatch_pending ──────────────────────────────────────────────────────────
|
||
|
||
/// Flush queued work to available agents.
|
||
fn dispatch_pending(
|
||
pool: &mut AgentPool,
|
||
queue: &mut EventQueue,
|
||
ctx: &Arc<PromptContext>,
|
||
) -> Vec<(Uuid, ThreadTags)> {
|
||
let mut dispatched_channels = Vec::new();
|
||
loop {
|
||
let batch = match queue.flush_next() {
|
||
Some(b) => b,
|
||
None => break,
|
||
};
|
||
let channel_id = batch.channel_id;
|
||
let typing_scope = batch
|
||
.events
|
||
.last()
|
||
.map(|event| queue::parse_thread_tags(&event.event))
|
||
.unwrap_or_default();
|
||
let affinity_hit = pool.has_session_for(channel_id);
|
||
let mut agent = match pool.try_claim(Some(channel_id)) {
|
||
Some(a) => a,
|
||
None => {
|
||
let pending = queue.pending_channels();
|
||
tracing::debug!(pending_channels = pending, "pool_exhausted");
|
||
queue.requeue_preserve_timestamps(batch);
|
||
queue.mark_complete(channel_id);
|
||
break;
|
||
}
|
||
};
|
||
tracing::debug!(agent = agent.index, channel = %channel_id, affinity_hit, "agent_claimed");
|
||
|
||
let recoverable_batch = match ctx.dedup_mode {
|
||
DedupMode::Queue => Some(batch.clone()),
|
||
DedupMode::Drop => None,
|
||
};
|
||
|
||
let result_tx = pool.result_tx();
|
||
let ctx_clone = Arc::clone(ctx);
|
||
let agent_index = agent.index;
|
||
|
||
// Goose-native non-cancelling steer seam: snapshot capability before
|
||
// the agent moves into `run_prompt_task`, and install the per-turn
|
||
// steer receiver on the read loop so the main loop's mode-gate fork
|
||
// (see the `if accepted && queue.is_channel_in_flight(...)` block
|
||
// in the relay event branch of the main `select!` loop) can drive
|
||
// it via the matching sender stored in `TaskMeta.steer_tx`.
|
||
// Install the steer channel for every prompt task — the supervisor
|
||
// uses try-and-tolerate: it attempts the steer for any agent and
|
||
// treats `-32601 method_not_found` as "fall back to cancel+merge".
|
||
let (tx, rx) = tokio::sync::mpsc::channel::<pool::SteerRequest>(1);
|
||
agent.acp.install_steer_rx(rx);
|
||
let steer_tx = Some(tx);
|
||
|
||
// Prompt text is now built inside run_prompt_task (needs async for
|
||
// context fetching). Pass None for prompt_text; batch carries the data.
|
||
let (control_tx, control_rx) = tokio::sync::oneshot::channel::<ControlSignal>();
|
||
|
||
let abort_handle = pool.join_set.spawn(async move {
|
||
pool::run_prompt_task(
|
||
agent,
|
||
Some(batch),
|
||
None,
|
||
ctx_clone,
|
||
result_tx,
|
||
Some(control_rx),
|
||
)
|
||
.await;
|
||
});
|
||
|
||
pool.task_map_mut().insert(
|
||
abort_handle.id(),
|
||
pool::TaskMeta {
|
||
agent_index,
|
||
channel_id: Some(channel_id),
|
||
recoverable_batch,
|
||
control_tx: Some(control_tx),
|
||
steer_tx,
|
||
},
|
||
);
|
||
dispatched_channels.push((channel_id, typing_scope));
|
||
}
|
||
tracing::debug!(
|
||
dispatched = dispatched_channels.len(),
|
||
queue_depth = queue.pending_channels(),
|
||
"dispatch_pending"
|
||
);
|
||
dispatched_channels
|
||
}
|
||
|
||
/// Spawn a task that posts a user-visible failure notice to the relay.
|
||
///
|
||
/// Shared by the hard-cap immediate dead-letter path and the retries-exhausted
|
||
/// dead-letter path so neither duplicates the tokio::spawn block.
|
||
fn spawn_failure_notice(
|
||
rest_client: Option<&relay::RestClient>,
|
||
batch: &FlushBatch,
|
||
content: String,
|
||
) {
|
||
if let Some(rest) = rest_client {
|
||
let thread_tags = batch
|
||
.events
|
||
.last()
|
||
.map(|be| queue::parse_thread_tags(&be.event))
|
||
.unwrap_or_default();
|
||
let rest = rest.clone();
|
||
let channel_id = batch.channel_id;
|
||
tokio::spawn(async move {
|
||
pool::post_failure_notice(&rest, channel_id, &thread_tags, &content).await;
|
||
});
|
||
}
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn handle_prompt_result(
|
||
pool: &mut AgentPool,
|
||
queue: &mut EventQueue,
|
||
config: &Config,
|
||
mut result: PromptResult,
|
||
heartbeat_in_flight: &mut bool,
|
||
removed_channels: &HashSet<Uuid>,
|
||
crash_history: &mut [SlotCircuit],
|
||
respawn_tx: &mpsc::Sender<RespawnResult>,
|
||
respawn_tasks: &mut tokio::task::JoinSet<()>,
|
||
observer: Option<observer::ObserverHandle>,
|
||
rest_client: Option<&relay::RestClient>,
|
||
) -> LoopAction {
|
||
let before = pool.task_map().len();
|
||
let agent_index = result.agent.index;
|
||
pool.task_map_mut()
|
||
.retain(|_, meta| meta.agent_index != agent_index);
|
||
debug_assert_eq!(before, pool.task_map().len() + 1);
|
||
|
||
// Requeue BEFORE mark_complete: requeue() sets retry_after with a future
|
||
// deadline, and mark_complete() checks for it to decide whether to preserve
|
||
// retry_counts. If mark_complete runs first, retry_counts is cleared and
|
||
// every retry starts at attempt 1 — defeating exponential backoff and
|
||
// dead-letter protection.
|
||
if let Some(batch) = result.batch.take() {
|
||
// Don't requeue batches for channels the agent was removed from —
|
||
// those events are stale and should be silently dropped.
|
||
if !removed_channels.contains(&batch.channel_id) {
|
||
if matches!(
|
||
result.outcome,
|
||
PromptOutcome::Cancelled | PromptOutcome::CancelDrainTimeout(_)
|
||
) {
|
||
// Cancel re-prompt: store as cancelled events so flush_next()
|
||
// merges them into the next FlushBatch.cancelled_events,
|
||
// enabling the annotated merged-prompt format. The batch's
|
||
// cancel_reason (set by the pool task per the control signal)
|
||
// selects steer vs interrupt framing. It is always set on this
|
||
// path; if somehow unset, fall back to the gentler Steer framing
|
||
// — consistent with MergeFraming::for_reason(None) and the
|
||
// system default — rather than telling the agent to supersede.
|
||
//
|
||
// CancelDrainTimeout shares this path with Cancelled: a failed
|
||
// 5s drain after a control-signal cancel is a cleanup-deadline
|
||
// problem, not the deterministic hard-cap death below — the
|
||
// original batch must survive with no retry/dead-letter
|
||
// accounting, same as a clean cancel.
|
||
let reason = batch.cancel_reason.unwrap_or(CancelReason::Steer);
|
||
queue.requeue_as_cancelled(batch, reason);
|
||
} else if matches!(result.outcome, PromptOutcome::Timeout(TimeoutKind::Hard)) {
|
||
// Hard-cap timeout is deterministic: re-running the same task
|
||
// from a fresh session will reproduce the same death. Dead-letter
|
||
// immediately without requeueing so the channel isn't subjected to
|
||
// up to 10 × 1-hour retry cycles.
|
||
tracing::error!(
|
||
channel_id = %batch.channel_id,
|
||
events = batch.events.len(),
|
||
"dead-lettering batch after hard-cap timeout — discarding {} events",
|
||
batch.events.len(),
|
||
);
|
||
let content = format!(
|
||
"⚠️ I couldn't process the last request (the turn exceeded the maximum duration ({}s)). Please re-send if it's still needed.",
|
||
config.max_turn_duration_secs
|
||
);
|
||
spawn_failure_notice(rest_client, &batch, content);
|
||
} else if let Some(dead) = queue.requeue(batch) {
|
||
// Dead-lettered: retries exhausted and the events are gone.
|
||
// Post a visible notice so the channel isn't left waiting on
|
||
// a turn that will never happen.
|
||
let reason = match &result.outcome {
|
||
PromptOutcome::Timeout(TimeoutKind::Idle) => "the turn timed out".to_string(),
|
||
// Unreachable today: Timeout(Hard) is consumed by the immediate
|
||
// dead-letter arm above before requeue() runs. Fail soft rather
|
||
// than panicking the main loop if that chain is ever reordered.
|
||
PromptOutcome::Timeout(TimeoutKind::Hard) => {
|
||
"the turn exceeded the maximum duration".to_string()
|
||
}
|
||
PromptOutcome::AgentExited => "the agent process exited".to_string(),
|
||
PromptOutcome::Error(e) => format!("{e}"),
|
||
_ => "repeated failures".to_string(),
|
||
};
|
||
let content = format!(
|
||
"⚠️ I couldn't process the last request after multiple retries ({reason}). Please re-send if it's still needed."
|
||
);
|
||
spawn_failure_notice(rest_client, &dead, content);
|
||
}
|
||
} else {
|
||
tracing::debug!(
|
||
channel_id = %batch.channel_id,
|
||
events = batch.events.len(),
|
||
"dropping failed batch for removed channel"
|
||
);
|
||
}
|
||
}
|
||
|
||
match &result.source {
|
||
PromptSource::Channel(ch) => queue.mark_complete(*ch),
|
||
PromptSource::Heartbeat => *heartbeat_in_flight = false,
|
||
}
|
||
|
||
// Strip sessions for channels the agent was removed from while this
|
||
// agent was checked out. This covers the gap where invalidate_channel_sessions
|
||
// only touches idle agents.
|
||
for ch in removed_channels {
|
||
result.agent.state.invalidate_channel(ch);
|
||
}
|
||
|
||
let outcome_label = match &result.outcome {
|
||
PromptOutcome::Ok(_) => "ok",
|
||
PromptOutcome::Error(_) => "error",
|
||
PromptOutcome::Timeout(TimeoutKind::Idle) => "idle_timeout",
|
||
PromptOutcome::Timeout(TimeoutKind::Hard) => "hard_timeout",
|
||
PromptOutcome::AgentExited => "exited",
|
||
PromptOutcome::Cancelled => "cancelled",
|
||
PromptOutcome::CancelDrainTimeout(_) => "cancel_drain_timeout",
|
||
};
|
||
let agent_index = result.agent.index;
|
||
// Capture the spawn-time configured model and our PID before the agent is
|
||
// moved into match arms below. `desired_model` reflects the config/persona
|
||
// model at spawn time — it does NOT reflect `session/set_model` overrides,
|
||
// which live in buzz-agent's session state and are what `llm: (model) …`
|
||
// errors carry. The two can legitimately differ; `configured_model=` is
|
||
// still valuable for identifying a stale orphan running an old model.
|
||
let harness_configured_model = result
|
||
.agent
|
||
.desired_model
|
||
.as_deref()
|
||
.unwrap_or("<none>")
|
||
.to_string();
|
||
let harness_pid = std::process::id();
|
||
|
||
let channel_id = match &result.source {
|
||
PromptSource::Channel(ch) => Some(*ch),
|
||
PromptSource::Heartbeat => None,
|
||
};
|
||
let emit_turn_error = |error_msg: &str, error_code: Option<i64>| {
|
||
if let Some(ref observer) = observer {
|
||
let mut payload = serde_json::json!({
|
||
"outcome": outcome_label,
|
||
"error": error_msg,
|
||
});
|
||
if let Some(code) = error_code {
|
||
payload["code"] = serde_json::json!(code);
|
||
}
|
||
observer.emit(
|
||
"turn_error",
|
||
Some(agent_index),
|
||
&observer::context_for(channel_id, None, None),
|
||
payload,
|
||
);
|
||
}
|
||
};
|
||
|
||
match result.outcome {
|
||
// Successful prompt — return agent to pool.
|
||
PromptOutcome::Ok(_) => {
|
||
tracing::debug!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
"agent_returned"
|
||
);
|
||
pool.return_agent(result.agent);
|
||
}
|
||
// Fatal outcomes: the agent subprocess is dead or poisoned — respawn it.
|
||
PromptOutcome::AgentExited | PromptOutcome::Timeout(_) => {
|
||
tracing::warn!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
configured_model = %harness_configured_model,
|
||
pid = harness_pid,
|
||
"agent_returned — respawning"
|
||
);
|
||
let death_message: String = match outcome_label {
|
||
"exited" => "Agent process exited unexpectedly".to_string(),
|
||
"hard_timeout" => format!(
|
||
"Agent turn exceeded the maximum duration ({}s)",
|
||
config.max_turn_duration_secs
|
||
),
|
||
_ => "Agent session timed out due to inactivity".to_string(),
|
||
};
|
||
emit_turn_error(&death_message, None);
|
||
|
||
let index = result.agent.index;
|
||
let slot_history = &mut crash_history[index];
|
||
if !spawn_respawn_task(
|
||
result.agent,
|
||
config,
|
||
slot_history,
|
||
respawn_tx,
|
||
respawn_tasks,
|
||
observer.clone(),
|
||
) {
|
||
// Circuit open — slot stays empty until maintenance refill.
|
||
if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) {
|
||
tracing::error!("all agents dead — exiting");
|
||
return LoopAction::Exit;
|
||
}
|
||
}
|
||
}
|
||
// Cancel-drain expiry: a control-signal cancel (steer fallback,
|
||
// interrupt, or explicit stop) did not drain within its bounded
|
||
// grace window. The process is poisoned/uncertain like a hard
|
||
// timeout — respawn it — but this is NOT the configured max-turn
|
||
// cap, so the message must name the actual grace, not
|
||
// `max_turn_duration_secs`. The triggering batch's fate (preserved
|
||
// for Steer/Interrupt, dropped for explicit Cancel/Rotate or a
|
||
// removed channel) is decided above — the message stays fate-neutral
|
||
// since it must be true in every case.
|
||
PromptOutcome::CancelDrainTimeout(grace) => {
|
||
tracing::warn!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
configured_model = %harness_configured_model,
|
||
pid = harness_pid,
|
||
grace = ?grace,
|
||
"agent_returned — respawning (cancel-drain timeout)"
|
||
);
|
||
let death_message = format!(
|
||
"Agent did not stop within {grace:?} after cancellation; the agent process is being replaced."
|
||
);
|
||
emit_turn_error(&death_message, None);
|
||
|
||
let index = result.agent.index;
|
||
let slot_history = &mut crash_history[index];
|
||
if !spawn_respawn_task(
|
||
result.agent,
|
||
config,
|
||
slot_history,
|
||
respawn_tx,
|
||
respawn_tasks,
|
||
observer.clone(),
|
||
) {
|
||
// Circuit open — slot stays empty until maintenance refill.
|
||
if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) {
|
||
tracing::error!("all agents dead — exiting");
|
||
return LoopAction::Exit;
|
||
}
|
||
}
|
||
}
|
||
// Errors fall into two categories:
|
||
//
|
||
// 1. Transport-class (Io, WriteTimeout, Timeout, Protocol): the stdio
|
||
// pipe may be corrupted or the agent desynchronized. These are fatal
|
||
// to the agent regardless of whether they occurred during session
|
||
// creation or an active prompt — respawn unconditionally.
|
||
//
|
||
// 2. Application-class (IdleTimeout, HardTimeout, Json): the pipe is
|
||
// intact but the prompt failed. Return the agent to the pool so it
|
||
// can be reused for the next event.
|
||
|
||
// Intentional cancel — agent is healthy, return it to the pool.
|
||
// No respawn, no retry penalty. The cancelled batch was already stored
|
||
// via requeue_as_cancelled() above and will be merged into the next
|
||
// FlushBatch by flush_next().
|
||
PromptOutcome::Cancelled => {
|
||
tracing::debug!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
configured_model = %harness_configured_model,
|
||
pid = harness_pid,
|
||
"agent_returned (cancelled)"
|
||
);
|
||
pool.return_agent(result.agent);
|
||
}
|
||
PromptOutcome::Error(ref e) => {
|
||
let is_transport_error = matches!(
|
||
e,
|
||
acp::AcpError::Io(_)
|
||
| acp::AcpError::WriteTimeout(_)
|
||
| acp::AcpError::Timeout(_)
|
||
| acp::AcpError::Protocol(_)
|
||
);
|
||
let error_code = match &e {
|
||
acp::AcpError::AgentError { code, .. } => Some(*code),
|
||
_ => None,
|
||
};
|
||
if is_transport_error {
|
||
tracing::warn!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
configured_model = %harness_configured_model,
|
||
pid = harness_pid,
|
||
error = %e,
|
||
"transport/protocol error — respawning agent"
|
||
);
|
||
emit_turn_error(&e.to_string(), error_code);
|
||
|
||
let index = result.agent.index;
|
||
let slot_history = &mut crash_history[index];
|
||
if !spawn_respawn_task(
|
||
result.agent,
|
||
config,
|
||
slot_history,
|
||
respawn_tx,
|
||
respawn_tasks,
|
||
observer,
|
||
) && pool.live_count() == 0
|
||
&& !any_respawn_in_flight(crash_history)
|
||
{
|
||
tracing::error!("all agents dead — exiting");
|
||
return LoopAction::Exit;
|
||
}
|
||
} else {
|
||
tracing::warn!(
|
||
agent = agent_index,
|
||
outcome = outcome_label,
|
||
configured_model = %harness_configured_model,
|
||
pid = harness_pid,
|
||
error = %e,
|
||
"agent_returned (application error — pipe intact)"
|
||
);
|
||
emit_turn_error(&e.to_string(), error_code);
|
||
pool.return_agent(result.agent);
|
||
}
|
||
}
|
||
}
|
||
LoopAction::Continue
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn recover_panicked_agent(
|
||
pool: &mut AgentPool,
|
||
queue: &mut EventQueue,
|
||
config: &Config,
|
||
join_error: tokio::task::JoinError,
|
||
heartbeat_in_flight: &mut bool,
|
||
removed_channels: &HashSet<Uuid>,
|
||
typing_channels: &mut HashMap<Uuid, ThreadTags>,
|
||
crash_history: &mut [SlotCircuit],
|
||
respawn_tx: &mpsc::Sender<RespawnResult>,
|
||
respawn_tasks: &mut tokio::task::JoinSet<()>,
|
||
observer: Option<observer::ObserverHandle>,
|
||
) {
|
||
let task_id = join_error.id();
|
||
let Some(meta) = pool.task_map_mut().remove(&task_id) else {
|
||
tracing::error!("panic for unknown task {task_id:?} — bug");
|
||
return;
|
||
};
|
||
let i = meta.agent_index;
|
||
|
||
// Requeue BEFORE mark_complete (same rationale as handle_prompt_result).
|
||
if let Some(batch) = meta.recoverable_batch {
|
||
if let Some(ch) = meta.channel_id {
|
||
if !removed_channels.contains(&ch) {
|
||
// Dead-letter on exhaustion is logged inside requeue(); a
|
||
// panic path has no outcome to report, so no notice here.
|
||
let _ = queue.requeue(batch);
|
||
tracing::warn!("requeued batch for panicked agent {i}");
|
||
} else {
|
||
tracing::debug!(
|
||
channel_id = %ch,
|
||
"dropping panicked batch for removed channel"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
if let Some(ch) = meta.channel_id {
|
||
queue.mark_complete(ch);
|
||
typing_channels.remove(&ch);
|
||
tracing::warn!("cleared wedged in-flight channel {ch} from panicked agent {i}");
|
||
} else {
|
||
*heartbeat_in_flight = false;
|
||
tracing::warn!("cleared wedged heartbeat_in_flight from panicked agent {i}");
|
||
}
|
||
|
||
if let Some(ref observer) = observer {
|
||
observer.emit(
|
||
"agent_panic",
|
||
Some(i),
|
||
&observer::context_for(meta.channel_id, None, None),
|
||
serde_json::json!({
|
||
"outcome": "panic",
|
||
"error": format!("Agent task panicked: {join_error}"),
|
||
}),
|
||
);
|
||
}
|
||
|
||
// Panics count as crashes for the circuit breaker.
|
||
// The panicked task already dropped the AcpClient, so we just need to
|
||
// check the circuit and spawn a fresh agent in the background.
|
||
let slot = &mut crash_history[i];
|
||
|
||
let delay = match slot.record_crash() {
|
||
CrashVerdict::CircuitOpen => {
|
||
tracing::error!(agent = i, "circuit open after panic — not respawning");
|
||
return;
|
||
}
|
||
CrashVerdict::HalfOpenProbe => {
|
||
tracing::info!(agent = i, "circuit half-open — probe respawn after panic");
|
||
Duration::ZERO
|
||
}
|
||
CrashVerdict::Respawn(d) => {
|
||
tracing::info!(
|
||
agent = i,
|
||
delay_ms = d.as_millis(),
|
||
"respawn backoff after panic"
|
||
);
|
||
d
|
||
}
|
||
};
|
||
|
||
// Spawn respawn work off the main loop.
|
||
slot.respawn_in_flight = true;
|
||
let cmd = config.agent_command.clone();
|
||
let args = config.agent_args.clone();
|
||
let env = config.persona_env_vars.clone();
|
||
let has_codex = config.has_generated_codex_config;
|
||
let guard = RespawnGuard::new(i, respawn_tx.clone());
|
||
respawn_tasks.spawn(async move {
|
||
if !delay.is_zero() {
|
||
tokio::time::sleep(delay).await;
|
||
}
|
||
let result = spawn_and_init(&cmd, &args, &env, has_codex, i, observer).await;
|
||
guard.send(result);
|
||
});
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn drain_ready_join_results(
|
||
pool: &mut AgentPool,
|
||
queue: &mut EventQueue,
|
||
config: &Config,
|
||
heartbeat_in_flight: &mut bool,
|
||
removed_channels: &HashSet<Uuid>,
|
||
typing_channels: &mut HashMap<Uuid, ThreadTags>,
|
||
crash_history: &mut [SlotCircuit],
|
||
respawn_tx: &mpsc::Sender<RespawnResult>,
|
||
respawn_tasks: &mut tokio::task::JoinSet<()>,
|
||
observer: Option<observer::ObserverHandle>,
|
||
) -> LoopAction {
|
||
while let Some(Some(join_result)) = pool.join_set.join_next().now_or_never() {
|
||
if let Err(join_error) = join_result {
|
||
tracing::error!("agent task panicked: {join_error}");
|
||
recover_panicked_agent(
|
||
pool,
|
||
queue,
|
||
config,
|
||
join_error,
|
||
heartbeat_in_flight,
|
||
removed_channels,
|
||
typing_channels,
|
||
crash_history,
|
||
respawn_tx,
|
||
respawn_tasks,
|
||
observer.clone(),
|
||
);
|
||
if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) {
|
||
return LoopAction::Exit;
|
||
}
|
||
}
|
||
}
|
||
LoopAction::Continue
|
||
}
|
||
|
||
fn dispatch_heartbeat(
|
||
pool: &mut AgentPool,
|
||
ctx: &Arc<PromptContext>,
|
||
heartbeat_in_flight: &mut bool,
|
||
) {
|
||
if *heartbeat_in_flight {
|
||
return;
|
||
}
|
||
let agent = match pool.try_claim(None) {
|
||
Some(a) => a,
|
||
None => return,
|
||
};
|
||
|
||
let prompt_text = ctx
|
||
.heartbeat_prompt
|
||
.clone()
|
||
.unwrap_or_else(default_heartbeat_prompt);
|
||
let result_tx = pool.result_tx();
|
||
let ctx_clone = Arc::clone(ctx);
|
||
let agent_index = agent.index;
|
||
|
||
let abort_handle = pool.join_set.spawn(async move {
|
||
pool::run_prompt_task(agent, None, Some(prompt_text), ctx_clone, result_tx, None).await;
|
||
});
|
||
|
||
pool.task_map_mut().insert(
|
||
abort_handle.id(),
|
||
pool::TaskMeta {
|
||
agent_index,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
*heartbeat_in_flight = true;
|
||
tracing::info!(agent = agent_index, "heartbeat_fired");
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod agent_draft_prompt_tests {
|
||
#[test]
|
||
fn shared_base_prompt_teaches_portable_agent_drafts() {
|
||
let prompt = include_str!("base_prompt.md");
|
||
assert!(prompt.contains("buzz agents draft-create"));
|
||
assert!(prompt.contains("ask for at most two things"));
|
||
assert!(prompt.contains("what it should do day-to-day"));
|
||
assert!(prompt.contains("owner saves it"));
|
||
assert!(prompt.contains("Do not ask about runtime, provider, model, credentials"));
|
||
}
|
||
}
|
||
|
||
fn default_heartbeat_prompt() -> String {
|
||
let now = chrono::Utc::now().to_rfc3339();
|
||
format!(
|
||
"[System: Heartbeat]\nTime: {now}\n\n\
|
||
You have been awakened for a routine heartbeat. You have NO incoming messages or\n\
|
||
active channel context for this turn.\n\n\
|
||
Your tasks:\n\
|
||
1. Run `buzz feed get --types needs_action` to check for pending workflow approvals or\n\
|
||
high-priority requests addressed to you.\n\
|
||
2. Run `buzz feed get --types mentions` to check for unanswered @mentions.\n\
|
||
3. If you find actionable items, address them using the appropriate CLI commands\n\
|
||
(e.g., `buzz workflows approve --token <UUID>`, `buzz messages send`,\n\
|
||
`buzz messages send --reply-to <event-id>`).\n\
|
||
4. If there are no pending actions or mentions, end your turn immediately.\n\n\
|
||
Do not run `buzz channels list` or `buzz messages search` unless you have a specific reason.\n\
|
||
Do not invent work — only act on items surfaced by the feed commands."
|
||
)
|
||
}
|
||
|
||
/// Spawn a background respawn task for a crashed agent slot.
|
||
///
|
||
/// Does the circuit breaker check synchronously (non-blocking), then spawns
|
||
/// the actual shutdown + backoff + spawn_and_init work into a background task.
|
||
/// The result comes back through `respawn_tx` so the main loop stays responsive.
|
||
///
|
||
/// Returns `true` if a respawn task was spawned, `false` if the circuit is open.
|
||
fn spawn_respawn_task(
|
||
old_agent: OwnedAgent,
|
||
config: &Config,
|
||
slot: &mut SlotCircuit,
|
||
respawn_tx: &mpsc::Sender<RespawnResult>,
|
||
respawn_tasks: &mut tokio::task::JoinSet<()>,
|
||
observer: Option<observer::ObserverHandle>,
|
||
) -> bool {
|
||
let index = old_agent.index;
|
||
|
||
// Circuit breaker: record crash, decide whether to respawn.
|
||
let delay = match slot.record_crash() {
|
||
CrashVerdict::CircuitOpen => {
|
||
tracing::error!(agent = index, "circuit open — not respawning");
|
||
return false;
|
||
}
|
||
CrashVerdict::HalfOpenProbe => {
|
||
tracing::info!(agent = index, "circuit half-open — probe respawn");
|
||
Duration::ZERO
|
||
}
|
||
CrashVerdict::Respawn(d) => {
|
||
tracing::info!(agent = index, delay_ms = d.as_millis(), "respawn backoff");
|
||
d
|
||
}
|
||
};
|
||
|
||
slot.respawn_in_flight = true;
|
||
|
||
// Spawn the actual work (shutdown + sleep + spawn + init) off the main loop.
|
||
let cmd = config.agent_command.clone();
|
||
let args = config.agent_args.clone();
|
||
let env = config.persona_env_vars.clone();
|
||
let has_codex = config.has_generated_codex_config;
|
||
let guard = RespawnGuard::new(index, respawn_tx.clone());
|
||
respawn_tasks.spawn(async move {
|
||
// Shutdown old agent (reap child, prevent zombie).
|
||
let mut agent = old_agent;
|
||
agent.acp.shutdown().await;
|
||
drop(agent);
|
||
|
||
if !delay.is_zero() {
|
||
tokio::time::sleep(delay).await;
|
||
}
|
||
|
||
let result = spawn_and_init(&cmd, &args, &env, has_codex, index, observer).await;
|
||
guard.send(result);
|
||
});
|
||
|
||
true
|
||
}
|
||
|
||
fn normalized_agent_name(init_result: &serde_json::Value) -> String {
|
||
init_result
|
||
.get("agentInfo")
|
||
.or_else(|| init_result.get("serverInfo"))
|
||
.and_then(|info| info.get("name"))
|
||
.and_then(|value| value.as_str())
|
||
.unwrap_or("unknown")
|
||
.trim()
|
||
.to_ascii_lowercase()
|
||
}
|
||
|
||
// ── spawn_and_init ────────────────────────────────────────────────────────────
|
||
/// Spawn an agent subprocess and run the MCP `initialize` handshake.
|
||
///
|
||
/// Takes owned args so it can run in a background `tokio::spawn` task without
|
||
/// borrowing `Config`. All respawn/refill paths use this.
|
||
async fn spawn_and_init(
|
||
command: &str,
|
||
args: &[String],
|
||
extra_env: &[(String, String)],
|
||
has_generated_codex_config: bool,
|
||
agent_index: usize,
|
||
observer: Option<observer::ObserverHandle>,
|
||
) -> Result<(AcpClient, u32, String)> {
|
||
let mut acp = AcpClient::spawn(command, args, extra_env, has_generated_codex_config)
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("failed to spawn agent: {e}"))?;
|
||
acp.set_observer(observer, agent_index);
|
||
|
||
match acp.initialize().await {
|
||
Ok(init_result) => {
|
||
tracing::info!("agent initialized: {init_result}");
|
||
let protocol_version = init_result["protocolVersion"].as_u64().unwrap_or(1) as u32;
|
||
acp.observe(
|
||
"agent_initialized",
|
||
serde_json::json!({
|
||
"agentIndex": agent_index,
|
||
"initializeResult": init_result,
|
||
}),
|
||
);
|
||
let agent_name = normalized_agent_name(&init_result);
|
||
Ok((acp, protocol_version, agent_name))
|
||
}
|
||
Err(e) => {
|
||
// Explicitly shut down the spawned child to prevent zombie/leak.
|
||
// Drop only does start_kill + try_wait (best-effort); shutdown()
|
||
// does start_kill + bounded wait (guaranteed reap).
|
||
acp.shutdown().await;
|
||
Err(anyhow::anyhow!("agent initialize failed: {e}"))
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn spawn_auth_client(agent: &AuthAgentArgs) -> Result<AcpClient, acp::AcpError> {
|
||
let agent_args = config::normalize_agent_args(&agent.agent_command, agent.agent_args.clone());
|
||
AcpClient::spawn(&agent.agent_command, &agent_args, &[], false).await
|
||
}
|
||
|
||
fn extract_auth_methods(init_result: &serde_json::Value) -> Vec<serde_json::Value> {
|
||
init_result
|
||
.get("authMethods")
|
||
.and_then(|methods| methods.as_array())
|
||
.cloned()
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
/// `buzz-acp auth-methods` — spawn an adapter, initialize it, print authMethods.
|
||
async fn run_auth_methods(args: AuthMethodsArgs) -> Result<()> {
|
||
let mut client = match spawn_auth_client(&args.agent).await {
|
||
Ok(c) => c,
|
||
Err(e) => {
|
||
eprintln!("error: failed to spawn agent: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
let init_result = match tokio::time::timeout(MODELS_TIMEOUT, client.initialize()).await {
|
||
Ok(Ok(result)) => result,
|
||
Ok(Err(e)) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent initialize failed: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
Err(_) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent timed out ({MODELS_TIMEOUT:?})");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
let methods = extract_auth_methods(&init_result);
|
||
client.shutdown().await;
|
||
|
||
if args.json {
|
||
let output = serde_json::json!({ "methods": methods });
|
||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||
} else if methods.is_empty() {
|
||
println!("No auth methods advertised.");
|
||
} else {
|
||
for method in methods {
|
||
let id = method
|
||
.get("id")
|
||
.and_then(|value| value.as_str())
|
||
.unwrap_or("unknown");
|
||
let name = method
|
||
.get("name")
|
||
.and_then(|value| value.as_str())
|
||
.unwrap_or(id);
|
||
println!("{id}\t{name}");
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// `buzz-acp authenticate` — invoke one adapter-owned auth method.
|
||
async fn run_authenticate(args: AuthenticateArgs) -> Result<()> {
|
||
let mut client = match spawn_auth_client(&args.agent).await {
|
||
Ok(c) => c,
|
||
Err(e) => {
|
||
eprintln!("error: failed to spawn agent: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
let init_result = match tokio::time::timeout(MODELS_TIMEOUT, client.initialize()).await {
|
||
Ok(Ok(result)) => result,
|
||
Ok(Err(e)) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent initialize failed: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
Err(_) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent initialize timed out ({MODELS_TIMEOUT:?})");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
let supports_method = extract_auth_methods(&init_result)
|
||
.iter()
|
||
.any(|method| method.get("id").and_then(|id| id.as_str()) == Some(args.method_id.as_str()));
|
||
if !supports_method {
|
||
client.shutdown().await;
|
||
eprintln!(
|
||
"error: auth method '{}' is not advertised by this adapter",
|
||
args.method_id
|
||
);
|
||
std::process::exit(1);
|
||
}
|
||
|
||
let result =
|
||
tokio::time::timeout(AUTHENTICATE_TIMEOUT, client.authenticate(&args.method_id)).await;
|
||
|
||
match result {
|
||
Ok(Ok(_)) => {
|
||
client.shutdown().await;
|
||
Ok(())
|
||
}
|
||
Ok(Err(e)) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: authenticate failed: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
Err(_) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: authenticate timed out ({AUTHENTICATE_TIMEOUT:?})");
|
||
std::process::exit(1);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Flow: spawn → initialize → session/new → print models → shutdown.
|
||
/// No relay connection, no MCP servers, no subscriptions. ~2-5s total.
|
||
async fn run_models(args: ModelsArgs) -> Result<()> {
|
||
use acp::{extract_model_config_options, extract_model_state};
|
||
|
||
let agent_args = config::normalize_agent_args(&args.agent.agent_command, args.agent.agent_args);
|
||
let cwd = std::env::current_dir()
|
||
.unwrap_or_else(|_| std::path::PathBuf::from("/"))
|
||
.to_string_lossy()
|
||
.to_string();
|
||
|
||
// Spawn outside the timeout so we always own the child for cleanup.
|
||
// `models` subcommand doesn't use persona packs — no extra env, no codex config.
|
||
let mut client =
|
||
match AcpClient::spawn(&args.agent.agent_command, &agent_args, &[], false).await {
|
||
Ok(c) => c,
|
||
Err(e) => {
|
||
eprintln!("error: failed to spawn agent: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
// Initialize + session/new under a timeout. Client is owned above,
|
||
// so shutdown() runs on all paths (success, error, timeout).
|
||
let protocol_result = tokio::time::timeout(MODELS_TIMEOUT, async {
|
||
let init = client.initialize().await?;
|
||
let session = client.session_new_full(&cwd, vec![], None).await?;
|
||
Ok::<_, acp::AcpError>((init, session))
|
||
})
|
||
.await;
|
||
|
||
let (init_result, session_resp) = match protocol_result {
|
||
Ok(Ok(tuple)) => tuple,
|
||
Ok(Err(e)) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent communication failed: {e}");
|
||
std::process::exit(1);
|
||
}
|
||
Err(_) => {
|
||
client.shutdown().await;
|
||
eprintln!("error: agent timed out ({MODELS_TIMEOUT:?})");
|
||
std::process::exit(1);
|
||
}
|
||
};
|
||
|
||
// Extract agent info from initialize response.
|
||
// ACP spec uses "serverInfo" (MCP heritage); some agents may use "agentInfo".
|
||
let info_obj = init_result
|
||
.get("serverInfo")
|
||
.or_else(|| init_result.get("agentInfo"));
|
||
let agent_name = info_obj
|
||
.and_then(|ai| ai.get("name"))
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("unknown");
|
||
let agent_version = info_obj
|
||
.and_then(|ai| ai.get("version"))
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("unknown");
|
||
|
||
// Extract model info from session/new response.
|
||
let config_options = extract_model_config_options(&session_resp.raw);
|
||
let model_state = extract_model_state(&session_resp.raw);
|
||
|
||
if args.json {
|
||
// Structured JSON output — consumed by Phase 3 `get_agent_models`.
|
||
let output = serde_json::json!({
|
||
"agent": {
|
||
"name": agent_name,
|
||
"version": agent_version,
|
||
},
|
||
"stable": {
|
||
"configOptions": config_options,
|
||
},
|
||
"unstable": model_state.as_ref().map(|ms| serde_json::json!({
|
||
"currentModelId": ms.get("currentModelId"),
|
||
"availableModels": ms.get("availableModels"),
|
||
})),
|
||
});
|
||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||
} else {
|
||
// Human-readable output.
|
||
println!("Agent: {} v{}", agent_name, agent_version);
|
||
println!();
|
||
|
||
let mut has_models = false;
|
||
|
||
if !config_options.is_empty() {
|
||
println!("Models (stable configOptions):");
|
||
for opt in &config_options {
|
||
let config_id = opt.get("configId").and_then(|v| v.as_str()).unwrap_or("?");
|
||
let display = opt
|
||
.get("displayName")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or(config_id);
|
||
println!(" {display} (configId: {config_id})");
|
||
if let Some(options) = opt.get("options").and_then(|v| v.as_array()) {
|
||
for o in options {
|
||
let val = o.get("value").and_then(|v| v.as_str()).unwrap_or("?");
|
||
let name = o.get("displayName").and_then(|v| v.as_str()).unwrap_or(val);
|
||
println!(" - {name} (value: {val})");
|
||
}
|
||
}
|
||
}
|
||
has_models = true;
|
||
}
|
||
|
||
if let Some(ref ms) = model_state {
|
||
let current = ms
|
||
.get("currentModelId")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("(none)");
|
||
println!("Models (unstable SessionModelState):");
|
||
println!(" Current: {current}");
|
||
if let Some(available) = ms.get("availableModels").and_then(|v| v.as_array()) {
|
||
println!(" Available:");
|
||
for m in available {
|
||
let id = m.get("modelId").and_then(|v| v.as_str()).unwrap_or("?");
|
||
let name = m.get("name").and_then(|v| v.as_str()).unwrap_or(id);
|
||
let desc = m.get("description").and_then(|v| v.as_str()).unwrap_or("");
|
||
if desc.is_empty() {
|
||
println!(" - {name} (id: {id})");
|
||
} else {
|
||
println!(" - {name} (id: {id}) — {desc}");
|
||
}
|
||
}
|
||
}
|
||
has_models = true;
|
||
}
|
||
|
||
if !has_models {
|
||
println!("No model information available from this agent.");
|
||
}
|
||
}
|
||
|
||
client.shutdown().await;
|
||
Ok(())
|
||
}
|
||
|
||
fn build_mcp_servers(config: &Config) -> Vec<McpServer> {
|
||
if config.mcp_command.is_empty() {
|
||
return vec![];
|
||
}
|
||
vec![McpServer {
|
||
name: std::path::Path::new(&config.mcp_command)
|
||
.file_stem()
|
||
.and_then(|s| s.to_str())
|
||
.unwrap_or("mcp")
|
||
.to_string(),
|
||
command: config.mcp_command.clone(),
|
||
args: vec![],
|
||
env: {
|
||
let mut env = vec![
|
||
EnvVar {
|
||
name: "BUZZ_RELAY_URL".into(),
|
||
value: config.relay_url.clone(),
|
||
},
|
||
EnvVar {
|
||
name: "BUZZ_PRIVATE_KEY".into(),
|
||
// bech32 encoding of a valid secret key is infallible.
|
||
// Panic here is correct: injecting a bogus secret would cause
|
||
// delayed, hard-to-diagnose agent failures downstream.
|
||
value: config
|
||
.keys
|
||
.secret_key()
|
||
.to_bech32()
|
||
.expect("secret key bech32 encoding should never fail"),
|
||
},
|
||
];
|
||
// Forward BUZZ_AUTH_TAG (NIP-OA owner attestation credential)
|
||
// so the MCP server can attach it to every signed event.
|
||
if let Ok(auth_tag) = std::env::var("BUZZ_AUTH_TAG") {
|
||
if !auth_tag.is_empty() {
|
||
env.push(EnvVar {
|
||
name: "BUZZ_AUTH_TAG".into(),
|
||
value: auth_tag,
|
||
});
|
||
}
|
||
}
|
||
env
|
||
},
|
||
}]
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod heartbeat_base_prompt_tests {
|
||
use super::*;
|
||
|
||
// Pins the heartbeat dispatch path (dispatch_heartbeat, ~line 2359): a
|
||
// legacy agent WITH a base_prompt must get [Base] prepended to the
|
||
// heartbeat user message, composed as `[Base]\n{bp}\n\n{prompt}`. This is
|
||
// the second half of the round-2 regression (the first being initial_message).
|
||
|
||
#[test]
|
||
fn test_heartbeat_legacy_agent_gets_base_prepended() {
|
||
// protocol_version 1 + Some(base_prompt): heartbeat prompt is prefixed
|
||
// with the [Base] section exactly as the legacy session/new path would.
|
||
let prompt = "[System: Heartbeat]\nrun feed get";
|
||
let composed = pool::prepend_base_for_legacy(1, Some("you are a helpful agent"), prompt);
|
||
assert_eq!(
|
||
composed,
|
||
"[Base]\nyou are a helpful agent\n\n[System: Heartbeat]\nrun feed get"
|
||
);
|
||
assert!(composed.starts_with("[Base]\nyou are a helpful agent\n\n"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_heartbeat_modern_agent_omits_base() {
|
||
// protocol_version 2 gets base_prompt via session/new; the heartbeat
|
||
// prompt is sent verbatim.
|
||
let prompt = "[System: Heartbeat]\nrun feed get";
|
||
let composed = pool::prepend_base_for_legacy(2, Some("you are a helpful agent"), prompt);
|
||
assert_eq!(composed, prompt);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod owner_control_command_tests {
|
||
use super::*;
|
||
use nostr::{EventBuilder, Keys, Kind, Tag};
|
||
|
||
fn make_event(kind: u32, content: &str, p_hex: Option<&str>) -> nostr::Event {
|
||
let keys = Keys::generate();
|
||
let tags = match p_hex {
|
||
Some(hex) => vec![Tag::parse(["p", hex]).expect("p tag")],
|
||
None => vec![],
|
||
};
|
||
EventBuilder::new(Kind::Custom(kind as u16), content)
|
||
.tags(tags)
|
||
.sign_with_keys(&keys)
|
||
.unwrap()
|
||
}
|
||
|
||
#[test]
|
||
fn owner_control_command_requires_kind_content_and_agent_mention() {
|
||
let agent = "ab".repeat(32);
|
||
|
||
let event = make_event(KIND_STREAM_MESSAGE, " !rotate ", Some(&agent));
|
||
assert!(is_owner_control_command(
|
||
&event,
|
||
KIND_STREAM_MESSAGE,
|
||
"!rotate",
|
||
&agent
|
||
));
|
||
|
||
let wrong_kind = make_event(1, "!rotate", Some(&agent));
|
||
assert!(!is_owner_control_command(&wrong_kind, 1, "!rotate", &agent));
|
||
|
||
let wrong_content = make_event(KIND_STREAM_MESSAGE, "!cancel", Some(&agent));
|
||
assert!(!is_owner_control_command(
|
||
&wrong_content,
|
||
KIND_STREAM_MESSAGE,
|
||
"!rotate",
|
||
&agent
|
||
));
|
||
|
||
let no_mention = make_event(KIND_STREAM_MESSAGE, "!rotate", None);
|
||
assert!(!is_owner_control_command(
|
||
&no_mention,
|
||
KIND_STREAM_MESSAGE,
|
||
"!rotate",
|
||
&agent
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn mode_gate_signal_maps_handling_to_control_signal() {
|
||
let owner = "a".repeat(64);
|
||
let other = "b".repeat(64);
|
||
|
||
// Queue: never signals — events wait for the turn to finish.
|
||
assert!(mode_gate_signal(MultipleEventHandling::Queue, &owner, Some(&owner)).is_none());
|
||
|
||
// Steer: always steers (eligibility already enforced upstream).
|
||
assert!(matches!(
|
||
mode_gate_signal(MultipleEventHandling::Steer, &other, Some(&owner)),
|
||
Some(ControlSignal::Steer)
|
||
));
|
||
// Steer even when owner is unknown — gate doesn't re-check authorship.
|
||
assert!(matches!(
|
||
mode_gate_signal(MultipleEventHandling::Steer, &other, None),
|
||
Some(ControlSignal::Steer)
|
||
));
|
||
|
||
// Interrupt: always interrupts for any eligible author.
|
||
assert!(matches!(
|
||
mode_gate_signal(MultipleEventHandling::Interrupt, &other, Some(&owner)),
|
||
Some(ControlSignal::Interrupt)
|
||
));
|
||
|
||
// OwnerInterrupt: interrupts only for the owner.
|
||
assert!(matches!(
|
||
mode_gate_signal(MultipleEventHandling::OwnerInterrupt, &owner, Some(&owner)),
|
||
Some(ControlSignal::Interrupt)
|
||
));
|
||
assert!(
|
||
mode_gate_signal(MultipleEventHandling::OwnerInterrupt, &other, Some(&owner)).is_none(),
|
||
"owner-interrupt must not fire for a non-owner author"
|
||
);
|
||
assert!(
|
||
mode_gate_signal(MultipleEventHandling::OwnerInterrupt, &owner, None).is_none(),
|
||
"owner-interrupt must not fire when the owner is unknown"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn signal_in_flight_task_sends_rotate_once() {
|
||
let mut pool = AgentPool::from_slots(vec![]);
|
||
let channel_id = Uuid::new_v4();
|
||
let other_channel_id = Uuid::new_v4();
|
||
let (control_tx, control_rx) = tokio::sync::oneshot::channel();
|
||
|
||
let abort_handle = pool.join_set.spawn(async {});
|
||
pool.task_map_mut().insert(
|
||
abort_handle.id(),
|
||
pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: Some(channel_id),
|
||
recoverable_batch: None,
|
||
control_tx: Some(control_tx),
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
|
||
assert!(!signal_in_flight_task(
|
||
&mut pool,
|
||
other_channel_id,
|
||
ControlSignal::Rotate
|
||
));
|
||
assert!(signal_in_flight_task(
|
||
&mut pool,
|
||
channel_id,
|
||
ControlSignal::Rotate
|
||
));
|
||
assert_eq!(control_rx.await.unwrap(), ControlSignal::Rotate);
|
||
assert!(!signal_in_flight_task(
|
||
&mut pool,
|
||
channel_id,
|
||
ControlSignal::Rotate
|
||
));
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod owner_cache_tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn new_with_some_caches_immediately() {
|
||
let cache = OwnerCache::new(Some("abcd".into()));
|
||
assert_eq!(cache.get(), Some("abcd"));
|
||
}
|
||
|
||
#[test]
|
||
fn new_with_none_returns_none() {
|
||
let cache = OwnerCache::new(None);
|
||
assert!(cache.get().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn get_returns_cached_value() {
|
||
let cache = OwnerCache::new(Some("ab".repeat(32)));
|
||
assert_eq!(cache.get(), Some("ab".repeat(32)).as_deref());
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod author_gate_tests {
|
||
use super::*;
|
||
|
||
/// A `RestClient` for tests. The author-gate decisions exercised here all
|
||
/// resolve from the owner pubkey or sibling cache before any HTTP call, so
|
||
/// this client is never actually used to make a request.
|
||
fn dummy_rest_client() -> relay::RestClient {
|
||
relay::RestClient {
|
||
http: reqwest::Client::new(),
|
||
base_url: "http://localhost:0".into(),
|
||
keys: nostr::Keys::generate(),
|
||
auth_tag_json: None,
|
||
}
|
||
}
|
||
|
||
const OWNER: &str = "00";
|
||
const SIBLING: &str = "11";
|
||
const EXTERNAL: &str = "22";
|
||
const STRANGER: &str = "33";
|
||
|
||
/// Owner + a known sibling, none of them on the explicit allowlist.
|
||
fn cache_with_sibling() -> OwnerCache {
|
||
let cache = OwnerCache::new(Some(OWNER.into()));
|
||
cache.cache_sibling(SIBLING.into(), true);
|
||
cache.cache_sibling(STRANGER.into(), false);
|
||
cache
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_allowlist_accepts_sibling_not_in_allowlist() {
|
||
let cache = cache_with_sibling();
|
||
let allowlist = HashSet::from([EXTERNAL.to_string()]);
|
||
assert!(
|
||
author_allowed(
|
||
&RespondTo::Allowlist,
|
||
&allowlist,
|
||
SIBLING,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"a same-owner sibling must fire a turn under Allowlist even when not listed"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_allowlist_accepts_explicit_external_pubkey() {
|
||
let cache = cache_with_sibling();
|
||
let allowlist = HashSet::from([EXTERNAL.to_string()]);
|
||
assert!(
|
||
author_allowed(
|
||
&RespondTo::Allowlist,
|
||
&allowlist,
|
||
EXTERNAL,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"an explicitly allowlisted external pubkey must still be accepted"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_allowlist_rejects_non_sibling_not_in_allowlist() {
|
||
let cache = cache_with_sibling();
|
||
let allowlist = HashSet::from([EXTERNAL.to_string()]);
|
||
assert!(
|
||
!author_allowed(
|
||
&RespondTo::Allowlist,
|
||
&allowlist,
|
||
STRANGER,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"a non-sibling absent from the allowlist must be dropped"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_allowlist_accepts_owner() {
|
||
let cache = cache_with_sibling();
|
||
let allowlist = HashSet::new();
|
||
assert!(
|
||
author_allowed(
|
||
&RespondTo::Allowlist,
|
||
&allowlist,
|
||
OWNER,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"the owner must always be accepted under Allowlist"
|
||
);
|
||
}
|
||
|
||
// The default `respond-to` is OwnerOnly. Under steering, "an ineligible
|
||
// author must NOT steer" is enforced *here* — author_allowed drops the
|
||
// event before it reaches the mode gate — not in the gate itself. These
|
||
// pin that invariant against the default mode.
|
||
#[tokio::test]
|
||
async fn test_owner_only_rejects_stranger_so_no_steer() {
|
||
let cache = cache_with_sibling();
|
||
assert!(
|
||
!author_allowed(
|
||
&RespondTo::OwnerOnly,
|
||
&HashSet::new(),
|
||
STRANGER,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"under the default OwnerOnly, a stranger must be dropped — so it can never reach the mode gate to steer"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_owner_only_admits_owner_and_sibling_to_steer() {
|
||
let cache = cache_with_sibling();
|
||
for (who, label) in [(OWNER, "owner"), (SIBLING, "sibling")] {
|
||
assert!(
|
||
author_allowed(
|
||
&RespondTo::OwnerOnly,
|
||
&HashSet::new(),
|
||
who,
|
||
&cache,
|
||
&dummy_rest_client()
|
||
)
|
||
.await,
|
||
"under default OwnerOnly, the {label} must be admitted so steering can fire"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod observer_chunk_coalescer_tests {
|
||
use super::*;
|
||
|
||
fn chunk_event(
|
||
seq: u64,
|
||
update_type: &str,
|
||
message_id: &str,
|
||
text: &str,
|
||
) -> observer::ObserverEvent {
|
||
observer::ObserverEvent {
|
||
seq,
|
||
timestamp: format!("2026-04-29T04:00:0{seq}Z"),
|
||
kind: "acp_read".to_string(),
|
||
agent_index: Some(0),
|
||
channel_id: Some("channel-1".to_string()),
|
||
session_id: Some("session-1".to_string()),
|
||
turn_id: Some("turn-1".to_string()),
|
||
payload: serde_json::json!({
|
||
"jsonrpc": "2.0",
|
||
"method": "session/update",
|
||
"params": {
|
||
"sessionId": "session-1",
|
||
"update": {
|
||
"sessionUpdate": update_type,
|
||
"messageId": message_id,
|
||
"content": {
|
||
"type": "text",
|
||
"text": text,
|
||
},
|
||
},
|
||
},
|
||
}),
|
||
}
|
||
}
|
||
|
||
fn non_chunk_event(seq: u64) -> observer::ObserverEvent {
|
||
observer::ObserverEvent {
|
||
seq,
|
||
timestamp: format!("2026-04-29T04:00:0{seq}Z"),
|
||
kind: "turn_started".to_string(),
|
||
agent_index: Some(0),
|
||
channel_id: Some("channel-1".to_string()),
|
||
session_id: Some("session-1".to_string()),
|
||
turn_id: Some("turn-1".to_string()),
|
||
payload: serde_json::json!({ "type": "turn_started" }),
|
||
}
|
||
}
|
||
|
||
fn chunk_text(event: &observer::ObserverEvent) -> &str {
|
||
event.payload["params"]["update"]["content"]["text"]
|
||
.as_str()
|
||
.expect("chunk text")
|
||
}
|
||
|
||
#[test]
|
||
fn coalesces_chunks_until_non_chunk_event() {
|
||
let mut coalescer = ObserverChunkCoalescer::default();
|
||
|
||
assert!(coalescer
|
||
.ingest(chunk_event(1, "agent_message_chunk", "message-1", "hello "))
|
||
.is_empty());
|
||
assert!(coalescer
|
||
.ingest(chunk_event(2, "agent_message_chunk", "message-1", "world"))
|
||
.is_empty());
|
||
|
||
let events = coalescer.ingest(non_chunk_event(3));
|
||
assert_eq!(events.len(), 2);
|
||
assert_eq!(events[0].seq, 2);
|
||
assert_eq!(chunk_text(&events[0]), "hello world");
|
||
assert_eq!(events[1].kind, "turn_started");
|
||
}
|
||
|
||
#[test]
|
||
fn keeps_independent_chunk_streams_separate() {
|
||
let mut coalescer = ObserverChunkCoalescer::default();
|
||
|
||
assert!(coalescer
|
||
.ingest(chunk_event(1, "agent_message_chunk", "message-1", "answer"))
|
||
.is_empty());
|
||
assert!(coalescer
|
||
.ingest(chunk_event(
|
||
2,
|
||
"agent_thought_chunk",
|
||
"thought-1",
|
||
"thinking"
|
||
))
|
||
.is_empty());
|
||
|
||
let events = coalescer.flush();
|
||
assert_eq!(events.len(), 2);
|
||
assert_eq!(chunk_text(&events[0]), "answer");
|
||
assert_eq!(chunk_text(&events[1]), "thinking");
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod build_mcp_servers_tests {
|
||
use super::*;
|
||
use std::sync::Mutex;
|
||
|
||
/// Env-var-touching tests must run serially — env vars are process-global.
|
||
static ENV_LOCK: Mutex<()> = Mutex::new(());
|
||
|
||
fn test_config() -> Config {
|
||
Config {
|
||
keys: nostr::Keys::generate(),
|
||
relay_url: "ws://localhost:3000".into(),
|
||
agent_command: "goose".into(),
|
||
agent_args: vec!["acp".into()],
|
||
mcp_command: "test-mcp-server".into(),
|
||
idle_timeout_secs: config::DEFAULT_IDLE_TIMEOUT_SECS,
|
||
max_turn_duration_secs: config::DEFAULT_MAX_TURN_DURATION_SECS,
|
||
agents: 1,
|
||
heartbeat_interval_secs: 0,
|
||
turn_liveness_secs: 10,
|
||
heartbeat_prompt: None,
|
||
system_prompt: None,
|
||
team_instructions: None,
|
||
initial_message: None,
|
||
subscribe_mode: config::SubscribeMode::All,
|
||
dedup_mode: config::DedupMode::Queue,
|
||
multiple_event_handling: config::MultipleEventHandling::Queue,
|
||
ignore_self: true,
|
||
kinds_override: None,
|
||
channels_override: None,
|
||
no_mention_filter: false,
|
||
config_path: std::path::PathBuf::from("./buzz-acp.toml"),
|
||
context_message_limit: 12,
|
||
max_turns_per_session: 0,
|
||
presence_enabled: true,
|
||
typing_enabled: true,
|
||
memory_enabled: false,
|
||
model: None,
|
||
permission_mode: config::PermissionMode::BypassPermissions,
|
||
respond_to: config::RespondTo::Anyone,
|
||
respond_to_allowlist: std::collections::HashSet::new(),
|
||
allowed_respond_to: vec![],
|
||
persona_env_vars: vec![],
|
||
has_generated_codex_config: false,
|
||
relay_observer: false,
|
||
agent_owner: None,
|
||
no_base_prompt: false,
|
||
base_prompt_content: None,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn session_new_mcp_server_has_required_fields() {
|
||
let config = test_config();
|
||
let servers = build_mcp_servers(&config);
|
||
assert_eq!(servers.len(), 1);
|
||
let server = &servers[0];
|
||
assert_eq!(server.name, "test-mcp-server");
|
||
|
||
let names: Vec<&str> = server.env.iter().map(|e| e.name.as_str()).collect();
|
||
assert!(
|
||
names.contains(&"BUZZ_RELAY_URL"),
|
||
"missing BUZZ_RELAY_URL; got {names:?}"
|
||
);
|
||
assert!(
|
||
names.contains(&"BUZZ_PRIVATE_KEY"),
|
||
"missing BUZZ_PRIVATE_KEY; got {names:?}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn session_new_mcp_server_forwards_buzz_auth_tag() {
|
||
let _guard = ENV_LOCK.lock().unwrap();
|
||
std::env::set_var("BUZZ_AUTH_TAG", "test-attestation-tag");
|
||
let config = test_config();
|
||
let servers = build_mcp_servers(&config);
|
||
std::env::remove_var("BUZZ_AUTH_TAG");
|
||
|
||
let server = &servers[0];
|
||
let auth_tag_env = server.env.iter().find(|e| e.name == "BUZZ_AUTH_TAG");
|
||
assert!(
|
||
auth_tag_env.is_some(),
|
||
"BUZZ_AUTH_TAG should be forwarded when set"
|
||
);
|
||
assert_eq!(auth_tag_env.unwrap().value, "test-attestation-tag");
|
||
}
|
||
|
||
#[test]
|
||
fn session_new_mcp_server_skips_empty_buzz_auth_tag() {
|
||
let _guard = ENV_LOCK.lock().unwrap();
|
||
std::env::set_var("BUZZ_AUTH_TAG", "");
|
||
let config = test_config();
|
||
let servers = build_mcp_servers(&config);
|
||
std::env::remove_var("BUZZ_AUTH_TAG");
|
||
|
||
let server = &servers[0];
|
||
let has_auth_tag = server.env.iter().any(|e| e.name == "BUZZ_AUTH_TAG");
|
||
assert!(!has_auth_tag, "empty BUZZ_AUTH_TAG should not be forwarded");
|
||
}
|
||
|
||
#[test]
|
||
fn empty_mcp_command_returns_no_servers() {
|
||
let mut config = test_config();
|
||
config.mcp_command = "".into();
|
||
let servers = build_mcp_servers(&config);
|
||
assert!(
|
||
servers.is_empty(),
|
||
"empty mcp_command should produce no MCP servers"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn absolute_path_mcp_command_uses_file_stem_as_name() {
|
||
let mut config = test_config();
|
||
config.mcp_command = "/opt/bin/my-mcp-server".into();
|
||
let servers = build_mcp_servers(&config);
|
||
assert_eq!(servers.len(), 1);
|
||
assert_eq!(servers[0].name, "my-mcp-server");
|
||
}
|
||
|
||
#[test]
|
||
fn mcp_command_with_no_stem_falls_back_to_mcp() {
|
||
// Path::new("").file_stem() returns None — exercises the unwrap_or("mcp") path.
|
||
let mut config = test_config();
|
||
config.mcp_command = "".into();
|
||
// Empty command returns no servers; test the stem logic directly.
|
||
assert_eq!(
|
||
std::path::Path::new("")
|
||
.file_stem()
|
||
.and_then(|s| s.to_str())
|
||
.unwrap_or("mcp"),
|
||
"mcp"
|
||
);
|
||
|
||
// Confirm a non-empty command with no stem (e.g. just a dot) also falls back.
|
||
config.mcp_command = ".".into();
|
||
let servers = build_mcp_servers(&config);
|
||
assert_eq!(servers.len(), 1);
|
||
assert_eq!(
|
||
servers[0].name, "mcp",
|
||
"Path::new(\".\").file_stem() is None — should fall back to \"mcp\""
|
||
);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod error_outcome_emission_tests {
|
||
//! Pins the policy that error-class outcomes surface to the activity feed
|
||
//! and never to the channel:
|
||
//!
|
||
//! - Channel silence is enforced *structurally* — `handle_prompt_result`
|
||
//! takes no relay handle, so it has no way to post a channel message. A
|
||
//! future re-introduction of channel notices would have to add the relay
|
||
//! parameter back, which these tests' construction would then refuse to
|
||
//! compile against.
|
||
//! - Feed coverage is the regression-prone half and is asserted at runtime:
|
||
//! each error outcome must emit exactly one `turn_error` observer event.
|
||
//! If any branch drops its `emit_turn_error` call, the matching test goes
|
||
//! red.
|
||
|
||
use super::*;
|
||
use crate::acp::{AcpClient, AcpError};
|
||
use crate::observer::ObserverHandle;
|
||
use crate::pool::{
|
||
AgentPool, OwnedAgent, PromptOutcome, PromptResult, PromptSource, TimeoutKind,
|
||
};
|
||
use crate::queue::{BatchEvent, FlushBatch};
|
||
use nostr::{EventBuilder, Keys, Kind};
|
||
use std::collections::HashSet;
|
||
|
||
fn test_config() -> Config {
|
||
Config {
|
||
keys: nostr::Keys::generate(),
|
||
relay_url: "ws://localhost:3000".into(),
|
||
// `true` exits cleanly, so the async respawn fails fast and
|
||
// harmlessly off the JoinSet — irrelevant to the synchronous
|
||
// feed emission under test.
|
||
agent_command: "true".into(),
|
||
agent_args: vec![],
|
||
mcp_command: "test-mcp-server".into(),
|
||
idle_timeout_secs: config::DEFAULT_IDLE_TIMEOUT_SECS,
|
||
max_turn_duration_secs: config::DEFAULT_MAX_TURN_DURATION_SECS,
|
||
agents: 1,
|
||
heartbeat_interval_secs: 0,
|
||
turn_liveness_secs: 10,
|
||
heartbeat_prompt: None,
|
||
system_prompt: None,
|
||
team_instructions: None,
|
||
initial_message: None,
|
||
subscribe_mode: config::SubscribeMode::All,
|
||
dedup_mode: config::DedupMode::Queue,
|
||
multiple_event_handling: config::MultipleEventHandling::Queue,
|
||
ignore_self: true,
|
||
kinds_override: None,
|
||
channels_override: None,
|
||
no_mention_filter: false,
|
||
config_path: std::path::PathBuf::from("./buzz-acp.toml"),
|
||
context_message_limit: 12,
|
||
max_turns_per_session: 0,
|
||
presence_enabled: true,
|
||
typing_enabled: true,
|
||
memory_enabled: false,
|
||
model: None,
|
||
permission_mode: config::PermissionMode::BypassPermissions,
|
||
respond_to: config::RespondTo::Anyone,
|
||
respond_to_allowlist: HashSet::new(),
|
||
allowed_respond_to: vec![],
|
||
persona_env_vars: vec![],
|
||
has_generated_codex_config: false,
|
||
relay_observer: false,
|
||
agent_owner: None,
|
||
no_base_prompt: false,
|
||
base_prompt_content: None,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn normalizes_agent_name_from_initialize_result() {
|
||
assert_eq!(
|
||
normalized_agent_name(&serde_json::json!({
|
||
"agentInfo": { "name": " Goose ", "version": "1.43.0" }
|
||
})),
|
||
"goose"
|
||
);
|
||
assert_eq!(
|
||
normalized_agent_name(&serde_json::json!({
|
||
"serverInfo": { "name": "buzz-agent" }
|
||
})),
|
||
"buzz-agent"
|
||
);
|
||
}
|
||
|
||
/// Spawn a real but inert agent subprocess (`cat`) so the error paths have
|
||
/// an `OwnedAgent` to move into respawn or return to the pool. The error
|
||
/// branches never talk to the subprocess.
|
||
async fn dummy_agent(index: usize) -> OwnedAgent {
|
||
OwnedAgent {
|
||
index,
|
||
acp: AcpClient::spawn("cat", &[], &[], false)
|
||
.await
|
||
.expect("spawn cat as inert agent"),
|
||
state: Default::default(),
|
||
model_capabilities: None,
|
||
desired_model: None,
|
||
model_overridden: false,
|
||
agent_name: "unknown".into(),
|
||
goose_system_prompt_supported: None,
|
||
// Error branches under test never read this; 1 is the legacy
|
||
// non-systemPrompt path, the simplest valid value.
|
||
protocol_version: 1,
|
||
}
|
||
}
|
||
|
||
/// Drive one error outcome through `handle_prompt_result` and return how
|
||
/// many `turn_error` events it emitted to the observer feed.
|
||
async fn turn_errors_emitted_for(outcome: PromptOutcome) -> usize {
|
||
let agent = dummy_agent(0).await;
|
||
let mut pool = AgentPool::from_slots(vec![None]);
|
||
|
||
// `handle_prompt_result` asserts it removes exactly one in-flight task
|
||
// for the completing agent (the slot was checked out, not idle). Mirror
|
||
// the real dispatch path by registering a TaskMeta keyed on a genuine
|
||
// `task::Id` — only obtainable from inside a spawned task.
|
||
let task_id = pool.join_set.spawn(async {}).id();
|
||
pool.task_map_mut().insert(
|
||
task_id,
|
||
crate::pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
|
||
let mut queue = EventQueue::new(config::DedupMode::Queue);
|
||
let config = test_config();
|
||
let mut heartbeat_in_flight = false;
|
||
let removed_channels = HashSet::new();
|
||
let mut crash_history = vec![SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
}];
|
||
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
|
||
let mut respawn_tasks = tokio::task::JoinSet::new();
|
||
let observer = ObserverHandle::in_process();
|
||
|
||
let result = PromptResult {
|
||
agent,
|
||
source: PromptSource::Channel(Uuid::new_v4()),
|
||
outcome,
|
||
batch: None,
|
||
};
|
||
|
||
handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
Some(observer.clone()),
|
||
None,
|
||
);
|
||
|
||
observer
|
||
.snapshot()
|
||
.iter()
|
||
.filter(|e| e.kind == "turn_error")
|
||
.count()
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn agent_exited_emits_exactly_one_feed_event() {
|
||
assert_eq!(turn_errors_emitted_for(PromptOutcome::AgentExited).await, 1);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn idle_timeout_emits_exactly_one_feed_event() {
|
||
assert_eq!(
|
||
turn_errors_emitted_for(PromptOutcome::Timeout(TimeoutKind::Idle)).await,
|
||
1
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn hard_timeout_emits_exactly_one_feed_event() {
|
||
assert_eq!(
|
||
turn_errors_emitted_for(PromptOutcome::Timeout(TimeoutKind::Hard)).await,
|
||
1
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn cancel_drain_timeout_emits_exactly_one_feed_event() {
|
||
assert_eq!(
|
||
turn_errors_emitted_for(PromptOutcome::CancelDrainTimeout(
|
||
std::time::Duration::from_secs(5)
|
||
))
|
||
.await,
|
||
1
|
||
);
|
||
}
|
||
|
||
/// idle_timeout outcome_label is "idle_timeout"; hard_timeout is "hard_timeout".
|
||
#[tokio::test]
|
||
async fn timeout_outcome_labels_differ() {
|
||
let check_label = |outcome: PromptOutcome, expected_label: &'static str| async move {
|
||
let agent = dummy_agent(0).await;
|
||
let mut pool = AgentPool::from_slots(vec![None]);
|
||
let task_id = pool.join_set.spawn(async {}).id();
|
||
pool.task_map_mut().insert(
|
||
task_id,
|
||
crate::pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
let mut queue = EventQueue::new(config::DedupMode::Queue);
|
||
let config = test_config();
|
||
let mut heartbeat_in_flight = false;
|
||
let removed_channels = HashSet::new();
|
||
let mut crash_history = vec![SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
}];
|
||
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
|
||
let mut respawn_tasks = tokio::task::JoinSet::new();
|
||
let observer = ObserverHandle::in_process();
|
||
let result = PromptResult {
|
||
agent,
|
||
source: PromptSource::Channel(Uuid::new_v4()),
|
||
outcome,
|
||
batch: None,
|
||
};
|
||
handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
Some(observer.clone()),
|
||
None,
|
||
);
|
||
let events = observer.snapshot();
|
||
let turn_error = events.iter().find(|e| e.kind == "turn_error").unwrap();
|
||
assert_eq!(
|
||
turn_error.payload["outcome"].as_str().unwrap(),
|
||
expected_label
|
||
);
|
||
};
|
||
check_label(PromptOutcome::Timeout(TimeoutKind::Idle), "idle_timeout").await;
|
||
check_label(PromptOutcome::Timeout(TimeoutKind::Hard), "hard_timeout").await;
|
||
check_label(
|
||
PromptOutcome::CancelDrainTimeout(std::time::Duration::from_secs(5)),
|
||
"cancel_drain_timeout",
|
||
)
|
||
.await;
|
||
}
|
||
|
||
/// hard-cap timeout dead-letters immediately (no requeue); idle timeout is requeued.
|
||
#[tokio::test]
|
||
async fn hard_timeout_not_requeued_idle_timeout_is_requeued() {
|
||
let make_batch = || {
|
||
let keys = Keys::generate();
|
||
let event = EventBuilder::new(Kind::Custom(9), "test")
|
||
.sign_with_keys(&keys)
|
||
.unwrap();
|
||
FlushBatch {
|
||
channel_id: Uuid::new_v4(),
|
||
events: vec![BatchEvent {
|
||
event,
|
||
prompt_tag: "test".into(),
|
||
received_at: std::time::Instant::now(),
|
||
}],
|
||
cancelled_events: vec![],
|
||
cancel_reason: None,
|
||
}
|
||
};
|
||
|
||
// Returns (pending_channels, queued_event_count_for_channel).
|
||
let run = |outcome: PromptOutcome, batch: FlushBatch| async move {
|
||
let channel_id = batch.channel_id;
|
||
let agent = dummy_agent(0).await;
|
||
let mut pool = AgentPool::from_slots(vec![None]);
|
||
let task_id = pool.join_set.spawn(async {}).id();
|
||
pool.task_map_mut().insert(
|
||
task_id,
|
||
crate::pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
let mut queue = EventQueue::new(config::DedupMode::Queue);
|
||
let config = test_config();
|
||
let mut heartbeat_in_flight = false;
|
||
let removed_channels = HashSet::new();
|
||
let mut crash_history = vec![SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
}];
|
||
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
|
||
let mut respawn_tasks = tokio::task::JoinSet::new();
|
||
let result = PromptResult {
|
||
agent,
|
||
source: PromptSource::Channel(channel_id),
|
||
outcome,
|
||
batch: Some(batch),
|
||
};
|
||
handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
None,
|
||
None,
|
||
);
|
||
(
|
||
queue.pending_channels(),
|
||
queue.queued_event_count(&channel_id),
|
||
)
|
||
};
|
||
|
||
// Hard timeout: batch must NOT be requeued (dead-lettered immediately).
|
||
let hard_batch = make_batch();
|
||
let (hard_channels, hard_events) =
|
||
run(PromptOutcome::Timeout(TimeoutKind::Hard), hard_batch).await;
|
||
assert_eq!(
|
||
hard_channels, 0,
|
||
"hard-cap timeout must not requeue the batch"
|
||
);
|
||
assert_eq!(hard_events, 0, "hard-cap timeout must drop all events");
|
||
|
||
// Idle timeout: batch IS requeued (first attempt, not yet dead-lettered).
|
||
let idle_batch = make_batch();
|
||
let (idle_channels, idle_events) =
|
||
run(PromptOutcome::Timeout(TimeoutKind::Idle), idle_batch).await;
|
||
assert_eq!(
|
||
idle_channels, 1,
|
||
"idle timeout must requeue the batch for retry"
|
||
);
|
||
assert_eq!(
|
||
idle_events, 1,
|
||
"idle timeout must preserve the event for retry"
|
||
);
|
||
}
|
||
|
||
/// Cancel-drain-timeout batches are requeued as cancelled (merge into the
|
||
/// next flush, `CancelReason` preserved) — never dead-lettered like a real
|
||
/// hard-cap. The agent itself is NOT returned to the idle pool: it is
|
||
/// handed to `spawn_respawn_task` instead, mirroring a fatal `Timeout`.
|
||
///
|
||
/// This reproduces the full steer-fallback incident, not just the
|
||
/// original batch in isolation: the steer ack handler already released
|
||
/// the new triggering event back to `queue` (`lib.rs`'s
|
||
/// `ExpectedRunIdMissing` path) before the cancel-drain expiry fires. The
|
||
/// next `flush_next()` must merge the surviving original event (via
|
||
/// `cancelled_events`) with that already-queued new event (via `events`)
|
||
/// exactly once each — proving no loss and no duplication.
|
||
#[tokio::test]
|
||
async fn cancel_drain_timeout_requeues_batch_and_does_not_return_agent() {
|
||
let keys = Keys::generate();
|
||
let original_event = EventBuilder::new(Kind::Custom(9), "original")
|
||
.sign_with_keys(&keys)
|
||
.unwrap();
|
||
let new_event = EventBuilder::new(Kind::Custom(9), "new")
|
||
.sign_with_keys(&keys)
|
||
.unwrap();
|
||
assert_ne!(
|
||
original_event.id, new_event.id,
|
||
"test fixture must use two distinct events"
|
||
);
|
||
let channel_id = Uuid::new_v4();
|
||
let batch = FlushBatch {
|
||
channel_id,
|
||
events: vec![BatchEvent {
|
||
event: original_event.clone(),
|
||
prompt_tag: "test".into(),
|
||
received_at: std::time::Instant::now(),
|
||
}],
|
||
cancelled_events: vec![],
|
||
cancel_reason: Some(CancelReason::Steer),
|
||
};
|
||
|
||
let agent = dummy_agent(0).await;
|
||
let mut pool = AgentPool::from_slots(vec![None]);
|
||
let task_id = pool.join_set.spawn(async {}).id();
|
||
pool.task_map_mut().insert(
|
||
task_id,
|
||
crate::pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
let mut queue = EventQueue::new(config::DedupMode::Queue);
|
||
// The steer ack handler releases the new event to the queue BEFORE
|
||
// signaling the fallback ControlSignal::Steer that ultimately times
|
||
// out on drain — so it is already queued by the time
|
||
// handle_prompt_result runs.
|
||
queue.push(QueuedEvent {
|
||
channel_id,
|
||
event: new_event.clone(),
|
||
received_at: std::time::Instant::now(),
|
||
prompt_tag: "test".into(),
|
||
});
|
||
let config = test_config();
|
||
let mut heartbeat_in_flight = false;
|
||
let removed_channels = HashSet::new();
|
||
let mut crash_history = vec![SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
}];
|
||
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
|
||
let mut respawn_tasks = tokio::task::JoinSet::new();
|
||
let observer = ObserverHandle::in_process();
|
||
let grace = std::time::Duration::from_secs(5);
|
||
let result = PromptResult {
|
||
agent,
|
||
source: PromptSource::Channel(channel_id),
|
||
outcome: PromptOutcome::CancelDrainTimeout(grace),
|
||
batch: Some(batch),
|
||
};
|
||
|
||
handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
Some(observer.clone()),
|
||
None,
|
||
);
|
||
|
||
// Batch preserved as a cancelled merge, not dead-lettered — same
|
||
// treatment as a normal `Cancelled` outcome. `handle_prompt_result`
|
||
// already called `mark_complete` internally, releasing the channel.
|
||
// `flush_next()` must merge the already-queued new event with the
|
||
// preserved original: each exactly once, in the correct bucket.
|
||
let requeued = queue.flush_next().expect("batch must be requeued");
|
||
assert_eq!(
|
||
requeued.events.len(),
|
||
1,
|
||
"exactly one new event must be in the regular events bucket"
|
||
);
|
||
assert_eq!(
|
||
requeued.events[0].event.id, new_event.id,
|
||
"the regular events bucket must hold the new (already-queued) event"
|
||
);
|
||
assert_eq!(
|
||
requeued.cancelled_events.len(),
|
||
1,
|
||
"exactly one original event must be in the cancelled_events bucket"
|
||
);
|
||
assert_eq!(
|
||
requeued.cancelled_events[0].event.id, original_event.id,
|
||
"the cancelled_events bucket must hold the original (interrupted) event"
|
||
);
|
||
assert_ne!(
|
||
requeued.events[0].event.id, requeued.cancelled_events[0].event.id,
|
||
"the new and original events must not be the same event"
|
||
);
|
||
assert_eq!(
|
||
requeued.cancel_reason,
|
||
Some(CancelReason::Steer),
|
||
"CancelReason must ride through to the requeued batch"
|
||
);
|
||
|
||
// Agent must NOT be back in the idle pool — it was handed to respawn.
|
||
assert_eq!(
|
||
pool.live_count(),
|
||
0,
|
||
"agent must not be returned to the pool after a cancel-drain timeout"
|
||
);
|
||
assert_eq!(
|
||
respawn_tasks.len(),
|
||
1,
|
||
"a respawn task must be spawned for the poisoned agent"
|
||
);
|
||
|
||
// The observer payload must be fate-neutral: it names the grace and
|
||
// the process replacement, and must NOT claim the batch was
|
||
// preserved — that claim is false for explicit Stop/removed-channel
|
||
// drops (see the sibling dropped-Stop test below), so the same
|
||
// wording is used regardless of fate.
|
||
let events = observer.snapshot();
|
||
let turn_error = events
|
||
.iter()
|
||
.find(|e| e.kind == "turn_error")
|
||
.expect("exactly one turn_error event must be emitted");
|
||
assert_eq!(
|
||
turn_error.payload["outcome"].as_str().unwrap(),
|
||
"cancel_drain_timeout"
|
||
);
|
||
assert_eq!(
|
||
turn_error.payload["error"].as_str().unwrap(),
|
||
format!("Agent did not stop within {grace:?} after cancellation; the agent process is being replaced."),
|
||
"observer message must name the actual grace and must not claim preservation"
|
||
);
|
||
assert_eq!(
|
||
events.iter().filter(|e| e.kind == "turn_error").count(),
|
||
1,
|
||
"exactly one turn_error event must be emitted"
|
||
);
|
||
}
|
||
|
||
/// Explicit Stop (`ControlSignal::Cancel`) on cancel-drain expiry drops
|
||
/// the triggering batch — `requeue_cancelled_batch` returns `None` for
|
||
/// `Cancel`/`Rotate`. The observer payload must be the SAME fate-neutral
|
||
/// text as the preserved-Steer case above: it must never claim work was
|
||
/// preserved when it was intentionally discarded. The poisoned agent is
|
||
/// still respawned exactly as in the preserved case.
|
||
#[tokio::test]
|
||
async fn cancel_drain_timeout_dropped_stop_batch_none_same_neutral_payload() {
|
||
let agent = dummy_agent(0).await;
|
||
let mut pool = AgentPool::from_slots(vec![None]);
|
||
let task_id = pool.join_set.spawn(async {}).id();
|
||
pool.task_map_mut().insert(
|
||
task_id,
|
||
crate::pool::TaskMeta {
|
||
agent_index: 0,
|
||
channel_id: None,
|
||
recoverable_batch: None,
|
||
control_tx: None,
|
||
steer_tx: None,
|
||
},
|
||
);
|
||
let mut queue = EventQueue::new(config::DedupMode::Queue);
|
||
let config = test_config();
|
||
let mut heartbeat_in_flight = false;
|
||
let removed_channels = HashSet::new();
|
||
let mut crash_history = vec![SlotCircuit {
|
||
crash_times: Vec::new(),
|
||
open_until: None,
|
||
respawn_in_flight: false,
|
||
}];
|
||
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
|
||
let mut respawn_tasks = tokio::task::JoinSet::new();
|
||
let observer = ObserverHandle::in_process();
|
||
let grace = std::time::Duration::from_secs(5);
|
||
let result = PromptResult {
|
||
agent,
|
||
source: PromptSource::Channel(Uuid::new_v4()),
|
||
outcome: PromptOutcome::CancelDrainTimeout(grace),
|
||
// Explicit Stop already dropped the batch upstream in
|
||
// `classify_control_cancel_failure` — `handle_prompt_result`
|
||
// never sees one to requeue.
|
||
batch: None,
|
||
};
|
||
|
||
handle_prompt_result(
|
||
&mut pool,
|
||
&mut queue,
|
||
&config,
|
||
result,
|
||
&mut heartbeat_in_flight,
|
||
&removed_channels,
|
||
&mut crash_history,
|
||
&respawn_tx,
|
||
&mut respawn_tasks,
|
||
Some(observer.clone()),
|
||
None,
|
||
);
|
||
|
||
// No batch to merge — the queue has nothing pending for any channel.
|
||
assert_eq!(
|
||
queue.pending_channels(),
|
||
0,
|
||
"a dropped Stop batch must not leave anything queued"
|
||
);
|
||
|
||
// Same respawn treatment as the preserved case: never returned idle.
|
||
assert_eq!(
|
||
pool.live_count(),
|
||
0,
|
||
"agent must not be returned to the pool after a cancel-drain timeout"
|
||
);
|
||
assert_eq!(
|
||
respawn_tasks.len(),
|
||
1,
|
||
"a respawn task must be spawned for the poisoned agent"
|
||
);
|
||
|
||
// The observer payload is byte-identical to the preserved-Steer case:
|
||
// fate-neutral, naming the grace, with no preservation claim.
|
||
let events = observer.snapshot();
|
||
let turn_error = events
|
||
.iter()
|
||
.find(|e| e.kind == "turn_error")
|
||
.expect("exactly one turn_error event must be emitted");
|
||
assert_eq!(
|
||
turn_error.payload["outcome"].as_str().unwrap(),
|
||
"cancel_drain_timeout"
|
||
);
|
||
assert_eq!(
|
||
turn_error.payload["error"].as_str().unwrap(),
|
||
format!("Agent did not stop within {grace:?} after cancellation; the agent process is being replaced."),
|
||
"observer message must be fate-neutral even though the batch was dropped"
|
||
);
|
||
assert_eq!(
|
||
events.iter().filter(|e| e.kind == "turn_error").count(),
|
||
1,
|
||
"exactly one turn_error event must be emitted"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn transport_error_emits_exactly_one_feed_event() {
|
||
let io = AcpError::Io(std::io::Error::other("pipe broke"));
|
||
assert_eq!(turn_errors_emitted_for(PromptOutcome::Error(io)).await, 1);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn application_error_emits_exactly_one_feed_event() {
|
||
let app = AcpError::IdleTimeout(std::time::Duration::from_secs(1));
|
||
assert_eq!(turn_errors_emitted_for(PromptOutcome::Error(app)).await, 1);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod observer_payload_trim_tests {
|
||
use super::*;
|
||
|
||
fn event_with_payload(kind: &str, payload: serde_json::Value) -> observer::ObserverEvent {
|
||
observer::ObserverEvent {
|
||
seq: 1,
|
||
timestamp: "2026-06-16T00:00:00Z".to_string(),
|
||
kind: kind.to_string(),
|
||
agent_index: Some(0),
|
||
channel_id: Some("11111111-1111-1111-1111-111111111111".to_string()),
|
||
session_id: Some("sess-1".to_string()),
|
||
turn_id: Some("turn-1".to_string()),
|
||
payload,
|
||
}
|
||
}
|
||
|
||
fn serialized(event: &observer::ObserverEvent) -> String {
|
||
serde_json::to_string(event).unwrap()
|
||
}
|
||
|
||
#[test]
|
||
fn test_under_budget_frame_passes_through_byte_identical() {
|
||
let mut event = event_with_payload("acp_read", serde_json::json!({ "body": "small" }));
|
||
let before = serialized(&event);
|
||
fit_observer_event_to_budget(&mut event);
|
||
assert_eq!(
|
||
serialized(&event),
|
||
before,
|
||
"under-budget frame must not be mutated"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_single_giant_leaf_is_elided_to_fit_with_envelope_intact() {
|
||
let big = "x".repeat(100_000);
|
||
let mut event = event_with_payload("acp_read", serde_json::json!({ "body": big }));
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(
|
||
serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN,
|
||
"frame must fit after trimming"
|
||
);
|
||
// Envelope intact.
|
||
assert_eq!(event.kind, "acp_read");
|
||
assert_eq!(event.turn_id.as_deref(), Some("turn-1"));
|
||
assert_eq!(
|
||
event.channel_id.as_deref(),
|
||
Some("11111111-1111-1111-1111-111111111111")
|
||
);
|
||
assert_eq!(event.seq, 1);
|
||
|
||
let leaf = event.payload["body"].as_str().unwrap();
|
||
assert!(
|
||
leaf.starts_with(&"x".repeat(OBSERVER_LEAF_RETAIN_BYTES)),
|
||
"head retained"
|
||
);
|
||
assert!(
|
||
leaf.ends_with(&"x".repeat(OBSERVER_LEAF_RETAIN_BYTES)),
|
||
"tail retained"
|
||
);
|
||
// N in the marker is RAW bytes removed: original len minus retained len.
|
||
let removed = 100_000 - leaf.chars().filter(|c| *c == 'x').count();
|
||
assert!(
|
||
leaf.contains(&format!("…[elided {removed} bytes]…")),
|
||
"marker reports raw bytes removed"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_multi_block_prompt_retains_every_section_header_after_elision() {
|
||
// The real session/prompt fix: format_prompt now emits one block per
|
||
// section, so the observer payload is params.prompt = [{text: "[Base]…"},
|
||
// {text: "[Agent Memory — core]…"}, … {text: "[Buzz event: …]…<huge>"}].
|
||
// An oversized section is its own leaf, so eliding its body keeps the
|
||
// leaf's head-3000 (which begins with the section's [Header] line) — every
|
||
// header survives, so the desktop "Prompt context" panel counts them all.
|
||
// This is the regression the single-fat-leaf shape caused (the trailing
|
||
// [Buzz event] header fell into the elided middle and the count collapsed
|
||
// to 1).
|
||
let sections = [
|
||
"[Base]\nyou are a helpful agent".to_string(),
|
||
"[System]\npersona text".to_string(),
|
||
"[Agent Memory — core]\nremember this".to_string(),
|
||
"[Context]\nScope: thread".to_string(),
|
||
// The triggering event body, oversized on its own.
|
||
format!("[Buzz event: @mention]\nContent: {}", "E".repeat(90_000)),
|
||
];
|
||
let block_refs: Vec<&str> = sections.iter().map(String::as_str).collect();
|
||
// Mirror the wire shape build_prompt_params produces: each block is its
|
||
// own {type:"text", text} leaf under params.prompt.
|
||
let prompt_blocks: Vec<serde_json::Value> = block_refs
|
||
.iter()
|
||
.map(|text| serde_json::json!({ "type": "text", "text": text }))
|
||
.collect();
|
||
let mut event = event_with_payload(
|
||
"acp_write",
|
||
serde_json::json!({
|
||
"method": "session/prompt",
|
||
"params": { "sessionId": "sess-1", "prompt": prompt_blocks },
|
||
}),
|
||
);
|
||
assert!(
|
||
serialized(&event).len() > OBSERVER_MAX_PLAINTEXT_LEN,
|
||
"precondition: oversized event body pushes the frame over the cap"
|
||
);
|
||
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(
|
||
serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN,
|
||
"frame must fit after trimming"
|
||
);
|
||
let blocks = event.payload["params"]["prompt"]
|
||
.as_array()
|
||
.expect("prompt array survives");
|
||
let texts: Vec<&str> = blocks.iter().map(|b| b["text"].as_str().unwrap()).collect();
|
||
for header in [
|
||
"[Base]",
|
||
"[System]",
|
||
"[Agent Memory — core]",
|
||
"[Context]",
|
||
"[Buzz event: @mention]",
|
||
] {
|
||
assert!(
|
||
texts.iter().any(|t| t.starts_with(header)),
|
||
"section header {header} must survive at the head of its own block"
|
||
);
|
||
}
|
||
// The oversized event body was elided in place (header kept, middle cut).
|
||
let event_block = texts
|
||
.iter()
|
||
.find(|t| t.starts_with("[Buzz event: @mention]"))
|
||
.unwrap();
|
||
assert!(
|
||
event_block.contains("…[elided"),
|
||
"the oversized event body is elided, not dropped"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_multi_leaf_elides_largest_shrinkable_first_and_stops_when_it_fits() {
|
||
// One leaf alone over the cap; a second smaller-but-still-large leaf.
|
||
// Eliding the biggest should suffice, leaving the smaller intact.
|
||
let mut event = event_with_payload(
|
||
"acp_write",
|
||
serde_json::json!({
|
||
"huge": "a".repeat(90_000),
|
||
"medium": "b".repeat(20_000),
|
||
}),
|
||
);
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN);
|
||
assert!(
|
||
event.payload["huge"].as_str().unwrap().contains("…[elided"),
|
||
"the largest leaf is elided"
|
||
);
|
||
assert_eq!(
|
||
event.payload["medium"].as_str().unwrap().len(),
|
||
20_000,
|
||
"the smaller leaf is left untouched once the frame fits"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_coalesced_chunk_nested_leaf_is_reached_by_recursive_walk() {
|
||
// The coalesced-chunk big leaf lives at params.update.content.text,
|
||
// not a top-level field — the walk must recurse to reach it.
|
||
let big = "z".repeat(80_000);
|
||
let mut event = event_with_payload(
|
||
"session_update",
|
||
serde_json::json!({
|
||
"params": {
|
||
"update": {
|
||
"sessionUpdate": "agent_message_chunk",
|
||
"content": { "text": big }
|
||
}
|
||
}
|
||
}),
|
||
);
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN);
|
||
let text = event.payload["params"]["update"]["content"]["text"]
|
||
.as_str()
|
||
.unwrap();
|
||
assert!(text.contains("…[elided"), "nested leaf was elided");
|
||
}
|
||
|
||
#[test]
|
||
fn test_many_medium_leaves_terminate_via_stub() {
|
||
// Many leaves each too small to shrink on their own (below 2x retain),
|
||
// collectively over the cap. No leaf can strictly shrink, so the trimmer
|
||
// must terminate via the stub rather than loop forever.
|
||
let leaf = "m".repeat(OBSERVER_LEAF_RETAIN_BYTES); // shorter than head+tail → cannot shrink
|
||
let items: Vec<serde_json::Value> = (0..40)
|
||
.map(|_| serde_json::Value::String(leaf.clone()))
|
||
.collect();
|
||
let mut event = event_with_payload("acp_read", serde_json::json!({ "items": items }));
|
||
assert!(
|
||
serialized(&event).len() > OBSERVER_MAX_PLAINTEXT_LEN,
|
||
"precondition: frame is over the cap"
|
||
);
|
||
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN);
|
||
assert_eq!(
|
||
event.payload["elided"].as_str().unwrap(),
|
||
"acp_read payload too large",
|
||
"fell back to the stub"
|
||
);
|
||
assert!(event.payload.get("originalBytes").is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn test_leaf_too_small_to_shrink_is_not_mutated() {
|
||
// A frame already under budget whose only leaf is below the shrink floor:
|
||
// nothing should change. (Under-budget short-circuits, and even if forced,
|
||
// leaf_shrinks would reject it.)
|
||
let short = "s".repeat(OBSERVER_LEAF_RETAIN_BYTES); // == head; cannot strictly shrink
|
||
assert!(
|
||
!leaf_shrinks(&short),
|
||
"a leaf at the retain floor must not shrink"
|
||
);
|
||
let longer = "L".repeat(OBSERVER_LEAF_RETAIN_BYTES * 2 + 100);
|
||
assert!(leaf_shrinks(&longer), "a clearly larger leaf must shrink");
|
||
}
|
||
|
||
#[test]
|
||
fn test_utf8_multibyte_leaf_elides_on_char_boundary() {
|
||
// A leaf of 3-byte chars (… = U+2026) — eliding must land on char
|
||
// boundaries and never panic or produce invalid UTF-8.
|
||
let big: String = "…".repeat(40_000); // 120_000 bytes
|
||
let mut event = event_with_payload("acp_read", serde_json::json!({ "body": big }));
|
||
fit_observer_event_to_budget(&mut event);
|
||
|
||
assert!(serialized(&event).len() <= OBSERVER_MAX_PLAINTEXT_LEN);
|
||
let leaf = event.payload["body"].as_str().unwrap();
|
||
// Valid UTF-8 by construction (it's a &str); confirm head/tail are whole
|
||
// multi-byte chars and the marker is present.
|
||
assert!(leaf.starts_with('…'));
|
||
assert!(leaf.ends_with('…'));
|
||
assert!(leaf.contains("[elided"));
|
||
}
|
||
}
|