ollie-core/session/agent.go

package session

import (
	"context"
	"os"
	"sync"
	"sync/atomic"

	"github.com/simonfxr/pubsub"
	"ollie/backend"
	olog "ollie/log"
	"ollie/tools"
)

// 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
	newDispatcher  func() tools.Dispatcher
	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 {
	if ag.runtime == nil {
		return nil
	}
	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 == nil || 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 == nil || 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
}