diff --git a/crates/buzz-acp/README.md b/crates/buzz-acp/README.md index e6164b02dd..e5861e2548 100644 --- a/crates/buzz-acp/README.md +++ b/crates/buzz-acp/README.md @@ -129,6 +129,17 @@ All configuration is via environment variables (or CLI flags — every env var h | `--heartbeat-prompt` | `BUZZ_ACP_HEARTBEAT_PROMPT` | (built-in) | Custom heartbeat prompt text. Conflicts with `--heartbeat-prompt-file`. | | `--heartbeat-prompt-file` | `BUZZ_ACP_HEARTBEAT_PROMPT_FILE` | — | Read heartbeat prompt from a file. Conflicts with `--heartbeat-prompt`. | +### Session Persistence + +| Flag | Env Var | Default | Description | +|------|---------|---------|-------------| +| `--state-dir` | `BUZZ_ACP_STATE_DIR` | `/.buzz-acp` | Directory for harness state that must outlive the process — today the channel → ACP session map (`sessions-.json`). On restart the harness issues `session/load` for each remembered channel instead of opening a fresh session, so the agent keeps its context. | +| `--no-resume-sessions` | `BUZZ_ACP_NO_RESUME_SESSIONS` | `false` | Do not remember channel sessions across restarts: every restart opens a fresh `session/new` per channel, as before the session store existed. | + +A resume that cannot happen never costs the turn: if the file is missing, the entry is stale, the agent does not advertise `loadSession`, or the agent rejects `session/load`, the harness forgets the mapping and falls back to `session/new` for that channel. + +> **Operators running in containers:** the default state directory lives under the working directory, which is ephemeral in most container and pod filesystems — it silently vanishes on every deploy, so every restart starts every channel from scratch. Set `BUZZ_ACP_STATE_DIR` to a path on mounted persistent storage (a volume, PVC, or bind mount, e.g. `/var/lib/buzz-acp`) when the harness does not live on a laptop. Gateways that supervise hosted harnesses can additionally key sessions by `(agent pubkey, channelId)` in their own store — the harness sends the channel on `session/new` as `_meta.channelId` — to resume across a node move even when the file did not survive. + ### Inbound Author Gate Controls which authors' events the harness forwards to the agent. Events from disallowed authors are silently dropped before reaching subscription rules. @@ -256,6 +267,7 @@ Forum event kinds: 4. **Prompting** — When events are pending and no prompt is in flight for that channel, drains all queued events for the oldest channel into a single batched prompt via ACP `session/prompt`. 5. **Agent response** — The agent processes the prompt and uses the Buzz CLI (`send_message`, `get_messages`, etc.) to interact with Buzz. 6. **Recovery** — If the agent crashes, the harness respawns it. If the relay disconnects, the harness reconnects with a `since` filter to avoid missing events. +7. **Restart** — The channel → session map is written to the state dir (see [Session Persistence](#session-persistence)); a restarted harness `session/load`s each remembered session and falls back to `session/new` when it cannot. Each channel has at most one prompt in flight. Multiple channels can be processed concurrently when agents > 1. diff --git a/crates/buzz-acp/src/acp.rs b/crates/buzz-acp/src/acp.rs index f8373bd66d..8b1d82e216 100644 --- a/crates/buzz-acp/src/acp.rs +++ b/crates/buzz-acp/src/acp.rs @@ -138,6 +138,15 @@ fn build_initialize_params() -> serde_json::Value { /// /// One `AcpClient` per agent process. Multiple sessions can be created on the /// same client via repeated calls to [`session_new`](AcpClient::session_new). +/// The Buzz channel a `session/new` is being opened for. See +/// [`AcpClient::session_new_with_origin`]. +#[derive(Debug, Clone, Copy)] +pub struct SessionOrigin<'a> { + pub channel_id: uuid::Uuid, + /// `"dm"`, `"stream"`, … as the relay reports it; `None` when unresolved. + pub channel_type: Option<&'a str>, +} + pub struct AcpClient { /// The agent child process (kept alive to prevent zombie). child: Child, @@ -200,6 +209,10 @@ pub struct AcpClient { /// a JSON-RPC *success*, not `-32601` — which the main loop would read as /// a delivered steer and drop the user's message from the queue. steering_supported: bool, + /// Whether the agent advertised `agentCapabilities.loadSession` at + /// `initialize`. Gates [`session_load`](Self::session_load): a harness + /// restart resumes a channel's remembered session only when the agent can. + load_session_supported: bool, /// Per-turn channel for receiving goose-native non-cancelling steer /// requests from the main loop. Installed by /// [`install_steer_rx`](Self::install_steer_rx) at dispatch and @@ -559,6 +572,7 @@ impl AcpClient { observer_context: ObserverContext::default(), active_run_id: None, steering_supported: false, + load_session_supported: false, steer_rx: None, goose_usage: UsageTracker::default(), standard_usage: StandardUsageTracker::default(), @@ -617,10 +631,48 @@ impl AcpClient { .pointer("/_meta/steering/supported") .and_then(|v| v.as_bool()) .unwrap_or(false); + self.load_session_supported = result + .pointer("/agentCapabilities/loadSession") + .and_then(|v| v.as_bool()) + .unwrap_or(false); tracing::debug!(target: "acp::init", "initialize response: {result}"); Ok(result) } + /// Whether the agent advertised `loadSession` at `initialize`. + pub fn load_session_supported(&self) -> bool { + self.load_session_supported + } + + /// Send `session/load` for a session this harness created before a restart. + /// + /// ACP has the agent replay the conversation as `session/update` + /// notifications before it answers; the read loop dispatches those exactly + /// as it does mid-turn, so nothing is persisted or re-published here. On + /// `Ok` the session is live again under the same id; on `Err` the caller + /// falls back to `session/new` — a resume that fails must never cost the + /// turn. + /// + /// Only meaningful when [`load_session_supported`](Self::load_session_supported) + /// is true; sending it to an agent that did not advertise the capability + /// gets a method-not-found error, which the caller treats like any other + /// resume failure. + pub async fn session_load( + &mut self, + session_id: &str, + cwd: &str, + mcp_servers: Vec, + ) -> Result { + let params = serde_json::json!({ + "sessionId": session_id, + "cwd": cwd, + "mcpServers": mcp_servers, + }); + let result = self.send_request("session/load", params).await?; + tracing::info!(target: "acp::session", "session resumed: {session_id}"); + Ok(result) + } + /// Send the ACP `authenticate` request for an adapter-advertised method. pub async fn authenticate(&mut self, method_id: &str) -> Result { let params = serde_json::json!({ @@ -654,6 +706,28 @@ impl AcpClient { mcp_servers: Vec, system_prompt: Option>, session_title: Option<&str>, + ) -> Result { + self.session_new_with_origin(cwd, mcp_servers, system_prompt, session_title, None) + .await + } + + /// [`session_new_full`](Self::session_new_full) plus the channel the session + /// is for, sent out of band as `_meta.channelId` / `_meta.channelType`. + /// + /// The prompt already names the channel in its `[Context]` block, but that + /// is text for the model. An agent — or a gateway in front of many agents, + /// like `fountain acp` — that wants to key sessions by channel (to resume + /// the same conversation after this harness restarts, or to attach + /// per-channel resources) has had nothing machine-readable to key on; + /// `sessionTitle` carries the channel *name* and is empty for DMs. Absent + /// for heartbeat sessions, which have no channel. + pub async fn session_new_with_origin( + &mut self, + cwd: &str, + mcp_servers: Vec, + system_prompt: Option>, + session_title: Option<&str>, + origin: Option>, ) -> Result { let mut params = serde_json::json!({ "cwd": cwd, @@ -673,6 +747,12 @@ impl AcpClient { // Merge — _meta may already carry systemPrompt from ClaudeMeta above. params["_meta"]["sessionTitle"] = serde_json::Value::String(title.to_owned()); } + if let Some(origin) = origin { + params["_meta"]["channelId"] = serde_json::Value::String(origin.channel_id.to_string()); + if let Some(channel_type) = origin.channel_type { + params["_meta"]["channelType"] = serde_json::Value::String(channel_type.to_owned()); + } + } let result = self.send_request("session/new", params).await?; let session_id = result["sessionId"] .as_str() @@ -3532,6 +3612,150 @@ mod tests { ); } + #[tokio::test] + async fn session_new_with_origin_sends_channel_id_and_type_in_meta() { + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{}}}' + read -t 2 REQ + echo '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"ses_test","_receivedRequest":'"$REQ"'}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + client + .initialize() + .await + .expect("initialize should succeed"); + + let channel_id = uuid::Uuid::new_v4(); + let resp = client + .session_new_with_origin( + "/tmp", + vec![], + None, + Some("Fizz · #buzz-dev"), + Some(SessionOrigin { + channel_id, + channel_type: Some("dm"), + }), + ) + .await + .expect("session_new_with_origin should succeed"); + + let meta = &resp.raw["_receivedRequest"]["params"]["_meta"]; + assert_eq!( + meta["channelId"].as_str(), + Some(channel_id.to_string().as_str()) + ); + assert_eq!(meta["channelType"].as_str(), Some("dm")); + // The existing member is untouched: the origin merges into _meta. + assert_eq!(meta["sessionTitle"].as_str(), Some("Fizz · #buzz-dev")); + } + + #[tokio::test] + async fn session_new_full_sends_no_channel_origin() { + // Heartbeat sessions and the plain constructor: no channel, no member. + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{}}}' + read -t 2 REQ + echo '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"ses_test","_receivedRequest":'"$REQ"'}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + client + .initialize() + .await + .expect("initialize should succeed"); + let resp = client + .session_new_full("/tmp", vec![], None, Some("t")) + .await + .expect("session_new_full should succeed"); + let meta = &resp.raw["_receivedRequest"]["params"]["_meta"]; + assert!(meta.get("channelId").is_none()); + assert!(meta.get("channelType").is_none()); + } + + #[tokio::test] + async fn initialize_records_load_session_capability() { + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true}}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + assert!(!client.load_session_supported()); + client + .initialize() + .await + .expect("initialize should succeed"); + assert!(client.load_session_supported()); + } + + #[tokio::test] + async fn initialize_leaves_load_session_unsupported_when_absent() { + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{}}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + client + .initialize() + .await + .expect("initialize should succeed"); + assert!(!client.load_session_supported()); + } + + #[tokio::test] + async fn session_load_sends_session_id_cwd_and_servers_and_tolerates_replay() { + // ACP: the agent replays history as session/update notifications + // BEFORE answering session/load. The read loop must ride through them. + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true}}}' + read -t 2 REQ + echo '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"ses_old","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"replayed"}}}}' + echo '{"jsonrpc":"2.0","id":1,"result":{"_receivedRequest":'"$REQ"'}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + client + .initialize() + .await + .expect("initialize should succeed"); + + let result = client + .session_load("ses_old", "/work", vec![]) + .await + .expect("session_load should succeed"); + let received = &result["_receivedRequest"]; + assert_eq!(received["method"].as_str(), Some("session/load")); + assert_eq!(received["params"]["sessionId"].as_str(), Some("ses_old")); + assert_eq!(received["params"]["cwd"].as_str(), Some("/work")); + assert!(received["params"]["mcpServers"].is_array()); + } + + #[tokio::test] + async fn session_load_error_is_an_error_not_a_hang() { + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true}}}' + read -t 2 REQ + echo '{"jsonrpc":"2.0","id":1,"error":{"code":-32602,"message":"session not found"}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + client + .initialize() + .await + .expect("initialize should succeed"); + assert!(client + .session_load("ses_gone", "/work", vec![]) + .await + .is_err()); + } + #[tokio::test] async fn session_new_full_omits_meta_when_session_title_none() { let script = r#" diff --git a/crates/buzz-acp/src/config.rs b/crates/buzz-acp/src/config.rs index f9e7bf1ed8..e8886c2ae3 100644 --- a/crates/buzz-acp/src/config.rs +++ b/crates/buzz-acp/src/config.rs @@ -409,6 +409,20 @@ pub struct CliArgs { #[arg(long, env = "BUZZ_ACP_NO_BASE_PROMPT")] pub no_base_prompt: bool, + /// Directory for harness state that must outlive the process — today the + /// channel → ACP session map (`sessions-.json`) that lets a restart + /// `session/load` each channel's session instead of opening a new one. + /// Defaults to `.buzz-acp/` under the working directory. Point it at a + /// persistent location when the working directory is not (a hosted + /// harness in a container, for example). + #[arg(long, env = "BUZZ_ACP_STATE_DIR")] + pub state_dir: Option, + + /// Do not remember channel sessions across restarts: every restart opens a + /// fresh `session/new` per channel, as before the session store existed. + #[arg(long, env = "BUZZ_ACP_NO_RESUME_SESSIONS")] + pub no_resume_sessions: bool, + /// Path to a custom base prompt file. Overrides the compiled-in default. /// Mutually exclusive with --no-base-prompt. #[arg( @@ -538,6 +552,10 @@ pub struct Config { /// `[Agent Memory — core]` section. On by default; disabled via the /// `--no-memory` / `BUZZ_ACP_NO_MEMORY` opt-out. pub memory_enabled: bool, + /// Where the channel → session map lives across restarts. `None` disables + /// resumption (`--no-resume-sessions`); otherwise resolved from + /// `--state-dir` / `BUZZ_ACP_STATE_DIR`, defaulting to `/.buzz-acp`. + pub state_dir: Option, /// Desired LLM model ID. Applied after every `session_new_full()`. pub model: Option, /// Sanitized session title, sent as `_meta.sessionTitle` on `session/new`. @@ -1104,6 +1122,15 @@ impl Config { presence_enabled: !args.no_presence, typing_enabled: !args.no_typing, memory_enabled: args.memory && !args.no_memory, + state_dir: if args.no_resume_sessions { + None + } else { + Some(args.state_dir.clone().unwrap_or_else(|| { + std::env::current_dir() + .unwrap_or_else(|_| PathBuf::from(".")) + .join(".buzz-acp") + })) + }, model, session_title: args .session_title @@ -1479,6 +1506,7 @@ mod tests { presence_enabled: true, typing_enabled: true, memory_enabled: true, + state_dir: None, model: None, session_title: None, permission_mode: PermissionMode::BypassPermissions, @@ -2205,6 +2233,37 @@ channels = "ALL" assert_eq!(configured.exit_after_inactivity, 120); } + #[test] + fn state_dir_defaults_under_cwd_and_no_resume_disables_it() { + let key = nostr::Keys::generate().secret_key().to_secret_hex(); + let default = Config::from_args(CliArgs::parse_from(["buzz-acp", "--private-key", &key])) + .expect("default config should build"); + let dir = default.state_dir.expect("resumption is on by default"); + assert!(dir.ends_with(".buzz-acp"), "got {}", dir.display()); + + let explicit = Config::from_args(CliArgs::parse_from([ + "buzz-acp", + "--private-key", + &key, + "--state-dir", + "/var/lib/buzz-acp", + ])) + .expect("config should build"); + assert_eq!( + explicit.state_dir.as_deref(), + Some(std::path::Path::new("/var/lib/buzz-acp")) + ); + + let off = Config::from_args(CliArgs::parse_from([ + "buzz-acp", + "--private-key", + &key, + "--no-resume-sessions", + ])) + .expect("config should build"); + assert!(off.state_dir.is_none()); + } + #[test] fn lazy_pool_defaults_off() { let key = "0".repeat(64); diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index 2a41ea7342..223ef587b9 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -9,6 +9,7 @@ mod pool; mod pool_lifecycle; mod queue; mod relay; +mod session_store; mod setup_mode; mod usage; @@ -2207,8 +2208,21 @@ async fn tokio_main() -> Result<()> { memory_enabled: config.memory_enabled, harness_name: crate::config::normalize_agent_command_identity(&config.agent_command), relay_url: config.relay_url.clone(), + session_store: Arc::new(match &config.state_dir { + Some(dir) => { + session_store::SessionStore::open(dir.join(format!("sessions-{pubkey_hex}.json"))) + } + None => session_store::SessionStore::disabled(), + }), }); + if !ctx.session_store.is_enabled() { + tracing::info!( + target: "session_store", + "channel session resumption disabled (--no-resume-sessions / BUZZ_ACP_NO_RESUME_SESSIONS): every restart opens fresh sessions" + ); + } + if !config.memory_enabled { tracing::info!( target: "engram::core", @@ -2693,6 +2707,8 @@ async fn tokio_main() -> Result<()> { // complete normally (the relay may reject actions if // the agent lost access). let drained_ids = queue.drain_channel(ch); + // A channel we left must not be resumed after a restart. + ctx.session_store.remove(ch); let invalidated = if pool_ready { pool.invalidate_channel_sessions(ch) } else { @@ -2825,6 +2841,8 @@ async fn tokio_main() -> Result<()> { ); } else { let invalidated = pool.invalidate_channel_sessions(buzz_event.channel_id); + // The owner asked for a fresh session; do not bring the old one back. + ctx.session_store.remove(buzz_event.channel_id); tracing::info!( channel_id = %buzz_event.channel_id, invalidated, @@ -7138,6 +7156,7 @@ mod build_mcp_servers_tests { presence_enabled: true, typing_enabled: true, memory_enabled: false, + state_dir: None, model: None, session_title: None, permission_mode: config::PermissionMode::BypassPermissions, @@ -7361,6 +7380,7 @@ mod error_outcome_emission_tests { presence_enabled: true, typing_enabled: true, memory_enabled: false, + state_dir: None, model: None, session_title: None, permission_mode: config::PermissionMode::BypassPermissions, diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 2efacce2b1..15dbb55a11 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -21,6 +21,9 @@ use std::cmp::Reverse; use std::collections::{HashMap, HashSet}; + +use crate::acp::SessionOrigin; +use crate::session_store::SessionStore; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -606,6 +609,12 @@ pub struct PromptContext { /// the desktop keys per (agent, relay) pair, e.g. `session_config_captured`, /// mirroring the `managed_agent_runtime_lifecycle` frames. pub relay_url: String, + /// Channel → session map that outlives the process. Consulted before + /// `session/new` for a channel this agent has no live session for: a + /// remembered session is `session/load`ed instead (when the agent + /// advertises `loadSession`), so a restart continues each channel's + /// conversation rather than starting over. See `session_store`. + pub session_store: Arc, } impl AgentPool { @@ -952,6 +961,57 @@ async fn resolve_new_session_channel_context( (is_dm, title_channel, Some(info.channel_type)) } +/// Resume the session the store remembers for `channel_id`, if there is one and +/// the agent can. Returns the live session id on success. +/// +/// Every failure path returns `None` and forgets the remembered id, so the +/// caller opens a fresh session and the same dead id is not retried on every +/// mention: an agent that did not advertise `loadSession`, a `session/load` +/// error (the agent no longer has it — a sandbox reclaimed, a machine change), +/// or an agent that exited (surfaced separately by the create path). +async fn resume_remembered_session( + agent: &mut OwnedAgent, + ctx: &PromptContext, + channel_id: Uuid, +) -> Option { + let remembered = ctx.session_store.get(channel_id)?; + if !agent.acp.load_session_supported() { + tracing::info!( + target: "pool::session", + "channel {channel_id} remembers session {remembered} but the agent does not support session/load; starting fresh" + ); + ctx.session_store.remove(channel_id); + return None; + } + let mcp_servers = mcp_servers_with_git_origin( + &ctx.mcp_servers, + Some(channel_id), + None, + ctx.session_title.as_deref(), + ); + match agent + .acp + .session_load(&remembered, &ctx.cwd, mcp_servers) + .await + { + Ok(_) => { + tracing::info!( + target: "pool::session", + "resumed session {remembered} for channel {channel_id} after restart" + ); + Some(remembered) + } + Err(e) => { + tracing::warn!( + target: "pool::session", + "could not resume session {remembered} for channel {channel_id} ({e}); starting fresh" + ); + ctx.session_store.remove(channel_id); + None + } + } +} + /// Create a new ACP session via `session_new_full()`, populate model capabilities /// on the agent (first session only), and apply `desired_model` if set. /// @@ -1006,7 +1066,7 @@ async fn create_session_and_apply_model( let resp = agent .acp - .session_new_full( + .session_new_with_origin( &ctx.cwd, mcp_servers, session_new_system_prompt( @@ -1016,6 +1076,10 @@ async fn create_session_and_apply_model( combined_system_prompt.as_deref(), ), session_title.as_deref(), + channel.id.map(|channel_id| SessionOrigin { + channel_id, + channel_type: channel.channel_type, + }), ) .await?; @@ -1689,6 +1753,20 @@ pub async fn run_prompt_task( PromptSource::Channel(cid) => { if let Some(sid) = agent.state.sessions.get(cid) { (sid.clone(), false) + } else if let Some(sid) = resume_remembered_session(&mut agent, &ctx, *cid).await { + // A session this harness opened before it last restarted, live + // again under the same id. Delivery state starts fresh — it is + // per-process bookkeeping — and no usage baseline is seeded, + // because a resumed session's prior usage is not zero. + agent.state.sessions.insert(*cid, sid.clone()); + agent + .state + .deliveries + .insert(*cid, ChannelDeliveryState::default()); + if let Some((pending_cid, section)) = pending_canvas.take() { + agent.state.canvas_sections.insert(pending_cid, section); + } + (sid, false) } else { // The title is channel-qualified (`Agent · #channel`) so one // agent in several channels doesn't produce identical session @@ -1714,6 +1792,7 @@ pub async fn run_prompt_task( "created session {sid} for channel {cid}" ); agent.state.sessions.insert(*cid, sid.clone()); + ctx.session_store.put(*cid, &sid); agent .state .deliveries @@ -7600,6 +7679,7 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" memory_enabled: false, harness_name: "goose".to_string(), relay_url: "ws://127.0.0.1:3000".to_string(), + session_store: Arc::new(SessionStore::disabled()), } } diff --git a/crates/buzz-acp/src/session_store.rs b/crates/buzz-acp/src/session_store.rs new file mode 100644 index 0000000000..c968fc2aeb --- /dev/null +++ b/crates/buzz-acp/src/session_store.rs @@ -0,0 +1,188 @@ +//! Durable channel → ACP session map, so a harness restart resumes instead of +//! forgetting. +//! +//! `AgentPool` keeps `SessionState.sessions` in memory. That is correct while +//! the process lives, but every restart — a desktop relaunch, a config-change +//! restart, or a hosted harness being moved between nodes on a deploy — emptied +//! it, and the next mention in every channel got a fresh `session/new`. For an +//! agent whose session *is* its workspace (a sandbox per session, as with the +//! `fountain acp` gateway) that discards the channel's memory and files each +//! time. +//! +//! The store is a small JSON file: `{ "": "", ... }`. +//! Writes are atomic (temp file + rename). Every failure is logged and +//! swallowed — the store is an optimisation on top of the in-memory map, never +//! a reason to refuse a turn. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Mutex; + +use uuid::Uuid; + +#[derive(Debug)] +pub struct SessionStore { + /// `None` disables persistence: `get` is always empty, `put`/`remove` no-op. + path: Option, + map: Mutex>, +} + +impl SessionStore { + /// A store that remembers nothing across restarts. + pub fn disabled() -> Self { + Self { + path: None, + map: Mutex::new(HashMap::new()), + } + } + + /// Open (or create on first write) the store at `path`, loading whatever is + /// there. An unreadable or malformed file is treated as empty and logged. + pub fn open(path: impl Into) -> Self { + let path = path.into(); + let map = match std::fs::read(&path) { + Ok(bytes) => match serde_json::from_slice::>(&bytes) { + Ok(map) => map, + Err(e) => { + tracing::warn!( + target: "session_store", + "ignoring malformed session store {}: {e}", + path.display() + ); + HashMap::new() + } + }, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => HashMap::new(), + Err(e) => { + tracing::warn!( + target: "session_store", + "cannot read session store {}: {e}", + path.display() + ); + HashMap::new() + } + }; + tracing::info!( + target: "session_store", + "session store {}: {} channel session(s) remembered", + path.display(), + map.len() + ); + Self { + path: Some(path), + map: Mutex::new(map), + } + } + + pub fn is_enabled(&self) -> bool { + self.path.is_some() + } + + /// The session last recorded for `channel_id`, if any. + pub fn get(&self, channel_id: Uuid) -> Option { + self.map.lock().ok()?.get(&channel_id).cloned() + } + + /// Remember `session_id` for `channel_id` and flush. + pub fn put(&self, channel_id: Uuid, session_id: &str) { + if self.path.is_none() { + return; + } + if let Ok(mut map) = self.map.lock() { + map.insert(channel_id, session_id.to_owned()); + self.flush(&map); + } + } + + /// Forget `channel_id` (the session was invalidated, or a resume of it + /// failed) and flush. + pub fn remove(&self, channel_id: Uuid) { + if self.path.is_none() { + return; + } + if let Ok(mut map) = self.map.lock() { + if map.remove(&channel_id).is_some() { + self.flush(&map); + } + } + } + + fn flush(&self, map: &HashMap) { + let Some(path) = &self.path else { return }; + let write = || -> std::io::Result<()> { + if let Some(dir) = path.parent() { + std::fs::create_dir_all(dir)?; + } + let tmp = path.with_extension("json.tmp"); + let bytes = serde_json::to_vec_pretty(map).map_err(std::io::Error::other)?; + std::fs::write(&tmp, bytes)?; + std::fs::rename(&tmp, path) + }; + if let Err(e) = write() { + tracing::warn!( + target: "session_store", + "cannot write session store {}: {e}", + path.display() + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tmp_path(name: &str) -> PathBuf { + std::env::temp_dir().join(format!( + "buzz-acp-session-store-{}-{name}.json", + std::process::id() + )) + } + + #[test] + fn disabled_store_remembers_nothing() { + let store = SessionStore::disabled(); + let ch = Uuid::new_v4(); + store.put(ch, "sess"); + assert_eq!(store.get(ch), None); + assert!(!store.is_enabled()); + } + + #[test] + fn put_survives_reopen_and_remove_forgets() { + let path = tmp_path("roundtrip"); + let _ = std::fs::remove_file(&path); + let ch_a = Uuid::new_v4(); + let ch_b = Uuid::new_v4(); + + let store = SessionStore::open(&path); + store.put(ch_a, "sess-a"); + store.put(ch_b, "sess-b"); + drop(store); + + // "The process restarted." + let store = SessionStore::open(&path); + assert_eq!(store.get(ch_a).as_deref(), Some("sess-a")); + assert_eq!(store.get(ch_b).as_deref(), Some("sess-b")); + + store.remove(ch_a); + drop(store); + let store = SessionStore::open(&path); + assert_eq!(store.get(ch_a), None); + assert_eq!(store.get(ch_b).as_deref(), Some("sess-b")); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn malformed_file_is_treated_as_empty() { + let path = tmp_path("malformed"); + std::fs::write(&path, b"{ not json").unwrap(); + let store = SessionStore::open(&path); + assert_eq!(store.get(Uuid::new_v4()), None); + // And it is writable afterwards. + let ch = Uuid::new_v4(); + store.put(ch, "sess"); + assert_eq!(SessionStore::open(&path).get(ch).as_deref(), Some("sess")); + let _ = std::fs::remove_file(&path); + } +}