nanocodex 0.2.0 → 0.4.0

package/browser/host.mjs

@@ -3,26 +3,41 @@ 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 toolMode = options.toolMode ?? "code";
+  if (toolMode !== "code" && toolMode !== "direct") {
+    throw new TypeError("toolMode must be code or direct");
+  }
   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) {
+  async function connect(endpoint, apiKey, sessionId, metadata = {}) {
     if (options.mpp) return connectMpp(endpoint);
+    const authorization = options.hostAuth
+      ? { authorization: "host_managed" }
+      : { authorization: "bearer", bearerToken: apiKey };
+    const request = { ...metadata };
+    delete request.authorization;
+    delete request.bearerToken;
+    Object.assign(request, authorization);
+    const opened = await createWebSocket(endpoint, sessionId, request);
+    const { socket, ...handshake } = normalizeWebSocketConnection(opened);
     return new Promise((resolve, reject) => {
       let settled = false;
-      const handle = nextHandle++;
-      const socket = createWebSocket(endpoint, sessionId);
       const connection = {
         socket,
         queue: [],
@@ -31,11 +46,21 @@ export function createBrowserHost(options = {}) {
         intentionallyClosed: false,
         overflowed: false,
       };
-      socket.addEventListener("open", () => {
+      const resolveOpen = () => {
+        if (settled) return;
         settled = true;
+        const handle = nextHandle++;
         connections.set(handle, connection);
-        resolve(JSON.stringify({ handle, status: 101, reasoning_included: false }));
-      }, { once: true });
+        resolve(JSON.stringify({
+          handle,
+          status: handshake.status ?? 101,
+          request_id: handshake.requestId,
+          server_model: handshake.serverModel,
+          reasoning_included: handshake.reasoningIncluded ?? false,
+          turn_state: handshake.turnState,
+        }));
+      };
+      socket.addEventListener("open", resolveOpen, { once: true });
       socket.addEventListener("message", (event) => {
         enqueue(connection, typeof event.data === "string"
           ? { kind: "text", text: event.data }
@@ -57,6 +82,11 @@ export function createBrowserHost(options = {}) {
           enqueue(connection, { kind: "error", detail: "WebSocket connection failed" });
         }
       });
+      if (socket.readyState === WEBSOCKET_OPEN) resolveOpen();
+      else if (socket.readyState > WEBSOCKET_OPEN) {
+        settled = true;
+        reject(new Error("WebSocket closed during connection"));
+      }
     });
   }
 
@@ -97,7 +127,7 @@ export function createBrowserHost(options = {}) {
 
   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,
@@ -123,7 +153,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),
       }));
     }
@@ -194,8 +224,20 @@ export function createBrowserHost(options = {}) {
     close,
     sleep: (milliseconds) => new Promise((resolve) => setTimeout(resolve, milliseconds)),
     executeCode: code.executeCode,
+    executeTool: code.executeTool,
+    toolMode: () => toolMode,
     toolDefinitions: code.toolDefinitions,
     emitEvent: onEvent,
     reset: code.reset,
   });
 }
+
+function normalizeWebSocketConnection(opened) {
+  if (opened?.socket && typeof opened.socket.addEventListener === "function") {
+    return opened;
+  }
+  if (!opened || typeof opened.addEventListener !== "function") {
+    throw new TypeError("createWebSocket must return a WebSocket or a connection descriptor");
+  }
+  return { socket: opened };
+}

package/browser/index.d.mts

@@ -1,6 +1,8 @@
 export { Actions } from "../index.mjs";
 export type {
   AgentEvent,
+  CostStatus,
+  EstimatedUsdCost,
   PromptInput,
   PromptItem,
   ReasoningMode,
@@ -10,5 +12,7 @@ export type {
   ToolContext,
   ToolMap,
   Turn,
+  TurnResult,
+  TurnUsage,
 } from "../types.mjs";
 export * as Agent from "./Agent.mjs";

package/index.d.mts

@@ -4,8 +4,10 @@ export type {
   AgentActions,
   AgentEvent,
   AgentOptions,
+  CostStatus,
   DefaultAgent,
   EventWatcher,
+  EstimatedUsdCost,
   ForkOptions,
   PromptInput,
   PromptItem,
@@ -16,5 +18,7 @@ export type {
   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,12 +42,18 @@ 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(turn) {
-  return JSON.parse(turnState(turn).raw.snapshot());
+export function getTurnSnapshot(result) {
+  return resultState(result).snapshot();
+}
+
+export function getTurnUsage(result) {
+  return resultState(result).usage();
 }
 
 export function steer(turn, options) {
@@ -66,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);
 }
 
@@ -83,7 +90,23 @@ export function setFastMode(agent, enabled) {
   return agentState(agent).raw.setFastMode(enabled);
 }
 
-export function subscribeAgentEvents(agent, listener, options = {}) {
+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");
@@ -91,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 = {}) {
@@ -104,6 +141,7 @@ export function toWasmConfig(options = {}) {
     throw new TypeError("apiKey must be a non-empty string");
   }
   const config = { api_key: apiKey };
+  copy(config, "model", options.model);
   copy(config, "thinking", options.thinking);
   copy(config, "reasoning_mode", options.reasoningMode);
   copy(config, "fast_mode", options.fastMode);
@@ -111,6 +149,7 @@ export function toWasmConfig(options = {}) {
   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;
 }
@@ -119,8 +158,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);
@@ -150,9 +191,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);
@@ -176,12 +226,23 @@ 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);
   },
+  executeTool(name, input, sessionId, callId) {
+    return requiredSessionHost(sessionId).executeTool(name, input, sessionId, callId);
+  },
+  toolMode(sessionId) {
+    // The WASM constructor asks before its session is adopted.
+    return (hostSessions.get(sessionId) ?? requiredActiveHost()).toolMode();
+  },
   toolDefinitions(sessionId) {
     // ModelRun builds its stable tool prefix inside the WASM constructor,
     // immediately before the returned session can be adopted. Runtime
@@ -202,6 +263,9 @@ function createAgent(raw, runtime) {
     raw,
     runtime,
     disposed: false,
+    released: false,
+    shutdownPromise: undefined,
+    subscriptions: new Set(),
     sessionId: raw.sessionId,
     uid: `agent-${nextAgentUid++}`,
   };
@@ -234,14 +298,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;
 }
 
@@ -251,11 +313,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");
@@ -264,7 +402,6 @@ function createTurn(raw, agent) {
   const turn = {
     get agent() { return state.agent; },
     result: () => getTurnResult(turn),
-    snapshot: () => getTurnSnapshot(turn),
     steer: (input) => steer(turn, input),
     cancel: () => cancel(turn),
     dispose() {
@@ -273,24 +410,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)) {
@@ -316,3 +496,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

@@ -12,13 +12,9 @@ export function create(options: create.Options): Promise<create.ReturnType>;
 export declare namespace create {
   type Options = AgentOptions & ({ apiKey: string; mpp?: never } | { apiKey?: never; mpp: MppSession }) & {
     apiBaseUrl?: string | undefined;
-    connectTimeoutMs?: number | undefined;
-    maxFrameBytes?: number | undefined;
-    maxQueuedBytes?: number | undefined;
-    maxQueuedMessages?: number | undefined;
     module?: unknown;
-    sendTimeoutMs?: number | undefined;
     tools?: ToolMap | undefined;
+    toolMode?: "code" | "direct" | undefined;
     websocketUrl?: string | undefined;
   };
   type ReturnType = Agent;

package/node/Agent.mjs

@@ -18,24 +18,33 @@ let NodeNanocodex;
 
 export function create(options = {}) {
   const {
+    model,
     thinking,
     reasoningMode,
     fastMode,
     instructions,
     sessionId,
+    workspace,
     resume,
     apiKey,
     mpp,
     websocketUrl,
     apiBaseUrl,
     module,
-    ...hostOptions
+    tools,
+    toolMode,
   } = options;
   const events = createEventChannel();
   if (mpp !== undefined && apiKey !== undefined) {
     throw new TypeError("apiKey and mpp are mutually exclusive");
   }
-  const host = createNodeHost({ ...hostOptions, mpp, onEvent: events.emit });
+  const host = createNodeHost({
+    mpp,
+    onEvent: events.emit,
+    tools,
+    toolMode,
+    workspace: workspace ?? resume?.workspace,
+  });
   activateHost(host);
   const runtime = defineRuntime({
     key: "node-wasm",
@@ -62,11 +71,13 @@ export function create(options = {}) {
     decorate: (agent) => agent.extend(agentActions()),
   });
   return createAgentClient(runtime, {
+    model,
     thinking,
     reasoningMode,
     fastMode,
     instructions,
     sessionId,
+    workspace,
     resume,
   });
 }