Review policies content starter

content/60_mcp_review_aws_iam_policies/.aillyrc

@@ -0,0 +1,12 @@
+---
+isolated: true
+mcp:
+  awslabs.aws-documentation-mcp-server:
+    command: "uvx"
+    args:
+      - awslabs.aws-documentation-mcp-server@latest
+    env:
+      FASTMCP_LOG_LEVEL: error
+---
+
+You are an AWS Docs Writer, preparing technical materials for analysis and summary by an LLM.

core/package.json

@@ -32,6 +32,7 @@
     "@aws-sdk/credential-providers": "^3.572.0",
     "@davidsouther/jiffies": "^2.2.5",
     "@dqbd/tiktoken": "^1.0.7",
+    "@modelcontextprotocol/sdk": "^1.12.1",
     "esbuild": "0.25.4",
     "gitignore-parser": "^0.0.2",
     "gray-matter": "^4.0.3",

core/src/actions/prompt_thread.test.ts

@@ -7,10 +7,12 @@ import { range } from "@davidsouther/jiffies/lib/cjs/range.js";
 import { cleanState } from "@davidsouther/jiffies/lib/cjs/scope/state";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { getPlugin, makePipelineSettings } from "..";
-import { loadContent } from "../content/content.js";
+import { type Content, loadContent } from "../content/content.js";
 import { getEngine } from "../engine/index.js";
 import { TIMEOUT } from "../engine/noop.js";
+import type { Tool } from "../engine/tool";
 import { LOGGER } from "../index.js";
+import { MCPClient, type MCPServersConfig } from "../mcp";
 import { withResolvers } from "../util.js";
 import {
   PromptThread,
@@ -107,9 +109,7 @@ describe("generateOne", () => {
     );
     expect(state.logger.info).toHaveBeenCalledWith("Skipping /b.txt");
     state.logger.info.mockClear();
-    //   });
 
-    //   it("generates others", async () => {
     const content = state.context["/c.txt"];
     expect(content.response).toBeUndefined();
     await generateOne(
@@ -168,8 +168,6 @@ describe("PromptThread", () => {
       system: [],
       meta: { isolated: true },
     });
-    state.logger.debug.mockClear();
-    state.logger.info.mockClear();
     const content = [...Object.values(context)];
     const plugin = await (await getPlugin("none")).default(
       state.engine,
@@ -186,6 +184,11 @@ describe("PromptThread", () => {
     expect(thread.finished).toBe(0);
     expect(thread.errors.length).toBe(0);
 
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
     await Promise.resolve();
     await Promise.resolve();
     await Promise.resolve();
@@ -208,8 +211,6 @@ describe("PromptThread", () => {
   it("runs sequence", async () => {
     const settings = await makePipelineSettings({ root: "/" });
     const context = await loadContent(state.fs);
-    state.logger.debug.mockClear();
-    state.logger.info.mockClear();
     const content = [...Object.values(context)];
     const plugin = await (await getPlugin("none")).default(
       state.engine,
@@ -232,4 +233,74 @@ describe("PromptThread", () => {
     expect(thread.finished).toBe(3);
     expect(thread.errors.length).toBe(0);
   });
+
+  it("runs with MCP", async () => {
+    const settings = await makePipelineSettings({
+      root: "/",
+      isolated: true,
+      combined: true,
+    });
+    const fs = new FileSystem(
+      new ObjectFileSystemAdapter({
+        ".ailly.md": "---\nmcp:\n  mock:\n    type: mock\n---\n",
+        "a.txt": "USE add WITH 40 7",
+      }),
+    );
+    const context = await loadContent(fs);
+    const content = [...Object.values(context)];
+    const client = new (class MockClient extends MCPClient {
+      initialize(_config?: MCPServersConfig): Promise<void> {
+        return Promise.resolve();
+      }
+      getAllTools(): Tool[] {
+        return [
+          {
+            name: "add",
+            parameters: {
+              type: "object",
+              properties: { args: { type: "array" } },
+            },
+          },
+        ];
+      }
+
+      async invokeTool(
+        toolName: string,
+        parameters: Record<string, unknown>,
+        _context?: string,
+      ): Promise<string> {
+        if (toolName === "add") {
+          const { args } = parameters;
+          const nums = (args as string[]).map(Number);
+          const sum = nums.reduce((a, b) => a + b, 0);
+          return `${sum}`;
+        }
+        return "";
+      }
+    })();
+    for (const f of content) {
+      f.context.mcpClient = client;
+    }
+    const plugin = await (await getPlugin("none")).default(
+      state.engine,
+      settings,
+    );
+    const thread = PromptThread.run(
+      content,
+      context,
+      settings,
+      state.engine,
+      plugin,
+    );
+
+    await thread.allSettled();
+
+    expect(thread.isDone).toBe(true);
+    expect(thread.finished).toBe(1);
+    expect(thread.errors.length).toBe(0);
+
+    expect(content.at(-1)?.response).toBe(
+      'USING TOOL add WITH ARGS [40, 7]\nTOOL RETURNED "47"\n',
+    );
+  });
 });

core/src/actions/prompt_thread.ts

@@ -14,6 +14,7 @@ import {
   LOGGER,
   type PipelineSettings,
 } from "../index.js";
+import { MCPClient } from "../mcp";
 import type { Plugin } from "../plugin/index.js";
 
 export interface PromptThreadsSummary {
@@ -117,6 +118,20 @@ export class PromptThread {
         this.settings.templateView,
       );
     }
+
+    if (c.meta?.mcp && !c.context.mcpClient) {
+      const mcpClient = new MCPClient();
+
+      // Extract MCP information from meta server and attach MCP Clients to Context
+      await mcpClient.initialize({ servers: c.meta.mcp });
+
+      // Assign all the tools to meta.tools
+      c.meta.tools = mcpClient.getAllTools();
+
+      // Attach MCP Clients to context
+      c.context.mcpClient = mcpClient;
+    }
+
     try {
       await this.template(c, this.view);
       await this.plugin.augment(c);
@@ -227,13 +242,39 @@ export function generateOne(
     },
   };
 
-  try {
+  async function runWithTools() {
     const generator = engine.generate(c, settings);
     assertExists(c.responseStream).resolve(generator.stream);
-    return generator.done.finally(() => {
-      c.response = generator.message();
+    await generator.done.finally(() => {
+      c.response = (c.response ?? "") + generator.message();
       assertExists(c.meta).debug = { ...c.meta?.debug, ...generator.debug() };
     });
+    const { toolUse } = generator.debug();
+    if (toolUse && c.meta && c.context.mcpClient) {
+      const result = await c.context.mcpClient.invokeTool(
+        toolUse.name,
+        toolUse.input,
+      );
+
+      c.meta.messages ??= [];
+
+      c.meta.messages.push({
+        role: "assistant",
+        content: c.response ?? "",
+        toolUse: {
+          name: toolUse.name,
+          input: toolUse.input,
+          result: JSON.stringify(result),
+          id: toolUse.id,
+        },
+      });
+      c.meta.continue = true;
+      return runWithTools();
+    }
+  }
+
+  try {
+    return runWithTools();
   } catch (err) {
     LOGGER.error(`Uncaught error in ${engine.name} generator`, { err });
     if (c.meta?.debug) {

core/src/content/content.test.ts

@@ -364,6 +364,65 @@ describe("Loading aillyrc", () => {
       },
     });
   });
+
+  test("it loads MCP server config", async () => {
+    const fs = new FileSystem(
+      new ObjectFileSystemAdapter({
+        ".aillyrc": `---
+mcp:
+  base:
+    type: stdio
+    command: base
+---
+`,
+        root: {
+          ".aillyrc": `---
+mcp:
+  root:
+    type: http
+    url: localhost:8080
+---
+`,
+
+          a: "a",
+        },
+      }),
+    );
+
+    const content = await loadContent(fs, {
+      meta: { context: "conversation" },
+    });
+
+    expect(content).toMatchObject({
+      "/root/a": {
+        name: "a",
+        path: "/root/a",
+        outPath: "/root/a.ailly.md",
+        prompt: "a",
+        context: {
+          system: [],
+          view: {},
+        },
+        meta: {
+          combined: false,
+          context: "conversation",
+          parent: "root",
+          root: "/root",
+          text: "a",
+          mcp: {
+            base: {
+              type: "stdio",
+              command: "base",
+            },
+            root: {
+              type: "http",
+              url: "localhost:8080",
+            },
+          },
+        },
+      },
+    });
+  });
 });
 
 describe("Loading", () => {
@@ -732,6 +791,7 @@ describe("Writing", () => {
     });
   });
 });
+
 describe("combined vs separate", () => {
   test("it loads combined prompt and responses", async () => {
     const fs = new FileSystem(

core/src/content/content.ts

@@ -1,6 +1,7 @@
 import { dirname, join } from "node:path";
 import matter, { type GrayMatterFile } from "gray-matter";
 import * as YAML from "yaml";
+import type { MCPClient, MCPConfig } from "../mcp";
 
 import {
   type FileSystem,
@@ -12,11 +13,9 @@ import {
   isAbsolute,
 } from "@davidsouther/jiffies/lib/cjs/fs.js";
 
-import {
-  assertExists,
-  checkExhaustive,
-} from "@davidsouther/jiffies/lib/cjs/assert";
+import { checkExhaustive } from "@davidsouther/jiffies/lib/cjs/assert";
 import type { EngineDebug, Message } from "../engine/index.js";
+import type { Tool } from "../engine/tool.js";
 import { LOGGER } from "../index.js";
 import {
   type PromiseWithResolvers,
@@ -60,6 +59,7 @@ export interface Context {
   edit?:
     | { start: number; end: number; file: string }
     | { after: number; file: string };
+  mcpClient?: MCPClient;
 }
 
 // Additional useful metadata.
@@ -80,6 +80,8 @@ export interface ContentMeta {
   prompt?: string;
   temperature?: number;
   maxTokens?: number;
+  tools?: Tool[];
+  mcp?: MCPConfig;
 }
 
 export type AillyEditReplace = { start: number; end: number; file: string };
@@ -359,13 +361,20 @@ export async function loadAillyRc(
 
 async function mergeAillyRc(
   content_meta: ContentMeta,
-  data: { [key: string]: unknown },
+  data: Record<string, unknown>,
   content: string,
   view: View,
   system: System[],
   fs: FileSystem,
 ): Promise<[System[], ContentMeta]> {
+  const mcp: MCPConfig = {
+    ...(content_meta.mcp ?? {}),
+    ...((data.mcp as MCPConfig) ?? {}),
+  };
   const meta = { ...{ parent: "root" as const }, ...content_meta, ...data };
+  if (Object.keys(mcp)) {
+    meta.mcp = mcp;
+  }
   switch (meta.parent) {
     case "root":
       if (!(content === "" && Object.keys(view).length === 0))