ollie-core/agent/agent.go

package agent

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

	"github.com/simonfxr/pubsub"
	"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
	cfg            agentConfig // per-turn config built from runtime
	agentName      string
	agentsDir      string
	baseLayers     []string // system prompt layers 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

	// Injected session-level dependencies (set at creation, stable for agent lifetime).
	bus            *pubsub.Bus
	log            *olog.Logger
	auditLog       *olog.Logger
	sessionID      string       // the owning session's ID
	startupMessages []string
	readPlanStep   func() string
	saveSession    func()       // trigger debounced persistence
	flushSave      func()       // immediately flush persistence
	turnError      func(ctx context.Context, errType, errMsg string) HookResult // overridable for tests
}

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

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

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

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

// 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 publishes an event on the agent's bus.
func (ag *Agent) emit(ev Event) {
	ag.bus.Publish("event", ev)
}

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

// SetCWD sets the agent's working directory (no validation — caller must validate).
func (ag *Agent) SetCwd(dir string) {
	ag.stateMu.Lock()
	ag.cwd = dir
	ag.stateMu.Unlock()
}

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

// 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.agentName, 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)
			}
		}
	}
}

// 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.SetCwd(dir)
	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 tool_list detects directory changes.
func (ag *Agent) refreshToolListing(listing string) {
	if ag.runtime == nil {
		return
	}
	const marker = "\n# Available Tools\n\n"
	if idx := strings.Index(ag.runtime.Preamble, marker); idx >= 0 {
		ag.runtime.Preamble = ag.runtime.Preamble[:idx]
	}
	if listing != "" {
		ag.runtime.Preamble += marker + listing
	}
}

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

// SaveFull persists the full session state (history + metadata) to the given path.
func (ag *Agent) SaveFull(path, sessionID, cwd, remote string) error {
	if ag.history == nil {
		return fmt.Errorf("no active session")
	}
	return ag.history.saveToFull(path, sessionID, ag.agentName,
		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())
}

// Reactions returns a map of response ID → emoji for all recorded reactions.
func (ag *Agent) Reactions() map[string]string {
	result := make(map[string]string)
	if ag.history == nil {
		return result
	}
	for _, reaction := range ag.history.Reactions {
		result[reaction.ResponseID] = reaction.Emoji
	}
	return result
}

// React records a reaction emoji against the most recent (or specified) response.
func (ag *Agent) React(responseID, emoji string) error {
	if ag.history == nil {
		return fmt.Errorf("no active session")
	}
	category, _, _ := classifyReaction(emoji)
	if category == "unknown" {
		return fmt.Errorf("unsupported reaction: %s", emoji)
	}
	if responseID == "" {
		for i := len(ag.history.messages) - 1; i >= 0; i-- {
			if ag.history.messages[i].Role == "assistant" {
				responseID = ag.history.messages[i].ID
				break
			}
		}
	}
	if responseID == "" {
		return fmt.Errorf("no assistant response to react to")
	}
	found := false
	for i := range ag.history.messages {
		if ag.history.messages[i].Role == "assistant" && ag.history.messages[i].ID == responseID {
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("assistant response not found: %s", responseID)
	}
	reaction := Reaction{ID: NewReactionID(), ResponseID: responseID, Emoji: emoji, Category: category, CreatedAt: time.Now()}
	replaced := false
	for i := range ag.history.Reactions {
		if ag.history.Reactions[i].ResponseID == responseID {
			if ag.history.Reactions[i].Emoji == emoji {
				return nil
			}
			ag.history.Reactions[i] = reaction
			replaced = true
			break
		}
	}
	if !replaced {
		ag.history.Reactions = append(ag.history.Reactions, reaction)
	}
	ag.history.recomputeReactionCounts()
	ag.saveSession()
	return nil
}

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

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