Vectrum
Sign In
85 lines · 2.1 KB · Go
// Package agent holds the turn loop and the conversation state it runs on.
//
// It must not import the ui package: the eval runner in Phase 3 drives this
// loop headless, and anything terminal-shaped here would make that impossible.
package agent

import (
	"sync"
	"time"

	"harness/internal/llm"
	"harness/internal/tools"
)

// ReadRecord lives in tools so the edit tool can query the tracker without
// importing this package.
type ReadRecord = tools.ReadRecord

// State is the conversation history plus the per-session file-read tracking.
type State struct {
	mu       sync.Mutex
	messages []llm.Message
	reads    map[string]ReadRecord
	usage    llm.Usage
	turns    int
}

func NewState() *State {
	return &State{reads: map[string]ReadRecord{}}
}

// MarkRead implements tools.ReadTracker.
func (s *State) MarkRead(path string, mod time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.reads[path] = ReadRecord{At: time.Now(), Mod: mod}
}

// LastRead returns the read record for path, if the file has been read.
func (s *State) LastRead(path string) (ReadRecord, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	r, ok := s.reads[path]
	return r, ok
}

// Append adds a message to the history.
func (s *State) Append(m llm.Message) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.messages = append(s.messages, m)
}

// History returns a copy of the conversation so far.
func (s *State) History() []llm.Message {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]llm.Message(nil), s.messages...)
}

// Len is the number of messages in the history.
func (s *State) Len() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.messages)
}

// AddUsage accumulates the token counts of one response.
func (s *State) AddUsage(u llm.Usage) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.usage.InputTokens += u.InputTokens
	s.usage.OutputTokens += u.OutputTokens
	s.usage.CacheReadInputTokens += u.CacheReadInputTokens
	s.usage.CacheCreationInputTokens += u.CacheCreationInputTokens
	s.turns++
}

// Usage is the running total for the session, and the number of model turns it
// took to get there.
func (s *State) Usage() (llm.Usage, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.usage, s.turns
}