@tanstack/ai-code-mode 0.1.0

package/package.json

@@ -0,0 +1,62 @@
+{
+  "name": "@tanstack/ai-code-mode",
+  "version": "0.1.0",
+  "description": "Code Mode for TanStack AI - LLM-driven code execution in secure sandboxes",
+  "author": "",
+  "publishConfig": {
+    "access": "public"
+  },
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/TanStack/ai.git",
+    "directory": "packages/typescript/ai-code-mode"
+  },
+  "type": "module",
+  "module": "./dist/esm/index.js",
+  "types": "./dist/esm/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/esm/index.d.ts",
+      "import": "./dist/esm/index.js"
+    }
+  },
+  "sideEffects": false,
+  "engines": {
+    "node": ">=18"
+  },
+  "files": [
+    "dist",
+    "src"
+  ],
+  "scripts": {
+    "build": "vite build",
+    "clean": "premove ./build ./dist",
+    "lint:fix": "eslint ./src --fix",
+    "test:build": "publint --strict",
+    "test:eslint": "eslint ./src",
+    "test:lib": "vitest",
+    "test:lib:dev": "pnpm test:lib --watch",
+    "test:types": "tsc"
+  },
+  "keywords": [
+    "ai",
+    "tanstack",
+    "code-mode",
+    "llm",
+    "sandbox",
+    "isolate"
+  ],
+  "dependencies": {
+    "esbuild": "^0.25.12"
+  },
+  "peerDependencies": {
+    "@tanstack/ai": "workspace:*",
+    "zod": "^3.0.0 || ^4.0.0"
+  },
+  "devDependencies": {
+    "@tanstack/ai": "workspace:*",
+    "@vitest/coverage-v8": "4.0.14",
+    "zod": "^4.2.0"
+  }
+}

package/src/agent-store.ts

@@ -0,0 +1,43 @@
+export interface AgentSession {
+  name: string
+  systemPrompt: string
+  memory: Record<string, unknown>
+  createdAt: number
+  lastUsedAt: number
+}
+
+export interface AgentStore {
+  get: (name: string) => Promise<AgentSession | null>
+  set: (name: string, session: AgentSession) => Promise<void>
+  delete: (name: string) => Promise<void>
+  list: () => Promise<Array<string>>
+}
+
+export class InMemoryAgentStore implements AgentStore {
+  private sessions = new Map<string, AgentSession>()
+
+  get(name: string): Promise<AgentSession | null> {
+    return Promise.resolve(this.sessions.get(name) ?? null)
+  }
+
+  set(name: string, session: AgentSession): Promise<void> {
+    this.sessions.set(name, session)
+    return Promise.resolve()
+  }
+
+  delete(name: string): Promise<void> {
+    this.sessions.delete(name)
+    return Promise.resolve()
+  }
+
+  list(): Promise<Array<string>> {
+    return Promise.resolve(Array.from(this.sessions.keys()))
+  }
+}
+
+export function generateAgentName(): string {
+  const hex = Array.from({ length: 8 }, () =>
+    Math.floor(Math.random() * 16).toString(16),
+  ).join('')
+  return `agent_${hex}`
+}

package/src/bindings/tool-to-binding.ts

@@ -0,0 +1,132 @@
+import { convertSchemaToJsonSchema } from '@tanstack/ai'
+import type { ToolExecutionContext } from '@tanstack/ai'
+import type { CodeModeTool, ToolBinding } from '../types'
+
+/**
+ * Convert an array of TanStack AI tools to a Record of ToolBindings
+ *
+ * @param tools - Array of tools to convert
+ * @param prefix - Optional prefix to add to binding names (e.g., 'external_')
+ */
+export function toolsToBindings(
+  tools: Array<CodeModeTool>,
+  prefix: string = '',
+): Record<string, ToolBinding> {
+  const bindings: Record<string, ToolBinding> = {}
+
+  for (const tool of tools) {
+    const bindingName = `${prefix}${tool.name}`
+    bindings[bindingName] = toolToBinding(tool, prefix)
+  }
+
+  return bindings
+}
+
+/**
+ * Convert a single TanStack AI tool to a ToolBinding
+ *
+ * @param tool - Tool to convert
+ * @param prefix - Optional prefix to add to binding name (e.g., 'external_')
+ * @throws Error if the tool doesn't have an execute function
+ */
+export function toolToBinding(
+  tool: CodeModeTool,
+  prefix: string = '',
+): ToolBinding {
+  // Convert schemas (Zod or Standard Schema) to JSON Schema
+  const inputSchema = convertSchemaToJsonSchema(tool.inputSchema) || {
+    type: 'object',
+    properties: {},
+  }
+
+  const outputSchema = tool.outputSchema
+    ? convertSchemaToJsonSchema(tool.outputSchema)
+    : undefined
+
+  // Get execute function
+  // ServerTool has execute, ToolDefinition (without .server()) does not
+  let execute: (
+    args: unknown,
+    context?: ToolExecutionContext,
+  ) => Promise<unknown>
+
+  if ('execute' in tool && typeof tool.execute === 'function') {
+    const toolExecute = tool.execute
+    execute = (args: unknown, context?: ToolExecutionContext) => {
+      // Pass context to the underlying tool so it can emit custom events
+      return Promise.resolve(toolExecute(args, context))
+    }
+  } else if ('__toolSide' in tool && tool.__toolSide === 'definition') {
+    throw new Error(
+      `Tool "${tool.name}" is a ToolDefinition without an execute function. ` +
+        `Call .server(fn) to provide an implementation before using with Code Mode.`,
+    )
+  } else {
+    throw new Error(
+      `Tool "${tool.name}" does not have an execute function. ` +
+        `Code Mode requires tools with implementations.`,
+    )
+  }
+
+  return {
+    name: `${prefix}${tool.name}`,
+    description: tool.description,
+    inputSchema,
+    outputSchema,
+    execute,
+  }
+}
+
+/**
+ * Create event-aware bindings that emit custom events for each external function call.
+ * Wraps each binding's execute function to emit events before and after execution.
+ *
+ * @param bindings - Original tool bindings
+ * @param emitCustomEvent - Callback to emit custom events to the stream
+ */
+export function createEventAwareBindings(
+  bindings: Record<string, ToolBinding>,
+  emitCustomEvent: (eventName: string, data: Record<string, any>) => void,
+): Record<string, ToolBinding> {
+  const wrapped: Record<string, ToolBinding> = {}
+
+  for (const [name, binding] of Object.entries(bindings)) {
+    wrapped[name] = {
+      ...binding,
+      execute: async (args: unknown) => {
+        // Emit call event
+        emitCustomEvent('code_mode:external_call', {
+          function: name,
+          args,
+          timestamp: Date.now(),
+        })
+
+        const startTime = Date.now()
+        try {
+          // Create context for the underlying tool so it can also emit events
+          const toolContext: ToolExecutionContext = { emitCustomEvent }
+          const result = await binding.execute(args, toolContext)
+
+          // Emit result event
+          emitCustomEvent('code_mode:external_result', {
+            function: name,
+            result,
+            duration: Date.now() - startTime,
+          })
+
+          return result
+        } catch (error) {
+          // Emit error event
+          emitCustomEvent('code_mode:external_error', {
+            function: name,
+            error: error instanceof Error ? error.message : String(error),
+            duration: Date.now() - startTime,
+          })
+          throw error
+        }
+      },
+    }
+  }
+
+  return wrapped
+}

package/src/code-wrapper.ts

@@ -0,0 +1,24 @@
+/**
+ * Wrap user code in an async IIFE to support top-level await
+ *
+ * The code is wrapped so that:
+ * 1. Top-level await works
+ * 2. Return values are captured and serialized (for isolate boundary transfer)
+ * 3. Errors propagate correctly
+ */
+export function wrapCode(code: string): string {
+  // Results must be serialized to JSON because objects can't be transferred
+  // directly across the isolate boundary in isolated-vm
+  return `
+(async function() {
+  try {
+    const __userResult = await (async function() {
+${code}
+    })();
+    return JSON.stringify(__userResult);
+  } catch (__error) {
+    throw __error;
+  }
+})()
+`
+}

package/src/create-code-mode-tool.ts

@@ -0,0 +1,254 @@
+import { z } from 'zod'
+import { toolDefinition } from '@tanstack/ai'
+import {
+  createEventAwareBindings,
+  toolsToBindings,
+} from './bindings/tool-to-binding'
+import { stripTypeScript } from './strip-typescript'
+import type { ServerTool, ToolExecutionContext } from '@tanstack/ai'
+import type {
+  CodeModeTool,
+  CodeModeToolConfig,
+  CodeModeToolResult,
+  IsolateContext,
+} from './types'
+
+/**
+ * Schema for the execute_typescript tool input
+ */
+const executeTypescriptInputSchema = z.object({
+  typescriptCode: z
+    .string()
+    .describe(
+      'TypeScript code to execute in the sandbox. ' +
+        'Use external_* functions to call available APIs. ' +
+        'Return a value to pass results back.',
+    ),
+})
+
+/**
+ * Schema for the execute_typescript tool output
+ */
+const executeTypescriptOutputSchema = z.object({
+  success: z.boolean().describe('Whether execution completed without errors'),
+  result: z
+    .unknown()
+    .optional()
+    .describe('Return value from the executed code'),
+  logs: z
+    .array(z.string())
+    .optional()
+    .describe('Console output captured during execution'),
+  error: z
+    .object({
+      message: z.string(),
+      name: z.string().optional(),
+      line: z.number().optional(),
+    })
+    .optional()
+    .describe('Error details if execution failed'),
+})
+
+export type ExecuteTypescriptInput = z.infer<
+  typeof executeTypescriptInputSchema
+>
+export type ExecuteTypescriptOutput = z.infer<
+  typeof executeTypescriptOutputSchema
+>
+
+/**
+ * Create an execute_typescript tool that can be used alongside other agent tools.
+ *
+ * This tool allows an LLM to execute TypeScript code in a secure sandbox.
+ * Tools passed in the config become `external_*` functions available inside the sandbox.
+ *
+ * @example
+ * ```typescript
+ * import { createCodeMode } from '@tanstack/ai-code-mode'
+ * import { createNodeIsolateDriver } from '@tanstack/ai-isolate-node'
+ *
+ * const { tool, systemPrompt } = createCodeMode({
+ *   driver: createNodeIsolateDriver(),
+ *   tools: [weatherTool, dbTool],  // Become external_fetchWeather, external_dbQuery
+ *   timeout: 30000,
+ * })
+ *
+ * chat({
+ *   systemPrompts: [myPrompt, systemPrompt],
+ *   tools: [tool, searchTool, emailTool],
+ *   messages,
+ * })
+ * ```
+ */
+export function createCodeModeTool(
+  config: CodeModeToolConfig,
+): ServerTool<
+  typeof executeTypescriptInputSchema,
+  typeof executeTypescriptOutputSchema,
+  'execute_typescript'
+> {
+  const {
+    driver,
+    tools,
+    timeout = 30000,
+    memoryLimit = 128,
+    getSkillBindings,
+  } = config
+
+  // Validate tools
+  if (tools.length === 0) {
+    throw new Error('At least one tool must be provided to createCodeModeTool')
+  }
+
+  // Transform tools to bindings with external_ prefix (static bindings)
+  const staticBindings = toolsToBindings(tools, 'external_')
+
+  // Create the tool definition
+  const definition = toolDefinition({
+    name: 'execute_typescript' as const,
+    description: buildToolDescription(tools),
+    inputSchema: executeTypescriptInputSchema,
+    outputSchema: executeTypescriptOutputSchema,
+  })
+
+  // Return server tool with execute function that accepts context
+  return definition.server(
+    async (
+      input,
+      toolContext?: ToolExecutionContext,
+    ): Promise<CodeModeToolResult> => {
+      const { typescriptCode } = input
+
+      // Get emitCustomEvent from context or use no-op
+      const emitCustomEvent = toolContext?.emitCustomEvent || (() => {})
+
+      if (!typescriptCode || typeof typescriptCode !== 'string') {
+        return {
+          success: false,
+          error: {
+            message: 'typescriptCode must be a non-empty string',
+            name: 'ValidationError',
+          },
+        }
+      }
+
+      // Create a fresh sandbox context for this execution
+      let isolateContext: IsolateContext | null = null
+
+      // Emit execution started event immediately
+      emitCustomEvent('code_mode:execution_started', {
+        timestamp: Date.now(),
+        codeLength: typescriptCode.length,
+      })
+
+      try {
+        // Step 1: Strip TypeScript (also serves as syntax validation via esbuild)
+        let strippedCode: string
+        try {
+          strippedCode = await stripTypeScript(typescriptCode)
+        } catch (error) {
+          // Type/syntax error from esbuild
+          return {
+            success: false,
+            error: {
+              message: error instanceof Error ? error.message : String(error),
+              name: 'TypeScriptError',
+            },
+          }
+        }
+
+        // Step 2: Get dynamic skill bindings if available
+        const skillBindings = getSkillBindings ? await getSkillBindings() : {}
+
+        // Step 3: Merge static and dynamic bindings, then wrap with event awareness
+        const allBindings = { ...staticBindings, ...skillBindings }
+        const eventAwareBindings = createEventAwareBindings(
+          allBindings,
+          emitCustomEvent,
+        )
+
+        // Step 4: Create sandbox context with event-aware bindings
+        isolateContext = await driver.createContext({
+          bindings: eventAwareBindings,
+          timeout,
+          memoryLimit,
+        })
+
+        // Step 5: Execute the code in the sandbox
+        const executionResult = await isolateContext.execute(strippedCode)
+
+        // Emit console logs as custom events
+        if (executionResult.logs && executionResult.logs.length > 0) {
+          for (const log of executionResult.logs) {
+            // Parse log level from prefix (added by sandbox console implementation)
+            let level: 'log' | 'warn' | 'error' | 'info' = 'log'
+            let message = log
+
+            if (log.startsWith('ERROR: ')) {
+              level = 'error'
+              message = log.slice(7)
+            } else if (log.startsWith('WARN: ')) {
+              level = 'warn'
+              message = log.slice(6)
+            } else if (log.startsWith('INFO: ')) {
+              level = 'info'
+              message = log.slice(6)
+            }
+
+            emitCustomEvent('code_mode:console', {
+              level,
+              message,
+              timestamp: Date.now(),
+            })
+          }
+        }
+
+        if (executionResult.success) {
+          return {
+            success: true,
+            result: executionResult.value,
+            logs: executionResult.logs,
+          }
+        } else {
+          return {
+            success: false,
+            error: executionResult.error
+              ? {
+                  message: executionResult.error.message,
+                  name: executionResult.error.name,
+                }
+              : { message: 'Unknown execution error' },
+            logs: executionResult.logs,
+          }
+        }
+      } catch (error) {
+        return {
+          success: false,
+          error: {
+            message: error instanceof Error ? error.message : String(error),
+            name: error instanceof Error ? error.name : 'Error',
+          },
+        }
+      } finally {
+        // Always clean up the sandbox context
+        if (isolateContext) {
+          await isolateContext.dispose()
+        }
+      }
+    },
+  )
+}
+
+/**
+ * Build the tool description including available external functions
+ */
+function buildToolDescription(tools: Array<CodeModeTool>): string {
+  const externalFunctions = tools.map((t) => `external_${t.name}`).join(', ')
+
+  return (
+    `Execute TypeScript code in a secure sandbox environment. ` +
+    `The code can use these external API functions: ${externalFunctions}. ` +
+    `All external_* calls are async and must be awaited. ` +
+    `Return a value to pass results back. Use console.log() for debugging.`
+  )
+}

package/src/create-code-mode.ts

@@ -0,0 +1,35 @@
+import { createCodeModeTool } from './create-code-mode-tool'
+import { createCodeModeSystemPrompt } from './create-system-prompt'
+import type { CodeModeToolConfig } from './types'
+
+/**
+ * Create both the `execute_typescript` tool and its matching system prompt
+ * from a single config object.
+ *
+ * This is the recommended way to set up Code Mode — it ensures the tool and
+ * system prompt always stay in sync.
+ *
+ * @example
+ * ```typescript
+ * import { createCodeMode } from '@tanstack/ai-code-mode'
+ * import { createNodeIsolateDriver } from '@tanstack/ai-isolate-node'
+ *
+ * const { tool, systemPrompt } = createCodeMode({
+ *   driver: createNodeIsolateDriver(),
+ *   tools: [weatherTool, dbTool],
+ *   timeout: 30000,
+ * })
+ *
+ * chat({
+ *   systemPrompts: [myPrompt, systemPrompt],
+ *   tools: [tool, ...otherTools],
+ *   messages,
+ * })
+ * ```
+ */
+export function createCodeMode(config: CodeModeToolConfig) {
+  return {
+    tool: createCodeModeTool(config),
+    systemPrompt: createCodeModeSystemPrompt(config),
+  }
+}

package/src/create-system-prompt.ts

@@ -0,0 +1,95 @@
+import { toolsToBindings } from './bindings/tool-to-binding'
+import { generateTypeStubs } from './type-generator/json-schema-to-ts'
+import type { CodeModeToolConfig } from './types'
+
+/**
+ * Create a system prompt snippet that documents the execute_typescript tool
+ * and all available external_* functions.
+ *
+ * Add this to your system prompts array when using createCodeModeTool.
+ *
+ * @example
+ * ```typescript
+ * import { createCodeMode } from '@tanstack/ai-code-mode'
+ * import { createNodeIsolateDriver } from '@tanstack/ai-isolate-node'
+ *
+ * const { tool, systemPrompt } = createCodeMode({
+ *   driver: createNodeIsolateDriver(),
+ *   tools: [weatherTool, dbTool],
+ * })
+ *
+ * chat({
+ *   systemPrompts: ['You are a helpful assistant.', systemPrompt],
+ *   tools: [tool, ...otherTools],
+ * })
+ * ```
+ */
+export function createCodeModeSystemPrompt(config: CodeModeToolConfig): string {
+  const { tools } = config
+
+  // Transform tools to bindings with external_ prefix to generate correct type stubs
+  const bindings = toolsToBindings(tools, 'external_')
+
+  // Generate TypeScript type stubs for the external functions
+  const typeStubs = generateTypeStubs(bindings)
+
+  // Build function documentation
+  const functionDocs = Object.entries(bindings)
+    .map(([name, binding]) => {
+      const doc = `- \`${name}(input)\`: ${binding.description}`
+      return doc
+    })
+    .join('\n')
+
+  return `## Code Execution Tool
+
+You have access to \`execute_typescript\` which runs TypeScript code in a sandboxed environment.
+
+### When to Use
+
+Use \`execute_typescript\` when you need to:
+- Process data with loops, conditionals, or complex logic
+- Make multiple API calls in parallel (Promise.all)
+- Transform, filter, or aggregate data
+- Perform calculations or data analysis
+
+For simple operations, prefer calling tools directly.
+
+### Available External APIs
+
+Inside your TypeScript code, you can call these async functions:
+
+${functionDocs}
+
+### Type Definitions
+
+\`\`\`typescript
+${typeStubs}
+\`\`\`
+
+### Example
+
+\`\`\`typescript
+// Fetch weather for multiple cities in parallel
+const cities = ["Tokyo", "Paris", "NYC"];
+const results = await Promise.all(
+  cities.map(city => external_fetchWeather({ location: city }))
+);
+
+// Find the warmest city
+const warmest = results.reduce((prev, curr) => 
+  curr.temperature > prev.temperature ? curr : prev
+);
+
+return { warmestCity: warmest.location, temperature: warmest.temperature };
+\`\`\`
+
+### Important Notes
+
+- All \`external_*\` calls are async - always use \`await\`
+- Return a value to pass results back to you
+- Use \`console.log()\` for debugging (logs are captured)
+- The sandbox is isolated - no network access or file system
+- Each execution is independent (no shared state between calls)
+`
+}

package/src/index.ts

@@ -0,0 +1,55 @@
+// Code Mode Tool API
+export { createCodeModeTool } from './create-code-mode-tool'
+export type {
+  ExecuteTypescriptInput,
+  ExecuteTypescriptOutput,
+} from './create-code-mode-tool'
+
+export { createCodeModeSystemPrompt } from './create-system-prompt'
+export { createCodeMode } from './create-code-mode'
+
+export {
+  InMemoryAgentStore,
+  generateAgentName,
+  type AgentSession,
+  type AgentStore,
+} from './agent-store'
+
+// Bindings utilities (useful for custom implementations)
+export {
+  toolToBinding,
+  toolsToBindings,
+  createEventAwareBindings,
+} from './bindings/tool-to-binding'
+
+// Type generator (useful for custom system prompts)
+export {
+  generateTypeStubs,
+  jsonSchemaToTypeScript,
+  type TypeGeneratorOptions,
+} from './type-generator/json-schema-to-ts'
+
+// TypeScript stripping utility
+export { stripTypeScript } from './strip-typescript'
+
+// Code wrapper utility (used by isolate drivers)
+export { wrapCode } from './code-wrapper'
+
+// All types
+export type {
+  // Tool-based API types
+  CodeModeToolConfig,
+  CodeModeToolResult,
+  // Isolate driver interfaces (used by driver packages)
+  IsolateDriver,
+  IsolateConfig,
+  IsolateContext,
+  ExecutionResult,
+  NormalizedError,
+  // Tool binding (used by driver packages)
+  ToolBinding,
+  // Tool input types
+  CodeModeTool,
+  // Re-exported from @tanstack/ai
+  ToolExecutionContext,
+} from './types'