From c02b8b37e35f1181e204a8f88c659f81c5fdde4c Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Mon, 11 May 2026 12:03:39 +0200 Subject: [PATCH] feat: use managed runtime for MCP semantic compute --- packages/cli/src/commands/serve-commands.ts | 5 ++ packages/cli/src/index.test.ts | 57 ++++++++++++++++++++ packages/cli/src/serve.test.ts | 60 +++++++++++++++++++++ packages/cli/src/serve.ts | 49 ++++++++++++++--- 4 files changed, 163 insertions(+), 8 deletions(-) diff --git a/packages/cli/src/commands/serve-commands.ts b/packages/cli/src/commands/serve-commands.ts index b7e659fb..28acf32d 100644 --- a/packages/cli/src/commands/serve-commands.ts +++ b/packages/cli/src/commands/serve-commands.ts @@ -1,5 +1,6 @@ import { type Command, InvalidArgumentError } from '@commander-js/extra-typings'; import { type KtxCliCommandContext, resolveCommandProjectDir } from '../cli-program.js'; +import { runtimeInstallPolicyFromFlags } from '../managed-python-command.js'; import type { KtxServeArgs } from '../serve.js'; import { profileMark } from '../startup-profile.js'; @@ -20,6 +21,8 @@ export function registerServeCommands(program: Command, context: KtxCliCommandCo .option('--user-id ', 'Local user id', 'local') .option('--semantic-compute', 'Enable semantic-layer compute', false) .option('--semantic-compute-url ', 'HTTP semantic-layer compute URL') + .option('--yes', 'Install the managed Python runtime without prompting when required', false) + .option('--no-input', 'Disable interactive managed runtime installation') .option('--database-introspection-url ', 'Daemon URL for live-database introspection') .option('--execute-queries', 'Allow semantic-layer query execution', false) .option('--memory-capture', 'Enable memory capture', false) @@ -40,6 +43,8 @@ export function registerServeCommands(program: Command, context: KtxCliCommandCo executeQueries: options.executeQueries === true, memoryCapture: options.memoryCapture === true, memoryModel: options.memoryModel, + cliVersion: context.packageInfo.version, + runtimeInstallPolicy: runtimeInstallPolicyFromFlags(options), }; const runner = context.deps.serveStdio ?? (await import('../serve.js')).runKtxServeStdio; context.setExitCode(await runner(args)); diff --git a/packages/cli/src/index.test.ts b/packages/cli/src/index.test.ts index ff61d91b..b22dbfab 100644 --- a/packages/cli/src/index.test.ts +++ b/packages/cli/src/index.test.ts @@ -599,6 +599,8 @@ describe('runKtxCli', () => { executeQueries: false, memoryCapture: false, memoryModel: undefined, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'prompt', }); }); @@ -2098,10 +2100,65 @@ describe('runKtxCli', () => { executeQueries: true, memoryCapture: true, memoryModel: 'openai/gpt-5.2', + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'prompt', }); expect(serveIo.stderr()).toBe(''); }); + it('routes serve managed runtime install policies', async () => { + const autoIo = makeIo(); + const neverIo = makeIo(); + const conflictIo = makeIo(); + const serveStdio = vi.fn(async () => 0); + + await expect( + runKtxCli(['serve', '--mcp', 'stdio', '--project-dir', tempDir, '--semantic-compute', '--yes'], autoIo.io, { + serveStdio, + }), + ).resolves.toBe(0); + await expect( + runKtxCli(['serve', '--mcp', 'stdio', '--project-dir', tempDir, '--semantic-compute', '--no-input'], neverIo.io, { + serveStdio, + }), + ).resolves.toBe(0); + await expect( + runKtxCli( + ['serve', '--mcp', 'stdio', '--project-dir', tempDir, '--semantic-compute', '--yes', '--no-input'], + conflictIo.io, + { serveStdio }, + ), + ).resolves.toBe(1); + + expect(serveStdio).toHaveBeenNthCalledWith(1, { + mcp: 'stdio', + projectDir: tempDir, + userId: 'local', + semanticCompute: true, + semanticComputeUrl: undefined, + databaseIntrospectionUrl: undefined, + executeQueries: false, + memoryCapture: false, + memoryModel: undefined, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'auto', + }); + expect(serveStdio).toHaveBeenNthCalledWith(2, { + mcp: 'stdio', + projectDir: tempDir, + userId: 'local', + semanticCompute: true, + semanticComputeUrl: undefined, + databaseIntrospectionUrl: undefined, + executeQueries: false, + memoryCapture: false, + memoryModel: undefined, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'never', + }); + expect(conflictIo.stderr()).toContain('Choose only one runtime install mode: --yes or --no-input'); + }); + it('prints dev help for bare dev commands', async () => { const testIo = makeIo(); diff --git a/packages/cli/src/serve.test.ts b/packages/cli/src/serve.test.ts index ca328d07..42c25c84 100644 --- a/packages/cli/src/serve.test.ts +++ b/packages/cli/src/serve.test.ts @@ -6,6 +6,19 @@ import { initKtxProject } from '@ktx/context/project'; import { describe, expect, it, vi } from 'vitest'; import { runKtxServeStdio } from './serve.js'; +function makeManagedRuntimeIo() { + let stdout = ''; + let stderr = ''; + return { + io: { + stdout: { write: (chunk: string) => (stdout += chunk) }, + stderr: { write: (chunk: string) => (stderr += chunk) }, + }, + stdout: () => stdout, + stderr: () => stderr, + }; +} + describe('runKtxServeStdio', () => { it('loads the project, creates local ports, and connects the server to stdio', async () => { const connect = vi.fn().mockResolvedValue(undefined); @@ -241,6 +254,53 @@ describe('runKtxServeStdio', () => { } }); + it('uses managed semantic compute when MCP semantic compute has no explicit HTTP URL', async () => { + const project = { projectDir: '/tmp/ktx-project', config: { connections: {} } } as never; + const semanticLayerCompute = { query: vi.fn(), validateSources: vi.fn(), generateSources: vi.fn() }; + const createManagedSemanticLayerCompute = vi.fn(async () => semanticLayerCompute); + const createContextTools = vi.fn(() => ({ connections: { list: async () => [] } })); + const managedRuntimeIo = makeManagedRuntimeIo(); + + await expect( + runKtxServeStdio( + { + mcp: 'stdio', + projectDir: '/tmp/ktx-project', + userId: 'agent', + semanticCompute: true, + semanticComputeUrl: undefined, + databaseIntrospectionUrl: undefined, + executeQueries: false, + memoryCapture: false, + memoryModel: undefined, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'auto', + }, + { + loadProject: async () => project, + createContextTools, + createManagedSemanticLayerCompute, + managedRuntimeIo: managedRuntimeIo.io, + createServer: vi.fn(() => ({ connect: vi.fn(async () => undefined) }) as never), + createTransport: vi.fn(() => ({}) as never), + stderr: { write: vi.fn() }, + }, + ), + ).resolves.toBe(0); + + expect(createManagedSemanticLayerCompute).toHaveBeenCalledWith({ + cliVersion: '0.2.0', + installPolicy: 'auto', + io: managedRuntimeIo.io, + }); + expect(createContextTools).toHaveBeenCalledWith( + project, + expect.objectContaining({ + semanticLayerCompute, + }), + ); + }); + it('uses the HTTP semantic compute port when a daemon URL is provided', async () => { const project = { projectDir: '/tmp/ktx-project', config: { connections: {} } } as never; const semanticLayerCompute = { query: vi.fn(), validateSources: vi.fn(), generateSources: vi.fn() }; diff --git a/packages/cli/src/serve.ts b/packages/cli/src/serve.ts index 7857e7d6..6e26ecf0 100644 --- a/packages/cli/src/serve.ts +++ b/packages/cli/src/serve.ts @@ -2,7 +2,6 @@ import { createLocalKtxLlmProviderFromConfig } from '@ktx/context'; import { createDefaultLocalQueryExecutor, type KtxSqlQueryExecutorPort } from '@ktx/context/connections'; import { createHttpSemanticLayerComputePort, - createPythonSemanticLayerComputePort, type KtxSemanticLayerComputePort, } from '@ktx/context/daemon'; import { createDefaultLocalIngestAdapters, type LocalIngestMcpOptions } from '@ktx/context/ingest'; @@ -15,8 +14,13 @@ import { createLocalProjectMemoryCapture, type MemoryCaptureService } from '@ktx import { type KtxLocalProject, loadKtxProject } from '@ktx/context/project'; import type { LocalScanMcpOptions } from '@ktx/context/scan'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; +import type { KtxCliIo } from './cli-runtime.js'; import { createKtxCliLocalIngestAdapters } from './local-adapters.js'; import { createKtxCliScanConnector } from './local-scan-connectors.js'; +import { + createManagedPythonSemanticLayerComputePort, + type KtxManagedPythonInstallPolicy, +} from './managed-python-command.js'; import { profileMark } from './startup-profile.js'; profileMark('module:serve'); @@ -31,6 +35,8 @@ export interface KtxServeArgs { executeQueries: boolean; memoryCapture: boolean; memoryModel?: string; + cliVersion?: string; + runtimeInstallPolicy?: KtxManagedPythonInstallPolicy; } interface KtxServeIo { @@ -48,6 +54,8 @@ interface KtxServeDeps { loadProject?: typeof loadKtxProject; createContextTools?: (project: KtxLocalProject, options?: LocalProjectContextToolOptions) => KtxMcpContextPorts; createSemanticLayerCompute?: () => KtxSemanticLayerComputePort; + createManagedSemanticLayerCompute?: typeof createManagedPythonSemanticLayerComputePort; + managedRuntimeIo?: KtxCliIo; createHttpSemanticLayerCompute?: (baseUrl: string) => KtxSemanticLayerComputePort; createIngestAdapters?: typeof createDefaultLocalIngestAdapters; createQueryExecutor?: () => KtxSqlQueryExecutorPort; @@ -57,6 +65,37 @@ interface KtxServeDeps { stderr?: KtxServeIo['stderr']; } +function requiredManagedRuntimeCliVersion(args: KtxServeArgs): string { + if (!args.cliVersion) { + throw new Error('Managed Python semantic compute requires a CLI version.'); + } + return args.cliVersion; +} + +async function createServeSemanticLayerCompute( + args: KtxServeArgs, + deps: KtxServeDeps, +): Promise { + if (!args.semanticCompute) { + return undefined; + } + if (args.semanticComputeUrl) { + return (deps.createHttpSemanticLayerCompute ?? ((baseUrl) => createHttpSemanticLayerComputePort({ baseUrl })))( + args.semanticComputeUrl, + ); + } + if (deps.createSemanticLayerCompute) { + return deps.createSemanticLayerCompute(); + } + const createManagedSemanticLayerCompute = + deps.createManagedSemanticLayerCompute ?? createManagedPythonSemanticLayerComputePort; + return createManagedSemanticLayerCompute({ + cliVersion: requiredManagedRuntimeCliVersion(args), + installPolicy: args.runtimeInstallPolicy ?? 'prompt', + io: deps.managedRuntimeIo ?? process, + }); +} + export async function runKtxServeStdio(args: KtxServeArgs, deps: KtxServeDeps = {}): Promise { const loadProjectFn = deps.loadProject ?? loadKtxProject; const createContextToolsFn = deps.createContextTools ?? createLocalProjectMcpContextPorts; @@ -65,13 +104,7 @@ export async function runKtxServeStdio(args: KtxServeArgs, deps: KtxServeDeps = const stderr = deps.stderr ?? process.stderr; const project = await loadProjectFn({ projectDir: args.projectDir }); - const semanticLayerCompute = args.semanticCompute - ? args.semanticComputeUrl - ? (deps.createHttpSemanticLayerCompute ?? ((baseUrl) => createHttpSemanticLayerComputePort({ baseUrl })))( - args.semanticComputeUrl, - ) - : (deps.createSemanticLayerCompute ?? createPythonSemanticLayerComputePort)() - : undefined; + const semanticLayerCompute = await createServeSemanticLayerCompute(args, deps); const queryExecutor = args.executeQueries ? (deps.createQueryExecutor ?? createDefaultLocalQueryExecutor)() : undefined;