1use std::{
2 collections::HashSet,
3 path::{Path, PathBuf},
4 sync::atomic::{AtomicU64, Ordering},
5 time::{Duration, SystemTime, UNIX_EPOCH},
6};
7
8use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
9use serde::{Deserialize, Serialize, de::DeserializeOwned};
10
11use crate::{
12 agent::{AgentConfig, AgentStatus, SpawnedAgentSummary, TeammateIdentity},
13 background::{
14 BackgroundNotification, BackgroundStore, BackgroundTaskStatus, BackgroundTaskSummary,
15 },
16 memory::journal::AgentMemoryState,
17 memory::{MemoryCursor, MemoryRecord, MemorySearchRequest, MemoryStore},
18 provider::ProviderId,
19 runtime::TaskItem,
20 session::PermissionRuleScope,
21 session::permission::{RememberedRule, RuleKey},
22 team::{TeamMemberSummary, TeamMessage, TeamProtocolRequestSummary, TeamStore},
23};
24
25use super::error::RuntimeError;
26
27static NEXT_STORE_ID: AtomicU64 = AtomicU64::new(1);
28#[cfg(test)]
29static NEXT_TEST_STORE_ID: AtomicU64 = AtomicU64::new(1);
30
31const DELIVERY_PENDING: i64 = 0;
32const DELIVERY_INFLIGHT: i64 = 1;
33const DELIVERY_ACKED: i64 = 2;
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PersistedAgentRecord {
37 pub(crate) id: String,
38 pub(crate) runtime_identifier: String,
39 pub(crate) name: String,
40 pub(crate) model: String,
41 pub(crate) provider_id: ProviderId,
42 pub(crate) config: AgentConfig,
43 pub(crate) hidden_tools: HashSet<String>,
44 pub(crate) max_rounds: Option<usize>,
45 pub(crate) teammate_identity: Option<TeammateIdentity>,
46 pub(crate) rounds_since_task: usize,
47 pub(crate) idle_requested: bool,
48 pub(crate) status: AgentStatus,
49 pub(crate) subagents: Vec<SpawnedAgentSummary>,
50}
51
52#[derive(Debug, Clone)]
53pub struct LoadedAgentState {
54 pub(crate) record: PersistedAgentRecord,
55 pub(crate) memory: AgentMemoryState,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct TaskStateSnapshot {
60 pub(crate) tasks: Vec<TaskItem>,
61}
62
63pub trait AgentStore: Send + Sync {
68 fn allows_disk_artifacts(&self) -> bool {
75 true
76 }
77
78 fn prepare_recovery(&self) -> Result<(), RuntimeError>;
79 fn create_agent(
80 &self,
81 record: &PersistedAgentRecord,
82 memory: &AgentMemoryState,
83 ) -> Result<(), RuntimeError>;
84 fn save_agent_record(&self, record: &PersistedAgentRecord) -> Result<(), RuntimeError>;
85 fn save_agent_memory(
86 &self,
87 agent_id: &str,
88 memory: &AgentMemoryState,
89 ) -> Result<(), RuntimeError>;
90 fn load_agent(&self, agent_id: &str) -> Result<Option<LoadedAgentState>, RuntimeError>;
91 fn list_agents(&self) -> Result<Vec<LoadedAgentState>, RuntimeError>;
92 fn list_agents_by_runtime(
93 &self,
94 runtime_identifier: &str,
95 ) -> Result<Vec<LoadedAgentState>, RuntimeError>;
96}
97
98pub trait RunStore: Send + Sync {
102 fn start_run(&self, agent_id: &str) -> Result<String, RuntimeError>;
103 fn update_run_state(
104 &self,
105 run_id: &str,
106 state: &str,
107 error: Option<&str>,
108 ) -> Result<(), RuntimeError>;
109 fn finish_run(&self, run_id: &str) -> Result<(), RuntimeError>;
110 fn fail_run(&self, run_id: &str, error: &str) -> Result<(), RuntimeError>;
111}
112
113pub trait TaskStore: Send + Sync {
118 fn load_tasks(&self, namespace: &Path) -> Result<Vec<TaskItem>, RuntimeError>;
119 fn capture_tasks(&self, namespace: &Path) -> Result<TaskStateSnapshot, RuntimeError>;
120 fn restore_tasks(
121 &self,
122 namespace: &Path,
123 snapshot: &TaskStateSnapshot,
124 ) -> Result<(), RuntimeError>;
125 fn replace_tasks(&self, namespace: &Path, tasks: &[TaskItem]) -> Result<(), RuntimeError>;
126
127 fn mutate(
139 &self,
140 namespace: &Path,
141 mutation: &mut dyn FnMut(&mut Vec<TaskItem>) -> Result<(), RuntimeError>,
142 ) -> Result<(), RuntimeError> {
143 let mut tasks = self.load_tasks(namespace)?;
144 mutation(&mut tasks)?;
145 self.replace_tasks(namespace, &tasks)
146 }
147}
148
149pub trait AuditStore: Send + Sync {
151 fn record_audit_event(
152 &self,
153 scope: &str,
154 event_type: &str,
155 payload: serde_json::Value,
156 ) -> Result<(), RuntimeError>;
157}
158
159pub trait LeaseStore: Send + Sync {
164 fn acquire_lease(&self, key: &str, owner: &str, ttl: Duration) -> Result<bool, RuntimeError>;
165 fn release_lease(&self, key: &str, owner: &str) -> Result<(), RuntimeError>;
166}
167
168pub trait PermissionRuleStore: Send + Sync {
176 fn save_rules(
181 &self,
182 session_id: &str,
183 project_id: Option<&str>,
184 rules: &[RememberedRule],
185 ) -> Result<(), RuntimeError>;
186
187 fn load_rules(
194 &self,
195 session_id: &str,
196 project_id: Option<&str>,
197 ) -> Result<Vec<RememberedRule>, RuntimeError>;
198
199 fn clear_rules(&self, session_id: &str) -> Result<(), RuntimeError>;
201}
202
203pub trait RuntimeStore:
209 AgentStore
210 + RunStore
211 + TaskStore
212 + AuditStore
213 + LeaseStore
214 + PermissionRuleStore
215 + TeamStore
216 + BackgroundStore
217 + MemoryStore
218 + Send
219 + Sync
220{
221}
222
223impl<T> RuntimeStore for T where
224 T: AgentStore
225 + RunStore
226 + TaskStore
227 + AuditStore
228 + LeaseStore
229 + PermissionRuleStore
230 + TeamStore
231 + BackgroundStore
232 + MemoryStore
233 + Send
234 + Sync
235{
236}
237
238impl TeamStore for SqliteRuntimeStore {
239 fn unread_team_count(&self, team_dir: &Path, agent_name: &str) -> Result<usize, RuntimeError> {
240 let conn = self.open()?;
241 let count = conn
242 .query_row(
243 "SELECT COUNT(*) FROM team_inbox WHERE team_dir = ?1 AND recipient = ?2 AND delivery_state = ?3",
244 params![Self::team_key(team_dir), agent_name, DELIVERY_PENDING],
245 |row| row.get::<_, i64>(0),
246 )
247 .map_err(sqlite_error)?;
248 Ok(count as usize)
249 }
250
251 fn load_team_members(&self, team_dir: &Path) -> Result<Vec<TeamMemberSummary>, RuntimeError> {
252 let conn = self.open()?;
253 let mut stmt = conn
254 .prepare("SELECT summary_json FROM team_members WHERE team_dir = ?1 ORDER BY name")
255 .map_err(sqlite_error)?;
256 let rows = stmt
257 .query_map(params![Self::team_key(team_dir)], |row| {
258 row.get::<_, String>(0)
259 })
260 .map_err(sqlite_error)?;
261 let mut members = Vec::new();
262 for row in rows {
263 members.push(from_json(&row.map_err(sqlite_error)?)?);
264 }
265 Ok(members)
266 }
267
268 fn upsert_team_member(
269 &self,
270 team_dir: &Path,
271 summary: &TeamMemberSummary,
272 ) -> Result<(), RuntimeError> {
273 let conn = self.open()?;
274 conn.execute(
275 r#"
276 INSERT INTO team_members (team_dir, name, summary_json)
277 VALUES (?1, ?2, ?3)
278 ON CONFLICT(team_dir, name) DO UPDATE SET summary_json = excluded.summary_json
279 "#,
280 params![Self::team_key(team_dir), summary.name, to_json(summary)?],
281 )
282 .map_err(sqlite_error)?;
283 Ok(())
284 }
285
286 fn read_team_inbox(
287 &self,
288 team_dir: &Path,
289 agent_name: &str,
290 ) -> Result<Vec<TeamMessage>, RuntimeError> {
291 let mut conn = self.open()?;
292 let tx = conn
293 .transaction_with_behavior(TransactionBehavior::Immediate)
294 .map_err(sqlite_error)?;
295 let team_key = Self::team_key(team_dir);
296 let ids_and_payloads = {
297 let mut stmt = tx
298 .prepare(
299 "SELECT id, payload_json FROM team_inbox WHERE team_dir = ?1 AND recipient = ?2 AND delivery_state = ?3 ORDER BY created_at, id",
300 )
301 .map_err(sqlite_error)?;
302 stmt.query_map(params![team_key, agent_name, DELIVERY_PENDING], |row| {
303 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
304 })
305 .map_err(sqlite_error)?
306 .collect::<Result<Vec<_>, _>>()
307 .map_err(sqlite_error)?
308 };
309
310 for (id, _) in &ids_and_payloads {
311 tx.execute(
312 "UPDATE team_inbox SET delivery_state = ?2 WHERE id = ?1",
313 params![id, DELIVERY_INFLIGHT],
314 )
315 .map_err(sqlite_error)?;
316 }
317 tx.commit().map_err(sqlite_error)?;
318
319 ids_and_payloads
320 .into_iter()
321 .map(|(_, payload)| from_json(&payload))
322 .collect()
323 }
324
325 fn ack_team_inbox(&self, team_dir: &Path, agent_name: &str) -> Result<(), RuntimeError> {
326 let conn = self.open()?;
327 conn.execute(
328 "UPDATE team_inbox SET delivery_state = ?3 WHERE team_dir = ?1 AND recipient = ?2 AND delivery_state = ?4",
329 params![Self::team_key(team_dir), agent_name, DELIVERY_ACKED, DELIVERY_INFLIGHT],
330 )
331 .map_err(sqlite_error)?;
332 Ok(())
333 }
334
335 fn requeue_team_inbox(&self, team_dir: &Path, agent_name: &str) -> Result<(), RuntimeError> {
336 let conn = self.open()?;
337 conn.execute(
338 "UPDATE team_inbox SET delivery_state = ?3 WHERE team_dir = ?1 AND recipient = ?2 AND delivery_state = ?4",
339 params![Self::team_key(team_dir), agent_name, DELIVERY_PENDING, DELIVERY_INFLIGHT],
340 )
341 .map_err(sqlite_error)?;
342 Ok(())
343 }
344
345 fn append_team_message(
346 &self,
347 team_dir: &Path,
348 recipient: &str,
349 message: &TeamMessage,
350 ) -> Result<(), RuntimeError> {
351 let conn = self.open()?;
352 conn.execute(
353 "INSERT INTO team_inbox (id, team_dir, recipient, payload_json, delivery_state, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
354 params![
355 next_id("teammsg"),
356 Self::team_key(team_dir),
357 recipient,
358 to_json(message)?,
359 DELIVERY_PENDING,
360 now_secs(),
361 ],
362 )
363 .map_err(sqlite_error)?;
364 Ok(())
365 }
366
367 fn load_team_requests(
368 &self,
369 team_dir: &Path,
370 ) -> Result<Vec<TeamProtocolRequestSummary>, RuntimeError> {
371 let conn = self.open()?;
372 let mut stmt = conn
373 .prepare(
374 "SELECT payload_json FROM team_requests WHERE team_dir = ?1 ORDER BY created_at, request_id",
375 )
376 .map_err(sqlite_error)?;
377 let rows = stmt
378 .query_map(params![Self::team_key(team_dir)], |row| {
379 row.get::<_, String>(0)
380 })
381 .map_err(sqlite_error)?;
382 let mut requests = Vec::new();
383 for row in rows {
384 requests.push(from_json(&row.map_err(sqlite_error)?)?);
385 }
386 Ok(requests)
387 }
388
389 fn upsert_team_request(
390 &self,
391 team_dir: &Path,
392 request: &TeamProtocolRequestSummary,
393 ) -> Result<(), RuntimeError> {
394 let conn = self.open()?;
395 conn.execute(
396 r#"
397 INSERT INTO team_requests (request_id, team_dir, payload_json, created_at)
398 VALUES (?1, ?2, ?3, ?4)
399 ON CONFLICT(request_id) DO UPDATE SET
400 team_dir = excluded.team_dir,
401 payload_json = excluded.payload_json
402 "#,
403 params![
404 request.request_id,
405 Self::team_key(team_dir),
406 to_json(request)?,
407 request.created_at as i64,
408 ],
409 )
410 .map_err(sqlite_error)?;
411 Ok(())
412 }
413
414 fn list_team_agent_names(&self, team_dir: &Path) -> Result<Vec<String>, RuntimeError> {
415 let conn = self.open()?;
416 let mut stmt = conn
417 .prepare("SELECT name FROM agents WHERE team_dir = ?1 ORDER BY name")
418 .map_err(sqlite_error)?;
419 stmt.query_map(params![Self::team_key(team_dir)], |row| {
420 row.get::<_, String>(0)
421 })
422 .map_err(sqlite_error)?
423 .collect::<Result<Vec<_>, _>>()
424 .map_err(sqlite_error)
425 }
426}
427
428impl BackgroundStore for SqliteRuntimeStore {
429 fn load_background_tasks(
430 &self,
431 agent_id: &str,
432 ) -> Result<Vec<BackgroundTaskSummary>, RuntimeError> {
433 let conn = self.open()?;
434 let mut stmt = conn
435 .prepare(
436 "SELECT payload_json FROM background_jobs WHERE agent_id = ?1 ORDER BY created_at, id",
437 )
438 .map_err(sqlite_error)?;
439 let rows = stmt
440 .query_map(params![agent_id], |row| row.get::<_, String>(0))
441 .map_err(sqlite_error)?;
442 let mut tasks = Vec::new();
443 for row in rows {
444 tasks.push(from_json(&row.map_err(sqlite_error)?)?);
445 }
446 Ok(tasks)
447 }
448
449 fn upsert_background_task(
450 &self,
451 agent_id: &str,
452 task: &BackgroundTaskSummary,
453 notification_state: i64,
454 ) -> Result<(), RuntimeError> {
455 let conn = self.open()?;
456 conn.execute(
457 r#"
458 INSERT INTO background_jobs (agent_id, id, payload_json, notification_state, created_at, updated_at)
459 VALUES (?1, ?2, ?3, ?4, ?5, ?5)
460 ON CONFLICT(agent_id, id) DO UPDATE SET
461 payload_json = excluded.payload_json,
462 notification_state = excluded.notification_state,
463 updated_at = excluded.updated_at
464 "#,
465 params![agent_id, task.id, to_json(task)?, notification_state, now_secs()],
466 )
467 .map_err(sqlite_error)?;
468 Ok(())
469 }
470
471 fn drain_background_notifications(
472 &self,
473 agent_id: &str,
474 ) -> Result<Vec<BackgroundNotification>, RuntimeError> {
475 let mut conn = self.open()?;
476 let tx = conn
477 .transaction_with_behavior(TransactionBehavior::Immediate)
478 .map_err(sqlite_error)?;
479 let jobs = {
480 let mut stmt = tx
481 .prepare(
482 "SELECT id, payload_json FROM background_jobs WHERE agent_id = ?1 AND notification_state = ?2 ORDER BY updated_at, id",
483 )
484 .map_err(sqlite_error)?;
485 stmt.query_map(params![agent_id, DELIVERY_PENDING], |row| {
486 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
487 })
488 .map_err(sqlite_error)?
489 .collect::<Result<Vec<_>, _>>()
490 .map_err(sqlite_error)?
491 };
492 for (id, _) in &jobs {
493 tx.execute(
494 "UPDATE background_jobs SET notification_state = ?3 WHERE agent_id = ?1 AND id = ?2",
495 params![agent_id, id, DELIVERY_INFLIGHT],
496 )
497 .map_err(sqlite_error)?;
498 }
499 tx.commit().map_err(sqlite_error)?;
500
501 jobs.into_iter()
502 .map(|(_, payload)| {
503 let task: BackgroundTaskSummary = from_json(&payload)?;
504 Ok(BackgroundNotification {
505 task_id: task.id,
506 command: task.command,
507 cwd: task.cwd,
508 status: task.status,
509 output_preview: task
510 .output_preview
511 .unwrap_or_else(|| "(no output)".to_string()),
512 })
513 })
514 .collect()
515 }
516
517 fn has_pending_background_notifications(&self, agent_id: &str) -> Result<bool, RuntimeError> {
518 let conn = self.open()?;
519 let exists = conn
520 .query_row(
521 "SELECT EXISTS(SELECT 1 FROM background_jobs WHERE agent_id = ?1 AND notification_state IN (?2, ?3))",
522 params![agent_id, DELIVERY_PENDING, DELIVERY_INFLIGHT],
523 |row| row.get::<_, i64>(0),
524 )
525 .map_err(sqlite_error)?;
526 Ok(exists != 0)
527 }
528
529 fn has_deliverable_background_notifications(
530 &self,
531 agent_id: &str,
532 ) -> Result<bool, RuntimeError> {
533 let conn = self.open()?;
534 let exists = conn
535 .query_row(
536 "SELECT EXISTS(SELECT 1 FROM background_jobs WHERE agent_id = ?1 AND notification_state = ?2)",
537 params![agent_id, DELIVERY_PENDING],
538 |row| row.get::<_, i64>(0),
539 )
540 .map_err(sqlite_error)?;
541 Ok(exists != 0)
542 }
543
544 fn ack_background_notifications(&self, agent_id: &str) -> Result<(), RuntimeError> {
545 let conn = self.open()?;
546 conn.execute(
547 "UPDATE background_jobs SET notification_state = ?2 WHERE agent_id = ?1 AND notification_state = ?3",
548 params![agent_id, DELIVERY_ACKED, DELIVERY_INFLIGHT],
549 )
550 .map_err(sqlite_error)?;
551 Ok(())
552 }
553
554 fn requeue_background_notifications(&self, agent_id: &str) -> Result<(), RuntimeError> {
555 let conn = self.open()?;
556 conn.execute(
557 "UPDATE background_jobs SET notification_state = ?2 WHERE agent_id = ?1 AND notification_state = ?3",
558 params![agent_id, DELIVERY_PENDING, DELIVERY_INFLIGHT],
559 )
560 .map_err(sqlite_error)?;
561 Ok(())
562 }
563}
564
565#[derive(Clone)]
566pub struct SqliteRuntimeStore {
568 path: PathBuf,
569}
570
571impl Default for SqliteRuntimeStore {
572 fn default() -> Self {
573 Self::new(Self::default_path())
574 }
575}
576
577impl SqliteRuntimeStore {
578 pub fn default_path() -> PathBuf {
580 default_store_dir().join("runtime.sqlite")
581 }
582
583 pub fn default_directory() -> PathBuf {
585 default_store_dir()
586 }
587
588 pub fn for_runtime_identifier(runtime_identifier: &str) -> Self {
590 Self::new(Self::path_for_runtime_identifier(runtime_identifier))
591 }
592
593 pub fn path_for_runtime_identifier(runtime_identifier: &str) -> PathBuf {
595 Self::default_directory().join(format!(
596 "runtime-{}.sqlite",
597 encode_runtime_identifier(runtime_identifier)
598 ))
599 }
600
601 pub fn list_persisted_runtime_identifiers() -> Result<Vec<String>, RuntimeError> {
603 let base = Self::default_directory();
604 let Ok(entries) = std::fs::read_dir(&base) else {
605 return Ok(Vec::new());
606 };
607
608 let mut runtime_identifiers = entries
609 .filter_map(|entry| entry.ok())
610 .filter_map(|entry| entry.file_name().into_string().ok())
611 .filter_map(|filename| decode_runtime_store_filename(&filename))
612 .collect::<Vec<_>>();
613 runtime_identifiers.sort();
614 runtime_identifiers.dedup();
615 Ok(runtime_identifiers)
616 }
617
618 pub fn new(path: impl Into<PathBuf>) -> Self {
620 Self { path: path.into() }
621 }
622
623 pub fn path(&self) -> &Path {
625 self.path.as_path()
626 }
627
628 fn open(&self) -> Result<Connection, RuntimeError> {
629 if let Some(parent) = self.path.parent() {
630 std::fs::create_dir_all(parent)
631 .map_err(|error| RuntimeError::Store(error.to_string()))?;
632 }
633 let conn = Connection::open(&self.path).map_err(sqlite_error)?;
634 conn.busy_timeout(Duration::from_secs(5))
635 .map_err(sqlite_error)?;
636 conn.pragma_update(None, "journal_mode", "WAL")
637 .map_err(sqlite_error)?;
638 conn.pragma_update(None, "foreign_keys", "ON")
639 .map_err(sqlite_error)?;
640 self.ensure_schema(&conn)?;
641 Ok(conn)
642 }
643
644 fn ensure_schema(&self, conn: &Connection) -> Result<(), RuntimeError> {
645 conn.execute_batch(
646 r#"
647 CREATE TABLE IF NOT EXISTS agents (
648 id TEXT PRIMARY KEY,
649 runtime_identifier TEXT NOT NULL,
650 name TEXT NOT NULL,
651 model TEXT NOT NULL,
652 provider_id TEXT NOT NULL,
653 team_dir TEXT NOT NULL,
654 tasks_namespace TEXT NOT NULL,
655 is_teammate INTEGER NOT NULL,
656 config_json TEXT NOT NULL,
657 hidden_tools_json TEXT NOT NULL,
658 max_rounds INTEGER,
659 teammate_identity_json TEXT,
660 rounds_since_task INTEGER NOT NULL,
661 idle_requested INTEGER NOT NULL,
662 status_json TEXT NOT NULL,
663 subagents_json TEXT NOT NULL,
664 created_at INTEGER NOT NULL,
665 updated_at INTEGER NOT NULL
666 );
667 CREATE TABLE IF NOT EXISTS agent_memory (
668 agent_id TEXT PRIMARY KEY,
669 revision INTEGER NOT NULL,
670 state_json TEXT NOT NULL,
671 updated_at INTEGER NOT NULL
672 );
673 CREATE TABLE IF NOT EXISTS agent_runs (
674 id TEXT PRIMARY KEY,
675 agent_id TEXT NOT NULL,
676 state TEXT NOT NULL,
677 error TEXT,
678 created_at INTEGER NOT NULL,
679 updated_at INTEGER NOT NULL
680 );
681 CREATE TABLE IF NOT EXISTS tasks (
682 namespace TEXT NOT NULL,
683 id INTEGER NOT NULL,
684 payload_json TEXT NOT NULL,
685 PRIMARY KEY (namespace, id)
686 );
687 CREATE TABLE IF NOT EXISTS task_edges (
688 namespace TEXT NOT NULL,
689 blocker_id INTEGER NOT NULL,
690 dependent_id INTEGER NOT NULL,
691 PRIMARY KEY (namespace, blocker_id, dependent_id)
692 );
693 CREATE TABLE IF NOT EXISTS team_members (
694 team_dir TEXT NOT NULL,
695 name TEXT NOT NULL,
696 summary_json TEXT NOT NULL,
697 PRIMARY KEY (team_dir, name)
698 );
699 CREATE TABLE IF NOT EXISTS team_inbox (
700 id TEXT PRIMARY KEY,
701 team_dir TEXT NOT NULL,
702 recipient TEXT NOT NULL,
703 payload_json TEXT NOT NULL,
704 delivery_state INTEGER NOT NULL,
705 created_at INTEGER NOT NULL
706 );
707 CREATE TABLE IF NOT EXISTS team_requests (
708 request_id TEXT PRIMARY KEY,
709 team_dir TEXT NOT NULL,
710 payload_json TEXT NOT NULL,
711 created_at INTEGER NOT NULL
712 );
713 CREATE TABLE IF NOT EXISTS background_jobs (
714 agent_id TEXT NOT NULL,
715 id TEXT NOT NULL,
716 payload_json TEXT NOT NULL,
717 notification_state INTEGER NOT NULL,
718 created_at INTEGER NOT NULL,
719 updated_at INTEGER NOT NULL,
720 PRIMARY KEY (agent_id, id)
721 );
722 CREATE TABLE IF NOT EXISTS audit_events (
723 id TEXT PRIMARY KEY,
724 scope TEXT NOT NULL,
725 event_type TEXT NOT NULL,
726 payload_json TEXT NOT NULL,
727 created_at INTEGER NOT NULL
728 );
729 CREATE TABLE IF NOT EXISTS leases (
730 key TEXT PRIMARY KEY,
731 owner TEXT NOT NULL,
732 expires_at INTEGER NOT NULL
733 );
734 CREATE TABLE IF NOT EXISTS permission_rules (
735 session_id TEXT NOT NULL,
736 project_id TEXT,
737 tool_name TEXT NOT NULL,
738 pattern TEXT,
739 allow INTEGER NOT NULL,
740 scope TEXT NOT NULL
741 );
742 CREATE TABLE IF NOT EXISTS long_term_memory (
743 record_id TEXT PRIMARY KEY,
744 agent_id TEXT NOT NULL,
745 kind TEXT NOT NULL,
746 content TEXT NOT NULL,
747 source_revision INTEGER NOT NULL,
748 created_at INTEGER NOT NULL,
749 metadata_json TEXT NOT NULL
750 );
751 CREATE VIRTUAL TABLE IF NOT EXISTS long_term_memory_fts USING fts5(
752 record_id UNINDEXED,
753 agent_id UNINDEXED,
754 content
755 );
756 CREATE TABLE IF NOT EXISTS long_term_memory_cursor (
757 agent_id TEXT PRIMARY KEY,
758 cursor_json TEXT NOT NULL,
759 updated_at INTEGER NOT NULL
760 );
761 "#,
762 )
763 .map_err(sqlite_error)?;
764 self.migrate_background_jobs_schema(conn)?;
765 self.migrate_permission_rules_schema(conn)
766 }
767
768 fn migrate_permission_rules_schema(&self, conn: &Connection) -> Result<(), RuntimeError> {
769 let Some(schema_sql) = conn
770 .query_row(
771 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'permission_rules'",
772 [],
773 |row| row.get::<_, String>(0),
774 )
775 .optional()
776 .map_err(sqlite_error)?
777 else {
778 return Ok(());
779 };
780
781 if !schema_sql.contains("project_id") {
782 conn.execute_batch("ALTER TABLE permission_rules ADD COLUMN project_id TEXT;")
783 .map_err(sqlite_error)?;
784 }
785
786 conn.execute_batch(
788 r#"
789 CREATE INDEX IF NOT EXISTS idx_perm_session ON permission_rules (session_id);
790 CREATE INDEX IF NOT EXISTS idx_perm_project ON permission_rules (project_id);
791 CREATE INDEX IF NOT EXISTS idx_perm_global ON permission_rules (scope);
792 "#,
793 )
794 .map_err(sqlite_error)?;
795
796 Ok(())
797 }
798
799 fn migrate_background_jobs_schema(&self, conn: &Connection) -> Result<(), RuntimeError> {
800 let Some(schema_sql) = conn
801 .query_row(
802 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'background_jobs'",
803 [],
804 |row| row.get::<_, String>(0),
805 )
806 .optional()
807 .map_err(sqlite_error)?
808 else {
809 return Ok(());
810 };
811
812 if schema_sql.contains("PRIMARY KEY (agent_id, id)")
813 || schema_sql.contains("PRIMARY KEY(agent_id, id)")
814 {
815 return Ok(());
816 }
817
818 conn.execute_batch(
819 r#"
820 ALTER TABLE background_jobs RENAME TO background_jobs_legacy;
821 CREATE TABLE background_jobs (
822 agent_id TEXT NOT NULL,
823 id TEXT NOT NULL,
824 payload_json TEXT NOT NULL,
825 notification_state INTEGER NOT NULL,
826 created_at INTEGER NOT NULL,
827 updated_at INTEGER NOT NULL,
828 PRIMARY KEY (agent_id, id)
829 );
830 INSERT INTO background_jobs (agent_id, id, payload_json, notification_state, created_at, updated_at)
831 SELECT agent_id, id, payload_json, notification_state, created_at, updated_at
832 FROM background_jobs_legacy;
833 DROP TABLE background_jobs_legacy;
834 "#,
835 )
836 .map_err(sqlite_error)?;
837
838 Ok(())
839 }
840
841 fn write_agent(
842 &self,
843 conn: &Connection,
844 record: &PersistedAgentRecord,
845 ) -> Result<(), RuntimeError> {
846 let now = now_secs();
847 conn.execute(
848 r#"
849 INSERT INTO agents (
850 id, runtime_identifier, name, model, provider_id, team_dir, tasks_namespace, is_teammate, config_json,
851 hidden_tools_json, max_rounds, teammate_identity_json, rounds_since_task,
852 idle_requested, status_json, subagents_json, created_at, updated_at
853 )
854 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
855 ON CONFLICT(id) DO UPDATE SET
856 runtime_identifier = excluded.runtime_identifier,
857 name = excluded.name,
858 model = excluded.model,
859 provider_id = excluded.provider_id,
860 team_dir = excluded.team_dir,
861 tasks_namespace = excluded.tasks_namespace,
862 is_teammate = excluded.is_teammate,
863 config_json = excluded.config_json,
864 hidden_tools_json = excluded.hidden_tools_json,
865 max_rounds = excluded.max_rounds,
866 teammate_identity_json = excluded.teammate_identity_json,
867 rounds_since_task = excluded.rounds_since_task,
868 idle_requested = excluded.idle_requested,
869 status_json = excluded.status_json,
870 subagents_json = excluded.subagents_json,
871 updated_at = excluded.updated_at
872 "#,
873 params![
874 record.id,
875 record.runtime_identifier,
876 record.name,
877 record.model,
878 record.provider_id.as_str(),
879 record.config.team.team_dir.to_string_lossy().into_owned(),
880 record.config.task.tasks_dir.to_string_lossy().into_owned(),
881 i64::from(record.teammate_identity.is_some()),
882 to_json(&record.config)?,
883 to_json(&record.hidden_tools)?,
884 record.max_rounds.map(|value| value as i64),
885 maybe_json(&record.teammate_identity)?,
886 record.rounds_since_task as i64,
887 i64::from(record.idle_requested),
888 to_json(&record.status)?,
889 to_json(&record.subagents)?,
890 now,
891 now,
892 ],
893 )
894 .map_err(sqlite_error)?;
895 Ok(())
896 }
897
898 fn write_agent_memory(
899 &self,
900 conn: &Connection,
901 agent_id: &str,
902 memory: &AgentMemoryState,
903 ) -> Result<(), RuntimeError> {
904 conn.execute(
905 r#"
906 INSERT INTO agent_memory (agent_id, revision, state_json, updated_at)
907 VALUES (?1, ?2, ?3, ?4)
908 ON CONFLICT(agent_id) DO UPDATE SET
909 revision = excluded.revision,
910 state_json = excluded.state_json,
911 updated_at = excluded.updated_at
912 "#,
913 params![
914 agent_id,
915 memory.revision as i64,
916 to_json(memory)?,
917 now_secs()
918 ],
919 )
920 .map_err(sqlite_error)?;
921 Ok(())
922 }
923
924 fn team_key(path: &Path) -> String {
925 path.to_string_lossy().into_owned()
926 }
927
928 fn task_namespace(path: &Path) -> String {
929 path.to_string_lossy().into_owned()
930 }
931}
932
933impl AgentStore for SqliteRuntimeStore {
934 fn prepare_recovery(&self) -> Result<(), RuntimeError> {
935 let mut conn = self.open()?;
936 let tx = conn
937 .transaction_with_behavior(TransactionBehavior::Immediate)
938 .map_err(sqlite_error)?;
939
940 tx.execute(
941 "UPDATE team_inbox SET delivery_state = ?1 WHERE delivery_state = ?2",
942 params![DELIVERY_PENDING, DELIVERY_INFLIGHT],
943 )
944 .map_err(sqlite_error)?;
945 tx.execute(
946 "UPDATE background_jobs SET notification_state = ?1 WHERE notification_state = ?2",
947 params![DELIVERY_PENDING, DELIVERY_INFLIGHT],
948 )
949 .map_err(sqlite_error)?;
950
951 {
952 let mut stmt = tx
953 .prepare("SELECT agent_id, id, payload_json FROM background_jobs")
954 .map_err(sqlite_error)?;
955 let rows = stmt
956 .query_map([], |row| {
957 Ok((
958 row.get::<_, String>(0)?,
959 row.get::<_, String>(1)?,
960 row.get::<_, String>(2)?,
961 ))
962 })
963 .map_err(sqlite_error)?;
964 for row in rows {
965 let (agent_id, id, payload) = row.map_err(sqlite_error)?;
966 let mut task: BackgroundTaskSummary = from_json(&payload)?;
967 if task.status == BackgroundTaskStatus::Running {
968 task.status = BackgroundTaskStatus::Interrupted;
969 tx.execute(
970 "UPDATE background_jobs SET payload_json = ?3, notification_state = ?4, updated_at = ?5 WHERE agent_id = ?1 AND id = ?2",
971 params![agent_id, id, to_json(&task)?, DELIVERY_PENDING, now_secs()],
972 )
973 .map_err(sqlite_error)?;
974 }
975 }
976 }
977
978 tx.execute(
979 "DELETE FROM leases WHERE expires_at <= ?1",
980 params![now_secs()],
981 )
982 .map_err(sqlite_error)?;
983 prune_stale_runtime_leases(&tx)?;
984 tx.commit().map_err(sqlite_error)
985 }
986
987 fn create_agent(
988 &self,
989 record: &PersistedAgentRecord,
990 memory: &AgentMemoryState,
991 ) -> Result<(), RuntimeError> {
992 let mut conn = self.open()?;
993 let tx = conn
994 .transaction_with_behavior(TransactionBehavior::Immediate)
995 .map_err(sqlite_error)?;
996 self.write_agent(&tx, record)?;
997 self.write_agent_memory(&tx, &record.id, memory)?;
998 tx.commit().map_err(sqlite_error)
999 }
1000
1001 fn save_agent_record(&self, record: &PersistedAgentRecord) -> Result<(), RuntimeError> {
1002 let mut conn = self.open()?;
1003 let tx = conn
1004 .transaction_with_behavior(TransactionBehavior::Immediate)
1005 .map_err(sqlite_error)?;
1006 self.write_agent(&tx, record)?;
1007 tx.commit().map_err(sqlite_error)
1008 }
1009
1010 fn save_agent_memory(
1011 &self,
1012 agent_id: &str,
1013 memory: &AgentMemoryState,
1014 ) -> Result<(), RuntimeError> {
1015 let mut conn = self.open()?;
1016 let tx = conn
1017 .transaction_with_behavior(TransactionBehavior::Immediate)
1018 .map_err(sqlite_error)?;
1019 self.write_agent_memory(&tx, agent_id, memory)?;
1020 tx.commit().map_err(sqlite_error)
1021 }
1022
1023 fn load_agent(&self, agent_id: &str) -> Result<Option<LoadedAgentState>, RuntimeError> {
1024 let conn = self.open()?;
1025 let record = conn
1026 .query_row(
1027 r#"
1028 SELECT
1029 id, runtime_identifier, name, model, provider_id, config_json,
1030 hidden_tools_json, max_rounds, teammate_identity_json, rounds_since_task,
1031 idle_requested, status_json, subagents_json
1032 FROM agents WHERE id = ?1
1033 "#,
1034 params![agent_id],
1035 |row| {
1036 let provider_id: String = row.get(4)?;
1037 let config_json: String = row.get(5)?;
1038 let hidden_tools_json: String = row.get(6)?;
1039 let teammate_identity_json: Option<String> = row.get(8)?;
1040 let status_json: String = row.get(11)?;
1041 let subagents_json: String = row.get(12)?;
1042 Ok(PersistedAgentRecord {
1043 id: row.get(0)?,
1044 runtime_identifier: row.get(1)?,
1045 name: row.get(2)?,
1046 model: row.get(3)?,
1047 provider_id: ProviderId::from(provider_id),
1048 config: from_json(&config_json).map_err(to_sql_error)?,
1049 hidden_tools: from_json(&hidden_tools_json).map_err(to_sql_error)?,
1050 max_rounds: row.get::<_, Option<i64>>(7)?.map(|value| value as usize),
1051 teammate_identity: teammate_identity_json
1052 .map(|json| from_json(&json))
1053 .transpose()
1054 .map_err(to_sql_error)?,
1055 rounds_since_task: row.get::<_, i64>(9)? as usize,
1056 idle_requested: row.get::<_, i64>(10)? != 0,
1057 status: from_json(&status_json).map_err(to_sql_error)?,
1058 subagents: from_json(&subagents_json).map_err(to_sql_error)?,
1059 })
1060 },
1061 )
1062 .optional()
1063 .map_err(sqlite_error)?;
1064 let Some(record) = record else {
1065 return Ok(None);
1066 };
1067
1068 let memory = conn
1069 .query_row(
1070 "SELECT state_json FROM agent_memory WHERE agent_id = ?1",
1071 params![agent_id],
1072 |row| {
1073 let state_json: String = row.get(0)?;
1074 from_json(&state_json).map_err(to_sql_error)
1075 },
1076 )
1077 .optional()
1078 .map_err(sqlite_error)?;
1079 let Some(memory) = memory else {
1080 return Err(RuntimeError::Store(format!(
1081 "Agent '{agent_id}' is missing persisted memory"
1082 )));
1083 };
1084
1085 Ok(Some(LoadedAgentState { record, memory }))
1086 }
1087
1088 fn list_agents(&self) -> Result<Vec<LoadedAgentState>, RuntimeError> {
1089 let conn = self.open()?;
1090 let mut stmt = conn
1091 .prepare("SELECT id FROM agents ORDER BY created_at, id")
1092 .map_err(sqlite_error)?;
1093 let ids = stmt
1094 .query_map([], |row| row.get::<_, String>(0))
1095 .map_err(sqlite_error)?
1096 .collect::<Result<Vec<_>, _>>()
1097 .map_err(sqlite_error)?;
1098 ids.into_iter()
1099 .map(|id| {
1100 self.load_agent(&id)?
1101 .ok_or_else(|| RuntimeError::Store(format!("Agent '{id}' disappeared")))
1102 })
1103 .collect()
1104 }
1105
1106 fn list_agents_by_runtime(
1107 &self,
1108 runtime_identifier: &str,
1109 ) -> Result<Vec<LoadedAgentState>, RuntimeError> {
1110 let conn = self.open()?;
1111 let mut stmt = conn
1112 .prepare("SELECT id FROM agents WHERE runtime_identifier = ?1 ORDER BY created_at, id")
1113 .map_err(sqlite_error)?;
1114 let ids = stmt
1115 .query_map(params![runtime_identifier], |row| row.get::<_, String>(0))
1116 .map_err(sqlite_error)?
1117 .collect::<Result<Vec<_>, _>>()
1118 .map_err(sqlite_error)?;
1119 ids.into_iter()
1120 .map(|id| {
1121 self.load_agent(&id)?
1122 .ok_or_else(|| RuntimeError::Store(format!("Agent '{id}' disappeared")))
1123 })
1124 .collect()
1125 }
1126}
1127
1128impl RunStore for SqliteRuntimeStore {
1129 fn start_run(&self, agent_id: &str) -> Result<String, RuntimeError> {
1130 let run_id = next_id("run");
1131 let conn = self.open()?;
1132 conn.execute(
1133 "INSERT INTO agent_runs (id, agent_id, state, error, created_at, updated_at) VALUES (?1, ?2, 'running', NULL, ?3, ?3)",
1134 params![run_id, agent_id, now_secs()],
1135 )
1136 .map_err(sqlite_error)?;
1137 Ok(run_id)
1138 }
1139
1140 fn update_run_state(
1141 &self,
1142 run_id: &str,
1143 state: &str,
1144 error: Option<&str>,
1145 ) -> Result<(), RuntimeError> {
1146 let conn = self.open()?;
1147 conn.execute(
1148 "UPDATE agent_runs SET state = ?2, error = ?3, updated_at = ?4 WHERE id = ?1",
1149 params![run_id, state, error, now_secs()],
1150 )
1151 .map_err(sqlite_error)?;
1152 Ok(())
1153 }
1154
1155 fn finish_run(&self, run_id: &str) -> Result<(), RuntimeError> {
1156 self.update_run_state(run_id, "finished", None)
1157 }
1158
1159 fn fail_run(&self, run_id: &str, error: &str) -> Result<(), RuntimeError> {
1160 self.update_run_state(run_id, "failed", Some(error))
1161 }
1162}
1163
1164impl TaskStore for SqliteRuntimeStore {
1165 fn load_tasks(&self, namespace: &Path) -> Result<Vec<TaskItem>, RuntimeError> {
1166 let conn = self.open()?;
1167 self.load_tasks_from_conn(&conn, namespace)
1168 }
1169
1170 fn capture_tasks(&self, namespace: &Path) -> Result<TaskStateSnapshot, RuntimeError> {
1171 Ok(TaskStateSnapshot {
1172 tasks: self.load_tasks(namespace)?,
1173 })
1174 }
1175
1176 fn restore_tasks(
1177 &self,
1178 namespace: &Path,
1179 snapshot: &TaskStateSnapshot,
1180 ) -> Result<(), RuntimeError> {
1181 self.replace_tasks(namespace, &snapshot.tasks)
1182 }
1183
1184 fn replace_tasks(&self, namespace: &Path, tasks: &[TaskItem]) -> Result<(), RuntimeError> {
1185 let mut conn = self.open()?;
1186 let tx = conn
1187 .transaction_with_behavior(TransactionBehavior::Immediate)
1188 .map_err(sqlite_error)?;
1189 self.replace_tasks_in_conn(&tx, namespace, tasks)?;
1190 tx.commit().map_err(sqlite_error)
1191 }
1192
1193 fn mutate(
1194 &self,
1195 namespace: &Path,
1196 mutation: &mut dyn FnMut(&mut Vec<TaskItem>) -> Result<(), RuntimeError>,
1197 ) -> Result<(), RuntimeError> {
1198 let mut conn = self.open()?;
1199 let tx = conn
1200 .transaction_with_behavior(TransactionBehavior::Immediate)
1201 .map_err(sqlite_error)?;
1202 let mut tasks = self.load_tasks_from_conn(&tx, namespace)?;
1203 mutation(&mut tasks)?;
1204 self.replace_tasks_in_conn(&tx, namespace, &tasks)?;
1205 tx.commit().map_err(sqlite_error)
1206 }
1207}
1208
1209impl AuditStore for SqliteRuntimeStore {
1210 fn record_audit_event(
1211 &self,
1212 scope: &str,
1213 event_type: &str,
1214 payload: serde_json::Value,
1215 ) -> Result<(), RuntimeError> {
1216 let conn = self.open()?;
1217 conn.execute(
1218 "INSERT INTO audit_events (id, scope, event_type, payload_json, created_at) VALUES (?1, ?2, ?3, ?4, ?5)",
1219 params![next_id("audit"), scope, event_type, payload.to_string(), now_secs()],
1220 )
1221 .map_err(sqlite_error)?;
1222 Ok(())
1223 }
1224}
1225
1226impl LeaseStore for SqliteRuntimeStore {
1227 fn acquire_lease(&self, key: &str, owner: &str, ttl: Duration) -> Result<bool, RuntimeError> {
1228 let mut conn = self.open()?;
1229 let tx = conn
1230 .transaction_with_behavior(TransactionBehavior::Immediate)
1231 .map_err(sqlite_error)?;
1232 let now = now_secs();
1233 tx.execute("DELETE FROM leases WHERE expires_at <= ?1", params![now])
1234 .map_err(sqlite_error)?;
1235 prune_stale_runtime_leases(&tx)?;
1236 let inserted = tx
1237 .execute(
1238 "INSERT OR IGNORE INTO leases (key, owner, expires_at) VALUES (?1, ?2, ?3)",
1239 params![key, owner, now + ttl.as_secs() as i64],
1240 )
1241 .map_err(sqlite_error)?;
1242 tx.commit().map_err(sqlite_error)?;
1243 Ok(inserted == 1)
1244 }
1245
1246 fn release_lease(&self, key: &str, owner: &str) -> Result<(), RuntimeError> {
1247 let conn = self.open()?;
1248 conn.execute(
1249 "DELETE FROM leases WHERE key = ?1 AND owner = ?2",
1250 params![key, owner],
1251 )
1252 .map_err(sqlite_error)?;
1253 Ok(())
1254 }
1255}
1256
1257impl PermissionRuleStore for SqliteRuntimeStore {
1258 fn save_rules(
1259 &self,
1260 session_id: &str,
1261 project_id: Option<&str>,
1262 rules: &[RememberedRule],
1263 ) -> Result<(), RuntimeError> {
1264 let mut conn = self.open()?;
1265 let tx = conn
1266 .transaction_with_behavior(TransactionBehavior::Immediate)
1267 .map_err(sqlite_error)?;
1268
1269 let session_scope = to_json(&PermissionRuleScope::Session)?;
1272 tx.execute(
1273 "DELETE FROM permission_rules WHERE session_id = ?1 AND scope = ?2",
1274 params![session_id, session_scope],
1275 )
1276 .map_err(sqlite_error)?;
1277
1278 for rule in rules {
1279 let scope_str = to_json(&rule.scope)?;
1280 tx.execute(
1281 r#"
1282 INSERT INTO permission_rules (session_id, project_id, tool_name, pattern, allow, scope)
1283 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
1284 "#,
1285 params![
1286 session_id,
1287 project_id,
1288 rule.key.tool_name,
1289 rule.key.pattern,
1290 rule.allow as i32,
1291 scope_str,
1292 ],
1293 )
1294 .map_err(sqlite_error)?;
1295 }
1296
1297 tx.commit().map_err(sqlite_error)?;
1298 Ok(())
1299 }
1300
1301 fn load_rules(
1302 &self,
1303 session_id: &str,
1304 project_id: Option<&str>,
1305 ) -> Result<Vec<RememberedRule>, RuntimeError> {
1306 let conn = self.open()?;
1307
1308 let session_scope = to_json(&PermissionRuleScope::Session)?;
1309 let project_scope = to_json(&PermissionRuleScope::Project)?;
1310 let global_scope = to_json(&PermissionRuleScope::Global)?;
1311
1312 let sql = r#"
1315 SELECT tool_name, pattern, allow, scope
1316 FROM permission_rules
1317 WHERE session_id = ?1 AND scope = ?2
1318 UNION
1319 SELECT tool_name, pattern, allow, scope
1320 FROM permission_rules
1321 WHERE project_id IS NOT NULL AND project_id = ?3 AND scope = ?4
1322 UNION
1323 SELECT tool_name, pattern, allow, scope
1324 FROM permission_rules
1325 WHERE scope = ?5
1326 "#;
1327
1328 let mut stmt = conn.prepare(sql).map_err(sqlite_error)?;
1329
1330 let project_id_param = project_id.unwrap_or("");
1333
1334 let rows = stmt
1335 .query_map(
1336 params![
1337 session_id,
1338 session_scope,
1339 project_id_param,
1340 project_scope,
1341 global_scope,
1342 ],
1343 |row| {
1344 Ok((
1345 row.get::<_, String>(0)?,
1346 row.get::<_, Option<String>>(1)?,
1347 row.get::<_, i32>(2)?,
1348 row.get::<_, String>(3)?,
1349 ))
1350 },
1351 )
1352 .map_err(sqlite_error)?;
1353
1354 let mut rules = Vec::new();
1355 for row in rows {
1356 let (tool_name, pattern, allow, scope_str) = row.map_err(sqlite_error)?;
1357 let scope: PermissionRuleScope = from_json(&scope_str)?;
1358 rules.push(RememberedRule {
1359 key: RuleKey { tool_name, pattern },
1360 allow: allow != 0,
1361 scope,
1362 });
1363 }
1364 Ok(rules)
1365 }
1366
1367 fn clear_rules(&self, session_id: &str) -> Result<(), RuntimeError> {
1368 let conn = self.open()?;
1369 conn.execute(
1370 "DELETE FROM permission_rules WHERE session_id = ?1",
1371 params![session_id],
1372 )
1373 .map_err(sqlite_error)?;
1374 Ok(())
1375 }
1376}
1377
1378impl MemoryStore for SqliteRuntimeStore {
1379 fn upsert_records(&self, records: &[MemoryRecord]) -> Result<(), RuntimeError> {
1380 if records.is_empty() {
1381 return Ok(());
1382 }
1383
1384 let mut conn = self.open()?;
1385 let tx = conn
1386 .transaction_with_behavior(TransactionBehavior::Immediate)
1387 .map_err(sqlite_error)?;
1388
1389 for record in records {
1390 tx.execute(
1391 r#"
1392 INSERT INTO long_term_memory (
1393 record_id, agent_id, kind, content, source_revision, created_at, metadata_json
1394 )
1395 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
1396 ON CONFLICT(record_id) DO UPDATE SET
1397 agent_id = excluded.agent_id,
1398 kind = excluded.kind,
1399 content = excluded.content,
1400 source_revision = excluded.source_revision,
1401 created_at = excluded.created_at,
1402 metadata_json = excluded.metadata_json
1403 "#,
1404 params![
1405 record.record_id,
1406 record.agent_id,
1407 format!("{:?}", record.kind).to_lowercase(),
1408 record.content,
1409 record.source_revision as i64,
1410 record.created_at,
1411 record.metadata_json,
1412 ],
1413 )
1414 .map_err(sqlite_error)?;
1415 tx.execute(
1416 "DELETE FROM long_term_memory_fts WHERE record_id = ?1",
1417 params![record.record_id],
1418 )
1419 .map_err(sqlite_error)?;
1420 tx.execute(
1421 "INSERT INTO long_term_memory_fts (record_id, agent_id, content) VALUES (?1, ?2, ?3)",
1422 params![record.record_id, record.agent_id, record.content],
1423 )
1424 .map_err(sqlite_error)?;
1425 }
1426
1427 tx.commit().map_err(sqlite_error)
1428 }
1429
1430 fn search_records_with_options(
1431 &self,
1432 request: &MemorySearchRequest,
1433 ) -> Result<Vec<MemoryRecord>, RuntimeError> {
1434 if request.query.trim().is_empty() || request.limit == 0 {
1435 return Ok(Vec::new());
1436 }
1437 let Some(query) = fts_query(&request.query) else {
1438 return Ok(Vec::new());
1439 };
1440
1441 let conn = self.open()?;
1442 let mut stmt = conn
1443 .prepare(
1444 r#"
1445 SELECT
1446 memory.record_id,
1447 memory.agent_id,
1448 memory.kind,
1449 memory.content,
1450 memory.source_revision,
1451 memory.created_at,
1452 memory.metadata_json,
1453 bm25(long_term_memory_fts) AS rank
1454 FROM long_term_memory_fts
1455 JOIN long_term_memory AS memory ON memory.record_id = long_term_memory_fts.record_id
1456 WHERE long_term_memory_fts.agent_id = ?1
1457 AND long_term_memory_fts.content MATCH ?2
1458 ORDER BY rank, memory.created_at DESC
1459 LIMIT ?3
1460 "#,
1461 )
1462 .map_err(sqlite_error)?;
1463
1464 stmt.query_map(
1465 params![request.agent_id, query, request.limit as i64],
1466 |row| {
1467 let kind = row.get::<_, String>(2)?;
1468 Ok(MemoryRecord {
1469 record_id: row.get(0)?,
1470 agent_id: row.get(1)?,
1471 kind: parse_memory_kind(&kind),
1472 content: row.get(3)?,
1473 source_revision: row.get::<_, i64>(4)? as u64,
1474 created_at: row.get(5)?,
1475 metadata_json: row.get(6)?,
1476 source: None,
1477 pinned: false,
1478 score: row.get::<_, Option<f64>>(7)?,
1479 })
1480 },
1481 )
1482 .map_err(sqlite_error)?
1483 .collect::<Result<Vec<_>, _>>()
1484 .map_err(sqlite_error)
1485 }
1486
1487 fn delete_records(&self, record_ids: &[String]) -> Result<(), RuntimeError> {
1488 if record_ids.is_empty() {
1489 return Ok(());
1490 }
1491
1492 let mut conn = self.open()?;
1493 let tx = conn
1494 .transaction_with_behavior(TransactionBehavior::Immediate)
1495 .map_err(sqlite_error)?;
1496 for record_id in record_ids {
1497 tx.execute(
1498 "DELETE FROM long_term_memory_fts WHERE record_id = ?1",
1499 params![record_id],
1500 )
1501 .map_err(sqlite_error)?;
1502 tx.execute(
1503 "DELETE FROM long_term_memory WHERE record_id = ?1",
1504 params![record_id],
1505 )
1506 .map_err(sqlite_error)?;
1507 }
1508 tx.commit().map_err(sqlite_error)
1509 }
1510
1511 fn tombstone_records(
1512 &self,
1513 agent_id: &str,
1514 record_ids: &[String],
1515 ) -> Result<usize, RuntimeError> {
1516 if record_ids.is_empty() {
1517 return Ok(0);
1518 }
1519
1520 let mut conn = self.open()?;
1521 let tx = conn
1522 .transaction_with_behavior(TransactionBehavior::Immediate)
1523 .map_err(sqlite_error)?;
1524 let mut affected = 0usize;
1525 for record_id in record_ids {
1526 tx.execute(
1527 "DELETE FROM long_term_memory_fts WHERE record_id = ?1",
1528 params![record_id],
1529 )
1530 .map_err(sqlite_error)?;
1531 affected += tx
1532 .execute(
1533 "DELETE FROM long_term_memory WHERE record_id = ?1 AND agent_id = ?2",
1534 params![record_id, agent_id],
1535 )
1536 .map_err(sqlite_error)?;
1537 }
1538 tx.commit().map_err(sqlite_error)?;
1539 Ok(affected)
1540 }
1541
1542 fn load_agent_memory_cursor(
1543 &self,
1544 agent_id: &str,
1545 ) -> Result<Option<MemoryCursor>, RuntimeError> {
1546 let conn = self.open()?;
1547 conn.query_row(
1548 "SELECT cursor_json FROM long_term_memory_cursor WHERE agent_id = ?1",
1549 params![agent_id],
1550 |row| row.get::<_, String>(0),
1551 )
1552 .optional()
1553 .map_err(sqlite_error)?
1554 .map(|json| from_json(&json))
1555 .transpose()
1556 }
1557
1558 fn save_agent_memory_cursor(
1559 &self,
1560 agent_id: &str,
1561 cursor: &MemoryCursor,
1562 ) -> Result<(), RuntimeError> {
1563 let conn = self.open()?;
1564 conn.execute(
1565 r#"
1566 INSERT INTO long_term_memory_cursor (agent_id, cursor_json, updated_at)
1567 VALUES (?1, ?2, ?3)
1568 ON CONFLICT(agent_id) DO UPDATE SET
1569 cursor_json = excluded.cursor_json,
1570 updated_at = excluded.updated_at
1571 "#,
1572 params![agent_id, to_json(cursor)?, now_secs()],
1573 )
1574 .map_err(sqlite_error)?;
1575 Ok(())
1576 }
1577}
1578
1579impl SqliteRuntimeStore {
1580 fn load_tasks_from_conn(
1581 &self,
1582 conn: &Connection,
1583 namespace: &Path,
1584 ) -> Result<Vec<TaskItem>, RuntimeError> {
1585 let mut stmt = conn
1586 .prepare("SELECT payload_json FROM tasks WHERE namespace = ?1 ORDER BY id")
1587 .map_err(sqlite_error)?;
1588 let rows = stmt
1589 .query_map(params![Self::task_namespace(namespace)], |row| {
1590 row.get::<_, String>(0)
1591 })
1592 .map_err(sqlite_error)?;
1593 let mut tasks = Vec::new();
1594 for row in rows {
1595 tasks.push(from_json(&row.map_err(sqlite_error)?)?);
1596 }
1597 Ok(tasks)
1598 }
1599
1600 fn replace_tasks_in_conn(
1601 &self,
1602 conn: &Connection,
1603 namespace: &Path,
1604 tasks: &[TaskItem],
1605 ) -> Result<(), RuntimeError> {
1606 let namespace = Self::task_namespace(namespace);
1607 conn.execute(
1608 "DELETE FROM tasks WHERE namespace = ?1",
1609 params![namespace.clone()],
1610 )
1611 .map_err(sqlite_error)?;
1612 conn.execute(
1613 "DELETE FROM task_edges WHERE namespace = ?1",
1614 params![namespace.clone()],
1615 )
1616 .map_err(sqlite_error)?;
1617 for task in tasks {
1618 conn.execute(
1619 "INSERT INTO tasks (namespace, id, payload_json) VALUES (?1, ?2, ?3)",
1620 params![namespace.clone(), task.id as i64, to_json(task)?],
1621 )
1622 .map_err(sqlite_error)?;
1623 for blocker in &task.blocked_by {
1624 conn.execute(
1625 "INSERT OR IGNORE INTO task_edges (namespace, blocker_id, dependent_id) VALUES (?1, ?2, ?3)",
1626 params![namespace.clone(), *blocker as i64, task.id as i64],
1627 )
1628 .map_err(sqlite_error)?;
1629 }
1630 }
1631 Ok(())
1632 }
1633}
1634
1635fn to_json<T: Serialize>(value: &T) -> Result<String, RuntimeError> {
1636 serde_json::to_string(value).map_err(|error| RuntimeError::Store(error.to_string()))
1637}
1638
1639fn maybe_json<T: Serialize>(value: &Option<T>) -> Result<Option<String>, RuntimeError> {
1640 value.as_ref().map(to_json).transpose()
1641}
1642
1643fn from_json<T: DeserializeOwned>(value: &str) -> Result<T, RuntimeError> {
1644 serde_json::from_str(value).map_err(|error| RuntimeError::Store(error.to_string()))
1645}
1646
1647fn sqlite_error(error: rusqlite::Error) -> RuntimeError {
1648 RuntimeError::Store(error.to_string())
1649}
1650
1651fn parse_memory_kind(kind: &str) -> crate::memory::MemoryRecordKind {
1652 match kind {
1653 "summary" => crate::memory::MemoryRecordKind::Summary,
1654 "fact" => crate::memory::MemoryRecordKind::Fact,
1655 _ => crate::memory::MemoryRecordKind::Episode,
1656 }
1657}
1658
1659fn fts_query(query: &str) -> Option<String> {
1660 let tokens = query
1661 .split(|ch: char| !ch.is_alphanumeric())
1662 .filter(|token| !token.is_empty())
1663 .map(|token| format!("\"{token}\""))
1664 .collect::<Vec<_>>();
1665
1666 if tokens.is_empty() {
1667 None
1668 } else {
1669 Some(tokens.join(" OR "))
1670 }
1671}
1672
1673fn to_sql_error(error: RuntimeError) -> rusqlite::Error {
1674 rusqlite::Error::FromSqlConversionFailure(
1675 0,
1676 rusqlite::types::Type::Text,
1677 Box::new(std::io::Error::other(error.to_string())),
1678 )
1679}
1680
1681fn next_id(prefix: &str) -> String {
1682 let counter = NEXT_STORE_ID.fetch_add(1, Ordering::Relaxed);
1683 format!("{prefix}-{:x}-{:x}", now_nanos(), counter)
1684}
1685
1686fn now_secs() -> i64 {
1687 SystemTime::now()
1688 .duration_since(UNIX_EPOCH)
1689 .unwrap_or_default()
1690 .as_secs() as i64
1691}
1692
1693fn now_nanos() -> u128 {
1694 SystemTime::now()
1695 .duration_since(UNIX_EPOCH)
1696 .unwrap_or_default()
1697 .as_nanos()
1698}
1699
1700fn prune_stale_runtime_leases(tx: &rusqlite::Transaction<'_>) -> Result<(), RuntimeError> {
1701 let mut stmt = tx
1702 .prepare("SELECT key, owner FROM leases")
1703 .map_err(sqlite_error)?;
1704 let leases = stmt
1705 .query_map([], |row| {
1706 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
1707 })
1708 .map_err(sqlite_error)?
1709 .collect::<Result<Vec<_>, _>>()
1710 .map_err(sqlite_error)?;
1711 drop(stmt);
1712
1713 for (key, owner) in leases {
1714 if runtime_owner_is_stale(&owner) {
1715 tx.execute("DELETE FROM leases WHERE key = ?1", params![key])
1716 .map_err(sqlite_error)?;
1717 }
1718 }
1719
1720 Ok(())
1721}
1722
1723fn runtime_owner_is_stale(owner: &str) -> bool {
1724 let Some(pid) = owner
1725 .strip_prefix("runtime-")
1726 .and_then(|value| value.parse::<u32>().ok())
1727 else {
1728 return false;
1729 };
1730
1731 #[cfg(unix)]
1732 {
1733 let pid = pid as i32;
1734 let result = unsafe { libc::kill(pid, 0) };
1735 if result == 0 {
1736 return false;
1737 }
1738
1739 match std::io::Error::last_os_error().raw_os_error() {
1740 Some(code) if code == libc::ESRCH => true,
1741 Some(code) if code == libc::EPERM => false,
1742 _ => false,
1743 }
1744 }
1745
1746 #[cfg(windows)]
1747 {
1748 use windows_sys::Win32::{
1749 Foundation::{CloseHandle, STILL_ACTIVE},
1750 System::Threading::{
1751 GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
1752 },
1753 };
1754
1755 const ERROR_ACCESS_DENIED: i32 = 5;
1756 const ERROR_INVALID_PARAMETER: i32 = 87;
1757
1758 unsafe {
1759 let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
1760 if handle.is_null() {
1761 return match std::io::Error::last_os_error().raw_os_error() {
1762 Some(ERROR_INVALID_PARAMETER) => true,
1763 Some(ERROR_ACCESS_DENIED) => false,
1764 _ => false,
1765 };
1766 }
1767
1768 let mut exit_code = 0u32;
1769 let result = GetExitCodeProcess(handle, &mut exit_code);
1770 let close_result = CloseHandle(handle);
1771 debug_assert_ne!(close_result, 0, "process handle should close");
1772
1773 if result == 0 {
1774 return false;
1775 }
1776
1777 exit_code != STILL_ACTIVE as u32
1778 }
1779 }
1780
1781 #[cfg(not(any(unix, windows)))]
1782 {
1783 false
1784 }
1785}
1786
1787fn encode_runtime_identifier(runtime_identifier: &str) -> String {
1788 let mut encoded = String::with_capacity(runtime_identifier.len() * 2);
1789 for byte in runtime_identifier.as_bytes() {
1790 use std::fmt::Write as _;
1791 let _ = write!(&mut encoded, "{byte:02x}");
1792 }
1793 encoded
1794}
1795
1796fn decode_runtime_store_filename(filename: &str) -> Option<String> {
1797 let encoded = filename.strip_prefix("runtime-")?.strip_suffix(".sqlite")?;
1798 if encoded.len() % 2 != 0 || encoded.is_empty() {
1799 return None;
1800 }
1801
1802 let mut bytes = Vec::with_capacity(encoded.len() / 2);
1803 let mut index = 0;
1804 while index < encoded.len() {
1805 let byte = u8::from_str_radix(&encoded[index..index + 2], 16).ok()?;
1806 bytes.push(byte);
1807 index += 2;
1808 }
1809 String::from_utf8(bytes).ok()
1810}
1811
1812#[cfg(not(test))]
1813fn default_store_dir() -> PathBuf {
1814 crate::default_paths::workspace_default_paths().root_dir
1815}
1816
1817#[cfg(test)]
1818fn default_store_dir() -> PathBuf {
1819 let suffix = NEXT_TEST_STORE_ID.fetch_add(1, Ordering::Relaxed);
1820 std::env::temp_dir()
1821 .join("mentra-test-runtime")
1822 .join(format!("process-{}-{suffix}", std::process::id()))
1823}
1824
1825#[cfg(test)]
1826mod tests {
1827 use super::*;
1828 use crate::memory::{MemoryRecord, MemoryRecordKind, MemoryStore};
1829
1830 #[test]
1831 fn runtime_identifier_round_trips_through_filename_encoding() {
1832 let runtime_identifier = "chat/example 01";
1833 let filename = format!(
1834 "runtime-{}.sqlite",
1835 encode_runtime_identifier(runtime_identifier)
1836 );
1837 assert_eq!(
1838 decode_runtime_store_filename(&filename).as_deref(),
1839 Some(runtime_identifier)
1840 );
1841 }
1842
1843 #[test]
1844 fn path_for_runtime_identifier_uses_runtime_specific_filename() {
1845 let path = SqliteRuntimeStore::path_for_runtime_identifier("session-a");
1846 assert!(
1847 path.file_name()
1848 .and_then(|name| name.to_str())
1849 .is_some_and(|name| name.starts_with("runtime-"))
1850 );
1851 assert!(
1852 path.file_name()
1853 .and_then(|name| name.to_str())
1854 .is_some_and(|name| name.ends_with(".sqlite"))
1855 );
1856 }
1857
1858 #[test]
1859 fn stale_runtime_owner_can_be_reclaimed() {
1860 let store = SqliteRuntimeStore::new(
1861 std::env::temp_dir().join(format!("mentra-store-lease-{}.sqlite", now_nanos())),
1862 );
1863 let conn = Connection::open(store.path()).expect("open store");
1864 store.ensure_schema(&conn).expect("ensure schema");
1865 conn.execute(
1866 "INSERT INTO leases (key, owner, expires_at) VALUES (?1, ?2, ?3)",
1867 params!["agent:test", "runtime-999999", now_secs() + 3600],
1868 )
1869 .expect("insert stale lease");
1870
1871 let acquired = store
1872 .acquire_lease("agent:test", "runtime-123", Duration::from_secs(60))
1873 .expect("acquire lease");
1874 assert!(acquired);
1875 }
1876
1877 #[test]
1878 fn background_tasks_are_scoped_per_agent() {
1879 let store = SqliteRuntimeStore::new(
1880 std::env::temp_dir().join(format!("mentra-store-background-{}.sqlite", now_nanos())),
1881 );
1882
1883 store
1884 .upsert_background_task(
1885 "agent-a",
1886 &BackgroundTaskSummary {
1887 id: "bg-1".to_string(),
1888 command: "echo a".to_string(),
1889 cwd: std::env::temp_dir().join("a"),
1890 status: BackgroundTaskStatus::Running,
1891 output_preview: None,
1892 },
1893 DELIVERY_ACKED,
1894 )
1895 .expect("seed agent a background task");
1896 store
1897 .upsert_background_task(
1898 "agent-b",
1899 &BackgroundTaskSummary {
1900 id: "bg-1".to_string(),
1901 command: "echo b".to_string(),
1902 cwd: std::env::temp_dir().join("b"),
1903 status: BackgroundTaskStatus::Finished,
1904 output_preview: Some("done".to_string()),
1905 },
1906 DELIVERY_PENDING,
1907 )
1908 .expect("seed agent b background task");
1909
1910 let agent_a_tasks = store
1911 .load_background_tasks("agent-a")
1912 .expect("load agent a background tasks");
1913 let agent_b_tasks = store
1914 .load_background_tasks("agent-b")
1915 .expect("load agent b background tasks");
1916
1917 assert_eq!(agent_a_tasks.len(), 1);
1918 assert_eq!(agent_a_tasks[0].command, "echo a");
1919 assert_eq!(agent_a_tasks[0].status, BackgroundTaskStatus::Running);
1920 assert_eq!(agent_b_tasks.len(), 1);
1921 assert_eq!(agent_b_tasks[0].command, "echo b");
1922 assert_eq!(agent_b_tasks[0].status, BackgroundTaskStatus::Finished);
1923 }
1924
1925 #[test]
1926 fn fts_query_returns_none_when_input_has_no_searchable_terms() {
1927 assert_eq!(fts_query("... --- \"\""), None);
1928 }
1929
1930 #[test]
1931 fn sqlite_memory_search_sanitizes_punctuation_heavy_queries() {
1932 let store = SqliteRuntimeStore::new(
1933 std::env::temp_dir().join(format!("mentra-store-memory-{}.sqlite", now_nanos())),
1934 );
1935 store
1936 .upsert_records(&[MemoryRecord {
1937 record_id: "episode:agent:1".to_string(),
1938 agent_id: "agent-1".to_string(),
1939 kind: MemoryRecordKind::Episode,
1940 content: "shared phrase alpha".to_string(),
1941 source_revision: 1,
1942 created_at: 1,
1943 metadata_json: "{}".to_string(),
1944 source: Some("seed".to_string()),
1945 pinned: false,
1946 score: None,
1947 }])
1948 .expect("seed records");
1949
1950 let records = store
1951 .search_records("agent-1", "(shared) alpha!!!", 10)
1952 .expect("search records");
1953 assert_eq!(records.len(), 1);
1954 assert_eq!(records[0].record_id, "episode:agent:1");
1955 }
1956
1957 #[test]
1958 fn sqlite_memory_search_ignores_non_searchable_queries() {
1959 let store = SqliteRuntimeStore::new(
1960 std::env::temp_dir().join(format!("mentra-store-empty-query-{}.sqlite", now_nanos())),
1961 );
1962 store
1963 .upsert_records(&[MemoryRecord {
1964 record_id: "episode:agent:1".to_string(),
1965 agent_id: "agent-1".to_string(),
1966 kind: MemoryRecordKind::Episode,
1967 content: "shared phrase alpha".to_string(),
1968 source_revision: 1,
1969 created_at: 1,
1970 metadata_json: "{}".to_string(),
1971 source: Some("seed".to_string()),
1972 pinned: false,
1973 score: None,
1974 }])
1975 .expect("seed records");
1976
1977 let records = store
1978 .search_records("agent-1", "... ---", 10)
1979 .expect("search records");
1980 assert!(records.is_empty());
1981 }
1982
1983 fn permission_store() -> SqliteRuntimeStore {
1986 SqliteRuntimeStore::new(
1987 std::env::temp_dir().join(format!("mentra-store-perm-{}.sqlite", now_nanos())),
1988 )
1989 }
1990
1991 #[test]
1992 fn permission_rules_save_and_load_round_trip() {
1993 use crate::session::PermissionRuleScope;
1994
1995 let store = permission_store();
1996
1997 let session_rule = RememberedRule {
1999 key: RuleKey {
2000 tool_name: "shell".to_string(),
2001 pattern: None,
2002 },
2003 allow: true,
2004 scope: PermissionRuleScope::Session,
2005 };
2006 let project_rule = RememberedRule {
2008 key: RuleKey {
2009 tool_name: "read".to_string(),
2010 pattern: Some("/tmp/*".to_string()),
2011 },
2012 allow: false,
2013 scope: PermissionRuleScope::Project,
2014 };
2015
2016 store
2017 .save_rules("session-1", Some("proj-x"), &[session_rule, project_rule])
2018 .expect("save rules");
2019
2020 let loaded = store
2022 .load_rules("session-1", Some("proj-x"))
2023 .expect("load rules");
2024
2025 assert_eq!(loaded.len(), 2, "expected 2 rules, got {loaded:?}");
2026
2027 let shell_rule = loaded
2028 .iter()
2029 .find(|r| r.key.tool_name == "shell")
2030 .expect("shell rule present");
2031 assert!(shell_rule.allow);
2032 assert_eq!(shell_rule.scope, PermissionRuleScope::Session);
2033 assert_eq!(shell_rule.key.pattern, None);
2034
2035 let read_rule = loaded
2036 .iter()
2037 .find(|r| r.key.tool_name == "read")
2038 .expect("read rule present");
2039 assert!(!read_rule.allow);
2040 assert_eq!(read_rule.scope, PermissionRuleScope::Project);
2041 assert_eq!(read_rule.key.pattern, Some("/tmp/*".to_string()));
2042 }
2043
2044 #[test]
2045 fn permission_rules_clear_removes_all_for_session() {
2046 use crate::session::PermissionRuleScope;
2047
2048 let store = permission_store();
2049 let rules = vec![RememberedRule {
2050 key: RuleKey {
2051 tool_name: "shell".to_string(),
2052 pattern: None,
2053 },
2054 allow: true,
2055 scope: PermissionRuleScope::Session,
2056 }];
2057
2058 store
2059 .save_rules("session-1", None, &rules)
2060 .expect("save rules");
2061 store.clear_rules("session-1").expect("clear rules");
2062
2063 let loaded = store
2064 .load_rules("session-1", None)
2065 .expect("load rules after clear");
2066 assert!(loaded.is_empty());
2067 }
2068
2069 #[test]
2070 fn permission_rules_are_scoped_per_session() {
2071 use crate::session::PermissionRuleScope;
2072
2073 let store = permission_store();
2076 let rules_a = vec![RememberedRule {
2077 key: RuleKey {
2078 tool_name: "shell".to_string(),
2079 pattern: None,
2080 },
2081 allow: true,
2082 scope: PermissionRuleScope::Session,
2083 }];
2084 let rules_b = vec![RememberedRule {
2087 key: RuleKey {
2088 tool_name: "read".to_string(),
2089 pattern: None,
2090 },
2091 allow: false,
2092 scope: PermissionRuleScope::Project,
2093 }];
2094
2095 store
2096 .save_rules("session-a", None, &rules_a)
2097 .expect("save rules a");
2098 store
2099 .save_rules("session-b", Some("proj-b"), &rules_b)
2100 .expect("save rules b");
2101
2102 let loaded_a = store.load_rules("session-a", None).expect("load rules a");
2104 let loaded_b = store
2106 .load_rules("session-b", Some("proj-b"))
2107 .expect("load rules b");
2108
2109 assert_eq!(loaded_a.len(), 1, "session-a should have 1 rule");
2110 assert_eq!(loaded_a[0].key.tool_name, "shell");
2111 assert_eq!(loaded_b.len(), 1, "session-b should have 1 rule");
2112 assert_eq!(loaded_b[0].key.tool_name, "read");
2113 }
2114
2115 #[test]
2116 fn permission_rules_save_replaces_existing_session_rules() {
2117 use crate::session::PermissionRuleScope;
2118
2119 let store = permission_store();
2120 let initial = vec![RememberedRule {
2121 key: RuleKey {
2122 tool_name: "shell".to_string(),
2123 pattern: None,
2124 },
2125 allow: true,
2126 scope: PermissionRuleScope::Session,
2127 }];
2128
2129 store
2130 .save_rules("session-1", None, &initial)
2131 .expect("save initial");
2132
2133 let updated = vec![RememberedRule {
2135 key: RuleKey {
2136 tool_name: "write".to_string(),
2137 pattern: None,
2138 },
2139 allow: false,
2140 scope: PermissionRuleScope::Session,
2141 }];
2142
2143 store
2144 .save_rules("session-1", None, &updated)
2145 .expect("save updated");
2146
2147 let loaded = store
2148 .load_rules("session-1", None)
2149 .expect("load after replace");
2150 assert_eq!(loaded.len(), 1, "should have exactly 1 rule after replace");
2151 assert_eq!(loaded[0].key.tool_name, "write");
2152 assert!(!loaded[0].allow);
2153 }
2154
2155 #[test]
2156 fn permission_rules_load_returns_empty_for_unknown_session() {
2157 let store = permission_store();
2158 let loaded = store
2159 .load_rules("nonexistent", None)
2160 .expect("load unknown session");
2161 assert!(loaded.is_empty());
2162 }
2163}