DEV Community

Cover image for Rewriting Claude Code in Rust
Josh Mo
Josh Mo

Posted on

Rewriting Claude Code in Rust

It's late. You're getting the builder itch - the kind of itch that never truly leaves you until you've actually built something. Whether you're actually going to use it in a week or so is another matter entirely, but that's besides the point. Your computer is turned on, you have Claude Code open in the terminal tab for your editor. You've been using it for a little while now.

And this gets you thinking. Surely it can't be that hard to build a coding agent, right? It's basically just an LLM with some tools and access to your filesystem?

In this post, we'll build a small demo that will be able to act as a minimal demo version of Claude Code written in Rust using the Rig AI framework. It will be able to use bash commands, read and write files, as well as stream output back to you.

Interested in just trying out the demo for yourself? Here it is.

Getting started

Before we get started, you'll need the Rust programming language installed. You will also need an API key for the model provider of your choice, although for the purposes of simplicity we'll be using an Anthropic API key.

To get started, we'll create our project and call it rig-cc:

cargo init rig-cc
cd rig-cc
Enter fullscreen mode Exit fullscreen mode

Next, we'll add our dependencies:

cargo add rig-core@0.30 tokio@1 anyhow \
serde serde_json thiserror futures -F tokio/full,serde/derive
Enter fullscreen mode Exit fullscreen mode

What did we add?

  • rig-core: The core Rig library.
  • serde/serde_json: A library for converting to and from Rust structs and various data formats, with serde_json for JSON interop.
  • tokio: the most popular async Rust runtime.
  • thiserror: Derive error handling
  • futures: A crate for dealing with streams and other async-related functionality. We need to use the Stream/StreamExt traits to be able to utilise the item stream returned from the agent.

Building our coding agent

Before we start, let's break down the components of a coding agent. It should be able to:

  • Hold a conversation and use context from previous messages
  • Read from and write to files
  • Use Bash commands, and additionally allow the user to manually interrupt the bash command if it takes too long
  • Allow users to exit the conversation

First, we'll build the ReadFile tool, for reading files. Tools in Rig are simply defined as types that implement the rig::tool::Tool type.

#[derive(Deserialize, Serialize)]
pub struct ReadFile;

#[derive(serde::Deserialize)]
pub struct ReadFileArgs {
    path: String,
}
Enter fullscreen mode Exit fullscreen mode

You can see below that we get the following:

  • The tool definition (through fn definition())
  • The relevant types (arguments, outputs, the tool name and error type)
  • What should actually happen when the tool is called (through fn call())
impl rig::tool::Tool for ReadFile {
    const NAME: &'static str = "read_file";
    type Error = rig::tool::ToolError;
    type Args = ReadFileArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "read_file".to_string(),
            description: "Read the contents of a file at the specified path. Returns the file contents as a string.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The path to the file to read"
                    }
                },
                "required": ["path"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        std::fs::read_to_string(&args.path).map_err(|e| {
            ToolError(format!(
                "Failed to read file '{path}': {e}",
                path = args.path
            ))
        })
    }
}
Enter fullscreen mode Exit fullscreen mode

The WriteFile tool is largely the same, as it's primarily just writing to a given file (and a check to see if the directory exists):

#[derive(Deserialize)]
pub struct WriteFileArgs {
    path: String,
    content: String,
}

#[derive(Deserialize, Serialize)]
pub struct WriteFile;

impl Tool for WriteFile {
    const NAME: &'static str = "write_file";
    type Error = ToolError;
    type Args = WriteFileArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "write_file".to_string(),
            description: "Write content to a file at the specified path. Creates parent directories if they don't exist. Overwrites the file if it already exists.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The path to the file to write"
                    },
                    "content": {
                        "type": "string",
                        "description": "The content to write to the file"
                    }
                },
                "required": ["path", "content"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        if let Some(parent) = std::path::Path::new(&args.path).parent()
            && !parent.as_os_str().is_empty()
        {
            std::fs::create_dir_all(parent)
                .map_err(|e| ToolError(format!("Failed to create directories: {}", e)))?;
        }

        std::fs::write(&args.path, &args.content)
            .map_err(|e| ToolError(format!("Failed to write file '{}': {}", args.path, e)))?;

        Ok(format!(
            "Successfully wrote {} bytes to '{}'",
            args.content.len(),
            args.path
        ))
    }
}
Enter fullscreen mode Exit fullscreen mode

Finally, we're going to write our bash tool. This has some somewhat complicated logic in it, so we'll split this into two parts: the function to call the bash command, and then implementing the tool.

First, we need to call the Bash command using std::process::Command - and pipe both stderr and stdout. This creates a child process that we can programmatically control (and more specifically, waiting for the process to finish so we can collect the results).

use tokio::process::Command;

let mut child = Command::new("bash")
    .arg("-c")
    .arg(&args.command)
    .stdout(std::process::Stdio::piped())
    .stderr(std::process::Stdio::piped())
    .spawn()
    .map_err(|e| ToolError(format!("Failed to spawn command: {}", e)))?;
Enter fullscreen mode Exit fullscreen mode

Next, we need to use the tokio timeout function which will take a std::time::Duration and a future. If the timeout is reached, it will return an error; if not, the result will be given as usual.

/// Timeout. If the execution time for a command exceeds the timeout, users should be able to manually interrupt it.
const WARNING_TIMEOUT_SECS: u64 = 60;

let warning_duration = Duration::from_secs(WARNING_TIMEOUT_SECS);

let status = match tokio::time::timeout(warning_duration, child.wait()).await {
    // If OK, automatically try to get the result
    Ok(result) => result.map_err(|e| ToolError(format!("Command failed: {}", e)))?,
    // An error here represents that the timeout was reached
    Err(_) => {
        eprintln!(
            "\n[Command running for >{WARNING_TIMEOUT_SECS}s. Press Ctrl+C to interrupt]");
        io::stderr().flush().ok();
        child.wait()
             .await
             .map_err(|e| ToolError(format!("Command failed: {}", e)))?
        }
};
Enter fullscreen mode Exit fullscreen mode

After that, we need to take the stdout/stderr pipes and get their entire contents. Because the take() function results in an Option<T>, we'll use the if let-some pattern which allows for simpler matching without needing to deal with a None value (if the value is None we can probably assume there's nothing there anyway).

let stdout = child.stdout.take();
let stderr = child.stderr.take();

let mut stdout_content = String::new();
let mut stderr_content = String::new();

if let Some(mut stdout) = stdout {
    use tokio::io::AsyncReadExt;
    let mut buf = Vec::new();
    stdout.read_to_end(&mut buf).await.ok();
    stdout_content = String::from_utf8_lossy(&buf).to_string();
}

if let Some(mut stderr) = stderr {
    use tokio::io::AsyncReadExt;
    let mut buf = Vec::new();
    stderr.read_to_end(&mut buf).await.ok();
    stderr_content = String::from_utf8_lossy(&buf).to_string();
}
Enter fullscreen mode Exit fullscreen mode

Finally, we'll create our tool output. If the command status is a success, we'll re-assign our contents of stdout as mutable and then push whatever's in stderr to the message. This will let the model know if there are any issues while running, since even on a successful run there may be errors causing issues that may hamper things like performance, latency and other things. If it fails, then we provide the exit code and the stdout/stderr contents following the previous method (appending a newline, then the contents of each).

Regardless of the status of the command, the final text will get truncated according to the maximum output (to avoid flooding the LLM with context and therefore increasing the context window usage by accident) - and then return the output.

/// Max output bytes. Prevents infinite context being stuffed into the response.
const MAX_OUTPUT_BYTES: usize = 50 * 1024;

let mut output = if status.success() {
    let mut out = stdout_content;
    if !stderr_content.is_empty() {
        if !out.is_empty() {
            out.push('\n');
        }
        out.push_str("stderr:\n");
        out.push_str(&stderr_content);
        }
    out
    } else {
        let mut out = format!("Exit code: {}\n", status.code().unwrap_or(-1));
        if !stdout_content.is_empty() {
            out.push_str("stdout:\n");
            out.push_str(&stdout_content);
            out.push('\n');
        }
        if !stderr_content.is_empty() {
            out.push_str("stderr:\n");
            out.push_str(&stderr_content);
        }
        out
};

let total_bytes = output.len();
if total_bytes > MAX_OUTPUT_BYTES {
    output.truncate(MAX_OUTPUT_BYTES);
    while !output.is_char_boundary(output.len()) {
        output.pop();
    }
    output.push_str(&format!(
        "\n... [output truncated, {total_bytes} bytes total]",
            ));
    }

Ok(output)
Enter fullscreen mode Exit fullscreen mode

OK, so now if we set up our Bash tool and arguments structs it should look like this (combine all the previous snippets into a single function like in the next snippet!):

#[derive(Deserialize)]
pub struct BashArgs {
    command: String,
}

#[derive(Deserialize, Serialize)]
pub struct Bash;

impl Bash {
    async fn run(&self, args: BashArgs) -> Result<String, ToolError> {
        use tokio::process::Command;

        let mut child = Command::new("bash")
            .arg("-c")
            .arg(&args.command)
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn()
            .map_err(|e| ToolError(format!("Failed to spawn command: {}", e)))?;

        let warning_duration = Duration::from_secs(WARNING_TIMEOUT_SECS);
        let status = match timeout(warning_duration, child.wait()).await {
            Ok(result) => result.map_err(|e| ToolError(format!("Command failed: {}", e)))?,
            Err(_) => {
                eprintln!(
                    "\n[Command running for >{}s. Press Ctrl+C to interrupt]",
                    WARNING_TIMEOUT_SECS
                );
                io::stderr().flush().ok();
                child
                    .wait()
                    .await
                    .map_err(|e| ToolError(format!("Command failed: {}", e)))?
            }
        };

        let stdout = child.stdout.take();
        let stderr = child.stderr.take();

        let mut stdout_content = String::new();
        let mut stderr_content = String::new();

        if let Some(mut stdout) = stdout {
            use tokio::io::AsyncReadExt;
            let mut buf = Vec::new();
            stdout.read_to_end(&mut buf).await.ok();
            stdout_content = String::from_utf8_lossy(&buf).to_string();
        }

        if let Some(mut stderr) = stderr {
            use tokio::io::AsyncReadExt;
            let mut buf = Vec::new();
            stderr.read_to_end(&mut buf).await.ok();
            stderr_content = String::from_utf8_lossy(&buf).to_string();
        }

        let mut output = if status.success() {
            let mut out = stdout_content;
            if !stderr_content.is_empty() {
                if !out.is_empty() {
                    out.push('\n');
                }
                out.push_str("stderr:\n");
                out.push_str(&stderr_content);
            }
            out
        } else {
            let mut out = format!("Exit code: {}\n", status.code().unwrap_or(-1));
            if !stdout_content.is_empty() {
                out.push_str("stdout:\n");
                out.push_str(&stdout_content);
                out.push('\n');
            }
            if !stderr_content.is_empty() {
                out.push_str("stderr:\n");
                out.push_str(&stderr_content);
            }
            out
        };

        let total_bytes = output.len();
        if total_bytes > MAX_OUTPUT_BYTES {
            output.truncate(MAX_OUTPUT_BYTES);
            while !output.is_char_boundary(output.len()) {
                output.pop();
            }
            output.push_str(&format!(
                "\n... [output truncated, {} bytes total]",
                total_bytes
            ));
        }

        Ok(output)
    }
}
Enter fullscreen mode Exit fullscreen mode

To finish this section off, we'll write the actual rig::tool::Tool implementation. Because we've already written the required function, instead of writing huge amounts of code in the async fn call() section we can write a nice little one liner.

impl Tool for Bash {
    const NAME: &'static str = "bash";
    type Error = ToolError;
    type Args = BashArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "bash".to_string(),
            description: "Execute a bash command and return its output. Use this for running shell commands, git operations, running tests, installing packages, etc. The command runs in the current working directory.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute"
                    }
                },
                "required": ["command"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        self.run(args).await
    }
}
Enter fullscreen mode Exit fullscreen mode

Writing the coding agent loop

With all the tools written, now we simply need to write our agent loop. First though we'll need to write our system prompt - for now we'll call our program Rig Code, and tell it what tools it has access to as well as some basic guidelines.

const SYSTEM_PROMPT: &str = r#"You are Rig Code, an interactive AI coding assistant running in the terminal.

You have access to these tools:
- bash: Execute shell commands (runs in current working directory)
- read_file: Read file contents
- write_file: Create or modify files

Guidelines:
- Use bash to explore projects, run tests, git operations, etc.
- Read files before modifying them to understand context
- Be concise and focused on solving the user's problem
- When making changes, explain what you're doing briefly
"#;
Enter fullscreen mode Exit fullscreen mode

Finally, we need to write our main function. The agent part is quite simple - we need to create an Anthropic client, then create an agent using the agent builder in Rig. Then we'll instantiate a message history, as well as an input/output buffer.

#[tokio::main]
async fn main() -> Result<()> {
    let client = Client::from_env();

    let agent = client
        .agent("claude-sonnet-4-5")
        .preamble(SYSTEM_PROMPT)
        .tool(ReadFile)
        .tool(WriteFile)
        .tool(Bash)
        .max_tokens(8192)
        .build();

    println!("Rig Code v0.1.0");
    println!("Type 'exit' or 'quit' to exit.\n");

    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut history: Vec<Message> = Vec::new();

    loop {
        // .. loop goes here
    }

    Ok(())
}
Enter fullscreen mode Exit fullscreen mode

Now for the loop! This should print a marker that makes it obvious it's the user's turn to type, then try to read the input whenever the user finishes inputting (ie, they press Return).

Once done, we then check for the following in the trimmed input:

  • If the user writes "/exit" or "/quit", break the loop and exit
  • If nothing's been inputted, just continue
  • If input exists, attempt to send a streaming prompt request to the model provider (we'll deal with this specifically later).

After the stream has finished, we should also show how many tokens were used.

print!("> ");
stdout.flush()?;

let mut input = String::new();

match stdin.read_line(&mut input) {
    Ok(0) => {
        println!("\nGoodbye!");
        break;
    }

    Ok(_) => {
        let input = input.trim();

        if input.eq_ignore_ascii_case("exit") || input.eq_ignore_ascii_case("quit") {
            println!("Goodbye!");
            break;
        }

        if input.is_empty() {
            continue;
        }

        println!();

        let mut stream = agent
            .stream_prompt(input)
            .with_history(history.clone())
            .multi_turn(100)
            .await;

        let mut response_text = String::new();
        let mut input_tokens = 0u64;
        let mut output_tokens = 0u64;

        while let Some(chunk) = stream.next().await {
            // .. handle stream (see below)
        }

        println!("\n");
        println!(
            "[Tokens: {} in / {} out]",
            // .. see below for format_number function
            format_number(input_tokens),
            format_number(output_tokens)
        );
        println!();

        history.push(Message::user(input));

        if !response_text.is_empty() {
            history.push(Message::assistant(response_text));
        }
    }

    Err(e) => {
        eprintln!("Error reading input: {}", e);
    }
}
Enter fullscreen mode Exit fullscreen mode

One small detail in the implementation is that if you're burning through a large number of tokens very quickly, it can be difficult to immediately tell at a glance what the number is. We'll implement a function below that makes it look slightly nicer and easier to look at.

fn format_number(n: u64) -> String {
    let s = n.to_string();
    let mut result = String::new();
    for (i, c) in s.chars().rev().enumerate() {
        if i > 0 && i % 3 == 0 {
            result.push(',');
        }
        result.push(c);
    }
    result.chars().rev().collect()
}
Enter fullscreen mode Exit fullscreen mode

The stream itself isn't difficult to deal with per se as you're just matching on item variants. You can see that we get a number of different variants - tool calls, text, tool results, and the final stream item (ie the aggregated text and token usage).

while let Some(chunk) = stream.next().await {
    match chunk {
        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::Text(text),
        )) => {
            print!("{}", text.text);
            stdout.flush()?;
            response_text.push_str(&text.text);
        }

        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::ToolCall { tool_call, .. },
        )) => {
            println!("\n[Calling tool: {}]", tool_call.function.name);
            stdout.flush()?;
        }

        Ok(MultiTurnStreamItem::StreamUserItem(
            rig::streaming::StreamedUserContent::ToolResult { tool_result, .. },
        )) => {
            println!("[Tool result received for: {}]", tool_result.id);
            stdout.flush()?;
        }

        Ok(MultiTurnStreamItem::FinalResponse(final_response)) => {
            let usage = final_response.usage();
            input_tokens = usage.input_tokens;
            output_tokens = usage.output_tokens;
        }

        Err(e) => {
            eprintln!("\nError: {}", e);
            break;
        }

        _ => {}
    }
}
Enter fullscreen mode Exit fullscreen mode

Beyond this tutorial

OK, OK, so maybe we didn't quite build out Claude Code. But that's okay! Now you know what the bare minimum is when it comes to building a coding agent in Rust.

Beyond this though if you want to productionise this, there are some additional features you may wish to consider:

  • Cost calculations (derived from token usage)
  • Auto compaction of context
  • Skills/MCP integrations
  • Improving the system prompt

If you're interested in learning more about Rig, check us out:

Top comments (0)