nanocodex 0.2.0 → 0.3.0

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 = {}) {
@@ -111,6 +148,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 +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);
@@ -150,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);
@@ -176,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);
@@ -202,6 +255,9 @@ function createAgent(raw, runtime) {
     raw,
     runtime,
     disposed: false,
+    released: false,
+    shutdownPromise: undefined,
+    subscriptions: new Set(),
     sessionId: raw.sessionId,
     uid: `agent-${nextAgentUid++}`,
   };
@@ -234,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;
 }
 
@@ -251,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");
@@ -264,7 +394,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 +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)) {
@@ -316,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

@@ -12,12 +12,7 @@ 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;
     websocketUrl?: string | undefined;
   };

package/node/Agent.mjs

@@ -23,19 +23,25 @@ export function create(options = {}) {
     fastMode,
     instructions,
     sessionId,
+    workspace,
     resume,
     apiKey,
     mpp,
     websocketUrl,
     apiBaseUrl,
     module,
-    ...hostOptions
+    tools,
   } = 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,
+    workspace: workspace ?? resume?.workspace,
+  });
   activateHost(host);
   const runtime = defineRuntime({
     key: "node-wasm",
@@ -67,6 +73,7 @@ export function create(options = {}) {
     fastMode,
     instructions,
     sessionId,
+    workspace,
     resume,
   });
 }

package/node/host.mjs

@@ -1,10 +1,13 @@
 import { Console } from "node:console";
 import { createRequire } from "node:module";
+import { resolve } from "node:path";
 import WebSocket from "ws";
+import packageMetadata from "../package.json" with { type: "json" };
 
 import { createCodeRuntime } from "../runtime/code-runtime.mjs";
 
 const RESPONSES_WEBSOCKETS_BETA = "responses_websockets=2026-02-06";
+const USER_AGENT = `nanocodex-wasm/${packageMetadata.version}`;
 const DEFAULT_MAX_QUEUED_MESSAGES = 4_096;
 const DEFAULT_MAX_QUEUED_BYTES = 32 * 1024 * 1024;
 const DEFAULT_MAX_FRAME_BYTES = 16 * 1024 * 1024;
@@ -12,7 +15,7 @@ const DEFAULT_MAX_FRAME_BYTES = 16 * 1024 * 1024;
 export function createNodeHost(options = {}) {
   const connections = new Map();
   const code = createCodeRuntime(options.tools, {
-    require: createRequire(`${process.cwd()}/`),
+    require: createRequire(resolve(options.workspace ?? process.cwd(), ".nanocodex-code-mode.cjs")),
     console: new Console({ stdout: process.stderr, stderr: process.stderr }),
   });
   const onEvent = options.onEvent || (() => {});
@@ -23,29 +26,47 @@ export function createNodeHost(options = {}) {
   const maxFrameBytes = options.maxFrameBytes ?? DEFAULT_MAX_FRAME_BYTES;
   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;
       let upgradeResponse;
+      const headers = {
+        Authorization: `Bearer ${apiKey}`,
+        "OpenAI-Beta": RESPONSES_WEBSOCKETS_BETA,
+        "x-openai-internal-codex-responses-lite": "true",
+        "session-id": sessionId,
+        "thread-id": sessionId,
+        "x-client-request-id": sessionId,
+        "x-responsesapi-include-timing-metrics": "true",
+        "User-Agent": USER_AGENT,
+      };
+      if (metadata.accountId) headers["ChatGPT-Account-ID"] = metadata.accountId;
+      if (metadata.fedramp) headers["X-OpenAI-Fedramp"] = "true";
+      if (metadata.turnState) headers["x-codex-turn-state"] = metadata.turnState;
       const socket = new WebSocket(endpoint, {
         handshakeTimeout: connectTimeoutMs,
         maxPayload: maxFrameBytes,
-        headers: {
-          Authorization: `Bearer ${apiKey}`,
-          "OpenAI-Beta": RESPONSES_WEBSOCKETS_BETA,
-          "x-openai-internal-codex-responses-lite": "true",
-          "session-id": sessionId,
-          "thread-id": sessionId,
-          "x-client-request-id": sessionId,
-          "x-responsesapi-include-timing-metrics": "true",
-          "User-Agent": "nanocodex-wasm/0.1.0",
-        },
+        headers,
       });
       const handle = nextHandle++;
       const connection = queueState(socket);
 
       socket.on("upgrade", (response) => { upgradeResponse = response; });
+      socket.on("unexpected-response", (_request, response) => {
+        if (settled) return;
+        settled = true;
+        const chunks = [];
+        response.on("data", (chunk) => chunks.push(Buffer.from(chunk)));
+        response.on("end", () => {
+          const error = new Error(`WebSocket handshake was rejected with HTTP ${response.statusCode}`);
+          error.status = response.statusCode;
+          error.body = chunks.length ? Buffer.concat(chunks).toString("utf8") : "empty response body";
+          const retryAfter = Number(header(response.headers, "retry-after"));
+          if (Number.isFinite(retryAfter) && retryAfter >= 0) error.retryAfter = retryAfter;
+          reject(error);
+        });
+      });
       socket.on("open", () => {
         settled = true;
         connections.set(handle, connection);

package/node/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/package.json

@@ -1,6 +1,6 @@
 {
   "name": "nanocodex",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "Viem-style Node and browser bindings for the Nanocodex WASM agent",
   "type": "module",
   "sideEffects": false,
@@ -25,7 +25,6 @@
     "pkg-node",
     "pkg-web",
     "runtime",
-    "scripts/check-package.mjs",
     "internal.mjs",
     "index.d.mts",
     "index.mjs",
@@ -39,7 +38,8 @@
   "scripts": {
     "check:package": "node scripts/check-package.mjs",
     "prepack": "npm run check:package",
-    "test": "node --test test/*.test.mjs && npm run test:typecheck && npm run check:package",
+    "test": "node --test test/*.test.mjs && npm run test:performance && npm run test:typecheck && npm run check:package",
+    "test:performance": "node --test test/performance.bench.mjs",
     "test:typecheck": "tsc -p test/tsconfig.json"
   },
   "dependencies": {