From 73d0f91d3c358c480eabf8f5edd336ea44d304a9 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Fri, 15 May 2026 12:49:51 +0200 Subject: [PATCH] feat: add backend-neutral agent tools --- packages/context/src/agent/agent-tool.test.ts | 39 ++++++++++++ packages/context/src/agent/agent-tool.ts | 61 +++++++++++++++++++ packages/context/src/agent/index.ts | 2 + .../src/ingest/local-bundle-runtime.ts | 9 ++- packages/context/src/ingest/ports.ts | 5 +- packages/context/src/memory/local-memory.ts | 4 +- packages/context/src/memory/types.ts | 5 +- packages/context/src/tools/base-tool.ts | 31 ++++------ 8 files changed, 123 insertions(+), 33 deletions(-) create mode 100644 packages/context/src/agent/agent-tool.test.ts create mode 100644 packages/context/src/agent/agent-tool.ts diff --git a/packages/context/src/agent/agent-tool.test.ts b/packages/context/src/agent/agent-tool.test.ts new file mode 100644 index 00000000..ad341ff6 --- /dev/null +++ b/packages/context/src/agent/agent-tool.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it, vi } from 'vitest'; +import { z } from 'zod'; +import { createAgentTool, toAiSdkTool, toAiSdkToolSet } from './agent-tool.js'; + +describe('agent tools', () => { + it('converts an agent tool to an AI SDK tool and preserves markdown output', async () => { + const execute = vi.fn(async (input: { name: string }) => ({ + markdown: `hello ${input.name}`, + structured: { ok: true }, + })); + const agentTool = createAgentTool({ + name: 'greet', + description: 'Greet someone', + inputSchema: z.object({ name: z.string() }), + execute, + }); + + const aiTool = toAiSdkTool(agentTool); + const output = await aiTool.execute?.({ name: 'Ada' }, { toolCallId: 'call-1', messages: [] } as never); + const modelOutput = aiTool.toModelOutput?.({ output } as never); + + expect(execute).toHaveBeenCalledWith({ name: 'Ada' }, { toolCallId: 'call-1' }); + expect(modelOutput).toEqual({ type: 'content', value: [{ type: 'text', text: 'hello Ada' }] }); + }); + + it('converts a named map of agent tools to an AI SDK tool set', () => { + const toolSet = toAiSdkToolSet({ + ping: createAgentTool({ + name: 'ping', + description: 'Ping', + inputSchema: z.object({}), + execute: async () => 'pong', + }), + }); + + expect(Object.keys(toolSet)).toEqual(['ping']); + expect(toolSet.ping?.description).toBe('Ping'); + }); +}); diff --git a/packages/context/src/agent/agent-tool.ts b/packages/context/src/agent/agent-tool.ts new file mode 100644 index 00000000..e454f8b7 --- /dev/null +++ b/packages/context/src/agent/agent-tool.ts @@ -0,0 +1,61 @@ +import { tool as aiTool, type Tool, type ToolSet } from 'ai'; +import { z, type ZodObject, type ZodRawShape } from 'zod'; + +export interface AgentToolCallOptions { + toolCallId?: string; +} + +export type AgentToolOutput = string | { markdown: string; structured?: unknown }; + +export interface AgentToolDefinition = ZodObject> { + name: string; + description: string; + inputSchema: TInputSchema; + execute(input: z.infer, options: AgentToolCallOptions): Promise; +} + +export type AgentToolSet = Record; + +export function createAgentTool>( + definition: AgentToolDefinition, +): AgentToolDefinition { + return definition; +} + +export function assertAgentToolSet(toolSet: AgentToolSet): void { + for (const [name, definition] of Object.entries(toolSet)) { + if (definition.name !== name) { + throw new Error(`Agent tool map key "${name}" does not match definition name "${definition.name}"`); + } + if (!(definition.inputSchema instanceof z.ZodObject)) { + throw new Error(`Agent tool "${name}" must use a Zod object input schema`); + } + } +} + +export function agentToolOutputToText(output: AgentToolOutput): string { + if (output && typeof output === 'object' && 'markdown' in output) { + return output.markdown; + } + return String(output); +} + +export function toAiSdkTool(definition: AgentToolDefinition): Tool { + return aiTool({ + description: definition.description, + inputSchema: definition.inputSchema, + execute: async (params, options) => + definition.execute(definition.inputSchema.parse(params), { + ...(options.toolCallId ? { toolCallId: options.toolCallId } : {}), + }), + toModelOutput: ({ output }) => ({ + type: 'content', + value: [{ type: 'text', text: agentToolOutputToText(output as AgentToolOutput) }], + }), + }); +} + +export function toAiSdkToolSet(toolSet: AgentToolSet): ToolSet { + assertAgentToolSet(toolSet); + return Object.fromEntries(Object.entries(toolSet).map(([name, definition]) => [name, toAiSdkTool(definition)])); +} diff --git a/packages/context/src/agent/index.ts b/packages/context/src/agent/index.ts index b4b94167..e442413b 100644 --- a/packages/context/src/agent/index.ts +++ b/packages/context/src/agent/index.ts @@ -1,3 +1,5 @@ +export type { AgentToolCallOptions, AgentToolDefinition, AgentToolOutput, AgentToolSet } from './agent-tool.js'; +export { agentToolOutputToText, assertAgentToolSet, createAgentTool, toAiSdkTool, toAiSdkToolSet } from './agent-tool.js'; export type { AgentRunnerServiceDeps, AgentTelemetryPort, diff --git a/packages/context/src/ingest/local-bundle-runtime.ts b/packages/context/src/ingest/local-bundle-runtime.ts index 047b7ee6..3372e7b8 100644 --- a/packages/context/src/ingest/local-bundle-runtime.ts +++ b/packages/context/src/ingest/local-bundle-runtime.ts @@ -2,9 +2,8 @@ import { mkdirSync } from 'node:fs'; import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; import type { KtxLlmProvider } from '@ktx/llm'; -import type { Tool } from 'ai'; import YAML from 'yaml'; -import type { AgentRunnerService } from '../agent/index.js'; +import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; import { AgentRunnerService as DefaultAgentRunnerService } from '../agent/index.js'; import { localConnectionInfoFromConfig, type KtxSqlQueryExecutorPort } from '../connections/index.js'; import type { KtxEmbeddingPort, KtxLogger } from '../core/index.js'; @@ -456,12 +455,12 @@ class NoopKnowledgeEventPort implements KnowledgeEventPort { class LocalIngestToolSet implements IngestToolsetLike { constructor( private readonly tools: BaseTool[], - private readonly sourceTools: Record = {}, + private readonly sourceTools: AgentToolSet = {}, ) {} - toAiSdkTools(context: ToolContext) { + toAgentTools(context: ToolContext) { return { - ...Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAiSdkTool(context)])), + ...Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAgentTool(context)])), ...this.sourceTools, }; } diff --git a/packages/context/src/ingest/ports.ts b/packages/context/src/ingest/ports.ts index fbb2451e..d6f93211 100644 --- a/packages/context/src/ingest/ports.ts +++ b/packages/context/src/ingest/ports.ts @@ -1,6 +1,5 @@ -import type { ToolSet } from 'ai'; import type { KtxModelRole } from '@ktx/llm'; -import type { AgentRunnerService } from '../agent/index.js'; +import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; import type { KtxEmbeddingPort } from '../core/embedding.js'; import type { GitService, KtxFileStorePort, KtxLogger, SessionOutcome } from '../core/index.js'; import type { CaptureSession, MemoryAction, MemoryKnowledgeSlRefsPort } from '../memory/index.js'; @@ -163,7 +162,7 @@ export interface IngestCommitMessagePort { } export interface IngestToolsetLike { - toAiSdkTools(context: ToolContext): ToolSet; + toAgentTools(context: ToolContext): AgentToolSet; } export interface IngestToolsetFactoryPort { diff --git a/packages/context/src/memory/local-memory.ts b/packages/context/src/memory/local-memory.ts index 3cc9d324..978898d3 100644 --- a/packages/context/src/memory/local-memory.ts +++ b/packages/context/src/memory/local-memory.ts @@ -386,8 +386,8 @@ class LocalShapeOnlySlValidator implements SlValidatorPort { class LocalMemoryToolSet implements MemoryToolSetLike { constructor(private readonly tools: BaseTool[]) {} - toAiSdkTools(context: ToolContext) { - return Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAiSdkTool(context)])); + toAgentTools(context: ToolContext) { + return Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAgentTool(context)])); } } diff --git a/packages/context/src/memory/types.ts b/packages/context/src/memory/types.ts index 207eb238..133b816f 100644 --- a/packages/context/src/memory/types.ts +++ b/packages/context/src/memory/types.ts @@ -1,5 +1,4 @@ -import type { Tool } from 'ai'; -import type { AgentRunnerService } from '../agent/index.js'; +import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; import type { GitService, KtxFileStorePort, KtxLogger, SessionWorktreeService } from '../core/index.js'; import type { PromptService } from '../prompts/index.js'; import type { SkillsRegistryService } from '../skills/index.js'; @@ -118,7 +117,7 @@ export interface MemoryCommitMessagePort { export interface MemoryFileStorePort extends KtxFileStorePort, MemoryCommitMessagePort {} export interface MemoryToolSetLike { - toAiSdkTools(context: ToolContext): Record; + toAgentTools(context: ToolContext): AgentToolSet; } export interface MemoryToolsetFactoryPort { diff --git a/packages/context/src/tools/base-tool.ts b/packages/context/src/tools/base-tool.ts index 0566a0ca..a76ca735 100644 --- a/packages/context/src/tools/base-tool.ts +++ b/packages/context/src/tools/base-tool.ts @@ -1,5 +1,5 @@ -import { tool } from 'ai'; import { z, type ZodType } from 'zod'; +import { createAgentTool, toAiSdkTool, type AgentToolDefinition } from '../agent/agent-tool.js'; import { noopLogger, type KtxLogger } from '../core/index.js'; import type { IngestToolMetadata, ToolSession } from './tool-session.js'; @@ -114,18 +114,16 @@ export abstract class BaseTool { }; } - toAiSdkTool(context: ToolContext): any { + toAgentTool(context: ToolContext): AgentToolDefinition { const toolName = this.name; - const logger = this.logger; - return tool({ + return createAgentTool({ + name: toolName, description: this.description, - inputSchema: this.inputSchema, + inputSchema: this.inputSchema as any, execute: async (params, { toolCallId }) => { - // Create context copy with current toolCallId (safe for parallel execution) - const callContext = { ...context, toolCallId }; + const callContext = { ...context, ...(toolCallId ? { toolCallId } : {}) }; - // Record tool execution start (input generation has already been tracked via onChunk) if (callContext.timingTracker && toolCallId) { callContext.timingTracker.recordToolExecutionStart(callContext.messageId, toolName, toolCallId); } @@ -136,8 +134,7 @@ export abstract class BaseTool { throw new Error('Authentication required: userId must be provided in ToolContext'); } const parsedInput = this.parseInput(params as Record); - const result = await this.call(parsedInput, callContext); - return result; + return await this.call(parsedInput, callContext); } catch (error) { state = 'error'; this.logger.error( @@ -151,19 +148,13 @@ export abstract class BaseTool { } } }, - // Send only markdown to the LLM; tool callers still receive the structured output. - toModelOutput: ({ output }) => { - if (output && typeof output === 'object' && 'markdown' in output) { - return { type: 'content', value: [{ type: 'text', text: output.markdown as string }] }; - } - if (typeof output !== 'string') { - logger.warn(`Tool ${toolName} returned unexpected output type: ${typeof output}. Coercing to string.`); - } - return { type: 'content', value: [{ type: 'text', text: String(output) }] }; - }, }); } + toAiSdkTool(context: ToolContext): any { + return toAiSdkTool(this.toAgentTool(context)); + } + parseInput(input: Record): z.infer { return this.inputSchema.parse(input); }