Skip to main content

mati_core/mcp/
server.rs

1//! MCP stdio server entry point (M-07).
2//!
3//! `serve()` is the entry point. It opens the store, loads the graph,
4//! constructs `MatiServer`, and runs the rmcp stdio transport. After the
5//! client disconnects, the process auto-promotes to a headless daemon and
6//! waits for an idle timeout or signal before shutting down (a panic hook
7//! is installed at startup; lifecycle events are recorded throughout;
8//! a boot-time auto-drain bounded by `AUTO_DRAIN_TIMEOUT` runs the dirty
9//! gotcha-index repair).
10//!
11//! Also binds the Unix daemon socket (`~/.mati/<slug>/mati.sock`) so that hook
12//! scripts using `mati get`/`mati ping` can route through the daemon protocol
13//! instead of trying to open the SurrealKV store directly (which would fail with
14//! a lock error while the MCP server holds the exclusive handle). The socket
15//! task is supervised: a watcher signals graceful shutdown if it dies, and
16//! a `SHUTDOWN_DRAIN_TIMEOUT` ceiling falls back to `abort_handle` so a
17//! wedged handler can never block exit.
18//!
19//! Public surface: `serve`, `socket_handle_connection`, `Shutdown` (+
20//! methods), and the policy constants `AUTO_DRAIN_TIMEOUT`,
21//! `MAX_CONCURRENT_CONNECTIONS`, `IDLE_SHUTDOWN_SECS`,
22//! `IDLE_CHECK_INTERVAL_SECS`, `UNIX_SOCK_PATH_MAX` — all shared with
23//! `cli::daemon` so both daemon paths use identical operational policy.
24
25use std::path::Path;
26use std::path::PathBuf;
27use std::sync::Arc;
28use std::time::Duration;
29
30use anyhow::Result;
31use rmcp::model::{ServerCapabilities, ServerInfo};
32use rmcp::{tool_handler, ServerHandler, ServiceExt};
33use serde::{Deserialize, Serialize};
34use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
35use tokio::net::UnixStream;
36
37use crate::graph::edges::EdgeKind;
38use crate::graph::Graph;
39
40use super::tools::MatiServer;
41use super::types::{MemBootstrapParams, MemGetParams, MemQueryParams, MemSetParams};
42
43#[derive(Debug)]
44pub(crate) enum ProxyDaemonResult {
45    Ok(serde_json::Value),
46    NotRunning,
47    StaleSocket,
48    Unresponsive,
49}
50
51#[tool_handler(router = self.tool_router)]
52impl ServerHandler for MatiServer {
53    fn get_info(&self) -> ServerInfo {
54        ServerInfo::new(
55            ServerCapabilities::builder()
56                .enable_tools()
57                .enable_tool_list_changed()
58                .build(),
59        )
60        .with_instructions(
61            "mati is a persistent engineering knowledge store for the current \
62                 codebase. Use mem_get for direct record lookup, mem_query for \
63                 search and graph traversal, mem_bootstrap for session context, \
64                 and mem_set for writing knowledge records.",
65        )
66    }
67}
68
69/// Start the MCP stdio proxy for the project rooted at `repo_root`.
70///
71/// After γ-C4, `mati serve` is a thin MCP-stdio ↔ UDS forwarder: every
72/// tool call is proxied over the Unix domain socket to a separate daemon
73/// process which owns the store, the graph, the socket listener, the
74/// idle-shutdown loop, signal handling, and the auto-drain pipeline.
75///
76/// On startup:
77/// 1. Resolve `~/.mati/<slug>/` from `repo_root`.
78/// 2. Ensure a daemon is running (auto-spawning one if necessary via the
79///    state-aware readiness machinery in `daemon_lifecycle::ensure_daemon`).
80/// 3. Bind the rmcp stdio transport and forward every request to the
81///    daemon via `MatiServer::with_socket_root`.
82///
83/// On client disconnect, this process exits cleanly — the daemon (separate
84/// process) is unaffected and remains available for the next `mati serve`
85/// invocation that Codex / Claude Code spawns.
86///
87/// Lifecycle events (`serve_start`, `serve_failed`, `serve_shutdown`,
88/// `startup`) are appended throughout so `mati doctor` can observe the
89/// proxy's cold-start path.
90pub async fn serve(repo_root: &Path) -> Result<()> {
91    let startup_t0 = std::time::Instant::now();
92
93    // Resolve the daemon root so we can emit lifecycle events even before
94    // the daemon is reachable.
95    let mati_root: PathBuf = crate::store::mati_home_opt()
96        .map(|h| h.join(crate::store::derive_slug(repo_root)))
97        .ok_or_else(|| anyhow::anyhow!("cannot resolve home directory for mati_root"))?;
98
99    super::metadata::record_lifecycle_event(&mati_root, "startup", "phase=ensure_daemon");
100
101    // The daemon owns the store. `ensure_daemon` spawns a daemon if needed
102    // and waits for it to be ready via the state-aware readiness machinery
103    // (`daemon_lifecycle::wait_for_ready`).
104    if !super::daemon_lifecycle::ensure_daemon(&mati_root).await {
105        super::metadata::record_lifecycle_event(
106            &mati_root,
107            "serve_failed",
108            "daemon unreachable after auto-spawn",
109        );
110        anyhow::bail!(
111            "mati serve: daemon unreachable. \
112             Run `mati daemon start` manually and check the lifecycle.log."
113        );
114    }
115
116    super::metadata::record_lifecycle_event(
117        &mati_root,
118        "serve_start",
119        &format!("pid={} owner=proxy", std::process::id()),
120    );
121
122    // Initialize the metrics handle so any local recording is no-op rather
123    // than panicking. The daemon owns the authoritative metrics surface.
124    super::metrics::init();
125
126    super::metadata::record_lifecycle_event(
127        &mati_root,
128        "startup",
129        &format!(
130            "phase=ready elapsed_ms={}",
131            startup_t0.elapsed().as_millis()
132        ),
133    );
134
135    // MCP stdio proxy: every tool call forwards over UDS to the daemon.
136    let transport = rmcp::transport::io::stdio();
137    let service = MatiServer::with_socket_root(mati_root.clone())
138        .serve(transport)
139        .await
140        .map_err(|e| anyhow::anyhow!("MCP proxy initialization failed: {e}"))
141        .inspect_err(|e| {
142            super::metadata::record_lifecycle_event(
143                &mati_root,
144                "serve_failed",
145                &format!("proxy init: {e:#}"),
146            )
147        })?;
148
149    let shutdown_reason: &'static str = match service.waiting().await {
150        Ok(_) => "client_disconnect",
151        Err(e) => {
152            super::metadata::record_lifecycle_event(
153                &mati_root,
154                "serve_failed",
155                &format!("proxy waiting: {e}"),
156            );
157            "mcp_waiting_error"
158        }
159    };
160    super::metadata::record_lifecycle_event(
161        &mati_root,
162        "serve_shutdown",
163        &format!("reason={shutdown_reason}"),
164    );
165    Ok(())
166}
167
168pub(crate) async fn proxy_daemon_result(
169    root: &Path,
170    cmd: &str,
171    args: serde_json::Value,
172) -> ProxyDaemonResult {
173    // Daemon-restart resilience: when `mati daemon stop` followed by
174    // `mati daemon start` happens during an active MCP-stdio session, the
175    // first call after the restart can fail in three ways:
176    //   1. Socket file transiently absent (NotRunning)
177    //   2. Connection refused before the new daemon's accept loop is up
178    //      (StaleSocket / Unresponsive depending on metadata state)
179    //   3. Connection succeeds but the request carries a stale session UUID
180    //      (cached by the rmcp tool dispatcher) → daemon returns
181    //      "session_mismatch" via the v2 fence in `dispatch_v2`.
182    //
183    // Without retry, every subsequent MCP tool call returns a structured
184    // error to Claude/Codex — a P9 violation in spirit since the agent's
185    // entire MCP session becomes unusable until restart.
186    //
187    // The retry is bounded: at most one re-connect after a brief delay,
188    // re-reading daemon metadata so the new session UUID is picked up.
189    // We do NOT retry indefinitely — a hard-down daemon must surface an
190    // error eventually so the caller can fall back.
191    let result = proxy_daemon_result_no_spawn(root, cmd, &args).await;
192
193    // Pass-33: if both retries failed because the daemon is gone (not
194    // because of a session mismatch or a transient stall), auto-spawn a
195    // fresh daemon and try one final time. Phase 3's `mati daemon stop`
196    // cycles for repair/init left the daemon unrun, breaking every MCP
197    // tool call until manual restart — this closes that hole.
198    //
199    // Only NotRunning/StaleSocket are eligible: Unresponsive means
200    // ensure_daemon has its own SIGTERM-and-cleanup recovery path that
201    // would conflict with our retry, and Ok / session-mismatch don't
202    // need a spawn.
203    if matches!(
204        &result,
205        ProxyDaemonResult::NotRunning | ProxyDaemonResult::StaleSocket
206    ) && super::daemon_lifecycle::ensure_daemon(root).await
207    {
208        match proxy_daemon_result_once(root, cmd, &args).await {
209            AttemptOutcome::Final(r) | AttemptOutcome::Retryable(r) => return r,
210        }
211    }
212
213    result
214}
215
216/// Inner: the original two-attempt retry without auto-spawn. Extracted so
217/// `daemon_lifecycle::ensure_daemon`'s probe can call this without
218/// triggering its own auto-spawn (which would loop indefinitely).
219pub(crate) async fn proxy_daemon_result_no_spawn(
220    root: &Path,
221    cmd: &str,
222    args: &serde_json::Value,
223) -> ProxyDaemonResult {
224    match proxy_daemon_result_once(root, cmd, args).await {
225        AttemptOutcome::Final(result) => result,
226        AttemptOutcome::Retryable(_) => {
227            // Brief settle — give the new daemon time to bind socket and
228            // publish metadata. 100ms is generous; daemon startup is ~50ms.
229            tokio::time::sleep(Duration::from_millis(100)).await;
230            match proxy_daemon_result_once(root, cmd, args).await {
231                AttemptOutcome::Final(result) | AttemptOutcome::Retryable(result) => result,
232            }
233        }
234    }
235}
236
237/// Outcome of a single `proxy_daemon_result` attempt.
238///
239/// `Retryable` carries the result the caller would have returned if no
240/// retry were attempted — used as the fallback if the second attempt also
241/// fails. This keeps the original error shape stable for callers that
242/// distinguish StaleSocket vs Unresponsive vs structured session_mismatch.
243enum AttemptOutcome {
244    Final(ProxyDaemonResult),
245    Retryable(ProxyDaemonResult),
246}
247
248async fn proxy_daemon_result_once(
249    root: &Path,
250    cmd: &str,
251    args: &serde_json::Value,
252) -> AttemptOutcome {
253    // Build v2 request from v1-style (cmd, args) using the same mapping
254    // as cli::daemon::daemon_result. Pure-reads only — mutating callers
255    // must use [`proxy_daemon_v2`] with a typed Command (see pass-29).
256    let v2_cmd = super::protocol::v1_to_v2_command(cmd, args);
257    proxy_daemon_send_v2(root, v2_cmd).await
258}
259
260/// Send a typed v2 [`super::protocol::Command`] to the daemon socket.
261///
262/// Mirrors [`proxy_daemon_result`] for callers (currently the MCP Socket-
263/// backend `mem_set` path) that have moved to typed commands and would
264/// otherwise have to round-trip through the legacy v1 mapper, which has
265/// no entries for mutating commands and panics on them.
266///
267/// Bounded auto-reconnect mirrors `proxy_daemon_result` so a daemon
268/// restart during an active session is recovered transparently.
269pub(crate) async fn proxy_daemon_v2(
270    root: &Path,
271    cmd: super::protocol::Command,
272) -> ProxyDaemonResult {
273    // Serialize once — every retry uses the same wire bytes.
274    let v2_cmd = match serde_json::to_value(&cmd) {
275        Ok(v) => v,
276        Err(_) => return ProxyDaemonResult::Unresponsive,
277    };
278
279    let result = match proxy_daemon_send_v2(root, v2_cmd.clone()).await {
280        AttemptOutcome::Final(result) => result,
281        AttemptOutcome::Retryable(_) => {
282            tokio::time::sleep(Duration::from_millis(100)).await;
283            match proxy_daemon_send_v2(root, v2_cmd.clone()).await {
284                AttemptOutcome::Final(result) | AttemptOutcome::Retryable(result) => result,
285            }
286        }
287    };
288
289    // Pass-33: parallel auto-spawn for the typed-Command path. Same
290    // policy as `proxy_daemon_result`: if the two retries failed because
291    // the daemon is gone, ensure_daemon spawns one and we try once more.
292    if matches!(
293        &result,
294        ProxyDaemonResult::NotRunning | ProxyDaemonResult::StaleSocket
295    ) && super::daemon_lifecycle::ensure_daemon(root).await
296    {
297        match proxy_daemon_send_v2(root, v2_cmd).await {
298            AttemptOutcome::Final(r) | AttemptOutcome::Retryable(r) => return r,
299        }
300    }
301
302    result
303}
304
305/// Inner socket transaction: connect, send a pre-built v2 command JSON,
306/// read the response. Shared between v1-style and typed-Command callers
307/// so the connect/refused/session-mismatch policy stays identical.
308async fn proxy_daemon_send_v2(root: &Path, v2_cmd: serde_json::Value) -> AttemptOutcome {
309    let sock_path = root.join("mati.sock");
310
311    if sock_path.as_os_str().len() > UNIX_SOCK_PATH_MAX {
312        tracing::warn!(
313            path = %sock_path.display(),
314            "mcp proxy: socket path exceeds Unix limit"
315        );
316        // Path-length violation is not transient — never retry.
317        return AttemptOutcome::Final(ProxyDaemonResult::NotRunning);
318    }
319
320    if !sock_path.exists() {
321        // Socket missing — daemon may be mid-restart. Retry once.
322        return AttemptOutcome::Retryable(ProxyDaemonResult::NotRunning);
323    }
324
325    let stream = match UnixStream::connect(&sock_path).await {
326        Ok(s) => s,
327        Err(e) => {
328            let is_refused = e.kind() == std::io::ErrorKind::ConnectionRefused;
329            if is_refused {
330                // Socket refused — use the metadata + PID liveness protocol
331                // to decide whether to clean up. Never blindly remove.
332                use super::metadata::{self as meta, StaleCheckResult};
333                match meta::check_and_cleanup_stale(root) {
334                    StaleCheckResult::StaleRemoved | StaleCheckResult::Clean => {
335                        return AttemptOutcome::Retryable(ProxyDaemonResult::StaleSocket);
336                    }
337                    StaleCheckResult::OrphanSocket => {
338                        // No metadata + ECONNREFUSED → stale
339                        let _ = std::fs::remove_file(&sock_path);
340                        return AttemptOutcome::Retryable(ProxyDaemonResult::StaleSocket);
341                    }
342                    StaleCheckResult::LiveDaemon { .. } => {
343                        // PID alive but socket refused — daemon is starting or broken
344                        return AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive);
345                    }
346                }
347            }
348            return AttemptOutcome::Retryable(ProxyDaemonResult::NotRunning);
349        }
350    };
351
352    // Read daemon metadata fresh per attempt so a session UUID rotated by
353    // a daemon restart between attempt 1 and attempt 2 is picked up.
354    let daemon_session = super::metadata::read_metadata(root)
355        .map(|m| m.session)
356        .unwrap_or_else(uuid::Uuid::nil);
357    let request = serde_json::json!({
358        "v": super::protocol::PROTOCOL_VERSION,
359        "id": uuid::Uuid::new_v4(),
360        "session": daemon_session,
361        "cmd": v2_cmd,
362    });
363
364    let (reader, mut writer) = stream.into_split();
365    let mut bytes = match serde_json::to_vec(&request) {
366        Ok(b) => b,
367        Err(_) => return AttemptOutcome::Final(ProxyDaemonResult::Unresponsive),
368    };
369    bytes.push(b'\n');
370
371    if writer.write_all(&bytes).await.is_err() {
372        return AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive);
373    }
374    if writer.shutdown().await.is_err() {
375        return AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive);
376    }
377
378    let mut buf_reader = BufReader::new(reader);
379    let mut line = String::new();
380    match tokio::time::timeout(Duration::from_secs(2), buf_reader.read_line(&mut line)).await {
381        Ok(Ok(n)) if n > 0 => {}
382        _ => return AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive),
383    }
384
385    // Parse v2 Response and convert to v1-compatible envelope for callers.
386    let resp: serde_json::Value = match serde_json::from_str(line.trim()) {
387        Ok(v) => v,
388        Err(_) => return AttemptOutcome::Final(ProxyDaemonResult::Unresponsive),
389    };
390
391    match resp.get("status").and_then(|s| s.as_str()) {
392        Some("ok") => {
393            let data = resp.get("data").cloned().unwrap_or(serde_json::Value::Null);
394            AttemptOutcome::Final(ProxyDaemonResult::Ok(
395                serde_json::json!({"ok": true, "v": 2, "data": data}),
396            ))
397        }
398        Some("err") => {
399            let code = resp
400                .get("code")
401                .and_then(|c| c.as_str())
402                .unwrap_or("internal");
403            let message = resp
404                .get("message")
405                .and_then(|m| m.as_str())
406                .unwrap_or("unknown error");
407            let envelope = serde_json::json!({
408                "ok": false, "v": 2, "error": message, "code": code
409            });
410            // Session mismatch is the canonical "daemon restarted, your
411            // cached session is stale" signal — see dispatch_v2.rs's fence
412            // and the symmetric handling in cli::daemon::send_v2_raw. The
413            // retry will re-read metadata and pick up the new session UUID.
414            if code == "session_mismatch" {
415                tracing::debug!(
416                    "mcp proxy: session mismatch — daemon may have restarted, will retry"
417                );
418                AttemptOutcome::Retryable(ProxyDaemonResult::Ok(envelope))
419            } else {
420                AttemptOutcome::Final(ProxyDaemonResult::Ok(envelope))
421            }
422        }
423        _ => AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive),
424    }
425}
426
427// cleanup_stale_pid and local is_pid_alive removed — callers now use
428// metadata::check_and_cleanup_stale which centralizes PID liveness checks.
429
430// ── Daemon socket — hook script bridge ───────────────────────────────────────
431
432/// Unix domain socket path length limit (macOS-compatible).
433///
434/// Public so the parallel daemon path in `cli::daemon` shares the same
435/// value — preventing one path's bound from drifting from the other's.
436pub const UNIX_SOCK_PATH_MAX: usize = 104;
437
438/// Max wait for a complete request line per connection.
439const READ_TIMEOUT: Duration = Duration::from_secs(3);
440
441/// Maximum number of daemon-socket connections handled concurrently.
442///
443/// A flood beyond this limit blocks at `accept` (TCP backlog absorbs the
444/// surplus); this gives natural backpressure rather than unbounded memory
445/// use. 64 is generous for a per-user daemon — typical hook traffic is
446/// O(1) concurrent. Public so `cli::daemon` shares the same bound.
447pub const MAX_CONCURRENT_CONNECTIONS: usize = 64;
448
449/// Maximum time the boot-time auto-drain (dirty-marker queue) can run
450/// before we give up and proceed to serve. Prevents a pathological dirty
451/// queue from blocking daemon startup. The dirty marker stays set; the
452/// user can run `mati repair` manually.
453///
454/// Public so `cli::daemon::run_daemon_start` can share the same ceiling.
455pub const AUTO_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
456
457/// Race-free shutdown signal for daemon-socket loops.
458///
459/// `signal()` is idempotent and `wait()` resolves immediately if the signal
460/// has already fired. The `enable()` pattern on `Notify::notified()`
461/// registers the future before the flag check, so a `signal()` race between
462/// flag-set and notify-fire cannot strand a waiter.
463///
464/// Shared with `cli::daemon` so both the embedded MCP-server socket loop
465/// and the headless `mati daemon start` loop use identical shutdown
466/// semantics.
467#[derive(Default)]
468pub struct Shutdown {
469    flag: std::sync::atomic::AtomicBool,
470    notify: tokio::sync::Notify,
471}
472
473impl Shutdown {
474    pub fn new() -> Self {
475        Self::default()
476    }
477
478    /// Idempotent — safe to call multiple times. Wakes every active waiter.
479    pub fn signal(&self) {
480        self.flag.store(true, std::sync::atomic::Ordering::SeqCst);
481        self.notify.notify_waiters();
482    }
483
484    pub fn is_set(&self) -> bool {
485        // SeqCst (matching the store): defense-in-depth correctness on
486        // weakly-ordered architectures (ARM/POWER). Without it, the load
487        // would rely on Notify's internal mutex acquire to synchronize
488        // with `signal()`'s store — which is the pattern in our `wait()`
489        // body and works in practice, but depends on Notify's
490        // implementation detail. Explicit SC pairing is cheap (one
491        // memory barrier at most) and removes the implicit dependency.
492        self.flag.load(std::sync::atomic::Ordering::SeqCst)
493    }
494
495    /// Future resolves once `signal()` has been called. Safe to call
496    /// repeatedly; safe to race with concurrent `signal()`.
497    pub async fn wait(&self) {
498        let notified = self.notify.notified();
499        tokio::pin!(notified);
500        // Register the receiver BEFORE the flag check so a `signal()` that
501        // fires between check and notify cannot be missed.
502        notified.as_mut().enable();
503        if self.is_set() {
504            return;
505        }
506        notified.await;
507    }
508}
509
510/// Daemon protocol version (must match `cli::daemon::PROTOCOL_VERSION`).
511const PROTOCOL_VERSION: u32 = 1;
512
513#[derive(Debug, Deserialize)]
514pub(crate) struct SocketRequest {
515    pub cmd: String,
516    #[allow(dead_code)] // Wire protocol field — must exist for deserialization
517    #[serde(default, rename = "v")]
518    pub version: Option<u32>,
519    #[serde(default)]
520    pub args: serde_json::Value,
521}
522
523#[derive(Debug, Serialize)]
524pub(crate) struct SocketResponse {
525    pub(crate) ok: bool,
526    #[serde(rename = "v")]
527    version: u32,
528    #[serde(skip_serializing_if = "Option::is_none")]
529    pub(crate) data: Option<serde_json::Value>,
530    #[serde(skip_serializing_if = "Option::is_none")]
531    pub(crate) error: Option<String>,
532}
533
534impl SocketResponse {
535    pub(crate) fn ok(data: serde_json::Value) -> Self {
536        Self {
537            ok: true,
538            version: PROTOCOL_VERSION,
539            data: Some(data),
540            error: None,
541        }
542    }
543    pub(crate) fn err(msg: impl Into<String>) -> Self {
544        Self {
545            ok: false,
546            version: PROTOCOL_VERSION,
547            data: None,
548            error: Some(msg.into()),
549        }
550    }
551}
552
553pub async fn socket_handle_connection(
554    graph: Arc<tokio::sync::RwLock<Graph>>,
555    repo_root: &Path,
556    stream: UnixStream,
557    peer: super::metadata::PeerContext,
558    daemon_session: uuid::Uuid,
559) -> Result<()> {
560    use super::protocol::MAX_FRAME_SIZE;
561    use tokio::io::AsyncReadExt;
562
563    let (reader, mut writer) = stream.into_split();
564    let mut buf = String::new();
565
566    // Cap the read at MAX_FRAME_SIZE + 1 bytes so the allocation is bounded
567    // before any JSON parsing occurs. If the client sends more data than
568    // MAX_FRAME_SIZE before the newline delimiter, `read_line` will stop at
569    // the take limit and the size check below will reject the request.
570    let limited = reader.take(MAX_FRAME_SIZE as u64 + 1);
571    let mut buf_reader = BufReader::new(limited);
572    match tokio::time::timeout(READ_TIMEOUT, buf_reader.read_line(&mut buf)).await {
573        Ok(Ok(0)) => return Ok(()),
574        Ok(Ok(_)) => {}
575        Ok(Err(e)) => anyhow::bail!("read error: {e}"),
576        Err(_) => anyhow::bail!("read timeout"),
577    }
578
579    if buf.len() > MAX_FRAME_SIZE {
580        let resp = super::protocol::Response::err(
581            uuid::Uuid::nil(),
582            super::protocol::ErrorCode::FrameTooLarge,
583            format!("request exceeds {MAX_FRAME_SIZE} byte limit"),
584        );
585        let json = serde_json::to_string(&resp)?;
586        writer.write_all(json.as_bytes()).await?;
587        writer.write_all(b"\n").await?;
588        writer.flush().await?;
589        return Ok(());
590    }
591
592    let trimmed = buf.trim();
593
594    // V2 protocol ONLY — no v1 fallback on the public wire.
595    // The v2 format requires `id` (UUID), `session` (UUID), and `cmd` as
596    // a tagged object with `type`. If decode fails, the request is rejected
597    // with a protocol error — there is no legacy v1 dispatch path.
598    let v2_req = match serde_json::from_str::<super::protocol::Request>(trimmed) {
599        Ok(r) => r,
600        Err(e) => {
601            // Return a v2-shaped error. Use nil UUID since we can't extract
602            // the request ID from a malformed payload.
603            let resp = super::protocol::Response::err(
604                uuid::Uuid::nil(),
605                super::protocol::ErrorCode::MalformedRequest,
606                format!("invalid v2 request: {e}"),
607            );
608            let json = serde_json::to_string(&resp)?;
609            writer.write_all(json.as_bytes()).await?;
610            writer.write_all(b"\n").await?;
611            writer.flush().await?;
612            return Ok(());
613        }
614    };
615
616    let ctx = super::dispatch_v2::RequestContext {
617        peer,
618        daemon_session,
619        repo_root: repo_root.to_path_buf(),
620    };
621    let resp = super::dispatch_v2::dispatch_v2(&graph, &ctx, v2_req).await;
622    let json = serde_json::to_string(&resp)?;
623    writer.write_all(json.as_bytes()).await?;
624    writer.write_all(b"\n").await?;
625    writer.flush().await?;
626    Ok(())
627}
628
629/// Build a `RequestContext` for the in-process v1 socket_dispatch path.
630///
631/// The wire layer (`socket_handle_connection`) carries authentic peer
632/// credentials and the daemon session UUID; v1 callers are in-process
633/// (e.g. tests), so they synthesize a context with the current process'
634/// identity. Used by the mem_* arms which now delegate to native handlers.
635fn build_v1_dispatch_ctx(repo_root: &Path) -> super::dispatch_v2::RequestContext {
636    super::dispatch_v2::RequestContext {
637        peer: super::metadata::PeerContext {
638            uid: super::metadata::current_euid(),
639            pid: Some(std::process::id()),
640        },
641        daemon_session: uuid::Uuid::nil(),
642        repo_root: repo_root.to_path_buf(),
643    }
644}
645
646pub(crate) async fn socket_dispatch(
647    graph: &Arc<tokio::sync::RwLock<Graph>>,
648    repo_root: &Path,
649    req: &SocketRequest,
650) -> SocketResponse {
651    use crate::store::session as sess;
652
653    match req.cmd.as_str() {
654        "ping" => SocketResponse::ok(serde_json::Value::String("pong".into())),
655
656        // Live daemon metrics snapshot — per-command counters, error rates,
657        // and p50/p95/p99 latencies. Pure read, no side effects, no audit.
658        // Returns `null` if the global metrics handle was never initialized
659        // (which only happens in tests that bypass `serve`).
660        "metrics" => match super::metrics::snapshot() {
661            Some(snap) => match serde_json::to_value(&snap) {
662                Ok(v) => SocketResponse::ok(v),
663                Err(e) => SocketResponse::err(format!("metrics serialize: {e}")),
664            },
665            None => SocketResponse::ok(serde_json::Value::Null),
666        },
667
668        // ── MCP tool commands ────────────────────────────────────────────
669        //
670        // γ-C4: the wire layer (`socket_handle_connection`) accepts only v2
671        // requests, which route MemGet / MemQuery / MemBootstrap / MemSet
672        // through `dispatch_v2` to the native handlers in `mcp::handlers`.
673        // These v1 arms are reachable only via in-process callers — they
674        // dispatch directly to the same canonical handlers so v1 and v2
675        // paths cannot drift.
676        "mem_get" => {
677            let params = match serde_json::from_value::<MemGetParams>(req.args.clone()) {
678                Ok(p) => p,
679                Err(e) => return SocketResponse::err(format!("invalid mem_get args: {e}")),
680            };
681            let input = super::protocol::MemGetInput { key: params.key };
682            let ctx = build_v1_dispatch_ctx(repo_root);
683            let g = graph.read().await;
684            match super::handlers::handle_mem_get(
685                g.store(),
686                graph,
687                &ctx,
688                uuid::Uuid::new_v4(),
689                &input,
690            )
691            .await
692            {
693                Ok(v) => SocketResponse::ok(serde_json::Value::String(
694                    serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".into()),
695                )),
696                Err((_code, msg)) => SocketResponse::err(msg),
697            }
698        }
699
700        "mem_query" => {
701            let params = match serde_json::from_value::<MemQueryParams>(req.args.clone()) {
702                Ok(p) => p,
703                Err(e) => return SocketResponse::err(format!("invalid mem_query args: {e}")),
704            };
705            let mode = match params.mode.as_str() {
706                "text" => super::protocol::QueryMode::Text,
707                "tag" => super::protocol::QueryMode::Tag,
708                "graph" => super::protocol::QueryMode::Graph,
709                "semantic" => super::protocol::QueryMode::Semantic,
710                other => {
711                    return SocketResponse::err(format!(
712                        "unknown mode: {other}. Valid modes: text, tag, graph, semantic"
713                    ));
714                }
715            };
716            let input = super::protocol::MemQueryInput {
717                query: params.query,
718                mode,
719                limit: params.limit as u32,
720            };
721            let g = graph.read().await;
722            match super::handlers::handle_mem_query(g.store(), &g, &input).await {
723                Ok(v) => SocketResponse::ok(serde_json::Value::String(
724                    serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".into()),
725                )),
726                Err((_code, msg)) => SocketResponse::err(msg),
727            }
728        }
729
730        "mem_bootstrap" => {
731            let params = match serde_json::from_value::<MemBootstrapParams>(req.args.clone()) {
732                Ok(p) => p,
733                Err(e) => return SocketResponse::err(format!("invalid mem_bootstrap args: {e}")),
734            };
735            let input = super::protocol::MemBootstrapInput {
736                context_files: params.context_files,
737            };
738            let ctx = build_v1_dispatch_ctx(repo_root);
739            let g = graph.read().await;
740            match super::handlers::handle_mem_bootstrap(
741                g.store(),
742                &g,
743                graph,
744                &ctx,
745                uuid::Uuid::new_v4(),
746                &input,
747            )
748            .await
749            {
750                Ok(s) => SocketResponse::ok(serde_json::Value::String(s)),
751                Err((_code, msg)) => SocketResponse::err(msg),
752            }
753        }
754
755        "mem_set" => {
756            let params = match serde_json::from_value::<MemSetParams>(req.args.clone()) {
757                Ok(p) => p,
758                Err(e) => return SocketResponse::err(format!("invalid mem_set args: {e}")),
759            };
760            let ctx = build_v1_dispatch_ctx(repo_root);
761            let response =
762                super::handlers::handle_mem_set(graph, &ctx, uuid::Uuid::new_v4(), &params).await;
763            SocketResponse::ok(serde_json::Value::String(response))
764        }
765
766        // ── Hook commands (store-only) ─────────────────────────────────
767        // Acquire a short-lived read lock for store access. The lock is
768        // released at the end of each arm — no risk of deadlock.
769        "get" => {
770            let key = match req.args.get("key").and_then(|v| v.as_str()) {
771                Some(k) => k,
772                None => return SocketResponse::err("missing args.key"),
773            };
774            let g = graph.read().await;
775            let store = g.store();
776            match store.get(key).await {
777                Ok(Some(record)) => {
778                    let confirmed = record
779                        .payload_as::<crate::store::GotchaRecord>()
780                        .map(|g| g.confirmed)
781                        .unwrap_or(false);
782                    match serde_json::to_value(&record) {
783                        Ok(mut val) => {
784                            if let Some(obj) = val.as_object_mut() {
785                                obj.insert(
786                                    "confirmed".to_string(),
787                                    serde_json::Value::Bool(confirmed),
788                                );
789                            }
790                            SocketResponse::ok(val)
791                        }
792                        Err(e) => SocketResponse::err(format!("serialize: {e}")),
793                    }
794                }
795                Ok(None) => SocketResponse::ok(serde_json::Value::Null),
796                Err(e) => SocketResponse::err(format!("store: {e}")),
797            }
798        }
799
800        // ── Internal hook-decide bulk command ────────────────────────────
801        // Returns file record + all linked gotcha records + consultation
802        // status in a single round-trip. NOT an MCP tool.
803        "hook_evaluate" => {
804            let file_key = match req.args.get("file_key").and_then(|v| v.as_str()) {
805                Some(k) => k,
806                None => return SocketResponse::err("missing args.file_key"),
807            };
808            let include_recent = req
809                .args
810                .get("include_recent")
811                .and_then(|v| v.as_bool())
812                .unwrap_or(false);
813            let actor = req.args.get("actor").and_then(|v| v.as_str());
814
815            let g = graph.read().await;
816            let store = g.store();
817
818            // 1. Fetch file record. Distinguish Ok(None) from Err.
819            let (file_record, store_error) = match store.get(file_key).await {
820                Ok(Some(r)) => (serde_json::to_value(&r).ok(), false),
821                Ok(None) => (None, false),
822                Err(e) => {
823                    tracing::warn!("hook_evaluate: store.get({file_key}) failed: {e}");
824                    (None, true)
825                }
826            };
827
828            // 2. Fetch all linked gotcha records.
829            //
830            // The canonical link is `file_record.payload.gotcha_keys`, written
831            // atomically by `compute_file_link_updates`. But CLAUDE.md flags
832            // this field as a *derived* index that can drift from the
833            // canonical `gotcha:*` records (e.g. if a gotcha was created
834            // before the file record existed, or if a partial-write left the
835            // file link stale). To make enforcement robust against that
836            // drift, we union three sources:
837            //   (a) `file_record.payload.gotcha_keys`               (primary)
838            //   (b) in-memory graph edges `HasGotcha` from file_key  (secondary)
839            //   (c) reverse scan of `gotcha:*` records whose
840            //       `affected_files` contains the relative path     (fallback)
841            // Source (c) is bounded by the active-gotcha count and
842            // short-circuited when (a) or (b) already produced results, so
843            // it does not add cost on the hot path.
844            let mut gotcha_records = serde_json::Map::new();
845            let mut gotcha_error = false;
846            let mut linked_keys: std::collections::BTreeSet<String> =
847                std::collections::BTreeSet::new();
848
849            if let Some(ref fr) = file_record {
850                if let Some(keys) = fr
851                    .pointer("/payload/gotcha_keys")
852                    .and_then(|v| v.as_array())
853                {
854                    for gk in keys {
855                        if let Some(key_str) = gk.as_str() {
856                            linked_keys.insert(key_str.to_string());
857                        }
858                    }
859                }
860            }
861
862            // (b) Graph-edge fallback. Loaded at boot from `graph:edge:*`,
863            // independent of the file record's denormalized list.
864            for nkey in g.neighbors(file_key, &crate::graph::EdgeKind::HasGotcha) {
865                linked_keys.insert(nkey);
866            }
867
868            // (c) Canonical reverse-lookup fallback. Only run when both
869            // derived indexes were empty AND a file record exists — covers
870            // the "file record present but gotcha_keys never synced" drift
871            // path that CLAUDE.md documents as a known trap. Bounded scan
872            // strips the relative path from the file_key once.
873            if linked_keys.is_empty() && file_record.is_some() {
874                let rel_path = file_key.strip_prefix("file:").unwrap_or(file_key);
875                if let Ok(all_gotchas) = store.scan_prefix("gotcha:").await {
876                    for r in all_gotchas {
877                        if !matches!(r.lifecycle, crate::store::RecordLifecycle::Active) {
878                            continue;
879                        }
880                        if let Some(g) = r.payload_as::<crate::store::GotchaRecord>() {
881                            if g.affected_files.iter().any(|af| af == rel_path) {
882                                linked_keys.insert(r.key.clone());
883                            }
884                        }
885                    }
886                }
887            }
888
889            for key_str in &linked_keys {
890                match store.get(key_str).await {
891                    Ok(Some(grec)) => {
892                        // Skip tombstoned gotchas so they never feed into enforcement.
893                        if !matches!(grec.lifecycle, crate::store::RecordLifecycle::Active) {
894                            continue;
895                        }
896                        // Inline confirmed flag (same as "get" handler).
897                        let confirmed = grec
898                            .payload_as::<crate::store::GotchaRecord>()
899                            .map(|g| g.confirmed)
900                            .unwrap_or(false);
901                        if let Ok(mut val) = serde_json::to_value(&grec) {
902                            if let Some(obj) = val.as_object_mut() {
903                                obj.insert(
904                                    "confirmed".to_string(),
905                                    serde_json::Value::Bool(confirmed),
906                                );
907                            }
908                            gotcha_records.insert(key_str.clone(), val);
909                        }
910                    }
911                    Ok(None) => {}
912                    Err(e) => {
913                        tracing::warn!("hook_evaluate: store.get({key_str}) failed: {e}");
914                        gotcha_error = true;
915                    }
916                }
917            }
918
919            // Project the unified gotcha_keys back into the returned
920            // file_record so decide.rs::evaluate (which iterates
921            // `payload.gotcha_keys`) sees every gotcha we just unioned —
922            // not just the ones the canonical link recorded.
923            let file_record = if let Some(mut fr) = file_record {
924                if !gotcha_records.is_empty() {
925                    if let Some(payload) = fr.pointer_mut("/payload") {
926                        if let Some(obj) = payload.as_object_mut() {
927                            let keys: Vec<serde_json::Value> = gotcha_records
928                                .keys()
929                                .map(|k| serde_json::Value::String(k.clone()))
930                                .collect();
931                            obj.insert("gotcha_keys".to_string(), serde_json::Value::Array(keys));
932                        }
933                    }
934                }
935                Some(fr)
936            } else {
937                None
938            };
939
940            // 3. Consultation status.
941            let consulted = sess::check_consulted(store, file_key, actor)
942                .await
943                .unwrap_or(false);
944            let consulted_recent = if include_recent {
945                sess::check_consulted_recent(store, file_key, 900, actor)
946                    .await
947                    .unwrap_or(false)
948            } else {
949                false
950            };
951
952            SocketResponse::ok(serde_json::json!({
953                "file_key": file_key,
954                "file_record": file_record,
955                "gotcha_records": gotcha_records,
956                "consulted": consulted,
957                "consulted_recent": consulted_recent,
958                "store_error": store_error,
959                "gotcha_error": gotcha_error,
960            }))
961        }
962
963        "log_hit" => {
964            let key = match req.args.get("key").and_then(|v| v.as_str()) {
965                Some(k) => k,
966                None => return SocketResponse::err("missing args.key"),
967            };
968            let g = graph.read().await;
969            if let Err(e) = sess::log_hit(g.store(), key).await {
970                tracing::warn!("daemon socket log_hit: {e}");
971            }
972            SocketResponse::ok(serde_json::Value::Null)
973        }
974
975        "log_miss" => {
976            let key = match req.args.get("key").and_then(|v| v.as_str()) {
977                Some(k) => k,
978                None => return SocketResponse::err("missing args.key"),
979            };
980            let g = graph.read().await;
981            if let Err(e) = sess::log_miss(g.store(), key).await {
982                tracing::warn!("daemon socket log_miss: {e}");
983            }
984            SocketResponse::ok(serde_json::Value::Null)
985        }
986
987        "log_compliance_miss" => {
988            let key = match req.args.get("key").and_then(|v| v.as_str()) {
989                Some(k) => k,
990                None => return SocketResponse::err("missing args.key"),
991            };
992            let g = graph.read().await;
993            let store = g.store();
994            if let Err(e) = sess::log_compliance_miss(store, key).await {
995                tracing::warn!("daemon socket log_compliance_miss: {e}");
996            }
997            // Record Deny enforcement event — best-effort
998            let _ = crate::store::enforcement::record_event(
999                store,
1000                crate::store::enforcement::EnforcementEventType::Deny,
1001                crate::store::enforcement::SubjectKind::File,
1002                key.to_string(),
1003                "claude".to_string(),
1004                None,
1005                "gotcha_above_threshold".to_string(),
1006                None,
1007            )
1008            .await;
1009            SocketResponse::ok(serde_json::Value::Null)
1010        }
1011
1012        "log_compliance_hit" => {
1013            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1014                Some(k) => k,
1015                None => return SocketResponse::err("missing args.key"),
1016            };
1017            let g = graph.read().await;
1018            let store = g.store();
1019            if let Err(e) = sess::log_compliance_hit(store, key).await {
1020                tracing::warn!("daemon socket log_compliance_hit: {e}");
1021            }
1022            // Record AllowAfterReceipt enforcement event — best-effort
1023            let _ = crate::store::enforcement::record_event(
1024                store,
1025                crate::store::enforcement::EnforcementEventType::AllowAfterReceipt,
1026                crate::store::enforcement::SubjectKind::File,
1027                key.to_string(),
1028                "claude".to_string(),
1029                None,
1030                "receipt_valid".to_string(),
1031                None,
1032            )
1033            .await;
1034            SocketResponse::ok(serde_json::Value::Null)
1035        }
1036
1037        "log_codex_shell_miss" => {
1038            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1039                Some(k) => k,
1040                None => return SocketResponse::err("missing args.key"),
1041            };
1042            let g = graph.read().await;
1043            if let Err(e) = sess::log_codex_shell_miss(g.store(), key).await {
1044                tracing::warn!("daemon socket log_codex_shell_miss: {e}");
1045            }
1046            SocketResponse::ok(serde_json::Value::Null)
1047        }
1048
1049        "log_bootstrap" => {
1050            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1051                Some(k) => k,
1052                None => return SocketResponse::err("missing args.key"),
1053            };
1054            let g = graph.read().await;
1055            if let Err(e) = sess::log_bootstrap(g.store(), key).await {
1056                tracing::warn!("daemon socket log_bootstrap: {e}");
1057            }
1058            SocketResponse::ok(serde_json::Value::Null)
1059        }
1060
1061        "log_prompt_nudge" => {
1062            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1063                Some(k) => k,
1064                None => return SocketResponse::err("missing args.key"),
1065            };
1066            let g = graph.read().await;
1067            if let Err(e) = sess::log_prompt_nudge(g.store(), key).await {
1068                tracing::warn!("daemon socket log_prompt_nudge: {e}");
1069            }
1070            SocketResponse::ok(serde_json::Value::Null)
1071        }
1072
1073        "session_check_consulted" => {
1074            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1075                Some(k) => k,
1076                None => return SocketResponse::err("missing args.key"),
1077            };
1078            let g = graph.read().await;
1079            match sess::check_consulted(g.store(), key, None).await {
1080                Ok(found) => SocketResponse::ok(serde_json::Value::Bool(found)),
1081                Err(e) => SocketResponse::err(format!("store: {e}")),
1082            }
1083        }
1084
1085        "session_check_consulted_recent" => {
1086            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1087                Some(k) => k,
1088                None => return SocketResponse::err("missing args.key"),
1089            };
1090            let ttl_secs = req
1091                .args
1092                .get("ttl_secs")
1093                .and_then(|v| v.as_u64())
1094                .unwrap_or(900);
1095            let g = graph.read().await;
1096            match sess::check_consulted_recent(g.store(), key, ttl_secs, None).await {
1097                Ok(found) => SocketResponse::ok(serde_json::Value::Bool(found)),
1098                Err(e) => SocketResponse::err(format!("store: {e}")),
1099            }
1100        }
1101
1102        "session_flush" => {
1103            let g = graph.read().await;
1104            if let Err(e) = sess::session_flush(g.store()).await {
1105                tracing::warn!("daemon socket session_flush: {e}");
1106            }
1107            SocketResponse::ok(serde_json::Value::Null)
1108        }
1109
1110        "session_harvest" => {
1111            // Note: uses no-staleness variant because StalenessAnalyzer (git2) is !Send.
1112            // Git-based staleness analysis runs on the next CLI-path harvest.
1113            let g = graph.read().await;
1114            if let Err(e) = sess::session_harvest_no_staleness(g.store()).await {
1115                tracing::warn!("daemon socket session_harvest: {e}");
1116            }
1117            SocketResponse::ok(serde_json::Value::Null)
1118        }
1119
1120        "reparse" => {
1121            let path = match req.args.get("path").and_then(|v| v.as_str()) {
1122                Some(p) => p,
1123                None => return SocketResponse::err("missing args.path"),
1124            };
1125            let g = graph.read().await;
1126            if let Err(e) = crate::analysis::reparse::reparse_impl(g.store(), repo_root, path).await
1127            {
1128                tracing::warn!("daemon socket reparse: {e}");
1129            }
1130            SocketResponse::ok(serde_json::Value::Null)
1131        }
1132
1133        "edit_hook" => {
1134            let path = match req.args.get("path").and_then(|v| v.as_str()) {
1135                Some(p) => p,
1136                None => return SocketResponse::err("missing args.path"),
1137            };
1138            let file_key = format!("file:{path}");
1139            let g = graph.read().await;
1140            let store = g.store();
1141            if let Err(e) = sess::log_hit(store, &file_key).await {
1142                tracing::warn!("daemon socket edit_hook: log_hit failed: {e}");
1143            }
1144            if let Err(e) = crate::analysis::reparse::reparse_impl(store, repo_root, path).await {
1145                tracing::warn!("daemon socket edit_hook: reparse failed (non-fatal): {e}");
1146            }
1147
1148            // Incremental blast radius update: recompute for the modified file,
1149            // its direct importers, and the files it imports.
1150            {
1151                use crate::analysis::blast_radius::BlastRadius;
1152                use crate::graph::edges::EdgeKind;
1153
1154                let mut keys_to_update = vec![file_key.clone()];
1155                // Files that import this file (their blast radius may change if
1156                // this file's import list changed).
1157                keys_to_update.extend(g.neighbors_incoming(&file_key, &EdgeKind::Imports));
1158                // Files this file imports (this file now counts as an importer).
1159                keys_to_update.extend(g.neighbors(&file_key, &EdgeKind::Imports));
1160
1161                for key in keys_to_update {
1162                    let br = BlastRadius::compute(&key, &g);
1163                    if let Ok(Some(mut rec)) = store.get(&key).await {
1164                        if let Some(mut fr) = rec.payload_as::<crate::store::record::FileRecord>() {
1165                            fr.blast_radius = Some(br);
1166                            rec.payload = serde_json::to_value(&fr).ok();
1167                            let _ = store.put(&key, &rec).await;
1168                        }
1169                    }
1170                }
1171            }
1172
1173            // Incremental staleness propagation: recompute for the edited
1174            // file's direct importers and their importers (depth 2 only).
1175            // Does NOT recompute the full repo — keeps the hook fast.
1176            {
1177                let mut affected_keys = vec![file_key.clone()];
1178                let d1 = g.neighbors_incoming(&file_key, &EdgeKind::Imports);
1179                for d1k in &d1 {
1180                    affected_keys.push(d1k.clone());
1181                    affected_keys.extend(g.neighbors_incoming(d1k, &EdgeKind::Imports));
1182                }
1183                // Collect records for just the affected neighborhood
1184                let mut neighborhood_recs = Vec::new();
1185                for key in &affected_keys {
1186                    if let Ok(Some(rec)) = store.get(key).await {
1187                        neighborhood_recs.push(rec);
1188                    }
1189                }
1190                // Also include the edited file itself as a potential source
1191                if let Ok(Some(rec)) = store.get(&file_key).await {
1192                    if !neighborhood_recs.iter().any(|r| r.key == file_key) {
1193                        neighborhood_recs.push(rec);
1194                    }
1195                }
1196                let propagation =
1197                    crate::analysis::propagation::compute_propagation(&neighborhood_recs, &g);
1198                for (key, prop) in &propagation {
1199                    if let Ok(Some(mut rec)) = store.get(key).await {
1200                        if let Some(mut fr) = rec.payload_as::<crate::store::record::FileRecord>() {
1201                            fr.propagated_staleness = Some(prop.clone());
1202                            rec.payload = serde_json::to_value(&fr).ok();
1203                            let _ = store.put(key, &rec).await;
1204                        }
1205                    }
1206                }
1207            }
1208
1209            SocketResponse::ok(serde_json::Value::Null)
1210        }
1211
1212        "doc_capture" => {
1213            let path = match req.args.get("path").and_then(|v| v.as_str()) {
1214                Some(p) => p,
1215                None => return SocketResponse::err("missing args.path"),
1216            };
1217            let content = req
1218                .args
1219                .get("content")
1220                .and_then(|v| v.as_str())
1221                .unwrap_or("");
1222            let g = graph.read().await;
1223            if let Err(e) = sess::doc_capture(g.store(), path, content).await {
1224                tracing::warn!("daemon socket doc_capture: {e}");
1225            }
1226            SocketResponse::ok(serde_json::Value::Null)
1227        }
1228
1229        "scan_prefix" => {
1230            let prefix = match req.args.get("prefix").and_then(|v| v.as_str()) {
1231                Some(p) => p,
1232                None => return SocketResponse::err("missing args.prefix"),
1233            };
1234            let g = graph.read().await;
1235            match g.store().scan_prefix(prefix).await {
1236                Ok(records) => match serde_json::to_value(&records) {
1237                    Ok(val) => SocketResponse::ok(val),
1238                    Err(e) => SocketResponse::err(format!("serialize: {e}")),
1239                },
1240                Err(e) => SocketResponse::err(format!("store: {e}")),
1241            }
1242        }
1243
1244        "scan_keys" => {
1245            let prefix = match req.args.get("prefix").and_then(|v| v.as_str()) {
1246                Some(p) => p,
1247                None => return SocketResponse::err("missing args.prefix"),
1248            };
1249            let g = graph.read().await;
1250            match g.store().scan_keys(prefix).await {
1251                Ok(keys) => match serde_json::to_value(&keys) {
1252                    Ok(val) => SocketResponse::ok(val),
1253                    Err(e) => SocketResponse::err(format!("serialize: {e}")),
1254                },
1255                Err(e) => SocketResponse::err(format!("store: {e}")),
1256            }
1257        }
1258
1259        "scan_enforcement_events" => {
1260            let since_seq = req
1261                .args
1262                .get("since_seq")
1263                .and_then(|v| v.as_u64())
1264                .unwrap_or(0);
1265            let until_seq = req
1266                .args
1267                .get("until_seq")
1268                .and_then(|v| v.as_u64())
1269                .unwrap_or(u64::MAX);
1270            let g = graph.read().await;
1271            match crate::store::enforcement::scan_enforcement_events(
1272                g.store(),
1273                since_seq,
1274                until_seq,
1275            )
1276            .await
1277            {
1278                Ok(events) => match serde_json::to_value(&events) {
1279                    Ok(val) => SocketResponse::ok(val),
1280                    Err(e) => SocketResponse::err(format!("serialize: {e}")),
1281                },
1282                Err(e) => SocketResponse::err(format!("store: {e}")),
1283            }
1284        }
1285
1286        "put" => {
1287            use crate::store::Record;
1288            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1289                Some(k) => k,
1290                None => return SocketResponse::err("missing args.key"),
1291            };
1292            let record: Record = match req
1293                .args
1294                .get("record")
1295                .and_then(|v| serde_json::from_value(v.clone()).ok())
1296            {
1297                Some(r) => r,
1298                None => return SocketResponse::err("put: invalid record"),
1299            };
1300            let g = graph.read().await;
1301            match g.store().put(key, &record).await {
1302                Ok(()) => SocketResponse::ok(serde_json::Value::Null),
1303                Err(e) => SocketResponse::err(format!("store put: {e}")),
1304            }
1305        }
1306
1307        "delete" => {
1308            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1309                Some(k) => k,
1310                None => return SocketResponse::err("missing args.key"),
1311            };
1312            let g = graph.read().await;
1313            match g.store().delete(key).await {
1314                Ok(()) => SocketResponse::ok(serde_json::Value::Null),
1315                Err(e) => SocketResponse::err(format!("delete: {e}")),
1316            }
1317        }
1318
1319        "history" => {
1320            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1321                Some(k) => k,
1322                None => return SocketResponse::err("missing args.key"),
1323            };
1324            let limit = req.args.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as usize;
1325            let g = graph.read().await;
1326            match g.store().history(key, limit) {
1327                Ok(entries) => match serde_json::to_value(&entries) {
1328                    Ok(val) => SocketResponse::ok(val),
1329                    Err(e) => SocketResponse::err(format!("serialize: {e}")),
1330                },
1331                Err(e) => SocketResponse::err(format!("history: {e}")),
1332            }
1333        }
1334
1335        "history_since" => {
1336            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1337                Some(k) => k,
1338                None => return SocketResponse::err("missing args.key"),
1339            };
1340            let since_ts = req
1341                .args
1342                .get("since_ts")
1343                .and_then(|v| v.as_u64())
1344                .unwrap_or(0);
1345            let limit = req.args.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as usize;
1346            let g = graph.read().await;
1347            match g.store().history_since(key, since_ts, limit) {
1348                Ok(entries) => match serde_json::to_value(&entries) {
1349                    Ok(val) => SocketResponse::ok(val),
1350                    Err(e) => SocketResponse::err(format!("serialize: {e}")),
1351                },
1352                Err(e) => SocketResponse::err(format!("history_since: {e}")),
1353            }
1354        }
1355
1356        "gotcha_write" => {
1357            use crate::store::gotcha_ops::apply_gotcha_write;
1358            use crate::store::Record;
1359
1360            let record: Record = match req
1361                .args
1362                .get("record")
1363                .and_then(|v| serde_json::from_value(v.clone()).ok())
1364            {
1365                Some(r) => r,
1366                None => return SocketResponse::err("missing or invalid args.record"),
1367            };
1368            let new_files: Vec<String> = req
1369                .args
1370                .get("new_files")
1371                .and_then(|v| serde_json::from_value(v.clone()).ok())
1372                .unwrap_or_default();
1373            let old_files: Vec<String> = req
1374                .args
1375                .get("old_files")
1376                .and_then(|v| serde_json::from_value(v.clone()).ok())
1377                .unwrap_or_default();
1378            let is_new = req
1379                .args
1380                .get("is_new")
1381                .and_then(|v| v.as_bool())
1382                .unwrap_or(false);
1383
1384            {
1385                let g = graph.read().await;
1386                match apply_gotcha_write(g.store(), &record, &old_files, &new_files, is_new).await {
1387                    Ok(()) => {}
1388                    Err(e) => return SocketResponse::err(format!("{e}")),
1389                }
1390            }
1391
1392            // Sync the in-memory graph: add HasGotcha edges for newly-affected files,
1393            // remove edges for files no longer affected. The persistent store was already
1394            // updated by apply_gotcha_write above; this keeps the in-memory adjacency list
1395            // in sync so that assemble_context_packet (bootstrap) sees the edges immediately
1396            // without requiring a daemon restart.
1397            let record_key = record.key.clone();
1398            let old_set: std::collections::HashSet<&str> =
1399                old_files.iter().map(String::as_str).collect();
1400            let new_set: std::collections::HashSet<&str> =
1401                new_files.iter().map(String::as_str).collect();
1402            {
1403                let mut g = graph.write().await;
1404                for file_path in new_set.difference(&old_set) {
1405                    let file_key = format!("file:{file_path}");
1406                    let _ = g
1407                        .add_edge(&file_key, EdgeKind::HasGotcha, &record_key)
1408                        .await;
1409                }
1410                for file_path in old_set.difference(&new_set) {
1411                    let file_key = format!("file:{file_path}");
1412                    let _ = g
1413                        .remove_edge(&file_key, &EdgeKind::HasGotcha, &record_key)
1414                        .await;
1415                }
1416            }
1417
1418            SocketResponse::ok(serde_json::Value::String("written".into()))
1419        }
1420
1421        "gotcha_tombstone" => {
1422            use crate::store::gotcha_ops::apply_gotcha_tombstone;
1423
1424            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1425                Some(k) => k,
1426                None => return SocketResponse::err("missing args.key"),
1427            };
1428            if !key.starts_with("gotcha:") {
1429                return SocketResponse::err("delete action only applies to gotcha: keys");
1430            }
1431            // Read affected_files from args if provided, otherwise look up the
1432            // record to get them. The MCP proxy sends delete without affected_files.
1433            let mut affected_files: Vec<String> = req
1434                .args
1435                .get("affected_files")
1436                .and_then(|v| serde_json::from_value(v.clone()).ok())
1437                .unwrap_or_default();
1438
1439            let g = graph.read().await;
1440            if affected_files.is_empty() {
1441                if let Ok(Some(record)) = g.store().get(key).await {
1442                    if let Some(gotcha) = record.payload_as::<crate::store::GotchaRecord>() {
1443                        affected_files = gotcha.affected_files;
1444                    }
1445                }
1446            }
1447            match apply_gotcha_tombstone(g.store(), key, &affected_files).await {
1448                Ok(()) => SocketResponse::ok(serde_json::Value::String("tombstoned".into())),
1449                Err(e) => SocketResponse::err(format!("{e}")),
1450            }
1451        }
1452
1453        "gotcha_confirm" => {
1454            let key = match req.args.get("key").and_then(|v| v.as_str()) {
1455                Some(k) => k,
1456                None => return SocketResponse::err("missing args.key"),
1457            };
1458
1459            // Read record
1460            let g = graph.read().await;
1461            let store = g.store();
1462            let mut record = match store.get(key).await {
1463                Ok(Some(r)) => r,
1464                Ok(None) => return SocketResponse::err(format!("record not found: {key}")),
1465                Err(e) => return SocketResponse::err(format!("store get: {e}")),
1466            };
1467
1468            if record.category != crate::store::record::Category::Gotcha {
1469                return SocketResponse::err(format!("{key} is not a gotcha record"));
1470            }
1471
1472            if !matches!(
1473                record.lifecycle,
1474                crate::store::record::RecordLifecycle::Active
1475            ) {
1476                return SocketResponse::err(format!(
1477                    "{key} is tombstoned — cannot confirm a deleted record"
1478                ));
1479            }
1480
1481            // Set confirmed + normalize severity
1482            if let Some(ref mut payload) = record.payload {
1483                if let Some(obj) = payload.as_object_mut() {
1484                    if let Some(sev) = obj
1485                        .get("severity")
1486                        .and_then(|v| v.as_str())
1487                        .map(|s| s.to_lowercase())
1488                    {
1489                        obj.insert("severity".to_string(), serde_json::Value::String(sev));
1490                    }
1491                    obj.insert("confirmed".to_string(), serde_json::Value::Bool(true));
1492                }
1493            }
1494
1495            record.source = crate::store::record::RecordSource::DeveloperManual;
1496            record.confidence.value = crate::store::record::ConfidenceScore::base_for_source(
1497                &crate::store::record::RecordSource::DeveloperManual,
1498            );
1499            record.confidence.confirmation_count += 1;
1500            record.quality = crate::health::quality::analyze(&record);
1501
1502            let now = std::time::SystemTime::now()
1503                .duration_since(std::time::UNIX_EPOCH)
1504                .unwrap_or_default()
1505                .as_secs();
1506            record.updated_at = now;
1507            record.version.logical_clock += 1;
1508            record.version.wall_clock = now;
1509
1510            // Extract affected_files for file-link sync
1511            let affected_files: Vec<String> = record
1512                .payload_as::<crate::store::record::GotchaRecord>()
1513                .map(|g| g.affected_files)
1514                .unwrap_or_default();
1515
1516            if let Err(e) = store.put(key, &record).await {
1517                return SocketResponse::err(format!("store put: {e}"));
1518            }
1519
1520            // Sync file:*.gotcha_keys — best-effort
1521            for file_path in &affected_files {
1522                let file_key = format!("file:{file_path}");
1523                if let Ok(Some(mut file_record)) = store.get(&file_key).await {
1524                    let needs_link = file_record
1525                        .payload
1526                        .as_ref()
1527                        .and_then(|p| p.get("gotcha_keys"))
1528                        .and_then(|v| v.as_array())
1529                        .map(|arr| !arr.iter().any(|v| v.as_str() == Some(key)))
1530                        .unwrap_or(true);
1531                    if needs_link {
1532                        if let Some(ref mut payload) = file_record.payload {
1533                            if let Some(obj) = payload.as_object_mut() {
1534                                let arr = obj.entry("gotcha_keys").or_insert(serde_json::json!([]));
1535                                if let Some(arr) = arr.as_array_mut() {
1536                                    arr.push(serde_json::Value::String(key.to_string()));
1537                                }
1538                            }
1539                        }
1540                        let _ = store.put(&file_key, &file_record).await;
1541                    }
1542                }
1543            }
1544
1545            // Propagate confirmation_count to linked file records
1546            crate::store::gotcha_ops::propagate_confirmation_to_files(store, &affected_files).await;
1547
1548            // Record ControlChanged::Confirmed enforcement event — best-effort.
1549            let _ = crate::store::enforcement::record_event(
1550                store,
1551                crate::store::enforcement::EnforcementEventType::ControlChanged {
1552                    change_kind: crate::store::enforcement::ControlChangeKind::Confirmed,
1553                },
1554                crate::store::enforcement::SubjectKind::Control,
1555                key.to_string(),
1556                "developer".to_string(),
1557                None,
1558                "control_confirmed".to_string(),
1559                None,
1560            )
1561            .await;
1562
1563            SocketResponse::ok(serde_json::json!({"confirmed": true, "key": key}))
1564        }
1565
1566        other => SocketResponse::err(format!("unknown command: {other}")),
1567    }
1568}
1569
1570// ── Auto-promotion: MCP server → headless daemon ─────────────────────────────
1571
1572/// Idle shutdown threshold — wall-clock seconds with no daemon socket requests.
1573///
1574/// Shared with `cli::daemon` so both daemon paths use the same idle policy.
1575pub const IDLE_SHUTDOWN_SECS: u64 = 30 * 60; // 30 min
1576
1577/// How often to check wall-clock idle time. Shared with `cli::daemon`.
1578pub const IDLE_CHECK_INTERVAL_SECS: u64 = 5 * 60; // 5 min
1579
1580// ── Tests ────────────────────────────────────────────────────────────────────
1581
1582#[cfg(test)]
1583mod shutdown_tests {
1584    use super::*;
1585    use std::sync::Arc;
1586    use std::time::Duration;
1587
1588    #[tokio::test]
1589    async fn shutdown_signal_before_wait_returns_immediately() {
1590        // Pre-signal: subsequent wait must NOT block. Tests the flag-check
1591        // arm of `wait()` before the notified.await.
1592        let s = Shutdown::new();
1593        s.signal();
1594        // Should return well under timeout — generous bound to avoid CI flake.
1595        tokio::time::timeout(Duration::from_millis(100), s.wait())
1596            .await
1597            .expect("wait must return immediately when already signaled");
1598        assert!(s.is_set());
1599    }
1600
1601    #[tokio::test]
1602    async fn shutdown_wait_then_signal_wakes_waiter() {
1603        let s = Arc::new(Shutdown::new());
1604        let s_clone = Arc::clone(&s);
1605        let waiter = tokio::spawn(async move { s_clone.wait().await });
1606
1607        // Give the waiter a moment to register on `notified()`.
1608        tokio::time::sleep(Duration::from_millis(20)).await;
1609        assert!(!s.is_set());
1610
1611        s.signal();
1612
1613        tokio::time::timeout(Duration::from_millis(200), waiter)
1614            .await
1615            .expect("waiter must wake within timeout")
1616            .expect("waiter task should not panic");
1617        assert!(s.is_set());
1618    }
1619
1620    #[tokio::test]
1621    async fn shutdown_multiple_concurrent_waiters_all_wake() {
1622        // The notify_waiters() in signal() must wake every active waiter.
1623        let s = Arc::new(Shutdown::new());
1624        let mut handles = Vec::new();
1625        for _ in 0..16 {
1626            let s = Arc::clone(&s);
1627            handles.push(tokio::spawn(async move { s.wait().await }));
1628        }
1629        // Let waiters register.
1630        tokio::time::sleep(Duration::from_millis(20)).await;
1631
1632        s.signal();
1633
1634        for h in handles {
1635            tokio::time::timeout(Duration::from_millis(200), h)
1636                .await
1637                .expect("each waiter must wake within timeout")
1638                .expect("waiter task should not panic");
1639        }
1640    }
1641
1642    #[tokio::test]
1643    async fn shutdown_signal_is_idempotent() {
1644        // Second signal must be a no-op. Subsequent waits still return.
1645        let s = Shutdown::new();
1646        s.signal();
1647        s.signal();
1648        s.signal();
1649        tokio::time::timeout(Duration::from_millis(100), s.wait())
1650            .await
1651            .expect("wait must still return on idempotent re-signal");
1652    }
1653
1654    /// Contract test: the bounded-drain pattern in `serve_loop_graceful`
1655    /// (and the caller-side hammer for `serve_daemon_socket`) relies on
1656    /// `JoinSet::abort_all()` actually causing in-flight tasks to wake
1657    /// with a cancellation error, so a subsequent `join_next` loop
1658    /// completes. If tokio ever changes this — e.g., requires polling
1659    /// each task explicitly — our drain-timeout fallback silently
1660    /// regresses to "wait forever after abort_all".
1661    #[tokio::test]
1662    async fn joinset_abort_all_makes_drain_finite() {
1663        let mut set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
1664        // Spawn a task that would otherwise run for a long time.
1665        set.spawn(async {
1666            tokio::time::sleep(Duration::from_secs(60)).await;
1667        });
1668
1669        // First drain attempt: time out (task is mid-sleep).
1670        let primary = tokio::time::timeout(Duration::from_millis(100), async {
1671            while set.join_next().await.is_some() {}
1672        })
1673        .await;
1674        assert!(
1675            primary.is_err(),
1676            "primary drain should time out while task is still sleeping"
1677        );
1678
1679        // Now abort and drain again — must complete promptly.
1680        set.abort_all();
1681        let secondary = tokio::time::timeout(Duration::from_millis(500), async {
1682            while set.join_next().await.is_some() {}
1683        })
1684        .await;
1685        assert!(
1686            secondary.is_ok(),
1687            "drain after abort_all must complete quickly"
1688        );
1689        assert!(set.is_empty(), "JoinSet should be empty after drain");
1690    }
1691
1692    /// Contract test: the panic-detection logic in `serve_daemon_socket`
1693    /// (and `cli::daemon::serve_loop_graceful`) relies on tokio's `JoinSet`
1694    /// reporting panicked tasks via `try_join_next() -> Some(Err(e))` with
1695    /// `e.is_panic() == true`. If tokio ever changes that, our handler-
1696    /// panic-is-terminal property silently regresses. Lock it down here.
1697    #[tokio::test]
1698    async fn joinset_panics_are_observable_via_try_join_next() {
1699        let mut set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
1700        set.spawn(async {
1701            panic!("simulated handler panic");
1702        });
1703
1704        // Wait until the panicked task has been catch_unwind'd at the
1705        // tokio spawn boundary and parked on the JoinSet's completion queue.
1706        // Poll try_join_next briefly; assert we see the panic.
1707        let deadline = std::time::Instant::now() + Duration::from_millis(500);
1708        loop {
1709            if let Some(res) = set.try_join_next() {
1710                let err = res.expect_err("panicked task should yield Err");
1711                assert!(
1712                    err.is_panic(),
1713                    "JoinError must report is_panic for panicking task; got: {err:?}"
1714                );
1715                return;
1716            }
1717            if std::time::Instant::now() >= deadline {
1718                panic!("try_join_next never reported the panic within 500ms");
1719            }
1720            tokio::time::sleep(Duration::from_millis(10)).await;
1721        }
1722    }
1723
1724    /// Race contract — exercises the enable() pattern. A waiter that is
1725    /// JUST being constructed (between the `notified()` call and the flag
1726    /// check) must NOT miss a `signal()` that fires concurrently.
1727    ///
1728    /// Probabilistic: runs many trials and asserts every one wakes.
1729    #[tokio::test]
1730    async fn shutdown_no_lost_signal_under_race() {
1731        for trial in 0..50 {
1732            let s = Arc::new(Shutdown::new());
1733            let s_waiter = Arc::clone(&s);
1734            let s_signaler = Arc::clone(&s);
1735
1736            let waiter = tokio::spawn(async move { s_waiter.wait().await });
1737
1738            // Yield briefly so the waiter has a chance to start `wait()`.
1739            tokio::task::yield_now().await;
1740
1741            // Signal at the moment the waiter is racing to register.
1742            s_signaler.signal();
1743
1744            tokio::time::timeout(Duration::from_millis(500), waiter)
1745                .await
1746                .unwrap_or_else(|_| panic!("trial {trial}: waiter stranded by lost signal"))
1747                .expect("waiter task should not panic");
1748        }
1749    }
1750}
1751
1752#[cfg(test)]
1753mod tests {
1754    use super::*;
1755    use crate::store::record::{
1756        Category, ConfidenceScore, FileRecord, GotchaRecord, Priority, QualityScore, Record,
1757        RecordLifecycle, RecordSource, RecordVersion, StalenessScore,
1758    };
1759    use crate::store::Store;
1760
1761    fn make_gotcha_record(key: &str, files: &[&str]) -> Record {
1762        let gotcha = GotchaRecord {
1763            rule: "test rule".into(),
1764            reason: "test reason".into(),
1765            severity: Priority::High,
1766            affected_files: files.iter().map(|s| s.to_string()).collect(),
1767            ref_url: None,
1768            discovered_session: 1_000_000,
1769            confirmed: true,
1770        };
1771        Record {
1772            key: key.to_string(),
1773            value: "test rule because test reason".into(),
1774            payload: serde_json::to_value(&gotcha).ok(),
1775            category: Category::Gotcha,
1776            priority: Priority::High,
1777            tags: vec![],
1778            created_at: 1_000_000,
1779            updated_at: 1_000_000,
1780            ref_url: None,
1781            staleness: StalenessScore::fresh(),
1782            lifecycle: RecordLifecycle::Active,
1783            version: RecordVersion {
1784                device_id: uuid::Uuid::new_v4(),
1785                logical_clock: 1,
1786                wall_clock: 1_000_000,
1787            },
1788            quality: QualityScore::layer0_default(),
1789            access_count: 0,
1790            last_accessed: 0,
1791            source: RecordSource::DeveloperManual,
1792            confidence: ConfidenceScore::for_new_record(&RecordSource::DeveloperManual),
1793            gap_analysis_score: 0.0,
1794        }
1795    }
1796
1797    fn make_file_record(path: &str) -> Record {
1798        let file = FileRecord {
1799            path: path.to_string(),
1800            purpose: String::new(),
1801            entry_points: vec![],
1802            imports: vec![],
1803            gotcha_keys: vec![],
1804            decision_keys: vec![],
1805            todos: vec![],
1806            unsafe_count: 0,
1807            unwrap_count: 0,
1808            change_frequency: 0,
1809            last_author: None,
1810            is_hotspot: false,
1811            token_cost_estimate: 0,
1812            last_modified_session: 0,
1813            content_hash: None,
1814            line_count: 0,
1815            blast_radius: None,
1816            propagated_staleness: None,
1817        };
1818        Record {
1819            key: format!("file:{path}"),
1820            value: String::new(),
1821            payload: serde_json::to_value(&file).ok(),
1822            category: Category::File,
1823            priority: Priority::Normal,
1824            tags: vec![],
1825            created_at: 1_000_000,
1826            updated_at: 1_000_000,
1827            ref_url: None,
1828            staleness: StalenessScore::fresh(),
1829            lifecycle: RecordLifecycle::Active,
1830            version: RecordVersion {
1831                device_id: uuid::Uuid::new_v4(),
1832                logical_clock: 1,
1833                wall_clock: 1_000_000,
1834            },
1835            quality: QualityScore::layer0_default(),
1836            access_count: 0,
1837            last_accessed: 0,
1838            source: RecordSource::StaticAnalysis,
1839            confidence: ConfidenceScore::for_new_record(&RecordSource::StaticAnalysis),
1840            gap_analysis_score: 0.0,
1841        }
1842    }
1843
1844    fn file_gotcha_keys(record: &Record) -> Vec<String> {
1845        record
1846            .payload
1847            .as_ref()
1848            .and_then(|p| p.get("gotcha_keys"))
1849            .and_then(|v| v.as_array())
1850            .map(|arr| {
1851                arr.iter()
1852                    .filter_map(|v| v.as_str().map(String::from))
1853                    .collect()
1854            })
1855            .unwrap_or_default()
1856    }
1857
1858    /// Test helper: wraps a Store in a Graph + Arc for socket_dispatch.
1859    ///
1860    /// Consumes the Store (Graph owns it). Returns the Arc and a reference
1861    /// to access the store through the graph for assertions.
1862    async fn make_test_graph(store: Store) -> Arc<tokio::sync::RwLock<Graph>> {
1863        let graph = Graph::load(store).await.expect("failed to load test graph");
1864        Arc::new(tokio::sync::RwLock::new(graph))
1865    }
1866
1867    async fn dispatch_with_graph(
1868        graph: &Arc<tokio::sync::RwLock<Graph>>,
1869        cmd: &str,
1870        args: serde_json::Value,
1871    ) -> SocketResponse {
1872        let req = SocketRequest {
1873            cmd: cmd.to_string(),
1874            version: Some(PROTOCOL_VERSION),
1875            args,
1876        };
1877        socket_dispatch(graph, Path::new("/tmp/mati-test"), &req).await
1878    }
1879
1880    // ── Regression: gotcha_write via socket syncs file links ─────────────
1881
1882    #[tokio::test]
1883    async fn socket_gotcha_write_adds_keys_to_file_records() {
1884        let dir = tempfile::TempDir::new().unwrap();
1885        let store = Store::open(dir.path()).await.unwrap();
1886        store
1887            .put("file:src/a.rs", &make_file_record("src/a.rs"))
1888            .await
1889            .unwrap();
1890        store
1891            .put("file:src/b.rs", &make_file_record("src/b.rs"))
1892            .await
1893            .unwrap();
1894        let graph = make_test_graph(store).await;
1895
1896        let record = make_gotcha_record("gotcha:socket-test", &["src/a.rs", "src/b.rs"]);
1897        let resp = dispatch_with_graph(&graph, "gotcha_write", serde_json::json!({
1898            "record": record, "new_files": ["src/a.rs", "src/b.rs"], "old_files": [], "is_new": true,
1899        })).await;
1900        assert!(resp.ok, "gotcha_write failed: {:?}", resp.error);
1901
1902        let g = graph.read().await;
1903        let a = g.store().get("file:src/a.rs").await.unwrap().unwrap();
1904        let b = g.store().get("file:src/b.rs").await.unwrap().unwrap();
1905        assert!(file_gotcha_keys(&a).contains(&"gotcha:socket-test".into()));
1906        assert!(file_gotcha_keys(&b).contains(&"gotcha:socket-test".into()));
1907    }
1908
1909    #[tokio::test]
1910    async fn socket_gotcha_write_edit_removes_key_from_old_file() {
1911        let dir = tempfile::TempDir::new().unwrap();
1912        let store = Store::open(dir.path()).await.unwrap();
1913        store
1914            .put("file:src/a.rs", &make_file_record("src/a.rs"))
1915            .await
1916            .unwrap();
1917        store
1918            .put("file:src/b.rs", &make_file_record("src/b.rs"))
1919            .await
1920            .unwrap();
1921        let graph = make_test_graph(store).await;
1922
1923        let record = make_gotcha_record("gotcha:edit-socket", &["src/a.rs"]);
1924        let resp = dispatch_with_graph(
1925            &graph,
1926            "gotcha_write",
1927            serde_json::json!({
1928                "record": record, "new_files": ["src/a.rs"], "old_files": [], "is_new": true,
1929            }),
1930        )
1931        .await;
1932        assert!(resp.ok);
1933
1934        let record2 = make_gotcha_record("gotcha:edit-socket", &["src/b.rs"]);
1935        let resp2 = dispatch_with_graph(&graph, "gotcha_write", serde_json::json!({
1936            "record": record2, "new_files": ["src/b.rs"], "old_files": ["src/a.rs"], "is_new": false,
1937        })).await;
1938        assert!(resp2.ok);
1939
1940        let g = graph.read().await;
1941        let a = g.store().get("file:src/a.rs").await.unwrap().unwrap();
1942        let b = g.store().get("file:src/b.rs").await.unwrap().unwrap();
1943        assert!(!file_gotcha_keys(&a).contains(&"gotcha:edit-socket".into()));
1944        assert!(file_gotcha_keys(&b).contains(&"gotcha:edit-socket".into()));
1945    }
1946
1947    #[tokio::test]
1948    async fn socket_gotcha_tombstone_removes_keys_from_file_records() {
1949        let dir = tempfile::TempDir::new().unwrap();
1950        let store = Store::open(dir.path()).await.unwrap();
1951        store
1952            .put("file:src/a.rs", &make_file_record("src/a.rs"))
1953            .await
1954            .unwrap();
1955        store
1956            .put("file:src/b.rs", &make_file_record("src/b.rs"))
1957            .await
1958            .unwrap();
1959        let graph = make_test_graph(store).await;
1960
1961        let record = make_gotcha_record("gotcha:tomb-socket", &["src/a.rs", "src/b.rs"]);
1962        let resp = dispatch_with_graph(&graph, "gotcha_write", serde_json::json!({
1963            "record": record, "new_files": ["src/a.rs", "src/b.rs"], "old_files": [], "is_new": true,
1964        })).await;
1965        assert!(resp.ok);
1966
1967        let resp2 = dispatch_with_graph(
1968            &graph,
1969            "gotcha_tombstone",
1970            serde_json::json!({
1971                "key": "gotcha:tomb-socket", "affected_files": ["src/a.rs", "src/b.rs"],
1972            }),
1973        )
1974        .await;
1975        assert!(resp2.ok, "gotcha_tombstone failed: {:?}", resp2.error);
1976
1977        let g = graph.read().await;
1978        let rec = g.store().get("gotcha:tomb-socket").await.unwrap().unwrap();
1979        assert!(matches!(rec.lifecycle, RecordLifecycle::Tombstoned { .. }));
1980        let a = g.store().get("file:src/a.rs").await.unwrap().unwrap();
1981        let b = g.store().get("file:src/b.rs").await.unwrap().unwrap();
1982        assert!(file_gotcha_keys(&a).is_empty());
1983        assert!(file_gotcha_keys(&b).is_empty());
1984    }
1985
1986    #[tokio::test]
1987    async fn socket_gotcha_write_rejects_duplicate_key() {
1988        let dir = tempfile::TempDir::new().unwrap();
1989        let store = Store::open(dir.path()).await.unwrap();
1990        let record1 = make_gotcha_record("gotcha:dup-socket", &["src/a.rs"]);
1991        store.put("gotcha:dup-socket", &record1).await.unwrap();
1992        let graph = make_test_graph(store).await;
1993
1994        let record2 = make_gotcha_record("gotcha:dup-socket", &["src/b.rs"]);
1995        let resp = dispatch_with_graph(
1996            &graph,
1997            "gotcha_write",
1998            serde_json::json!({
1999                "record": record2, "new_files": ["src/b.rs"], "old_files": [], "is_new": true,
2000            }),
2001        )
2002        .await;
2003        assert!(!resp.ok, "duplicate key should be rejected");
2004        assert!(resp
2005            .error
2006            .as_deref()
2007            .unwrap_or("")
2008            .contains("already exists"));
2009
2010        let g = graph.read().await;
2011        let original = g.store().get("gotcha:dup-socket").await.unwrap().unwrap();
2012        let payload = original.payload_as::<GotchaRecord>().unwrap();
2013        assert_eq!(payload.affected_files, vec!["src/a.rs"]);
2014    }
2015
2016    // ── Wire-level size enforcement ────────────────────────────────────
2017
2018    #[tokio::test]
2019    async fn oversized_request_returns_frame_too_large_with_response() {
2020        use super::super::protocol::MAX_FRAME_SIZE;
2021        use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
2022
2023        let dir = tempfile::TempDir::new().unwrap();
2024        let store = Store::open(dir.path()).await.unwrap();
2025        let graph = make_test_graph(store).await;
2026
2027        let (client, server) = UnixStream::pair().unwrap();
2028        let peer = super::super::metadata::PeerContext {
2029            uid: 501,
2030            pid: None,
2031        };
2032
2033        // Payload larger than MAX_FRAME_SIZE.
2034        let oversized = "x".repeat(MAX_FRAME_SIZE + 100);
2035        let payload = format!("{oversized}\n");
2036
2037        // Split client: write oversized request, then read response.
2038        let (client_read, client_write) = client.into_split();
2039
2040        let write_handle = tokio::spawn(async move {
2041            let mut w = client_write;
2042            w.write_all(payload.as_bytes()).await.unwrap();
2043            w.shutdown().await.unwrap();
2044        });
2045
2046        let handle_result =
2047            socket_handle_connection(graph, dir.path(), server, peer, uuid::Uuid::nil()).await;
2048        assert!(handle_result.is_ok());
2049
2050        write_handle.await.unwrap();
2051
2052        // Read the error response from the server.
2053        let mut reader = tokio::io::BufReader::new(client_read);
2054        let mut line = String::new();
2055        reader.read_line(&mut line).await.unwrap();
2056        let resp: serde_json::Value = serde_json::from_str(line.trim()).unwrap();
2057
2058        assert_eq!(resp["status"], "err");
2059        assert_eq!(resp["code"], "frame_too_large");
2060        assert!(
2061            resp["message"]
2062                .as_str()
2063                .unwrap()
2064                .contains(&MAX_FRAME_SIZE.to_string()),
2065            "error message should mention the size limit"
2066        );
2067    }
2068
2069    #[tokio::test]
2070    async fn normal_sized_request_is_not_rejected_by_size_check() {
2071        use super::super::protocol::MAX_FRAME_SIZE;
2072        use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
2073
2074        let dir = tempfile::TempDir::new().unwrap();
2075        let store = Store::open(dir.path()).await.unwrap();
2076        let graph = make_test_graph(store).await;
2077
2078        let (client, server) = UnixStream::pair().unwrap();
2079        let peer = super::super::metadata::PeerContext {
2080            uid: 501,
2081            pid: None,
2082        };
2083
2084        // A valid v2 ping request — well under MAX_FRAME_SIZE.
2085        let request = serde_json::json!({
2086            "v": 2,
2087            "id": uuid::Uuid::new_v4(),
2088            "session": uuid::Uuid::nil(),
2089            "cmd": { "type": "ping" }
2090        });
2091        let payload = format!("{}\n", serde_json::to_string(&request).unwrap());
2092        assert!(
2093            payload.len() < MAX_FRAME_SIZE,
2094            "test payload should be small"
2095        );
2096
2097        let (client_read, client_write) = client.into_split();
2098
2099        let write_handle = tokio::spawn(async move {
2100            let mut w = client_write;
2101            w.write_all(payload.as_bytes()).await.unwrap();
2102            w.shutdown().await.unwrap();
2103        });
2104
2105        let handle_result =
2106            socket_handle_connection(graph, dir.path(), server, peer, uuid::Uuid::nil()).await;
2107        assert!(handle_result.is_ok());
2108
2109        write_handle.await.unwrap();
2110
2111        // Read response — should be a successful pong, not FrameTooLarge.
2112        let mut reader = tokio::io::BufReader::new(client_read);
2113        let mut line = String::new();
2114        reader.read_line(&mut line).await.unwrap();
2115        let resp: serde_json::Value = serde_json::from_str(line.trim()).unwrap();
2116
2117        assert_eq!(resp["status"], "ok", "ping should succeed, got: {resp}");
2118    }
2119
2120    // ── Daemon-restart resilience ──────────────────────────────────────
2121    //
2122    // Regression for the smoke-test failure: after a daemon stop+start,
2123    // the MCP-stdio bridge sees `session_mismatch` (or transient
2124    // `Unresponsive`) on the first call because its cached daemon session
2125    // UUID predates the restart. Without retry, every subsequent
2126    // mem_get/mem_query/mem_bootstrap/mem_set returns a structured error
2127    // that effectively wedges the agent's MCP session.
2128    //
2129    // The fix in `proxy_daemon_result` is one bounded auto-reconnect: the
2130    // helper re-reads daemon metadata fresh (picking up the new session
2131    // UUID) and re-issues the request. This test asserts the reconnect
2132    // succeeds end-to-end and DOES NOT propagate the session_mismatch
2133    // error envelope to the caller.
2134
2135    /// Spawn a tiny daemon-substitute that binds the given socket and
2136    /// answers each connection with the supplied JSON response (one line),
2137    /// then closes the connection. Returns the JoinHandle so the test can
2138    /// await it.
2139    async fn spawn_canned_responder(
2140        sock_path: std::path::PathBuf,
2141        responses: Vec<serde_json::Value>,
2142    ) -> tokio::task::JoinHandle<()> {
2143        // Bind in this task synchronously so the caller can issue
2144        // requests immediately without a sleep race.
2145        let listener = tokio::net::UnixListener::bind(&sock_path).expect("bind responder socket");
2146        tokio::spawn(async move {
2147            for resp in responses {
2148                let (stream, _) = match listener.accept().await {
2149                    Ok(s) => s,
2150                    Err(_) => return,
2151                };
2152                let (reader, mut writer) = stream.into_split();
2153                // Drain the request line so the peer's `shutdown()` returns Ok.
2154                let mut buf_reader = tokio::io::BufReader::new(reader);
2155                let mut line = String::new();
2156                let _ = tokio::io::AsyncBufReadExt::read_line(&mut buf_reader, &mut line).await;
2157                let mut bytes = serde_json::to_vec(&resp).unwrap();
2158                bytes.push(b'\n');
2159                let _ = tokio::io::AsyncWriteExt::write_all(&mut writer, &bytes).await;
2160                let _ = tokio::io::AsyncWriteExt::shutdown(&mut writer).await;
2161            }
2162        })
2163    }
2164
2165    #[tokio::test]
2166    async fn mcp_call_after_daemon_restart_does_not_kill_transport() {
2167        // Scenario: the proxy's first attempt hits a daemon whose session
2168        // UUID does not match (simulating a daemon restart between two
2169        // tool calls). The fix retries once, re-reads metadata, and the
2170        // second attempt succeeds.
2171
2172        let dir = tempfile::TempDir::new().unwrap();
2173        let root = dir.path().to_path_buf();
2174        let sock_path = root.join("mati.sock");
2175
2176        // Initial daemon session "before restart". The proxy will read
2177        // this UUID, but our canned responder pretends not to recognize
2178        // it (returning session_mismatch). After the retry delay, we
2179        // rotate metadata to a new UUID — exactly what `mati daemon stop`
2180        // + `mati daemon start` would do in production.
2181        let session_before = uuid::Uuid::new_v4();
2182        let session_after = uuid::Uuid::new_v4();
2183
2184        let meta_before = super::super::metadata::DaemonMetadata {
2185            pid: std::process::id(),
2186            session: session_before,
2187            owner: super::super::metadata::DaemonOwner::Daemon,
2188        };
2189        super::super::metadata::publish_metadata(&root, &meta_before).unwrap();
2190
2191        // Stage two responses on the same socket: the first is a
2192        // SessionMismatch err (pre-restart daemon view), the second is a
2193        // successful pong (post-restart daemon view).
2194        let responder_handle = spawn_canned_responder(
2195            sock_path.clone(),
2196            vec![
2197                serde_json::json!({
2198                    "v": 2,
2199                    "id": uuid::Uuid::new_v4(),
2200                    "status": "err",
2201                    "code": "session_mismatch",
2202                    "message": "session mismatch: re-read daemon metadata and retry",
2203                }),
2204                serde_json::json!({
2205                    "v": 2,
2206                    "id": uuid::Uuid::new_v4(),
2207                    "status": "ok",
2208                    "data": "pong",
2209                }),
2210            ],
2211        )
2212        .await;
2213
2214        // Concurrent metadata rotation — fires during the retry delay.
2215        // Mirrors what a real daemon restart does: writes fresh metadata.
2216        let root_for_rotate = root.clone();
2217        let rotate_handle = tokio::spawn(async move {
2218            // Sleep just less than the proxy's 100ms retry settle so the
2219            // metadata rewrite is committed before the second attempt.
2220            tokio::time::sleep(Duration::from_millis(20)).await;
2221            let meta_after = super::super::metadata::DaemonMetadata {
2222                pid: std::process::id(),
2223                session: session_after,
2224                owner: super::super::metadata::DaemonOwner::Daemon,
2225            };
2226            super::super::metadata::publish_metadata(&root_for_rotate, &meta_after).unwrap();
2227        });
2228
2229        // Wrap in a tokio timeout: if the retry path is missing, the
2230        // proxy returns the first attempt's envelope without ever
2231        // dialing the second responder, which would leave the test
2232        // hanging on the spare canned response. The timeout converts
2233        // that latent hang into a deterministic failure with a clear
2234        // error message.
2235        let result = tokio::time::timeout(
2236            Duration::from_secs(5),
2237            super::proxy_daemon_result(&root, "ping", serde_json::json!({})),
2238        )
2239        .await
2240        .expect("proxy_daemon_result should resolve within 5s — retry path appears wedged");
2241
2242        rotate_handle.await.unwrap();
2243        // Drop the responder task — the second canned response may go
2244        // unconsumed in failure modes. Aborting prevents the test from
2245        // hanging on `responder_handle.await` in failure mode.
2246        responder_handle.abort();
2247
2248        // The proxy must transparently recover: caller sees Ok, not the
2249        // session_mismatch error envelope from the first attempt.
2250        match result {
2251            super::ProxyDaemonResult::Ok(v) => {
2252                let ok = v.get("ok") == Some(&serde_json::Value::Bool(true));
2253                let code = v.get("code").and_then(|c| c.as_str()).unwrap_or("");
2254                assert!(
2255                    ok,
2256                    "second attempt should succeed after metadata rotation, \
2257                     but caller saw the first attempt's session_mismatch envelope: \
2258                     ok={ok} code={code:?} v={v}"
2259                );
2260            }
2261            other => panic!(
2262                "expected Ok(true) after auto-reconnect, got {other:?}; \
2263                 the daemon-restart retry path is not engaging"
2264            ),
2265        }
2266    }
2267
2268    #[tokio::test]
2269    async fn mcp_call_session_mismatch_no_retry_target_returns_envelope() {
2270        // Negative-side guard: if the second attempt also fails with the
2271        // same error (e.g. the daemon was not actually restarted), the
2272        // proxy still returns the structured error envelope to the
2273        // caller — it does NOT panic, hang, or close the rmcp transport.
2274        // This preserves the per-call structured-error discipline that
2275        // keeps Claude's MCP session alive.
2276
2277        let dir = tempfile::TempDir::new().unwrap();
2278        let root = dir.path().to_path_buf();
2279        let sock_path = root.join("mati.sock");
2280
2281        let session = uuid::Uuid::new_v4();
2282        let meta = super::super::metadata::DaemonMetadata {
2283            pid: std::process::id(),
2284            session,
2285            owner: super::super::metadata::DaemonOwner::Daemon,
2286        };
2287        super::super::metadata::publish_metadata(&root, &meta).unwrap();
2288
2289        // Both attempts get a session_mismatch — emulates a daemon that
2290        // truly cannot be reconciled (wedged in a state the proxy can't
2291        // recover from).
2292        let responder_handle = spawn_canned_responder(
2293            sock_path.clone(),
2294            vec![
2295                serde_json::json!({
2296                    "v": 2,
2297                    "id": uuid::Uuid::new_v4(),
2298                    "status": "err",
2299                    "code": "session_mismatch",
2300                    "message": "session mismatch (1)",
2301                }),
2302                serde_json::json!({
2303                    "v": 2,
2304                    "id": uuid::Uuid::new_v4(),
2305                    "status": "err",
2306                    "code": "session_mismatch",
2307                    "message": "session mismatch (2)",
2308                }),
2309            ],
2310        )
2311        .await;
2312
2313        let result = tokio::time::timeout(
2314            Duration::from_secs(5),
2315            super::proxy_daemon_result(&root, "ping", serde_json::json!({})),
2316        )
2317        .await
2318        .expect("proxy_daemon_result must resolve within 5s");
2319        responder_handle.abort();
2320
2321        // The caller MUST get a structured Ok envelope with ok:false +
2322        // the session_mismatch code, never a panic or transport-killing
2323        // surprise. socket_call (in tools.rs) renders this to a JSON
2324        // error string — which is exactly the contract the rmcp loop
2325        // expects: a String response, not a Result::Err.
2326        match result {
2327            super::ProxyDaemonResult::Ok(v) => {
2328                assert_eq!(v.get("ok"), Some(&serde_json::Value::Bool(false)));
2329                assert_eq!(
2330                    v.get("code").and_then(|c| c.as_str()),
2331                    Some("session_mismatch")
2332                );
2333            }
2334            other => panic!("expected structured Ok envelope, got {other:?}"),
2335        }
2336    }
2337
2338    // ── Pass-29 regression: proxy_daemon_result handles side-effecting reads ──
2339    //
2340    // Pre-fix: every Socket-backed `mem_get` and `mem_bootstrap` MCP call
2341    // panicked the rmcp task at `v1_to_v2_command` (no match arm), which
2342    // surfaced to the client as `Transport closed` and wedged Phases 6–17
2343    // of the smoke. The translation layer is the load-bearing artifact
2344    // — pass 27's mock-UnixListener test bypassed it entirely, so the
2345    // bug shipped.
2346    //
2347    // These tests drive `proxy_daemon_result` with the exact strings
2348    // tools.rs sends today. Without the new arms in v1_to_v2_command,
2349    // both panic. With the fix, both return a clean `NotRunning` because
2350    // the socket doesn't exist — proving the translation succeeded
2351    // before the connect attempt.
2352
2353    #[tokio::test]
2354    async fn proxy_daemon_result_handles_mem_get_translation_no_panic() {
2355        let dir = tempfile::TempDir::new().unwrap();
2356        // No socket file present — the call must reach the
2357        // sock_path.exists() guard, which it cannot do if v1_to_v2_command
2358        // panics first.
2359        let result = super::proxy_daemon_result(
2360            dir.path(),
2361            "mem_get",
2362            serde_json::json!({ "key": "file:src/main.rs" }),
2363        )
2364        .await;
2365        assert!(
2366            matches!(result, super::ProxyDaemonResult::NotRunning),
2367            "mem_get without daemon must return NotRunning, got {result:?}"
2368        );
2369    }
2370
2371    #[tokio::test]
2372    async fn proxy_daemon_result_handles_mem_bootstrap_translation_no_panic() {
2373        let dir = tempfile::TempDir::new().unwrap();
2374        let result = super::proxy_daemon_result(
2375            dir.path(),
2376            "mem_bootstrap",
2377            serde_json::json!({ "context_files": ["src/lib.rs"] }),
2378        )
2379        .await;
2380        assert!(
2381            matches!(result, super::ProxyDaemonResult::NotRunning),
2382            "mem_bootstrap without daemon must return NotRunning, got {result:?}"
2383        );
2384    }
2385
2386    #[tokio::test]
2387    async fn proxy_daemon_v2_typed_path_handles_mem_set_mutations_no_panic() {
2388        // The Socket-backend mem_set now takes the typed path. With no
2389        // daemon present, the typed-Command serialize→connect path must
2390        // surface as a clean NotRunning, never a panic. This is the
2391        // load-bearing fence: any future caller that accidentally routes
2392        // gotcha_upsert through the v1 mapper would fail
2393        // v1_to_v2_command_no_mutations_silently_accepted in protocol.rs;
2394        // here we make sure the typed path itself is wired end-to-end.
2395        let dir = tempfile::TempDir::new().unwrap();
2396        let cmd = super::super::protocol::Command::GotchaConfirm(
2397            super::super::protocol::GotchaConfirmInput {
2398                key: "gotcha:test".into(),
2399            },
2400        );
2401        let result = super::proxy_daemon_v2(dir.path(), cmd).await;
2402        assert!(
2403            matches!(result, super::ProxyDaemonResult::NotRunning),
2404            "typed proxy_daemon_v2 must return NotRunning when daemon is absent, got {result:?}"
2405        );
2406    }
2407}