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        -- 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        // Older or newer schema on disk: caller must run `agentdb migrate`
245        // (or call `schema::migrate(conn)` programmatically).
246        Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
247        // Missing version key indicates a corrupt or pre-v0.1 database.
248        None => Err(crate::error::AgentDbError::SchemaMigration),
249    }
250}
251
252/// Idempotent migration runner.
253///
254/// Re-runs `bootstrap()` (all DDL uses `CREATE … IF NOT EXISTS` so existing
255/// tables/columns are left intact), then stamps the current schema version.
256/// Safe to call on databases created by any prior version of AgentDB.
257///
258/// In addition to adding new tables, this applies additive `ALTER TABLE …
259/// ADD COLUMN` statements for columns introduced in later schema versions.
260/// SQLite ignores duplicate columns via `IF NOT EXISTS` semantics.
261pub fn migrate(conn: &Connection) -> Result<()> {
262    // Re-run bootstrap to create any tables introduced after the DB was first opened.
263    bootstrap(conn)?;
264
265    // v2 → v3: add `error` column to _adb_workflows (was missing before v0.5.0).
266    let _ = conn.execute_batch("ALTER TABLE _adb_workflows ADD COLUMN error TEXT;");
267
268    // v2 → v3: add `updated_at` column to _adb_vectors.
269    let _ = conn.execute_batch(
270        "ALTER TABLE _adb_vectors ADD COLUMN updated_at INTEGER NOT NULL DEFAULT 0;",
271    );
272
273    // v3 → v4: message FTS virtual table (bootstrap already uses IF NOT EXISTS).
274    // No ALTER TABLE needed — new table is created by bootstrap() above.
275
276    // v4 → v5: add embedding model provenance to vectors.
277    let _ = conn.execute_batch("ALTER TABLE _adb_vectors ADD COLUMN model TEXT;");
278
279    // v4 → v5: add token count to messages for context budgeting.
280    let _ = conn.execute_batch("ALTER TABLE _adb_messages ADD COLUMN token_count INTEGER;");
281
282    // Stamp the new version.
283    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}