1use crate::error::Result;
2use rusqlite::Connection;
3
4pub const SCHEMA_VERSION: &str = "5";
5
6pub fn bootstrap(conn: &Connection) -> Result<()> {
7 conn.execute_batch(
8 "
9 PRAGMA journal_mode=WAL;
10 PRAGMA foreign_keys=ON;
11 PRAGMA synchronous=NORMAL;
12
13 CREATE TABLE IF NOT EXISTS _adb_meta (
14 key TEXT PRIMARY KEY,
15 value TEXT NOT NULL
16 );
17
18 CREATE TABLE IF NOT EXISTS _adb_collections (
19 id TEXT PRIMARY KEY,
20 name TEXT UNIQUE NOT NULL,
21 dim INTEGER NOT NULL,
22 metric TEXT NOT NULL DEFAULT 'cosine',
23 count INTEGER NOT NULL DEFAULT 0,
24 created_at INTEGER NOT NULL
25 );
26
27 CREATE TABLE IF NOT EXISTS _adb_vectors (
28 id TEXT NOT NULL,
29 collection_id TEXT NOT NULL,
30 vector BLOB NOT NULL,
31 metadata TEXT,
32 created_at INTEGER NOT NULL,
33 updated_at INTEGER NOT NULL DEFAULT 0,
34 PRIMARY KEY (id, collection_id),
35 FOREIGN KEY (collection_id)
36 REFERENCES _adb_collections(id) ON DELETE CASCADE
37 );
38
39 CREATE TABLE IF NOT EXISTS _adb_hnsw_index (
40 collection_id TEXT PRIMARY KEY,
41 index_blob BLOB NOT NULL,
42 built_at INTEGER NOT NULL,
43 is_dirty INTEGER NOT NULL DEFAULT 0
44 );
45
46 CREATE TABLE IF NOT EXISTS _adb_nodes (
47 id TEXT PRIMARY KEY,
48 kind TEXT NOT NULL,
49 data TEXT,
50 created_at INTEGER NOT NULL,
51 updated_at INTEGER NOT NULL
52 );
53
54 CREATE TABLE IF NOT EXISTS _adb_edges (
55 src TEXT NOT NULL,
56 dst TEXT NOT NULL,
57 relation TEXT NOT NULL,
58 weight REAL NOT NULL DEFAULT 1.0,
59 created_at INTEGER NOT NULL,
60 PRIMARY KEY (src, dst, relation),
61 FOREIGN KEY (src) REFERENCES _adb_nodes(id) ON DELETE CASCADE,
62 FOREIGN KEY (dst) REFERENCES _adb_nodes(id) ON DELETE CASCADE
63 );
64
65 CREATE INDEX IF NOT EXISTS idx_edges_src ON _adb_edges(src);
66 CREATE INDEX IF NOT EXISTS idx_edges_dst ON _adb_edges(dst);
67 CREATE INDEX IF NOT EXISTS idx_vectors_col ON _adb_vectors(collection_id);
68 CREATE INDEX IF NOT EXISTS idx_nodes_kind ON _adb_nodes(kind);
69
70 -- Conversations / message threading
71 CREATE TABLE IF NOT EXISTS _adb_conversations (
72 id TEXT PRIMARY KEY,
73 title TEXT,
74 metadata TEXT,
75 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
76 updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
77 );
78
79 CREATE TABLE IF NOT EXISTS _adb_messages (
80 id TEXT PRIMARY KEY,
81 conversation_id TEXT NOT NULL REFERENCES _adb_conversations(id) ON DELETE CASCADE,
82 role TEXT NOT NULL,
83 content TEXT NOT NULL,
84 metadata TEXT,
85 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
86 );
87 CREATE INDEX IF NOT EXISTS idx_messages_conv ON _adb_messages(conversation_id, created_at);
88
89 -- Workflow persistence
90 CREATE TABLE IF NOT EXISTS _adb_workflows (
91 id TEXT PRIMARY KEY,
92 name TEXT NOT NULL,
93 status TEXT NOT NULL DEFAULT 'pending',
94 input TEXT,
95 output TEXT,
96 error TEXT,
97 metadata TEXT,
98 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
99 updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
100 );
101
102 CREATE TABLE IF NOT EXISTS _adb_workflow_steps (
103 id TEXT PRIMARY KEY,
104 workflow_id TEXT NOT NULL REFERENCES _adb_workflows(id) ON DELETE CASCADE,
105 step_index INTEGER NOT NULL,
106 name TEXT NOT NULL,
107 status TEXT NOT NULL DEFAULT 'pending',
108 input TEXT,
109 output TEXT,
110 error TEXT,
111 started_at INTEGER,
112 completed_at INTEGER
113 );
114 CREATE INDEX IF NOT EXISTS idx_wf_steps ON _adb_workflow_steps(workflow_id, step_index);
115
116 -- Reasoning traces
117 CREATE TABLE IF NOT EXISTS _adb_traces (
118 id TEXT PRIMARY KEY,
119 session_id TEXT,
120 parent_id TEXT REFERENCES _adb_traces(id),
121 trace_type TEXT NOT NULL,
122 content TEXT NOT NULL,
123 metadata TEXT,
124 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
125 );
126 CREATE INDEX IF NOT EXISTS idx_traces_session ON _adb_traces(session_id, created_at);
127 CREATE INDEX IF NOT EXISTS idx_traces_parent ON _adb_traces(parent_id);
128
129 -- Message full-text search (schema v4)
130 CREATE VIRTUAL TABLE IF NOT EXISTS _adb_messages_fts
131 USING fts5(
132 message_id UNINDEXED,
133 conversation_id UNINDEXED,
134 content,
135 tokenize='porter ascii'
136 );
137
138 -- Tool registry (schema v5)
139 CREATE TABLE IF NOT EXISTS _adb_tools (
140 id TEXT PRIMARY KEY,
141 name TEXT UNIQUE NOT NULL,
142 description TEXT,
143 parameters_schema TEXT,
144 version TEXT NOT NULL DEFAULT '1.0.0',
145 created_at INTEGER NOT NULL,
146 updated_at INTEGER NOT NULL
147 );
148 CREATE INDEX IF NOT EXISTS idx_tools_name ON _adb_tools(name);
149
150 -- Structured tool call log (schema v5)
151 CREATE TABLE IF NOT EXISTS _adb_tool_calls (
152 id TEXT PRIMARY KEY,
153 session_id TEXT,
154 tool_name TEXT NOT NULL,
155 arguments TEXT,
156 result TEXT,
157 error TEXT,
158 latency_ms INTEGER,
159 created_at INTEGER NOT NULL
160 );
161 CREATE INDEX IF NOT EXISTS idx_tool_calls_session ON _adb_tool_calls(session_id, created_at);
162 CREATE INDEX IF NOT EXISTS idx_tool_calls_tool_name ON _adb_tool_calls(tool_name, created_at);
163
164 -- Immutable audit log (schema v5)
165 CREATE TABLE IF NOT EXISTS _adb_audit_log (
166 id TEXT PRIMARY KEY,
167 timestamp INTEGER NOT NULL,
168 actor TEXT,
169 action TEXT NOT NULL,
170 table_name TEXT NOT NULL,
171 record_id TEXT NOT NULL,
172 old_value TEXT,
173 new_value TEXT,
174 reason TEXT
175 );
176 CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON _adb_audit_log(timestamp);
177 CREATE INDEX IF NOT EXISTS idx_audit_table_rec ON _adb_audit_log(table_name, record_id);
178 CREATE INDEX IF NOT EXISTS idx_audit_actor ON _adb_audit_log(actor, timestamp);
179
180 -- Token-budgeted context window entries (schema v5)
181 CREATE TABLE IF NOT EXISTS _adb_context_entries (
182 id TEXT PRIMARY KEY,
183 session_id TEXT NOT NULL,
184 source_type TEXT NOT NULL,
185 source_id TEXT NOT NULL,
186 content_preview TEXT,
187 token_count INTEGER NOT NULL DEFAULT 0,
188 relevance_score REAL NOT NULL DEFAULT 0.0,
189 priority INTEGER NOT NULL DEFAULT 0,
190 included_at INTEGER NOT NULL
191 );
192 CREATE INDEX IF NOT EXISTS idx_ctx_session ON _adb_context_entries(session_id, priority DESC, relevance_score DESC);
193 CREATE INDEX IF NOT EXISTS idx_ctx_source ON _adb_context_entries(source_type, source_id);
194
195 -- Versioned prompt templates (schema v5)
196 CREATE TABLE IF NOT EXISTS _adb_prompt_templates (
197 id TEXT PRIMARY KEY,
198 name TEXT NOT NULL,
199 version INTEGER NOT NULL DEFAULT 1,
200 template TEXT NOT NULL,
201 model_hint TEXT,
202 max_tokens INTEGER,
203 metadata TEXT,
204 created_at INTEGER NOT NULL,
205 UNIQUE (name, version)
206 );
207 CREATE INDEX IF NOT EXISTS idx_prompt_name_ver ON _adb_prompt_templates(name, version DESC);
208
209 -- Privacy / data classification labels (schema v5)
210 CREATE TABLE IF NOT EXISTS _adb_data_labels (
211 table_name TEXT NOT NULL,
212 record_id TEXT NOT NULL,
213 label TEXT NOT NULL,
214 tagged_by TEXT,
215 tagged_at INTEGER NOT NULL,
216 PRIMARY KEY (table_name, record_id, label)
217 );
218 CREATE INDEX IF NOT EXISTS idx_data_labels_table ON _adb_data_labels(table_name, label);
219 CREATE INDEX IF NOT EXISTS idx_data_labels_label ON _adb_data_labels(label);
220 ",
221 )?;
222
223 conn.execute(
224 "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('schema_version', ?1)",
225 rusqlite::params![SCHEMA_VERSION],
226 )?;
227 conn.execute(
228 "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('created_at', ?1)",
229 rusqlite::params![now_ms()],
230 )?;
231 Ok(())
232}
233
234pub fn check_version(conn: &Connection) -> Result<()> {
235 let version: Option<String> = conn
236 .query_row(
237 "SELECT value FROM _adb_meta WHERE key = 'schema_version'",
238 [],
239 |row| row.get(0),
240 )
241 .ok();
242 match version.as_deref() {
243 Some(v) if v == SCHEMA_VERSION => Ok(()),
244 Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
247 None => Err(crate::error::AgentDbError::SchemaMigration),
249 }
250}
251
252pub fn migrate(conn: &Connection) -> Result<()> {
262 bootstrap(conn)?;
264
265 let _ = conn.execute_batch("ALTER TABLE _adb_workflows ADD COLUMN error TEXT;");
267
268 let _ = conn.execute_batch(
270 "ALTER TABLE _adb_vectors ADD COLUMN updated_at INTEGER NOT NULL DEFAULT 0;",
271 );
272
273 let _ = conn.execute_batch("ALTER TABLE _adb_vectors ADD COLUMN model TEXT;");
278
279 let _ = conn.execute_batch("ALTER TABLE _adb_messages ADD COLUMN token_count INTEGER;");
281
282 conn.execute(
284 "INSERT INTO _adb_meta (key, value) VALUES ('schema_version', ?1)
285 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
286 rusqlite::params![SCHEMA_VERSION],
287 )?;
288 Ok(())
289}
290
291pub fn now_ms() -> i64 {
292 std::time::SystemTime::now()
293 .duration_since(std::time::UNIX_EPOCH)
294 .unwrap_or_default()
295 .as_millis() as i64
296}