initial commit

1 files changed, 611 insertions, 0 deletions

diff --git a/main.go b/main.go
new file mode 100644
index 0000000..df82ac9
--- /dev/null
+++ b/main.go
@@ -0,0 +1,611 @@
+// pico — a minimal single-binary CLI coding agent.
+package main
+
+import (
+        "bufio"
+        "bytes"
+        "encoding/json"
+        "fmt"
+        "io"
+        "net/http"
+        "os"
+        "os/exec"
+        "path/filepath"
+        "regexp"
+        "sort"
+        "strconv"
+        "strings"
+        "syscall"
+        "time"
+        "unicode/utf8"
+)
+
+// ------------------------------------------------------------------- palette
+
+const (
+        reset = "\x1b[0m"
+        bold  = "\x1b[1m"
+        dim   = "\x1b[2m"
+        blue  = "\x1b[34m"
+        cyan  = "\x1b[36m"
+        green = "\x1b[32m"
+        red   = "\x1b[31m"
+)
+
+// -------------------------------------------------------------------- config
+
+var (
+        apiURL = firstEnv("API_URL", "ANTHROPIC_BASE_URL", "http://localhost:8000")
+        apiKey = firstEnv("API_KEY", "ANTHROPIC_API_KEY", "")
+        model  = firstEnv("MODEL", "ANTHROPIC_MODEL", "gpt-oss-120b")
+)
+
+func firstEnv(keys ...string) string {
+        def := keys[len(keys)-1]
+        for _, k := range keys[:len(keys)-1] {
+                if k == "" {
+                        continue
+                }
+                if v := os.Getenv(k); v != "" {
+                        return strings.TrimRight(v, "/")
+                }
+        }
+        return def
+}
+
+// --------------------------------------------------------------------- tools
+
+type param struct {
+        name string
+        typ  string // trailing "?" marks optional
+}
+
+type tool struct {
+        name   string
+        desc   string
+        params []param
+        run    func(in map[string]any) (string, error)
+}
+
+var tools = []tool{
+        {"read", "Read file with line numbers (file path, not directory)",
+                []param{{"path", "string"}, {"offset", "integer?"}, {"limit", "integer?"}}, toolRead},
+        {"write", "Write content to file",
+                []param{{"path", "string"}, {"content", "string"}}, toolWrite},
+        {"edit", "Replace old with new in file (old must be unique unless all=true)",
+                []param{{"path", "string"}, {"old", "string"}, {"new", "string"}, {"all", "boolean?"}}, toolEdit},
+        {"glob", "Find files by pattern, sorted by mtime",
+                []param{{"pat", "string"}, {"path", "string?"}}, toolGlob},
+        {"grep", "Search files for regex pattern",
+                []param{{"pat", "string"}, {"path", "string?"}}, toolGrep},
+        {"bash", "Run shell command",
+                []param{{"cmd", "string"}}, toolBash},
+}
+
+var byName = func() map[string]*tool {
+        m := map[string]*tool{}
+        for i := range tools {
+                m[tools[i].name] = &tools[i]
+        }
+        return m
+}()
+
+// schemas is built once at startup and reused verbatim on every request, so
+// the serialized prefix stays byte-identical and KV cache holds.
+// Properties are emitted in declaration order, which map marshaling would not
+// preserve. Each property carries "type" and nothing else.
+var schemas = func() []json.RawMessage {
+        out := make([]json.RawMessage, len(tools))
+        for i, t := range tools {
+                var props, required []string
+                for _, p := range t.params {
+                        props = append(props, fmt.Sprintf(`%s:{"type":%s}`, jstr(p.name), jstr(strings.TrimSuffix(p.typ, "?"))))
+                        if !strings.HasSuffix(p.typ, "?") {
+                                required = append(required, jstr(p.name))
+                        }
+                }
+                out[i] = json.RawMessage(fmt.Sprintf(
+                        `{"name":%s,"description":%s,"input_schema":{"type":"object","properties":{%s},"required":[%s]}}`,
+                        jstr(t.name), jstr(t.desc), strings.Join(props, ","), strings.Join(required, ",")))
+        }
+        return out
+}()
+
+func jstr(s string) string {
+        enc, _ := json.Marshal(s)
+        return string(enc)
+}
+
+// ------------------------------------------------------------- arg coercion
+
+func argStr(in map[string]any, k, def string) string {
+        if s, ok := in[k].(string); ok {
+                return s
+        }
+        return def
+}
+
+func argInt(in map[string]any, k string, def int) int {
+        switch v := in[k].(type) {
+        case float64:
+                return int(v)
+        case string:
+                if n, err := strconv.Atoi(v); err == nil {
+                        return n
+                }
+        }
+        return def
+}
+
+func argBool(in map[string]any, k string, def bool) bool {
+        if b, ok := in[k].(bool); ok {
+                return b
+        }
+        return def
+}
+
+// ---------------------------------------------------------------------- read
+
+func toolRead(in map[string]any) (string, error) {
+        data, err := os.ReadFile(argStr(in, "path", ""))
+        if err != nil {
+                return "", err
+        }
+        // SplitAfter keeps the newline on each line; drop the empty tail element.
+        lines := strings.SplitAfter(string(data), "\n")
+        if lines[len(lines)-1] == "" {
+                lines = lines[:len(lines)-1]
+        }
+        offset := argInt(in, "offset", 0)
+        limit := argInt(in, "limit", len(lines))
+        start := min(max(offset, 0), len(lines))
+        end := min(max(offset+limit, start), len(lines))
+
+        var b strings.Builder
+        for i, ln := range lines[start:end] {
+                fmt.Fprintf(&b, "%4d| %s", start+i+1, ln)
+        }
+        return b.String(), nil
+}
+
+// ---------------------------------------------------------------- write, edit
+
+func toolWrite(in map[string]any) (string, error) {
+        if err := os.WriteFile(argStr(in, "path", ""), []byte(argStr(in, "content", "")), 0644); err != nil {
+                return "", err
+        }
+        return "ok", nil
+}
+
+func toolEdit(in map[string]any) (string, error) {
+        path := argStr(in, "path", "")
+        data, err := os.ReadFile(path)
+        if err != nil {
+                return "", err
+        }
+        s, old, new := string(data), argStr(in, "old", ""), argStr(in, "new", "")
+
+        // Both of these are ordinary results, not errors: the model reads them
+        // and retries with a longer `old`.
+        switch n := strings.Count(s, old); {
+        case n == 0:
+                return "error: old_string not found", nil
+        case n > 1 && !argBool(in, "all", false):
+                return fmt.Sprintf("error: old_string appears %d times, must be unique (use all=true)", n), nil
+        }
+
+        if argBool(in, "all", false) {
+                s = strings.ReplaceAll(s, old, new)
+        } else {
+                s = strings.Replace(s, old, new, 1)
+        }
+        if err := os.WriteFile(path, []byte(s), 0644); err != nil {
+                return "", err
+        }
+        return "ok", nil
+}
+
+// ---------------------------------------------------------------- glob, grep
+
+func argPath(in map[string]any) string {
+        if p := argStr(in, "path", "."); p != "" {
+                return p
+        }
+        return "."
+}
+
+func toolGlob(in map[string]any) (string, error) {
+        matches := globExpand(strings.ReplaceAll(argPath(in)+"/"+argStr(in, "pat", ""), "//", "/"))
+        if len(matches) == 0 {
+                return "none", nil
+        }
+        mtime := make(map[string]int64, len(matches))
+        for _, m := range matches {
+                if fi, err := os.Stat(m); err == nil && !fi.IsDir() {
+                        mtime[m] = fi.ModTime().UnixNano()
+                }
+        }
+        sort.SliceStable(matches, func(i, j int) bool { return mtime[matches[i]] > mtime[matches[j]] })
+        return strings.Join(matches, "\n"), nil
+}
+
+func toolGrep(in map[string]any) (string, error) {
+        re, err := regexp.Compile(argStr(in, "pat", ""))
+        if err != nil {
+                return "", err
+        }
+        var hits []string
+        for _, p := range globExpand(argPath(in) + "/**") {
+                f, err := os.Open(p)
+                if err != nil {
+                        continue
+                }
+                // Directories appear in the ** expansion; scanning one yields nothing.
+                sc := bufio.NewScanner(f)
+                sc.Buffer(make([]byte, 0, 64*1024), 8<<20)
+                for n := 1; sc.Scan(); n++ {
+                        if line := sc.Text(); re.MatchString(line) {
+                                hits = append(hits, fmt.Sprintf("%s:%d:%s", p, n, strings.TrimRight(line, " \t\r")))
+                        }
+                }
+                f.Close()
+                if len(hits) >= 50 {
+                        break
+                }
+        }
+        if len(hits) == 0 {
+                return "none", nil
+        }
+        if len(hits) > 50 {
+                hits = hits[:50]
+        }
+        return strings.Join(hits, "\n"), nil
+}
+
+// ---------------------------------------------------------------------- bash
+
+func toolBash(in map[string]any) (string, error) {
+        pr, pw, err := os.Pipe()
+        if err != nil {
+                return "", err
+        }
+        cmd := exec.Command("sh", "-c", argStr(in, "cmd", ""))
+        cmd.Stdout, cmd.Stderr = pw, pw // one pipe: interleaved in real order
+        // Own process group, so the timeout below can kill the whole tree.
+        cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+        if err := cmd.Start(); err != nil {
+                pw.Close()
+                pr.Close()
+                return "", err
+        }
+        // The parent MUST drop its own copy of the write end, or the pipe never
+        // reaches EOF and the reader below blocks forever -- even for `true`.
+        pw.Close()
+
+        done := make(chan string, 1)
+        go func() {
+                defer pr.Close()
+                var acc strings.Builder
+                sc := bufio.NewScanner(pr)
+                sc.Buffer(make([]byte, 0, 64*1024), 8<<20)
+                for sc.Scan() {
+                        line := sc.Text()
+                        fmt.Printf("  %s│ %s%s\n", dim, strings.TrimRight(line, " \t\r"), reset)
+                        acc.WriteString(line + "\n")
+                }
+                done <- acc.String()
+        }()
+
+        var out string
+        select {
+        case out = <-done:
+        case <-time.After(30 * time.Second):
+                // Kill the GROUP, not just sh: `sh -c "sleep 60"` forks, and the
+                // grandchild inherits the pipe's write end. Killing sh alone leaves
+                // the pipe open and the drain below blocks until the grandchild
+                // exits on its own -- a 30s timeout that takes 60s.
+                if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
+                        cmd.Process.Kill()
+                }
+                out = <-done // let the reader drain before reporting
+                out += "\n(timed out after 30s)"
+        }
+        cmd.Wait() // always reap
+
+        if out = strings.TrimSpace(out); out == "" {
+                return "(empty)", nil
+        }
+        return out, nil
+}
+
+// ------------------------------------------------------------------ dispatch
+
+func runTool(name string, input json.RawMessage) string {
+        t, ok := byName[name]
+        if !ok {
+                return fmt.Sprintf("error: unknown tool %q", name)
+        }
+        var in map[string]any
+        json.Unmarshal(input, &in)
+        out, err := t.run(in)
+        if err != nil {
+                return "error: " + err.Error()
+        }
+        return out
+}
+
+// --------------------------------------------------------------- glob engine
+
+func globExpand(pattern string) []string {
+        root, rest := ".", pattern
+        if strings.HasPrefix(pattern, "/") {
+                root, rest = "/", pattern[1:]
+        }
+        var out []string
+        globWalk(root, strings.Split(rest, "/"), &out)
+        if root == "." {
+                for i := range out {
+                        out[i] = strings.TrimPrefix(out[i], "./")
+                }
+        }
+        return out
+}
+
+func globWalk(dir string, segs []string, out *[]string) {
+        if len(segs) == 0 {
+                *out = append(*out, dir)
+                return
+        }
+        seg, rest := segs[0], segs[1:]
+
+        if seg == "" || seg == "." {
+                globWalk(dir, rest, out)
+                return
+        }
+
+        if seg == "**" {
+                globWalk(dir, rest, out) // ** matches zero path segments
+                for _, e := range readDir(dir) {
+                        if strings.HasPrefix(e.Name(), ".") {
+                                continue
+                        }
+                        switch p := join(dir, e.Name()); {
+                        case e.IsDir():
+                                globWalk(p, segs, out) // keep ** at the head
+                        case len(rest) == 0:
+                                *out = append(*out, p)
+                        }
+                }
+                return
+        }
+
+        for _, e := range readDir(dir) {
+                name := e.Name()
+                if strings.HasPrefix(name, ".") && !strings.HasPrefix(seg, ".") {
+                        continue
+                }
+                if ok, _ := filepath.Match(seg, name); !ok {
+                        continue
+                }
+                p := join(dir, name)
+                if len(rest) == 0 {
+                        *out = append(*out, p) // file or directory
+                } else if e.IsDir() {
+                        globWalk(p, rest, out)
+                }
+        }
+}
+
+// readDir returns entries sorted by name; unreadable directories yield none.
+func readDir(dir string) []os.DirEntry {
+        es, err := os.ReadDir(dir)
+        if err != nil {
+                return nil
+        }
+        return es
+}
+
+func join(dir, name string) string {
+        if strings.HasSuffix(dir, "/") {
+                return dir + name
+        }
+        return dir + "/" + name
+}
+
+// ------------------------------------------------------------------ API call
+
+type message struct {
+        Role    string `json:"role"`
+        Content any    `json:"content"`
+}
+
+type toolResult struct {
+        Type      string `json:"type"`
+        ToolUseID string `json:"tool_use_id"`
+        Content   string `json:"content"`
+}
+
+type block struct {
+        Type  string          `json:"type"`
+        Text  string          `json:"text"`
+        ID    string          `json:"id"`
+        Name  string          `json:"name"`
+        Input json.RawMessage `json:"input"`
+}
+
+var systemPrompt = func() string {
+        cwd, _ := os.Getwd()
+        return "Concise coding assistant. cwd: " + cwd
+}()
+
+func callAPI(messages []message) ([]json.RawMessage, error) {
+        body, err := json.Marshal(map[string]any{
+                "model":      model,
+                "max_tokens": 8192,
+                "system":     systemPrompt,
+                "messages":   messages,
+                "tools":      schemas,
+        })
+        if err != nil {
+                return nil, err
+        }
+        req, err := http.NewRequest("POST", apiURL+"/v1/messages", bytes.NewReader(body))
+        if err != nil {
+                return nil, err
+        }
+        // Every header is byte-identical across turns: a varying one would
+        // invalidate llama-swap's KV cache on every request.
+        req.Header.Set("Content-Type", "application/json")
+        req.Header.Set("anthropic-version", "2023-06-01")
+        req.Header.Set("x-api-key", apiKey)
+
+        // No timeout: a cold llama-swap model load costs ~10-25s before prefill
+        // even begins, and llama-swap's own responseHeader timeout is 1440s.
+        resp, err := http.DefaultClient.Do(req)
+        if err != nil {
+                return nil, err
+        }
+        defer resp.Body.Close()
+        raw, err := io.ReadAll(resp.Body)
+        if err != nil {
+                return nil, err
+        }
+        if resp.StatusCode >= 400 {
+                return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(raw), 500))
+        }
+        // Content blocks stay raw: thinking/text/tool_use and anything unknown
+        // all round-trip verbatim into the next request.
+        var parsed struct {
+                Content []json.RawMessage `json:"content"`
+        }
+        if err := json.Unmarshal(raw, &parsed); err != nil {
+                return nil, err
+        }
+        return parsed.Content, nil
+}
+
+// ---------------------------------------------------------------- agent loop
+
+func agentLoop(messages *[]message) error {
+        for {
+                blocks, err := callAPI(*messages)
+                if err != nil {
+                        return err
+                }
+                var results []toolResult
+                for _, raw := range blocks {
+                        var b block
+                        if err := json.Unmarshal(raw, &b); err != nil {
+                                return err
+                        }
+                        switch b.Type {
+                        case "text":
+                                fmt.Printf("\n%s⏺%s %s\n", cyan, reset, renderMarkdown(b.Text))
+                        case "tool_use":
+                                fmt.Printf("\n%s⏺ %s%s(%s%s%s)\n",
+                                        green, capitalize(b.Name), reset, dim, argPreview(b.Name, b.Input), reset)
+                                out := runTool(b.Name, b.Input)
+                                printResult(out)
+                                results = append(results, toolResult{"tool_result", b.ID, out})
+                        }
+                        // Other block types (thinking) are not displayed but still echoed.
+                }
+                *messages = append(*messages, message{"assistant", blocks})
+                if len(results) == 0 {
+                        return nil
+                }
+                *messages = append(*messages, message{"user", results})
+        }
+}
+
+// ------------------------------------------------------------------- display
+
+var boldRe = regexp.MustCompile(`(?s)\*\*(.*?)\*\*`)
+
+func renderMarkdown(s string) string { return boldRe.ReplaceAllString(s, bold+"$1"+reset) }
+
+func capitalize(s string) string {
+        if s == "" {
+                return s
+        }
+        return strings.ToUpper(s[:1]) + s[1:]
+}
+
+// truncate cuts to n code points, never bytes, and never adds an ellipsis.
+func truncate(s string, n int) string {
+        if utf8.RuneCountInString(s) <= n {
+                return s
+        }
+        return string([]rune(s)[:n])
+}
+
+func argPreview(name string, input json.RawMessage) string {
+        t, ok := byName[name]
+        if !ok {
+                return ""
+        }
+        var in map[string]any
+        json.Unmarshal(input, &in)
+        for _, p := range t.params {
+                v, present := in[p.name]
+                if !present {
+                        continue
+                }
+                if s, isStr := v.(string); isStr {
+                        return truncate(s, 50)
+                }
+                enc, _ := json.Marshal(v)
+                return truncate(string(enc), 50)
+        }
+        return ""
+}
+
+func printResult(out string) {
+        lines := strings.Split(out, "\n")
+        preview := truncate(lines[0], 60)
+        if len(lines) > 1 {
+                preview += fmt.Sprintf(" ... +%d lines", len(lines)-1)
+        } else if utf8.RuneCountInString(lines[0]) > 60 {
+                preview += "..."
+        }
+        fmt.Printf("  %s⎿  %s%s\n", dim, preview, reset)
+}
+
+// ---------------------------------------------------------------------- REPL
+
+func main() {
+        cwd, _ := os.Getwd()
+        fmt.Printf("%spico%s | %s%s | %s%s\n\n", bold, reset, dim, model, cwd, reset)
+
+        sep := dim + strings.Repeat("─", 80) + reset
+        stdin := bufio.NewScanner(os.Stdin)
+        stdin.Buffer(make([]byte, 0, 64*1024), 8<<20)
+
+        var messages []message
+        for {
+                fmt.Println(sep)
+                fmt.Printf("%s%s❯%s ", bold, blue, reset)
+                if !stdin.Scan() {
+                        fmt.Println()
+                        return
+                }
+                fmt.Println(sep)
+
+                switch line := strings.TrimSpace(stdin.Text()); line {
+                case "":
+                        continue
+                case "/q", "exit":
+                        return
+                case "/c":
+                        messages = nil
+                        fmt.Printf("%s⏺ Cleared conversation%s\n", green, reset)
+                default:
+                        messages = append(messages, message{"user", line})
+                        if err := agentLoop(&messages); err != nil {
+                                fmt.Printf("%s⏺ Error: %v%s\n", red, err, reset)
+                                continue
+                        }
+                        fmt.Println()
+                }
+        }
+}