ollie/session/registry.go

package session

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

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

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

	// Package config, set via Init.
	pkgCtx         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.
// Must be called before any other session functions.
func Init(cfg InitConfig) {
	pkgCtx = 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()

	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.
func BuildIndex() []byte {
	var sb strings.Builder
	mu.RLock()
	names := make([]string, 0, len(sessions))
	for name := range sessions {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		sess := sessions[name]
		paused := sess.IsPaused()
		connected := sess.IsConnected()
		cwd := sess.CWD()
		var be, model string
		var agentList []string
		for _, ag := range sess.Agents() {
			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",
			name, sess.ID(), paused, connected, cwd, be, model, agents)
	}
	mu.RUnlock()
	return []byte(sb.String())
}

// CreateEmpty creates an empty session (no agent).
func CreateEmpty(name 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(pkgCtx)
	sess := NewEmpty(sessID, ctx, cancel)

	// Spawn tool server
	cwd, _ := os.Getwd()
	var dialOpts []toolsrv.Option
	if pkgYolo {
		dialOpts = append(dialOpts, toolsrv.WithYolo())
	}
	proc, err := toolsrv.Spawn(ctx, cwd, dialOpts...)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("spawn tool server: %w", err)
	}
	sess.SetProc(proc)
	sess.SetKeeper(toolsrv.NewProcessKeeper(ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
		return toolsrv.Spawn(ctx, cwd, dialOpts...)
	}))
	conn, err := toolsrv.Dial(ctx, toolsrv.LocalAddr(proc.Socket))
	if err != nil {
		cancel()
		proc.Close()
		return nil, fmt.Errorf("dial tool server: %w", err)
	}
	sess.SetToolsConn(conn)

	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 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 backendName, modelName, remoteTarget string
	for _, arg := range args {
		if k, v, ok := strings.Cut(arg, "="); ok {
			switch k {
			case "backend":
				backendName = v
			case "model":
				modelName = v
			case "remote":
				remoteTarget = v
			}
		} else if profile == "default" {
			profile = arg
		}
	}

	cfg := LoadAgentConfig(pkgAgentsDir, profile, nil)

	uname := NextUname()
	sessID := sess.ID()
	cwd := sess.CWD()

	// 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.OpModel, layers.EnvBlock)

	// Apply backend/model overrides to runtime
	if backendName != "" {
		if be, err := backend.NewWithName(backendName); err == nil {
			rt.Backend = be
		}
	}
	if modelName != "" {
		rt.Backend.SetModel(modelName)
	}

	// 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,
	})

	sess.AddAgent(ag)
	pkgLog.Info("created agent %s in session %s (profile=%s)", uname, sessName, profile)
	return ag, nil
}