ollie/cmd/olliesrv/internal/session/registry.go

// Package session manages the lifecycle of agent sessions.
//
// A session owns one or more agents, a tool server connection, and persistence
// state. The package-level registry (initialized via Init) tracks all active
// sessions and provides lookup, creation, and teardown.
//
// Typical flow: Init → Create → Register → (use) → Kill or Shutdown.
package session

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

	"ollie/cmd/olliesrv/internal/agent"
	"ollie/cmd/olliesrv/internal/backend"
	olog "ollie/log"
	"ollie/util"
)

// Package-level session registry with thread-safe access.
var (
	mu       sync.RWMutex
	sessions = make(map[string]*Session)

	// Package config, set via Init.
	serverCtx      context.Context
	pkgLog         *olog.Logger
	pkgSink        *olog.Sink
	pkgAgentsDir   string
	pkgSessionsDir string
	pkgYolo        bool
)

// InitConfig configures the session package.
type InitConfig struct {
	Ctx         context.Context
	Log         *olog.Logger
	Sink        *olog.Sink
	AgentsDir   string
	SessionsDir string
	Yolo        bool
}

// Init initializes the session package with the given configuration.
func Init(cfg InitConfig) {
	serverCtx = cfg.Ctx
	pkgLog = cfg.Log
	pkgSink = cfg.Sink
	pkgAgentsDir = cfg.AgentsDir
	pkgSessionsDir = cfg.SessionsDir
	pkgYolo = cfg.Yolo
}

// Sessions returns a snapshot of all sessions.
func Sessions() map[string]*Session {
	mu.RLock()
	defer mu.RUnlock()
	result := make(map[string]*Session, len(sessions))
	for k, v := range sessions {
		result[k] = v
	}
	return result
}

// Lookup returns the session by name or ID, or nil.
func Lookup(nameOrID string) *Session {
	mu.RLock()
	defer mu.RUnlock()
	if s := sessions[nameOrID]; s != nil {
		return s
	}
	for _, sess := range sessions {
		if sess.ID == nameOrID {
			return sess
		}
	}
	return nil
}

// IsSubagent returns true if the given agent ID belongs to a sub-agent (depth > 0).
func IsSubagent(agentID string) bool {
	mu.RLock()
	defer mu.RUnlock()
	for _, sess := range sessions {
		if ag := sess.FindAgent(agentID); ag != nil {
			return ag.Depth() > 0
		}
	}
	return false
}

// Register adds a session to the registry.
func Register(name string, sess *Session) {
	mu.Lock()
	sessions[name] = sess
	mu.Unlock()
}

// Kill terminates a session and removes it from the registry.
func Kill(name string) error {
	mu.Lock()
	sess := sessions[name]
	if sess == nil {
		mu.Unlock()
		return fmt.Errorf("session not found: %s", name)
	}
	delete(sessions, name)
	mu.Unlock()

	RemovePersistedSession(sess.ID)

	var agentIDs []string
	for _, ag := range sess.Agents() {
		agentIDs = append(agentIDs, ag.ID())
	}
	sess.Close()
	if sess.Cancel != nil {
		sess.Cancel()
	}
	if sess.Keeper != nil {
		sess.Keeper.Close()
	} else if sess.Proc != nil {
		sess.Proc.Close()
	}

	pkgLog.Info("killed session %s", name)
	for _, agentID := range agentIDs {
		PublishEvent("session."+sess.ID+".agent."+agentID+".kill", "")
	}
	PublishEvent("session."+sess.ID+".kill", "")
	return nil
}

// Rename renames a session.
func Rename(oldName, newName string) error {
	mu.Lock()
	sess, ok := sessions[oldName]
	if !ok {
		mu.Unlock()
		return fmt.Errorf("session not found: %s", oldName)
	}
	if _, exists := sessions[newName]; exists {
		mu.Unlock()
		return fmt.Errorf("session already exists: %s", newName)
	}
	for _, ag := range sess.Agents() {
		if ag.IsRunning() {
			mu.Unlock()
			return fmt.Errorf("cannot rename while agent is running")
		}
	}
	sess.SetName(newName)
	sessions[newName] = sess
	delete(sessions, oldName)
	mu.Unlock()

	PersistSession(newName)

	pkgLog.Info("renamed session %s -> %s", oldName, newName)
	PublishEvent("session."+sess.ID+".rename", oldName+" "+newName)
	return nil
}

// Shutdown pauses all sessions (to allow later resume), waits for idle, then closes them.
func Shutdown() {
	// First, pause all sessions to persist their state.
	PauseAll()
	waitIdle(100*time.Millisecond, 5*time.Second)

	mu.Lock()
	names := make([]string, 0, len(sessions))
	for name := range sessions {
		names = append(names, name)
	}
	mu.Unlock()

	for _, name := range names {
		mu.Lock()
		sess := sessions[name]
		delete(sessions, name)
		mu.Unlock()
		if sess != nil {
			// Session is already paused, so Cancel/Keeper/Proc should be nil.
			// But close anything still open defensively.
			if sess.Cancel != nil {
				sess.Cancel()
			}
			sess.Close()
			if sess.Keeper != nil {
				sess.Keeper.Close()
			} else if sess.Proc != nil {
				sess.Proc.Close()
			}
			pkgLog.Info("shutdown session %s", name)
		}
	}
}

// PauseAll pauses every active session. Agents are interrupted first.
func PauseAll() {
	InterruptAll()
	mu.RLock()
	var toProcess []*Session
	for _, sess := range sessions {
		if !sess.IsPaused() {
			toProcess = append(toProcess, sess)
		}
	}
	mu.RUnlock()

	for _, sess := range toProcess {
		if err := sess.Pause(); err != nil {
			pkgLog.Warn("pause session %s: %v", sess.ID, err)
		}
	}
}

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

func waitIdle(poll, timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		allIdle := true
		mu.RLock()
		for _, sess := range sessions {
			for _, ag := range sess.Agents() {
				if ag.State() != "idle" {
					allIdle = false
					break
				}
			}
			if !allIdle {
				break
			}
		}
		mu.RUnlock()
		if allIdle {
			return
		}
		time.Sleep(poll)
	}
	pkgLog.Warn("shutdown: timed out waiting for idle")
}

// BuildIndex builds the session index content.
// Format: session-id\tsession-name\tpaused\tconnected\tremote\tcwd
func BuildIndex() []byte {
	var sb strings.Builder
	mu.RLock()
	defer mu.RUnlock()

	type entry struct {
		name   string
		sess   *Session
		paused bool
	}
	entries := make([]entry, 0, len(sessions))
	for name, sess := range sessions {
		entries = append(entries, entry{name: name, sess: sess, paused: sess.IsPaused()})
	}
	sort.Slice(entries, func(i, j int) bool {
		if entries[i].paused != entries[j].paused {
			return !entries[i].paused
		}
		return entries[i].name < entries[j].name
	})

	for _, e := range entries {
		sess := e.sess
		sess.mu.RLock()
		connected := sess.toolsConn != nil && !sess.paused
		sess.mu.RUnlock()
		fmt.Fprintf(&sb, "%s\t%s\t%t\t%t\t%s\t%s\n",
			sess.ID, e.name, e.paused, connected, sess.Remote, sess.Cwd())
	}
	return []byte(sb.String())
}

// BuildAgentIndex builds the agent index for a session.
// Format: session-id\tagent-id\tagent-name\tparent-id\tdepth\tstate
func BuildAgentIndex(sess *Session) []byte {
	var sb strings.Builder
	for _, ag := range sess.Agents() {
		fmt.Fprintf(&sb, "%s\t%s\t%s\t%s\t%d\t%s\n",
			sess.ID, ag.ID(), ag.Name(), ag.ParentID(), ag.Depth(), ag.State())
	}
	return []byte(sb.String())
}

// CreateEmpty creates an empty session (no agent) and registers it.
// CreateEmpty returns the session named `name`, creating it if it does not yet
// exist (get-or-create). created reports whether a new session was made; when
// false the returned session already existed and its configuration is left
// untouched — creation never mutates an existing session.
func CreateEmpty(name, remote string, yolo ...bool) (sess *Session, created bool, err error) {
	sessID := util.NewUUID()
	if name == "" {
		name = sessID
		if idx := strings.IndexByte(sessID, '-'); idx > 0 {
			name = sessID[:idx]
		}
	}

	// Early path: return the existing session without building a tool server.
	mu.RLock()
	existing := sessions[name]
	mu.RUnlock()
	if existing != nil {
		return existing, false, nil
	}

	ctx, cancel := context.WithCancel(serverCtx)
	sess = NewEmpty(sessID, ctx, cancel)
	sess.SetName(name)
	sess.Remote = remote
	sess.Yolo = len(yolo) > 0 && yolo[0]

	cwd, _ := os.Getwd()
	infra, err := SetupToolServer(ToolServerConfig{
		Ctx:          ctx,
		CWD:          cwd,
		RemoteTarget: remote,
		SessionID:    sessID,
		Yolo:         sess.Yolo,
	})
	if err != nil {
		cancel()
		return nil, false, fmt.Errorf("setup tool server: %w", err)
	}
	sess.Proc = infra.Proc
	sess.Keeper = infra.Keeper
	sess.SetToolsConn(infra.ToolsConn)

	// Start bypass approval loop
	sess.StartBypassLoop()

	// Atomic check-and-insert. If another caller won the race, discard our
	// freshly-built session and return theirs — still get-or-create.
	mu.Lock()
	if won := sessions[name]; won != nil {
		mu.Unlock()
		cancel()
		return won, false, nil
	}
	sessions[name] = sess
	mu.Unlock()

	pkgLog.Info("created empty session %s (id=%s)", name, sessID)
	PublishEvent("session."+sessID+".new", name)
	return sess, true, nil
}

// NextUname generates a unique agent name (12 hex chars).
func NextUname() string {
	b := make([]byte, 6)
	rand.Read(b)
	return fmt.Sprintf("%06x%06x", b[:3], b[3:6])
}

// --- Shared agent construction ---

// AgentParams describes the parameters for building an agent within a session.
type AgentParams struct {
	Profile  string
	ID       string         // uname; if empty, generate one
	Name     string         // display name; optional
	Backend  string         // backend name; empty = config default
	Model    string         // model override; empty = config default
	Remote   string         // remote target override
	ParentID string         // immutable ID of the agent that spawned this agent
	History  *agent.History // nil = fresh agent
	CwdOverride string      // per-agent cwd override; empty = inherit session cwd
}

// buildAgent constructs a fully-wired agent within a session.
// Uses the session's CWD.
func buildAgent(sess *Session, p AgentParams) (*agent.Agent, error) {
	if p.ID == "" {
		p.ID = NextUname()
	}
	cwd := sess.Cwd()
	if cwd == "" {
		return nil, fmt.Errorf("session cwd is required")
	}

	// effectiveCwd is where this agent's tools run and what its runtime sees.
	// It is the per-agent override when set, otherwise the session cwd. The
	// shared tool server process stays configured with the session cwd; the
	// override is applied per-agent on the tool server below.
	effectiveCwd := cwd
	if p.CwdOverride != "" {
		effectiveCwd = util.ExpandHome(p.CwdOverride)
	}

	cfg, err := agent.LoadConfig(pkgAgentsDir, p.Profile, nil)
	if err != nil {
		return nil, fmt.Errorf("load agent config %q: %w", p.Profile, err)
	}
	sessID := sess.ID

	// Set up tool server — reuse existing session infrastructure
	var reuseFrom *InfraConfig
	if sess.Proc != nil {
		reuseFrom = &InfraConfig{
			Proc:      sess.Proc,
			Keeper:    sess.Keeper,
			ToolsConn: sess.ToolsConn(),
		}
	}

	remote := p.Remote
	if remote == "" {
		remote = sess.Remote
	}

	infra, err := SetupToolServer(ToolServerConfig{
		Ctx:          sess.Ctx,
		CWD:          cwd,
		RemoteTarget: remote,
		SessionID:    sessID,
		Yolo:         sess.Yolo,
		ReuseFrom:    reuseFrom,
	})
	if err != nil {
		return nil, fmt.Errorf("setup tool server: %w", err)
	}

	// First agent in a session sets the shared infra
	if sess.Proc == nil {
		sess.Proc = infra.Proc
		sess.Keeper = infra.Keeper
		sess.SetToolsConn(infra.ToolsConn)
	}

	// Build prompt layers
	layers := BuildPromptLayers(cfg, effectiveCwd, sessID, p.ID, infra.Platform, infra.IsGitRepo, "")

	// Load tools
	LoadTools(cfg, infra.ToolsConn, sessID, p.ID, func(f string, a ...any) {
		pkgLog.Error("session %s agent %s: "+f, append([]any{sessID, p.ID}, a...)...)
	})

	// Build runtime
	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + p.ID}
	rt := agent.BuildRuntime(cfg, infra.ToolsConn, effectiveCwd, env, layers.SystemPrompt, layers.EnvBlock)

	// Create backend
	beName := p.Backend
	if beName == "" && cfg != nil {
		beName = cfg.Backend
	}
	be, err := backend.NewWithName(beName)
	if err != nil {
		return nil, fmt.Errorf("create backend: %w", err)
	}
	rt.Backend = be

	if p.Model != "" {
		rt.Backend.SetModel(p.Model)
	} else if cfg != nil && cfg.Model != "" {
		rt.Backend.SetModel(cfg.Model)
	}

	ag := agent.NewAgent(agent.AgentParams{
		History:       p.History,
		Runtime:       rt,
		Profile:       p.Profile,
		AgentsDir:     pkgAgentsDir,
		ID:            p.ID,
		GetCwd:        sess.Cwd,
		ParentID:      p.ParentID,
		SystemPrompt:  layers.SystemPrompt,
		EnvBlock:      layers.EnvBlock,
		NewToolServer: infra.NewToolServer,
		NewBackend:    backend.NewWithName,
		Log:           pkgSink.NewLogger("agent"),
		SessionID:     sessID,
		Save:          sess.saveSession,
		Flush:         sess.flushSave,
	})

	if p.Name != "" {
		ag.SetName(p.Name)
	}
	if p.CwdOverride != "" {
		ag.SetCwdOverride(p.CwdOverride)
	}
	return ag, nil
}

// CreateAgentWithParams creates a new agent within an existing session using
// a pre-built AgentParams struct.
func CreateAgentWithParams(sessName string, p AgentParams) (*agent.Agent, error) {
	sess := Lookup(sessName)
	if sess == nil {
		return nil, fmt.Errorf("session not found: %s", sessName)
	}

	if p.Profile == "" {
		p.Profile = "default"
	}

	if sess.Cwd() == "" {
		return nil, fmt.Errorf("session cwd is required")
	}

	ag, err := buildAgent(sess, p)
	if err != nil {
		return nil, err
	}

	if err := sess.AddAgent(ag); err != nil {
		ag.Close()
		return nil, err
	}
	pkgLog.Info("created agent %s in session %s (profile=%s)", ag.Name(), sessName, p.Profile)
	return ag, nil
}