Move shared test helpers into helper files

[onechat] / terminal.go

package main

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

// terminalResultLines is the most lines of a tool result that the client shows.
const terminalResultLines = 10

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

// terminal 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 terminal(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 terminalSession(ctx, conn, in, out)
}

// terminalSession 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 terminalSession(ctx context.Context, conn net.Conn, in io.Reader, out io.Writer) error {
        sent := make(chan error, 1)
        go func() {
                sent <- terminalSendLines(ctx, conn, in)
        }()

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

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

        return err
}

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

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

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

        return scanner.Err()
}

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

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

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

        return scanner.Err()
}

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

        return err
}

// terminalEventText returns the text that the client shows for event, or "" for
// none. The client shows a user message only from another source, such as
// telegram, because the user typed the others. A tool result shows its first
// terminalResultLines lines.
func terminalEventText(event Event) string {
        switch event.Kind {
        case "ready":
                return "[connected to " + event.Text + "]"
        case "user":
                if event.Source == terminalSource {
                        return ""
                }

                return "[" + event.Source + "] " + event.Text
        case "reply":
                return event.Text
        case "tool":
                return "[tool] " + event.Text
        case "echo":
                return "[result] " + terminalFirstLines(event.Text, terminalResultLines)
        case "busy":
                return "[working]"
        case "idle":
                return "[done]"
        case "error":
                return "[error] " + event.Text
        default:
                return ""
        }
}

// terminalFirstLines returns the first count lines of text, with the number of
// the lines left out.
func terminalFirstLines(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)
}