feat!: Support Claude 4.6 generation models

Author: niallthomson

AGENTS.md

@@ -15,6 +15,7 @@ src/                    # Source code
 ├── ai/                # AI model integration and prompts
 │   ├── model/         # Bedrock client and rate limiting
 │   ├── prompts/       # Prompt templates for various operations
+│   ├── tools/         # Tool definitions for model interactions
 │   └── mcp/           # Model Context Protocol server
 ├── check/             # Content validation (lint, links, images)
 ├── content/           # Content processing

README.md

@@ -398,6 +398,17 @@ Toolkit for Markdown uses AWS Bedrock for AI processing. Ensure the following is
 - AWS credentials configured
 - Access to Bedrock models in the appropriate AWS account
 
+**Supported Models:**
+
+The tool requires Anthropic Claude models via Amazon Bedrock. The following model families are supported:
+
+- Claude Opus 4 (including 4.1, 4.5, 4.6)
+- Claude Sonnet 4 (including 4.5, 4.6)
+- Claude Haiku 4.5
+- Claude 3.7 Sonnet
+
+Model IDs are matched using a prefix, so both standard and cross-region inference profiles (e.g. `us.anthropic.claude-sonnet-4-5-20250929-v1:0`) are accepted. Other model providers and older Claude versions (such as Claude 3.5 Haiku) are not supported.
+
 ## Commands
 
 ### `review`

package.json

@@ -33,7 +33,7 @@
   "author": "",
   "dependencies": {
     "@aws-sdk/client-bedrock-runtime": "^3.985.0",
-    "@modelcontextprotocol/sdk": "^1.18.2",
+    "@modelcontextprotocol/sdk": "^1.27.1",
     "chalk": "^5.3.0",
     "commander": "^14.0.0",
     "diff": "^7.0.0",
@@ -50,7 +50,7 @@
     "unified": "^11.0.5",
     "unist-util-visit": "^5.1.0",
     "yaml": "^2.8.0",
-    "zod": "^3.25.75"
+    "zod": "^4.3.6"
   },
   "devDependencies": {
     "@biomejs/biome": "^2.0.6",

src/ai/index.ts

@@ -16,3 +16,4 @@
 
 export * from "./model/index.js";
 export * from "./prompts/index.js";
+export * from "./tools/index.js";

src/ai/model/client.ts

@@ -22,8 +22,11 @@ import {
   CountTokensCommand,
   type ImageFormat,
   type Message,
+  type ToolConfiguration,
 } from "@aws-sdk/client-bedrock-runtime";
+import { z } from "zod";
 import type { Prompt } from "../prompts/index.js";
+import type { ToolDefinition } from "../tools/types.js";
 import {
   NoopRateLimiter,
   type RateLimiter,
@@ -140,76 +143,102 @@ export class DefaultBedrockClient implements BedrockClient {
    */
   async generate(
     prompt: Prompt,
+    tools: ToolDefinition[] = [],
     cacheEnabled: boolean = false,
   ): Promise<BedrockClientGenerateResponse> {
     const tokenUsageCounter = new TokenUsageCounter();
 
     const promptText = prompt.prompt;
     const promptContext = prompt.context;
-    let prefill = prompt.prefill ?? "";
 
     let iterations = 0;
 
-    while (true) {
-      // TODO: Emit event about iteration counter
+    const toolConfig: ToolConfiguration | undefined =
+      tools.length > 0
+        ? {
+            tools: tools.map((tool) => ({
+              toolSpec: {
+                name: tool.name,
+                description: tool.description,
+                inputSchema: {
+                  json: z.toJSONSchema(tool.parameters) as Record<
+                    string,
+                    unknown
+                  >,
+                },
+              },
+            })) as ToolConfiguration["tools"],
+          }
+        : undefined;
+
+    const toolMap = new Map(tools.map((t) => [t.name, t]));
+
+    const messages: Message[] = [
+      {
+        role: "user" as ConversationRole,
+        content: [
+          ...(promptContext ? [{ text: promptContext }] : []),
+          ...(cacheEnabled
+            ? [
+                {
+                  cachePoint: {
+                    type: "default" as CachePointType,
+                  },
+                },
+              ]
+            : []),
+          {
+            text: promptText,
+          },
+          ...(prompt.images || []).map((img) => ({
+            image: {
+              format: img.format as ImageFormat,
+              source: { bytes: img.bytes },
+            },
+          })),
+        ],
+      },
+    ];
 
+    while (true) {
       if (iterations > this.maxIterations) {
         throw new Error("Maximum iterations breached");
       }
 
-      const conversation = [
-        {
-          role: "user" as ConversationRole,
-          content: [
-            ...(promptContext ? [{ text: promptContext }] : []),
-
-            ...(cacheEnabled && this.isCachingSupported(this.modelId)
-              ? [
-                  {
-                    cachePoint: {
-                      type: "default" as CachePointType,
-                    },
-                  },
-                ]
-              : []),
-            {
-              text: promptText,
-            },
-            ...(prompt.images || []).map((img) => ({
-              image: {
-                format: img.format as ImageFormat,
-                source: { bytes: img.bytes },
-              },
-            })),
-          ],
-        },
-        ...(prefill
-          ? [
-              {
-                role: "assistant" as ConversationRole,
-                content: [{ text: prefill }],
-              },
-            ]
-          : []),
-      ];
-
       const estimatedTokens = await this.estimateTokens(
-        conversation,
+        messages,
         prompt.sampleOutput,
       );
-      /*const estimatedTokens = estimateTokens(
-        promptContext ?? "",
-        promptText,
-        prefill,
-        prompt.sampleOutput ?? "",
-      );*/
 
       tokenUsageCounter.addEstimated(estimatedTokens);
 
       const command = new ConverseCommand({
         modelId: this.modelId,
-        messages: conversation,
+        messages,
+        toolConfig,
+        additionalModelRequestFields: {
+          reasoning_config: {
+            type: "enabled",
+            budget_tokens: 2048,
+          },
+        },
         inferenceConfig: { maxTokens: this.maxTokens },
+        ...(prompt.outputSchema && {
+          outputConfig: {
+            textFormat: {
+              type: "json_schema",
+              structure: {
+                jsonSchema: {
+                  schema: JSON.stringify(
+                    z.toJSONSchema(prompt.outputSchema.schema),
+                  ),
+                  name: prompt.outputSchema.name,
+                  description: prompt.outputSchema.description,
+                },
+              },
+            },
+          },
+        }),
       });
 
       await Promise.all([
@@ -221,9 +250,6 @@ export class DefaultBedrockClient implements BedrockClient {
 
       const responseObject = await this.client.send(command);
 
-      // biome-ignore lint/style/noNonNullAssertion: Need to see if this needs better checks
-      let response = prefill + responseObject.output!.message!.content![0].text;
-
       this.tokenRateLimiter.consume(
         responseObject.usage?.totalTokens || 0,
         timestamp,
@@ -233,13 +259,52 @@ export class DefaultBedrockClient implements BedrockClient {
         tokenUsageCounter.addUsage(responseObject.usage);
       }
 
-      if (responseObject.stopReason === "max_tokens") {
-        prefill = response.trimEnd();
-        iterations++;
+      const responseContent = responseObject.output?.message?.content ?? [];
+
+      if (responseObject.stopReason === "tool_use") {
+        messages.push({
+          role: "assistant" as ConversationRole,
+          content: responseContent,
+        });
+
+        const toolResultContent = [];
+
+        for (const block of responseContent) {
+          if (!block.toolUse) continue;
 
+          const toolUseId = block.toolUse.toolUseId;
+          const toolName = block.toolUse.name ?? "";
+          const input = (block.toolUse.input ?? {}) as Record<string, unknown>;
+          const tool = toolMap.get(toolName);
+
+          const resultText = tool
+            ? tool.execute(input)
+            : `Unknown tool: ${toolName}`;
+
+          toolResultContent.push({
+            toolResult: {
+              toolUseId,
+              content: [{ text: resultText }],
+            },
+          });
+        }
+
+        messages.push({
+          role: "user" as ConversationRole,
+          content: toolResultContent,
+        });
+
+        iterations++;
         continue;
       }
 
+      const textBlock = responseContent.find((block) => block.text);
+      let response = textBlock?.text ?? "";
+
+      if (responseObject.stopReason === "max_tokens") {
+        throw new Error("Response exceeded maximum token limit");
+      }
+
       if (responseObject.stopReason !== "end_turn") {
         throw new Error(`Unexpected stop reason: ${responseObject.stopReason}`);
       }
@@ -286,43 +351,4 @@ export class DefaultBedrockClient implements BedrockClient {
 
     return responseObject.inputTokens ?? 0;
   }
-
-  /**
-   * Determines whether prompt caching is supported for the given model ID.
-   *
-   * Caching can improve performance and reduce costs by reusing previously
-   * processed context across multiple requests. This feature is only available
-   * for specific model versions.
-   *
-   * @param modelId - The AWS Bedrock model identifier to check
-   * @returns True if the model supports prompt caching, false otherwise
-   *
-   * @private
-   */
-  private isCachingSupported(modelId: string) {
-    const validModels = [
-      "anthropic.claude-opus-4-6",
-      "anthropic.claude-opus-4-5",
-      "anthropic.claude-opus-4-1",
-      "anthropic.claude-opus-4",
-      "anthropic.claude-sonnet-4-6",
-      "anthropic.claude-sonnet-4-5",
-      "anthropic.claude-haiku-4-5",
-      "anthropic.claude-opus-4",
-      "anthropic.claude-sonnet-4",
-      "anthropic.claude-3-7-sonnet",
-      "anthropic.claude-3-5-haiku",
-      "amazon.nova-micro-v1:0",
-      "amazon.nova-lite-v1:0",
-      "amazon.nova-pro-v1:0",
-    ];
-
-    for (const validModel of validModels) {
-      if (modelId.includes(validModel)) {
-        return true;
-      }
-    }
-
-    return false;
-  }
 }

src/ai/model/types.ts

@@ -15,6 +15,7 @@
  */
 
 import type { Prompt } from "../prompts/index.js";
+import type { ToolDefinition } from "../tools/types.js";
 
 export interface TokenUsage {
   inputTokens: number;
@@ -33,6 +34,7 @@ export interface BedrockClientGenerateResponse {
 export interface BedrockClient {
   generate(
     prompt: Prompt,
+    tools: ToolDefinition[],
     cacheEnabled: boolean,
   ): Promise<BedrockClientGenerateResponse>;
 }

src/ai/prompts/reviewPrompt.ts

@@ -22,11 +22,7 @@ import { loadImage } from "../../content/utils/markdownUtils.js";
 import type { Language } from "../../languages/index.js";
 import { buildContextPrompt } from "./contextPrompt.js";
 import type { Exemplar, Prompt } from "./types.js";
-import {
-  type ContextStrategy,
-  extractFileSection,
-  getContext,
-} from "./utils.js";
+import { type ContextStrategy, getContext } from "./utils.js";
 
 const template = `Your task is to review the content provided for file "{{file}}" and update it to improve it in terms of style, grammar and syntax.
 
@@ -56,9 +52,17 @@ For any finding which cannot be reliably remediated, such as missing images or b
 [Some screenshot](/pods1.png)
 </example_comment>
 
-Write the output as markdown in a similar style to the example content. Respond with the resulting file enclosed in <file></file> including the path to the file as an attribute.
+Write the translated content in a similar style to the example content. Use the write_file tool to output the result to "{{currentNode.filePath}}" in chunks:
+- Each chunk is a separate call to write_file
+- You MUST NOT write more than ~3000 tokens per chunk (roughly 2000-2500 words)
+- Break at natural boundaries: section headers, major paragraphs
+- First call: mode="create"
+- Subsequent calls: mode="append"
+- Continue until complete
 
-ONLY respond with the content between the "<file></file>" tags.`;
+Write substantial chunks to minimize tool calls while staying well under the output limit.
+
+You're final response to the user MUST simply be "Success".`;
 
 export async function buildReviewPrompt(
   tree: ContentTree,
@@ -95,16 +99,6 @@ export async function buildReviewPrompt(
       checkIssues: checkIssues && checkIssues.length > 0 ? checkIssues : null,
     }),
     sampleOutput: currentNode.content || undefined,
-    prefill: `<file path="${currentNode.filePath}">`,
-    transform: (input) => {
-      const fileSection = extractFileSection(input);
-
-      if (fileSection.path !== currentNode.filePath) {
-        throw new Error(`Unexpected file path in output: ${fileSection.path}`);
-      }
-
-      return fileSection.content;
-    },
   };
 
   if (includeImages && currentNode.content) {