@desplega.ai/agent-swarm 2.4.3 → 2.4.5

package/src/be/memory/providers/openai-embedding.ts

@@ -3,12 +3,40 @@ import { getOpenRouterAttributionHeaders } from "../../../utils/openrouter-base-
 import { DEFAULT_EMBEDDING_MODEL, EMBEDDING_DIMENSIONS } from "../constants";
 import type { EmbeddingProvider } from "../types";
 
+// text-embedding-3-* reject any input over 8,192 tokens with a 400, and one
+// such input fails every input in its request. There is no tokenizer here, and
+// dense text (UUIDs, hashes, CJK) measures 1.2-1.8 chars per token, so no char
+// cap is safe by itself: start at ~3 chars per token and halve on rejection.
+export const MAX_EMBEDDING_INPUT_CHARS = 24_000;
+const MIN_EMBEDDING_INPUT_CHARS = 1_000;
+// The API also caps one request at 300,000 tokens. Even at one char per token,
+// this keeps a request under it.
+const MAX_BATCH_REQUEST_CHARS = 200_000;
+
 interface OpenAIEmbeddingConfig {
   apiKey?: string;
   model?: string;
   dimensions?: number;
 }
 
+function cleanInput(text: string): string {
+  return text
+    .replace(/[\n\r]/g, " ")
+    .trim()
+    .slice(0, MAX_EMBEDDING_INPUT_CHARS);
+}
+
+/** A 400 is about the request's inputs, so the other inputs can still embed. */
+function isInputRejection(err: unknown): boolean {
+  return err instanceof OpenAI.APIError && err.status === 400;
+}
+
+function isInputTooLong(err: unknown): boolean {
+  return (
+    isInputRejection(err) && /maximum (context|input) length/i.test((err as Error).message ?? "")
+  );
+}
+
 export class OpenAIEmbeddingProvider implements EmbeddingProvider {
   readonly name: string;
   readonly dimensions: number;
@@ -47,83 +75,109 @@ export class OpenAIEmbeddingProvider implements EmbeddingProvider {
     return this.client;
   }
 
-  async embed(text: string): Promise<Float32Array | null> {
-    const client = this.getClient();
-    if (!client) return null;
-
-    try {
-      const cleaned = text.replace(/[\n\r]/g, " ").trim();
-      if (!cleaned) return null;
-
-      const response = await client.embeddings.create({
-        model: this.model,
-        input: cleaned,
-        dimensions: this.dimensions,
-        // Without it the SDK asks for base64 and decodes any reply as base64. A
-        // gateway that answers with plain floats would then yield garbage.
-        encoding_format: "float",
-      });
+  private async request(client: OpenAI, input: string | string[]) {
+    return client.embeddings.create({
+      model: this.model,
+      input,
+      dimensions: this.dimensions,
+      // Without it the SDK asks for base64 and decodes any reply as base64. A
+      // gateway that answers with plain floats would then yield garbage.
+      encoding_format: "float",
+    });
+  }
 
-      const values = response.data[0]?.embedding;
-      if (!values) return null;
+  private toVector(values: number[] | undefined, label: string): Float32Array | null {
+    if (!values) return null;
+    if (values.length !== this.dimensions) {
+      console.error(
+        `[memory] ${label} dimension mismatch: expected=${this.dimensions} got=${values.length}. Provider may not support the 'dimensions' parameter.`,
+      );
+      return null;
+    }
+    return new Float32Array(values);
+  }
 
-      if (values.length !== this.dimensions) {
-        console.error(
-          `[memory] Embedding dimension mismatch: expected=${this.dimensions} got=${values.length}. Provider may not support the 'dimensions' parameter.`,
-        );
+  /** Embed one cleaned input, cutting it in half each time the provider says it is too long. */
+  private async embedOne(client: OpenAI, cleaned: string): Promise<Float32Array | null> {
+    let input = cleaned;
+    while (true) {
+      try {
+        const response = await this.request(client, input);
+        return this.toVector(response.data[0]?.embedding, "Embedding");
+      } catch (err) {
+        const message = (err as Error).message;
+        if (isInputTooLong(err) && input.length > MIN_EMBEDDING_INPUT_CHARS) {
+          const next = Math.max(MIN_EMBEDDING_INPUT_CHARS, Math.floor(input.length / 2));
+          console.warn(
+            `[memory] Embedding input over the token cap at ${input.length} chars, retrying at ${next}: ${message}`,
+          );
+          input = input.slice(0, next);
+          continue;
+        }
+        console.error("[memory] Embedding failed:", message);
         return null;
       }
-
-      return new Float32Array(values);
-    } catch (err) {
-      console.error("[memory] Embedding failed:", (err as Error).message);
-      return null;
     }
   }
 
-  async embedBatch(texts: string[]): Promise<(Float32Array | null)[]> {
+  async embed(text: string): Promise<Float32Array | null> {
     const client = this.getClient();
-    if (!client) return texts.map(() => null);
+    if (!client) return null;
 
-    const cleaned = texts.map((t) => t.replace(/[\n\r]/g, " ").trim());
-    const nonEmptyIndices: number[] = [];
-    const nonEmptyTexts: string[] = [];
+    const cleaned = cleanInput(text);
+    if (!cleaned) return null;
+    return this.embedOne(client, cleaned);
+  }
+
+  async embedBatch(texts: string[]): Promise<(Float32Array | null)[]> {
+    const results: (Float32Array | null)[] = texts.map(() => null);
+    const client = this.getClient();
+    if (!client) return results;
 
+    const cleaned = texts.map(cleanInput);
+    const groups: number[][] = [];
+    let group: number[] = [];
+    let groupChars = 0;
     for (let i = 0; i < cleaned.length; i++) {
-      if (cleaned[i]) {
-        nonEmptyIndices.push(i);
-        nonEmptyTexts.push(cleaned[i]!);
+      const length = cleaned[i]!.length;
+      if (length === 0) continue;
+      if (group.length > 0 && groupChars + length > MAX_BATCH_REQUEST_CHARS) {
+        groups.push(group);
+        group = [];
+        groupChars = 0;
       }
+      group.push(i);
+      groupChars += length;
     }
+    if (group.length > 0) groups.push(group);
 
-    if (nonEmptyTexts.length === 0) return texts.map(() => null);
-
-    try {
-      const response = await client.embeddings.create({
-        model: this.model,
-        input: nonEmptyTexts,
-        dimensions: this.dimensions,
-        encoding_format: "float",
-      });
-
-      const results: (Float32Array | null)[] = texts.map(() => null);
-      for (const item of response.data) {
-        const originalIndex = nonEmptyIndices[item.index];
-        if (originalIndex !== undefined && item.embedding) {
-          if (item.embedding.length !== this.dimensions) {
-            console.error(
-              `[memory] Batch embedding dimension mismatch: expected=${this.dimensions} got=${item.embedding.length}. Provider may not support the 'dimensions' parameter.`,
-            );
-            continue;
-          }
-          results[originalIndex] = new Float32Array(item.embedding);
+    for (const indices of groups) {
+      try {
+        const response = await this.request(
+          client,
+          indices.map((i) => cleaned[i]!),
+        );
+        for (const item of response.data) {
+          const originalIndex = indices[item.index];
+          if (originalIndex === undefined) continue;
+          results[originalIndex] = this.toVector(item.embedding, "Batch embedding");
+        }
+      } catch (err) {
+        if (!isInputRejection(err)) {
+          console.error("[memory] Batch embedding failed:", (err as Error).message);
+          continue;
+        }
+        // One rejected input fails the whole request: embed each input alone so
+        // only that input is cut down, and the rest embed as they are.
+        console.warn(
+          `[memory] Batch embedding rejected, embedding ${indices.length} inputs one by one: ${(err as Error).message}`,
+        );
+        for (const i of indices) {
+          results[i] = await this.embedOne(client, cleaned[i]!);
         }
       }
-
-      return results;
-    } catch (err) {
-      console.error("[memory] Batch embedding failed:", (err as Error).message);
-      return texts.map(() => null);
     }
+
+    return results;
   }
 }

package/src/be/memory/providers/sqlite-store.ts

@@ -12,8 +12,10 @@ import {
 } from "../constants";
 import { isLongtermKey, LONGTERM_ROOT, tierSource } from "../key-paths";
 import { recencyDecay } from "../reranker";
+import { MAX_CANDIDATE_FETCH } from "../search-candidates";
 import type {
   MemoryCandidate,
+  MemoryDocumentRef,
   MemoryEditInput,
   MemoryEditResult,
   MemoryHealth,
@@ -29,6 +31,16 @@ import type {
 
 const VECTOR_BYTES = EMBEDDING_DIMENSIONS * Float32Array.BYTES_PER_ELEMENT;
 
+/**
+ * Rows one hybrid or FTS arm fetches for a `limit`-row search: 4x overfetch up
+ * to 100, never fewer than `limit`. A fixed 100 ceiling returned the same rows
+ * for every wider limit, so searchMemoryCandidates read it as exhaustion while
+ * other documents sat behind 100 chunks of one task.
+ */
+function armFetchLimit(limit: number): number {
+  return Math.min(Math.max(Math.min(limit * 4, 100), limit), MAX_CANDIDATE_FETCH);
+}
+
 // memory_fts.memory_id is UNINDEXED (FTS5 cannot index it), so a CORRELATED
 // lookup against it is a full table scan per outer row. The old backfill used
 // `WHERE NOT EXISTS (SELECT 1 FROM memory_fts f WHERE f.memory_id = m.id)`,
@@ -512,6 +524,32 @@ export class SqliteMemoryStore implements MemoryStore {
     return results;
   }
 
+  async replaceDocument(
+    document: MemoryDocumentRef,
+    inputs: MemoryInput[],
+  ): Promise<AgentMemory[]> {
+    const results: AgentMemory[] = [];
+    await getDbClient().transaction(async (tx) => {
+      // Only this owner's rows of this source, in this scope, under this key.
+      // Another source holding the key is left alone and the insert fails.
+      const rows = await tx.query<{ id: string }>(
+        `SELECT id FROM agent_memory
+          WHERE key = ? AND scope = ? AND COALESCE(agentId, '') = COALESCE(?, '') AND source = ?`,
+        [document.key, document.scope, document.agentId ?? null, document.source],
+      );
+      const ids = rows.map((row) => row.id);
+      if (ids.length > 0) {
+        await this.purgeByIds(ids);
+        const placeholders = ids.map(() => "?").join(",");
+        await tx.run(`DELETE FROM agent_memory WHERE id IN (${placeholders})`, ids);
+      }
+      for (const input of inputs) {
+        results.push(await this.store(input));
+      }
+    });
+    return results;
+  }
+
   async get(id: string): Promise<AgentMemory | null> {
     const db = getDbClient();
     const row = await db.get<AgentMemoryRow>("SELECT * FROM agent_memory WHERE id = ?", [id]);
@@ -624,7 +662,7 @@ export class SqliteMemoryStore implements MemoryStore {
       keyPrefix?: string;
     },
   ): Promise<MemoryCandidate[]> {
-    const overfetchLimit = Math.min(Math.max(options.limit * 4, options.limit), 100);
+    const overfetchLimit = armFetchLimit(options.limit);
     const vectorCandidates = await this.searchWithVec(queryEmbedding, agentId, {
       ...options,
       limit: overfetchLimit,
@@ -725,7 +763,7 @@ export class SqliteMemoryStore implements MemoryStore {
     }
 
     try {
-      const sqlLimit = Math.min(Math.max(limit * 4, limit), 100);
+      const sqlLimit = armFetchLimit(limit);
       const rows = await getDbClient().query<AgentMemoryRow & { rank: number }>(
         `SELECT m.*, bm25(memory_fts) AS rank
          FROM memory_fts
@@ -1254,29 +1292,37 @@ export class SqliteMemoryStore implements MemoryStore {
   }
 
   async updateEmbedding(id: string, embedding: Float32Array, model: string): Promise<void> {
-    const db = getDbClient();
     const buffer = serializeEmbedding(embedding);
-    await db.run("UPDATE agent_memory SET embedding = ?, embeddingModel = ? WHERE id = ?", [
-      buffer,
-      model,
-      id,
-    ]);
-
-    if (this.vecInitialized && this.getVecTableSchema()) {
-      const vecBuffer = this.toVecBuffer(embedding);
-      if (!vecBuffer) {
-        console.warn(
-          `[memory-vec] update skipped memory_id=${id} reason=invalid_dimensions dimensions=${embedding.length} expected=${EMBEDDING_DIMENSIONS}`,
-        );
-        return;
-      }
-      try {
-        await db.run("DELETE FROM memory_vec WHERE memory_id = ?", [id]);
-        await db.run("INSERT INTO memory_vec(memory_id, embedding) VALUES (?, ?)", [id, vecBuffer]);
-      } catch (err) {
-        console.error(`[memory-vec] update failed memory_id=${id}: ${(err as Error).message}`);
+    // One transaction: a delete or replaceDocument queued behind this write
+    // cannot land between the row UPDATE and the vec insert and orphan it.
+    await getDbClient().transaction(async (tx) => {
+      const updated = await tx.run(
+        "UPDATE agent_memory SET embedding = ?, embeddingModel = ? WHERE id = ?",
+        [buffer, model, id],
+      );
+      // The row was deleted (or replaced) while its embedding was in flight:
+      // a vec row now would point at nothing.
+      if (updated.changes === 0) return;
+
+      if (this.vecInitialized && this.getVecTableSchema()) {
+        const vecBuffer = this.toVecBuffer(embedding);
+        if (!vecBuffer) {
+          console.warn(
+            `[memory-vec] update skipped memory_id=${id} reason=invalid_dimensions dimensions=${embedding.length} expected=${EMBEDDING_DIMENSIONS}`,
+          );
+          return;
+        }
+        try {
+          await tx.run("DELETE FROM memory_vec WHERE memory_id = ?", [id]);
+          await tx.run("INSERT INTO memory_vec(memory_id, embedding) VALUES (?, ?)", [
+            id,
+            vecBuffer,
+          ]);
+        } catch (err) {
+          console.error(`[memory-vec] update failed memory_id=${id}: ${(err as Error).message}`);
+        }
       }
-    }
+    });
   }
 
   async rewriteForScrub(

package/src/be/memory/reranker.ts

@@ -160,5 +160,35 @@ export function rerank(candidates: MemoryCandidate[], options: RerankOptions): M
   });
 
   scored.sort((a, b) => b.similarity - a.similarity || b.rawSimilarity - a.rawSimilarity);
-  return scored.slice(0, limit);
+  return collapseTaskCompletionChunks(scored).slice(0, limit);
+}
+
+/**
+ * The document rerank() collapses a candidate into, or null when the candidate
+ * is ranked on its own (anything but a chunk of a multi-chunk task memory).
+ */
+export function collapsedDocumentKey(
+  candidate: Pick<MemoryCandidate, "source" | "totalChunks" | "key" | "agentId" | "scope">,
+): string | null {
+  if (candidate.source !== "task_completion" || candidate.totalChunks <= 1 || !candidate.key) {
+    return null;
+  }
+  return JSON.stringify([candidate.agentId, candidate.scope, candidate.key]);
+}
+
+/**
+ * Keep only the best-ranked chunk of each multi-chunk task_completion memory.
+ * One long task output splits into a hundred chunks or more, which would
+ * otherwise fill every result slot. Chunks of one task share a key (see
+ * taskCompletionMemoryKey), scoped by owner and visibility.
+ */
+function collapseTaskCompletionChunks<T extends MemoryCandidate>(ranked: T[]): T[] {
+  const seen = new Set<string>();
+  return ranked.filter((candidate) => {
+    const document = collapsedDocumentKey(candidate);
+    if (document === null) return true;
+    if (seen.has(document)) return false;
+    seen.add(document);
+    return true;
+  });
 }

package/src/be/memory/search-candidates.ts

@@ -0,0 +1,47 @@
+/**
+ * Candidate fetch for memory search, sized in documents rather than rows.
+ *
+ * rerank() collapses the chunks of one task_completion document into its best
+ * chunk. When those chunks fill the candidate window, the collapse leaves
+ * fewer documents than the caller asked for while other matches sit just past
+ * the cutoff. This fetch widens the window until it holds `limit` documents,
+ * or the store has nothing more to return.
+ */
+import { collapsedDocumentKey } from "./reranker";
+import type { MemoryCandidate, MemorySearchOptions, MemoryStore } from "./types";
+
+/** Upper bound on one widened fetch. */
+export const MAX_CANDIDATE_FETCH = 4096;
+const WIDEN_FACTOR = 4;
+
+export function countCandidateDocuments(candidates: MemoryCandidate[]): number {
+  const documents = new Set<string>();
+  let standalone = 0;
+  for (const candidate of candidates) {
+    const document = collapsedDocumentKey(candidate);
+    if (document === null) standalone++;
+    else documents.add(document);
+  }
+  return standalone + documents.size;
+}
+
+export async function searchMemoryCandidates(
+  store: MemoryStore,
+  embedding: Float32Array,
+  agentId: string,
+  options: MemorySearchOptions & { limit: number },
+): Promise<MemoryCandidate[]> {
+  let fetchLimit = options.limit;
+  let rows = await store.search(embedding, agentId, { ...options, limit: fetchLimit });
+  while (fetchLimit < MAX_CANDIDATE_FETCH) {
+    const documents = countCandidateDocuments(rows);
+    // Enough documents, or no chunk duplicates for the collapse to remove.
+    if (documents >= options.limit || documents === rows.length) break;
+    const nextLimit = Math.min(fetchLimit * WIDEN_FACTOR, MAX_CANDIDATE_FETCH);
+    const wider = await store.search(embedding, agentId, { ...options, limit: nextLimit });
+    if (wider.length <= rows.length) break; // the store is exhausted
+    rows = wider;
+    fetchLimit = nextLimit;
+  }
+  return rows;
+}

package/src/be/memory/types.ts

@@ -20,6 +20,12 @@ export interface EmbeddingProvider {
 export interface MemoryStore {
   store(input: MemoryInput): Promise<AgentMemory>;
   storeBatch(inputs: MemoryInput[]): Promise<AgentMemory[]>;
+  /**
+   * In one transaction, delete every row of `document` (same key, scope,
+   * owner and source) and insert `inputs`. Rows of other owners, scopes or
+   * sources are never touched.
+   */
+  replaceDocument(document: MemoryDocumentRef, inputs: MemoryInput[]): Promise<AgentMemory[]>;
   get(id: string): Promise<AgentMemory | null>;
   peek(id: string): Promise<AgentMemory | null>;
   search(
@@ -79,6 +85,14 @@ export interface MemoryInput {
   key?: string | null;
 }
 
+/** The rows of one keyed document: every chunk shares these four fields. */
+export interface MemoryDocumentRef {
+  key: string;
+  scope: AgentMemoryScope;
+  agentId: string | null;
+  source: AgentMemorySource;
+}
+
 export interface MemoryCandidate extends AgentMemory {
   similarity: number;
   /** Raw cosine similarity before reranking (preserved for diagnostics). */

package/src/be/migrations/232_agent_session_tokens.sql

@@ -0,0 +1,41 @@
+PRAGMA defer_foreign_keys = ON;
+
+-- Add agent-bound worker tokens: a fourth binding that names only the agent.
+-- parentTokenId links a child task token to the agent token that minted it,
+-- so revoking the agent token revokes its children.
+CREATE TABLE session_tokens_new (
+  id TEXT PRIMARY KEY,
+  tokenHash TEXT NOT NULL UNIQUE,
+  tokenPreview TEXT NOT NULL,
+  agentId TEXT NOT NULL,
+  taskId TEXT,
+  runtimeId TEXT,
+  parentRuntimeId TEXT,
+  controllerId TEXT,
+  parentTokenId TEXT,
+  expiresAt TEXT NOT NULL,
+  revokedAt TEXT,
+  createdAt TEXT NOT NULL DEFAULT (datetime('now')),
+  lastUsedAt TEXT,
+  CHECK (
+    (taskId IS NOT NULL AND runtimeId IS NULL AND controllerId IS NULL
+      AND (parentRuntimeId IS NULL OR parentTokenId IS NULL))
+    OR (taskId IS NULL AND runtimeId IS NOT NULL AND parentRuntimeId IS NULL AND controllerId IS NULL AND parentTokenId IS NULL)
+    OR (controllerId IS NOT NULL AND taskId IS NULL AND runtimeId IS NULL AND parentRuntimeId IS NULL AND parentTokenId IS NULL)
+    OR (taskId IS NULL AND runtimeId IS NULL AND controllerId IS NULL AND parentRuntimeId IS NULL AND parentTokenId IS NULL)
+  )
+);
+
+INSERT INTO session_tokens_new
+  (id, tokenHash, tokenPreview, agentId, taskId, runtimeId, parentRuntimeId, controllerId,
+   expiresAt, revokedAt, createdAt, lastUsedAt)
+SELECT id, tokenHash, tokenPreview, agentId, taskId, runtimeId, parentRuntimeId, controllerId,
+  expiresAt, revokedAt, createdAt, lastUsedAt
+FROM session_tokens;
+
+DROP TABLE session_tokens;
+ALTER TABLE session_tokens_new RENAME TO session_tokens;
+CREATE INDEX idx_session_tokens_agentId ON session_tokens(agentId);
+CREATE INDEX idx_session_tokens_runtimeId ON session_tokens(runtimeId);
+CREATE INDEX idx_session_tokens_parentRuntimeId ON session_tokens(parentRuntimeId);
+CREATE INDEX idx_session_tokens_parentTokenId ON session_tokens(parentTokenId);

package/src/be/scripts/typecheck.ts

@@ -55,6 +55,8 @@ export type AgentTaskStepConfig = {
   routingNote?: string;
   tags?: string[];
   priority?: number;
+  /** Task IDs that must complete before this task starts. */
+  dependsOn?: string[];
   offerMode?: boolean;
   dir?: string;
   vcsRepo?: string;