Check the socket path before the daemon writes

[onechat] / chat.go

package main

import (
        "bufio"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net"
        "path/filepath"
        "strings"
)

const (
        // stopCommand is the input line that interrupts the running turn.
        stopCommand = "/stop"

        // resultLines is the most lines of a tool result that the client shows.
        resultLines = 10
)

// chat runs "onechat chat": the terminal client of the daemon. Each line of
// input is a user message, and the line /stop interrupts the running turn. The
// client shows the replies, the tool calls, and their results as they happen.
// It ends at the end of the input or when the daemon closes the connection.
func chat(ctx context.Context, in io.Reader, out io.Writer) error {
        dir, err := dataDir()
        if err != nil {
                return err
        }

        path := filepath.Join(dir, socketFileName)
        err = checkSocketPath(path)
        if err != nil {
                return err
        }

        var dialer net.Dialer
        conn, err := dialer.DialContext(ctx, "unix", path)
        if err != nil {
                return fmt.Errorf(`connect to the daemon (start it with "onechat daemon"): %w`, err)
        }

        return chatSession(ctx, conn, in, out)
}

// errDaemonClosed tells that the daemon closed the connection.
var errDaemonClosed = errors.New("the daemon closed the connection")

// chatSession runs the client on conn until the input ends, the daemon closes
// the connection, or ctx ends, and closes conn. A read of the input that waits
// for a line can stay behind when the daemon closes the connection.
func chatSession(ctx context.Context, conn net.Conn, in io.Reader, out io.Writer) error {
        sent := make(chan error, 1)
        go func() {
                sent <- sendLines(ctx, conn, in)
        }()

        shown := make(chan error, 1)
        go func() {
                shown <- showEvents(conn, out)
        }()

        var err error
        select {
        case err = <-sent:
                err = errors.Join(err, conn.Close())
                <-shown
        case err = <-shown:
                err = errors.Join(cmpError(err, errDaemonClosed), conn.Close())
        case <-ctx.Done():
                err = conn.Close()
                <-shown
        }

        return err
}

// cmpError returns err, or fallback when err is nil.
func cmpError(err, fallback error) error {
        if err == nil {
                return fallback
        }

        return err
}

// sendLines sends each line of in as a request until the input ends.
func sendLines(ctx context.Context, conn net.Conn, in io.Reader) error {
        encoder := json.NewEncoder(conn)
        scanner := bufio.NewScanner(in)
        scanner.Buffer(make([]byte, 64*1024), requestBytes)
        for ctx.Err() == nil && scanner.Scan() {
                line := scanner.Text()
                request := clientRequest{
                        Type: "send",
                        Text: line,
                }

                switch strings.TrimSpace(line) {
                case "":
                        continue
                case stopCommand:
                        request = clientRequest{
                                Type: "interrupt",
                        }
                }

                err := encoder.Encode(request)
                if err != nil {
                        return fmt.Errorf("send to the daemon: %w", err)
                }
        }

        return scanner.Err()
}

// showEvents writes each event from the daemon to out until the connection
// closes.
func showEvents(conn net.Conn, out io.Writer) error {
        scanner := bufio.NewScanner(conn)
        scanner.Buffer(make([]byte, 64*1024), requestBytes)
        for scanner.Scan() {
                var event Event
                err := json.Unmarshal(scanner.Bytes(), &event)
                if err != nil {
                        return fmt.Errorf("decode event: %w", err)
                }

                text := eventText(event)
                if text == "" {
                        continue
                }

                _, err = io.WriteString(out, text+"\n")
                if err != nil {
                        return err
                }
        }

        return scanner.Err()
}

// eventText returns the text that the client shows for event, or "" for none.
// The client shows no user message, because the user typed it. A tool result
// shows its first resultLines lines.
func eventText(event Event) string {
        switch event.Kind {
        case "ready":
                return "[connected to " + event.Text + "]"
        case "reply":
                return event.Text
        case "tool":
                return "[tool] " + event.Text
        case "echo":
                return "[result] " + firstLines(event.Text, resultLines)
        case "busy":
                return "[working]"
        case "idle":
                return "[done]"
        case "error":
                return "[error] " + event.Text
        default:
                return ""
        }
}

// firstLines returns the first count lines of text, with the number of the
// lines left out.
func firstLines(text string, count int) string {
        text = strings.TrimRight(text, "\n")
        lines := strings.Split(text, "\n")
        if len(lines) <= count {
                return text
        }

        return strings.Join(lines[:count], "\n") + fmt.Sprintf("\n[%d more lines]", len(lines)-count)
}