ollie/fs/persist.go

package fs

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"ollie/agent"
	"ollie/backend"
	"ollie/prompts"
	coresession "ollie/session"
	"ollie/toolsrv"
)

// KillFromRoot is the public API for callers with *Tree.
func KillFromRoot(root *Tree, id string) {
	Kill(state(root), id)
}

// RenameFromRoot is the public API for callers with *Tree.
func RenameFromRoot(root *Tree, old, new string) error {
	return Rename(state(root), old, new)
}

// --- Persistence ---

func activeSessionsDir(rs *rootState) string {
	return filepath.Join(rs.cfg.SessionsDir, "active")
}

func persistSession(rs *rootState, id string) {
	rs.mu.RLock()
	sess, ok := rs.sessions[id]
	rs.mu.RUnlock()
	if !ok {
		return
	}
	if sess.Core == nil || sess.Core.Agent() == nil {
		// Empty session — no agent to persist. Silently skip.
		return
	}
	dir := activeSessionsDir(rs)
	os.MkdirAll(dir, 0700)
	path := filepath.Join(dir, sess.id+".json")
	ag := sess.Core.Agent()
	if err := ag.SaveFull(path, sess.id, sess.Name(), sess.Core.CWD(), ""); err != nil {
		rs.cfg.Log.Error("persist session %s: %v", id, err)
	}
}

func removePersistedAgent(rs *rootState, id string) {
	path := filepath.Join(activeSessionsDir(rs), id+".json")
	os.Remove(path)
}

func saveAllSessions(rs *rootState) {
	rs.mu.RLock()
	ids := make([]string, 0, len(rs.sessions))
	for id := range rs.sessions {
		ids = append(ids, id)
	}
	rs.mu.RUnlock()
	for _, id := range ids {
		persistSession(rs, id)
	}
}

func restoreAllSessions(rs *rootState) {
	dir := activeSessionsDir(rs)
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}

	type loadedSession struct {
		ps   *agent.PersistedAgent
		name string
	}
	var loaded []loadedSession
	for _, e := range entries {
		if !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		path := filepath.Join(dir, e.Name())
		ps, err := agent.LoadPersistedAgent(path)
		if err != nil {
			rs.cfg.Log.Error("restore session %s: %v", e.Name(), err)
			continue
		}
		loaded = append(loaded, loadedSession{ps: ps, name: e.Name()})
	}

	if len(loaded) == 0 {
		return
	}

	var wg sync.WaitGroup
	for _, ls := range loaded {
		wg.Add(1)
		go func(ls loadedSession) {
			defer wg.Done()
			if err := restoreSession(rs, ls.ps); err != nil {
				rs.cfg.Log.Error("restore session %s: %v", ls.ps.ID, err)
			}
		}(ls)
	}
	wg.Wait()
}

func restoreSession(rs *rootState, ps *agent.PersistedAgent) error {
	cwd := ps.CWD
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	// Profile and agent display name are independent persisted fields. Never
	// use the agent name (or its immutable ID) as a profile fallback: doing so
	// makes restored sessions report an ID such as "eda74c4e8fb4" as the
	// active profile.
	profile := ps.Profile
	if profile == "" {
		profile = "default"
	}
	sessID := ps.ID

	cfg := LoadAgentConfig(rs.cfg.AgentsDir, profile, nil)

	backendName := ps.Backend
	if backendName == "" && cfg != nil && cfg.Backend != "" {
		backendName = cfg.Backend
	}
	be, err := backend.NewWithName(backendName)
	if err != nil {
		return fmt.Errorf("backend: %w", err)
	}
	modelName := ps.Model
	if modelName == "" && cfg != nil && cfg.Model != "" {
		modelName = cfg.Model
	}
	if modelName == "" {
		modelName = os.Getenv("OLLIE_MODEL")
	}
	if modelName != "" {
		be.SetModel(modelName)
	}

	uname := ps.ID
	if uname == "" {
		uname = nextUname(rs)
	}
	var newToolServer func() toolsrv.Runner
	var promptEnv []string
	var proc *toolsrv.Process
	remoteTarget := ps.Remote

	if remoteTarget != "" {
		var spawnErr error
		proc, spawnErr = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{
			SSHTarget: remoteTarget,
			CWD:       cwd,
		})
		if spawnErr != nil {
			return fmt.Errorf("remote spawn: %w", spawnErr)
		}
		newToolServer = func() toolsrv.Runner {
			conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket))
			return conn
		}
		promptEnv = []string{
			"PRIME_CWD=" + cwd,
			"PRIME_PLATFORM=" + proc.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo),
		}
	} else {
		var dialOpts []toolsrv.Option
		if rs.cfg.Yolo {
			dialOpts = append(dialOpts, toolsrv.WithYolo())
		}
		var spawnErr error
		proc, spawnErr = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...)
		if spawnErr != nil {
			return fmt.Errorf("local spawn: %w", spawnErr)
		}
		newToolServer = func() toolsrv.Runner {
			conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket))
			return conn
		}
		promptEnv = agent.PromptEnv(cwd)
	}

	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
	env = append(env, promptEnv...)

	var spOverride string
	if cfg != nil {
		spOverride = cfg.SystemPrompt
	}
	sysPrompt := prompts.ResolveSystemPrompt(spOverride)
	envMap := make(map[string]string)
	for _, e := range env {
		if k, v, ok := strings.Cut(e, "="); ok {
			envMap[k] = v
		}
	}
	opModel := prompts.OperationalModel(true, false, envMap)
	platform := "linux"
	isGitRepo := false
	for _, e := range promptEnv {
		if k, v, ok := strings.Cut(e, "="); ok {
			switch k {
			case "PRIME_PLATFORM":
				platform = v
			case "PRIME_IS_GIT_REPO":
				isGitRepo = v == "true"
			}
		}
	}
	envBlock := prompts.Environment(cwd, platform, isGitRepo, "")

	toolSrv := newToolServer()
	// Load configured tools before building the runtime. BuildRuntime snapshots
	// the loaded-tool list to construct the restored agent's tool definitions.
	if cfg != nil {
		for _, tl := range cfg.AutoLoad {
			if err := loadToolOnRunner(toolSrv, tl); err != nil {
				rs.cfg.Log.Error("session %s: autoLoad tool %q: %v", sessID, tl, err)
			}
		}
	}

	rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, sysPrompt, opModel, envBlock)

	restoredSession := agent.RestoreHistory(ps)

	var sessPtr *Session
	al := NewAgentLog(remoteTarget)
	core := coresession.New(coresession.Config{
		Backend:        be,
		Profile:        profile,
		AgentsDir:      rs.cfg.AgentsDir,
		SessionsDir:    rs.cfg.SessionsDir,
		SessionID:      sessID,
		AgentID:        uname,
		CWD:            cwd,
		Remote:         remoteTarget,
		History:        restoredSession,
		Runtime:        rt,
		NewToolServer:  newToolServer,
		PromptEnvExtra: promptEnv,
		BaseLayers:     []string{sysPrompt, opModel, envBlock},
		Log:            rs.cfg.Sink.NewLogger("core"),
		Output:         NewEventHandler(al),
		ReadPlanStep: func() string {
			if sessPtr == nil || sessPtr.AgentLog() == nil {
				return ""
			}
			al := sessPtr.AgentLog()
			al.mu.RLock()
			data := make([]byte, len(al.plan))
			copy(data, al.plan)
			al.mu.RUnlock()
			return coresession.NextUncheckedStep(data)
		},
	})

	sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
	sess := NewSession(sessID, core, sessionCtx, sessionCancel)
	if ps.SessionName != "" {
		sess.SetName(ps.SessionName)
	}
	sessPtr = sess
	sess.uname = uname
	sess.SetAgentLog(uname, al)
	sess.proc = proc
	// Use the multiplexed connection for both agent execution and 9P management.
	sess.toolsConn = toolSrv

	// Restore agent display name from persisted session.
	ag := sess.Core.Agent()
	if ps.Agent != "" {
		ag.SetName(ps.Agent)
	}

	replayMessagesToLog(sess.AgentLog(), ps.Messages)

	// Pre-load autoLoad tools into the remote.
	if cfg != nil {
		ag := sess.Core.Agent()
		for _, tl := range cfg.AutoLoad {
			if err := sess.loadTool(tl, ag); err != nil {
				rs.cfg.Log.Error("session %s: autoLoad tool %q: %v", sessID, tl, err)
			}
		}
	}

	name := sess.Name()
	rs.mu.Lock()
	rs.sessions[name] = sess
	rs.mu.Unlock()

	toolCount := 0
	if toolSrv != nil {
		if infos, err := toolSrv.ListTools(); err == nil {
			toolCount = len(infos)
		}
	}
	rs.cfg.Log.Info("restored session %s: cfg=%v autoload=%d tools=%d (backend=%s model=%s agent=%s)",
		sessID, cfg != nil, func() int {
			if cfg != nil {
				return len(cfg.AutoLoad)
			}
			return -1
		}(), toolCount, backendName, modelName, profile)
	return nil
}