feat(genkit-tools/mcp): MCP fix for Antigravity

genkit-tools/cli/src/commands/mcp.ts

@@ -24,18 +24,35 @@ import { startMcpServer } from '../mcp/server';
 
 interface McpOptions {
   projectRoot?: string;
-  debug?: boolean;
+  debug?: boolean | string;
+  explicitProjectRoot?: boolean;
+  timeout?: string;
 }
 
 /** Command to run MCP server. */
 export const mcp = new Command('mcp')
   .option('--project-root [projectRoot]', 'Project root')
-  .option('-d, --debug', 'debug to file', false)
+  .option('-d, --debug [path]', 'debug to file')
+  .option(
+    '--timeout [timeout]',
+    'Timeout for runtime to start (ms). Default 30000.'
+  )
+  .option(
+    '--explicitProjectRoot',
+    'Whether runtime dependent tools need projectRoot specified. Needed for use with Google Antigravity',
+    false
+  )
   .description('run MCP stdio server (EXPERIMENTAL, subject to change)')
   .action(async (options: McpOptions) => {
     forceStderr();
     if (options.debug) {
-      debugToFile();
+      debugToFile(
+        typeof options.debug === 'string' ? options.debug : undefined
+      );
     }
-    await startMcpServer(options.projectRoot ?? (await findProjectRoot()));
+    await startMcpServer({
+      projectRoot: options.projectRoot ?? (await findProjectRoot()),
+      explicitProjectRoot: options.explicitProjectRoot ?? false,
+      timeout: options.timeout ? parseInt(options.timeout, 10) : undefined,
+    });
   });

genkit-tools/cli/src/mcp/flows.ts

@@ -18,19 +18,31 @@ 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 {
+  McpToolOptions,
+  getCommonSchema,
+  resolveProjectRoot,
+} from './utils.js';
 
-export function defineFlowTools(server: McpServer, manager: McpRuntimeManager) {
+export function defineFlowTools(server: McpServer, options: McpToolOptions) {
   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(options.explicitProjectRoot),
     },
-    async () => {
+    async (opts) => {
       await record(new McpRunToolEvent('list_flows'));
-      const runtimeManager = await manager.getManager();
+      const rootOrError = resolveProjectRoot(
+        options.explicitProjectRoot,
+        opts,
+        options.projectRoot
+      );
+      if (typeof rootOrError !== 'string') return rootOrError;
+
+      const runtimeManager = await options.manager.getManager(rootOrError);
       const actions = await runtimeManager.listActions();
 
       let flows = '';
@@ -56,21 +68,28 @@ export function defineFlowTools(server: McpServer, manager: McpRuntimeManager) {
     {
       title: 'Run Flow',
       description: 'Runs the flow with the provided input',
-      inputSchema: {
+      inputSchema: getCommonSchema(options.explicitProjectRoot, {
         flowName: z.string().describe('name of the flow'),
         input: z
           .string()
           .describe(
             '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(
+        options.explicitProjectRoot,
+        opts,
+        options.projectRoot
+      );
+      if (typeof rootOrError !== 'string') return rootOrError;
+      const { flowName, input } = opts;
 
       try {
-        const runtimeManager = await manager.getManager();
+        const runtimeManager = await options.manager.getManager(rootOrError);
         const response = await runtimeManager.runAction({
           key: `/flow/${flowName}`,
           input: input !== undefined ? JSON.parse(input) : undefined,

genkit-tools/cli/src/mcp/runtime.ts

@@ -18,89 +18,121 @@ 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 {
+  McpToolOptions,
+  getCommonSchema,
+  resolveProjectRoot,
+} from './utils.js';
 
-export function defineRuntimeTools(
-  server: McpServer,
-  manager: McpRuntimeManager
-) {
+export function defineRuntimeTools(server: McpServer, options: McpToolOptions) {
   server.registerTool(
     'start_runtime',
     {
       title: 'Starts a Genkit runtime process',
       description: `Use this to start a Genkit runtime process (This is typically the entry point to the users app). Once started, the runtime will be picked up by the \`genkit start\` command to power the Dev UI features like model and flow playgrounds. The inputSchema for this tool matches the function prototype for \`NodeJS.child_process.spawn\`.
 
       Examples: 
-        {command: 'go', args: ['run', 'main.go']}
-        {command: 'npm', args: ['run', 'dev']}`,
-      inputSchema: {
+        {command: "go", args: ["run", "main.go"]}
+        {command: "npm", args: ["run", "dev"]}
+        {command: "npm", args: ["run", "dev"], projectRoot: "path/to/project"}`,
+      inputSchema: getCommonSchema(options.explicitProjectRoot, {
         command: z.string().describe('The command to run'),
         args: z
           .array(z.string())
           .describe(
-            'List of command line arguments. IMPORTANT: This must be an array of strings, not a single string.'
+            'The array of string args for the command to run. Eg: `["run", "dev"]`.'
           ),
-      },
+      }),
     },
-    async ({ command, args }) => {
+    async (opts) => {
       await record(new McpRunToolEvent('start_runtime'));
-      await manager.getManagerWithDevProcess(command, args);
+      const rootOrError = resolveProjectRoot(
+        options.explicitProjectRoot,
+        opts,
+        options.projectRoot
+      );
+      if (typeof rootOrError !== 'string') return rootOrError;
 
-      return {
-        content: [{ type: 'text', text: `Done.` }],
-      };
-    }
-  );
-
-  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) {
+      try {
+        await options.manager.getManagerWithDevProcess({
+          projectRoot: rootOrError,
+          command: opts.command,
+          args: opts.args,
+          explicitProjectRoot: options.explicitProjectRoot,
+          timeout: options.timeout,
+        });
+      } catch (err) {
         return {
-          isError: true,
           content: [
-            { type: 'text', text: `No runtime process currently running.` },
+            {
+              type: 'text',
+              text:
+                'Error creating runtime manager: ' +
+                (err instanceof Error ? err.stack : JSON.stringify(err)),
+            },
           ],
+          isError: true,
         };
       }
 
-      await runtimeManager.processManager?.kill();
       return {
-        content: [{ type: 'text', text: `Done.` }],
+        content: [
+          {
+            type: 'text',
+            text: `Done.`,
+          },
+        ],
       };
     }
   );
 
-  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) {
+  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(options.explicitProjectRoot),
+      },
+      async (opts) => {
+        await record(new McpRunToolEvent(name));
+        const rootOrError = resolveProjectRoot(
+          options.explicitProjectRoot,
+          opts,
+          options.projectRoot
+        );
+        if (typeof rootOrError !== 'string') return rootOrError;
+
+        const runtimeManager = await options.manager.getManager(rootOrError);
+        if (!runtimeManager.processManager) {
+          return {
+            isError: true,
+            content: [
+              { type: 'text', text: `No runtime process currently running.` },
+            ],
+          };
+        }
+
+        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'
   );
 }

genkit-tools/cli/src/mcp/server.ts

@@ -23,24 +23,36 @@ import { defineInitPrompt } from './prompts/init';
 import { defineRuntimeTools } from './runtime';
 import { defineTraceTools } from './trace';
 import { defineUsageGuideTool } from './usage';
-import { McpRuntimeManager } from './util';
+import { McpRuntimeManager, McpToolOptions } from './utils';
 
-export async function startMcpServer(projectRoot: string) {
+export async function startMcpServer(params: {
+  projectRoot: string;
+  explicitProjectRoot: boolean;
+  timeout?: number;
+}) {
+  const { projectRoot, explicitProjectRoot, timeout } = params;
   logger.info(`Starting MCP server in: ${projectRoot}`);
+
   const server = new McpServer({
     name: 'Genkit MCP',
     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);
+  const manager = new McpRuntimeManager();
+  const options: McpToolOptions = {
+    projectRoot,
+    explicitProjectRoot,
+    timeout,
+    manager,
+  };
+
+  defineFlowTools(server, options);
+  defineTraceTools(server, options);
+  defineRuntimeTools(server, options);
 
   return new Promise(async (resolve) => {
     const transport = new StdioServerTransport();

genkit-tools/cli/src/mcp/trace.ts

@@ -18,30 +18,38 @@ 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 {
+  McpToolOptions,
+  getCommonSchema,
+  resolveProjectRoot,
+} from './utils.js';
 
-export function defineTraceTools(
-  server: McpServer,
-  manager: McpRuntimeManager
-) {
+export function defineTraceTools(server: McpServer, options: McpToolOptions) {
   server.registerTool(
     'get_trace',
     {
       title: 'Get Genkit Trace',
       description: 'Returns the trace details',
-      inputSchema: {
+      inputSchema: getCommonSchema(options.explicitProjectRoot, {
         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(
+        options.explicitProjectRoot,
+        opts,
+        options.projectRoot
+      );
+      if (typeof rootOrError !== 'string') return rootOrError;
+      const { traceId } = opts;
 
       try {
-        const runtimeManager = await manager.getManager();
+        const runtimeManager = await options.manager.getManager(rootOrError);
         const response = await runtimeManager.getTrace({ traceId });
         return {
           content: [