ollie-core/session/session.go

package session

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/simonfxr/pubsub"
	"ollie/backend"
	olog "ollie/log"
	"ollie/paths"
	"ollie/tools"
)

// toolClassifier reports whether a named tool is safe to run concurrently
// with other read-class tools. nil means treat all tools as serial.
type toolClassifier func(name string) bool

// BuildRuntime constructs a Runtime from a pre-configured Dispatcher and
// optional agent config. cwd sets the working directory reported in the
// system prompt; if empty, the process working directory is used.
// env provides additional environment variables injected into prompt resolution
// subprocesses (e.g. OLLIE_SESSION_ID=xxx).
// The caller is responsible for registering all servers on d before calling this.
func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string, baseLayers ...string) *Runtime {
	var messages []string

	var allToolInfos []tools.ToolInfo
	var allTools []backend.Tool
	serverOf := make(map[string]string)

	if cfg == nil || cfg.ToolsEnabled() {
		var listErr error
		allToolInfos, listErr = d.ListTools()
		if listErr != nil {
			messages = append(messages, fmt.Sprintf("list tools: %v", listErr))
		}
		for _, t := range allToolInfos {
			serverOf[t.Name] = t.Server
		}
		// Only built-in executors (with InputSchema) become backend tools.
		// Named tool scripts are promoted via the tool registry and appear in the preamble.
		allTools = toolInfosToBackend(allToolInfos)

		// Append named tool scripts for preamble listing only.
		allToolInfos = append(allToolInfos, tools.DiscoverTools()...)
	}

	hooks := Hooks{}
	var preamble string
	var genParams backend.GenerationParams
	var maxSteps int
	if cfg != nil {
		for k, v := range cfg.Hooks {
			hooks[k] = []string(v)
		}
		if resolved, err := resolvePrompt(cfg.Prompt, cwd, env); err != nil {
			fmt.Fprintf(os.Stderr, "resolve prompt: %v\n", err)
		} else {
			preamble = resolved
		}
		genParams = backend.GenerationParams{
			MaxTokens:           cfg.MaxTokens,
			MaxCompletionTokens: cfg.MaxCompletionTokens,
			Temperature:         cfg.Temperature,
			TopP:                cfg.TopP,
			TopK:                cfg.TopK,
			MinP:                cfg.MinP,
			TopA:                cfg.TopA,
			FrequencyPenalty:    cfg.FrequencyPenalty,
			PresencePenalty:     cfg.PresencePenalty,
			RepetitionPenalty:   cfg.RepetitionPenalty,
			ThinkingBudget:      cfg.Reasoning,
			ReasoningEffort:     cfg.ReasoningEffort,
			IncludeReasoning:    cfg.IncludeReasoning,
			ResponseFormat:      cfg.ResponseFormat,
			Stop:                cfg.Stop,
			Verbosity:           cfg.Verbosity,
		}
		maxSteps = cfg.MaxSteps
		if len(cfg.AllowTools) > 0 {
			if srv, ok := d.GetServer("execute"); ok {
				if rs, ok := srv.(tools.ToolRestrictionSetter); ok {
					rs.SetAllowTools(cfg.AllowTools)
				}
			}
		}
	}

	exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) {
		// Dynamically resolve the server for this tool — handles
		// lazy-promoted tools that appeared after BuildRuntime.
		infos, listErr := d.ListTools()
		if listErr != nil {
			return "", nil, listErr
		}
		server := ""
		for _, t := range infos {
			if t.Name == name {
				server = t.Server
				break
			}
		}
		if server == "" {
			return "", nil, fmt.Errorf("unknown tool: %s", name)
		}
		raw, err := d.Dispatch(ctx, server, name, args)
		if err != nil {
			return "", nil, err
		}
		text, blocks, isErr := extractToolResult(raw)
		if isErr {
			return "", nil, fmt.Errorf("%s", text)
		}
		return text, blocks, nil
	}

	var classify toolClassifier
	if srv, ok := d.GetServer("execute"); ok {
		if pc, ok := srv.(tools.ParallelClassifier); ok {
			classify = pc.IsParallelRead
		}
	}

	var tierFn func(string, json.RawMessage) ResultTier
	if srv, ok := d.GetServer("execute"); ok {
		if tc, ok := srv.(tools.TierClassifier); ok {
			tierFn = func(name string, args json.RawMessage) ResultTier {
				switch tc.ResultTierArgs(name, args) {
				case "cold":
					return TierCold
				case "warm":
					return TierWarm
				default:
					return TierHot
				}
			}
		}
	}

	var backendName, modelName, compactionModel string
	if cfg != nil {
		backendName = cfg.Backend
		modelName = cfg.Model
		compactionModel = cfg.CompactionModel
	}

	// Prepend base layers (system prompt, operational model, environment)
	// before the agent-specific preamble.
	if len(baseLayers) > 0 {
		var prefix strings.Builder
		for _, layer := range baseLayers {
			if layer != "" {
				prefix.WriteString(layer)
				prefix.WriteByte('\n')
			}
		}
		if prefix.Len() > 0 {
			preamble = prefix.String() + preamble
		}
	}

	// Append compact tool surface listing (name + description).
	// Full prompts are available on-demand via /tools write (net/dns pattern).
	var toolListing strings.Builder
	for _, ti := range allToolInfos {
		if ti.Description != "" && ti.Server == "" {
			// Only named tool scripts (Server==""), not built-in executors
			fmt.Fprintf(&toolListing, "- **%s** — %s\n", ti.Name, ti.Description)
		}
	}
	if toolListing.Len() > 0 {
		preamble += "\n# Available Tools\n\n" + toolListing.String()
	}

	return &Runtime{
		Dispatcher:      d,
		Tools:           allTools,
		Exec:            exec,
		ClassifyTool:    classify,
		ClassifyTier:    tierFn,
		Hooks:           hooks,
		Preamble:        preamble,
		GenParams:       genParams,
		MaxSteps:        maxSteps,
		CfgBackend:      backendName,
		CfgModel:        modelName,
		CompactionModel: compactionModel,
		Messages:        messages,
	}
}

// DefaultPromptsDir returns the default directory for prompt templates.
func DefaultPromptsDir() string {
	return paths.CfgDir() + "/prompts"
}

// PromptsDirs returns all prompt directories from OLLIE_PROMPTS_PATH (colon-separated).
func PromptsDirs() []string {
	if p := os.Getenv("OLLIE_PROMPTS_PATH"); p != "" {
		return strings.Split(p, ":")
	}
	return []string{DefaultPromptsDir()}
}

// AgentsDirs returns all agent directories from OLLIE_AGENTS_PATH (colon-separated).
func AgentsDirs() []string {
	if p := os.Getenv("OLLIE_AGENTS_PATH"); p != "" {
		return strings.Split(p, ":")
	}
	return []string{paths.CfgDir() + "/agents"}
}

// AgentConfigPath resolves the config file path for a named agent.
// It searches all directories in OLLIE_AGENTS_PATH, falling back to agentsDir.
func AgentConfigPath(agentsDir, name string) string {
	for _, dir := range AgentsDirs() {
		p := dir + "/" + name + ".json"
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}
	return agentsDir + "/" + name + ".json"
}

// NewSessionID generates a unique, lexicographically sortable session identifier.
// Format: <unix-nanoseconds>-<random-hex> — sortable by creation time, unique
// even if two sessions are created within the same nanosecond.
func NewSessionID() string {
	b := make([]byte, 3)
	rand.Read(b) //nolint:errcheck
	return strconv.FormatInt(time.Now().UnixNano(), 10) + "-" + fmt.Sprintf("%06x", b)
}

// NewResponseID generates a unique identifier for a single assistant response.
func NewResponseID() string {
	b := make([]byte, 3)
	rand.Read(b) //nolint:errcheck
	return "resp_" + strconv.FormatInt(time.Now().UnixNano(), 10) + "_" + fmt.Sprintf("%06x", b)
}

// NewReactionID generates a unique identifier for a user reaction.
func NewReactionID() string {
	b := make([]byte, 3)
	rand.Read(b) //nolint:errcheck
	return "react_" + strconv.FormatInt(time.Now().UnixNano(), 10) + "_" + fmt.Sprintf("%06x", b)
}

// infoEvent wraps a plain-text message as an info Event.
func infoEvent(text string) Event {
	return Event{Role: "info", Content: text + "\n"}
}

// actionHandle holds the cancel function for the current agent turn.
type actionHandle struct {
	cancel context.CancelCauseFunc
}

// Config is the configuration for creating an agent.
type Config struct {
	Backend       backend.Backend
	ModelName     string // if non-empty, overrides backend's default model
	AgentName     string
	AgentsDir     string
	SessionsDir   string
	SessionID     string
	AgentID       string // unique agent identity
	CWD           string // working directory for tool execution and system prompt
	History       *History
	Runtime       *Runtime
	NewDispatcher func() tools.Dispatcher
	NewBackend    func(string) (backend.Backend, error) // if nil, defaults to backend.NewWithName
	Log           *olog.Logger                          // if nil, logging is disabled
	// MaxSteps overrides the agent JSON maxSteps when non-zero.
	// 0 means use the value from the agent config (or unlimited if absent).
	MaxSteps int
	// ReadPlanStep, if non-nil, is called to get the next unchecked plan step.
	// When provided, it is used directly instead of reading from the filesystem.
	ReadPlanStep func() string
	// ListHandlers provides session-specific list implementations for slash
	// commands such as /skills, /tools, /agents. Keyed by command name without
	// the leading slash. Commands with no entry emit nothing.
	ListHandlers map[string]func() []string
	// PromptEnvExtra holds PRIME_* env vars for prompt resolution.
	// For remote sessions these come from HostInfo; for local, from PromptEnv().
	// Stored so /agent reloads use the correct values.
	PromptEnvExtra []string
	// Remote is the SSH target for remote execution (empty = local).
	Remote string
	// BaseLayers are the system prompt, operational model, and environment
	// block prepended to the agent preamble. Stored so /agent reloads
	// preserve the server-injected context.
	BaseLayers []string
}

// Session is the concrete session type. It owns all session and agent state
// but has no knowledge of how output is rendered.
type Session struct {
	// Session-level state
	id       string
	bus      *pubsub.Bus
	envMu    sync.RWMutex
	env      map[string]string
	plan     []byte
	prevPrompt string

	// Agent
	r            *Agent
	log             *olog.Logger
	sessionsDir     string
	listHandlers    map[string]func() []string
	remote          string   // SSH target for remote execution
	mu              sync.RWMutex
	auditLog        *olog.Logger

	// Debounced session persistence
	saveMu    sync.Mutex
	saveDirty bool
	saveTimer *time.Timer

}

// ToolCallCount returns the total number of tool calls executed in this
// session. The counter is monotonically increasing and never resets.
// Blocked calls (pre-tool hook exit 2) are not counted.
func (a *Session) ToolCallCount() int64 {
	return a.r.toolCallCount.Load()
}

// SetEnv stores a session-scoped variable and propagates it to the execute server.
func (a *Session) SetEnv(key, value string) {
	a.envMu.Lock()
	a.env[key] = value
	a.envMu.Unlock()
	if a.r.runtime == nil || a.r.runtime.Dispatcher == nil {
		return
	}
	if srv, ok := a.r.runtime.Dispatcher.GetServer("execute"); ok {
		if es, ok := srv.(tools.EnvSetter); ok {
			es.SetEnv(key, value)
		}
	}
}

// pushSessionEnv injects OLLIE_SESSION_ID into the execute server subprocess env.
func (a *Session) pushSessionEnv() {
	if a.r.runtime == nil || a.r.runtime.Dispatcher == nil || a.id == "" {
		return
	}
	if srv, ok := a.r.runtime.Dispatcher.GetServer("execute"); ok {
		if es, ok := srv.(tools.EnvSetter); ok {
			es.SetEnv("OLLIE_SESSION_ID", a.id)
			if a.r.id != "" {
				es.SetEnv("OLLIE_UNAME", a.r.id)
			}
		}
	}
}

// pushLockDir sets the flock directory on the execute server to the session tmpdir.
func (a *Session) pushLockDir() {
	if a.r.runtime == nil || a.r.runtime.Dispatcher == nil || a.id == "" {
		return
	}

}


var sweepTmpOnce sync.Once

// ollieTmpDir returns the base temp directory for session tmpdirs.
func ollieTmpDir() string {
	if p := os.Getenv("OLLIE_TMP_PATH"); p != "" {
		return p
	}
	return filepath.Join(os.TempDir(), "ollie")
}

// readNextUnchecked reads a plan file and returns the first unchecked step.
func readNextUnchecked(path string) string {
	data, err := os.ReadFile(path)
	if err != nil || len(data) == 0 {
		return ""
	}
	return NextUncheckedStep(data)
}

// NextUncheckedStep returns the first unchecked step from plan bytes.
func NextUncheckedStep(data []byte) string {
	for _, line := range strings.Split(string(data), "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "- [ ]") {
			return strings.TrimSpace(trimmed[5:])
		}
	}
	return ""
}

// sweepStaleTmpDirs removes session tmpdirs left by a previous crash.
// All session tmpdirs are owned by a single process, so anything
// present at startup is stale.
func sweepStaleTmpDirs() {
	sweepTmpOnce.Do(func() {
		base := ollieTmpDir()
		os.RemoveAll(base)      //nolint:errcheck
		os.MkdirAll(base, 0700) //nolint:errcheck
	})
}

// New creates an agent from the given configuration.
func New(cfg Config) *Session {
	sweepStaleTmpDirs()
	if cfg.ModelName != "" {
		cfg.Backend.SetModel(cfg.ModelName)
	}
	if cfg.NewBackend == nil {
		cfg.NewBackend = backend.NewWithName
	}
	rt := cfg.Runtime
	if rt == nil {
		rt = &Runtime{}
	}
	// Store the backend on the runtime so it's the single source of truth.
	rt.Backend = cfg.Backend
	// A non-zero MaxSteps in Config takes precedence over the
	// value loaded from the agent JSON.
	if cfg.MaxSteps > 0 {
		rt.MaxSteps = cfg.MaxSteps
	}
	var readPlanStep func() string
	if cfg.SessionID != "" {
		os.MkdirAll(filepath.Join(ollieTmpDir(), cfg.SessionID), 0700) //nolint:errcheck
	}
	if cfg.ReadPlanStep != nil {
		readPlanStep = cfg.ReadPlanStep
	}

	log := cfg.Log
	if log == nil {
		log = olog.NewWriter("core", olog.LevelError+1, io.Discard, io.Discard)
	}

	bus := pubsub.NewBus()
	auditLog := log.Sub("audit")

	a := &Session{
		id:  cfg.SessionID,
		bus: bus,
		env: make(map[string]string),
		r: &Agent{
			history:         cfg.History,
			runtime:         rt,
			cwd:             paths.ExpandHome(cfg.CWD),
			id:              cfg.AgentID,
			agentName:       cfg.AgentName,
			agentsDir:       cfg.AgentsDir,
			promptEnvExtra:  cfg.PromptEnvExtra,
			baseLayers:      cfg.BaseLayers,
			newDispatcher:   cfg.NewDispatcher,
			newBackend:      cfg.NewBackend,
			bus:             bus,
			log:             log,
			auditLog:        auditLog,
			sessionID:       cfg.SessionID,
			startupMessages: rt.Messages,
			readPlanStep:    readPlanStep,
		},
		log:          log,
		auditLog:     auditLog,
		sessionsDir:  cfg.SessionsDir,
		remote:       cfg.Remote,
		listHandlers: cfg.ListHandlers,
	}
	a.r.InitCond()
	a.r.state = "idle"
	a.r.flushSave = a.flushSave
	a.r.saveSession = a.saveSession
	a.r.turnError = a.r.defaultTurnError
	a.pushSessionEnv()
	a.pushLockDir()
	return a
}

// Close releases resources for this session, including its tmpdir.
func (a *Session) Close() {
	a.log.Debug("Close() session=%q", a.id)
	a.flushSave()
	if a.r.runtime != nil && a.r.runtime.Dispatcher != nil {
		if srv, ok := a.r.runtime.Dispatcher.GetServer("execute"); ok {
			if c, ok := srv.(interface{ Close() }); ok {
				a.log.Debug("Close() calling execute.Close()")
				c.Close()
			}
		}
	}
	if a.id != "" {
		os.RemoveAll(filepath.Join(ollieTmpDir(), a.id)) //nolint:errcheck
	}
}

// execServer returns the execute server if available, or nil.
func (a *Session) execServer() interface{} {
	if a.r.runtime == nil || a.r.runtime.Dispatcher == nil {
		return nil
	}
	srv, _ := a.r.runtime.Dispatcher.GetServer("execute")
	return srv
}

func (a *Session) Detach() bool {
	if srv := a.execServer(); srv != nil {
		if d, ok := srv.(interface{ Detach() bool }); ok {
			return d.Detach()
		}
	}
	return false
}

func (a *Session) ListDetached() []DetachedInfo {
	if srv := a.execServer(); srv != nil {
		type listDetacher interface {
			ListDetachedRaw() []any
		}
		if ld, ok := srv.(listDetacher); ok {
			raw := ld.ListDetachedRaw()
			out := make([]DetachedInfo, 0, len(raw))
			for _, r := range raw {
				if m, ok := r.(map[string]any); ok {
					di := DetachedInfo{}
					if v, ok := m["pid"].(int); ok {
						di.PID = v
					}
					if v, ok := m["command"].(string); ok {
						di.Command = v
					}
					if v, ok := m["started"].(int64); ok {
						di.Started = v
					}
					if v, ok := m["exited"].(bool); ok {
						di.Exited = v
					}
					if v, ok := m["exit_code"].(int); ok {
						di.ExitCode = v
					}
					out = append(out, di)
				}
			}
			return out
		}
	}
	return nil
}

func (a *Session) SignalDetached(pid, signal int) error {
	if srv := a.execServer(); srv != nil {
		type signaler interface {
			SignalDetached(int, syscall.Signal) error
		}
		if sg, ok := srv.(signaler); ok {
			return sg.SignalDetached(pid, syscall.Signal(signal))
		}
	}
	return fmt.Errorf("no execute server available")
}

func (a *Session) GetDetachedOutput(pid int) (string, error) {
	if srv := a.execServer(); srv != nil {
		type outputGetter interface {
			GetDetachedOutput(int) (string, error)
		}
		if og, ok := srv.(outputGetter); ok {
			return og.GetDetachedOutput(pid)
		}
	}
	return "", fmt.Errorf("no execute server available")
}

func (a *Session) DismissDetached(pid int) bool {
	if srv := a.execServer(); srv != nil {
		type dismisser interface {
			DismissDetached(int) bool
		}
		if d, ok := srv.(dismisser); ok {
			return d.DismissDetached(pid)
		}
	}
	return false
}

func (a *Session) InjectSystemEvent(content string) {
	a.Queue("<detached-process-result>\n" + content + "\n</detached-process-result>")
}

// classifyReaction returns a category and description for a reaction emoji.
func classifyReaction(emoji string) (category, description string, positive bool) {
	switch emoji {
	case "👍", "✅":
		return "positive", "The response was good. Keep doing what you're doing.", true
	case "🚀", "🎉":
		return "excellent", "The response was exactly what was wanted.", true
	case "👎", "❌":
		return "negative", "The response was wrong or unhelpful.", false
	case "💩", "🤬":
		return "terrible", "The response was fundamentally wrong. Stop this approach entirely and reassess from scratch.", false
	case "🤔":
		return "confused", "The response was unclear or confusing.", false
	default:
		return "unknown", "", false
	}
}

func (a *Session) Reactions() map[string]string {
	result := make(map[string]string)
	if a.r.history == nil {
		return result
	}
	for _, reaction := range a.r.history.Reactions {
		result[reaction.ResponseID] = reaction.Emoji
	}
	return result
}

func (a *Session) React(emoji string) {
	_ = a.ReactTo("", emoji)
}

func (a *Session) ReactTo(responseID, emoji string) error {
	if a.r.history == nil {
		return fmt.Errorf("no active session")
	}
	category, _, _ := classifyReaction(emoji)
	if category == "unknown" {
		return fmt.Errorf("unsupported reaction: %s", emoji)
	}
	if responseID == "" {
		for i := len(a.r.history.messages) - 1; i >= 0; i-- {
			if a.r.history.messages[i].Role == "assistant" {
				responseID = a.r.history.messages[i].ID
				break
			}
		}
	}
	if responseID == "" {
		return fmt.Errorf("no assistant response to react to")
	}
	found := false
	for i := range a.r.history.messages {
		if a.r.history.messages[i].Role == "assistant" && a.r.history.messages[i].ID == responseID {
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("assistant response not found: %s", responseID)
	}

	reaction := Reaction{ID: NewReactionID(), ResponseID: responseID, Emoji: emoji, Category: category, CreatedAt: time.Now()}
	replaced := false
	for i := range a.r.history.Reactions {
		if a.r.history.Reactions[i].ResponseID == responseID {
			if a.r.history.Reactions[i].Emoji == emoji {
				return nil
			}
			a.r.history.Reactions[i] = reaction
			replaced = true
			break
		}
	}
	if !replaced {
		a.r.history.Reactions = append(a.r.history.Reactions, reaction)
	}
	a.r.history.recomputeReactionCounts()
	a.saveSession()
	return nil
}

func (a *Session) AgentName() string {
	v := a.r.Name()
	a.log.Debug("AgentName() = %q", v)
	return v
}
func (a *Session) BackendName() string {
	v := a.r.BackendName()
	a.log.Debug("BackendName() = %q", v)
	return v
}
func (a *Session) ModelName() string {
	v := a.r.ModelName()
	a.log.Debug("ModelName() = %q", v)
	return v
}

// Agent returns the active agent.
func (a *Session) Agent() *Agent { return a.r }

func (a *Session) State() string {
	return a.r.State()
}

func (a *Session) setState(state string) {
	a.log.Debug("state -> %q", state)
	a.r.SetState(state)
}

// WaitChange blocks until the named field changes from current, then returns
// the new value. Returns ("", false) if ctx is cancelled.
func (a *Session) WaitChange(ctx context.Context, field, current string) (string, bool) {
	// For state changes, delegate directly to the agent.
	if field == WatchState {
		return a.r.WaitChange(ctx, field, current)
	}

	// Other watched fields (usage, ctxsz, cwd, agent) are session-level;
	// they change as a side effect of agent activity, so we still listen
	// on the agent's change signal.
	read := func() string {
		switch field {
		case WatchUsage:
			return a.Usage()
		case WatchCtxSz:
			return a.CtxSz()
		case WatchCWD:
			return a.CWD()
		case WatchAgent:
			return a.AgentName()
		}
		return ""
	}

	stop := context.AfterFunc(ctx, func() {
		a.r.changeMu.Lock()
		a.r.changeCond.Broadcast()
		a.r.changeMu.Unlock()
	})
	defer stop()

	a.r.changeMu.Lock()
	defer a.r.changeMu.Unlock()
	for ctx.Err() == nil {
		if v := read(); v != current {
			return v, true
		}
		a.r.changeCond.Wait()
	}
	return "", false
}

func (a *Session) Reply() string {
	r := a.r.Reply()
	a.log.Debug("Reply() len=%d", len(r))
	return r
}

// CWD returns the current working directory for tool execution.
func (a *Session) CWD() string {
	cwd := a.r.Cwd()
	if cwd != "" {
		a.log.Debug("CWD() = %q", cwd)
		return cwd
	}
	wd, _ := os.Getwd()
	a.log.Debug("CWD() = %q (from getwd)", wd)
	return wd
}

// SetCWD changes the working directory for tool execution and updates the
// system prompt. Returns an error if the path does not exist.
func (a *Session) SetCWD(dir string) error {
	a.log.Debug("SetCWD(%q)", dir)
	dir = paths.ExpandHome(dir)
	if dir != "" {
		if _, err := os.Stat(dir); err != nil {
			return fmt.Errorf("cwd: %w", err)
		}
	}
	oldCwd := a.r.Cwd()
	a.r.SetCwd(dir)
	// Update cwd references in the system prompt.
	if oldCwd != "" && dir != "" && oldCwd != dir {
		a.r.runtime.Preamble = strings.ReplaceAll(a.r.runtime.Preamble, oldCwd, dir)
	}
	// Propagate to any tool server that knows how to handle it (e.g. execute).
	if a.r.runtime != nil && a.r.runtime.Dispatcher != nil {
		if srv, ok := a.r.runtime.Dispatcher.GetServer("execute"); ok {
			if ws, ok := srv.(tools.CWDSetter); ok {
				ws.SetCWD(dir)
			}
		}
	}
	a.r.notifyChange()
	return nil
}

// SetSessionID renames the session. It updates the in-memory ID, renames
// persisted files on disk, and propagates to the execute server env.
func (a *Session) SetSessionID(newID string) error {
	a.log.Debug("SetSessionID(%q) old=%q", newID, a.id)
	oldID := a.id
	if oldID == newID {
		return nil
	}
	// Rename active persisted files on disk.
	if a.sessionsDir != "" && oldID != "" {
		for _, suffix := range []string{".json", ".compaction.jsonl"} {
			oldPath := a.activeSessionPath(oldID, suffix)
			if _, err := os.Stat(oldPath); err == nil {
				if err := os.Rename(oldPath, a.activeSessionPath(newID, suffix)); err != nil {
					return fmt.Errorf("rename %s: %w", suffix, err)
				}
			}
		}
	}
	a.id = newID
	// Update session ID references in the system prompt.
	a.r.runtime.Preamble = strings.ReplaceAll(a.r.runtime.Preamble, oldID, newID)
	// Rename tmpdir so isread markers remain valid after rename.
	oldTemp := filepath.Join(ollieTmpDir(), oldID)
	newTemp := filepath.Join(ollieTmpDir(), newID)
	if _, err := os.Stat(oldTemp); err == nil {
		os.Rename(oldTemp, newTemp) //nolint:errcheck
	}
	a.pushSessionEnv()
	return nil
}

// defaultContextLength is used when the backend cannot report the model's
// actual context window (e.g. CodeWhisperer). 128k tokens is a safe default
// for modern models.
const defaultContextLength = 128000
const defaultToolResultMaxBytes = 131072

func (a *Session) activeSessionPath(id, suffix string) string {
	return filepath.Join(a.sessionsDir, "active", id+suffix)
}

func (a *Session) saveSession() {
	a.saveMu.Lock()
	a.saveDirty = true
	if a.saveTimer == nil {
		a.saveTimer = time.AfterFunc(2*time.Second, a.flushSave)
	}
	a.saveMu.Unlock()
}

// flushSave immediately persists the session if dirty.
func (a *Session) flushSave() {
	a.saveMu.Lock()
	dirty := a.saveDirty
	a.saveDirty = false
	if a.saveTimer != nil {
		a.saveTimer.Stop()
		a.saveTimer = nil
	}
	a.saveMu.Unlock()
	if !dirty {
		return
	}
	if a.r.history == nil || a.id == "" || a.sessionsDir == "" {
		return
	}
	path := a.activeSessionPath(a.id, ".json")
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		a.log.Error("session save: %v", err)
		return
	}
	if err := a.r.history.saveToFull(path, a.id, a.r.agentName,
		a.r.runtime.Backend.Name(), a.r.runtime.Backend.Model(), a.CWD(), a.remote); err != nil {
		a.log.Error("session save: %v", err)
	}
}

// SaveSession writes the current session state to the given path, including
// backend and model metadata for external restore.
func (a *Session) SaveSession(path string) error {
	a.mu.RLock()
	defer a.mu.RUnlock()
	if a.r.history == nil {
		return fmt.Errorf("no active session")
	}
	return a.r.history.saveToFull(path, a.id, a.r.agentName,
		a.r.runtime.Backend.Name(), a.r.runtime.Backend.Model(), a.CWD(), a.remote)
}

// Interrupt cancels the current in-progress agent turn.
// Returns true if an action was running and was cancelled.
func (a *Session) Interrupt(cause error) bool {
	a.log.Debug("Interrupt() cause=%v", cause)
	return a.r.Interrupt(cause)
}

func (a *Session) Inject(prompt string) {
	// If an inject is already pending, fall back to the normal FIFO so nothing
	// is lost. Use CompareAndSwap to avoid a race between the nil check and store.
	if !a.r.pendingInject.CompareAndSwap(nil, &prompt) {
		a.r.fifo.Push(prompt)
		return
	}
	a.emit(Event{Role: "info", Content: "\n"})
	a.emit(Event{Role: "user", Content: prompt})
}

func (a *Session) injectRewrite(prompt string) {
	a.r.pendingInject.Store(&prompt)
	a.emit(Event{Role: "info", Content: "\n"})
	a.emit(Event{Role: "user", Content: prompt})
}

func (a *Session) Queue(prompt string) {
	a.r.fifo.Push(prompt)
	a.bus.Publish("queued", prompt)
}

func (a *Session) drainQueue() {
	if prompt, ok := a.r.fifo.Pop(); ok {
		a.Submit(context.Background(), prompt)
	}
}

func (a *Session) Bus() *pubsub.Bus {
	return a.bus
}

func (a *Session) emit(ev Event) {
	a.bus.Publish("event", ev)
}

func (a *Session) PopQueue() (string, bool) {
	return a.r.fifo.Pop()
}

func (a *Session) IsRunning() bool {
	v := a.r.IsRunning()
	a.log.Debug("IsRunning() = %v", v)
	return v
}

func (a *Session) CtxSz() string {
	if a.r.history == nil {
		a.log.Debug("CtxSz() no session")
		return "no active session"
	}
	ctxLen := a.r.runtime.Backend.ContextLength(context.Background())
	if ctxLen <= 0 {
		ctxLen = defaultContextLength
	}
	estimated := a.r.history.estimateTokens()
	pct := estimated * 100 / ctxLen
	v := fmt.Sprintf("%d / %d (%d%%)", estimated, ctxLen, pct)
	a.log.Debug("CtxSz() = %q", v)
	return v
}

func (a *Session) Cost() string {
	if a.r.history == nil {
		return "no active session"
	}
	return fmt.Sprintf("costLast=$%.4f\ncostSession=$%.4f\n",
		a.r.history.LastTurnCostUSD, a.r.history.SessionCostUSD)
}

func (a *Session) Usage() string {
	if a.r.history == nil {
		a.log.Debug("Usage() no session")
		return "no active session"
	}
	str := fmt.Sprintf("%d in, %d out, %d requests",
		a.r.history.TotalInputTokens, a.r.history.TotalOutputTokens,
		a.r.history.TotalRequests)
	if a.r.history.TotalCachedInputTokens > 0 {
		str += fmt.Sprintf(", %d cached", a.r.history.TotalCachedInputTokens)
	}
	if a.r.history.Estimated {
		str += " [estimated]"
	}
	a.log.Debug("Usage() = %q", str)
	return str
}

func (a *Session) Context() []backend.Message {
	a.mu.RLock()
	var msgs []backend.Message
	if a.r.history != nil {
		msgs = slices.Clone(a.r.history.history())
	}
	a.mu.RUnlock()
	if a.r.runtime.Preamble != "" {
		msgs = append([]backend.Message{{Role: "system", Content: a.r.runtime.Preamble}}, msgs...)
	}
	return msgs
}

func (a *Session) SystemPrompt() string {
	a.log.Debug("SystemPrompt() len=%d", len(a.r.runtime.Preamble))
	return a.r.runtime.Preamble
}

func (a *Session) GenerationParams() backend.GenerationParams {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.r.runtime.GenParams
}

func (a *Session) CompactionModel() string {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.r.runtime.CompactionModel
}

func (a *Session) SetCompactionModel(model string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.r.runtime.CompactionModel = model
}

func (a *Session) SetGenerationParams(params backend.GenerationParams) error {
	if a.IsRunning() {
		return fmt.Errorf("cannot change params while agent is running")
	}
	a.mu.Lock()
	a.r.runtime.GenParams = params
	a.mu.Unlock()
	return nil
}

func (a *Session) ListModels() string {
	a.log.Debug("ListModels()")
	models := a.r.runtime.Backend.Models(context.Background())
	slices.Sort(models)
	return strings.Join(models, "\n")
}

// firstSentence returns the first sentence of s (up to the first period or
// newline), trimmed. Falls back to s truncated at 80 chars if no sentence end
// is found.
func firstSentence(s string) string {
	for i, r := range s {
		if r == '.' || r == '\n' {
			return strings.TrimSpace(s[:i+1])
		}
	}
	if len(s) > 80 {
		return s[:77] + "..."
	}
	return s
}

// Submit processes one line of user input: slash commands are dispatched
// immediately; any other input is delegated to the agent for turn execution.
func (a *Session) Submit(ctx context.Context, input string) {
	a.log.Debug("Submit() input_len=%d running=%v", len(input), a.IsRunning())
	if input == "" {
		return
	}

	// Commands are handled at the session level (they reference sessionsDir,
	// listHandlers, etc.). If the input is a command, handle it and return.
	if a.IsRunning() {
		if a.handleCommand(ctx, input) {
			return
		}
	} else {
		// Serialize commands with turns — acquire submitMu to check.
		a.r.submitMu.Lock()
		if a.handleCommand(ctx, input) {
			a.r.submitMu.Unlock()
			return
		}
		a.r.submitMu.Unlock()
	}

	// Delegate to the agent for turn execution (or FIFO queueing if running).
	a.r.Submit(ctx, input)
}

func toolInfosToBackend(infos []tools.ToolInfo) []backend.Tool {
	out := make([]backend.Tool, len(infos))
	for i, t := range infos {
		out[i] = backend.Tool{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.InputSchema,
		}
	}
	return out
}

func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) {
	var result struct {
		IsError bool `json:"isError"`
		Content []struct {
			Type      string `json:"type"`
			Text      string `json:"text"`
			MediaType string `json:"media_type"`
			Data      string `json:"data"`
		} `json:"content"`
	}
	if err := json.Unmarshal(raw, &result); err != nil {
		return string(raw), nil, false
	}
	var parts []string
	for _, c := range result.Content {
		switch c.Type {
		case "text":
			parts = append(parts, c.Text)
		case "image":
			contentBlocks = append(contentBlocks, backend.ContentBlock{
				Type: "image",
				ImageSource: &backend.ImageSource{
					Type:      "base64",
					MediaType: c.MediaType,
					Data:      c.Data,
				},
			})
		}
	}
	return strings.Join(parts, "\n"), contentBlocks, result.IsError
}

// WatchField names supported by Session.WaitChange.
const (
	WatchState = "state"
	WatchUsage = "usage"
	WatchCtxSz = "ctxsz"
	WatchCWD   = "cwd"
	WatchAgent = "agent"
)

// ErrInterrupted is returned when the user cancels an agent turn (Ctrl-C).
var ErrInterrupted = errors.New("interrupted")

// Event is a typed output event emitted during an agent turn or in response
// to a command.
type Event struct {
	Role       string
	Name       string
	Content    string
	ResponseID string
}

// EventHandler receives events from the agent.
type EventHandler func(Event)

// DetachedInfo describes a detached process for external consumers.
type DetachedInfo struct {
	PID      int
	Command  string
	Started  int64 // unix timestamp
	Exited   bool
	ExitCode int
}