mirror of
https://github.com/block/buzz.git
synced 2026-11-01 06:51:17 +00:00
a1cf1db67f
Signed-off-by: Will Pfleger <pfleger.will@gmail.com> Co-authored-by: npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7 <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@sprout-oss.stage.blox.sqprod.co>
3056 lines
117 KiB
Rust
3056 lines
117 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, resolve_model_switch_method, AcpClient,
|
||
AcpError, McpServer, ModelSwitchMethod, StopReason,
|
||
};
|
||
use crate::config::{DedupMode, PermissionMode};
|
||
use crate::observer;
|
||
use crate::queue::{
|
||
ContextMessage, ConversationContext, FlushBatch, PromptChannelInfo, PromptProfile,
|
||
PromptProfileLookup,
|
||
};
|
||
use crate::relay::{ChannelInfo, RestClient};
|
||
|
||
// ── FlushBatch Clone note ─────────────────────────────────────────────────────
|
||
// FlushBatch and BatchEvent derive Clone (added in queue.rs) so we can store
|
||
// a recoverable copy in TaskMeta for panic recovery in Queue mode.
|
||
|
||
// ── Types ─────────────────────────────────────────────────────────────────────
|
||
|
||
/// 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>>,
|
||
}
|
||
|
||
/// 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>,
|
||
/// 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,
|
||
) {
|
||
if control_signal == ControlSignal::Rotate {
|
||
state.invalidate(source);
|
||
}
|
||
}
|
||
|
||
/// Control signal for an in-flight channel turn.
|
||
#[derive(Clone, Copy, 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.
|
||
Interrupt,
|
||
/// Stop the current turn and drop its triggering batch. The session is
|
||
/// invalidated just like cancel; the next turn creates a fresh session.
|
||
Rotate,
|
||
}
|
||
|
||
/// 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,
|
||
}
|
||
|
||
// ── AgentPool impl ────────────────────────────────────────────────────────────
|
||
|
||
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
|
||
}
|
||
|
||
// ── Accessors ─────────────────────────────────────────────────────────
|
||
|
||
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
|
||
}
|
||
|
||
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
|
||
}
|
||
}
|
||
|
||
// ── run_prompt_task ───────────────────────────────────────────────────────────
|
||
|
||
/// 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.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.
|
||
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?;
|
||
}
|
||
None => {
|
||
tracing::warn!(
|
||
target: "pool::model",
|
||
"desired model {desired} not found in agent's available models — proceeding with agent default"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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]`.
|
||
fn framed_system_prompt(base_prompt: Option<&str>, system_prompt: Option<&str>) -> Option<String> {
|
||
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,
|
||
}
|
||
}
|
||
|
||
/// 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>>,
|
||
) {
|
||
// ── Determine source and resolve/create session ───────────────────────
|
||
|
||
// 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,
|
||
}),
|
||
);
|
||
|
||
// ── Reaction cleanup guard ────────────────────────────────────────────
|
||
// Collects event IDs up front. On drop (any exit path — normal, early
|
||
// return, or panic), spawns best-effort cleanup of both 👀 and 💬.
|
||
// See `ReactionGuard` docs for ordering guarantees and known edge cases.
|
||
let reaction_ids: Vec<String> = batch
|
||
.as_ref()
|
||
.map(|b| b.events.iter().map(|be| be.event.id.to_hex()).collect())
|
||
.unwrap_or_default();
|
||
let _reaction_guard = ReactionGuard::new(ctx.rest_client.clone(), reaction_ids.clone());
|
||
|
||
// ── 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(),
|
||
);
|
||
|
||
// ── NIP-AE: fetch core engram before session creation ───────────────
|
||
//
|
||
// 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,
|
||
}),
|
||
);
|
||
|
||
// ── Send initial_message on new channel sessions ──────────────────────
|
||
|
||
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;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── Build prompt text (with optional context fetch) ──────────────────
|
||
|
||
// 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;
|
||
};
|
||
|
||
// 💬 — fire-and-forget so the prompt fires immediately.
|
||
// The guard's cleanup (spawned on drop) removes 💬 after the turn completes.
|
||
// A brief race where 💬 appears slightly after the agent starts is acceptable.
|
||
if !reaction_ids.is_empty() {
|
||
let rest = ctx.rest_client.clone();
|
||
let ids = reaction_ids.clone();
|
||
tokio::spawn(async move {
|
||
react_working(&rest, &ids).await;
|
||
});
|
||
}
|
||
|
||
// ── 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(),
|
||
};
|
||
|
||
// ── Control-aware prompt dispatch ─────────────────────────────────────
|
||
// 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);
|
||
// 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 = match control_signal {
|
||
ControlSignal::Interrupt => requeue_batch_if_queue(&ctx, batch),
|
||
ControlSignal::Cancel | ControlSignal::Rotate => None,
|
||
};
|
||
let _ = result_tx.send(PromptResult {
|
||
agent,
|
||
source,
|
||
outcome: PromptOutcome::Cancelled,
|
||
batch: retry_batch,
|
||
});
|
||
return;
|
||
}
|
||
Err(AcpError::AgentExited) => {
|
||
agent.state.invalidate_all();
|
||
let retry_batch = match control_signal {
|
||
ControlSignal::Interrupt => requeue_batch_if_queue(&ctx, batch),
|
||
ControlSignal::Cancel | ControlSignal::Rotate => None,
|
||
};
|
||
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 = match control_signal {
|
||
ControlSignal::Interrupt => requeue_batch_if_queue(&ctx, batch),
|
||
ControlSignal::Cancel | ControlSignal::Rotate => None,
|
||
};
|
||
let _ = result_tx.send(PromptResult {
|
||
agent,
|
||
source,
|
||
outcome: PromptOutcome::Timeout,
|
||
batch: retry_batch,
|
||
});
|
||
return;
|
||
}
|
||
Err(e) => {
|
||
agent.state.invalidate(&source);
|
||
let retry_batch = match control_signal {
|
||
ControlSignal::Interrupt => requeue_batch_if_queue(&ctx, batch),
|
||
ControlSignal::Cancel | ControlSignal::Rotate => None,
|
||
};
|
||
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 control_signal == ControlSignal::Rotate {
|
||
tracing::debug!(
|
||
target: "pool::prompt",
|
||
"rotate 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);
|
||
|
||
// ── Session rotation on context exhaustion ────────────────
|
||
let should_rotate = matches!(
|
||
stop_reason,
|
||
StopReason::MaxTokens | StopReason::MaxTurnRequests
|
||
);
|
||
|
||
// ── Proactive turn-based rotation ─────────────────────────
|
||
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),
|
||
});
|
||
}
|
||
}
|
||
// _reaction_guard drops here → spawns clear_reactions for all exit paths.
|
||
}
|
||
|
||
// ── Context fetching ──────────────────────────────────────────────────────────
|
||
|
||
/// 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
|
||
}
|
||
|
||
/// 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);
|
||
lookup.insert(
|
||
pk.to_ascii_lowercase(),
|
||
PromptProfile {
|
||
display_name,
|
||
nip05_handle,
|
||
},
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
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,
|
||
})
|
||
}
|
||
|
||
// ── Internal helpers ──────────────────────────────────────────────────────────
|
||
|
||
/// 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,
|
||
}
|
||
}
|
||
|
||
/// 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}");
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── Reaction indicators ───────────────────────────────────────────────────────
|
||
//
|
||
// Two-phase lifecycle visible to users:
|
||
// 👀 "seen" — event was queued and an agent will handle it
|
||
// 💬 "working" — agent is actively prompting
|
||
//
|
||
// 💬 is awaited inline in `run_prompt_task` before the prompt fires, so
|
||
// add-before-remove ordering is structural. 👀 is fire-and-forget from
|
||
// `main.rs` at queue-push time for immediate responsiveness; on rare
|
||
// fast-failure paths the guard's cleanup may race with the 👀 add,
|
||
// leaving a cosmetic stale 👀 (see `ReactionGuard` docs).
|
||
//
|
||
// Cleanup is fire-and-forget via `ReactionGuard` (spawned on drop).
|
||
// Failures are debug-logged and ignored — reactions are cosmetic.
|
||
|
||
/// Drop guard that spawns reaction cleanup on any exit path.
|
||
///
|
||
/// Created at the top of `run_prompt_task`. On drop — normal return, early
|
||
/// return, or panic — spawns fire-and-forget removal of both 👀 and 💬.
|
||
///
|
||
/// ## Ordering
|
||
///
|
||
/// 💬 (`react_working`) is fire-and-forget (spawned before the prompt fires).
|
||
/// A brief race where 💬 appears slightly after the agent starts is acceptable.
|
||
///
|
||
/// 👀 (`react_seen`) is fire-and-forget from `main.rs` at queue-push time.
|
||
/// On rare fast-failure paths (e.g., `session_new` error on an idle agent),
|
||
/// the cleanup spawn may race with the 👀 add, leaving a stale 👀. This is
|
||
/// accepted as a cosmetic edge case — the message will be retried and the
|
||
/// stale 👀 is harmless.
|
||
struct ReactionGuard {
|
||
rest: Option<crate::relay::RestClient>,
|
||
ids: Vec<String>,
|
||
}
|
||
|
||
impl ReactionGuard {
|
||
fn new(rest: crate::relay::RestClient, ids: Vec<String>) -> Self {
|
||
Self {
|
||
rest: if ids.is_empty() { None } else { Some(rest) },
|
||
ids,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Drop for ReactionGuard {
|
||
fn drop(&mut self) {
|
||
// Guard against drop outside a tokio runtime (e.g., in unit tests or
|
||
// during process teardown before the runtime is fully initialized).
|
||
// `run_prompt_task` is always spawned via `JoinSet::spawn`, so a
|
||
// runtime handle is normally available; `try_current` is the safe
|
||
// fallback for the rare cases it isn't.
|
||
if let Some(rest) = self.rest.take() {
|
||
let ids = std::mem::take(&mut self.ids);
|
||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||
handle.spawn(clear_reactions(rest, ids));
|
||
}
|
||
// If no runtime is available, reactions are left as-is — they are
|
||
// cosmetic indicators and the stale state is harmless.
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── 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!({}),
|
||
);
|
||
}
|
||
}
|
||
|
||
// ── Turn completion scope guard ──────────────────────────────────────────────
|
||
// 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!({}),
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
const REACTION_SEEN: &str = "👀";
|
||
const REACTION_WORKING: &str = "💬";
|
||
|
||
/// Best-effort timeout for a single reaction REST call.
|
||
const REACTION_TIMEOUT: Duration = Duration::from_millis(500);
|
||
|
||
/// Percent-encode a string for use in a URL path segment (used in tests only).
|
||
#[cfg(test)]
|
||
fn pct_encode(s: &str) -> String {
|
||
let mut out = String::with_capacity(s.len() * 3);
|
||
for byte in s.bytes() {
|
||
match byte {
|
||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
|
||
out.push(byte as char);
|
||
}
|
||
_ => {
|
||
use std::fmt::Write;
|
||
let _ = write!(out, "%{byte:02X}");
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Best-effort: add a reaction via a signed Nostr kind-7 event (NIP-25).
|
||
///
|
||
/// Builds a reaction event with `buzz_sdk::build_reaction`, signs it with
|
||
/// the keys already stored in `RestClient`, and submits via `POST /events`.
|
||
/// Returns immediately on timeout or any error — reactions are cosmetic.
|
||
pub(crate) async fn reaction_add(rest: &crate::relay::RestClient, event_id: &str, emoji: &str) {
|
||
let target_id = match nostr::EventId::from_hex(event_id) {
|
||
Ok(id) => id,
|
||
Err(e) => {
|
||
tracing::debug!(event_id, emoji, "reaction add: invalid event ID: {e}");
|
||
return;
|
||
}
|
||
};
|
||
let builder = match buzz_sdk::build_reaction(target_id, emoji) {
|
||
Ok(b) => b,
|
||
Err(e) => {
|
||
tracing::warn!(event_id, emoji, "reaction add: build failed: {e}");
|
||
return;
|
||
}
|
||
};
|
||
let event = match builder.sign_with_keys(&rest.keys) {
|
||
Ok(e) => e,
|
||
Err(e) => {
|
||
tracing::warn!(event_id, emoji, "reaction add: sign failed: {e}");
|
||
return;
|
||
}
|
||
};
|
||
match tokio::time::timeout(REACTION_TIMEOUT, rest.submit_event(&event)).await {
|
||
Ok(Ok(_)) => {}
|
||
Ok(Err(e)) => tracing::debug!(event_id, emoji, "reaction add failed: {e}"),
|
||
Err(_) => tracing::debug!(event_id, emoji, "reaction add timed out"),
|
||
}
|
||
}
|
||
|
||
/// Best-effort: remove a reaction via a signed kind:5 (NIP-09) deletion event.
|
||
///
|
||
/// Queries kind:7 reactions by our pubkey targeting the event, finds the matching
|
||
/// emoji, then submits a signed kind:5 deletion via `POST /events`.
|
||
/// Returns immediately on timeout or any error — reactions are cosmetic.
|
||
pub(crate) async fn reaction_remove(rest: &crate::relay::RestClient, event_id: &str, emoji: &str) {
|
||
use nostr::{Alphabet, SingleLetterTag};
|
||
|
||
// Step 1: query our kind:7 reactions targeting this event.
|
||
let my_pubkey = rest.keys.public_key();
|
||
let e_tag = SingleLetterTag::lowercase(Alphabet::E);
|
||
let filter = nostr::Filter::new()
|
||
.kind(nostr::Kind::Reaction)
|
||
.author(my_pubkey)
|
||
.custom_tags(e_tag, [event_id]);
|
||
|
||
let resp = match tokio::time::timeout(Duration::from_millis(1_000), rest.query(&[filter])).await
|
||
{
|
||
Ok(Ok(v)) => v,
|
||
Ok(Err(e)) => {
|
||
tracing::debug!(event_id, emoji, "reaction remove: query failed: {e}");
|
||
return;
|
||
}
|
||
Err(_) => {
|
||
tracing::debug!(event_id, emoji, "reaction remove: query timed out");
|
||
return;
|
||
}
|
||
};
|
||
|
||
// Find our reaction event with matching emoji content.
|
||
let reid = resp.as_array().and_then(|events| {
|
||
events.iter().find_map(|ev| {
|
||
let content = ev.get("content")?.as_str()?;
|
||
if content != emoji {
|
||
return None;
|
||
}
|
||
ev.get("id")?.as_str().map(|s| s.to_string())
|
||
})
|
||
});
|
||
|
||
let reid = match reid {
|
||
Some(id) => id,
|
||
None => {
|
||
tracing::debug!(event_id, emoji, "reaction remove: no reaction event found");
|
||
return;
|
||
}
|
||
};
|
||
|
||
// Step 2: build and submit a signed kind:5 deletion for the reaction event.
|
||
let target_id = match nostr::EventId::from_hex(&reid) {
|
||
Ok(id) => id,
|
||
Err(e) => {
|
||
tracing::debug!(
|
||
event_id,
|
||
emoji,
|
||
"reaction remove: invalid reaction event ID: {e}"
|
||
);
|
||
return;
|
||
}
|
||
};
|
||
let builder = match buzz_sdk::build_remove_reaction(target_id) {
|
||
Ok(b) => b,
|
||
Err(e) => {
|
||
tracing::warn!(event_id, emoji, "reaction remove: build failed: {e}");
|
||
return;
|
||
}
|
||
};
|
||
let event = match builder.sign_with_keys(&rest.keys) {
|
||
Ok(e) => e,
|
||
Err(e) => {
|
||
tracing::warn!(event_id, emoji, "reaction remove: sign failed: {e}");
|
||
return;
|
||
}
|
||
};
|
||
match tokio::time::timeout(Duration::from_millis(1_000), rest.submit_event(&event)).await {
|
||
Ok(Ok(_)) => {}
|
||
Ok(Err(e)) => tracing::debug!(event_id, emoji, "reaction remove failed: {e}"),
|
||
Err(_) => tracing::debug!(event_id, emoji, "reaction remove timed out"),
|
||
}
|
||
}
|
||
|
||
/// Maximum concurrent reaction HTTP requests per fan-out call.
|
||
/// Prevents unbounded parallelism when a large batch of events arrives.
|
||
const REACTION_CONCURRENCY: usize = 10;
|
||
|
||
/// Add 💬 to all events, capped at `REACTION_CONCURRENCY` concurrent requests.
|
||
/// Awaited inline before the prompt fires.
|
||
async fn react_working(rest: &crate::relay::RestClient, event_ids: &[String]) {
|
||
for chunk in event_ids.chunks(REACTION_CONCURRENCY) {
|
||
futures_util::future::join_all(
|
||
chunk
|
||
.iter()
|
||
.map(|eid| reaction_add(rest, eid, REACTION_WORKING)),
|
||
)
|
||
.await;
|
||
}
|
||
}
|
||
|
||
/// Fire-and-forget: remove both 👀 and 💬 from all events. Spawned on turn complete.
|
||
/// Capped at `REACTION_CONCURRENCY` concurrent requests per chunk to avoid
|
||
/// unbounded HTTP fan-out on large batches.
|
||
async fn clear_reactions(rest: crate::relay::RestClient, event_ids: Vec<String>) {
|
||
// Each event needs two removals (👀 and 💬); pair them and chunk by
|
||
// REACTION_CONCURRENCY pairs so the total concurrent requests stay bounded.
|
||
for chunk in event_ids.chunks(REACTION_CONCURRENCY) {
|
||
futures_util::future::join_all(chunk.iter().flat_map(|eid| {
|
||
[
|
||
reaction_remove(&rest, eid, REACTION_SEEN),
|
||
reaction_remove(&rest, eid, REACTION_WORKING),
|
||
]
|
||
}))
|
||
.await;
|
||
}
|
||
}
|
||
|
||
// ─── Unit Tests ──────────────────────────────────────────────────────────────
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use nostr::{EventBuilder, Keys, Kind, Tag};
|
||
use serde_json::json;
|
||
|
||
// ── prepend_base_for_legacy regression tests ─────────────────────────────
|
||
// 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");
|
||
}
|
||
|
||
// ── framed_system_prompt tests ───────────────────────────────────────────
|
||
// 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());
|
||
}
|
||
|
||
// ── with_core tests ──────────────────────────────────────────────────────
|
||
|
||
#[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());
|
||
}
|
||
|
||
// ── parse_thread_response tests ──────────────────────────────────────────
|
||
|
||
#[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());
|
||
}
|
||
|
||
// ── parse_dm_response tests ──────────────────────────────────────────────
|
||
|
||
#[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());
|
||
}
|
||
|
||
// ── json_to_context_message tests ────────────────────────────────────────
|
||
|
||
#[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![],
|
||
};
|
||
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()),
|
||
})
|
||
);
|
||
}
|
||
|
||
#[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");
|
||
}
|
||
|
||
// ── pct_encode tests ─────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_pct_encode_hex_passthrough() {
|
||
let hex = "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2";
|
||
assert_eq!(pct_encode(hex), hex);
|
||
}
|
||
|
||
#[test]
|
||
fn test_pct_encode_emoji() {
|
||
// 👀 = U+1F440 = F0 9F 91 80 in UTF-8
|
||
assert_eq!(pct_encode("👀"), "%F0%9F%91%80");
|
||
}
|
||
|
||
#[test]
|
||
fn test_pct_encode_emoji_speech_balloon() {
|
||
// 💬 = U+1F4AC = F0 9F 92 AC in UTF-8
|
||
assert_eq!(pct_encode("💬"), "%F0%9F%92%AC");
|
||
}
|
||
|
||
#[test]
|
||
fn test_pct_encode_empty() {
|
||
assert_eq!(pct_encode(""), "");
|
||
}
|
||
|
||
#[test]
|
||
fn test_pct_encode_unreserved_passthrough() {
|
||
assert_eq!(pct_encode("AZaz09-_.~"), "AZaz09-_.~");
|
||
}
|
||
|
||
#[test]
|
||
fn test_pct_encode_reserved_chars() {
|
||
assert_eq!(pct_encode("/"), "%2F");
|
||
assert_eq!(pct_encode("+"), "%2B");
|
||
assert_eq!(pct_encode(" "), "%20");
|
||
}
|
||
|
||
// ── 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");
|
||
}
|
||
|
||
// ── 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.
|
||
}
|
||
}
|