diff --git a/genkit-tools/cli/src/mcp/flows.ts b/genkit-tools/cli/src/mcp/flows.ts index 5b4f5a8d96..8fbcf2648e 100644 --- a/genkit-tools/cli/src/mcp/flows.ts +++ b/genkit-tools/cli/src/mcp/flows.ts @@ -18,19 +18,27 @@ import { record } from '@genkit-ai/tools-common/utils'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp'; import z from 'zod'; import { McpRunToolEvent } from './analytics.js'; -import { McpRuntimeManager } from './util.js'; +import { + McpRuntimeManager, + getCommonSchema, + resolveProjectRoot, +} from './util.js'; -export function defineFlowTools(server: McpServer, manager: McpRuntimeManager) { +export function defineFlowTools(server: McpServer, projectRoot: string) { server.registerTool( 'list_flows', { title: 'List Genkit Flows', description: 'Use this to discover available Genkit flows or inspect the input schema of Genkit flows to know how to successfully call them.', + inputSchema: getCommonSchema(), }, - async () => { + async (opts) => { await record(new McpRunToolEvent('list_flows')); - const runtimeManager = await manager.getManager(); + const rootOrError = resolveProjectRoot(opts, projectRoot); + if (typeof rootOrError !== 'string') return rootOrError; + + const runtimeManager = await McpRuntimeManager.getManager(rootOrError); const actions = await runtimeManager.listActions(); let flows = ''; @@ -56,7 +64,7 @@ export function defineFlowTools(server: McpServer, manager: McpRuntimeManager) { { title: 'Run Flow', description: 'Runs the flow with the provided input', - inputSchema: { + inputSchema: getCommonSchema({ flowName: z.string().describe('name of the flow'), input: z .string() @@ -64,13 +72,16 @@ export function defineFlowTools(server: McpServer, manager: McpRuntimeManager) { 'Flow input as JSON object encoded as string (it will be passed through `JSON.parse`). Must conform to the schema.' ) .optional(), - }, + }), }, - async ({ flowName, input }) => { + async (opts) => { await record(new McpRunToolEvent('run_flow')); + const rootOrError = resolveProjectRoot(opts, projectRoot); + if (typeof rootOrError !== 'string') return rootOrError; + const { flowName, input } = opts; try { - const runtimeManager = await manager.getManager(); + const runtimeManager = await McpRuntimeManager.getManager(rootOrError); const response = await runtimeManager.runAction({ key: `/flow/${flowName}`, input: input !== undefined ? JSON.parse(input) : undefined, diff --git a/genkit-tools/cli/src/mcp/runtime.ts b/genkit-tools/cli/src/mcp/runtime.ts index daa52e622d..c3bf6ea704 100644 --- a/genkit-tools/cli/src/mcp/runtime.ts +++ b/genkit-tools/cli/src/mcp/runtime.ts @@ -18,12 +18,13 @@ import { record } from '@genkit-ai/tools-common/utils'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp'; import { z } from 'zod'; import { McpRunToolEvent } from './analytics.js'; -import { McpRuntimeManager } from './util.js'; +import { + McpRuntimeManager, + getCommonSchema, + resolveProjectRoot, +} from './util.js'; -export function defineRuntimeTools( - server: McpServer, - manager: McpRuntimeManager -) { +export function defineRuntimeTools(server: McpServer, projectRoot: string) { server.registerTool( 'start_runtime', { @@ -33,14 +34,21 @@ export function defineRuntimeTools( Examples: {command: 'go', args: ['run', 'main.go']} {command: 'npm', args: ['run', 'dev']}`, - inputSchema: { + inputSchema: getCommonSchema({ command: z.string(), args: z.array(z.string()), - }, + }), }, - async ({ command, args }) => { + async (opts) => { await record(new McpRunToolEvent('start_runtime')); - await manager.getManagerWithDevProcess(command, args); + const rootOrError = resolveProjectRoot(opts, projectRoot); + if (typeof rootOrError !== 'string') return rootOrError; + + await McpRuntimeManager.getManagerWithDevProcess( + rootOrError, + opts.command, + opts.args + ); return { content: [{ type: 'text', text: `Done.` }], @@ -48,55 +56,49 @@ export function defineRuntimeTools( } ); - server.registerTool( - 'kill_runtime', - { - title: 'Kills any existing Genkit runtime process', - description: - 'Use this to stop an existing runtime that was started using the `start_runtime` tool', - }, - async () => { - await record(new McpRunToolEvent('kill_runtime')); - const runtimeManager = await manager.getManager(); - if (!runtimeManager.processManager) { - return { - isError: true, - content: [ - { type: 'text', text: `No runtime process currently running.` }, - ], - }; - } + const registerControlTool = ( + name: string, + title: string, + action: 'kill' | 'restart' + ) => { + server.registerTool( + name, + { + title, + description: `Use this to ${action} an existing runtime that was started using the \`start_runtime\` tool`, + inputSchema: getCommonSchema(), + }, + async (opts) => { + await record(new McpRunToolEvent(name)); + const rootOrError = resolveProjectRoot(opts, projectRoot); + if (typeof rootOrError !== 'string') return rootOrError; - await runtimeManager.processManager?.kill(); - return { - content: [{ type: 'text', text: `Done.` }], - }; - } - ); + const runtimeManager = await McpRuntimeManager.getManager(rootOrError); + if (!runtimeManager.processManager) { + return { + isError: true, + content: [ + { type: 'text', text: `No runtime process currently running.` }, + ], + }; + } - server.registerTool( - 'restart_runtime', - { - title: 'Restarts any existing Genkit runtime process', - description: - 'Use this to restart an existing runtime that was started using the `start_runtime` tool', - }, - async () => { - await record(new McpRunToolEvent('restart_runtime')); - const runtimeManager = await manager.getManager(); - if (!runtimeManager.processManager) { + await runtimeManager.processManager[action](); return { - isError: true, - content: [ - { type: 'text', text: `No runtime process currently running.` }, - ], + content: [{ type: 'text', text: `Done.` }], }; } + ); + }; - await runtimeManager.processManager?.restart(); - return { - content: [{ type: 'text', text: `Done.` }], - }; - } + registerControlTool( + 'kill_runtime', + 'Kills any existing Genkit runtime process', + 'kill' + ); + registerControlTool( + 'restart_runtime', + 'Restarts any existing Genkit runtime process', + 'restart' ); } diff --git a/genkit-tools/cli/src/mcp/server.ts b/genkit-tools/cli/src/mcp/server.ts index 6aecf2a218..7d8d663c3e 100644 --- a/genkit-tools/cli/src/mcp/server.ts +++ b/genkit-tools/cli/src/mcp/server.ts @@ -23,7 +23,7 @@ import { defineInitPrompt } from './prompts/init'; import { defineRuntimeTools } from './runtime'; import { defineTraceTools } from './trace'; import { defineUsageGuideTool } from './usage'; -import { McpRuntimeManager } from './util'; +import { isAntigravity, McpRuntimeManager } from './util'; export async function startMcpServer(projectRoot: string) { const server = new McpServer({ @@ -31,21 +31,23 @@ export async function startMcpServer(projectRoot: string) { version: '0.0.2', }); - const manager = new McpRuntimeManager(projectRoot); - await defineDocsTool(server); await defineUsageGuideTool(server); defineInitPrompt(server); - defineRuntimeTools(server, manager); - defineFlowTools(server, manager); - defineTraceTools(server, manager); + defineFlowTools(server, projectRoot); + defineTraceTools(server, projectRoot); + // Disable runtime tools in AGY. Something about AGY env is messing with + // runtime discoverability. + if (!isAntigravity) { + defineRuntimeTools(server, projectRoot); + } return new Promise(async (resolve) => { const transport = new StdioServerTransport(); const cleanup = async () => { try { - await manager.kill(); + await McpRuntimeManager.kill(); } catch (e) { // ignore } @@ -54,7 +56,7 @@ export async function startMcpServer(projectRoot: string) { }; transport.onclose = async () => { try { - await manager.kill(); + await McpRuntimeManager.kill(); } catch (e) { // ignore } diff --git a/genkit-tools/cli/src/mcp/trace.ts b/genkit-tools/cli/src/mcp/trace.ts index 7e8520cf40..7ad30c4c20 100644 --- a/genkit-tools/cli/src/mcp/trace.ts +++ b/genkit-tools/cli/src/mcp/trace.ts @@ -18,30 +18,34 @@ import { record } from '@genkit-ai/tools-common/utils'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp'; import z from 'zod'; import { McpRunToolEvent } from './analytics.js'; -import { McpRuntimeManager } from './util.js'; +import { + McpRuntimeManager, + getCommonSchema, + resolveProjectRoot, +} from './util.js'; -export function defineTraceTools( - server: McpServer, - manager: McpRuntimeManager -) { +export function defineTraceTools(server: McpServer, projectRoot: string) { server.registerTool( 'get_trace', { title: 'Get Genkit Trace', description: 'Returns the trace details', - inputSchema: { + inputSchema: getCommonSchema({ traceId: z .string() .describe( 'trace id (typically returned after running a flow or other actions)' ), - }, + }), }, - async ({ traceId }) => { + async (opts) => { await record(new McpRunToolEvent('get_trace')); + const rootOrError = resolveProjectRoot(opts, projectRoot); + if (typeof rootOrError !== 'string') return rootOrError; + const { traceId } = opts; try { - const runtimeManager = await manager.getManager(); + const runtimeManager = await McpRuntimeManager.getManager(rootOrError); const response = await runtimeManager.getTrace({ traceId }); return { content: [ diff --git a/genkit-tools/cli/src/mcp/util.ts b/genkit-tools/cli/src/mcp/util.ts index f39ac569bf..0425da33e4 100644 --- a/genkit-tools/cli/src/mcp/util.ts +++ b/genkit-tools/cli/src/mcp/util.ts @@ -15,39 +15,74 @@ */ import { RuntimeManager } from '@genkit-ai/tools-common/manager'; +import { z } from 'zod'; import { startDevProcessManager, startManager } from '../utils/manager-utils'; +export const isAntigravity = !!process.env.ANTIGRAVITY_ENV; + +export function getCommonSchema(shape: z.ZodRawShape = {}): z.ZodRawShape { + return !isAntigravity + ? shape + : { + projectRoot: z + .string() + .describe( + 'The path to the current project root (a.k.a workspace directory or project directory)' + ), + ...shape, + }; +} + +export function resolveProjectRoot( + opts: { + [x: string]: any; + }, + fallback: string +): string | { content: any[]; isError: boolean } { + if (isAntigravity && !opts?.projectRoot) { + return { + content: [ + { type: 'text', text: 'Project root is required for this tool.' }, + ], + isError: true, + }; + } + return opts?.projectRoot ?? fallback; +} + /** Genkit Runtime manager specifically for the MCP server. Allows lazy * initialization and dev process manangement. */ export class McpRuntimeManager { - private manager: RuntimeManager | undefined; - - constructor(private projectRoot: string) {} + private static manager: RuntimeManager | undefined; + private static currentProjectRoot: string | undefined; - async getManager() { - if (!this.manager) { - this.manager = await startManager( - this.projectRoot, - true /* manageHealth */ - ); + static async getManager(projectRoot: string) { + if (this.manager && this.currentProjectRoot === projectRoot) { + return this.manager; + } + if (this.manager) { + await this.manager.stop(); } + this.manager = await startManager(projectRoot, true /* manageHealth */); + this.currentProjectRoot = projectRoot; return this.manager; } - async getManagerWithDevProcess(command: string, args: string[]) { + static async getManagerWithDevProcess( + projectRoot: string, + command: string, + args: string[] + ) { if (this.manager) { await this.manager.stop(); } - const devManager = await startDevProcessManager( - this.projectRoot, - command, - args - ); + const devManager = await startDevProcessManager(projectRoot, command, args); this.manager = devManager.manager; + this.currentProjectRoot = projectRoot; return this.manager; } - async kill() { + static async kill() { if (this.manager) { await this.manager.stop(); } diff --git a/genkit-tools/cli/src/utils/manager-utils.ts b/genkit-tools/cli/src/utils/manager-utils.ts index 7badf39a11..feecc73514 100644 --- a/genkit-tools/cli/src/utils/manager-utils.ts +++ b/genkit-tools/cli/src/utils/manager-utils.ts @@ -82,7 +82,9 @@ export async function startDevProcessManager( projectRoot, processManager, }); - const processPromise = processManager.start(); + const processPromise = processManager.start({ + cwd: projectRoot, + }); return { manager, processPromise }; } diff --git a/genkit-tools/common/src/manager/manager.ts b/genkit-tools/common/src/manager/manager.ts index 82ee715b17..14761b7746 100644 --- a/genkit-tools/common/src/manager/manager.ts +++ b/genkit-tools/common/src/manager/manager.ts @@ -183,7 +183,7 @@ export class RuntimeManager { throw new Error( input?.runtimeId ? `No runtime found with ID ${input.runtimeId}.` - : 'No runtimes found. Make sure your app is running using `genkit start -- ...`. See getting started documentation.' + : 'No runtimes found. Make sure your app is running using the `start_runtime` MCP tool or the CLI: `genkit start -- ...`. See getting started documentation.' ); } const response = await axios @@ -206,7 +206,7 @@ export class RuntimeManager { throw new Error( input.runtimeId ? `No runtime found with ID ${input.runtimeId}.` - : 'No runtimes found. Make sure your app is running using `genkit start -- ...`. See getting started documentation.' + : 'No runtimes found. Make sure your app is running using the `start_runtime` MCP tool or the CLI: `genkit start -- ...`. See getting started documentation.' ); } if (streamingCallback) { diff --git a/genkit-tools/common/src/manager/process-manager.ts b/genkit-tools/common/src/manager/process-manager.ts index 3c4f905fb7..3a68bc5825 100644 --- a/genkit-tools/common/src/manager/process-manager.ts +++ b/genkit-tools/common/src/manager/process-manager.ts @@ -25,6 +25,7 @@ export interface AppProcessStatus { } export interface ProcessManagerStartOptions { + cwd?: string; nonInteractive?: boolean; } @@ -50,6 +51,7 @@ export class ProcessManager { return new Promise((resolve, reject) => { this._status = 'running'; this.appProcess = spawn(this.command, this.args, { + cwd: options?.cwd, env: { ...process.env, ...this.env,