@vellumai/assistant 0.11.2 → 0.11.3-staging.2

package/src/persistence/migrations/361-normalize-managed-connection-rows.test.ts

@@ -0,0 +1,168 @@
+import { Database } from "bun:sqlite";
+import { describe, expect, test } from "bun:test";
+
+import { drizzle } from "drizzle-orm/bun-sqlite";
+
+import * as schema from "../schema.js";
+import { migrateNormalizeManagedConnectionRows } from "./361-normalize-managed-connection-rows.js";
+
+function createTestDb() {
+  const sqlite = new Database(":memory:");
+  sqlite.exec(/*sql*/ `
+    CREATE TABLE provider_connections (
+      name        TEXT PRIMARY KEY,
+      provider    TEXT NOT NULL,
+      auth        TEXT NOT NULL,
+      label       TEXT,
+      base_url    TEXT,
+      models      TEXT,
+      created_at  INTEGER NOT NULL,
+      updated_at  INTEGER NOT NULL
+    );
+  `);
+  return { sqlite, db: drizzle(sqlite, { schema }) };
+}
+
+function insertRow(
+  sqlite: Database,
+  name: string,
+  provider: string,
+  auth: string,
+): void {
+  sqlite
+    .query(
+      /*sql*/ `INSERT INTO provider_connections (name, provider, auth, created_at, updated_at)
+       VALUES (?, ?, ?, 1, 1)`,
+    )
+    .run(name, provider, auth);
+}
+
+function readRows(
+  sqlite: Database,
+): Array<{ name: string; provider: string; auth: string }> {
+  return sqlite
+    .query(
+      `SELECT name, provider, auth FROM provider_connections ORDER BY name`,
+    )
+    .all() as Array<{ name: string; provider: string; auth: string }>;
+}
+
+function readRow(
+  sqlite: Database,
+  name: string,
+): { provider: string; auth: string } {
+  const row = readRows(sqlite).find((r) => r.name === name);
+  if (!row) {
+    throw new Error(`row "${name}" not found`);
+  }
+  return row;
+}
+
+describe("migration 361: normalize managed connection rows", () => {
+  test("rewrites a platform-auth row with a concrete provider to provider vellum", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(sqlite, "managed-openai", "openai", '{"type":"platform"}');
+
+    migrateNormalizeManagedConnectionRows(db);
+
+    const row = readRow(sqlite, "managed-openai");
+    expect(row.provider).toBe("vellum");
+    expect(JSON.parse(row.auth)).toEqual({ type: "platform" });
+  });
+
+  test("heals a stuck legacy canonical row (name vellum, concrete provider, platform auth)", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(sqlite, "vellum", "anthropic", '{"type":"platform"}');
+
+    migrateNormalizeManagedConnectionRows(db);
+
+    expect(readRow(sqlite, "vellum").provider).toBe("vellum");
+  });
+
+  test("rewrites a vellum-provider row with non-platform auth to platform auth", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(
+      sqlite,
+      "keyed-vellum",
+      "vellum",
+      '{"type":"api_key","credential":"vault/x"}',
+    );
+    insertRow(sqlite, "none-vellum", "vellum", '{"type":"none"}');
+
+    migrateNormalizeManagedConnectionRows(db);
+
+    expect(JSON.parse(readRow(sqlite, "keyed-vellum").auth)).toEqual({
+      type: "platform",
+    });
+    expect(JSON.parse(readRow(sqlite, "none-vellum").auth)).toEqual({
+      type: "platform",
+    });
+  });
+
+  test("leaves consistent rows untouched", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(sqlite, "vellum", "vellum", '{"type":"platform"}');
+    insertRow(
+      sqlite,
+      "anthropic-key",
+      "anthropic",
+      '{"type":"api_key","credential":"vault/anthropic"}',
+    );
+    insertRow(
+      sqlite,
+      "chatgpt-subscription",
+      "openai",
+      '{"type":"oauth_subscription"}',
+    );
+    insertRow(sqlite, "ollama-local", "ollama", '{"type":"none"}');
+
+    const before = readRows(sqlite);
+    migrateNormalizeManagedConnectionRows(db);
+
+    expect(readRows(sqlite)).toEqual(before);
+  });
+
+  test("leaves a BYOK row claiming the canonical name untouched", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(
+      sqlite,
+      "vellum",
+      "anthropic",
+      '{"type":"api_key","credential":"vault/anthropic"}',
+    );
+
+    migrateNormalizeManagedConnectionRows(db);
+
+    const row = readRow(sqlite, "vellum");
+    expect(row.provider).toBe("anthropic");
+    expect(JSON.parse(row.auth).type).toBe("api_key");
+  });
+
+  test("leaves a row with unparseable auth untouched", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(sqlite, "broken", "vellum", "not json");
+
+    migrateNormalizeManagedConnectionRows(db);
+
+    expect(readRow(sqlite, "broken").auth).toBe("not json");
+  });
+
+  test("is idempotent", () => {
+    const { sqlite, db } = createTestDb();
+    insertRow(sqlite, "managed-openai", "openai", '{"type":"platform"}');
+    insertRow(sqlite, "keyed-vellum", "vellum", '{"type":"api_key"}');
+
+    migrateNormalizeManagedConnectionRows(db);
+    const afterFirst = readRows(sqlite);
+    migrateNormalizeManagedConnectionRows(db);
+
+    expect(readRows(sqlite)).toEqual(afterFirst);
+  });
+
+  test("no-ops on a database from before the provider_connections table", () => {
+    const sqlite = new Database(":memory:");
+    const db = drizzle(sqlite, { schema });
+
+    expect(() => migrateNormalizeManagedConnectionRows(db)).not.toThrow();
+  });
+});

package/src/persistence/migrations/361-normalize-managed-connection-rows.ts

@@ -0,0 +1,76 @@
+import { type DrizzleDb, getSqliteFrom } from "../db-connection.js";
+
+const VELLUM_PROVIDER = "vellum";
+const PLATFORM_AUTH_JSON = JSON.stringify({ type: "platform" });
+
+/**
+ * Reconcile `provider_connections` rows so platform auth and provider
+ * "vellum" always pair. The two columns record the same fact (the managed,
+ * platform-billed route); auth derivation sets them together and the
+ * connection routes reject writes that split them, but rows written before
+ * that guard can disagree. The provider column is the authoritative managed
+ * signal (`isVellumManagedConnection`), so mismatched rows must converge:
+ *
+ * - Platform auth with any other provider: the provider becomes "vellum".
+ *   These rows dispatch through the managed proxy (transport selection keys
+ *   on the managed identity), so normalizing the provider preserves billing
+ *   and transport behavior, while normalizing the auth instead would
+ *   silently turn a platform-billed row into a keyless BYOK row. This also
+ *   heals a stuck pre-consolidation canonical row (name "vellum", concrete
+ *   provider, platform auth) that boot seeding skips because its provider
+ *   differs from the sentinel.
+ * - Provider "vellum" with any other auth: the auth becomes
+ *   `{"type":"platform"}`. A vellum row means the managed route by
+ *   definition, and non-platform auth on it never yields a usable adapter.
+ *
+ * Rows whose auth column is not valid JSON are left alone (the DB loaders
+ * already treat them as invalid). Idempotent: a reconciled row matches
+ * neither rule, and the table-exists guard covers databases from before
+ * migration 243.
+ */
+export function migrateNormalizeManagedConnectionRows(
+  database: DrizzleDb,
+): void {
+  const raw = getSqliteFrom(database);
+
+  const tableExists = raw
+    .query(
+      `SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'provider_connections'`,
+    )
+    .get();
+  if (!tableExists) {
+    return;
+  }
+
+  const rows = raw
+    .query(`SELECT name, provider, auth FROM provider_connections`)
+    .all() as Array<{ name: string; provider: string; auth: string }>;
+
+  const now = Date.now();
+  for (const row of rows) {
+    let authType: unknown;
+    try {
+      const parsed: unknown = JSON.parse(row.auth);
+      authType =
+        parsed !== null && typeof parsed === "object"
+          ? (parsed as { type?: unknown }).type
+          : undefined;
+    } catch {
+      continue;
+    }
+
+    if (authType === "platform" && row.provider !== VELLUM_PROVIDER) {
+      raw
+        .prepare(
+          `UPDATE provider_connections SET provider = ?, updated_at = ? WHERE name = ?`,
+        )
+        .run(VELLUM_PROVIDER, now, row.name);
+    } else if (row.provider === VELLUM_PROVIDER && authType !== "platform") {
+      raw
+        .prepare(
+          `UPDATE provider_connections SET auth = ?, updated_at = ? WHERE name = ?`,
+        )
+        .run(PLATFORM_AUTH_JSON, now, row.name);
+    }
+  }
+}

package/src/persistence/migrations/362-add-conversation-subagent-kind.test.ts

@@ -0,0 +1,187 @@
+import { Database } from "bun:sqlite";
+import { describe, expect, test } from "bun:test";
+
+import { drizzle } from "drizzle-orm/bun-sqlite";
+
+import * as schema from "../schema.js";
+import { migrateAddConversationSubagentKind } from "./362-add-conversation-subagent-kind.js";
+
+function createTestDb() {
+  const sqlite = new Database(":memory:");
+  // Pre-360 shape: only the columns the migration and tests touch. The
+  // `subagents` table (migration 311) always precedes 360 in the step order
+  // and is the backfill source.
+  sqlite.exec(/*sql*/ `
+    CREATE TABLE conversations (
+      id TEXT PRIMARY KEY,
+      title TEXT,
+      created_at INTEGER NOT NULL,
+      updated_at INTEGER NOT NULL
+    );
+    CREATE TABLE subagents (
+      id TEXT PRIMARY KEY,
+      parent_conversation_id TEXT NOT NULL,
+      conversation_id TEXT NOT NULL,
+      role TEXT NOT NULL,
+      is_fork INTEGER NOT NULL
+    );
+  `);
+  return { sqlite, db: drizzle(sqlite, { schema }) };
+}
+
+function kindOf(
+  sqlite: Database,
+  conversationId: string,
+): { role: unknown; mode: unknown } {
+  const row = sqlite
+    .query(
+      "SELECT subagent_role, subagent_spawn_mode FROM conversations WHERE id = ?",
+    )
+    .get(conversationId) as {
+    subagent_role: unknown;
+    subagent_spawn_mode: unknown;
+  };
+  return { role: row.subagent_role, mode: row.subagent_spawn_mode };
+}
+
+function columnInfo(sqlite: Database) {
+  return sqlite.query("PRAGMA table_info(conversations)").all() as Array<{
+    name: string;
+    notnull: number;
+  }>;
+}
+
+describe("migration 362: conversations subagent role / spawn mode", () => {
+  test("adds both columns, nullable", () => {
+    const { sqlite, db } = createTestDb();
+    const before = columnInfo(sqlite).map((c) => c.name);
+    expect(before).not.toContain("subagent_role");
+    expect(before).not.toContain("subagent_spawn_mode");
+
+    migrateAddConversationSubagentKind(db);
+
+    for (const name of ["subagent_role", "subagent_spawn_mode"]) {
+      const column = columnInfo(sqlite).find((c) => c.name === name);
+      expect(column).toBeDefined();
+      expect(column?.notnull).toBe(0);
+    }
+  });
+
+  test("existing non-subagent rows read back null", () => {
+    const { sqlite, db } = createTestDb();
+    sqlite.exec(/*sql*/ `
+      INSERT INTO conversations (id, created_at, updated_at)
+      VALUES ('conv-1', 1000, 1000)
+    `);
+
+    migrateAddConversationSubagentKind(db);
+
+    expect(kindOf(sqlite, "conv-1")).toEqual({ role: null, mode: null });
+  });
+
+  test("round-trips an insert that sets the new columns", () => {
+    const { sqlite, db } = createTestDb();
+    migrateAddConversationSubagentKind(db);
+
+    sqlite.exec(/*sql*/ `
+      INSERT INTO conversations (id, created_at, updated_at, subagent_role, subagent_spawn_mode)
+      VALUES ('child-1', 2000, 2000, 'researcher', 'regular')
+    `);
+
+    expect(kindOf(sqlite, "child-1")).toEqual({
+      role: "researcher",
+      mode: "regular",
+    });
+  });
+
+  test("backfills role and derives spawn mode from surviving subagents rows", () => {
+    const { sqlite, db } = createTestDb();
+    sqlite.exec(/*sql*/ `
+      INSERT INTO conversations (id, created_at, updated_at) VALUES
+        ('sub-regular', 1000, 1000),
+        ('sub-fork', 2000, 2000),
+        ('sub-advisor', 3000, 3000),
+        ('plain-conv', 4000, 4000);
+      INSERT INTO subagents (id, parent_conversation_id, conversation_id, role, is_fork) VALUES
+        ('s1', 'parent-a', 'sub-regular', 'coder', 0),
+        ('s2', 'parent-b', 'sub-fork', 'general', 1),
+        ('s3', 'parent-c', 'sub-advisor', 'advisor', 1);
+    `);
+
+    migrateAddConversationSubagentKind(db);
+
+    expect(kindOf(sqlite, "sub-regular")).toEqual({
+      role: "coder",
+      mode: "regular",
+    });
+    expect(kindOf(sqlite, "sub-fork")).toEqual({
+      role: "general",
+      mode: "fork",
+    });
+    // The advisor is a role, so the derived mode must be `advisor_consult`
+    // rather than the bare `fork` its `is_fork` flag would otherwise imply.
+    expect(kindOf(sqlite, "sub-advisor")).toEqual({
+      role: "advisor",
+      mode: "advisor_consult",
+    });
+    // Conversations with no surviving subagents row stay unattributed.
+    expect(kindOf(sqlite, "plain-conv")).toEqual({ role: null, mode: null });
+  });
+
+  test("backfill does not overwrite already-stamped values", () => {
+    const { sqlite, db } = createTestDb();
+    migrateAddConversationSubagentKind(db);
+    sqlite.exec(/*sql*/ `
+      INSERT INTO conversations (id, created_at, updated_at, subagent_role, subagent_spawn_mode)
+      VALUES ('stamped', 1000, 1000, 'general', 'voice_continuation');
+      INSERT INTO subagents (id, parent_conversation_id, conversation_id, role, is_fork)
+      VALUES ('s1', 'parent-stale', 'stamped', 'general', 1);
+    `);
+
+    migrateAddConversationSubagentKind(db);
+
+    expect(kindOf(sqlite, "stamped")).toEqual({
+      role: "general",
+      mode: "voice_continuation",
+    });
+  });
+
+  test("indexes subagents.conversation_id so the backfill is not quadratic", () => {
+    const { sqlite, db } = createTestDb();
+
+    migrateAddConversationSubagentKind(db);
+
+    const indexes = (
+      sqlite.query("PRAGMA index_list(subagents)").all() as Array<{
+        name: string;
+      }>
+    ).map((i) => i.name);
+    expect(indexes).toContain("idx_subagents_conversation_id");
+
+    // Both backfill statements look `subagents` up once per `conversations`
+    // row. Without the index that lookup is a full scan of `subagents`, which
+    // makes a step that blocks daemon startup quadratic.
+    const plan = (
+      sqlite
+        .query(
+          "EXPLAIN QUERY PLAN SELECT 1 FROM subagents s WHERE s.conversation_id = 'c1'",
+        )
+        .all() as Array<{ detail: string }>
+    )
+      .map((row) => row.detail)
+      .join(" ");
+    expect(plan).toContain("idx_subagents_conversation_id");
+    expect(plan).not.toContain("SCAN");
+  });
+
+  test("is idempotent: re-run is a no-op", () => {
+    const { sqlite, db } = createTestDb();
+
+    migrateAddConversationSubagentKind(db);
+    expect(() => migrateAddConversationSubagentKind(db)).not.toThrow();
+
+    const names = columnInfo(sqlite).map((c) => c.name);
+    expect(names.filter((n) => n === "subagent_role")).toHaveLength(1);
+    expect(names.filter((n) => n === "subagent_spawn_mode")).toHaveLength(1);
+  });
+});

package/src/persistence/migrations/362-add-conversation-subagent-kind.ts

@@ -0,0 +1,93 @@
+import type { DrizzleDb } from "../db-connection.js";
+import { tableHasColumn } from "./schema-introspection.js";
+
+const ROLE_COLUMN = "subagent_role";
+const SPAWN_MODE_COLUMN = "subagent_spawn_mode";
+
+/**
+ * Add `subagent_role` and `subagent_spawn_mode` columns to `conversations`.
+ *
+ * Both are nullable, free-form strings describing the subagent that owns the
+ * conversation: the role it was spawned with and how it was spawned. The
+ * accepted values, and what each one means, are documented once on
+ * `SubagentRole` / `SubagentSpawnMode` in `subagent/types.ts`; the columns
+ * take whatever the row carried, including role names retired since. NULL for
+ * every conversation that is not a subagent.
+ *
+ * Together they make delegated LLM spend separable. Every subagent variety
+ * (a fire-and-forget spawn, a context-inheriting fork, and an advisor
+ * consult) emits usage under `llm_call_site = "subagentSpawn"` in a child
+ * conversation, so the largest descendant cost bucket could not be
+ * decomposed. The advisor in particular is a ROLE, not a call site.
+ *
+ * Denormalized onto `conversations` rather than joined from `subagents` at
+ * telemetry-flush time because `subagents` rows are deleted on dispose (TTL
+ * sweep ~30 minutes after the run goes terminal) while usage telemetry
+ * flushes on a watermark that can trail arbitrarily far behind after an
+ * ingest outage. Same rationale that put `parent_conversation_id` here
+ * (migration 342). No index: these are projected alongside an existing
+ * primary-key join, never filtered on.
+ *
+ * Backfills from the `subagents` table (migration 311) for subagents that
+ * have not yet been disposed, which links unflushed usage rows of
+ * pre-migration subagents. The spawn mode is derived from the surviving
+ * `role` / `is_fork` columns; long-disposed subagents have no surviving
+ * linkage anywhere and correctly stay NULL. The derivation cannot recover
+ * `voice_continuation` (the live-voice continuation is an unlabelled fork
+ * pre-migration), so those rows backfill as `fork`, a small, bounded
+ * inaccuracy confined to the backfill window.
+ *
+ * The backfill needs `subagents(conversation_id)` indexed, which migration 311
+ * did not create: both statements look the `subagents` table up once per
+ * `conversations` row, so unindexed they are a full scan of `subagents` per
+ * row and the whole step is quadratic. Migrations run during startup and block
+ * readiness, so on a long-lived assistant that is a multi-minute stall at
+ * boot. The index is created here rather than appended as a later migration
+ * because it has to exist BEFORE these statements run; it also serves the
+ * by-conversation lookup in `subagent-store.ts`, which scanned until now.
+ *
+ * Idempotent: guarded with `tableHasColumn` so a crash between the `ALTER
+ * TABLE` and the checkpoint write doesn't cause a duplicate-column error on
+ * the next boot, `IF NOT EXISTS` on the index, and the backfill only fills
+ * NULL rows for the same reason.
+ */
+export function migrateAddConversationSubagentKind(database: DrizzleDb): void {
+  database.run(
+    `CREATE INDEX IF NOT EXISTS idx_subagents_conversation_id ON subagents(conversation_id)`,
+  );
+  if (!tableHasColumn(database, "conversations", ROLE_COLUMN)) {
+    database.run(`ALTER TABLE conversations ADD COLUMN ${ROLE_COLUMN} TEXT`);
+  }
+  if (!tableHasColumn(database, "conversations", SPAWN_MODE_COLUMN)) {
+    database.run(
+      `ALTER TABLE conversations ADD COLUMN ${SPAWN_MODE_COLUMN} TEXT`,
+    );
+  }
+  database.run(`
+    UPDATE conversations
+    SET subagent_role = (
+      SELECT s.role FROM subagents s
+      WHERE s.conversation_id = conversations.id
+    )
+    WHERE subagent_role IS NULL
+      AND EXISTS (
+        SELECT 1 FROM subagents s WHERE s.conversation_id = conversations.id
+      )
+  `);
+  database.run(`
+    UPDATE conversations
+    SET subagent_spawn_mode = (
+      SELECT CASE
+        WHEN s.role = 'advisor' THEN 'advisor_consult'
+        WHEN s.is_fork = 1 THEN 'fork'
+        ELSE 'regular'
+      END
+      FROM subagents s
+      WHERE s.conversation_id = conversations.id
+    )
+    WHERE subagent_spawn_mode IS NULL
+      AND EXISTS (
+        SELECT 1 FROM subagents s WHERE s.conversation_id = conversations.id
+      )
+  `);
+}