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

import (
	"context"
	"strings"
	"sync"
	"time"

	"codey/internal/acp"
	"codey/internal/agent"
	"codey/internal/approval"
)

// flushAfter and flushBytes are when buffered text goes out. One notification
// per token would be correct and unusable: T3 Code buffers inbound
// notifications in a queue 32 deep and drops the oldest when it overflows, so
// a fast model would lose the middle of its own sentence. Coalescing is what
// keeps us clear of that, and 50 ms is still below what reads as lag.
const (
	flushAfter = 50 * time.Millisecond
	flushBytes = 512
)

// emitter turns one turn's [agent.Event] stream into session/update
// notifications.
//
// It is called from the agent's own goroutine, one event at a time, and holds
// a mutex anyway: the final flush comes from whichever goroutine ended the
// turn, and that is a different one.
type emitter struct {
	send      func(acp.Update)
	now       func() time.Time
	afterTime time.Duration
	afterSize int

	mu    sync.Mutex
	buf   strings.Builder
	since time.Time

	// anon names tool calls the model did not give an id to — the hosted
	// tools, whose result blocks can arrive after a pause split them from
	// their call. Dispatch is sequential, so the last one started is the one a
	// result with no id belongs to.
	anonN    int
	lastAnon string

	// current is the call most recently announced, which is the one an
	// approval prompt is about: the loop emits tool_start and then asks, one
	// call at a time.
	current string
}

// currentCall is the id of the tool call an approval prompt belongs to.
func (e *emitter) currentCall() string {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.current
}

func newEmitter(send func(acp.Update)) *emitter {
	return &emitter{send: send, now: time.Now, afterTime: flushAfter, afterSize: flushBytes}
}

// handle is the [agent.Agent] observer.
func (e *emitter) handle(ev agent.Event) {
	switch ev.Kind {
	case agent.EventTextDelta:
		e.appendText(ev.Text)

	case agent.EventText:
		// The non-streaming path: the whole message at once.
		e.Flush()
		e.chunk(acp.UpdateAgentMessageChunk, ev.Text)

	case agent.EventToolStart:
		// Anything buffered is said before the tool runs. Otherwise a sentence
		// written just before a thirty-second command would sit in the buffer
		// until the command finished, and the user would watch a tool run for
		// no stated reason.
		e.Flush()
		e.toolCall(ev)

	case agent.EventToolResult:
		e.Flush()
		e.toolResult(ev)

	case agent.EventNotice, agent.EventCompact:
		// The loop talking about the run rather than about the task. It goes
		// in the transcript because there is nowhere else for it to go, and
		// silence here is how a compaction becomes a mystery.
		if ev.Text != "" {
			e.Flush()
			e.chunk(acp.UpdateAgentMessageChunk, ev.Text)
		}
	}
}

func (e *emitter) appendText(text string) {
	if text == "" {
		return
	}
	e.mu.Lock()
	if e.buf.Len() == 0 {
		e.since = e.now()
	}
	e.buf.WriteString(text)
	ready := e.buf.Len() >= e.afterSize || e.now().Sub(e.since) >= e.afterTime
	var out string
	if ready {
		out = e.takeLocked()
	}
	e.mu.Unlock()

	if out != "" {
		e.chunk(acp.UpdateAgentMessageChunk, out)
	}
}

// Flush sends whatever text is buffered. It is safe to call when there is
// none, and must be called once more when the turn ends: the last delta of a
// message is usually short and recent, so nothing else would push it out.
func (e *emitter) Flush() {
	e.mu.Lock()
	out := e.takeLocked()
	e.mu.Unlock()
	if out != "" {
		e.chunk(acp.UpdateAgentMessageChunk, out)
	}
}

func (e *emitter) takeLocked() string {
	out := e.buf.String()
	e.buf.Reset()
	return out
}

func (e *emitter) chunk(kind, text string) {
	e.send(acp.Update{SessionUpdate: kind, Content: acp.TextContent(text)})
}

func (e *emitter) toolCall(ev agent.Event) {
	desc := approval.Describe(ev.Tool, ev.Input)
	id := e.callID(ev, true)

	e.mu.Lock()
	e.current = id
	e.mu.Unlock()

	e.send(acp.Update{
		SessionUpdate: acp.UpdateToolCall,
		ToolCallID:    id,
		Title:         title(desc),
		Kind:          toolKind(ev.Tool),
		Status:        acp.StatusInProgress,
		RawInput:      ev.Input,
	})
}

func (e *emitter) toolResult(ev agent.Event) {
	status := acp.StatusCompleted
	if ev.IsError {
		status = acp.StatusFailed
	}
	e.send(acp.Update{
		SessionUpdate: acp.UpdateToolCallUpdate,
		ToolCallID:    e.callID(ev, false),
		Status:        status,
		Content:       []acp.ToolCallContent{acp.ToolCallText(ev.Result)},
	})
}

func (e *emitter) callID(ev agent.Event, start bool) string {
	if ev.ID != "" {
		return ev.ID
	}
	e.mu.Lock()
	defer e.mu.Unlock()
	if start {
		e.anonN++
		e.lastAnon = "call-" + itoa(e.anonN)
	}
	if e.lastAnon == "" {
		e.lastAnon = "call-0"
	}
	return e.lastAnon
}

// title is the one line a client shows beside the tool row. The summary
// already says what will happen; the detail is the argument that makes it
// specific, and a call without one is better off saying nothing than saying
// an empty pair of quotes.
func title(desc approval.Request) string {
	detail := strings.TrimSpace(desc.Detail)
	if detail == "" {
		return desc.Summary
	}
	if i := strings.IndexByte(detail, '\n'); i >= 0 {
		detail = detail[:i] + " …"
	}
	const max = 120
	if len(detail) > max {
		detail = detail[:max] + " …"
	}
	return desc.Summary + ": " + detail
}

// toolKind maps Codey's tools onto ACP's vocabulary, which a client uses to
// choose an icon and decide what to fold away. An unmapped tool is "other",
// which is a fine thing to be.
func toolKind(tool string) string {
	switch tool {
	case "read_file", "skill":
		return acp.KindRead
	case "glob", "grep":
		return acp.KindSearch
	case "edit_file", "write_file":
		return acp.KindEdit
	case "bash":
		return acp.KindExecute
	case "web_search", "web_fetch":
		return acp.KindFetch
	default:
		return acp.KindOther
	}
}

// itoa avoids pulling strconv in for one call site.
func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

// notifier is what an emitter sends through: the connection, bound to one
// session id.
func notifier(conn *acp.Conn, sessionID string) func(acp.Update) {
	return func(u acp.Update) {
		// A notification that cannot be written means the connection is gone,
		// and the turn will find that out on its next move. There is nowhere
		// useful to report it from here.
		_ = conn.Notify(context.Background(), acp.MethodSessionUpdate,
			acp.SessionNotification{SessionID: sessionID, Update: u})
	}
}