ollie/session/session.go

package session

import (
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// Config is the configuration for creating a session.
type Config struct {
	Backend        backend.Backend
	ModelName      string
	AgentName      string
	AgentsDir      string
	SessionsDir    string
	SessionID      string
	AgentID        string
	CWD            string
	History        *agent.History
	Runtime        *agent.Runtime
	NewToolServer  func() toolsrv.Runner
	NewBackend     func(string) (backend.Backend, error)
	Log            *olog.Logger
	MaxSteps       int
	ReadPlanStep   func() string
	ListHandlers   map[string]func() []string
	PromptEnvExtra []string
	Remote         string
	BaseLayers     []string
	Output         agent.EventHandler
}

// Session is the concrete session type. It owns session-level state and
// delegates agent operations to its owned agents. Operations that act on
// "the agent" (Agent(), CWD, etc.) use the agent identified by activeAgentName.
type Session struct {
	id         string
	envMu      sync.RWMutex
	env        map[string]string
	plan       []byte
	prevPrompt string

	agents           []*agent.Agent
	activeAgentName  string // name of the agent returned by Agent()
	log              *olog.Logger
	sessionsDir      string
	listHandlers     map[string]func() []string
	remote           string
	mu               sync.RWMutex
	auditLog         *olog.Logger

	saveMu    sync.Mutex
	saveDirty bool
	saveTimer *time.Timer
}

// NewSessionID generates a UUIDv4 session identifier.
func NewSessionID() string {
	b := make([]byte, 16)
	rand.Read(b) //nolint:errcheck
	b[6] = (b[6] & 0x0f) | 0x40 // version 4
	b[8] = (b[8] & 0x3f) | 0x80 // variant 10xx
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
}

// NextUncheckedStep returns the first unchecked step from plan bytes.
func NextUncheckedStep(data []byte) string {
	for _, line := range strings.Split(string(data), "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "- [ ]") {
			return strings.TrimSpace(trimmed[5:])
		}
	}
	return ""
}

var sweepTmpOnce sync.Once

func ollieTmpDir() string {
	if p := os.Getenv("OLLIE_TMP_PATH"); p != "" {
		return p
	}
	return filepath.Join(os.TempDir(), "ollie")
}

func sweepStaleTmpDirs() {
	sweepTmpOnce.Do(func() {
		base := ollieTmpDir()
		os.RemoveAll(base)      //nolint:errcheck
		os.MkdirAll(base, 0700) //nolint:errcheck
	})
}

// New creates a session with an owned agent from the given configuration.
func New(cfg Config) *Session {
	sweepStaleTmpDirs()
	if cfg.ModelName != "" {
		cfg.Backend.SetModel(cfg.ModelName)
	}
	if cfg.NewBackend == nil {
		cfg.NewBackend = backend.NewWithName
	}
	rt := cfg.Runtime
	if rt == nil {
		rt = &agent.Runtime{}
	}
	rt.Backend = cfg.Backend
	if cfg.MaxSteps > 0 {
		rt.MaxSteps = cfg.MaxSteps
	}
	if cfg.SessionID != "" {
		os.MkdirAll(filepath.Join(ollieTmpDir(), cfg.SessionID), 0700) //nolint:errcheck
	}

	log := cfg.Log
	if log == nil {
		log = olog.NewWriter("core", olog.LevelError+1, io.Discard, io.Discard)
	}

	auditLog := log.Sub("audit")

	a := &Session{
		id:           cfg.SessionID,
		env:          make(map[string]string),
		agents:       make([]*agent.Agent, 0, 1),
		log:          log,
		auditLog:     auditLog,
		sessionsDir:  cfg.SessionsDir,
		remote:       cfg.Remote,
		listHandlers: cfg.ListHandlers,
	}

	ag := agent.NewAgent(agent.AgentCfg{
		History:        cfg.History,
		Runtime:        rt,
		AgentName:      cfg.AgentName,
		AgentsDir:      cfg.AgentsDir,
		AgentID:        cfg.AgentID,
		CWD:            paths.ExpandHome(cfg.CWD),
		BaseLayers:     cfg.BaseLayers,
		PromptEnvExtra: cfg.PromptEnvExtra,
		NewToolServer:  cfg.NewToolServer,
		NewBackend:     cfg.NewBackend,
		Output:         cfg.Output,
		Log:            log,
		AuditLog:       auditLog,
		SessionID:      cfg.SessionID,
		StartupMsgs:    rt.Messages,
		ReadPlanStep:   cfg.ReadPlanStep,
		SaveSession:    a.saveSession,
		FlushSave:      a.flushSave,
})

	a.agents = append(a.agents, ag)
	a.activeAgentName = ag.Name()
	ag.SetSessionEnv(a.id)
	return a
}

// Agent returns the currently active agent by name.
// If no active agent is set, returns the first agent (backward compatible).
func (a *Session) Agent() *agent.Agent {
	if len(a.agents) == 0 {
		return nil
	}
	if a.activeAgentName != "" {
		for _, ag := range a.agents {
			if ag.Name() == a.activeAgentName {
				return ag
			}
		}
	}
	// Fallback: first agent or match by ID
	return a.agents[0]
}

// SetActiveAgent sets the name of the active agent. Returns nil if found,
// or an error if no agent with that name exists.
func (a *Session) SetActiveAgent(name string) error {
	for _, ag := range a.agents {
		if ag.Name() == name {
			a.activeAgentName = name
			return nil
		}
	}
	return fmt.Errorf("agent %q not found", name)
}

// FindAgent returns the agent matching the given name or ID, or nil.
func (a *Session) FindAgent(nameOrID string) *agent.Agent {
	for _, ag := range a.agents {
		if ag.Name() == nameOrID || ag.ID() == nameOrID {
			return ag
		}
	}
	return nil
}

// Agents returns the full agent slice.
func (a *Session) Agents() []*agent.Agent { return a.agents }

// AgentAt returns the agent at the given index, or nil.
func (a *Session) AgentAt(idx int) *agent.Agent {
	if idx < 0 || idx >= len(a.agents) {
		return nil
	}
	return a.agents[idx]
}

// AgentCount returns the number of agents.
func (a *Session) AgentCount() int { return len(a.agents) }

// SaveFuncs returns save and flush callbacks suitable for use as
// agent.AgentCfg.SaveSession and FlushSave when adding agents
// to this session.
func (a *Session) SaveFuncs() (save func(), flush func()) {
	return a.saveSession, a.flushSave
}

// AddAgent appends an agent to the session and sets it as the active agent.
func (a *Session) AddAgent(ag *agent.Agent) {
	a.agents = append(a.agents, ag)
	if a.activeAgentName == "" {
		a.activeAgentName = ag.Name()
	}
	ag.SetSessionEnv(a.id)
}

// Close releases resources for this session.
func (a *Session) Close() {
	a.log.Debug("Close() session=%q", a.id)
	a.flushSave()
	for _, ag := range a.agents {
		ag.Close()
	}
	if a.id != "" {
		os.RemoveAll(filepath.Join(ollieTmpDir(), a.id)) //nolint:errcheck
	}
}

// SetEnv stores a session-scoped variable and propagates to all agents.
func (a *Session) SetEnv(key, value string) {
	a.envMu.Lock()
	a.env[key] = value
	a.envMu.Unlock()
	for _, ag := range a.agents {
		ag.SetEnv(key, value)
	}
}

// CWD returns the current working directory for tool execution
// (from the primary agent).
func (a *Session) CWD() string {
	if ag := a.Agent(); ag != nil {
		if c := ag.Cwd(); c != "" {
			return c
		}
	}
	wd, _ := os.Getwd()
	return wd
}

// SetCWD validates and sets the working directory on the primary agent.
func (a *Session) SetCWD(dir string) error {
	dir = paths.ExpandHome(dir)
	if dir != "" {
		if _, err := os.Stat(dir); err != nil {
			return fmt.Errorf("cwd: %w", err)
		}
	}
	if ag := a.Agent(); ag != nil {
		ag.SetCWD(dir)
	}
	return nil
}

// SetSessionID renames the session.
func (a *Session) SetSessionID(newID string) error {
	oldID := a.id
	if oldID == newID {
		return nil
	}
	if a.sessionsDir != "" && oldID != "" {
		for _, suffix := range []string{".json", ".compaction.jsonl"} {
			oldPath := a.activeSessionPath(oldID, suffix)
			if _, err := os.Stat(oldPath); err == nil {
				if err := os.Rename(oldPath, a.activeSessionPath(newID, suffix)); err != nil {
					return fmt.Errorf("rename %s: %w", suffix, err)
				}
			}
		}
	}
	a.id = newID
	for _, ag := range a.agents {
		ag.RenamePreamble(oldID, newID)
	}
	oldTemp := filepath.Join(ollieTmpDir(), oldID)
	newTemp := filepath.Join(ollieTmpDir(), newID)
	if _, err := os.Stat(oldTemp); err == nil {
		os.Rename(oldTemp, newTemp) //nolint:errcheck
	}
	if ag := a.Agent(); ag != nil {
		ag.SetSessionEnv(newID)
	}
	return nil
}

// WaitChange blocks until the named field changes from current.
func (a *Session) WaitChange(ctx context.Context, field, current string) (string, bool) {
	ag := a.Agent()
	if ag == nil {
		return "", false
	}
	if field == agent.WatchState {
		return ag.WaitChange(ctx, field, current)
	}
	// Other fields — read via agent methods, use agent's change signal.
	read := func() string {
		switch field {
		case "usage":
			return ag.UsageStr()
		case "ctxsz":
			return ag.CtxSz()
		case "cwd":
			return a.CWD()
		case "agent":
			return ag.Name()
		}
		return ""
	}
	stop := context.AfterFunc(ctx, func() { ag.BroadcastChange() })
	defer stop()
	for ctx.Err() == nil {
		if v := read(); v != current {
			return v, true
		}
		ag.WaitForChange(ctx)
	}
	return "", false
}

func (a *Session) activeSessionPath(id, suffix string) string {
	return filepath.Join(a.sessionsDir, "active", id+suffix)
}

func (a *Session) saveSession() {
	a.saveMu.Lock()
	a.saveDirty = true
	if a.saveTimer == nil {
		a.saveTimer = time.AfterFunc(2*time.Second, a.flushSave)
	}
	a.saveMu.Unlock()
}

func (a *Session) flushSave() {
	a.saveMu.Lock()
	dirty := a.saveDirty
	a.saveDirty = false
	if a.saveTimer != nil {
		a.saveTimer.Stop()
		a.saveTimer = nil
	}
	a.saveMu.Unlock()
	if !dirty || a.id == "" || a.sessionsDir == "" {
		return
	}
	path := a.activeSessionPath(a.id, ".json")
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		a.log.Error("session save: %v", err)
		return
	}
	if ag := a.Agent(); ag != nil {
		if err := ag.SaveFull(path, a.id, a.CWD(), a.remote); err != nil {
			a.log.Error("session save: %v", err)
		}
	}
}

// SaveSession writes the current session state to the given path.
// Returns nil if there is no agent (empty session).
func (a *Session) SaveSession(path string) error {
	if ag := a.Agent(); ag != nil {
		return ag.SaveFull(path, a.id, a.CWD(), a.remote)
	}
	return nil
}