@tanstack/ai-code-mode-snippets 0.3.14

package/src/create-snippet-management-tools.ts

@@ -0,0 +1,298 @@
+import { toolDefinition } from '@tanstack/ai'
+import { toolsToBindings } from '@tanstack/ai-code-mode'
+import { z } from 'zod'
+import { createDefaultTrustStrategy } from './trust-strategies'
+import { snippetToTool } from './snippets-to-tools'
+import type { SchemaInput, ServerTool, ToolRegistry } from '@tanstack/ai'
+import type { CodeModeToolConfig, ToolBinding } from '@tanstack/ai-code-mode'
+import type { SnippetStorage } from './types'
+import type { TrustStrategy } from './trust-strategies'
+
+interface CreateSnippetManagementToolsOptions {
+  /**
+   * Storage implementation for snippets
+   */
+  storage: SnippetStorage
+
+  /**
+   * Trust strategy for determining initial trust level.
+   * If not provided, uses the storage's trustStrategy or falls back to default.
+   */
+  trustStrategy?: TrustStrategy
+
+  /**
+   * Tool registry for adding newly registered snippets immediately.
+   * When provided, register_snippet will add the new snippet to this registry
+   * so it's available as a direct tool in the current chat session.
+   */
+  registry?: ToolRegistry
+
+  /**
+   * Code mode config for creating snippet tools.
+   * Required when registry is provided.
+   */
+  config?: CodeModeToolConfig
+
+  /**
+   * Pre-computed bindings for external_* functions.
+   * Required when registry is provided.
+   */
+  baseBindings?: Record<string, ToolBinding>
+}
+
+/**
+ * Create tools for searching, retrieving, and registering snippets.
+ * These tools allow the LLM to interact with the snippet library at runtime.
+ *
+ * When registry, config, and baseBindings are provided, newly registered snippets
+ * will be immediately added to the registry and available as direct tools.
+ */
+export function createSnippetManagementTools({
+  storage,
+  trustStrategy,
+  registry,
+  config,
+  baseBindings,
+}: CreateSnippetManagementToolsOptions): Array<
+  ServerTool<SchemaInput, SchemaInput, string>
+> {
+  // Use provided strategy, or storage's strategy, or default
+  const strategy =
+    trustStrategy ?? storage.trustStrategy ?? createDefaultTrustStrategy()
+
+  // Compute bindings if not provided but config is available
+  const bindings =
+    baseBindings ?? (config ? toolsToBindings(config.tools, 'external_') : {})
+  return [
+    // Search for snippets
+    toolDefinition({
+      name: 'search_snippets',
+      description:
+        'Search the snippet library for reusable snippets. Use this to find snippets that can help accomplish a task. Returns matching snippets with their descriptions.',
+      inputSchema: z.object({
+        query: z
+          .string()
+          .describe('Search query describing what you want to accomplish'),
+        limit: z
+          .number()
+          .optional()
+          .default(5)
+          .describe('Maximum number of results (default: 5)'),
+      }),
+      outputSchema: z.array(
+        z.object({
+          name: z.string(),
+          description: z.string(),
+          usageHints: z.array(z.string()),
+          trustLevel: z.enum(['untrusted', 'provisional', 'trusted']),
+        }),
+      ),
+    }).server(async ({ query, limit }) => {
+      const results = await storage.search(query, { limit: limit ?? 5 })
+      return results.map((s) => ({
+        name: s.name,
+        description: s.description,
+        usageHints: s.usageHints,
+        trustLevel: s.trustLevel,
+      }))
+    }),
+
+    // Get full snippet details
+    toolDefinition({
+      name: 'get_snippet',
+      description:
+        'Get the full implementation details of a snippet, including its code. Use this after search_snippets to see how a snippet works before using it.',
+      inputSchema: z.object({
+        name: z.string().describe('The snippet name (without snippet_ prefix)'),
+      }),
+      outputSchema: z.object({
+        name: z.string().optional(),
+        description: z.string().optional(),
+        code: z.string().optional(),
+        inputSchema: z.string().optional().describe('JSON Schema as string'),
+        outputSchema: z.string().optional().describe('JSON Schema as string'),
+        usageHints: z.array(z.string()).optional(),
+        dependsOn: z.array(z.string()).optional(),
+        trustLevel: z.enum(['untrusted', 'provisional', 'trusted']).optional(),
+        stats: z
+          .object({
+            executions: z.number(),
+            successRate: z.number(),
+          })
+          .optional(),
+        error: z.string().optional(),
+      }),
+    }).server(async ({ name }) => {
+      const snippet = await storage.get(name)
+      if (!snippet) {
+        return { error: `Snippet '${name}' not found` }
+      }
+      return {
+        name: snippet.name,
+        description: snippet.description,
+        code: snippet.code,
+        inputSchema: JSON.stringify(snippet.inputSchema),
+        outputSchema: JSON.stringify(snippet.outputSchema),
+        usageHints: snippet.usageHints,
+        dependsOn: snippet.dependsOn,
+        trustLevel: snippet.trustLevel,
+        stats: snippet.stats,
+      }
+    }),
+
+    // Register a new snippet
+    toolDefinition({
+      name: 'register_snippet',
+      description:
+        'Save working TypeScript code as a reusable snippet for future use. Only register code that has been tested and works correctly. The snippet becomes available as a callable tool immediately.',
+      inputSchema: z.object({
+        name: z
+          .string()
+          .regex(
+            /^[a-z][a-z0-9_]*$/,
+            'Must be snake_case starting with a letter',
+          )
+          .describe(
+            'Unique snippet name in snake_case (e.g., fetch_github_stats)',
+          ),
+        description: z
+          .string()
+          .describe('Clear description of what the snippet does'),
+        code: z
+          .string()
+          .describe(
+            'The TypeScript code. Receives `input` variable, can call external_* and snippet_* functions, should return a value.',
+          ),
+        inputSchema: z
+          .string()
+          .describe(
+            'JSON Schema as a JSON string describing the input parameter, e.g. {"type":"object","properties":{"a":{"type":"number"}},"required":["a"]}',
+          ),
+        outputSchema: z
+          .string()
+          .describe(
+            'JSON Schema as a JSON string describing the return value, e.g. {"type":"object","properties":{"result":{"type":"number"}}}',
+          ),
+        usageHints: z
+          .array(z.string())
+          .describe(
+            'Hints about when to use this snippet, e.g. "Use when user asks about..."',
+          ),
+        dependsOn: z
+          .array(z.string())
+          .optional()
+          .default([])
+          .describe('Names of other snippets this snippet calls'),
+      }),
+      outputSchema: z.object({
+        success: z.boolean().optional(),
+        snippetId: z.string().optional(),
+        name: z.string().optional(),
+        message: z.string().optional(),
+        error: z.string().optional(),
+      }),
+    }).server(async (rawSnippetDef, context) => {
+      // Parse the JSON string schemas
+      let inputSchema: Record<string, unknown>
+      let outputSchema: Record<string, unknown>
+      try {
+        inputSchema = JSON.parse(rawSnippetDef.inputSchema) as Record<
+          string,
+          unknown
+        >
+      } catch {
+        return { error: 'inputSchema must be a valid JSON string' }
+      }
+      try {
+        outputSchema = JSON.parse(rawSnippetDef.outputSchema) as Record<
+          string,
+          unknown
+        >
+      } catch {
+        return { error: 'outputSchema must be a valid JSON string' }
+      }
+
+      const snippetDef = {
+        ...rawSnippetDef,
+        inputSchema,
+        outputSchema,
+      }
+      try {
+        // Validate the snippet name isn't reserved
+        if (snippetDef.name.startsWith('external_')) {
+          return { error: "Snippet names cannot start with 'external_'" }
+        }
+        if (snippetDef.name.startsWith('snippet_')) {
+          return {
+            error:
+              "Snippet names should not include the 'snippet_' prefix - it will be added automatically",
+          }
+        }
+
+        // Check if snippet already exists
+        const existing = await storage.get(snippetDef.name)
+        if (existing) {
+          return {
+            error: `Snippet '${snippetDef.name}' already exists. Use a different name or update the existing snippet.`,
+          }
+        }
+
+        // Generate a unique ID
+        const id = crypto.randomUUID()
+
+        // Get initial trust level from strategy
+        const initialTrustLevel = strategy.getInitialTrustLevel()
+
+        // Save the snippet
+        const snippet = await storage.save({
+          id,
+          name: snippetDef.name,
+          description: snippetDef.description,
+          code: snippetDef.code,
+          inputSchema: snippetDef.inputSchema,
+          outputSchema: snippetDef.outputSchema,
+          usageHints: snippetDef.usageHints,
+          dependsOn: snippetDef.dependsOn ?? [],
+          trustLevel: initialTrustLevel,
+          stats: { executions: 0, successRate: 0 },
+        })
+
+        // If registry and config are available, add the snippet as a tool immediately
+        if (registry && config) {
+          const snippetTool = snippetToTool({
+            snippet,
+            driver: config.driver,
+            bindings,
+            storage,
+            timeout: config.timeout,
+            memoryLimit: config.memoryLimit,
+          })
+          registry.add(snippetTool)
+          console.log(
+            `[register_snippet] Added snippet '${snippet.name}' to registry immediately`,
+          )
+        }
+
+        // Emit event for UI notification
+        context?.emitCustomEvent('snippet:registered', {
+          id: snippet.id,
+          name: snippet.name,
+          description: snippet.description,
+          timestamp: Date.now(),
+        })
+
+        return {
+          success: true,
+          snippetId: snippet.id,
+          name: snippet.name,
+          message: `Snippet '${snippet.name}' registered successfully and is now available as the '${snippet.name}' tool.`,
+        }
+      } catch (error) {
+        console.error('[register_snippet] Error:', error)
+        return {
+          error: `Failed to register snippet: ${error instanceof Error ? error.message : String(error)}`,
+        }
+      }
+    }),
+  ]
+}

package/src/create-snippets-system-prompt.ts

@@ -0,0 +1,289 @@
+import { generateSnippetTypes } from './generate-snippet-types'
+import type { Snippet } from './types'
+
+interface CreateSnippetsSystemPromptOptions {
+  /**
+   * Snippets that were selected for this request
+   */
+  selectedSnippets: Array<Snippet>
+
+  /**
+   * Total number of snippets in the library
+   */
+  totalSnippetCount: number
+
+  /**
+   * Whether snippets are exposed as direct tools (not just sandbox bindings)
+   * @default true
+   */
+  snippetsAsTools?: boolean
+}
+
+/**
+ * Generate example input from a JSON Schema
+ */
+function generateExampleFromSchema(schema: Record<string, unknown>): string {
+  if (schema.type === 'object' && schema.properties) {
+    const props = schema.properties as Record<string, { type: string }>
+    const example: Record<string, unknown> = {}
+
+    for (const [key, value] of Object.entries(props)) {
+      if (value.type === 'string') example[key] = `'example_${key}'`
+      else if (value.type === 'number') example[key] = 0
+      else if (value.type === 'boolean') example[key] = true
+      else if (value.type === 'array') example[key] = []
+      else example[key] = null
+    }
+
+    return JSON.stringify(example).replace(/"/g, '')
+  }
+  return '{}'
+}
+
+/**
+ * Create system prompt documentation for the snippet library.
+ * This is appended to the Code Mode system prompt.
+ */
+export function createSnippetsSystemPrompt({
+  selectedSnippets,
+  totalSnippetCount,
+  snippetsAsTools = true,
+}: CreateSnippetsSystemPromptOptions): string {
+  // No snippets in library
+  if (totalSnippetCount === 0) {
+    return `## Snippet Library
+
+You have access to a snippet library for storing reusable code. The library is currently empty.
+
+### Snippet Management Tools
+
+- \`search_snippets(query, limit?)\` - Search for snippets (currently empty)
+- \`get_snippet(name)\` - Get full snippet details including code
+- \`register_snippet(...)\` - Save working code as a reusable snippet
+
+When you write useful, reusable code, consider registering it as a snippet for future use.
+
+**Important**: Newly registered snippets become available as tools on the **next message**, not immediately in the current conversation turn.
+`
+  }
+
+  // No snippets selected for this conversation
+  if (selectedSnippets.length === 0) {
+    return `## Snippet Library
+
+You have access to a persistent snippet library with ${totalSnippetCount} snippet${totalSnippetCount === 1 ? '' : 's'}. No snippets were pre-loaded for this conversation based on context.
+
+### Snippet Management Tools
+
+- \`search_snippets(query, limit?)\` - Search for relevant snippets
+- \`get_snippet(name)\` - Get full snippet details including code
+- \`register_snippet(...)\` - Save working code as a reusable snippet
+
+When you write useful, reusable code, consider registering it as a snippet for future use.
+
+**Important**: Newly registered snippets become available as tools on the **next message**, not immediately in the current conversation turn.
+`
+  }
+
+  if (snippetsAsTools) {
+    // Snippets are available as direct tools
+    const snippetToolDocs = selectedSnippets
+      .map((snippet) => {
+        const inputExample = generateExampleFromSchema(snippet.inputSchema)
+        const trustBadge =
+          snippet.trustLevel === 'trusted'
+            ? '✓ trusted'
+            : snippet.trustLevel === 'provisional'
+              ? '◐ provisional'
+              : '○ untrusted'
+
+        return `
+### ${snippet.name} [${trustBadge}]
+
+${snippet.description}
+
+${snippet.usageHints.map((h) => `- ${h}`).join('\n')}
+
+**Input Schema:**
+\`\`\`json
+${JSON.stringify(snippet.inputSchema, null, 2)}
+\`\`\`
+
+**Output Schema:**
+\`\`\`json
+${JSON.stringify(snippet.outputSchema, null, 2)}
+\`\`\`
+
+**Example:**
+Call the \`${snippet.name}\` tool with: ${inputExample}
+`
+      })
+      .join('\n---\n')
+
+    return `## Snippet Library
+
+${selectedSnippets.length} snippet${selectedSnippets.length === 1 ? '' : 's'} pre-loaded for this conversation (${totalSnippetCount} total in library).
+
+### Available Snippet Tools
+
+These snippets are available as **direct tools** you can call (marked with [SNIPPET] in description):
+
+${snippetToolDocs}
+
+### Snippet Management Tools
+
+- \`search_snippets(query, limit?)\` - Find additional snippets not pre-loaded
+- \`get_snippet(name)\` - Get full details of any snippet
+- \`register_snippet(...)\` - Save working code as a new snippet
+
+### Using Snippets
+
+Snippets are **regular tools** - call them directly like any other tool. No need to use \`execute_typescript\`.
+
+### Creating New Snippets
+
+When you write useful, reusable code with \`execute_typescript\`, register it:
+
+\`\`\`typescript
+// After verifying code works, call the register_snippet tool
+register_snippet({
+  name: 'compare_npm_packages',
+  description: 'Compare download counts for multiple NPM packages',
+  code: \`
+    const { packages } = input;
+    const results = await Promise.all(
+      packages.map(pkg => external_getNpmDownloads({ package: pkg }))
+    );
+    return packages.map((pkg, i) => ({ package: pkg, downloads: results[i].downloads }))
+      .sort((a, b) => b.downloads - a.downloads);
+  \`,
+  inputSchema: { 
+    type: 'object', 
+    properties: { packages: { type: 'array', items: { type: 'string' } } },
+    required: ['packages']
+  },
+  outputSchema: {
+    type: 'array',
+    items: { type: 'object', properties: { package: { type: 'string' }, downloads: { type: 'number' } } }
+  },
+  usageHints: ['Use when comparing popularity of NPM packages'],
+  dependsOn: [],
+});
+\`\`\`
+
+**Important**: Newly registered snippets become available as tools on the **next message**, not immediately in the current conversation turn.
+`
+  }
+
+  // Snippets as sandbox bindings (legacy mode)
+  const snippetDocs = selectedSnippets
+    .map((snippet) => {
+      const inputExample = generateExampleFromSchema(snippet.inputSchema)
+      const trustBadge =
+        snippet.trustLevel === 'trusted'
+          ? '✓ trusted'
+          : snippet.trustLevel === 'provisional'
+            ? '◐ provisional'
+            : '○ untrusted'
+
+      return `
+### snippet_${snippet.name} [${trustBadge}]
+
+${snippet.description}
+
+${snippet.usageHints.map((h) => `- ${h}`).join('\n')}
+
+**Input Schema:**
+\`\`\`json
+${JSON.stringify(snippet.inputSchema, null, 2)}
+\`\`\`
+
+**Output Schema:**
+\`\`\`json
+${JSON.stringify(snippet.outputSchema, null, 2)}
+\`\`\`
+
+**Example:**
+\`\`\`typescript
+const result = await snippet_${snippet.name}(${inputExample});
+\`\`\`
+`
+    })
+    .join('\n---\n')
+
+  // Generate type stubs for selected snippets
+  const typeStubs = generateSnippetTypes(selectedSnippets)
+
+  return `## Snippet Library
+
+${selectedSnippets.length} snippet${selectedSnippets.length === 1 ? '' : 's'} pre-loaded for this conversation (${totalSnippetCount} total in library).
+
+### Pre-loaded Snippets
+
+These are available as \`snippet_*\` functions in your TypeScript code:
+
+${snippetDocs}
+
+### Type Definitions
+
+\`\`\`typescript
+${typeStubs}
+\`\`\`
+
+### Snippet Management Tools
+
+- \`search_snippets(query, limit?)\` - Find additional snippets not pre-loaded
+- \`get_snippet(name)\` - Get full details of any snippet
+- \`register_snippet(...)\` - Save working code as a new snippet
+
+### Using Snippets
+
+Snippets work just like \`external_*\` functions inside \`execute_typescript\`:
+
+\`\`\`typescript
+// Call a pre-loaded snippet
+const stats = await snippet_fetch_github_stats({ owner: 'tanstack', repo: 'query' });
+
+// Compose snippets with external tools
+const repos = await external_searchRepositories({ query: 'react state' });
+const detailed = await Promise.all(
+  repos.items.slice(0, 5).map(r => 
+    snippet_fetch_github_stats({ owner: r.owner.login, repo: r.name })
+  )
+);
+\`\`\`
+
+### Creating New Snippets
+
+When you write useful, reusable code, register it:
+
+\`\`\`typescript
+// After verifying code works, call the register_snippet tool
+register_snippet({
+  name: 'compare_npm_packages',
+  description: 'Compare download counts for multiple NPM packages',
+  code: \`
+    const { packages } = input;
+    const results = await Promise.all(
+      packages.map(pkg => external_getNpmDownloads({ package: pkg }))
+    );
+    return packages.map((pkg, i) => ({ package: pkg, downloads: results[i].downloads }))
+      .sort((a, b) => b.downloads - a.downloads);
+  \`,
+  inputSchema: { 
+    type: 'object', 
+    properties: { packages: { type: 'array', items: { type: 'string' } } },
+    required: ['packages']
+  },
+  outputSchema: {
+    type: 'array',
+    items: { type: 'object', properties: { package: { type: 'string' }, downloads: { type: 'number' } } }
+  },
+  usageHints: ['Use when comparing popularity of NPM packages'],
+  dependsOn: [],
+});
+\`\`\`
+
+**Important**: Newly registered snippets become available as tools on the **next message**, not immediately in the current conversation turn.
+`
+}