simple-claude-code.py

edit-file-tool.ts

import fs from "node:fs/promises";
import z from "zod";
import createTool from "./_create.js";
import { execa } from "execa";

export default createTool({
  schema: {
    name: "edit_file",
    description: "Edit a file with various modes",
    input_schema: {
      type: "object",
      properties: {
        path: {
          type: "string",
          description: "Path to the file to edit",
        },
        mode: {
          type: "string",
          enum: ["replace", "insertLinesBefore", "insertLinesAfter", "deleteLines", "replaceLines"],
          description:
            "Editing mode: 'replace' (default), 'insertLinesBefore', 'insertLinesAfter', 'deleteLines', or 'replaceLines'",
          default: "replace",
        },
        content: {
          type: "string",
          description:
            "Content to write (complete file for 'replace' or new content for 'insertLinesBefore'/'insertLinesAfter'/'replaceLines')",
        },
        line: {
          type: "integer",
          description:
            "Line number for the operation (1-based): for 'insertLinesBefore' the new content will be placed before this line; for 'insertLinesAfter' the new content will be placed after this line; for 'deleteLines'/'replaceLines' this is the starting line",
        },
        count: {
          type: "integer",
          description: "Number of lines to delete or replace (for 'deleteLines' and 'replaceLines' modes)",
          default: 1,
        },
        prettier: {
          type: "boolean",
          description: "Whether to format the file with prettier (defaults to true)",
          default: true,
        },
      },
      required: ["path"],
    },
  },

  validator: z.object({
    path: z.string(),
    mode: z
      .enum(["replace", "insertLinesBefore", "insertLinesAfter", "deleteLines", "replaceLines"])
      .default("replace"),
    content: z.string().optional(),
    line: z.number().int().positive().optional(),
    count: z.number().int().positive().default(1),
    prettier: z.boolean().optional().default(true),
  }),

  async exec({ path: filePath, mode, content, line, count, prettier = true }) {
    // Check if the file exists
    let fileExists = true;
    try {
      await fs.access(filePath);
    } catch {
      fileExists = false;
      // For non-replace modes, the file needs to exist
      if (mode !== "replace") {
        return { content: `File ${filePath} does not exist`, isError: true };
      }
    }

    // Check required parameters based on mode
    if (
      (mode === "replace" || mode === "insertLinesBefore" || mode === "insertLinesAfter" || mode === "replaceLines") &&
      !content
    ) {
      return { content: `Content is required for ${mode} mode`, isError: true };
    }

    if (
      (mode === "insertLinesBefore" ||
        mode === "insertLinesAfter" ||
        mode === "deleteLines" ||
        mode === "replaceLines") &&
      !line
    ) {
      return { content: `Line number is required for ${mode} mode`, isError: true };
    }

    // Handle different edit modes
    switch (mode) {
      case "replace": {
        // Simply replace the entire file content
        await fs.writeFile(filePath, content!);
        break;
      }

      case "insertLinesBefore":
      case "insertLinesAfter":
      case "deleteLines":
      case "replaceLines": {
        // These modes require reading the file content first
        if (!fileExists) {
          return { content: `Cannot modify non-existent file ${filePath}`, isError: true };
        }

        const fileContent = await fs.readFile(filePath, "utf-8");
        const lines = fileContent.split("\n");

        // Validate line number
        if (line! > lines.length + 1) {
          return {
            content: `Line number ${line} is out of range (file has ${lines.length} lines)`,
            isError: true,
          };
        }

        // Line numbers in the file are 1-based, arrays are 0-based
        const lineIndex = line! - 1;

        if (mode === "insertLinesBefore") {
          // Insert new content before specified line
          const newLines = content!.split("\n");
          lines.splice(lineIndex, 0, ...newLines);
        } else if (mode === "insertLinesAfter") {
          // Insert new content after specified line
          const newLines = content!.split("\n");
          lines.splice(lineIndex + 1, 0, ...newLines);
        } else if (mode === "deleteLines") {
          // Validate delete range
          if (lineIndex + count > lines.length) {
            return {
              content: `Cannot delete ${count} lines starting at line ${line} (out of range)`,
              isError: true,
            };
          }

          // Delete specified number of lines
          lines.splice(lineIndex, count);
        } else if (mode === "replaceLines") {
          // Validate replace range
          if (lineIndex + count > lines.length) {
            return {
              content: `Cannot replace ${count} lines starting at line ${line} (out of range)`,
              isError: true,
            };
          }

          // Replace specified number of lines with new content
          const newLines = content!.split("\n");
          lines.splice(lineIndex, count, ...newLines);
        }

        // Write the modified content back to the file
        await fs.writeFile(filePath, lines.join("\n"));
        break;
      }

      default:
        return { content: `Unknown mode: ${mode}`, isError: true };
    }

    // Format with prettier if requested
    if (prettier) {
      const prettierResult = await execa("yarn", ["prettier", "--write", filePath], { reject: false, all: true });
      if (prettierResult.failed) {
        return { content: `File updated but prettier failed: ${prettierResult.all}`, isError: true };
      }
    }

    // Run linting
    const lintResult = await execa("yarn", ["lint"], { reject: false, all: true });
    if (lintResult.failed) {
      return { content: `File updated but lint failed: ${lintResult.all}`, isError: true };
    }

    return { content: `File updated successfully at ${filePath} using mode: ${mode}` };
  },
});

edit-file.tests.ts

import fs from "node:fs/promises";
import { beforeEach, describe, expect, it, vi } from "vitest";
import editFileTool from "../edit-file.js";
import { execa } from "execa";

// Mock dependencies
vi.mock("node:fs/promises");
vi.mock("execa");

describe("edit-file tool", () => {
  const mockFsAccess = vi.mocked(fs.access);
  const mockFsReadFile = vi.mocked(fs.readFile);
  const mockFsWriteFile = vi.mocked(fs.writeFile);
  const mockExeca = vi.mocked(execa);

  beforeEach(() => {
    vi.resetAllMocks();

    // Default mock for prettier and lint to succeed
    mockExeca.mockResolvedValue({
      failed: false,
      exitCode: 0,
      stdout: "",
      stderr: "",
      all: "",
      command: "",
      escapedCommand: "",
      killed: false,
      timedOut: false,
      isCanceled: false,
      signal: null,
    } as any);
  });

  describe("input validation", () => {
    it("should reject invalid input", async () => {
      const result = await editFileTool.exec({
        // missing required path property
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Invalid input");
    });

    it("should require content for replace mode", async () => {
      mockFsAccess.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replace",
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Content is required for replace mode");
    });

    it("should require line number for deleteLines mode", async () => {
      mockFsAccess.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "deleteLines",
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Line number is required for deleteLines mode");
    });
  });

  describe("file existence checks", () => {
    it("should error for non-existent file in non-replace modes", async () => {
      mockFsAccess.mockRejectedValue(new Error("File not found"));

      const result = await editFileTool.exec({
        path: "nonexistent.txt",
        mode: "deleteLines",
        line: 1,
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("does not exist");
    });

    it("should create file if it does not exist in replace mode", async () => {
      mockFsAccess.mockRejectedValue(new Error("File not found"));
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "new-file.txt",
        mode: "replace",
        content: "New content",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("new-file.txt", "New content");
      expect(result.content).toContain("File updated successfully");
    });
  });

  describe("replace mode", () => {
    it("should replace entire file content", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replace",
        content: "New file content",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "New file content");
      expect(result.content).toContain("File updated successfully");
    });
  });

  describe("deleteLines mode", () => {
    it("should delete a single line from file", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "deleteLines",
        line: 2,
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should delete multiple lines from file", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3\nLine 4");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "deleteLines",
        line: 2,
        count: 2,
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nLine 4");
      expect(result.content).toContain("File updated successfully");
    });

    it("should error when trying to delete lines out of range", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "deleteLines",
        line: 3,
        count: 2,
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Cannot delete 2 lines starting at line 3 (out of range)");
      expect(mockFsWriteFile).not.toHaveBeenCalled();
    });

    it("should error when line number exceeds file length", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "deleteLines",
        line: 5,
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Line number 5 is out of range");
      expect(mockFsWriteFile).not.toHaveBeenCalled();
    });
  });

  describe("replaceLines mode", () => {
    it("should replace a single line in file", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replaceLines",
        line: 2,
        content: "Replaced Line",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nReplaced Line\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should replace multiple lines with single line", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3\nLine 4");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replaceLines",
        line: 2,
        count: 2,
        content: "Replaced Content",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nReplaced Content\nLine 4");
      expect(result.content).toContain("File updated successfully");
    });

    it("should replace a single line with multiple lines", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replaceLines",
        line: 2,
        content: "New Line 1\nNew Line 2",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nNew Line 1\nNew Line 2\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should error when trying to replace lines out of range", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replaceLines",
        line: 3,
        count: 2,
        content: "Replacement Text",
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Cannot replace 2 lines starting at line 3 (out of range)");
      expect(mockFsWriteFile).not.toHaveBeenCalled();
    });

    it("should require content for replaceLines mode", async () => {
      mockFsAccess.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replaceLines",
        line: 1,
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Content is required for replaceLines mode");
    });
  });

  describe("insertLinesBefore mode", () => {
    it("should insert a single line before specified line", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesBefore",
        line: 2,
        content: "New Line",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nNew Line\nLine 2\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should insert multiple lines before specified line", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesBefore",
        line: 2,
        content: "New Line 1\nNew Line 2",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nNew Line 1\nNew Line 2\nLine 2\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should insert at the beginning of the file", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesBefore",
        line: 1,
        content: "New First Line",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "New First Line\nLine 1\nLine 2\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });
  });

  describe("insertLinesAfter mode", () => {
    it("should insert a single line after specified line", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesAfter",
        line: 2,
        content: "New Line",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nLine 2\nNew Line\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should insert multiple lines after specified line", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesAfter",
        line: 2,
        content: "New Line 1\nNew Line 2",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nLine 2\nNew Line 1\nNew Line 2\nLine 3");
      expect(result.content).toContain("File updated successfully");
    });

    it("should insert at the end of the file", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsReadFile.mockResolvedValue("Line 1\nLine 2\nLine 3");
      mockFsWriteFile.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesAfter",
        line: 3,
        content: "New Last Line",
        prettier: false, // Skip prettier to simplify test
      });

      expect(result.isError).toBe(false);
      expect(mockFsWriteFile).toHaveBeenCalledWith("test.txt", "Line 1\nLine 2\nLine 3\nNew Last Line");
      expect(result.content).toContain("File updated successfully");
    });

    it("should require content for insertLinesAfter mode", async () => {
      mockFsAccess.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesAfter",
        line: 1,
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Content is required for insertLinesAfter mode");
    });

    it("should require line number for insertLinesAfter mode", async () => {
      mockFsAccess.mockResolvedValue(undefined);

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "insertLinesAfter",
        content: "Test content",
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Line number is required for insertLinesAfter mode");
    });
  });

  describe("error handling", () => {
    it("should handle filesystem errors", async () => {
      mockFsAccess.mockResolvedValue(undefined);
      mockFsWriteFile.mockRejectedValue(new Error("Permission denied"));

      const result = await editFileTool.exec({
        path: "test.txt",
        mode: "replace",
        content: "New content",
      });

      expect(result.isError).toBe(true);
      expect(result.content).toContain("Error: Error: Permission denied");
    });
  });
});

version1.py

import anthropic

import sys
import pathlib

SYSTEM_PROMPT = f"""
You are Claude, an AI assistant focused on helping write high-quality code.

Here is your current source code:

```python
{pathlib.Path(sys.argv[0]).read_text()}
```
""".lstrip()

client = anthropic.Anthropic()

def get_user_input():
    first_line = input()
    
    # Check if this is a heredoc-style input
    if first_line.startswith('<<'):
        try:
            delimiter = first_line[2:].strip()
            if not delimiter:
                return first_line  # If it's just <<, treat as regular input
                
            lines = []
            while True:
                line = input()
                if line.strip() == delimiter:
                    break
                lines.append(line)
            
            return '\n'.join(lines)
        except EOFError:
            # If EOF occurs during multi-line input, return what we have
            return '\n'.join(lines) if lines else first_line
    else:
        # Regular single-line input
        return first_line

messages = []

while True:
    if len(messages) == 0:
        print("You: ", end="")
    else:
        print("\n\nYou: ", end="")

    try:
        user_input = get_user_input()
    except EOFError:
        break

    if user_input.lower() in ["exit", "quit", "bye"]:
        break
        
    messages.append({"role": "user", "content": user_input})

    with client.messages.stream(
        model="claude-3-7-sonnet-20250219",
        max_tokens=4000,
        system=[{"type": "text", "text": SYSTEM_PROMPT}],
        messages=messages
    ) as stream:
        for event in stream:
            if event.type == "text":
                sys.stdout.write(event.text)
                sys.stdout.flush()
            elif event.type == "content_block_stop":
                messages.append({"role": "assistant", "content": event.content_block})
            elif event.type == "message_stop":
                print(f"\n[Input tokens: {event.message.usage.input_tokens}, Output tokens: {event.message.usage.output_tokens}

version2.py

import anthropic

import sys
import pathlib
import argparse

SYSTEM_PROMPT = f"""
You are Claude, an AI assistant focused on helping write high-quality code.

Here is your current source code:

```python
{pathlib.Path(sys.argv[0]).read_text()}
```
""".lstrip()

client = anthropic.Anthropic()

def update_source_code(source_path, new_content):
    """Update the source code file with new content"""
    try:
        pathlib.Path(source_path).write_text(new_content)
        return {"success": True, "message": f"Source code updated successfully at {source_path}"}
    except Exception as e:
        return {"success": False, "message": f"Error updating source code: {str(e)}"}

def get_user_input():
    first_line = input()
    
    # Check if this is a heredoc-style input
    if first_line.startswith('<<'):
        try:
            delimiter = first_line[2:].strip()
            if not delimiter:
                return first_line  # If it's just <<, treat as regular input
                
            lines = []
            while True:
                line = input()
                if line.strip() == delimiter:
                    break
                lines.append(line)
            
            return '\n'.join(lines)
        except EOFError:
            # If EOF occurs during multi-line input, return what we have
            return '\n'.join(lines) if lines else first_line
    else:
        # Regular single-line input
        return first_line

def parse_cmdline():
    parser = argparse.ArgumentParser(
        description="An interactive CLI for Claude, focused on code assistance.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument(
        "--model", 
        default="claude-3-7-sonnet-20250219",
        help="The Claude model to use for responses"
    )
    
    parser.add_argument(
        "--max-tokens", 
        type=int, 
        default=8000,
        help="Maximum number of tokens in Claude's response"
    )
    
    parser.add_argument(
        "initial_prompt", 
        nargs="?", 
        default=None,
        help="Initial prompt to send to Claude (optional)"
    )
    
    return parser.parse_args()

args = parse_cmdline()
messages = []

source_code_tool = {
    "name": "update_source_code",
    "description": "Update this script's source code with new content",
    "input_schema": {
        "type": "object",
        "properties": {
            "new_content": {
                "type": "string",
                "description": "The new source code content to write to the file"
            }
        },
        "required": ["new_content"]
    }
}

while True:
    if len(messages) == 0 and len(sys.argv) > 1:
        messages.append({"role": "user", "content": args.initial_prompt})
    else:
        if len(messages) == 0:
            print("You: ", end="")
        else:
            print("\n\nYou: ", end="")

        try:
            user_input = get_user_input()
        except EOFError:
            break
        
        if user_input.lower() in ["exit", "quit", "bye"]:
            break
            
        messages.append({"role": "user", "content": user_input})

    with client.messages.stream(
        model=args.model,
        max_tokens=args.max_tokens,
        system=[{"type": "text", "text": SYSTEM_PROMPT}],
        messages=messages,
        tools=[source_code_tool]
    ) as stream:
        for event in stream:
            if event.type == "text":
                sys.stdout.write(event.text)
                sys.stdout.flush()
            elif event.type == "input_json":
                sys.stdout.write(event.partial_json)
                sys.stdout.flush()
            elif event.type == "content_block_start" and event.content_block.type == "tool_use":
                print(f"\n[Using tool: {event.content_block.name}]\n")
            elif event.type == "content_block_stop":
                messages.append({"role": "assistant", "content": event.content_block})
                if event.content_block.type == "tool_use":
                    if event.content_block.name == "update_source_code":
                        result = update_source_code(sys.argv[0], event.content_block.input["new_content"])
                    print(f"\n[Tool result: {result['message']}]\n")
            elif event.type == "message_stop":
            	print(f"\n[Input tokens: {event.message.usage.input_tokens}, Output tokens: {event.message.usage.output_tokens}]\n")