From a66e68a653b5fb48e14aa6150d6a0fd343b452d7 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Fri, 15 May 2026 12:55:46 +0200 Subject: [PATCH] feat: add agent runner port --- .../src/agent/agent-runner.service.test.ts | 49 ++++++++++++++++--- .../context/src/agent/agent-runner.service.ts | 16 ++++-- packages/context/src/agent/index.ts | 1 + .../curator-pagination.service.ts | 4 +- .../src/ingest/local-bundle-runtime.ts | 6 +-- packages/context/src/ingest/local-ingest.ts | 6 +-- packages/context/src/ingest/ports.ts | 4 +- .../src/ingest/stages/stage-3-work-units.ts | 4 +- .../ingest/stages/stage-4-reconciliation.ts | 4 +- packages/context/src/memory/local-memory.ts | 4 +- packages/context/src/memory/types.ts | 4 +- 11 files changed, 73 insertions(+), 29 deletions(-) diff --git a/packages/context/src/agent/agent-runner.service.test.ts b/packages/context/src/agent/agent-runner.service.test.ts index 3208bda7..d1af53f5 100644 --- a/packages/context/src/agent/agent-runner.service.test.ts +++ b/packages/context/src/agent/agent-runner.service.test.ts @@ -7,6 +7,8 @@ vi.mock('ai', () => ({ })); import { generateText } from 'ai'; +import { z } from 'zod'; +import { createAgentTool } from './agent-tool.js'; import { AgentRunnerService, type RunLoopStepInfo } from './agent-runner.service.js'; describe('AgentRunnerService.runLoop', () => { @@ -42,7 +44,14 @@ describe('AgentRunnerService.runLoop', () => { (generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] }); const repairHandler = vi.fn(); llmProvider.repairToolCallHandler.mockReturnValueOnce(repairHandler); - const tools = { noop: { description: 'noop', inputSchema: {}, execute: vi.fn() } }; + const tools = { + noop: createAgentTool({ + name: 'noop', + description: 'noop', + inputSchema: z.object({}), + execute: vi.fn(async () => 'ok'), + }), + }; await runner.runLoop({ modelRole: 'candidateExtraction', systemPrompt: 'SYS', @@ -55,7 +64,7 @@ describe('AgentRunnerService.runLoop', () => { expect(call.system).toEqual({ role: 'system', content: 'SYS' }); expect(call.messages).toEqual([{ role: 'user', content: 'USR' }]); expect(call.prompt).toBeUndefined(); - expect(call.tools).toEqual(tools); + expect(call.tools).toMatchObject({ noop: { description: 'noop' } }); expect(call.stopWhen).toBe(17); expect(call.temperature).toBe(0); expect(call.experimental_repairToolCall).toBe(repairHandler); @@ -63,6 +72,33 @@ describe('AgentRunnerService.runLoop', () => { expect(llmProvider.repairToolCallHandler).toHaveBeenCalledWith({ source: 'ktx-agent-runner' }); }); + it('converts AgentToolSet to AI SDK tools before generateText', async () => { + (generateText as any).mockResolvedValue({} as never); + await runner.runLoop({ + modelRole: 'default', + systemPrompt: 'system', + userPrompt: 'user', + toolSet: { + emit_candidate: createAgentTool({ + name: 'emit_candidate', + description: 'Emit candidate', + inputSchema: z.object({ key: z.string() }), + execute: async ({ key }) => ({ markdown: key, structured: { key } }), + }), + }, + stepBudget: 3, + telemetryTags: {}, + }); + + expect(generateText).toHaveBeenCalledWith( + expect.objectContaining({ + tools: expect.objectContaining({ + emit_candidate: expect.objectContaining({ description: 'Emit candidate' }), + }), + }), + ); + }); + it('returns stopReason=natural when the loop completes without error', async () => { (generateText as any).mockResolvedValue({ text: 'done', toolCalls: [], steps: [] }); const result = await runner.runLoop({ @@ -289,11 +325,12 @@ describe('AgentRunnerService.runLoop', () => { systemPrompt: 'SECRET SYSTEM PROMPT', userPrompt: 'SECRET USER PROMPT', toolSet: { - emit_candidate: { + emit_candidate: createAgentTool({ + name: 'emit_candidate', description: 'SECRET TOOL DESCRIPTION', - inputSchema: {}, - execute: vi.fn(), - } as any, + inputSchema: z.object({}), + execute: vi.fn(async () => 'ok'), + }), }, stepBudget: 10, telemetryTags: { operationName: 'ingest-bundle-wu', source: 'metabase', jobId: 'job-1', unitKey: 'cards/1' }, diff --git a/packages/context/src/agent/agent-runner.service.ts b/packages/context/src/agent/agent-runner.service.ts index 128818f9..aec2b9cf 100644 --- a/packages/context/src/agent/agent-runner.service.ts +++ b/packages/context/src/agent/agent-runner.service.ts @@ -1,7 +1,8 @@ import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider, type KtxModelRole } from '@ktx/llm'; -import { generateText, stepCountIs, type TelemetrySettings, type Tool } from 'ai'; +import { generateText, stepCountIs, type TelemetrySettings, type ToolSet } from 'ai'; import { noopLogger, type KtxLogger } from '../core/index.js'; import { summarizeKtxLlmDebugRequest, type KtxLlmDebugRequestRecorder } from '../llm/index.js'; +import { toAiSdkToolSet, type AgentToolSet } from './agent-tool.js'; export type RunLoopStopReason = 'budget' | 'natural' | 'error'; @@ -14,7 +15,7 @@ export interface RunLoopParams { modelRole: KtxModelRole; systemPrompt: string; userPrompt: string; - toolSet: Record; + toolSet: AgentToolSet; stepBudget: number; telemetryTags: Record; onStepFinish?: (info: RunLoopStepInfo) => void | Promise; @@ -25,6 +26,10 @@ export interface RunLoopResult { error?: Error; } +export interface AgentRunnerPort { + runLoop(params: RunLoopParams): Promise; +} + export interface AgentTelemetryPort { createTelemetry(tags: Record): TelemetrySettings; } @@ -36,7 +41,7 @@ export interface AgentRunnerServiceDeps { logger?: KtxLogger; } -export class AgentRunnerService { +export class AgentRunnerService implements AgentRunnerPort { private readonly logger: KtxLogger; constructor(private readonly deps: AgentRunnerServiceDeps) { @@ -47,11 +52,12 @@ export class AgentRunnerService { let stepIndex = 0; try { const model = this.deps.llmProvider.getModel(params.modelRole); + const aiToolSet = toAiSdkToolSet(params.toolSet); const builder = new KtxMessageBuilder(this.deps.llmProvider); const built = builder.wrapSimple({ system: params.systemPrompt, messages: [{ role: 'user', content: params.userPrompt }], - tools: params.toolSet, + tools: aiToolSet, model, }); const promptMessages = splitKtxSystemMessages(built.messages); @@ -79,7 +85,7 @@ export class AgentRunnerService { }), ...(promptMessages.system ? { system: promptMessages.system } : {}), messages: promptMessages.messages, - tools: built.tools as Record, + tools: built.tools as ToolSet, onStepFinish: async () => { stepIndex += 1; if (!params.onStepFinish) { diff --git a/packages/context/src/agent/index.ts b/packages/context/src/agent/index.ts index e442413b..675828e5 100644 --- a/packages/context/src/agent/index.ts +++ b/packages/context/src/agent/index.ts @@ -1,6 +1,7 @@ export type { AgentToolCallOptions, AgentToolDefinition, AgentToolOutput, AgentToolSet } from './agent-tool.js'; export { agentToolOutputToText, assertAgentToolSet, createAgentTool, toAiSdkTool, toAiSdkToolSet } from './agent-tool.js'; export type { + AgentRunnerPort, AgentRunnerServiceDeps, AgentTelemetryPort, RunLoopParams, diff --git a/packages/context/src/ingest/context-candidates/curator-pagination.service.ts b/packages/context/src/ingest/context-candidates/curator-pagination.service.ts index 95ba5a1b..cd05f90e 100644 --- a/packages/context/src/ingest/context-candidates/curator-pagination.service.ts +++ b/packages/context/src/ingest/context-candidates/curator-pagination.service.ts @@ -1,5 +1,5 @@ import type { KtxModelRole } from '@ktx/llm'; -import type { AgentRunnerService, AgentToolSet } from '../../agent/index.js'; +import type { AgentRunnerPort, AgentToolSet } from '../../agent/index.js'; import { type KtxLogger, noopLogger } from '../../core/index.js'; import type { MemoryAction } from '../../memory/index.js'; import type { ContextCandidateForDedup, CuratorPaginationPort, CuratorPaginationReport } from '../ports.js'; @@ -49,7 +49,7 @@ interface CuratorPaginationResult extends ReconciliationOutcome { export interface CuratorPaginationServiceDeps { store: ContextCandidateStorePort; - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; settings: CuratorPaginationSettings; logger?: KtxLogger; } diff --git a/packages/context/src/ingest/local-bundle-runtime.ts b/packages/context/src/ingest/local-bundle-runtime.ts index 8165deee..ddbf1b0c 100644 --- a/packages/context/src/ingest/local-bundle-runtime.ts +++ b/packages/context/src/ingest/local-bundle-runtime.ts @@ -3,7 +3,7 @@ import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; import type { KtxLlmProvider } from '@ktx/llm'; import YAML from 'yaml'; -import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; +import type { AgentRunnerPort, 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'; @@ -99,7 +99,7 @@ const LOCAL_SHAPE_WARNING = 'Local ingest validates semantic-layer YAML shape on export interface CreateLocalBundleIngestRuntimeOptions { project: KtxLocalProject; adapters: SourceAdapter[]; - agentRunner?: AgentRunnerService; + agentRunner?: AgentRunnerPort; llmProvider?: KtxLlmProvider; llmDebugRequestFile?: string; memoryModel?: string; @@ -577,7 +577,7 @@ function localIngestLlmProviderGuardMessage(projectDir: string): string { } function resolveAgentRunner(options: CreateLocalBundleIngestRuntimeOptions): { - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; llmProvider?: KtxLlmProvider; } { const llmProvider = diff --git a/packages/context/src/ingest/local-ingest.ts b/packages/context/src/ingest/local-ingest.ts index 6056f6ed..526183ec 100644 --- a/packages/context/src/ingest/local-ingest.ts +++ b/packages/context/src/ingest/local-ingest.ts @@ -2,7 +2,7 @@ import { randomUUID } from 'node:crypto'; import { cp, mkdir, rm } from 'node:fs/promises'; import { isAbsolute, resolve } from 'node:path'; import type { KtxLlmProvider } from '@ktx/llm'; -import type { AgentRunnerService } from '../agent/index.js'; +import type { AgentRunnerPort } from '../agent/index.js'; import type { KtxSqlQueryExecutorPort } from '../connections/index.js'; import type { KtxLogger } from '../core/index.js'; import type { KtxSemanticLayerComputePort } from '../daemon/index.js'; @@ -28,7 +28,7 @@ export interface RunLocalIngestOptions { trigger?: IngestTrigger; jobId?: string; memoryFlow?: MemoryFlowEventSink; - agentRunner?: AgentRunnerService; + agentRunner?: AgentRunnerPort; llmProvider?: KtxLlmProvider; llmDebugRequestFile?: string; memoryModel?: string; @@ -167,7 +167,7 @@ async function runScheduledPullJob(options: { trigger?: IngestTrigger; jobId?: string; memoryFlow?: MemoryFlowEventSink; - agentRunner?: AgentRunnerService; + agentRunner?: AgentRunnerPort; llmProvider?: KtxLlmProvider; memoryModel?: string; semanticLayerCompute?: KtxSemanticLayerComputePort; diff --git a/packages/context/src/ingest/ports.ts b/packages/context/src/ingest/ports.ts index a1b75bc0..15e387e9 100644 --- a/packages/context/src/ingest/ports.ts +++ b/packages/context/src/ingest/ports.ts @@ -1,5 +1,5 @@ import type { KtxModelRole } from '@ktx/llm'; -import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; +import type { AgentRunnerPort, 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'; @@ -349,7 +349,7 @@ export interface IngestBundleRunnerDeps { registry: SourceAdapterRegistryPort; diffSetService: DiffSetComputerPort; sessionWorktreeService: IngestSessionWorktreePort; - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; gitService: GitService; lockingService: IngestLockPort; storage: IngestStoragePort; diff --git a/packages/context/src/ingest/stages/stage-3-work-units.ts b/packages/context/src/ingest/stages/stage-3-work-units.ts index 3e5c7405..28c88829 100644 --- a/packages/context/src/ingest/stages/stage-3-work-units.ts +++ b/packages/context/src/ingest/stages/stage-3-work-units.ts @@ -1,4 +1,4 @@ -import type { AgentRunnerService, AgentToolSet } from '@ktx/context/agent'; +import type { AgentRunnerPort, AgentToolSet } from '@ktx/context/agent'; import type { KtxModelRole } from '@ktx/llm'; import type { CaptureSession, MemoryAction } from '../../memory/index.js'; import { listTouchedSlSources, type TouchedSlSource } from '../../tools/index.js'; @@ -13,7 +13,7 @@ export interface TouchedValidationResult { export interface WorkUnitExecutionDeps { sessionWorktreeGit: { revParseHead(): Promise }; - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; validateTouchedSources: (touched: TouchedSlSource[]) => Promise; resetHardTo: (targetSha: string) => Promise; buildSystemPrompt: (wu: WorkUnit) => string; diff --git a/packages/context/src/ingest/stages/stage-4-reconciliation.ts b/packages/context/src/ingest/stages/stage-4-reconciliation.ts index 0f43f3c4..09658d26 100644 --- a/packages/context/src/ingest/stages/stage-4-reconciliation.ts +++ b/packages/context/src/ingest/stages/stage-4-reconciliation.ts @@ -1,4 +1,4 @@ -import type { AgentRunnerService, AgentToolSet } from '@ktx/context/agent'; +import type { AgentRunnerPort, AgentToolSet } from '@ktx/context/agent'; import type { KtxModelRole } from '@ktx/llm'; import type { EvictionUnit } from '../types.js'; import type { StageIndex } from './stage-index.types.js'; @@ -6,7 +6,7 @@ import type { StageIndex } from './stage-index.types.js'; export interface ReconciliationContext { stageIndex: StageIndex; evictionUnit: EvictionUnit | undefined; - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; buildSystemPrompt: (idx: StageIndex, ev: EvictionUnit | undefined) => string; buildUserPrompt: (idx: StageIndex, ev: EvictionUnit | undefined) => string; buildToolSet: () => AgentToolSet; diff --git a/packages/context/src/memory/local-memory.ts b/packages/context/src/memory/local-memory.ts index 978898d3..4c15082c 100644 --- a/packages/context/src/memory/local-memory.ts +++ b/packages/context/src/memory/local-memory.ts @@ -2,7 +2,7 @@ import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; import type { KtxLlmProvider } from '@ktx/llm'; import YAML from 'yaml'; -import { AgentRunnerService } from '../agent/index.js'; +import { AgentRunnerService, type AgentRunnerPort } from '../agent/index.js'; import { localConnectionInfoFromConfig } from '../connections/index.js'; import type { KtxEmbeddingPort, KtxFileStorePort, KtxFileWriteResult } from '../core/index.js'; import { type KtxLogger, noopLogger, SessionWorktreeService } from '../core/index.js'; @@ -64,7 +64,7 @@ const LOCAL_SHAPE_WARNING = 'Local memory capture validates semantic-layer YAML export interface CreateLocalProjectMemoryCaptureOptions { llmProvider?: KtxLlmProvider; - agentRunner?: AgentRunnerService; + agentRunner?: AgentRunnerPort; memoryModel?: string; semanticLayerCompute?: KtxSemanticLayerComputePort; queryExecutor?: { execute(input: { connectionId: string; sql: string; maxRows?: number }): Promise }; diff --git a/packages/context/src/memory/types.ts b/packages/context/src/memory/types.ts index 133b816f..ef622a8f 100644 --- a/packages/context/src/memory/types.ts +++ b/packages/context/src/memory/types.ts @@ -1,4 +1,4 @@ -import type { AgentRunnerService, AgentToolSet } from '../agent/index.js'; +import type { AgentRunnerPort, 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'; @@ -149,7 +149,7 @@ export interface MemoryAgentServiceDeps { slSourcesRepository: SlSourcesIndexPort; sessionWorktreeService: SessionWorktreeService; semanticLayerSourceReconciler: MemorySlSourceReconcilerPort; - agentRunner: AgentRunnerService; + agentRunner: AgentRunnerPort; slValidator: SlValidatorPort; toolsetFactory: MemoryToolsetFactoryPort; telemetry?: MemoryTelemetryPort;