ollie/fs/lifecycle.go

package fs

import (
	"context"
	"crypto/rand"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

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

// ── Event ring ──────────────────────────────────────────────────

const eventRingCap = 100

// globalRootState is set by NewRoot so handler functions can access
// the session map without going through ctx.Root (which is nil at
// build time when closures are created).
var globalRootState *rootState

// rootState is the package-level state stored in root.Data.
type rootState struct {
	cfg      Config
	mu       sync.RWMutex
	sessions map[string]*Session
}

func state(root *Tree) *rootState {
	return root.Data.(*rootState)
}

type eventRing struct {
	mu    sync.Mutex
	cond  *sync.Cond
	buf   []string
	cap   int
	start int
	count int
	total int64
}

var evRing = func() *eventRing {
	e := &eventRing{cap: eventRingCap, buf: make([]string, eventRingCap)}
	e.cond = sync.NewCond(&e.mu)
	return e
}()

// PostEvent appends a structured event line and wakes all waiters.
func PostEvent(line string) {
	evRing.mu.Lock()
	idx := (evRing.start + evRing.count) % evRing.cap
	evRing.buf[idx] = line
	if evRing.count < evRing.cap {
		evRing.count++
	} else {
		evRing.start = (evRing.start + 1) % evRing.cap
	}
	evRing.total++
	evRing.mu.Unlock()
	evRing.cond.Broadcast()
}

// WaitEvent blocks until events are available since the given offset.
func WaitEvent(ctx context.Context, offsetStr string) ([]byte, string, error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	context.AfterFunc(ctx, func() { evRing.cond.Broadcast() })

	evRing.mu.Lock()
	var offset int64
	if offsetStr != "" {
		fmt.Sscanf(offsetStr, "%d", &offset)
	} else {
		offset = evRing.total
	}
	for evRing.total <= offset {
		if ctx.Err() != nil {
			evRing.mu.Unlock()
			return nil, offsetStr, nil
		}
		evRing.cond.Wait()
	}
	base := evRing.total - int64(evRing.count)
	if offset < base {
		offset = base
	}
	skip := int(offset - base)
	var sb strings.Builder
	for i := 0; i < evRing.count-skip; i++ {
		idx := (evRing.start + skip + i) % evRing.cap
		sb.WriteString(evRing.buf[idx])
		sb.WriteByte('\n')
	}
	next := fmt.Sprintf("%d", evRing.total)
	evRing.mu.Unlock()
	return []byte(sb.String()), next, nil
}

// ── buildIndex — session/idx content ────────────────────────────

func buildIndex(rs *rootState) []byte {
	var sb strings.Builder
	rs.mu.RLock()
	ids := make([]string, 0, len(rs.sessions))
	for id := range rs.sessions {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	for _, id := range ids {
		sess := rs.sessions[id]
		sess.mu.RLock()
		name := sess.name
		sid := sess.id
		if sess.Core != nil {
			for _, ag := range sess.Core.Agents() {
				st := ag.State()
				cwd := sess.Core.CWD()
				be := ag.BackendName()
				model := ag.ModelName()
				agName := ag.Name()
				fmt.Fprintf(&sb, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", name, st, cwd, be, model, agName, sid)
			}
		} else {
			fmt.Fprintf(&sb, "%s\t\t\t\t\t\t%s\n", name, sid)
		}
		sess.mu.RUnlock()
	}
	rs.mu.RUnlock()
	return []byte(sb.String())
}

// ── Session lifecycle ───────────────────────────────────────────

// Kill terminates a session and removes it from the tree.
func Kill(rs *rootState, id string) {
	rs.mu.Lock()
	sess := rs.sessions[id]
	if sess == nil {
		rs.mu.Unlock()
		rs.cfg.Log.Error("Kill: session %q not found in sessions map", id)
		return
	}
	delete(rs.sessions, id)
	rs.mu.Unlock()

	var agentNames []string
	if sess.Core != nil {
		for _, ag := range sess.Core.Agents() {
			agentNames = append(agentNames, ag.ID())
		}
		sess.Core.Close()
	}
	sess.Cancel()
	if sess.keeper != nil {
		sess.keeper.Close()
	} else if sess.proc != nil {
		sess.proc.Close()
	}
	removePersistedAgent(rs, sess.id)
	rs.cfg.Log.Info("killed session %s", id)
	for _, name := range agentNames {
		PostEvent("A kill session/" + id + "/agent/" + name)
	}
	PostEvent("S kill session/" + id)
}

// Rename renames a session by updating its Name and the sessions map.
func Rename(rs *rootState, old, new string) error {
	rs.mu.Lock()
	sess, ok := rs.sessions[old]
	if !ok {
		rs.mu.Unlock()
		return fmt.Errorf("session not found: %s", old)
	}
	if _, exists := rs.sessions[new]; exists {
		rs.mu.Unlock()
		return fmt.Errorf("session already exists: %s", new)
	}
	if sess.Core != nil {
		for _, ag := range sess.Core.Agents() {
			if ag.IsRunning() {
				rs.mu.Unlock()
				return fmt.Errorf("cannot rename while agent is running")
			}
		}
	}
	sess.name = new
	rs.sessions[new] = sess
	delete(rs.sessions, old)
	rs.mu.Unlock()
	// Persist the updated session and agent display names immediately. The
	// agent's own autosave does not know the fs-level session display name.
	persistSession(rs, new)
	rs.cfg.Log.Info("renamed session %s -> %s", old, new)
	PostEvent("S rename session/" + old + " session/" + new)
	return nil
}

// Shutdown interrupts all sessions, waits for idle, persists, and closes.
func Shutdown(root *Tree) {
	rs := root.Data.(*rootState)
	InterruptAll(root)
	waitIdle(rs, 100*time.Millisecond, 5*time.Second)
	saveAllSessions(rs)
	rs.mu.Lock()
	ids := make([]string, 0, len(rs.sessions))
	for id := range rs.sessions {
		ids = append(ids, id)
	}
	rs.mu.Unlock()
	for _, id := range ids {
		rs.mu.Lock()
		sess := rs.sessions[id]
		delete(rs.sessions, id)
		rs.mu.Unlock()
		if sess != nil {
			sess.Cancel()
			if sess.Core != nil {
				sess.Core.Close()
			}
			if sess.keeper != nil {
				sess.keeper.Close()
			} else if sess.proc != nil {
				sess.proc.Close()
			}
			rs.cfg.Log.Info("shutdown session %s", id)
		}
	}
}

func waitIdle(rs *rootState, poll, timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		allIdle := true
		rs.mu.RLock()
		for _, sess := range rs.sessions {
			if sess.Core != nil {
				for _, ag := range sess.Core.Agents() {
					if ag.State() != "idle" {
						allIdle = false
						break
					}
				}
			}
			if !allIdle {
				break
			}
		}
		rs.mu.RUnlock()
		if allIdle {
			return
		}
		time.Sleep(poll)
	}
	rs.cfg.Log.Warn("shutdown: timed out waiting for sessions to become idle, saving anyway")
}

// InterruptAll interrupts every active session.
func InterruptAll(root *Tree) {
	rs := root.Data.(*rootState)
	rs.mu.RLock()
	defer rs.mu.RUnlock()
	for _, sess := range rs.sessions {
		if sess.Core != nil {
			for _, ag := range sess.Core.Agents() {
				ag.Interrupt(agent.ErrInterrupted)
			}
		}
	}
}

// Lookup returns the session for the given ID, or nil.
func Lookup(root *Tree, id string) *Session {
	rs := root.Data.(*rootState)
	rs.mu.RLock()
	defer rs.mu.RUnlock()
	// Try by display name first (map key).
	if s := rs.sessions[id]; s != nil {
		return s
	}
	// Fall back to immutable session ID.
	for _, sess := range rs.sessions {
		if sess.id == id {
			return sess
		}
	}
	return nil
}

// LookupByUname returns the session with the given uname, or nil.
func LookupByUname(root *Tree, uname string) *Session {
	rs := root.Data.(*rootState)
	rs.mu.RLock()
	defer rs.mu.RUnlock()
	for _, sess := range rs.sessions {
		if sess.uname == uname {
			return sess
		}
	}
	return nil
}

// AddSession inserts a pre-built session (used by tests).
func AddSession(root *Tree, sess *Session) {
	rs := root.Data.(*rootState)
	rs.mu.Lock()
	rs.sessions[sess.RunnableID()] = sess
	rs.mu.Unlock()
}

// ── Create session/agent ────────────────────────────────────────

func nextUname(rs *rootState) string {
	b := make([]byte, 6)
	rand.Read(b)
	return fmt.Sprintf("%06x%06x", b[:3], b[3:6])
}

// Create creates a new empty session from key=value args and registers it.
func Create(rs *rootState, args []string) (string, error) {
	name := ""
	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
		default:
			return "", fmt.Errorf("unknown option %q (valid: name)", k)
		}
	}

	sessID := coresession.NewSessionID()
	if name == "" {
		name = sessID
		if idx := strings.IndexByte(sessID, '-'); idx > 0 {
			name = sessID[:idx]
		}
	}

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

	sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
	sess := NewSession(sessID, nil, sessionCtx, sessionCancel)
	sess.SetName(name)

	// Spawn the tool server so the session is immediately operational.
	cwd, _ := os.Getwd()
	var dialOpts []toolsrv.Option
	if rs.cfg.Yolo {
		dialOpts = append(dialOpts, toolsrv.WithYolo())
	}
	proc, err := toolsrv.Spawn(sessionCtx, cwd, dialOpts...)
	if err != nil {
		sessionCancel()
		return "", fmt.Errorf("spawn tool server: %w", err)
	}
	sess.proc = proc
	sess.keeper = toolsrv.NewProcessKeeper(sessionCtx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
		return toolsrv.Spawn(ctx, cwd, dialOpts...)
	})
	conn, err := toolsrv.Dial(sessionCtx, toolsrv.LocalAddr(proc.Socket))
	if err != nil {
		sessionCancel()
		proc.Close()
		return "", fmt.Errorf("dial tool server: %w", err)
	}
	sess.toolsConn = conn

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

	rs.cfg.Log.Info("created empty session %s (id=%s)", name, sessID)
	PostEvent("S new session/" + name)
	return name, nil
}

// CreateFromRoot is the public API for callers that have a *Tree.
func CreateFromRoot(root *Tree, args []string) (string, error) {
	return Create(root.Data.(*rootState), args)
}

// CreateAgent creates and configures an agent within a session.
func CreateAgent(rs *rootState, sessName string, args []string) error {
	backendOverride := ""
	modelOverride := ""
	profile := ""
	name := ""
	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":
			profile = 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)
		}
	}

	rs.mu.RLock()
	sess, ok := rs.sessions[sessName]
	rs.mu.RUnlock()
	if !ok {
		sessID := coresession.NewSessionID()
		sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
		sess = NewSession(sessID, nil, sessionCtx, sessionCancel)
		sess.SetName(sessName)
	}

	if profile == "" {
		if v := os.Getenv("OLLIE_DEFAULT_AGENT"); v != "" {
			profile = v
		} else {
			profile = "default"
		}
	}

	cwd = paths.ExpandHome(os.ExpandEnv(cwd))
	if cwd == "" {
		return fmt.Errorf("cwd is required (set cwd=<path> when creating an agent)")
	}
	if remoteTarget == "" {
		if info, err := os.Stat(cwd); err != nil {
			return fmt.Errorf("cwd %q does not exist", cwd)
		} else if !info.IsDir() {
			return fmt.Errorf("cwd %q is not a directory", cwd)
		}
	}

	uname := nextUname(rs)
	sessID := sess.id

	cfg := LoadAgentConfig(rs.cfg.AgentsDir, profile, 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 := rs.cfg.MkdirAll(rs.cfg.SessionsDir, PermMkdirPrivate); err != nil {
		return fmt.Errorf("sessions dir: %w", err)
	}

	var dialOpts []toolsrv.Option
	if rs.cfg.Yolo {
		dialOpts = append(dialOpts, toolsrv.WithYolo())
	}

	var proc *toolsrv.Process
	var newToolServer func() toolsrv.Runner
	var promptEnv []string

	if remoteTarget != "" {
		proc, err = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{
			SSHTarget: remoteTarget,
			CWD:       cwd,
		})
		if err != nil {
			return fmt.Errorf("remote spawn: %w", err)
		}
		keeper := toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
			return toolsrv.SpawnRemote(ctx, toolsrv.RemoteConfig{SSHTarget: remoteTarget, CWD: cwd})
		})
		sess.keeper = keeper
		newToolServer = func() toolsrv.Runner {
			conn, err := keeper.Dial()
			if err != nil {
				return nil
			}
			return conn
		}
		promptEnv = []string{
			"PRIME_CWD=" + cwd,
			"PRIME_PLATFORM=" + proc.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo),
		}
	} else if sess.proc != nil {
		// Reuse the tool server spawned during Create.
		proc = sess.proc
		if sess.keeper != nil {
			newToolServer = func() toolsrv.Runner {
				conn, err := sess.keeper.Dial()
				if err != nil {
					rs.cfg.Log.Error("toolsrv dial: %v", err)
					return nil
				}
				return conn
			}
		} else {
			newToolServer = func() toolsrv.Runner {
				return sess.toolsConn
			}
		}
		if cwc, ok := sess.toolsConn.(interface{ SetCWD(string) }); ok {
			cwc.SetCWD(cwd)
		}
		promptEnv = agent.PromptEnv(cwd)
	} else {
		proc, err = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...)
		if err != nil {
			return fmt.Errorf("local spawn: %w", err)
		}
		keeper := toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
			return toolsrv.Spawn(ctx, cwd, dialOpts...)
		})
		sess.keeper = keeper
		newToolServer = func() toolsrv.Runner {
			conn, err := keeper.Dial()
			if err != nil {
				rs.cfg.Log.Error("toolsrv dial: %v", err)
				return nil
			}
			return conn
		}
		promptEnv = agent.PromptEnv(cwd)
	}

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

	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(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()
	if toolSrv == nil {
		// Tool server is vital infrastructure — retry once with a fresh spawn.
		rs.cfg.Log.Error("session %s: initial tool server dial failed, retrying...", sessName)
		if sess.keeper != nil {
			// Force respawn via keeper
			if conn, err := sess.keeper.Dial(); err == nil {
				toolSrv = conn
			}
		}
		if toolSrv == nil {
			return fmt.Errorf("tool server unavailable after retry (session infrastructure broken — try again)")
		}
	}
	// Load configured tools before building the runtime. BuildRuntime snapshots
	// the loaded-tool list to construct the agent's initial 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", sessName, tl, err)
			}
		}
	}

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

	al := NewAgentLog(remoteTarget)

	if sess.Core == nil {
		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,
			Runtime:        rt,
			NewToolServer:  newToolServer,
			PromptEnvExtra: promptEnv,
			BaseLayers:     []string{sysPrompt, opModel, envBlock},
			Log:            rs.cfg.Sink.NewLogger("core"),
			Output:         NewEventHandler(al),
			ReadPlanStep: func() string {
				return ""
			},
		})
		sess.mu.Lock()
		sess.Core = core
		sess.uname = uname
		sess.setAgentLogLocked(uname, al)
		sess.proc = proc
		sess.mu.Unlock()
	} else {
		saveFn, flushFn := sess.Core.SaveFuncs()
		ag := agent.NewAgent(agent.AgentCfg{
			History:        nil,
			Runtime:        rt,
			Profile:        profile,
			AgentsDir:      rs.cfg.AgentsDir,
			AgentID:        uname,
			CWD:            cwd,
			BaseLayers:     []string{sysPrompt, opModel, envBlock},
			PromptEnvExtra: promptEnv,
			NewToolServer:  newToolServer,
			NewBackend:     backend.NewWithName,
			Output:         NewEventHandler(al),
			Log:            rs.cfg.Sink.NewLogger("core"),
			AuditLog:       rs.cfg.Sink.NewLogger("core").Sub("audit"),
			SessionID:      sessID,
			ReadPlanStep: func() string {
				return ""
			},
			SaveSession: saveFn,
			FlushSave:   flushFn,
		})
		sess.Core.AddAgent(ag)
		sess.mu.Lock()
		sess.setAgentLogLocked(uname, al)
		sess.mu.Unlock()
	}

	rs.mu.Lock()
	if _, exists := rs.sessions[sessName]; !exists {
		rs.sessions[sessName] = sess
	}
	rs.mu.Unlock()

	ag := sess.Core.FindAgent(uname)
	ag.Submit(sess.SessionCtx, "/agent "+profile)
	if name != "" {
		ag.SetName(name)
	}

	rs.cfg.Log.Info("created agent in session %s (backend=%s model=%s agent=%s)",
		sessName, ag.BackendName(), ag.ModelName(), ag.Name())
	PostEvent("A new session/" + sessName + "/agent/" + ag.ID())
	return nil
}

// CreateAgentFromRoot is the public API for callers that have a *Tree.
func CreateAgentFromRoot(root *Tree, sessName string, args []string) (string, error) {
	rs := root.Data.(*rootState)
	if err := CreateAgent(rs, sessName, args); err != nil {
		return "", err
	}
	rs.mu.RLock()
	sess := rs.sessions[sessName]
	rs.mu.RUnlock()
	if sess == nil || sess.Core == nil || sess.Core.AgentCount() == 0 {
		return "", fmt.Errorf("agent not created")
	}
	agents := sess.Core.Agents()
	return agents[len(agents)-1].ID(), nil
}