@assistant-ui/core 0.3.17 → 0.3.18

package/src/tests/tools-scope-migration.test.tsx

@@ -23,6 +23,34 @@ afterEach(() => {
 });
 
 describe("Tools model-context registration", () => {
+  it("registers and removes a tool UI named __proto__", async () => {
+    let aui!: AnyClient;
+    const Harness = () => {
+      aui = useAui({ tools: Tools({}) } as never);
+      return null;
+    };
+    render(<Harness />);
+
+    let remove!: () => void;
+    await act(async () => {
+      remove = aui.tools().setToolUI("__proto__", () => null);
+      await vi.waitFor(() => {
+        const toolUIs = aui.tools().getState().toolUIs;
+        expect(Object.hasOwn(toolUIs, "__proto__")).toBe(true);
+        expect(toolUIs.__proto__).toHaveLength(1);
+      });
+    });
+
+    await act(async () => {
+      remove();
+      await vi.waitFor(() =>
+        expect(Object.hasOwn(aui.tools().getState().toolUIs, "__proto__")).toBe(
+          false,
+        ),
+      );
+    });
+  });
+
   it("migrates to a structurally replaced modelContext scope", async () => {
     let aui!: AnyClient;
     const Harness = ({ generation }: { generation: number }) => {

package/src/types/attachment.ts

@@ -57,3 +57,7 @@ export const isCreateAttachment = (
   attachment: File | CreateAttachment,
 ): attachment is CreateAttachment =>
   "content" in attachment && !("lastModified" in attachment);
+
+export const isAttachmentComplete = (
+  attachment: Attachment,
+): attachment is CompleteAttachment => attachment.status.type === "complete";

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

@@ -0,0 +1,213 @@
+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("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,93 @@
+const done = <T>(): IteratorResult<T> => ({ done: true, value: undefined });
+
+/** Resolves at cancellation, so an await on a stream can stop being the only way out. */
+const whenAborted = (signal: AbortSignal): Promise<undefined> =>
+  new Promise((resolve) => {
+    if (signal.aborted) return resolve(undefined);
+    signal.addEventListener("abort", () => resolve(undefined), { once: true });
+  });
+
+export const openAbortableIterable = async <T>(
+  source: AsyncIterable<T> | Promise<AsyncIterable<T>>,
+  signal: AbortSignal,
+): Promise<AsyncIterable<T> | undefined> => {
+  const opened = Promise.resolve(source);
+  const iterable = await Promise.race([opened, whenAborted(signal)]);
+  if (!iterable) {
+    void opened
+      .then((late) => late[Symbol.asyncIterator]().return?.(undefined))
+      .catch(() => {});
+  }
+  return iterable;
+};
+
+/**
+ * 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]();
+    const finalize = () => {
+      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 (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);
+
+          iterator.next().then(
+            (result) => {
+              release();
+              resolve(result);
+            },
+            (error: unknown) => {
+              release();
+              finalize();
+              reject(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/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();
+  });
+});

package/src/utils/record.ts

@@ -0,0 +1,11 @@
+/**
+ * Builds a dictionary with no prototype from the own entries of its arguments.
+ * Bracket assignment on a plain object routes a `__proto__` key through the
+ * `Object.prototype` setter instead of defining an own property, so a record
+ * assigned into by a key the app does not control is built this way.
+ */
+export const nullProtoRecord = <T>(
+  base?: Readonly<Record<string, T>>,
+  ...rest: readonly (Readonly<Record<string, NoInfer<T>>> | undefined)[]
+): Record<string, T> =>
+  Object.assign(Object.create(null) as Record<string, T>, base, ...rest);