ollie/cmd/olliesrv/internal/agent/agent.go

package agent

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"slices"
	"sync"
	"sync/atomic"

	"ollie/cmd/olliesrv/internal/backend"
	lib9p "ollie/lib9p/client"
	olog "ollie/log"
	toolsrvclient "ollie/toolsrv/client"
	"ollie/util"
)

// Agent holds the state of the current Agent entity (the "agent"
// in the traditional sense). It is swappable: when the user runs /agent,
// a new Agent is built from the new agent config while the session
// host remains stable.
type Agent struct {
	history       *History
	runtime       *Runtime
	name          string // display name (defaults to id)
	profile       string // config profile name (e.g. "default")
	agentsDir     string
	systemPrompt  string // system prompt for /agent reloads
	envBlock      string // environment block for /agent reloads
	newToolServer func() *toolsrvclient.Conn
	newBackend    func(string) (backend.Backend, error)
	currentAction atomic.Pointer[actionHandle]
	warnedContext bool
	resultCache   sync.Map

	// Execution state — owned by the agent, protected by stateMu.
	state         string // "idle", "thinking", "calling: <tool>"
	reply         string // last assistant response
	cwd           string // working directory for tool execution
	id            string // agent identity (unique principal)
	parentID      string // immutable ID of the agent that spawned this agent
	fifo          Fifo   // prompt queue
	Feed          Feed   // streaming input gate (change-detecting)
	toolCallCount atomic.Int64
	pendingInject atomic.Pointer[string]
	submitMu      sync.Mutex // serializes Submit calls (commands + turns)
	stateMu       sync.RWMutex
	signalMu      sync.Mutex
	signalCh      chan struct{} // closed on state change; replaced with a fresh channel

	// Event handler — receives events from this agent.
	output          EventHandler
	log             *olog.Logger
	sessionID       string // the owning session's ID
	startupMessages []string
	save            func()                      // trigger debounced persistence
	flush           func()                      // immediately flush persistence
	onStateChange   func(agentID, state string) // optional callback for state changes

	// Chat log — the rendered output stream, exposed via 9P.
	chatMu       sync.RWMutex
	chatLog      []byte
	chatVers     uint32
	chatCond     *sync.Cond
	chatSignalMu sync.Mutex
	chatSignalCh chan struct{}
	plan         []byte
}

// Backend returns the active backend from the runtime.
func (ag *Agent) Backend() backend.Backend {
	return ag.runtime.Backend
}

// Name returns the agent's display name.
func (ag *Agent) Name() string { return ag.name }

// Profile returns the config profile name (e.g. "default").
func (ag *Agent) Profile() string { return ag.profile }

// Messages returns the conversation history messages.
func (ag *Agent) Messages() []backend.Message {
	if ag.history == nil {
		return nil
	}
	return ag.history.history()
}

// UsageStats holds token usage and cost information.
type UsageStats struct {
	TotalInputTokens         int
	TotalCachedInputTokens   int
	TotalCacheCreationTokens int
	TotalOutputTokens        int
	TotalRequests            int
	Estimated                bool
	LastTurnCostUSD          float64
	SessionCostUSD           float64
	CacheHitRatio            float64
}

// Usage returns the agent's usage statistics.
func (ag *Agent) Usage() *UsageStats {
	if ag.history == nil {
		return nil
	}
	return &UsageStats{
		TotalInputTokens:         ag.history.TotalInputTokens,
		TotalCachedInputTokens:   ag.history.TotalCachedInputTokens,
		TotalCacheCreationTokens: ag.history.TotalCacheCreationTokens,
		TotalOutputTokens:        ag.history.TotalOutputTokens,
		TotalRequests:            ag.history.TotalRequests,
		Estimated:                ag.history.Estimated,
		LastTurnCostUSD:          ag.history.LastTurnCostUSD,
		SessionCostUSD:           ag.history.SessionCostUSD,
		CacheHitRatio:            ag.history.cacheHitRatio(),
	}
}

// SetName changes the agent's display name without reloading config.
func (ag *Agent) SetName(name string) {
	ag.stateMu.Lock()
	ag.name = name
	ag.stateMu.Unlock()
}

// ID returns the agent's unique identity.
func (ag *Agent) ID() string { return ag.id }

// ParentID returns the immutable ID of the agent that spawned this agent.
func (ag *Agent) ParentID() string { return ag.parentID }

// --- Chat log methods ---

// AppendChat appends data to the chat log and notifies stream readers.
func (ag *Agent) AppendChat(data []byte) {
	if len(data) == 0 {
		return
	}
	ag.chatMu.Lock()
	ag.chatLog = append(ag.chatLog, data...)
	ag.chatVers++
	ag.chatMu.Unlock()
	ag.chatCond.Broadcast()
	ag.chatSignalMu.Lock()
	close(ag.chatSignalCh)
	ag.chatSignalCh = make(chan struct{})
	ag.chatSignalMu.Unlock()
}

// EnsureTrailingNewline appends a newline if the log doesn't end with one.
func (ag *Agent) EnsureTrailingNewline() {
	ag.chatMu.Lock()
	if len(ag.chatLog) > 0 && ag.chatLog[len(ag.chatLog)-1] != '\n' {
		ag.chatLog = append(ag.chatLog, '\n')
	}
	ag.chatVers++
	ag.chatMu.Unlock()
}

// ChatMu returns the chat log mutex for external locking (streaming).
func (ag *Agent) ChatMu() *sync.RWMutex { return &ag.chatMu }

// ChatCond returns the condvar for blocking chat readers.
func (ag *Agent) ChatCond() *sync.Cond { return ag.chatCond }

// ChatSignal returns the current chat signal channel (closed on new chat data).
func (ag *Agent) ChatSignal() <-chan struct{} {
	ag.chatSignalMu.Lock()
	ch := ag.chatSignalCh
	ag.chatSignalMu.Unlock()
	return ch
}

// ChatLog returns the raw chat log bytes (caller must hold ChatMu.RLock).
func (ag *Agent) ChatLog() []byte { return ag.chatLog }

// ChatRead returns new chat data since the given offset (base).
// Returns (data, nextBase, error). If no new data, returns empty data.
func (ag *Agent) ChatRead(base string) ([]byte, string, error) {
	var offset int
	if base != "" {
		fmt.Sscanf(base, "%d", &offset)
	} else {
		ag.chatMu.RLock()
		offset = len(ag.chatLog)
		ag.chatMu.RUnlock()
	}
	ag.chatMu.RLock()
	log := ag.chatLog
	if len(log) <= offset {
		ag.chatMu.RUnlock()
		return nil, fmt.Sprintf("%d", offset), nil
	}
	data := make([]byte, len(log)-offset)
	copy(data, log[offset:])
	newOffset := len(log)
	ag.chatMu.RUnlock()
	return data, fmt.Sprintf("%d", newOffset), nil
}

// Plan returns a copy of the plan.
func (ag *Agent) Plan() []byte {
	ag.chatMu.RLock()
	p := make([]byte, len(ag.plan))
	copy(p, ag.plan)
	ag.chatMu.RUnlock()
	return p
}

// SetPlan replaces the plan.
func (ag *Agent) SetPlan(data []byte) {
	ag.chatMu.Lock()
	ag.plan = make([]byte, len(data))
	copy(ag.plan, data)
	ag.chatMu.Unlock()
}

// BackendName returns the name of the active backend.
func (ag *Agent) BackendName() string {
	if ag.runtime.Backend == nil {
		return ""
	}
	return ag.runtime.Backend.Name()
}

// ModelName returns the name of the active model.
func (ag *Agent) ModelName() string {
	if ag.runtime.Backend == nil {
		return ""
	}
	return ag.runtime.Backend.Model()
}

// State returns the agent's current execution state.
func (ag *Agent) State() string {
	ag.stateMu.RLock()
	s := ag.state
	ag.stateMu.RUnlock()
	return s
}

// SetState sets the agent's execution state and notifies waiters.
func (ag *Agent) SetState(state string) {
	ag.stateMu.Lock()
	ag.state = state
	ag.stateMu.Unlock()
	ag.notifyChange()
	if ag.onStateChange != nil {
		ag.onStateChange(ag.id, state)
	}
}

// Reply returns the agent's last assistant response.
func (ag *Agent) Reply() string {
	ag.stateMu.RLock()
	r := ag.reply
	ag.stateMu.RUnlock()
	return r
}

// SetOnStateChange sets a callback invoked whenever agent state changes.
func (ag *Agent) SetOnStateChange(fn func(agentID, state string)) {
	ag.onStateChange = fn
}

// SetReply sets the agent's last response.
func (ag *Agent) setReply(reply string) {
	ag.stateMu.Lock()
	ag.reply = reply
	ag.stateMu.Unlock()
}

// notifyChange wakes all goroutines waiting on state changes.
func (ag *Agent) notifyChange() {
	ag.signalMu.Lock()
	close(ag.signalCh)
	ag.signalCh = make(chan struct{})
	ag.signalMu.Unlock()
}

// SignalCh returns the current signal channel (closed on any change).
func (ag *Agent) SignalCh() <-chan struct{} {
	ag.signalMu.Lock()
	ch := ag.signalCh
	ag.signalMu.Unlock()
	return ch
}

// WaitChange blocks until the agent's state differs from current.
// Returns the new value and true, or ("", false) if ctx is cancelled.
func (ag *Agent) WaitChange(ctx context.Context, field, current string) (string, bool) {
	for {
		// Snapshot the signal channel before reading state.
		ag.signalMu.Lock()
		ch := ag.signalCh
		ag.signalMu.Unlock()

		var val string
		switch field {
		case WatchState:
			val = ag.State()
		case WatchFeed:
			val = ag.Feed.Hash()
			if val == "" {
				val = current // no data yet — block
			}
		default:
			return "", false
		}
		if val != current {
			return val, true
		}
		// Wait for either a state change or context cancellation.
		select {
		case <-ch:
			// Changed — loop to re-check.
		case <-ctx.Done():
			return "", false
		}
	}
}

// emit sends an event to the agent's output handler.
func (ag *Agent) emit(ev Event) {
	if ag.output != nil {
		ag.output(ev)
	}
}

// SetToolServer updates the tool server connection and factory.
// Used when resuming a paused session that was restored without infra.
func (ag *Agent) SetToolServer(newToolServer func() *toolsrvclient.Conn, conn *toolsrvclient.Conn) {
	ag.newToolServer = newToolServer
	if ag.runtime != nil {
		ag.runtime.ToolServer = conn
	}
}

// CWD returns the agent's working directory.
func (ag *Agent) Cwd() string {
	ag.stateMu.RLock()
	c := ag.cwd
	ag.stateMu.RUnlock()
	return c
}

// effectiveCwd returns the agent's cwd, falling back to os.Getwd() if empty.
func (ag *Agent) effectiveCwd() string {
	if c := ag.Cwd(); c != "" {
		return c
	}
	wd, _ := os.Getwd()
	return wd
}

// IsRunning returns true if the agent has an active turn in progress.
func (ag *Agent) IsRunning() bool {
	return ag.currentAction.Load() != nil
}

// Interrupt cancels the current in-progress agent turn.
// Returns true if an action was running and was cancelled.
func (ag *Agent) Interrupt(cause error) bool {
	if h := ag.currentAction.Load(); h != nil {
		h.cancel(cause)
		return true
	}
	return false
}

// Compact runs manual context compaction. Returns an error if the agent is
// running or compaction fails.
func (ag *Agent) Compact(ctx context.Context) error {
	if ag.IsRunning() {
		return fmt.Errorf("cannot compact while agent is running")
	}
	if ag.history == nil {
		return nil
	}
	ag.SetState("compacting")
	_, err := ag.runCompact(ctx, "manual")
	ag.SetState("idle")
	if err != nil {
		return err
	}
	ag.save()
	return nil
}

// Clear resets the agent history. Returns an error if the agent is running.
func (ag *Agent) Clear() error {
	if ag.IsRunning() {
		return fmt.Errorf("cannot clear while agent is running")
	}
	ag.history = nil
	return nil
}

// SwitchProfile loads a new agent profile by name, replacing the runtime
// and clearing history. Returns an error if the agent is running or the
// profile cannot be loaded.
func (ag *Agent) SwitchProfile(name string) error {
	if ag.IsRunning() {
		return fmt.Errorf("cannot switch agent while agent is running")
	}
	cfgPath := AgentConfigPath(ag.agentsDir, name)
	f, err := os.Open(cfgPath)
	if err != nil {
		return fmt.Errorf("agent %q: %w", name, err)
	}
	cfg, err := Load(f)
	f.Close()
	if err != nil {
		return fmt.Errorf("agent %q: %w", name, err)
	}
	disp := ag.newToolServer()
	if disp != nil && ag.id != "" {
		disp.SetAgentID(ag.id)
	}
	env := []string{"OLLIE_SESSION_ID=" + ag.sessionID, "OLLIE_UNAME=" + ag.id}
	rt := BuildRuntime(cfg, disp, ag.cwd, env, ag.systemPrompt, ag.envBlock)
	if cfg.Backend != "" {
		newBe, err := ag.newBackend(cfg.Backend)
		if err != nil {
			return fmt.Errorf("backend %q: %w", cfg.Backend, err)
		}
		if cfg.Model != "" {
			newBe.SetModel(cfg.Model)
		}
		rt.Backend = newBe
	} else {
		rt.Backend = ag.runtime.Backend
		if cfg.Model != "" {
			rt.Backend.SetModel(cfg.Model)
		}
	}
	ag.runtime = rt
	ag.profile = name
	ag.history = nil
	ag.save()
	ag.flush()
	ag.notifyChange()
	return nil
}

// Event is a typed output event emitted during an agent turn or in response
// to a command.
// Event carries a single piece of output from the agent loop to consumers
// (frontends, loggers, the turn orchestrator in turn.go).
//
// The Role field determines the event semantics:
//
//	Role          Name       Content                     Emitted by
//	────────────  ─────────  ──────────────────────────  ──────────
//	"user"        —          user input text             turn.go (before run)
//	"assistant"   —          streamed LLM text chunk     loop.go (streamResponse)
//	"reasoning"   —          <think>…</think> chunks    loop.go (streamResponse)
//	"call"        tool name  JSON args                   loop.go (execOne)
//	"tool"        tool name  result text (may stream)    loop.go (execOne)
//	"usage"       —          "in out est cost cached creation" loop.go (streamResponse)
//	"state"       —          "thinking"|"compacting"|…   loop.go / turn.go
//	"limitretry"  —          —                           loop.go (rate limit hit)
//	"retry"       —          "HH:MM:SS" countdown       loop.go (retryCountdown)
//	"maxsteps"    —          step count as string        loop.go (budget exhausted)
//	"error"       —          error message               turn.go / loop.go
//	"info"        —          informational text          turn.go (compaction, cost)
//
// turn.go intercepts events before forwarding to the external handler:
//   - "assistant" → accumulates reply text
//   - "call"      → updates agent state to "calling: <name>"
//   - "tool"      → logs result
//   - "state"     → calls ag.SetState()
//   - "limitretry"→ sets state to "limitretry"
//   - "usage"     → parses token counts, updates history
//   - "error"     → logs
type Event struct {
	Role         string
	Name         string
	Content      string
	ResponseID   string
	OutputFormat string
}

// EventHandler receives events from the agent.
type EventHandler func(Event)

// actionHandle holds the cancel function for the current agent turn.
type actionHandle struct {
	cancel context.CancelCauseFunc
}

// WatchField names supported by Agent.WaitChange.
const (
	WatchState = "state"
	WatchFeed  = "feed"
)

// SetSessionEnv injects session env vars into the execute server.
func (ag *Agent) SetSessionEnv(sessionID string) {
	if ag.runtime == nil || ag.runtime.ToolServer == nil {
		return
	}
	ag.runtime.ToolServer.SetEnv("OLLIE_SESSION_ID", sessionID)
	if ag.id != "" {
		ag.runtime.ToolServer.SetAgentID(ag.id)
	}
	if u := os.Getenv("USER"); u != "" {
		ag.runtime.ToolServer.SetEnv("USER", u)
	}
}

// SetEnv stores an environment variable on the execute server.
func (ag *Agent) SetEnv(key, value string) {
	ag.runtime.ToolServer.SetEnv(key, value)
}

// Close releases agent resources (dispatcher, execute server).
func (ag *Agent) Close() {
	if ag.runtime.ToolServer != nil {
		ag.runtime.ToolServer.Close()
	}
}

// SetCWD updates the agent's working directory, preamble environment section, and dispatcher.
func (ag *Agent) SetCWD(dir string) {
	ag.stateMu.Lock()
	ag.cwd = dir
	ag.stateMu.Unlock()
	isGitRepo := util.IsGitRepo(dir)
	ag.envBlock = EnvironmentBlock(dir, runtime.GOOS, isGitRepo, "")
	ag.runtime.Preamble.Set(SectionEnv, ag.envBlock)
	if ag.runtime.ToolServer != nil {
		ag.runtime.ToolServer.SetCWD(dir)
	}
	ag.notifyChange()
}

// RefreshTools forces a refresh of the tool listing in the agent's preamble.
func (ag *Agent) RefreshTools() {
	if infos, err := ag.runtime.ToolServer.ListTools(); err == nil {
		ag.runtime.Preamble.Set(SectionTools, RenderTools(infos))
	}
}

// SetToolsPreamble updates the tools section of the preamble.
func (ag *Agent) SetToolsPreamble(content string) {
	ag.runtime.Preamble.Set(SectionTools, content)
}

// CtxSz returns a human-readable context size string.
func (ag *Agent) CtxSz() string {
	if ag.history == nil {
		return "no active session"
	}
	ctxLen := ag.runtime.Backend.ContextLength(context.Background())
	if ctxLen <= 0 {
		ctxLen = defaultContextLength
	}
	estimated := ag.history.estimateTokens()
	pct := estimated * 100 / ctxLen
	return fmt.Sprintf("%d / %d (%d%%)", estimated, ctxLen, pct)
}

// CostStr returns formatted cost information.
func (ag *Agent) CostStr() string {
	if ag.history == nil {
		return "no active session"
	}
	return fmt.Sprintf("costLast=$%.4f\ncostSession=$%.4f\n",
		ag.history.LastTurnCostUSD, ag.history.SessionCostUSD)
}

// UsageStr returns formatted usage information.
func (ag *Agent) UsageStr() string {
	if ag.history == nil {
		return "no active session"
	}
	str := fmt.Sprintf("%d in, %d out, %d requests",
		ag.history.TotalInputTokens, ag.history.TotalOutputTokens,
		ag.history.TotalRequests)
	if ag.history.TotalCachedInputTokens > 0 {
		str += fmt.Sprintf(", %d cached", ag.history.TotalCachedInputTokens)
	}
	if ag.history.Estimated {
		str += " [estimated]"
	}
	return str
}

// Context returns the full message context (system prompt + history).
func (ag *Agent) Context() []backend.Message {
	var msgs []backend.Message
	if ag.history != nil {
		msgs = slices.Clone(ag.history.history())
	}
	if preamble := ag.runtime.PreambleString(); preamble != "" {
		msgs = append([]backend.Message{{Role: "system", Content: preamble}}, msgs...)
	}
	return msgs
}

// SystemPrompt returns the rendered system prompt.
func (ag *Agent) SystemPrompt() string {
	return ag.runtime.PreambleString()
}

// GenParams returns the current generation parameters.
func (ag *Agent) GenParams() backend.GenerationParams {
	return ag.runtime.GenParams
}

// SetGenParams sets the generation parameters.
func (ag *Agent) SetGenParams(params backend.GenerationParams) {
	ag.runtime.GenParams = params
}

// ListModels returns available models from the backend.
func (ag *Agent) ListModels() []string {
	return ag.runtime.Backend.Models(context.Background())
}

// ToolServer returns the tool execution server, or nil if unavailable.
// Exported for use by the 9P filesystem layer to sync tool registries.
func (ag *Agent) ToolServer() *toolsrvclient.Conn {
	return ag.runtime.ToolServer
}

// SetRuntime replaces the agent's runtime (used on session resume).
func (ag *Agent) SetRuntime(rt *Runtime) {
	ag.runtime = rt
}

// Queue pushes a prompt onto the agent's FIFO.
func (ag *Agent) Queue(prompt string) {
	ag.fifo.Push(prompt)
}

// PopQueue pops the next prompt from the FIFO.
func (ag *Agent) PopQueue() (string, bool) {
	return ag.fifo.Pop()
}

// FeedWrite stores data in the feed and signals waiters.
func (ag *Agent) FeedWrite(data []byte) {
	ag.Feed.Set(data)
	ag.notifyChange()
}

// ConsumeFeed reads from the agent's feed file in a loop via 9P.
// Blocks until feed changes, submits the new data as a prompt.
// Exits when ctx is cancelled or on read error.
func ConsumeFeed(ctx context.Context, ag *Agent) {
	path := fmt.Sprintf("session/%s/agent/%s/feed", ag.sessionID, ag.id)
	client, err := lib9p.DialDefault()
	if err != nil {
		ag.log.Error("feed consumer: dial: %v", err)
		return
	}
	go func() {
		<-ctx.Done()
		client.Close()
	}()
	for ctx.Err() == nil {
		data, err := client.Read(path)
		if err != nil {
			return
		}
		if len(data) == 0 {
			continue
		}
		ag.Submit(ctx, string(data))
		ag.EnsureTrailingNewline()
	}
}

// AgentParams holds the runtime dependencies for constructing a new Agent.
// Not to be confused with AgentConfig, which is the on-disk JSON schema.
type AgentParams struct {
	ID            string // unique agent identity (uname)
	SessionID     string // id of session the agent belongs to
	Profile       string // config profile name (e.g. "default" → agents/default.json)
	ParentID      string // immutable ID of the agent that spawned this agent
	History       *History
	Runtime       *Runtime
	AgentsDir     string
	Cwd           string // working directory for tool execution
	SystemPrompt  string
	EnvBlock      string
	NewToolServer func() *toolsrvclient.Conn
	NewBackend    func(string) (backend.Backend, error)
	Log           *olog.Logger
	StartupMsgs   []string
	Save          func()
	Flush         func()
}

// NewAgent constructs an Agent from the given configuration.
// Post-construction wiring (SetSessionEnv) is handled internally —
// no additional calls are required after construction.
func NewAgent(cfg AgentParams) *Agent {
	ag := &Agent{
		history:         cfg.History,
		runtime:         cfg.Runtime,
		profile:         cfg.Profile,
		name:            cfg.ID, // display name defaults to uname
		agentsDir:       cfg.AgentsDir,
		id:              cfg.ID,
		parentID:        cfg.ParentID,
		cwd:             cfg.Cwd,
		systemPrompt:    cfg.SystemPrompt,
		envBlock:        cfg.EnvBlock,
		newToolServer:   cfg.NewToolServer,
		newBackend:      cfg.NewBackend,
		log:             cfg.Log,
		sessionID:       cfg.SessionID,
		startupMessages: cfg.StartupMsgs,
		save:            cfg.Save,
		flush:           cfg.Flush,
		state:           "idle",
	}
	ag.signalCh = make(chan struct{})
	ag.chatSignalCh = make(chan struct{})
	ag.chatCond = sync.NewCond(ag.chatMu.RLocker())
	ag.initChatHandler()
	ag.SetSessionEnv(cfg.SessionID)
	return ag
}