assistant-ui 0.0.119 → 0.0.121

package/src/lib/cloud-provision.ts

@@ -0,0 +1,93 @@
+import { randomUUID } from "node:crypto";
+import path from "node:path";
+import { readOptionalFile, writeCloudFile } from "./cloud-config";
+import { validateCloudUrl } from "./cloud-auth";
+
+export type CloudProvisionRequest = {
+  org_id: string;
+  project_id: string;
+  harness_name: string;
+  backend_urls: string[];
+  allow_localhost: boolean;
+  idempotency_key: string;
+  access_code?: string;
+};
+
+export type CloudProvisionResult = {
+  project_id: string;
+  workspace_id: string;
+  harness_id: string;
+  harness_origin: string;
+  api_key: string;
+  api_origin: string;
+};
+
+export const resolveProvisionRequest = async (
+  directory: string,
+  input: Omit<CloudProvisionRequest, "idempotency_key" | "access_code">,
+  accessCode?: string,
+): Promise<CloudProvisionRequest> => {
+  const file = path.join(directory, ".assistant-ui", "provision.json");
+  const previous = await readOptionalFile(file);
+  let request: CloudProvisionRequest;
+  if (previous) {
+    request = JSON.parse(previous) as CloudProvisionRequest;
+    const { idempotency_key: _key, ...savedInput } = request;
+    if (JSON.stringify(savedInput) !== JSON.stringify(input)) {
+      throw new Error(
+        "This directory has a pending setup for another project or backend. Resume that setup or use a new directory.",
+      );
+    }
+  } else {
+    request = { ...input, idempotency_key: randomUUID() };
+    await writeCloudFile(file, `${JSON.stringify(request, null, 2)}\n`, {
+      exclusive: true,
+    });
+  }
+  return { ...request, ...(accessCode ? { access_code: accessCode } : {}) };
+};
+
+export const provisionCloudHarness = async (
+  managementUrl: string,
+  accessToken: string,
+  request: CloudProvisionRequest,
+): Promise<CloudProvisionResult> => {
+  const response = await fetch(
+    `${validateCloudUrl(managementUrl)}/api/cli/provision`,
+    {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify(request),
+      signal: AbortSignal.timeout(30_000),
+      redirect: "error",
+    },
+  );
+  if (!response.ok) {
+    const body = (await response.json().catch(() => null)) as {
+      error?: string;
+      message?: string;
+    } | null;
+    const detail = body?.message ?? body?.error ?? response.statusText;
+    throw new Error(`Cloud setup failed (${response.status}): ${detail}`);
+  }
+  const result = (await response.json()) as CloudProvisionResult;
+  if (
+    !result ||
+    [
+      result.api_key,
+      result.harness_id,
+      result.workspace_id,
+      result.harness_origin,
+      result.api_origin,
+    ].some((value) => typeof value !== "string" || value.length === 0) ||
+    result.project_id !== request.project_id
+  ) {
+    throw new Error(
+      "Cloud returned an incomplete setup response. Rerun the same command to resume safely.",
+    );
+  }
+  return result;
+};

package/src/lib/cloud-setup-login.test.ts

@@ -0,0 +1,361 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { connectCloudLoginSetup, setupLoginUrl } from "./cloud-setup-login";
+
+const fixture = vi.hoisted(() => ({
+  state: {
+    version: 2,
+    id: "session-a",
+    status: "planning",
+    inputs: [] as {
+      id: string;
+      status: string;
+      preset?: string;
+      help?: { href?: string };
+    }[],
+  } as
+    | {
+        version: number;
+        id: string | null;
+        status: string;
+        inputs: {
+          id: string;
+          status: string;
+          preset?: string;
+          help?: { href?: string };
+        }[];
+      }
+    | undefined,
+  status: "connected",
+  listeners: new Set<() => void>(),
+  ask: vi.fn(),
+  answer: vi.fn(),
+  dispose: vi.fn(),
+  transport: vi.fn(),
+}));
+vi.mock("statewire", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("statewire")>()),
+  StatewireHttp: (...args: unknown[]) => fixture.transport(...args),
+  StatewireClient: class {
+    get state() {
+      return fixture.state;
+    }
+    get connection() {
+      return { status: fixture.status };
+    }
+    commands = { "agent/ask": fixture.ask, "checkout/answer": fixture.answer };
+    subscribe = (listener: () => void) => {
+      fixture.listeners.add(listener);
+      return () => fixture.listeners.delete(listener);
+    };
+    dispose = fixture.dispose;
+  },
+}));
+
+const issuer = "https://accounts.assistant-ui.com";
+const authorization = () => ({
+  deviceCode: "private-device-code",
+  userCode: "ABCD-EFGH",
+  verificationUri: `${issuer}/device`,
+  verificationUriComplete: `${issuer}/device?user_code=ABCD-EFGH&device_code=do-not-forward`,
+  expiresAt: Date.now() + 60_000,
+  intervalMs: 5_000,
+  resource: `${issuer}/api/auth`,
+});
+const update = () => {
+  for (const listener of fixture.listeners) listener();
+};
+afterEach(() => {
+  fixture.state = {
+    version: 2,
+    id: "session-a",
+    status: "planning",
+    inputs: [],
+  };
+  fixture.status = "connected";
+  fixture.listeners.clear();
+  vi.useRealTimers();
+});
+
+describe("cloud setup device login", () => {
+  it("shares only the public code, trusted approval URL, expiration and attempt", () => {
+    const url = new URL(setupLoginUrl(authorization(), issuer));
+    expect(url.origin).toBe(issuer);
+    expect([...url.searchParams.keys()]).toEqual([
+      "user_code",
+      "setup_attempt",
+      "setup_expires",
+    ]);
+    expect(url.searchParams.get("user_code")).toBe("ABCD-EFGH");
+    expect(url.searchParams.get("setup_attempt")).toMatch(/^[0-9a-f-]{36}$/);
+    expect(url.href).not.toContain("private-device-code");
+    expect(url.href).not.toContain("device_code");
+  });
+  it.each([
+    {
+      verificationUriComplete:
+        "https://accounts.assistant-ui.com.evil.test/device?user_code=ABCD-EFGH",
+    },
+    { verificationUriComplete: `${issuer}/device?user_code=WRONG` },
+    {
+      verificationUriComplete: `${issuer}/device?user_code=ABCD-EFGH&user_code=SECOND`,
+    },
+    { verificationUriComplete: `${issuer}/device?user_code=ABCD-EFGH#token` },
+    { expiresAt: 1 },
+    { expiresAt: Date.now() + 60 * 60_000 },
+  ])("rejects unexpected approval metadata %j", (patch) => {
+    expect(() =>
+      setupLoginUrl({ ...authorization(), ...patch }, issuer),
+    ).toThrow("invalid device approval link");
+  });
+  it("rejects arbitrary issuers and malformed attempt identifiers", () => {
+    expect(() =>
+      setupLoginUrl(authorization(), "https://accounts.other.test"),
+    ).toThrow();
+    expect(() =>
+      setupLoginUrl(authorization(), issuer, "../redirect"),
+    ).toThrow();
+  });
+  it("publishes a reserved input and sends only terminal status after CLI approval", async () => {
+    fixture.ask.mockResolvedValue({ inputId: "q1" });
+    fixture.answer.mockResolvedValue(undefined);
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    await bridge.publish(authorization(), issuer);
+    expect(fixture.ask).toHaveBeenLastCalledWith(
+      expect.objectContaining({
+        kind: "text",
+        preset: "assistant-ui-cli-login",
+      }),
+    );
+    const call = fixture.ask.mock.lastCall![0];
+    expect(JSON.stringify(call)).not.toContain("private-device-code");
+    await bridge.complete("signed-in");
+    expect(fixture.answer).toHaveBeenLastCalledWith({
+      inputId: "q1",
+      answer: "signed-in",
+    });
+    await bridge.dispose();
+    expect(fixture.listeners.size).toBe(0);
+  });
+  it("ignores a forged answer, but cancels a dismissed input", async () => {
+    fixture.ask.mockResolvedValue({ inputId: "q2" });
+    const bridge = await connectCloudLoginSetup(
+      "http://127.0.0.1:8791/session-a",
+    );
+    await bridge.publish(authorization(), issuer);
+    fixture.state!.inputs = [{ id: "q2", status: "answered" }];
+    update();
+    expect(bridge.signal.aborted).toBe(false);
+    fixture.state!.inputs[0]!.status = "dismissed";
+    update();
+    expect(bridge.signal.aborted).toBe(true);
+    const count = fixture.ask.mock.calls.length;
+    await expect(bridge.publish(authorization(), issuer)).rejects.toThrow(
+      "cancelled",
+    );
+    expect(fixture.ask.mock.calls.length).toBe(count);
+    await bridge.dispose();
+  });
+  it.each(["cancelled", "done"])(
+    "cancels when the checkout becomes %s",
+    async (status) => {
+      const bridge = await connectCloudLoginSetup(
+        "https://checkout.test/session-a",
+      );
+      fixture.state!.status = status;
+      update();
+      expect(bridge.signal.aborted).toBe(true);
+      await bridge.dispose();
+    },
+  );
+  it("cancels when the original session is replaced", async () => {
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    fixture.state!.id = "session-b";
+    update();
+    expect(bridge.signal.aborted).toBe(true);
+    await bridge.dispose();
+  });
+  it.each([null, 1, "input", {}, { id: "q1" }, { status: "pending" }])(
+    "rejects malformed initial input entries: %j",
+    async (input) => {
+      fixture.state!.inputs = [input] as unknown as NonNullable<
+        typeof fixture.state
+      >["inputs"];
+      await expect(
+        connectCloudLoginSetup("https://checkout.test/session-a"),
+      ).rejects.toThrow("Start the setup");
+      expect(fixture.ask).not.toHaveBeenCalled();
+      expect(fixture.listeners.size).toBe(0);
+    },
+  );
+  it.each([null, 1, "input", {}, { id: "q1" }, { status: "pending" }])(
+    "cancels without throwing or answering malformed updated entries: %j",
+    async (input) => {
+      fixture.ask.mockResolvedValueOnce({ inputId: "q1" });
+      const bridge = await connectCloudLoginSetup(
+        "https://checkout.test/session-a",
+      );
+      await bridge.publish(authorization(), issuer);
+      fixture.state!.inputs = [input] as unknown as NonNullable<
+        typeof fixture.state
+      >["inputs"];
+      expect(update).not.toThrow();
+      expect(bridge.signal.aborted).toBe(true);
+      await bridge.complete("cancelled");
+      expect(fixture.answer).not.toHaveBeenCalled();
+      await bridge.dispose();
+    },
+  );
+  it("disposes a connection that never supplies a snapshot before its deadline", async () => {
+    fixture.state = undefined;
+    fixture.status = "connecting";
+    const deadline = new AbortController();
+    const timeout = vi
+      .spyOn(AbortSignal, "timeout")
+      .mockReturnValueOnce(deadline.signal);
+    const count = fixture.dispose.mock.calls.length;
+    try {
+      const connection = connectCloudLoginSetup(
+        "https://checkout.test/session-a",
+      );
+      deadline.abort(new Error("snapshot timeout"));
+      await expect(connection).rejects.toThrow("snapshot timeout");
+      expect(fixture.dispose.mock.calls.length).toBe(count + 1);
+      expect(fixture.listeners.size).toBe(0);
+    } finally {
+      timeout.mockRestore();
+    }
+  });
+  it("observes dismissal even when it arrives before the ask response", async () => {
+    let finish!: (result: { inputId: string }) => void;
+    fixture.ask.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+        }),
+    );
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    const publishing = bridge.publish(authorization(), issuer);
+    fixture.state!.inputs = [{ id: "q3", status: "dismissed" }];
+    update();
+    finish({ inputId: "q3" });
+    await expect(publishing).rejects.toThrow("cancelled");
+    expect(bridge.signal.aborted).toBe(true);
+    await bridge.dispose();
+  });
+  it("reconciles an accepted prompt from its snapshot when the ask acknowledgment times out", async () => {
+    let finish!: (result: { inputId: string }) => void;
+    fixture.ask.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+        }),
+    );
+    fixture.answer.mockResolvedValue(undefined);
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    const deadline = new AbortController();
+    const timeout = vi
+      .spyOn(AbortSignal, "timeout")
+      .mockReturnValueOnce(deadline.signal);
+    const publishing = bridge.publish(authorization(), issuer);
+    const href = fixture.ask.mock.lastCall![0].help.href;
+    fixture.state!.inputs = [
+      {
+        id: "late",
+        status: "pending",
+        preset: "assistant-ui-cli-login",
+        help: { href },
+      },
+    ];
+    update();
+    deadline.abort(new Error("ask timeout"));
+    await expect(publishing).rejects.toThrow("ask timeout");
+    expect(fixture.answer).toHaveBeenLastCalledWith({
+      inputId: "late",
+      answer: "failed",
+    });
+    const disposing = bridge.dispose();
+    finish({ inputId: "late" });
+    await disposing;
+    expect(fixture.ask).toHaveBeenCalledOnce();
+    expect(fixture.answer).toHaveBeenCalledOnce();
+    expect(fixture.listeners.size).toBe(0);
+    timeout.mockRestore();
+  });
+  it("reconciles a late accepted prompt after cancellation without replaying the ask", async () => {
+    let finish!: (result: { inputId: string }) => void;
+    fixture.ask.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+        }),
+    );
+    fixture.answer.mockResolvedValue(undefined);
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    const deadline = new AbortController();
+    const timeout = vi
+      .spyOn(AbortSignal, "timeout")
+      .mockReturnValueOnce(deadline.signal);
+    const publishing = bridge.publish(authorization(), issuer);
+    deadline.abort(new Error("ask timeout"));
+    await expect(publishing).rejects.toThrow("ask timeout");
+    await bridge.complete("cancelled");
+    const disposing = bridge.dispose();
+    finish({ inputId: "late" });
+    await disposing;
+    expect(fixture.answer).toHaveBeenLastCalledWith({
+      inputId: "late",
+      answer: "cancelled",
+    });
+    expect(fixture.ask).toHaveBeenCalledOnce();
+    timeout.mockRestore();
+  });
+  it("never reconciles a late acknowledgment into a replacement session with reused input IDs", async () => {
+    let finish!: (result: { inputId: string }) => void;
+    fixture.ask.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+        }),
+    );
+    const bridge = await connectCloudLoginSetup(
+      "https://checkout.test/session-a",
+    );
+    const publishing = bridge.publish(authorization(), issuer);
+    fixture.state!.id = "session-b";
+    fixture.state!.inputs = [{ id: "q1", status: "pending" }];
+    update();
+    await expect(publishing).rejects.toThrow("cancelled");
+    await bridge.complete("cancelled");
+    const disposing = bridge.dispose();
+    finish({ inputId: "q1" });
+    await disposing;
+    expect(fixture.answer).not.toHaveBeenCalled();
+  });
+  it("disposes a stopped connection before requesting a device grant", async () => {
+    fixture.state = undefined;
+    fixture.status = "stopped";
+    const count = fixture.dispose.mock.calls.length;
+    await expect(
+      connectCloudLoginSetup("https://checkout.test/session-a"),
+    ).rejects.toThrow("stopped");
+    expect(fixture.dispose.mock.calls.length).toBe(count + 1);
+    expect(fixture.listeners.size).toBe(0);
+  });
+  it.each([
+    "https://user:pass@checkout.test/session",
+    "http://remote.test/session",
+    "https://checkout.test/session?token=x",
+  ])("rejects unsafe setup targets %s", async (url) => {
+    await expect(connectCloudLoginSetup(url)).rejects.toThrow("HTTPS");
+  });
+});

package/src/lib/cloud-setup-login.ts

@@ -0,0 +1,240 @@
+import { randomUUID } from "node:crypto";
+import type { DeviceAuthorization } from "aui-auth/device";
+import { StatewireClient, StatewireHttp } from "statewire";
+import { abortable } from "./cloud-abort";
+import { validateCloudUrl } from "./cloud-url";
+
+type SetupState = {
+  version: number;
+  id: string | null;
+  status: string;
+  inputs: {
+    id: string;
+    status: string;
+    preset?: string;
+    help?: { href?: string };
+  }[];
+};
+
+type SetupCommands = {
+  "agent/ask": (input: {
+    kind: "text";
+    preset: string;
+    prompt: string;
+    help: { summary: string; href: string };
+  }) => { inputId: string };
+  "checkout/answer": (answer: { inputId: string; answer: string }) => void;
+};
+
+const validInputs = (inputs: unknown): inputs is SetupState["inputs"] =>
+  Array.isArray(inputs) &&
+  inputs.every(
+    (input) =>
+      input !== null &&
+      typeof input === "object" &&
+      typeof input.id === "string" &&
+      typeof input.status === "string",
+  );
+
+export const setupLoginUrl = (
+  authorization: DeviceAuthorization,
+  issuer: string,
+  attempt: string = randomUUID(),
+) => {
+  const url = new URL(authorization.verificationUriComplete);
+  const issuerOrigin = new URL(issuer).origin;
+  if (
+    !["https://accounts.assistant-ui.com", "https://accounts.aui.dev"].includes(
+      issuerOrigin,
+    ) ||
+    url.origin !== issuerOrigin ||
+    !/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i.test(
+      attempt,
+    ) ||
+    !/^[A-Z0-9-]{4,32}$/.test(authorization.userCode) ||
+    url.pathname !== "/device" ||
+    url.username ||
+    url.password ||
+    url.hash ||
+    url.searchParams.get("user_code") !== authorization.userCode ||
+    url.searchParams.getAll("user_code").length !== 1 ||
+    !Number.isFinite(authorization.expiresAt) ||
+    authorization.expiresAt <= Date.now() ||
+    authorization.expiresAt > Date.now() + 10 * 60 * 1000
+  ) {
+    throw new Error("Accounts returned an invalid device approval link.");
+  }
+  const publicUrl = new URL("/device", url.origin);
+  publicUrl.searchParams.set("user_code", authorization.userCode);
+  publicUrl.searchParams.set("setup_attempt", attempt);
+  publicUrl.searchParams.set("setup_expires", String(authorization.expiresAt));
+  return publicUrl.href;
+};
+
+export const connectCloudLoginSetup = async (url: string) => {
+  const target = validateCloudUrl(url);
+  const controller = new AbortController();
+  const client = new StatewireClient<SetupState | undefined, SetupCommands>({
+    transport: StatewireHttp({ url: target }),
+    onError: (error) => controller.abort(error),
+  });
+  let unsubscribe = () => {};
+  let inputId: string | undefined;
+  let sessionId: string;
+  let pendingHref: string | undefined;
+  let pendingAsk: Promise<{ inputId: string }> | undefined;
+  let terminalResult: "signed-in" | "cancelled" | "failed" | undefined;
+  let completion: Promise<void> | undefined;
+  let disposed = false;
+  const reconcile = () => {
+    const state = client.state;
+    if (state && !validInputs(state.inputs)) {
+      controller.abort(new Error("The setup returned invalid inputs."));
+      return Promise.resolve();
+    }
+    if (
+      disposed ||
+      inputId === undefined ||
+      terminalResult === undefined ||
+      state?.id !== sessionId ||
+      !Array.isArray(state.inputs) ||
+      !["planning", "installing"].includes(state.status) ||
+      ["dismissed", "answered"].includes(
+        state.inputs.find((input) => input.id === inputId)?.status ?? "",
+      )
+    )
+      return Promise.resolve();
+    completion ??= abortable(
+      client.commands["checkout/answer"]({ inputId, answer: terminalResult }),
+      AbortSignal.timeout(10_000),
+    );
+    return completion;
+  };
+  const checkSession = () => {
+    const next = client.state;
+    if (next && !validInputs(next.inputs)) {
+      controller.abort(new Error("The setup returned invalid inputs."));
+      return;
+    }
+    if (next?.id === sessionId && Array.isArray(next.inputs) && pendingHref) {
+      inputId ??= next.inputs.find(
+        (input) =>
+          input.preset === "assistant-ui-cli-login" &&
+          input.help?.href === pendingHref,
+      )?.id;
+    }
+    if (
+      next?.version !== 2 ||
+      next.id !== sessionId ||
+      !Array.isArray(next.inputs) ||
+      !["planning", "installing"].includes(next.status) ||
+      (inputId !== undefined &&
+        next.inputs.find((input) => input.id === inputId)?.status ===
+          "dismissed")
+    ) {
+      controller.abort(new Error("The setup login was cancelled."));
+    }
+    if (client.connection.status === "stopped")
+      controller.abort(new Error("The setup connection stopped."));
+    void reconcile().catch(() => {});
+  };
+  const dispose = async () => {
+    try {
+      if (pendingAsk && terminalResult) {
+        await abortable(pendingAsk, AbortSignal.timeout(10_000)).catch(
+          () => {},
+        );
+        await reconcile().catch(() => {});
+      }
+    } finally {
+      disposed = true;
+      unsubscribe();
+      client.dispose();
+    }
+  };
+  try {
+    let resolveReady!: () => void;
+    const ready = new Promise<void>((resolve) => {
+      resolveReady = resolve;
+    });
+    unsubscribe = client.subscribe(() => {
+      if (client.state !== undefined) resolveReady();
+      if (client.connection.status === "stopped") {
+        controller.abort(new Error("The setup connection stopped."));
+      }
+    });
+    if (client.state !== undefined) resolveReady();
+    if (client.connection.status === "stopped") {
+      controller.abort(new Error("The setup connection stopped."));
+    }
+    await abortable(
+      ready,
+      AbortSignal.any([controller.signal, AbortSignal.timeout(10_000)]),
+    );
+    const state = client.state!;
+    if (
+      state?.version !== 2 ||
+      !state.id ||
+      !validInputs(state.inputs) ||
+      !["planning", "installing"].includes(state.status)
+    ) {
+      throw new Error(
+        "Start the setup in the wizard before signing in the CLI.",
+      );
+    }
+    sessionId = state.id;
+    unsubscribe();
+    unsubscribe = client.subscribe(checkSession);
+    checkSession();
+  } catch (error) {
+    await dispose();
+    throw error;
+  }
+  const command = <T>(run: () => Promise<T>) => {
+    controller.signal.throwIfAborted();
+    return abortable(
+      run(),
+      AbortSignal.any([controller.signal, AbortSignal.timeout(10_000)]),
+    );
+  };
+  return {
+    signal: controller.signal,
+    async publish(authorization: DeviceAuthorization, issuer: string) {
+      controller.signal.throwIfAborted();
+      const href = setupLoginUrl(authorization, issuer);
+      pendingHref = href;
+      pendingAsk = client.commands["agent/ask"]({
+        kind: "text",
+        preset: "assistant-ui-cli-login",
+        prompt: "Sign in to Assistant Cloud",
+        help: {
+          summary:
+            "Authorize the CLI in your browser, then return to this setup.",
+          href,
+        },
+      }).then(async (result) => {
+        if (!disposed && client.state?.id === sessionId) {
+          inputId = result.inputId;
+          checkSession();
+          await reconcile().catch(() => {});
+        }
+        return result;
+      });
+      try {
+        await command(() => pendingAsk!);
+        checkSession();
+        controller.signal.throwIfAborted();
+      } catch (error) {
+        terminalResult = controller.signal.aborted ? "cancelled" : "failed";
+        await reconcile().catch(() => {});
+        throw error;
+      }
+      return href;
+    },
+    async complete(result: "signed-in" | "cancelled" | "failed") {
+      terminalResult = result;
+      await reconcile();
+    },
+    dispose,
+  };
+};