@assistant-ui/vscode 0.0.0-stage → 0.0.0

package/src/host/router.ts

@@ -0,0 +1,253 @@
+import {
+  isWebviewToHostMessage,
+  VSCODE_BRIDGE_CHANNEL,
+  VSCODE_VIRTUAL_ORIGIN,
+  type FetchRequestMessage,
+  type HostToWebviewMessage,
+} from "../protocol";
+import { isForeignHttpUrl } from "../route-url";
+
+export type RouteHandler = (request: Request) => Response | Promise<Response>;
+
+export type HttpMethod =
+  | "GET"
+  | "HEAD"
+  | "POST"
+  | "PUT"
+  | "PATCH"
+  | "DELETE"
+  | "OPTIONS";
+
+export type RouteHandlers = { [M in HttpMethod]?: RouteHandler };
+
+export type WebviewRoutes = Record<string, RouteHandlers>;
+
+export type WebviewLike = {
+  postMessage(message: any): PromiseLike<boolean>;
+  onDidReceiveMessage(listener: (message: any) => unknown): {
+    dispose(): unknown;
+  };
+};
+
+export type ServeWebviewRoutesOptions = {
+  /** Milliseconds to coalesce response chunks before posting them. */
+  flushInterval?: number;
+  /** Buffered byte count that posts a chunk without waiting. */
+  flushSize?: number;
+  onError?: (error: unknown) => void;
+};
+
+export type Disposable = { dispose(): void };
+
+const BODYLESS_METHODS = new Set(["GET", "HEAD"]);
+
+const concat = (parts: Uint8Array[], size: number) => {
+  const out = new Uint8Array(size);
+  let offset = 0;
+  for (const part of parts) {
+    out.set(part, offset);
+    offset += part.byteLength;
+  }
+  return out;
+};
+
+const errorMessage = (error: unknown) =>
+  error instanceof Error ? error.message : String(error);
+
+export function serveRoutes(
+  webview: WebviewLike,
+  routes: WebviewRoutes,
+  {
+    flushInterval = 16,
+    flushSize = 64 * 1024,
+    onError = (error) => console.error(error),
+  }: ServeWebviewRoutesOptions = {},
+): Disposable {
+  const inflight = new Map<string, AbortController>();
+
+  const post = (message: HostToWebviewMessage, controller: AbortController) => {
+    Promise.resolve(webview.postMessage(message)).then(
+      (delivered) => {
+        if (!delivered) controller.abort();
+      },
+      () => controller.abort(),
+    );
+  };
+
+  const dispatch = async (
+    request: Request,
+    pathname: string,
+  ): Promise<Response> => {
+    const route = Object.hasOwn(routes, pathname) ? routes[pathname] : null;
+    if (!route) return new Response("Not Found", { status: 404 });
+    const handler = Object.hasOwn(route, request.method)
+      ? route[request.method as HttpMethod]
+      : null;
+    if (!handler) {
+      return new Response("Method Not Allowed", {
+        status: 405,
+        headers: { Allow: Object.keys(route).join(", ") },
+      });
+    }
+    try {
+      return await handler(request);
+    } catch (error) {
+      if (request.signal.aborted) throw error;
+      onError(error);
+      return new Response("Internal Server Error", { status: 500 });
+    }
+  };
+
+  const pump = async (
+    id: string,
+    body: ReadableStream<Uint8Array>,
+    controller: AbortController,
+  ) => {
+    const { signal } = controller;
+    const reader = body.getReader();
+    const onAbort = () => {
+      reader.cancel(signal.reason).catch(() => undefined);
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+
+    let buffered: Uint8Array[] = [];
+    let size = 0;
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const flush = () => {
+      if (timer !== undefined) clearTimeout(timer);
+      timer = undefined;
+      if (size === 0 || signal.aborted) return;
+      const chunk = concat(buffered, size);
+      buffered = [];
+      size = 0;
+      post(
+        { channel: VSCODE_BRIDGE_CHANNEL, kind: "fetch:chunk", id, chunk },
+        controller,
+      );
+    };
+
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done || signal.aborted) break;
+        if (value.byteLength === 0) continue;
+        buffered.push(value);
+        size += value.byteLength;
+        if (size >= flushSize || flushInterval <= 0) flush();
+        else timer ??= setTimeout(flush, flushInterval);
+      }
+      flush();
+    } catch (error) {
+      flush();
+      throw error;
+    } finally {
+      if (timer !== undefined) clearTimeout(timer);
+      signal.removeEventListener("abort", onAbort);
+      reader.releaseLock();
+    }
+  };
+
+  const handle = async (message: FetchRequestMessage) => {
+    const { id } = message;
+    if (inflight.has(id)) {
+      onError(new Error(`Ignored a request reusing the in-flight id ${id}`));
+      return;
+    }
+    const controller = new AbortController();
+    inflight.set(id, controller);
+    try {
+      const url = new URL(message.url);
+      if (isForeignHttpUrl(url)) {
+        throw new TypeError(`Refused a request for ${url.origin}`);
+      }
+      // The webview chooses this URL, so a route sees only its path and query.
+      const requestUrl = new URL(VSCODE_VIRTUAL_ORIGIN);
+      requestUrl.pathname = url.pathname;
+      requestUrl.search = url.search;
+      const request = new Request(requestUrl, {
+        method: message.method,
+        headers: message.headers,
+        body: BODYLESS_METHODS.has(message.method.toUpperCase())
+          ? null
+          : message.body,
+        signal: controller.signal,
+      });
+      const response = await dispatch(request, url.pathname);
+      if (controller.signal.aborted) {
+        void response.body?.cancel().catch(() => undefined);
+        return;
+      }
+      post(
+        {
+          channel: VSCODE_BRIDGE_CHANNEL,
+          kind: "fetch:head",
+          id,
+          status: response.status,
+          statusText: response.statusText,
+          headers: [...response.headers],
+        },
+        controller,
+      );
+      if (request.method === "HEAD") {
+        void response.body?.cancel().catch(() => undefined);
+      } else if (response.body) {
+        await pump(id, response.body, controller);
+      }
+      if (!controller.signal.aborted) {
+        post(
+          { channel: VSCODE_BRIDGE_CHANNEL, kind: "fetch:end", id },
+          controller,
+        );
+      }
+    } catch (error) {
+      if (controller.signal.aborted) return;
+      onError(error);
+      post(
+        {
+          channel: VSCODE_BRIDGE_CHANNEL,
+          kind: "fetch:error",
+          id,
+          message: errorMessage(error),
+        },
+        controller,
+      );
+    } finally {
+      if (inflight.get(id) === controller) inflight.delete(id);
+    }
+  };
+
+  const subscription = webview.onDidReceiveMessage((message: unknown) => {
+    if (!isWebviewToHostMessage(message)) return;
+    if (message.kind === "fetch:request") void handle(message);
+    else if (message.kind === "fetch:reset") {
+      const controllers = [...inflight.values()];
+      inflight.clear();
+      for (const controller of controllers) controller.abort();
+    } else if (message.kind === "fetch:abort") {
+      const controller = inflight.get(message.id);
+      if (controller) {
+        inflight.delete(message.id);
+        controller.abort();
+      }
+    }
+  });
+
+  return {
+    dispose: () => {
+      subscription.dispose();
+      for (const [id, controller] of inflight) {
+        post(
+          {
+            channel: VSCODE_BRIDGE_CHANNEL,
+            kind: "fetch:error",
+            id,
+            message: "Webview routes disposed",
+          },
+          controller,
+        );
+        controller.abort();
+      }
+      inflight.clear();
+    },
+  };
+}

package/src/host/serve.ts

@@ -0,0 +1,150 @@
+import { DEFAULT_EXTERNAL_SCHEMES, parseExternalUrl } from "../external-url";
+import {
+  isWebviewToHostMessage,
+  VSCODE_BRIDGE_CHANNEL,
+  type RpcRequestMessage,
+  type RpcResponseMessage,
+} from "../protocol";
+import {
+  serveRoutes,
+  type Disposable,
+  type ServeWebviewRoutesOptions,
+  type WebviewLike,
+  type WebviewRoutes,
+} from "./router";
+
+/** Matches `vscode.Memento`, whose `update(key, undefined)` removes the key. */
+export type MementoLike = {
+  get(key: string): unknown;
+  update(key: string, value: unknown): PromiseLike<void>;
+};
+
+export type ServeWebviewHostOptions = ServeWebviewRoutesOptions & {
+  routes?: WebviewRoutes;
+  /**
+   * Opens a URL forwarded by `installLinkInterceptor`, after the host has
+   * checked its scheme against `externalSchemes`.
+   */
+  openExternal?: (url: string) => PromiseLike<boolean> | boolean | void;
+  /** URL schemes `openExternal` may receive, such as `"https:"`. */
+  externalSchemes?: readonly string[];
+  /** Backs `createVSCodeStorage`, such as `context.globalState`. */
+  storage?: MementoLike;
+  /** Namespaces the webview's keys inside `storage`. */
+  storagePrefix?: string;
+};
+
+type RpcHandler = (...params: unknown[]) => unknown;
+
+const errorMessage = (error: unknown) =>
+  error instanceof Error ? error.message : String(error);
+
+const stringParam = (value: unknown, name: string): string => {
+  if (typeof value !== "string")
+    throw new TypeError(`${name} must be a string`);
+  return value;
+};
+
+/**
+ * Serves `vscodeFetch` routes and the webview's host calls, such as
+ * `installLinkInterceptor` opening links and `createVSCodeStorage`, for one
+ * webview. Host calls without a handler reject in the webview.
+ */
+export function serveWebviewHost(
+  webview: WebviewLike,
+  {
+    routes = {},
+    openExternal,
+    externalSchemes = DEFAULT_EXTERNAL_SCHEMES,
+    storage,
+    storagePrefix = "@assistant-ui/vscode:",
+    ...routeOptions
+  }: ServeWebviewHostOptions = {},
+): Disposable {
+  const handlers: Record<string, RpcHandler> = {};
+
+  if (openExternal) {
+    handlers.openExternal = async (value) => {
+      const url = parseExternalUrl(value, externalSchemes);
+      if (!url) throw new Error(`Refused to open ${String(value)}`);
+      return (await openExternal(url)) !== false;
+    };
+  }
+
+  if (storage) {
+    const storageKey = (key: unknown) =>
+      storagePrefix + stringParam(key, "key");
+    handlers["storage.getItem"] = (key) => {
+      const value = storage.get(storageKey(key));
+      return typeof value === "string" ? value : null;
+    };
+    handlers["storage.setItem"] = async (key, value) => {
+      await storage.update(storageKey(key), stringParam(value, "value"));
+    };
+    handlers["storage.removeItem"] = async (key) => {
+      await storage.update(storageKey(key), undefined);
+    };
+  }
+
+  const respond = (
+    response:
+      | { id: string; ok: true; result: unknown }
+      | { id: string; ok: false; message: string },
+  ) => {
+    const message: RpcResponseMessage = {
+      channel: VSCODE_BRIDGE_CHANNEL,
+      kind: "rpc:response",
+      ...response,
+    };
+    Promise.resolve(webview.postMessage(message)).catch(() => undefined);
+  };
+
+  const pending = new Set<string>();
+
+  const handle = async ({ id, method, params }: RpcRequestMessage) => {
+    if (pending.has(id)) return;
+    pending.add(id);
+    try {
+      const handler = Object.hasOwn(handlers, method) ? handlers[method] : null;
+      if (!handler) throw new Error(`${method} is not served by this host`);
+      const result = await handler(...params);
+      if (pending.delete(id)) respond({ id, ok: true, result });
+    } catch (error) {
+      if (pending.delete(id)) {
+        respond({ id, ok: false, message: errorMessage(error) });
+      }
+    }
+  };
+
+  const server = serveRoutes(webview, routes, routeOptions);
+  const subscription = webview.onDidReceiveMessage((message: unknown) => {
+    if (isWebviewToHostMessage(message) && message.kind === "rpc:request") {
+      void handle(message);
+    }
+  });
+
+  return {
+    dispose: () => {
+      subscription.dispose();
+      for (const id of pending) {
+        respond({ id, ok: false, message: "Webview host disposed" });
+      }
+      pending.clear();
+      server.dispose();
+    },
+  };
+}
+
+/**
+ * Serves `routes` to `createVSCodeFetch` / `vscodeFetch` calls made inside
+ * `webview`, streaming each response back as it is produced. Host calls such
+ * as `installLinkInterceptor` or `createVSCodeStorage` reject in the webview;
+ * use `serveWebviewHost` to serve them.
+ */
+export function serveWebviewRoutes(
+  webview: WebviewLike,
+  routes: WebviewRoutes,
+  options?: ServeWebviewRoutesOptions,
+): Disposable {
+  return serveWebviewHost(webview, { ...options, routes });
+}

package/src/host.ts

@@ -0,0 +1,27 @@
+export {
+  createCspNonce,
+  createWebviewCsp,
+  renderWebviewHtml,
+  type RenderWebviewHtmlOptions,
+  type WebviewCspMode,
+  type WebviewCspOptions,
+  type WebviewCspSource,
+  type WebviewResourceLike,
+  type WebviewSurface,
+} from "./host/html";
+export {
+  type Disposable,
+  type HttpMethod,
+  type RouteHandler,
+  type RouteHandlers,
+  type ServeWebviewRoutesOptions,
+  type WebviewLike,
+  type WebviewRoutes,
+} from "./host/router";
+export {
+  serveWebviewHost,
+  serveWebviewRoutes,
+  type MementoLike,
+  type ServeWebviewHostOptions,
+} from "./host/serve";
+export * from "./protocol";

package/src/protocol.ts

@@ -0,0 +1,131 @@
+export const VSCODE_BRIDGE_CHANNEL = "aui-vscode";
+
+export const VSCODE_VIRTUAL_ORIGIN = "https://webview.vscode.invalid";
+
+export type BridgeHeaders = [name: string, value: string][];
+
+type BridgeEnvelope<TKind extends string> = {
+  channel: typeof VSCODE_BRIDGE_CHANNEL;
+  kind: TKind;
+  id: string;
+};
+
+export type FetchRequestMessage = BridgeEnvelope<"fetch:request"> & {
+  url: string;
+  method: string;
+  headers: BridgeHeaders;
+  body: Uint8Array<ArrayBuffer> | null;
+};
+
+export type FetchAbortMessage = BridgeEnvelope<"fetch:abort">;
+
+export type FetchResetMessage = BridgeEnvelope<"fetch:reset">;
+
+export type FetchHeadMessage = BridgeEnvelope<"fetch:head"> & {
+  status: number;
+  statusText: string;
+  headers: BridgeHeaders;
+};
+
+export type FetchChunkMessage = BridgeEnvelope<"fetch:chunk"> & {
+  chunk: Uint8Array<ArrayBuffer>;
+};
+
+export type FetchEndMessage = BridgeEnvelope<"fetch:end">;
+
+export type FetchErrorMessage = BridgeEnvelope<"fetch:error"> & {
+  message: string;
+};
+
+export type RpcRequestMessage = BridgeEnvelope<"rpc:request"> & {
+  method: string;
+  params: unknown[];
+};
+
+export type RpcResponseMessage = BridgeEnvelope<"rpc:response"> &
+  ({ ok: true; result: unknown } | { ok: false; message: string });
+
+export type WebviewToHostMessage =
+  | FetchRequestMessage
+  | FetchAbortMessage
+  | FetchResetMessage
+  | RpcRequestMessage;
+
+export type HostToWebviewMessage =
+  | FetchHeadMessage
+  | FetchChunkMessage
+  | FetchEndMessage
+  | FetchErrorMessage
+  | RpcResponseMessage;
+
+export type BridgeMessage = WebviewToHostMessage | HostToWebviewMessage;
+
+type Fields = Record<string, unknown>;
+
+const isString = (value: unknown): value is string => typeof value === "string";
+
+const isBytes = (value: unknown) => value instanceof Uint8Array;
+
+const isHeaders = (value: unknown) =>
+  Array.isArray(value) &&
+  value.every(
+    (entry) =>
+      Array.isArray(entry) &&
+      entry.length === 2 &&
+      isString(entry[0]) &&
+      isString(entry[1]),
+  );
+
+const hasNoPayload = () => true;
+
+const WEBVIEW_TO_HOST: Record<
+  WebviewToHostMessage["kind"],
+  (message: Fields) => boolean
+> = {
+  "fetch:request": (m) =>
+    isString(m.url) &&
+    isString(m.method) &&
+    isHeaders(m.headers) &&
+    (m.body === null || isBytes(m.body)),
+  "fetch:abort": hasNoPayload,
+  "fetch:reset": hasNoPayload,
+  "rpc:request": (m) => isString(m.method) && Array.isArray(m.params),
+};
+
+const HOST_TO_WEBVIEW: Record<
+  HostToWebviewMessage["kind"],
+  (message: Fields) => boolean
+> = {
+  "fetch:head": (m) =>
+    typeof m.status === "number" &&
+    isString(m.statusText) &&
+    isHeaders(m.headers),
+  "fetch:chunk": (m) => isBytes(m.chunk),
+  "fetch:end": hasNoPayload,
+  "fetch:error": (m) => isString(m.message),
+  "rpc:response": (m) =>
+    m.ok === true || (m.ok === false && isString(m.message)),
+};
+
+const matches = (
+  validators: Record<string, (message: Fields) => boolean>,
+  value: unknown,
+) => {
+  if (typeof value !== "object" || value === null) return false;
+  const message = value as Fields;
+  return (
+    message.channel === VSCODE_BRIDGE_CHANNEL &&
+    isString(message.id) &&
+    isString(message.kind) &&
+    Object.hasOwn(validators, message.kind) &&
+    validators[message.kind]!(message)
+  );
+};
+
+export const isWebviewToHostMessage = (
+  value: unknown,
+): value is WebviewToHostMessage => matches(WEBVIEW_TO_HOST, value);
+
+export const isHostToWebviewMessage = (
+  value: unknown,
+): value is HostToWebviewMessage => matches(HOST_TO_WEBVIEW, value);

package/src/route-url.ts

@@ -0,0 +1,6 @@
+import { VSCODE_VIRTUAL_ORIGIN } from "./protocol";
+
+/** Routes are served on the virtual origin, so an `http:` or `https:` URL on any other origin reaches none of them. */
+export const isForeignHttpUrl = (url: URL) =>
+  (url.protocol === "http:" || url.protocol === "https:") &&
+  url.origin !== VSCODE_VIRTUAL_ORIGIN;