ollie/session/registry.go

package session

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

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

// 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 // server lifetime context; parent of all sessions
	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.
// Must be called before any other session functions.
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
}

// 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()

	// Remove from persistent store — kill is unconditional
	RemovePersistedSession(sess.ID())

	var agentIDs []string
	for _, ag := range sess.Agents() {
		agentIDs = append(agentIDs, ag.ID())
	}
	sess.Close()
	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()

	// Persist immediately so the rename survives restart
	PersistSession(newName)

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

// Shutdown interrupts all sessions, waits for idle, and closes them.
func Shutdown() {
	InterruptAll()
	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 {
			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)
		}
	}
}

// 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.
// Active sessions are listed first, paused sessions last.
func BuildIndex() []byte {
	var sb strings.Builder
	mu.RLock()
	defer mu.RUnlock()

	// Collect sessions and sort: active first (alphabetically), then paused (alphabetically)
	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 {
		// Active (not paused) comes before paused
		if entries[i].paused != entries[j].paused {
			return !entries[i].paused // active first
		}
		return entries[i].name < entries[j].name
	})

	for _, e := range entries {
		sess := e.sess
		connected := sess.IsConnected()
		var cwd, be, model string
		var agentList []string
		for _, ag := range sess.Agents() {
			if cwd == "" {
				cwd = ag.Cwd()
			}
			if be == "" {
				be = ag.BackendName()
				model = ag.ModelName()
			}
			agentList = append(agentList, fmt.Sprintf("%s:%s:%s", ag.Name(), ag.ID(), ag.State()))
		}
		agents := strings.Join(agentList, ";")
		fmt.Fprintf(&sb, "%s\t%s\t%t\t%t\t%s\t%s\t%s\t%s\n",
			e.name, sess.ID(), e.paused, connected, cwd, be, model, agents)
	}
	return []byte(sb.String())
}

// BuildAgentIndex builds the agent index for a session.
// Format: name<tab>id<tab>cwd<tab>backend<tab>model per line.
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%s\n",
			ag.Name(), ag.ID(), ag.Cwd(), ag.BackendName(), ag.ModelName())
	}
	return []byte(sb.String())
}

// CreateEmpty creates an empty session (no agent).
// If remote is non-empty, spawns a remote tool server via SSH.
func CreateEmpty(name, remote string) (*Session, error) {
	sessID := NewSessionID()
	if name == "" {
		name = sessID
		if idx := strings.IndexByte(sessID, '-'); idx > 0 {
			name = sessID[:idx]
		}
	}

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

	ctx, cancel := context.WithCancel(serverCtx)
	sess := NewEmpty(sessID, ctx, cancel)
	sess.SetName(name)
	sess.SetRemote(remote)

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

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

	pkgLog.Info("created empty session %s (id=%s)", name, sessID)
	PublishEvent("session."+sessID+".new", name)
	return sess, 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])
}

// CreateAgent creates a new agent within an existing session.
// Args: [profile] [key=value...] where keys include name, backend, model, remote.
func CreateAgent(sessName string, args []string) (*agent.Agent, error) {
	sess := Lookup(sessName)
	if sess == nil {
		return nil, fmt.Errorf("session not found: %s", sessName)
	}

	// Parse args
	profile := "default"
	var agentName, backendName, modelName, remoteTarget, cwdOverride string
	for _, arg := range args {
		if k, v, ok := strings.Cut(arg, "="); ok {
			switch k {
			case "name":
				agentName = v
			case "backend":
				backendName = v
			case "model":
				modelName = v
			case "remote":
				remoteTarget = v
			case "cwd":
				cwdOverride = v
			}
		} else if profile == "default" {
			profile = arg
		}
	}

	cfg := agent.LoadConfig(pkgAgentsDir, profile, nil)

	uname := NextUname()
	sessID := sess.ID()
	if cwdOverride == "" {
		return nil, fmt.Errorf("cwd= is required when creating an agent")
	}
	cwd := cwdOverride

	// Set up tool server - reuse existing session infrastructure
	infra, err := SetupToolServer(ToolServerConfig{
		Ctx:          sess.Ctx(),
		CWD:          cwd,
		RemoteTarget: remoteTarget,
		Yolo:         pkgYolo,
		ReuseFrom: &InfraConfig{
			Proc:      sess.Proc(),
			Keeper:    sess.Keeper(),
			ToolsConn: sess.ToolsConn(),
		},
	})
	if err != nil {
		return nil, fmt.Errorf("setup tool server: %w", err)
	}

	// Build prompt layers
	layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "")

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

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

	// Create backend - use override if provided, otherwise use config or env default
	beName := backendName
	if beName == "" {
		beName = rt.CfgBackend
	}
	be, err := backend.NewWithName(beName)
	if err != nil {
		return nil, fmt.Errorf("create backend: %w", err)
	}
	rt.Backend = be

	// Apply model override if provided
	if modelName != "" {
		rt.Backend.SetModel(modelName)
	} else if rt.CfgModel != "" {
		rt.Backend.SetModel(rt.CfgModel)
	}

	// Create agent
	ag := agent.NewAgent(agent.AgentCfg{
		History:        nil, // new agent starts with empty history
		Runtime:        rt,
		Profile:        profile,
		AgentsDir:      pkgAgentsDir,
		AgentID:        uname,
		CWD:            cwd,
		BaseLayers:     layers.BaseLayers(),
		PromptEnvExtra: infra.PromptEnv,
		NewToolServer:  infra.NewToolServer,
		NewBackend:     backend.NewWithName,
		Log:            pkgSink.NewLogger("agent"),
		SessionID:      sessID,
		Save:           sess.saveSession,
		Flush:          sess.flushSave,
	})

	sess.AddAgent(ag)
	if agentName != "" {
		ag.SetName(agentName)
	}
	pkgLog.Info("created agent %s in session %s (profile=%s)", ag.Name(), sessName, profile)
	return ag, nil
}