@botpress/adk-cli
Version:
Command-line interface for the Botpress Agent Development Kit (ADK)
585 lines (577 loc) • 20.7 kB
JavaScript
// @bun
import {
createCliLogger
} from "./chunk-gzwt1qdr.js";
// src/utils/sqlite-trace-store.ts
import { Database } from "bun:sqlite";
import { mkdirSync } from "fs";
import { join } from "path";
var logger = createCliLogger({ tag: "sqlite-trace-store" });
var FLUSH_DELAY_MS = 16;
var FLUSH_MAX_BATCH = 250;
var CHECKPOINT_INTERVAL_MS = 5000;
function dataAttributeJsonPath(name) {
return `$."${name.replace(/\\/g, "\\\\").replace(/"/g, "\\\"")}"`;
}
var SPAN_PAYLOAD_TTL_MS = 7 * 24 * 60 * 60 * 1000;
var SPAN_PAYLOAD_PRUNE_INTERVAL_MS = 5 * 60 * 1000;
var SCHEMA = `
CREATE TABLE IF NOT EXISTS spans (
trace_id TEXT NOT NULL,
span_id TEXT NOT NULL,
parent_id TEXT,
name TEXT NOT NULL,
label TEXT NOT NULL,
status TEXT NOT NULL CHECK(status IN ('ok','error','running')),
error TEXT,
tier TEXT NOT NULL CHECK(tier IN ('concise','standard','verbose')),
started_at REAL NOT NULL,
ended_at REAL,
duration REAL,
bot_id TEXT,
conversation_id TEXT,
user_id TEXT,
message_id TEXT,
workflow_id TEXT,
event_id TEXT,
integration TEXT,
channel TEXT,
data TEXT NOT NULL DEFAULT '{}',
resource TEXT NOT NULL DEFAULT '{}',
context_json TEXT NOT NULL DEFAULT '{}',
updated_at REAL NOT NULL,
PRIMARY KEY (trace_id, span_id)
);
CREATE INDEX IF NOT EXISTS idx_spans_conversation ON spans(conversation_id) WHERE conversation_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_spans_workflow ON spans(workflow_id) WHERE workflow_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_spans_user ON spans(user_id) WHERE user_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_spans_message ON spans(message_id) WHERE message_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_spans_event ON spans(event_id) WHERE event_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_spans_started ON spans(started_at);
CREATE INDEX IF NOT EXISTS idx_spans_running ON spans(status) WHERE status = 'running';
CREATE INDEX IF NOT EXISTS idx_spans_updated ON spans(updated_at);
CREATE TABLE IF NOT EXISTS conversation_users (
bot_id TEXT NOT NULL,
conversation_id TEXT NOT NULL,
user_token TEXT NOT NULL,
created_at REAL NOT NULL,
PRIMARY KEY (bot_id, conversation_id)
);
CREATE TABLE IF NOT EXISTS span_payloads (
trace_id TEXT NOT NULL,
span_id TEXT NOT NULL,
attr_key TEXT NOT NULL,
content_type TEXT NOT NULL,
value TEXT NOT NULL,
size_bytes INTEGER NOT NULL,
created_at REAL NOT NULL,
updated_at REAL NOT NULL,
expires_at REAL NOT NULL,
PRIMARY KEY (trace_id, span_id, attr_key)
);
CREATE INDEX IF NOT EXISTS idx_span_payloads_expires ON span_payloads(expires_at);
`;
function rowToSpan(row) {
return {
id: {
trace: row.trace_id,
span: row.span_id,
parent: row.parent_id
},
name: row.name,
label: row.label,
status: row.status,
error: row.error ?? undefined,
timing: {
startedAt: row.started_at,
endedAt: row.ended_at ?? undefined,
duration: row.duration ?? undefined
},
context: JSON.parse(row.context_json || "{}"),
tier: row.tier,
data: JSON.parse(row.data || "{}"),
resource: JSON.parse(row.resource || "{}")
};
}
function rowToSpanPayload(row) {
return {
traceId: row.trace_id,
spanId: row.span_id,
key: row.attr_key,
contentType: row.content_type,
value: row.value,
sizeBytes: row.size_bytes,
createdAt: row.created_at,
updatedAt: row.updated_at,
expiresAt: row.expires_at
};
}
class SqliteTraceStore {
db = null;
dbPath;
tracesDir;
insertStmt = null;
updateStmt = null;
nextPayloadPruneAt = 0;
writeQueue = [];
flushTimer = null;
checkpointTimer = null;
flushListener = null;
constructor(agentPath) {
this.tracesDir = agentPath ? join(agentPath, ".adk", "bot", "traces") : process.env.TRACE_DIR ?? "traces";
this.dbPath = join(this.tracesDir, "traces.db");
}
getDb() {
if (this.db)
return this.db;
mkdirSync(this.tracesDir, { recursive: true });
try {
this.db = new Database(this.dbPath);
this.db.exec("PRAGMA journal_mode=WAL");
this.db.exec("PRAGMA synchronous=NORMAL");
this.db.exec("PRAGMA busy_timeout=5000");
this.db.exec("PRAGMA wal_autocheckpoint=0");
this.db.exec(SCHEMA);
this.ensureConversationUsersBotId(this.db);
this.startCheckpointTimer();
this.insertStmt = this.db.prepare(`
INSERT OR REPLACE INTO spans (
trace_id, span_id, parent_id, name, label, status, error, tier,
started_at, ended_at, duration,
bot_id, conversation_id, user_id, message_id, workflow_id, event_id, integration, channel,
data, resource, context_json, updated_at
) VALUES (
?, ?, ?, ?, ?, ?, ?, ?,
?, ?, ?,
?, ?, ?, ?, ?, ?, ?, ?,
?, ?, ?, ?
)
`);
this.updateStmt = this.db.prepare(`
UPDATE spans SET
label = ?, status = ?, error = ?,
ended_at = ?, duration = ?,
bot_id = ?, conversation_id = ?, user_id = ?, message_id = ?,
workflow_id = ?, event_id = ?, integration = ?, channel = ?,
data = ?, resource = ?, context_json = ?, updated_at = ?
WHERE trace_id = ? AND span_id = ?
`);
return this.db;
} catch (err) {
logger.warn(`Failed to initialize database: ${err instanceof Error ? err.message : String(err)}`);
return null;
}
}
ensureConversationUsersBotId(db) {
const columns = db.prepare("PRAGMA table_info(conversation_users)").all();
if (columns.some((c) => c.name === "bot_id")) {
db.exec("CREATE INDEX IF NOT EXISTS idx_conversation_users_lookup ON conversation_users(bot_id, conversation_id)");
return;
}
db.exec(`
BEGIN;
ALTER TABLE conversation_users RENAME TO conversation_users_legacy;
CREATE TABLE conversation_users (
bot_id TEXT NOT NULL,
conversation_id TEXT NOT NULL,
user_token TEXT NOT NULL,
created_at REAL NOT NULL,
PRIMARY KEY (bot_id, conversation_id)
);
CREATE INDEX IF NOT EXISTS idx_conversation_users_lookup ON conversation_users(bot_id, conversation_id);
COMMIT;
`);
}
setFlushListener(fn) {
this.flushListener = fn;
}
enqueueInsert(span, payloads = []) {
this.writeQueue.push({ op: "insert", span, payloads });
this.scheduleFlush();
}
enqueueUpdate(span, payloads = []) {
this.writeQueue.push({ op: "update", span, payloads });
this.scheduleFlush();
}
scheduleFlush() {
if (this.writeQueue.length >= FLUSH_MAX_BATCH) {
this.flush();
return;
}
if (this.flushTimer)
return;
this.flushTimer = setTimeout(() => {
this.flushTimer = null;
this.flush();
}, FLUSH_DELAY_MS);
this.flushTimer.unref?.();
}
flush() {
if (this.writeQueue.length === 0)
return;
if (this.flushTimer) {
clearTimeout(this.flushTimer);
this.flushTimer = null;
}
const batch = this.writeQueue;
this.writeQueue = [];
const db = this.getDb();
if (!db)
return;
try {
if (batch.some((item) => item.payloads.length > 0)) {
this.pruneExpiredSpanPayloads();
}
db.transaction((items) => {
for (const { op, span, payloads } of items) {
try {
if (op === "insert")
this.insertSpan(span);
else
this.updateSpan(span);
this.writeSpanPayloads(span.id.trace, span.id.span, payloads);
} catch (err) {
logger.warn(`Failed to write span ${span.id.span}: ${err instanceof Error ? err.message : String(err)}`);
}
}
})(batch);
if (this.flushListener) {
try {
this.flushListener(batch.map((b) => b.span));
} catch {}
}
} catch (err) {
logger.warn(`Span batch flush failed: ${err instanceof Error ? err.message : String(err)}`);
}
}
startCheckpointTimer() {
if (this.checkpointTimer)
return;
this.checkpointTimer = setInterval(() => {
try {
this.db?.exec("PRAGMA wal_checkpoint(PASSIVE)");
} catch {}
}, CHECKPOINT_INTERVAL_MS);
this.checkpointTimer.unref?.();
}
insertSpan(span) {
const db = this.getDb();
if (!db || !this.insertStmt)
return;
const now = Date.now();
this.insertStmt.run(span.id.trace, span.id.span, span.id.parent, span.name, span.label, span.status, span.error ?? null, span.tier, span.timing.startedAt, span.timing.endedAt ?? null, span.timing.duration ?? null, span.context.botId ?? null, span.context.conversationId ?? null, span.context.userId ?? null, span.context.messageId ?? null, span.context.workflowId ?? null, span.context.eventId ?? null, span.context.integration ?? null, span.context.channel ?? null, JSON.stringify(span.data), JSON.stringify(span.resource), JSON.stringify(span.context), now);
}
updateSpan(span) {
const db = this.getDb();
if (!db || !this.updateStmt)
return;
const now = Date.now();
this.updateStmt.run(span.label, span.status, span.error ?? null, span.timing.endedAt ?? null, span.timing.duration ?? null, span.context.botId ?? null, span.context.conversationId ?? null, span.context.userId ?? null, span.context.messageId ?? null, span.context.workflowId ?? null, span.context.eventId ?? null, span.context.integration ?? null, span.context.channel ?? null, JSON.stringify(span.data), JSON.stringify(span.resource), JSON.stringify(span.context), now, span.id.trace, span.id.span);
if (db.changes === 0) {
this.insertSpan(span);
}
}
upsertSpanPayloads(traceId, spanId, payloads) {
if (payloads.length === 0)
return;
const db = this.getDb();
if (!db)
return;
this.pruneExpiredSpanPayloads();
db.transaction((items) => {
this.writeSpanPayloads(traceId, spanId, items);
})(payloads);
}
writeSpanPayloads(traceId, spanId, payloads) {
if (payloads.length === 0)
return;
const db = this.getDb();
if (!db)
return;
const now = Date.now();
const expiresAt = now + SPAN_PAYLOAD_TTL_MS;
const stmt = db.prepare(`
INSERT INTO span_payloads (
trace_id, span_id, attr_key, content_type, value, size_bytes, created_at, updated_at, expires_at
) VALUES (
?, ?, ?, ?, ?, ?, ?, ?, ?
)
ON CONFLICT(trace_id, span_id, attr_key) DO UPDATE SET
content_type = excluded.content_type,
value = excluded.value,
size_bytes = excluded.size_bytes,
updated_at = excluded.updated_at,
expires_at = excluded.expires_at
`);
for (const payload of payloads) {
stmt.run(traceId, spanId, payload.key, payload.contentType, payload.value, payload.sizeBytes, now, now, expiresAt);
}
}
getSpanPayload(traceId, spanId, key) {
this.flush();
const db = this.getDb();
if (!db)
return null;
const now = Date.now();
const row = db.prepare("SELECT * FROM span_payloads WHERE trace_id = ? AND span_id = ? AND attr_key = ? AND expires_at > ?").get(traceId, spanId, key, now);
if (row) {
return rowToSpanPayload(row);
}
db.prepare("DELETE FROM span_payloads WHERE trace_id = ? AND span_id = ? AND attr_key = ? AND expires_at <= ?").run(traceId, spanId, key, now);
return null;
}
pruneExpiredSpanPayloads(now = Date.now()) {
if (now < this.nextPayloadPruneAt)
return 0;
const db = this.getDb();
if (!db)
return 0;
db.prepare("DELETE FROM span_payloads WHERE expires_at <= ?").run(now);
this.nextPayloadPruneAt = now + SPAN_PAYLOAD_PRUNE_INTERVAL_MS;
return db.changes;
}
getTraceById(traceId) {
this.flush();
const db = this.getDb();
if (!db)
return [];
const rows = db.prepare("SELECT * FROM spans WHERE trace_id = ? ORDER BY started_at").all(traceId);
return rows.map(rowToSpan);
}
getRecentTraces(limit = 100) {
this.flush();
const db = this.getDb();
if (!db)
return [];
const rows = db.prepare("SELECT * FROM spans ORDER BY started_at DESC LIMIT ?").all(limit);
return rows.map(rowToSpan);
}
getTraceByAttribute(name, value, startTs, endTs) {
this.flush();
const db = this.getDb();
if (!db)
return [];
const indexedColumns = {
conversationId: "conversation_id",
messageId: "message_id",
eventId: "event_id",
userId: "user_id",
workflowId: "workflow_id",
botId: "bot_id",
integration: "integration",
channel: "channel"
};
const column = indexedColumns[name];
let traceIds;
if (column) {
let traceQuery = `SELECT DISTINCT trace_id FROM spans WHERE ${column} = ?`;
const params = [value];
if (startTs) {
traceQuery += " AND started_at >= ?";
params.push(startTs);
}
if (endTs) {
traceQuery += " AND started_at <= ?";
params.push(endTs);
}
const traceRows = db.prepare(traceQuery).all(...params);
traceIds = traceRows.map((r) => r.trace_id);
} else {
let traceQuery = `SELECT DISTINCT trace_id FROM spans WHERE CAST(json_extract(data, ?) AS TEXT) = ?`;
const params = [dataAttributeJsonPath(name), value];
if (startTs) {
traceQuery += " AND started_at >= ?";
params.push(startTs);
}
if (endTs) {
traceQuery += " AND started_at <= ?";
params.push(endTs);
}
const traceRows = db.prepare(traceQuery).all(...params);
traceIds = traceRows.map((r) => r.trace_id);
}
if (traceIds.length === 0)
return [];
const placeholders = traceIds.map(() => "?").join(",");
const rows = db.prepare(`SELECT * FROM spans WHERE trace_id IN (${placeholders}) ORDER BY started_at`).all(...traceIds);
return rows.map(rowToSpan);
}
getChangedSince(ts, traceIds) {
this.flush();
const db = this.getDb();
if (!db)
return [];
if (traceIds && traceIds.length > 0) {
const placeholders = traceIds.map(() => "?").join(",");
const rows2 = db.prepare(`SELECT * FROM spans WHERE updated_at >= ? AND trace_id IN (${placeholders}) ORDER BY started_at`).all(ts, ...traceIds);
return rows2.map(rowToSpan);
}
const rows = db.prepare("SELECT * FROM spans WHERE updated_at >= ? ORDER BY started_at").all(ts);
return rows.map(rowToSpan);
}
getAllTodayTraces() {
this.flush();
const db = this.getDb();
if (!db)
return [];
const startOfDay = new Date;
startOfDay.setHours(0, 0, 0, 0);
const rows = db.prepare("SELECT * FROM spans WHERE started_at >= ? ORDER BY started_at").all(startOfDay.getTime());
return rows.map(rowToSpan);
}
setConversationUser(botId, conversationId, userToken) {
const db = this.getDb();
if (!db)
return;
db.prepare("INSERT OR REPLACE INTO conversation_users (bot_id, conversation_id, user_token, created_at) VALUES (?, ?, ?, ?)").run(botId, conversationId, userToken, Date.now());
}
getConversationUser(botId, conversationId) {
const db = this.getDb();
if (!db)
return null;
const row = db.prepare("SELECT user_token FROM conversation_users WHERE bot_id = ? AND conversation_id = ?").get(botId, conversationId);
return row?.user_token ?? null;
}
listConversations(options) {
this.flush();
const db = this.getDb();
if (!db)
return [];
const limit = options.limit ?? 20;
const params = [];
let sql = `
SELECT
conversation_id,
MIN(started_at) as first_seen,
MAX(started_at) as last_seen,
COUNT(DISTINCT trace_id) as trace_count,
MAX(integration) as integration,
MAX(channel) as channel,
MAX(CASE WHEN status = 'error' THEN 1 ELSE 0 END) as has_errors
FROM spans
WHERE conversation_id IS NOT NULL
`;
if (options.botId) {
sql += " AND bot_id = ?";
params.push(options.botId);
}
sql += " GROUP BY conversation_id";
if (options.since != null) {
sql += " HAVING MAX(started_at) >= ?";
params.push(options.since);
}
sql += " ORDER BY last_seen DESC LIMIT ?";
params.push(limit);
const rows = db.prepare(sql).all(...params);
return rows.map((r) => ({
conversationId: r.conversation_id,
firstSeen: r.first_seen,
lastSeen: r.last_seen,
traceCount: r.trace_count,
integration: r.integration,
channel: r.channel,
hasErrors: r.has_errors === 1
}));
}
sweepRunningSpans() {
const db = this.getDb();
if (!db)
return 0;
const now = Date.now();
db.prepare(`
UPDATE spans
SET status = 'error',
error = 'Process terminated',
ended_at = COALESCE(ended_at, ?),
duration = COALESCE(ended_at, ?) - started_at,
updated_at = ?
WHERE status = 'running'
`).run(now, now, now);
return db.changes;
}
queryTraceIds(filters) {
this.flush();
const db = this.getDb();
if (!db)
return [];
let sql = "SELECT trace_id FROM spans WHERE parent_id IS NULL";
const params = [];
if (filters.since != null) {
sql += " AND started_at >= ?";
params.push(filters.since);
}
if (filters.until != null) {
sql += " AND started_at <= ?";
params.push(filters.until);
}
if (filters.status === "error") {
sql += " AND trace_id IN (SELECT DISTINCT trace_id FROM spans WHERE status = 'error')";
}
if (filters.conversations && filters.conversations.length > 0) {
const placeholders = filters.conversations.map(() => "?").join(",");
sql += ` AND trace_id IN (SELECT DISTINCT trace_id FROM spans WHERE conversation_id IN (${placeholders}))`;
params.push(...filters.conversations);
}
if (filters.workflows && filters.workflows.length > 0) {
const placeholders = filters.workflows.map(() => "?").join(",");
sql += ` AND trace_id IN (SELECT DISTINCT trace_id FROM spans WHERE name = 'handler.workflow' AND json_extract(data, '$.workflow.name') IN (${placeholders}))`;
params.push(...filters.workflows);
}
if (filters.actions && filters.actions.length > 0) {
const placeholders = filters.actions.map(() => "?").join(",");
sql += ` AND trace_id IN (SELECT DISTINCT trace_id FROM spans WHERE (name = 'handler.action' AND json_extract(data, '$.action.name') IN (${placeholders})) OR (name = 'autonomous.tool' AND json_extract(data, '$.autonomous.tool.name') IN (${placeholders})))`;
params.push(...filters.actions, ...filters.actions);
}
if (filters.triggers && filters.triggers.length > 0) {
const placeholders = filters.triggers.map(() => "?").join(",");
sql += ` AND trace_id IN (SELECT DISTINCT trace_id FROM spans WHERE name = 'handler.trigger' AND json_extract(data, '$.trigger.name') IN (${placeholders}))`;
params.push(...filters.triggers);
}
const limit = filters.limit ?? 20;
sql += " ORDER BY started_at DESC";
sql += " LIMIT ?";
params.push(limit);
const rows = db.prepare(sql).all(...params);
return rows.map((r) => r.trace_id);
}
getCompletedRootSpansSince(ts) {
this.flush();
const db = this.getDb();
if (!db)
return [];
const rows = db.prepare("SELECT * FROM spans WHERE parent_id IS NULL AND status != 'running' AND updated_at >= ? ORDER BY started_at").all(ts);
return rows.map(rowToSpan);
}
findTraceIdsByPrefix(prefix, limit = 10) {
this.flush();
const db = this.getDb();
if (!db)
return [];
if (prefix.length === 0)
return [];
const upperBound = prefix.slice(0, -1) + String.fromCharCode(prefix.charCodeAt(prefix.length - 1) + 1);
const rows = db.prepare("SELECT DISTINCT trace_id FROM spans WHERE trace_id >= ? AND trace_id < ? ORDER BY started_at DESC LIMIT ?").all(prefix, upperBound, limit);
return rows.map((r) => r.trace_id);
}
close() {
if (this.flushTimer) {
clearTimeout(this.flushTimer);
this.flushTimer = null;
}
if (this.checkpointTimer) {
clearInterval(this.checkpointTimer);
this.checkpointTimer = null;
}
try {
this.flush();
this.db?.exec("PRAGMA wal_checkpoint(TRUNCATE)");
} catch {}
try {
this.db?.close();
this.db = null;
this.insertStmt = null;
this.updateStmt = null;
} catch {}
}
}
export { SqliteTraceStore };