mirror of
https://github.com/block/buzz.git
synced 2026-11-04 00:12:02 +00:00
3169 lines
123 KiB
Rust
3169 lines
123 KiB
Rust
//! Agent pool — owns N AcpClient instances and dispatches prompt tasks.
|
||
//!
|
||
//! # Mental model
|
||
//!
|
||
//! ```text
|
||
//! AgentPool
|
||
//! ├── agents: Vec<Option<OwnedAgent>> ← idle agents sit here
|
||
//! ├── join_set: JoinSet<()> ← in-flight tasks
|
||
//! ├── task_map: HashMap<Id, TaskMeta> ← 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<Uuid>,
|
||
/// Clone of batch for Queue mode panic recovery.
|
||
pub recoverable_batch: Option<FlushBatch>,
|
||
/// Control signal for the in-flight prompt task.
|
||
/// `None` for heartbeat tasks (not controllable) and after signal is consumed.
|
||
pub control_tx: Option<tokio::sync::oneshot::Sender<ControlSignal>>,
|
||
/// 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<tokio::sync::mpsc::Sender<SteerRequest>>,
|
||
}
|
||
|
||
/// 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<serde_json::Value>,
|
||
/// Unstable: SessionModelState from session/new.
|
||
pub available_models_raw: Option<serde_json::Value>,
|
||
}
|
||
|
||
/// 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<Uuid, String>,
|
||
pub heartbeat_session: Option<String>,
|
||
/// Per-channel turn counters for proactive session rotation.
|
||
/// Incremented on each successful prompt; reset when the session is rotated.
|
||
pub turn_counts: HashMap<Uuid, u32>,
|
||
/// 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<Uuid, String>,
|
||
}
|
||
|
||
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<AgentModelCapabilities>,
|
||
/// Desired model ID (from `Config.model`). Applied after every `session_new_full()`.
|
||
pub desired_model: Option<String>,
|
||
/// 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<Option<OwnedAgent>>,
|
||
result_tx: mpsc::UnboundedSender<PromptResult>,
|
||
result_rx: mpsc::UnboundedReceiver<PromptResult>,
|
||
pub join_set: JoinSet<()>,
|
||
task_map: HashMap<tokio::task::Id, TaskMeta>,
|
||
}
|
||
|
||
/// 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<FlushBatch>,
|
||
}
|
||
|
||
/// 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<String>,
|
||
/// Oneshot for the read loop to report the outcome.
|
||
pub ack_tx: tokio::sync::oneshot::Sender<SteerAck>,
|
||
}
|
||
|
||
/// 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<McpServer>,
|
||
pub initial_message: Option<String>,
|
||
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<String>,
|
||
pub heartbeat_prompt: Option<String>,
|
||
/// 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<Uuid, ChannelInfo>,
|
||
/// 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<nostr::PublicKey>,
|
||
/// 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<Option<OwnedAgent>>) -> 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<Uuid>) -> Option<OwnedAgent> {
|
||
// 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<tokio::task::Id, TaskMeta> {
|
||
&self.task_map
|
||
}
|
||
|
||
pub fn task_map_mut(&mut self) -> &mut HashMap<tokio::task::Id, TaskMeta> {
|
||
&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<PromptResult> {
|
||
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<PromptResult>, &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<PromptResult, mpsc::error::TryRecvError> {
|
||
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<Option<OwnedAgent>> {
|
||
&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<String, AcpError> {
|
||
// 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<String> = 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<String> {
|
||
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<String> {
|
||
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<String>, core: Option<&str>) -> Option<String> {
|
||
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<FlushBatch>,
|
||
prompt_text: Option<String>,
|
||
ctx: Arc<PromptContext>,
|
||
result_tx: mpsc::UnboundedSender<PromptResult>,
|
||
control_rx: Option<tokio::sync::oneshot::Receiver<ControlSignal>>,
|
||
) {
|
||
// 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<String> = 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<String> = 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<String> = None;
|
||
let prompt_sections: Vec<String> = 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, Fut, T>(f: F) -> Option<T>
|
||
where
|
||
F: Fn() -> Fut,
|
||
Fut: std::future::Future<Output = Option<T>>,
|
||
{
|
||
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<PromptChannelInfo> {
|
||
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<PromptChannelInfo>,
|
||
ctx: &PromptContext,
|
||
) -> Option<ConversationContext> {
|
||
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<String> {
|
||
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<String> {
|
||
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<String> = 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<PromptProfileLookup> {
|
||
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::<serde_json::Value>(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<PromptProfileLookup> {
|
||
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<nostr::PublicKey> = 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<ConversationContext> {
|
||
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<ConversationContext> {
|
||
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<ConversationContext> {
|
||
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<ConversationContext> {
|
||
let arr = json.get("messages").and_then(|v| v.as_array())?;
|
||
|
||
let mut messages: Vec<ContextMessage> =
|
||
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<ContextMessage> {
|
||
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<ConversationContext> {
|
||
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<ConversationContext> {
|
||
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<ContextMessage> = 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<FlushBatch>) -> Option<FlushBatch> {
|
||
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<FlushBatch>,
|
||
) -> Option<FlushBatch> {
|
||
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<observer::ObserverHandle>,
|
||
agent_index: Option<usize>,
|
||
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<observer::ObserverHandle>,
|
||
agent_index: Option<usize>,
|
||
channel_id: Option<uuid::Uuid>,
|
||
turn_id: String,
|
||
}
|
||
|
||
impl TurnCompletionGuard {
|
||
fn new(
|
||
observer: Option<observer::ObserverHandle>,
|
||
agent_index: Option<usize>,
|
||
channel_id: Option<uuid::Uuid>,
|
||
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.
|
||
}
|
||
}
|