open-standards/wai/wai-app/web/model.js

// Where a program comes from, each an async iterator of text chunks:
//
//   ownModel   the visitor's own model: an endpoint that streams the common
//              `chat/completions` server-sent-events format (a model server on their own
//              machine, or a provider they hold a key for), called from this page with
//              their key. The key is held by the caller in memory and sent to that
//              endpoint alone; nothing here stores it.
//   bridge     the visitor's coding agent, through the bridge `wai mcp --toolset app` runs on 127.0.0.1:
//              each program the agent streams, read with the pairing code.
//   replay     a recorded stream: chunks and the pauses between them.
//
// No source calls a model of this site's own: there is none.

// Lines of a server-sent-events body, as they arrive.
async function* sseLines(body) {
  const dec = new TextDecoder();
  let buf = "";
  const reader = body.getReader();
  for (;;) {
    const { value, done } = await reader.read();
    if (done) break;
    buf += dec.decode(value, { stream: true });
    let nl;
    while ((nl = buf.indexOf("\n")) >= 0) {
      const line = buf.slice(0, nl).replace(/\r$/, "");
      buf = buf.slice(nl + 1);
      yield line;
    }
  }
  if (buf) yield buf;
}

// The text of each `data:` event of a `chat/completions` stream: `choices[0].delta.content`.
export async function* completionChunks(lines) {
  for await (const line of lines) {
    if (!line.startsWith("data:")) continue;
    const data = line.slice(5).trim();
    if (!data || data === "[DONE]") continue;
    let j;
    try { j = JSON.parse(data); } catch { continue; }
    if (j.error) throw new Error(`the model answered: ${j.error.message || JSON.stringify(j.error)}`);
    const c = j.choices && j.choices[0] && j.choices[0].delta && j.choices[0].delta.content;
    if (typeof c === "string" && c) yield c;
  }
}

// The visitor's own model. `endpoint` is the full URL of its `chat/completions` route.
export async function* ownModel({ endpoint, model, key, system, prompt, signal }) {
  const u = new URL(endpoint);
  if (u.protocol !== "https:" && !(u.protocol === "http:" && ["127.0.0.1", "localhost", "[::1]"].includes(u.hostname))) {
    throw new Error("the endpoint must be https, or http on this machine (127.0.0.1)");
  }
  const headers = { "content-type": "application/json" };
  if (key) headers.authorization = `Bearer ${key}`;
  const r = await fetch(u.toString(), {
    method: "POST",
    headers,
    signal,
    credentials: "omit",
    referrerPolicy: "no-referrer",
    cache: "no-store",
    body: JSON.stringify({ model, stream: true, messages: [{ role: "system", content: system }, { role: "user", content: prompt }] }),
  });
  if (!r.ok || !r.body) throw new Error(`the model's endpoint answered ${r.status}`);
  yield* completionChunks(sseLines(r.body));
}

// The coding agent's bridge: `onProgram(iterator)` is called for each program the agent
// streams. Returns a function that disconnects.
export function bridge({ url, code, onProgram, onState }) {
  const u = new URL(url);
  if (!["127.0.0.1", "localhost", "[::1]"].includes(u.hostname)) throw new Error("the bridge is on this machine: http://127.0.0.1:<port>");
  const es = new EventSource(`${u.origin}/create/stream?code=${encodeURIComponent(code)}`);
  let queue = null;
  es.addEventListener("hello", () => onState && onState("connected"));
  es.addEventListener("start", () => {
    if (queue) queue.end();
    queue = channel();
    onProgram(queue.iter);
  });
  es.addEventListener("chunk", (e) => { if (queue) queue.put(JSON.parse(e.data).text); });
  es.addEventListener("end", () => { if (queue) { queue.end(); queue = null; } });
  es.onerror = () => onState && onState(es.readyState === EventSource.CLOSED ? "closed" : "retrying");
  return () => { es.close(); if (queue) queue.end(); };
}

function channel() {
  const items = [];
  let wake = null;
  let done = false;
  return {
    put(x) { items.push(x); if (wake) { wake(); wake = null; } },
    end() { done = true; if (wake) { wake(); wake = null; } },
    iter: (async function* () {
      for (;;) {
        if (items.length) { yield items.shift(); continue; }
        if (done) return;
        await new Promise((r) => { wake = r; });
      }
    })(),
  };
}

// A recorded stream: `{ chunks: [...], pause_ms }` (or `pauses: [...]`), replayed with its
// pauses, `speed` times faster.
export async function* replay(rec, { speed = 1, signal } = {}) {
  for (let i = 0; i < rec.chunks.length; i++) {
    if (signal && signal.aborted) return;
    const ms = rec.pauses ? rec.pauses[i] || 0 : rec.pause_ms || 0;
    if (ms > 0) await new Promise((r) => setTimeout(r, ms / speed));
    yield rec.chunks[i];
  }
}