@vellumai/assistant 0.12.2 → 0.12.3-staging.2

package/src/persistence/__tests__/client-connection-events-store.test.ts

@@ -0,0 +1,322 @@
+import { Database } from "bun:sqlite";
+import { describe, expect, test } from "bun:test";
+
+import { drizzle } from "drizzle-orm/bun-sqlite";
+
+import {
+  CLIENT_CONNECTION_EVENT_RETENTION_MS,
+  CLIENT_CONNECTION_FLAP_WINDOW_MS,
+  type ClientConnectionEvent,
+  coalesceClientConnectionSessions,
+  listClientConnectionHistory,
+  recordClientConnectionEvent,
+} from "../client-connection-events-store.js";
+import { migrateCreateClientConnectionEvents } from "../migrations/382-create-client-connection-events.js";
+import * as schema from "../schema.js";
+
+function createTestDb() {
+  const sqlite = new Database(":memory:");
+  sqlite.exec("PRAGMA journal_mode=WAL");
+  const db = drizzle(sqlite, { schema });
+  migrateCreateClientConnectionEvents(db);
+  return db;
+}
+
+function event(
+  overrides: Partial<ClientConnectionEvent> &
+    Pick<ClientConnectionEvent, "reason" | "occurredAt">,
+): ClientConnectionEvent {
+  return {
+    id: overrides.id ?? `evt-${overrides.occurredAt}`,
+    clientId: overrides.clientId ?? "client-123",
+    interfaceId: overrides.interfaceId ?? "chrome-extension",
+    connectionId: overrides.connectionId ?? `conn-${overrides.occurredAt}`,
+    reason: overrides.reason,
+    occurredAt: overrides.occurredAt,
+    actorPrincipalId: overrides.actorPrincipalId ?? "user-123",
+    clientVersion: overrides.clientVersion ?? "0.12.1",
+    sseWatchdog: overrides.sseWatchdog ?? true,
+    machineName: overrides.machineName ?? null,
+  };
+}
+
+describe("coalesceClientConnectionSessions", () => {
+  test("merges a reconnect shorter than the flap window", () => {
+    const sessions = coalesceClientConnectionSessions(
+      [
+        event({
+          reason: "sse_open",
+          occurredAt: 1_000,
+          connectionId: "conn-1",
+        }),
+        event({
+          reason: "stale_replaced",
+          occurredAt: 2_000,
+          connectionId: "conn-1",
+        }),
+        event({
+          reason: "sse_open",
+          occurredAt: 2_000 + CLIENT_CONNECTION_FLAP_WINDOW_MS - 1,
+          connectionId: "conn-2",
+        }),
+      ],
+      10_000 + CLIENT_CONNECTION_FLAP_WINDOW_MS,
+    );
+
+    expect(sessions).toHaveLength(1);
+    expect(sessions[0].flapCount).toBe(1);
+    expect(sessions[0].endedAt).toBeNull();
+    expect(sessions[0].connectionIds).toHaveLength(2);
+  });
+
+  test("keeps a gap at or beyond the flap window as a separate session", () => {
+    const sessions = coalesceClientConnectionSessions([
+      event({ reason: "sse_open", occurredAt: 1_000 }),
+      event({ reason: "sse_close", occurredAt: 2_000 }),
+      event({
+        reason: "sse_open",
+        occurredAt: 2_000 + CLIENT_CONNECTION_FLAP_WINDOW_MS,
+      }),
+    ]);
+
+    expect(sessions).toHaveLength(2);
+    expect(sessions.every((session) => session.flapCount === 0)).toBe(true);
+  });
+
+  test("leaves an unmatched open as a still-connected session", () => {
+    const now = 8_000;
+    const sessions = coalesceClientConnectionSessions(
+      [event({ reason: "sse_open", occurredAt: 1_000 })],
+      now,
+    );
+
+    expect(sessions).toHaveLength(1);
+    expect(sessions[0].endedAt).toBeNull();
+    expect(sessions[0].durationMs).toBe(7_000);
+    expect(sessions[0].closeReason).toBeNull();
+  });
+
+  test("ends an orphaned open when a later open arrives with no close", () => {
+    const sessions = coalesceClientConnectionSessions(
+      [
+        event({
+          reason: "sse_open",
+          occurredAt: 1_000,
+          connectionId: "conn-1",
+        }),
+        event({
+          reason: "sse_open",
+          occurredAt: 5_000,
+          connectionId: "conn-2",
+        }),
+      ],
+      8_000,
+    );
+
+    expect(sessions).toHaveLength(2);
+    expect(sessions[1].startedAt).toBe(1_000);
+    expect(sessions[1].endedAt).toBe(5_000);
+    expect(sessions[1].closeReason).toBe("implicit_replaced");
+    expect(sessions[1].durationMs).toBe(4_000);
+    expect(sessions[0].startedAt).toBe(5_000);
+    expect(sessions[0].endedAt).toBeNull();
+    expect(sessions[0].durationMs).toBe(3_000);
+  });
+});
+
+describe("recordClientConnectionEvent / listClientConnectionHistory", () => {
+  test("writes events and returns coalesced sessions", () => {
+    const db = createTestDb();
+    const t0 = 1_700_000_000_000;
+
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-1",
+        reason: "sse_open",
+        occurredAt: t0,
+        actorPrincipalId: "user-123",
+        clientVersion: "0.12.1",
+        sseWatchdog: true,
+      },
+      db,
+    );
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-1",
+        reason: "stale_replaced",
+        occurredAt: t0 + 5_000,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-2",
+        reason: "sse_open",
+        occurredAt: t0 + 6_000,
+        actorPrincipalId: "user-123",
+        clientVersion: "0.12.1",
+        sseWatchdog: true,
+      },
+      db,
+    );
+
+    const history = listClientConnectionHistory(
+      { clientId: "client-123", now: t0 + 10_000 },
+      db,
+    );
+
+    expect(history.events).toHaveLength(3);
+    expect(history.sessions).toHaveLength(1);
+    expect(history.sessions[0].flapCount).toBe(1);
+    expect(history.sessions[0].endedAt).toBeNull();
+    expect(history.sessions[0].sseWatchdog).toBe(true);
+    expect(history.sessions[0].clientVersion).toBe("0.12.1");
+  });
+
+  test("prunes rows older than the retention window", () => {
+    const db = createTestDb();
+    const now = 1_700_000_000_000;
+
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-old",
+        reason: "sse_open",
+        occurredAt: now - CLIENT_CONNECTION_EVENT_RETENTION_MS - 1,
+      },
+      db,
+    );
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-new",
+        reason: "sse_open",
+        occurredAt: now,
+      },
+      db,
+    );
+
+    const history = listClientConnectionHistory({ now }, db);
+    expect(history.events).toHaveLength(1);
+    expect(history.events[0].connectionId).toBe("conn-new");
+  });
+
+  test("does not throw when listing without an injected database", () => {
+    const history = listClientConnectionHistory();
+    expect(Array.isArray(history.events)).toBe(true);
+    expect(Array.isArray(history.sessions)).toBe(true);
+  });
+
+  test("preserves a session that opened before the since bound", () => {
+    const db = createTestDb();
+    const t0 = 1_700_000_000_000;
+
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-1",
+        reason: "sse_open",
+        occurredAt: t0,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-1",
+        reason: "sse_close",
+        occurredAt: t0 + 8_000,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+
+    const history = listClientConnectionHistory(
+      { clientId: "client-123", since: t0 + 5_000, now: t0 + 10_000 },
+      db,
+    );
+
+    expect(history.sessions).toHaveLength(1);
+    expect(history.sessions[0].startedAt).toBe(t0);
+    expect(history.sessions[0].endedAt).toBe(t0 + 8_000);
+    expect(history.sessions[0].durationMs).toBe(8_000);
+    expect(history.events).toHaveLength(1);
+    expect(history.events[0].reason).toBe("sse_close");
+  });
+
+  test("keeps a still-open session whose only event is before since", () => {
+    const db = createTestDb();
+    const t0 = 1_700_000_000_000;
+
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-1",
+        reason: "sse_open",
+        occurredAt: t0,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+
+    const history = listClientConnectionHistory(
+      { clientId: "client-123", since: t0 + 5_000, now: t0 + 10_000 },
+      db,
+    );
+
+    expect(history.sessions).toHaveLength(1);
+    expect(history.sessions[0].startedAt).toBe(t0);
+    expect(history.sessions[0].endedAt).toBeNull();
+    expect(history.sessions[0].durationMs).toBe(10_000);
+    expect(history.events).toHaveLength(0);
+  });
+
+  test("omits a session that ended before the since bound", () => {
+    const db = createTestDb();
+    const t0 = 1_700_000_000_000;
+
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-old",
+        reason: "sse_open",
+        occurredAt: t0,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+    recordClientConnectionEvent(
+      {
+        clientId: "client-123",
+        interfaceId: "chrome-extension",
+        connectionId: "conn-old",
+        reason: "sse_close",
+        occurredAt: t0 + 1_000,
+        actorPrincipalId: "user-123",
+      },
+      db,
+    );
+
+    const history = listClientConnectionHistory(
+      { clientId: "client-123", since: t0 + 5_000, now: t0 + 10_000 },
+      db,
+    );
+
+    expect(history.sessions).toHaveLength(0);
+    expect(history.events).toHaveLength(0);
+  });
+});

package/src/persistence/attachments-store.ts

@@ -17,7 +17,7 @@ import {
 } from "node:fs";
 import { basename, dirname, extname, join } from "node:path";
 
-import { eq } from "drizzle-orm";
+import { and, eq, inArray } from "drizzle-orm";
 import { v4 as uuid } from "uuid";
 
 import {
@@ -31,7 +31,7 @@ import { getWorkspaceDir } from "../util/platform.js";
 import { getConversationAttachmentsDirPath } from "./conversation-directories.js";
 import { getDb } from "./db-connection.js";
 import { rawAll, rawGet, rawRun } from "./raw-query.js";
-import { attachments, messageAttachments } from "./schema.js";
+import { attachments, messageAttachments, messages } from "./schema.js";
 
 export interface StoredAttachment {
   id: string;
@@ -206,16 +206,23 @@ function insertMessageAttachmentLink(
   attachmentId: string,
   position: number,
 ): void {
-  const db = getDb();
-  db.insert(messageAttachments)
-    .values({
-      id: uuid(),
-      messageId,
-      attachmentId,
-      position,
-      createdAt: Date.now(),
-    })
-    .run();
+  rawRun(
+    "attachments:insertMessageAttachmentLink",
+    `INSERT INTO message_attachments
+       (id, message_id, attachment_id, position, created_at)
+     SELECT ?, ?, ?, ?, ?
+     WHERE NOT EXISTS (
+       SELECT 1 FROM message_attachments
+       WHERE message_id = ? AND attachment_id = ?
+     )`,
+    uuid(),
+    messageId,
+    attachmentId,
+    position,
+    Date.now(),
+    messageId,
+    attachmentId,
+  );
 }
 
 function persistAttachmentFilePath(
@@ -1256,6 +1263,20 @@ export function linkAttachmentToMessage(
     throw new Error(`Message not found: ${messageId}`);
   }
 
+  const existing = getDb()
+    .select({ attachmentId: messageAttachments.attachmentId })
+    .from(messageAttachments)
+    .where(
+      and(
+        eq(messageAttachments.messageId, messageId),
+        eq(messageAttachments.attachmentId, attachmentId),
+      ),
+    )
+    .get();
+  if (existing) {
+    return existing.attachmentId;
+  }
+
   const scopedAttachmentId = scopeAttachmentToConversation(
     attachmentId,
     ctx.conversationId,
@@ -1301,41 +1322,60 @@ export function getAttachmentsForMessage(
 export function getAttachmentMetadataForMessage(
   messageId: string,
 ): StoredAttachment[] {
-  const db = getDb();
-  const links = db
-    .select({ attachmentId: messageAttachments.attachmentId })
-    .from(messageAttachments)
-    .where(eq(messageAttachments.messageId, messageId))
-    .orderBy(messageAttachments.position)
-    .all();
+  return getAttachmentMetadataForMessages([messageId]).map(
+    (row) => row.attachment,
+  );
+}
 
-  if (links.length === 0) {
-    return [];
-  }
+export interface MessageAttachmentMetadata {
+  messageId: string;
+  messageMetadata: string | null;
+  attachment: StoredAttachment;
+}
 
-  const results: StoredAttachment[] = [];
-  for (const link of links) {
-    if (!link.attachmentId) {
-      continue;
-    }
-    const row = db
-      .select({
-        id: attachments.id,
-        originalFilename: attachments.originalFilename,
-        mimeType: attachments.mimeType,
-        sizeBytes: attachments.sizeBytes,
-        kind: attachments.kind,
-        thumbnailBase64: attachments.thumbnailBase64,
-        createdAt: attachments.createdAt,
-      })
-      .from(attachments)
-      .where(eq(attachments.id, link.attachmentId))
-      .get();
-    if (row) {
-      results.push(row);
-    }
+/**
+ * Return linked attachment metadata together with the metadata of each source
+ * message. Rows follow the requested message order, then link position.
+ */
+export function getAttachmentMetadataForMessages(
+  messageIds: string[],
+): MessageAttachmentMetadata[] {
+  if (messageIds.length === 0) {
+    return [];
   }
-  return results;
+  const db = getDb();
+  const rows = db
+    .select({
+      messageId: messageAttachments.messageId,
+      messageMetadata: messages.metadata,
+      position: messageAttachments.position,
+      id: attachments.id,
+      originalFilename: attachments.originalFilename,
+      mimeType: attachments.mimeType,
+      sizeBytes: attachments.sizeBytes,
+      kind: attachments.kind,
+      thumbnailBase64: attachments.thumbnailBase64,
+      createdAt: attachments.createdAt,
+    })
+    .from(messageAttachments)
+    .innerJoin(messages, eq(messages.id, messageAttachments.messageId))
+    .innerJoin(attachments, eq(attachments.id, messageAttachments.attachmentId))
+    .where(inArray(messageAttachments.messageId, messageIds))
+    .all();
+  const messageOrder = new Map(messageIds.map((id, index) => [id, index]));
+  rows.sort((a, b) => {
+    const messageDelta =
+      (messageOrder.get(a.messageId) ?? 0) -
+      (messageOrder.get(b.messageId) ?? 0);
+    return messageDelta !== 0 ? messageDelta : a.position - b.position;
+  });
+  return rows.map(
+    ({ messageId, messageMetadata, position: _position, ...a }) => ({
+      messageId,
+      messageMetadata,
+      attachment: a,
+    }),
+  );
 }
 
 /**