ollie/fs/lifecycle.go

package fs

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

	"github.com/simonfxr/pubsub"

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

// 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.
// Exported for use by fshandlers package.
type RootState struct {
	Cfg      Config
	Mu       sync.RWMutex
	Sessions map[string]*Session
}

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

var eventBus = pubsub.NewBus()

// Event represents a published event with topic and payload.
type Event struct {
	Topic   string
	Payload string
}

// PublishEvent publishes an event to a specific topic.
// Topic hierarchy:
//   session.{sid}.new/kill/pause/resume/rename
//   session.{sid}.agent.{aid}.new/kill/state
//
// Wildcard routing publishes to all ancestor wildcards:
//   session.{sid}.agent.{aid}.state -> also publishes to:
//     session.{sid}.agent.{aid}.*
//     session.{sid}.agent.*
//     session.{sid}.*
//     session.*
//     *
func PublishEvent(topic, payload string) {
	ev := Event{Topic: topic, Payload: payload}
	eventBus.Publish(topic, ev)
	
	// Publish to hierarchical wildcards
	parts := strings.Split(topic, ".")
	for i := len(parts) - 1; i >= 1; i-- {
		wildcard := strings.Join(parts[:i], ".") + ".*"
		eventBus.Publish(wildcard, ev)
	}
	eventBus.Publish("*", ev)
}

// SubscribeEvents returns a channel that receives all events matching the topic.
// Use "*" for all events, "session.*" for all session events, etc.
func SubscribeEvents(ctx context.Context, topic string) <-chan Event {
	ch := make(chan Event, 64)
	sub := eventBus.SubscribeChan(topic, ch, pubsub.KeepOpen)
	go func() {
		<-ctx.Done()
		eventBus.Unsubscribe(sub)
		close(ch)
	}()
	return ch
}

// WaitEvent blocks until an event is available. Returns the event as "topic payload\n".
// For backwards compatibility with the old eventwait interface.
func WaitEvent(ctx context.Context, _ string) ([]byte, string, error) {
	ch := SubscribeEvents(ctx, "*")
	select {
	case ev, ok := <-ch:
		if !ok {
			return nil, "", nil
		}
		line := ev.Topic
		if ev.Payload != "" {
			line += " " + ev.Payload
		}
		return []byte(line + "\n"), "", nil
	case <-ctx.Done():
		return nil, "", nil
	}
}

// BuildIndex builds the session/idx content.
//
// Format: one line per session (not per agent)
// Fields (tab-separated):
//   name, id, paused, connected, cwd, backend, model, agents
// Where agents is a semicolon-separated list of: agentName:agentId:state
//
// Example:
//   myproject\tf60604bf-...\tfalse\ttrue\t/home/user/proj\tkiro\tclaude-opus-4.5\tdev:7c36c9e4:idle;test:abc123:thinking

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
		paused := sess.paused
		
		// Check connected status
		connected := false
		if conn := sess.toolsConn; conn != nil {
			if pinger, ok := conn.(interface{ Ping() error }); ok {
				connected = pinger.Ping() == nil
			} else {
				connected = true
			}
		}
		
		var cwd, be, model string
		var agentList []string
		if sess.Core != nil {
			cwd = sess.Core.CWD()
			for _, ag := range sess.Core.Agents() {
				if be == "" {
					be = ag.BackendName()
					model = ag.ModelName()
				}
				// Format: name:id:state
				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, sid, paused, connected, cwd, be, model, agents)
		sess.mu.RUnlock()
	}
	rs.Mu.RUnlock()
	return []byte(sb.String())
}

// 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 _, agentID := range agentNames {
		PublishEvent("session."+sess.id+".agent."+agentID+".kill", "")
	}
	PublishEvent("session."+sess.id+".kill", "")
}

// 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)
	PublishEvent("session."+sess.id+".rename", old+" "+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
}

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)
	PublishEvent("session."+sessID+".new", 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)
	}

	// Set up tool server infrastructure.
	// If session already has a proc from Create(), reuse it.
	var reuseFrom *InfraConfig
	if sess.proc != nil && remoteTarget == "" {
		reuseFrom = &InfraConfig{
			Proc:      sess.proc,
			Keeper:    sess.keeper,
			ToolsConn: sess.toolsConn,
		}
	}

	infra, err := SetupToolServer(ToolServerConfig{
		Ctx:          rs.Cfg.Ctx,
		CWD:          cwd,
		RemoteTarget: remoteTarget,
		Yolo:         rs.Cfg.Yolo,
		ReuseFrom:    reuseFrom,
	})
	if err != nil {
		return err
	}

	// Update session's keeper if we spawned new infrastructure.
	if reuseFrom == nil {
		sess.keeper = infra.Keeper
	}

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

	toolSrv := infra.ToolsConn
	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 {
			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.
	LoadAutoLoadTools(cfg, toolSrv, func(f string, a ...any) {
		rs.Cfg.Log.Error("session %s: "+f, append([]any{sessName}, a...)...)
	})

	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
	env = append(env, infra.PromptEnv...)
	rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, layers.SystemPrompt, layers.OpModel, layers.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:  infra.NewToolServer,
			PromptEnvExtra: infra.PromptEnv,
			BaseLayers:     layers.BaseLayers(),
			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 = infra.Proc
		sess.mu.Unlock()
	} else {
		rt.Backend = be
		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:     layers.BaseLayers(),
			PromptEnvExtra: infra.PromptEnv,
			NewToolServer:  infra.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)
		// Wire up state change events - capture sessID for hierarchical topic
		sid := sessID
		ag.SetOnStateChange(func(agentID, state string) {
			PublishEvent("session."+sid+".agent."+agentID+".state", state)
		})
		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())
	PublishEvent("session."+sessID+".agent."+ag.ID()+".new", ag.Name())
	return nil
}