apunkt/ktx
1
0
Fork 0
mirror of https://github.com/Kaelio/ktx.git synced 2026-06-25 08:48:08 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat: add claude-code llm backend with runtime port (#115)

Browse source 
* docs: revise claude-code ingest backend spec

* docs: keep claude-code spec focused on ingest

* docs: expand claude-code spec to full llm parity

* Refine claude-code backend spec after adversarial review iteration 1

* Refine claude-code backend spec after adversarial review iteration 2

* Refine claude-code backend spec after adversarial review iteration 3

* feat: recognize claude-code llm backend

* feat: add ktx llm runtime port

* feat: add claude-code llm runtime

* feat: route non-agent llm calls through runtime

* feat: run ingest agents through llm runtime

* feat: support claude-code setup and status

* test: verify claude-code backend runtime

* docs: add claude-code backend v1 runtime plan

* fix: close claude-code runtime isolation checks

* fix: warn on claude-code prompt caching during setup

* chore: verify claude-code v1 closure

* docs: add claude-code backend v1 isolation closure plan

* fix: update claude-code ingest setup guidance

* docs: add claude-code backend v1 ingest guidance closure plan

* docs: align claude-code isolation spec with sdk metadata

* test: cover claude-code host discovery metadata

* fix: tolerate claude-code host discovery metadata

* docs: clarify claude-code host discovery metadata

* docs: add claude-code auth-probe isolation fix plan

* chore: prepare kaelio ktx rc1 release

* chore: add semantic release workflow

* fix: unblock ci checks

* chore(release): 0.1.0-rc.1

* feat: add Claude Code model selection to setup

* fix: keep git maintenance attached in local repos
This commit is contained in:
Andrey Avtomonov 2026-05-16 12:06:34 +02:00 • committed by GitHub
parent e6d578c03f
commit b565e44a22
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
109 changed files with 10218 additions and 1093 deletions
Download patch file Download diff file Expand all files Collapse all files

9
packages/context/src/agent/agent-runner.service.test.ts
View file

@ -55,7 +55,14 @@ 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.noop).toEqual(
expect.objectContaining({
description: 'noop',
inputSchema: {},
execute: expect.any(Function),
toModelOutput: expect.any(Function),
}),
);
expect(call.stopWhen).toBe(17);
expect(call.temperature).toBe(0);
expect(call.experimental_repairToolCall).toBe(repairHandler);

111
packages/context/src/agent/agent-runner.service.ts
View file

@ -1,33 +1,15 @@
import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider, type KtxModelRole } from '@ktx/llm';
import { generateText, stepCountIs, type TelemetrySettings, type Tool } from 'ai';
import { noopLogger, type KtxLogger } from '../core/index.js';
import { summarizeKtxLlmDebugRequest, type KtxLlmDebugRequestRecorder } from '../llm/index.js';
export type RunLoopStopReason = 'budget' | 'natural' | 'error';
export interface RunLoopStepInfo {
stepIndex: number;
stepBudget: number;
}
export interface RunLoopParams {
modelRole: KtxModelRole;
systemPrompt: string;
userPrompt: string;
toolSet: Record<string, Tool>;
stepBudget: number;
telemetryTags: Record<string, string>;
onStepFinish?: (info: RunLoopStepInfo) => void | Promise<void>;
}
export interface RunLoopResult {
stopReason: RunLoopStopReason;
error?: Error;
}
export interface AgentTelemetryPort {
createTelemetry(tags: Record<string, string>): TelemetrySettings;
}
import type { KtxLlmProvider } from '@ktx/llm';
import type { KtxLogger } from '../core/index.js';
import { AiSdkKtxLlmRuntime, type AgentTelemetryPort } from '../llm/ai-sdk-runtime.js';
import type { KtxLlmDebugRequestRecorder } from '../llm/debug-request-recorder.js';
import type { AgentRunnerPort, RunLoopParams, RunLoopResult } from '../llm/runtime-port.js';
export type {
RunLoopParams,
RunLoopResult,
RunLoopStepInfo,
RunLoopStopReason,
} from '../llm/runtime-port.js';
export type { AgentTelemetryPort } from '../llm/ai-sdk-runtime.js';
export interface AgentRunnerServiceDeps {
llmProvider: KtxLlmProvider;
@ -36,71 +18,14 @@ export interface AgentRunnerServiceDeps {
logger?: KtxLogger;
}
export class AgentRunnerService {
private readonly logger: KtxLogger;
export class AgentRunnerService implements AgentRunnerPort {
private readonly runtime: AiSdkKtxLlmRuntime;
constructor(private readonly deps: AgentRunnerServiceDeps) {
this.logger = deps.logger ?? noopLogger;
constructor(deps: AgentRunnerServiceDeps) {
this.runtime = new AiSdkKtxLlmRuntime(deps);
}
async runLoop(params: RunLoopParams): Promise<RunLoopResult> {
let stepIndex = 0;
try {
const model = this.deps.llmProvider.getModel(params.modelRole);
const builder = new KtxMessageBuilder(this.deps.llmProvider);
const built = builder.wrapSimple({
system: params.systemPrompt,
messages: [{ role: 'user', content: params.userPrompt }],
tools: params.toolSet,
model,
});
const promptMessages = splitKtxSystemMessages(built.messages);
await this.deps.debugRequestRecorder?.record(
summarizeKtxLlmDebugRequest({
operationName: params.telemetryTags.operationName ?? 'ktx-agent-runner',
source: params.telemetryTags.source,
jobId: params.telemetryTags.jobId,
unitKey: params.telemetryTags.unitKey,
modelRole: params.modelRole,
modelId: (model as { modelId?: string }).modelId ?? params.modelRole,
messages: built.messages,
tools: built.tools as Record<string, { providerOptions?: unknown }>,
}),
);
await generateText({
model,
temperature: 0,
stopWhen: stepCountIs(params.stepBudget),
experimental_telemetry: this.deps.telemetry?.createTelemetry(params.telemetryTags),
experimental_repairToolCall: this.deps.llmProvider.repairToolCallHandler({
source: params.telemetryTags.operationName ?? 'ktx-agent-runner',
}),
...(promptMessages.system ? { system: promptMessages.system } : {}),
messages: promptMessages.messages,
tools: built.tools as Record<string, Tool>,
onStepFinish: async () => {
stepIndex += 1;
if (!params.onStepFinish) {
return;
}
try {
await params.onStepFinish({ stepIndex, stepBudget: params.stepBudget });
} catch (err) {
this.logger.warn(
`[agent-runner] onStepFinish callback threw; ignoring: ${
err instanceof Error ? err.message : String(err)
}`,
);
}
},
});
return { stopReason: 'natural' };
} catch (error) {
const err = error instanceof Error ? error : new Error(String(error));
this.logger.warn(`[agent-runner] loop failed: ${err.message}`);
return { stopReason: 'error', error: err };
}
runLoop(params: RunLoopParams): Promise<RunLoopResult> {
return this.runtime.runAgentLoop(params);
}
}

9
packages/context/src/core/git.service.test.ts
View file

@ -1,4 +1,4 @@
import { mkdtemp, realpath, rm, writeFile } from 'node:fs/promises';
import { mkdtemp, readFile, realpath, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
@ -52,6 +52,13 @@ describe('GitService', () => {
const after = await service.revParseHead();
expect(after).toBe(before);
});
it('keeps git auto-maintenance attached for deterministic cleanup', async () => {
const config = await readFile(join(tempDir, '.git', 'config'), 'utf-8');
expect(config).toMatch(/\[gc]\n\s+autoDetach = false/);
expect(config).toMatch(/\[maintenance]\n\s+autoDetach = false/);
});
});
describe('commitFile `created` flag', () => {

6
packages/context/src/core/git.service.ts
View file

@ -105,6 +105,12 @@ export class GitService {
this.logger.log('Initialized git repository');
}
// Keep any auto-maintenance triggered by writes in-process. Detached maintenance can
// keep object-pack directories alive briefly after awaited git commands complete,
// which makes temp-project cleanup flaky in CI.
await this.git.addConfig('gc.autoDetach', 'false');
await this.git.addConfig('maintenance.autoDetach', 'false');
// Ensure HEAD always resolves to a commit so callers (e.g., the memory-agent squash flow)
// can rely on `revParseHead()` returning a SHA. Idempotent: skip if HEAD already exists.
const head = await this.revParseHead();

105
packages/context/src/ingest/adapters/historic-sql/local-ingest-acceptance.test.ts
View file

@ -2,7 +2,7 @@ import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import YAML from 'yaml';
import { AgentRunnerService } from '../../../agent/index.js';
import type { AgentRunnerPort, RunLoopParams } from '../../../llm/index.js';
import { initKtxProject, loadKtxProject, type KtxLocalProject } from '../../../project/index.js';
import {
type SqlAnalysisBatchItem,
@ -47,8 +47,8 @@ class AcceptanceHistoricSqlReader implements HistoricSqlReader {
}
}
class HistoricSqlAcceptanceAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: any) => {
class HistoricSqlAcceptanceAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (params.telemetryTags?.operationName !== 'ingest-bundle-wu') {
return { stopReason: 'natural' as const };
}
@ -59,78 +59,65 @@ class HistoricSqlAcceptanceAgentRunner extends AgentRunnerService {
}
if (params.telemetryTags.unitKey === 'historic-sql-table-public-orders') {
const result = await emitEvidence.execute(
{
kind: 'table_usage',
table: 'public.orders',
rawPath: 'tables/public.orders.json',
usage: {
narrative: 'Analysts repeatedly inspect paid order lifecycle by customer segment.',
frequencyTier: 'high',
commonFilters: ['status'],
commonGroupBys: ['status', 'segment'],
commonJoins: [{ table: 'public.customers', on: ['customer_id', 'id'] }],
staleSince: null,
},
const result = await emitEvidence.execute({
kind: 'table_usage',
table: 'public.orders',
rawPath: 'tables/public.orders.json',
usage: {
narrative: 'Analysts repeatedly inspect paid order lifecycle by customer segment.',
frequencyTier: 'high',
commonFilters: ['status'],
commonGroupBys: ['status', 'segment'],
commonJoins: [{ table: 'public.customers', on: ['customer_id', 'id'] }],
staleSince: null,
},
{ toolCallId: 'historic-sql-orders-usage' },
);
if (!String(result).includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected orders evidence result: ${String(result)}`);
});
if (!result.markdown.includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected orders evidence result: ${result.markdown}`);
}
}
if (params.telemetryTags.unitKey === 'historic-sql-table-public-customers') {
const result = await emitEvidence.execute(
{
kind: 'table_usage',
table: 'public.customers',
rawPath: 'tables/public.customers.json',
usage: {
narrative: 'Customers provide segment context for paid order lifecycle analysis.',
frequencyTier: 'mid',
commonFilters: [],
commonGroupBys: ['segment'],
commonJoins: [{ table: 'public.orders', on: ['id', 'customer_id'] }],
staleSince: null,
},
const result = await emitEvidence.execute({
kind: 'table_usage',
table: 'public.customers',
rawPath: 'tables/public.customers.json',
usage: {
narrative: 'Customers provide segment context for paid order lifecycle analysis.',
frequencyTier: 'mid',
commonFilters: [],
commonGroupBys: ['segment'],
commonJoins: [{ table: 'public.orders', on: ['id', 'customer_id'] }],
staleSince: null,
},
{ toolCallId: 'historic-sql-customers-usage' },
);
if (!String(result).includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected customers evidence result: ${String(result)}`);
});
if (!result.markdown.includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected customers evidence result: ${result.markdown}`);
}
}
if (params.telemetryTags.unitKey === 'historic-sql-patterns-part-0001') {
const result = await emitEvidence.execute(
{
kind: 'pattern',
rawPath: 'patterns-input/part-0001.json',
pattern: {
slug: 'paid-order-lifecycle',
title: 'Paid Order Lifecycle',
narrative: 'Analysts join orders and customers to compare paid order lifecycle by segment.',
definitionSql:
'select o.status, c.segment, count(*) from public.orders o join public.customers c on c.id = o.customer_id group by o.status, c.segment',
tablesInvolved: ['public.orders', 'public.customers'],
slRefs: ['orders', 'customers'],
constituentTemplateIds: ['pg:orders-lifecycle'],
},
const result = await emitEvidence.execute({
kind: 'pattern',
rawPath: 'patterns-input/part-0001.json',
pattern: {
slug: 'paid-order-lifecycle',
title: 'Paid Order Lifecycle',
narrative: 'Analysts join orders and customers to compare paid order lifecycle by segment.',
definitionSql:
'select o.status, c.segment, count(*) from public.orders o join public.customers c on c.id = o.customer_id group by o.status, c.segment',
tablesInvolved: ['public.orders', 'public.customers'],
slRefs: ['orders', 'customers'],
constituentTemplateIds: ['pg:orders-lifecycle'],
},
{ toolCallId: 'historic-sql-pattern' },
);
if (!String(result).includes('Recorded historic-SQL pattern evidence')) {
throw new Error(`Unexpected pattern evidence result: ${String(result)}`);
});
if (!result.markdown.includes('Recorded historic-SQL pattern evidence')) {
throw new Error(`Unexpected pattern evidence result: ${result.markdown}`);
}
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
function acceptanceSqlAnalysis(): SqlAnalysisPort {

7
packages/context/src/ingest/context-candidates/curator-pagination.service.ts
View file

@ -1,7 +1,6 @@
import type { KtxModelRole } from '@ktx/llm';
import type { ToolSet } from 'ai';
import type { AgentRunnerService } from '../../agent/index.js';
import { type KtxLogger, noopLogger } from '../../core/index.js';
import type { AgentRunnerPort, KtxRuntimeToolSet } from '../../llm/index.js';
import type { MemoryAction } from '../../memory/index.js';
import type { ContextCandidateForDedup, CuratorPaginationPort, CuratorPaginationReport } from '../ports.js';
import type {
@ -38,7 +37,7 @@ export interface CuratorPaginationInput {
modelRole: KtxModelRole;
buildSystemPrompt: () => string;
buildUserPrompt: (input: CuratorPaginationPromptInput) => string;
buildToolSet: (passNumber: number) => ToolSet;
buildToolSet: (passNumber: number) => KtxRuntimeToolSet;
getReconciliationActions: () => MemoryAction[];
onStepFinish?: (info: { passNumber: number; stepIndex: number; stepBudget: number }) => void;
}
@ -50,7 +49,7 @@ interface CuratorPaginationResult extends ReconciliationOutcome {
export interface CuratorPaginationServiceDeps {
store: ContextCandidateStorePort;
agentRunner: AgentRunnerService;
agentRunner: AgentRunnerPort;
settings: CuratorPaginationSettings;
logger?: KtxLogger;
}

14
packages/context/src/ingest/ingest-bundle.runner.test.ts
View file

@ -200,7 +200,7 @@ const makeDeps = () => {
const slValidator = { validateSingleSource: vi.fn().mockResolvedValue({ errors: [], warnings: [] }) };
const toolsetFactory = {
createIngestWuToolset: vi.fn().mockReturnValue({
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
}),
@ -419,7 +419,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
deps.toolsetFactory.createIngestWuToolset.mockImplementation((toolSession: any) => {
sessions.push(toolSession);
return {
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
};
@ -591,7 +591,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
deps.toolsetFactory.createIngestWuToolset.mockImplementation((toolSession: any) => {
currentToolSession = toolSession;
return {
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
};
@ -663,7 +663,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
deps.toolsetFactory.createIngestWuToolset.mockImplementation((toolSession: any) => {
currentToolSession = toolSession;
return {
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
};
@ -834,7 +834,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
it('stores memory-flow provenance and transcript summaries in the ingest report body', async () => {
const deps = makeDeps();
deps.toolsetFactory.createIngestWuToolset.mockReturnValue({
toAiSdkTools: vi.fn().mockReturnValue({
toRuntimeTools: vi.fn().mockReturnValue({
read_raw_span: {
description: 'read a raw span',
inputSchema: {},
@ -1376,7 +1376,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
deps.toolsetFactory.createIngestWuToolset.mockImplementation((toolSession: any) => {
currentToolSession = toolSession;
return {
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
};
@ -1933,7 +1933,7 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
deps.toolsetFactory.createIngestWuToolset.mockImplementation((toolSession: any) => {
currentToolSession = toolSession;
return {
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
getToolNames: vi.fn().mockReturnValue([]),
};

120
packages/context/src/ingest/ingest-bundle.runner.ts
View file

@ -1,9 +1,9 @@
import { mkdir, readFile, rm, writeFile } from 'node:fs/promises';
import { dirname, join } from 'node:path';
import { type Tool, tool } from 'ai';
import pLimit from 'p-limit';
import { z } from 'zod';
import { type KtxLogger, noopLogger } from '../core/index.js';
import { createRuntimeToolDescriptorFromAiTool, type KtxRuntimeToolSet } from '../llm/index.js';
import type { CaptureSession, MemoryAction } from '../memory/index.js';
import type { SemanticLayerService, SemanticLayerSource, SlValidationDeps } from '../sl/index.js';
import { createTouchedSlSources, type ToolContext, type ToolSession } from '../tools/index.js';
@ -694,8 +694,9 @@ export class IngestBundleRunner {
};
const skillsLoadedPerWu: string[] = [];
const loadSkillTool: Record<string, Tool> = {
load_skill: tool({
const loadSkillTool: KtxRuntimeToolSet = {
load_skill: {
name: 'load_skill',
description:
'Load a skill to get specialized instructions. Call this when a skill listed in the system prompt matches the current task.',
inputSchema: z.object({ name: z.string() }),
@ -705,19 +706,23 @@ export class IngestBundleRunner {
const available =
(await this.deps.skillsRegistry.listSkills('memory_agent')).map((s) => s.name).join(', ') ||
'(none)';
return `Skill "${name}" not available. Available: ${available}`;
return { markdown: `Skill "${name}" not available. Available: ${available}` };
}
const body = await readFile(join(skill.path, 'SKILL.md'), 'utf-8');
if (!skillsLoadedPerWu.includes(skill.name)) {
skillsLoadedPerWu.push(skill.name);
}
return {
const structured = {
name: skill.name,
skillDirectory: skill.path,
content: this.deps.skillsRegistry.stripFrontmatter(body),
};
return {
markdown: `# ${structured.name}\n\n${structured.content}`,
structured,
};
},
}),
},
};
const priorProvenance = await this.deps.provenance.findLatestArtifactsForRawPaths(
@ -726,12 +731,15 @@ export class IngestBundleRunner {
wu.rawFiles,
);
const wuEmitUnmappedFallbackTool = {
emit_unmapped_fallback: createEmitUnmappedFallbackTool({
stageIndex,
allowedPaths: new Set(wu.rawFiles),
tableRefExists: (tableRef) =>
this.tableRefExistsInSemanticLayer(scopedSemanticLayerService, slConnectionIds, tableRef),
}),
emit_unmapped_fallback: createRuntimeToolDescriptorFromAiTool(
'emit_unmapped_fallback',
createEmitUnmappedFallbackTool({
stageIndex,
allowedPaths: new Set(wu.rawFiles),
tableRefExists: (tableRef) =>
this.tableRefExistsInSemanticLayer(scopedSemanticLayerService, slConnectionIds, tableRef),
}),
),
};
const systemPrompt = buildWuSystemPrompt({
@ -765,7 +773,7 @@ export class IngestBundleRunner {
wu: wuInner,
loadSkillTool,
emitUnmappedFallbackTool: wuEmitUnmappedFallbackTool,
toolsetTools: wuToolset.toAiSdkTools(wuToolContext),
toolsetTools: wuToolset.toRuntimeTools(wuToolContext),
}),
join(transcriptDir, `${wuInner.unitKey}.jsonl`),
wuInner.unitKey,
@ -921,53 +929,79 @@ export class IngestBundleRunner {
ingest: ingestToolMetadata,
session: rcToolSession,
};
const rcLoadSkill: Record<string, Tool> = {
load_skill: tool({
const rcLoadSkill: KtxRuntimeToolSet = {
load_skill: {
name: 'load_skill',
description: 'Load a skill.',
inputSchema: z.object({ name: z.string() }),
execute: async ({ name }) => {
const skill = await this.deps.skillsRegistry.getSkill(name, 'memory_agent');
if (!skill) {
return `Skill "${name}" not found`;
return { markdown: `Skill "${name}" not found` };
}
const body = await readFile(join(skill.path, 'SKILL.md'), 'utf-8');
return { name: skill.name, content: this.deps.skillsRegistry.stripFrontmatter(body) };
const structured = { name: skill.name, content: this.deps.skillsRegistry.stripFrontmatter(body) };
return { markdown: `# ${structured.name}\n\n${structured.content}`, structured };
},
}),
},
};
const allStagedPaths = new Set<string>([...currentHashes.keys()]);
const rcRawSpanTool = { read_raw_span: createReadRawSpanTool({ stagedDir, allowedPaths: allStagedPaths }) };
const rcStageListTool = { stage_list: createStageListTool({ stageIndex }) };
const rcStageDiffTool = { stage_diff: createStageDiffTool({ stageIndex }) };
const rcRawSpanTool = {
read_raw_span: createRuntimeToolDescriptorFromAiTool(
'read_raw_span',
createReadRawSpanTool({ stagedDir, allowedPaths: allStagedPaths }),
),
};
const rcStageListTool = {
stage_list: createRuntimeToolDescriptorFromAiTool('stage_list', createStageListTool({ stageIndex })),
};
const rcStageDiffTool = {
stage_diff: createRuntimeToolDescriptorFromAiTool('stage_diff', createStageDiffTool({ stageIndex })),
};
const rcEvictionListTool = {
eviction_list: createEvictionListTool({
provenance: this.deps.provenance,
connectionId: job.connectionId,
sourceKey: job.sourceKey,
deletedRawPaths: eviction?.deletedRawPaths ?? [],
}),
eviction_list: createRuntimeToolDescriptorFromAiTool(
'eviction_list',
createEvictionListTool({
provenance: this.deps.provenance,
connectionId: job.connectionId,
sourceKey: job.sourceKey,
deletedRawPaths: eviction?.deletedRawPaths ?? [],
}),
),
};
const rcEmitConflictResolutionTool = {
emit_conflict_resolution: createEmitConflictResolutionTool({ stageIndex }),
emit_conflict_resolution: createRuntimeToolDescriptorFromAiTool(
'emit_conflict_resolution',
createEmitConflictResolutionTool({ stageIndex }),
),
};
const rcEmitEvictionDecisionTool = {
emit_eviction_decision: createEmitEvictionDecisionTool({
stageIndex,
deletedRawPaths: eviction?.deletedRawPaths ?? [],
}),
emit_eviction_decision: createRuntimeToolDescriptorFromAiTool(
'emit_eviction_decision',
createEmitEvictionDecisionTool({
stageIndex,
deletedRawPaths: eviction?.deletedRawPaths ?? [],
}),
),
};
const rcEmitArtifactResolutionTool = {
emit_artifact_resolution: createEmitArtifactResolutionTool({
stageIndex,
allowedPaths: allStagedPaths,
}),
emit_artifact_resolution: createRuntimeToolDescriptorFromAiTool(
'emit_artifact_resolution',
createEmitArtifactResolutionTool({
stageIndex,
allowedPaths: allStagedPaths,
}),
),
};
const rcEmitUnmappedFallbackTool = {
emit_unmapped_fallback: createEmitUnmappedFallbackTool({
stageIndex,
allowedPaths: allStagedPaths,
tableRefExists: (tableRef) => this.tableRefExistsInSemanticLayer(rcScopedSl, slConnectionIds, tableRef),
}),
emit_unmapped_fallback: createRuntimeToolDescriptorFromAiTool(
'emit_unmapped_fallback',
createEmitUnmappedFallbackTool({
stageIndex,
allowedPaths: allStagedPaths,
tableRefExists: (tableRef) => this.tableRefExistsInSemanticLayer(rcScopedSl, slConnectionIds, tableRef),
}),
),
};
const reconcileBaseFraming = await this.deps.promptService.loadPrompt('memory_agent_bundle_ingest_reconcile');
@ -1026,7 +1060,7 @@ export class IngestBundleRunner {
emitArtifactResolutionTool: rcEmitArtifactResolutionTool,
emitUnmappedFallbackTool: rcEmitUnmappedFallbackTool,
readRawSpanTool: rcRawSpanTool,
toolsetTools: rcToolset.toAiSdkTools(rcToolContext),
toolsetTools: rcToolset.toRuntimeTools(rcToolContext),
}),
join(transcriptDir, 'reconcile.jsonl'),
'reconcile',
@ -1075,7 +1109,7 @@ export class IngestBundleRunner {
emitArtifactResolutionTool: rcEmitArtifactResolutionTool,
emitUnmappedFallbackTool: rcEmitUnmappedFallbackTool,
readRawSpanTool: rcRawSpanTool,
toolsetTools: rcToolset.toAiSdkTools(rcToolContext),
toolsetTools: rcToolset.toRuntimeTools(rcToolContext),
}),
join(transcriptDir, 'reconcile.jsonl'),
'reconcile',

177
packages/context/src/ingest/local-bundle-ingest.test.ts
View file

@ -3,7 +3,7 @@ import { tmpdir } from 'node:os';
import { join } from 'node:path';
import Database from 'better-sqlite3';
import YAML from 'yaml';
import { AgentRunnerService } from '../agent/index.js';
import type { AgentRunnerPort, RunLoopParams } from '../llm/index.js';
import { initKtxProject, type KtxLocalProject, loadKtxProject } from '../project/index.js';
import { makeLocalGitRepo } from '../test/make-local-git-repo.js';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
@ -13,16 +13,12 @@ import { createDefaultLocalIngestAdapters, localPullConfigForAdapter } from './l
import { getLocalIngestStatus, runLocalIngest } from './local-ingest.js';
import type { ChunkResult, DiffSet, SourceAdapter } from './types.js';
class TestAgentRunner extends AgentRunnerService {
override runLoop = vi.fn().mockResolvedValue({ stopReason: 'natural' as const });
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
class TestAgentRunner implements AgentRunnerPort {
runLoop = vi.fn().mockResolvedValue({ stopReason: 'natural' as const });
}
class LookerSlWritingAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: any) => {
class LookerSlWritingAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (
params.telemetryTags?.operationName === 'ingest-bundle-wu' &&
params.telemetryTags?.unitKey === 'looker-explore-ecommerce-orders'
@ -31,130 +27,100 @@ class LookerSlWritingAgentRunner extends AgentRunnerService {
if (!ledger?.execute) {
throw new Error('record_verification_ledger tool was not available to the Looker WorkUnit');
}
await ledger.execute(
{
summary: 'Test fixture verified Looker explore target identifiers before writing SL.',
verifiedIdentifiers: ['prod-warehouse', 'public.orders'],
unverifiedIdentifiers: [],
},
{ toolCallId: 'looker-verification-ledger', messages: [] },
);
await ledger.execute({
summary: 'Test fixture verified Looker explore target identifiers before writing SL.',
verifiedIdentifiers: ['prod-warehouse', 'public.orders'],
unverifiedIdentifiers: [],
});
const slWrite = params.toolSet.sl_write_source;
if (!slWrite?.execute) {
throw new Error('sl_write_source tool was not available to the Looker WorkUnit');
}
const result = await slWrite.execute(
{
connectionId: 'prod-warehouse',
sourceName: 'looker__ecommerce__orders',
source: {
name: 'looker__ecommerce__orders',
table: 'public.orders',
grain: ['id'],
columns: [
{ name: 'id', type: 'number' },
{ name: 'revenue', type: 'number' },
],
measures: [{ name: 'total_revenue', expr: 'sum(revenue)' }],
},
const result = await slWrite.execute({
connectionId: 'prod-warehouse',
sourceName: 'looker__ecommerce__orders',
source: {
name: 'looker__ecommerce__orders',
table: 'public.orders',
grain: ['id'],
columns: [
{ name: 'id', type: 'number' },
{ name: 'revenue', type: 'number' },
],
measures: [{ name: 'total_revenue', expr: 'sum(revenue)' }],
},
{ toolCallId: 'looker-sl-write' },
);
if (!result.structured.success) {
});
if (!(result.structured as { success?: boolean } | undefined)?.success) {
throw new Error(result.markdown);
}
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
class WikiWritingAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: any) => {
class WikiWritingAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (params.telemetryTags?.operationName === 'ingest-bundle-wu') {
const ledger = params.toolSet.record_verification_ledger;
if (!ledger?.execute) {
throw new Error('record_verification_ledger tool was not available to the WorkUnit');
}
await ledger.execute(
{
summary: 'Test fixture writes wiki-only context with no warehouse identifiers.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
},
{ toolCallId: 'wiki-verification-ledger', messages: [] },
);
await ledger.execute({
summary: 'Test fixture writes wiki-only context with no warehouse identifiers.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
});
const wikiWrite = params.toolSet.wiki_write;
if (!wikiWrite?.execute) {
throw new Error('wiki_write tool was not available to the WorkUnit');
}
const result = await wikiWrite.execute(
{
key: 'orders_context',
summary: 'Orders source context',
content: 'Orders are purchase records used for revenue analysis.',
tags: ['orders'],
},
{ toolCallId: 'wiki-write' },
);
if (!result.structured.success) {
const result = await wikiWrite.execute({
key: 'orders_context',
summary: 'Orders source context',
content: 'Orders are purchase records used for revenue analysis.',
tags: ['orders'],
});
if (!(result.structured as { success?: boolean } | undefined)?.success) {
throw new Error(result.markdown);
}
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
class WikiWritingWithRawPathAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: any) => {
class WikiWritingWithRawPathAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (params.telemetryTags?.operationName === 'ingest-bundle-wu') {
const ledger = params.toolSet.record_verification_ledger;
if (!ledger?.execute) {
throw new Error('record_verification_ledger tool was not available to the WorkUnit');
}
await ledger.execute(
{
summary: 'Test fixture writes wiki-only context with explicit raw provenance and no warehouse identifiers.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
},
{ toolCallId: 'wiki-raw-path-verification-ledger', messages: [] },
);
await ledger.execute({
summary: 'Test fixture writes wiki-only context with explicit raw provenance and no warehouse identifiers.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
});
const wikiWrite = params.toolSet.wiki_write;
if (!wikiWrite?.execute) {
throw new Error('wiki_write tool was not available to the WorkUnit');
}
const result = await wikiWrite.execute(
{
key: 'orders_context',
summary: 'Orders source context',
content: 'Orders are purchase records used for revenue analysis.',
tags: ['orders'],
rawPaths: ['orders/orders.json'],
},
{ toolCallId: 'wiki-write' },
);
if (!result.structured.success) {
const result = await wikiWrite.execute({
key: 'orders_context',
summary: 'Orders source context',
content: 'Orders are purchase records used for revenue analysis.',
tags: ['orders'],
rawPaths: ['orders/orders.json'],
});
if (!(result.structured as { success?: boolean } | undefined)?.success) {
throw new Error(result.markdown);
}
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
class HistoricSqlEvidenceAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: any) => {
class HistoricSqlEvidenceAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (
params.telemetryTags?.operationName === 'ingest-bundle-wu' &&
params.telemetryTags?.unitKey === 'historic-sql-table-public-orders'
@ -163,31 +129,24 @@ class HistoricSqlEvidenceAgentRunner extends AgentRunnerService {
if (!emitEvidence?.execute) {
throw new Error('emit_historic_sql_evidence tool was not available to the historic-SQL WorkUnit');
}
const result = await emitEvidence.execute(
{
kind: 'table_usage',
table: 'public.orders',
rawPath: 'tables/public.orders.json',
usage: {
narrative: 'Orders are repeatedly queried by lifecycle status.',
frequencyTier: 'high',
commonFilters: ['status'],
commonJoins: [],
staleSince: null,
},
const result = await emitEvidence.execute({
kind: 'table_usage',
table: 'public.orders',
rawPath: 'tables/public.orders.json',
usage: {
narrative: 'Orders are repeatedly queried by lifecycle status.',
frequencyTier: 'high',
commonFilters: ['status'],
commonJoins: [],
staleSince: null,
},
{ toolCallId: 'historic-sql-evidence' },
);
if (!String(result).includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected historic-SQL evidence result: ${String(result)}`);
});
if (!result.markdown.includes('Recorded historic-SQL table_usage evidence')) {
throw new Error(`Unexpected historic-SQL evidence result: ${result.markdown}`);
}
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
class HistoricSqlEvidenceTestAdapter implements SourceAdapter {

45
packages/context/src/ingest/local-bundle-runtime.test.ts
View file

@ -1,7 +1,7 @@
import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { AgentRunnerService } from '../agent/index.js';
import type { AgentRunnerPort } from '../llm/index.js';
import { initKtxProject, type KtxLocalProject, loadKtxProject } from '../project/index.js';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { FakeSourceAdapter } from './adapters/fake/fake.adapter.js';
@ -17,6 +17,10 @@ type RuntimeWithConnectionDeps = {
};
};
function testAgentRunner(): AgentRunnerPort {
return { runLoop: vi.fn().mockResolvedValue({ stopReason: 'natural' as const }) };
}
describe('createLocalBundleIngestRuntime', () => {
let tempDir: string;
let project: KtxLocalProject;
@ -55,15 +59,42 @@ describe('createLocalBundleIngestRuntime', () => {
}),
).toThrow(
[
'ktx ingest requires llm.provider.backend: anthropic, vertex, or gateway, or an injected agentRunner.',
`Configure an Anthropic provider, then rerun ingest:`,
` ktx setup --project-dir ${project.projectDir} --anthropic-api-key-env ANTHROPIC_API_KEY --anthropic-model claude-sonnet-4-6 --no-input`,
'ktx ingest requires llm.provider.backend: anthropic, vertex, gateway, or claude-code, or an injected agentRunner.',
'Configure a local Claude Code session or API-backed LLM, then rerun ingest:',
` ktx setup --project-dir ${project.projectDir} --llm-backend claude-code --no-input`,
` ktx setup --project-dir ${project.projectDir} --llm-backend anthropic --anthropic-api-key-env ANTHROPIC_API_KEY --anthropic-model claude-sonnet-4-6 --no-input`,
].join('\n'),
);
});
it('uses a runtime-backed agent runner when claude-code is configured', () => {
const runtime = {
generateText: vi.fn(),
generateObject: vi.fn(),
runAgentLoop: vi.fn(async () => ({ stopReason: 'natural' as const })),
};
project.config.llm = {
provider: { backend: 'claude-code' },
models: { default: 'sonnet' },
promptCaching: { enabled: false },
};
const createLlmRuntime = vi.fn(() => runtime);
const created = createLocalBundleIngestRuntime({
project,
adapters: [new FakeSourceAdapter()],
createLlmRuntime,
});
expect(created).toBeDefined();
expect(createLlmRuntime).toHaveBeenCalledWith(
project.config.llm,
expect.objectContaining({ projectDir: project.projectDir }),
);
});
it('builds runner deps with local SQLite stores and context tools enabled', async () => {
const agentRunner = new AgentRunnerService({ llmProvider: { getModel: () => ({}) as never } as any });
const agentRunner = testAgentRunner();
const runtime = createLocalBundleIngestRuntime({
project,
@ -94,7 +125,7 @@ describe('createLocalBundleIngestRuntime', () => {
project_id: 'acme',
dataset_id: 'warehouse',
};
const agentRunner = new AgentRunnerService({ llmProvider: { getModel: () => ({}) as never } as any });
const agentRunner = testAgentRunner();
const runtime = createLocalBundleIngestRuntime({
project,
@ -114,7 +145,7 @@ describe('createLocalBundleIngestRuntime', () => {
});
it('passes project connection config to local ingest query executors', async () => {
const agentRunner = new AgentRunnerService({ llmProvider: { getModel: () => ({}) as never } as any });
const agentRunner = testAgentRunner();
const queryExecutor = {
execute: vi.fn(async () => ({
headers: ['answer'],

77
packages/context/src/ingest/local-bundle-runtime.ts
View file

@ -1,20 +1,20 @@
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 { AgentRunnerService as DefaultAgentRunnerService } from '../agent/index.js';
import { localConnectionInfoFromConfig, type KtxSqlQueryExecutorPort } from '../connections/index.js';
import type { KtxEmbeddingPort, KtxLogger } from '../core/index.js';
import { noopLogger, SessionWorktreeService } from '../core/index.js';
import type { KtxSemanticLayerComputePort } from '../daemon/index.js';
import {
createJsonlKtxLlmDebugRequestRecorder,
createRuntimeToolDescriptorFromAiTool,
createLocalKtxEmbeddingProviderFromConfig,
createLocalKtxLlmProviderFromConfig,
createLocalKtxLlmRuntimeFromConfig,
KtxIngestEmbeddingPortAdapter,
RuntimeAgentRunner,
type AgentRunnerPort,
type KtxLlmRuntimePort,
type KtxRuntimeToolSet,
} from '../llm/index.js';
import type { KtxLocalProject } from '../project/index.js';
import { ktxLocalStateDbPath } from '../project/index.js';
@ -100,8 +100,9 @@ const LOCAL_SHAPE_WARNING = 'Local ingest validates semantic-layer YAML shape on
export interface CreateLocalBundleIngestRuntimeOptions {
project: KtxLocalProject;
adapters: SourceAdapter[];
agentRunner?: AgentRunnerService;
llmProvider?: KtxLlmProvider;
agentRunner?: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
createLlmRuntime?: typeof createLocalKtxLlmRuntimeFromConfig;
llmDebugRequestFile?: string;
memoryModel?: string;
semanticLayerCompute?: KtxSemanticLayerComputePort;
@ -456,12 +457,12 @@ class NoopKnowledgeEventPort implements KnowledgeEventPort {
class LocalIngestToolSet implements IngestToolsetLike {
constructor(
private readonly tools: BaseTool[],
private readonly sourceTools: Record<string, Tool> = {},
private readonly sourceTools: KtxRuntimeToolSet = {},
) {}
toAiSdkTools(context: ToolContext) {
toRuntimeTools(context: ToolContext): KtxRuntimeToolSet {
return {
...Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAiSdkTool(context)])),
...Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toRuntimeTool(context)])),
...this.sourceTools,
};
}
@ -541,13 +542,16 @@ class LocalIngestToolsetFactory implements IngestToolsetFactoryPort {
}
createIngestWuToolset(session: ToolSession, options?: { includeContextEvidenceTools?: boolean }): IngestToolsetLike {
const sourceTools: Record<string, Tool> =
const sourceTools: KtxRuntimeToolSet =
session.ingest?.sourceKey === 'historic-sql'
? {
emit_historic_sql_evidence: createEmitHistoricSqlEvidenceTool({
connectionId: session.connectionId,
session,
}),
emit_historic_sql_evidence: createRuntimeToolDescriptorFromAiTool(
'emit_historic_sql_evidence',
createEmitHistoricSqlEvidenceTool({
connectionId: session.connectionId,
session,
}),
),
}
: {};
return new LocalIngestToolSet(
@ -571,36 +575,36 @@ function nextLocalJobId(): string {
function localIngestLlmProviderGuardMessage(projectDir: string): string {
return [
'ktx ingest requires llm.provider.backend: anthropic, vertex, or gateway, or an injected agentRunner.',
'Configure an Anthropic provider, then rerun ingest:',
` ktx setup --project-dir ${projectDir} --anthropic-api-key-env ANTHROPIC_API_KEY --anthropic-model claude-sonnet-4-6 --no-input`,
'ktx ingest requires llm.provider.backend: anthropic, vertex, gateway, or claude-code, or an injected agentRunner.',
'Configure a local Claude Code session or API-backed LLM, then rerun ingest:',
` ktx setup --project-dir ${projectDir} --llm-backend claude-code --no-input`,
` ktx setup --project-dir ${projectDir} --llm-backend anthropic --anthropic-api-key-env ANTHROPIC_API_KEY --anthropic-model claude-sonnet-4-6 --no-input`,
].join('\n');
}
function resolveAgentRunner(options: CreateLocalBundleIngestRuntimeOptions): {
agentRunner: AgentRunnerService;
llmProvider?: KtxLlmProvider;
agentRunner: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
} {
const llmProvider =
options.llmProvider ?? createLocalKtxLlmProviderFromConfig(options.project.config.llm) ?? undefined;
const llmRuntime =
options.llmRuntime ??
(options.createLlmRuntime ?? createLocalKtxLlmRuntimeFromConfig)(options.project.config.llm, {
projectDir: options.project.projectDir,
env: process.env,
}) ??
undefined;
if (options.agentRunner) {
return { agentRunner: options.agentRunner, ...(llmProvider ? { llmProvider } : {}) };
return { agentRunner: options.agentRunner, ...(llmRuntime ? { llmRuntime } : {}) };
}
if (!llmProvider) {
if (!llmRuntime) {
throw new Error(localIngestLlmProviderGuardMessage(options.project.projectDir));
}
return {
agentRunner: new DefaultAgentRunnerService({
llmProvider,
logger: options.logger ?? noopLogger,
...(options.llmDebugRequestFile
? { debugRequestRecorder: createJsonlKtxLlmDebugRequestRecorder(options.llmDebugRequestFile) }
: {}),
}),
llmProvider,
agentRunner: new RuntimeAgentRunner(llmRuntime),
llmRuntime,
};
}
@ -627,7 +631,7 @@ export function createLocalBundleIngestRuntime(
const knowledgeIndex = new LocalKnowledgeIndex(options.project, embedding);
const knowledgeEvents = new NoopKnowledgeEventPort();
const wikiService = new KnowledgeWikiService(rootFileStore, embedding, knowledgeIndex, options.project.git, logger);
const { agentRunner, llmProvider } = resolveAgentRunner(options);
const { agentRunner, llmRuntime } = resolveAgentRunner(options);
const promptService = new PromptService({ promptsDir, partials: [], logger });
const storage = new LocalIngestStorage(options.project);
const registry = registerAdapters(options.adapters);
@ -681,10 +685,11 @@ export function createLocalBundleIngestRuntime(
commitMessages: new LocalCommitMessagePort(),
embedding,
contextEvidenceIndex: new ContextEvidenceIndexService({ store: contextStore, embeddings: embedding, logger }),
pageTriage: llmProvider
llmRuntime,
pageTriage: llmRuntime
? new PageTriageService({
store: contextStore,
llmProvider,
llmRuntime,
settings: {
enabled: true,
maxConcurrency: 2,

17
packages/context/src/ingest/local-ingest.ts
View file

@ -1,11 +1,10 @@
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 { KtxSqlQueryExecutorPort } from '../connections/index.js';
import type { KtxLogger } from '../core/index.js';
import type { KtxSemanticLayerComputePort } from '../daemon/index.js';
import type { AgentRunnerPort, KtxLlmRuntimePort } from '../llm/index.js';
import type { KtxLocalProject } from '../project/index.js';
import { ktxLocalStateDbPath } from '../project/index.js';
import { planMetabaseFanoutChildren } from './adapters/metabase/fanout-planner.js';
@ -28,8 +27,8 @@ export interface RunLocalIngestOptions {
trigger?: IngestTrigger;
jobId?: string;
memoryFlow?: MemoryFlowEventSink;
agentRunner?: AgentRunnerService;
llmProvider?: KtxLlmProvider;
agentRunner?: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
llmDebugRequestFile?: string;
memoryModel?: string;
semanticLayerCompute?: KtxSemanticLayerComputePort;
@ -41,7 +40,7 @@ export interface LocalIngestMcpOptions
extends Pick<
RunLocalIngestOptions,
| 'agentRunner'
| 'llmProvider'
| 'llmRuntime'
| 'memoryModel'
| 'semanticLayerCompute'
| 'queryExecutor'
@ -167,8 +166,8 @@ async function runScheduledPullJob(options: {
trigger?: IngestTrigger;
jobId?: string;
memoryFlow?: MemoryFlowEventSink;
agentRunner?: AgentRunnerService;
llmProvider?: KtxLlmProvider;
agentRunner?: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
memoryModel?: string;
semanticLayerCompute?: KtxSemanticLayerComputePort;
queryExecutor?: KtxSqlQueryExecutorPort;
@ -221,7 +220,7 @@ export async function runLocalIngest(options: RunLocalIngestOptions): Promise<Lo
jobId,
memoryFlow: options.memoryFlow,
agentRunner: options.agentRunner,
llmProvider: options.llmProvider,
llmRuntime: options.llmRuntime,
memoryModel: options.memoryModel,
semanticLayerCompute: options.semanticLayerCompute,
queryExecutor: options.queryExecutor,
@ -406,7 +405,7 @@ export async function runLocalMetabaseIngest(
jobId: childJobId,
memoryFlow: options.memoryFlow,
agentRunner: options.agentRunner,
llmProvider: options.llmProvider,
llmRuntime: options.llmRuntime,
memoryModel: options.memoryModel,
semanticLayerCompute: options.semanticLayerCompute,
queryExecutor: options.queryExecutor,

10
packages/context/src/ingest/local-metabase-ingest.test.ts
View file

@ -1,24 +1,20 @@
import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { AgentRunnerService } from '../agent/index.js';
import type { AgentRunnerPort, RunLoopParams } from '../llm/index.js';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { initKtxProject, type KtxLocalProject } from '../project/index.js';
import { LocalMetabaseDiscoveryCache } from './adapters/metabase/local-source-state-store.js';
import { getLocalIngestStatus, runLocalMetabaseIngest } from './local-ingest.js';
import type { ChunkResult, FetchContext, SourceAdapter } from './types.js';
class TestAgentRunner extends AgentRunnerService {
override runLoop = vi.fn(async (params: Parameters<AgentRunnerService['runLoop']>[0]) => {
class TestAgentRunner implements AgentRunnerPort {
runLoop = vi.fn(async (params: RunLoopParams) => {
if (params.userPrompt.includes('metabase-db-2')) {
return { stopReason: 'error' as const, error: new Error('database 2 failed') };
}
return { stopReason: 'natural' as const };
});
constructor() {
super({ llmProvider: { getModel: () => ({}) as never } as never });
}
}
class FakeMetabaseSourceAdapter implements SourceAdapter {

73
packages/context/src/ingest/page-triage/page-triage.service.test.ts
View file

@ -21,7 +21,11 @@ describe('PageTriageService', () => {
};
let promptService: { loadPrompt: ReturnType<typeof vi.fn<(name: string) => Promise<string>>> };
let adapter: { triageSupported: true; getTriageSignals: ReturnType<typeof vi.fn> };
let generateTextMock: ReturnType<typeof vi.fn>;
let llmRuntime: {
generateText: ReturnType<typeof vi.fn>;
generateObject: ReturnType<typeof vi.fn>;
runAgentLoop: ReturnType<typeof vi.fn>;
};
beforeEach(async () => {
stagedDir = await mkdtemp(join(tmpdir(), 'page-triage-'));
@ -88,31 +92,16 @@ describe('PageTriageService', () => {
.fn<(name: string) => Promise<string>>()
.mockImplementation((name) => Promise.resolve(`prompt:${name}`)),
};
generateTextMock = vi.fn();
llmRuntime = {
generateText: vi.fn(),
generateObject: vi.fn(),
runAgentLoop: vi.fn(),
};
service = new PageTriageService({
store: repository as any,
llmProvider: {
getModel: vi.fn().mockReturnValue('model'),
getModelByName: vi.fn(),
cacheMarker: vi.fn(),
repairToolCallHandler: vi.fn(),
thinkingProviderOptions: vi.fn(),
telemetryConfig: vi.fn(),
promptCachingConfig: vi.fn(() => ({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
})),
activeBackend: vi.fn(() => 'anthropic'),
} as any,
llmRuntime: llmRuntime as any,
settings: triageSettings,
promptService: promptService as any,
generateText: generateTextMock as any,
});
});
@ -121,10 +110,10 @@ describe('PageTriageService', () => {
});
it('writes light-lane candidates and keeps the page out of full WorkUnits', async () => {
generateTextMock
.mockResolvedValueOnce({ text: JSON.stringify({ lane: 'light', reason: 'short durable policy' }) } as any)
.mockResolvedValueOnce({
text: JSON.stringify({
llmRuntime.generateText
.mockResolvedValueOnce(JSON.stringify({ lane: 'light', reason: 'short durable policy' }))
.mockResolvedValueOnce(
JSON.stringify({
candidates: [
{
candidateKey: 'support-handoff-owner',
@ -142,7 +131,7 @@ describe('PageTriageService', () => {
},
],
}),
} as any);
);
const result = await service.triageRun({
stagedDir,
@ -171,6 +160,7 @@ describe('PageTriageService', () => {
});
expect(result.fullRawPaths.has('pages/page-1/page.md')).toBe(false);
expect(adapter.getTriageSignals).toHaveBeenCalledWith(stagedDir, 'page-1');
expect(llmRuntime.generateText).toHaveBeenCalledWith(expect.objectContaining({ role: 'triage' }));
expect(repository.setDocumentTriageLane).toHaveBeenCalledWith('run-1', 'pages/page-1/page.md', 'light');
expect(repository.insertCandidate).toHaveBeenCalledWith(
expect.objectContaining({
@ -225,23 +215,20 @@ describe('PageTriageService', () => {
}
return Promise.resolve(`prompt:${name}`);
});
generateTextMock
llmRuntime.generateText
.mockImplementationOnce((args: any) => {
const systemMessage = args.system ?? args.messages.find((m: { role: string }) => m.role === 'system');
const userMessage = args.messages.find((m: { role: string }) => m.role === 'user');
const systemText =
typeof systemMessage === 'string' ? systemMessage : (systemMessage.content as string);
const userText = userMessage.content as string;
const systemText = args.system as string;
const userText = args.prompt as string;
expect(systemText).toContain(
'Reusable templates and scripts are durable knowledge regardless of subject matter.',
);
expect(systemText).toContain('Date-titled standups are still skip; named templates and scripts are not.');
expect(userText).toContain('Cold Call Script');
expect(userText).not.toContain('Reusable templates and scripts are durable knowledge');
return { text: JSON.stringify({ lane: 'light', reason: 'reusable sales script' }) } as any;
return JSON.stringify({ lane: 'light', reason: 'reusable sales script' });
})
.mockResolvedValueOnce({
text: JSON.stringify({
.mockResolvedValueOnce(
JSON.stringify({
candidates: [
{
candidateKey: 'cold-call-script',
@ -259,7 +246,7 @@ describe('PageTriageService', () => {
},
],
}),
} as any);
);
const result = await service.triageRun({
stagedDir,
@ -312,9 +299,7 @@ describe('PageTriageService', () => {
'utf-8',
);
generateTextMock.mockResolvedValue({
text: JSON.stringify({ lane: 'full', reason: 'durable policy page' }),
} as any);
llmRuntime.generateText.mockResolvedValue(JSON.stringify({ lane: 'full', reason: 'durable policy page' }));
const result = await service.triageRun({
stagedDir,
@ -351,7 +336,7 @@ describe('PageTriageService', () => {
});
it('falls back to full when classifier output is malformed', async () => {
generateTextMock.mockResolvedValueOnce({ text: 'not-json' } as any);
llmRuntime.generateText.mockResolvedValueOnce('not-json');
const result = await service.triageRun({
stagedDir,
@ -370,8 +355,8 @@ describe('PageTriageService', () => {
});
it('promotes a light page to full when light extraction fails', async () => {
generateTextMock
.mockResolvedValueOnce({ text: JSON.stringify({ lane: 'light', reason: 'short durable policy' }) } as any)
llmRuntime.generateText
.mockResolvedValueOnce(JSON.stringify({ lane: 'light', reason: 'short durable policy' }))
.mockRejectedValueOnce(new Error('provider unavailable'));
const result = await service.triageRun({
@ -405,7 +390,7 @@ describe('PageTriageService', () => {
});
expect(result).toEqual({ enabled: false, report: undefined, fullRawPaths: new Set<string>(), warnings: [] });
expect(generateTextMock).not.toHaveBeenCalled();
expect(llmRuntime.generateText).not.toHaveBeenCalled();
expect(repository.setDocumentTriageLane).not.toHaveBeenCalled();
});
});

24
packages/context/src/ingest/page-triage/page-triage.service.ts
View file

@ -1,11 +1,10 @@
import { createHash } from 'node:crypto';
import { readdir, readFile } from 'node:fs/promises';
import { dirname, join, relative } from 'node:path';
import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider } from '@ktx/llm';
import { generateText, type ToolSet } from 'ai';
import pLimit from 'p-limit';
import { z } from 'zod';
import { type KtxLogger, noopLogger } from '../../core/index.js';
import type { KtxLlmRuntimePort } from '../../llm/index.js';
import type { PromptService } from '../../prompts/index.js';
import type { InsertContextCandidateInput } from '../context-candidates/index.js';
import type { JsonValue } from '../ports.js';
@ -100,20 +99,17 @@ export interface PageTriageSettings {
export interface PageTriageServiceDeps {
store: PageTriageStorePort;
llmProvider: KtxLlmProvider;
llmRuntime: KtxLlmRuntimePort;
settings: PageTriageSettings;
promptService: PromptService;
logger?: KtxLogger;
generateText?: typeof generateText;
}
export class PageTriageService {
private readonly logger: KtxLogger;
private readonly runGenerateText: typeof generateText;
constructor(private readonly deps: PageTriageServiceDeps) {
this.logger = deps.logger ?? noopLogger;
this.runGenerateText = deps.generateText ?? generateText;
}
async triageRun(args: PageTriageRunArgs): Promise<PageTriageRunResult> {
@ -339,22 +335,12 @@ export class PageTriageService {
jobId: string;
unitKey: string;
}): Promise<string> {
const model = this.deps.llmProvider.getModel('triage');
const built = new KtxMessageBuilder(this.deps.llmProvider).wrapSimple({
return this.deps.llmRuntime.generateText({
role: 'triage',
system: params.system,
messages: [{ role: 'user', content: params.prompt }],
tools: {},
model,
});
const split = splitKtxSystemMessages(built.messages);
const result = await this.runGenerateText({
model,
prompt: params.prompt,
temperature: 0,
...(split.system ? { system: split.system } : {}),
messages: split.messages,
tools: built.tools as ToolSet,
});
return result.text;
}
private async buildClassifierSystem(): Promise<string> {

10
packages/context/src/ingest/ports.ts
View file

@ -1,8 +1,7 @@
import type { ToolSet } from 'ai';
import type { KtxModelRole } from '@ktx/llm';
import type { AgentRunnerService } from '../agent/index.js';
import type { KtxEmbeddingPort } from '../core/embedding.js';
import type { GitService, KtxFileStorePort, KtxLogger, SessionOutcome } from '../core/index.js';
import type { AgentRunnerPort, KtxLlmRuntimePort, KtxRuntimeToolSet } from '../llm/index.js';
import type { CaptureSession, MemoryAction, MemoryKnowledgeSlRefsPort } from '../memory/index.js';
import type { PromptService } from '../prompts/index.js';
import type { SkillsRegistryService } from '../skills/index.js';
@ -163,7 +162,7 @@ export interface IngestCommitMessagePort {
}
export interface IngestToolsetLike {
toAiSdkTools(context: ToolContext): ToolSet;
toRuntimeTools(context: ToolContext): KtxRuntimeToolSet;
}
export interface IngestToolsetFactoryPort {
@ -315,7 +314,7 @@ export interface CuratorPaginationPort {
items: ReconcileCandidateForPrompt[];
runState: ReconcilePromptRunState;
}) => string;
buildToolSet: (passNumber: number) => ToolSet;
buildToolSet: (passNumber: number) => KtxRuntimeToolSet;
getReconciliationActions: () => MemoryAction[];
onStepFinish?: (info: { passNumber: number; stepIndex: number; stepBudget: number }) => void;
}): Promise<ReconciliationOutcome & { report: CuratorPaginationReport; warnings: string[] }>;
@ -350,7 +349,8 @@ export interface IngestBundleRunnerDeps {
registry: SourceAdapterRegistryPort;
diffSetService: DiffSetComputerPort;
sessionWorktreeService: IngestSessionWorktreePort;
agentRunner: AgentRunnerService;
agentRunner: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
gitService: GitService;
lockingService: IngestLockPort;
storage: IngestStoragePort;

23
packages/context/src/ingest/stages/build-reconcile-context.test.ts
View file

@ -141,26 +141,17 @@ describe('buildReconcileToolSet', () => {
toolsetTools: { sl_write_source: { description: 'sl write', inputSchema: {} as any, execute: slWrite } as any },
});
const correction = await toolSet.sl_write_source.execute?.(
{ connectionId: 'warehouse', sourceName: 'accounts' },
{ toolCallId: 't1' } as any,
);
const correction = await toolSet.sl_write_source.execute?.({ connectionId: 'warehouse', sourceName: 'accounts' });
expect(slWrite).not.toHaveBeenCalled();
expect(correction).toMatchObject({ structured: { success: false, reason: 'verification_ledger_required' } });
await toolSet.record_verification_ledger.execute?.(
{
summary: 'Verified warehouse.accounts with entity_details.',
verifiedIdentifiers: ['warehouse.accounts'],
unverifiedIdentifiers: [],
},
{ toolCallId: 't2' } as any,
);
const written = await toolSet.sl_write_source.execute?.(
{ connectionId: 'warehouse', sourceName: 'accounts' },
{ toolCallId: 't3' } as any,
);
await toolSet.record_verification_ledger.execute?.({
summary: 'Verified warehouse.accounts with entity_details.',
verifiedIdentifiers: ['warehouse.accounts'],
unverifiedIdentifiers: [],
});
const written = await toolSet.sl_write_source.execute?.({ connectionId: 'warehouse', sourceName: 'accounts' });
expect(slWrite).toHaveBeenCalledTimes(1);
expect(written).toMatchObject({ structured: { success: true } });

24
packages/context/src/ingest/stages/build-reconcile-context.ts
View file

@ -1,5 +1,5 @@
import type { Tool, ToolSet } from 'ai';
import { buildCanonicalPinsPromptBlock, type CanonicalPin } from '../canonical-pins.js';
import type { KtxRuntimeToolSet } from '../../llm/index.js';
import {
createVerificationLedgerState,
VERIFICATION_LEDGER_PROMPT,
@ -181,19 +181,19 @@ export function buildReconcileUserPrompt(
}
export interface ReconcileToolSetInput {
loadSkillTool: Record<string, Tool>;
stageListTool: Record<string, Tool>;
stageDiffTool: Record<string, Tool>;
evictionListTool: Record<string, Tool>;
emitConflictResolutionTool: Record<string, Tool>;
emitEvictionDecisionTool: Record<string, Tool>;
emitArtifactResolutionTool: Record<string, Tool>;
emitUnmappedFallbackTool: Record<string, Tool>;
readRawSpanTool: Record<string, Tool>;
toolsetTools: ToolSet;
loadSkillTool: KtxRuntimeToolSet;
stageListTool: KtxRuntimeToolSet;
stageDiffTool: KtxRuntimeToolSet;
evictionListTool: KtxRuntimeToolSet;
emitConflictResolutionTool: KtxRuntimeToolSet;
emitEvictionDecisionTool: KtxRuntimeToolSet;
emitArtifactResolutionTool: KtxRuntimeToolSet;
emitUnmappedFallbackTool: KtxRuntimeToolSet;
readRawSpanTool: KtxRuntimeToolSet;
toolsetTools: KtxRuntimeToolSet;
}
export function buildReconcileToolSet(input: ReconcileToolSetInput): ToolSet {
export function buildReconcileToolSet(input: ReconcileToolSetInput): KtxRuntimeToolSet {
const state = createVerificationLedgerState();
return withVerificationLedger(
{

17
packages/context/src/ingest/stages/build-wu-context.test.ts
View file

@ -87,21 +87,18 @@ describe('buildWuToolSet', () => {
toolsetTools: { wiki_write: { description: 'write', inputSchema: {} as any, execute: wikiWrite } as any },
});
const correction = await toolSet.wiki_write.execute?.({ key: 'customer-rules' }, { toolCallId: 't1' } as any);
const correction = await toolSet.wiki_write.execute?.({ key: 'customer-rules' });
expect(wikiWrite).not.toHaveBeenCalled();
expect(correction).toMatchObject({ structured: { success: false, reason: 'verification_ledger_required' } });
expect(String((correction as any).markdown)).toContain('record_verification_ledger');
await toolSet.record_verification_ledger.execute?.(
{
summary: 'No warehouse identifiers will be emitted in this wiki write.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
},
{ toolCallId: 't2' } as any,
);
const written = await toolSet.wiki_write.execute?.({ key: 'customer-rules' }, { toolCallId: 't3' } as any);
await toolSet.record_verification_ledger.execute?.({
summary: 'No warehouse identifiers will be emitted in this wiki write.',
verifiedIdentifiers: [],
unverifiedIdentifiers: [],
});
const written = await toolSet.wiki_write.execute?.({ key: 'customer-rules' });
expect(wikiWrite).toHaveBeenCalledTimes(1);
expect(written).toMatchObject({ structured: { success: true } });

27
packages/context/src/ingest/stages/build-wu-context.ts
View file

@ -1,6 +1,6 @@
import type { Tool, ToolSet } from 'ai';
import { buildCanonicalPinsPromptBlock, type CanonicalPin } from '../canonical-pins.js';
import { createLookerQueryToSlTool } from '../adapters/looker/tools/looker-query-to-sl.tool.js';
import { createRuntimeToolDescriptorFromAiTool, type KtxRuntimeToolSet } from '../../llm/index.js';
import type { IngestProvenanceRow } from '../ports.js';
import { createReadRawFileTool } from '../tools/read-raw-file.tool.js';
import { createReadRawSpanTool } from '../tools/read-raw-span.tool.js';
@ -88,12 +88,12 @@ export interface BuildWuToolSetInput {
sourceKey?: string;
stagedDir: string;
wu: WorkUnit;
loadSkillTool: Record<string, Tool>;
emitUnmappedFallbackTool: Record<string, Tool>;
toolsetTools: ToolSet;
loadSkillTool: KtxRuntimeToolSet;
emitUnmappedFallbackTool: KtxRuntimeToolSet;
toolsetTools: KtxRuntimeToolSet;
}
function withoutWriteSlTools(toolset: ToolSet, wu: WorkUnit): ToolSet {
function withoutWriteSlTools(toolset: KtxRuntimeToolSet, wu: WorkUnit): KtxRuntimeToolSet {
if (!wu.slDisallowed) {
return toolset;
}
@ -103,9 +103,12 @@ function withoutWriteSlTools(toolset: ToolSet, wu: WorkUnit): ToolSet {
return next;
}
export function buildWuToolSet(input: BuildWuToolSetInput): ToolSet {
export function buildWuToolSet(input: BuildWuToolSetInput): KtxRuntimeToolSet {
const allowedPaths = new Set<string>([...input.wu.rawFiles, ...input.wu.dependencyPaths]);
const lookerTools: ToolSet = input.sourceKey === 'looker' ? { looker_query_to_sl: createLookerQueryToSlTool() } : {};
const lookerTools: KtxRuntimeToolSet =
input.sourceKey === 'looker'
? { looker_query_to_sl: createRuntimeToolDescriptorFromAiTool('looker_query_to_sl', createLookerQueryToSlTool()) }
: {};
const state = createVerificationLedgerState();
return withVerificationLedger(
withoutWriteSlTools(
@ -114,8 +117,14 @@ export function buildWuToolSet(input: BuildWuToolSetInput): ToolSet {
...lookerTools,
...input.loadSkillTool,
...input.emitUnmappedFallbackTool,
read_raw_file: createReadRawFileTool({ stagedDir: input.stagedDir, allowedPaths }),
read_raw_span: createReadRawSpanTool({ stagedDir: input.stagedDir, allowedPaths }),
read_raw_file: createRuntimeToolDescriptorFromAiTool(
'read_raw_file',
createReadRawFileTool({ stagedDir: input.stagedDir, allowedPaths }),
),
read_raw_span: createRuntimeToolDescriptorFromAiTool(
'read_raw_span',
createReadRawSpanTool({ stagedDir: input.stagedDir, allowedPaths }),
),
},
input.wu,
),

7
packages/context/src/ingest/stages/stage-3-work-units.ts
View file

@ -1,6 +1,5 @@
import type { AgentRunnerService } from '@ktx/context/agent';
import type { KtxModelRole } from '@ktx/llm';
import type { Tool } from 'ai';
import type { AgentRunnerPort, KtxRuntimeToolSet } from '@ktx/context';
import type { CaptureSession, MemoryAction } from '../../memory/index.js';
import { listTouchedSlSources, type TouchedSlSource } from '../../tools/index.js';
import type { WorkUnit } from '../types.js';
@ -14,12 +13,12 @@ export interface TouchedValidationResult {
export interface WorkUnitExecutionDeps {
sessionWorktreeGit: { revParseHead(): Promise<string | null> };
agentRunner: AgentRunnerService;
agentRunner: AgentRunnerPort;
validateTouchedSources: (touched: TouchedSlSource[]) => Promise<TouchedValidationResult>;
resetHardTo: (targetSha: string) => Promise<void>;
buildSystemPrompt: (wu: WorkUnit) => string;
buildUserPrompt: (wu: WorkUnit) => string;
buildToolSet: (wu: WorkUnit) => Record<string, Tool>;
buildToolSet: (wu: WorkUnit) => KtxRuntimeToolSet;
captureSession: CaptureSession;
sessionActions: MemoryAction[];
modelRole: KtxModelRole;

7
packages/context/src/ingest/stages/stage-4-reconciliation.ts
View file

@ -1,16 +1,15 @@
import type { AgentRunnerService } from '@ktx/context/agent';
import type { AgentRunnerPort, KtxRuntimeToolSet } from '@ktx/context';
import type { KtxModelRole } from '@ktx/llm';
import type { ToolSet } from 'ai';
import type { EvictionUnit } from '../types.js';
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: () => ToolSet;
buildToolSet: () => KtxRuntimeToolSet;
modelRole: KtxModelRole;
stepBudget: number;
sourceKey: string;

13
packages/context/src/ingest/tools/tool-call-logger.ts
View file

@ -1,6 +1,6 @@
import { appendFile, mkdir } from 'node:fs/promises';
import { dirname } from 'node:path';
import type { ToolExecuteFunction, ToolExecutionOptions, ToolSet } from 'ai';
import type { KtxRuntimeToolSet } from '../../llm/index.js';
export interface ToolCallLogEntry {
ts: string;
@ -31,7 +31,7 @@ interface ToolCallLoggerOptions {
* sequential (`generateText` awaits each tool result), so per-WU files are
* effectively single-writer and lines land in call order.
*/
export function wrapToolsWithLogger<T extends ToolSet>(
export function wrapToolsWithLogger<T extends KtxRuntimeToolSet>(
tools: T,
logFilePath: string,
wuKey: string,
@ -44,17 +44,13 @@ export function wrapToolsWithLogger<T extends ToolSet>(
wrapped[name] = original;
continue;
}
const wrappedExecute: ToolExecuteFunction<unknown, unknown> = async (
input: unknown,
opts: ToolExecutionOptions,
) => {
const wrappedExecute = async (input: unknown) => {
const start = Date.now();
try {
const output = await (originalExecute as ToolExecuteFunction<unknown, unknown>)(input, opts);
const output = await originalExecute(input);
const entry: ToolCallLogEntry = {
ts: new Date().toISOString(),
wuKey,
toolCallId: opts.toolCallId,
toolName: name,
durationMs: Date.now() - start,
input,
@ -67,7 +63,6 @@ export function wrapToolsWithLogger<T extends ToolSet>(
const entry: ToolCallLogEntry = {
ts: new Date().toISOString(),
wuKey,
toolCallId: opts.toolCallId,
toolName: name,
durationMs: Date.now() - start,
input,

20
packages/context/src/ingest/tools/verification-ledger.tool.ts
View file

@ -1,5 +1,5 @@
import { tool, type ToolExecuteFunction, type ToolExecutionOptions, type ToolSet } from 'ai';
import { z } from 'zod';
import type { KtxRuntimeToolDescriptor, KtxRuntimeToolSet } from '../../llm/index.js';
const verificationLedgerInputSchema = z.object({
summary: z.string().min(1).max(2000),
@ -37,22 +37,19 @@ export function createVerificationLedgerState(): VerificationLedgerState {
return { entries: [] };
}
export function withVerificationLedger(tools: ToolSet, state: VerificationLedgerState): ToolSet {
const wrapped: ToolSet = {};
export function withVerificationLedger(tools: KtxRuntimeToolSet, state: VerificationLedgerState): KtxRuntimeToolSet {
const wrapped: KtxRuntimeToolSet = {};
for (const [name, original] of Object.entries(tools)) {
if (!WRITE_TOOL_NAMES.has(name) || typeof original.execute !== 'function') {
wrapped[name] = original;
continue;
}
const originalExecute = original.execute;
const guardedExecute: ToolExecuteFunction<unknown, unknown> = async (
input: unknown,
opts: ToolExecutionOptions,
) => {
const guardedExecute = async (input: unknown) => {
if (state.entries.length === 0) {
return verificationRequiredOutput(name);
}
return (originalExecute as ToolExecuteFunction<unknown, unknown>)(input, opts);
return originalExecute(input);
};
wrapped[name] = { ...original, execute: guardedExecute };
}
@ -60,8 +57,9 @@ export function withVerificationLedger(tools: ToolSet, state: VerificationLedger
return wrapped;
}
function createRecordVerificationLedgerTool(state: VerificationLedgerState) {
return tool({
function createRecordVerificationLedgerTool(state: VerificationLedgerState): KtxRuntimeToolDescriptor {
return {
name: 'record_verification_ledger',
description:
'Record the pre-write verification ledger required by loaded ingest skills. Call this before wiki/SL/fallback writes to state what was verified, which tool calls support it, and what remains intentionally unverified.',
inputSchema: verificationLedgerInputSchema,
@ -78,7 +76,7 @@ function createRecordVerificationLedgerTool(state: VerificationLedgerState) {
structured: { success: true, entry },
};
},
});
};
}
function verificationRequiredOutput(toolName: string) {

164
packages/context/src/llm/ai-sdk-runtime.ts Normal file
View file

@ -0,0 +1,164 @@
import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider } from '@ktx/llm';
import { generateText, Output, stepCountIs, type FlexibleSchema, type TelemetrySettings, type ToolSet } from 'ai';
import type { z } from 'zod';
import { noopLogger, type KtxLogger } from '../core/index.js';
import { summarizeKtxLlmDebugRequest, type KtxLlmDebugRequestRecorder } from './debug-request-recorder.js';
import { createAiSdkToolSet } from './runtime-tools.js';
import type {
KtxGenerateObjectInput,
KtxGenerateTextInput,
KtxLlmRuntimePort,
RunLoopParams,
RunLoopResult,
} from './runtime-port.js';
export interface AgentTelemetryPort {
createTelemetry(tags: Record<string, string>): TelemetrySettings;
}
export interface AiSdkKtxLlmRuntimeDeps {
llmProvider: KtxLlmProvider;
telemetry?: AgentTelemetryPort;
logger?: KtxLogger;
debugRequestRecorder?: KtxLlmDebugRequestRecorder;
}
function hasTools(tools: Record<string, unknown>): boolean {
return Object.keys(tools).length > 0;
}
export class AiSdkKtxLlmRuntime implements KtxLlmRuntimePort {
private readonly logger: KtxLogger;
constructor(private readonly deps: AiSdkKtxLlmRuntimeDeps) {
this.logger = deps.logger ?? noopLogger;
}
async generateText(input: KtxGenerateTextInput): Promise<string> {
const model = this.deps.llmProvider.getModel(input.role);
if ((model as { provider?: string }).provider === 'deterministic') {
return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`;
}
const tools = createAiSdkToolSet(input.tools ?? {});
const built = new KtxMessageBuilder(this.deps.llmProvider).wrapSimple({
system: input.system,
messages: [{ role: 'user', content: input.prompt }],
tools,
model,
});
const split = splitKtxSystemMessages(built.messages);
const result = await generateText({
model,
temperature: input.temperature ?? 0,
...(split.system ? { system: split.system } : {}),
messages: split.messages,
tools: built.tools as ToolSet,
...(hasTools(tools)
? {
experimental_repairToolCall: this.deps.llmProvider.repairToolCallHandler({
source: `ktx-${input.role}`,
}),
}
: {}),
});
if (typeof result.text !== 'string') {
throw new Error('KTX LLM text generation returned no text');
}
return result.text;
}
async generateObject<TOutput, TSchema extends z.ZodType<TOutput>>(
input: KtxGenerateObjectInput<TOutput, TSchema>,
): Promise<TOutput> {
const model = this.deps.llmProvider.getModel(input.role);
const tools = createAiSdkToolSet(input.tools ?? {});
const built = new KtxMessageBuilder(this.deps.llmProvider).wrapSimple({
system: input.system,
messages: [{ role: 'user', content: input.prompt }],
tools,
model,
});
const split = splitKtxSystemMessages(built.messages);
const result = await generateText({
model,
temperature: input.temperature ?? 0,
...(split.system ? { system: split.system } : {}),
messages: split.messages,
tools: built.tools as ToolSet,
...(hasTools(tools)
? {
experimental_repairToolCall: this.deps.llmProvider.repairToolCallHandler({
source: `ktx-${input.role}`,
}),
}
: {}),
output: Output.object({ schema: input.schema as unknown as FlexibleSchema<TOutput> }),
});
if (result.output == null) {
throw new Error('KTX LLM object generation returned no output');
}
return result.output as TOutput;
}
async runAgentLoop(params: RunLoopParams): Promise<RunLoopResult> {
let stepIndex = 0;
try {
const model = this.deps.llmProvider.getModel(params.modelRole);
const tools = createAiSdkToolSet(params.toolSet);
const builder = new KtxMessageBuilder(this.deps.llmProvider);
const built = builder.wrapSimple({
system: params.systemPrompt,
messages: [{ role: 'user', content: params.userPrompt }],
tools,
model,
});
const promptMessages = splitKtxSystemMessages(built.messages);
await this.deps.debugRequestRecorder?.record(
summarizeKtxLlmDebugRequest({
operationName: params.telemetryTags.operationName ?? 'ktx-agent-runner',
source: params.telemetryTags.source,
jobId: params.telemetryTags.jobId,
unitKey: params.telemetryTags.unitKey,
modelRole: params.modelRole,
modelId: (model as { modelId?: string }).modelId ?? params.modelRole,
messages: built.messages,
tools: built.tools as Record<string, { providerOptions?: unknown }>,
}),
);
await generateText({
model,
temperature: 0,
stopWhen: stepCountIs(params.stepBudget),
experimental_telemetry: this.deps.telemetry?.createTelemetry(params.telemetryTags) ?? this.deps.llmProvider.telemetryConfig(),
experimental_repairToolCall: this.deps.llmProvider.repairToolCallHandler({
source: params.telemetryTags.operationName ?? 'ktx-agent-runner',
}),
...(promptMessages.system ? { system: promptMessages.system } : {}),
messages: promptMessages.messages,
tools: built.tools as ToolSet,
onStepFinish: async () => {
stepIndex += 1;
if (!params.onStepFinish) {
return;
}
try {
await params.onStepFinish({ stepIndex, stepBudget: params.stepBudget });
} catch (err) {
this.logger.warn(
`[agent-runner] onStepFinish callback threw; ignoring: ${
err instanceof Error ? err.message : String(err)
}`,
);
}
},
});
return { stopReason: 'natural' };
} catch (error) {
const err = error instanceof Error ? error : new Error(String(error));
this.logger.warn(`[agent-runner] loop failed: ${err.message}`);
return { stopReason: 'error', error: err };
}
}
}

19
packages/context/src/llm/claude-code-env.test.ts Normal file
View file

@ -0,0 +1,19 @@
import { describe, expect, it } from 'vitest';
import { CLAUDE_CODE_PROVIDER_ENV_DENYLIST, createKtxClaudeCodeEnv } from './claude-code-env.js';
describe('createKtxClaudeCodeEnv', () => {
it('strips provider-routing credentials from the Claude Code child environment', () => {
const seeded = Object.fromEntries(CLAUDE_CODE_PROVIDER_ENV_DENYLIST.map((key) => [key, `${key}-value`]));
const env = createKtxClaudeCodeEnv({
...seeded,
PATH: '/usr/bin',
HOME: '/Users/test',
});
for (const key of CLAUDE_CODE_PROVIDER_ENV_DENYLIST) {
expect(env).not.toHaveProperty(key);
}
expect(env.PATH).toBe('/usr/bin');
expect(env.HOME).toBe('/Users/test');
});
});

23
packages/context/src/llm/claude-code-env.ts Normal file
View file

@ -0,0 +1,23 @@
export const CLAUDE_CODE_PROVIDER_ENV_DENYLIST = [
'ANTHROPIC_API_KEY',
'ANTHROPIC_AUTH_TOKEN',
'ANTHROPIC_BASE_URL',
'ANTHROPIC_MODEL',
'ANTHROPIC_VERTEX_PROJECT_ID',
'CLOUD_ML_REGION',
'GOOGLE_APPLICATION_CREDENTIALS',
'GOOGLE_CLOUD_PROJECT',
'AWS_ACCESS_KEY_ID',
'AWS_SECRET_ACCESS_KEY',
'AWS_SESSION_TOKEN',
'AWS_REGION',
'AWS_PROFILE',
'CLAUDE_CODE_USE_BEDROCK',
'CLAUDE_CODE_USE_VERTEX',
] as const;
const DENYLIST = new Set<string>(CLAUDE_CODE_PROVIDER_ENV_DENYLIST);
export function createKtxClaudeCodeEnv(env: NodeJS.ProcessEnv = process.env): Record<string, string | undefined> {
return Object.fromEntries(Object.entries(env).filter(([key]) => !DENYLIST.has(key)));
}

17
packages/context/src/llm/claude-code-models.test.ts Normal file
View file

@ -0,0 +1,17 @@
import { describe, expect, it } from 'vitest';
import { resolveClaudeCodeModel } from './claude-code-models.js';
describe('resolveClaudeCodeModel', () => {
it.each([
['sonnet', 'claude-sonnet-4-6'],
['opus', 'claude-opus-4-7'],
['haiku', 'claude-haiku-4-5'],
['claude-sonnet-4-6', 'claude-sonnet-4-6'],
])('maps %s to %s', (input, expected) => {
expect(resolveClaudeCodeModel(input)).toBe(expected);
});
it('rejects unsupported aliases', () => {
expect(() => resolveClaudeCodeModel('gpt-5')).toThrow('Unsupported Claude Code model');
});
});

19
packages/context/src/llm/claude-code-models.ts Normal file
View file

@ -0,0 +1,19 @@
const CLAUDE_CODE_MODEL_ALIASES: Record<string, string> = {
sonnet: 'claude-sonnet-4-6',
opus: 'claude-opus-4-7',
haiku: 'claude-haiku-4-5',
};
const FULL_MODEL_ID = /^claude-(sonnet|opus|haiku)-[0-9]+-[0-9]+$/;
export function resolveClaudeCodeModel(model: string): string {
const normalized = model.trim();
const alias = CLAUDE_CODE_MODEL_ALIASES[normalized];
if (alias) {
return alias;
}
if (FULL_MODEL_ID.test(normalized)) {
return normalized;
}
throw new Error(`Unsupported Claude Code model "${model}". Use sonnet, opus, haiku, or a claude-* model id.`);
}

464
packages/context/src/llm/claude-code-runtime.test.ts Normal file
View file

@ -0,0 +1,464 @@
import { describe, expect, it, vi } from 'vitest';
import { z } from 'zod';
import type { SDKMessage } from '@anthropic-ai/claude-agent-sdk';
import { ClaudeCodeKtxLlmRuntime, mapClaudeCodeStopReason, runClaudeCodeAuthProbe } from './claude-code-runtime.js';
async function* stream(messages: SDKMessage[]): AsyncGenerator<SDKMessage, void> {
for (const message of messages) {
yield message;
}
}
function initMessage(overrides: Partial<Extract<SDKMessage, { type: 'system'; subtype: 'init' }>> = {}): Extract<
SDKMessage,
{ type: 'system'; subtype: 'init' }
> {
return {
type: 'system',
subtype: 'init',
apiKeySource: 'none' as never, // pragma: allowlist secret
claude_code_version: '0.3.142',
cwd: '/tmp/project',
tools: [],
mcp_servers: [],
model: 'claude-sonnet-4-6',
permissionMode: 'dontAsk',
slash_commands: [],
output_style: 'default',
skills: [],
plugins: [],
uuid: '00000000-0000-4000-8000-000000000001',
session_id: 'session-id',
...overrides,
};
}
function resultMessage(overrides: Partial<Extract<SDKMessage, { type: 'result' }>> = {}): Extract<
SDKMessage,
{ type: 'result' }
> {
return {
type: 'result',
subtype: 'success',
duration_ms: 1,
duration_api_ms: 1,
is_error: false,
num_turns: 1,
result: 'ok',
stop_reason: null,
total_cost_usd: 0,
usage: {} as never,
modelUsage: {},
permission_denials: [],
errors: [],
uuid: '00000000-0000-4000-8000-000000000002',
session_id: 'session-id',
...overrides,
} as Extract<SDKMessage, { type: 'result' }>;
}
describe('ClaudeCodeKtxLlmRuntime', () => {
it('passes isolation options and scrubbed env to text generation', async () => {
const query = vi.fn((_input: any) => stream([initMessage(), resultMessage({ result: 'hello' })]));
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: { ANTHROPIC_API_KEY: 'sk-ant-test', PATH: '/usr/bin' }, // pragma: allowlist secret
});
await expect(runtime.generateText({ role: 'default', prompt: 'say hello' })).resolves.toBe('hello');
expect(query).toHaveBeenCalledWith({
prompt: 'say hello',
options: expect.objectContaining({
cwd: '/tmp/project',
model: 'claude-sonnet-4-6',
maxTurns: 1,
settingSources: [],
skills: [],
plugins: [],
tools: [],
allowedTools: [],
permissionMode: 'dontAsk',
persistSession: false,
env: expect.not.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test' }),
}),
});
});
it('validates structured output with the caller schema', async () => {
const schema = z.object({ answer: z.string() });
const query = vi.fn((_input: any) => stream([initMessage(), resultMessage({ structured_output: { answer: 'yes' } })]));
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
});
await expect(runtime.generateObject({ role: 'default', prompt: 'json', schema })).resolves.toEqual({ answer: 'yes' });
expect(query.mock.calls[0][0].options.outputFormat).toMatchObject({
type: 'json_schema',
schema: expect.objectContaining({ type: 'object' }),
});
});
it('registers only exact KTX MCP tool ids and denies non-KTX tools', async () => {
const query = vi.fn((_input: any) =>
stream([
initMessage({ tools: ['mcp__ktx__load_skill'], mcp_servers: [{ name: 'ktx', status: 'connected' }] }),
{
type: 'assistant',
message: { role: 'assistant', content: [] },
parent_tool_use_id: null,
uuid: '00000000-0000-4000-8000-000000000003',
session_id: 'session-id',
} as unknown as SDKMessage,
resultMessage({ subtype: 'error_max_turns', is_error: true }),
]),
);
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
});
const onStepFinish = vi.fn();
await runtime.runAgentLoop({
modelRole: 'default',
systemPrompt: 'system',
userPrompt: 'user',
toolSet: {
load_skill: {
name: 'load_skill',
description: 'Load skill.',
inputSchema: z.object({ name: z.string() }),
execute: async () => ({ markdown: 'loaded' }),
},
},
stepBudget: 1,
telemetryTags: { operationName: 'test' },
onStepFinish,
});
const options = query.mock.calls[0][0].options;
expect(options.allowedTools).toEqual(['mcp__ktx__load_skill']);
expect(await options.canUseTool('mcp__ktx__load_skill', {}, { signal: new AbortController().signal, toolUseID: '1' })).toEqual({
behavior: 'allow',
toolUseID: '1',
});
expect(await options.canUseTool('Bash', {}, { signal: new AbortController().signal, toolUseID: '2' })).toMatchObject({
behavior: 'deny',
toolUseID: '2',
});
expect(onStepFinish).toHaveBeenCalledWith({ stepIndex: 1, stepBudget: 1 });
});
it('treats host-discovered commands skills and agents as non-fatal init metadata for text and auth probe', async () => {
const hostDiscoveredInit = initMessage({
slash_commands: ['/help', '/compact', '/clear', '/user-command'],
skills: ['pdf', 'docx'],
agents: ['claude', 'Explore', 'general-purpose'],
});
const textQuery = vi.fn((_input: any) => stream([hostDiscoveredInit, resultMessage({ result: 'hello' })]));
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query: textQuery,
env: { ANTHROPIC_API_KEY: 'sk-ant-test', PATH: '/usr/bin' }, // pragma: allowlist secret
});
await expect(runtime.generateText({ role: 'default', prompt: 'say hello' })).resolves.toBe('hello');
const textOptions = textQuery.mock.calls[0][0].options;
expect(textOptions).toMatchObject({
settingSources: [],
skills: [],
plugins: [],
tools: [],
allowedTools: [],
permissionMode: 'dontAsk',
persistSession: false,
env: expect.not.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test' }),
});
expect(textOptions.disallowedTools).toEqual(expect.arrayContaining(['Agent', 'Task', 'Bash']));
expect(await textOptions.canUseTool('Agent', {}, { signal: new AbortController().signal, toolUseID: 'agent' })).toMatchObject({
behavior: 'deny',
toolUseID: 'agent',
});
expect(await textOptions.canUseTool('Skill', {}, { signal: new AbortController().signal, toolUseID: 'skill' })).toMatchObject({
behavior: 'deny',
toolUseID: 'skill',
});
expect(
await textOptions.canUseTool('SlashCommand', {}, { signal: new AbortController().signal, toolUseID: 'slash' }),
).toMatchObject({
behavior: 'deny',
toolUseID: 'slash',
});
const probeQuery = vi.fn((_input: any) => stream([hostDiscoveredInit, resultMessage({ result: 'ok' })]));
await expect(
runClaudeCodeAuthProbe({
projectDir: '/tmp/project',
model: 'sonnet',
query: probeQuery,
env: { ANTHROPIC_AUTH_TOKEN: 'token', HOME: '/Users/test' },
}),
).resolves.toEqual({ ok: true });
expect(probeQuery.mock.calls[0][0].options).toMatchObject({
settingSources: [],
skills: [],
plugins: [],
tools: [],
allowedTools: [],
permissionMode: 'dontAsk',
persistSession: false,
env: expect.objectContaining({ HOME: '/Users/test' }),
});
expect(probeQuery.mock.calls[0][0].options.env).not.toEqual(
expect.objectContaining({ ANTHROPIC_AUTH_TOKEN: 'token' }),
);
});
it('allows host-discovered context during agent loops while requiring exact KTX MCP tools and servers', async () => {
const query = vi.fn((_input: any) =>
stream([
initMessage({
tools: ['mcp__ktx__load_skill'],
mcp_servers: [{ name: 'ktx', status: 'connected' }],
slash_commands: ['/help', '/compact', '/clear'],
skills: ['memory-agent', 'doc-reader'],
agents: ['claude', 'Plan', 'Explore'],
}),
{
type: 'assistant',
message: { role: 'assistant', content: [] },
parent_tool_use_id: null,
uuid: '00000000-0000-4000-8000-000000000006',
session_id: 'session-id',
} as unknown as SDKMessage,
resultMessage({ subtype: 'error_max_turns', is_error: true }),
]),
);
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
});
await expect(
runtime.runAgentLoop({
modelRole: 'default',
systemPrompt: 'system',
userPrompt: 'user',
toolSet: {
load_skill: {
name: 'load_skill',
description: 'Load skill.',
inputSchema: z.object({ name: z.string() }),
execute: async () => ({ markdown: 'loaded' }),
},
},
stepBudget: 1,
telemetryTags: { operationName: 'test' },
}),
).resolves.toEqual({ stopReason: 'budget' });
const options = query.mock.calls[0][0].options;
expect(options.allowedTools).toEqual(['mcp__ktx__load_skill']);
expect(await options.canUseTool('mcp__ktx__load_skill', {}, { signal: new AbortController().signal, toolUseID: '1' })).toEqual({
behavior: 'allow',
toolUseID: '1',
});
expect(await options.canUseTool('Task', {}, { signal: new AbortController().signal, toolUseID: '2' })).toMatchObject({
behavior: 'deny',
toolUseID: '2',
});
expect(await options.canUseTool('Skill', {}, { signal: new AbortController().signal, toolUseID: '3' })).toMatchObject({
behavior: 'deny',
toolUseID: '3',
});
});
it('still rejects unexpected tools, missing KTX tools, plugins, and non-KTX MCP servers from init messages', async () => {
const query = vi.fn((_input: any) =>
stream([
initMessage({
tools: ['Bash'],
mcp_servers: [{ name: 'filesystem', status: 'connected' }],
plugins: [{ name: 'host-plugin', path: '/tmp/plugin' }],
}),
resultMessage({ result: 'hello' }),
]),
);
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
});
await expect(
runtime.generateText({
role: 'default',
prompt: 'say hello',
tools: {
load_skill: {
name: 'load_skill',
description: 'Load skill.',
inputSchema: z.object({ name: z.string() }),
execute: async () => ({ markdown: 'loaded' }),
},
},
}),
).rejects.toThrow(
/Claude Code runtime isolation failed: .*tools=Bash.*missing_tools=mcp__ktx__load_skill.*mcp_servers=filesystem.*plugins=host-plugin/,
);
});
it('passes scrubbed env to object generation and agent loops', async () => {
const schema = z.object({ answer: z.string() });
const objectQuery = vi.fn((_input: any) =>
stream([initMessage(), resultMessage({ structured_output: { answer: 'yes' } })]),
);
const objectRuntime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query: objectQuery,
env: { ANTHROPIC_API_KEY: 'sk-ant-test', AWS_PROFILE: 'prod', PATH: '/usr/bin' }, // pragma: allowlist secret
});
await expect(objectRuntime.generateObject({ role: 'default', prompt: 'json', schema })).resolves.toEqual({
answer: 'yes',
});
expect(objectQuery.mock.calls[0][0].options.env).toEqual(expect.objectContaining({ PATH: '/usr/bin' }));
expect(objectQuery.mock.calls[0][0].options.env).not.toEqual(
expect.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test', AWS_PROFILE: 'prod' }), // pragma: allowlist secret
);
const agentQuery = vi.fn((_input: any) =>
stream([
initMessage({ tools: ['mcp__ktx__load_skill'], mcp_servers: [{ name: 'ktx', status: 'connected' }] }),
{
type: 'assistant',
message: { role: 'assistant', content: [] },
parent_tool_use_id: null,
uuid: '00000000-0000-4000-8000-000000000004',
session_id: 'session-id',
} as unknown as SDKMessage,
resultMessage({ subtype: 'error_max_turns', is_error: true }),
]),
);
const agentRuntime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query: agentQuery,
env: { ANTHROPIC_AUTH_TOKEN: 'token', CLAUDE_CODE_USE_VERTEX: '1', HOME: '/Users/test' },
});
await agentRuntime.runAgentLoop({
modelRole: 'default',
systemPrompt: 'system',
userPrompt: 'user',
toolSet: {
load_skill: {
name: 'load_skill',
description: 'Load skill.',
inputSchema: z.object({ name: z.string() }),
execute: async () => ({ markdown: 'loaded' }),
},
},
stepBudget: 1,
telemetryTags: { operationName: 'test' },
});
expect(agentQuery.mock.calls[0][0].options.env).toEqual(expect.objectContaining({ HOME: '/Users/test' }));
expect(agentQuery.mock.calls[0][0].options.env).not.toEqual(
expect.objectContaining({ ANTHROPIC_AUTH_TOKEN: 'token', CLAUDE_CODE_USE_VERTEX: '1' }),
);
});
it('logs and ignores onStepFinish callback errors', async () => {
const query = vi.fn((_input: any) =>
stream([
initMessage(),
{
type: 'assistant',
message: { role: 'assistant', content: [] },
parent_tool_use_id: null,
uuid: '00000000-0000-4000-8000-000000000005',
session_id: 'session-id',
} as unknown as SDKMessage,
resultMessage({ subtype: 'success', terminal_reason: 'completed' }),
]),
);
const logger = {
debug: vi.fn(),
log: vi.fn(),
warn: vi.fn(),
error: vi.fn(),
};
const runtime = new ClaudeCodeKtxLlmRuntime({
projectDir: '/tmp/project',
modelSlots: { default: 'sonnet' },
query,
env: {},
logger,
});
await expect(
runtime.runAgentLoop({
modelRole: 'default',
systemPrompt: 'system',
userPrompt: 'user',
toolSet: {},
stepBudget: 1,
telemetryTags: { operationName: 'test' },
onStepFinish: async () => {
throw new Error('callback exploded');
},
}),
).resolves.toEqual({ stopReason: 'natural' });
expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('callback exploded'));
});
it('maps max-turn terminal reasons to budget', () => {
expect(mapClaudeCodeStopReason(resultMessage({ subtype: 'error_max_turns' }))).toBe('budget');
expect(mapClaudeCodeStopReason(resultMessage({ terminal_reason: 'max_turns' }))).toBe('budget');
expect(mapClaudeCodeStopReason(resultMessage({ stop_reason: 'max_turns' }))).toBe('budget');
expect(mapClaudeCodeStopReason(resultMessage({ subtype: 'success', terminal_reason: 'completed' }))).toBe('natural');
expect(mapClaudeCodeStopReason(resultMessage({ subtype: 'error_during_execution' }))).toBe('error');
});
it('auth probe uses isolation options and a scrubbed env', async () => {
const query = vi.fn((_input: any) => stream([initMessage(), resultMessage({ result: 'ok' })]));
await expect(
runClaudeCodeAuthProbe({ projectDir: '/tmp/project', model: 'sonnet', query, env: { ANTHROPIC_API_KEY: 'sk-ant-test' } }), // pragma: allowlist secret
).resolves.toEqual({ ok: true });
expect(query.mock.calls[0][0].options).toMatchObject({
settingSources: [],
skills: [],
plugins: [],
tools: [],
allowedTools: [],
persistSession: false,
env: expect.not.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test' }),
});
});
it('reports unsupported Claude Code models without framing them as auth failures', async () => {
await expect(
runClaudeCodeAuthProbe({
projectDir: '/tmp/project',
model: 'gpt-5',
query: vi.fn(),
env: {},
}),
).resolves.toEqual({
ok: false,
message: 'Unsupported Claude Code model "gpt-5". Use sonnet, opus, haiku, or a claude-* model id.',
});
});
});

327
packages/context/src/llm/claude-code-runtime.ts Normal file
View file

@ -0,0 +1,327 @@
import {
createSdkMcpServer,
query as defaultQuery,
type Options,
type SDKMessage,
type SDKResultMessage,
} from '@anthropic-ai/claude-agent-sdk';
import { z } from 'zod';
import { noopLogger, type KtxLogger } from '../core/index.js';
import { createKtxClaudeCodeEnv } from './claude-code-env.js';
import { resolveClaudeCodeModel } from './claude-code-models.js';
import { createClaudeSdkTools, mcpToolIds } from './runtime-tools.js';
import type {
KtxGenerateObjectInput,
KtxGenerateTextInput,
KtxLlmRuntimePort,
KtxRuntimeToolSet,
RunLoopParams,
RunLoopResult,
RunLoopStopReason,
} from './runtime-port.js';
type QueryFn = (params: Parameters<typeof defaultQuery>[0]) => AsyncIterable<SDKMessage>;
export interface ClaudeCodeKtxLlmRuntimeDeps {
projectDir: string;
modelSlots: { default: string } & Partial<Record<string, string>>;
query?: QueryFn;
env?: NodeJS.ProcessEnv;
logger?: KtxLogger;
}
const BUILTIN_TOOLS = [
'Agent',
'Task',
'AskUserQuestion',
'Bash',
'Read',
'Edit',
'Write',
'Glob',
'Grep',
'WebFetch',
'WebSearch',
'TodoWrite',
];
function isResult(message: SDKMessage): message is SDKResultMessage {
return message.type === 'result';
}
function resultError(result: SDKResultMessage): Error | undefined {
if (result.subtype === 'success') {
return undefined;
}
const details = result.errors.length > 0 ? `: ${result.errors.join('; ')}` : '';
return new Error(`Claude Code query failed (${result.subtype})${details}`);
}
export function mapClaudeCodeStopReason(result: SDKResultMessage): RunLoopStopReason {
if (result.subtype === 'error_max_turns') {
return 'budget';
}
if (result.terminal_reason === 'max_turns' || result.stop_reason === 'max_turns') {
return 'budget';
}
if (result.subtype === 'success') {
return result.terminal_reason && result.terminal_reason !== 'completed' ? 'error' : 'natural';
}
return 'error';
}
function jsonSchema(schema: z.ZodType): Record<string, unknown> {
return z.toJSONSchema(schema, { target: 'draft-7' }) as Record<string, unknown>;
}
function modelForRole(modelSlots: ClaudeCodeKtxLlmRuntimeDeps['modelSlots'], role: string): string {
return resolveClaudeCodeModel(modelSlots[role] ?? modelSlots.default);
}
function assertInitIsolation(
message: SDKMessage,
allowedToolIds: Set<string>,
expectedMcpServerNames: Set<string>,
): void {
if (message.type !== 'system' || message.subtype !== 'init') {
return;
}
const activeToolIds = new Set(message.tools);
const unexpectedTools = message.tools.filter((toolName) => !allowedToolIds.has(toolName));
const missingTools = [...allowedToolIds].filter((toolName) => !activeToolIds.has(toolName));
const activeMcpServerNames = message.mcp_servers.map((server) => server.name);
const unexpectedMcpServers = activeMcpServerNames.filter((name) => !expectedMcpServerNames.has(name));
const missingMcpServers = [...expectedMcpServerNames].filter((name) => !activeMcpServerNames.includes(name));
const unexpectedPlugins = message.plugins.map((plugin) => plugin.name);
if (
unexpectedTools.length > 0 ||
missingTools.length > 0 ||
unexpectedMcpServers.length > 0 ||
missingMcpServers.length > 0 ||
unexpectedPlugins.length > 0
) {
throw new Error(
`Claude Code runtime isolation failed: tools=${unexpectedTools.join(',') || '(none)'} missing_tools=${
missingTools.join(',') || '(none)'
} mcp_servers=${unexpectedMcpServers.join(',') || '(none)'} missing_mcp_servers=${
missingMcpServers.join(',') || '(none)'
} plugins=${unexpectedPlugins.join(',') || '(none)'} host_slash_commands=${
message.slash_commands.length
} host_skills=${message.skills.length} host_agents=${message.agents?.join(',') || '(none)'}`,
);
}
}
function expectedMcpServerNames(tools: KtxRuntimeToolSet | undefined): Set<string> {
return tools && Object.keys(tools).length > 0 ? new Set(['ktx']) : new Set();
}
function baseOptions(input: {
projectDir: string;
model: string;
env: NodeJS.ProcessEnv | undefined;
maxTurns: number;
tools?: KtxRuntimeToolSet;
}): Options {
const toolIds = mcpToolIds(input.tools ?? {});
const allowedToolIds = new Set(toolIds);
return {
cwd: input.projectDir,
model: input.model,
maxTurns: input.maxTurns,
settingSources: [],
skills: [],
plugins: [],
tools: [],
allowedTools: toolIds,
disallowedTools: BUILTIN_TOOLS,
canUseTool: async (toolName, _toolInput, options) =>
allowedToolIds.has(toolName)
? { behavior: 'allow', toolUseID: options.toolUseID }
: {
behavior: 'deny',
message: `KTX claude-code runtime only permits current KTX MCP tools; denied ${toolName}.`,
toolUseID: options.toolUseID,
},
permissionMode: 'dontAsk',
persistSession: false,
env: createKtxClaudeCodeEnv(input.env),
...(input.tools && Object.keys(input.tools).length > 0
? { mcpServers: { ktx: createSdkMcpServer({ name: 'ktx', tools: createClaudeSdkTools(input.tools) }) } }
: {}),
};
}
async function collectResult(params: {
query: QueryFn;
prompt: string;
options: Options;
allowedToolIds: Set<string>;
expectedMcpServerNames: Set<string>;
onAssistantTurn?: () => Promise<void>;
}): Promise<SDKResultMessage> {
let result: SDKResultMessage | undefined;
for await (const message of params.query({ prompt: params.prompt, options: params.options })) {
assertInitIsolation(message, params.allowedToolIds, params.expectedMcpServerNames);
if (message.type === 'assistant' && message.parent_tool_use_id === null) {
await params.onAssistantTurn?.();
}
if (isResult(message)) {
result = message;
}
}
if (!result) {
throw new Error('Claude Code query returned no result message');
}
return result;
}
export class ClaudeCodeKtxLlmRuntime implements KtxLlmRuntimePort {
private readonly runQuery: QueryFn;
private readonly logger: KtxLogger;
constructor(private readonly deps: ClaudeCodeKtxLlmRuntimeDeps) {
this.runQuery = deps.query ?? defaultQuery;
this.logger = deps.logger ?? noopLogger;
}
async generateText(input: KtxGenerateTextInput): Promise<string> {
const options = baseOptions({
projectDir: this.deps.projectDir,
model: modelForRole(this.deps.modelSlots, input.role),
env: this.deps.env,
maxTurns: 1,
tools: input.tools,
});
const result = await collectResult({
query: this.runQuery,
prompt: [input.system, input.prompt].filter(Boolean).join('\n\n'),
options,
allowedToolIds: new Set(mcpToolIds(input.tools ?? {})),
expectedMcpServerNames: expectedMcpServerNames(input.tools),
});
const error = resultError(result);
if (error) {
throw error;
}
if (result.subtype !== 'success') {
throw new Error(`Claude Code query failed (${result.subtype})`);
}
return result.result;
}
async generateObject<TOutput, TSchema extends z.ZodType<TOutput>>(
input: KtxGenerateObjectInput<TOutput, TSchema>,
): Promise<TOutput> {
const options = {
...baseOptions({
projectDir: this.deps.projectDir,
model: modelForRole(this.deps.modelSlots, input.role),
env: this.deps.env,
maxTurns: 1,
tools: input.tools,
}),
outputFormat: { type: 'json_schema' as const, schema: jsonSchema(input.schema as z.ZodType) },
};
const result = await collectResult({
query: this.runQuery,
prompt: [input.system, input.prompt].filter(Boolean).join('\n\n'),
options,
allowedToolIds: new Set(mcpToolIds(input.tools ?? {})),
expectedMcpServerNames: expectedMcpServerNames(input.tools),
});
const error = resultError(result);
if (error) {
throw error;
}
if (result.subtype !== 'success') {
throw new Error(`Claude Code query failed (${result.subtype})`);
}
return (input.schema as z.ZodType<TOutput>).parse(result.structured_output);
}
async runAgentLoop(params: RunLoopParams): Promise<RunLoopResult> {
let stepIndex = 0;
try {
const options = baseOptions({
projectDir: this.deps.projectDir,
model: modelForRole(this.deps.modelSlots, params.modelRole),
env: this.deps.env,
maxTurns: params.stepBudget,
tools: params.toolSet,
});
const result = await collectResult({
query: this.runQuery,
prompt: params.userPrompt,
options: { ...options, systemPrompt: params.systemPrompt },
allowedToolIds: new Set(mcpToolIds(params.toolSet)),
expectedMcpServerNames: expectedMcpServerNames(params.toolSet),
onAssistantTurn: async () => {
stepIndex += 1;
if (!params.onStepFinish) {
return;
}
try {
await params.onStepFinish({ stepIndex, stepBudget: params.stepBudget });
} catch (error) {
this.logger.warn(
`[claude-code-runner] onStepFinish callback threw; ignoring: ${
error instanceof Error ? error.message : String(error)
}`,
);
}
},
});
const stopReason = mapClaudeCodeStopReason(result);
const error = resultError(result);
return { stopReason, ...(stopReason === 'error' && error ? { error } : {}) };
} catch (error) {
const err = error instanceof Error ? error : new Error(String(error));
return { stopReason: 'error', error: err };
}
}
}
export async function runClaudeCodeAuthProbe(input: {
projectDir: string;
model: string;
query?: QueryFn;
env?: NodeJS.ProcessEnv;
}): Promise<{ ok: true } | { ok: false; message: string }> {
let model: string;
try {
model = resolveClaudeCodeModel(input.model);
} catch (error) {
return {
ok: false,
message: error instanceof Error ? error.message : String(error),
};
}
try {
const options = baseOptions({
projectDir: input.projectDir,
model,
env: input.env,
maxTurns: 1,
});
const result = await collectResult({
query: input.query ?? defaultQuery,
prompt: 'Reply with exactly: ok',
options,
allowedToolIds: new Set(),
expectedMcpServerNames: new Set(),
});
const error = resultError(result);
if (error) {
throw error;
}
return { ok: true };
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
return {
ok: false,
message: `Claude Code authentication is not usable. Authenticate Claude Code locally with the Claude Code CLI, then rerun setup or the command. ${message}`,
};
}
}

87
packages/context/src/llm/generation.ts
View file

@ -1,85 +1,12 @@
import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider, type KtxModelRole } from '@ktx/llm';
import { generateText, Output, type FlexibleSchema, type ToolSet } from 'ai';
import type { z } from 'zod';
import type { KtxGenerateObjectInput, KtxGenerateTextInput, KtxLlmRuntimePort } from './runtime-port.js';
type GenerateTextInput = Parameters<typeof generateText>[0];
type GenerateTextFn = (input: GenerateTextInput) => Promise<{ text?: string; output?: unknown }>;
function hasTools(tools: ToolSet): boolean {
return Object.keys(tools).length > 0;
export async function generateKtxText(input: KtxGenerateTextInput & { runtime: KtxLlmRuntimePort }): Promise<string> {
return input.runtime.generateText(input);
}
interface GenerateKtxTextInput {
llmProvider: KtxLlmProvider;
role: KtxModelRole;
prompt: string;
system?: string;
tools?: ToolSet;
temperature?: number;
generateText?: GenerateTextFn;
}
export async function generateKtxText(input: GenerateKtxTextInput): Promise<string> {
const model = input.llmProvider.getModel(input.role);
if ((model as { provider?: string }).provider === 'deterministic') {
return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`;
}
const built = new KtxMessageBuilder(input.llmProvider).wrapSimple({
system: input.system,
messages: [{ role: 'user', content: input.prompt }],
tools: input.tools ?? {},
model,
});
const split = splitKtxSystemMessages(built.messages);
const result = await (input.generateText ?? generateText)({
model,
temperature: input.temperature ?? 0,
...(split.system ? { system: split.system } : {}),
messages: split.messages,
tools: built.tools as ToolSet,
...(hasTools(built.tools as ToolSet)
? {
experimental_repairToolCall: input.llmProvider.repairToolCallHandler({
source: `ktx-${input.role}`,
}),
}
: {}),
});
if (typeof result.text !== 'string') {
throw new Error('KTX LLM text generation returned no text');
}
return result.text;
}
export async function generateKtxObject<TOutput, TSchema>(
input: GenerateKtxTextInput & { schema: TSchema },
export async function generateKtxObject<TOutput, TSchema extends z.ZodType<TOutput>>(
input: KtxGenerateObjectInput<TOutput, TSchema> & { runtime: KtxLlmRuntimePort },
): Promise<TOutput> {
const model = input.llmProvider.getModel(input.role);
const built = new KtxMessageBuilder(input.llmProvider).wrapSimple({
system: input.system,
messages: [{ role: 'user', content: input.prompt }],
tools: input.tools ?? {},
model,
});
const split = splitKtxSystemMessages(built.messages);
const result = await (input.generateText ?? generateText)({
model,
temperature: input.temperature ?? 0,
...(split.system ? { system: split.system } : {}),
messages: split.messages,
tools: built.tools as ToolSet,
...(hasTools(built.tools as ToolSet)
? {
experimental_repairToolCall: input.llmProvider.repairToolCallHandler({
source: `ktx-${input.role}`,
}),
}
: {}),
output: Output.object({
schema: input.schema as FlexibleSchema<TOutput>,
}),
});
if (result.output == null) {
throw new Error('KTX LLM object generation returned no output');
}
return result.output as TOutput;
return input.runtime.generateObject(input);
}

27
packages/context/src/llm/index.ts
View file

@ -1,5 +1,31 @@
export { KtxIngestEmbeddingPortAdapter, KtxScanEmbeddingPortAdapter } from './embedding-port.js';
export { AiSdkKtxLlmRuntime } from './ai-sdk-runtime.js';
export type { AgentTelemetryPort, AiSdkKtxLlmRuntimeDeps } from './ai-sdk-runtime.js';
export { createKtxClaudeCodeEnv, CLAUDE_CODE_PROVIDER_ENV_DENYLIST } from './claude-code-env.js';
export { resolveClaudeCodeModel } from './claude-code-models.js';
export { ClaudeCodeKtxLlmRuntime, mapClaudeCodeStopReason, runClaudeCodeAuthProbe } from './claude-code-runtime.js';
export { generateKtxObject, generateKtxText } from './generation.js';
export type {
AgentRunnerPort,
KtxGenerateObjectInput,
KtxGenerateTextInput,
KtxLlmRuntimePort,
KtxRuntimeToolDescriptor,
KtxRuntimeToolOutput,
KtxRuntimeToolSet,
RunLoopParams,
RunLoopResult,
RunLoopStepInfo,
RunLoopStopReason,
} from './runtime-port.js';
export { RuntimeAgentRunner } from './runtime-port.js';
export {
createAiSdkToolSet,
createClaudeSdkTools,
createRuntimeToolDescriptorFromAiTool,
createRuntimeToolSetFromAiSdkTools,
normalizeKtxRuntimeToolOutput,
} from './runtime-tools.js';
export type {
KtxLlmDebugProviderOptionsEntry,
KtxLlmDebugRequest,
@ -15,6 +41,7 @@ export {
MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV,
createLocalKtxEmbeddingProviderFromConfig,
createLocalKtxLlmProviderFromConfig,
createLocalKtxLlmRuntimeFromConfig,
resolveLocalKtxEmbeddingConfig,
resolveLocalKtxLlmConfig,
} from './local-config.js';

34
packages/context/src/llm/local-config.ts
View file

@ -9,11 +9,17 @@ import {
} from '@ktx/llm';
import { resolveKtxConfigReference } from '../core/config-reference.js';
import type { KtxProjectEmbeddingConfig, KtxProjectLlmConfig } from '../project/config.js';
import { AiSdkKtxLlmRuntime } from './ai-sdk-runtime.js';
import { ClaudeCodeKtxLlmRuntime } from './claude-code-runtime.js';
import type { KtxLlmRuntimePort } from './runtime-port.js';
interface LocalConfigDeps {
env?: NodeJS.ProcessEnv;
projectDir?: string;
createKtxLlmProvider?: typeof createKtxLlmProvider;
createKtxEmbeddingProvider?: typeof createKtxEmbeddingProvider;
createClaudeCodeRuntime?: (deps: ConstructorParameters<typeof ClaudeCodeKtxLlmRuntime>[0]) => KtxLlmRuntimePort;
createAiSdkRuntime?: (deps: { llmProvider: KtxLlmProvider }) => KtxLlmRuntimePort;
}
export const MANAGED_SENTENCE_TRANSFORMERS_BASE_URL = 'managed:local-embeddings';
@ -106,7 +112,33 @@ export function createLocalKtxLlmProviderFromConfig(
deps: LocalConfigDeps = {},
): KtxLlmProvider | null {
const resolved = resolveLocalKtxLlmConfig(config, deps.env ?? process.env);
return resolved ? (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved) : null;
if (!resolved || resolved.backend === 'claude-code') {
return null;
}
return (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved);
}
export function createLocalKtxLlmRuntimeFromConfig(
config: KtxProjectLlmConfig,
deps: LocalConfigDeps = {},
): KtxLlmRuntimePort | null {
const resolved = resolveLocalKtxLlmConfig(config, deps.env ?? process.env);
if (!resolved) {
return null;
}
if (resolved.backend === 'claude-code') {
const projectDir = deps.projectDir;
if (!projectDir) {
throw new Error('projectDir is required when creating the claude-code LLM runtime');
}
return (deps.createClaudeCodeRuntime ?? ((runtimeDeps) => new ClaudeCodeKtxLlmRuntime(runtimeDeps)))({
projectDir,
modelSlots: resolved.modelSlots,
env: deps.env,
});
}
const llmProvider = (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved);
return (deps.createAiSdkRuntime ?? ((runtimeDeps) => new AiSdkKtxLlmRuntime(runtimeDeps)))({ llmProvider });
}
function resolveSentenceTransformersBaseUrl(

25
packages/context/src/llm/runtime-local-config.test.ts Normal file
View file

@ -0,0 +1,25 @@
import { describe, expect, it, vi } from 'vitest';
import { createLocalKtxLlmProviderFromConfig, createLocalKtxLlmRuntimeFromConfig } from './local-config.js';
describe('local KTX LLM runtime config', () => {
it('creates a Claude Code runtime for claude-code backend without creating an AI SDK provider', () => {
const runtime = createLocalKtxLlmRuntimeFromConfig(
{
provider: { backend: 'claude-code' },
models: { default: 'sonnet', triage: 'haiku' },
},
{ env: {}, projectDir: '/tmp/project', createClaudeCodeRuntime: vi.fn((deps) => ({ deps }) as never) },
);
expect(runtime).toMatchObject({ deps: expect.objectContaining({ projectDir: '/tmp/project' }) });
});
it('returns null from the AI SDK provider factory for claude-code backend', () => {
expect(
createLocalKtxLlmProviderFromConfig({
provider: { backend: 'claude-code' },
models: { default: 'sonnet' },
}),
).toBeNull();
});
});

75
packages/context/src/llm/runtime-port.ts Normal file
View file

@ -0,0 +1,75 @@
import type { KtxModelRole } from '@ktx/llm';
import type { z } from 'zod';
export interface KtxRuntimeToolOutput<TOutput = unknown> {
markdown: string;
structured?: TOutput;
}
export interface KtxRuntimeToolDescriptor<TInput = unknown, TOutput = unknown> {
name: string;
description: string;
inputSchema: z.ZodObject<z.ZodRawShape>;
execute(input: TInput): Promise<KtxRuntimeToolOutput<TOutput>>;
}
export type KtxRuntimeToolSet = Record<string, KtxRuntimeToolDescriptor>;
export type RunLoopStopReason = 'budget' | 'natural' | 'error';
export interface RunLoopStepInfo {
stepIndex: number;
stepBudget: number;
}
export interface RunLoopParams {
modelRole: KtxModelRole;
systemPrompt: string;
userPrompt: string;
toolSet: KtxRuntimeToolSet;
stepBudget: number;
telemetryTags: Record<string, string>;
onStepFinish?: (info: RunLoopStepInfo) => void | Promise<void>;
}
export interface RunLoopResult {
stopReason: RunLoopStopReason;
error?: Error;
}
export interface KtxGenerateTextInput {
role: KtxModelRole;
prompt: string;
system?: string;
tools?: KtxRuntimeToolSet;
temperature?: number;
}
export interface KtxGenerateObjectInput<TOutput, TSchema extends z.ZodType<TOutput>> {
role: KtxModelRole;
prompt: string;
system?: string;
tools?: KtxRuntimeToolSet;
temperature?: number;
schema: TSchema;
}
export interface KtxLlmRuntimePort {
generateText(input: KtxGenerateTextInput): Promise<string>;
generateObject<TOutput, TSchema extends z.ZodType<TOutput>>(
input: KtxGenerateObjectInput<TOutput, TSchema>,
): Promise<TOutput>;
runAgentLoop(params: RunLoopParams): Promise<RunLoopResult>;
}
export interface AgentRunnerPort {
runLoop(params: RunLoopParams): Promise<RunLoopResult>;
}
export class RuntimeAgentRunner implements AgentRunnerPort {
constructor(private readonly runtime: KtxLlmRuntimePort) {}
runLoop(params: RunLoopParams): Promise<RunLoopResult> {
return this.runtime.runAgentLoop(params);
}
}

43
packages/context/src/llm/runtime-tools.test.ts Normal file
View file

@ -0,0 +1,43 @@
import { describe, expect, it, vi } from 'vitest';
import { z } from 'zod';
import { createAiSdkToolSet, createClaudeSdkTools, normalizeKtxRuntimeToolOutput } from './runtime-tools.js';
import type { KtxRuntimeToolDescriptor } from './runtime-port.js';
describe('runtime tool descriptors', () => {
const descriptor: KtxRuntimeToolDescriptor<{ id: string }, { ok: boolean }> = {
name: 'read_thing',
description: 'Read one thing.',
inputSchema: z.object({ id: z.string() }),
execute: vi.fn(async (input) => ({
markdown: `Read ${input.id}`,
structured: { ok: true },
})),
};
it('normalizes string and object tool outputs into markdown plus optional structured payload', () => {
expect(normalizeKtxRuntimeToolOutput('plain text')).toEqual({ markdown: 'plain text' });
expect(normalizeKtxRuntimeToolOutput({ markdown: 'shown', structured: { id: 1 } })).toEqual({
markdown: 'shown',
structured: { id: 1 },
});
expect(normalizeKtxRuntimeToolOutput({ name: 'skill', content: 'body' })).toEqual({
markdown: '```json\n{\n "name": "skill",\n "content": "body"\n}\n```',
structured: { name: 'skill', content: 'body' },
});
});
it('builds AI SDK tools that expose markdown to the model', async () => {
const tools = createAiSdkToolSet({ read_thing: descriptor });
const output = await tools.read_thing.execute?.({ id: 'a' }, { toolCallId: 'call-1', messages: [] } as never);
const modelOutput = tools.read_thing.toModelOutput?.({ output } as never);
expect(modelOutput).toEqual({ type: 'text', value: 'Read a' });
});
it('builds Claude SDK tools that return text content only', async () => {
const tools = createClaudeSdkTools({ read_thing: descriptor });
const result = await tools[0].handler({ id: 'b' } as never, {});
expect(result).toEqual({ content: [{ type: 'text', text: 'Read b' }] });
});
});

91
packages/context/src/llm/runtime-tools.ts Normal file
View file

@ -0,0 +1,91 @@
import { tool as aiTool, type Tool, type ToolSet } from 'ai';
import { tool as claudeTool, type SdkMcpToolDefinition } from '@anthropic-ai/claude-agent-sdk';
import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
import { z } from 'zod';
import type { KtxRuntimeToolDescriptor, KtxRuntimeToolOutput, KtxRuntimeToolSet } from './runtime-port.js';
function isRuntimeOutput(value: unknown): value is KtxRuntimeToolOutput {
return Boolean(
value &&
typeof value === 'object' &&
'markdown' in value &&
typeof (value as { markdown?: unknown }).markdown === 'string',
);
}
export function normalizeKtxRuntimeToolOutput(value: unknown): KtxRuntimeToolOutput {
if (isRuntimeOutput(value)) {
return 'structured' in value ? { markdown: value.markdown, structured: value.structured } : { markdown: value.markdown };
}
if (typeof value === 'string') {
return { markdown: value };
}
return {
markdown: `\`\`\`json\n${JSON.stringify(value, null, 2)}\n\`\`\``,
structured: value,
};
}
function assertObjectSchema(name: string, schema: z.ZodType): asserts schema is z.ZodObject<z.ZodRawShape> {
if (!(schema instanceof z.ZodObject)) {
throw new Error(`KTX runtime tool "${name}" must use z.object input schema for claude-code`);
}
}
export function createAiSdkToolSet(tools: KtxRuntimeToolSet = {}): ToolSet {
return Object.fromEntries(
Object.entries(tools).map(([name, descriptor]) => [
name,
aiTool({
description: descriptor.description,
inputSchema: descriptor.inputSchema,
execute: async (input) => descriptor.execute(input),
toModelOutput: ({ output }) => {
const normalized = normalizeKtxRuntimeToolOutput(output);
return { type: 'text', value: normalized.markdown };
},
}),
]),
);
}
export function createClaudeSdkTools(tools: KtxRuntimeToolSet = {}): Array<SdkMcpToolDefinition<z.ZodRawShape>> {
return Object.values(tools).map((descriptor) => {
assertObjectSchema(descriptor.name, descriptor.inputSchema);
return claudeTool(
descriptor.name,
descriptor.description,
descriptor.inputSchema.shape,
async (input): Promise<CallToolResult> => {
const normalized = normalizeKtxRuntimeToolOutput(await descriptor.execute(input));
return { content: [{ type: 'text', text: normalized.markdown }] };
},
);
});
}
export function mcpToolIds(tools: KtxRuntimeToolSet = {}): string[] {
return Object.keys(tools).map((name) => `mcp__ktx__${name}`);
}
export function createRuntimeToolDescriptorFromAiTool(name: string, aiSdkTool: Tool): KtxRuntimeToolDescriptor {
return {
name,
description: aiSdkTool.description ?? '',
inputSchema: aiSdkTool.inputSchema as KtxRuntimeToolDescriptor['inputSchema'],
execute: async (input) => {
if (typeof aiSdkTool.execute !== 'function') {
throw new Error(`KTX runtime tool "${name}" has no execute function`);
}
return normalizeKtxRuntimeToolOutput(
await aiSdkTool.execute(input as never, { toolCallId: `runtime-${name}` } as never),
);
},
};
}
export function createRuntimeToolSetFromAiSdkTools(tools: ToolSet = {}): KtxRuntimeToolSet {
return Object.fromEntries(
Object.entries(tools).map(([name, aiSdkTool]) => [name, createRuntimeToolDescriptorFromAiTool(name, aiSdkTool as Tool)]),
);
}

32
packages/context/src/memory/local-memory.ts
View file

@ -1,13 +1,17 @@
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 { localConnectionInfoFromConfig } from '../connections/index.js';
import type { KtxEmbeddingPort, KtxFileStorePort, KtxFileWriteResult } from '../core/index.js';
import { type KtxLogger, noopLogger, SessionWorktreeService } from '../core/index.js';
import type { KtxSemanticLayerComputePort } from '../daemon/index.js';
import { createLocalKtxLlmProviderFromConfig } from '../llm/index.js';
import {
createLocalKtxLlmRuntimeFromConfig,
RuntimeAgentRunner,
type AgentRunnerPort,
type KtxLlmRuntimePort,
type KtxRuntimeToolSet,
} from '../llm/index.js';
import type { KtxLocalProject } from '../project/index.js';
import { PromptService } from '../prompts/index.js';
import { SkillsRegistryService } from '../skills/index.js';
@ -63,8 +67,8 @@ const LOCAL_AUTHOR = { name: 'KTX Local', email: 'local@ktx.local' };
const LOCAL_SHAPE_WARNING = 'Local memory ingest validates semantic-layer YAML shape only.';
export interface CreateLocalProjectMemoryIngestOptions {
llmProvider?: KtxLlmProvider;
agentRunner?: AgentRunnerService;
llmRuntime?: KtxLlmRuntimePort;
agentRunner?: AgentRunnerPort;
memoryModel?: string;
semanticLayerCompute?: KtxSemanticLayerComputePort;
queryExecutor?: { execute(input: { connectionId: string; sql: string; maxRows?: number }): Promise<KtxQueryResult> };
@ -89,7 +93,8 @@ export function createLocalProjectMemoryIngest(
const slSearchService = new SlSearchService(embedding, slSourcesRepository, logger);
const wikiService = new KnowledgeWikiService(rootFileStore, embedding, knowledgeIndex, project.git, logger);
const authorResolver = new LocalAuthorResolver();
const llmProvider = options.llmProvider ?? createLocalKtxLlmProviderFromConfig(project.config.llm);
const llmRuntime =
options.llmRuntime ?? createLocalKtxLlmRuntimeFromConfig(project.config.llm, { projectDir: project.projectDir });
const toolsetFactory = new LocalMemoryToolsetFactory({
project,
embedding,
@ -104,10 +109,7 @@ export function createLocalProjectMemoryIngest(
});
const agentRunner =
options.agentRunner ??
new AgentRunnerService({
llmProvider: requireLlmProvider(llmProvider),
logger,
});
new RuntimeAgentRunner(requireLlmRuntime(llmRuntime));
const memoryAgent = new MemoryAgentService({
settings: {
knowledge: { userScopedKnowledgeEnabled: false },
@ -143,11 +145,11 @@ export function createLocalProjectMemoryIngest(
});
}
function requireLlmProvider(provider: KtxLlmProvider | null | undefined): KtxLlmProvider {
if (!provider) {
function requireLlmRuntime(runtime: KtxLlmRuntimePort | null | undefined): KtxLlmRuntimePort {
if (!runtime) {
throw new Error('createLocalProjectMemoryIngest requires llm.provider.backend or an injected agentRunner');
}
return provider;
return runtime;
}
class LocalMemoryFileStore implements MemoryFileStorePort {
@ -386,8 +388,8 @@ class LocalShapeOnlySlValidator implements SlValidatorPort<SlValidationDeps> {
class LocalMemoryToolSet implements MemoryToolSetLike {
constructor(private readonly tools: BaseTool[]) {}
toAiSdkTools(context: ToolContext) {
return Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toAiSdkTool(context)]));
toRuntimeTools(context: ToolContext): KtxRuntimeToolSet {
return Object.fromEntries(this.tools.map((tool) => [tool.name, tool.toRuntimeTool(context)]));
}
}

42
packages/context/src/memory/memory-agent.service.ingest.test.ts
View file

@ -1,3 +1,6 @@
import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
// Module-level mock for 'ai' so generateText is a stub. This file is separate from
@ -15,7 +18,6 @@ import { MemoryAgentService } from './memory-agent.service.js';
interface BuiltMocks {
appSettings: any;
llmProvider: any;
prompt: any;
eventTracker: any;
telemetry: any;
@ -63,7 +65,6 @@ const buildMocks = (overrides: Partial<BuiltMocks> = {}): BuiltMocks => {
llm: { memoryIngestionModel: 'test-model' },
},
},
llmProvider: { getModel: vi.fn().mockReturnValue({}) },
prompt: { loadPrompt: vi.fn().mockResolvedValue('base framing') },
eventTracker: { trackEvent: vi.fn(), createTelemetryIntegration: vi.fn().mockReturnValue(undefined) },
telemetry: {
@ -124,11 +125,11 @@ const buildMocks = (overrides: Partial<BuiltMocks> = {}): BuiltMocks => {
slValidator: { validateSingleSource: vi.fn().mockResolvedValue({ errors: [], warnings: [] }) },
toolsetFactory: {
createIngestWuToolset: vi.fn().mockReturnValue({
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
}),
createToolset: vi.fn().mockReturnValue({
toAiSdkTools: vi.fn().mockReturnValue({}),
toRuntimeTools: vi.fn().mockReturnValue({}),
getAllTools: vi.fn().mockReturnValue([]),
}),
},
@ -241,6 +242,39 @@ describe('MemoryAgentService.ingest — session-branch orchestration', () => {
expect(result.commitHash).toBe('cafebabe');
});
it('normalizes load_skill output to markdown while preserving structured payload', async () => {
const tempDir = await mkdtemp(join(tmpdir(), 'ktx-memory-skill-'));
const skillDir = join(tempDir, 'memory_agent');
await mkdir(skillDir, { recursive: true });
await writeFile(join(skillDir, 'SKILL.md'), '---\nname: memory_agent\n---\nSkill body', 'utf-8');
try {
const agentRunner = {
runLoop: vi.fn(async (params: any) => {
const result = await params.toolSet.load_skill.execute({ name: 'memory_agent' });
expect(result.markdown).toContain('memory_agent');
expect(result.structured).toMatchObject({ name: 'memory_agent' });
return { stopReason: 'natural' as const };
}),
};
const mocks = buildMocks({
agentRunner,
skillsRegistry: {
listSkills: vi.fn().mockResolvedValue([{ name: 'memory_agent', path: skillDir }]),
buildSkillsPrompt: vi.fn().mockReturnValue(''),
getSkill: vi.fn().mockResolvedValue({ name: 'memory_agent', path: skillDir }),
stripFrontmatter: vi.fn().mockReturnValue('Skill body'),
},
});
const svc = buildService(mocks);
await svc.ingest(baseInput);
expect(agentRunner.runLoop).toHaveBeenCalled();
} finally {
await rm(tempDir, { recursive: true, force: true });
}
});
it('logs prompt debug output when KTX_MEMORY_AGENT_DEBUG_PROMPTS is enabled', async () => {
const previousDebugPrompts = process.env.KTX_MEMORY_AGENT_DEBUG_PROMPTS;
const mocks = buildMocks();

21
packages/context/src/memory/memory-agent.service.ts
View file

@ -1,10 +1,10 @@
import { createHash } from 'node:crypto';
import { readFile } from 'node:fs/promises';
import { join } from 'node:path';
import { tool } from 'ai';
import * as YAML from 'yaml';
import { z } from 'zod';
import { type KtxLogger, noopLogger } from '../core/index.js';
import type { KtxRuntimeToolSet } from '../llm/index.js';
import {
revertSourceToPreHead,
type SemanticLayerSource,
@ -125,8 +125,9 @@ export class MemoryAgentService {
session: toolSession,
};
const loadSkillTool = {
load_skill: tool({
const loadSkillTool: KtxRuntimeToolSet = {
load_skill: {
name: 'load_skill',
description:
'Load a skill to get specialized instructions. Call this when a skill listed in the system prompt matches the current task.',
inputSchema: z.object({
@ -137,23 +138,27 @@ export class MemoryAgentService {
if (!skill) {
const available =
(await this.deps.skillsRegistry.listSkills('memory_agent')).map((s) => s.name).join(', ') || '(none)';
return `Skill "${name}" not available to the memory agent. Available: ${available}`;
return { markdown: `Skill "${name}" not available to the memory agent. Available: ${available}` };
}
try {
const body = await readFile(join(skill.path, 'SKILL.md'), 'utf-8');
if (!skillsLoaded.includes(skill.name)) {
skillsLoaded.push(skill.name);
}
return {
const structured = {
name: skill.name,
skillDirectory: skill.path,
content: this.deps.skillsRegistry.stripFrontmatter(body),
};
return {
markdown: `# ${structured.name}\n\n${structured.content}`,
structured,
};
} catch (e) {
return `Error loading skill "${name}": ${e instanceof Error ? e.message : String(e)}`;
return { markdown: `Error loading skill "${name}": ${e instanceof Error ? e.message : String(e)}` };
}
},
}),
},
};
const skillNames: string[] = [...DEFAULT_SKILL_NAMES];
@ -212,7 +217,7 @@ export class MemoryAgentService {
modelRole: 'candidateExtraction',
systemPrompt,
userPrompt: prompt,
toolSet: { ...toolset.toAiSdkTools(toolContext), ...loadSkillTool },
toolSet: { ...toolset.toRuntimeTools(toolContext), ...loadSkillTool },
stepBudget,
telemetryTags: {
operationName: 'memory-agent-ingest',

7
packages/context/src/memory/types.ts
View file

@ -1,5 +1,4 @@
import type { Tool } from 'ai';
import type { AgentRunnerService } from '../agent/index.js';
import type { AgentRunnerPort, KtxRuntimeToolSet } from '../llm/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<MemoryFileStorePort>, MemoryCommitMessagePort {}
export interface MemoryToolSetLike {
toAiSdkTools(context: ToolContext): Record<string, Tool>;
toRuntimeTools(context: ToolContext): KtxRuntimeToolSet;
}
export interface MemoryToolsetFactoryPort {
@ -150,7 +149,7 @@ export interface MemoryAgentServiceDeps {
slSourcesRepository: SlSourcesIndexPort;
sessionWorktreeService: SessionWorktreeService<MemoryFileStorePort>;
semanticLayerSourceReconciler: MemorySlSourceReconcilerPort;
agentRunner: AgentRunnerService;
agentRunner: AgentRunnerPort;
slValidator: SlValidatorPort<SlValidationDeps>;
toolsetFactory: MemoryToolsetFactoryPort;
telemetry?: MemoryTelemetryPort;

27
packages/context/src/project/config.test.ts
View file

@ -180,6 +180,31 @@ llm:
});
});
it('parses Claude Code as a first-class LLM backend', () => {
const config = parseKtxProjectConfig(`
llm:
provider:
backend: claude-code
models:
default: sonnet
triage: haiku
candidateExtraction: sonnet
curator: sonnet
reconcile: sonnet
repair: opus
`);
expect(config.llm.provider.backend).toBe('claude-code');
expect(config.llm.models).toEqual({
default: 'sonnet',
triage: 'haiku',
candidateExtraction: 'sonnet',
curator: 'sonnet',
reconcile: 'sonnet',
repair: 'opus',
});
});
it('parses gateway LLM, OpenAI scan embeddings, and sentence-transformers ingest embeddings', () => {
const config = parseKtxProjectConfig(`
llm:
@ -497,7 +522,7 @@ describe('generateKtxProjectConfigJsonSchema', () => {
const llm = (schema.properties as Record<string, { properties?: Record<string, unknown> }>).llm;
const provider = llm?.properties?.provider as { properties?: Record<string, unknown> };
const backend = provider?.properties?.backend as { enum?: readonly string[] };
expect(backend?.enum).toEqual(['none', 'anthropic', 'vertex', 'gateway']);
expect(backend?.enum).toEqual(['none', 'anthropic', 'vertex', 'gateway', 'claude-code']);
const storage = (schema.properties as Record<string, { properties?: Record<string, unknown> }>).storage;
const state = storage?.properties?.state as { enum?: readonly string[] };

6
packages/context/src/project/config.ts
View file

@ -3,7 +3,7 @@ import YAML from 'yaml';
import * as z from 'zod';
import { connectionConfigSchema } from './driver-schemas.js';
const KTX_LLM_BACKENDS = ['none', 'anthropic', 'vertex', 'gateway'] as const;
const KTX_LLM_BACKENDS = ['none', 'anthropic', 'vertex', 'gateway', 'claude-code'] as const;
const KTX_EMBEDDING_BACKENDS = ['none', 'deterministic', 'openai', 'sentence-transformers'] as const;
const KTX_PROMPT_CACHE_TTLS = ['5m', '1h'] as const;
const KTX_ENRICHMENT_MODES = ['none', 'deterministic', 'llm'] as const;
@ -46,7 +46,9 @@ const llmProviderSchema = z
backend: z
.enum(KTX_LLM_BACKENDS)
.default('none')
.describe('LLM provider backend. "none" disables LLM features; "anthropic" / "vertex" / "gateway" require the matching nested credentials block.'),
.describe(
'LLM provider backend. "none" disables LLM features; "anthropic" / "vertex" / "gateway" require the matching nested credentials block; "claude-code" uses the local Claude Code session.',
),
vertex: vertexProviderSchema.optional().describe('Vertex AI credentials, used when backend is "vertex".'),
anthropic: apiCredentialsSchema.optional().describe('Anthropic API credentials, used when backend is "anthropic".'),
gateway: apiCredentialsSchema.optional().describe('AI Gateway credentials, used when backend is "gateway".'),

82
packages/context/src/scan/description-generation.test.ts
View file

@ -31,46 +31,32 @@ function createCache(initial: Record<string, string> = {}): KtxDescriptionCacheP
function createLlmProvider(text = 'generated description') {
vi.mocked(generateText).mockResolvedValue({ text } as never);
return {
getModel: vi.fn().mockReturnValue({ modelId: 'claude-sonnet-4-6', provider: 'anthropic' }),
getModelByName: vi.fn(),
cacheMarker: vi.fn(),
repairToolCallHandler: vi.fn(),
thinkingProviderOptions: vi.fn(),
telemetryConfig: vi.fn(),
promptCachingConfig: vi.fn(() => ({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
})),
activeBackend: vi.fn(() => 'anthropic'),
generateText: vi.fn(async (input) => {
const result = await generateText({
system: input.system ? { role: 'system', content: input.system } : undefined,
messages: [{ role: 'user', content: input.prompt }],
temperature: input.temperature,
} as never);
return result.text;
}),
generateObject: vi.fn(),
runAgentLoop: vi.fn(),
} as any;
}
function createFailingLlmProvider(message = 'timeout exceeded when trying to connect') {
vi.mocked(generateText).mockRejectedValue(new Error(message) as never);
return {
getModel: vi.fn().mockReturnValue({ modelId: 'claude-sonnet-4-6', provider: 'anthropic' }),
getModelByName: vi.fn(),
cacheMarker: vi.fn(),
repairToolCallHandler: vi.fn(),
thinkingProviderOptions: vi.fn(),
telemetryConfig: vi.fn(),
promptCachingConfig: vi.fn(() => ({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
})),
activeBackend: vi.fn(() => 'anthropic'),
generateText: vi.fn(async (input) => {
const result = await generateText({
system: input.system ? { role: 'system', content: input.system } : undefined,
messages: [{ role: 'user', content: input.prompt }],
temperature: input.temperature,
} as never);
return result.text;
}),
generateObject: vi.fn(),
runAgentLoop: vi.fn(),
} as any;
}
@ -158,10 +144,10 @@ describe('KTX description prompt builders', () => {
describe('KtxDescriptionGenerator', () => {
it('generates column descriptions with pre-fetched values, cache hits, and word-limit metadata', async () => {
const cache = createCache({ 'warehouse.public.orders.cached_status': 'Cached status description' });
const llmProvider = createLlmProvider('Payment state');
const llmRuntime = createLlmProvider('Payment state');
const connector = createConnector();
const generator = new KtxDescriptionGenerator({
llmProvider,
llmRuntime,
cache,
settings: {
columnMaxWords: 12,
@ -222,7 +208,7 @@ describe('KtxDescriptionGenerator', () => {
it('samples through the connector when column values are not pre-fetched', async () => {
const connector = createConnector();
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Current order state'),
llmRuntime: createLlmProvider('Current order state'),
settings: {
columnMaxWords: 12,
tableMaxWords: 18,
@ -271,7 +257,7 @@ describe('KtxDescriptionGenerator', () => {
})),
};
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Generated through sampler'),
llmRuntime: createLlmProvider('Generated through sampler'),
settings: {
columnMaxWords: 12,
tableMaxWords: 18,
@ -310,7 +296,7 @@ describe('KtxDescriptionGenerator', () => {
const cache = createCache();
const connector = createConnector();
const generator = new KtxDescriptionGenerator({
llmProvider: createFailingLlmProvider(),
llmRuntime: createFailingLlmProvider(),
cache,
settings: {
columnMaxWords: 12,
@ -355,7 +341,7 @@ describe('KtxDescriptionGenerator', () => {
const cache = createCache();
const connector = createConnector();
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Commerce orders'),
llmRuntime: createLlmProvider('Commerce orders'),
cache,
settings: {
columnMaxWords: 12,
@ -424,7 +410,7 @@ describe('KtxDescriptionGenerator resilience', () => {
const logger = createLogger();
const warnings: Array<{ code: string; table?: string }> = [];
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Commerce orders'),
llmRuntime: createLlmProvider('Commerce orders'),
logger,
onWarning: (warning) => warnings.push({ code: warning.code, ...(warning.table ? { table: warning.table } : {}) }),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, concurrencyLimit: 2 },
@ -455,7 +441,7 @@ describe('KtxDescriptionGenerator resilience', () => {
const logger = createLogger();
const warnings: Array<{ code: string; table?: string; metadata?: Record<string, unknown> }> = [];
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Customer reference data'),
llmRuntime: createLlmProvider('Customer reference data'),
logger,
onWarning: (warning) =>
warnings.push({
@ -503,7 +489,7 @@ describe('KtxDescriptionGenerator resilience', () => {
};
const warnings: string[] = [];
const generator = new KtxDescriptionGenerator({
llmProvider: createFailingLlmProvider(),
llmRuntime: createFailingLlmProvider(),
onWarning: (warning) => warnings.push(warning.code),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});
@ -528,7 +514,7 @@ describe('KtxDescriptionGenerator resilience', () => {
};
const warnings: string[] = [];
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Orders mart'),
llmRuntime: createLlmProvider('Orders mart'),
onWarning: (warning) => warnings.push(warning.code),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});
@ -562,7 +548,7 @@ describe('KtxDescriptionGenerator resilience', () => {
};
const warnings: string[] = [];
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('should not be called'),
llmRuntime: createLlmProvider('should not be called'),
onWarning: (warning) => warnings.push(warning.code),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});
@ -588,7 +574,7 @@ describe('KtxDescriptionGenerator resilience', () => {
};
const logger = createLogger();
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Payment lifecycle state'),
llmRuntime: createLlmProvider('Payment lifecycle state'),
logger,
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});
@ -625,7 +611,7 @@ describe('KtxDescriptionGenerator resilience', () => {
sampleColumn,
};
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('Customer reference identifier'),
llmRuntime: createLlmProvider('Customer reference identifier'),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});
@ -657,7 +643,7 @@ describe('KtxDescriptionGenerator resilience', () => {
};
vi.mocked(generateText).mockClear();
const generator = new KtxDescriptionGenerator({
llmProvider: createLlmProvider('should not be called'),
llmRuntime: createLlmProvider('should not be called'),
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
});

12
packages/context/src/scan/description-generation.ts
View file

@ -1,5 +1,4 @@
import type { KtxLlmProvider } from '@ktx/llm';
import { generateKtxText } from '../llm/index.js';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import type {
KtxColumnSampleInput,
KtxColumnSampleResult,
@ -120,7 +119,7 @@ export interface KtxGenerateDataSourceDescriptionInput {
}
export interface KtxDescriptionGeneratorOptions {
llmProvider: KtxLlmProvider;
llmRuntime: KtxLlmRuntimePort;
cache?: KtxDescriptionCachePort;
logger?: KtxScanLoggerPort;
onWarning?: (warning: KtxScanWarning) => void;
@ -400,14 +399,14 @@ Data source type: ${input.dataSourceType}`;
}
export class KtxDescriptionGenerator {
private readonly llmProvider: KtxLlmProvider;
private readonly llmRuntime: KtxLlmRuntimePort;
private readonly cache?: KtxDescriptionCachePort;
private readonly logger?: KtxScanLoggerPort;
private readonly onWarning?: (warning: KtxScanWarning) => void;
private readonly settings: ResolvedKtxDescriptionGenerationSettings;
constructor(options: KtxDescriptionGeneratorOptions) {
this.llmProvider = options.llmProvider;
this.llmRuntime = options.llmRuntime;
this.cache = options.cache;
this.logger = options.logger;
this.onWarning = options.onWarning;
@ -779,8 +778,7 @@ export class KtxDescriptionGenerator {
private async generateAiDescription(prompt: KtxDescriptionPrompt, _operationName: string): Promise<string | null> {
try {
const text = await generateKtxText({
llmProvider: this.llmProvider,
const text = await this.llmRuntime.generateText({
role: 'candidateExtraction',
system: prompt.system,
prompt: prompt.user,

1
packages/context/src/scan/index.ts
View file

@ -264,7 +264,6 @@ export type {
} from './relationship-graph-resolver.js';
export { resolveKtxRelationshipGraph } from './relationship-graph-resolver.js';
export type {
KtxRelationshipLlmProposalGenerateText,
KtxRelationshipLlmProposalResult,
KtxRelationshipLlmProposalSettings,
ProposeKtxRelationshipCandidatesWithLlmInput,

20
packages/context/src/scan/local-enrichment.test.ts
View file

@ -356,7 +356,7 @@ describe('local scan enrichment', () => {
it('honors scan relationship config when LLM proposals are disabled', async () => {
const providers = createDeterministicLocalScanEnrichmentProviders({ embeddingDimensions: 3 });
const getModel = vi.fn(() => ({ modelId: 'provider/language-model', provider: 'gateway' }));
const generateObject = vi.fn();
const result = await runLocalScanEnrichment({
connectionId: 'warehouse',
mode: 'relationships',
@ -365,9 +365,9 @@ describe('local scan enrichment', () => {
context: { runId: 'scan-run-llm-disabled' },
providers: {
...providers,
llm: {
...providers.llm,
getModel: getModel as never,
llmRuntime: {
...providers.llmRuntime,
generateObject: generateObject as never,
},
},
relationshipSettings: {
@ -378,7 +378,7 @@ describe('local scan enrichment', () => {
});
expect(result.summary.llmRelationshipValidation).toBe('skipped');
expect(getModel).not.toHaveBeenCalledWith('candidateExtraction');
expect(generateObject).not.toHaveBeenCalled();
});
it('skips relationship detection when scan relationships are disabled', async () => {
@ -628,7 +628,7 @@ describe('local scan enrichment', () => {
connector: scanConnector,
context: { runId: 'scan-run-batched-embeddings' },
providers: {
llm: deterministicProviders.llm,
llmRuntime: deterministicProviders.llmRuntime,
embedding: {
dimensions: 3,
maxBatchSize: 2,
@ -658,7 +658,7 @@ describe('local scan enrichment', () => {
providerIdentity: { provider: 'deterministic', embeddingDimensions: 6 },
});
const getModel = vi.spyOn(providers.llm, 'getModel');
const generateText = vi.spyOn(providers.llmRuntime, 'generateText');
const embedBatch = vi.spyOn(providers.embedding, 'embedBatch');
const second = await runLocalScanEnrichment({
connectionId: 'warehouse',
@ -676,7 +676,7 @@ describe('local scan enrichment', () => {
expect(first.state.resumedStages).toEqual([]);
expect(second.state.resumedStages).toEqual(['descriptions', 'embeddings', 'relationships']);
expect(second.state.completedStages).toEqual(['descriptions', 'embeddings', 'relationships']);
expect(getModel).not.toHaveBeenCalled();
expect(generateText).not.toHaveBeenCalled();
expect(embedBatch).not.toHaveBeenCalled();
expect(second.descriptionUpdates).toEqual(first.descriptionUpdates);
expect(second.embeddingUpdates).toEqual(first.embeddingUpdates);
@ -711,7 +711,7 @@ describe('local scan enrichment', () => {
tables: [{ ...firstTable, name: 'customers' }],
})),
};
const getModel = vi.spyOn(providers.llm, 'getModel');
const generateText = vi.spyOn(providers.llmRuntime, 'generateText');
const result = await runLocalScanEnrichment({
connectionId: 'warehouse',
@ -727,7 +727,7 @@ describe('local scan enrichment', () => {
expect(result.state.resumedStages).toEqual([]);
expect(result.state.completedStages).toEqual(['descriptions', 'embeddings', 'relationships']);
expect(getModel).toHaveBeenCalled();
expect(generateText).toHaveBeenCalled();
});
it('runs providerless enriched scans as relationship-only discovery enrichment', async () => {

49
packages/context/src/scan/local-enrichment.ts
View file

@ -1,5 +1,5 @@
import type { KtxLlmProvider } from '@ktx/llm';
import pLimit from 'p-limit';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import { buildDefaultKtxProjectConfig, type KtxScanRelationshipConfig } from '../project/config.js';
import { type KtxDescriptionColumnTable, KtxDescriptionGenerator } from './description-generation.js';
import { buildKtxColumnEmbeddingText } from './embedding-text.js';
@ -49,7 +49,7 @@ export interface DeterministicLocalScanEnrichmentProviderOptions {
}
export interface KtxLocalScanEnrichmentProviders {
llm: KtxLlmProvider;
llmRuntime: KtxLlmRuntimePort;
embedding: KtxEmbeddingPort;
}
@ -190,7 +190,7 @@ export function createDeterministicLocalScanEnrichmentProviders(
const dimensions = options.embeddingDimensions ?? 8;
const maxBatchSize = options.maxBatchSize ?? 64;
return {
llm: deterministicLlmProvider(),
llmRuntime: deterministicLlmRuntime(),
embedding: {
dimensions,
maxBatchSize,
@ -201,41 +201,16 @@ export function createDeterministicLocalScanEnrichmentProviders(
};
}
function deterministicLlmProvider(): KtxLlmProvider {
const model = { modelId: 'deterministic-scan', provider: 'deterministic' };
function deterministicLlmRuntime(): KtxLlmRuntimePort {
return {
getModel() {
return model as ReturnType<KtxLlmProvider['getModel']>;
async generateText(input) {
return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`;
},
getModelByName() {
return model as ReturnType<KtxLlmProvider['getModelByName']>;
async generateObject() {
return { pkCandidates: [], fkCandidates: [] } as never;
},
cacheMarker() {
return undefined;
},
repairToolCallHandler() {
throw new Error('deterministic scan provider does not support tool-call repair');
},
thinkingProviderOptions() {
return {};
},
telemetryConfig() {
return undefined;
},
promptCachingConfig() {
return {
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
};
},
activeBackend() {
return 'gateway';
async runAgentLoop() {
return { stopReason: 'natural' };
},
};
}
@ -324,7 +299,7 @@ async function generateDescriptions(input: {
}): Promise<KtxLocalScanEnrichmentResult['descriptionUpdates']> {
const warningSink = input.warnings;
const generator = new KtxDescriptionGenerator({
llmProvider: input.providers.llm,
llmRuntime: input.providers.llmRuntime,
...(input.context.logger ? { logger: input.context.logger } : {}),
...(warningSink
? {
@ -643,7 +618,7 @@ export async function runLocalScanEnrichment(
schema,
context: input.context,
settings: relationshipSettings,
llmProvider: input.providers?.llm ?? null,
llmRuntime: input.providers?.llmRuntime ?? null,
});
await relationshipProgress?.update(

48
packages/context/src/scan/local-scan.test.ts
View file

@ -1,10 +1,10 @@
import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import type { KtxLlmProvider } from '@ktx/llm';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import YAML from 'yaml';
import type { SourceAdapter } from '../ingest/index.js';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import { initKtxProject, type KtxLocalProject, loadKtxProject } from '../project/index.js';
import { filterSnapshotTables, getLocalScanReport, getLocalScanStatus, resolveEnabledTables, runLocalScan } from './local-scan.js';
import type { KtxQueryResult, KtxReadOnlyQueryInput, KtxSchemaSnapshot, KtxSchemaTable } from './types.js';
@ -79,25 +79,11 @@ function relationshipSqlResult(
throw new Error(`Unexpected relationship SQL: ${input.sql}`);
}
function deterministicLlmProvider(): KtxLlmProvider {
function deterministicLlmRuntime(): KtxLlmRuntimePort {
return {
getModel: () => ({ provider: 'deterministic', modelId: 'deterministic' }) as never,
getModelByName: () => ({ provider: 'deterministic', modelId: 'deterministic' }) as never,
cacheMarker: () => undefined,
repairToolCallHandler: (() => undefined) as never,
thinkingProviderOptions: () => ({}),
telemetryConfig: () => undefined,
promptCachingConfig: () => ({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
}),
activeBackend: () => 'gateway',
generateText: vi.fn(async (input) => `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`),
generateObject: vi.fn(async () => ({ pkCandidates: [], fkCandidates: [] }) as never),
runAgentLoop: vi.fn(),
};
}
@ -571,7 +557,7 @@ describe('local scan', () => {
llmProposals: false,
maxLlmTablesPerBatch: 7,
};
const getModel = vi.fn(() => ({ modelId: 'provider/language-model', provider: 'gateway' }));
const generateObject = vi.fn(async () => ({ pkCandidates: [], fkCandidates: [] }));
const connector = {
id: 'test:warehouse',
driver: 'postgres' as const,
@ -650,9 +636,9 @@ describe('local scan', () => {
detectRelationships: true,
connector,
enrichmentProviders: {
llm: {
...deterministicLlmProvider(),
getModel: getModel as never,
llmRuntime: {
...deterministicLlmRuntime(),
generateObject: generateObject as never,
},
embedding: {
dimensions: 8,
@ -668,7 +654,7 @@ describe('local scan', () => {
expect(result.report.relationships.accepted).toBe(1);
expect(result.report.enrichment.llmRelationshipValidation).toBe('skipped');
expect(getModel).not.toHaveBeenCalledWith('candidateExtraction');
expect(generateObject).not.toHaveBeenCalled();
});
it('accepts no-declared-constraint relationships and writes relationship artifacts', async () => {
@ -1206,7 +1192,7 @@ describe('local scan', () => {
mode: 'enriched',
connector,
enrichmentProviders: {
llm: deterministicLlmProvider(),
llmRuntime: deterministicLlmRuntime(),
embedding: {
dimensions: 8,
maxBatchSize: 64,
@ -1314,7 +1300,7 @@ describe('local scan', () => {
return { values: ['1'], nullCount: 0, distinctCount: 1 };
},
};
const llm = deterministicLlmProvider();
const llmRuntime = deterministicLlmRuntime();
const first = await runLocalScan({
project,
@ -1323,7 +1309,7 @@ describe('local scan', () => {
mode: 'enriched',
connector,
enrichmentProviders: {
llm,
llmRuntime,
embedding: {
dimensions: 8,
maxBatchSize: 64,
@ -1344,7 +1330,7 @@ describe('local scan', () => {
});
expect(first.report.enrichment.embeddings).toBe('failed');
const getModel = vi.spyOn(llm, 'getModel');
const generateObject = vi.spyOn(llmRuntime, 'generateObject');
const retry = await runLocalScan({
project,
adapters: [fetchOnlyAdapter()],
@ -1352,7 +1338,7 @@ describe('local scan', () => {
mode: 'enriched',
connector,
enrichmentProviders: {
llm,
llmRuntime,
embedding: {
dimensions: 8,
maxBatchSize: 64,
@ -1373,8 +1359,8 @@ describe('local scan', () => {
failedStages: [],
});
expect(retry.report.enrichment.embeddings).toBe('completed');
expect(getModel).toHaveBeenCalledTimes(1);
expect(getModel).toHaveBeenCalledWith('candidateExtraction');
expect(generateObject).toHaveBeenCalledTimes(1);
expect(generateObject).toHaveBeenCalledWith(expect.objectContaining({ role: 'candidateExtraction' }));
expect(embeddingAttempts).toBe(2);
const reportPath = retry.report.artifactPaths.reportPath;

16
packages/context/src/scan/local-scan.ts
View file

@ -8,7 +8,7 @@ import {
} from '../ingest/index.js';
import {
createLocalKtxEmbeddingProviderFromConfig,
createLocalKtxLlmProviderFromConfig,
createLocalKtxLlmRuntimeFromConfig,
KtxScanEmbeddingPortAdapter,
} from '../llm/index.js';
import type { KtxProjectLlmConfig, KtxScanEnrichmentConfig, KtxScanRelationshipConfig } from '../project/config.js';
@ -150,6 +150,7 @@ interface LocalScanEnrichmentProviderDeps {
createKtxLlmProvider?: typeof createKtxLlmProvider;
createKtxEmbeddingProvider?: typeof createKtxEmbeddingProvider;
env?: NodeJS.ProcessEnv;
projectDir?: string;
}
export function createLocalScanEnrichmentProvidersFromConfig(
@ -165,14 +166,17 @@ export function createLocalScanEnrichmentProvidersFromConfig(
return null;
}
const llm = createLocalKtxLlmProviderFromConfig(llmConfig, deps);
const llmRuntime = createLocalKtxLlmRuntimeFromConfig(llmConfig, {
...deps,
projectDir: deps.projectDir,
});
const embeddingProvider = createLocalKtxEmbeddingProviderFromConfig(config.embeddings, deps);
if (!llm || !embeddingProvider) {
if (!llmRuntime || !embeddingProvider) {
return null;
}
return {
llm,
llmRuntime,
embedding: new KtxScanEmbeddingPortAdapter(embeddingProvider),
};
}
@ -378,7 +382,9 @@ export async function runLocalScan(options: RunLocalScanOptions): Promise<LocalS
connector && (mode !== 'structural' || options.detectRelationships)
? options.enrichmentProviders !== undefined
? options.enrichmentProviders
: createLocalScanEnrichmentProvidersFromConfig(options.project.config.scan.enrichment, options.project.config.llm)
: createLocalScanEnrichmentProvidersFromConfig(options.project.config.scan.enrichment, options.project.config.llm, {
projectDir: options.project.projectDir,
})
: null;
await options.progress?.update(0.15, 'Inspecting database schema');

6
packages/context/src/scan/relationship-benchmarks.ts
View file

@ -6,6 +6,7 @@ import { gunzipSync } from 'node:zlib';
import Database from 'better-sqlite3';
import YAML from 'yaml';
import { z } from 'zod';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import type { KtxEnrichedRelationship, KtxEnrichedSchema, KtxRelationshipType } from './enrichment-types.js';
import { snapshotToKtxEnrichedSchema } from './local-enrichment.js';
import type { KtxRelationshipDiscoveryCandidate } from './relationship-candidates.js';
@ -13,7 +14,6 @@ import {
generateKtxRelationshipDiscoveryCandidates,
mergeKtxRelationshipDiscoveryCandidates,
} from './relationship-candidates.js';
import type { KtxLlmProvider } from '@ktx/llm';
import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
import {
discoverKtxCompositeRelationships,
@ -527,7 +527,7 @@ export function isKtxRelationshipBenchmarkTuningEligible(input: {
}
export function ktxRelationshipBenchmarkDetectorWithLlm(
llmProvider: KtxLlmProvider,
llmRuntime: KtxLlmRuntimePort,
): KtxRelationshipBenchmarkDetector {
return {
async detect(input) {
@ -566,7 +566,7 @@ export function ktxRelationshipBenchmarkDetectorWithLlm(
connectionId: input.snapshot.connectionId,
schema: input.schema,
profile: profiles,
llmProvider,
llmRuntime,
});
const candidates = mergeKtxRelationshipDiscoveryCandidates([
...broadRelationshipCandidates,

59
packages/context/src/scan/relationship-discovery.test.ts
View file

@ -1,6 +1,6 @@
import type { KtxLlmProvider } from '@ktx/llm';
import Database from 'better-sqlite3';
import { afterEach, describe, expect, it, vi } from 'vitest';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import { buildDefaultKtxProjectConfig } from '../project/config.js';
import { snapshotToKtxEnrichedSchema } from './local-enrichment.js';
import {
@ -216,29 +216,11 @@ function connector(executor: InMemorySqliteExecutor | null): KtxScanConnector {
};
}
function llmProvider(): KtxLlmProvider {
const model = { modelId: 'claude-sonnet-4-6', provider: 'anthropic' };
function llmRuntime(output: unknown): KtxLlmRuntimePort {
return {
getModel: vi.fn(() => model as ReturnType<KtxLlmProvider['getModel']>),
getModelByName: vi.fn(() => model as ReturnType<KtxLlmProvider['getModelByName']>),
cacheMarker: vi.fn(),
repairToolCallHandler: vi.fn(),
thinkingProviderOptions: vi.fn(() => ({})),
telemetryConfig: vi.fn(() => undefined),
promptCachingConfig: vi.fn(
() =>
({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
}) as ReturnType<KtxLlmProvider['promptCachingConfig']>,
),
activeBackend: vi.fn(() => 'anthropic' as ReturnType<KtxLlmProvider['activeBackend']>),
generateText: vi.fn(),
generateObject: vi.fn(async () => output) as KtxLlmRuntimePort['generateObject'],
runAgentLoop: vi.fn(),
};
}
@ -505,21 +487,19 @@ describe('production relationship discovery', () => {
INSERT INTO customers (id) VALUES (1), (2);
INSERT INTO orders (id, buyer_ref) VALUES (10, 1), (11, 2);
`);
const generateText = vi.fn(async () => ({
output: {
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.91, rationale: 'Unique customer key.' }],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'buyer_ref',
toTable: 'customers',
toColumn: 'id',
confidence: 0.89,
rationale: 'Buyer reference values align with customer identifiers.',
},
],
},
}));
const llmOutput = {
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.91, rationale: 'Unique customer key.' }],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'buyer_ref',
toTable: 'customers',
toColumn: 'id',
confidence: 0.89,
rationale: 'Buyer reference values align with customer identifiers.',
},
],
};
const result = await discoverKtxRelationships({
connectionId: 'warehouse',
@ -528,8 +508,7 @@ describe('production relationship discovery', () => {
schema: snapshotToKtxEnrichedSchema(llmOnlyRelationshipSnapshot()),
context: { runId: 'llm-relationship-orchestrator' },
settings: relationshipSettings(),
llmProvider: llmProvider(),
generateText,
llmRuntime: llmRuntime(llmOutput),
});
expect(result.llmRelationshipValidation).toBe('completed');

13
packages/context/src/scan/relationship-discovery.ts
View file

@ -1,4 +1,4 @@
import type { KtxLlmProvider } from '@ktx/llm';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import type { KtxScanRelationshipConfig } from '../project/config.js';
import type { KtxEnrichedRelationship, KtxEnrichedSchema, KtxRelationshipUpdate } from './enrichment-types.js';
import {
@ -15,10 +15,7 @@ import {
type KtxResolvedRelationshipDiscoveryCandidate,
resolveKtxRelationshipGraph,
} from './relationship-graph-resolver.js';
import {
type KtxRelationshipLlmProposalGenerateText,
proposeKtxRelationshipCandidatesWithLlm,
} from './relationship-llm-proposal.js';
import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
import {
createKtxRelationshipProfileCache,
type KtxRelationshipProfileArtifact,
@ -42,8 +39,7 @@ export interface DiscoverKtxRelationshipsInput {
schema: KtxEnrichedSchema;
context: KtxScanContext;
settings: KtxScanRelationshipConfig;
llmProvider?: KtxLlmProvider | null;
generateText?: KtxRelationshipLlmProposalGenerateText;
llmRuntime?: KtxLlmRuntimePort | null;
}
export interface DiscoverKtxRelationshipsResult {
@ -246,11 +242,10 @@ export async function discoverKtxRelationships(
connectionId: input.connectionId,
schema: input.schema,
profile,
llmProvider: input.llmProvider ?? null,
llmRuntime: input.llmRuntime ?? null,
settings: {
maxTablesPerBatch: input.settings.maxLlmTablesPerBatch,
},
generateText: input.generateText,
})
: { candidates: [], warnings: [], llmCalls: 0, summary: 'skipped' as const };
const candidates = mergeKtxRelationshipDiscoveryCandidates([

130
packages/context/src/scan/relationship-llm-proposal.test.ts
View file

@ -1,32 +1,14 @@
import type { KtxLlmProvider } from '@ktx/llm';
import { describe, expect, it, vi } from 'vitest';
import type { KtxLlmRuntimePort } from '../llm/index.js';
import type { KtxEnrichedColumn, KtxEnrichedSchema, KtxEnrichedTable } from './enrichment-types.js';
import type { KtxRelationshipProfileArtifact } from './relationship-profiling.js';
import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
function llmProvider(provider = 'anthropic'): KtxLlmProvider {
const model = { modelId: 'claude-sonnet-4-6', provider };
function llmRuntime(output?: unknown): KtxLlmRuntimePort {
return {
getModel: vi.fn(() => model as ReturnType<KtxLlmProvider['getModel']>),
getModelByName: vi.fn(() => model as ReturnType<KtxLlmProvider['getModelByName']>),
cacheMarker: vi.fn(),
repairToolCallHandler: vi.fn(),
thinkingProviderOptions: vi.fn(() => ({})),
telemetryConfig: vi.fn(() => undefined),
promptCachingConfig: vi.fn(
() =>
({
enabled: false,
systemTtl: '1h',
toolsTtl: '1h',
historyTtl: '5m',
cacheSystem: true,
cacheTools: true,
cacheHistory: true,
vertexFallbackTo5m: false,
}) as ReturnType<KtxLlmProvider['promptCachingConfig']>,
),
activeBackend: vi.fn(() => provider as ReturnType<KtxLlmProvider['activeBackend']>),
generateText: vi.fn(),
generateObject: vi.fn(async () => output) as KtxLlmRuntimePort['generateObject'],
runAgentLoop: vi.fn(),
};
}
@ -125,28 +107,25 @@ function profile(): KtxRelationshipProfileArtifact {
describe('relationship LLM proposals', () => {
it('maps valid structured FK proposals into review candidates with rationale evidence', async () => {
const generateText = vi.fn(async () => ({
output: {
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.94, rationale: 'Unique customer identifier.' }],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'buyer_ref',
toTable: 'customers',
toColumn: 'id',
confidence: 0.88,
rationale: 'Buyer reference values match customer identifiers.',
},
],
},
}));
const runtime = llmRuntime({
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.94, rationale: 'Unique customer identifier.' }],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'buyer_ref',
toTable: 'customers',
toColumn: 'id',
confidence: 0.88,
rationale: 'Buyer reference values match customer identifiers.',
},
],
});
const result = await proposeKtxRelationshipCandidatesWithLlm({
connectionId: 'warehouse',
schema: schema(),
profile: profile(),
llmProvider: llmProvider(),
generateText,
llmRuntime: runtime,
});
expect(result.summary).toBe('completed');
@ -164,42 +143,27 @@ describe('relationship LLM proposals', () => {
reasons: ['llm_proposal', 'llm_pk_proposal'],
},
});
expect(generateText).toHaveBeenCalledWith(
expect(runtime.generateObject).toHaveBeenCalledWith(
expect.objectContaining({
system: expect.objectContaining({
role: 'system',
content: expect.stringContaining('You are helping KTX review possible SQL relationships'),
}),
messages: expect.arrayContaining([
expect.objectContaining({
role: 'user',
content: expect.stringContaining('"tables"'),
}),
]),
role: 'candidateExtraction',
system: expect.stringContaining('You are helping KTX review possible SQL relationships'),
prompt: expect.stringContaining('"tables"'),
}),
);
const call = (
generateText.mock.calls as unknown as Array<[{ messages: Array<{ role: string; content: string }> }]>
)[0]?.[0];
const userMessage = call?.messages.find((m) => m.role === 'user');
expect(userMessage?.content).not.toContain('You are helping KTX review possible SQL relationships');
expect(call?.messages.some((m) => m.role === 'system')).toBe(false);
const call = vi.mocked(runtime.generateObject).mock.calls[0]?.[0];
expect(call?.prompt).not.toContain('You are helping KTX review possible SQL relationships');
});
it('skips deterministic providers without calling generateText', async () => {
const generateText = vi.fn();
it('skips when no runtime is configured', async () => {
const result = await proposeKtxRelationshipCandidatesWithLlm({
connectionId: 'warehouse',
schema: schema(),
profile: profile(),
llmProvider: llmProvider('deterministic'),
generateText,
llmRuntime: null,
});
expect(result).toMatchObject({ candidates: [], llmCalls: 0, summary: 'skipped' });
expect(result.warnings).toEqual([]);
expect(generateText).not.toHaveBeenCalled();
});
it('returns recoverable warnings for invalid references and generation failures', async () => {
@ -207,22 +171,19 @@ describe('relationship LLM proposals', () => {
connectionId: 'warehouse',
schema: schema(),
profile: profile(),
llmProvider: llmProvider(),
generateText: vi.fn(async () => ({
output: {
pkCandidates: [],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'missing_column',
toTable: 'customers',
toColumn: 'id',
confidence: 0.7,
rationale: 'Invalid source column.',
},
],
},
})),
llmRuntime: llmRuntime({
pkCandidates: [],
fkCandidates: [
{
fromTable: 'orders',
fromColumn: 'missing_column',
toTable: 'customers',
toColumn: 'id',
confidence: 0.7,
rationale: 'Invalid source column.',
},
],
}),
});
expect(invalidReference.candidates).toEqual([]);
expect(invalidReference.summary).toBe('completed');
@ -235,10 +196,13 @@ describe('relationship LLM proposals', () => {
connectionId: 'warehouse',
schema: schema(),
profile: profile(),
llmProvider: llmProvider(),
generateText: vi.fn(async () => {
throw new Error('model unavailable');
}),
llmRuntime: {
generateText: vi.fn(),
generateObject: vi.fn(async () => {
throw new Error('model unavailable');
}),
runAgentLoop: vi.fn(),
},
});
expect(failed).toMatchObject({ candidates: [], llmCalls: 1, summary: 'failed' });
expect(failed.warnings[0]).toMatchObject({

21
packages/context/src/scan/relationship-llm-proposal.ts
View file

@ -1,7 +1,5 @@
import type { KtxLlmProvider } from '@ktx/llm';
import type { generateText } from 'ai';
import { z } from 'zod';
import { generateKtxObject } from '../llm/index.js';
import { generateKtxObject, type KtxLlmRuntimePort } from '../llm/index.js';
import type { KtxEnrichedColumn, KtxEnrichedSchema, KtxEnrichedTable } from './enrichment-types.js';
import {
normalizeKtxRelationshipName,
@ -32,10 +30,6 @@ const relationshipLlmProposalSchema = z.object({
});
type KtxRelationshipLlmProposalOutput = z.infer<typeof relationshipLlmProposalSchema>;
type GenerateTextInput = Parameters<typeof generateText>[0];
export type KtxRelationshipLlmProposalGenerateText = (
input: GenerateTextInput,
) => Promise<{ text?: string; output?: unknown }>;
export interface KtxRelationshipLlmProposalSettings {
maxTablesPerBatch: number;
@ -48,9 +42,8 @@ export interface ProposeKtxRelationshipCandidatesWithLlmInput {
connectionId: string;
schema: KtxEnrichedSchema;
profile: KtxRelationshipProfileArtifact;
llmProvider: KtxLlmProvider | null;
llmRuntime: KtxLlmRuntimePort | null;
settings?: Partial<KtxRelationshipLlmProposalSettings>;
generateText?: KtxRelationshipLlmProposalGenerateText;
}
export interface KtxRelationshipLlmProposalResult {
@ -77,11 +70,6 @@ function clampConfidence(value: number): number {
return Number(Math.max(0, Math.min(1, value)).toFixed(3));
}
function modelIsDeterministic(llmProvider: KtxLlmProvider): boolean {
const model = llmProvider.getModel('candidateExtraction');
return (model as { provider?: string }).provider === 'deterministic';
}
function findTable(schema: KtxEnrichedSchema, name: string): KtxEnrichedTable | null {
const normalized = name.toLowerCase();
return schema.tables.find((table) => table.ref.name.toLowerCase() === normalized) ?? null;
@ -238,7 +226,7 @@ function generationFailureWarning(error: unknown): KtxScanWarning {
export async function proposeKtxRelationshipCandidatesWithLlm(
input: ProposeKtxRelationshipCandidatesWithLlmInput,
): Promise<KtxRelationshipLlmProposalResult> {
if (!input.llmProvider || modelIsDeterministic(input.llmProvider)) {
if (!input.llmRuntime) {
return { candidates: [], warnings: [], llmCalls: 0, summary: 'skipped' };
}
@ -256,12 +244,11 @@ export async function proposeKtxRelationshipCandidatesWithLlm(
KtxRelationshipLlmProposalOutput,
typeof relationshipLlmProposalSchema
>({
llmProvider: input.llmProvider,
runtime: input.llmRuntime,
role: 'candidateExtraction',
system,
prompt,
schema: relationshipLlmProposalSchema,
generateText: input.generateText,
});
const output = relationshipLlmProposalSchema.parse(generated);
const mapped = mapValidProposals(input.schema, output, settings);

19
packages/context/src/tools/base-tool.ts
View file

@ -1,6 +1,8 @@
import { tool } from 'ai';
import { z, type ZodType } from 'zod';
import { noopLogger, type KtxLogger } from '../core/index.js';
import type { KtxRuntimeToolDescriptor } from '../llm/runtime-port.js';
import { normalizeKtxRuntimeToolOutput } from '../llm/runtime-tools.js';
import type { IngestToolMetadata, ToolSession } from './tool-session.js';
export interface ToolOutput<T = unknown> {
@ -164,6 +166,23 @@ export abstract class BaseTool<TInput extends ZodType = ZodType> {
});
}
toRuntimeTool(context: ToolContext): KtxRuntimeToolDescriptor {
const toolName = this.name;
return {
name: toolName,
description: this.description,
inputSchema: this.inputSchema as unknown as KtxRuntimeToolDescriptor['inputSchema'],
execute: async (params) => {
const callContext = { ...context };
if (!callContext.userId) {
throw new Error('Authentication required: userId must be provided in ToolContext');
}
const parsedInput = this.parseInput(params as Record<string, any>);
return normalizeKtxRuntimeToolOutput(await this.call(parsedInput, callContext));
},
};
}
parseInput(input: Record<string, any>): z.infer<TInput> {
return this.inputSchema.parse(input);
}