Closes the #3110 split — lib.rs is now just the crate doc comment and the pub mod list. journal.rs's new doc comment fixes a pre-existing bug: the old JournalPriority doc text in lib.rs was actually half Capability's doc (a leftover from an earlier reorder that moved the code but not the comment above it).
1244 lines
53 KiB
Rust
1244 lines
53 KiB
Rust
//! Sqlite-backed message broker. Survives `hive-c0re` restart, and taps every
|
|
//! send/recv onto a broadcast channel so the dashboard can stream it.
|
|
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::path::Path;
|
|
use std::sync::Mutex;
|
|
|
|
use anyhow::{Context, Result};
|
|
use chrono::Utc;
|
|
|
|
use hive_sh4re::inbox::{InboxRow, Message};
|
|
|
|
use crate::db::Migration;
|
|
use rusqlite::{Connection, OptionalExtension, params};
|
|
use serde::Serialize;
|
|
use tokio::sync::broadcast;
|
|
|
|
const SCHEMA: &str = r"
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
sender TEXT NOT NULL,
|
|
recipient TEXT NOT NULL,
|
|
body TEXT NOT NULL,
|
|
sent_at INTEGER NOT NULL,
|
|
delivered_at INTEGER,
|
|
in_reply_to INTEGER,
|
|
priority INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_messages_undelivered
|
|
ON messages (recipient, priority DESC, id) WHERE delivered_at IS NULL;
|
|
|
|
CREATE TABLE IF NOT EXISTS kv (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
";
|
|
|
|
/// Capacity of the live event channel. Slow subscribers (e.g. an idle browser)
|
|
/// may drop events past this; we send a `lagged` notice in their stream.
|
|
const EVENT_CHANNEL: usize = 256;
|
|
|
|
/// A single message hand-off from broker to recipient. Carries the
|
|
/// broker's row id (so the harness can drive `ack_turn` later) and
|
|
/// the redelivery flag (so the harness can prepend the
|
|
/// "may already be handled" hint to the wake prompt). The
|
|
/// `Message` itself is identical to a pristine `Send` payload.
|
|
#[derive(Debug, Clone)]
|
|
pub struct Delivery {
|
|
pub id: i64,
|
|
pub redelivered: bool,
|
|
pub message: Message,
|
|
}
|
|
|
|
/// Intra-process broker event. `recv_blocking_batch` listens on the
|
|
/// same channel as the dashboard forwarder; the forwarder re-emits
|
|
/// each event as a `DashboardEvent` with a freshly-stamped seq from
|
|
/// the Coordinator. The broker itself doesn't stamp seqs — that's a
|
|
/// wire concern, not a storage concern.
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "snake_case", tag = "kind")]
|
|
pub enum MessageEvent {
|
|
Sent {
|
|
/// Broker row id — used by the dashboard to track thread parents.
|
|
id: i64,
|
|
from: String,
|
|
to: String,
|
|
body: String,
|
|
at: i64,
|
|
in_reply_to: Option<i64>,
|
|
},
|
|
Delivered {
|
|
/// Broker row id — used by the dashboard to track thread parents.
|
|
id: i64,
|
|
from: String,
|
|
to: String,
|
|
body: String,
|
|
at: i64,
|
|
in_reply_to: Option<i64>,
|
|
},
|
|
}
|
|
|
|
/// Per-recipient in-memory bookkeeping for the deliver-then-ack
|
|
/// flow. Source of truth is the DB columns `delivered_at` +
|
|
/// `acked_at`; the in-memory state here is purely an optimisation
|
|
/// (avoids scanning the messages table on `AckTurn`) plus the
|
|
/// redelivery-hint marker.
|
|
#[derive(Default)]
|
|
struct RecipientInflight {
|
|
/// Message ids the broker has handed to this recipient since the
|
|
/// last `AckTurn`. Drained on `ack_turn`, which then runs a
|
|
/// single `UPDATE … WHERE id IN (…)` to set `acked_at`.
|
|
unacked_ids: Vec<i64>,
|
|
/// Message ids resurfaced by the most recent `requeue_inflight`
|
|
/// call. The next `recv_batch` pop of any id in this set tags
|
|
/// the response with `redelivered: true` so the harness can
|
|
/// prepend the "may already be handled" hint to the wake prompt;
|
|
/// successful pops drain the id from the set.
|
|
requeued_ids: HashSet<i64>,
|
|
}
|
|
|
|
pub struct Broker {
|
|
conn: Mutex<Connection>,
|
|
events: broadcast::Sender<MessageEvent>,
|
|
/// Per-recipient deliver/ack tracking. Lost on hive-c0re restart
|
|
/// (harmless — the harness fires `RequeueInflight` on its own
|
|
/// boot, which rebuilds the `requeued_ids` set from the DB and
|
|
/// clears any stale `unacked_ids`).
|
|
inflight: Mutex<HashMap<String, RecipientInflight>>,
|
|
}
|
|
|
|
/// Ordered schema migrations for the broker. Tracked in `schema_versions`
|
|
/// under key `"broker"`. Each migration declares the column it adds, so a
|
|
/// legacy DB (fully or partially migrated via the old per-column approach)
|
|
/// is converged by skipping the migrations whose column already exists.
|
|
const BROKER_MIGRATIONS: &[Migration] = &[
|
|
// v1: acked_at on messages, with backfill so existing delivered rows
|
|
// are not phantom-requeued on the next open. The BEGIN/COMMIT block
|
|
// makes ALTER + UPDATE atomic — either both land or neither, so the
|
|
// guard column (acked_at) faithfully marks the whole step as done.
|
|
Migration {
|
|
sql: "BEGIN;\
|
|
ALTER TABLE messages ADD COLUMN acked_at INTEGER;\
|
|
UPDATE messages SET acked_at = delivered_at \
|
|
WHERE delivered_at IS NOT NULL AND acked_at IS NULL;\
|
|
COMMIT;",
|
|
adds_column: Some(("messages", "acked_at")),
|
|
},
|
|
// v2: in_reply_to for thread-parent tracking. NULL = root of a thread.
|
|
Migration {
|
|
sql: "ALTER TABLE messages ADD COLUMN in_reply_to INTEGER",
|
|
adds_column: Some(("messages", "in_reply_to")),
|
|
},
|
|
// v3: priority for operator-message fast-path. Rebuild the delivery
|
|
// index to include priority as a secondary sort key. Atomic BEGIN/COMMIT
|
|
// so priority existing implies the index rebuild also committed.
|
|
Migration {
|
|
sql: "BEGIN;\
|
|
ALTER TABLE messages ADD COLUMN \
|
|
priority INTEGER NOT NULL DEFAULT 0;\
|
|
DROP INDEX IF EXISTS idx_messages_undelivered;\
|
|
CREATE INDEX IF NOT EXISTS idx_messages_undelivered \
|
|
ON messages (recipient, priority DESC, id) WHERE delivered_at IS NULL;\
|
|
COMMIT;",
|
|
adds_column: Some(("messages", "priority")),
|
|
},
|
|
];
|
|
|
|
impl Broker {
|
|
pub fn open(path: &Path) -> Result<Self> {
|
|
let conn = crate::db::open(path, "broker")?;
|
|
conn.execute_batch(SCHEMA).context("apply broker schema")?;
|
|
crate::db::apply_versioned_migrations(&conn, "broker", BROKER_MIGRATIONS)
|
|
.context("broker migrations")?;
|
|
let (events, _) = broadcast::channel(EVENT_CHANNEL);
|
|
Ok(Self {
|
|
conn: Mutex::new(conn),
|
|
events,
|
|
inflight: Mutex::new(HashMap::new()),
|
|
})
|
|
}
|
|
|
|
pub fn subscribe(&self) -> broadcast::Receiver<MessageEvent> {
|
|
self.events.subscribe()
|
|
}
|
|
|
|
/// Set a small persistent key/value pair (upsert). The `kv` table is a
|
|
/// general single-value store for state that must survive a hive-c0re
|
|
/// restart but isn't worth a dedicated table — currently just the
|
|
/// `hivectl start` running-agents snapshot (see
|
|
/// `Coordinator::set_last_stopped_running`).
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error if the sqlite upsert fails.
|
|
pub fn kv_set(&self, key: &str, value: &str) -> Result<()> {
|
|
let conn = self.conn.lock().unwrap();
|
|
conn.execute(
|
|
"INSERT INTO kv (key, value) VALUES (?1, ?2)
|
|
ON CONFLICT(key) DO UPDATE SET value = excluded.value",
|
|
params![key, value],
|
|
)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Read a persistent key/value pair. `None` when the key is absent.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error if the sqlite query fails.
|
|
pub fn kv_get(&self, key: &str) -> Result<Option<String>> {
|
|
let conn = self.conn.lock().unwrap();
|
|
let value: Option<String> = conn
|
|
.query_row("SELECT value FROM kv WHERE key = ?1", params![key], |row| {
|
|
row.get(0)
|
|
})
|
|
.optional()?;
|
|
Ok(value)
|
|
}
|
|
|
|
/// Delete a persistent key/value pair. Idempotent — deleting an absent
|
|
/// key is a no-op.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error if the sqlite delete fails.
|
|
pub fn kv_delete(&self, key: &str) -> Result<()> {
|
|
let conn = self.conn.lock().unwrap();
|
|
conn.execute("DELETE FROM kv WHERE key = ?1", params![key])?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn send(&self, message: &Message) -> Result<()> {
|
|
let conn = self.conn.lock().unwrap();
|
|
let now = Utc::now().timestamp();
|
|
// Operator messages get elevated priority so they surface before
|
|
// queued wakes (bash completions, forge events, etc.) when the
|
|
// harness pops the next turn driver. All other senders stay at 0.
|
|
let priority: i64 =
|
|
i64::from(message.from.as_str() == hive_sh4re::manager::OPERATOR_RECIPIENT);
|
|
conn.execute(
|
|
"INSERT INTO messages (sender, recipient, body, sent_at, in_reply_to, priority) \
|
|
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
params![
|
|
message.from.as_str(),
|
|
message.to,
|
|
message.body,
|
|
now,
|
|
message.in_reply_to,
|
|
priority
|
|
],
|
|
)?;
|
|
let row_id = conn.last_insert_rowid();
|
|
drop(conn);
|
|
let _ = self.events.send(MessageEvent::Sent {
|
|
id: row_id,
|
|
from: message.from.to_string(),
|
|
to: message.to.clone(),
|
|
body: message.body.clone(),
|
|
at: now,
|
|
in_reply_to: message.in_reply_to,
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
/// Unread (unacked) messages addressed to `recipient`, newest-first.
|
|
/// Filters to `acked_at IS NULL` so the agent inbox view clears after
|
|
/// "mark all read" — mirroring exactly what `mark_all_read` will drain.
|
|
pub fn recent_for(&self, recipient: &str, limit: u64) -> Result<Vec<InboxRow>> {
|
|
let conn = self.conn.lock().unwrap();
|
|
let limit_i = i64::try_from(limit.min(i64::MAX as u64)).unwrap_or(i64::MAX);
|
|
let mut stmt = conn.prepare(
|
|
"SELECT id, sender, body, sent_at, in_reply_to
|
|
FROM messages
|
|
WHERE recipient = ?1
|
|
AND acked_at IS NULL
|
|
ORDER BY id DESC
|
|
LIMIT ?2",
|
|
)?;
|
|
let rows = stmt.query_map(params![recipient, limit_i], |row| {
|
|
Ok(InboxRow {
|
|
id: row.get(0)?,
|
|
from: row.get(1)?,
|
|
body: row.get(2)?,
|
|
at: row.get(3)?,
|
|
in_reply_to: row.get(4)?,
|
|
})
|
|
})?;
|
|
rows.collect::<rusqlite::Result<Vec<_>>>()
|
|
.map_err(Into::into)
|
|
}
|
|
|
|
/// Latest `limit` messages across every recipient, newest-first.
|
|
/// Backs the dashboard's message-flow backfill so a reload doesn't
|
|
/// blank the operator's view of recent traffic. Returns each row as
|
|
/// a [`MessageEvent::Sent`] so the dashboard's live renderer (which
|
|
/// already speaks `MessageEvent`) can replay history through the
|
|
/// same code path. We don't synthesise `Delivered` events here —
|
|
/// the recv-side acks live in a different table column and would
|
|
/// double-render on backfill; the live stream picks them up
|
|
/// immediately on the first new `recv`.
|
|
pub fn recent_all(&self, limit: u64) -> Result<Vec<MessageEvent>> {
|
|
let conn = self.conn.lock().unwrap();
|
|
let limit_i = i64::try_from(limit.min(i64::MAX as u64)).unwrap_or(i64::MAX);
|
|
let mut stmt = conn.prepare(
|
|
"SELECT id, sender, recipient, body, sent_at, in_reply_to
|
|
FROM messages
|
|
ORDER BY id DESC
|
|
LIMIT ?1",
|
|
)?;
|
|
let rows = stmt.query_map(params![limit_i], |row| {
|
|
Ok(MessageEvent::Sent {
|
|
id: row.get(0)?,
|
|
from: row.get(1)?,
|
|
to: row.get(2)?,
|
|
body: row.get(3)?,
|
|
at: row.get(4)?,
|
|
in_reply_to: row.get(5)?,
|
|
})
|
|
})?;
|
|
rows.collect::<rusqlite::Result<Vec<_>>>()
|
|
.map_err(Into::into)
|
|
}
|
|
|
|
/// Unacknowledged messages addressed to `recipient`, newest-first.
|
|
/// Backs the dashboard's operator inbox: the operator never
|
|
/// `recv`s over an agent socket, so messages to `"operator"` sit in
|
|
/// the broker with `acked_at IS NULL` until the operator hits "mark
|
|
/// all read" (which calls [`Broker::mark_all_read`]). This read
|
|
/// mirrors that filter EXACTLY — `recipient = ?1 AND acked_at IS
|
|
/// NULL`, with no `delivered_at` condition — so everything listed
|
|
/// here is precisely what `mark_all_read` will clear (operator rows
|
|
/// may never get `delivered_at` set). Returned as
|
|
/// [`MessageEvent::Sent`] so the dashboard reuses its live renderer.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `Err` if the `SQLite` prepare or query fails.
|
|
pub fn unread_for_recipient(&self, recipient: &str, limit: u64) -> Result<Vec<MessageEvent>> {
|
|
let conn = self.conn.lock().unwrap();
|
|
let limit_i = i64::try_from(limit.min(i64::MAX as u64)).unwrap_or(i64::MAX);
|
|
let mut stmt = conn.prepare(
|
|
"SELECT id, sender, recipient, body, sent_at, in_reply_to
|
|
FROM messages
|
|
WHERE recipient = ?1 AND acked_at IS NULL
|
|
ORDER BY id DESC
|
|
LIMIT ?2",
|
|
)?;
|
|
let rows = stmt.query_map(params![recipient, limit_i], |row| {
|
|
Ok(MessageEvent::Sent {
|
|
id: row.get(0)?,
|
|
from: row.get(1)?,
|
|
to: row.get(2)?,
|
|
body: row.get(3)?,
|
|
at: row.get(4)?,
|
|
in_reply_to: row.get(5)?,
|
|
})
|
|
})?;
|
|
rows.collect::<rusqlite::Result<Vec<_>>>()
|
|
.map_err(Into::into)
|
|
}
|
|
|
|
/// Number of undelivered messages addressed to `recipient`. Non-mutating
|
|
/// — used by the harness to surface "N unread" in tool-result status
|
|
/// lines without popping the queue.
|
|
pub fn count_pending(&self, recipient: &str) -> Result<u64> {
|
|
let conn = self.conn.lock().unwrap();
|
|
// Skip rows closed by `ack_until` while still pending — they
|
|
// will never pop, so counting them would show phantom unread.
|
|
let n: i64 = conn.query_row(
|
|
"SELECT COUNT(*) FROM messages
|
|
WHERE recipient = ?1 AND delivered_at IS NULL AND acked_at IS NULL",
|
|
params![recipient],
|
|
|row| row.get(0),
|
|
)?;
|
|
Ok(u64::try_from(n.max(0)).unwrap_or(0))
|
|
}
|
|
|
|
/// Send a "your parent changed from X to Y" notification to `child`,
|
|
/// coalescing with any existing undelivered one so that multiple moves
|
|
/// while the agent is offline collapse into a single message spanning
|
|
/// the full arc (e.g. A→B then B→C becomes "your parent changed from A to C").
|
|
///
|
|
/// If an undelivered system reparent notification for `child` already
|
|
/// exists, its body is updated in-place preserving the original "from"
|
|
/// label. If none exists, a fresh message is inserted with `old_label`
|
|
/// as the source.
|
|
pub fn send_coalescing_reparent(
|
|
&self,
|
|
child: &str,
|
|
old_label: &str,
|
|
new_label: &str,
|
|
) -> Result<()> {
|
|
const PREFIX: &str = "your parent changed from ";
|
|
const SEPARATOR: &str = " to ";
|
|
let conn = self.conn.lock().unwrap();
|
|
let now = Utc::now().timestamp();
|
|
let existing: Option<(i64, String)> = conn
|
|
.query_row(
|
|
"SELECT id, body FROM messages
|
|
WHERE recipient = ?1
|
|
AND sender = ?2
|
|
AND body LIKE ?3
|
|
AND delivered_at IS NULL
|
|
AND acked_at IS NULL
|
|
LIMIT 1",
|
|
params![
|
|
child,
|
|
hive_sh4re::manager::SYSTEM_SENDER,
|
|
format!("{PREFIX}%")
|
|
],
|
|
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)),
|
|
)
|
|
.optional()?;
|
|
if let Some((row_id, old_body)) = existing {
|
|
// Preserve the original "from" label from the earlier notification.
|
|
let original_from = old_body
|
|
.strip_prefix(PREFIX)
|
|
.and_then(|rest| rest.split(SEPARATOR).next())
|
|
.unwrap_or(old_label);
|
|
let new_body = format!("{PREFIX}{original_from}{SEPARATOR}{new_label}");
|
|
conn.execute(
|
|
"UPDATE messages SET body = ?1, sent_at = ?2 WHERE id = ?3",
|
|
params![new_body, now, row_id],
|
|
)?;
|
|
drop(conn);
|
|
let _ = self.events.send(MessageEvent::Sent {
|
|
id: row_id,
|
|
from: hive_sh4re::manager::SYSTEM_SENDER.to_owned(),
|
|
to: child.to_owned(),
|
|
body: new_body,
|
|
at: now,
|
|
in_reply_to: None,
|
|
});
|
|
} else {
|
|
let body = format!("{PREFIX}{old_label}{SEPARATOR}{new_label}");
|
|
conn.execute(
|
|
"INSERT INTO messages (sender, recipient, body, sent_at, in_reply_to)
|
|
VALUES (?1, ?2, ?3, ?4, NULL)",
|
|
params![hive_sh4re::manager::SYSTEM_SENDER, child, body, now],
|
|
)?;
|
|
let row_id = conn.last_insert_rowid();
|
|
drop(conn);
|
|
let _ = self.events.send(MessageEvent::Sent {
|
|
id: row_id,
|
|
from: hive_sh4re::manager::SYSTEM_SENDER.to_owned(),
|
|
to: child.to_owned(),
|
|
body,
|
|
at: now,
|
|
in_reply_to: None,
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Long-poll variant of `recv_batch`: returns immediately if any
|
|
/// row is pending (popping up to `max`); otherwise waits up to
|
|
/// `timeout` for the broker to emit a `Sent { to: recipient }`
|
|
/// event and re-tries the pop. Lets agents react to new mail
|
|
/// without polling their socket on a fixed interval AND lets a
|
|
/// single round-trip drain a burst of messages.
|
|
///
|
|
/// **Subscribe-before-check order matters.** If we polled the
|
|
/// sqlite row first and only then called `subscribe()`, a
|
|
/// concurrent `send` landing in that window would commit +
|
|
/// broadcast its event *before* our receiver existed — and we'd
|
|
/// then sit on the long-poll until the timeout (or another,
|
|
/// unrelated send) fired. That looked externally like "the agent
|
|
/// processed one wake then went deaf until the operator poked it
|
|
/// again". Subscribing first guarantees any post-subscribe send
|
|
/// notifies us; the redundant `recv_batch()` catches the message
|
|
/// either way.
|
|
///
|
|
/// `max == 0` returns an empty vec without subscribing or waiting.
|
|
pub async fn recv_blocking_batch(
|
|
&self,
|
|
recipient: &str,
|
|
timeout: std::time::Duration,
|
|
max: usize,
|
|
) -> Result<Vec<Delivery>> {
|
|
if max == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut rx = self.subscribe();
|
|
// Immediate check: any pending sqlite messages.
|
|
let batch = self.recv_batch(recipient, max)?;
|
|
if !batch.is_empty() {
|
|
return Ok(batch);
|
|
}
|
|
let deadline = tokio::time::Instant::now() + timeout;
|
|
loop {
|
|
let Some(remaining) = deadline.checked_duration_since(tokio::time::Instant::now())
|
|
else {
|
|
return Ok(Vec::new());
|
|
};
|
|
match tokio::time::timeout(remaining, rx.recv()).await {
|
|
Err(_) => return Ok(Vec::new()),
|
|
// Channel lagged or closed — fall back to a direct recv_batch
|
|
// (in case we missed our notification while behind).
|
|
Ok(Err(_)) => return self.recv_batch(recipient, max),
|
|
// A message was sent to this recipient (or a wake injected).
|
|
// Re-poll sqlite; the event payload is ignored — recv_batch is
|
|
// the source of truth and prevents double-delivery.
|
|
Ok(Ok(MessageEvent::Sent { to, .. })) if to == recipient => {
|
|
let batch = self.recv_batch(recipient, max)?;
|
|
if !batch.is_empty() {
|
|
return Ok(batch);
|
|
}
|
|
// Lost a race (concurrent recv drained it). Keep waiting.
|
|
}
|
|
Ok(Ok(_)) => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Delete fully-acked messages older than `older_than_secs`.
|
|
/// Unacked rows (delivered but not yet acknowledged by a clean
|
|
/// turn-end, plus undelivered rows) are always kept regardless of
|
|
/// age — the former because they're recoverable via
|
|
/// `requeue_inflight`, the latter because they're still in flight
|
|
/// from the broker's POV. Returns the number of rows removed.
|
|
pub fn vacuum_delivered(&self, older_than_secs: i64) -> Result<u64> {
|
|
let cutoff = Utc::now().timestamp() - older_than_secs;
|
|
let conn = self.conn.lock().unwrap();
|
|
let n = conn.execute(
|
|
"DELETE FROM messages
|
|
WHERE acked_at IS NOT NULL
|
|
AND acked_at < ?1",
|
|
params![cutoff],
|
|
)?;
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
|
|
/// Pop up to `max` pending messages for `recipient` in one
|
|
/// round-trip. Every popped row is marked `delivered_at = NOW`,
|
|
/// pushed onto the per-recipient `unacked_ids` list (so the next
|
|
/// `ack_turn` closes them out), and tagged with
|
|
/// `redelivered = true` if it was resurfaced by the most recent
|
|
/// `requeue_inflight`. Emits one `MessageEvent::Delivered` per
|
|
/// popped row so the dashboard forwarder stream sees one event
|
|
/// per message regardless of batch size.
|
|
///
|
|
/// `max == 0` short-circuits to an empty vec (no DB hit); any
|
|
/// positive value caps the batch at `max`. Ordering: higher
|
|
/// priority first; FIFO within the same priority level. Operator
|
|
/// messages (priority = 1) therefore surface before queued bash
|
|
/// completions and forge events (priority = 0).
|
|
///
|
|
/// Lock order: `inflight` FIRST, then `conn`. `requeue_inflight`
|
|
/// and `ack_turn` follow the same order so a concurrent pop can't
|
|
/// race the requeue's DB update vs in-memory populate and miss
|
|
/// the redelivered tag.
|
|
pub fn recv_batch(&self, recipient: &str, max: usize) -> Result<Vec<Delivery>> {
|
|
if max == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
// Same lock order as `recv` / `ack_turn` / `requeue_inflight`.
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
let conn = self.conn.lock().unwrap();
|
|
let max_i = i64::try_from(max).unwrap_or(i64::MAX);
|
|
// `acked_at IS NULL` matters for rows closed by `ack_until`
|
|
// while still pending (never delivered): they carry an ack but
|
|
// no `delivered_at`, and must not pop.
|
|
let mut stmt = conn.prepare(
|
|
"SELECT id, sender, recipient, body, in_reply_to
|
|
FROM messages
|
|
WHERE recipient = ?1 AND delivered_at IS NULL AND acked_at IS NULL
|
|
ORDER BY priority DESC, id ASC
|
|
LIMIT ?2",
|
|
)?;
|
|
let rows: Vec<(i64, String, String, String, Option<i64>)> = stmt
|
|
.query_map(params![recipient, max_i], |row| {
|
|
Ok((
|
|
row.get(0)?,
|
|
row.get(1)?,
|
|
row.get(2)?,
|
|
row.get(3)?,
|
|
row.get(4)?,
|
|
))
|
|
})?
|
|
.collect::<rusqlite::Result<_>>()?;
|
|
drop(stmt);
|
|
if rows.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
// Stamp all popped rows in a single UPDATE — under the broker
|
|
// mutex, well within sqlite's 999-param default.
|
|
let now = Utc::now().timestamp();
|
|
let ids: Vec<i64> = rows.iter().map(|(id, _, _, _, _)| *id).collect();
|
|
let placeholders = std::iter::repeat_n("?", ids.len())
|
|
.collect::<Vec<_>>()
|
|
.join(",");
|
|
let sql = format!("UPDATE messages SET delivered_at = ? WHERE id IN ({placeholders})");
|
|
let mut params_vec: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(ids.len() + 1);
|
|
params_vec.push(&now);
|
|
for id in &ids {
|
|
params_vec.push(id);
|
|
}
|
|
conn.execute(&sql, params_vec.as_slice())?;
|
|
drop(conn);
|
|
// Bookkeeping + assemble the Delivery list. Per-row
|
|
// `requeued_ids` lookup runs once per pop, same as `recv`.
|
|
let slot = inflight.entry(recipient.to_owned()).or_default();
|
|
let mut deliveries = Vec::with_capacity(rows.len());
|
|
for (id, from, to, body, in_reply_to) in rows {
|
|
slot.unacked_ids.push(id);
|
|
let redelivered = slot.requeued_ids.remove(&id);
|
|
deliveries.push(Delivery {
|
|
id,
|
|
redelivered,
|
|
message: Message {
|
|
from: hive_sh4re::manager::trusted_sender(&from),
|
|
to,
|
|
body,
|
|
in_reply_to,
|
|
},
|
|
});
|
|
}
|
|
drop(inflight);
|
|
// Mirror the per-row Delivered emit `recv` does so the
|
|
// dashboard forwarder sees one event per message regardless of
|
|
// which surface the harness used.
|
|
for d in &deliveries {
|
|
let _ = self.events.send(MessageEvent::Delivered {
|
|
id: d.id,
|
|
from: d.message.from.to_string(),
|
|
to: d.message.to.clone(),
|
|
body: d.message.body.clone(),
|
|
at: now,
|
|
in_reply_to: d.message.in_reply_to,
|
|
});
|
|
}
|
|
Ok(deliveries)
|
|
}
|
|
|
|
/// Drain the per-recipient unacked-id list and mark every row
|
|
/// `acked_at = NOW`. Fired by the harness after `TurnOutcome::Ok`.
|
|
/// Returns the number of rows acked (zero is normal — claude
|
|
/// may have not called recv during the turn). Tolerant of ids
|
|
/// that no longer exist in the DB (vacuumed, manually deleted)
|
|
/// — `UPDATE … WHERE id IN (…)` simply matches zero rows.
|
|
pub fn ack_turn(&self, recipient: &str) -> Result<u64> {
|
|
// Same lock order as `recv` and `requeue_inflight`.
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
let ids: Vec<i64> = inflight
|
|
.get_mut(recipient)
|
|
.map(|s| std::mem::take(&mut s.unacked_ids))
|
|
.unwrap_or_default();
|
|
if ids.is_empty() {
|
|
return Ok(0);
|
|
}
|
|
let now = Utc::now().timestamp();
|
|
let conn = self.conn.lock().unwrap();
|
|
// Bind every id explicitly. Caps in the hundreds in the worst
|
|
// case (a single very chatty turn); well under sqlite's 999
|
|
// default param limit and we're already serialising on the
|
|
// broker mutex.
|
|
let placeholders = std::iter::repeat_n("?", ids.len())
|
|
.collect::<Vec<_>>()
|
|
.join(",");
|
|
let sql = format!("UPDATE messages SET acked_at = ? WHERE id IN ({placeholders})");
|
|
let mut params_vec: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(ids.len() + 1);
|
|
params_vec.push(&now);
|
|
for id in &ids {
|
|
params_vec.push(id);
|
|
}
|
|
let n = conn.execute(&sql, params_vec.as_slice())?;
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
|
|
/// Bulk-ack every message addressed to `recipient` with row id
|
|
/// `<= up_to` that isn't acked yet — pending AND delivered rows
|
|
/// alike. The agent-facing triage tool behind `AckUntil`: after a
|
|
/// redelivered flood the agent acks everything up to the highest id
|
|
/// it has seen instead of re-popping each row. Recipient-scoped by
|
|
/// the WHERE clause, so an agent can never touch another agent's
|
|
/// rows. Also drains the in-memory `unacked_ids` / `requeued_ids`
|
|
/// bookkeeping below the cutoff so a later `ack_turn` doesn't
|
|
/// re-update rows this call already closed and a stale redelivery
|
|
/// tag doesn't outlive its row. Returns the number of rows newly
|
|
/// acked.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Propagates sqlite errors from the `UPDATE`.
|
|
pub fn ack_until(&self, recipient: &str, up_to: i64) -> Result<u64> {
|
|
// Same lock order as `recv` / `ack_turn` / `requeue_inflight`:
|
|
// `inflight` FIRST, then `conn`.
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
if let Some(state) = inflight.get_mut(recipient) {
|
|
state.unacked_ids.retain(|&id| id > up_to);
|
|
state.requeued_ids.retain(|&id| id > up_to);
|
|
}
|
|
let conn = self.conn.lock().unwrap();
|
|
let n = conn.execute(
|
|
"UPDATE messages SET acked_at = ?1
|
|
WHERE recipient = ?2 AND id <= ?3 AND acked_at IS NULL",
|
|
params![Utc::now().timestamp(), recipient, up_to],
|
|
)?;
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
|
|
/// Resurface every message the broker previously handed to this
|
|
/// recipient that never got `acked_at` set. Used by the harness at
|
|
/// boot to recover from the crashed-mid-turn / OOM-killed /
|
|
/// container-restarted cases. Three steps:
|
|
///
|
|
/// 1. Clear any stale in-memory state for this recipient (the
|
|
/// previous harness session's `unacked_ids` are irrelevant —
|
|
/// the new session will repopulate from fresh pops).
|
|
/// 2. Find every row where `recipient = me`, `delivered_at IS NOT
|
|
/// NULL`, `acked_at IS NULL`. Reset `delivered_at = NULL` so
|
|
/// the next `Recv` pops them again.
|
|
/// 3. Remember each id in the per-recipient `requeued_ids` set so
|
|
/// the next pop tags the response with `redelivered: true`.
|
|
///
|
|
/// Returns the number of rows requeued. Safe to call when there's
|
|
/// nothing in flight (returns 0). Safe to call multiple times
|
|
/// (idempotent — the second call finds nothing because the rows
|
|
/// are now back in the pending state).
|
|
pub fn requeue_inflight(&self, recipient: &str) -> Result<u64> {
|
|
// Hold inflight + conn together so a concurrent `recv` can't
|
|
// pop a just-requeued row between our DB update and our
|
|
// in-memory populate and miss the redelivered tag.
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
let conn = self.conn.lock().unwrap();
|
|
let mut stmt = conn.prepare(
|
|
"SELECT id FROM messages
|
|
WHERE recipient = ?1
|
|
AND delivered_at IS NOT NULL
|
|
AND acked_at IS NULL",
|
|
)?;
|
|
let ids: Vec<i64> = stmt
|
|
.query_map(params![recipient], |row| row.get(0))?
|
|
.collect::<rusqlite::Result<_>>()?;
|
|
drop(stmt);
|
|
if !ids.is_empty() {
|
|
let placeholders = std::iter::repeat_n("?", ids.len())
|
|
.collect::<Vec<_>>()
|
|
.join(",");
|
|
let sql =
|
|
format!("UPDATE messages SET delivered_at = NULL WHERE id IN ({placeholders})");
|
|
let params_vec: Vec<&dyn rusqlite::ToSql> =
|
|
ids.iter().map(|id| id as &dyn rusqlite::ToSql).collect();
|
|
conn.execute(&sql, params_vec.as_slice())?;
|
|
}
|
|
let slot = inflight.entry(recipient.to_owned()).or_default();
|
|
slot.unacked_ids.clear();
|
|
slot.requeued_ids.clear();
|
|
slot.requeued_ids.extend(ids.iter().copied());
|
|
Ok(u64::try_from(ids.len()).unwrap_or(0))
|
|
}
|
|
|
|
/// Operator-driven "clear the inbox": mark every message addressed
|
|
/// to `recipient` as acked. Backfills `delivered_at = NOW` for any
|
|
/// row that was still pending (undelivered), so the row doesn't
|
|
/// become impossible to vacuum later — `vacuum_delivered` requires
|
|
/// both timestamps set. Also clears the in-memory inflight state
|
|
/// for the recipient so a subsequent `ack_turn` doesn't try to
|
|
/// re-mark ids that are already acked. Returns the number of rows
|
|
/// affected (zero is normal — inbox already empty).
|
|
///
|
|
/// Distinct from `ack_turn` (which acks only the per-turn unacked
|
|
/// ids the harness pulled via `recv_batch`) and `requeue_inflight`
|
|
/// (which puts inflight-but-unacked rows BACK on the queue). This
|
|
/// is the destructive "drain everything for this agent" path the
|
|
/// dashboard surfaces as the side-panel "mark all read" button.
|
|
/// Backs `POST /api/agent/{name}/mark-all-read`.
|
|
pub fn mark_all_read(&self, recipient: &str) -> Result<u64> {
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
let conn = self.conn.lock().unwrap();
|
|
let now = Utc::now().timestamp();
|
|
// Two-axis update in one statement: set acked_at on every
|
|
// row for the recipient that doesn't have it yet, AND backfill
|
|
// delivered_at if it was NULL so the row is fully consumed
|
|
// (vacuum_delivered's `acked_at IS NOT NULL AND acked_at < ?`
|
|
// predicate then collects it on the normal hourly sweep).
|
|
let n = conn.execute(
|
|
"UPDATE messages
|
|
SET delivered_at = COALESCE(delivered_at, ?1),
|
|
acked_at = ?1
|
|
WHERE recipient = ?2
|
|
AND acked_at IS NULL",
|
|
params![now, recipient],
|
|
)?;
|
|
// Drop in-memory inflight bookkeeping for this recipient: the
|
|
// ids we just acked might still be in `unacked_ids` from a
|
|
// prior `recv_batch`; leaving them would cause the next
|
|
// `ack_turn` to re-issue an UPDATE against rows that no longer
|
|
// need it (correct but wasteful) and the requeue path would
|
|
// see stale ids. Cleanest to reset.
|
|
if let Some(slot) = inflight.get_mut(recipient) {
|
|
slot.unacked_ids.clear();
|
|
slot.requeued_ids.clear();
|
|
}
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
|
|
/// Per-process counter so each test gets a unique sqlite path even
|
|
/// when threads run concurrently. Avoids pulling in a `tempfile`
|
|
/// dep just for this one module.
|
|
static TEST_COUNTER: AtomicU64 = AtomicU64::new(0);
|
|
|
|
struct TmpBroker {
|
|
path: std::path::PathBuf,
|
|
pub broker: Broker,
|
|
}
|
|
|
|
impl Drop for TmpBroker {
|
|
fn drop(&mut self) {
|
|
let _ = std::fs::remove_file(&self.path);
|
|
}
|
|
}
|
|
|
|
fn open_broker() -> TmpBroker {
|
|
let n = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
|
|
let pid = std::process::id();
|
|
let path = std::env::temp_dir().join(format!("hive-broker-test-{pid}-{n}.sqlite"));
|
|
let _ = std::fs::remove_file(&path);
|
|
let broker = Broker::open(&path).expect("open broker");
|
|
TmpBroker { path, broker }
|
|
}
|
|
|
|
fn msg(from: &str, to: &str, body: &str) -> Message {
|
|
Message {
|
|
from: hive_types::Ident::parse(from).expect("test sender must be a valid ident"),
|
|
to: to.to_owned(),
|
|
body: body.to_owned(),
|
|
in_reply_to: None,
|
|
}
|
|
}
|
|
|
|
/// Convenience wrapper for tests that want single-pop semantics
|
|
/// — the broker only exposes `recv_batch` publicly, so tests
|
|
/// funnel single pops through here.
|
|
fn pop_one(broker: &Broker, recipient: &str) -> Option<Delivery> {
|
|
let mut batch = broker.recv_batch(recipient, 1).unwrap();
|
|
batch.pop()
|
|
}
|
|
|
|
#[test]
|
|
fn kv_set_get_delete_roundtrip() {
|
|
let tb = open_broker();
|
|
let b = &tb.broker;
|
|
// Absent key -> None.
|
|
assert_eq!(b.kv_get("k").unwrap(), None);
|
|
// Set then get.
|
|
b.kv_set("k", "v1").unwrap();
|
|
assert_eq!(b.kv_get("k").unwrap(), Some("v1".to_owned()));
|
|
// Upsert overwrites.
|
|
b.kv_set("k", "v2").unwrap();
|
|
assert_eq!(b.kv_get("k").unwrap(), Some("v2".to_owned()));
|
|
// Delete clears; deleting again is a no-op.
|
|
b.kv_delete("k").unwrap();
|
|
assert_eq!(b.kv_get("k").unwrap(), None);
|
|
b.kv_delete("k").unwrap();
|
|
}
|
|
|
|
/// Happy path: send → recv → `ack_turn` drains the in-memory list
|
|
/// and marks the row `acked_at IS NOT NULL`. A second recv finds
|
|
/// nothing pending (the row stays in the table for vacuum).
|
|
#[test]
|
|
fn ack_turn_marks_delivered_rows_acked() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "hi")).unwrap();
|
|
let d = pop_one(broker, "b").expect("popped");
|
|
assert_eq!(d.message.body, "hi");
|
|
assert!(!d.redelivered);
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 1);
|
|
// ack_turn drained the unacked list; calling again is a no-op.
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 0);
|
|
// Recv finds nothing — the row is now delivered + acked.
|
|
assert!(pop_one(broker, "b").is_none());
|
|
}
|
|
|
|
/// Bulk triage: three queued messages, agent acks up to the second
|
|
/// id — the first two never deliver again, the third still pops.
|
|
#[test]
|
|
fn ack_until_acks_only_rows_at_or_below_cutoff() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "one")).unwrap();
|
|
broker.send(&msg("a", "b", "two")).unwrap();
|
|
broker.send(&msg("a", "b", "three")).unwrap();
|
|
// Pop the first two so we know their ids (FIFO).
|
|
let d1 = pop_one(broker, "b").expect("popped one");
|
|
let d2 = pop_one(broker, "b").expect("popped two");
|
|
assert_eq!(broker.ack_until("b", d2.id).unwrap(), 2);
|
|
// The cutoff also drained the in-memory unacked list, so a
|
|
// turn-level ack right after finds nothing left to do.
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 0);
|
|
// A restart-style requeue finds nothing below the cutoff …
|
|
assert_eq!(broker.requeue_inflight("b").unwrap(), 0);
|
|
// … and the third message (id above the cutoff) still pops.
|
|
let d3 = pop_one(broker, "b").expect("third still pending");
|
|
assert_eq!(d3.message.body, "three");
|
|
assert!(d1.id < d2.id && d2.id < d3.id);
|
|
}
|
|
|
|
/// Pending (never-delivered) rows below the cutoff are acked too —
|
|
/// that's the whole point for a stale backlog the agent never
|
|
/// popped individually.
|
|
#[test]
|
|
fn ack_until_covers_pending_rows() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "stale-1")).unwrap();
|
|
broker.send(&msg("a", "b", "stale-2")).unwrap();
|
|
// Learn the highest id by peeking via recent_for (non-mutating).
|
|
let rows = broker.recent_for("b", 10).unwrap();
|
|
let max_id = rows.iter().map(|r| r.id).max().expect("rows");
|
|
assert_eq!(broker.ack_until("b", max_id).unwrap(), 2);
|
|
assert!(pop_one(broker, "b").is_none(), "backlog cleared");
|
|
}
|
|
|
|
/// Recipient scoping: acking b's inbox never touches c's rows,
|
|
/// even when c's ids fall below the cutoff.
|
|
#[test]
|
|
fn ack_until_is_recipient_scoped() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "c", "for-c")).unwrap();
|
|
broker.send(&msg("a", "b", "for-b")).unwrap();
|
|
let rows = broker.recent_for("b", 10).unwrap();
|
|
let max_id = rows.iter().map(|r| r.id).max().expect("rows");
|
|
assert_eq!(broker.ack_until("b", max_id).unwrap(), 1);
|
|
let d = pop_one(broker, "c").expect("c's message untouched");
|
|
assert_eq!(d.message.body, "for-c");
|
|
}
|
|
|
|
/// Crash-recovery: send → recv → (no ack) → `requeue_inflight`
|
|
/// resets `delivered_at` + tags the next pop as redelivered. After
|
|
/// that `ack_turn` closes it out cleanly.
|
|
#[test]
|
|
fn requeue_inflight_resurfaces_unacked_with_redelivered_flag() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "hi")).unwrap();
|
|
let d1 = pop_one(broker, "b").expect("popped");
|
|
assert!(!d1.redelivered);
|
|
// Simulate harness crash: never call ack_turn. Now boot the
|
|
// new harness — requeue_inflight resurfaces the row.
|
|
assert_eq!(broker.requeue_inflight("b").unwrap(), 1);
|
|
let d2 = pop_one(broker, "b").expect("popped again");
|
|
assert_eq!(d2.message.body, "hi");
|
|
assert!(d2.redelivered, "second pop should be tagged redelivered");
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 1);
|
|
}
|
|
|
|
/// Idempotency: a second `requeue_inflight` on the same recipient
|
|
/// finds nothing because the prior call already reset
|
|
/// `delivered_at` (the row is back in the pending state, not
|
|
/// inflight).
|
|
#[test]
|
|
fn requeue_inflight_is_idempotent() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "hi")).unwrap();
|
|
pop_one(broker, "b").expect("popped");
|
|
assert_eq!(broker.requeue_inflight("b").unwrap(), 1);
|
|
// Second call: the row is pending (delivered_at IS NULL) so
|
|
// nothing matches the inflight filter.
|
|
assert_eq!(broker.requeue_inflight("b").unwrap(), 0);
|
|
}
|
|
|
|
/// Multiple messages, partial drain: pop two, `ack_turn` covers
|
|
/// both even though one was popped before the other.
|
|
#[test]
|
|
fn ack_turn_handles_batch() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "one")).unwrap();
|
|
broker.send(&msg("a", "b", "two")).unwrap();
|
|
broker.send(&msg("a", "b", "three")).unwrap();
|
|
pop_one(broker, "b").expect("popped 1");
|
|
pop_one(broker, "b").expect("popped 2");
|
|
pop_one(broker, "b").expect("popped 3");
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 3);
|
|
assert!(pop_one(broker, "b").is_none());
|
|
}
|
|
|
|
/// Vacuum filter respects the new `acked_at` semantics — a
|
|
/// delivered-but-not-acked row is NOT vacuumed regardless of
|
|
/// age (the requeue path needs it).
|
|
#[test]
|
|
fn vacuum_preserves_unacked_inflight_rows() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "stuck")).unwrap();
|
|
pop_one(broker, "b").expect("popped");
|
|
// Wide window — should still skip unacked rows.
|
|
let removed = broker.vacuum_delivered(-i64::from(u8::MAX)).unwrap();
|
|
assert_eq!(removed, 0, "unacked inflight row must survive vacuum");
|
|
// After ack_turn the row is fair game.
|
|
broker.ack_turn("b").unwrap();
|
|
let removed = broker.vacuum_delivered(-i64::from(u8::MAX)).unwrap();
|
|
assert_eq!(removed, 1, "acked row is now vacuumable");
|
|
}
|
|
|
|
/// Recv ordering: requeued rows go back into FIFO position
|
|
/// (they keep their original id). New sends added after the
|
|
/// requeue arrive after them.
|
|
#[test]
|
|
fn requeued_rows_come_back_in_original_order() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "first")).unwrap();
|
|
broker.send(&msg("a", "b", "second")).unwrap();
|
|
// Pop both, ack neither.
|
|
pop_one(broker, "b").expect("popped 1");
|
|
pop_one(broker, "b").expect("popped 2");
|
|
broker.requeue_inflight("b").unwrap();
|
|
// Now add a brand new message AFTER the requeue.
|
|
broker.send(&msg("a", "b", "third")).unwrap();
|
|
let d1 = pop_one(broker, "b").expect("re-pop 1");
|
|
assert_eq!(d1.message.body, "first");
|
|
assert!(d1.redelivered);
|
|
let d2 = pop_one(broker, "b").expect("re-pop 2");
|
|
assert_eq!(d2.message.body, "second");
|
|
assert!(d2.redelivered);
|
|
let d3 = pop_one(broker, "b").expect("re-pop 3");
|
|
assert_eq!(d3.message.body, "third");
|
|
assert!(
|
|
!d3.redelivered,
|
|
"fresh-send-after-requeue must NOT be tagged redelivered"
|
|
);
|
|
}
|
|
|
|
/// Happy path for `recv_batch`: pops in FIFO order, respects
|
|
/// `max`, leaves the rest pending for the next call.
|
|
#[test]
|
|
fn recv_batch_pops_fifo_capped_at_max() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
for i in 0..5 {
|
|
broker.send(&msg("a", "b", &format!("m{i}"))).unwrap();
|
|
}
|
|
let batch = broker.recv_batch("b", 3).unwrap();
|
|
let bodies: Vec<_> = batch.iter().map(|d| d.message.body.as_str()).collect();
|
|
assert_eq!(bodies, vec!["m0", "m1", "m2"]);
|
|
// Remaining two stay pending; a second batch drains them.
|
|
let next = broker.recv_batch("b", 10).unwrap();
|
|
let bodies: Vec<_> = next.iter().map(|d| d.message.body.as_str()).collect();
|
|
assert_eq!(bodies, vec!["m3", "m4"]);
|
|
// ack_turn closes out all five popped rows in one go.
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 5);
|
|
}
|
|
|
|
/// Wake messages go through sqlite and get real row ids that the
|
|
/// agent can pass to `ack_until`. A wake sent when no recv is parked
|
|
/// is not lost — it sits in sqlite until the next recv pops it.
|
|
#[test]
|
|
fn wake_persisted_with_real_id_and_ackable() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
// Two wakes arrive (simulating bash-task completions).
|
|
broker
|
|
.send(&msg("bash-task-1", "b", "bash task `1` finished: exit=0"))
|
|
.unwrap();
|
|
broker
|
|
.send(&msg("bash-task-2", "b", "bash task `2` finished: exit=0"))
|
|
.unwrap();
|
|
let batch = broker.recv_batch("b", 10).unwrap();
|
|
assert_eq!(batch.len(), 2);
|
|
// Every delivery has a real (non-zero) id.
|
|
assert!(batch[0].id > 0);
|
|
assert!(batch[1].id > 0);
|
|
// Agent can ack them explicitly via ack_until.
|
|
let max_id = batch.iter().map(|d| d.id).max().unwrap();
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 2);
|
|
// Confirming the rows are now closed.
|
|
assert!(pop_one(broker, "b").is_none());
|
|
// ack_until on already-acked ids is a no-op.
|
|
assert_eq!(broker.ack_until("b", max_id).unwrap(), 0);
|
|
}
|
|
|
|
/// Wakes and regular mail both go through sqlite: FIFO by insertion order.
|
|
#[test]
|
|
fn wake_and_mail_ordered_fifo() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "mail")).unwrap();
|
|
broker.send(&msg("matrix", "b", "matrix wake")).unwrap();
|
|
let batch = broker.recv_batch("b", 10).unwrap();
|
|
assert_eq!(batch.len(), 2);
|
|
// FIFO within same priority: mail was inserted first, comes out first.
|
|
assert_eq!(batch[0].message.body, "mail");
|
|
assert_eq!(batch[1].message.body, "matrix wake");
|
|
assert!(batch[0].id > 0);
|
|
assert!(batch[1].id > 0);
|
|
}
|
|
|
|
/// Operator messages (priority = 1) surface before queued bash
|
|
/// completions and forge events (priority = 0) even when inserted
|
|
/// later, so an interrupt arrives before the task wake drives a turn.
|
|
#[test]
|
|
fn operator_message_jumps_queue_over_lower_priority() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
// Low-priority wake arrives first.
|
|
broker
|
|
.send(&msg("bash-task-123", "agent", "task done"))
|
|
.unwrap();
|
|
broker.send(&msg("forge", "agent", "new pr")).unwrap();
|
|
// Operator interrupt arrives after both.
|
|
broker
|
|
.send(&msg("operator", "agent", "stop what you're doing"))
|
|
.unwrap();
|
|
let batch = broker.recv_batch("agent", 10).unwrap();
|
|
assert_eq!(batch.len(), 3);
|
|
// Operator message surfaces first despite arriving last.
|
|
assert_eq!(batch[0].message.body, "stop what you're doing");
|
|
assert_eq!(batch[0].message.from.as_str(), "operator");
|
|
// Remaining two in FIFO order.
|
|
assert_eq!(batch[1].message.body, "task done");
|
|
assert_eq!(batch[2].message.body, "new pr");
|
|
}
|
|
|
|
/// `recv_batch` with no pending traffic returns an empty vec
|
|
/// (the "(empty)" path), not an error.
|
|
#[test]
|
|
fn recv_batch_returns_empty_when_idle() {
|
|
let h = open_broker();
|
|
let batch = h.broker.recv_batch("ghost", 5).unwrap();
|
|
assert!(batch.is_empty());
|
|
}
|
|
|
|
/// `max = 0` short-circuits without touching the DB (covered by
|
|
/// asserting we don't accidentally pop a pending row).
|
|
#[test]
|
|
fn recv_batch_zero_max_pops_nothing() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "stay")).unwrap();
|
|
assert!(broker.recv_batch("b", 0).unwrap().is_empty());
|
|
// The pending row is still in flight for the next real recv.
|
|
let d = pop_one(broker, "b").expect("still pending");
|
|
assert_eq!(d.message.body, "stay");
|
|
}
|
|
|
|
/// `recv_batch` tags requeued rows with `redelivered: true` and
|
|
/// drains them from the per-recipient `requeued_ids` set so a
|
|
/// fresh follow-up recv after the batch doesn't double-tag.
|
|
#[test]
|
|
fn recv_batch_propagates_redelivered_flag() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "one")).unwrap();
|
|
broker.send(&msg("a", "b", "two")).unwrap();
|
|
pop_one(broker, "b").expect("popped 1");
|
|
pop_one(broker, "b").expect("popped 2");
|
|
broker.requeue_inflight("b").unwrap();
|
|
let batch = broker.recv_batch("b", 5).unwrap();
|
|
assert_eq!(batch.len(), 2);
|
|
assert!(batch.iter().all(|d| d.redelivered));
|
|
// Fresh send after the batch is NOT tagged redelivered.
|
|
broker.send(&msg("a", "b", "three")).unwrap();
|
|
let d = pop_one(broker, "b").expect("re-pop 3");
|
|
assert_eq!(d.message.body, "three");
|
|
assert!(!d.redelivered);
|
|
}
|
|
|
|
/// `mark_all_read` covers a mix of pending + delivered + acked rows:
|
|
/// pending rows get both `delivered_at` and `acked_at` backfilled,
|
|
/// delivered-but-unacked rows just get `acked_at` set, already-acked
|
|
/// rows pass through untouched. Returns the count of rows mutated.
|
|
#[test]
|
|
fn mark_all_read_drains_all_states_for_recipient() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
// Set up three rows in three different states:
|
|
// r1 — pending (never popped, both timestamps NULL)
|
|
// r2 — delivered + unacked (popped, harness didn't ack yet)
|
|
// r3 — delivered + acked (popped + ack_turn ran)
|
|
broker.send(&msg("a", "b", "pending")).unwrap();
|
|
broker.send(&msg("a", "b", "delivered")).unwrap();
|
|
broker.send(&msg("a", "b", "acked")).unwrap();
|
|
// Pop both deliverable rows, then ack only the last so r2 stays
|
|
// delivered-but-unacked. After pop r3 is still in the unacked
|
|
// list; recv pops in FIFO order so first pop = "pending", but
|
|
// we want THAT row to remain undelivered. Workaround: pop two
|
|
// rows (so "pending" and "delivered" come off the queue) and
|
|
// requeue the first to put "pending" back. Then send a fourth
|
|
// "acked" and pop+ack just that.
|
|
//
|
|
// Simpler approach: bypass the queue helpers and craft the row
|
|
// states directly via send + recv + ack. FIFO order is by
|
|
// insertion; we pop two, ack only the second.
|
|
let _ = pop_one(broker, "b").expect("pop 1: pending → now delivered");
|
|
let _ = pop_one(broker, "b").expect("pop 2: delivered");
|
|
let _ = pop_one(broker, "b").expect("pop 3: acked-soon");
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 3);
|
|
// Now reshape: requeue first two so they're pending again.
|
|
// (Hack — easier: just call mark_all_read on the state we
|
|
// have, which is "three rows already acked". Should return
|
|
// zero because no row has acked_at IS NULL.)
|
|
assert_eq!(broker.mark_all_read("b").unwrap(), 0);
|
|
// Add a fresh pending row + a delivered-but-unacked row.
|
|
broker.send(&msg("a", "b", "new pending")).unwrap();
|
|
broker.send(&msg("a", "b", "new delivered")).unwrap();
|
|
let _ = pop_one(broker, "b").expect("pop new pending → delivered");
|
|
// Don't ack — leaves it delivered+unacked.
|
|
// Now: one row is pending (delivered_at IS NULL), one is
|
|
// delivered+unacked. mark_all_read should hit both.
|
|
assert_eq!(broker.mark_all_read("b").unwrap(), 2);
|
|
// Second call: nothing pending now.
|
|
assert_eq!(broker.mark_all_read("b").unwrap(), 0);
|
|
// Confirm via recv — inbox is empty.
|
|
assert!(pop_one(broker, "b").is_none());
|
|
}
|
|
|
|
/// Per-recipient isolation: marking alice doesn't touch bob's inbox.
|
|
#[test]
|
|
fn mark_all_read_is_per_recipient() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("x", "alice", "for alice")).unwrap();
|
|
broker.send(&msg("x", "bob", "for bob")).unwrap();
|
|
assert_eq!(broker.mark_all_read("alice").unwrap(), 1);
|
|
// bob's row still pending — pop succeeds.
|
|
let d = pop_one(broker, "bob").expect("bob pop");
|
|
assert_eq!(d.message.body, "for bob");
|
|
}
|
|
|
|
/// After `mark_all_read`, a subsequent `ack_turn` from a stale
|
|
/// in-memory unacked list MUST NOT panic or double-ack. The
|
|
/// inflight bookkeeping is cleared by `mark_all_read`.
|
|
#[test]
|
|
fn mark_all_read_clears_inflight_so_ack_turn_is_noop() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("a", "b", "hi")).unwrap();
|
|
let _ = pop_one(broker, "b").expect("popped");
|
|
// Now there's an id in the unacked_ids list. mark_all_read
|
|
// should clear it so the harness's next ack_turn (which still
|
|
// thinks the id is unacked) returns 0 cleanly.
|
|
assert_eq!(broker.mark_all_read("b").unwrap(), 1);
|
|
assert_eq!(broker.ack_turn("b").unwrap(), 0);
|
|
}
|
|
|
|
/// Per-recipient isolation: `requeue_inflight("a")` doesn't touch
|
|
/// b's inflight rows.
|
|
#[test]
|
|
fn requeue_inflight_is_per_recipient() {
|
|
let h = open_broker();
|
|
let broker = &h.broker;
|
|
broker.send(&msg("x", "alice", "for alice")).unwrap();
|
|
broker.send(&msg("x", "bob", "for bob")).unwrap();
|
|
pop_one(broker, "alice").expect("popped alice");
|
|
pop_one(broker, "bob").expect("popped bob");
|
|
// Requeue only alice. Bob's row stays inflight.
|
|
assert_eq!(broker.requeue_inflight("alice").unwrap(), 1);
|
|
let d = pop_one(broker, "alice").expect("re-pop alice");
|
|
assert!(d.redelivered);
|
|
// Bob has nothing pending (his row is still delivered, not requeued).
|
|
assert!(pop_one(broker, "bob").is_none());
|
|
}
|
|
}
|