@posthog/agent 1.24.2 → 1.25.0

package/src/adapters/claude/tools.ts

@@ -0,0 +1,609 @@
+import type {
+  PlanEntry,
+  ToolCallContent,
+  ToolCallLocation,
+  ToolKind,
+} from "@agentclientprotocol/sdk";
+import type { HookCallback, HookInput } from "@anthropic-ai/claude-agent-sdk";
+import type {
+  ToolResultBlockParam,
+  WebSearchToolResultBlockParam,
+} from "@anthropic-ai/sdk/resources";
+import type {
+  BetaBashCodeExecutionToolResultBlockParam,
+  BetaCodeExecutionToolResultBlockParam,
+  BetaRequestMCPToolResultBlockParam,
+  BetaTextEditorCodeExecutionToolResultBlockParam,
+  BetaToolSearchToolResultBlockParam,
+  BetaWebFetchToolResultBlockParam,
+  BetaWebSearchToolResultBlockParam,
+} from "@anthropic-ai/sdk/resources/beta.mjs";
+import { Logger } from "@/utils/logger.js";
+import {
+  replaceAndCalculateLocation,
+  SYSTEM_REMINDER,
+  toolNames,
+} from "./mcp-server.js";
+
+interface ToolInfo {
+  title: string;
+  kind: ToolKind;
+  content: ToolCallContent[];
+  locations?: ToolCallLocation[];
+}
+
+interface ToolUpdate {
+  title?: string;
+  content?: ToolCallContent[];
+  locations?: ToolCallLocation[];
+}
+
+export function toolInfoFromToolUse(
+  toolUse: any,
+  cachedFileContent: { [key: string]: string },
+  logger: Logger = new Logger({ debug: false, prefix: "[ClaudeTools]" }),
+): ToolInfo {
+  const name = toolUse.name;
+  const input = toolUse.input;
+
+  switch (name) {
+    case "Task":
+      return {
+        title: input?.description ? input.description : "Task",
+        kind: "think",
+        content: input?.prompt
+          ? [
+              {
+                type: "content",
+                content: { type: "text", text: input.prompt },
+              },
+            ]
+          : [],
+      };
+
+    case "NotebookRead":
+      return {
+        title: input?.notebook_path
+          ? `Read Notebook ${input.notebook_path}`
+          : "Read Notebook",
+        kind: "read",
+        content: [],
+        locations: input?.notebook_path ? [{ path: input.notebook_path }] : [],
+      };
+
+    case "NotebookEdit":
+      return {
+        title: input?.notebook_path
+          ? `Edit Notebook ${input.notebook_path}`
+          : "Edit Notebook",
+        kind: "edit",
+        content: input?.new_source
+          ? [
+              {
+                type: "content",
+                content: { type: "text", text: input.new_source },
+              },
+            ]
+          : [],
+        locations: input?.notebook_path ? [{ path: input.notebook_path }] : [],
+      };
+
+    case "Bash":
+    case toolNames.bash:
+      return {
+        title: input?.command
+          ? `\`${input.command.replaceAll("`", "\\`")}\``
+          : "Terminal",
+        kind: "execute",
+        content: input?.description
+          ? [
+              {
+                type: "content",
+                content: { type: "text", text: input.description },
+              },
+            ]
+          : [],
+      };
+
+    case "BashOutput":
+    case toolNames.bashOutput:
+      return {
+        title: "Tail Logs",
+        kind: "execute",
+        content: [],
+      };
+
+    case "KillShell":
+    case toolNames.killShell:
+      return {
+        title: "Kill Process",
+        kind: "execute",
+        content: [],
+      };
+
+    case toolNames.read: {
+      let limit = "";
+      if (input.limit) {
+        limit =
+          " (" +
+          ((input.offset ?? 0) + 1) +
+          " - " +
+          ((input.offset ?? 0) + input.limit) +
+          ")";
+      } else if (input.offset) {
+        limit = ` (from line ${input.offset + 1})`;
+      }
+      return {
+        title: `Read ${input.file_path ?? "File"}${limit}`,
+        kind: "read",
+        locations: input.file_path
+          ? [
+              {
+                path: input.file_path,
+                line: input.offset ?? 0,
+              },
+            ]
+          : [],
+        content: [],
+      };
+    }
+
+    case "Read":
+      return {
+        title: "Read File",
+        kind: "read",
+        content: [],
+        locations: [{ path: input.file_path, line: input.offset ?? 0 }],
+      };
+
+    case "LS":
+      return {
+        title: `List the ${input?.path ? `\`${input.path}\`` : "current"} directory's contents`,
+        kind: "search",
+        content: [],
+        locations: [],
+      };
+
+    case toolNames.edit:
+    case "Edit": {
+      const path = input?.file_path ?? input?.file_path;
+      let oldText = input.old_string ?? null;
+      let newText = input.new_string ?? "";
+      let affectedLines: number[] = [];
+
+      if (path && oldText) {
+        try {
+          const oldContent = cachedFileContent[path] || "";
+          const newContent = replaceAndCalculateLocation(oldContent, [
+            {
+              oldText,
+              newText,
+              replaceAll: false,
+            },
+          ]);
+          oldText = oldContent;
+          newText = newContent.newContent;
+          affectedLines = newContent.lineNumbers;
+        } catch (e) {
+          logger.error("Failed to edit file", e);
+        }
+      }
+      return {
+        title: path ? `Edit \`${path}\`` : "Edit",
+        kind: "edit",
+        content:
+          input && path
+            ? [
+                {
+                  type: "diff",
+                  path,
+                  oldText,
+                  newText,
+                },
+              ]
+            : [],
+        locations: path
+          ? affectedLines.length > 0
+            ? affectedLines.map((line) => ({ line, path }))
+            : [{ path }]
+          : [],
+      };
+    }
+
+    case toolNames.write: {
+      let content: ToolCallContent[] = [];
+      if (input?.file_path) {
+        content = [
+          {
+            type: "diff",
+            path: input.file_path,
+            oldText: null,
+            newText: input.content,
+          },
+        ];
+      } else if (input?.content) {
+        content = [
+          {
+            type: "content",
+            content: { type: "text", text: input.content },
+          },
+        ];
+      }
+      return {
+        title: input?.file_path ? `Write ${input.file_path}` : "Write",
+        kind: "edit",
+        content,
+        locations: input?.file_path ? [{ path: input.file_path }] : [],
+      };
+    }
+
+    case "Write":
+      return {
+        title: input?.file_path ? `Write ${input.file_path}` : "Write",
+        kind: "edit",
+        content: input?.file_path
+          ? [
+              {
+                type: "diff",
+                path: input.file_path,
+                oldText: null,
+                newText: input.content,
+              },
+            ]
+          : [],
+        locations: input?.file_path ? [{ path: input.file_path }] : [],
+      };
+
+    case "Glob": {
+      let label = "Find";
+      if (input.path) {
+        label += ` \`${input.path}\``;
+      }
+      if (input.pattern) {
+        label += ` \`${input.pattern}\``;
+      }
+      return {
+        title: label,
+        kind: "search",
+        content: [],
+        locations: input.path ? [{ path: input.path }] : [],
+      };
+    }
+
+    case "Grep": {
+      let label = "grep";
+
+      if (input["-i"]) {
+        label += " -i";
+      }
+      if (input["-n"]) {
+        label += " -n";
+      }
+
+      if (input["-A"] !== undefined) {
+        label += ` -A ${input["-A"]}`;
+      }
+      if (input["-B"] !== undefined) {
+        label += ` -B ${input["-B"]}`;
+      }
+      if (input["-C"] !== undefined) {
+        label += ` -C ${input["-C"]}`;
+      }
+
+      if (input.output_mode) {
+        switch (input.output_mode) {
+          case "FilesWithMatches":
+            label += " -l";
+            break;
+          case "Count":
+            label += " -c";
+            break;
+          default:
+            break;
+        }
+      }
+
+      if (input.head_limit !== undefined) {
+        label += ` | head -${input.head_limit}`;
+      }
+
+      if (input.glob) {
+        label += ` --include="${input.glob}"`;
+      }
+
+      if (input.type) {
+        label += ` --type=${input.type}`;
+      }
+
+      if (input.multiline) {
+        label += " -P";
+      }
+
+      label += ` "${input.pattern}"`;
+
+      if (input.path) {
+        label += ` ${input.path}`;
+      }
+
+      return {
+        title: label,
+        kind: "search",
+        content: [],
+      };
+    }
+
+    case "WebFetch":
+      return {
+        title: input?.url ? `Fetch ${input.url}` : "Fetch",
+        kind: "fetch",
+        content: input?.prompt
+          ? [
+              {
+                type: "content",
+                content: { type: "text", text: input.prompt },
+              },
+            ]
+          : [],
+      };
+
+    case "WebSearch": {
+      let label = `"${input.query}"`;
+
+      if (input.allowed_domains && input.allowed_domains.length > 0) {
+        label += ` (allowed: ${input.allowed_domains.join(", ")})`;
+      }
+
+      if (input.blocked_domains && input.blocked_domains.length > 0) {
+        label += ` (blocked: ${input.blocked_domains.join(", ")})`;
+      }
+
+      return {
+        title: label,
+        kind: "fetch",
+        content: [],
+      };
+    }
+
+    case "TodoWrite":
+      return {
+        title: Array.isArray(input?.todos)
+          ? `Update TODOs: ${input.todos.map((todo: any) => todo.content).join(", ")}`
+          : "Update TODOs",
+        kind: "think",
+        content: [],
+      };
+
+    case "ExitPlanMode":
+      return {
+        title: "Ready to code?",
+        kind: "switch_mode",
+        content: input?.plan
+          ? [{ type: "content", content: { type: "text", text: input.plan } }]
+          : [],
+      };
+
+    case "Other": {
+      let output: string;
+      try {
+        output = JSON.stringify(input, null, 2);
+      } catch {
+        output = typeof input === "string" ? input : "{}";
+      }
+      return {
+        title: name || "Unknown Tool",
+        kind: "other",
+        content: [
+          {
+            type: "content",
+            content: {
+              type: "text",
+              text: `\`\`\`json\n${output}\`\`\``,
+            },
+          },
+        ],
+      };
+    }
+
+    default:
+      return {
+        title: name || "Unknown Tool",
+        kind: "other",
+        content: [],
+      };
+  }
+}
+
+export function toolUpdateFromToolResult(
+  toolResult:
+    | ToolResultBlockParam
+    | BetaWebSearchToolResultBlockParam
+    | BetaWebFetchToolResultBlockParam
+    | WebSearchToolResultBlockParam
+    | BetaCodeExecutionToolResultBlockParam
+    | BetaBashCodeExecutionToolResultBlockParam
+    | BetaTextEditorCodeExecutionToolResultBlockParam
+    | BetaRequestMCPToolResultBlockParam
+    | BetaToolSearchToolResultBlockParam,
+  toolUse: any | undefined,
+): ToolUpdate {
+  switch (toolUse?.name) {
+    case "Read":
+    case toolNames.read:
+      if (Array.isArray(toolResult.content) && toolResult.content.length > 0) {
+        return {
+          content: toolResult.content.map((content: any) => ({
+            type: "content",
+            content:
+              content.type === "text"
+                ? {
+                    type: "text",
+                    text: markdownEscape(
+                      content.text.replace(SYSTEM_REMINDER, ""),
+                    ),
+                  }
+                : content,
+          })),
+        };
+      } else if (
+        typeof toolResult.content === "string" &&
+        toolResult.content.length > 0
+      ) {
+        return {
+          content: [
+            {
+              type: "content",
+              content: {
+                type: "text",
+                text: markdownEscape(
+                  toolResult.content.replace(SYSTEM_REMINDER, ""),
+                ),
+              },
+            },
+          ],
+        };
+      }
+      return {};
+
+    case toolNames.bash:
+    case "edit":
+    case "Edit":
+    case toolNames.edit:
+    case toolNames.write:
+    case "Write": {
+      if (
+        "is_error" in toolResult &&
+        toolResult.is_error &&
+        toolResult.content &&
+        toolResult.content.length > 0
+      ) {
+        // Only return errors
+        return toAcpContentUpdate(toolResult.content, true);
+      }
+      return {};
+    }
+
+    case "ExitPlanMode": {
+      return { title: "Exited Plan Mode" };
+    }
+    default: {
+      return toAcpContentUpdate(
+        toolResult.content,
+        "is_error" in toolResult ? toolResult.is_error : false,
+      );
+    }
+  }
+}
+
+function toAcpContentUpdate(
+  content: any,
+  isError: boolean = false,
+): { content?: ToolCallContent[] } {
+  if (Array.isArray(content) && content.length > 0) {
+    return {
+      content: content.map((content: any) => ({
+        type: "content",
+        content:
+          isError && content.type === "text"
+            ? {
+                ...content,
+                text: `\`\`\`\n${content.text}\n\`\`\``,
+              }
+            : content,
+      })),
+    };
+  } else if (typeof content === "string" && content.length > 0) {
+    return {
+      content: [
+        {
+          type: "content",
+          content: {
+            type: "text",
+            text: isError ? `\`\`\`\n${content}\n\`\`\`` : content,
+          },
+        },
+      ],
+    };
+  }
+  return {};
+}
+
+export type ClaudePlanEntry = {
+  content: string;
+  status: "pending" | "in_progress" | "completed";
+  activeForm: string;
+};
+
+export function planEntries(input: { todos: ClaudePlanEntry[] }): PlanEntry[] {
+  return input.todos.map((input) => ({
+    content: input.content,
+    status: input.status,
+    priority: "medium",
+  }));
+}
+
+export function markdownEscape(text: string): string {
+  let escapedText = "```";
+  for (const [m] of text.matchAll(/^```+/gm)) {
+    while (m.length >= escapedText.length) {
+      escapedText += "`";
+    }
+  }
+  return `${escapedText}\n${text}${text.endsWith("\n") ? "" : "\n"}${escapedText}`;
+}
+
+/* A global variable to store callbacks that should be executed when receiving hooks from Claude Code */
+const toolUseCallbacks: {
+  [toolUseId: string]: {
+    onPostToolUseHook?: (
+      toolUseID: string,
+      toolInput: unknown,
+      toolResponse: unknown,
+    ) => Promise<void>;
+  };
+} = {};
+
+/* Setup callbacks that will be called when receiving hooks from Claude Code */
+export const registerHookCallback = (
+  toolUseID: string,
+  {
+    onPostToolUseHook,
+  }: {
+    onPostToolUseHook?: (
+      toolUseID: string,
+      toolInput: unknown,
+      toolResponse: unknown,
+    ) => Promise<void>;
+  },
+) => {
+  toolUseCallbacks[toolUseID] = {
+    onPostToolUseHook,
+  };
+};
+
+/* A callback for Claude Code that is called when receiving a PostToolUse hook */
+export const createPostToolUseHook =
+  (
+    logger: Logger = new Logger({ prefix: "[createPostToolUseHook]" }),
+  ): HookCallback =>
+  async (
+    input: HookInput,
+    toolUseID: string | undefined,
+  ): Promise<{ continue: boolean }> => {
+    if (input.hook_event_name === "PostToolUse" && toolUseID) {
+      const onPostToolUseHook = toolUseCallbacks[toolUseID]?.onPostToolUseHook;
+      if (onPostToolUseHook) {
+        await onPostToolUseHook(
+          toolUseID,
+          input.tool_input,
+          input.tool_response,
+        );
+        delete toolUseCallbacks[toolUseID]; // Cleanup after execution
+      } else {
+        logger.error(
+          `No onPostToolUseHook found for tool use ID: ${toolUseID}`,
+        );
+        delete toolUseCallbacks[toolUseID];
+      }
+    }
+    return { continue: true };
+  };