ollie-9p/session/session.go

package session

import (
	"context"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"ollie/pkg/agent"
	"ollie/pkg/backend"
	"ollie/pkg/config"
	olog "ollie/pkg/log"
	"ollie/pkg/paths"
	"ollie/pkg/remote"
	"ollie/pkg/tools"
	"ollie/pkg/tools/execute"
	"olliesrv/fs"
	"olliesrv/prompts"
)

// Session holds all state for one agent session.
type Session struct {
	mu         sync.RWMutex
	id         string
	uname      string // immutable user principal (numeric UID), set at creation
	Core       agent.Core
	Ctx        context.Context
	cancel     context.CancelFunc
	log        []byte
	logVers    uint32
	ChatOffset int
	plan       []byte
	prevPrompt []byte          // last submitted prompt; overwritten on each new submission
	peers      map[string]bool // peer session IDs (bidirectional links)
	remote     string          // SSH target for remote execution (empty = local)
	mount      *mountState     // non-nil for remote sessions that need a local mount

	// Cached model list (expensive API call; refreshed every 24h).
	modelsMu      sync.Mutex
	modelsCache   string
	modelsCacheAt time.Time
}

func NewSession(id string, core agent.Core, ctx context.Context, cancel context.CancelFunc) *Session {
	sess := &Session{id: id, Core: core, Ctx: ctx, cancel: cancel}
	sess.startEventLog()
	return sess
}

const modelsCacheTTL = 24 * time.Hour

// CachedListModels returns the model list, using a 24h cache to avoid
// repeated expensive API calls.
// NOTE: Must NOT be called while holding sess.mu (caller content() holds RLock).
// Uses its own modelsMu to avoid deadlock.
func (sess *Session) CachedListModels() string {
	sess.modelsMu.Lock()
	if sess.modelsCache != "" && time.Since(sess.modelsCacheAt) < modelsCacheTTL {
		result := sess.modelsCache
		sess.modelsMu.Unlock()
		return result
	}
	sess.modelsMu.Unlock()

	// Cache miss — fetch and store.
	result := sess.Core.ListModels()
	sess.modelsMu.Lock()
	sess.modelsCache = result
	sess.modelsCacheAt = time.Now()
	sess.modelsMu.Unlock()
	return result
}

// InvalidateModelsCache clears the cached model list.
func (sess *Session) InvalidateModelsCache() {
	sess.modelsMu.Lock()
	sess.modelsCache = ""
	sess.modelsCacheAt = time.Time{}
	sess.modelsMu.Unlock()
}

// HasPeer reports whether peerID is a registered peer of this session.
func (sess *Session) HasPeer(peerID string) bool {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return sess.peers[peerID]
}

func (sess *Session) RunnableID() string { return sess.id }
func (sess *Session) Uname() string      { return sess.uname }

func (sess *Session) Cancel() {
	sess.cancel()
}

func (sess *Session) Interrupt() {
	sess.Core.Interrupt(agent.ErrInterrupted)
}

// AppendLog appends data to the session's log and bumps the version.
func (sess *Session) AppendLog(data []byte) {
	if len(data) == 0 {
		return
	}
	sess.mu.Lock()
	sess.log = append(sess.log, data...)
	sess.logVers++
	sess.mu.Unlock()
}

// EnsureTrailingNewline appends a newline if the log doesn't already end with one.
func (sess *Session) EnsureTrailingNewline() {
	sess.mu.Lock()
	if len(sess.log) > 0 && sess.log[len(sess.log)-1] != '\n' {
		sess.log = append(sess.log, '\n')
	}
	sess.logVers++
	sess.mu.Unlock()
}

// startEventLog subscribes to the agent's "event" bus topic and writes
// all events to the session chat log.
func (sess *Session) startEventLog() {
	streamingRole := "" // tracks current streaming role ("assistant", "reasoning", or "tool")

	sess.Core.Bus().Subscribe("event", func(ev agent.Event) {
		switch ev.Role {
		case "assistant", "reasoning":
			if streamingRole != ev.Role {
				if streamingRole != "" {
					sess.AppendLog([]byte("\n"))
				}
				sess.AppendLog([]byte("[" + ev.Role + "]\n"))
				if ev.Role == "assistant" {
					sess.mu.Lock()
					sess.ChatOffset = len(sess.log)
					sess.mu.Unlock()
				}
				streamingRole = ev.Role
			}
			sess.AppendLog(FormatEvent(ev))

		case "tool":
			// Tool events: empty content = header-only (stream start),
			// non-empty while streaming = chunk, non-empty without streaming = full result.
			if streamingRole == "tool" {
				if ev.Content == "" {
					// End of stream.
					sess.AppendLog([]byte("\n"))
					streamingRole = ""
				} else {
					sess.AppendLog([]byte(ev.Content))
				}
				return
			}
			if streamingRole != "" {
				sess.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			if ev.Content == "" {
				// Stream start: write header, enter streaming mode.
				sess.AppendLog([]byte("[tool:" + ev.Name + "]\n"))
				streamingRole = "tool"
			} else {
				// Non-streamed tool result (fast tool, no streaming happened).
				sess.AppendLog(FormatEvent(ev))
			}

		case "call":
			if streamingRole != "" {
				sess.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			sess.AppendLog(FormatEvent(ev))

		default:
			if streamingRole != "" {
				sess.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			sess.AppendLog(FormatEvent(ev))
		}
	})
}

// LogInfo returns the current log length and version atomically.
func (sess *Session) LogInfo() (length int, vers uint32) {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return len(sess.log), sess.logVers
}

// Mu returns the session's mutex for external synchronization (e.g. D-Bus adapter).
func (sess *Session) Mu() *sync.RWMutex { return &sess.mu }

// Log returns the raw chat log bytes. Caller must hold Mu().RLock().
func (sess *Session) Log() []byte { return sess.log }

// Plan returns the plan bytes. Caller must hold Mu().RLock().
func (sess *Session) Plan() []byte { return sess.plan }

// SetPlan sets the plan. Caller must hold Mu().Lock().
func (sess *Session) SetPlan(p []byte) { sess.plan = p }

// PrevPrompt returns the last submitted prompt. Caller must hold Mu().RLock().
func (sess *Session) PrevPrompt() []byte { return sess.prevPrompt }

var sessionStoreOrder = []string{"new", "idx", "ls", "kill", "sh", "b", "bfg", "bbg", "cleanup"}

// FileMode returns the mode for a fixed session file,
// or 0 and false if the name is not a fixed file.
func FileMode(name string) (os.FileMode, bool) {
	m, ok := fs.Perms[fs.PathSessions].Files[name]
	return m, ok
}

// ManagerConfig holds the dependencies for a Manager.
type ManagerConfig struct {
	AgentsDir   string
	SessionsDir string
	Log         *olog.Logger
	Sink        *olog.Sink
	ReadFile    func(string) ([]byte, error)
	MkdirAll    func(string, os.FileMode) error
	// NewCore, if non-nil, replaces the default backend.New + agent.NewAgentCore
	// path. It receives the session ID, agent name, and cwd, and returns a Core.
	NewCore func(sessionID, agentName, cwd string) (agent.Core, error)
	// Strict rejects inline code steps; only tool steps are allowed.
	Strict bool
	// Yolo skips the landrun sandbox.
	Yolo bool
	// NoMount disables the per-session FUSE mount (e.g. when listening on TCP).
	NoMount bool
	// Enable9P indicates the 9P listener is active (for operational model injection).
	Enable9P bool
	// EnableDBus indicates the D-Bus adapter is active (for operational model injection).
	EnableDBus bool
	// InvalidateModels clears the model cache, forcing a refresh.
	InvalidateModels func()
	// OnSessionCreated is called after a new session is added to the manager.
	// Receives the session ID and the Session pointer.
	OnSessionCreated func(id string, sess *Session)
	// OnSessionKilled is called after a session is removed from the manager.
	OnSessionKilled func(id string)
	// OnSessionRenamed is called after a session is renamed.
	OnSessionRenamed func(oldID, newID string)
}

// Manager manages session lifecycle and exposes sessions as a Tree.
type Manager struct {
	tree     *fs.Tree
	cfg      ManagerConfig
	mu       sync.RWMutex
	sessions map[string]*Session
	nextUID  atomic.Uint32 // incrementing principal counter
}

// Tree returns the Tree view of the session namespace.
func (s *Manager) Tree() *fs.Tree { return s.tree }

// nextUname generates the next uname atomically.
func (s *Manager) nextUname() string {
	return fmt.Sprintf("%d", s.nextUID.Add(1))
}

func NewManager(cfg ManagerConfig) *Manager {
	if cfg.ReadFile == nil {
		cfg.ReadFile = os.ReadFile
	}
	if cfg.MkdirAll == nil {
		cfg.MkdirAll = os.MkdirAll
	}
	ss := &Manager{
		cfg:      cfg,
		sessions: make(map[string]*Session),
	}
	ss.nextUID.Store(9999)
	ss.restoreAllSessions()
	ss.tree = fs.NewTree(nil, 0,
		fs.WithStat(func(_ []string, name string) (os.FileInfo, error) { return ss.stat(name) }),
		fs.WithOpener(func(_ []string, name string) (fs.File, error) { return ss.openEntry(name) }),
		fs.WithLister(func(_ []string) ([]os.DirEntry, error) { return ss.list() }),
		fs.WithReaddir(func(_ []string, name string) ([]os.DirEntry, error) { return ss.Readdir(name) }),
		fs.WithDeleter(func(_ []string, name string) error { return ss.del(name) }),
		fs.WithCreator(func(_ []string, name string, _ os.FileMode) error { return ss.create(name) }),
		fs.WithRenamer(func(_ []string, old, new string) error { return ss.renameSession(old, new) }),
	)
	return ss
}

// AddSession inserts a pre-built session into the fs.
func (s *Manager) AddSession(sess *Session) {
	s.mu.Lock()
	s.sessions[sess.RunnableID()] = sess
	s.mu.Unlock()
}

// List returns all root-level entries (fixed files + session directories).
func (s *Manager) List() ([]os.DirEntry, error) {
	return s.list()
}

func (s *Manager) list() ([]os.DirEntry, error) {
	entries := make([]os.DirEntry, 0, len(sessionStoreOrder))
	for _, name := range sessionStoreOrder {
		entries = append(entries, fs.FileEntry(name, fs.Perms[fs.PathSessions].Files[name]))
	}
	s.mu.RLock()
	for id := range s.sessions {
		entries = append(entries, fs.DirEntry(id, fs.Perms[fs.PathSessionDir].DirMode))
	}
	s.mu.RUnlock()
	return entries, nil
}

// Readdir lists entries in a subdirectory (e.g. "{id}", "{id}/peer").
func (s *Manager) Readdir(name string) ([]os.DirEntry, error) {
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	sess := s.Session(sessID)
	if sess == nil {
		return nil, fmt.Errorf("session not found: %s", sessID)
	}
	// {id} — list session files + peer/
	if len(parts) == 1 {
		sfs, err := s.openStore(sess)
		if err != nil {
			return nil, err
		}
		entries, err := sfs.List()
		if err != nil {
			return nil, err
		}
		// Append peer/ and proc/ directories
		entries = append(entries, fs.DirEntry("peer", 0755))
		entries = append(entries, fs.DirEntry("proc", 0755))
		return entries, nil
	}
	// {id}/peer — list peer session IDs
	if len(parts) == 2 && parts[1] == "peer" {
		peers, _ := s.PeerList(sessID)
		entries := make([]os.DirEntry, len(peers))
		for i, p := range peers {
			entries[i] = fs.FileEntry(p, 0666)
		}
		return entries, nil
	}
	// {id}/proc — list detached process PIDs
	if len(parts) == 2 && parts[1] == "proc" {
		procs := sess.Core.ListDetached()
		entries := make([]os.DirEntry, len(procs))
		for i, p := range procs {
			entries[i] = fs.FileEntry(fmt.Sprintf("%d", p.PID), 0666)
		}
		return entries, nil
	}
	return nil, fmt.Errorf("%s: not a directory", name)
}

// --- Peer Management ---

// PeerAdd creates a bidirectional peer link between two sessions.
func (s *Manager) PeerAdd(sessID, peerID string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}
	if !peerOk {
		return fmt.Errorf("peer session not found: %s", peerID)
	}
	if sessID == peerID {
		return fmt.Errorf("cannot peer a session with itself")
	}

	sess.mu.Lock()
	if sess.peers == nil {
		sess.peers = make(map[string]bool)
	}
	sess.peers[peerID] = true
	sess.mu.Unlock()

	peer.mu.Lock()
	if peer.peers == nil {
		peer.peers = make(map[string]bool)
	}
	peer.peers[sessID] = true
	peer.mu.Unlock()
	return nil
}

// PeerRemove removes a bidirectional peer link between two sessions.
func (s *Manager) PeerRemove(sessID, peerID string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}

	sess.mu.Lock()
	delete(sess.peers, peerID)
	sess.mu.Unlock()

	if peerOk {
		peer.mu.Lock()
		delete(peer.peers, sessID)
		peer.mu.Unlock()
	}
	return nil
}

// PeerList returns the peer session IDs for a session.
func (s *Manager) PeerList(sessID string) ([]string, error) {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	s.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("session not found: %s", sessID)
	}
	sess.mu.RLock()
	result := make([]string, 0, len(sess.peers))
	for id := range sess.peers {
		result = append(result, id)
	}
	sess.mu.RUnlock()
	return result, nil
}

// PeerSubmit sends a prompt to a peer session.
func (s *Manager) PeerSubmit(sessID, peerID, prompt string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}
	if !peerOk {
		return fmt.Errorf("peer session not found: %s", peerID)
	}
	sess.mu.RLock()
	isPeer := sess.peers[peerID]
	sess.mu.RUnlock()
	if !isPeer {
		return fmt.Errorf("%s is not a peer of %s", peerID, sessID)
	}
	go peer.Core.Submit(peer.Ctx, prompt)
	return nil
}

func (s *Manager) stat(name string) (os.FileInfo, error) {
	// Top-level fixed files (new, idx, sh, etc.)
	if mode, ok := fs.Perms[fs.PathSessions].Files[name]; ok {
		return &fs.SyntheticFileInfo{Name_: name, Mode_: mode}, nil
	}
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	s.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("%s: not found", name)
	}
	// Session directory: {id}
	if len(parts) == 1 {
		return &fs.SyntheticFileInfo{Name_: sessID, Mode_: fs.Perms[fs.PathSessionDir].DirMode, IsDir_: true}, nil
	}
	// Peer directory: {id}/peer
	if parts[1] == "peer" {
		if len(parts) == 2 {
			return &fs.SyntheticFileInfo{Name_: "peer", Mode_: 0755, IsDir_: true}, nil
		}
		// Peer file: {id}/peer/{peer-id}
		peerID := parts[2]
		sess.mu.RLock()
		isPeer := sess.peers[peerID]
		sess.mu.RUnlock()
		if !isPeer {
			return nil, fmt.Errorf("%s: not found", name)
		}
		return &fs.SyntheticFileInfo{Name_: peerID, Mode_: 0666}, nil
	}
	// Proc directory: {id}/proc
	if parts[1] == "proc" {
		if len(parts) == 2 {
			return &fs.SyntheticFileInfo{Name_: "proc", Mode_: 0755, IsDir_: true}, nil
		}
		// Proc file: {id}/proc/{pid} — read-only output
		return &fs.SyntheticFileInfo{Name_: parts[2], Mode_: 0444}, nil
	}
	// Session file: {id}/{file}
	sfs, err := s.openStore(sess)
	if err != nil {
		return nil, err
	}
	return sfs.Stat(parts[1])
}

func (s *Manager) openEntry(name string) (fs.File, error) {
	notBlocking := func(context.Context, string) ([]byte, string, error) {
		return nil, "", fmt.Errorf("blocking read not supported")
	}

	// Peer file: {id}/peer/{peer-id} — write submits prompt to peer
	parts := strings.SplitN(name, "/", 3)

	// Proc file: {id}/proc/{pid} — read returns ring buffer output
	if len(parts) == 3 && parts[1] == "proc" {
		sessID := parts[0]
		sess := s.Session(sessID)
		if sess == nil {
			return nil, fmt.Errorf("%s: not found", name)
		}
		pidStr := parts[2]
		pid, err := strconv.Atoi(pidStr)
		if err != nil {
			return nil, fmt.Errorf("invalid pid: %s", pidStr)
		}
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: pidStr, Mode_: 0444}, nil
			},
			ReadFn: func() ([]byte, error) {
				output, err := sess.Core.GetDetachedOutput(pid)
				if err != nil {
					return nil, err
				}
				return []byte(output), nil
			},
			WriteFn:        func([]byte) error { return fmt.Errorf("read-only") },
			BlockingReadFn: notBlocking,
		}, nil
	}

	if len(parts) == 3 && parts[1] == "peer" {
		sessID := parts[0]
		peerID := parts[2]
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: peerID, Mode_: 0666}, nil
			},
			ReadFn: func() ([]byte, error) {
				return nil, nil // empty read
			},
			WriteFn: func(data []byte) error {
				prompt := strings.TrimSpace(string(data))
				if prompt == "" {
					return nil
				}
				return s.PeerSubmit(sessID, peerID, prompt)
			},
			BlockingReadFn: notBlocking,
		}, nil
	}

	// Top-level fixed files.
	switch name {
	case "new":
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: "new", Mode_: fs.Perms[fs.PathSessions].Files["new"]}, nil
			},
			ReadFn: func() ([]byte, error) {
				return []byte("name=\ncwd=\nremote=\nbackend=\nmodel=\nagent=\nmaxTokens=\nmaxCompletionTokens=\ntemperature=\ntopP=\ntopK=\nminP=\ntopA=\nfrequencyPenalty=\npresencePenalty=\nrepetitionPenalty=\nreasoning=\nreasoningEffort=\nincludeReasoning=\nresponseFormat=\nstop=\nverbosity=\n"), nil
			},
			WriteFn: func(data []byte) error {
				_, err := s.CreateSession(strings.Fields(strings.TrimSpace(string(data))))
				return err
			},
			BlockingReadFn: notBlocking,
		}, nil
	case "idx":
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: "idx", Mode_: fs.Perms[fs.PathSessions].Files["idx"]}, nil
			},
			ReadFn:         func() ([]byte, error) { return s.index(), nil },
			WriteFn:        func([]byte) error { return fmt.Errorf("idx: read-only") },
			BlockingReadFn: notBlocking,
		}, nil
	}
	if _, ok := fs.Perms[fs.PathSessions].Files[name]; ok {
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathSessions].Files[name]}, nil
			},
			ReadFn: func() ([]byte, error) {
				return s.cfg.ReadFile(paths.CfgDir() + "/scripts/s/" + name)
			},
			WriteFn:        func([]byte) error { return fmt.Errorf("%s: not writable", name) },
			BlockingReadFn: notBlocking,
		}, nil
	}

	// Hierarchical paths: {id}/{file} or {id}/t/{tool}
	parts = strings.SplitN(name, "/", 3)
	sessID := parts[0]
	sess := s.Session(sessID)
	if sess == nil {
		return nil, fmt.Errorf("%s: not found", name)
	}
	if len(parts) == 1 {
		return nil, fmt.Errorf("%s: is a directory", name)
	}
	// Session file: {id}/{file}
	sfs, err := s.openStore(sess)
	if err != nil {
		return nil, err
	}
	return sfs.Open(parts[1])
}

func (s *Manager) create(name string) error {
	// Peer creation: {id}/peer/{peer-id}
	parts := strings.SplitN(name, "/", 3)
	if len(parts) == 3 && parts[1] == "peer" {
		return s.PeerAdd(parts[0], parts[2])
	}
	return fmt.Errorf("create not supported: %s", name)
}

func (s *Manager) del(name string) error {
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	// Delete session directory itself.
	if len(parts) == 1 {
		s.mu.RLock()
		_, ok := s.sessions[sessID]
		s.mu.RUnlock()
		if ok {
			s.KillSession(sessID)
			return nil
		}
		return fmt.Errorf("session not found: %s", sessID)
	}
	// Peer removal: {id}/peer/{peer-id}
	if len(parts) == 3 && parts[1] == "peer" {
		return s.PeerRemove(sessID, parts[2])
	}
	// Proc dismiss: {id}/proc/{pid}
	if len(parts) == 3 && parts[1] == "proc" {
		sess := s.Session(sessID)
		if sess == nil {
			return fmt.Errorf("session not found: %s", sessID)
		}
		pid, err := strconv.Atoi(parts[2])
		if err != nil {
			return fmt.Errorf("invalid pid: %s", parts[2])
		}
		if !sess.Core.DismissDetached(pid) {
			return fmt.Errorf("process %d not found or still running", pid)
		}
		return nil
	}
	// Session files are synthetic; allow rm -r to continue.
	return nil
}

// Session returns the session for the given ID, or nil.
func (s *Manager) Session(id string) *Session {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.sessions[id]
}

// SessionByUname returns the session with the given uname (principal), or nil.
func (s *Manager) SessionByUname(uname string) *Session {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sess := range s.sessions {
		if sess.uname == uname {
			return sess
		}
	}
	return nil
}

// OpenStore returns a Tree for the given session ID.
func (s *Manager) OpenStore(id string) (*fs.Tree, error) {
	if sess := s.Session(id); sess != nil {
		return s.openStore(sess)
	}
	return nil, fmt.Errorf("session not found: %s", id)
}

func (s *Manager) openStore(sess *Session) (*fs.Tree, error) {
	return NewSessionTree(
		sess,
		s.cfg.Log,
		func() { s.KillSession(sess.id) },
		func(newID string) error { return s.renameSession(sess.id, newID) },
		nil, // no transcript saving
		s.cfg.InvalidateModels,
	), nil
}

// InterruptAll interrupts every active session.
func (s *Manager) InterruptAll() {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sess := range s.sessions {
		sess.Core.Interrupt(agent.ErrInterrupted)
	}
}

// Shutdown kills all active sessions.
// --- Session Persistence ---

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

func (s *Manager) persistSession(id string) {
	s.mu.RLock()
	sess, ok := s.sessions[id]
	s.mu.RUnlock()
	if !ok {
		return
	}
	dir := s.activeSessionsDir()
	os.MkdirAll(dir, 0700)
	path := filepath.Join(dir, id+".json")
	if err := sess.Core.SaveSession(path); err != nil {
		s.cfg.Log.Error("persist session %s: %v", id, err)
	}
}

func (s *Manager) removePersistedSession(id string) {
	path := filepath.Join(s.activeSessionsDir(), id+".json")
	os.Remove(path)
}

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

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

	// Load all persisted session JSONs (fast, sequential disk reads)
	type loadedSession struct {
		ps   *agent.PersistedSession
		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.LoadPersistedSession(path)
		if err != nil {
			s.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
	}

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

func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
	cwd := ps.CWD
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	agentName := ps.Agent
	if agentName == "" {
		agentName = "default"
	}
	sessID := ps.ID

	cfg := LoadAgentConfig(s.cfg.AgentsDir, agentName, 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 := s.nextUname()
	var newDisp func() tools.Dispatcher
	var promptEnv []string
	remoteTarget := ps.Remote

	if remoteTarget != "" {
		rsrv, dialErr := remote.Dial(context.Background(), remote.Config{
			SSHTarget: remoteTarget,
			CWD:       cwd,
		})
		if dialErr != nil {
			return fmt.Errorf("remote dial: %w", dialErr)
		}
		newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
			"execute": func() tools.Server { return rsrv },
		})
		promptEnv = []string{
			"PRIME_CWD=" + cwd,
			"PRIME_PLATFORM=" + rsrv.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo),
		}
	} else {
		var execOpts []execute.Option
		if !s.cfg.NoMount {
			execOpts = append(execOpts, WithMount9P())
		}
		if s.cfg.Strict {
			execOpts = append(execOpts, execute.WithStrict())
		}
		if s.cfg.Yolo {
			execOpts = append(execOpts, execute.WithYolo())
		}
		newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
			"execute": execute.Decl(cwd, execOpts...),
		})
		promptEnv = agent.PromptEnv(cwd)
	}

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

	// Compute base layers (same as new-session path).
	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(s.cfg.Enable9P, s.cfg.EnableDBus, 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, "")

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

	restoredSession := agent.RestoreSession(ps)

	var sessPtr *Session
	core := agent.NewAgentCore(agent.AgentCoreConfig{
		Backend:        be,
		AgentName:      agentName,
		AgentsDir:      s.cfg.AgentsDir,
		SessionsDir:    s.cfg.SessionsDir,
		SessionID:      sessID,
		Uname:          uname,
		CWD:            cwd,
		Remote:         remoteTarget,
		Session:        restoredSession,
		Runtime:        rt,
		NewDispatcher:  newDisp,
		PromptEnvExtra: promptEnv,
		BaseLayers:     []string{sysPrompt, opModel, envBlock},
		Log:            s.cfg.Sink.NewLogger("core"),
		ReadPlanStep: func() string {
			if sessPtr == nil {
				return ""
			}
			sessPtr.mu.RLock()
			data := make([]byte, len(sessPtr.plan))
			copy(data, sessPtr.plan)
			sessPtr.mu.RUnlock()
			return agent.NextUncheckedStep(data)
		},
	})

	ctx, cancel := context.WithCancel(context.Background())
	sess := NewSession(sessID, core, ctx, cancel)
	sessPtr = sess
	sess.uname = uname
	sess.remote = remoteTarget



	// Replay tail of persisted messages into the chat log so the GUI
	// and `chat` file show recent history on restore.
	replayMessagesToLog(sess, ps.Messages)

	s.mu.Lock()
	s.sessions[sessID] = sess
	s.mu.Unlock()

	s.cfg.Log.Info("restored session %s (backend=%s model=%s agent=%s)", sessID, backendName, modelName, agentName)
	if s.cfg.OnSessionCreated != nil {
		s.cfg.OnSessionCreated(sessID, sess)
	}
	return nil
}

func (s *Manager) Shutdown() {
	// Interrupt all in-progress turns and wait for them to finish
	// before persisting state, so we capture the latest messages.
	s.InterruptAll()
	s.waitIdle(100*time.Millisecond, 5*time.Second)
	s.saveAllSessions()
	s.mu.Lock()
	ids := make([]string, 0, len(s.sessions))
	for id := range s.sessions {
		ids = append(ids, id)
	}
	s.mu.Unlock()
	for _, id := range ids {
		s.mu.Lock()
		sess := s.sessions[id]
		delete(s.sessions, id)
		s.mu.Unlock()
		if sess != nil {
			sess.Cancel()
			sess.Core.Close()
			if sess.mount != nil {
				teardownMount(sess.mount)
			}
			s.cfg.Log.Info("shutdown session %s", id)
		}
	}
}

// waitIdle polls all sessions until they are idle or the timeout expires.
func (s *Manager) waitIdle(poll, timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		allIdle := true
		s.mu.RLock()
		for _, sess := range s.sessions {
			if sess.Core.State() != "idle" {
				allIdle = false
				break
			}
		}
		s.mu.RUnlock()
		if allIdle {
			return
		}
		time.Sleep(poll)
	}
	s.cfg.Log.Warn("shutdown: timed out waiting for sessions to become idle, saving anyway")
}

func (s *Manager) KillSession(id string) {
	s.mu.Lock()
	sess := s.sessions[id]
	delete(s.sessions, id)
	s.mu.Unlock()
	if sess != nil {
		sess.Cancel()
		sess.Core.Close()
		if sess.mount != nil {
			teardownMount(sess.mount)
		}
		s.removePersistedSession(id)
		s.cfg.Log.Info("killed session %s", id)
		if s.cfg.OnSessionKilled != nil {
			s.cfg.OnSessionKilled(id)
		}
	}
}

func (s *Manager) index() []byte {
	var sb strings.Builder
	s.mu.RLock()
	ids := make([]string, 0, len(s.sessions))
	for id := range s.sessions {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	for _, id := range ids {
		sess := s.sessions[id]
		sess.mu.RLock()
		state := sess.Core.State()
		cwd := sess.Core.CWD()
		be := sess.Core.BackendName()
		model := sess.Core.ModelName()
		agent := sess.Core.AgentName()
		sess.mu.RUnlock()
		fmt.Fprintf(&sb, "%s\t%s\t%s\t%s\t%s\t%s\n", id, state, cwd, be, model, agent)
	}
	s.mu.RUnlock()
	return []byte(sb.String())
}

// CreateSession creates a new agent session from key=value args.
// Returns the session ID on success.
func (s *Manager) CreateSession(args []string) (string, error) {
	name := ""
	backendOverride := ""
	modelOverride := ""
	agentName := "default"
	cwd := ""
	remoteTarget := ""
	systemPromptOverride := ""
	for _, arg := range args {
		k, v, ok := strings.Cut(arg, "=")
		if !ok {
			return "", fmt.Errorf("invalid option %q (expected key=value)", arg)
		}
		if v == "" {
			continue
		}
		switch k {
		case "name":
			name = v
		case "backend":
			backendOverride = v
		case "model":
			modelOverride = v
		case "agent":
			agentName = v
		case "cwd":
			cwd = v
		case "remote":
			remoteTarget = v
		case "system_prompt":
			systemPromptOverride = v
		default:
			return "", fmt.Errorf("unknown option %q (valid: name, backend, model, agent, cwd, remote, system_prompt)", k)
		}
	}
	if agentName == "default" {
		if v := os.Getenv("OLLIE_DEFAULT_AGENT"); v != "" {
			agentName = v
		}
	}

	cwd = paths.ExpandHome(os.ExpandEnv(cwd))
	if cwd == "" {
		return "", fmt.Errorf("cwd is required (e.g. new cwd=/path/to/project)")
	}

	sessID := name
	if sessID == "" {
		sessID = agent.NewSessionID()
	}

	s.mu.RLock()
	_, exists := s.sessions[sessID]
	s.mu.RUnlock()
	if exists {
		return "", fmt.Errorf("session already exists: %s", sessID)
	}

	var core agent.Core
	var sessPtr *Session
	uname := s.nextUname()
	if s.cfg.NewCore != nil {
		var err error
		core, err = s.cfg.NewCore(sessID, agentName, cwd)
		if err != nil {
			return "", err
		}
	} else {
		cfg := LoadAgentConfig(s.cfg.AgentsDir, agentName, nil)
		if cfg != nil {
			if backendOverride == "" && cfg.Backend != "" {
				backendOverride = cfg.Backend
			}
			if modelOverride == "" && cfg.Model != "" {
				modelOverride = cfg.Model
			}
		}

		be, err := backend.NewWithName(backendOverride)
		if err != nil {
			return "", fmt.Errorf("backend: %w", err)
		}

		if modelOverride == "" {
			modelOverride = os.Getenv("OLLIE_MODEL")
		}
		if modelOverride != "" {
			be.SetModel(modelOverride)
		}

		if err := s.cfg.MkdirAll(s.cfg.SessionsDir, fs.PermMkdirPrivate); err != nil {
			return "", fmt.Errorf("sessions dir: %w", err)
		}

		var execOpts []execute.Option
		if !s.cfg.NoMount {
			execOpts = append(execOpts, WithMount9P())
		}
		if s.cfg.Strict {
			execOpts = append(execOpts, execute.WithStrict())
		}
		if s.cfg.Yolo {
			execOpts = append(execOpts, execute.WithYolo())
		}
		if cfg != nil && len(cfg.AllowExecutors) > 0 {
			execOpts = append(execOpts, execute.WithAllowExecutors(cfg.AllowExecutors))
		}
		if cfg != nil && len(cfg.AllowTools) > 0 {
			execOpts = append(execOpts, execute.WithAllowTools(cfg.AllowTools))
		}
		var newDisp func() tools.Dispatcher
		var remoteEnv []string // extra env vars from remote host info
		if remoteTarget != "" {
			// Dial eagerly so we get HostInfo for prompt resolution.
			// FIXME: Eager dial means session creation blocks on SSH connect.
			// Eventually make this async or move host info into the prompt
			// scripts themselves (running remotely via the RPC).
			rsrv, dialErr := remote.Dial(context.Background(), remote.Config{
				SSHTarget: remoteTarget,
				CWD:       cwd,
			})
			if dialErr != nil {
				return "", fmt.Errorf("remote dial: %w", dialErr)
			}
			newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
				"execute": func() tools.Server { return rsrv },
			})
			remoteEnv = []string{
				"PRIME_CWD=" + cwd,
				"PRIME_PLATFORM=" + rsrv.Info.Platform,
				"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo),
			}
		} else {
			newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
				"execute": execute.Decl(cwd, execOpts...),
			})
		}

		var promptEnv []string
		if len(remoteEnv) > 0 {
			promptEnv = remoteEnv
		} else {
			promptEnv = agent.PromptEnv(cwd)
		}
		env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
		env = append(env, promptEnv...)

		// Compose base layers: system prompt + operational model + environment.
		spOverride := systemPromptOverride
		if spOverride == "" && 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(s.cfg.Enable9P, s.cfg.EnableDBus, envMap)

		// Determine environment values for the environment layer.
		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, "")

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

		// sessPtr is set after NewSession; the ReadPlanStep closure captures it.
		core = agent.NewAgentCore(agent.AgentCoreConfig{
			Backend:        be,
			AgentName:      agentName,
			AgentsDir:      s.cfg.AgentsDir,
			SessionsDir:    s.cfg.SessionsDir,
			SessionID:      sessID,
			Uname:          uname,
			CWD:            cwd,
			Remote:         remoteTarget,
			Runtime:        rt,
			NewDispatcher:  newDisp,
			PromptEnvExtra: promptEnv,
			BaseLayers:     []string{sysPrompt, opModel, envBlock},
			Log:            s.cfg.Sink.NewLogger("core"),
			ReadPlanStep: func() string {
				if sessPtr == nil {
					return ""
				}
				sessPtr.mu.RLock()
				data := make([]byte, len(sessPtr.plan))
				copy(data, sessPtr.plan)
				sessPtr.mu.RUnlock()
				return agent.NextUncheckedStep(data)
			},
		})
	}

	ctx, cancel := context.WithCancel(context.Background())
	sess := NewSession(sessID, core, ctx, cancel)
	sessPtr = sess
	sess.remote = remoteTarget

	// For remote sessions, the mount isn't managed by a local execute.Server
	// (which doesn't exist). Start it directly so $OLLIE is available for
	s.mu.Lock()
	sess.uname = uname
	s.sessions[sessID] = sess
	s.mu.Unlock()

	// HACK: Load the agent now that the session mount is active.
	core.Submit(ctx, "/agent "+agentName)

	s.cfg.Log.Info("new session %s (backend=%s model=%s agent=%s)",
		sessID, core.BackendName(), core.ModelName(), core.AgentName())
	if s.cfg.OnSessionCreated != nil {
		s.cfg.OnSessionCreated(sessID, sess)
	}
	return sessID, nil
}

// RenameSession renames a session (exported for D-Bus adapter).
func (s *Manager) RenameSession(old, new string) error {
	return s.renameSession(old, new)
}

func (s *Manager) renameSession(old, new string) error {
	// Session rename: {oldID} -> {newID}
	oldID := old
	newID := new
	s.mu.Lock()

	sess, ok := s.sessions[oldID]
	if !ok {
		s.mu.Unlock()
		return fmt.Errorf("session not found: %s", oldID)
	}
	if _, exists := s.sessions[newID]; exists {
		s.mu.Unlock()
		return fmt.Errorf("session already exists: %s", newID)
	}
	if sess.Core.IsRunning() {
		s.mu.Unlock()
		return fmt.Errorf("cannot rename while agent is running")
	}

	if err := sess.Core.SetSessionID(newID); err != nil {
		s.mu.Unlock()
		return err
	}

	sess.id = newID
	s.sessions[newID] = sess
	delete(s.sessions, oldID)
	s.mu.Unlock()

	sess.AppendLog([]byte(fmt.Sprintf(":: session renamed: %s -> %s\n", oldID, newID)))
	s.cfg.Log.Info("renamed session %s -> %s", oldID, newID)
	if s.cfg.OnSessionRenamed != nil {
		s.cfg.OnSessionRenamed(oldID, newID)
	}
	return nil
}

// LoadAgentConfig resolves and loads the config for a named agent.
// Returns nil if the config file does not exist; BuildRuntime handles nil configs.
func LoadAgentConfig(agentsDir, name string, open func(string) (*os.File, error)) *config.Config {
	if open == nil {
		open = os.Open
	}
	path := agent.AgentConfigPath(agentsDir, name)
	f, err := open(path)
	if err != nil {
		return nil
	}
	defer f.Close()
	cfg, _ := config.Load(f)
	return cfg
}

// FormatEvent converts an agent Event to bytes for appending to a chat log.
// Streaming roles (assistant, reasoning) return only the content chunk;
// the caller (startEventLog) is responsible for writing the [role] header.
func FormatEvent(ev agent.Event) []byte {
	switch ev.Role {
	case "user":
		return []byte("[user]\n" + ev.Content + "\n")
	case "assistant":
		return []byte(ev.Content)
	case "reasoning":
		return []byte(ev.Content)
	case "call":
		args := squashWhitespace(ev.Content)
		return []byte("[call:" + ev.Name + "]\n" + args + "\n")
	case "tool":
		return []byte("[tool:" + ev.Name + "]\n" + strings.TrimRight(ev.Content, "\n") + "\n")
	case "error":
		return []byte("[error]\n" + ev.Content + "\n")
	case "info":
		return []byte("[info]\n" + ev.Content)
	case "retry":
		return []byte("[retry]\n" + ev.Content + "s\n")
	case "stalled":
		return []byte("[stalled]\n")
	case "maxsteps":
		return []byte("[maxsteps]\n" + ev.Content + "\n")
	default:
		return nil
	}
}

func squashWhitespace(s string) string {
	return strings.Join(strings.Fields(s), " ")
}

// replayMessagesToLog renders the tail of a persisted message list into the
// session's chat log so that restored sessions show recent history.
func replayMessagesToLog(sess *Session, messages []backend.Message) {
	const maxReplay = 20
	start := len(messages) - maxReplay
	if start < 0 {
		start = 0
	}
	// Skip system messages at the start of the replay window.
	for start < len(messages) && messages[start].Role == "system" {
		start++
	}
	for _, m := range messages[start:] {
		switch m.Role {
		case "system":
			// skip
		case "user":
			sess.AppendLog([]byte("[user]\n" + m.Content + "\n"))
		case "assistant":
			sess.AppendLog([]byte("[assistant]\n"))
			if m.Content != "" {
				sess.AppendLog([]byte(m.Content + "\n"))
			}
			for _, tc := range m.ToolCalls {
				args := squashWhitespace(string(tc.Arguments))
				sess.AppendLog([]byte("[call:" + tc.Name + "]\n" + args + "\n"))
			}
		case "tool":
			sess.AppendLog([]byte("[tool]\n" + strings.TrimRight(m.Content, "\n") + "\n"))
		}
	}
}

// --- Code Completion ---

// Complete performs a stateless one-shot code completion via backend.Generate.
func (s *Manager) Complete(cwd, filePath, prefix, suffix, extraContext string) (string, error) {
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	cwd = paths.ExpandHome(os.ExpandEnv(cwd))

	model := os.Getenv("OLLIE_COMPLETE_MODEL")
	backendName := os.Getenv("OLLIE_COMPLETE_BACKEND")
	if model == "" || backendName == "" {
		return "", fmt.Errorf("OLLIE_COMPLETE_MODEL and OLLIE_COMPLETE_BACKEND must be set")
	}

	const prefixMax = 12000
	const suffixMax = 1000
	if len(prefix) > prefixMax {
		prefix = prefix[len(prefix)-prefixMax:]
	}
	if len(suffix) > suffixMax {
		suffix = suffix[:suffixMax]
	}

	fileHint := ""
	if filePath != "" {
		fileHint = " in " + filePath
	}
	contextBlock := ""
	if extraContext != "" {
		contextBlock = "\n" + extraContext
	}
	prompt := fmt.Sprintf(`Implement the code at the cursor%s. The prefix ends at the point where new code is needed. Write the implementation — do not echo stubs, TODOs, or placeholder returns from the prefix. Output ONLY raw code. No reasoning, no shell commands, no explanations, no markdown fences, no backticks, no preamble. Your entire response must be valid code that can be inserted directly into the file.
%s
<prefix>
%s
</prefix>
<suffix>
%s
</suffix>`, fileHint, contextBlock, prefix, suffix)

	result, err := backend.Generate(context.Background(), backend.GenerateRequest{
		Prompt:  prompt,
		Backend: backendName,
		Model:   model,
	})
	if err != nil {
		return "", err
	}

	result = stripCompletionNoise(result)
	result = stripPrefixEcho(prefix, result)
	return result, nil
}

func stripCompletionNoise(s string) string {
	var lines []string
	for _, line := range strings.Split(s, "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "```") {
			continue
		}
		if trimmed == "<prefix>" || trimmed == "</prefix>" ||
			trimmed == "<suffix>" || trimmed == "</suffix>" {
			continue
		}
		if strings.HasPrefix(trimmed, ":: ") {
			continue
		}
		lines = append(lines, line)
	}
	return strings.Join(lines, "\n")
}

func stripPrefixEcho(prefix, result string) string {
	tailMax := 200
	if len(prefix) < tailMax {
		tailMax = len(prefix)
	}
	for i := tailMax; i > 0; i-- {
		tail := prefix[len(prefix)-i:]
		if strings.HasPrefix(result, tail) {
			return result[len(tail):]
		}
	}
	return result
}

func crc32Str(s string) uint32 {
	return crc32.ChecksumIEEE([]byte(s))
}