use anyhow::Result;
use std::path::{Path, PathBuf};
pub const MAX_HOOK_PAYLOAD_CHARS: usize = 10_000;
#[derive(Debug, Clone, PartialEq)]
pub struct InboxMessage {
pub id: String,
pub text: String,
pub sent_at: i64,
}
pub fn write_message(dir: &Path, session_id: &str, text: &str) -> Result<InboxMessage> {
use std::sync::atomic::{AtomicU32, Ordering};
static SEQ: AtomicU32 = AtomicU32::new(0);
let sent_at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
let message = InboxMessage {
id: format!(
"msg-{sent_at}-{}-{:010}",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed),
),
text: text.to_string(),
sent_at,
};
let dir = inbox_dir(dir, session_id);
std::fs::create_dir_all(&dir)?;
let path = dir.join(format!("{}.json", message.id));
let body = serde_json::json!({
"id": message.id,
"text": message.text,
"sentAt": message.sent_at,
});
let tmp = path.with_extension("tmp");
std::fs::write(&tmp, serde_json::to_string_pretty(&body)?)?;
std::fs::rename(&tmp, &path)?;
Ok(message)
}
pub fn read_pending_messages(dir: &Path, session_id: &str) -> Result<Vec<InboxMessage>> {
let dir = inbox_dir(dir, session_id);
let entries = match std::fs::read_dir(&dir) {
Ok(e) => e,
Err(_) => return Ok(Vec::new()), };
let mut messages: Vec<InboxMessage> = entries
.filter_map(|e| e.ok())
.filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
.filter_map(|e| {
let raw = std::fs::read_to_string(e.path()).ok()?;
let v: serde_json::Value = serde_json::from_str(&raw).ok()?;
Some(InboxMessage {
id: v.get("id")?.as_str()?.to_string(),
text: v.get("text")?.as_str()?.to_string(),
sent_at: v.get("sentAt").and_then(|t| t.as_i64()).unwrap_or(0),
})
})
.collect();
messages.sort_by(|a, b| a.sent_at.cmp(&b.sent_at).then(a.id.cmp(&b.id)));
Ok(messages)
}
pub fn mark_messages_delivered(dir: &Path, session_id: &str, ids: &[String]) -> Result<()> {
let dir = inbox_dir(dir, session_id);
let mut first_err: Option<std::io::Error> = None;
for id in ids {
let path = dir.join(format!("{id}.json"));
if path.exists() {
let delivered = dir.join(format!("{id}.json.delivered"));
if let Err(e) = std::fs::rename(&path, &delivered) {
first_err.get_or_insert(e);
}
}
}
match first_err {
Some(e) => Err(e.into()),
None => Ok(()),
}
}
fn inbox_dir(dir: &Path, session_id: &str) -> PathBuf {
dir.join(format!("{session_id}.inbox"))
}
fn drain_batch(pending: Vec<InboxMessage>, max_chars: usize) -> (String, Vec<String>) {
let mut joined = String::new();
let mut ids = Vec::new();
for msg in pending {
let mut candidate = joined.clone();
if !candidate.is_empty() {
candidate.push_str("\n\n");
}
candidate.push_str(&msg.text);
if candidate.chars().count() > max_chars {
if ids.is_empty() {
joined = candidate.chars().take(max_chars).collect();
ids.push(msg.id);
}
break;
}
joined = candidate;
ids.push(msg.id);
}
(joined, ids)
}
pub fn drain_for_stop(dir: &Path, session_id: &str) -> Result<Option<serde_json::Value>> {
let pending = read_pending_messages(dir, session_id)?;
if pending.is_empty() {
return Ok(None);
}
let (reason, ids) = drain_batch(pending, MAX_HOOK_PAYLOAD_CHARS);
if let Err(e) = mark_messages_delivered(dir, session_id, &ids) {
tracing::warn!("failed to mark inbox messages delivered for {session_id}: {e}");
}
Ok(Some(serde_json::json!({ "decision": "block", "reason": reason })))
}
pub fn drain_for_prompt_submit(dir: &Path, session_id: &str) -> Result<Option<serde_json::Value>> {
let pending = read_pending_messages(dir, session_id)?;
if pending.is_empty() {
return Ok(None);
}
let (context, ids) = drain_batch(pending, MAX_HOOK_PAYLOAD_CHARS);
if let Err(e) = mark_messages_delivered(dir, session_id, &ids) {
tracing::warn!("failed to mark inbox messages delivered for {session_id}: {e}");
}
Ok(Some(serde_json::json!({
"hookSpecificOutput": {
"hookEventName": "UserPromptSubmit",
"additionalContext": context,
}
})))
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn write_then_read_returns_the_message() {
let dir = tempdir().unwrap();
let written = write_message(dir.path(), "sess-1", "hello worker").unwrap();
let pending = read_pending_messages(dir.path(), "sess-1").unwrap();
assert_eq!(pending, vec![written]);
}
#[test]
fn read_pending_on_missing_dir_returns_empty() {
let dir = tempdir().unwrap();
assert!(read_pending_messages(dir.path(), "no-such-session").unwrap().is_empty());
}
#[test]
fn messages_are_returned_oldest_first() {
let dir = tempdir().unwrap();
let a = write_message(dir.path(), "sess-1", "first").unwrap();
let b = write_message(dir.path(), "sess-1", "second").unwrap();
let pending = read_pending_messages(dir.path(), "sess-1").unwrap();
assert_eq!(pending, vec![a, b]);
}
#[test]
fn delivered_messages_are_no_longer_pending() {
let dir = tempdir().unwrap();
let msg = write_message(dir.path(), "sess-1", "hello").unwrap();
mark_messages_delivered(dir.path(), "sess-1", std::slice::from_ref(&msg.id)).unwrap();
assert!(read_pending_messages(dir.path(), "sess-1").unwrap().is_empty());
let dir_path = inbox_dir(dir.path(), "sess-1");
assert!(dir_path.join(format!("{}.json.delivered", msg.id)).exists());
}
#[test]
fn messages_for_different_sessions_do_not_collide() {
let dir = tempdir().unwrap();
write_message(dir.path(), "sess-1", "for one").unwrap();
write_message(dir.path(), "sess-2", "for two").unwrap();
assert_eq!(read_pending_messages(dir.path(), "sess-1").unwrap().len(), 1);
assert_eq!(read_pending_messages(dir.path(), "sess-2").unwrap().len(), 1);
}
#[test]
fn drain_for_stop_returns_none_when_nothing_pending() {
let dir = tempdir().unwrap();
assert!(drain_for_stop(dir.path(), "sess-1").unwrap().is_none());
}
#[test]
fn drain_for_stop_blocks_with_joined_reason_and_delivers() {
let dir = tempdir().unwrap();
write_message(dir.path(), "sess-1", "first").unwrap();
write_message(dir.path(), "sess-1", "second").unwrap();
let response = drain_for_stop(dir.path(), "sess-1").unwrap().unwrap();
assert_eq!(response["decision"], "block");
assert_eq!(response["reason"], "first\n\nsecond");
assert!(read_pending_messages(dir.path(), "sess-1").unwrap().is_empty());
}
#[test]
fn drain_for_stop_is_self_limiting_across_repeat_invocations() {
let dir = tempdir().unwrap();
write_message(dir.path(), "sess-1", "only message").unwrap();
assert!(drain_for_stop(dir.path(), "sess-1").unwrap().is_some());
assert!(drain_for_stop(dir.path(), "sess-1").unwrap().is_none());
}
#[test]
fn drain_for_prompt_submit_returns_none_when_nothing_pending() {
let dir = tempdir().unwrap();
assert!(drain_for_prompt_submit(dir.path(), "sess-1").unwrap().is_none());
}
#[test]
fn drain_for_prompt_submit_shapes_hook_specific_output_and_delivers() {
let dir = tempdir().unwrap();
write_message(dir.path(), "sess-1", "context for the next turn").unwrap();
let response = drain_for_prompt_submit(dir.path(), "sess-1").unwrap().unwrap();
assert_eq!(response["hookSpecificOutput"]["hookEventName"], "UserPromptSubmit");
assert_eq!(
response["hookSpecificOutput"]["additionalContext"],
"context for the next turn"
);
assert!(read_pending_messages(dir.path(), "sess-1").unwrap().is_empty());
}
#[test]
fn drain_for_prompt_submit_truncates_a_single_oversized_message_and_delivers_it() {
let dir = tempdir().unwrap();
let long = "x".repeat(MAX_HOOK_PAYLOAD_CHARS + 500);
write_message(dir.path(), "sess-1", &long).unwrap();
let response = drain_for_prompt_submit(dir.path(), "sess-1").unwrap().unwrap();
let context = response["hookSpecificOutput"]["additionalContext"].as_str().unwrap();
assert_eq!(context.chars().count(), MAX_HOOK_PAYLOAD_CHARS);
assert!(read_pending_messages(dir.path(), "sess-1").unwrap().is_empty());
}
#[test]
fn drain_for_prompt_submit_leaves_overflow_messages_pending_instead_of_discarding_them() {
let dir = tempdir().unwrap();
let big = "x".repeat(MAX_HOOK_PAYLOAD_CHARS);
write_message(dir.path(), "sess-1", &big).unwrap();
write_message(dir.path(), "sess-1", "second message").unwrap();
let response = drain_for_prompt_submit(dir.path(), "sess-1").unwrap().unwrap();
let context = response["hookSpecificOutput"]["additionalContext"].as_str().unwrap();
assert_eq!(context, big, "only the first message should be in this batch");
let still_pending = read_pending_messages(dir.path(), "sess-1").unwrap();
assert_eq!(still_pending.len(), 1, "the overflow message must stay pending, not be lost");
assert_eq!(still_pending[0].text, "second message");
}
#[test]
fn drain_for_stop_leaves_overflow_messages_pending_instead_of_discarding_them() {
let dir = tempdir().unwrap();
let big = "x".repeat(MAX_HOOK_PAYLOAD_CHARS);
write_message(dir.path(), "sess-1", &big).unwrap();
write_message(dir.path(), "sess-1", "second message").unwrap();
let response = drain_for_stop(dir.path(), "sess-1").unwrap().unwrap();
assert_eq!(response["reason"], big);
let still_pending = read_pending_messages(dir.path(), "sess-1").unwrap();
assert_eq!(still_pending.len(), 1, "the overflow message must stay pending, not be lost");
assert_eq!(still_pending[0].text, "second message");
}
#[test]
fn drain_batch_includes_everything_that_fits() {
let pending = vec![
InboxMessage { id: "a".into(), text: "first".into(), sent_at: 1 },
InboxMessage { id: "b".into(), text: "second".into(), sent_at: 2 },
];
let (joined, ids) = drain_batch(pending, 100);
assert_eq!(joined, "first\n\nsecond");
assert_eq!(ids, vec!["a".to_string(), "b".to_string()]);
}
#[test]
fn drain_batch_stops_before_a_message_that_would_overflow() {
let pending = vec![
InboxMessage { id: "a".into(), text: "12345".into(), sent_at: 1 },
InboxMessage { id: "b".into(), text: "67890".into(), sent_at: 2 },
];
let (joined, ids) = drain_batch(pending, 5);
assert_eq!(joined, "12345");
assert_eq!(ids, vec!["a".to_string()]);
}
#[test]
fn drain_batch_truncates_a_single_message_that_alone_exceeds_the_cap() {
let pending = vec![InboxMessage { id: "a".into(), text: "1234567890".into(), sent_at: 1 }];
let (joined, ids) = drain_batch(pending, 5);
assert_eq!(joined, "12345");
assert_eq!(ids, vec!["a".to_string()]);
}
}