Skip to main content

agentdb/
schema.rs

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        -- Enforce immutability: prevent DELETE and UPDATE on the audit log
181        CREATE TRIGGER IF NOT EXISTS _adb_audit_no_delete
182            BEFORE DELETE ON _adb_audit_log
183            BEGIN SELECT RAISE(ABORT, 'audit log is immutable: DELETE not permitted'); END;
184        CREATE TRIGGER IF NOT EXISTS _adb_audit_no_update
185            BEFORE UPDATE ON _adb_audit_log
186            BEGIN SELECT RAISE(ABORT, 'audit log is immutable: UPDATE not permitted'); END;
187
188        -- Token-budgeted context window entries (schema v5)
189        CREATE TABLE IF NOT EXISTS _adb_context_entries (
190            id              TEXT PRIMARY KEY,
191            session_id      TEXT NOT NULL,
192            source_type     TEXT NOT NULL,
193            source_id       TEXT NOT NULL,
194            content_preview TEXT,
195            token_count     INTEGER NOT NULL DEFAULT 0,
196            relevance_score REAL NOT NULL DEFAULT 0.0,
197            priority        INTEGER NOT NULL DEFAULT 0,
198            included_at     INTEGER NOT NULL
199        );
200        CREATE INDEX IF NOT EXISTS idx_ctx_session ON _adb_context_entries(session_id, priority DESC, relevance_score DESC);
201        CREATE INDEX IF NOT EXISTS idx_ctx_source  ON _adb_context_entries(source_type, source_id);
202
203        -- Versioned prompt templates (schema v5)
204        CREATE TABLE IF NOT EXISTS _adb_prompt_templates (
205            id         TEXT PRIMARY KEY,
206            name       TEXT NOT NULL,
207            version    INTEGER NOT NULL DEFAULT 1,
208            template   TEXT NOT NULL,
209            model_hint TEXT,
210            max_tokens INTEGER,
211            metadata   TEXT,
212            created_at INTEGER NOT NULL,
213            UNIQUE (name, version)
214        );
215        CREATE INDEX IF NOT EXISTS idx_prompt_name_ver ON _adb_prompt_templates(name, version DESC);
216
217        -- Privacy / data classification labels (schema v5)
218        CREATE TABLE IF NOT EXISTS _adb_data_labels (
219            table_name TEXT NOT NULL,
220            record_id  TEXT NOT NULL,
221            label      TEXT NOT NULL,
222            tagged_by  TEXT,
223            tagged_at  INTEGER NOT NULL,
224            PRIMARY KEY (table_name, record_id, label)
225        );
226        CREATE INDEX IF NOT EXISTS idx_data_labels_table ON _adb_data_labels(table_name, label);
227        CREATE INDEX IF NOT EXISTS idx_data_labels_label ON _adb_data_labels(label);
228        ",
229    )?;
230
231    conn.execute(
232        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('schema_version', ?1)",
233        rusqlite::params![SCHEMA_VERSION],
234    )?;
235    conn.execute(
236        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('created_at', ?1)",
237        rusqlite::params![now_ms()],
238    )?;
239    Ok(())
240}
241
242pub fn check_version(conn: &Connection) -> Result<()> {
243    let version: Option<String> = conn
244        .query_row(
245            "SELECT value FROM _adb_meta WHERE key = 'schema_version'",
246            [],
247            |row| row.get(0),
248        )
249        .ok();
250    match version.as_deref() {
251        Some(v) if v == SCHEMA_VERSION => Ok(()),
252        // Older or newer schema on disk: caller must run `agentdb migrate`
253        // (or call `schema::migrate(conn)` programmatically).
254        Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
255        // Missing version key indicates a corrupt or pre-v0.1 database.
256        None => Err(crate::error::AgentDbError::SchemaMigration),
257    }
258}
259
260/// Idempotent migration runner.
261///
262/// Re-runs `bootstrap()` (all DDL uses `CREATE … IF NOT EXISTS` so existing
263/// tables/columns are left intact), then stamps the current schema version.
264/// Safe to call on databases created by any prior version of AgentDB.
265///
266/// In addition to adding new tables, this applies additive `ALTER TABLE …
267/// ADD COLUMN` statements for columns introduced in later schema versions.
268/// SQLite ignores duplicate columns via `IF NOT EXISTS` semantics.
269pub fn migrate(conn: &Connection) -> Result<()> {
270    // Re-run bootstrap to create any tables introduced after the DB was first opened.
271    bootstrap(conn)?;
272
273    // v2 → v3: add `error` column to _adb_workflows (was missing before v0.5.0).
274    let _ = conn.execute_batch("ALTER TABLE _adb_workflows ADD COLUMN error TEXT;");
275
276    // v2 → v3: add `updated_at` column to _adb_vectors.
277    let _ = conn.execute_batch(
278        "ALTER TABLE _adb_vectors ADD COLUMN updated_at INTEGER NOT NULL DEFAULT 0;",
279    );
280
281    // v3 → v4: message FTS virtual table (bootstrap already uses IF NOT EXISTS).
282    // No ALTER TABLE needed — new table is created by bootstrap() above.
283
284    // v4 → v5: add embedding model provenance to vectors.
285    let _ = conn.execute_batch("ALTER TABLE _adb_vectors ADD COLUMN model TEXT;");
286
287    // v4 → v5: add token count to messages for context budgeting.
288    let _ = conn.execute_batch("ALTER TABLE _adb_messages ADD COLUMN token_count INTEGER;");
289
290    // Stamp the new version.
291    conn.execute(
292        "INSERT INTO _adb_meta (key, value) VALUES ('schema_version', ?1)
293         ON CONFLICT(key) DO UPDATE SET value = excluded.value",
294        rusqlite::params![SCHEMA_VERSION],
295    )?;
296    Ok(())
297}
298
299pub fn now_ms() -> i64 {
300    std::time::SystemTime::now()
301        .duration_since(std::time::UNIX_EPOCH)
302        .unwrap_or_default()
303        .as_millis() as i64
304}