ollie/fs/types.go

package fs

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

	"ollie/agent"
	"ollie/backend"
	"ollie/format"
	"ollie/session"
	"ollie/toolsrv"
)

// SessionNode wraps session.Session for 9P handlers.
// It adds AgentLogs (9P chat streaming state) and models cache.
type SessionNode struct {
	Core *session.Session

	// AgentLogs for chat streaming, keyed by agent ID.
	mu        sync.RWMutex
	AgentLogs map[string]*AgentLog

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

// NewSessionNode creates a SessionNode wrapping a session.
func NewSessionNode(core *session.Session) *SessionNode {
	return &SessionNode{Core: core, AgentLogs: make(map[string]*AgentLog)}
}

// --- Delegation to Core ---

func (s *SessionNode) ID() string        { return s.Core.ID() }
func (s *SessionNode) Name() string      { return s.Core.Name() }
func (s *SessionNode) SetName(n string)  { s.Core.SetName(n) }
func (s *SessionNode) Uname() string     { return s.Core.Uname() }
func (s *SessionNode) Cancel()           { s.Core.Cancel() }
func (s *SessionNode) IsPaused() bool    { return s.Core.IsPaused() }
func (s *SessionNode) IsConnected() bool { return s.Core.IsConnected() }

func (s *SessionNode) Pause() error {
	err := s.Core.Pause()
	if err == nil {
		session.PublishEvent("session."+s.Core.ID()+".pause", "")
	}
	return err
}

func (s *SessionNode) Resume() error {
	err := s.Core.Resume()
	if err == nil {
		session.PublishEvent("session."+s.Core.ID()+".resume", "")
	}
	return err
}

func (s *SessionNode) LoadTool(name string, ag *agent.Agent) error {
	return s.Core.LoadTool(name, ag)
}

func (s *SessionNode) Ctx() context.Context {
	return s.Core.Ctx()
}

func (s *SessionNode) ToolsConn() toolsrv.Runner {
	return s.Core.ToolsConn()
}

// --- AgentLog management ---

func (s *SessionNode) AgentLog() *AgentLog {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, al := range s.AgentLogs {
		return al
	}
	return nil
}

func (s *SessionNode) AgentLogFor(agentID string) *AgentLog {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.AgentLogs[agentID]
}

func (s *SessionNode) SetAgentLog(agentID string, al *AgentLog) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.AgentLogs == nil {
		s.AgentLogs = make(map[string]*AgentLog)
	}
	s.AgentLogs[agentID] = al
}

// --- Models cache ---

const modelsCacheTTL = 24 * time.Hour

func (s *SessionNode) CachedListModels() string {
	s.modelsMu.Lock()
	if s.modelsCache != "" && time.Since(s.modelsCacheAt) < modelsCacheTTL {
		result := s.modelsCache
		s.modelsMu.Unlock()
		return result
	}
	s.modelsMu.Unlock()

	agents := s.Core.Agents()
	if len(agents) == 0 {
		return ""
	}
	models := agents[0].ListModels()
	result := strings.Join(models, "\n")
	s.modelsMu.Lock()
	s.modelsCache = result
	s.modelsCacheAt = time.Now()
	s.modelsMu.Unlock()
	return result
}

func (s *SessionNode) InvalidateModelsCache() {
	s.modelsMu.Lock()
	s.modelsCache = ""
	s.modelsCacheAt = time.Time{}
	s.modelsMu.Unlock()
}

// AgentLog holds per-agent 9P streaming state.
type AgentLog struct {
	mu         sync.RWMutex
	log        []byte
	logVers    uint32
	ChatOffset int
	plan       []byte
	prevPrompt []byte
	remote     string
	chatCond   *sync.Cond
}

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

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

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

func (as *AgentLog) LogInfo() (length int, vers uint32) {
	as.mu.RLock()
	defer as.mu.RUnlock()
	return len(as.log), as.logVers
}

func (as *AgentLog) Log() []byte       { return as.log }
func (as *AgentLog) Plan() []byte      { return as.plan }
func (as *AgentLog) SetPlan(p []byte)  { as.plan = p }
func (as *AgentLog) PrevPrompt() []byte { return as.prevPrompt }
func (as *AgentLog) Mu() *sync.RWMutex { return &as.mu }
func (as *AgentLog) Remote() string    { return as.remote }

// NewEventHandler returns an event handler that writes to an AgentLog.
func NewEventHandler(as *AgentLog) agent.EventHandler {
	streamingRole := ""
	streamingResponseID := ""
	streamingFence := ""

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

		case "reasoning":
			// Absorbed silently

		case "tool":
			if ev.Name == "reasoning_think" {
				return
			}
			if streamingRole == "tool" {
				if ev.Content == "" {
					if streamingFence != "" {
						as.AppendLog([]byte("\n" + streamingFence + "\n"))
						streamingFence = ""
					}
					as.AppendLog([]byte("\n" + format.EndTag + "\n"))
					streamingRole = ""
				} else {
					as.AppendLog([]byte(ev.Content))
				}
				return
			}
			if streamingRole != "" {
				if streamingFence != "" {
					as.AppendLog([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				as.AppendLog([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			if ev.Content == "" {
				as.AppendLog([]byte(format.ToolDelim(ev.Name) + "```" + ev.OutputFormat + "\n"))
				streamingFence = format.ToolFence
				streamingRole = "tool"
			} else {
				as.AppendLog(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))
			}

		case "retry":
			if streamingRole == "retry" {
				as.AppendLog([]byte("\n" + ev.Content))
				return
			}
			if streamingRole != "" {
				as.AppendLog([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			as.AppendLog([]byte(format.RetryTag + "\n"))
			streamingRole = "retry"
			as.AppendLog(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		case "call":
			if streamingRole != "" {
				if streamingFence != "" {
					as.AppendLog([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				as.AppendLog([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			as.AppendLog(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		default:
			if streamingRole != "" {
				if streamingFence != "" {
					as.AppendLog([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				as.AppendLog([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			as.AppendLog(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))
		}
	}
}

// replayMessagesToLog renders persisted messages into an AgentLog.
// It produces the same block structure as the live NewEventHandler by
// using format.FormatEvent for call/tool formatting.
func replayMessagesToLog(as *AgentLog, messages []backend.Message) {
	const maxReplay = 20
	start := len(messages) - maxReplay
	if start < 0 {
		start = 0
	}
	for start < len(messages) && messages[start].Role == "system" {
		start++
	}

	// Build a lookup from ToolCallID → tool name so tool result blocks
	// can reference the tool name (matching the live EventHandler output).
	callNames := make(map[string]string)
	for _, m := range messages[start:] {
		for _, tc := range m.ToolCalls {
			if tc.ID != "" {
				callNames[tc.ID] = tc.Name
			}
		}
	}

	for _, m := range messages[start:] {
		switch m.Role {
		case "system":
		case "user":
			as.AppendLog(format.FormatEvent("user", "", m.Content, ""))
		case "assistant":
			if m.ID != "" {
				as.AppendLog([]byte(format.AssistantHeaderPrefix + m.ID + "]]]\n"))
			} else {
				as.AppendLog([]byte(format.AssistantTag + "\n"))
			}
			if m.Content != "" {
				as.AppendLog([]byte(m.Content + "\n"))
			}
			as.AppendLog([]byte(format.EndTag + "\n"))
			for _, tc := range m.ToolCalls {
				as.AppendLog(format.FormatEvent("call", tc.Name, string(tc.Arguments), ""))
			}
		case "tool":
			name := callNames[m.ToolCallID]
			if name == "" {
				name = m.ToolCallID
			}
			as.AppendLog(format.FormatEvent("tool", name, m.Content, ""))
		}
	}
}