//! Agent pool — owns N AcpClient instances and dispatches prompt tasks. //! //! # Mental model //! //! ```text //! AgentPool //! ├── agents: Vec> ← idle agents sit here //! ├── join_set: JoinSet<()> ← in-flight tasks //! ├── task_map: HashMap ← panic recovery metadata //! └── result_tx/rx: mpsc channel ← tasks return agents here //! //! Dispatch: //! try_claim() → OwnedAgent (removed from slot) //! spawn run_prompt_task(agent, ...) into join_set //! task sends PromptResult { agent, outcome } via result_tx //! rx_and_join_set() → poll result_rx for PromptResult //! return_agent(agent) → puts agent back in slot //! ``` //! //! `AcpClient` is NOT Clone — ownership moves out on claim and back on return. use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Duration; use tokio::sync::mpsc; use tokio::task::JoinSet; use tokio::time::timeout; use uuid::Uuid; use crate::acp::{ extract_model_config_options, extract_model_state, model_in_catalog, resolve_model_switch_method, AcpClient, AcpError, McpServer, ModelSwitchMethod, StopReason, }; use crate::config::{DedupMode, PermissionMode}; use crate::observer; use crate::queue::{ CancelReason, ContextMessage, ConversationContext, FlushBatch, PromptChannelInfo, PromptProfile, PromptProfileLookup, }; use crate::relay::{ChannelInfo, RestClient}; // FlushBatch and BatchEvent derive Clone (added in queue.rs) so we can store // a recoverable copy in TaskMeta for panic recovery in Queue mode. /// Metadata stored per in-flight task for panic recovery. pub struct TaskMeta { pub agent_index: usize, pub channel_id: Option, /// Clone of batch for Queue mode panic recovery. pub recoverable_batch: Option, /// Control signal for the in-flight prompt task. /// `None` for heartbeat tasks (not controllable) and after signal is consumed. pub control_tx: Option>, /// Steer request channel for non-cancelling mid-turn delivery. /// Capacity-1; `try_send` from the main loop fails on `Full`/`Closed`, /// in which case the caller must fall back to the universal /// `ControlSignal::Steer` cancel+merge path. `None` for heartbeat /// tasks only — all prompt tasks install a steer channel regardless /// of the agent's name. pub steer_tx: Option>, } /// Agent-level model capabilities. Populated on first session creation. /// The catalog is the same across all sessions for a given agent process. /// Fields are read by the desktop's `get_agent_models` Tauri command (Phase 3). #[allow(dead_code)] // Scaffolding for desktop integration — fields read via serde. pub struct AgentModelCapabilities { /// Stable: configOptions with category "model" from session/new. pub config_options_raw: Vec, /// Unstable: SessionModelState from session/new. pub available_models_raw: Option, } /// Per-channel session IDs and turn counters. /// /// Separated from `OwnedAgent` so the state machine is testable without /// spawning a real agent subprocess. #[derive(Default)] pub struct SessionState { /// channel_id → session_id pub sessions: HashMap, pub heartbeat_session: Option, /// Per-channel turn counters for proactive session rotation. /// Incremented on each successful prompt; reset when the session is rotated. pub turn_counts: HashMap, /// Turn counter for the heartbeat session. pub heartbeat_turn_count: u32, /// channel_id → rendered NIP-AE core prompt section, populated once at /// session creation per Tyler's spec (no mid-session refresh). pub core_sections: HashMap, } impl SessionState { /// Invalidate the session (and turn counter) for a specific prompt source. pub fn invalidate(&mut self, source: &PromptSource) { match source { PromptSource::Channel(cid) => { self.invalidate_channel(cid); } PromptSource::Heartbeat => { self.heartbeat_session = None; self.heartbeat_turn_count = 0; } } } /// Invalidate a single channel's session and turn counter. /// Returns `true` if the channel had an active session. pub fn invalidate_channel(&mut self, channel_id: &Uuid) -> bool { self.turn_counts.remove(channel_id); self.core_sections.remove(channel_id); self.sessions.remove(channel_id).is_some() } /// Invalidate all sessions and turn counters (e.g. after agent exit). pub fn invalidate_all(&mut self) { self.sessions.clear(); self.turn_counts.clear(); self.heartbeat_session = None; self.heartbeat_turn_count = 0; self.core_sections.clear(); } #[cfg(test)] fn has_channel_state(&self, channel_id: &Uuid) -> bool { self.sessions.contains_key(channel_id) || self.turn_counts.contains_key(channel_id) || self.core_sections.contains_key(channel_id) } } /// An agent with its session state, owned by the pool or a running task. pub struct OwnedAgent { pub index: usize, pub acp: AcpClient, pub state: SessionState, /// Model catalog from first session/new. None until first session created. pub model_capabilities: Option, /// Desired model ID (from `Config.model`). Applied after every `session_new_full()`. pub desired_model: Option, /// Whether `desired_model` was set by a live `SwitchModel` control signal /// (as opposed to being derived from config/persona at spawn). Used by the /// desktop reader to distinguish a genuine runtime override from a stale /// session whose persona model was edited. Reset on spawn/restart. pub model_overridden: bool, /// Protocol version reported by the agent in its initialize response. /// Agents declaring >= 2 support `systemPrompt` in session/new. pub protocol_version: u32, } /// Pool of agents with take-and-return ownership semantics. /// /// Agents are either idle (sitting in `agents[i]`) or checked out /// (running inside a spawned task). The `task_map` tracks in-flight /// tasks for panic recovery. pub struct AgentPool { agents: Vec>, result_tx: mpsc::UnboundedSender, result_rx: mpsc::UnboundedReceiver, pub join_set: JoinSet<()>, task_map: HashMap, } /// Result returned by a completed prompt task. pub struct PromptResult { pub agent: OwnedAgent, pub source: PromptSource, pub outcome: PromptOutcome, /// Present on failure in Queue mode, for requeue. pub batch: Option, } /// Whether the prompt came from a channel event or a heartbeat. #[derive(Debug)] pub enum PromptSource { Channel(Uuid), Heartbeat, } /// Apply state effects for Race 1, where a control signal arrives just after the /// prompt completed naturally. The prompt result has already been consumed by /// `select!`, so the harness must synthesize a successful result while still /// honoring any load-bearing control signal semantics. fn apply_completed_before_control_signal( state: &mut SessionState, source: &PromptSource, control_signal: &ControlSignal, ) { // Rotate and SwitchModel both invalidate so the next turn creates a fresh // session. For SwitchModel the caller has already set `desired_model`, so // the fresh session applies the new model on its next creation. if matches!( control_signal, ControlSignal::Rotate | ControlSignal::SwitchModel(_) ) { state.invalidate(source); } } /// Control signal for an in-flight channel turn. /// /// Not `Copy`: `SwitchModel` carries an owned `String`. Callers must clone when /// a value is needed after a move, or match by reference. #[derive(Clone, Debug, Eq, PartialEq)] pub enum ControlSignal { /// Stop the current turn and drop its triggering batch. Cancel, /// Stop the current turn and requeue its triggering batch for a merged /// re-prompt framed as a **supersede**: the new request replaces the old. Interrupt, /// Stop the current turn and requeue its triggering batch for a merged /// re-prompt framed as a **steer**: a message arrived while the agent was /// working; it should continue its work and incorporate the message if /// relevant, not treat it as a replacement task. This is the default /// mid-turn delivery path (see [`MultipleEventHandling::Steer`]). Steer, /// Stop the current turn and drop its triggering batch. The session is /// invalidated just like cancel; the next turn creates a fresh session. Rotate, /// Switch the agent's model, then requeue the triggering batch so it /// re-runs on a fresh session under the new model. The model lands by /// setting `OwnedAgent::desired_model` before invalidation; the requeued /// turn re-creates the session and re-applies `desired_model`. Runtime-only /// — never persisted, gone on restart/respawn. SwitchModel(String), } /// Goose-native non-cancelling steer request, sent from the main loop to an /// in-flight prompt task's read loop via a capacity-1 mpsc channel. /// /// The read loop owns the `AcpClient`'s reader/writer for the duration of the /// turn, so we cannot drive a steer write from the main thread directly. The /// main loop carries the steer prompt body (already framed by /// `queue::native_steer_framing()` + `queue::format_event_block`); the read /// loop completes `sessionId` (lexical) and `expectedRunId` /// (`AcpClient::active_run_id` at write time) when it actually emits the /// JSON-RPC request. The main loop awaits a `SteerAck` on the `ack_tx` /// oneshot. /// /// ## Why the read loop fills params, not the main loop /// /// `expectedRunId` is a *moving target*: the read loop updates /// `self.active_run_id` as goose emits `session/update` notifications, and /// the steer is rejected if the supplied id doesn't match the *current* run. /// A snapshot taken at dispatch (or at mode-gate time) can be stale by the /// time the read loop actually writes the steer line. Filling params at /// write time uses the freshest possible run id and is correct-by- /// construction on the one field whose freshness the protocol checks. /// `sessionId` is in lexical scope inside the read loop's caller /// (`session_prompt_blocks_with_idle_timeout`), so no plumbing is required /// for that — only a function parameter pass-through. /// /// If `active_run_id` is `None` at write time (no `session/update` seen yet /// — e.g. agents that never emit run-id metadata), the steer cannot form a /// valid `expectedRunId` and the read loop acks /// [`SteerError::ExpectedRunIdMissing`]. The main loop maps this to the /// "Err-before-pending" bucket: no withhold/mark was established at /// `pool::send_steer` time because the request was rejected before any /// write, so the watcher only needs to release nothing and fall back to the /// universal `ControlSignal::Steer` cancel+merge path. pub struct SteerRequest { /// Prompt body text blocks. Each entry becomes one `text` content /// block in `params.prompt`. Built by the main loop via /// `queue::native_steer_framing()` + `queue::format_event_block` so /// the wording cannot drift from the cancel+merge fallback path. pub prompt_blocks: Vec, /// Oneshot for the read loop to report the outcome. pub ack_tx: tokio::sync::oneshot::Sender, } /// Why a goose-native steer failed. /// /// String and integer fields are intentionally `Debug`-only — read by /// `tracing` macros in the main loop's `PoolEvent::SteerAck` arm via /// `?ack`. The dead-code lint can't see that path because it doesn't /// trace through `Debug` derives, hence the `#[allow]`. #[allow(dead_code)] #[derive(Debug)] pub enum SteerError { /// The agent returned a JSON-RPC error response to the steer request. /// /// `code` is the JSON-RPC error code: /// - `-32601` (`method_not_found`): the agent does not implement the /// steer extension. The main loop should fire the cancel+merge /// fallback so the message still reaches the agent. /// - Any other code: the write landed and the agent rejected it at the /// application level (e.g. wrong run id). Release the withheld event /// for normal dispatch; do NOT fire the fallback — the turn is still /// running or just ended. AgentError { code: i64, message: String }, /// Transport-level failure: write error, read EOF, JSON-RPC framing /// violation, etc. The string carries the underlying `AcpError`'s display. Transport(String), /// At steer-write time `AcpClient::active_run_id` was `None`, so the /// read loop couldn't form a valid `expectedRunId`. The read loop drops /// the request without writing anything; the main loop should release /// any withheld event and fall back to the universal cancel+merge /// `ControlSignal::Steer` path. This is in the same "Err-before-pending" /// bucket as `Transport` write failures: no in-process state was /// established, so no in-process cleanup is needed. ExpectedRunIdMissing, /// The read loop never got to dispatch the steer because the prompt /// completed first. Delivery state for the underlying message is /// unknown after prompt completion — the main loop must treat this as /// "release the withheld event so normal dispatch handles it" with no /// claims that the agent did or did not incorporate it. /// /// Returned synchronously by `send_steer` when no task is in flight /// for the channel. Never sent through the ack channel — the ack /// watcher is only spawned on `send_steer` success. PromptCompleted, } /// Outcome of a goose-native steer, sent from the read loop back to the /// main loop's ack watcher. #[derive(Debug)] pub enum SteerAck { /// The agent returned a successful response to the steer request. /// The main loop must drop the withheld event (`remove_event`) — it /// has been delivered via the non-cancelling path. Success, /// The steer was attempted but failed. Delivery state for the /// underlying message is unknown after prompt completion; the main /// loop must release the withheld event and fall back to the /// universal `Steer` cancel+merge path so the message still reaches /// the agent. Err(SteerError), /// The prompt completed before the read loop selected the steer arm. /// Treated as a benign no-op: release the withheld event for normal /// dispatch. Do not fire the fallback `Steer` signal — there is no /// in-flight turn to signal, and normal dispatch handles delivery. PromptCompletedNeutral, } /// Outcome of a prompt task. #[allow(dead_code)] pub enum PromptOutcome { Ok(StopReason), Error(AcpError), AgentExited, Timeout, /// Intentional cancel via `!cancel` command or interrupt mode. /// Agent is healthy — no respawn, no retry penalty. Cancelled, } /// Immutable config subset shared (via `Arc`) by all spawned prompt tasks. /// /// Built once from `Config` at startup. Avoids cloning the full config /// into every task. pub struct PromptContext { pub mcp_servers: Vec, pub initial_message: Option, pub idle_timeout: Duration, pub max_turn_duration: Duration, /// Interval between per-turn `turn_liveness` observer pings. `Duration::ZERO` /// disables emission. This is the desktop crash-backstop signal — distinct /// from `heartbeat_prompt` (agent self-prompting). pub turn_liveness_interval: Duration, pub dedup_mode: DedupMode, pub system_prompt: Option, pub heartbeat_prompt: Option, /// Base prompt content, or `None` if `--no-base-prompt` was passed. /// /// `'static` because `PromptContext` is `Arc`-shared across async tasks. /// Content from `--base-prompt-file` is promoted via `Box::leak` in `main.rs` /// after validated file read in `Config::from_cli()`. The compiled-in default /// (`include_str!`) is inherently `'static`. pub base_prompt: Option<&'static str>, pub cwd: String, /// REST client for pre-prompt context fetches (thread/DM history). pub rest_client: RestClient, /// Channel metadata from discovery (name, type). Read-only after startup. pub channel_info: std::collections::HashMap, /// Max messages to include in thread/DM context. 0 = disabled. pub context_message_limit: u32, /// Max turns per session before proactive rotation. 0 = disabled. pub max_turns_per_session: u32, /// Permission mode to apply after session creation. `Default` = skip. pub permission_mode: PermissionMode, /// Agent identity — used to derive the NIP-AE conversation key at /// session creation for core injection. pub agent_keys: nostr::Keys, /// Owner pubkey (hex), if resolved at startup. When unset, NIP-AE core /// injection is skipped entirely (no owner = no `(agent, owner)` pair). pub agent_owner_pubkey: Option, /// Whether NIP-AE agent core memory injection is enabled. When false, /// the per-session core engram fetch is skipped and `core_sections` /// remains empty for every channel, so `format_prompt` renders no /// `[Agent Memory — core]` section. On by default; disabled via /// `--no-memory` / `BUZZ_ACP_NO_MEMORY`. pub memory_enabled: bool, } impl AgentPool { /// Create a pool from pre-indexed slots (may contain None for failed startups). /// /// Slot positions are preserved so that `agent.index` always matches the /// index into `self.agents`. Use this instead of `new()` when the startup /// loop skips failed agents — `new()` would pack agents densely and break /// the index invariant. pub fn from_slots(slots: Vec>) -> Self { let (result_tx, result_rx) = mpsc::unbounded_channel(); Self { agents: slots, result_tx, result_rx, join_set: JoinSet::new(), task_map: HashMap::new(), } } /// Try to claim an idle agent for the given channel (or heartbeat if `None`). /// /// Pass 1: prefer an agent that already has a session for `channel_id`. /// Pass 2: any idle agent. /// /// Returns `None` if all agents are checked out. pub fn try_claim(&mut self, channel_id: Option) -> Option { // Pass 1: prefer agent with existing session for this channel. if let Some(cid) = channel_id { let idx = self.agents.iter().position(|slot| { slot.as_ref() .map(|a| a.state.sessions.contains_key(&cid)) .unwrap_or(false) }); if let Some(i) = idx { return self.agents[i].take(); } } // Pass 2: first idle agent. let idx = self.agents.iter().position(|slot| slot.is_some()); idx.map(|i| self.agents[i].take().unwrap()) } /// Return an agent to its slot after a task completes. pub fn return_agent(&mut self, agent: OwnedAgent) { let idx = agent.index; if self.agents[idx].is_some() { // This is a bug: two tasks returned the same agent index. Log it // loudly so it shows up in production logs, then overwrite — the // alternative (dropping the incoming agent) would permanently leak // the slot. tracing::error!( idx, "BUG: return_agent called for slot {idx} which is already occupied — overwriting" ); } self.agents[idx] = Some(agent); } /// Whether any agent is currently idle (sitting in its slot). pub fn any_idle(&self) -> bool { self.agents.iter().any(|slot| slot.is_some()) } /// Whether any idle agent already has a session for `channel_id`. /// Used to compute `affinity_hit` before calling `try_claim`. pub fn has_session_for(&self, channel_id: Uuid) -> bool { self.agents.iter().any(|slot| { slot.as_ref() .map(|a| a.state.sessions.contains_key(&channel_id)) .unwrap_or(false) }) } /// Count of agents that are alive: idle OR checked out (have a task_map entry). /// /// Used to detect when all agents have exited so the caller can respawn. pub fn live_count(&self) -> usize { let idle = self.agents.iter().filter(|s| s.is_some()).count(); let checked_out = self.task_map.len(); idle + checked_out } pub fn task_map(&self) -> &HashMap { &self.task_map } pub fn task_map_mut(&mut self) -> &mut HashMap { &mut self.task_map } /// Try to send a goose-native steer request to the in-flight task for /// `channel_id`. /// /// Returns `Ok(())` if the request was accepted by the read loop's /// receiver (capacity-1 mpsc; one slot is the single in-flight steer /// write). Returns `Err(SteerError::Transport(_))` on `Full`/`Closed` /// (already-in-flight write, or read loop torn down). Callers must /// fall back to the universal `ControlSignal::Steer` cancel+merge path /// on `Err`. /// /// This does **not** spawn the ack watcher — the caller owns the /// oneshot `ack_tx` inside `SteerRequest` and is responsible for /// awaiting it and applying the locked Success / Err / PromptCompletedNeutral /// semantics. Caller is also responsible for the synchronous /// `queue.mark_native_steer_pending(...)` *before* spawning the /// watcher, to close the result-vs-ack race. /// /// Returns `Err(SteerError::PromptCompleted)` if no task is in flight /// for `channel_id` (the prompt completed between the mode-gate check /// and this call, or the channel was never in flight). This is /// semantically a soft no-op — the caller should release any withheld /// event and let normal dispatch handle delivery. pub fn send_steer( &mut self, channel_id: Uuid, request: SteerRequest, ) -> Result<(), SteerError> { let meta = self .task_map .values_mut() .find(|m| m.channel_id == Some(channel_id)) .ok_or(SteerError::PromptCompleted)?; let tx = meta .steer_tx .as_ref() .ok_or_else(|| SteerError::Transport("steer_tx not installed".into()))?; tx.try_send(request) .map_err(|e| SteerError::Transport(e.to_string())) } pub fn result_tx(&self) -> mpsc::UnboundedSender { self.result_tx.clone() } /// Split-borrow: returns mutable refs to `result_rx` and `join_set` /// simultaneously. This lets callers poll both in a single `select!` /// without a double-borrow error on `&mut AgentPool`. pub fn rx_and_join_set( &mut self, ) -> (&mut mpsc::UnboundedReceiver, &mut JoinSet<()>) { (&mut self.result_rx, &mut self.join_set) } /// Non-blocking drain of the result channel. Used during shutdown to /// collect agents that completed while join_set was being drained. pub fn result_rx_try_recv(&mut self) -> Result { self.result_rx.try_recv() } /// Check whether a slot is alive: either idle in the pool or checked out /// for an in-flight task. Returns `false` only when the slot is truly /// empty and available for refill. pub fn slot_alive(&self, index: usize) -> bool { let idle = self.agents.get(index).is_some_and(|s| s.is_some()); if idle { return true; } // Check if the agent is checked out (in-flight on a task). self.task_map.values().any(|m| m.agent_index == index) } pub fn agents_mut(&mut self) -> &mut Vec> { &mut self.agents } /// Remove the session for `channel_id` from all idle agents. /// /// Called when the agent is removed from a channel — stale sessions /// should not be reused. Checked-out agents (in-flight) are not /// modified; their sessions will fail naturally on the next prompt /// if the relay rejects the request. /// /// Returns the number of sessions invalidated. pub fn invalidate_channel_sessions(&mut self, channel_id: Uuid) -> usize { let mut count = 0; for slot in &mut self.agents { if let Some(agent) = slot.as_mut() { if agent.state.invalidate_channel(&channel_id) { count += 1; } } } count } /// Idle-path model switch: set `desired_model` on the idle agent for /// `channel_id` and invalidate its session so the next turn re-creates the /// session under the new model. /// /// Pre-cancel guard: the desired model is validated against the agent's /// cached catalog *before* the session is invalidated, so an unsupported /// pick is rejected without disturbing the existing session. /// /// Returns [`IdleSwitchResult`] describing what happened. The model does not /// take effect — and the panel does not reflect it — until the agent next /// runs a turn (no live session exists to re-emit `session_config_captured` /// from an idle agent). This lag is intentional: faking the emit would /// surface an override the session has not actually applied. pub fn switch_idle_agent_model( &mut self, channel_id: Uuid, model_id: &str, ) -> IdleSwitchResult { let Some(agent) = self .agents .iter_mut() .flatten() .find(|a| a.state.sessions.contains_key(&channel_id)) else { return IdleSwitchResult::NoIdleAgent; }; // Pre-cancel guard against the cached catalog. None = catalog not yet // populated (no session ever created); defer validation to apply time. if let Some(caps) = agent.model_capabilities.as_ref() { if !model_in_catalog( &caps.config_options_raw, caps.available_models_raw.as_ref(), model_id, ) { return IdleSwitchResult::UnsupportedModel; } } agent.desired_model = Some(model_id.to_string()); agent.model_overridden = true; agent.state.invalidate_channel(&channel_id); IdleSwitchResult::Switched } } /// Outcome of [`AgentPool::switch_idle_agent_model`]. #[derive(Debug, PartialEq, Eq)] pub enum IdleSwitchResult { /// `desired_model` set and the channel session invalidated. Switched, /// Desired model is not in the agent's cached catalog — pick rejected, /// session untouched. UnsupportedModel, /// No idle agent available (all checked out / none spawned). NoIdleAgent, } /// Timeout for a single pre-prompt context fetch attempt (thread/DM history). /// Each call gets this budget; with one retry the total worst-case is /// 2 × CONTEXT_FETCH_TIMEOUT + CONTEXT_FETCH_RETRY_DELAY ≈ 6.5 s. const CONTEXT_FETCH_TIMEOUT: Duration = Duration::from_millis(3_000); /// Delay between the first failed context fetch and the single retry. const CONTEXT_FETCH_RETRY_DELAY: Duration = Duration::from_millis(500); /// Timeout for model-switch requests (`session/set_config_option`, `session/set_model`). const MODEL_SWITCH_TIMEOUT: Duration = Duration::from_secs(5); /// Timeout for permission-mode requests (`session/set_config_option` with `configId: "mode"`). const PERMISSION_MODE_TIMEOUT: Duration = Duration::from_secs(5); /// Create a new ACP session via `session_new_full()`, populate model capabilities /// on the agent (first session only), and apply `desired_model` if set. /// /// On error from `session_new_full()`, returns the `AcpError` — caller handles /// error reporting. Model-switch failures are logged and gracefully ignored /// (the agent proceeds with its default model). async fn create_session_and_apply_model( agent: &mut OwnedAgent, ctx: &PromptContext, agent_core: Option<&str>, ) -> Result { // Combine base_prompt + system_prompt + agent core into a single // systemPrompt value for the session/new request. Only sent when the agent // declares protocol version >= 2 (supports systemPrompt); legacy agents // ignore it and receive the same content as user-message sections via // `format_prompt`. Core already carries its own `[Agent Memory — core]` // header from `engram_fetch::build_core_section`, so we just append it. let combined_system_prompt: Option = if agent.protocol_version >= 2 { with_core( framed_system_prompt(&ctx.cwd, ctx.base_prompt, ctx.system_prompt.as_deref()), agent_core, ) } else { None }; let resp = agent .acp .session_new_full( &ctx.cwd, ctx.mcp_servers.clone(), combined_system_prompt.as_deref(), ) .await?; // Populate model capabilities on first session creation. if agent.model_capabilities.is_none() { agent.model_capabilities = Some(AgentModelCapabilities { config_options_raw: extract_model_config_options(&resp.raw), available_models_raw: extract_model_state(&resp.raw), }); } // Apply desired_model if set, matching against the fresh session/new response. // Track whether the switch succeeded so session_config_captured reflects // the post-switch state (not the pre-switch desired state). let switch_succeeded = if let Some(ref desired) = agent.desired_model { match resolve_model_switch_method(&resp.raw, desired) { Some(method) => { apply_model_switch(&mut agent.acp, &resp.session_id, desired, &method).await?; true } None => { tracing::warn!( target: "pool::model", "desired model {desired} not found in agent's available models — proceeding with agent default" ); // Surface the miss so the desktop ModelPicker can reject a live // pick rather than silently no-op. On the busy path the turn has // already been cancelled+requeued by the time we get here, so the // turn restarts on the unchanged model and the user is told no. agent.acp.observe( "control_result", serde_json::json!({ "type": "switch_model", "status": "unsupported_model", "modelId": desired, }), ); false } } } else { false }; // Emit session config for desktop consumption (config bridge tier 1b). // Emitted AFTER desired_model resolution so the desktop caches the // post-switch state. modelOverridden reflects whether the switch actually // applied — false on the unsupported arm so the panel doesn't show a // stale override badge. agent.acp.observe( "session_config_captured", serde_json::json!({ "configOptions": resp.raw.get("configOptions").cloned().unwrap_or(serde_json::Value::Null), "modes": resp.raw.get("modes").cloned().unwrap_or(serde_json::Value::Null), "models": resp.raw.get("models").cloned().unwrap_or(serde_json::Value::Null), "modelOverridden": agent.model_overridden && switch_succeeded, }), ); // Apply permission mode if not the agent's built-in default AND the agent // advertises the requested mode in session/new. Agents that don't support // the mode (e.g., goose crashes on unrecognized set_config_option values) // are safely skipped — the harness auto-approves via handle_permission_request. if !ctx.permission_mode.is_default() && agent_supports_mode(&resp.raw, ctx.permission_mode.as_wire_str()) { apply_permission_mode(&mut agent.acp, &resp.session_id, &ctx.permission_mode).await?; } Ok(resp.session_id) } /// Send the appropriate ACP model-switch request with a timeout. /// /// On timeout or error, logs a warning and returns — the caller proceeds /// with the agent's default model. This is intentionally non-fatal: a stale /// response from a timed-out request is safely ignored by `read_until_response` /// (non-matching JSON-RPC IDs are skipped). async fn apply_model_switch( acp: &mut AcpClient, session_id: &str, desired: &str, method: &ModelSwitchMethod, ) -> Result<(), AcpError> { let method_label = match method { ModelSwitchMethod::ConfigOption { config_id, .. } => { format!("configOption (configId={config_id})") } ModelSwitchMethod::SetModel { .. } => "set_model".to_string(), }; let result = tokio::time::timeout(MODEL_SWITCH_TIMEOUT, async { match method { ModelSwitchMethod::ConfigOption { config_id, option_value, } => { acp.session_set_config_option(session_id, config_id, option_value) .await } ModelSwitchMethod::SetModel { model_id } => { acp.session_set_model(session_id, model_id).await } } }) .await; match result { Ok(Ok(_)) => { tracing::info!( target: "pool::model", "applied model {desired} via {method_label} on session {session_id}" ); } // Transport-class errors may have corrupted the stdio stream — propagate // so the caller can respawn the agent instead of reusing a poisoned one. Ok(Err(e @ AcpError::Io(_))) | Ok(Err(e @ AcpError::WriteTimeout(_))) | Ok(Err(e @ AcpError::Timeout(_))) | Ok(Err(e @ AcpError::Protocol(_))) | Ok(Err(e @ AcpError::AgentExited)) => { tracing::error!( target: "pool::model", "fatal error setting model {desired} via {method_label}: {e}" ); return Err(e); } // Application-level errors (Json, etc.) — agent is fine, just uses default model. Ok(Err(e)) => { tracing::warn!( target: "pool::model", "failed to set model {desired} via {method_label}: {e} — proceeding with agent default" ); } Err(_) => { // Outer timeout fired — the inner send_request may have left the // stream in an unknown state. Treat as transport error. tracing::error!( target: "pool::model", "model set via {method_label} timed out ({MODEL_SWITCH_TIMEOUT:?}) — treating as fatal" ); return Err(AcpError::Timeout(MODEL_SWITCH_TIMEOUT)); } } Ok(()) } /// Set the session permission mode via `session/set_config_option`. /// /// Non-fatal for most errors: logs and proceeds. The agent falls back /// to its default permission mode (`"default"`), which still works via /// Check if the agent's `session/new` response advertises a given mode ID /// in `result.modes.availableModes[].id`. Returns `false` if the modes /// field is absent or the mode isn't listed. fn agent_supports_mode(session_new_result: &serde_json::Value, mode_wire: &str) -> bool { session_new_result .get("modes") .and_then(|m| m.get("availableModes")) .and_then(|a| a.as_array()) .map(|modes| { modes .iter() .any(|m| m.get("id").and_then(|v| v.as_str()) == Some(mode_wire)) }) .unwrap_or(false) } /// per-tool auto-approval in `handle_permission_request`. /// /// **Fatal exception:** if the agent process exits (e.g., goose crashes on /// unrecognized methods), returns `Err(AgentExited)` so the caller can respawn. async fn apply_permission_mode( acp: &mut AcpClient, session_id: &str, mode: &PermissionMode, ) -> Result<(), AcpError> { let wire = mode.as_wire_str(); let result = tokio::time::timeout(PERMISSION_MODE_TIMEOUT, async { acp.session_set_config_option(session_id, "mode", wire) .await }) .await; match result { Ok(Ok(_)) => { tracing::info!( target: "pool::permission", "applied permission mode {wire:?} on session {session_id}" ); } // Transport-class errors may have corrupted the stdio stream — propagate // so the caller can respawn the agent. Ok(Err(e @ AcpError::Io(_))) | Ok(Err(e @ AcpError::WriteTimeout(_))) | Ok(Err(e @ AcpError::Timeout(_))) | Ok(Err(e @ AcpError::Protocol(_))) | Ok(Err(e @ AcpError::AgentExited)) => { tracing::error!( target: "pool::permission", "fatal error setting permission mode {wire:?}: {e}" ); return Err(e); } // Application-level errors — agent is fine, just uses default permission mode. Ok(Err(e)) => { tracing::warn!( target: "pool::permission", "failed to set permission mode {wire:?}: {e} — falling back to per-tool auto-approval" ); } Err(_) => { // Outer timeout fired — stream may be in unknown state. tracing::error!( target: "pool::permission", "permission mode set timed out ({PERMISSION_MODE_TIMEOUT:?}) — treating as fatal" ); return Err(AcpError::Timeout(PERMISSION_MODE_TIMEOUT)); } } Ok(()) } /// Prepend the `[Base]` section to a user-message body for legacy agents. /// /// Legacy agents (`protocol_version < 2`) don't receive `base_prompt` via the /// system role in `session/new`, so it must ride along in the user message. /// Agents with `protocol_version >= 2`, or any agent without a `base_prompt`, /// get `body` unchanged. The gate lives here so the heartbeat and /// initial-message dispatch paths can't drift apart again. pub(crate) fn prepend_base_for_legacy( protocol_version: u32, base_prompt: Option<&str>, body: &str, ) -> String { match base_prompt { Some(bp) if protocol_version < 2 => { format!("{}\n\n{body}", crate::queue::base_section(bp)) } _ => body.to_string(), } } /// Frame the `session/new` `systemPrompt` so each present prompt carries its own /// header, keeping the base/persona boundary recoverable downstream. /// /// The header framing matches the legacy per-turn path (`queue::base_section` /// for `[Base]`, `[System]\n{...}` for the persona) so the desktop observer can /// split the combined value into labeled sub-sections. Each prompt is wrapped /// only when present, so a persona-only agent yields `[System]\n{persona}` /// rather than an unlabeled blob that would be mislabeled as `[Base]`. /// /// Prepends a `[Workspace]` section naming the agent's absolute working /// directory. The base prompt describes the workspace layout but never its /// absolute root, so without this anchor a model fills the gap by searching /// `$HOME` (triggering macOS TCC prompts) or by inventing its own workspace /// directory. The line is emitted only when a real base prompt is present and /// `cwd` is an absolute path other than the `/` fallback — naming `/` as the /// workspace would itself invite a `$HOME`-wide scan. fn framed_system_prompt( cwd: &str, base_prompt: Option<&str>, system_prompt: Option<&str>, ) -> Option { let body = match (base_prompt, system_prompt) { (Some(bp), Some(sp)) => Some(format!( "{}\n\n[System]\n{sp}", crate::queue::base_section(bp) )), (Some(bp), None) => Some(crate::queue::base_section(bp)), (None, Some(sp)) => Some(format!("[System]\n{sp}")), (None, None) => None, }?; // Anchor the workspace only when a base prompt is present — the workspace // section grounds the base prompt's layout description, so it is meaningless // for a persona-only (`[System]`-only) agent that never received that layout. match (base_prompt, workspace_section(cwd)) { (Some(_), Some(workspace)) => Some(format!("{workspace}\n\n{body}")), _ => Some(body), } } /// Render the `[Workspace]` grounding section, or `None` when `cwd` is unusable. /// /// Skips relative paths and the `/` fallback (`std::env::current_dir()` resolves /// to `/` on failure): a `/`-rooted workspace line would actively encourage the /// `$HOME`-wide scan this section exists to prevent. fn workspace_section(cwd: &str) -> Option { if cwd != "/" && cwd.starts_with('/') { Some(format!( "[Workspace]\nYour absolute working directory is `{cwd}`. All workspace \ files — `AGENTS.md`, `RESEARCH/`, `PLANS/`, `GUIDES/`, `WORK_LOGS/`, \ `OUTBOX/` — and any repositories you clone (under `{cwd}/REPOS/`) live \ here. This is where you already are; do not search `$HOME` or other \ directories for them." )) } else { None } } /// Append the agent's core memory section onto the framed system prompt. /// /// Core already carries its own `[Agent Memory — core]` header from /// `engram_fetch::build_core_section`, so it is joined with a blank-line /// separator and never re-labeled. Either side may be absent. fn with_core(framed: Option, core: Option<&str>) -> Option { match (framed, core) { (Some(framed), Some(core)) => Some(format!("{framed}\n\n{core}")), (Some(framed), None) => Some(framed), (None, Some(core)) => Some(core.to_string()), (None, None) => None, } } /// Core async function spawned for each prompt. /// /// Lifecycle: /// 1. Resolve or create a session (channel or heartbeat). /// 2. Send `initial_message` on new channel sessions (if configured). /// 3. Fetch conversation context if needed (thread reply or DM). /// 4. Build the prompt text from batch + context. /// 5. Send the actual prompt with turn timeout. /// 6. Handle all error paths, always returning the agent via `result_tx`. /// /// The agent is ALWAYS returned — even on panic the `JoinSet` detects the /// abort and the caller uses `task_map` to recover the agent index. pub async fn run_prompt_task( mut agent: OwnedAgent, batch: Option, prompt_text: Option, ctx: Arc, result_tx: mpsc::UnboundedSender, control_rx: Option>, ) { // Is this a channel prompt or a heartbeat? let source = match &batch { Some(b) => PromptSource::Channel(b.channel_id), None => PromptSource::Heartbeat, }; let turn_id = uuid::Uuid::new_v4().to_string(); let observer_channel_id = match &source { PromptSource::Channel(channel_id) => Some(*channel_id), PromptSource::Heartbeat => None, }; agent.acp.set_observer_context(observer::context_for( observer_channel_id, None, Some(turn_id.clone()), )); let triggering_event_ids: Vec = batch .as_ref() .map(|b| b.events.iter().map(|be| be.event.id.to_hex()).collect()) .unwrap_or_default(); agent.acp.observe( "turn_started", serde_json::json!({ "source": match &source { PromptSource::Channel(_) => "channel", PromptSource::Heartbeat => "heartbeat", }, "triggeringEventIds": triggering_event_ids, }), ); // ── Turn completion guard ───────────────────────────────────────────── // Emits `turn_completed` on any exit path. Captures observer handle and // metadata now, before the agent is moved into PromptResult. let _turn_guard = TurnCompletionGuard::new( agent.acp.observer_handle(), agent.acp.observer_agent_index(), observer_channel_id, turn_id.clone(), ); // // Core memory is delivered inside the system prompt the harness already // builds (system role for protocol >= 2, the `[System]` user-message // section for legacy agents). To put it on the wire at `session/new` for // modern agents, the fetch must run *before* the session is created — so // we do it here and cache the rendered section in `state.core_sections`. // // Core is keyed by (agent_keys, owner) — both fixed for the process — so // it is identical across channels; the per-channel cache just avoids a // re-fetch on each new session and is cleared on session invalidation. // // Failure modes (all fail open — no crash, no block): // * no owner configured → skip (no NIP-AE namespace exists) // * confirmed absence → cache the onboarding nudge so the agent // learns how to bootstrap itself. // * transport / decrypt / parse error → inject nothing. We never // mistake "relay slow or broken" for "no core" — that would invite // the agent to overwrite real, just-unreachable memory. // * fetch exceeds CORE_FETCH_TIMEOUT → inject nothing, same reason. // // Per Tyler's locked spec: NO mid-session refreshes. Re-fetch only // happens when a session is invalidated and recreated (see // `SessionState::invalidate_channel`). // // Operator opt-out: `--no-memory` / `BUZZ_ACP_NO_MEMORY` skips the fetch. if ctx.memory_enabled { if let (PromptSource::Channel(cid), Some(owner_pk)) = (&source, ctx.agent_owner_pubkey.as_ref()) { let is_new_channel_session = !agent.state.sessions.contains_key(cid); if is_new_channel_session && !agent.state.core_sections.contains_key(cid) { // Bounded — we'd rather start the session with no core hint // than block session creation on a stalled relay. const CORE_FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3); let fetch = crate::engram_fetch::build_core_section( &ctx.rest_client, &ctx.agent_keys, owner_pk, ); let section = match tokio::time::timeout(CORE_FETCH_TIMEOUT, fetch).await { Ok(s) => s, Err(_) => { tracing::warn!( target: "engram::core", channel = %cid, timeout_ms = CORE_FETCH_TIMEOUT.as_millis() as u64, "core fetch timed out — emitting no section" ); None } }; if let Some(rendered) = section { tracing::info!( target: "engram::core", channel = %cid, section_len = rendered.len(), "injected NIP-AE core section into system prompt" ); agent.state.core_sections.insert(*cid, rendered); } } } } // The core section to fold into the system prompt for this turn's session. // Channel-scoped; heartbeats carry no owner core. let agent_core: Option = match &source { PromptSource::Channel(cid) => agent.state.core_sections.get(cid).cloned(), PromptSource::Heartbeat => None, }; let (session_id, is_new_session) = match &source { PromptSource::Channel(cid) => { if let Some(sid) = agent.state.sessions.get(cid) { (sid.clone(), false) } else { // Create new session with model application. match create_session_and_apply_model(&mut agent, &ctx, agent_core.as_deref()).await { Ok(sid) => { tracing::info!( target: "pool::session", "created session {sid} for channel {cid}" ); agent.state.sessions.insert(*cid, sid.clone()); (sid, true) } Err(AcpError::AgentExited) => { agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: requeue_batch_if_queue(&ctx, batch), }); return; } Err(e) => { let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(e), batch: requeue_batch_if_queue(&ctx, batch), }); return; } } } } PromptSource::Heartbeat => { if let Some(sid) = &agent.state.heartbeat_session { (sid.clone(), false) } else { match create_session_and_apply_model(&mut agent, &ctx, None).await { Ok(sid) => { tracing::info!( target: "pool::session", "created heartbeat session {sid} for agent {}", agent.index ); agent.state.heartbeat_session = Some(sid.clone()); (sid, true) } Err(AcpError::AgentExited) => { agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: None, }); return; } Err(e) => { let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(e), batch: None, }); return; } } } } }; agent.acp.set_observer_context(observer::context_for( observer_channel_id, Some(session_id.clone()), Some(turn_id.clone()), )); agent.acp.observe( "session_resolved", serde_json::json!({ "sessionId": session_id, "isNewSession": is_new_session, }), ); if is_new_session { if let (PromptSource::Channel(cid), Some(ref initial_msg)) = (&source, &ctx.initial_message) { tracing::info!( target: "pool::session", "sending initial_message to session {session_id} for channel {cid}" ); // For agents with systemPrompt support (protocol_version >= 2), // base_prompt is delivered via the system role in session/new. // Legacy agents receive it via [Base] in the user message instead. let init_msg = prepend_base_for_legacy(agent.protocol_version, ctx.base_prompt, initial_msg); let init_result = agent .acp .session_prompt_with_idle_timeout( &session_id, &init_msg, ctx.idle_timeout, ctx.max_turn_duration, ) .await; match init_result { Ok(stop_reason) => { tracing::info!( target: "pool::session", "initial_message complete for channel {cid}: {stop_reason:?}" ); } Err(AcpError::AgentExited) => { agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: requeue_batch_if_queue(&ctx, batch), }); return; } Err(AcpError::IdleTimeout(_)) => { tracing::warn!( target: "pool::session", "initial_message idle timeout ({}s) for channel {cid} — cancelling", ctx.idle_timeout.as_secs() ); match agent .acp .cancel_with_cleanup(&session_id, ctx.idle_timeout) .await { Ok(_) => { agent.state.invalidate(&source); } Err(AcpError::AgentExited) => { agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: requeue_batch_if_queue(&ctx, batch), }); return; } Err(e) => { tracing::error!( target: "pool::session", "cancel_with_cleanup failed during initial_message timeout: {e}" ); agent.state.invalidate(&source); } } let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: requeue_batch_if_queue(&ctx, batch), }); return; } Err(AcpError::HardTimeout) => { tracing::error!( target: "pool::session", "hard timeout ({}s cap) during initial_message for channel {cid} — agent process is unrecoverable", ctx.max_turn_duration.as_secs() ); agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: requeue_batch_if_queue(&ctx, batch), }); return; } Err(e) => { tracing::error!( target: "pool::session", "initial_message failed for channel {cid}: {e} — invalidating session" ); agent.state.invalidate(&source); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(e), batch: requeue_batch_if_queue(&ctx, batch), }); return; } } } } // When the batch is a single slash-command message (e.g. "@Eva /goal …"), // `slash_command` holds the bare command. It is sent as the FIRST prompt // content block so ACP connectors' slash-command detection // (`prompt[0].text.startsWith("/")`) fires; the wrapped Buzz context // follows as a second block. let mut slash_command: Option = None; let prompt_sections: Vec = if let Some(text) = prompt_text { // Pre-built prompt (heartbeat or legacy path) — a single block. vec![text] } else if let Some(ref b) = batch { // Build prompt from batch with context enrichment. // Try startup cache first; lazy-fetch via REST for dynamic channels. let channel_info = match ctx.channel_info.get(&b.channel_id) { Some(ci) => Some(PromptChannelInfo { name: ci.name.clone(), channel_type: ci.channel_type.clone(), }), None => fetch_channel_info(b.channel_id, &ctx.rest_client).await, }; let conversation_context = if ctx.context_message_limit > 0 { fetch_conversation_context(b, &channel_info, &ctx).await } else { None }; let profile_lookup = fetch_prompt_profile_lookup(b, conversation_context.as_ref(), &ctx.rest_client).await; let known_names: Vec<&str> = profile_lookup .iter() .flat_map(|lookup| lookup.values()) .flat_map(|p| [p.display_name.as_deref(), p.nip05_handle.as_deref()]) .flatten() .collect(); slash_command = crate::queue::slash_command_for_batch(b, &known_names); if let Some(ref cmd) = slash_command { tracing::info!( target: "pool::prompt", channel = %b.channel_id, command = %cmd, "slash-command pass-through" ); } crate::queue::format_prompt( b, &crate::queue::FormatPromptArgs { agent_core: agent_core.as_deref(), channel_info: channel_info.as_ref(), conversation_context: conversation_context.as_ref(), profile_lookup: profile_lookup.as_ref(), has_system_prompt_support: agent.protocol_version >= 2, base_prompt: ctx.base_prompt, system_prompt: ctx.system_prompt.as_deref(), }, ) } else { // Should not happen — batch is None only for heartbeats which have prompt_text. // Return the agent to the pool to prevent a permanent slot leak. tracing::error!("run_prompt_task: no batch and no prompt_text — returning agent"); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(AcpError::Protocol("no batch and no prompt_text".into())), batch: None, }); return; }; // ── Send the actual prompt ──────────────────────────────────────────── // Slash-command pass-through sends the bare command as the first text // block (so connector detection fires), then each prompt section as its // own block. Per-section blocks let the observer size trimmer elide a // section body in place while every `[Header]` line survives at the head // of its own leaf — so the "Prompt context" panel counts every section. let prompt_blocks: Vec<&str> = match slash_command { Some(ref cmd) => std::iter::once(cmd.as_str()) .chain(prompt_sections.iter().map(String::as_str)) .collect(), None => prompt_sections.iter().map(String::as_str).collect(), }; // When control_rx is Some (channel tasks), wrap the prompt in select! so // the main loop can cancel, interrupt, or rotate it. Heartbeats // (control_rx=None) take the simple await path — they are not controllable. // // The liveness future emits `turn_liveness` pings on an interval and never // resolves; it rides every prompt-await path as a non-winning select arm so // a turn stays alive on the desktop while it runs. Built from a captured // observer handle (not `&agent.acp`) because the prompt holds `&mut agent.acp`. let liveness = run_turn_liveness( agent.acp.observer_handle(), agent.acp.observer_agent_index(), observer::context_for( observer_channel_id, Some(session_id.clone()), Some(turn_id.clone()), ), ctx.turn_liveness_interval, ); tokio::pin!(liveness); let prompt_result = match control_rx { None => { // Heartbeat / non-cancellable path. tokio::select! { biased; result = agent.acp.session_prompt_blocks_with_idle_timeout( &session_id, &prompt_blocks, ctx.idle_timeout, ctx.max_turn_duration, ) => result, _ = &mut liveness => unreachable!("liveness future never resolves"), } } Some(rx) => { tokio::select! { biased; result = agent.acp.session_prompt_blocks_with_idle_timeout( &session_id, &prompt_blocks, ctx.idle_timeout, ctx.max_turn_duration, ) => result, _ = &mut liveness => unreachable!("liveness future never resolves"), mode = rx => { let control_signal = mode.unwrap_or(ControlSignal::Cancel); // Land the model switch before any cancel/requeue work: setting // `desired_model` here means the fresh session created by the // requeued turn (busy) or the next turn (already-completed) // applies the new model. Runtime-only — never persisted. if let ControlSignal::SwitchModel(ref model_id) = control_signal { agent.desired_model = Some(model_id.clone()); agent.model_overridden = true; } // Control signal received. Guard against Race 1: the turn may // have completed naturally just as cancel fired. if agent.acp.has_in_flight_prompt() { // Prompt is genuinely in-flight — cancel it. match agent .acp .cancel_with_cleanup_grace( &session_id, std::time::Duration::from_secs(5), ) .await { Ok(stop_reason) => { log_stop_reason(&source, &stop_reason); agent.state.invalidate(&source); let retry_batch = requeue_cancelled_batch(&ctx, control_signal, batch); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Cancelled, batch: retry_batch, }); return; } Err(AcpError::AgentExited) => { agent.state.invalidate_all(); let retry_batch = requeue_cancelled_batch(&ctx, control_signal, batch); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: retry_batch, }); return; } Err(AcpError::IdleTimeout(_) | AcpError::HardTimeout) => { // Cancel drain timed out — agent state uncertain. agent.state.invalidate(&source); let retry_batch = requeue_cancelled_batch(&ctx, control_signal, batch); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: retry_batch, }); return; } Err(e) => { agent.state.invalidate(&source); let retry_batch = requeue_cancelled_batch(&ctx, control_signal, batch); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(e), batch: retry_batch, }); return; } } } else { // Race 1 resolution: turn completed naturally before cancel // could fire. last_prompt_id is None — cleared by // session_prompt_with_idle_timeout() on success. The prompt // future was dropped by select! — its Ok result is gone. // // Note: this `else` branch (last_prompt_id is None) cannot // fire during the pre-prompt phase because `biased` select! // polls the prompt arm first. That arm sets last_prompt_id // synchronously before its first yield point, so by the time // the cancel arm can win, last_prompt_id is already Some. // This branch only fires when the turn genuinely completed // and last_prompt_id was cleared by the success path. // // MUST send a PromptResult or the main loop deadlocks. if matches!( control_signal, ControlSignal::Rotate | ControlSignal::SwitchModel(_) ) { tracing::debug!( target: "pool::prompt", "rotate/switch signal arrived but turn already completed — invalidating session" ); } else { tracing::debug!( target: "pool::prompt", "control signal arrived but turn already completed — treating as success" ); } apply_completed_before_control_signal( &mut agent.state, &source, &control_signal, ); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Ok(StopReason::EndTurn), batch: None, // turn succeeded — batch was processed, no requeue }); return; } } } } }; match prompt_result { Ok(stop_reason) => { log_stop_reason(&source, &stop_reason); let should_rotate = matches!( stop_reason, StopReason::MaxTokens | StopReason::MaxTurnRequests ); let should_rotate = should_rotate || { let limit = ctx.max_turns_per_session; if limit > 0 { match &source { PromptSource::Channel(cid) => { let count = agent.state.turn_counts.entry(*cid).or_insert(0); *count += 1; *count >= limit } PromptSource::Heartbeat => { agent.state.heartbeat_turn_count += 1; agent.state.heartbeat_turn_count >= limit } } } else { false } }; if should_rotate { tracing::info!( target: "pool::session", "rotating session for {source:?} after {stop_reason:?}", ); agent.state.invalidate(&source); } let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Ok(stop_reason), batch: None, }); } Err(AcpError::AgentExited) => { tracing::error!(target: "pool::prompt", "agent {} exited during prompt", agent.index); agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: requeue_batch_if_queue(&ctx, batch), }); } Err(AcpError::IdleTimeout(_)) => { tracing::warn!( target: "pool::prompt", "idle timeout ({}s) — cancelling session {session_id}", ctx.idle_timeout.as_secs() ); match agent .acp .cancel_with_cleanup(&session_id, ctx.idle_timeout) .await { Ok(stop_reason) => { log_stop_reason(&source, &stop_reason); // Timeout triggers respawn in handle_prompt_result — // session state will be discarded with the old agent. let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: requeue_batch_if_queue(&ctx, batch), }); } Err(AcpError::AgentExited) => { tracing::error!( target: "pool::prompt", "agent {} exited during cancel_with_cleanup", agent.index ); agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::AgentExited, batch: requeue_batch_if_queue(&ctx, batch), }); } Err(e) => { tracing::error!( target: "pool::prompt", "cancel_with_cleanup error: {e} — invalidating session" ); agent.state.invalidate(&source); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: requeue_batch_if_queue(&ctx, batch), }); } } } Err(AcpError::HardTimeout) => { tracing::error!( target: "pool::prompt", "hard timeout ({}s cap) — agent process is unrecoverable, invalidating all sessions", ctx.max_turn_duration.as_secs() ); agent.state.invalidate_all(); let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Timeout, batch: requeue_batch_if_queue(&ctx, batch), }); } Err(e) => { tracing::error!(target: "pool::prompt", "session_prompt error: {e}"); // AgentError means the agent caught a problem before mutating // session state (e.g. bad LLM response). The session is healthy — // don't invalidate it. Other errors may have corrupted state. if !matches!(e, AcpError::AgentError(_)) { agent.state.invalidate(&source); } let _ = result_tx.send(PromptResult { agent, source, outcome: PromptOutcome::Error(e), batch: requeue_batch_if_queue(&ctx, batch), }); } } } /// Retry wrapper for context fetches: one retry with `CONTEXT_FETCH_RETRY_DELAY` /// on any `None` result. The closure is called twice at most. /// /// Using a closure (not a `Future`) so the retry can construct a fresh `Future` /// each attempt without requiring `Clone` or re-boxing. async fn fetch_with_retry(f: F) -> Option where F: Fn() -> Fut, Fut: std::future::Future>, { if let Some(result) = f().await { return Some(result); } tokio::time::sleep(CONTEXT_FETCH_RETRY_DELAY).await; f().await } /// Lazy-fetch channel metadata for a channel not in the startup discovery cache. /// /// Handles channels added dynamically via membership notifications after startup. /// Uses `CONTEXT_FETCH_TIMEOUT` with one retry on failure. Returns `None` on /// persistent failure (graceful degradation — prompt will lack channel name and /// DM detection). async fn fetch_channel_info(channel_id: Uuid, rest: &RestClient) -> Option { use nostr::{Alphabet, SingleLetterTag}; let d_tag = SingleLetterTag::lowercase(Alphabet::D); let filter = nostr::Filter::new() .kind(nostr::Kind::Custom( buzz_core::kind::KIND_NIP29_GROUP_METADATA as u16, )) .custom_tags(d_tag, [channel_id.to_string()]); fetch_with_retry(|| async { match timeout( CONTEXT_FETCH_TIMEOUT, rest.query(std::slice::from_ref(&filter)), ) .await { Ok(Ok(json)) => { let events = json.as_array()?; let ev = events.first()?; let tags = ev.get("tags")?.as_array()?; let mut name = None; let mut is_hidden = false; let mut is_private = false; for tag in tags { if let Some(arr) = tag.as_array() { match arr.first().and_then(|v| v.as_str()) { Some("name") => name = arr.get(1).and_then(|v| v.as_str()), Some("hidden") => is_hidden = true, Some("private") => is_private = true, _ => {} } } } let channel_type = if is_hidden { "dm".to_string() } else if is_private { "private".to_string() } else { "stream".to_string() }; Some(PromptChannelInfo { name: name.unwrap_or("unknown").to_string(), channel_type, }) } Ok(Err(e)) => { tracing::debug!( channel_id = %channel_id, "channel info fetch failed: {e} — will retry" ); None } Err(_) => { tracing::debug!( channel_id = %channel_id, "channel info fetch timed out — will retry" ); None } } }) .await } /// Fetch conversation context (thread or DM) for a batch before prompting. /// /// Returns `None` if: /// - The event is a plain channel message (not a thread reply, not a DM) /// - The REST fetch fails or times out (graceful degradation) /// - `context_message_limit` is 0 /// /// For batches with multiple events, thread context is fetched for the **last** /// reply event only (most recent = most likely to need a response). async fn fetch_conversation_context( batch: &FlushBatch, channel_info: &Option, ctx: &PromptContext, ) -> Option { let limit = ctx.context_message_limit; let is_dm = channel_info .as_ref() .map(|ci| ci.channel_type == "dm") .unwrap_or(false); // Check thread tags on the last event first — this applies to both // channels and DMs. A DM reply needs thread context (not channel history) // because /api/channels/{id}/messages excludes thread replies. let last_event = batch.events.last()?; let tags = crate::queue::parse_thread_tags(&last_event.event); if let Some(root_id) = tags.root_event_id { return fetch_thread_context(batch.channel_id, &root_id, limit, &ctx.rest_client).await; } // DM non-reply: fetch recent conversation history. if is_dm { return fetch_dm_context(batch.channel_id, limit, &ctx.rest_client).await; } None } /// Normalize AND validate a pubkey for the batch profile API request. /// Returns `None` for malformed input — only valid 64-char hex passes. /// See also: `normalize_lookup_key` in queue.rs (normalize-only, no validation). fn normalize_prompt_pubkey(pubkey: &str) -> Option { let normalized = pubkey.trim().to_ascii_lowercase(); if normalized.len() == 64 && normalized.chars().all(|c| c.is_ascii_hexdigit()) { Some(normalized) } else { None } } fn collect_prompt_pubkeys( batch: &FlushBatch, conversation_context: Option<&ConversationContext>, ) -> Vec { let mut pubkeys = HashSet::new(); for event in &batch.events { pubkeys.insert(event.event.pubkey.to_hex().to_ascii_lowercase()); for mentioned in crate::queue::parse_thread_tags(&event.event).mentioned_pubkeys { if let Some(normalized) = normalize_prompt_pubkey(&mentioned) { pubkeys.insert(normalized); } } } let context_messages = match conversation_context { Some(ConversationContext::Thread { messages, .. }) | Some(ConversationContext::Dm { messages, .. }) => Some(messages), None => None, }; if let Some(messages) = context_messages { for message in messages { if let Some(normalized) = normalize_prompt_pubkey(&message.pubkey) { pubkeys.insert(normalized); } } } let mut pubkeys: Vec = pubkeys.into_iter().collect(); pubkeys.sort(); pubkeys } /// Detect whether a kind:0 profile event belongs to an owned agent. /// /// Agents carry a NIP-OA `["auth", owner_pk, conditions, sig]` tag in their /// profile; humans do not. This checks for the tag's presence/shape only — a /// cheap routing heuristic for reply anchoring, not a verified security gate /// (the signing path in `lib.rs::check_sibling_via_profile` does full /// verification where it matters). fn profile_event_is_agent(ev: &serde_json::Value) -> bool { ev.get("tags") .and_then(|t| t.as_array()) .is_some_and(|tags| { tags.iter().any(|tag| { tag.as_array() .is_some_and(|parts| parts.len() == 4 && parts[0].as_str() == Some("auth")) }) }) } /// Parse kind:0 profile events into a `PromptProfileLookup`. /// /// Each kind:0 event has `pubkey` and JSON `content` with optional fields: /// `display_name` (or `name`), `nip05`. fn parse_kind0_profile_lookup(json: serde_json::Value) -> Option { let events = json.as_array()?; let mut lookup = PromptProfileLookup::new(); for ev in events { let pubkey = ev.get("pubkey").and_then(|v| v.as_str()); let content_str = ev.get("content").and_then(|v| v.as_str()); if let (Some(pk), Some(content)) = (pubkey, content_str) { if let Ok(profile) = serde_json::from_str::(content) { let display_name = profile .get("display_name") .or_else(|| profile.get("name")) .and_then(|v| v.as_str()) .map(str::to_string); let nip05_handle = profile .get("nip05") .and_then(|v| v.as_str()) .map(str::to_string); let is_agent = profile_event_is_agent(ev); lookup.insert( pk.to_ascii_lowercase(), PromptProfile { display_name, nip05_handle, is_agent, }, ); } } } if lookup.is_empty() { None } else { Some(lookup) } } async fn fetch_prompt_profile_lookup( batch: &FlushBatch, conversation_context: Option<&ConversationContext>, rest: &RestClient, ) -> Option { let pubkeys = collect_prompt_pubkeys(batch, conversation_context); if pubkeys.is_empty() { return None; } // Query kind:0 (NIP-01 profile metadata) for all pubkeys. let authors: Vec = pubkeys .iter() .filter_map(|s| nostr::PublicKey::from_hex(s).ok()) .collect(); if authors.is_empty() { return None; } let filter = nostr::Filter::new() .kind(nostr::Kind::Metadata) .authors(authors); fetch_with_retry(|| async { match timeout( CONTEXT_FETCH_TIMEOUT, rest.query(std::slice::from_ref(&filter)), ) .await { Ok(Ok(json)) => parse_kind0_profile_lookup(json), Ok(Err(e)) => { tracing::debug!("prompt profile lookup failed: {e} — will retry"); None } Err(_) => { tracing::debug!("prompt profile lookup timed out — will retry"); None } } }) .await } /// Fetch thread context via Nostr query: root event by ID + replies by `#e` tag. async fn fetch_thread_context( channel_id: Uuid, root_event_id: &str, limit: u32, rest: &RestClient, ) -> Option { use nostr::{Alphabet, SingleLetterTag}; // Defense-in-depth: validate hex event ID. if root_event_id.is_empty() || root_event_id.len() != 64 || !root_event_id.chars().all(|c| c.is_ascii_hexdigit()) { tracing::warn!( channel_id = %channel_id, "invalid root_event_id (expected 64 hex chars) — skipping thread context fetch" ); return None; } let e_tag = SingleLetterTag::lowercase(Alphabet::E); let h_tag = SingleLetterTag::lowercase(Alphabet::H); let ch_str = channel_id.to_string(); // Two filters: (1) root event by ID, (2) replies with #e=root + #h=channel. let root_filter = nostr::Filter::new().id(nostr::EventId::from_hex(root_event_id).ok()?); let replies_filter = nostr::Filter::new() .kinds([ nostr::Kind::Custom(buzz_core::kind::KIND_STREAM_MESSAGE as u16), nostr::Kind::Custom(buzz_core::kind::KIND_STREAM_MESSAGE_V2 as u16), ]) .custom_tags(e_tag, [root_event_id]) .custom_tags(h_tag, [ch_str.as_str()]) .limit(limit as usize); fetch_with_retry(|| async { match timeout( CONTEXT_FETCH_TIMEOUT, rest.query(&[root_filter.clone(), replies_filter.clone()]), ) .await { Ok(Ok(json)) => parse_nostr_thread_response(json, root_event_id), Ok(Err(e)) => { tracing::warn!( channel_id = %channel_id, root = root_event_id, "thread context fetch failed: {e} — will retry" ); None } Err(_) => { tracing::warn!( channel_id = %channel_id, root = root_event_id, "thread context fetch timed out — will retry" ); None } } }) .await } /// Fetch DM context via Nostr query: recent messages in channel by `#h` tag. async fn fetch_dm_context( channel_id: Uuid, limit: u32, rest: &RestClient, ) -> Option { use nostr::{Alphabet, SingleLetterTag}; let h_tag = SingleLetterTag::lowercase(Alphabet::H); let ch_str = channel_id.to_string(); let filter = nostr::Filter::new() .kinds([ nostr::Kind::Custom(buzz_core::kind::KIND_STREAM_MESSAGE as u16), nostr::Kind::Custom(buzz_core::kind::KIND_STREAM_MESSAGE_V2 as u16), ]) .custom_tags(h_tag, [ch_str.as_str()]) .limit(limit as usize); fetch_with_retry(|| async { match timeout( CONTEXT_FETCH_TIMEOUT, rest.query(std::slice::from_ref(&filter)), ) .await { Ok(Ok(json)) => parse_nostr_dm_response(json, limit), Ok(Err(e)) => { tracing::warn!( channel_id = %channel_id, "DM context fetch failed: {e} — will retry" ); None } Err(_) => { tracing::warn!( channel_id = %channel_id, "DM context fetch timed out — will retry" ); None } } }) .await } /// Parse the legacy REST thread response (used in tests only). #[cfg(test)] fn parse_thread_response(json: serde_json::Value) -> Option { let mut messages = Vec::new(); // Root message. if let Some(root) = json.get("root") { if let Some(msg) = json_to_context_message(root) { messages.push(msg); } } // Replies. if let Some(replies) = json.get("replies").and_then(|v| v.as_array()) { for reply in replies { if let Some(msg) = json_to_context_message(reply) { messages.push(msg); } } } let total_replies = json .get("total_replies") .and_then(|v| v.as_u64()) .unwrap_or(0) as usize; let total = total_replies + 1; // +1 for root let truncated = total > messages.len(); if messages.is_empty() { return None; } Some(ConversationContext::Thread { messages, total, truncated, }) } /// Parse the DM messages REST response into a `ConversationContext::Dm`. /// /// Parse the legacy REST DM response (used in tests only). #[cfg(test)] fn parse_dm_response(json: serde_json::Value, limit: u32) -> Option { let arr = json.get("messages").and_then(|v| v.as_array())?; let mut messages: Vec = arr.iter().filter_map(json_to_context_message).collect(); // API returns newest-first; reverse to chronological for the prompt. messages.reverse(); // The relay's next_cursor is always set when the page is non-empty (not // just when more pages exist), so we can't use it for truncation detection. // Instead, compare returned count against the requested limit. let truncated = messages.len() >= limit as usize; let total = if truncated { messages.len() + 1 // indicate there are more } else { messages.len() }; if messages.is_empty() { return None; } Some(ConversationContext::Dm { messages, total, truncated, }) } /// Extract a `ContextMessage` from a JSON message object. /// /// Works with both thread reply objects and channel message objects. fn json_to_context_message(obj: &serde_json::Value) -> Option { let content = obj.get("content").and_then(|v| v.as_str())?; let pubkey = obj .get("pubkey") .and_then(|v| v.as_str()) .unwrap_or("unknown"); let timestamp = obj .get("created_at") .and_then(|v| { // Handle both string timestamps and integer timestamps. v.as_str().map(|s| s.to_string()).or_else(|| { v.as_i64().map(|ts| { chrono::DateTime::from_timestamp(ts, 0) .map(|dt| dt.to_rfc3339()) .unwrap_or_else(|| ts.to_string()) }) }) }) .unwrap_or_else(|| "unknown".to_string()); Some(ContextMessage { pubkey: pubkey.to_string(), timestamp, content: content.to_string(), }) } /// Parse a Nostr query response (array of events) into thread context. /// /// Separates the root event (matching `root_event_id`) from replies, sorts /// chronologically by `created_at`. fn parse_nostr_thread_response( json: serde_json::Value, root_event_id: &str, ) -> Option { let events = json.as_array()?; let mut root_msg = None; let mut reply_msgs = Vec::new(); for ev in events { let ev_id = ev.get("id").and_then(|v| v.as_str()).unwrap_or(""); if let Some(msg) = json_to_context_message(ev) { if ev_id == root_event_id { root_msg = Some(msg); } else { reply_msgs.push(( ev.get("created_at").and_then(|v| v.as_u64()).unwrap_or(0), msg, )); } } } // Sort replies chronologically. reply_msgs.sort_by_key(|(ts, _)| *ts); let mut messages = Vec::new(); if let Some(root) = root_msg { messages.push(root); } messages.extend(reply_msgs.into_iter().map(|(_, msg)| msg)); let total = messages.len(); if messages.is_empty() { return None; } Some(ConversationContext::Thread { messages, total, truncated: false, // query returns all within limit }) } /// Parse a Nostr query response (array of events) into DM context. /// /// Events arrive in relay order (newest first); reversed to chronological. fn parse_nostr_dm_response(json: serde_json::Value, limit: u32) -> Option { let events = json.as_array()?; let mut messages: Vec<(u64, ContextMessage)> = events .iter() .filter_map(|ev| { let ts = ev.get("created_at").and_then(|v| v.as_u64()).unwrap_or(0); json_to_context_message(ev).map(|msg| (ts, msg)) }) .collect(); // Sort chronologically (oldest first). messages.sort_by_key(|(ts, _)| *ts); let messages: Vec = messages.into_iter().map(|(_, msg)| msg).collect(); let truncated = messages.len() >= limit as usize; let total = if truncated { messages.len() + 1 } else { messages.len() }; if messages.is_empty() { return None; } Some(ConversationContext::Dm { messages, total, truncated, }) } /// Return the batch for requeue only in Queue mode; drop it in Drop mode. #[inline] fn requeue_batch_if_queue(ctx: &PromptContext, batch: Option) -> Option { match ctx.dedup_mode { DedupMode::Queue => batch, DedupMode::Drop => None, } } /// Map a cancelling [`ControlSignal`] to the [`CancelReason`] that should frame /// the merged re-prompt, then requeue the batch (in `Queue` dedup mode) with /// that reason stamped onto [`FlushBatch::cancel_reason`]. `Cancel`/`Rotate` /// drop the batch entirely. The reason is consumed by the main loop at requeue /// time (`requeue_as_cancelled`) and ultimately by `format_prompt`. #[inline] fn requeue_cancelled_batch( ctx: &PromptContext, signal: ControlSignal, batch: Option, ) -> Option { let reason = match signal { ControlSignal::Steer => CancelReason::Steer, ControlSignal::Interrupt | ControlSignal::SwitchModel(_) => CancelReason::Interrupt, // Cancel/Rotate discard the batch — no merged re-prompt. ControlSignal::Cancel | ControlSignal::Rotate => return None, }; requeue_batch_if_queue(ctx, batch).map(|mut b| { b.cancel_reason = Some(reason); b }) } /// Log a stop reason at the appropriate tracing level. fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) { let label = match source { PromptSource::Channel(cid) => format!("channel {cid}"), PromptSource::Heartbeat => "heartbeat".to_string(), }; match stop_reason { StopReason::EndTurn => { tracing::info!(target: "pool::prompt", "turn complete for {label}: end_turn"); } StopReason::Cancelled => { tracing::warn!(target: "pool::prompt", "turn cancelled for {label}"); } StopReason::MaxTokens => { tracing::warn!(target: "pool::prompt", "turn hit max_tokens for {label} — session will be rotated"); } StopReason::MaxTurnRequests => { tracing::warn!(target: "pool::prompt", "turn hit max_turn_requests for {label} — session will be rotated"); } StopReason::Refusal => { tracing::warn!(target: "pool::prompt", "turn refused for {label}"); } } } // ── Turn liveness emission ─────────────────────────────────────────────────── // Periodically emits a `turn_liveness` observer event while a turn is in-flight, // so the desktop can prune turns whose host died without unwinding (kill -9 / // crash) far sooner than the no-activity backstop. Runs as a non-resolving // `select!` arm in `run_prompt_task`: it lives and dies with the prompt future, // so emission stops on every exit path (complete / cancel / error / panic) with // no separate teardown to forget. // // Takes a captured `ObserverHandle` rather than `&agent.acp` because the prompt // future holds `&mut agent.acp` for its whole duration — a second borrow would // not compile. // // This future never resolves; callers must race it against the prompt and rely // on drop for teardown. When `interval` is zero, liveness is disabled and the // future parks forever without emitting. async fn run_turn_liveness( observer: Option, agent_index: Option, context: observer::ObserverContext, interval: Duration, ) { let Some(observer) = observer else { return std::future::pending::<()>().await; }; if interval.is_zero() { return std::future::pending::<()>().await; } let mut ticker = tokio::time::interval(interval); // The first tick completes immediately; skip it so the first liveness ping // fires one interval after the turn starts, not at t=0 (turn_started already // marks t=0). ticker.tick().await; loop { ticker.tick().await; observer.emit( "turn_liveness", agent_index, &context, serde_json::json!({}), ); } } // Emits a `turn_completed` observer event on drop, covering ALL exit paths // (success, error, timeout, cancel, panic) from `run_prompt_task`. Captures // observer handle and metadata at creation time so it remains valid even after // the agent is moved into `PromptResult`. struct TurnCompletionGuard { observer: Option, agent_index: Option, channel_id: Option, turn_id: String, } impl TurnCompletionGuard { fn new( observer: Option, agent_index: Option, channel_id: Option, turn_id: String, ) -> Self { Self { observer, agent_index, channel_id, turn_id, } } } impl Drop for TurnCompletionGuard { fn drop(&mut self) { if let Some(observer) = self.observer.take() { let context = observer::context_for(self.channel_id, None, Some(self.turn_id.clone())); observer.emit( "turn_completed", self.agent_index, &context, serde_json::json!({}), ); } } } // ─── Unit Tests ────────────────────────────────────────────────────────────── #[cfg(test)] mod tests { use super::*; use nostr::{EventBuilder, Keys, Kind, Tag}; use serde_json::json; // These pin the initial_message dispatch path (run_prompt_task, ~line 855): // a legacy agent WITH a base_prompt must get [Base] prepended to the user // message. This is the exact regression that shipped in the round-2 bug. #[test] fn test_initial_message_legacy_agent_gets_base_prepended() { // protocol_version 1 + Some(base_prompt): [Base] rides along in the // user message, composed as `[Base]\n{bp}\n\n{initial_msg}`. let composed = prepend_base_for_legacy(1, Some("you are a helpful agent"), "hello channel"); assert_eq!(composed, "[Base]\nyou are a helpful agent\n\nhello channel"); assert!(composed.starts_with("[Base]\nyou are a helpful agent\n\n")); } #[test] fn test_initial_message_modern_agent_omits_base() { // protocol_version 2 receives base_prompt via session/new, so the user // message is left untouched even when a base_prompt is present. let composed = prepend_base_for_legacy(2, Some("you are a helpful agent"), "hello channel"); assert_eq!(composed, "hello channel"); } #[test] fn test_initial_message_legacy_agent_without_base_is_unchanged() { // No base_prompt configured: nothing to prepend regardless of version. let composed = prepend_base_for_legacy(1, None, "hello channel"); assert_eq!(composed, "hello channel"); } // Pin the session/new systemPrompt framing: each present prompt carries its // own header so the desktop observer can split into labeled sub-sections. #[test] fn test_framed_system_prompt_both_present_carries_both_headers() { let framed = framed_system_prompt("/", Some("base text"), Some("persona text")) .expect("both present yields Some"); assert_eq!(framed, "[Base]\nbase text\n\n[System]\npersona text"); } #[test] fn test_framed_system_prompt_base_only_labels_base() { let framed = framed_system_prompt("/", Some("base text"), None).expect("base yields Some"); assert_eq!(framed, "[Base]\nbase text"); } #[test] fn test_framed_system_prompt_persona_only_labels_system() { // A bare persona would be mislabeled "Base" downstream — it must carry // its own [System] header even when no base prompt exists. let framed = framed_system_prompt("/", None, Some("persona text")).expect("persona yields Some"); assert_eq!(framed, "[System]\npersona text"); } #[test] fn test_framed_system_prompt_neither_is_none() { assert!(framed_system_prompt("/", None, None).is_none()); } #[test] fn test_framed_system_prompt_absolute_cwd_prepends_workspace_before_base() { let framed = framed_system_prompt("/Users/me/.buzz", Some("base text"), None) .expect("base yields Some"); assert!( framed.starts_with("[Workspace]\n"), "workspace section must lead: {framed}" ); assert!(framed.contains("`/Users/me/.buzz`")); assert!( framed.contains("\n\n[Base]\nbase text"), "base must follow the workspace section: {framed}" ); } #[test] fn test_framed_system_prompt_persona_only_omits_workspace() { // The workspace section grounds the base prompt's layout; a persona-only // agent never received that layout, so no [Workspace] anchor is emitted. let framed = framed_system_prompt("/Users/me/.buzz", None, Some("persona text")) .expect("persona yields Some"); assert_eq!(framed, "[System]\npersona text"); } #[test] fn test_framed_system_prompt_root_cwd_omits_workspace() { // The "/" fallback must never be named — it would invite a $HOME scan. let framed = framed_system_prompt("/", Some("base text"), None).expect("base yields Some"); assert_eq!(framed, "[Base]\nbase text"); } #[test] fn test_workspace_section_relative_cwd_is_none() { assert!(workspace_section("relative/path").is_none()); assert!(workspace_section("").is_none()); } #[test] fn test_with_core_appends_below_framed() { let framed = with_core( Some("[System]\npersona".to_string()), Some("[Agent Memory — core]\nbe helpful"), ) .expect("both present yields Some"); assert_eq!( framed, "[System]\npersona\n\n[Agent Memory — core]\nbe helpful" ); } #[test] fn test_with_core_framed_only_passes_through() { let framed = with_core(Some("[System]\npersona".to_string()), None) .expect("framed-only yields Some"); assert_eq!(framed, "[System]\npersona"); } #[test] fn test_with_core_core_only_is_just_core() { let framed = with_core(None, Some("[Agent Memory — core]\nbe helpful")) .expect("core-only yields Some"); assert_eq!(framed, "[Agent Memory — core]\nbe helpful"); } #[test] fn test_with_core_neither_is_none() { assert!(with_core(None, None).is_none()); } #[test] fn test_parse_thread_response_basic() { let json = json!({ "root": { "event_id": "abc123", "pubkey": "pub1", "content": "root message", "created_at": 1710518400 }, "replies": [ { "event_id": "def456", "pubkey": "pub2", "content": "first reply", "created_at": 1710518460 } ], "total_replies": 1 }); let ctx = parse_thread_response(json).expect("should parse"); match ctx { ConversationContext::Thread { messages, total, truncated, } => { assert_eq!(messages.len(), 2); // root + 1 reply assert_eq!(total, 2); // 1 reply + 1 root assert!(!truncated); assert_eq!(messages[0].content, "root message"); assert_eq!(messages[1].content, "first reply"); } _ => panic!("expected Thread context"), } } #[test] fn test_parse_thread_response_truncated() { let json = json!({ "root": { "event_id": "abc", "pubkey": "pub1", "content": "root", "created_at": 1710518400 }, "replies": [ { "event_id": "def", "pubkey": "pub2", "content": "reply1", "created_at": 1710518460 } ], "total_replies": 10 }); let ctx = parse_thread_response(json).expect("should parse"); match ctx { ConversationContext::Thread { messages, total, truncated, } => { assert_eq!(messages.len(), 2); assert_eq!(total, 11); // 10 replies + 1 root assert!(truncated); } _ => panic!("expected Thread context"), } } #[test] fn test_parse_thread_response_empty() { let json = json!({ "root": null, "replies": [], "total_replies": 0 }); assert!(parse_thread_response(json).is_none()); } #[test] fn test_parse_thread_response_missing_fields() { // Malformed JSON — no root, no replies key. let json = json!({ "something": "else" }); assert!(parse_thread_response(json).is_none()); } #[test] fn test_parse_dm_response_basic() { let json = json!({ "messages": [ { "event_id": "msg2", "pubkey": "pub2", "content": "newer message", "created_at": 1710518500 }, { "event_id": "msg1", "pubkey": "pub1", "content": "older message", "created_at": 1710518400 } ], "next_cursor": null }); // limit=12 > 2 messages → not truncated. let ctx = parse_dm_response(json, 12).expect("should parse"); match ctx { ConversationContext::Dm { messages, total, truncated, } => { // Should be reversed to chronological order. assert_eq!(messages.len(), 2); assert_eq!(messages[0].content, "older message"); assert_eq!(messages[1].content, "newer message"); assert!(!truncated); assert_eq!(total, 2); } _ => panic!("expected Dm context"), } } #[test] fn test_parse_dm_response_truncated() { let json = json!({ "messages": [ { "event_id": "msg1", "pubkey": "pub1", "content": "message", "created_at": 1710518400 } ], "next_cursor": "00000000660f5a80" }); // limit=1 == 1 message → truncated. let ctx = parse_dm_response(json, 1).expect("should parse"); match ctx { ConversationContext::Dm { truncated, total, .. } => { assert!(truncated); assert_eq!(total, 2); // 1 message + indicator } _ => panic!("expected Dm context"), } } #[test] fn test_parse_dm_response_not_truncated_despite_cursor() { // Relay always sets next_cursor when page is non-empty, but if // returned count < limit, the page is complete. let json = json!({ "messages": [ { "event_id": "msg1", "pubkey": "pub1", "content": "only message", "created_at": 1710518400 } ], "next_cursor": "00000000660f5a80" }); // limit=12 > 1 message → NOT truncated despite next_cursor being set. let ctx = parse_dm_response(json, 12).expect("should parse"); match ctx { ConversationContext::Dm { truncated, total, .. } => { assert!(!truncated, "should not be truncated when count < limit"); assert_eq!(total, 1); } _ => panic!("expected Dm context"), } } #[test] fn test_parse_dm_response_empty() { let json = json!({ "messages": [], "next_cursor": null }); assert!(parse_dm_response(json, 12).is_none()); } #[test] fn test_parse_dm_response_missing_messages_key() { let json = json!({ "data": [] }); assert!(parse_dm_response(json, 12).is_none()); } #[test] fn test_json_to_context_message_integer_timestamp() { let obj = json!({ "pubkey": "abc", "content": "hello", "created_at": 1710518400 }); let msg = json_to_context_message(&obj).expect("should parse"); assert_eq!(msg.pubkey, "abc"); assert_eq!(msg.content, "hello"); assert!(msg.timestamp.contains("2024")); // 1710518400 = 2024-03-15 } #[test] fn test_json_to_context_message_string_timestamp() { let obj = json!({ "pubkey": "abc", "content": "hello", "created_at": "2026-03-15T16:30:00+00:00" }); let msg = json_to_context_message(&obj).expect("should parse"); assert_eq!(msg.timestamp, "2026-03-15T16:30:00+00:00"); } #[test] fn test_json_to_context_message_missing_content() { let obj = json!({ "pubkey": "abc" }); assert!(json_to_context_message(&obj).is_none()); } #[test] fn test_collect_prompt_pubkeys_includes_authors_mentions_and_context() { let keys = Keys::generate(); let p_tag = Tag::parse([ "p", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", ]) .unwrap(); let event = EventBuilder::new(Kind::Custom(9), "hello") .tags([p_tag]) .sign_with_keys(&keys) .unwrap(); let author_hex = event.pubkey.to_hex(); let batch = FlushBatch { channel_id: Uuid::new_v4(), events: vec![crate::queue::BatchEvent { event, prompt_tag: "@mention".into(), received_at: std::time::Instant::now(), }], cancelled_events: vec![], cancel_reason: None, }; let context = ConversationContext::Thread { messages: vec![ContextMessage { pubkey: "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".into(), timestamp: "2026-03-25T05:51:25Z".into(), content: "follow up".into(), }], total: 1, truncated: false, }; let pubkeys = collect_prompt_pubkeys(&batch, Some(&context)); let mut expected = vec![ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(), author_hex, "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string(), ]; expected.sort(); assert_eq!(pubkeys, expected); } #[test] fn test_parse_kind0_profile_lookup_extracts_display_name_and_nip05() { let lookup = parse_kind0_profile_lookup(json!([ { "id": "0000000000000000000000000000000000000000000000000000000000000001", "pubkey": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "kind": 0, "content": "{\"display_name\":\"Wes\",\"nip05\":\"wes@example.com\"}", "created_at": 1000, "tags": [], "sig": "0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" } ])) .expect("lookup should parse"); assert_eq!( lookup.get("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), Some(&PromptProfile { display_name: Some("Wes".into()), nip05_handle: Some("wes@example.com".into()), is_agent: false, }) ); } #[test] fn test_profile_event_is_agent_detects_nip_oa_auth_tag() { // Agent profile carries a 4-element NIP-OA ["auth", owner, cond, sig] tag. let agent_ev = json!({ "pubkey": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "tags": [["auth", "owner_pk", "conditions", "sig"]], }); assert!(profile_event_is_agent(&agent_ev)); // Human profile: no auth tag. let human_ev = json!({ "pubkey": "bbbb", "tags": [["t", "topic"]] }); assert!(!profile_event_is_agent(&human_ev)); // Empty / missing tags → not an agent. assert!(!profile_event_is_agent(&json!({ "tags": [] }))); assert!(!profile_event_is_agent(&json!({}))); // Malformed auth tag (wrong arity) → not treated as an agent. let malformed = json!({ "tags": [["auth", "owner_pk"]] }); assert!(!profile_event_is_agent(&malformed)); } #[test] fn test_parse_kind0_profile_lookup_returns_none_for_empty() { assert!(parse_kind0_profile_lookup(json!([])).is_none()); assert!(parse_kind0_profile_lookup(json!({})).is_none()); } #[test] fn test_json_to_context_message_missing_pubkey_uses_default() { let obj = json!({ "content": "hello" }); let msg = json_to_context_message(&obj).expect("should parse"); assert_eq!(msg.pubkey, "unknown"); } // ── SessionState tests ─────────────────────────────────────────────── fn make_state() -> (SessionState, Uuid, Uuid) { let ch_a = Uuid::new_v4(); let ch_b = Uuid::new_v4(); let mut s = SessionState::default(); s.sessions.insert(ch_a, "sess-a".into()); s.sessions.insert(ch_b, "sess-b".into()); s.turn_counts.insert(ch_a, 5); s.turn_counts.insert(ch_b, 3); s.core_sections.insert(ch_a, "core-a".into()); s.core_sections.insert(ch_b, "core-b".into()); s.heartbeat_session = Some("sess-hb".into()); s.heartbeat_turn_count = 7; (s, ch_a, ch_b) } #[test] fn test_rotate_after_natural_completion_invalidates_channel_state() { let (mut s, ch_a, ch_b) = make_state(); apply_completed_before_control_signal( &mut s, &PromptSource::Channel(ch_a), &ControlSignal::Rotate, ); assert!(!s.sessions.contains_key(&ch_a)); assert!(!s.turn_counts.contains_key(&ch_a)); assert!(!s.core_sections.contains_key(&ch_a)); assert!(!s.has_channel_state(&ch_a)); assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); assert_eq!(s.heartbeat_session.as_deref(), Some("sess-hb")); assert_eq!(s.heartbeat_turn_count, 7); } #[test] fn test_cancel_after_natural_completion_preserves_channel_state() { let (mut s, ch_a, ch_b) = make_state(); apply_completed_before_control_signal( &mut s, &PromptSource::Channel(ch_a), &ControlSignal::Cancel, ); assert_eq!(s.sessions.get(&ch_a).unwrap(), "sess-a"); assert_eq!(*s.turn_counts.get(&ch_a).unwrap(), 5); assert_eq!(s.core_sections.get(&ch_a).unwrap(), "core-a"); assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); } #[test] fn test_invalidate_channel_clears_session_and_turn_count() { let (mut s, ch_a, ch_b) = make_state(); s.invalidate(&PromptSource::Channel(ch_a)); assert!(!s.sessions.contains_key(&ch_a)); assert!(!s.turn_counts.contains_key(&ch_a)); assert!(!s.core_sections.contains_key(&ch_a)); assert!(!s.has_channel_state(&ch_a)); // ch_b untouched assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); // heartbeat untouched assert_eq!(s.heartbeat_session.as_deref(), Some("sess-hb")); assert_eq!(s.heartbeat_turn_count, 7); } #[test] fn test_invalidate_heartbeat_clears_session_and_turn_count() { let (mut s, ch_a, ch_b) = make_state(); s.invalidate(&PromptSource::Heartbeat); assert!(s.heartbeat_session.is_none()); assert_eq!(s.heartbeat_turn_count, 0); // channels untouched assert_eq!(s.sessions.len(), 2); assert_eq!(*s.turn_counts.get(&ch_a).unwrap(), 5); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_a).unwrap(), "core-a"); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); } #[test] fn test_invalidate_all_clears_everything() { let (mut s, _ch_a, _ch_b) = make_state(); s.invalidate_all(); assert!(s.sessions.is_empty()); assert!(s.turn_counts.is_empty()); assert!(s.core_sections.is_empty()); assert!(s.heartbeat_session.is_none()); assert_eq!(s.heartbeat_turn_count, 0); } #[test] fn test_invalidate_nonexistent_channel_is_noop() { let (mut s, ch_a, ch_b) = make_state(); let ghost = Uuid::new_v4(); s.invalidate(&PromptSource::Channel(ghost)); // Everything still intact. assert_eq!(s.sessions.len(), 2); assert_eq!(s.turn_counts.len(), 2); assert_eq!(*s.turn_counts.get(&ch_a).unwrap(), 5); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_a).unwrap(), "core-a"); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); } #[test] fn test_invalidate_all_on_empty_state_is_noop() { let mut s = SessionState::default(); s.invalidate_all(); // should not panic assert!(s.sessions.is_empty()); assert!(s.turn_counts.is_empty()); assert!(s.core_sections.is_empty()); } #[test] fn test_invalidate_channel_returns_true_when_session_existed() { let (mut s, ch_a, ch_b) = make_state(); assert!(s.invalidate_channel(&ch_a)); assert!(!s.sessions.contains_key(&ch_a)); assert!(!s.turn_counts.contains_key(&ch_a)); assert!(!s.core_sections.contains_key(&ch_a)); assert!(!s.has_channel_state(&ch_a)); // ch_b untouched assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); // heartbeat untouched assert_eq!(s.heartbeat_session.as_deref(), Some("sess-hb")); assert_eq!(s.heartbeat_turn_count, 7); } #[test] fn test_invalidate_channel_returns_false_when_no_session() { let (mut s, _ch_a, _ch_b) = make_state(); let ghost = Uuid::new_v4(); assert!(!s.invalidate_channel(&ghost)); // Nothing changed. assert_eq!(s.sessions.len(), 2); assert_eq!(s.turn_counts.len(), 2); } #[test] fn test_removed_channels_cleaned_via_invalidate_channel() { // Simulates handle_prompt_result: channels removed while agent // was checked out should have both sessions and turn_counts stripped. let (mut s, ch_a, ch_b) = make_state(); let removed = vec![ch_a]; for ch in &removed { s.invalidate_channel(ch); } assert!(!s.sessions.contains_key(&ch_a)); assert!(!s.turn_counts.contains_key(&ch_a)); assert!(!s.core_sections.contains_key(&ch_a)); assert!(!s.has_channel_state(&ch_a)); assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); } // ── ControlSignal::SwitchModel (Phase 3a, Option ii) ───────────────────── #[test] fn test_switch_model_after_natural_completion_invalidates_channel_state() { let (mut s, ch_a, ch_b) = make_state(); // SwitchModel must invalidate just like Rotate so the requeued turn // re-creates a fresh session that re-applies the new desired_model. apply_completed_before_control_signal( &mut s, &PromptSource::Channel(ch_a), &ControlSignal::SwitchModel("gpt-5".into()), ); assert!(!s.has_channel_state(&ch_a)); // ch_b untouched — the switch is channel-scoped. assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); } // ── turn liveness emission ─────────────────────────────────────────────── // `run_turn_liveness` is raced against a "prompt" future the same way // `run_prompt_task` does it: the prompt wins the select and the liveness // future is dropped. We assert what the observer saw. fn liveness_count(handle: &observer::ObserverHandle) -> usize { handle .snapshot() .iter() .filter(|e| e.kind == "turn_liveness") .count() } #[tokio::test(start_paused = true)] async fn test_liveness_fires_while_prompt_pends_then_stops() { let observer = observer::ObserverHandle::in_process(); let context = observer::context_for(None, None, Some("t-1".into())); let liveness = run_turn_liveness( Some(observer.clone()), Some(0), context, Duration::from_secs(10), ); tokio::pin!(liveness); // Prompt pends for 25s, then completes — first liveness tick at 10s, // second at 20s, so the observer must see exactly two pings. tokio::select! { biased; () = tokio::time::sleep(Duration::from_secs(25)) => {} _ = &mut liveness => unreachable!("liveness future never resolves"), } assert_eq!(liveness_count(&observer), 2); // The turn carried the live turn_id on each ping. let pings: Vec<_> = observer .snapshot() .into_iter() .filter(|e| e.kind == "turn_liveness") .collect(); assert!(pings.iter().all(|e| e.turn_id.as_deref() == Some("t-1"))); // After the prompt wins the select, the liveness future is dropped — // advancing the clock further produces no new pings. tokio::time::sleep(Duration::from_secs(60)).await; assert_eq!(liveness_count(&observer), 2); } #[tokio::test(start_paused = true)] async fn test_liveness_disabled_when_interval_zero_emits_nothing() { let observer = observer::ObserverHandle::in_process(); let context = observer::context_for(None, None, Some("t-1".into())); let liveness = run_turn_liveness(Some(observer.clone()), Some(0), context, Duration::ZERO); tokio::pin!(liveness); tokio::select! { biased; () = tokio::time::sleep(Duration::from_secs(120)) => {} _ = &mut liveness => unreachable!("disabled liveness future never resolves"), } assert_eq!(liveness_count(&observer), 0); } #[tokio::test(start_paused = true)] async fn test_liveness_without_observer_emits_nothing() { // A turn that never started has no observer handle — the future must // park without emitting or panicking. let context = observer::context_for(None, None, Some("t-1".into())); let liveness = run_turn_liveness(None, None, context, Duration::from_secs(10)); tokio::pin!(liveness); tokio::select! { biased; () = tokio::time::sleep(Duration::from_secs(120)) => {} _ = &mut liveness => unreachable!("handle-less liveness future never resolves"), } // No observer to assert against — reaching here without panic is the test. } }