Vectrum
Sign In
247 lines · 6.4 KB · Go
package serve

import (
	"fmt"
	"strings"
	"time"

	"github.com/johannSo/codey/internal/acp"
	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/subagent"
)

// A subagent is shown as a tool call, because a tool call is the only
// container the protocol has for work with a beginning, a running middle and
// an end. The row opens when the child starts, fills in as it works, and
// closes with its report.
//
// The id is prefixed so it cannot collide with a real tool call: the parent's
// own spawn_agent call is a row too, and the two are different things.
const agentCallPrefix = "agent-"

func agentCallID(id string) string { return agentCallPrefix + id }

// watchAgents subscribes this session to its children. Both callbacks arrive
// on the children's own goroutines, so everything they touch is guarded.
func (s *session) watchAgents() {
	if s.agents == nil {
		return
	}
	s.agents.Subscribe(s.onAgentEvent, s.onAgentsChanged)
}

// onAgentEvent is one child saying something.
func (s *session) onAgentEvent(ev subagent.Event) {
	line := agentLine(ev.Agent)
	if line == "" {
		return
	}
	s.emit.agentProgress(ev, line)
}

// onAgentsChanged is the supervisor telling us something moved. Statuses are
// read from the snapshot rather than carried on the event, so this is also
// where a child that has finished gets its row closed.
func (s *session) onAgentsChanged() {
	if s.agents == nil {
		return
	}
	for _, info := range s.agents.Snapshot() {
		if info.Status.Live() {
			continue
		}
		s.emit.agentFinished(info)
	}

	// Wake a turn that is waiting for children to finish. Buffered and
	// non-blocking: a signal nobody is waiting for is not worth keeping, and
	// the waiter re-reads the snapshot anyway.
	select {
	case s.agentsMoved <- struct{}{}:
	default:
	}
}

// agentLine renders one of a child's events as a line for its row. Only the
// things worth watching from outside are rendered: what it said, and what it
// ran. Its usage and its turn boundaries are the parent's business, not the
// reader's.
func agentLine(ev agent.Event) string {
	switch ev.Kind {
	case agent.EventText:
		return strings.TrimSpace(ev.Text) + "\n"
	case agent.EventTextDelta:
		return ev.Text
	case agent.EventToolStart:
		return "\n· " + ev.Tool + "\n"
	default:
		return ""
	}
}

// agentProgress opens a child's row if it is new and appends to it.
func (e *emitter) agentProgress(ev subagent.Event, line string) {
	id := agentCallID(ev.ID)

	e.mu.Lock()
	buf, known := e.running[id]
	if !known {
		buf = &strings.Builder{}
		e.running[id] = buf
	}
	buf.WriteString(line)

	out := e.dueLocked(id, buf)
	e.mu.Unlock()

	if !known {
		e.send(acp.Update{
			SessionUpdate: acp.UpdateToolCall,
			ToolCallID:    id,
			Title:         agentTitle(ev),
			Kind:          acp.KindThink,
			Status:        acp.StatusInProgress,
		})
	}
	if out == "" {
		return
	}
	e.send(acp.Update{
		SessionUpdate: acp.UpdateToolCallUpdate,
		ToolCallID:    id,
		Status:        acp.StatusInProgress,
		Content:       []acp.ToolCallContent{acp.ToolCallText(out)},
	})
}

// agentFinished closes a child's row with its report. It is called on every
// change once the child has ended, so it has to be the one that decides it has
// already been called.
func (e *emitter) agentFinished(info subagent.Info) {
	id := agentCallID(info.ID)

	e.mu.Lock()
	if e.closed == nil {
		e.closed = make(map[string]bool)
	}
	if e.closed[id] {
		e.mu.Unlock()
		return
	}
	e.closed[id] = true
	_, opened := e.running[id]
	e.forgetLocked(id)
	e.mu.Unlock()

	// An agent that finished without ever saying anything — killed while
	// queued, or failed on its first request — has no row yet. Open one, so
	// what happened to it is visible rather than merely absent.
	if !opened {
		e.send(acp.Update{
			SessionUpdate: acp.UpdateToolCall,
			ToolCallID:    id,
			Title:         agentInfoTitle(info),
			Kind:          acp.KindThink,
			Status:        acp.StatusInProgress,
		})
	}

	status := acp.StatusCompleted
	if info.Status != subagent.Done {
		status = acp.StatusFailed
	}
	e.send(acp.Update{
		SessionUpdate: acp.UpdateToolCallUpdate,
		ToolCallID:    id,
		Title:         agentInfoTitle(info),
		Status:        status,
		Content:       []acp.ToolCallContent{acp.ToolCallText(agentOutcome(info))},
	})
}

func agentTitle(ev subagent.Event) string {
	title := fmt.Sprintf("subagent %s (%s)", ev.ID, ev.Preset)
	if ev.Model != "" {
		title += " on " + ev.Model
	}
	return title
}

func agentInfoTitle(info subagent.Info) string {
	title := fmt.Sprintf("subagent %s (%s)", info.ID, info.Preset)
	if info.Model != "" {
		title += " on " + info.Model
	}
	if d := info.Elapsed(); d > 0 {
		title += fmt.Sprintf(" — %s", d.Round(time.Second))
	}
	return title
}

// agentOutcome is what the closed row shows: the report when there is one, and
// why there is not when there is not.
func agentOutcome(info subagent.Info) string {
	var b strings.Builder
	fmt.Fprintf(&b, "%s (%s) %s after %d turns", info.ID, info.Preset, endedWord(info.Status), info.Turns)
	if cost, ok := info.Cost(); ok && cost > 0 {
		fmt.Fprintf(&b, ", $%.4f", cost)
	}
	b.WriteString("\n\n")
	switch {
	case info.Report != "":
		b.WriteString(info.Report)
	case info.Err != "":
		b.WriteString(info.Err)
	default:
		b.WriteString("It reported nothing.")
	}
	return b.String()
}

func endedWord(s subagent.Status) string {
	switch s {
	case subagent.Done:
		return "finished"
	case subagent.Failed:
		return "failed"
	case subagent.Killed:
		return "was stopped"
	default:
		return string(s)
	}
}

// reportBlocks folds the reports this session is still entitled to hand over
// into one message for the parent.
//
// The claim is made here rather than when the report was drained: agent_wait
// may have collected it in between, and putting it in again would charge the
// model twice for something already in its context.
func (s *session) reportBlocks() string {
	if s.agents == nil {
		return ""
	}
	var parts []string
	for _, r := range s.agents.TakeReports() {
		if !s.agents.ClaimReport(r.ID) {
			continue
		}
		parts = append(parts, r.Notification())
	}
	return strings.Join(parts, "\n\n")
}

// agentsLive reports whether any child is still going to do something. A turn
// stays open while one is: their updates belong inside the turn that started
// them, and a report that arrives after the turn closed has nowhere to go.
func (s *session) agentsLive() bool {
	if s.agents == nil {
		return false
	}
	for _, info := range s.agents.Snapshot() {
		if info.Status.Live() {
			return true
		}
	}
	return false
}