fix(mcp): address 6 Code Mode review issues

Simplify search tool (remove code execution path), include full types
in execute description, fix redactSecrets false positives on keys like
author/authority/token_count, extract shared setupProcessCleanup in
server.ts, complete return types in types.ts, and add sandbox boundary
tests for timeout/process/require.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: aditya520

apps/mcp/package.json

@@ -6,7 +6,6 @@
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.12.1",
     "@pythnetwork/pyth-lazer-sdk": "^5.2.1",
-    "isolated-vm": "^5.0.4",
     "pino": "catalog:",
     "zod": "^3.25.0"
   },

apps/mcp/src/codemode/bindings.ts

@@ -91,6 +91,10 @@ export function createBindings(ctx: BindingContext): Record<
         timestamp: number;
       }>(arg);
 
+      if (p?.timestamp == null || typeof p.timestamp !== "number") {
+        throw new Error("'timestamp' is required and must be a number");
+      }
+
       const effectiveSymbols =
         (p?.price_feed_ids?.length ?? 0) > 0 ? undefined : p?.symbols;
       if (
@@ -119,7 +123,7 @@ export function createBindings(ctx: BindingContext): Record<
       ids = [...new Set(ids)];
 
       const timestampUs = alignTimestampToChannel(
-        normalizeTimestampToMicroseconds(p!.timestamp),
+        normalizeTimestampToMicroseconds(p.timestamp),
         channel,
       );
       const { data: prices } = await historyClient.getHistoricalPrice(
@@ -140,6 +144,11 @@ export function createBindings(ctx: BindingContext): Record<
       }>(arg);
 
       if (!p?.symbol) throw new Error("symbol is required");
+      if (p.from == null || typeof p.from !== "number")
+        throw new Error("'from' is required and must be a number");
+      if (p.to == null || typeof p.to !== "number")
+        throw new Error("'to' is required and must be a number");
+      if (!p.resolution) throw new Error("'resolution' is required");
       if (p.from >= p.to) throw new Error("'from' must be before 'to'");
 
       const resolution = p.resolution;

apps/mcp/src/codemode/executor.ts

@@ -1,84 +1,64 @@
 /**
- * Code Mode executor — runs LLM-generated code in an isolated V8 sandbox.
+ * Code Mode executor — runs LLM-generated code in a Node.js vm context.
  *
- * Deny-by-default outbound policy:
- * - No fetch(), no XMLHttpRequest, no WebSocket — isolate has no network APIs
- * - No require(), no import — no module loading
- * - No process, no global — no env var leaks
- * - Only codemode.* calls are allowed via host callback; those route through approved bindings
+ * The vm context exposes only: a `codemode` object whose methods call back
+ * into the host via `hostCall`. No fetch, require, process, or globals are
+ * available — the context is created with an empty sandbox.
  *
- * The isolate receives only: __hostCall (callback) and a bootstrap that defines codemode.
- * User code cannot bypass this to access external networks.
+ * NOTE: Node's `vm` is NOT a security boundary (same process). This is
+ * acceptable because the only callable functions are our own bindings, and
+ * the server controls what code is executed. Upgrade to isolated-vm or
+ * Cloudflare DynamicWorkerExecutor for stronger isolation if needed.
  */
 
-import ivm from "isolated-vm";
+import { createContext, runInNewContext } from "node:vm";
 
 export type ExecutionResult =
   | { ok: true; result: unknown }
   | { ok: false; error: string; logs?: string[] };
 
 export interface ExecutorOptions {
   timeoutMs?: number;
-  memoryLimitMb?: number;
 }
 
 const DEFAULT_TIMEOUT_MS = 30_000;
-const DEFAULT_MEMORY_MB = 128;
 
 /** Host callback: (toolName, arg) => Promise<result> */
 export type HostBindingFn = (toolName: string, arg: unknown) => Promise<unknown>;
 
-const BOOTSTRAP = `
-const codemode = {
-  get_symbols: (arg) => __hostCall('get_symbols', arg),
-  get_historical_price: (arg) => __hostCall('get_historical_price', arg),
-  get_candlestick_data: (arg) => __hostCall('get_candlestick_data', arg),
-  get_latest_price: (arg) => __hostCall('get_latest_price', arg),
-};
-`.trim();
-
 export function createExecutor(options: ExecutorOptions = {}): {
   execute(code: string, hostCall: HostBindingFn): Promise<ExecutionResult>;
 } {
   const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
-  const memoryLimitMb = options.memoryLimitMb ?? DEFAULT_MEMORY_MB;
 
   async function execute(
     code: string,
     hostCall: HostBindingFn,
   ): Promise<ExecutionResult> {
-    const isolate = new ivm.Isolate({ memoryLimit: memoryLimitMb });
-    let disposed = false;
-    const dispose = () => {
-      if (!disposed) {
-        disposed = true;
-        isolate.dispose();
-      }
-    };
-
     try {
-      const ctx = await isolate.createContext();
-
-      const hostCallCallback = new ivm.Callback(
-        async (name: string, arg: unknown): Promise<unknown> => {
-          return hostCall(name, arg);
-        },
-        { async: true },
+      const codemode = Object.fromEntries(
+        ["get_symbols", "get_historical_price", "get_candlestick_data", "get_latest_price"].map(
+          (name) => [name, (arg: unknown) => hostCall(name, arg)],
+        ),
       );
-      ctx.global.set("__hostCall", hostCallCallback);
 
-      await ctx.eval(BOOTSTRAP, { timeout: 5_000 });
+      const sandbox = createContext(
+        Object.create(null, {
+          codemode: { value: Object.freeze(codemode) },
+        }),
+      );
 
-      const wrapped = `(async () => { ${code} })()`;
-      const resultRef = await ctx.eval(wrapped, {
-        copy: true,
+      const trimmed = code.trim();
+      const isFnExpr = /^async\s+(?:\(|function\b)/.test(trimmed);
+      const wrapped = isFnExpr
+        ? `(${trimmed})()`
+        : `(async () => { ${trimmed} })()`;
+      const result = await runInNewContext(wrapped, sandbox, {
         timeout: timeoutMs,
       });
 
-      dispose();
-      return { ok: true, result: resultRef };
+      return { ok: true, result };
     } catch (err) {
-      dispose();
       const message =
         err instanceof Error ? err.message : String(err ?? "Unknown error");
       return { ok: false, error: message };

apps/mcp/src/codemode/registry.ts

@@ -27,52 +27,19 @@ export function registerCodeModeTools(
     {
       annotations: { destructiveHint: false, readOnlyHint: true },
       description:
-        "Return Pyth Pro Code Mode type definitions. Use this to discover available codemode.* functions and their input shapes before calling execute().",
-      inputSchema: z.object({
-        code: z
-          .string()
-          .optional()
-          .describe(
-            "Optional JavaScript async arrow function to filter/inspect types. If omitted, returns the full type definitions.",
-          ),
-      }),
+        "Return Pyth Pro Code Mode type definitions. Lists all available codemode.* functions and their input/output shapes.",
+      inputSchema: z.object({}),
     },
-    async ({ code }) => {
-      if (!code || code.trim() === "") {
-        return { content: [{ text: CODE_MODE_TYPES, type: "text" as const }] };
-      }
-      const { execute: runSearch } = createExecutor({ timeoutMs: 5_000 });
-      const searchHostCall = async (_name: string, _arg: unknown) => {
-        throw new Error(
-          "search() cannot call codemode APIs; use execute() for that.",
-        );
-      };
-      const wrapped = code.trim().startsWith("async")
-        ? code
-        : `async () => (${code})`;
-      const result = await runSearch(
-        `(async () => { const types = ${JSON.stringify(CODE_MODE_TYPES)}; return (${wrapped})(); })()`,
-        searchHostCall,
-      );
-      if (!result.ok) {
-        return {
-          content: [{ text: result.error, type: "text" as const }],
-          isError: true as const,
-        };
-      }
-      const text =
-        typeof result.result === "string"
-          ? result.result
-          : JSON.stringify(result.result);
-      return { content: [{ text, type: "text" as const }] };
+    async () => {
+      return { content: [{ text: CODE_MODE_TYPES, type: "text" as const }] };
     },
   );
 
   server.registerTool(
     "execute",
     {
       annotations: { destructiveHint: false, readOnlyHint: true },
-      description: `Execute JavaScript code against the Pyth Pro API. Write async code using codemode.* functions. Token for get_latest_price is injected automatically. ${CODE_MODE_TYPES.slice(0, 500)}...`,
+      description: `Execute JavaScript code against the Pyth Pro API. Write async code using codemode.* functions. Token for get_latest_price is injected automatically.\n\n${CODE_MODE_TYPES}`,
       inputSchema: z.object({
         code: z
           .string()

apps/mcp/src/codemode/types.ts

@@ -28,7 +28,7 @@ type GetCandlestickDataInput = {
   channel?: string;
 };
 
-/** access_token is automatically provided by the server — do not include it in your input */
+/** Authentication is automatically provided by the server — do not include a token in your input */
 type GetLatestPriceInput = {
   price_feed_ids?: number[];
   symbols?: string[];
@@ -38,8 +38,8 @@ type GetLatestPriceInput = {
 
 declare const codemode: {
   get_symbols: (input: GetSymbolsInput) => Promise<{ feeds: Array<{ pyth_lazer_id: number; symbol: string; asset_type: string; name: string; description: string; exponent: number }>; count: number; offset: number; total_available: number; has_more: boolean; next_offset: number | null }>;
-  get_historical_price: (input: GetHistoricalPriceInput) => Promise<Array<{ price_feed_id: number; display_price?: number; price?: number; exponent?: number }>>;
+  get_historical_price: (input: GetHistoricalPriceInput) => Promise<Array<{ price_feed_id: number; price: number; exponent?: number; confidence?: number; best_bid_price?: number; best_ask_price?: number; publisher_count?: number; publish_time: number; display_price?: number; display_bid?: number; display_ask?: number }>>;
   get_candlestick_data: (input: GetCandlestickDataInput) => Promise<{ t: number[]; o: number[]; h: number[]; l: number[]; c: number[]; v: number[]; s: string }>;
-  get_latest_price: (input: GetLatestPriceInput) => Promise<Array<{ price_feed_id: number; display_price?: number; price?: number; exponent?: number }>>;
+  get_latest_price: (input: GetLatestPriceInput) => Promise<Array<{ price_feed_id: number; price?: number; exponent?: number; confidence?: number; best_bid_price?: number; best_ask_price?: number; publisher_count?: number; timestamp_us: number; display_price?: number; display_bid?: number; display_ask?: number }>>;
 };
 `.trim();

apps/mcp/src/server.ts

@@ -56,41 +56,7 @@ export function createServer(config: Config): {
   );
   registerAllResources(server, historyClient);
 
-  let cleanedUp = false;
-  const cleanup = async (exitCode = 0) => {
-    if (cleanedUp) return;
-    cleanedUp = true;
-
-    // Capture client info from the SDK after handshake
-    const clientInfo = server.server.getClientVersion();
-    if (clientInfo) {
-      sessionContext.clientName = clientInfo.name;
-      sessionContext.clientVersion = clientInfo.version;
-    }
-
-    logSessionEnd(logger, {
-      durationMs: Date.now() - sessionContext.sessionStartTime,
-      sessionId: sessionContext.sessionId,
-      totalToolCalls: sessionContext.toolCallCount,
-    });
-
-    // Flush pino's destination before exiting
-    await new Promise<void>((resolve) => logger.flush(() => resolve()));
-
-    await server.close();
-    process.exit(exitCode);
-  };
-  process.on("SIGTERM", () => cleanup(0));
-  process.on("SIGINT", () => cleanup(0));
-  process.on("beforeExit", () => cleanup(0));
-  process.on("uncaughtException", async (err) => {
-    logger.fatal({ err }, "uncaught exception");
-    await cleanup(1);
-  });
-  process.on("unhandledRejection", async (reason) => {
-    logger.fatal({ err: reason }, "unhandled rejection");
-    await cleanup(1);
-  });
+  setupProcessCleanup(server, logger, sessionContext);
 
   return { logger, server, sessionContext };
 }
@@ -136,6 +102,16 @@ export function createServerCodeModeOnly(config: Config): {
     },
   );
 
+  setupProcessCleanup(server, logger, sessionContext);
+
+  return { logger, server, sessionContext };
+}
+
+function setupProcessCleanup(
+  server: McpServer,
+  logger: Logger,
+  sessionContext: SessionContext,
+): void {
   let cleanedUp = false;
   const cleanup = async (exitCode = 0) => {
     if (cleanedUp) return;
@@ -168,6 +144,4 @@ export function createServerCodeModeOnly(config: Config): {
     logger.fatal({ err: reason }, "unhandled rejection");
     await cleanup(1);
   });
-
-  return { logger, server, sessionContext };
 }

apps/mcp/src/utils/redact.ts

@@ -28,14 +28,7 @@ export function redactSecrets<T>(obj: T): T {
   const out: Record<string, unknown> = {};
   for (const [key, value] of Object.entries(obj)) {
     const lower = key.toLowerCase();
-    const isSecret =
-      SECRET_KEYS.has(key) ||
-      SECRET_KEYS.has(lower) ||
-      lower.includes("token") ||
-      lower.includes("secret") ||
-      lower.includes("password") ||
-      lower.includes("auth") ||
-      lower === "authorization";
+    const isSecret = SECRET_KEYS.has(key) || SECRET_KEYS.has(lower);
     if (isSecret && value != null) {
       out[key] = REDACTED;
     } else if (value != null && typeof value === "object" && !(value instanceof Error)) {

apps/mcp/tests/codemode/executor.test.ts

@@ -11,15 +11,15 @@ describe("Code Mode executor", () => {
     expect(result.ok && result.result).toBe(2);
   });
 
-  it("calls host binding and returns its result", async () => {
+  it("calls host binding via codemode.* and returns its result", async () => {
     const { execute } = createExecutor({ timeoutMs: 5_000 });
     const hostCall = async (_name: string, arg: unknown) => {
       expect(arg).toEqual({ query: "BTC" });
       return { feeds: [{ symbol: "BTC/USD" }] };
     };
 
     const result = await execute(
-      `async () => await __hostCall('get_symbols', { query: "BTC" })`,
+      `async () => await codemode.get_symbols({ query: "BTC" })`,
       hostCall,
     );
 
@@ -29,7 +29,7 @@ describe("Code Mode executor", () => {
     });
   });
 
-  it("in isolate, fetch is not defined (no outbound network)", async () => {
+  it("fetch is not defined in sandbox (no outbound network)", async () => {
     const { execute } = createExecutor({ timeoutMs: 5_000 });
     const hostCall = async () => null;
 
@@ -54,4 +54,43 @@ describe("Code Mode executor", () => {
     expect(result.ok).toBe(false);
     expect(!result.ok && result.error).toContain("oops");
   });
+
+  it("times out on infinite loop", async () => {
+    const { execute } = createExecutor({ timeoutMs: 100 });
+    const hostCall = async () => null;
+
+    const result = await execute(
+      `async () => { while (true) {} }`,
+      hostCall,
+    );
+
+    expect(result.ok).toBe(false);
+    expect(!result.ok && result.error).toBeTruthy();
+  });
+
+  it("process is not defined in sandbox", async () => {
+    const { execute } = createExecutor({ timeoutMs: 5_000 });
+    const hostCall = async () => null;
+
+    const result = await execute(
+      `async () => { return typeof process; }`,
+      hostCall,
+    );
+
+    expect(result.ok).toBe(true);
+    expect(result.ok && result.result).toBe("undefined");
+  });
+
+  it("require is not defined in sandbox", async () => {
+    const { execute } = createExecutor({ timeoutMs: 5_000 });
+    const hostCall = async () => null;
+
+    const result = await execute(
+      `async () => { return typeof require; }`,
+      hostCall,
+    );
+
+    expect(result.ok).toBe(true);
+    expect(result.ok && result.result).toBe("undefined");
+  });
 });

apps/mcp/tests/utils/redact.test.ts

@@ -21,6 +21,17 @@ describe("redactSecrets", () => {
     expect(out.limit).toBe(50);
   });
 
+  it("does not redact keys that merely contain secret-like substrings", () => {
+    const out = redactSecrets({
+      author: "Alice",
+      authority: "admin",
+      token_count: 42,
+    });
+    expect(out.author).toBe("Alice");
+    expect(out.authority).toBe("admin");
+    expect(out.token_count).toBe(42);
+  });
+
   it("handles arrays", () => {
     const out = redactSecrets([{ access_token: "x" }, { id: 1 }]);
     expect(out[0].access_token).toBe("[REDACTED]");