main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/anthropics/anthropic-sdk-go"
	"github.com/invopop/jsonschema"
)

func main() {
	client := anthropic.NewClient()

	scanner := bufio.NewScanner(os.Stdin)
	getUserMessage := func() (string, bool) {
		if !scanner.Scan() {
			return "", false
		}
		return scanner.Text(), true
	}

	tools := []ToolDefinition{ReadFileDefinition, ListFilesDefinition, EditFileDefinition, PingDefinition, BashDefinition}
	agent := NewAgent(&client, getUserMessage, tools)
	err := agent.Run(context.TODO())
	if err != nil {
		fmt.Printf("Error: %s\n", err.Error())
	}
}

func NewAgent(client *anthropic.Client, getUserMessage func() (string, bool), tools []ToolDefinition) *Agent {
	return &Agent{
		client:         client,
		getUserMessage: getUserMessage,
		tools:          tools,
	}
}

type Agent struct {
	client         *anthropic.Client
	getUserMessage func() (string, bool)
	tools          []ToolDefinition
}

func (a *Agent) Run(ctx context.Context) error {
	conversation := []anthropic.MessageParam{}

	fmt.Printf("Chat with Kimi (use 'ctrl-c' to quit)\n\n")

	readUserInput := true
	for {
		if readUserInput {
			fmt.Print("\u001b[94mYou\u001b[0m: ")
			userInput, ok := a.getUserMessage()
			if !ok {
				break
			}

			userMessage := anthropic.NewUserMessage(anthropic.NewTextBlock(userInput))
			conversation = append(conversation, userMessage)
		}

		message, err := a.runInference(ctx, conversation)
		if err != nil {
			return err
		}
		conversation = append(conversation, message.ToParam())

		toolResults := []anthropic.ContentBlockParamUnion{}
		for _, content := range message.Content {
			switch content.Type {
			case "text":
				fmt.Printf("\u001b[93mKimi\u001b[0m: %s\n\n", content.Text)
			case "tool_use":
				result := a.executeTool(content.ID, content.Name, content.Input)
				toolResults = append(toolResults, result)
			}
		}
		if len(toolResults) == 0 {
			readUserInput = true
			continue
		}
		readUserInput = false
		conversation = append(conversation, anthropic.NewUserMessage(toolResults...))
	}

	return nil
}

func (a *Agent) runInference(ctx context.Context, conversation []anthropic.MessageParam) (*anthropic.Message, error) {
	anthropicTools := []anthropic.ToolUnionParam{}
	for _, tool := range a.tools {
		anthropicTools = append(anthropicTools, anthropic.ToolUnionParam{
			OfTool: &anthropic.ToolParam{
				Name:        tool.Name,
				Description: anthropic.String(tool.Description),
				InputSchema: tool.InputSchema,
			},
		})
	}

	message, err := a.client.Messages.New(ctx, anthropic.MessageNewParams{
		Model:     "kimi-k2-thinking",
		MaxTokens: int64(1024),
		System: []anthropic.TextBlockParam{{Text: func() string {
			b, _ := os.ReadFile("system_prompt.txt")
			return string(b)
		}()}},
		Messages: conversation,
		Tools:    anthropicTools,
	})
	return message, err
}

func GenerateSchema[T any]() anthropic.ToolInputSchemaParam {
	reflector := jsonschema.Reflector{
		AllowAdditionalProperties: false,
		DoNotReference:            true,
	}
	var v T

	schema := reflector.Reflect(v)

	return anthropic.ToolInputSchemaParam{
		Properties: schema.Properties,
	}
}

func (a *Agent) executeTool(id, name string, input json.RawMessage) anthropic.ContentBlockParamUnion {
	var toolDef ToolDefinition
	var found bool
	for _, tool := range a.tools {
		if tool.Name == name {
			toolDef = tool
			found = true
			break
		}
	}
	if !found {
		return anthropic.NewToolResultBlock(id, "tool not found", true)
	}

	fmt.Printf("\u001b[92mtool\u001b[0m: %s(%s)\n\n", name, input)
	response, err := toolDef.Function(input)
	if err != nil {
		return anthropic.NewToolResultBlock(id, err.Error(), true)
	}
	return anthropic.NewToolResultBlock(id, response, false)
}

// ****************************************************************************
// Tools
// ****************************************************************************

type ToolDefinition struct {
	Name        string                         `json:"name"`
	Description string                         `json:"description"`
	InputSchema anthropic.ToolInputSchemaParam `json:"input_schema"`
	Function    func(input json.RawMessage) (string, error)
}

// Read file
// ****************************************************************************

var ReadFileDefinition = ToolDefinition{
	Name:        "read_file",
	Description: "Read the contents of a given relative file path. Use this when you want to see what's inside a file. Do not use this with directory names.",
	InputSchema: ReadFileInputSchema,
	Function:    ReadFile,
}

type ReadFileInput struct {
	Path string `json:"path" jsonschema_description:"The relative path of a file in the working directory."`
}

var ReadFileInputSchema = GenerateSchema[ReadFileInput]()

func ReadFile(input json.RawMessage) (string, error) {
	readFileInput := ReadFileInput{}
	err := json.Unmarshal(input, &readFileInput)
	if err != nil {
		return "", err
	}

	content, err := os.ReadFile(readFileInput.Path)
	if err != nil {
		return "", err
	}

	return string(content), nil
}

// List files and directories at a given path.
// ****************************************************************************

var ListFilesDefinition = ToolDefinition{
	Name:        "list_files",
	Description: "List files and directories at a given path. If no path is provided, lists files in the current directory.",
	InputSchema: ListFilesInputSchema,
	Function:    ListFiles,
}

type ListFilesInput struct {
	Path string `json:"path,omitempty" jsonschema_description:"Optional relative path to list files from. Defaults to current directory if not provided."`
}

var ListFilesInputSchema = GenerateSchema[ListFilesInput]()

func ListFiles(input json.RawMessage) (string, error) {
	listFilesInput := ListFilesInput{}
	err := json.Unmarshal(input, &listFilesInput)
	if err != nil {
		panic(err)
	}

	dir := "."
	if listFilesInput.Path != "" {
		dir = listFilesInput.Path
	}

	var files []string
	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		relPath, err := filepath.Rel(dir, path)
		if err != nil {
			return err
		}

		if relPath != "." {
			if info.IsDir() {
				files = append(files, relPath+"/")
			} else {
				files = append(files, relPath)
			}
		}
		return nil
	})

	if err != nil {
		return "", err
	}

	result, err := json.Marshal(files)
	if err != nil {
		return "", err
	}

	return string(result), nil
}

// Edit a file.
// ****************************************************************************

var EditFileDefinition = ToolDefinition{
	Name: "edit_file",
	Description: `Make edits to a text file.

Replaces 'old_str' with 'new_str' in the given file. 'old_str' and 'new_str' MUST be different from each other.

If the file specified with path doesn't exist, it will be created.
`,
	InputSchema: EditFileInputSchema,
	Function:    EditFile,
}

type EditFileInput struct {
	Path   string `json:"path" jsonschema_description:"The path to the file"`
	OldStr string `json:"old_str" jsonschema_description:"Text to search for - must match exactly and must only have one match exactly"`
	NewStr string `json:"new_str" jsonschema_description:"Text to replace old_str with"`
}

var EditFileInputSchema = GenerateSchema[EditFileInput]()

func EditFile(input json.RawMessage) (string, error) {
	editFileInput := EditFileInput{}
	err := json.Unmarshal(input, &editFileInput)
	if err != nil {
		return "", err
	}

	if editFileInput.Path == "" || editFileInput.OldStr == editFileInput.NewStr {
		return "", fmt.Errorf("invalid input parameters")
	}

	content, err := os.ReadFile(editFileInput.Path)
	if err != nil {
		if os.IsNotExist(err) && editFileInput.OldStr == "" {
			return createNewFile(editFileInput.Path, editFileInput.NewStr)
		}
		return "", err
	}

	oldContent := string(content)
	newContent := strings.Replace(oldContent, editFileInput.OldStr, editFileInput.NewStr, -1)

	if oldContent == newContent && editFileInput.OldStr != "" {
		return "", fmt.Errorf("old_str not found in file")
	}

	err = os.WriteFile(editFileInput.Path, []byte(newContent), 0644)
	if err != nil {
		return "", err
	}

	return "OK", nil
}

func createNewFile(filePath, content string) (string, error) {
	dir := path.Dir(filePath)
	if dir != "." {
		err := os.MkdirAll(dir, 0755)
		if err != nil {
			return "", fmt.Errorf("failed to create directory: %w", err)
		}
	}

	err := os.WriteFile(filePath, []byte(content), 0644)
	if err != nil {
		return "", fmt.Errorf("failed to create file: %w", err)
	}

	return fmt.Sprintf("Successfully created file %s", filePath), nil
}

// Ping a host
// ****************************************************************************

var PingDefinition = ToolDefinition{
	Name:        "ping",
	Description: "Ping some host on the internet",
	InputSchema: PingInputSchema,
	Function:    Ping,
}

type PingInput struct {
	Host string `json:"host" jsonschema_description:"Hostname or IP address to ping"`
}

var PingInputSchema = GenerateSchema[PingInput]()

func Ping(input json.RawMessage) (string, error) {
	pingInput := PingInput{}
	err := json.Unmarshal(input, &pingInput)
	if err != nil {
		return "", err
	}

	if pingInput.Host == "" {
		return "", fmt.Errorf("host is required")
	}

	cmd := exec.Command("ping", "-c", "5", pingInput.Host)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Sprintf("error: %s\n%s", err.Error(), string(output)), nil
	}

	return string(output), nil
}

// Bash - Execute bash commands
// ****************************************************************************

var BashDefinition = ToolDefinition{
	Name: "bash",
	Description: func() string {
		b, _ := os.ReadFile("bash.txt")
		return string(b)
	}(),
	InputSchema: BashInputSchema,
	Function:    Bash,
}

type BashInput struct {
	Command     string `json:"command" jsonschema_description:"The command to execute"`
	Timeout     int    `json:"timeout,omitempty" jsonschema_description:"Optional timeout in milliseconds. Defaults to 120000ms (2 minutes)."`
	Workdir     string `json:"workdir,omitempty" jsonschema_description:"The working directory to run the command in. Defaults to current directory. Use this instead of 'cd' commands."`
	Description string `json:"description" jsonschema_description:"Clear, concise description of what this command does in 5-10 words. Examples:\nInput: ls\nOutput: Lists files in current directory\n\nInput: git status\nOutput: Shows working tree status\n\nInput: npm install\nOutput: Installs package dependencies\n\nInput: mkdir foo\nOutput: Creates directory 'foo'"`
}

var BashInputSchema = GenerateSchema[BashInput]()

const (
	DefaultTimeoutMs = 120000 // 2 minutes
	MaxOutputLength  = 30000
)

const SigkillTimeoutMs = 200

func Bash(input json.RawMessage) (string, error) {
	bashInput := BashInput{}
	err := json.Unmarshal(input, &bashInput)
	if err != nil {
		return "", err
	}

	if bashInput.Command == "" {
		return "", fmt.Errorf("command is required")
	}

	// Validate timeout
	if bashInput.Timeout < 0 {
		return "", fmt.Errorf("invalid timeout value: %d. Timeout must be a positive number", bashInput.Timeout)
	}

	// Set default timeout if not provided
	timeout := bashInput.Timeout
	if timeout <= 0 {
		timeout = DefaultTimeoutMs
	}

	// Set working directory
	workdir := bashInput.Workdir
	if workdir == "" {
		workdir, _ = os.Getwd()
	}

	// Execute command using bash with process group
	cmd := exec.Command("bash", "-c", bashInput.Command)
	cmd.Dir = workdir
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	// Create pipes for streaming output
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", fmt.Errorf("failed to create stdout pipe: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return "", fmt.Errorf("failed to create stderr pipe: %w", err)
	}

	// Start the command
	if err := cmd.Start(); err != nil {
		return "", fmt.Errorf("failed to start command: %w", err)
	}

	var output strings.Builder
	var outputMu sync.Mutex
	var exited bool
	var timedOut bool

	// Function to kill the process tree
	killTree := func() {
		pid := cmd.Process.Pid
		if exited {
			return
		}

		// Send SIGTERM to the process group
		pgid, err := syscall.Getpgid(pid)
		if err == nil {
			syscall.Kill(-pgid, syscall.SIGTERM)
		} else {
			cmd.Process.Signal(syscall.SIGTERM)
		}

		// Wait for SIGKILL_TIMEOUT_MS then send SIGKILL
		time.Sleep(time.Duration(SigkillTimeoutMs) * time.Millisecond)
		if !exited {
			if pgid, err := syscall.Getpgid(pid); err == nil {
				syscall.Kill(-pgid, syscall.SIGKILL)
			} else {
				cmd.Process.Kill()
			}
		}
	}

	// Append output from reader
	appendOutput := func(reader io.Reader) {
		buf := make([]byte, 4096)
		for {
			n, err := reader.Read(buf)
			if n > 0 {
				outputMu.Lock()
				if output.Len() <= MaxOutputLength {
					output.Write(buf[:n])
				}
				outputMu.Unlock()
			}
			if err != nil {
				break
			}
		}
	}

	// Read stdout and stderr concurrently
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		appendOutput(stdout)
	}()
	go func() {
		defer wg.Done()
		appendOutput(stderr)
	}()

	// Set up timeout timer
	timeoutTimer := time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
		timedOut = true
		killTree()
	})
	defer timeoutTimer.Stop()

	// Wait for output readers to finish
	wg.Wait()

	// Wait for the command to finish
	cmdErr := cmd.Wait()
	exited = true

	outputStr := output.String()

	// Build result metadata
	var resultMetadata []string

	if len(outputStr) > MaxOutputLength {
		outputStr = outputStr[:MaxOutputLength]
		resultMetadata = append(resultMetadata, fmt.Sprintf("bash tool truncated output as it exceeded %d char limit", MaxOutputLength))
	}

	if timedOut {
		resultMetadata = append(resultMetadata, fmt.Sprintf("bash tool terminated command after exceeding timeout %d ms", timeout))
	}

	// Append metadata if any
	if len(resultMetadata) > 0 {
		outputStr += "\n\n<bash_metadata>\n" + strings.Join(resultMetadata, "\n") + "\n</bash_metadata>"
	}

	if cmdErr != nil {
		// Return output with error info but don't fail the tool
		exitCode := -1
		if cmd.ProcessState != nil {
			exitCode = cmd.ProcessState.ExitCode()
		}
		return fmt.Sprintf("%s\n\nexit code: %d", outputStr, exitCode), nil
	}

	return outputStr, nil
}