jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-27 04:30:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(sprout-acp): Agent context enrichment — envelope, hints, thread/DM auto-fetch (#69)

Browse Source
This commit is contained in:
tlongwell-block
2026-03-15 21:18:35 -04:00
committed by GitHub
parent 076050c4f6
commit fcb1d37702
5 changed files with 1598 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/sprout-acp/src/config.rs
View File

@@ -147,6 +147,12 @@ pub struct CliArgs {
#[arg(long, env = "SPROUT_ACP_NO_IGNORE_SELF")]
pub no_ignore_self: bool,
/// Maximum number of context messages to include for thread replies and DMs.
/// Set to 0 to disable automatic context fetching. Max 100.
#[arg(long, env = "SPROUT_ACP_CONTEXT_MESSAGE_LIMIT", default_value_t = 12,
value_parser = clap::value_parser!(u32).range(0..=100))]
pub context_message_limit: u32,
}
// ── Merged NIP-01 filter ──────────────────────────────────────────────────────
@@ -183,6 +189,7 @@ pub struct Config {
pub channels_override: Option<Vec<String>>,
pub no_mention_filter: bool,
pub config_path: PathBuf,
pub context_message_limit: u32,
}
impl Config {
@@ -261,13 +268,14 @@ impl Config {
channels_override: args.channels,
no_mention_filter: args.no_mention_filter,
config_path: args.config,
context_message_limit: args.context_message_limit,
})
}
/// Human-readable summary (no secrets).
pub fn summary(&self) -> String {
format!(
"relay={} pubkey={} agent_cmd={} {} mcp_cmd={} timeout={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} ignore_self={}",
"relay={} pubkey={} agent_cmd={} {} mcp_cmd={} timeout={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} ignore_self={} context_limit={}",
self.relay_url,
self.keys.public_key().to_hex(),
self.agent_command,
@@ -279,6 +287,7 @@ impl Config {
self.subscribe_mode,
self.dedup_mode,
self.ignore_self,
self.context_message_limit,
)
}
}
@@ -572,6 +581,7 @@ mod tests {
channels_override: None,
no_mention_filter: false,
config_path: PathBuf::from("./sprout-acp.toml"),
context_message_limit: 12,
}
}

176
crates/sprout-acp/src/main.rs
View File

@@ -7,7 +7,7 @@ mod pool;
mod queue;
mod relay;
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
@@ -75,12 +75,13 @@ async fn main() -> Result<()> {
tracing::info!("subscribed to membership notifications");
// ── Step 3: Discover channels and build subscription rules ────────────────
let channels = relay
let channel_info_map = relay
.discover_channels()
.await
.map_err(|e| anyhow::anyhow!("channel discovery error: {e}"))?;
tracing::info!("discovered {} channel(s)", channels.len());
tracing::info!("discovered {} channel(s)", channel_info_map.len());
let channel_ids: Vec<Uuid> = channel_info_map.keys().copied().collect();
let rules: Vec<SubscriptionRule> = match config.subscribe_mode {
SubscribeMode::Mentions => {
@@ -122,7 +123,7 @@ async fn main() -> Result<()> {
};
// ── Step 4: Subscribe to channels ────────────────────────────────────────
let channel_filters = config::resolve_channel_filters(&config, &channels, &rules);
let channel_filters = config::resolve_channel_filters(&config, &channel_ids, &rules);
if channel_filters.is_empty() {
tracing::warn!("no channel subscriptions resolved — agent will sit idle");
}
@@ -149,6 +150,9 @@ async fn main() -> Result<()> {
.unwrap_or_else(|_| std::path::PathBuf::from("/"))
.to_string_lossy()
.to_string(),
rest_client: relay.rest_client(),
channel_info: channel_info_map,
context_message_limit: config.context_message_limit,
});
// ── Step 6: Heartbeat timer ───────────────────────────────────────────────
@@ -183,11 +187,35 @@ async fn main() -> Result<()> {
});
}
// Track the newest membership notification timestamp per channel so that
// replayed events (returned in DESC order on reconnect) don't override the
// correct final state. The first event seen per channel is the newest; any
// older duplicate for the same channel is skipped.
// Track the newest membership notification timestamp per channel.
// On reconnect the relay replays events newest-first, so the first event
// per channel is authoritative. Any later event with ts < newest is stale.
// Exact duplicates (same event ID) are caught by seen_membership_ids.
//
// Uses strict `<` (not `<=`) so that legitimate live events at the same
// second are both processed. The seen_membership_ids set handles exact
// replays that share the same timestamp.
let mut membership_newest_ts: HashMap<Uuid, u64> = HashMap::new();
// Dedup set for exact membership event replays (bounded, cleared at 2000).
let mut seen_membership_ids: HashSet<String> = HashSet::new();
// Channels the agent has been removed from. When a checked-out agent is
// returned to the pool, its sessions for these channels are stripped, and
// failed/panicked batches for these channels are dropped instead of requeued.
//
// Cleared on re-add (KIND_MEMBER_ADDED_NOTIFICATION) so re-joined channels
// regain session affinity.
//
// Known limitation: if a batch is in-flight when the channel is removed AND
// re-added before the batch returns, the stale batch may be requeued. This
// is acceptable because: (a) the agent is a member again and has access,
// (b) the events are from the agent's authorized history, (c) the window
// is extremely narrow (membership changes are rare, prompt turns are seconds),
// and (d) fixing this would require per-channel epoch tracking on TaskMeta
// and PromptResult — significant complexity for a benign edge case. If strict
// causal invalidation is needed, add a monotonic epoch counter per channel
// and capture it in TaskMeta at dispatch time.
let mut removed_channels: HashSet<Uuid> = HashSet::new();
// ── Step 8: Main orchestration loop ──────────────────────────────────────
//
@@ -225,29 +253,57 @@ async fn main() -> Result<()> {
{
let ch = sprout_event.channel_id;
let ts = sprout_event.event.created_at.as_u64();
let eid = sprout_event.event.id.to_hex();
// Skip stale membership events: on reconnect the relay
// replays events newest-first, so the first event per
// channel is authoritative. Any later (older) event for
// the same channel is outdated and must be ignored.
let dominated = membership_newest_ts
.get(&ch)
.is_some_and(|&newest| ts < newest);
if dominated {
// Two-layer membership dedup:
//
// 1. Exact duplicate rejection (seen_membership_ids):
// Catches the same event replayed on reconnect.
//
// 2. Timestamp watermark (membership_newest_ts):
// Uses strict `<` so that older events from reconnect
// replay are dropped, but legitimate live events at the
// same second are both processed. This is safe because
// exact duplicates are already caught by layer 1.
//
// Why not `<=`? That would suppress legitimate live
// add→remove (or remove→add) sequences in the same second,
// leaving the harness in the wrong membership state.
if !seen_membership_ids.insert(eid.clone()) {
tracing::debug!(
channel_id = %ch,
kind = kind_u32,
ts,
"skipping stale membership notification (newer already processed)"
"skipping duplicate membership notification (same event_id)"
);
continue;
}
membership_newest_ts
.entry(ch)
.and_modify(|v| *v = (*v).max(ts))
.or_insert(ts);
// Bound the dedup set to prevent unbounded growth.
// Re-insert the current ID after clearing so it stays
// protected against immediate replay (same pattern as
// relay.rs BgState::record_event).
if seen_membership_ids.len() > 2000 {
seen_membership_ids.clear();
seen_membership_ids.insert(eid);
}
if let Some(&newest) = membership_newest_ts.get(&ch) {
if ts < newest {
tracing::debug!(
channel_id = %ch,
kind = kind_u32,
ts,
newest,
"skipping stale membership notification (older than newest)"
);
continue;
}
}
membership_newest_ts.insert(ch, ts);
if kind_u32 == KIND_MEMBER_ADDED_NOTIFICATION {
// Clear removal tracking so sessions are not
// stripped for a legitimately re-added channel.
removed_channels.remove(&ch);
if let Some(filter) = config::resolve_dynamic_channel_filter(&config, ch, &rules) {
tracing::info!(channel_id = %ch, "membership notification: subscribing to new channel");
if let Err(e) = relay.subscribe_channel(ch, filter).await {
@@ -261,6 +317,23 @@ async fn main() -> Result<()> {
if let Err(e) = relay.unsubscribe_channel(ch).await {
tracing::warn!("failed to unsubscribe from channel {ch}: {e}");
}
// Drain queued events and invalidate sessions for the
// removed channel. Events already in-flight will
// complete normally (the relay may reject actions if
// the agent lost access).
let drained = queue.drain_channel(ch);
let invalidated = pool.invalidate_channel_sessions(ch);
// Track removed channels so checked-out agents get
// their sessions stripped when they return to the pool.
removed_channels.insert(ch);
if drained > 0 || invalidated > 0 {
tracing::info!(
channel_id = %ch,
drained,
invalidated,
"cleaned up after membership removal"
);
}
}
continue;
}
@@ -329,6 +402,7 @@ async fn main() -> Result<()> {
&config,
*result,
&mut heartbeat_in_flight,
&removed_channels,
)
.await
== LoopAction::Exit
@@ -340,6 +414,7 @@ async fn main() -> Result<()> {
&mut queue,
&config,
&mut heartbeat_in_flight,
&removed_channels,
)
.await
== LoopAction::Exit
@@ -356,6 +431,7 @@ async fn main() -> Result<()> {
&config,
join_error,
&mut heartbeat_in_flight,
&removed_channels,
)
.await;
if pool.live_count() == 0 {
@@ -420,7 +496,6 @@ fn dispatch_pending(pool: &mut AgentPool, queue: &mut EventQueue, ctx: &Arc<Prom
};
tracing::debug!(agent = agent.index, channel = %channel_id, affinity_hit, "agent_claimed");
let prompt_text = queue::format_prompt(&batch, ctx.system_prompt.as_deref());
let recoverable_batch = match ctx.dedup_mode {
DedupMode::Queue => Some(batch.clone()),
DedupMode::Drop => None,
@@ -430,8 +505,10 @@ fn dispatch_pending(pool: &mut AgentPool, queue: &mut EventQueue, ctx: &Arc<Prom
let ctx_clone = Arc::clone(ctx);
let agent_index = agent.index;
// Prompt text is now built inside run_prompt_task (needs async for
// context fetching). Pass None for prompt_text; batch carries the data.
let abort_handle = pool.join_set.spawn(async move {
pool::run_prompt_task(agent, Some(batch), prompt_text, ctx_clone, result_tx).await;
pool::run_prompt_task(agent, Some(batch), None, ctx_clone, result_tx).await;
});
pool.task_map_mut().insert(
@@ -457,8 +534,9 @@ async fn handle_prompt_result(
pool: &mut AgentPool,
queue: &mut EventQueue,
config: &Config,
result: PromptResult,
mut result: PromptResult,
heartbeat_in_flight: &mut bool,
removed_channels: &HashSet<Uuid>,
) -> LoopAction {
let before = pool.task_map().len();
let agent_index = result.agent.index;
@@ -472,7 +550,27 @@ async fn handle_prompt_result(
}
if let Some(batch) = result.batch {
queue.requeue(batch);
// Don't requeue batches for channels the agent was removed from —
// those events are stale and should be silently dropped.
if !removed_channels.contains(&batch.channel_id) {
queue.requeue(batch);
} else {
tracing::debug!(
channel_id = %batch.channel_id,
events = batch.events.len(),
"dropping failed batch for removed channel"
);
}
}
// Strip sessions for channels the agent was removed from while this
// agent was checked out. This covers the gap where invalidate_channel_sessions
// only touches idle agents.
if !removed_channels.is_empty() {
result
.agent
.sessions
.retain(|ch, _| !removed_channels.contains(ch));
}
let outcome_label = match &result.outcome {
@@ -522,6 +620,7 @@ async fn recover_panicked_agent(
config: &Config,
join_error: tokio::task::JoinError,
heartbeat_in_flight: &mut bool,
removed_channels: &HashSet<Uuid>,
) {
let task_id = join_error.id();
let Some(meta) = pool.task_map_mut().remove(&task_id) else {
@@ -539,8 +638,16 @@ async fn recover_panicked_agent(
}
if let Some(batch) = meta.recoverable_batch {
queue.requeue(batch);
tracing::warn!("requeued batch for panicked agent {i}");
// Don't requeue batches for removed channels.
if !removed_channels.contains(&batch.channel_id) {
queue.requeue(batch);
tracing::warn!("requeued batch for panicked agent {i}");
} else {
tracing::debug!(
channel_id = %batch.channel_id,
"dropping panicked batch for removed channel"
);
}
}
match spawn_and_init(config).await {
@@ -566,11 +673,20 @@ async fn drain_ready_join_results(
queue: &mut EventQueue,
config: &Config,
heartbeat_in_flight: &mut bool,
removed_channels: &HashSet<Uuid>,
) -> LoopAction {
while let Some(Some(join_result)) = pool.join_set.join_next().now_or_never() {
if let Err(join_error) = join_result {
tracing::error!("agent task panicked: {join_error}");
recover_panicked_agent(pool, queue, config, join_error, heartbeat_in_flight).await;
recover_panicked_agent(
pool,
queue,
config,
join_error,
heartbeat_in_flight,
removed_channels,
)
.await;
if pool.live_count() == 0 {
return LoopAction::Exit;
}
@@ -603,7 +719,7 @@ fn dispatch_heartbeat(
let agent_index = agent.index;
let abort_handle = pool.join_set.spawn(async move {
pool::run_prompt_task(agent, None, prompt_text, ctx_clone, result_tx).await;
pool::run_prompt_task(agent, None, Some(prompt_text), ctx_clone, result_tx).await;
});
pool.task_map_mut().insert(

591
crates/sprout-acp/src/pool.rs
View File

@@ -30,7 +30,8 @@ use uuid::Uuid;
use crate::acp::{AcpClient, AcpError, McpServer, StopReason};
use crate::config::DedupMode;
use crate::queue::FlushBatch;
use crate::queue::{ContextMessage, ConversationContext, FlushBatch, PromptChannelInfo};
use crate::relay::{ChannelInfo, RestClient};
// ── FlushBatch Clone note ─────────────────────────────────────────────────────
// FlushBatch and BatchEvent derive Clone (added in queue.rs) so we can store
@@ -104,6 +105,12 @@ pub struct PromptContext {
pub system_prompt: Option<String>,
pub heartbeat_prompt: Option<String>,
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,
}
// ── AgentPool impl ────────────────────────────────────────────────────────────
@@ -209,24 +216,49 @@ impl AgentPool {
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.sessions.remove(&channel_id).is_some() {
count += 1;
}
}
}
count
}
}
// ── run_prompt_task ───────────────────────────────────────────────────────────
/// Timeout for pre-prompt context fetches (thread/DM history).
const CONTEXT_FETCH_TIMEOUT: Duration = Duration::from_millis(500);
/// 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. Send the actual prompt with turn timeout.
/// 4. Handle all error paths, always returning the agent via `result_tx`.
/// 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: String,
prompt_text: Option<String>,
ctx: Arc<PromptContext>,
result_tx: mpsc::UnboundedSender<PromptResult>,
) {
@@ -410,6 +442,47 @@ pub async fn run_prompt_task(
}
}
// ── Build prompt text (with optional context fetch) ──────────────────
let prompt_text = if let Some(text) = prompt_text {
// Pre-built prompt (heartbeat or legacy path).
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
};
crate::queue::format_prompt(
b,
ctx.system_prompt.as_deref(),
channel_info.as_ref(),
conversation_context.as_ref(),
)
} else {
// Should not happen — batch is None only for heartbeats which have prompt_text.
// Return the agent to the pool to prevent a permanent slot leak.
tracing::error!("run_prompt_task: no batch and no prompt_text — returning agent");
let _ = result_tx.send(PromptResult {
agent,
source,
outcome: PromptOutcome::Error(AcpError::Protocol("no batch and no prompt_text".into())),
batch: None,
});
return;
};
// ── Send the actual prompt ────────────────────────────────────────────
let prompt_result = timeout(
@@ -514,6 +587,266 @@ pub async fn run_prompt_task(
}
}
// ── Context fetching ──────────────────────────────────────────────────────────
/// Lazy-fetch channel metadata for a channel not in the startup discovery cache.
///
/// Handles channels added dynamically via membership notifications after startup.
/// Uses the same 500ms timeout as context fetches. Returns `None` on failure
/// (graceful degradation — prompt will lack channel name and DM detection).
async fn fetch_channel_info(channel_id: Uuid, rest: &RestClient) -> Option<PromptChannelInfo> {
let path = format!("/api/channels/{}", channel_id);
let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await;
match result {
Ok(Ok(json)) => {
let name = json
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string();
let channel_type = json
.get("channel_type")
.and_then(|v| v.as_str())
.unwrap_or("stream")
.to_string();
Some(PromptChannelInfo { name, channel_type })
}
Ok(Err(e)) => {
tracing::debug!(
channel_id = %channel_id,
"channel info fetch failed: {e} — using defaults"
);
None
}
Err(_) => {
tracing::debug!(
channel_id = %channel_id,
"channel info fetch timed out — using defaults"
);
None
}
}
}
/// 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
}
/// Fetch thread context via REST: `GET /api/channels/{id}/threads/{event_id}?limit=N`
async fn fetch_thread_context(
channel_id: Uuid,
root_event_id: &str,
limit: u32,
rest: &RestClient,
) -> Option<ConversationContext> {
// Defense-in-depth: validate hex before interpolating into URL path.
// Nostr event IDs are 32-byte SHA-256 hashes = 64 hex chars.
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 path = format!(
"/api/channels/{}/threads/{}?limit={}",
channel_id, root_event_id, limit
);
let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await;
match result {
Ok(Ok(json)) => parse_thread_response(json),
Ok(Err(e)) => {
tracing::warn!(
channel_id = %channel_id,
root = root_event_id,
"thread context fetch failed: {e} — falling back to hints-only"
);
None
}
Err(_) => {
tracing::warn!(
channel_id = %channel_id,
root = root_event_id,
"thread context fetch timed out — falling back to hints-only"
);
None
}
}
}
/// Fetch DM context via REST: `GET /api/channels/{id}/messages?limit=N`
async fn fetch_dm_context(
channel_id: Uuid,
limit: u32,
rest: &RestClient,
) -> Option<ConversationContext> {
let path = format!("/api/channels/{}/messages?limit={}", channel_id, limit);
let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await;
match result {
Ok(Ok(json)) => parse_dm_response(json, limit),
Ok(Err(e)) => {
tracing::warn!(
channel_id = %channel_id,
"DM context fetch failed: {e} — falling back to hints-only"
);
None
}
Err(_) => {
tracing::warn!(
channel_id = %channel_id,
"DM context fetch timed out — falling back to hints-only"
);
None
}
}
}
/// Parse the thread REST response into a `ConversationContext::Thread`.
///
/// Expected shape: `{ "root": {...}, "replies": [...], "total_replies": N }`
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`.
///
/// Expected shape: `{ "messages": [...], "next_cursor": ... }`
/// Messages arrive newest-first from the API; we reverse to chronological order.
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(),
})
}
// ── Internal helpers ──────────────────────────────────────────────────────────
/// Return the batch for requeue only in Queue mode; drop it in Drop mode.
@@ -549,3 +882,253 @@ fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) {
}
}
}
// ─── Unit Tests ──────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
// ── 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_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");
}
}

839
crates/sprout-acp/src/queue.rs
View File

@@ -241,6 +241,32 @@ impl EventQueue {
pub fn pending_channels(&self) -> usize {
self.queues.len()
}
/// Drop all queued (non-in-flight) events for a channel.
///
/// Used when the agent is removed from a channel — any pending events
/// for that channel are stale and should not be prompted. Does NOT
/// affect in-flight prompts (those will complete normally; the agent
/// may fail to act if it lost access, but that's handled by the relay).
///
/// Also clears any `retry_after` throttle for the channel.
///
/// Returns the number of events dropped.
pub fn drain_channel(&mut self, channel_id: Uuid) -> usize {
let dropped = self
.queues
.remove(&channel_id)
.map(|q| q.len())
.unwrap_or(0);
self.retry_after.remove(&channel_id);
dropped
}
/// Whether a prompt is currently in-flight for the given channel.
#[allow(dead_code)]
pub fn is_channel_in_flight(&self, channel_id: Uuid) -> bool {
self.in_flight_channels.contains(&channel_id)
}
}
impl Default for EventQueue {
@@ -249,76 +275,322 @@ impl Default for EventQueue {
}
}
// ── NIP-10 tag parsing ────────────────────────────────────────────────────────
/// Parsed thread relationship from NIP-10 `e` tags.
#[derive(Debug, Clone, Default)]
pub struct ThreadTags {
/// Root event ID (hex). Present for all thread replies.
pub root_event_id: Option<String>,
/// Parent event ID (hex). For direct replies to root, equals root.
pub parent_event_id: Option<String>,
/// Mentioned pubkeys from `p` tags (hex).
pub mentioned_pubkeys: Vec<String>,
}
/// Parse NIP-10 thread tags from a Nostr event.
///
/// Detection logic (per research doc §4c):
/// - Find an `e` tag with `root` marker → its value is `root_event_id`
/// - Find an `e` tag with `reply` marker → its value is `parent_event_id`
/// - If only `reply` marker found (direct reply to root), root == parent
/// - `p` tags → mentioned pubkeys
///
/// NOTE: Only handles NIP-10 marker-based format (preferred). The deprecated
/// positional format (no markers, `["e", id, relay_url]`) is not supported —
/// Sprout always generates marker-based tags (see relay messages.rs:762-783).
pub fn parse_thread_tags(event: &Event) -> ThreadTags {
let mut root = None;
let mut reply = None;
let mut mentions = Vec::new();
for tag in event.tags.iter() {
let parts = tag.as_slice();
match parts.first().map(|s| s.as_str()) {
Some("e") if parts.len() >= 4 => {
let id = &parts[1];
let marker = &parts[3];
match marker.as_str() {
"root" => root = Some(id.clone()),
"reply" => reply = Some(id.clone()),
_ => {}
}
}
Some("p") if parts.len() >= 2 => {
mentions.push(parts[1].clone());
}
_ => {}
}
}
// For direct replies to root: single "reply" tag, no "root" tag.
// In that case, root == parent.
let (root_event_id, parent_event_id) = match (root, reply) {
(Some(r), Some(p)) => (Some(r), Some(p)),
(Some(r), None) => (Some(r.clone()), Some(r)),
(None, Some(p)) => (Some(p.clone()), Some(p)),
(None, None) => (None, None),
};
ThreadTags {
root_event_id,
parent_event_id,
mentioned_pubkeys: mentions,
}
}
// ── Prompt formatting ─────────────────────────────────────────────────────────
/// Stream-message kinds — these get the compact format (no raw Tags line).
const STREAM_MESSAGE_KINDS: &[u32] = &[
sprout_core::kind::KIND_STREAM_MESSAGE,
sprout_core::kind::KIND_STREAM_MESSAGE_V2,
];
/// Conversation context fetched by the harness before prompting.
#[derive(Debug, Clone)]
pub enum ConversationContext {
/// Thread context for a reply event.
Thread {
messages: Vec<ContextMessage>,
total: usize,
truncated: bool,
},
/// DM conversation history.
Dm {
messages: Vec<ContextMessage>,
total: usize,
truncated: bool,
},
}
/// Format the per-event block lines for a single [`BatchEvent`].
/// A single message in a conversation context section.
#[derive(Debug, Clone)]
pub struct ContextMessage {
pub pubkey: String,
pub timestamp: String,
pub content: String,
}
/// Channel metadata for prompt formatting.
#[derive(Debug, Clone)]
pub struct PromptChannelInfo {
pub name: String,
pub channel_type: String,
}
/// Format the per-event `[Event]` block for a single [`BatchEvent`].
///
/// Non-stream-message kinds (anything not in `[9, 40002]`) include a
/// `Tags:` line with the raw Nostr tags serialised as a JSON array-of-arrays.
fn format_event_block(channel_id: Uuid, be: &BatchEvent) -> String {
let npub = be
.event
.pubkey
.to_bech32()
.unwrap_or_else(|_| be.event.pubkey.to_hex());
/// Includes: event_id, channel (name + UUID), kind, sender (hex + npub),
/// time, content, all tags (never stripped), and parsed structural fields.
fn format_event_block(
channel_id: Uuid,
channel_info: Option<&PromptChannelInfo>,
be: &BatchEvent,
) -> String {
let hex = be.event.pubkey.to_hex();
let npub = be.event.pubkey.to_bech32().unwrap_or_else(|_| hex.clone());
let time = chrono::DateTime::from_timestamp(be.event.created_at.as_u64() as i64, 0)
.map(|dt| dt.to_rfc3339())
.unwrap_or_else(|| be.event.created_at.as_u64().to_string());
let kind = be.event.kind.as_u16() as u32;
let event_id = be.event.id.to_hex();
let channel_display = match channel_info {
Some(ci) => format!("{} (#{channel_id})", ci.name),
None => channel_id.to_string(),
};
let mut block = format!(
"Channel: {channel_id}\nKind: {kind}\nFrom: {npub}\nTime: {time}\nContent: {}",
"Event ID: {event_id}\n\
Channel: {channel_display}\n\
Kind: {kind}\n\
From: {npub} (hex: {hex})\n\
Time: {time}\n\
Content: {}",
be.event.content,
);
// Include raw tags for non-stream-message kinds.
if !STREAM_MESSAGE_KINDS.contains(&kind) {
let tags_json: Vec<&[String]> = be.event.tags.iter().map(|t| t.as_slice()).collect();
if let Ok(tags_str) = serde_json::to_string(&tags_json) {
block.push_str(&format!("\nTags: {tags_str}"));
}
// Always include tags — they carry structural information.
let tags_json: Vec<&[String]> = be.event.tags.iter().map(|t| t.as_slice()).collect();
if let Ok(tags_str) = serde_json::to_string(&tags_json) {
block.push_str(&format!("\nTags: {tags_str}"));
}
// Parsed structural fields.
let thread = parse_thread_tags(&be.event);
let mut parsed_parts = Vec::new();
if let Some(ref p) = thread.parent_event_id {
parsed_parts.push(format!("parent={p}"));
}
if let Some(ref r) = thread.root_event_id {
parsed_parts.push(format!("root={r}"));
}
if !thread.mentioned_pubkeys.is_empty() {
parsed_parts.push(format!(
"mentions=[{}]",
thread.mentioned_pubkeys.join(", ")
));
}
if !parsed_parts.is_empty() {
block.push_str(&format!("\nParsed: {}", parsed_parts.join(", ")));
}
block
}
/// Format a `[Context]` hints section based on event scope.
fn format_context_hints(
channel_id: Uuid,
channel_info: Option<&PromptChannelInfo>,
thread_tags: &ThreadTags,
is_dm: bool,
has_conversation_context: bool,
) -> String {
let channel_display = match channel_info {
Some(ci) => format!("{} (#{channel_id})", ci.name),
None => channel_id.to_string(),
};
// DM check comes first — a DM reply has both thread tags AND is_dm=true,
// and the scope should be "dm" (not "thread") because the agent is in a DM.
if is_dm {
let is_reply = thread_tags.root_event_id.is_some();
// DM replies use get_thread() because /messages excludes thread replies.
// DM non-replies use get_channel_history() for recent conversation.
let ctx_hint = if has_conversation_context && is_reply {
"Thread context included below. Use get_thread() for full history if truncated."
} else if has_conversation_context {
"Conversation context included below. Use get_channel_history() for full history if truncated."
} else if is_reply {
"Use get_thread() to fetch the reply chain."
} else {
"Use get_channel_history() for conversation context."
};
let mut s = format!(
"[Context]\n\
Scope: dm\n\
Channel: {channel_display}\n\
{ctx_hint}"
);
// If this is a DM reply, include thread structural info as supplementary.
if let Some(ref root) = thread_tags.root_event_id {
s.push_str(&format!("\nThread root: {root}"));
if let Some(ref parent) = thread_tags.parent_event_id {
if parent != root {
s.push_str(&format!("\nParent: {parent}"));
}
}
}
s
} else if let Some(ref root) = thread_tags.root_event_id {
let ctx_hint = if has_conversation_context {
"Thread context included below. Use get_thread() for full history if truncated."
} else {
"Use get_thread() to fetch thread context."
};
let mut s = format!(
"[Context]\n\
Scope: thread\n\
Channel: {channel_display}\n\
Thread root: {root}"
);
if let Some(ref parent) = thread_tags.parent_event_id {
if parent != root {
s.push_str(&format!("\nParent: {parent}"));
}
}
s.push_str(&format!("\n{ctx_hint}"));
s
} else {
format!(
"[Context]\n\
Scope: channel\n\
Channel: {channel_display}\n\
Hint: Use get_channel_history() for recent messages if needed."
)
}
}
/// Format a conversation context section (thread or DM).
fn format_conversation_context(ctx: &ConversationContext) -> String {
let (label, messages, total, truncated) = match ctx {
ConversationContext::Thread {
messages,
total,
truncated,
} => ("Thread Context", messages, total, truncated),
ConversationContext::Dm {
messages,
total,
truncated,
} => ("Conversation Context", messages, total, truncated),
};
let trunc_label = if *truncated { ", truncated" } else { "" };
let mut s = format!(
"[{label} ({} of {total} messages{trunc_label})]",
messages.len()
);
for (i, msg) in messages.iter().enumerate() {
s.push_str(&format!(
"\n[{}] {} ({}): {}",
i + 1,
msg.pubkey,
msg.timestamp,
msg.content,
));
}
s
}
/// Format a [`FlushBatch`] into a prompt string for the agent.
///
/// If `system_prompt` is `Some`, it is prepended as a `[System]` block.
///
/// Single-event format:
/// ```text
/// [Sprout event: <prompt_tag>]
/// Channel: ...
/// Kind: ...
/// From: ...
/// Time: ...
/// Content: ...
/// ```
///
/// Batch format (N > 1):
/// ```text
/// [Sprout events — N events]
///
/// --- Event 1 (<prompt_tag>) ---
/// Channel: ...
/// ...
/// ```
pub fn format_prompt(batch: &FlushBatch, system_prompt: Option<&str>) -> String {
let body = if batch.events.len() == 1 {
/// Produces a stable prompt with these sections (in order):
/// 1. `[System]` — system prompt (if configured)
/// 2. `[Context]` — scope, channel name, structural hints
/// 3. `[Thread Context]` or `[Conversation Context]` — if fetched
/// 4. `[Event]` / `[Sprout events]` — the triggering event(s)
pub fn format_prompt(
batch: &FlushBatch,
system_prompt: Option<&str>,
channel_info: Option<&PromptChannelInfo>,
conversation_context: Option<&ConversationContext>,
) -> String {
// Scope is always derived from the LAST event in the batch — that's the
// one the agent is responding to. Thread/DM context is supplementary info
// included alongside, not a scope override. This prevents mixed batches
// (thread reply + later plain message) from being mislabeled as "thread".
let last_event = batch.events.last().expect("batch must have ≥1 event");
let thread_tags = parse_thread_tags(&last_event.event);
let is_dm = channel_info
.map(|ci| ci.channel_type == "dm")
.unwrap_or(false);
let mut sections: Vec<String> = Vec::with_capacity(4);
// 1. System prompt.
if let Some(sp) = system_prompt {
sections.push(format!("[System]\n{sp}"));
}
// 2. Context hints.
sections.push(format_context_hints(
batch.channel_id,
channel_info,
&thread_tags,
is_dm,
conversation_context.is_some(),
));
// 3. Conversation context (thread or DM).
if let Some(ctx) = conversation_context {
sections.push(format_conversation_context(ctx));
}
// 4. Event block(s).
let event_section = if batch.events.len() == 1 {
let be = &batch.events[0];
format!(
"[Sprout event: {}]\n{}",
be.prompt_tag,
format_event_block(batch.channel_id, be)
format_event_block(batch.channel_id, channel_info, be)
)
} else {
let mut s = format!("[Sprout events — {} events]", batch.events.len());
@@ -327,16 +599,14 @@ pub fn format_prompt(batch: &FlushBatch, system_prompt: Option<&str>) -> String
"\n\n--- Event {} ({}) ---\n{}",
i + 1,
be.prompt_tag,
format_event_block(batch.channel_id, be)
format_event_block(batch.channel_id, channel_info, be)
));
}
s
};
sections.push(event_section);
match system_prompt {
Some(sp) => format!("[System]\n{sp}\n\n{body}"),
None => body,
}
sections.join("\n\n")
}
// ─── Unit Tests ──────────────────────────────────────────────────────────────
@@ -571,12 +841,17 @@ mod tests {
}],
};
let prompt = format_prompt(&batch, None);
let prompt = format_prompt(&batch, None, None, None);
assert!(prompt.starts_with("[Sprout event: @mention]\n"));
// Should contain [Context] section before the event.
assert!(prompt.contains("[Context]"));
assert!(prompt.contains("Scope: channel"));
assert!(prompt.contains("[Sprout event: @mention]\n"));
assert!(prompt.contains(&format!("Channel: {}", ch)));
assert!(prompt.contains(&format!("From: {}", npub)));
assert!(prompt.contains("Content: Hello @agent"));
// Event ID should be present.
assert!(prompt.contains("Event ID:"));
// Should NOT contain "--- Event 1 ---" (that's the multi-event format).
assert!(!prompt.contains("--- Event 1 ---"));
}
@@ -661,21 +936,16 @@ mod tests {
],
};
let prompt = format_prompt(&batch, None);
let prompt = format_prompt(&batch, None, None, None);
assert!(prompt.starts_with("[Sprout events — 3 events]"));
assert!(prompt.contains("[Context]"));
assert!(prompt.contains("[Sprout events — 3 events]"));
assert!(prompt.contains("--- Event 1 (tag-a) ---"));
assert!(prompt.contains("--- Event 2 (tag-b) ---"));
assert!(prompt.contains("--- Event 3 (tag-c) ---"));
assert!(prompt.contains("Content: first message"));
assert!(prompt.contains("Content: second message"));
assert!(prompt.contains("Content: third message"));
// All events reference the same channel.
assert_eq!(
prompt.matches(&format!("Channel: {}", ch)).count(),
3,
"each event block should include the channel id"
);
}
// ── Test 11: system prompt prepended ─────────────────────────────────────
@@ -694,8 +964,8 @@ mod tests {
}],
};
let prompt = format_prompt(&batch, Some("You are a triage bot."));
assert!(prompt.starts_with("[System]\nYou are a triage bot.\n\n[Sprout event: test]\n"));
let prompt = format_prompt(&batch, Some("You are a triage bot."), None, None);
assert!(prompt.starts_with("[System]\nYou are a triage bot.\n\n[Context]"));
}
// ── Test 12: drop mode discards in-flight channel events ─────────────────
@@ -999,4 +1269,451 @@ mod tests {
.expect("ch should be flushable after throttle expires");
assert_eq!(batch3.channel_id, ch);
}
// ── NIP-10 tag parsing tests ─────────────────────────────────────────────
/// Build an event with specific tags for thread testing.
fn make_event_with_tags(content: &str, tags: Vec<Vec<String>>) -> Event {
let keys = Keys::generate();
let nostr_tags: Vec<nostr::Tag> = tags
.iter()
.map(|t| {
let strs: Vec<&str> = t.iter().map(|s| s.as_str()).collect();
nostr::Tag::parse(&strs).unwrap()
})
.collect();
EventBuilder::new(Kind::Custom(9), content, nostr_tags)
.sign_with_keys(&keys)
.unwrap()
}
#[test]
fn test_parse_thread_tags_no_tags() {
let event = make_event("plain message");
let tags = parse_thread_tags(&event);
assert!(tags.root_event_id.is_none());
assert!(tags.parent_event_id.is_none());
assert!(tags.mentioned_pubkeys.is_empty());
}
#[test]
fn test_parse_thread_tags_direct_reply() {
// Direct reply to root: single "reply" tag.
let event = make_event_with_tags(
"reply to root",
vec![vec!["e".into(), "abc123".into(), "".into(), "reply".into()]],
);
let tags = parse_thread_tags(&event);
assert_eq!(tags.root_event_id.as_deref(), Some("abc123"));
assert_eq!(tags.parent_event_id.as_deref(), Some("abc123"));
}
#[test]
fn test_parse_thread_tags_nested_reply() {
// Nested reply: root + reply tags.
let event = make_event_with_tags(
"nested reply",
vec![
vec!["e".into(), "root123".into(), "".into(), "root".into()],
vec!["e".into(), "parent456".into(), "".into(), "reply".into()],
],
);
let tags = parse_thread_tags(&event);
assert_eq!(tags.root_event_id.as_deref(), Some("root123"));
assert_eq!(tags.parent_event_id.as_deref(), Some("parent456"));
}
#[test]
fn test_parse_thread_tags_with_mentions() {
let event = make_event_with_tags(
"hey @alice",
vec![
vec!["p".into(), "alice_pubkey".into()],
vec!["p".into(), "bob_pubkey".into()],
],
);
let tags = parse_thread_tags(&event);
assert!(tags.root_event_id.is_none());
assert_eq!(tags.mentioned_pubkeys, vec!["alice_pubkey", "bob_pubkey"]);
}
#[test]
fn test_parse_thread_tags_root_only() {
// Only root marker, no reply marker — root == parent.
let event = make_event_with_tags(
"reply",
vec![vec!["e".into(), "root123".into(), "".into(), "root".into()]],
);
let tags = parse_thread_tags(&event);
assert_eq!(tags.root_event_id.as_deref(), Some("root123"));
assert_eq!(tags.parent_event_id.as_deref(), Some("root123"));
}
// ── Context formatting tests ─────────────────────────────────────────────
#[test]
fn test_format_prompt_with_channel_info() {
let ch = Uuid::new_v4();
let event = make_event("hello");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
};
let ci = PromptChannelInfo {
name: "engineering".into(),
channel_type: "stream".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None);
assert!(prompt.contains("engineering (#"));
assert!(prompt.contains("Scope: channel"));
}
#[test]
fn test_format_prompt_dm_scope() {
let ch = Uuid::new_v4();
let event = make_event("hey");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "dm".into(),
received_at: Instant::now(),
}],
};
let ci = PromptChannelInfo {
name: "DM".into(),
channel_type: "dm".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None);
assert!(prompt.contains("Scope: dm"));
}
#[test]
fn test_format_prompt_thread_scope() {
let ch = Uuid::new_v4();
let event = make_event_with_tags(
"yes go ahead",
vec![vec![
"e".into(),
"root123".into(),
"".into(),
"reply".into(),
]],
);
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "@mention".into(),
received_at: Instant::now(),
}],
};
let prompt = format_prompt(&batch, None, None, None);
assert!(prompt.contains("Scope: thread"));
assert!(prompt.contains("Thread root: root123"));
}
#[test]
fn test_format_prompt_with_thread_context() {
let ch = Uuid::new_v4();
let event = make_event_with_tags(
"yes go ahead",
vec![vec![
"e".into(),
"root123".into(),
"".into(),
"reply".into(),
]],
);
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "@mention".into(),
received_at: Instant::now(),
}],
};
let ctx = ConversationContext::Thread {
messages: vec![
ContextMessage {
pubkey: "npub1xyz".into(),
timestamp: "2026-03-15T16:30:00Z".into(),
content: "Let's refactor auth".into(),
},
ContextMessage {
pubkey: "npub1def".into(),
timestamp: "2026-03-15T16:35:00Z".into(),
content: "yes go ahead".into(),
},
],
total: 5,
truncated: true,
};
let prompt = format_prompt(&batch, None, None, Some(&ctx));
assert!(prompt.contains("[Thread Context (2 of 5 messages, truncated)]"));
assert!(prompt.contains("Let's refactor auth"));
assert!(prompt.contains("Thread context included below"));
}
#[test]
fn test_format_prompt_with_dm_context() {
let ch = Uuid::new_v4();
let event = make_event("ok do that");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "dm".into(),
received_at: Instant::now(),
}],
};
let ci = PromptChannelInfo {
name: "DM".into(),
channel_type: "dm".into(),
};
let ctx = ConversationContext::Dm {
messages: vec![ContextMessage {
pubkey: "npub1abc".into(),
timestamp: "2026-03-15T16:00:00Z".into(),
content: "Can you deploy?".into(),
}],
total: 1,
truncated: false,
};
let prompt = format_prompt(&batch, None, Some(&ci), Some(&ctx));
assert!(prompt.contains("Scope: dm"));
assert!(prompt.contains("[Conversation Context (1 of 1 messages)]"));
assert!(prompt.contains("Can you deploy?"));
}
#[test]
fn test_format_prompt_dm_reply_hints_get_thread() {
let ch = Uuid::new_v4();
// DM reply event — has thread e-tags.
let event = make_event_with_tags(
"sounds good, do it",
vec![vec![
"e".into(),
"root123".into(),
"".into(),
"reply".into(),
]],
);
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "dm".into(),
received_at: Instant::now(),
}],
};
let ci = PromptChannelInfo {
name: "DM".into(),
channel_type: "dm".into(),
};
// Thread context fetched (as the fetch path does for DM replies).
let ctx = ConversationContext::Thread {
messages: vec![ContextMessage {
pubkey: "npub1xyz".into(),
timestamp: "2026-03-15T16:30:00Z".into(),
content: "Should I deploy?".into(),
}],
total: 1,
truncated: false,
};
let prompt = format_prompt(&batch, None, Some(&ci), Some(&ctx));
// Scope should be "dm", not "thread".
assert!(
prompt.contains("Scope: dm"),
"DM reply should have Scope: dm, got:\n{prompt}"
);
// Hint should point to get_thread(), not get_channel_history().
assert!(
prompt.contains("get_thread()"),
"DM reply hint should mention get_thread(), got:\n{prompt}"
);
// Thread structural info should be present.
assert!(
prompt.contains("Thread root: root123"),
"DM reply should include thread root"
);
// Thread context should be included.
assert!(prompt.contains("Should I deploy?"));
}
#[test]
fn test_format_prompt_dm_non_reply_hints_get_channel_history() {
let ch = Uuid::new_v4();
let event = make_event("hey there");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "dm".into(),
received_at: Instant::now(),
}],
};
let ci = PromptChannelInfo {
name: "DM".into(),
channel_type: "dm".into(),
};
// No context fetched — hints only.
let prompt = format_prompt(&batch, None, Some(&ci), None);
assert!(prompt.contains("Scope: dm"));
assert!(
prompt.contains("get_channel_history()"),
"DM non-reply hint should mention get_channel_history()"
);
assert!(
!prompt.contains("get_thread()"),
"DM non-reply should NOT mention get_thread()"
);
}
#[test]
fn test_format_event_block_includes_event_id() {
let ch = Uuid::new_v4();
let event = make_event("test");
let event_id = event.id.to_hex();
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
};
let prompt = format_prompt(&batch, None, None, None);
assert!(
prompt.contains(&format!("Event ID: {event_id}")),
"prompt should contain the event ID"
);
}
#[test]
fn test_format_event_block_includes_hex_and_npub() {
let ch = Uuid::new_v4();
let event = make_event("test");
let hex = event.pubkey.to_hex();
let npub = event.pubkey.to_bech32().unwrap();
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
};
let prompt = format_prompt(&batch, None, None, None);
assert!(
prompt.contains(&format!("From: {npub} (hex: {hex})")),
"prompt should contain both npub and hex"
);
}
#[test]
fn test_format_event_block_always_includes_tags() {
let ch = Uuid::new_v4();
// Kind 9 (stream message) — tags were previously stripped.
let event = make_event_with_tags("hello", vec![vec!["h".into(), ch.to_string()]]);
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
};
let prompt = format_prompt(&batch, None, None, None);
assert!(
prompt.contains("Tags:"),
"tags should always be included, even for stream messages"
);
}
// ── drain_channel tests ──────────────────────────────────────────────────
#[test]
fn test_drain_channel_removes_pending_events() {
let mut q = EventQueue::new(DedupMode::Queue);
let ch = Uuid::new_v4();
q.push(make_queued(ch, "msg1"));
q.push(make_queued(ch, "msg2"));
assert_eq!(q.pending_count(), 2);
let dropped = q.drain_channel(ch);
assert_eq!(dropped, 2);
assert_eq!(q.pending_count(), 0);
}
#[test]
fn test_drain_channel_does_not_affect_other_channels() {
let mut q = EventQueue::new(DedupMode::Queue);
let ch_a = Uuid::new_v4();
let ch_b = Uuid::new_v4();
q.push(make_queued(ch_a, "A"));
q.push(make_queued(ch_b, "B"));
let dropped = q.drain_channel(ch_a);
assert_eq!(dropped, 1);
assert_eq!(q.pending_count(), 1); // ch_b still has 1
}
#[test]
fn test_drain_channel_clears_retry_after() {
let mut q = EventQueue::new(DedupMode::Queue);
let ch = Uuid::new_v4();
q.push(make_queued(ch, "msg"));
let batch = q.flush_next().unwrap();
q.requeue(batch); // sets retry_after
q.mark_complete(ch);
// Channel is throttled — verify drain clears it.
assert!(!q.has_flushable_work());
let dropped = q.drain_channel(ch);
assert_eq!(dropped, 1);
assert_eq!(q.pending_count(), 0);
}
#[test]
fn test_drain_channel_empty_returns_zero() {
let mut q = EventQueue::new(DedupMode::Queue);
let ch = Uuid::new_v4();
assert_eq!(q.drain_channel(ch), 0);
}
#[test]
fn test_drain_channel_does_not_affect_in_flight() {
let mut q = EventQueue::new(DedupMode::Queue);
let ch = Uuid::new_v4();
q.push(make_queued(ch, "msg1"));
let _batch = q.flush_next().unwrap(); // now in-flight
assert!(q.is_in_flight());
// Push another event while in-flight.
q.push(make_queued(ch, "msg2"));
// drain_channel should only remove the queued event, not the in-flight one.
let dropped = q.drain_channel(ch);
assert_eq!(dropped, 1);
assert!(q.is_in_flight()); // in-flight unaffected
}
}

81
crates/sprout-acp/src/relay.rs
View File

@@ -36,6 +36,7 @@ fn event_channel_capacity() -> usize {
std::env::var("SPROUT_ACP_EVENT_BUFFER")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.map(|v| v.max(1)) // mpsc::channel panics on capacity 0
.unwrap_or(EVENT_CHANNEL_CAPACITY_DEFAULT)
}
/// Maximum number of seen event IDs before the dedup set is cleared.
@@ -61,6 +62,51 @@ use crate::config::ChannelFilter;
// ── Types ─────────────────────────────────────────────────────────────────────
/// Metadata about a channel, populated at discovery time.
#[derive(Debug, Clone)]
pub struct ChannelInfo {
pub name: String,
pub channel_type: String,
}
/// Lightweight REST client for pre-prompt context fetches.
///
/// Extracted from `HarnessRelay` fields so it can be shared (via `Arc`) with
/// spawned prompt tasks without giving them access to the WebSocket.
#[derive(Debug, Clone)]
pub struct RestClient {
pub http: reqwest::Client,
pub base_url: String,
pub api_token: Option<String>,
pub keys: Keys,
}
impl RestClient {
/// GET a JSON endpoint, returning the parsed value on success.
pub async fn get_json(&self, path: &str) -> Result<Value, RelayError> {
let url = format!("{}{}", self.base_url, path);
let builder = self.http.get(&url);
let builder = apply_auth(builder, &self.api_token, &self.keys);
let resp = builder
.send()
.await
.map_err(|e| RelayError::Http(e.to_string()))?;
if !resp.status().is_success() {
return Err(RelayError::Http(format!(
"GET {} returned HTTP {}",
path,
resp.status()
)));
}
resp.json()
.await
.map_err(|e| RelayError::Http(e.to_string()))
}
}
/// Events the harness cares about.
#[derive(Debug, Clone)]
pub struct SproutEvent {
@@ -240,7 +286,7 @@ impl HarnessRelay {
}
/// Discover channels the agent is a member of via `GET /api/channels?member=true`.
pub async fn discover_channels(&self) -> Result<Vec<Uuid>, RelayError> {
pub async fn discover_channels(&self) -> Result<HashMap<Uuid, ChannelInfo>, RelayError> {
let http_url = relay_ws_to_http(&self.relay_url);
let url = format!("{http_url}/api/channels?member=true");
@@ -268,11 +314,23 @@ impl HarnessRelay {
.as_array()
.ok_or_else(|| RelayError::Http("expected JSON array from /api/channels".into()))?;
let mut ids = Vec::with_capacity(channels.len());
let mut map = HashMap::with_capacity(channels.len());
for ch in channels {
if let Some(id_str) = ch.get("id").and_then(|v| v.as_str()) {
match id_str.parse::<Uuid>() {
Ok(uuid) => ids.push(uuid),
Ok(uuid) => {
let name = ch
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string();
let channel_type = ch
.get("channel_type")
.and_then(|v| v.as_str())
.unwrap_or("stream")
.to_string();
map.insert(uuid, ChannelInfo { name, channel_type });
}
Err(e) => {
warn!("skipping channel with unparseable id {id_str:?}: {e}");
}
@@ -280,8 +338,21 @@ impl HarnessRelay {
}
}
debug!("discovered {} channel(s)", ids.len());
Ok(ids)
debug!("discovered {} channel(s)", map.len());
Ok(map)
}
/// Build a [`RestClient`] that shares this relay's HTTP credentials.
///
/// The returned client is cheap to clone (wraps `reqwest::Client` which is
/// internally `Arc`-ed) and safe to share across spawned tasks via `Arc`.
pub fn rest_client(&self) -> RestClient {
RestClient {
http: self.http.clone(),
base_url: relay_ws_to_http(&self.relay_url),
api_token: self.api_token.clone(),
keys: self.keys.clone(),
}
}
/// Subscribe to events in a channel using the given filter.