ollie/agent/agent.go

package agent

import (
	"context"
	"fmt"
	"os"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"

	"ollie/backend"
	olog "ollie/log"
	"ollie/toolsrv"
)

// 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
	promptEnvExtra []string // PRIME_* vars for prompt resolution
	newToolServer  func() toolsrv.Runner
	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)
	fifo          Fifo   // prompt queue
	toolCallCount atomic.Int64
	pendingInject atomic.Pointer[string]
	submitMu      sync.Mutex // serializes Submit calls (commands + turns)
	stateMu       sync.RWMutex
	changeMu      sync.Mutex
	changeCond    *sync.Cond

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

// 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()
}

// TaskState returns the current task state, if any.
func (ag *Agent) TaskState() *TaskState {
	if ag.history == nil {
		return nil
	}
	return ag.history.TaskState
}

// 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
}

// 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,
	}
}

// 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 }

// 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.changeMu.Lock()
	ag.changeCond.Broadcast()
	ag.changeMu.Unlock()
}

// 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) {
	done := make(chan struct{})
	context.AfterFunc(ctx, func() {
		ag.changeMu.Lock()
		ag.changeCond.Broadcast()
		ag.changeMu.Unlock()
		close(done)
	})

	ag.changeMu.Lock()
	for {
		var val string
		switch field {
		case WatchState:
			val = ag.State()
		default:
			ag.changeMu.Unlock()
			return "", false
		}
		if val != current {
			ag.changeMu.Unlock()
			return val, true
		}
		if ctx.Err() != nil {
			ag.changeMu.Unlock()
			return "", false
		}
		ag.changeCond.Wait()
		if ctx.Err() != nil {
			ag.changeMu.Unlock()
			return "", false
		}
	}
}

// InitCond initializes the changeCond. Must be called once after construction.
func (ag *Agent) InitCond() {
	ag.changeCond = sync.NewCond(&ag.changeMu)
}

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

// SetOutput sets the agent's event output handler.
// This allows the fs layer to wire up the event handler after agent creation.
func (ag *Agent) SetOutput(h EventHandler) {
	ag.output = h
}

// 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() toolsrv.Runner, conn toolsrv.Runner) {
	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
}

// Event is a typed output event emitted during an agent turn or in response
// to a command.
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"
)

// HasHistory returns true if the agent has an active conversation history.
func (ag *Agent) HasHistory() bool {
	return ag.history != nil
}

// SaveTo saves the current history to the given path.
func (ag *Agent) SaveTo(path, name, cwd string) error {
	if ag.history == nil {
		return fmt.Errorf("no active session")
	}
	return ag.history.saveTo(path, name, ag.Name(), cwd)
}

// Restore restores agent history from a persisted session.
func (ag *Agent) Restore(ps *PersistedAgent) {
	ag.history = RestoreHistory(ps)
}

// ToolCallCount returns the total number of tool calls executed.
func (ag *Agent) ToolCallCount() int64 {
	return ag.toolCallCount.Load()
}

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

// SetEnv stores an environment variable on the execute server.
func (ag *Agent) SetEnv(key, value string) {
	if ag.runtime.ToolServer == nil {
		return
	}
	if srv := ag.runtime.ToolServer; srv != nil {
		if es, ok := srv.(interface{ SetEnv(string, string) }); ok {
			es.SetEnv(key, value)
		}
	}
}

// Close releases agent resources (dispatcher, execute server).
func (ag *Agent) Close() {
	if ag.runtime.ToolServer == nil {
		return
	}
	if srv := ag.runtime.ToolServer; srv != nil {
		if c, ok := srv.(interface{ Close() }); ok {
			c.Close()
		}
	}
}

// SetCWD updates the agent's working directory, preamble references, and dispatcher.
func (ag *Agent) SetCWD(dir string) {
	oldCwd := ag.Cwd()
	ag.stateMu.Lock()
	ag.cwd = dir
	ag.stateMu.Unlock()
	if oldCwd != "" && dir != "" && oldCwd != dir {
		ag.runtime.Preamble = strings.ReplaceAll(ag.runtime.Preamble, oldCwd, dir)
	}
	if ag.runtime.ToolServer != nil {
		if srv := ag.runtime.ToolServer; srv != nil {
			if ws, ok := srv.(interface{ SetCWD(string) }); ok {
				ws.SetCWD(dir)
			}
		}
	}
	ag.notifyChange()
}

// RenamePreamble replaces old references in the preamble with new ones.
func (ag *Agent) RenamePreamble(old, new string) {
	if ag.runtime != nil {
		ag.runtime.Preamble = strings.ReplaceAll(ag.runtime.Preamble, old, new)
	}
}

// refreshToolListing replaces the "# Available Tools" section in the preamble
// with an updated listing. Called when directory contents change.
// Note: This only updates the tool listing, not the full tool documentation.
// Tool documentation requires a full BuildPreamble call (on agent reload).
func (ag *Agent) refreshToolListing(listing string) {
	if ag.runtime == nil {
		return
	}

	// Find the "# Available Tools" section
	const availableToolsMarker = "\n# Available Tools\n\n"
	availableToolsIdx := strings.Index(ag.runtime.Preamble, availableToolsMarker)

	if availableToolsIdx < 0 {
		// No tool listing exists yet
		if listing != "" {
			// Find where to insert - before "# Tool Documentation" or "## Tool discovery"
			const toolDocsMarker = "\n# Tool Documentation\n\n"
			const discoveryMarker = "\n## Tool discovery\n"

			toolDocsIdx := strings.Index(ag.runtime.Preamble, toolDocsMarker)
			discoveryIdx := strings.Index(ag.runtime.Preamble, discoveryMarker)

			var insertIdx int
			var insertBefore string

			if toolDocsIdx >= 0 {
				insertIdx = toolDocsIdx
				insertBefore = toolDocsMarker
			} else if discoveryIdx >= 0 {
				insertIdx = discoveryIdx
				insertBefore = discoveryMarker
			} else {
				// Just append
				ag.runtime.Preamble += availableToolsMarker + listing
				return
			}

			// Insert before the marker
			ag.runtime.Preamble = ag.runtime.Preamble[:insertIdx] + availableToolsMarker + listing + insertBefore + ag.runtime.Preamble[insertIdx+len(insertBefore):]
		}
		return
	}

	// Tool listing exists - find where it ends
	// Look for next section marker after the listing
	afterMarker := ag.runtime.Preamble[availableToolsIdx+len(availableToolsMarker):]

	// Find next section marker
	const toolDocsMarker = "\n# Tool Documentation\n\n"
	const discoveryMarker = "\n## Tool discovery\n"

	var endIdx int
	toolDocsIdx := strings.Index(afterMarker, toolDocsMarker)
	discoveryIdx := strings.Index(afterMarker, discoveryMarker)

	if toolDocsIdx >= 0 {
		endIdx = toolDocsIdx
	} else if discoveryIdx >= 0 {
		endIdx = discoveryIdx
	} else {
		// No markers found - listing goes to end
		endIdx = len(afterMarker)
	}

	// Rebuild preamble with new listing
	beforeListing := ag.runtime.Preamble[:availableToolsIdx+len(availableToolsMarker)]
	afterListing := afterMarker[endIdx:]

	if listing == "" {
		// Remove the entire "# Available Tools" section
		// Need to also remove the preceding "\n" if it exists
		if availableToolsIdx > 0 && ag.runtime.Preamble[availableToolsIdx-1] == '\n' {
			beforeListing = ag.runtime.Preamble[:availableToolsIdx]
		}
		ag.runtime.Preamble = beforeListing + afterListing
	} else {
		ag.runtime.Preamble = beforeListing + listing + afterListing
	}
}

// wireToolsChanged sets the OnToolsChanged callback on the tool server
// so the preamble stays in sync when tools are added/removed from disk.
func (ag *Agent) wireToolsChanged() {
	if ag.runtime == nil || ag.runtime.ToolServer == nil {
		return
	}
	type hasHook interface {
		SetOnToolsChanged(func(string))
	}
	if srv, ok := ag.runtime.ToolServer.(hasHook); ok {
		srv.SetOnToolsChanged(ag.refreshToolListing)
	}
}

// toolsNeedRefresh reports whether the tool registry has changed since the
// last time Tools was populated. If the tool server doesn't support revision
// tracking (e.g. remote Conn), this always returns true (safe fallback).
func (ag *Agent) toolsNeedRefresh() bool {
	type revisionProvider interface {
		ToolRegistryRevision() uint64
	}
	srv, ok := ag.runtime.ToolServer.(revisionProvider)
	if !ok {
		return true // no revision support — always refetch
	}
	rev := srv.ToolRegistryRevision()
	if rev != ag.runtime.toolRegistryRevision {
		ag.runtime.toolRegistryRevision = rev
		return true
	}
	return false
}

// SaveFull persists the full session state (history + metadata) to the given path.
func (ag *Agent) SaveFull(path, sessionID, sessionName, cwd, remote string) error {
	// Switching profiles intentionally clears the conversation history, but
	// the profile and display name still need to be persisted. Save an empty
	// history in that case rather than dropping the entire metadata snapshot.
	history := ag.history
	if history == nil {
		history = &History{}
	}
	return history.saveToFull(path, sessionID, sessionName, ag.id, ag.Name(), ag.profile,
		ag.runtime.Backend.Name(), ag.runtime.Backend.Model(), cwd, remote)
}

// 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 ag.runtime.Preamble != "" {
		msgs = append([]backend.Message{{Role: "system", Content: ag.runtime.Preamble}}, msgs...)
	}
	return msgs
}

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

// 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
}

// CompactionModel returns the configured compaction model name.
func (ag *Agent) CompactionModel() string {
	return ag.runtime.CompactionModel
}

// SetCompactionModel sets the compaction model.
func (ag *Agent) SetCompactionModel(model string) {
	ag.runtime.CompactionModel = model
}

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

// execServer returns the execute server, or nil if unavailable.
func (ag *Agent) execServer() toolsrv.Runner {
	return ag.runtime.ToolServer
}

// 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() toolsrv.Runner {
	return ag.execServer()
}

// 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()
}

// BroadcastChange wakes all goroutines waiting on state changes.
func (ag *Agent) BroadcastChange() {
	ag.changeMu.Lock()
	ag.changeCond.Broadcast()
	ag.changeMu.Unlock()
}

// WaitForChange blocks until a state change is broadcast or ctx is cancelled.
func (ag *Agent) WaitForChange(ctx context.Context) {
	ag.changeMu.Lock()
	if ctx.Err() == nil {
		ag.changeCond.Wait()
	}
	ag.changeMu.Unlock()
}

// Detach detaches the current running process to background.
func (ag *Agent) Detach() bool {
	if srv := ag.execServer(); srv != nil {
		if d, ok := srv.(interface{ Detach() bool }); ok {
			return d.Detach()
		}
	}
	return false
}

// DetachedInfo describes a detached process.
type DetachedInfo struct {
	PID      int
	Command  string
	Started  int64
	Exited   bool
	ExitCode int
}

// ListDetached returns info about all detached processes.
func (ag *Agent) ListDetached() []DetachedInfo {
	srv := ag.execServer()
	if srv == nil {
		return nil
	}
	type listDetacher interface{ ListDetachedRaw() []any }
	ld, ok := srv.(listDetacher)
	if !ok {
		return nil
	}
	raw := ld.ListDetachedRaw()
	out := make([]DetachedInfo, 0, len(raw))
	for _, r := range raw {
		if m, ok := r.(map[string]any); ok {
			di := DetachedInfo{}
			if v, ok := m["pid"].(int); ok {
				di.PID = v
			}
			if v, ok := m["command"].(string); ok {
				di.Command = v
			}
			if v, ok := m["started"].(int64); ok {
				di.Started = v
			}
			if v, ok := m["exited"].(bool); ok {
				di.Exited = v
			}
			if v, ok := m["exit_code"].(int); ok {
				di.ExitCode = v
			}
			out = append(out, di)
		}
	}
	return out
}

// SignalDetached sends a signal to a detached process.
func (ag *Agent) SignalDetached(pid int, signal int) error {
	if srv := ag.execServer(); srv != nil {
		type signaler interface {
			SignalDetached(int, syscall.Signal) error
		}
		if sg, ok := srv.(signaler); ok {
			return sg.SignalDetached(pid, syscall.Signal(signal))
		}
	}
	return fmt.Errorf("no execute server available")
}

// GetDetachedOutput reads output from a detached process.
func (ag *Agent) GetDetachedOutput(pid int) (string, error) {
	if srv := ag.execServer(); srv != nil {
		type outputGetter interface{ GetDetachedOutput(int) (string, error) }
		if og, ok := srv.(outputGetter); ok {
			return og.GetDetachedOutput(pid)
		}
	}
	return "", fmt.Errorf("no execute server available")
}

// DismissDetached removes a finished detached process.
func (ag *Agent) DismissDetached(pid int) bool {
	if srv := ag.execServer(); srv != nil {
		type dismisser interface{ DismissDetached(int) bool }
		if d, ok := srv.(dismisser); ok {
			return d.DismissDetached(pid)
		}
	}
	return false
}