Tool context injection (#10)

* feat: inject server-only tool context to keep auth out of LLM inputs

* chore: bump version to 0.3.0-alpha.1

Author: dayhaysoos

README.md

@@ -1100,6 +1100,36 @@ enforce access control in your app.
 
 ---
 
+## Injecting trusted context into tools
+
+Some tool handlers need server-only context (orgId, externalId, userId) that
+should never appear in the tool schema or LLM prompt. Pass that data via
+`toolContext` in `chat.send` or `chat.sendForExternalId`. The context is merged
+into tool args after validation, and `toolContext` wins on conflicts.
+
+```typescript
+await ctx.runAction(components.databaseChat.chat.sendForExternalId, {
+  conversationId,
+  externalId,
+  message,
+  config: {
+    apiKey: process.env.OPENROUTER_API_KEY!,
+    systemPrompt: SYSTEM_PROMPT,
+    tools,
+    toolContext: {
+      orgId,
+      userId,
+      externalId,
+    },
+  },
+});
+```
+
+Note: Do not include fields like `orgId` in the tool schema so the LLM never
+sees them.
+
+---
+
 ## API Reference
 
 ### Component Functions

convex/component/chat.test.ts

@@ -3,6 +3,8 @@ import { describe, it, expect } from "vitest";
 import { convexTest } from "convex-test";
 import schema from "./schema";
 import { api } from "./_generated/api";
+import { executeToolWithContext } from "./toolExecution";
+import type { DatabaseChatTool } from "./tools";
 
 const modules = import.meta.glob("./**/*.ts");
 
@@ -134,6 +136,47 @@ describe("databaseChat chat", () => {
     });
   });
 
+  describe("tool context injection", () => {
+    it("merges toolContext into tool args before execution", async () => {
+      const ctx = {
+        runQuery: async (_handler: string, args: Record<string, unknown>) => args,
+        runMutation: async (
+          _handler: string,
+          args: Record<string, unknown>
+        ) => args,
+        runAction: async (_handler: string, args: Record<string, unknown>) =>
+          args,
+      };
+
+      const tool: DatabaseChatTool = {
+        name: "searchRecords",
+        description: "Search records",
+        parameters: {
+          type: "object",
+          properties: {
+            query: { type: "string" },
+          },
+          required: ["query"],
+        },
+        handler: "testHandler",
+      };
+
+      const { result, args } = await executeToolWithContext(
+        ctx,
+        tool,
+        { query: "react", orgId: "llm-org" },
+        { orgId: "org123", externalId: "user:1" }
+      );
+
+      expect(result).toEqual({
+        query: "react",
+        orgId: "org123",
+        externalId: "user:1",
+      });
+      expect(args).toEqual(result);
+    });
+  });
+
   // TODO: Add integration test with mocked fetch for chat.send
   // This would require setting up MSW or similar to mock OpenRouter responses
   describe.skip("chat.send integration", () => {

convex/component/chat.ts

@@ -10,6 +10,7 @@ import {
 } from "./tools";
 import type { DatabaseChatTool } from "./tools";
 import { DeltaStreamer } from "./deltaStreamer";
+import { executeToolWithContext } from "./toolExecution";
 
 const sendConfigValidator = v.object({
   apiKey: v.string(),
@@ -19,6 +20,7 @@ const sendConfigValidator = v.object({
   tools: v.optional(v.array(databaseChatToolValidator)),
   // Max messages to include in LLM context (default: 50)
   maxMessagesForLLM: v.optional(v.number()),
+  toolContext: v.optional(v.any()),
 });
 
 const sendReturnValidator = v.object({
@@ -98,6 +100,7 @@ async function sendInternal(
         handler: string;
       }>;
       maxMessagesForLLM?: number;
+      toolContext?: Record<string, unknown>;
     };
   }
 ) {
@@ -207,14 +210,19 @@ async function sendInternal(
 
         // Execute the tool
         try {
-          const result = await executeToolHandler(ctx, tool, parsedArgs);
+          const { result, args: mergedArgs } = await executeToolWithContext(
+            ctx,
+            tool,
+            parsedArgs,
+            config.toolContext
+          );
           toolResults.push({
             toolCallId: toolCall.id,
             result: JSON.stringify(result),
           });
           executedToolCalls.push({
             name: toolCall.name,
-            args: parsedArgs,
+            args: mergedArgs,
             result,
           });
         } catch (error) {
@@ -294,23 +302,6 @@ async function sendInternal(
   }
 }
 
-async function executeToolHandler(
-  ctx: any,
-  tool: DatabaseChatTool,
-  args: Record<string, unknown>
-) {
-  const handlerType = tool.handlerType ?? "query";
-  switch (handlerType) {
-    case "mutation":
-      return await ctx.runMutation(tool.handler as any, args);
-    case "action":
-      return await ctx.runAction(tool.handler as any, args);
-    case "query":
-    default:
-      return await ctx.runQuery(tool.handler as any, args);
-  }
-}
-
 const DEFAULT_SYSTEM_PROMPT = `You are a helpful assistant that can search and query a database.
 When users ask questions, use the available tools to find relevant information.
 If you don't have access to a tool that can answer the question, say so.

convex/component/client.ts

@@ -133,6 +133,8 @@ export interface SendMessageOptions {
   model?: string;
   /** Override system prompt for this message */
   systemPrompt?: string;
+  /** Server-side context merged into tool args (not exposed to LLM) */
+  toolContext?: Record<string, unknown>;
 }
 
 export interface SendMessageResult {
@@ -349,6 +351,7 @@ export class DatabaseChatClient {
         systemPrompt: options.systemPrompt ?? this.config.systemPrompt,
         tools: this.tools.length > 0 ? this.tools : undefined,
         maxMessagesForLLM: this.config.maxMessagesForLLM ?? 50,
+        toolContext: options.toolContext,
       },
     });
   }

convex/component/toolExecution.ts

@@ -0,0 +1,45 @@
+import type { DatabaseChatTool } from "./tools";
+
+export type ToolExecutionContext = {
+  runQuery: (handler: any, args: Record<string, unknown>) => Promise<unknown>;
+  runMutation: (handler: any, args: Record<string, unknown>) => Promise<unknown>;
+  runAction: (handler: any, args: Record<string, unknown>) => Promise<unknown>;
+};
+
+export function mergeToolArgs(
+  parsedArgs: Record<string, unknown>,
+  toolContext?: Record<string, unknown>
+): Record<string, unknown> {
+  if (!toolContext || Object.keys(toolContext).length === 0) {
+    return { ...parsedArgs };
+  }
+  return { ...parsedArgs, ...toolContext };
+}
+
+export async function executeToolWithContext(
+  ctx: ToolExecutionContext,
+  tool: DatabaseChatTool,
+  parsedArgs: Record<string, unknown>,
+  toolContext?: Record<string, unknown>
+): Promise<{ result: unknown; args: Record<string, unknown> }> {
+  const mergedArgs = mergeToolArgs(parsedArgs, toolContext);
+  const result = await executeToolHandler(ctx, tool, mergedArgs);
+  return { result, args: mergedArgs };
+}
+
+async function executeToolHandler(
+  ctx: ToolExecutionContext,
+  tool: DatabaseChatTool,
+  args: Record<string, unknown>
+) {
+  const handlerType = tool.handlerType ?? "query";
+  switch (handlerType) {
+    case "mutation":
+      return await ctx.runMutation(tool.handler as any, args);
+    case "action":
+      return await ctx.runAction(tool.handler as any, args);
+    case "query":
+    default:
+      return await ctx.runQuery(tool.handler as any, args);
+  }
+}

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "@dayhaysoos/convex-database-chat",
-  "version": "0.3.0-alpha.0",
+  "version": "0.3.0-alpha.1",
   "description": "A Convex component for adding natural language database queries to your app.",
   "type": "module",
   "main": "./dist/src/index.js",