[onechat] / turn.go

package main

import (
        "context"
        "errors"
        "fmt"
        "path/filepath"
        "sync"
        "time"
)

// newTurnLoop returns the loop of agent with its memory and its subagents. Its
// calls go through client, and notify gets each event. The clock gives the date
// of a turn.
func newTurnLoop(agent *Agent, memory *Memory, subagents *Subagents, client ChainClient, clock func() time.Time, notify func(Event)) (*TurnLoop, error) {
        layout, err := loadViewLayout(filepath.Join(agent.dir, viewLayoutFileName))
        if err != nil {
                return nil, err
        }

        l := &TurnLoop{
                agent:  agent,
                memory: memory,
                client: client,
                layout: layout,
                clock:  clock,
                notify: notify,
                wake:   make(chan struct{}, 1),
        }

        l.runs = &SubagentRuns{
                agent:     agent,
                memory:    memory,
                subagents: subagents,
                layout:    layout,
                client:    client,
                state:     l.state,
                report:    l.Report,
                note:      l.note,
                fail:      l.fail,
                notify:    notify,
                running:   make(map[string]*SubagentRun),
        }

        l.tools = &Toolbox{
                log:       l,
                chat:      memory,
                subagents: subagents,
                runs:      l.runs,
                workdir:   agent.config.Workdir,
        }

        return l, nil
}

// Event tells the clients what the agent does as it happens. Kind is the kind
// of a logged message (user, reply, tool, echo, note, work), or one of ready,
// busy, idle, error, and subagent, which tells that the subagent Text starts a
// run. ID is the ID of a logged message. Source names the client that sent a
// user message, such as terminal or telegram, so that each client can show the
// messages of the others, or the subagent of a report. Status tells how the run
// of a report ended: finished, stopped, or failed.
type Event struct {
        Kind   string `json:"kind"`
        ID     int    `json:"id,omitempty"`
        Text   string `json:"text,omitempty"`
        Source string `json:"source,omitempty"`
        Status string `json:"status,omitempty"`
}

// TurnLoop runs the turns of one agent. A user message goes to the log when it
// comes, and it wakes the loop. A turn waits until every message before it is
// summarized. Then the turn runs a Chain whose first request holds the prompt,
// the tools, and the input of turnInput.
//
// The view goes out as the messages of a ViewLayout, so the first request of a
// turn reads the view of the turn before from the prompt cache. User messages
// that come during a turn join the next request of its chain after the tool
// results, and so do the reports of the subagents. The subagents run until the
// loop ends.
type TurnLoop struct {
        agent  *Agent
        memory *Memory
        client ChainClient
        tools  *Toolbox
        runs   *SubagentRuns
        layout *ViewLayout
        clock  func() time.Time
        notify func(Event)
        wake   chan struct{}

        mutex   sync.Mutex
        pending []Message
        cancel  context.CancelFunc
        failure error
}

// Submit logs a user message from the client source and wakes the loop. The
// message joins the next request: the first one of a new turn, or the next one
// of a running turn.
func (l *TurnLoop) Submit(source, text string) (Message, error) {
        message, err := l.record("user", source, text)
        if err != nil {
                return Message{}, err
        }

        l.mutex.Lock()
        l.pending = append(l.pending, message)
        l.mutex.Unlock()

        select {
        case l.wake <- struct{}{}:
        default:
        }

        return message, nil
}

// Report logs the report text of the subagent name as a work message and tells
// the clients, with the status of the run. Unless the subagent was stopped, the
// report joins the waiting messages and wakes the loop, as a user message of
// Submit does. A failed write stops the loop, because the log refuses all
// writes after it.
func (l *TurnLoop) Report(name, text, status string) {
        message, err := l.Work(name, text, status)
        switch {
        case err != nil:
                l.fail(fmt.Errorf("%w: %w", errChainLogWrite, err))
        case status != subagentStopped:
                l.Wake(message)
        }
}

// Work logs text as a work message of source, such as the name of a subagent,
// and tells the clients, with status. The message starts no turn until it goes
// to Wake.
func (l *TurnLoop) Work(source, text, status string) (Message, error) {
        message, err := l.memory.Append("work", text)
        if err != nil {
                return Message{}, err
        }

        l.notify(Event{
                Kind:   "work",
                ID:     message.ID,
                Text:   text,
                Source: source,
                Status: status,
        })

        return message, nil
}

// Wake adds the logged messages to the waiting messages and wakes the loop
// once, as a user message of Submit does. They join the first request of a new
// turn together, or the next request of a running turn after the tool results.
func (l *TurnLoop) Wake(messages ...Message) {
        if len(messages) == 0 {
                return
        }

        l.mutex.Lock()
        l.pending = append(l.pending, messages...)
        l.mutex.Unlock()

        select {
        case l.wake <- struct{}{}:
        default:
        }
}

// StopSubagent ends the run of the subagent name and tells if it ran. Its
// report is logged, but it starts no turn.
func (l *TurnLoop) StopSubagent(name string) (bool, error) {
        return l.runs.Stop(name)
}

// Interrupt ends the running turn. A running tool is killed, and the turn sends
// no further request.
func (l *TurnLoop) Interrupt() bool {
        l.mutex.Lock()
        defer l.mutex.Unlock()

        if l.cancel == nil {
                return false
        }

        l.cancel()

        return true
}

// Busy tells if a turn runs.
func (l *TurnLoop) Busy() bool {
        l.mutex.Lock()
        defer l.mutex.Unlock()

        return l.cancel != nil
}

// Append logs a message and tells the clients. The toolbox logs its calls and
// results through it.
func (l *TurnLoop) Append(kind, text string) (Message, error) {
        return l.record(kind, "", text)
}

// Flush makes the log durable.
func (l *TurnLoop) Flush(ctx context.Context) error {
        return l.memory.Flush(ctx)
}

// Run runs a turn each time user messages or reports wait, until ctx ends. It
// returns the error of a failed log write, which stops the agent. At its start
// it resumes the subagents that a stop of the daemon cut, and at its end it
// ends the runs of the subagents and waits for them.
func (l *TurnLoop) Run(ctx context.Context) error {
        defer l.runs.Close()

        err := l.runs.Resume(ctx)
        if err != nil {
                return fmt.Errorf("resume subagents: %w", err)
        }

        for {
                for l.waiting() {
                        err := l.turn(ctx)
                        if err != nil {
                                return err
                        }
                }

                select {
                case <-l.wake:
                case <-ctx.Done():
                        return ctx.Err()
                }

                err := l.failed()
                if err != nil {
                        return err
                }
        }
}

// note logs a note and tells the clients. A note starts no turn.
func (l *TurnLoop) note(text string) error {
        _, err := l.record("note", "", text)

        return err
}

// record logs a message and tells the clients, with the client source of a user
// message.
func (l *TurnLoop) record(kind, source, text string) (Message, error) {
        message, err := l.memory.Append(kind, text)
        if err != nil {
                return Message{}, err
        }

        l.notify(Event{
                Kind:   kind,
                ID:     message.ID,
                Text:   text,
                Source: source,
        })

        return message, nil
}

// waiting tells if user messages or reports wait for a request.
func (l *TurnLoop) waiting() bool {
        l.mutex.Lock()
        defer l.mutex.Unlock()

        return len(l.pending) > 0
}

// fail stops the loop with err.
func (l *TurnLoop) fail(err error) {
        l.mutex.Lock()
        l.failure = errors.Join(l.failure, err)
        l.mutex.Unlock()

        select {
        case l.wake <- struct{}{}:
        default:
        }
}

// failed returns the error of a failed write of a report.
func (l *TurnLoop) failed() error {
        l.mutex.Lock()
        defer l.mutex.Unlock()

        return l.failure
}

// turn runs one turn. A failed or interrupted model call ends the turn and
// tells the clients. The messages of the turn stay in the log and the view, and
// the next turn starts only with a new message, so a lasting failure such as a
// needed login does not repeat. Only a failed log write is an error of turn.
func (l *TurnLoop) turn(ctx context.Context) error {
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()

        l.mutex.Lock()
        l.cancel = cancel
        l.mutex.Unlock()

        defer func() {
                l.mutex.Lock()
                l.cancel = nil
                l.mutex.Unlock()
        }()

        l.notify(Event{
                Kind: "busy",
        })

        err := l.chain(ctx)
        if err != nil && !errors.Is(err, errChainLogWrite) {
                l.notify(Event{
                        Kind: "error",
                        Text: turnError(ctx, err),
                })

                err = nil
        }

        l.notify(Event{
                Kind: "idle",
        })

        return err
}

// chain runs the request chain of a turn until the final answer.
func (l *TurnLoop) chain(ctx context.Context) error {
        messages := l.take()
        if len(messages) == 0 {
                return nil
        }

        err := l.memory.Wait(ctx)
        if err != nil {
                return err
        }

        instructions, err := l.agent.Prompt()
        if err != nil {
                return err
        }

        skip := make(map[int]bool, len(messages))
        for _, message := range messages {
                skip[message.ID] = true
        }

        view, err := l.layout.Split(l.memory.ViewLines(skip))
        if err != nil {
                return err
        }

        chain := &Chain{
                client: l.client,
                tools:  l.tools,
                request: ModelRequest{
                        Model:        l.agent.config.Model,
                        Effort:       l.agent.config.Effort,
                        Instructions: instructions,
                        Tools:        promptTools(),
                },
                take: l.take,
        }

        _, err = chain.Run(ctx, turnInput(view, l.state(), messages))

        return err
}

// take returns the user messages that wait for a request, and removes them.
func (l *TurnLoop) take() []Message {
        l.mutex.Lock()
        defer l.mutex.Unlock()

        messages := l.pending
        l.pending = nil

        return messages
}

// state returns the per-turn state: the end of the view, and the date and time
// of the turn.
func (l *TurnLoop) state() string {
        return "</chat>\n\nNow: " + l.clock().Format("Monday 2006-01-02 15:04 -07:00")
}

// turnError returns the text of the error that ended a turn.
func turnError(ctx context.Context, err error) string {
        if ctx.Err() != nil {
                return "interrupted"
        }

        return err.Error()
}

// turnInput returns the input of the first request of a turn, stable parts
// first: a user message for each view message, one with the per-turn state,
// and one for each new message, whole. The view leaves out the leaf lines of
// the new messages.
func turnInput(view []string, state string, messages []Message) []any {
        input := make([]any, 0, len(view)+1+len(messages))
        for _, text := range view {
                input = append(input, modelUserMessage(text))
        }

        input = append(input, modelUserMessage(state))
        for _, message := range messages {
                input = append(input, modelUserMessage(message.Text))
        }

        return input
}