@ar-ai/brink-cli 0.0.0-stage → 1.0.4

package/support/messages-bridge.js

@@ -0,0 +1,317 @@
+"use strict";
+
+const { randomUUID } = require("crypto");
+const { buildMessagesRequest } = require("./messages-request");
+const { cleanHeaders, requestWithRetries, upstreamFailure } = require("./transport");
+const { failure, writeChunk, sseEvents, decodedStream } = require("./bridge-stream");
+const { observeActivity } = require("./activity");
+
+const UPSTREAM = new URL("https://agentrouter.org/v1/responses");
+
+function messageUsage(usage) {
+  if (!usage) return {};
+  const result = {};
+  const cached = usage.input_tokens_details?.cached_tokens;
+  if (Number.isFinite(usage.input_tokens))
+    result.input_tokens = Math.max(0, usage.input_tokens - (cached || 0));
+  if (Number.isFinite(cached)) result.cache_read_input_tokens = cached;
+  if (Number.isFinite(usage.output_tokens)) result.output_tokens = usage.output_tokens;
+  return result;
+}
+
+function errorBody(error) {
+  return { type: "error", error: {
+    type: error.status === 400 ? "invalid_request_error" :
+      error.status === 401 ? "authentication_error" :
+      error.status === 403 ? "permission_error" :
+      error.status === 429 ? "rate_limit_error" : "api_error",
+    message: error.message,
+  } };
+}
+
+class MessagesWriter {
+  constructor(res, payload, tracker) {
+    this.res = res;
+    this.streaming = payload.stream === true;
+    this.showThinking = payload.thinking?.display !== "omitted";
+    this.tracker = tracker;
+    this.items = new Map();
+    this.blocks = new Map();
+    this.terminal = null;
+    this.refused = false;
+    this.toolNames = new Set((payload.tools || []).map(tool => tool.name));
+    this.message = { id: `msg_${randomUUID().replace(/-/g, "")}`, type: "message",
+      role: "assistant", model: payload.model, content: [], stop_reason: null,
+      stop_sequence: null, usage: {} };
+  }
+
+  async event(type, fields = {}) {
+    if (this.streaming)
+      await writeChunk(this.res, `event: ${type}\ndata: ${JSON.stringify({ type, ...fields })}\n\n`);
+  }
+
+  async start() {
+    await this.event("message_start", { message: { ...this.message,
+      usage: { input_tokens: 0, output_tokens: 0 } } });
+  }
+
+  async block(key, content) {
+    let state = this.blocks.get(key);
+    if (state) return state;
+    state = { index: this.message.content.length, content, text: "", closed: false };
+    this.blocks.set(key, state);
+    this.message.content.push(content);
+    await this.event("content_block_start", { index: state.index, content_block: { ...content } });
+    return state;
+  }
+
+  async append(state, text, final = false) {
+    if (typeof text !== "string") throw failure("AgentRouter returned non-text content");
+    if (final) {
+      if (!text.startsWith(state.text)) throw failure("AgentRouter changed already-streamed content");
+      text = text.slice(state.text.length);
+    }
+    if (!text) return;
+    if (state.closed) throw failure("AgentRouter changed a completed content block");
+    state.text += text;
+    const kind = state.content.type;
+    if (kind === "text") state.content.text = state.text;
+    else if (kind === "thinking") state.content.thinking = state.text;
+    await this.event("content_block_delta", { index: state.index, delta:
+      kind === "tool_use" ? { type: "input_json_delta", partial_json: text } :
+      kind === "thinking" ? { type: "thinking_delta", thinking: text } :
+      { type: "text_delta", text } });
+  }
+
+  async close(state) {
+    if (state.closed) return;
+    if (state.content.type === "tool_use") {
+      let input;
+      try { input = JSON.parse(state.text); } catch {
+        throw failure("AgentRouter returned incomplete or invalid tool arguments", 502, "incomplete_stream");
+      }
+      if (!input || typeof input !== "object" || Array.isArray(input))
+        throw failure("AgentRouter tool arguments must be a JSON object");
+      state.content.input = input;
+    }
+    state.closed = true;
+    await this.event("content_block_stop", { index: state.index });
+  }
+
+  async item(index, item, final = false) {
+    if (!item || !Number.isInteger(index)) throw failure("AgentRouter returned an invalid output item");
+    const previous = this.items.get(index);
+    if (previous && (previous.type !== item.type || (previous.id && item.id && previous.id !== item.id)))
+      throw failure("AgentRouter changed an output item");
+    this.items.set(index, item);
+    if (item.type === "function_call") {
+      if (!item.call_id || !this.toolNames.has(item.name))
+        throw failure("AgentRouter returned an unknown tool or missing call ID");
+      const state = await this.block(`${index}:tool`, {
+        type: "tool_use", id: item.call_id, name: item.name, input: {},
+      });
+      if (state.content.id !== item.call_id || state.content.name !== item.name)
+        throw failure("AgentRouter changed a tool call identity");
+      if (final || item.arguments) await this.append(state, item.arguments || "", true);
+      if (final) await this.close(state);
+    } else if (item.type === "message" || item.type === "reasoning") {
+      if (item.type === "message" && item.role !== "assistant")
+        throw failure("AgentRouter returned a non-assistant output message");
+      const parts = item.type === "reasoning" ? item.summary : item.content;
+      for (const [partIndex, part] of (parts || []).entries())
+        await this.part(index, partIndex, part, final);
+    } else {
+      throw failure(`Unsupported Responses output type: ${item.type}`);
+    }
+  }
+
+  async part(index, partIndex, part, final = false) {
+    const thinking = part?.type === "summary_text";
+    if (thinking && !this.showThinking) return;
+    if (!thinking && !["output_text", "refusal"].includes(part?.type))
+      throw failure(`Unsupported Responses content type: ${part?.type}`);
+    const refused = part.type === "refusal";
+    this.refused ||= refused;
+    const state = await this.block(`${index}:${thinking ? "summary" : "text"}:${partIndex}`,
+      thinking ? { type: "thinking", thinking: "", signature: "" } : { type: "text", text: "" });
+    await this.append(state, (refused ? part.refusal : part.text) || "", true);
+    if (final) await this.close(state);
+  }
+
+  async consume(value) {
+    if (!value || typeof value !== "object") throw failure("Invalid Responses payload");
+    if (this.terminal) throw failure("AgentRouter emitted data after its terminal response");
+    const type = value.type;
+    if (type === "error" || type === "response.failed" || value.error) {
+      const error = value.error || value.response?.error || value;
+      throw failure(typeof error === "string" ? error : error.message || "AgentRouter response failed",
+        502, error.code || "upstream_error");
+    }
+    observeActivity(value, this.tracker, { summaries: this.showThinking });
+    if (type === "response.output_item.added" || type === "response.output_item.done") {
+      await this.item(value.output_index, value.item, type.endsWith(".done"));
+    } else if (type === "response.content_part.added" || type === "response.content_part.done") {
+      await this.part(value.output_index, value.content_index, value.part, type.endsWith(".done"));
+    } else if (type === "response.reasoning_summary_part.added" || type === "response.reasoning_summary_part.done") {
+      await this.part(value.output_index, value.summary_index, value.part, type.endsWith(".done"));
+    } else if (["response.output_text.delta", "response.refusal.delta", "response.reasoning_summary_text.delta"].includes(type)) {
+      const thinking = type === "response.reasoning_summary_text.delta";
+      if (thinking && !this.showThinking) return;
+      this.refused ||= type === "response.refusal.delta";
+      const key = `${value.output_index}:${thinking ? "summary" : "text"}:${thinking ? value.summary_index : value.content_index}`;
+      const state = await this.block(key, thinking ? { type: "thinking", thinking: "", signature: "" } : { type: "text", text: "" });
+      await this.append(state, value.delta);
+    } else if (type === "response.function_call_arguments.delta" || type === "response.function_call_arguments.done") {
+      const state = this.blocks.get(`${value.output_index}:tool`);
+      if (!state) throw failure("AgentRouter emitted arguments before the tool identity");
+      await this.append(state, value.delta ?? value.arguments, type.endsWith(".done"));
+    } else if (["response.output_text.done", "response.refusal.done", "response.reasoning_summary_text.done"].includes(type)) {
+      const thinking = type === "response.reasoning_summary_text.done";
+      await this.part(value.output_index, thinking ? value.summary_index : value.content_index,
+        thinking ? { type: "summary_text", text: value.text } :
+        type === "response.refusal.done" ? { type: "refusal", refusal: value.refusal } :
+        { type: "output_text", text: value.text });
+    } else if (type === "response.completed" || type === "response.incomplete" || value.object === "response") {
+      const response = value.response || value;
+      if (!["completed", "incomplete"].includes(response.status))
+        throw failure(response.error?.message || "AgentRouter response did not complete", 502, "incomplete_stream");
+      if (!Array.isArray(response.output)) throw failure("AgentRouter response is missing output");
+      for (const [index, item] of response.output.entries()) await this.item(index, item, true);
+      const incomplete = response.status === "incomplete";
+      const reason = response.incomplete_details?.reason;
+      if (incomplete && !["max_output_tokens", "content_filter"].includes(reason))
+        throw failure("AgentRouter response is incomplete", 502, "incomplete_stream");
+      this.message.stop_reason = incomplete ? reason === "max_output_tokens" ? "max_tokens" : "refusal" :
+        this.message.content.some(part => part.type === "tool_use") ? "tool_use" :
+        this.refused ? "refusal" : "end_turn";
+      this.message.usage = messageUsage(response.usage);
+      if (response.usage) this.tracker?.usage(response.usage);
+      if (incomplete) this.tracker?.fail("AgentRouter reported an incomplete response (generation limit or content restriction)");
+      this.terminal = response.status;
+    }
+  }
+
+  async finish() {
+    if (!this.terminal) throw failure("AgentRouter stream ended before its terminal response", 502, "incomplete_stream");
+    for (const state of this.blocks.values()) await this.close(state);
+    await this.event("message_delta", { delta: { stop_reason: this.message.stop_reason, stop_sequence: null }, usage: this.message.usage });
+    if (this.terminal === "completed") this.tracker?.complete();
+    await this.event("message_stop");
+    return this.message;
+  }
+}
+
+async function readJSON(stream) {
+  const chunks = [];
+  let bytes = 0;
+  for await (const chunk of stream) {
+    bytes += chunk.length;
+    if (bytes > 32 * 1024 * 1024) throw failure("AgentRouter JSON response exceeds 32 MiB");
+    chunks.push(chunk);
+  }
+  try { return JSON.parse(Buffer.concat(chunks).toString("utf8")); }
+  catch { throw failure("AgentRouter returned invalid Responses JSON"); }
+}
+
+async function bridgeMessages(req, res, payload, { agent, tracker } = {}) {
+  const controller = new AbortController();
+  let result;
+  let stream;
+  let writer;
+  let passthrough = false;
+  let cancelled = false;
+  const destroyUpstream = () => {
+    controller.abort();
+    stream?.destroy();
+    result?.stream.destroy();
+    result?.response.destroy();
+    result?.request.destroy();
+  };
+  const cancel = () => {
+    if (res.writableFinished) return;
+    cancelled = true;
+    destroyUpstream();
+  };
+  req.once("aborted", cancel);
+  res.once("close", cancel);
+  res.once("error", cancel);
+  try {
+    const request = buildMessagesRequest(payload);
+    if (payload.output_config?.effort === "max")
+      tracker?.note("info", "Mapped Messages max effort to Responses xhigh");
+    const body = JSON.stringify(request);
+    const headers = cleanHeaders(req.headers);
+    if (!headers.authorization && typeof headers["x-api-key"] === "string")
+      headers.authorization = `Bearer ${headers["x-api-key"]}`;
+    for (const key of ["host", "content-length", "content-encoding", "expect", "accept-encoding", "x-api-key", "anthropic-version", "anthropic-beta"])
+      delete headers[key];
+    headers.host = UPSTREAM.host;
+    headers["content-type"] = "application/json";
+    headers["content-length"] = Buffer.byteLength(body);
+    headers["accept-encoding"] = "identity";
+    headers.accept = request.stream ? "text/event-stream" : "application/json";
+    if (cancelled || req.aborted || res.destroyed) return;
+    result = await requestWithRetries(UPSTREAM, { method: "POST", headers, agent }, body, {
+      signal: controller.signal, tracker,
+      clientOwnsRetries: req.headers["x-agentrouter-retry-owner"] === "client",
+    });
+    const status = result.response.statusCode || 502;
+    tracker?.upstream(status);
+    if (status < 200 || status >= 300 || result.retryReason) {
+      tracker?.fail(upstreamFailure(result.retryReason ? null : status, { code: result.retryReason }));
+      passthrough = true;
+      res.writeHead(status, cleanHeaders(result.response.headers));
+      for await (const chunk of result.stream) await writeChunk(res, chunk);
+      res.end();
+      return;
+    }
+    stream = decodedStream(result.response, result.stream);
+    writer = new MessagesWriter(res, payload, tracker);
+    const responseHeaders = cleanHeaders(result.response.headers);
+    for (const key of ["content-length", "content-encoding", "etag"])
+      delete responseHeaders[key];
+    responseHeaders["cache-control"] = "no-cache";
+    responseHeaders["x-accel-buffering"] = "no";
+    responseHeaders["content-type"] = writer.streaming ? "text/event-stream; charset=utf-8" : "application/json";
+    if (writer.streaming) {
+      res.writeHead(200, responseHeaders);
+      res.flushHeaders?.();
+      await writer.start();
+    }
+    if (String(result.response.headers["content-type"] || "").includes("text/event-stream")) {
+      for await (const event of sseEvents(stream)) {
+        if (event.data.trim() === "[DONE]") continue;
+        let value;
+        try { value = JSON.parse(event.data); }
+        catch { throw failure("AgentRouter returned malformed Responses SSE JSON"); }
+        if (!value.type && event.event) value.type = event.event;
+        await writer.consume(value);
+      }
+    } else await writer.consume(await readJSON(stream));
+    const message = await writer.finish();
+    if (writer.streaming) res.end();
+    else {
+      res.writeHead(200, responseHeaders);
+      res.end(JSON.stringify(message));
+    }
+  } catch (error) {
+    if (!cancelled && !res.destroyed && !res.writableEnded)
+      tracker?.fail(upstreamFailure(null, error));
+    destroyUpstream();
+    if (cancelled || res.destroyed || res.writableEnded) return;
+    if (passthrough) { res.destroy(); return; }
+    if (res.headersSent) {
+      try { await writer.event("error", { error: errorBody(error).error }); res.end(); }
+      catch { res.destroy(); }
+    } else {
+      res.writeHead(error.status || 502, { "content-type": "application/json" });
+      res.end(JSON.stringify(errorBody(error)));
+    }
+  } finally {
+    req.off("aborted", cancel);
+    res.off("close", cancel);
+    res.off("error", cancel);
+  }
+}
+
+module.exports = { bridgeMessages };