@vellumai/assistant 0.12.3-staging.2 → 0.12.4-staging.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. package/Dockerfile +7 -0
  2. package/README.md +12 -12
  3. package/docker/mousepad/90-vellum.gschema.override +4 -0
  4. package/openapi.yaml +6 -2
  5. package/package.json +1 -1
  6. package/src/__tests__/approval-routes-http.test.ts +4 -2
  7. package/src/__tests__/call-routes-http.test.ts +4 -2
  8. package/src/__tests__/conversation-resurface-done.test.ts +289 -0
  9. package/src/__tests__/gemini-provider.test.ts +31 -0
  10. package/src/__tests__/provider-error-scenarios.test.ts +46 -0
  11. package/src/__tests__/provider-platform-proxy-integration.test.ts +45 -15
  12. package/src/__tests__/retry-fallback-escalation.test.ts +22 -0
  13. package/src/__tests__/runtime-attachment-metadata.test.ts +4 -2
  14. package/src/__tests__/secret-routes-platform-proxy.test.ts +1 -0
  15. package/src/__tests__/send-endpoint-busy.test.ts +4 -2
  16. package/src/__tests__/steer-on-enqueue-question.test.ts +265 -34
  17. package/src/cli/commands/channel-verification-sessions.help.ts +12 -2
  18. package/src/cli/lib/bundled-marketplace.json +4 -4
  19. package/src/config/feature-flag-registry.json +8 -8
  20. package/src/config/sidebar-done-gate.ts +27 -0
  21. package/src/daemon/conversation-turn-finalize.ts +12 -0
  22. package/src/daemon/handlers/__tests__/conversation-history-search.test.ts +61 -0
  23. package/src/daemon/handlers/config-channels.ts +3 -1
  24. package/src/daemon/handlers/conversation-history.ts +11 -1
  25. package/src/daemon/handlers/conversations.ts +118 -10
  26. package/src/desktop/desktop-browser-endpoint.ts +1 -0
  27. package/src/ipc/routes/__tests__/conversation-sync-ipc-routes.test.ts +49 -1
  28. package/src/ipc/routes/conversation-sync-ipc-routes.ts +38 -2
  29. package/src/permissions/question-prompter.ts +14 -4
  30. package/src/persistence/conversation-crud.ts +114 -3
  31. package/src/persistence/conversation-queries.ts +48 -3
  32. package/src/persistence/embeddings/__tests__/embedding-local-lifecycle.test.ts +157 -32
  33. package/src/persistence/embeddings/embedding-backend-shutdown.test.ts +47 -0
  34. package/src/persistence/embeddings/embedding-backend.ts +23 -41
  35. package/src/persistence/embeddings/embedding-gemini.test.ts +17 -0
  36. package/src/persistence/embeddings/embedding-gemini.ts +3 -1
  37. package/src/persistence/embeddings/embedding-local.ts +72 -17
  38. package/src/persistence/embeddings/embedding-types.ts +0 -19
  39. package/src/plugins/defaults/memory/__tests__/memory-retrospective-skill-card.test.ts +2 -0
  40. package/src/plugins/defaults/memory/memory-retrospective-skill-card.ts +5 -0
  41. package/src/providers/gemini/client.ts +17 -1
  42. package/src/providers/malformed-tool-call.ts +55 -0
  43. package/src/providers/openai/__tests__/chat-completions-malformed-tool-call.test.ts +105 -0
  44. package/src/providers/openai/chat-completions-provider.ts +30 -0
  45. package/src/providers/platform-proxy/constants.ts +5 -0
  46. package/src/providers/retry.ts +68 -36
  47. package/src/providers/vellum-model-routing.test.ts +1 -0
  48. package/src/runtime/__tests__/agent-wake-archived-gate.test.ts +90 -0
  49. package/src/runtime/__tests__/question-resolution.test.ts +32 -0
  50. package/src/runtime/agent-wake.ts +9 -2
  51. package/src/runtime/auth/__tests__/middleware.test.ts +138 -0
  52. package/src/runtime/auth/middleware.ts +15 -11
  53. package/src/runtime/middleware/auth.ts +7 -1
  54. package/src/runtime/pending-interactions.ts +8 -8
  55. package/src/runtime/question-resolution.ts +54 -8
  56. package/src/runtime/routes/__tests__/conversation-list-routes.test.ts +146 -0
  57. package/src/runtime/routes/__tests__/global-search-routes.test.ts +74 -1
  58. package/src/runtime/routes/conversation-list-routes.ts +16 -14
  59. package/src/runtime/routes/conversation-query-routes.ts +4 -0
  60. package/src/runtime/routes/global-search-routes.ts +22 -7
  61. package/src/runtime/sync/documents-sidecar-publish.test.ts +50 -4
  62. package/src/runtime/sync/resource-sync-events.ts +11 -0
  63. package/src/runtime/sync/worker-daemon-notify.test.ts +30 -0
  64. package/src/runtime/sync/worker-daemon-notify.ts +44 -0
  65. package/src/runtime/verification-outbound-actions.ts +2 -2
  66. package/src/tools/ask-question/ask-question-tool.test.ts +38 -0
  67. package/src/tools/ask-question/ask-question-tool.ts +12 -1
  68. package/src/util/worker-pipe.test.ts +100 -0
  69. package/src/util/worker-pipe.ts +44 -0
@@ -27,9 +27,14 @@ import { LocalEmbeddingBackend } from "../embedding-local.js";
27
27
  type Internals = any;
28
28
 
29
29
  /**
30
- * A stand-in for a Bun subprocess whose stdin pipe is broken. `write` throws
31
- * EPIPE the way a real worker's pipe does once the child is gone.
30
+ * A backend of its own for one test. Production gets the per-model instance
31
+ * from `forModel`; constructing directly keeps tests from sharing state.
32
32
  */
33
+ function newBackend(): Internals {
34
+ return new (LocalEmbeddingBackend as Internals)("test-model");
35
+ }
36
+
37
+ /** A stand-in for a Bun subprocess whose stdin `write` fails synchronously. */
33
38
  function brokenPipeProc(pid = 4242) {
34
39
  return {
35
40
  pid,
@@ -196,7 +201,7 @@ describe("PID file ownership", () => {
196
201
  */
197
202
  test("releasePidFile leaves an entry that names someone else's worker", () => {
198
203
  writeFileSync(pidPath, "777");
199
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
204
+ const backend = newBackend();
200
205
 
201
206
  backend.releasePidFile(888);
202
207
 
@@ -206,7 +211,7 @@ describe("PID file ownership", () => {
206
211
 
207
212
  test("releasePidFile drops the entry when it names our worker", () => {
208
213
  writeFileSync(pidPath, "777");
209
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
214
+ const backend = newBackend();
210
215
 
211
216
  backend.releasePidFile(777);
212
217
 
@@ -219,7 +224,7 @@ describe("PID file ownership", () => {
219
224
  */
220
225
  test("disposing an instance with no worker keeps another's PID entry", () => {
221
226
  writeFileSync(pidPath, "777");
222
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
227
+ const backend = newBackend();
223
228
  backend.disposeRequested = true;
224
229
  backend.workerProc = null;
225
230
 
@@ -232,12 +237,12 @@ describe("PID file ownership", () => {
232
237
 
233
238
  describe("broken worker pipe", () => {
234
239
  /**
235
- * A write to a dead worker raises EPIPE. Escaping, it reached the daemon's
236
- * `unhandledRejection` handler and terminated the process; it must instead
237
- * come back as an ordinary failed embed the backend chain can fall back from.
240
+ * A broken pipe comes back as an ordinary failed embed the backend chain can
241
+ * fall back from. An escaping rejection reaches the daemon's
242
+ * `unhandledRejection` handler, which shuts the process down.
238
243
  */
239
244
  test("EPIPE on write resolves the request as an error", async () => {
240
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
245
+ const backend = newBackend();
241
246
  backend.workerProc = brokenPipeProc();
242
247
 
243
248
  const response = await backend.sendRequest(["hello"]);
@@ -247,7 +252,7 @@ describe("broken worker pipe", () => {
247
252
  });
248
253
 
249
254
  test("EPIPE surfaces to embed() as a rejection, not a process crash", async () => {
250
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
255
+ const backend = newBackend();
251
256
  backend.workerProc = brokenPipeProc();
252
257
  // Skip initialization: the worker is already (notionally) running.
253
258
  backend.ensureInitialized = async () => {};
@@ -257,8 +262,34 @@ describe("broken worker pipe", () => {
257
262
  );
258
263
  });
259
264
 
265
+ /**
266
+ * A real pipe, because the failure lives in Bun's behaviour rather than in
267
+ * ours: a batch larger than the pipe buffer, sent to a worker too busy to
268
+ * drain stdin, leaves the write pending, and the pending write rejects when
269
+ * the worker dies. A synchronous guard never sees that rejection.
270
+ */
271
+ test("a write still pending when the worker dies resolves the request as an error", async () => {
272
+ const backend = newBackend();
273
+ const proc = Bun.spawn({
274
+ cmd: [process.execPath, "-e", "setTimeout(() => {}, 60_000)"],
275
+ windowsHide: true,
276
+ stdin: "pipe",
277
+ stdout: "ignore",
278
+ stderr: "ignore",
279
+ });
280
+ spawned.push(proc);
281
+ backend.workerProc = proc;
282
+
283
+ const request = backend.sendRequest(["x".repeat(4 * 1024 * 1024)]);
284
+ proc.kill("SIGKILL");
285
+ const response = await request;
286
+
287
+ expect(response.error).toContain("worker pipe write failed");
288
+ expect(backend.pendingRequests.size).toBe(0);
289
+ });
290
+
260
291
  test("a failed write does not leak the pending request", async () => {
261
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
292
+ const backend = newBackend();
262
293
  backend.workerProc = brokenPipeProc();
263
294
 
264
295
  await backend.sendRequest(["hello"]);
@@ -274,7 +305,7 @@ describe("releasing ownership of a worker", () => {
274
305
  * strands it: no handle, no PID entry, invisible to every later reclaim.
275
306
  */
276
307
  test("stdout ending kills the child before clearing workerProc", async () => {
277
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
308
+ const backend = newBackend();
278
309
  const proc = liveProc();
279
310
  backend.workerProc = proc;
280
311
 
@@ -285,7 +316,7 @@ describe("releasing ownership of a worker", () => {
285
316
  });
286
317
 
287
318
  test("shutdown reaps the worker and resolves only once it exits", async () => {
288
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
319
+ const backend = newBackend();
289
320
  const proc = liveProc(555);
290
321
  backend.workerProc = proc;
291
322
  writeFileSync(getEmbedWorkerPidPath(), "555");
@@ -298,7 +329,7 @@ describe("releasing ownership of a worker", () => {
298
329
  });
299
330
 
300
331
  test("shutdown settles in-flight requests instead of hanging them", async () => {
301
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
332
+ const backend = newBackend();
302
333
  backend.workerProc = liveProc(556);
303
334
 
304
335
  const inFlight = backend.sendRequest(["hello"]);
@@ -315,7 +346,7 @@ describe("releasing ownership of a worker", () => {
315
346
  * pending requests are only settled after the wait.
316
347
  */
317
348
  test("a worker that ignores SIGTERM is escalated to SIGKILL, not waited on forever", async () => {
318
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
349
+ const backend = newBackend();
319
350
  backend.terminateGraceMs = 50;
320
351
  const signals: (string | undefined)[] = [];
321
352
  const wedged = {
@@ -354,7 +385,7 @@ describe("termination that cannot be confirmed", () => {
354
385
  }
355
386
 
356
387
  test("keeps the PID publication instead of unpublishing a possibly-live child", async () => {
357
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
388
+ const backend = newBackend();
358
389
  backend.terminateGraceMs = 20;
359
390
  const proc = wedgedProc();
360
391
  backend.workerProc = proc;
@@ -370,7 +401,7 @@ describe("termination that cannot be confirmed", () => {
370
401
  });
371
402
 
372
403
  test("refuses to start a replacement while the child may still be alive", async () => {
373
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
404
+ const backend = newBackend();
374
405
  // No workerPath recorded, so liveness falls back to a PID probe. Our own
375
406
  // PID is certainly alive, standing in for a child that will not die.
376
407
  backend.unconfirmedWorker = process.pid;
@@ -381,7 +412,7 @@ describe("termination that cannot be confirmed", () => {
381
412
  });
382
413
 
383
414
  test("surfaces the incomplete teardown through embed() rather than failing silently", async () => {
384
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
415
+ const backend = newBackend();
385
416
  backend.unconfirmedWorker = process.pid;
386
417
 
387
418
  await expect(backend.embed(["hello"])).rejects.toThrow(
@@ -401,7 +432,7 @@ describe("termination that cannot be confirmed", () => {
401
432
  });
402
433
  await Bun.sleep(400);
403
434
 
404
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
435
+ const backend = newBackend();
405
436
  backend.workerPath = scriptPath;
406
437
  backend.unconfirmedWorker = proc.pid;
407
438
  writeFileSync(getEmbedWorkerPidPath(), String(proc.pid));
@@ -440,7 +471,7 @@ describe("reclaiming before spawning a replacement", () => {
440
471
  await Bun.sleep(400);
441
472
  writeFileSync(getEmbedWorkerPidPath(), String(proc.pid));
442
473
 
443
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
474
+ const backend = newBackend();
444
475
  await backend.reclaimOwnedWorkers(scriptPath);
445
476
 
446
477
  // Confirmed gone by the time reclaim returns, not merely signalled.
@@ -464,7 +495,7 @@ describe("reclaiming before spawning a replacement", () => {
464
495
  spawned.push(proc);
465
496
  await Bun.sleep(400);
466
497
 
467
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
498
+ const backend = newBackend();
468
499
  backend.terminateGraceMs = 300;
469
500
  await backend.reclaimOwnedWorkers(scriptPath);
470
501
 
@@ -497,7 +528,7 @@ setTimeout(() => {}, 60_000);\n`,
497
528
  expect(workers.length).toBe(1);
498
529
  expect(workers[0].ppid).toBe(parent.pid);
499
530
 
500
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
531
+ const backend = newBackend();
501
532
  backend.terminateGraceMs = 300;
502
533
  await backend.reclaimOwnedWorkers(scriptPath);
503
534
 
@@ -521,7 +552,7 @@ describe("shutdown versus in-flight initialization", () => {
521
552
  * lets the initializer spawn a worker afterwards, with nobody left to reap it.
522
553
  */
523
554
  test("shutdown waits for an in-flight init and reaps the worker it creates", async () => {
524
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
555
+ const backend = newBackend();
525
556
  const proc = liveProc(600);
526
557
 
527
558
  let finishInit: () => void = () => {};
@@ -564,7 +595,7 @@ describe("parent exit orderings", () => {
564
595
  * process's and left it alone. So it kills its own child on the way out.
565
596
  */
566
597
  test("terminateNow kills the worker and releases the entry without waiting", () => {
567
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
598
+ const backend = newBackend();
568
599
  const signals: string[] = [];
569
600
  const proc = {
570
601
  pid: 909,
@@ -590,7 +621,7 @@ describe("parent exit orderings", () => {
590
621
  * process while the child was still being adopted.
591
622
  */
592
623
  test("shutdown does not inherit the model-load wait when init never settles", async () => {
593
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
624
+ const backend = newBackend();
594
625
  backend.terminateGraceMs = 50;
595
626
  // An initialization that never resolves, as during a cold model load.
596
627
  backend.initInFlight = new Promise<void>(() => {});
@@ -602,7 +633,7 @@ describe("parent exit orderings", () => {
602
633
  });
603
634
 
604
635
  test("a backend already shutting down refuses to spawn a replacement", async () => {
605
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
636
+ const backend = newBackend();
606
637
  backend.disposeRequested = true;
607
638
 
608
639
  await expect(backend.embed(["hello"])).rejects.toThrow(/shutting down/);
@@ -621,7 +652,7 @@ describe("transient failures stay transient", () => {
621
652
  const PERMANENT_FAILURE_SENTINEL = "Local embedding backend unavailable";
622
653
 
623
654
  test("an unconfirmed termination does not claim permanent unavailability", async () => {
624
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
655
+ const backend = newBackend();
625
656
  backend.unconfirmedWorker = process.pid;
626
657
 
627
658
  const err = await backend.ensureInitialized().catch((e: Error) => e);
@@ -631,7 +662,7 @@ describe("transient failures stay transient", () => {
631
662
  });
632
663
 
633
664
  test("shutting down does not claim permanent unavailability", async () => {
634
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
665
+ const backend = newBackend();
635
666
  backend.disposeRequested = true;
636
667
 
637
668
  const err = await backend.embed(["hello"]).catch((e: Error) => e);
@@ -652,7 +683,7 @@ describe("audit regressions", () => {
652
683
  * embed queued behind it.
653
684
  */
654
685
  test("a startup failure against a wedged worker does not hang", async () => {
655
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
686
+ const backend = newBackend();
656
687
  backend.terminateGraceMs = 30;
657
688
  const signals: (string | undefined)[] = [];
658
689
  backend.workerProc = {
@@ -680,7 +711,7 @@ describe("audit regressions", () => {
680
711
  * startup for the full two-minute timeout.
681
712
  */
682
713
  test("terminateNow clears the stdout buffer so a retry can parse ready", () => {
683
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
714
+ const backend = newBackend();
684
715
  backend.stdoutBuffer = '{"id":5,"vect';
685
716
  backend.workerProc = null;
686
717
 
@@ -707,7 +738,7 @@ describe("audit regressions", () => {
707
738
  spawned.push(proc);
708
739
  await Bun.sleep(400);
709
740
 
710
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
741
+ const backend = newBackend();
711
742
  backend.workerPath = scriptPath;
712
743
  backend.workerProc = null; // spawned, never adopted
713
744
 
@@ -722,7 +753,7 @@ describe("audit regressions", () => {
722
753
  * contract, so it must leave the instance able to start a replacement.
723
754
  */
724
755
  test("terminateNow leaves the instance reusable", () => {
725
- const backend = new LocalEmbeddingBackend("test-model") as Internals;
756
+ const backend = newBackend();
726
757
  backend.workerProc = null;
727
758
  backend.stdoutReaderActive = true;
728
759
 
@@ -766,3 +797,97 @@ test("orphan reclaim outlasts the embedding reap it must not truncate", () => {
766
797
  EMBEDDING_SHUTDOWN_BUDGET_MS + 1_000,
767
798
  );
768
799
  });
800
+
801
+ describe("one backend per model", () => {
802
+ test("forModel hands out one instance per model", () => {
803
+ const first = LocalEmbeddingBackend.forModel("per-model-a");
804
+
805
+ expect(LocalEmbeddingBackend.forModel("per-model-a")).toBe(first);
806
+ expect(LocalEmbeddingBackend.forModel("per-model-b")).not.toBe(first);
807
+ });
808
+
809
+ /**
810
+ * A backend-cache reset disposes the backend it evicts, and that disposal
811
+ * waits for in-flight embeds. The next embed needs the model again inside
812
+ * that window, and the embeds in flight keep their worker.
813
+ */
814
+ test("needing the model again while a disposal is pending keeps the busy worker", () => {
815
+ const backend = LocalEmbeddingBackend.forModel(
816
+ "per-model-reopen",
817
+ ) as Internals;
818
+ const proc = liveProc(7101);
819
+ backend.workerProc = proc;
820
+ backend.activeEmbeds = 1;
821
+ backend.dispose();
822
+
823
+ const again = LocalEmbeddingBackend.forModel(
824
+ "per-model-reopen",
825
+ ) as Internals;
826
+ backend.activeEmbeds = 0;
827
+ backend.disposeIfIdle();
828
+
829
+ expect(again).toBe(backend);
830
+ expect(again.workerProc).toBe(proc);
831
+ expect(proc.killed).toBe(false);
832
+ });
833
+
834
+ test("a pending disposal nobody withdraws still releases the worker once idle", () => {
835
+ const backend = LocalEmbeddingBackend.forModel(
836
+ "per-model-dispose",
837
+ ) as Internals;
838
+ const proc = liveProc(7102);
839
+ backend.workerProc = proc;
840
+ backend.activeEmbeds = 1;
841
+ backend.dispose();
842
+
843
+ backend.activeEmbeds = 0;
844
+ backend.disposeIfIdle();
845
+
846
+ expect(backend.workerProc).toBeNull();
847
+ expect(proc.killed).toBe(true);
848
+ });
849
+
850
+ /** The process is exiting, so nothing may spawn a worker behind the reap. */
851
+ test("a shut-down backend stays shut down", async () => {
852
+ const backend = LocalEmbeddingBackend.forModel("per-model-shutdown");
853
+ await backend.shutdown();
854
+
855
+ const again = LocalEmbeddingBackend.forModel("per-model-shutdown");
856
+
857
+ expect(again).toBe(backend);
858
+ await expect(again.embed(["hello"])).rejects.toThrow(/shutting down/);
859
+ });
860
+
861
+ /**
862
+ * The backend cache forgets a backend whose disposal is pending, so process
863
+ * teardown reaches it through the class's own list.
864
+ */
865
+ test("shutdownAll reaps a busy backend no cache still holds", async () => {
866
+ const backend = LocalEmbeddingBackend.forModel(
867
+ "per-model-evicted",
868
+ ) as Internals;
869
+ const proc = liveProc(7103);
870
+ backend.workerProc = proc;
871
+ backend.activeEmbeds = 1;
872
+ backend.dispose();
873
+
874
+ await LocalEmbeddingBackend.shutdownAll();
875
+
876
+ expect(proc.killed).toBe(true);
877
+ expect(backend.workerProc).toBeNull();
878
+ });
879
+
880
+ test("terminateAllNow kills a busy backend no cache still holds", () => {
881
+ const backend = LocalEmbeddingBackend.forModel(
882
+ "per-model-evicted-now",
883
+ ) as Internals;
884
+ const proc = liveProc(7104);
885
+ backend.workerProc = proc;
886
+ backend.activeEmbeds = 1;
887
+ backend.dispose();
888
+
889
+ LocalEmbeddingBackend.terminateAllNow();
890
+
891
+ expect(proc.killed).toBe(true);
892
+ });
893
+ });
@@ -0,0 +1,47 @@
1
+ /**
2
+ * Process teardown reaches every local embedding worker the process owns.
3
+ *
4
+ * Kept in its own file: `shutdownEmbeddingBackends` latches the module into its
5
+ * shut-down state for the rest of the process.
6
+ */
7
+ import { expect, test } from "bun:test";
8
+
9
+ import { AssistantConfigSchema } from "../../config/schema.js";
10
+ import {
11
+ clearEmbeddingBackendCache,
12
+ selectEmbeddingBackend,
13
+ shutdownEmbeddingBackends,
14
+ } from "./embedding-backend.js";
15
+
16
+ /** Reach past `private`, which is compile-time only, so tests drive real state. */
17
+ type Internals = any;
18
+
19
+ /**
20
+ * A backend-cache reset forgets a backend whose disposal is waiting on
21
+ * in-flight embeds, while that backend still holds a live worker.
22
+ */
23
+ test("shutdown reaps a busy local worker the cache has already forgotten", async () => {
24
+ const config = AssistantConfigSchema.parse({
25
+ memory: { embeddings: { provider: "local" } },
26
+ });
27
+ const { backend } = await selectEmbeddingBackend(config);
28
+ const delegate = await (backend as Internals).getDelegate();
29
+ let killed = false;
30
+ delegate.workerProc = {
31
+ pid: 7201,
32
+ exited: new Promise<number>(() => {}),
33
+ kill() {
34
+ killed = true;
35
+ },
36
+ stdin: { write: () => 0, flush: () => 0 },
37
+ };
38
+ delegate.terminateGraceMs = 10;
39
+ delegate.activeEmbeds = 1;
40
+
41
+ clearEmbeddingBackendCache();
42
+ expect(killed).toBe(false);
43
+
44
+ await shutdownEmbeddingBackends();
45
+
46
+ expect(killed).toBe(true);
47
+ });
@@ -16,6 +16,7 @@ import {
16
16
  recordBillingSuccess,
17
17
  } from "./embedding-billing-breaker.js";
18
18
  import { GeminiEmbeddingBackend } from "./embedding-gemini.js";
19
+ import type { LocalEmbeddingBackend } from "./embedding-local.js";
19
20
  import { OllamaEmbeddingBackend } from "./embedding-ollama.js";
20
21
  import {
21
22
  OpenAIEmbeddingBackend,
@@ -78,6 +79,13 @@ export const EMBEDDING_SHUTDOWN_BUDGET_MS = 8_000;
78
79
  * the failure is contained and other embedding backends can be used instead.
79
80
  */
80
81
 
82
+ /**
83
+ * The local backend class, once {@link LazyLocalEmbeddingBackend} has loaded it.
84
+ * It lists every local backend this process owns, which process teardown needs
85
+ * and {@link backendCache} cannot supply.
86
+ */
87
+ let localEmbeddingBackends: typeof LocalEmbeddingBackend | null = null;
88
+
81
89
  class LazyLocalEmbeddingBackend implements EmbeddingBackend {
82
90
  readonly provider = "local" as const;
83
91
  readonly model: string;
@@ -114,21 +122,9 @@ class LazyLocalEmbeddingBackend implements EmbeddingBackend {
114
122
  this.delegate?.dispose?.();
115
123
  }
116
124
 
117
- terminateNow(): void {
118
- this.delegate?.terminateNow?.();
119
- }
120
-
121
- async sweepOwnedWorkers(): Promise<void> {
122
- await this.delegate?.sweepOwnedWorkers?.();
123
- }
124
-
125
- async shutdown(): Promise<void> {
126
- // A delegate under construction still ends up owning a worker, so settle
127
- // the in-flight import before tearing down rather than skipping it.
128
- if (!this.delegate && this.initPromise) {
129
- await this.initPromise.catch(() => undefined);
130
- }
131
- await this.delegate?.shutdown?.();
125
+ /** Settles once a delegate under construction has been handed out or failed. */
126
+ async loaded(): Promise<void> {
127
+ await this.initPromise?.catch(() => undefined);
132
128
  }
133
129
 
134
130
  resetForRetry(): void {
@@ -146,7 +142,8 @@ class LazyLocalEmbeddingBackend implements EmbeddingBackend {
146
142
  try {
147
143
  const { LocalEmbeddingBackend } =
148
144
  await import("./embedding-local.js");
149
- this.delegate = new LocalEmbeddingBackend(this.model);
145
+ localEmbeddingBackends = LocalEmbeddingBackend;
146
+ this.delegate = LocalEmbeddingBackend.forModel(this.model);
150
147
  return this.delegate;
151
148
  } catch (err) {
152
149
  localBackendBroken = true;
@@ -270,7 +267,7 @@ export function clearEmbeddingBackendCache(): void {
270
267
  }
271
268
 
272
269
  /**
273
- * Tear down every cached backend's OS resources and empty the caches.
270
+ * Tear down every local embedding worker this process owns and empty the caches.
274
271
  *
275
272
  * Called on daemon shutdown so process-owned embedding workers exit with their
276
273
  * owner instead of being orphaned. `clearEmbeddingBackendCache()` is the
@@ -279,24 +276,18 @@ export function clearEmbeddingBackendCache(): void {
279
276
  */
280
277
  export async function shutdownEmbeddingBackends(): Promise<void> {
281
278
  embeddingBackendsShutDown = true;
282
- const backends = new Set(backendCache.values());
279
+ const loading = [...backendCache.values()].filter(
280
+ (backend) => backend instanceof LazyLocalEmbeddingBackend,
281
+ );
283
282
  backendCache.clear();
284
283
  vectorCache.clear();
285
284
  vectorCacheBytes = 0;
286
285
  backendDimCache.clear();
287
286
 
288
- const teardown = Promise.all(
289
- [...backends].map(async (backend) => {
290
- try {
291
- await backend.shutdown?.();
292
- await backend.sweepOwnedWorkers?.();
293
- } catch (err) {
294
- log.warn(
295
- { err, provider: backend.provider, model: backend.model },
296
- "Failed to shut down embedding backend",
297
- );
298
- }
299
- }),
287
+ // A delegate under construction still ends up owning a worker, so let it
288
+ // join the class's list before tearing that list down.
289
+ const teardown = Promise.all(loading.map((backend) => backend.loaded())).then(
290
+ () => localEmbeddingBackends?.shutdownAll(),
300
291
  );
301
292
 
302
293
  const timedOut = Symbol("timeout");
@@ -313,7 +304,7 @@ export async function shutdownEmbeddingBackends(): Promise<void> {
313
304
  }
314
305
 
315
306
  /**
316
- * SIGKILL every cached backend's worker synchronously.
307
+ * SIGKILL every local embedding worker this process owns, synchronously.
317
308
  *
318
309
  * For a process that must exit immediately and cannot run the graceful
319
310
  * teardown: the memory worker on PID-file eviction, where staying alive to reap
@@ -322,16 +313,7 @@ export async function shutdownEmbeddingBackends(): Promise<void> {
322
313
  * sweep has passed, leaving two workers alive (JARVIS-1125).
323
314
  */
324
315
  export function terminateEmbeddingWorkersNow(): void {
325
- for (const backend of new Set(backendCache.values())) {
326
- try {
327
- backend.terminateNow?.();
328
- } catch (err) {
329
- log.warn(
330
- { err, provider: backend.provider, model: backend.model },
331
- "Failed to terminate embedding worker",
332
- );
333
- }
334
- }
316
+ localEmbeddingBackends?.terminateAllNow();
335
317
  }
336
318
 
337
319
  /** Reset the sticky local-backend failure flag without evicting live backends. */
@@ -439,6 +439,23 @@ describe("GeminiEmbeddingBackend: batched text inputs", () => {
439
439
  expect(vectors.map((v) => v[0])).toEqual(texts(250).map(textIndex));
440
440
  });
441
441
 
442
+ test("managed text inputs use single calls from the first request and preserve vector order", async () => {
443
+ const fetchMock = routedFetch();
444
+ globalThis.fetch = fetchMock as unknown as typeof fetch;
445
+ const backend = new GeminiEmbeddingBackend("test-key", "test-model", {
446
+ managedBaseUrl: "https://proxy.example.com/v1/runtime-proxy/gemini",
447
+ interCallDelayMs: 0,
448
+ });
449
+
450
+ expect(await backend.embed(texts(3))).toEqual([[0], [1], [2]]);
451
+ expect(await backend.embed(["t4", "t5"])).toEqual([[4], [5]]);
452
+ expect(calledUrls(fetchMock)).toEqual(
453
+ Array(5).fill(
454
+ "https://proxy.example.com/v1/runtime-proxy/gemini/v1beta/models/test-model:embedContent",
455
+ ),
456
+ );
457
+ });
458
+
442
459
  test("each batched request names the model with the models/ prefix and carries taskType and outputDimensionality", async () => {
443
460
  const fetchMock = routedFetch();
444
461
  globalThis.fetch = fetchMock as unknown as typeof fetch;
@@ -71,7 +71,8 @@ export class GeminiEmbeddingBackend implements EmbeddingBackend {
71
71
  }
72
72
 
73
73
  /**
74
- * Embed `inputs` in order. Runs of text inputs go through
74
+ * Embed `inputs` in order. Managed requests use single `embedContent`
75
+ * calls because Vertex does not expose the batch route. Direct text runs use
75
76
  * `batchEmbedContents`, {@link GEMINI_EMBED_BATCH_SIZE} texts per round
76
77
  * trip, so a corpus re-embed costs one request per hundred sections rather
77
78
  * than one per section; a lone text and every multimodal input take the
@@ -99,6 +100,7 @@ export class GeminiEmbeddingBackend implements EmbeddingBackend {
99
100
  // Gather the run of text inputs starting here, up to one batch.
100
101
  let end = i;
101
102
  while (
103
+ !this.managed &&
102
104
  !this.batchRouteUnavailable &&
103
105
  end < normalized.length &&
104
106
  end - i < GEMINI_EMBED_BATCH_SIZE &&