1use 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
69pub async fn serve(repo_root: &Path) -> Result<()> {
91 let startup_t0 = std::time::Instant::now();
92
93 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 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 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 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 let result = proxy_daemon_result_no_spawn(root, cmd, &args).await;
192
193 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
216pub(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 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
237enum 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 let v2_cmd = super::protocol::v1_to_v2_command(cmd, args);
257 proxy_daemon_send_v2(root, v2_cmd).await
258}
259
260pub(crate) async fn proxy_daemon_v2(
270 root: &Path,
271 cmd: super::protocol::Command,
272) -> ProxyDaemonResult {
273 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 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
305async 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 return AttemptOutcome::Final(ProxyDaemonResult::NotRunning);
318 }
319
320 if !sock_path.exists() {
321 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 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 let _ = std::fs::remove_file(&sock_path);
340 return AttemptOutcome::Retryable(ProxyDaemonResult::StaleSocket);
341 }
342 StaleCheckResult::LiveDaemon { .. } => {
343 return AttemptOutcome::Retryable(ProxyDaemonResult::Unresponsive);
345 }
346 }
347 }
348 return AttemptOutcome::Retryable(ProxyDaemonResult::NotRunning);
349 }
350 };
351
352 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 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 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
427pub const UNIX_SOCK_PATH_MAX: usize = 104;
437
438const READ_TIMEOUT: Duration = Duration::from_secs(3);
440
441pub const MAX_CONCURRENT_CONNECTIONS: usize = 64;
448
449pub const AUTO_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
456
457#[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 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 self.flag.load(std::sync::atomic::Ordering::SeqCst)
493 }
494
495 pub async fn wait(&self) {
498 let notified = self.notify.notified();
499 tokio::pin!(notified);
500 notified.as_mut().enable();
503 if self.is_set() {
504 return;
505 }
506 notified.await;
507 }
508}
509
510const PROTOCOL_VERSION: u32 = 1;
512
513#[derive(Debug, Deserialize)]
514pub(crate) struct SocketRequest {
515 pub cmd: String,
516 #[allow(dead_code)] #[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 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 let v2_req = match serde_json::from_str::<super::protocol::Request>(trimmed) {
599 Ok(r) => r,
600 Err(e) => {
601 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
629fn 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 "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 "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(), ¶ms).await;
763 SocketResponse::ok(serde_json::Value::String(response))
764 }
765
766 "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 "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 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 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 for nkey in g.neighbors(file_key, &crate::graph::EdgeKind::HasGotcha) {
865 linked_keys.insert(nkey);
866 }
867
868 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 if !matches!(grec.lifecycle, crate::store::RecordLifecycle::Active) {
894 continue;
895 }
896 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 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 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 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 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 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 {
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 keys_to_update.extend(g.neighbors_incoming(&file_key, &EdgeKind::Imports));
1158 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 {
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 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 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 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 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 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 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 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 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 crate::store::gotcha_ops::propagate_confirmation_to_files(store, &affected_files).await;
1547
1548 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
1570pub const IDLE_SHUTDOWN_SECS: u64 = 30 * 60; pub const IDLE_CHECK_INTERVAL_SECS: u64 = 5 * 60; #[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 let s = Shutdown::new();
1593 s.signal();
1594 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 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 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 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 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 #[tokio::test]
1662 async fn joinset_abort_all_makes_drain_finite() {
1663 let mut set: tokio::task::JoinSet<()> = tokio::task::JoinSet::new();
1664 set.spawn(async {
1666 tokio::time::sleep(Duration::from_secs(60)).await;
1667 });
1668
1669 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 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 #[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 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 #[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 tokio::task::yield_now().await;
1740
1741 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 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 #[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 #[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 let oversized = "x".repeat(MAX_FRAME_SIZE + 100);
2035 let payload = format!("{oversized}\n");
2036
2037 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 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 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 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 async fn spawn_canned_responder(
2140 sock_path: std::path::PathBuf,
2141 responses: Vec<serde_json::Value>,
2142 ) -> tokio::task::JoinHandle<()> {
2143 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 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 let dir = tempfile::TempDir::new().unwrap();
2173 let root = dir.path().to_path_buf();
2174 let sock_path = root.join("mati.sock");
2175
2176 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 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 let root_for_rotate = root.clone();
2217 let rotate_handle = tokio::spawn(async move {
2218 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 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 responder_handle.abort();
2247
2248 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 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 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 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 #[tokio::test]
2354 async fn proxy_daemon_result_handles_mem_get_translation_no_panic() {
2355 let dir = tempfile::TempDir::new().unwrap();
2356 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 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}