nanocodex 0.1.1 → 0.3.0

package/browser/host.mjs

@@ -3,25 +3,30 @@ import { createCodeRuntime } from "../runtime/code-runtime.mjs";
 const DEFAULT_MAX_QUEUED_MESSAGES = 4_096;
 const DEFAULT_MAX_QUEUED_BYTES = 32 * 1024 * 1024;
 const DEFAULT_MAX_BUFFERED_SEND_BYTES = 16 * 1024 * 1024;
+const WEBSOCKET_OPEN = 1;
 
 export function createBrowserHost(options = {}) {
-  const WebSocketImpl = options.WebSocketImpl || globalThis.WebSocket;
-  if (!WebSocketImpl) throw new Error("WebSocket is unavailable in this runtime");
+  const WebSocketImpl = options.WebSocketImpl ?? globalThis.WebSocket;
+  const createWebSocket = options.createWebSocket
+    ?? (WebSocketImpl && ((endpoint) => new WebSocketImpl(endpoint)));
+  if (!options.mpp && !createWebSocket) {
+    throw new Error("WebSocket is unavailable in this runtime");
+  }
   const connections = new Map();
   const code = createCodeRuntime(options.tools);
   const onEvent = options.onEvent || (() => {});
-  const createWebSocket = options.createWebSocket || ((endpoint) => new WebSocketImpl(endpoint));
   const maxQueuedMessages = options.maxQueuedMessages ?? DEFAULT_MAX_QUEUED_MESSAGES;
   const maxQueuedBytes = options.maxQueuedBytes ?? DEFAULT_MAX_QUEUED_BYTES;
   const maxBufferedSendBytes = options.maxBufferedSendBytes ?? DEFAULT_MAX_BUFFERED_SEND_BYTES;
   const encoder = new TextEncoder();
   let nextHandle = 1;
 
-  function connect(endpoint, _apiKey, sessionId) {
+  function connect(endpoint, _apiKey, sessionId, metadata = {}) {
+    if (options.mpp) return connectMpp(endpoint);
     return new Promise((resolve, reject) => {
       let settled = false;
       const handle = nextHandle++;
-      const socket = createWebSocket(endpoint, sessionId);
+      const socket = createWebSocket(endpoint, sessionId, metadata);
       const connection = {
         socket,
         queue: [],
@@ -59,9 +64,44 @@ export function createBrowserHost(options = {}) {
     });
   }
 
+  async function connectMpp(endpoint) {
+    if (typeof options.mpp.ws !== "function") {
+      throw new TypeError("mpp must provide ws(endpoint)");
+    }
+    const socket = await options.mpp.ws(endpoint);
+    if (!socket || typeof socket.addEventListener !== "function") {
+      throw new TypeError("mpp.ws(endpoint) must return a WebSocket");
+    }
+    const handle = nextHandle++;
+    const connection = {
+      socket,
+      queue: [],
+      queuedBytes: 0,
+      waiter: undefined,
+      intentionallyClosed: false,
+      overflowed: false,
+      managed: true,
+    };
+    connections.set(handle, connection);
+    socket.addEventListener("message", (event) => {
+      enqueue(connection, typeof event.data === "string"
+        ? { kind: "text", text: event.data }
+        : { kind: "binary" });
+    });
+    socket.addEventListener("close", (event) => {
+      if (!connection.intentionallyClosed && !connection.overflowed) {
+        enqueue(connection, { kind: "closed", detail: `with code ${event.code ?? 1000}` });
+      }
+    });
+    socket.addEventListener("error", () => {
+      enqueue(connection, { kind: "error", detail: "MPP WebSocket connection failed" });
+    });
+    return JSON.stringify({ handle, status: 101, reasoning_included: false });
+  }
+
   function send(handle, message) {
     const connection = connections.get(handle);
-    if (!connection || connection.socket.readyState !== WebSocketImpl.OPEN) {
+    if (!connection || connection.socket.readyState !== WEBSOCKET_OPEN) {
       return Promise.resolve(JSON.stringify({
         ok: false,
         reconnectable: true,
@@ -69,6 +109,10 @@ export function createBrowserHost(options = {}) {
       }));
     }
     try {
+      if (connection.managed) {
+        connection.socket.send(JSON.stringify({ mpp: "message", data: message }));
+        return Promise.resolve(JSON.stringify({ ok: true }));
+      }
       const frameBytes = encoder.encode(message).byteLength;
       if (frameBytes > maxBufferedSendBytes
         || connection.socket.bufferedAmount + frameBytes > maxBufferedSendBytes) {
@@ -83,7 +127,7 @@ export function createBrowserHost(options = {}) {
     } catch (error) {
       return Promise.resolve(JSON.stringify({
         ok: false,
-        reconnectable: connection.socket.readyState !== WebSocketImpl.OPEN,
+        reconnectable: connection.socket.readyState !== WEBSOCKET_OPEN,
         error: error instanceof Error ? error.message : String(error),
       }));
     }

package/browser/index.d.mts

@@ -1,13 +1,18 @@
 export { Actions } from "../index.mjs";
 export type {
   AgentEvent,
+  CostStatus,
+  EstimatedUsdCost,
   PromptInput,
   PromptItem,
   ReasoningMode,
+  SessionSnapshot,
   Thinking,
   Tool,
   ToolContext,
   ToolMap,
   Turn,
+  TurnResult,
+  TurnUsage,
 } from "../types.mjs";
 export * as Agent from "./Agent.mjs";

package/index.d.mts

@@ -4,16 +4,21 @@ export type {
   AgentActions,
   AgentEvent,
   AgentOptions,
+  CostStatus,
   DefaultAgent,
   EventWatcher,
+  EstimatedUsdCost,
   ForkOptions,
   PromptInput,
   PromptItem,
   ReasoningMode,
+  SessionSnapshot,
   Thinking,
   Tool,
   ToolContext,
   ToolMap,
   Turn,
+  TurnResult,
+  TurnUsage,
   WatchEventsOptions,
 } from "./types.mjs";

package/internal.mjs

@@ -1,5 +1,6 @@
-const AGENT_STATE = Symbol("nanocodex.agent");
-const TURN_STATE = Symbol("nanocodex.turn");
+const agentStates = new WeakMap();
+const turnStates = new WeakMap();
+const resultStates = new WeakMap();
 const hostSessions = new Map();
 const hostConnections = new Map();
 let activeHost;
@@ -41,10 +42,20 @@ export function prompt(agent, options) {
 
 export function getTurnResult(turn) {
   const state = turnState(turn);
-  state.result ||= Promise.resolve().then(() => state.raw.result());
+  state.result ||= Promise.resolve()
+    .then(() => state.raw.result())
+    .then(createTurnResult);
   return state.result;
 }
 
+export function getTurnSnapshot(result) {
+  return resultState(result).snapshot();
+}
+
+export function getTurnUsage(result) {
+  return resultState(result).usage();
+}
+
 export function steer(turn, options) {
   const state = turnState(turn);
   const input = actionInput(options);
@@ -62,7 +73,7 @@ export async function fork(agent, options) {
   const at = options?.at;
   const raw = at === undefined
     ? await state.raw.fork()
-    : await state.raw.forkFrom(turnState(at).raw);
+    : await state.raw.forkFrom(resultState(at).raw);
   return createAgent(raw, state.runtime);
 }
 
@@ -71,7 +82,31 @@ export async function spawn(agent) {
   return createAgent(await state.raw.spawn(), state.runtime);
 }
 
-export function subscribeAgentEvents(agent, listener, options = {}) {
+export function setThinking(agent, thinking) {
+  return agentState(agent).raw.setThinking(thinking);
+}
+
+export function setFastMode(agent, enabled) {
+  return agentState(agent).raw.setFastMode(enabled);
+}
+
+export function compact(agent) {
+  return agentState(agent).raw.compact();
+}
+
+export async function shutdown(agent) {
+  const state = knownAgentState(agent);
+  if (state.shutdownPromise) return state.shutdownPromise;
+  if (state.disposed) throw new Error("the Nanocodex agent has been disposed");
+  if (typeof state.raw.shutdown !== "function") {
+    throw new Error("this Nanocodex runtime does not expose graceful shutdown");
+  }
+  state.disposed = true;
+  state.shutdownPromise = joinAgentShutdown(state);
+  return state.shutdownPromise;
+}
+
+export function subscribeAgentEvents(agent, listener, options = {}, onRelease) {
   const state = agentState(agent);
   if (typeof state.runtime.subscribe !== "function") {
     throw new Error("this Nanocodex runtime does not expose agent events");
@@ -79,11 +114,25 @@ export function subscribeAgentEvents(agent, listener, options = {}) {
   if (typeof listener !== "function") {
     throw new TypeError("watchAgentEvents requires a listener");
   }
-  return state.runtime.subscribe((event) => {
+  const unsubscribe = state.runtime.subscribe((event, encodedLength) => {
     if (options.includeAllSessions || !event?.request_id || event.request_id === agent.sessionId) {
-      listener(event);
+      listener(event, encodedLength);
     }
   });
+  let active = true;
+  const subscription = {
+    close(notify) {
+      if (!active) return;
+      active = false;
+      state.subscriptions.delete(subscription);
+      const errors = [];
+      runCleanup(errors, () => unsubscribe?.());
+      if (notify) runCleanup(errors, () => onRelease?.());
+      throwCleanupErrors(errors);
+    },
+  };
+  state.subscriptions.add(subscription);
+  return () => subscription.close(false);
 }
 
 export function toWasmConfig(options = {}) {
@@ -94,10 +143,13 @@ export function toWasmConfig(options = {}) {
   const config = { api_key: apiKey };
   copy(config, "thinking", options.thinking);
   copy(config, "reasoning_mode", options.reasoningMode);
+  copy(config, "fast_mode", options.fastMode);
   copy(config, "websocket_url", options.websocketUrl);
   copy(config, "api_base_url", options.apiBaseUrl);
   copy(config, "instructions", options.instructions);
   copy(config, "session_id", options.sessionId);
+  copy(config, "workspace", options.workspace);
+  copy(config, "resume", options.resume);
   return config;
 }
 
@@ -105,8 +157,10 @@ export function createEventChannel() {
   const listeners = new Set();
   return Object.freeze({
     emit(eventJson) {
+      if (!listeners.size) return;
       const event = typeof eventJson === "string" ? JSON.parse(eventJson) : eventJson;
-      for (const listener of listeners) listener(event);
+      const encodedLength = typeof eventJson === "string" ? eventJson.length : undefined;
+      for (const listener of listeners) listener(event, encodedLength);
     },
     subscribe(listener) {
       listeners.add(listener);
@@ -136,9 +190,18 @@ export function releaseHostSession(host, sessionId) {
 }
 
 const hostBridge = Object.freeze({
-  async connect(endpoint, apiKey, sessionId) {
+  async connect(endpoint, apiKey, accountId, fedramp, sessionId, turnState) {
     const host = requiredSessionHost(sessionId);
-    const result = JSON.parse(await host.connect(endpoint, apiKey, sessionId));
+    let result;
+    try {
+      result = JSON.parse(await host.connect(endpoint, apiKey, sessionId, {
+        accountId: accountId ?? undefined,
+        fedramp,
+        turnState: turnState ?? undefined,
+      }));
+    } catch (error) {
+      throw JSON.stringify(connectFailure(error));
+    }
     const handle = nextHostConnection++;
     hostConnections.set(handle, { host, handle: result.handle });
     hostSessions.set(sessionId, host);
@@ -162,8 +225,12 @@ const hostBridge = Object.freeze({
     hostConnections.delete(handle);
     connection.host.close(connection.handle);
   },
-  sleep(milliseconds) {
-    return new Promise((resolve) => setTimeout(resolve, milliseconds));
+  sleep(sessionId, milliseconds) {
+    const host = requiredSessionHost(sessionId);
+    if (typeof host.sleep !== "function") {
+      throw new TypeError("the selected Nanocodex host must define sleep(milliseconds)");
+    }
+    return host.sleep(milliseconds);
   },
   executeCode(source, sessionId, callId) {
     return requiredSessionHost(sessionId).executeCode(source, sessionId, callId);
@@ -188,6 +255,9 @@ function createAgent(raw, runtime) {
     raw,
     runtime,
     disposed: false,
+    released: false,
+    shutdownPromise: undefined,
+    subscriptions: new Set(),
     sessionId: raw.sessionId,
     uid: `agent-${nextAgentUid++}`,
   };
@@ -220,14 +290,12 @@ function agentView(state, extensions) {
       return agentView(state, deepMerge(extensions, extension));
     },
     dispose() {
-      if (state.disposed) return;
-      state.disposed = true;
-      state.runtime.release?.(state.raw);
-      state.runtime.dispose(state.raw);
+      if (state.shutdownPromise) return;
+      releaseAgentState(state);
     },
   };
   agent = Object.assign(base, extensions);
-  Object.defineProperty(agent, AGENT_STATE, { value: state });
+  agentStates.set(agent, state);
   return agent;
 }
 
@@ -237,11 +305,87 @@ function requiredSessionHost(sessionId) {
   return host;
 }
 
+function releaseAgentState(state) {
+  if (state.released) return;
+  state.disposed = true;
+  state.released = true;
+  const errors = [];
+  for (const subscription of [...state.subscriptions]) {
+    runCleanup(errors, () => subscription.close(true));
+  }
+  runCleanup(errors, () => state.runtime.release?.(state.raw));
+  runCleanup(errors, () => state.runtime.dispose(state.raw));
+  throwCleanupErrors(errors);
+}
+
+async function joinAgentShutdown(state) {
+  await Promise.resolve();
+  let shutdownFailed = false;
+  let shutdownError;
+  try {
+    await state.raw.shutdown();
+  } catch (error) {
+    shutdownFailed = true;
+    shutdownError = error;
+  }
+
+  let cleanupFailed = false;
+  let cleanupError;
+  try {
+    releaseAgentState(state);
+  } catch (error) {
+    cleanupFailed = true;
+    cleanupError = error;
+  }
+
+  if (shutdownFailed && cleanupFailed) {
+    const cleanupErrors = cleanupError instanceof AggregateError
+      ? cleanupError.errors
+      : [cleanupError];
+    throw new AggregateError(
+      [shutdownError, ...cleanupErrors],
+      "Nanocodex driver shutdown and resource release both failed",
+    );
+  }
+  if (shutdownFailed) throw shutdownError;
+  if (cleanupFailed) throw cleanupError;
+}
+
 function requiredActiveHost() {
   if (!activeHost) throw new Error("no Nanocodex host is active");
   return activeHost;
 }
 
+function connectFailure(error) {
+  if (typeof error === "string") {
+    try {
+      const encoded = JSON.parse(error);
+      if (encoded?.kind === "handshake_rejected" || encoded?.kind === "transport") {
+        return encoded;
+      }
+    } catch {}
+  }
+  const status = Number(error?.status);
+  if (Number.isInteger(status) && status >= 100 && status <= 599) {
+    const retryAfter = Number(error?.retryAfter);
+    return {
+      kind: "handshake_rejected",
+      status,
+      body: typeof error?.body === "string" ? error.body : errorDetail(error),
+      ...(Number.isFinite(retryAfter) && retryAfter >= 0 ? { retry_after: retryAfter } : {}),
+    };
+  }
+  return {
+    kind: "transport",
+    detail: errorDetail(error),
+    reconnectable: true,
+  };
+}
+
+function errorDetail(error) {
+  return error && (error.stack || error.message) || String(error);
+}
+
 function createTurn(raw, agent) {
   if (!raw || typeof raw.result !== "function") {
     throw new TypeError("the runtime returned an invalid Nanocodex turn handle");
@@ -258,24 +402,67 @@ function createTurn(raw, agent) {
       state.raw.free();
     },
   };
-  Object.defineProperty(turn, TURN_STATE, { value: state });
+  turnStates.set(turn, state);
   return Object.freeze(turn);
 }
 
+function createTurnResult(raw) {
+  if (
+    !raw
+    || typeof raw.finalMessage !== "string"
+    || typeof raw.snapshot !== "function"
+    || typeof raw.usage !== "function"
+  ) {
+    raw?.free?.();
+    throw new TypeError("the runtime returned an invalid Nanocodex turn result");
+  }
+  const state = {
+    raw,
+    snapshotValue: undefined,
+    usageValue: undefined,
+    snapshot() {
+      state.snapshotValue ||= freezeJson(JSON.parse(raw.snapshot()));
+      return state.snapshotValue;
+    },
+    usage() {
+      state.usageValue ||= freezeJson(JSON.parse(raw.usage()));
+      return state.usageValue;
+    },
+  };
+  const result = {
+    finalMessage: raw.finalMessage,
+    get snapshot() { return state.snapshot(); },
+    get usage() { return state.usage(); },
+  };
+  resultStates.set(result, state);
+  return Object.freeze(result);
+}
+
 function agentState(agent) {
-  const state = agent?.[AGENT_STATE];
-  if (!state) throw new TypeError("expected a Nanocodex agent");
+  const state = knownAgentState(agent);
   if (state.disposed) throw new Error("the Nanocodex agent has been disposed");
   return state;
 }
 
+function knownAgentState(agent) {
+  const state = agentStates.get(agent);
+  if (!state) throw new TypeError("expected a Nanocodex agent");
+  return state;
+}
+
 function turnState(turn) {
-  const state = turn?.[TURN_STATE];
+  const state = turnStates.get(turn);
   if (!state) throw new TypeError("expected a Nanocodex turn");
   if (state.disposed) throw new Error("the Nanocodex turn has been disposed");
   return state;
 }
 
+function resultState(result) {
+  const state = resultStates.get(result);
+  if (!state) throw new TypeError("expected a completed Nanocodex turn result");
+  return state;
+}
+
 function actionInput(options) {
   const input = options?.input;
   if (typeof input !== "string" && !Array.isArray(input)) {
@@ -301,3 +488,32 @@ function isObject(value) {
 function copy(target, key, value) {
   if (value !== undefined) target[key] = value;
 }
+
+function runCleanup(errors, cleanup) {
+  try {
+    cleanup();
+  } catch (error) {
+    errors.push(error);
+  }
+}
+
+function throwCleanupErrors(errors) {
+  if (errors.length === 1) throw errors[0];
+  if (errors.length > 1) {
+    throw new AggregateError(errors, "multiple Nanocodex resources failed to release");
+  }
+}
+
+function freezeJson(value) {
+  if (!value || typeof value !== "object") return value;
+  const pending = [value];
+  while (pending.length) {
+    const current = pending.pop();
+    if (Object.isFrozen(current)) continue;
+    for (const child of Object.values(current)) {
+      if (child && typeof child === "object" && !Object.isFrozen(child)) pending.push(child);
+    }
+    Object.freeze(current);
+  }
+  return value;
+}

package/node/Agent.d.mts

@@ -1,6 +1,7 @@
 import type {
   AgentOptions,
   DefaultAgent,
+  MppSession,
   ToolMap,
 } from "../types.mjs";
 
@@ -9,14 +10,9 @@ export type Agent = DefaultAgent;
 /** Creates a Node-hosted Rust/WASM Agent. */
 export function create(options: create.Options): Promise<create.ReturnType>;
 export declare namespace create {
-  type Options = AgentOptions & {
+  type Options = AgentOptions & ({ apiKey: string; mpp?: never } | { apiKey?: never; mpp: MppSession }) & {
     apiBaseUrl?: string | undefined;
-    apiKey: string;
-    connectTimeoutMs?: number | undefined;
-    maxFrameBytes?: number | undefined;
-    maxQueuedBytes?: number | undefined;
-    maxQueuedMessages?: number | undefined;
-    sendTimeoutMs?: number | undefined;
+    module?: unknown;
     tools?: ToolMap | undefined;
     websocketUrl?: string | undefined;
   };

package/node/Agent.mjs

@@ -1,4 +1,5 @@
 import { createRequire } from "node:module";
+import initWeb, { Nanocodex as WebNanocodex } from "../pkg-web/nanocodex.js";
 
 import { agentActions } from "../actions/index.mjs";
 import {
@@ -12,32 +13,51 @@ import {
 } from "../internal.mjs";
 import { createNodeHost } from "./host.mjs";
 
-const require = createRequire(import.meta.url);
-const { Nanocodex } = require("../pkg-node/nanocodex.js");
+let initializedWeb;
+let NodeNanocodex;
 
 export function create(options = {}) {
   const {
     thinking,
     reasoningMode,
+    fastMode,
     instructions,
     sessionId,
+    workspace,
+    resume,
     apiKey,
+    mpp,
     websocketUrl,
     apiBaseUrl,
-    ...hostOptions
+    module,
+    tools,
   } = options;
   const events = createEventChannel();
-  const host = createNodeHost({ ...hostOptions, onEvent: events.emit });
+  if (mpp !== undefined && apiKey !== undefined) {
+    throw new TypeError("apiKey and mpp are mutually exclusive");
+  }
+  const host = createNodeHost({
+    mpp,
+    onEvent: events.emit,
+    tools,
+    workspace: workspace ?? resume?.workspace,
+  });
   activateHost(host);
   const runtime = defineRuntime({
     key: "node-wasm",
     name: "Nanocodex Node WASM",
     type: "node",
-    create(config) {
+    async create(config) {
+      activateHost(host);
+      const Nanocodex = module === undefined
+        ? loadNodeNanocodex()
+        : await loadWebNanocodex(module);
       activateHost(host);
       return new Nanocodex(JSON.stringify(toWasmConfig({
-        apiKey,
-        websocketUrl,
+        apiKey: apiKey ?? (mpp === undefined ? undefined : "mpp-managed"),
+        websocketUrl: websocketUrl ?? (mpp === undefined
+          ? undefined
+          : "wss://openai.mpp.tempo.xyz/v1/responses"),
         apiBaseUrl,
         ...config,
       })));
@@ -47,5 +67,25 @@ export function create(options = {}) {
     release: (raw) => releaseHostSession(host, raw.sessionId),
     decorate: (agent) => agent.extend(agentActions()),
   });
-  return createAgentClient(runtime, { thinking, reasoningMode, instructions, sessionId });
+  return createAgentClient(runtime, {
+    thinking,
+    reasoningMode,
+    fastMode,
+    instructions,
+    sessionId,
+    workspace,
+    resume,
+  });
+}
+
+function loadNodeNanocodex() {
+  const require = createRequire(import.meta.url);
+  NodeNanocodex ||= require("../pkg-node/nanocodex.js").Nanocodex;
+  return NodeNanocodex;
+}
+
+async function loadWebNanocodex(module) {
+  initializedWeb ||= initWeb({ module_or_path: module });
+  await initializedWeb;
+  return WebNanocodex;
 }