@codecaine-ai/agent-kernel-db 0.0.0-stage → 0.0.2

package/src/actions/read-api.ts

@@ -0,0 +1,314 @@
+/**
+ * Read API helpers — container-first. All reads are keyed by containerId;
+ * there is no app-session identity in this package.
+ */
+import { and, asc, count, desc, eq, gt, inArray, max, type SQL } from "drizzle-orm";
+
+import type { KernelDatabase } from "../client";
+import { piAgentSessionSelection } from "../read-compat";
+import { agentRuns } from "../schema/agent-runs";
+import { containers } from "../schema/containers";
+import { piAgentSessions } from "../schema/pi-agent-sessions";
+import { traceEvents } from "../schema/trace-events";
+import type {
+  AgentRun,
+  Container,
+  PiAgentSession,
+  TraceEventRow,
+} from "../types";
+
+export interface KernelTraceReadOptions {
+  after?: string | null;
+  limit?: number;
+  maxContainers?: number;
+}
+
+/**
+ * `kind` is optional here so callers that build rows by hand (fixtures in
+ * linked consumers) keep compiling; reads always fill it, and a missing
+ * value means "pi".
+ */
+export interface PiAgentSessionWithEventCount
+  extends Omit<PiAgentSession, "kind"> {
+  kind?: PiAgentSession["kind"];
+  eventCount: number;
+}
+
+export interface KernelTraceReadRows {
+  rootContainer: Container;
+  containers: Container[];
+  piSessions: PiAgentSessionWithEventCount[];
+  agentRuns: AgentRun[];
+  events: TraceEventRow[];
+}
+
+export interface KernelTraceDeleteResult {
+  containerIds: string[];
+  piSessionIds: string[];
+  deleted: {
+    traceEvents: number;
+    agentRuns: number;
+    piAgentSessions: number;
+    containers: number;
+  };
+}
+
+function clampLimit(limit: number | undefined, fallback: number, max: number): number {
+  if (limit === undefined || !Number.isFinite(limit)) return fallback;
+  return Math.max(1, Math.min(Math.floor(limit), max));
+}
+
+export interface SessionContainerStatsRow {
+	container: Container;
+	piSessionCount: number;
+	eventCount: number;
+	latestEventAt: string | null;
+}
+
+export interface ListSessionContainersOptions {
+	/** Scope to one kernel's containers. */
+	kernelId?: string;
+	/** Container kind to list; defaults to "session". */
+	kind?: string;
+	limit?: number;
+}
+
+/**
+ * Containers of kind "session" (newest first) with the per-container pi
+ * session count, event count, and latest event timestamp — the rows behind
+ * the trace-sessions index.
+ */
+export async function listSessionContainersWithStats(
+	db: KernelDatabase,
+	opts: ListSessionContainersOptions = {},
+): Promise<SessionContainerStatsRow[]> {
+	const limit = clampLimit(opts.limit, 100, 500);
+	const kind = opts.kind ?? "session";
+	const conditions: SQL[] = [eq(containers.kind, kind)];
+	if (opts.kernelId) conditions.push(eq(containers.kernelId, opts.kernelId));
+
+	const rows = await db
+		.select()
+		.from(containers)
+		.where(and(...conditions))
+		.orderBy(desc(containers.createdAt))
+		.limit(limit);
+	if (rows.length === 0) return [];
+
+	const containerIds = rows.map((row) => row.id);
+	const piCounts = await db
+		.select({ containerId: piAgentSessions.containerId, count: count() })
+		.from(piAgentSessions)
+		.where(inArray(piAgentSessions.containerId, containerIds))
+		.groupBy(piAgentSessions.containerId);
+	const eventStats = await db
+		.select({
+			containerId: traceEvents.containerId,
+			count: count(),
+			latestEventAt: max(traceEvents.timestamp),
+		})
+		.from(traceEvents)
+		.where(inArray(traceEvents.containerId, containerIds))
+		.groupBy(traceEvents.containerId);
+
+	const piCountByContainer = new Map(
+		piCounts.map((row) => [row.containerId, Number(row.count ?? 0)]),
+	);
+	const eventStatsByContainer = new Map(
+		eventStats.map((row) => [
+			row.containerId,
+			{ count: Number(row.count ?? 0), latestEventAt: row.latestEventAt ?? null },
+		]),
+	);
+
+	return rows.map((container) => {
+		const stats = eventStatsByContainer.get(container.id);
+		return {
+			container,
+			piSessionCount: piCountByContainer.get(container.id) ?? 0,
+			eventCount: stats?.count ?? 0,
+			latestEventAt: stats?.latestEventAt ?? null,
+		};
+	});
+}
+
+/** Breadth-first walk of a container subtree, root first. */
+export async function listContainerTree(
+  db: KernelDatabase,
+  rootContainerId: string,
+  opts: { maxContainers?: number } = {},
+): Promise<Container[]> {
+  const maxContainers = clampLimit(opts.maxContainers, 500, 5000);
+  const [root] = await db
+    .select()
+    .from(containers)
+    .where(eq(containers.id, rootContainerId))
+    .limit(1);
+
+  if (!root) return [];
+
+  const rows: Container[] = [root];
+  const seen = new Set<string>([root.id]);
+  let frontier = [root.id];
+
+  while (frontier.length > 0 && rows.length < maxContainers) {
+    const remaining = maxContainers - rows.length;
+    const children = await db
+      .select()
+      .from(containers)
+      .where(inArray(containers.parentContainerId, frontier))
+      .orderBy(asc(containers.createdAt))
+      .limit(remaining);
+
+    frontier = [];
+    for (const child of children) {
+      if (seen.has(child.id)) continue;
+      seen.add(child.id);
+      rows.push(child);
+      frontier.push(child.id);
+    }
+  }
+
+  return rows;
+}
+
+/**
+ * Everything the viewer needs for one container subtree: the container rows,
+ * their sessions (with event counts), runs, and events — keyed by containerId
+ * only.
+ */
+export async function getKernelTraceReadRows(
+  db: KernelDatabase,
+  rootContainerId: string,
+  opts: KernelTraceReadOptions = {},
+): Promise<KernelTraceReadRows | undefined> {
+  const containerRows = await listContainerTree(db, rootContainerId, {
+    maxContainers: opts.maxContainers,
+  });
+  const rootContainer = containerRows[0];
+  if (!rootContainer) return undefined;
+
+  const containerIds = containerRows.map((container) => container.id);
+  const limit = clampLimit(opts.limit, 5000, 10000);
+
+  const piSessionRows: PiAgentSession[] = await db
+    .select(piAgentSessionSelection(db))
+    .from(piAgentSessions)
+    .where(inArray(piAgentSessions.containerId, containerIds))
+    .orderBy(asc(piAgentSessions.createdAt));
+
+  const piSessionIds = piSessionRows.map((pi) => pi.id);
+  const eventCountRows =
+    piSessionIds.length > 0
+      ? await db
+          .select({
+            piSessionId: traceEvents.piSessionId,
+            eventCount: count(),
+          })
+          .from(traceEvents)
+          .where(inArray(traceEvents.piSessionId, piSessionIds))
+          .groupBy(traceEvents.piSessionId)
+      : [];
+  const eventCountsByPiSession = new Map<string, number>();
+  for (const row of eventCountRows) {
+    if (row.piSessionId) {
+      eventCountsByPiSession.set(row.piSessionId, Number(row.eventCount ?? 0));
+    }
+  }
+  const piSessions: PiAgentSessionWithEventCount[] = piSessionRows.map(
+    (session) => ({
+      ...session,
+      eventCount: eventCountsByPiSession.get(session.id) ?? 0,
+    }),
+  );
+
+  const runRows: AgentRun[] = await db
+    .select()
+    .from(agentRuns)
+    .where(inArray(agentRuns.containerId, containerIds))
+    .orderBy(asc(agentRuns.startedAt));
+
+  const eventConditions: SQL[] = [
+    inArray(traceEvents.containerId, containerIds),
+  ];
+  if (opts.after) eventConditions.push(gt(traceEvents.timestamp, opts.after));
+
+  const eventRows: TraceEventRow[] = await db
+    .select()
+    .from(traceEvents)
+    .where(and(...eventConditions))
+    .orderBy(asc(traceEvents.timestamp), asc(traceEvents.eventId))
+    .limit(limit);
+
+  return {
+    rootContainer,
+    containers: containerRows,
+    piSessions,
+    agentRuns: runRows,
+    events: eventRows,
+  };
+}
+
+/** Delete a container subtree and all rows hanging off it. */
+export async function deleteKernelTraceRows(
+  db: KernelDatabase,
+  rootContainerId: string,
+  opts: Pick<KernelTraceReadOptions, "maxContainers"> = {},
+): Promise<KernelTraceDeleteResult | undefined> {
+  const containerRows = await listContainerTree(db, rootContainerId, {
+    maxContainers: opts.maxContainers,
+  });
+  const rootContainer = containerRows[0];
+  if (!rootContainer) return undefined;
+
+  const containerIds = containerRows.map((container) => container.id);
+
+  return db.transaction(async (tx) => {
+    const piSessionRows = await tx
+      .select({ id: piAgentSessions.id })
+      .from(piAgentSessions)
+      .where(inArray(piAgentSessions.containerId, containerIds));
+    const piSessionIds = piSessionRows.map((session) => session.id);
+
+    const eventConditions: SQL[] = [
+      inArray(traceEvents.containerId, containerIds),
+    ];
+    const deletedEvents = await tx
+      .delete(traceEvents)
+      .where(and(...eventConditions))
+      .returning({ eventId: traceEvents.eventId });
+
+    const deletedRuns = await tx
+      .delete(agentRuns)
+      .where(inArray(agentRuns.containerId, containerIds))
+      .returning({ id: agentRuns.id });
+
+    const deletedPiSessions =
+      piSessionIds.length > 0
+        ? await tx
+            .delete(piAgentSessions)
+            .where(inArray(piAgentSessions.id, piSessionIds))
+            .returning({ id: piAgentSessions.id })
+        : [];
+
+    const deletedContainers: Array<{ id: string }> = [];
+    for (const container of [...containerRows].reverse()) {
+      const rows = await tx
+        .delete(containers)
+        .where(eq(containers.id, container.id))
+        .returning({ id: containers.id });
+      deletedContainers.push(...rows);
+    }
+
+    return {
+      containerIds,
+      piSessionIds,
+      deleted: {
+        traceEvents: deletedEvents.length,
+        agentRuns: deletedRuns.length,
+        piAgentSessions: deletedPiSessions.length,
+        containers: deletedContainers.length,
+      },
+    };
+  });
+}

package/src/actions/trace-blobs.ts

@@ -0,0 +1,71 @@
+import { createHash } from "node:crypto";
+import { eq, inArray } from "drizzle-orm";
+import type { KernelDatabase } from "../client";
+import { traceBlobs } from "../schema/trace-blobs";
+import type { NewTraceBlob, TraceBlob } from "../types";
+
+/**
+ * Content address for trace blob bytes: "b1-" + sha256hex of the raw bytes.
+ * The kernel capture code hashes payloads with this before storing them.
+ */
+export function hashTraceBlobBytes(bytes: Uint8Array): string {
+  return `b1-${createHash("sha256").update(bytes).digest("hex")}`;
+}
+
+/**
+ * Idempotent multi-row insert keyed by hash. Blobs are content-addressed and
+ * immutable, so a conflicting insert is benign — the existing row (including
+ * its createdAt) is left untouched. Returns the number of rows actually
+ * inserted (0 when every blob already existed).
+ */
+export async function upsertTraceBlobs(
+  db: KernelDatabase,
+  blobs: NewTraceBlob[],
+): Promise<number> {
+  if (blobs.length === 0) return 0;
+
+  const inserted = await db
+    .insert(traceBlobs)
+    .values(blobs)
+    .onConflictDoNothing({ target: traceBlobs.hash })
+    .returning({ hash: traceBlobs.hash });
+  return inserted.length;
+}
+
+export async function getTraceBlob(
+  db: KernelDatabase,
+  hash: string,
+): Promise<TraceBlob | null> {
+  const [row] = await db
+    .select()
+    .from(traceBlobs)
+    .where(eq(traceBlobs.hash, hash))
+    .limit(1);
+  return row ?? null;
+}
+
+/** Blob row metadata — everything except the data payload. */
+export type TraceBlobMeta = Omit<TraceBlob, "data">;
+
+/**
+ * Metadata rows (no data column) for the given hashes, for existence and
+ * size checks without pulling payload bytes off disk. Missing hashes are
+ * simply absent from the result.
+ */
+export async function getTraceBlobMetas(
+  db: KernelDatabase,
+  hashes: string[],
+): Promise<TraceBlobMeta[]> {
+  if (hashes.length === 0) return [];
+
+  return db
+    .select({
+      hash: traceBlobs.hash,
+      kind: traceBlobs.kind,
+      mimeType: traceBlobs.mimeType,
+      byteLength: traceBlobs.byteLength,
+      createdAt: traceBlobs.createdAt,
+    })
+    .from(traceBlobs)
+    .where(inArray(traceBlobs.hash, hashes));
+}

package/src/actions/trace-events.ts

@@ -0,0 +1,158 @@
+import { and, asc, eq, gt, inArray, sql } from "drizzle-orm";
+import type { TraceEvent } from "@codecaine-ai/agent-kernel-protocol";
+import type { KernelDatabase } from "../client";
+import { traceEvents } from "../schema/trace-events";
+import type { TraceEventRow } from "../types";
+import { runImmediateTransaction, traceEventToRow } from "./tx-local";
+
+/**
+ * Idempotent batch insert keyed by event_id (INSERT OR IGNORE) — replaying
+ * a batch never duplicates rows. The transport-only piSessionUuid on the
+ * envelope is resolved into the pi_session_id column at write time.
+ */
+export async function insertTraceEventsBatch(
+  db: KernelDatabase,
+  events: TraceEvent[],
+): Promise<number> {
+  if (events.length === 0) return 0;
+
+  const inserted = await db
+    .insert(traceEvents)
+    .values(events.map(traceEventToRow))
+    .onConflictDoNothing({ target: traceEvents.eventId })
+    .returning({ eventId: traceEvents.eventId });
+  return inserted.length;
+}
+
+export interface ListTraceEventsOptions {
+  typeFilter?: string[];
+  after?: string;
+  limit?: number;
+}
+
+export async function listTraceEventsForContainer(
+  db: KernelDatabase,
+  containerId: string,
+  opts: ListTraceEventsOptions = {},
+): Promise<TraceEventRow[]> {
+  const { typeFilter, after, limit = 100 } = opts;
+
+  const conditions = [eq(traceEvents.containerId, containerId)];
+  if (typeFilter && typeFilter.length > 0) {
+    conditions.push(inArray(traceEvents.type, typeFilter));
+  }
+  if (after) {
+    conditions.push(gt(traceEvents.timestamp, after));
+  }
+
+  return db
+    .select()
+    .from(traceEvents)
+    .where(and(...conditions))
+    .orderBy(asc(traceEvents.timestamp), asc(traceEvents.eventId))
+    .limit(Math.min(limit, 1000));
+}
+
+export async function listTraceEventsForRun(
+  db: KernelDatabase,
+  runId: string,
+  opts: ListTraceEventsOptions = {},
+): Promise<TraceEventRow[]> {
+  const { typeFilter, after, limit = 100 } = opts;
+
+  const conditions = [eq(traceEvents.runId, runId)];
+  if (typeFilter && typeFilter.length > 0) {
+    conditions.push(inArray(traceEvents.type, typeFilter));
+  }
+  if (after) {
+    conditions.push(gt(traceEvents.timestamp, after));
+  }
+
+  return db
+    .select()
+    .from(traceEvents)
+    .where(and(...conditions))
+    .orderBy(asc(traceEvents.timestamp), asc(traceEvents.eventId))
+    .limit(Math.min(limit, 1000));
+}
+
+/** Envelope for a stored row; pi_session_id travels back as piSessionUuid. */
+function traceEventFromRow(row: TraceEventRow): TraceEvent {
+  return {
+    eventId: row.eventId,
+    containerId: row.containerId,
+    type: row.type,
+    source: row.source,
+    traceLevel: row.traceLevel as TraceEvent["traceLevel"],
+    eventData: row.eventData as TraceEvent["eventData"],
+    timestamp: row.timestamp,
+    ...(row.runId !== null && { runId: row.runId }),
+    ...(row.piSessionId !== null && { piSessionUuid: row.piSessionId }),
+    ...(row.agentId !== null && { agentId: row.agentId }),
+    ...(row.userId !== null && { userId: row.userId }),
+    ...(row.spanId !== null && { spanId: row.spanId }),
+    ...(row.parentEventId !== null && { parentEventId: row.parentEventId }),
+  };
+}
+
+/**
+ * Every event of one run as envelopes, oldest first (uses idx_events_run).
+ * Unlike listTraceEventsForRun there is no row limit: a model-node replay
+ * rebuilds its result from the whole run.
+ */
+export async function getTraceEventsForRun(
+  db: KernelDatabase,
+  runId: string,
+  types?: readonly string[],
+): Promise<TraceEvent[]> {
+  const conditions = [eq(traceEvents.runId, runId)];
+  if (types && types.length > 0) {
+    conditions.push(inArray(traceEvents.type, [...types]));
+  }
+  const rows = await db
+    .select()
+    .from(traceEvents)
+    .where(and(...conditions))
+    .orderBy(asc(traceEvents.timestamp), asc(traceEvents.eventId));
+  return rows.map(traceEventFromRow);
+}
+
+/**
+ * Promotable upsert for nested tool ends: inserts, or replaces an existing
+ * row with the same event_id only when the stored row's event_data.timing is
+ * 'approximate' and the new one is 'live'. Any other conflict keeps the
+ * stored row.
+ *
+ *   INSERT … ON CONFLICT(event_id) DO UPDATE
+ *     SET event_data = excluded.event_data, timestamp = excluded.timestamp
+ *     WHERE json_extract(trace_events.event_data, '$.timing') = 'approximate'
+ *       AND json_extract(excluded.event_data, '$.timing') = 'live'
+ */
+export async function upsertPromotableTraceEvent(
+  db: KernelDatabase,
+  event: TraceEvent,
+): Promise<"inserted" | "promoted" | "kept"> {
+  return runImmediateTransaction(db, (tx) => {
+    const existed =
+      tx
+        .select({ eventId: traceEvents.eventId })
+        .from(traceEvents)
+        .where(eq(traceEvents.eventId, event.eventId))
+        .get() !== undefined;
+    const written = tx
+      .insert(traceEvents)
+      .values(traceEventToRow(event))
+      .onConflictDoUpdate({
+        target: traceEvents.eventId,
+        set: {
+          eventData: sql`excluded.event_data`,
+          timestamp: sql`excluded.timestamp`,
+        },
+        setWhere: sql`json_extract(${traceEvents.eventData}, '$.timing') = 'approximate' AND json_extract(excluded.event_data, '$.timing') = 'live'`,
+      })
+      .returning({ eventId: traceEvents.eventId })
+      .all();
+    if (!existed) return "inserted" as const;
+    return written.length > 0 ? ("promoted" as const) : ("kept" as const);
+  });
+}