@vellumai/assistant 0.3.2 → 0.3.4

package/src/memory/channel-delivery-store.ts

@@ -15,7 +15,7 @@ import { eq, and, lte, isNotNull } from 'drizzle-orm';
 import { v4 as uuid } from 'uuid';
 import { getDb } from './db.js';
 import { channelInboundEvents, conversations } from './schema.js';
-import { getOrCreateConversation } from './conversation-key-store.js';
+import { getConversationByKey, getOrCreateConversation, setConversationKeyIfAbsent } from './conversation-key-store.js';
 import {
   classifyError,
   retryDelayForAttempt,
@@ -31,6 +31,7 @@ export interface InboundResult {
 
 export interface RecordInboundOptions {
   sourceMessageId?: string;
+  assistantId?: string;
 }
 
 /**
@@ -69,8 +70,30 @@ export function recordInbound(
     };
   }
 
-  const conversationKey = `${sourceChannel}:${externalChatId}`;
-  const mapping = getOrCreateConversation(conversationKey);
+  const assistantId = options?.assistantId;
+  const legacyKey = `${sourceChannel}:${externalChatId}`;
+  const scopedKey = assistantId ? `asst:${assistantId}:${sourceChannel}:${externalChatId}` : legacyKey;
+
+  // Resolve conversation mapping with assistant-scoped keying:
+  // 1. If scoped key exists, use it directly.
+  // 2. If assistantId is "self" and legacy key exists, reuse the legacy
+  //    conversation and create a scoped alias to prevent future bleed.
+  // 3. Otherwise, create/get conversation from the scoped key.
+  let mapping: { conversationId: string; created: boolean };
+  const scopedMapping = assistantId ? getConversationByKey(scopedKey) : null;
+  if (scopedMapping) {
+    mapping = { conversationId: scopedMapping.conversationId, created: false };
+  } else if (assistantId === 'self') {
+    const legacyMapping = getConversationByKey(legacyKey);
+    if (legacyMapping) {
+      mapping = { conversationId: legacyMapping.conversationId, created: false };
+      setConversationKeyIfAbsent(scopedKey, legacyMapping.conversationId);
+    } else {
+      mapping = getOrCreateConversation(scopedKey);
+    }
+  } else {
+    mapping = getOrCreateConversation(scopedKey);
+  }
   const now = Date.now();
   const eventId = uuid();
 
@@ -316,6 +339,54 @@ export function getDeadLetterEvents(): Array<{
     .all();
 }
 
+// ── Deliver-once guard for terminal reply idempotency ────────────────
+//
+// When both the main poll (processChannelMessageWithApprovals) and the
+// post-decision poll (schedulePostDecisionDelivery) race to deliver the
+// final assistant reply for the same run, this guard ensures only one
+// of them actually sends the message. The guard is run-scoped so old
+// assistant messages from previous runs are not affected.
+
+const deliveredRuns = new Set<string>();
+
+/** TTL for delivery claims — 10 minutes, well beyond the poll max-wait. */
+const CLAIM_TTL_MS = 10 * 60 * 1000;
+
+/**
+ * Atomically claim the right to deliver the final reply for a run.
+ * Returns `true` if this caller won the claim (and should proceed with
+ * delivery). Returns `false` if another caller already claimed it.
+ *
+ * This is an in-memory guard — sufficient because both racing pollers
+ * execute within the same process. The Set is never persisted; on restart
+ * there are no in-flight pollers to race.
+ *
+ * Claims are automatically evicted after CLAIM_TTL_MS to prevent
+ * unbounded Set growth over the lifetime of the process.
+ */
+export function claimRunDelivery(runId: string): boolean {
+  if (deliveredRuns.has(runId)) return false;
+  deliveredRuns.add(runId);
+  setTimeout(() => deliveredRuns.delete(runId), CLAIM_TTL_MS);
+  return true;
+}
+
+/**
+ * Reset the deliver-once guard for a run. Used to release a claim when
+ * delivery fails (so the other racing poller can retry) and in tests
+ * for isolation between test cases.
+ */
+export function resetRunDeliveryClaim(runId: string): void {
+  deliveredRuns.delete(runId);
+}
+
+/**
+ * Clear all delivery claims. Used in tests for full isolation.
+ */
+export function resetAllRunDeliveryClaims(): void {
+  deliveredRuns.clear();
+}
+
 /**
  * Reset dead-lettered events back to 'failed' so the sweep can retry
  * them. Resets attempt counter and sets an immediate retry_after.

package/src/memory/channel-guardian-store.ts

@@ -58,6 +58,7 @@ export interface GuardianApprovalRequest {
   id: string;
   runId: string;
   conversationId: string;
+  assistantId: string;
   channel: string;
   requesterExternalUserId: string;
   requesterChatId: string;
@@ -114,6 +115,7 @@ function rowToApprovalRequest(row: typeof channelGuardianApprovalRequests.$infer
     id: row.id,
     runId: row.runId,
     conversationId: row.conversationId,
+    assistantId: row.assistantId,
     channel: row.channel,
     requesterExternalUserId: row.requesterExternalUserId,
     requesterChatId: row.requesterChatId,
@@ -305,6 +307,7 @@ export function consumeChallenge(
 export function createApprovalRequest(params: {
   runId: string;
   conversationId: string;
+  assistantId?: string;
   channel: string;
   requesterExternalUserId: string;
   requesterChatId: string;
@@ -323,6 +326,7 @@ export function createApprovalRequest(params: {
     id,
     runId: params.runId,
     conversationId: params.conversationId,
+    assistantId: params.assistantId ?? 'self',
     channel: params.channel,
     requesterExternalUserId: params.requesterExternalUserId,
     requesterChatId: params.requesterChatId,
@@ -388,25 +392,32 @@ export function getUnresolvedApprovalForRun(runId: string): GuardianApprovalRequ
 /**
  * Find a pending guardian approval request by the guardian's chat ID.
  * Used when the guardian sends a decision from their chat.
+ *
+ * When `assistantId` is provided, the lookup is scoped to that assistant,
+ * preventing cross-assistant approval consumption in shared guardian chats.
  */
 export function getPendingApprovalByGuardianChat(
   channel: string,
   guardianChatId: string,
+  assistantId?: string,
 ): GuardianApprovalRequest | null {
   const db = getDb();
   const now = Date.now();
 
+  const conditions = [
+    eq(channelGuardianApprovalRequests.channel, channel),
+    eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
+    eq(channelGuardianApprovalRequests.status, 'pending'),
+    gt(channelGuardianApprovalRequests.expiresAt, now),
+  ];
+  if (assistantId) {
+    conditions.push(eq(channelGuardianApprovalRequests.assistantId, assistantId));
+  }
+
   const row = db
     .select()
     .from(channelGuardianApprovalRequests)
-    .where(
-      and(
-        eq(channelGuardianApprovalRequests.channel, channel),
-        eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
-        eq(channelGuardianApprovalRequests.status, 'pending'),
-        gt(channelGuardianApprovalRequests.expiresAt, now),
-      ),
-    )
+    .where(and(...conditions))
     .orderBy(desc(channelGuardianApprovalRequests.createdAt))
     .get();
 
@@ -418,27 +429,34 @@ export function getPendingApprovalByGuardianChat(
  * guardian chat, and channel. Used when a callback button provides a run ID,
  * so the decision is applied to exactly the right approval even when
  * multiple approvals target the same guardian chat.
+ *
+ * When `assistantId` is provided, the lookup is further scoped to that
+ * assistant to prevent cross-assistant approval consumption.
  */
 export function getPendingApprovalByRunAndGuardianChat(
   runId: string,
   channel: string,
   guardianChatId: string,
+  assistantId?: string,
 ): GuardianApprovalRequest | null {
   const db = getDb();
   const now = Date.now();
 
+  const conditions = [
+    eq(channelGuardianApprovalRequests.runId, runId),
+    eq(channelGuardianApprovalRequests.channel, channel),
+    eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
+    eq(channelGuardianApprovalRequests.status, 'pending'),
+    gt(channelGuardianApprovalRequests.expiresAt, now),
+  ];
+  if (assistantId) {
+    conditions.push(eq(channelGuardianApprovalRequests.assistantId, assistantId));
+  }
+
   const row = db
     .select()
     .from(channelGuardianApprovalRequests)
-    .where(
-      and(
-        eq(channelGuardianApprovalRequests.runId, runId),
-        eq(channelGuardianApprovalRequests.channel, channel),
-        eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
-        eq(channelGuardianApprovalRequests.status, 'pending'),
-        gt(channelGuardianApprovalRequests.expiresAt, now),
-      ),
-    )
+    .where(and(...conditions))
     .get();
 
   return row ? rowToApprovalRequest(row) : null;
@@ -448,25 +466,32 @@ export function getPendingApprovalByRunAndGuardianChat(
  * Return all pending (non-expired) guardian approval requests for a given
  * guardian chat and channel. Used to detect ambiguity when a guardian sends
  * a plain-text decision while multiple approvals are pending.
+ *
+ * When `assistantId` is provided, the results are scoped to that assistant
+ * to prevent cross-assistant approval consumption.
  */
 export function getAllPendingApprovalsByGuardianChat(
   channel: string,
   guardianChatId: string,
+  assistantId?: string,
 ): GuardianApprovalRequest[] {
   const db = getDb();
   const now = Date.now();
 
+  const conditions = [
+    eq(channelGuardianApprovalRequests.channel, channel),
+    eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
+    eq(channelGuardianApprovalRequests.status, 'pending'),
+    gt(channelGuardianApprovalRequests.expiresAt, now),
+  ];
+  if (assistantId) {
+    conditions.push(eq(channelGuardianApprovalRequests.assistantId, assistantId));
+  }
+
   const rows = db
     .select()
     .from(channelGuardianApprovalRequests)
-    .where(
-      and(
-        eq(channelGuardianApprovalRequests.channel, channel),
-        eq(channelGuardianApprovalRequests.guardianChatId, guardianChatId),
-        eq(channelGuardianApprovalRequests.status, 'pending'),
-        gt(channelGuardianApprovalRequests.expiresAt, now),
-      ),
-    )
+    .where(and(...conditions))
     .orderBy(desc(channelGuardianApprovalRequests.createdAt))
     .all();
 
@@ -525,7 +550,7 @@ export interface VerificationRateLimit {
   actorChatId: string;
   /** Individual attempt timestamps (epoch-ms) within the sliding window. */
   attemptTimestamps: number[];
-  /** Derived count of attempts currently inside the window (convenience). */
+  /** Total stored attempt count (may include expired timestamps; use lockedUntil for enforcement decisions). */
   invalidAttempts: number;
   lockedUntil: number | null;
   createdAt: number;
@@ -645,6 +670,9 @@ export function recordInvalidAttempt(
     channel,
     actorExternalUserId,
     actorChatId,
+    // Legacy columns kept for backward compatibility with upgraded databases
+    invalidAttempts: 0,
+    windowStartedAt: 0,
     attemptTimestampsJson: JSON.stringify(timestamps),
     lockedUntil,
     createdAt: now,

package/src/memory/conversation-key-store.ts

@@ -67,6 +67,26 @@ export function setConversationKey(conversationKey: string, conversationId: stri
     .run();
 }
 
+/**
+ * Insert a conversation-key mapping only if the key does not already exist.
+ *
+ * Uses `onConflictDoNothing` on the unique `conversationKey` column to
+ * avoid unique-constraint races when concurrent first messages attempt
+ * to migrate a legacy key to a new scoped alias.
+ */
+export function setConversationKeyIfAbsent(conversationKey: string, conversationId: string): void {
+  const db = getDb();
+  db.insert(conversationKeys)
+    .values({
+      id: uuid(),
+      conversationKey,
+      conversationId,
+      createdAt: Date.now(),
+    })
+    .onConflictDoNothing()
+    .run();
+}
+
 /**
  * Get or create a conversation for the given conversationKey.
  *

package/src/memory/conversation-store.ts

@@ -84,7 +84,7 @@ export function deleteConversation(id: string): void {
   });
 }
 
-export function listConversations(limit?: number, includeBackground = false) {
+export function listConversations(limit?: number, includeBackground = false, offset = 0) {
   const db = getDb();
   const where = includeBackground ? undefined : sql`${conversations.threadType} != 'background'`;
   const query = db
@@ -92,10 +92,22 @@ export function listConversations(limit?: number, includeBackground = false) {
     .from(conversations)
     .where(where)
     .orderBy(desc(conversations.updatedAt))
-    .limit(limit ?? 100);
+    .limit(limit ?? 100)
+    .offset(offset);
   return query.all();
 }
 
+export function countConversations(includeBackground = false): number {
+  const db = getDb();
+  const where = includeBackground ? undefined : sql`${conversations.threadType} != 'background'`;
+  const [{ total }] = db
+    .select({ total: count() })
+    .from(conversations)
+    .where(where)
+    .all();
+  return total;
+}
+
 export function getLatestConversation() {
   const db = getDb();
   const result = db

package/src/memory/db-connection.ts

@@ -0,0 +1,28 @@
+import { Database } from 'bun:sqlite';
+import { drizzle } from 'drizzle-orm/bun-sqlite';
+import * as schema from './schema.js';
+import { getDbPath, ensureDataDir, migrateToDataLayout, migrateToWorkspaceLayout } from '../util/platform.js';
+
+let db: ReturnType<typeof drizzle<typeof schema>> | null = null;
+
+export function getDb() {
+  if (!db) {
+    migrateToDataLayout();
+    migrateToWorkspaceLayout();
+    ensureDataDir();
+    const sqlite = new Database(getDbPath());
+    sqlite.exec('PRAGMA journal_mode=WAL');
+    sqlite.exec('PRAGMA foreign_keys = ON');
+    db = drizzle(sqlite, { schema });
+  }
+  return db;
+}
+
+/** Reset the db singleton. Used by tests to ensure isolation between test files. */
+export function resetDb(): void {
+  if (db) {
+    const raw = (db as unknown as { $client: Database }).$client;
+    raw.close();
+    db = null;
+  }
+}