@assistant-ui/core 0.3.17 → 0.3.19

package/src/utils/abortable-iterable.test.ts

@@ -0,0 +1,328 @@
+import { describe, expect, it, vi } from "vitest";
+import { abortableIterable, openAbortableIterable } from "./abortable-iterable";
+
+const deferred = <T>() => {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((r) => {
+    resolve = r;
+  });
+  return { promise, resolve };
+};
+
+const collect = async <T>(iterable: AsyncIterable<T>) => {
+  const seen: T[] = [];
+  for await (const value of iterable) seen.push(value);
+  return seen;
+};
+
+describe("abortableIterable", () => {
+  it.each([false, true])(
+    "releases the opening abort listener on rejection=%s",
+    async (fails) => {
+      const controller = new AbortController();
+      const added = vi.spyOn(controller.signal, "addEventListener");
+      const removed = vi.spyOn(controller.signal, "removeEventListener");
+      const source: AsyncIterable<number> = {
+        async *[Symbol.asyncIterator]() {
+          yield 1;
+        },
+      };
+      const opening = openAbortableIterable(
+        fails ? Promise.reject(new Error("opening failed")) : source,
+        controller.signal,
+      );
+      if (fails) await expect(opening).rejects.toThrow("opening failed");
+      else await expect(opening).resolves.toBe(source);
+      expect(added).toHaveBeenCalledTimes(1);
+      expect(removed).toHaveBeenCalledWith("abort", added.mock.calls[0]![1]);
+    },
+  );
+
+  it("finalizes once when an aborted read later rejects", async () => {
+    let rejectRead!: (error: unknown) => void;
+    const read = new Promise<IteratorResult<number>>((_, reject) => {
+      rejectRead = reject;
+    });
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const source: AsyncIterable<number> = {
+      [Symbol.asyncIterator]: () => ({ next: () => read, return: finalize }),
+    };
+    const controller = new AbortController();
+    const iterator = abortableIterable(source, controller.signal)[
+      Symbol.asyncIterator
+    ]();
+    const pending = iterator.next();
+    controller.abort();
+    await expect(pending).resolves.toEqual({ done: true, value: undefined });
+    expect(finalize).toHaveBeenCalledTimes(1);
+    rejectRead(new Error("late failure"));
+    await read.catch(() => {});
+    await iterator.return?.();
+    expect(finalize).toHaveBeenCalledTimes(1);
+  });
+
+  it("releases the listener and finalizes after a synchronous read failure", async () => {
+    const controller = new AbortController();
+    const added = vi.spyOn(controller.signal, "addEventListener");
+    const removed = vi.spyOn(controller.signal, "removeEventListener");
+    const error = new Error("read failed");
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const source: AsyncIterable<number> = {
+      [Symbol.asyncIterator]: () => ({
+        next: () => {
+          throw error;
+        },
+        return: finalize,
+      }),
+    };
+    await expect(
+      collect(abortableIterable(source, controller.signal)),
+    ).rejects.toBe(error);
+    expect(removed).toHaveBeenCalledWith("abort", added.mock.calls[0]![1]);
+    expect(finalize).toHaveBeenCalledTimes(1);
+    controller.abort();
+    expect(finalize).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not read or finalize again after the consumer returns", async () => {
+    const next = vi.fn(async () => ({ done: false as const, value: 1 }));
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const source: AsyncIterable<number> = {
+      [Symbol.asyncIterator]: () => ({ next, return: finalize }),
+    };
+    const iterator = abortableIterable(source, new AbortController().signal)[
+      Symbol.asyncIterator
+    ]();
+    await iterator.return?.();
+    await iterator.return?.();
+    await expect(iterator.next()).resolves.toEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(next).not.toHaveBeenCalled();
+    expect(finalize).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not finalize a naturally exhausted source", async () => {
+    const next = vi.fn(async () => ({ done: true as const, value: undefined }));
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const source: AsyncIterable<number> = {
+      [Symbol.asyncIterator]: () => ({ next, return: finalize }),
+    };
+    const iterator = abortableIterable(source, new AbortController().signal)[
+      Symbol.asyncIterator
+    ]();
+    await iterator.next();
+    await iterator.next();
+    await iterator.return?.();
+    expect(next).toHaveBeenCalledTimes(1);
+    expect(finalize).not.toHaveBeenCalled();
+  });
+
+  it("settles an opening stream on abort and finalizes a late iterable", async () => {
+    const opened = deferred<AsyncIterable<number>>();
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const controller = new AbortController();
+
+    const opening = openAbortableIterable(opened.promise, controller.signal);
+    controller.abort();
+
+    await expect(opening).resolves.toBeUndefined();
+    opened.resolve({
+      [Symbol.asyncIterator]: () => ({
+        next: () => Promise.resolve({ done: true, value: undefined }),
+        return: finalize,
+      }),
+    });
+    await vi.waitFor(() => expect(finalize).toHaveBeenCalledTimes(1));
+  });
+
+  it("passes a stream that completes on its own through untouched", async () => {
+    async function* source() {
+      yield 1;
+      yield 2;
+    }
+    const controller = new AbortController();
+
+    expect(
+      await collect(abortableIterable(source(), controller.signal)),
+    ).toEqual([1, 2]);
+  });
+
+  // the reason this wrapper exists: the source ignores the signal and parks on
+  // its own work, so nothing but the abort can end the loop
+  it("ends the loop while the source is parked mid-chunk", async () => {
+    const hang = deferred<void>();
+    const finallyRan = vi.fn();
+    async function* source() {
+      try {
+        yield 1;
+        await hang.promise;
+        yield 2;
+      } finally {
+        finallyRan();
+      }
+    }
+    const controller = new AbortController();
+
+    const seen: number[] = [];
+    const consuming = (async () => {
+      for await (const value of abortableIterable(
+        source(),
+        controller.signal,
+      )) {
+        seen.push(value);
+      }
+    })();
+    // let the loop take the first chunk and park awaiting the second
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    controller.abort();
+
+    // settles even though the source never does
+    await consuming;
+    expect(seen).toEqual([1]);
+    expect(finallyRan).not.toHaveBeenCalled();
+  });
+
+  it("finalizes a source still suspended at a yield when the signal aborts", async () => {
+    const finallyRan = vi.fn();
+    async function* source() {
+      try {
+        yield 1;
+        yield 2;
+      } finally {
+        finallyRan();
+      }
+    }
+    const controller = new AbortController();
+
+    const seen: number[] = [];
+    for await (const value of abortableIterable(source(), controller.signal)) {
+      seen.push(value);
+      controller.abort();
+    }
+
+    expect(seen).toEqual([1]);
+    expect(finallyRan).toHaveBeenCalledTimes(1);
+  });
+
+  it("reports a stream aborted before the first read as exhausted", async () => {
+    const next = vi.fn();
+    const controller = new AbortController();
+    controller.abort();
+
+    const seen = await collect(
+      abortableIterable(
+        {
+          [Symbol.asyncIterator]: () => ({ next }),
+        } as AsyncIterable<number>,
+        controller.signal,
+      ),
+    );
+
+    expect(seen).toEqual([]);
+    expect(next).not.toHaveBeenCalled();
+  });
+
+  it("propagates a source failure", async () => {
+    async function* source() {
+      yield 1;
+      throw new Error("stream failed");
+    }
+    const controller = new AbortController();
+
+    await expect(
+      collect(abortableIterable(source(), controller.signal)),
+    ).rejects.toThrow("stream failed");
+  });
+
+  it("finalizes a source after its next call rejects", async () => {
+    const finalize = vi.fn(async () => ({
+      done: true as const,
+      value: undefined,
+    }));
+    const source: AsyncIterable<number> = {
+      [Symbol.asyncIterator]: () => ({
+        next: () => Promise.reject(new Error("stream failed")),
+        return: finalize,
+      }),
+    };
+    const controller = new AbortController();
+
+    await expect(
+      collect(abortableIterable(source, controller.signal)),
+    ).rejects.toThrow("stream failed");
+    expect(finalize).toHaveBeenCalledTimes(1);
+  });
+
+  it("finalizes the source when the consumer breaks out", async () => {
+    const finallyRan = vi.fn();
+    async function* source() {
+      try {
+        yield 1;
+        yield 2;
+      } finally {
+        finallyRan();
+      }
+    }
+    const controller = new AbortController();
+
+    for await (const _ of abortableIterable(source(), controller.signal)) {
+      break;
+    }
+
+    // finalized without the break waiting on it
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(finallyRan).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not let a source that parks in its own cleanup stall the break", async () => {
+    const hang = deferred<void>();
+    async function* source() {
+      try {
+        yield 1;
+        yield 2;
+      } finally {
+        await hang.promise;
+      }
+    }
+    const controller = new AbortController();
+
+    // resolves even though the source's finally never does
+    for await (const _ of abortableIterable(source(), controller.signal)) {
+      break;
+    }
+  });
+
+  it("releases its abort listener after each chunk", async () => {
+    const controller = new AbortController();
+    const added = vi.spyOn(controller.signal, "addEventListener");
+    const removed = vi.spyOn(controller.signal, "removeEventListener");
+    async function* source() {
+      yield 1;
+      yield 2;
+      yield 3;
+    }
+
+    await collect(abortableIterable(source(), controller.signal));
+
+    expect(added).toHaveBeenCalledTimes(removed.mock.calls.length);
+    expect(added.mock.calls.length).toBeGreaterThan(0);
+  });
+});

package/src/utils/abortable-iterable.ts

@@ -0,0 +1,104 @@
+const done = <T>(): IteratorResult<T> => ({ done: true, value: undefined });
+
+export const openAbortableIterable = async <T>(
+  source: AsyncIterable<T> | Promise<AsyncIterable<T>>,
+  signal: AbortSignal,
+): Promise<AsyncIterable<T> | undefined> => {
+  const opened = Promise.resolve(source);
+  let release = () => {};
+  const aborted = new Promise<undefined>((resolve) => {
+    const onAbort = () => resolve(undefined);
+    if (signal.aborted) return onAbort();
+    signal.addEventListener("abort", onAbort, { once: true });
+    release = () => signal.removeEventListener("abort", onAbort);
+  });
+  try {
+    const iterable = await Promise.race([opened, aborted]);
+    if (!iterable) {
+      void opened
+        .then((late) => late[Symbol.asyncIterator]().return?.(undefined))
+        .catch(() => {});
+    }
+    return iterable;
+  } finally {
+    release();
+  }
+};
+
+/**
+ * Presents a stream as exhausted once the signal aborts, so a consumer stops
+ * waiting on it.
+ *
+ * A caller's stream is free to ignore the abort signal it is handed. One that
+ * also stops yielding, because it is awaiting its own work, leaves `for await`
+ * parked on `next()` with nothing left to wake it: the run never settles and
+ * whatever is serialized behind it never starts. Reporting the stream as
+ * exhausted ends the loop at cancellation instead.
+ *
+ * The source is finalized on the way out, but never awaited, since a stream
+ * that is already parked would not settle that either.
+ */
+export const abortableIterable = <T>(
+  source: AsyncIterable<T>,
+  signal: AbortSignal,
+): AsyncIterable<T> => ({
+  [Symbol.asyncIterator]: () => {
+    const iterator = source[Symbol.asyncIterator]();
+    let finished = false;
+    const finalize = () => {
+      if (finished) return;
+      finished = true;
+      try {
+        iterator.return?.(undefined)?.catch(() => {});
+      } catch {
+        // a source that throws on finalize has nothing left to clean up
+      }
+    };
+
+    return {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next: () => {
+        if (finished || signal.aborted) {
+          finalize();
+          return Promise.resolve(done<T>());
+        }
+
+        return new Promise<IteratorResult<T>>((resolve, reject) => {
+          // Registered per chunk and released as soon as the chunk settles, so
+          // a long stream does not accumulate one reaction per chunk.
+          const onAbort = () => {
+            finalize();
+            resolve(done<T>());
+          };
+          signal.addEventListener("abort", onAbort, { once: true });
+          const release = () => signal.removeEventListener("abort", onAbort);
+
+          const onError = (error: unknown) => {
+            release();
+            finalize();
+            reject(error);
+          };
+          try {
+            void iterator.next().then((result) => {
+              release();
+              if (finished) return resolve(done<T>());
+              if (result.done) finished = true;
+              resolve(result);
+            }, onError);
+          } catch (error) {
+            onError(error);
+          }
+        });
+      },
+      // The consumer breaks out of the loop to stop a cancelled run, so this is
+      // reached on the same path as an abort and finalizes the same way: a
+      // source whose own cleanup parks must not stall the break either.
+      return: async () => {
+        finalize();
+        return done<T>();
+      },
+    };
+  },
+});

package/src/utils/composite-context-provider.test.ts

@@ -114,6 +114,31 @@ describe("CompositeContextProvider", () => {
     expect(composite.getModelContext().system).toBeUndefined();
   });
 
+  it("finishes unregistering when the provider cleanup throws", () => {
+    const composite = new CompositeContextProvider();
+    const cleanupError = new Error("cleanup failed");
+    const subscriber = vi.fn();
+    const unsubscribe = vi.fn(() => {
+      throw cleanupError;
+    });
+
+    composite.subscribe(subscriber);
+    const unregister = composite.registerModelContextProvider({
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: () => unsubscribe,
+    });
+    subscriber.mockClear();
+
+    expect(() => unregister()).toThrow(cleanupError);
+    expect(composite.getModelContext().system).toBeUndefined();
+    expect(subscriber).toHaveBeenCalledOnce();
+    expect(unsubscribe).toHaveBeenCalledOnce();
+
+    expect(() => unregister()).not.toThrow();
+    expect(subscriber).toHaveBeenCalledOnce();
+    expect(unsubscribe).toHaveBeenCalledOnce();
+  });
+
   it("notifies every subscriber and rethrows on provider updates", () => {
     const composite = new CompositeContextProvider();
     const error = new Error("subscriber failed");

package/src/utils/composite-context-provider.ts

@@ -32,12 +32,33 @@ export class CompositeContextProvider implements ModelContextProvider {
     }
     this._providerUnsubscribes.set(id, unsubscribe);
     this.notifySubscribers();
+    let released = false;
     return () => {
+      if (released) return;
+      released = true;
       this._providers.delete(id);
       const unsubscribe = this._providerUnsubscribes.get(id);
-      unsubscribe?.();
       this._providerUnsubscribes.delete(id);
-      this.notifySubscribers();
+
+      let cleanupFailed = false;
+      let cleanupError: unknown;
+      const runCleanup = (cleanup: () => void) => {
+        try {
+          cleanup();
+        } catch (error) {
+          if (cleanupFailed) {
+            console.error(error);
+          } else {
+            cleanupFailed = true;
+            cleanupError = error;
+          }
+        }
+      };
+
+      if (unsubscribe) runCleanup(unsubscribe);
+      runCleanup(() => this.notifySubscribers());
+
+      if (cleanupFailed) throw cleanupError;
     };
   }
 

package/src/utils/getGroupStatus.test.ts

@@ -1,5 +1,13 @@
 import { describe, expect, it } from "vitest";
-import { getGroupStatus } from "./getGroupStatus";
+import type {
+  MessagePartStatus,
+  ToolCallMessagePartStatus,
+} from "../types/message";
+import { getGroupStatus, getGroupSummary } from "./getGroupStatus";
+
+type Status = MessagePartStatus | ToolCallMessagePartStatus;
+
+const part = (status: Status) => ({ status });
 
 describe("getGroupStatus", () => {
   it("reports running when the first member is running and the last is complete", () => {
@@ -10,7 +18,6 @@ describe("getGroupStatus", () => {
     const status = getGroupStatus(parts);
 
     expect(status).toEqual({ type: "running" });
-    expect(getGroupStatus(parts, [0, 1])).toBe(status);
     expect(status).toBe(getGroupStatus([{ status: { type: "running" } }]));
     expect(Object.isFrozen(status)).toBe(true);
   });
@@ -27,10 +34,69 @@ describe("getGroupStatus", () => {
 
   it("matches the positional outcome for statusless adapters", () => {
     expect(
-      getGroupStatus(
+      getGroupSummary(
         [{ status: { type: "complete" } }, { status: { type: "running" } }],
         [0, 1],
-      ),
+      ).status,
     ).toEqual({ type: "running" });
   });
 });
+
+describe("getGroupSummary", () => {
+  it("tallies every status over the requested indices and reports running when any part runs", () => {
+    const summary = getGroupSummary(
+      [
+        part({ type: "complete" }),
+        part({ type: "running" }),
+        part({ type: "incomplete", reason: "error" }),
+        part({ type: "requires-action", reason: "tool-calls" }),
+      ],
+      [0, 1, 2, 3],
+    );
+
+    expect(summary.counts).toEqual({
+      running: 1,
+      complete: 1,
+      incomplete: 1,
+      requiresAction: 1,
+    });
+    expect(
+      Object.values(summary.counts).reduce((sum, count) => sum + count),
+    ).toBe(4);
+    expect(summary.status).toEqual({ type: "running" });
+  });
+
+  it("uses the last status when no part runs and treats missing parts as complete", () => {
+    const summary = getGroupSummary(
+      [
+        undefined,
+        part({ type: "incomplete", reason: "cancelled" }),
+        part({ type: "requires-action", reason: "interrupt" }),
+      ],
+      [0, 1, 2],
+    );
+
+    expect(summary.counts).toEqual({
+      running: 0,
+      complete: 1,
+      incomplete: 1,
+      requiresAction: 1,
+    });
+    expect(summary.status).toEqual({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+  });
+
+  it("returns complete with zero counts for empty indices", () => {
+    expect(getGroupSummary([], [])).toEqual({
+      status: { type: "complete" },
+      counts: {
+        running: 0,
+        complete: 0,
+        incomplete: 0,
+        requiresAction: 0,
+      },
+    });
+  });
+});

package/src/utils/getGroupStatus.ts

@@ -10,22 +10,45 @@ type PartWithStatus = {
 
 export const getGroupStatus = (
   parts: readonly (PartWithStatus | undefined)[],
-  indices?: readonly number[],
 ): MessagePartStatus | ToolCallMessagePartStatus => {
-  if (indices) {
-    for (const index of indices) {
-      if (parts[index]?.status.type === "running") return RUNNING_STATUS;
-    }
-
-    const lastIndex = indices.at(-1);
-    return lastIndex === undefined
-      ? COMPLETE_STATUS
-      : (parts[lastIndex]?.status ?? COMPLETE_STATUS);
-  }
-
   for (const part of parts) {
     if (part?.status.type === "running") return RUNNING_STATUS;
   }
 
   return parts.at(-1)?.status ?? COMPLETE_STATUS;
 };
+
+export const getGroupSummary = (
+  parts: readonly (PartWithStatus | undefined)[],
+  indices: readonly number[],
+) => {
+  const counts = {
+    running: 0,
+    complete: 0,
+    incomplete: 0,
+    requiresAction: 0,
+  };
+  let status: MessagePartStatus | ToolCallMessagePartStatus = COMPLETE_STATUS;
+  let isRunning = false;
+
+  for (const index of indices) {
+    status = parts[index]?.status ?? COMPLETE_STATUS;
+    switch (status.type) {
+      case "running":
+        counts.running++;
+        isRunning = true;
+        break;
+      case "complete":
+        counts.complete++;
+        break;
+      case "incomplete":
+        counts.incomplete++;
+        break;
+      case "requires-action":
+        counts.requiresAction++;
+        break;
+    }
+  }
+
+  return { status: isRunning ? RUNNING_STATUS : status, counts };
+};

package/src/utils/record.test.ts

@@ -0,0 +1,48 @@
+import { describe, expect, it } from "vitest";
+import { nullProtoRecord } from "./record";
+
+describe("nullProtoRecord", () => {
+  it("has no prototype", () => {
+    expect(Object.getPrototypeOf(nullProtoRecord())).toBeNull();
+    expect(Object.getPrototypeOf(nullProtoRecord({ a: 1 }))).toBeNull();
+  });
+
+  it("copies own entries from every source, later sources winning", () => {
+    expect({ ...nullProtoRecord({ a: 1, b: 2 }, undefined, { b: 3 }) }).toEqual(
+      {
+        a: 1,
+        b: 3,
+      },
+    );
+  });
+
+  it.each(["__proto__", "constructor", "toString"])(
+    "assigns %s as an own property",
+    (key) => {
+      const record = nullProtoRecord<number>();
+      record[key] = 1;
+
+      expect(Object.hasOwn(record, key)).toBe(true);
+      expect(record[key]).toBe(1);
+      expect(Object.keys(record)).toEqual([key]);
+    },
+  );
+
+  it.each(["__proto__", "constructor", "toString"])(
+    "carries %s through a source",
+    (key) => {
+      const record = nullProtoRecord<number>({ [key]: 1, ok: 2 });
+
+      expect(Object.keys(record)).toEqual([key, "ok"]);
+      expect(record[key]).toBe(1);
+    },
+  );
+
+  it("reads an absent prototype-named key as undefined", () => {
+    const record = nullProtoRecord<number>({ ok: 1 });
+
+    expect(record["__proto__"]).toBeUndefined();
+    expect(record["constructor"]).toBeUndefined();
+    expect(record["toString"]).toBeUndefined();
+  });
+});