ollie/fs/session/types.go

package fs

import (
	"context"
	"strings"
	"sync"
	"time"

	"ollie/agent"
	"ollie/session"
	"ollie/toolsrv"
)

// Session holds session-level state exposed via 9P.
// Agent-scoped state (log, plan, etc.) lives in AgentLog.
type Session struct {
	mu         sync.RWMutex
	id         string
	uname      string // immutable user principal (numeric UID), set at creation
	Core       *session.Session
	SessionCtx context.Context
	cancel     context.CancelFunc
	proc       *toolsrv.Process // toolsrv subprocess (session owns lifecycle)

	// AgentLog for the session's agent (will become a map when multi-agent lands).
	AgentLog *AgentLog

	// Cached model list (expensive API call; refreshed every 24h).
	modelsMu      sync.Mutex
	modelsCache   string
	modelsCacheAt time.Time
}

// AgentLog holds all per-agent state exposed via 9P.
type AgentLog struct {
	mu         sync.RWMutex
	log        []byte
	logVers    uint32
	ChatOffset int
	plan       []byte
	prevPrompt []byte // last submitted prompt; overwritten on each new submission
	remote     string // SSH target for remote execution (empty = local)

	// Chat stream: condition variable broadcast on every log append.
	// Readers of the `chat` file block on this until new data arrives.
	chatCond *sync.Cond
}

func NewAgentLog(remote string) *AgentLog {
	al := &AgentLog{remote: remote}
	al.chatCond = sync.NewCond(al.mu.RLocker())
	return al
}

func NewSession(id string, core *session.Session, ctx context.Context, cancel context.CancelFunc) *Session {
	sess := &Session{id: id, Core: core, SessionCtx: ctx, cancel: cancel}
	return sess
}

const modelsCacheTTL = 24 * time.Hour

// CachedListModels returns the model list, using a 24h cache to avoid
// repeated expensive API calls.
// NOTE: Must NOT be called while holding sess.mu (caller content() holds RLock).
// Uses its own modelsMu to avoid deadlock.
func (sess *Session) CachedListModels() string {
	sess.modelsMu.Lock()
	if sess.modelsCache != "" && time.Since(sess.modelsCacheAt) < modelsCacheTTL {
		result := sess.modelsCache
		sess.modelsMu.Unlock()
		return result
	}
	sess.modelsMu.Unlock()

	// Cache miss — fetch and store.
	models := sess.Core.Agent().ListModels()
	result := strings.Join(models, "\n")
	sess.modelsMu.Lock()
	sess.modelsCache = result
	sess.modelsCacheAt = time.Now()
	sess.modelsMu.Unlock()
	return result
}

// InvalidateModelsCache clears the cached model list.
func (sess *Session) InvalidateModelsCache() {
	sess.modelsMu.Lock()
	sess.modelsCache = ""
	sess.modelsCacheAt = time.Time{}
	sess.modelsMu.Unlock()
}

func (sess *Session) RunnableID() string { return sess.id }
func (sess *Session) Uname() string      { return sess.uname }

func (sess *Session) Cancel() {
	sess.cancel()
}

func (sess *Session) Interrupt() {
	sess.Core.Agent().Interrupt(agent.ErrInterrupted)
}

// --- AgentLog methods ---

// AppendLog appends data to the agent's log and bumps the version.
func (as *AgentLog) AppendLog(data []byte) {
	if len(data) == 0 {
		return
	}
	as.mu.Lock()
	as.log = append(as.log, data...)
	as.logVers++
	as.mu.Unlock()
	as.chatCond.Broadcast()
}

// EnsureTrailingNewline appends a newline if the log doesn't already end with one.
func (as *AgentLog) EnsureTrailingNewline() {
	as.mu.Lock()
	if len(as.log) > 0 && as.log[len(as.log)-1] != '\n' {
		as.log = append(as.log, '\n')
	}
	as.logVers++
	as.mu.Unlock()
}

// LogInfo returns the current log length and version atomically.
func (as *AgentLog) LogInfo() (length int, vers uint32) {
	as.mu.RLock()
	defer as.mu.RUnlock()
	return len(as.log), as.logVers
}

// Log returns the raw chat log bytes. Caller must hold Mu().RLock().
func (as *AgentLog) Log() []byte { return as.log }

// Plan returns the plan bytes. Caller must hold Mu().RLock().
func (as *AgentLog) Plan() []byte { return as.plan }

// SetPlan sets the plan. Caller must hold Mu().Lock().
func (as *AgentLog) SetPlan(p []byte) { as.plan = p }

// PrevPrompt returns the last submitted prompt. Caller must hold Mu().RLock().
func (as *AgentLog) PrevPrompt() []byte { return as.prevPrompt }

// Mu returns the agent state's mutex.
func (as *AgentLog) Mu() *sync.RWMutex { return &as.mu }

// Remote returns the SSH target for remote execution.
func (as *AgentLog) Remote() string { return as.remote }

// NewEventHandler returns an agent.EventHandler that writes formatted events
// to this AgentLog. Use as the Output callback when creating an agent.
func NewEventHandler(as *AgentLog) agent.EventHandler {
	streamingRole := ""
	streamingResponseID := ""

	return func(ev agent.Event) {
		switch ev.Role {
		case "assistant", "reasoning":
			if streamingRole != ev.Role || (ev.Role == "assistant" && ev.ResponseID != streamingResponseID) {
				if streamingRole != "" {
					as.AppendLog([]byte("\n"))
				}
				if ev.Role == "assistant" && ev.ResponseID != "" {
					as.AppendLog([]byte("[assistant:" + ev.ResponseID + "]\n"))
					streamingResponseID = ev.ResponseID
				} else {
					as.AppendLog([]byte("[" + ev.Role + "]\n"))
				}
				if ev.Role == "assistant" {
					as.mu.Lock()
					as.ChatOffset = len(as.log)
					as.mu.Unlock()
				}
				streamingRole = ev.Role
			}
			as.AppendLog(FormatEvent(ev))

		case "tool":
			if streamingRole == "tool" {
				if ev.Content == "" {
					as.AppendLog([]byte("\n"))
					streamingRole = ""
				} else {
					as.AppendLog([]byte(ev.Content))
				}
				return
			}
			if streamingRole != "" {
				as.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			if ev.Content == "" {
				as.AppendLog([]byte("[tool:" + ev.Name + "]\n"))
				streamingRole = "tool"
			} else {
				as.AppendLog(FormatEvent(ev))
			}

		case "retry":
			if streamingRole == "retry" {
				as.AppendLog([]byte("\n" + ev.Content))
				return
			}
			if streamingRole != "" {
				as.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			as.AppendLog([]byte("[retry]\n"))
			streamingRole = "retry"
			as.AppendLog(FormatEvent(ev))

		case "call":
			if streamingRole != "" {
				as.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			as.AppendLog(FormatEvent(ev))

		default:
			if streamingRole != "" {
				as.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			as.AppendLog(FormatEvent(ev))
		}
	}
}