ollie/fs/types.go

package fs

import (
	"context"
	"encoding/json"
	"fmt"
	"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 // immutable UUID, set at creation, never changes
	name       string // mutable friendly name, defaults to id
	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)

	// toolsConn is the multiplexed connection to the remote tool server,
	// shared by agent execution and 9P tool-management handlers.
	toolsConn toolsrv.Runner

	// AgentLogs for the session's agents, keyed by agent ID.
	// Backward compat: AgentLog() returns the first agent's log (by map iteration order).
	AgentLogs map[string]*AgentLog

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

	// disallowTools is the set of tool names that cannot be loaded
	// for agents in this session. Set from agent config at creation.
	disallowTools map[string]struct{}
}

// 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 {
	name := id
	if idx := strings.IndexByte(id, '-'); idx > 0 {
		name = id[:idx]
	}
	sess := &Session{id: id, name: name, Core: core, SessionCtx: ctx, cancel: cancel}
	return sess
}

// AgentLog returns the first agent's log for backward compatibility.
// The iteration order over the map is non-deterministic; when multi-agent
// support is fully wired, callers should use AgentLogFor or AgentLogs directly.
func (sess *Session) AgentLog() *AgentLog {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	for _, al := range sess.AgentLogs {
		return al
	}
	return nil
}

// AgentLogFor returns the AgentLog for the given agent ID, or nil.
func (sess *Session) AgentLogFor(agentID string) *AgentLog {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return sess.AgentLogs[agentID]
}

// SetAgentLog stores the agent log for the given agent ID.
func (sess *Session) SetAgentLog(agentID string, al *AgentLog) {
	sess.mu.Lock()
	defer sess.mu.Unlock()
	sess.setAgentLogLocked(agentID, al)
}

// setAgentLogLocked is like SetAgentLog but assumes sess.mu is already held.
func (sess *Session) setAgentLogLocked(agentID string, al *AgentLog) {
	if sess.AgentLogs == nil {
		sess.AgentLogs = make(map[string]*AgentLog)
	}
	sess.AgentLogs[agentID] = al
}

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

// loadTool loads a tool into the remote ollie-remote process via the
// session's management toolsConn. Returns an error if the tool is in
// the session's disallow list or if the remote rejects it.
//
// This is the single convergent path for tool loading — both the 9P
// agent/{id}/tools write handler and the autoLoad initialization call this.
func (sess *Session) loadTool(name string, ag *agent.Agent) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return nil
	}

	// Check disallow list.
	sess.mu.RLock()
	_, blocked := sess.disallowTools[name]
	sess.mu.RUnlock()
	if blocked {
		return fmt.Errorf("tool %q is disallowed for this session", name)
	}

	// Load on the remote ollie-remote process (for agent execution).
	// Prefer the session's management toolsConn (avoids deadlocking with
	// the agent loop), fall back to the agent's ToolServer.
	var runner toolsrv.Runner
	if sess.toolsConn != nil {
		runner = sess.toolsConn
	} else if ag != nil {
		runner = ag.ToolServer()
	}
	return loadToolOnRunner(runner, name)
}

func loadToolOnRunner(runner toolsrv.Runner, name string) error {
	if runner == nil {
		return nil
	}
	ctxTO, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	args, _ := json.Marshal(map[string]string{"name": name})
	if _, err := runner.CallTool(ctxTO, "tool_load", json.RawMessage(args)); err != nil {
		return fmt.Errorf("remote load: %w", err)
	}
	return nil
}

// SetDisallowTools replaces the session's disallow tool list.
func (sess *Session) SetDisallowTools(names []string) {
	sess.mu.Lock()
	defer sess.mu.Unlock()
	if len(names) == 0 {
		sess.disallowTools = nil
		return
	}
	sess.disallowTools = make(map[string]struct{}, len(names))
	for _, n := range names {
		sess.disallowTools[n] = struct{}{}
	}
}

func (sess *Session) RunnableID() string { return sess.id }
func (sess *Session) Uname() string      { return sess.uname }
func (sess *Session) Name() string {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return sess.name
}
func (sess *Session) SetName(name string) {
	sess.mu.Lock()
	defer sess.mu.Unlock()
	sess.name = name
}

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 := ""
	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[[[end]]]\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 "reasoning":
			// Reasoning is included in the LLM context but not appended
			// to the chat log, so the log parser doesn't misinterpret
			// header-like text that reasoning might produce.
			//
			// Don't set streamingRole here either — since we never write
			// a [[[reasoning]]] header to the log, the role-transition
			// logic would emit a spurious [[[end]]] when the next role
			// (assistant) fires. Just absorb the event silently.

		case "tool":
			// reasoning_think writes reasoning content — suppress from log
			// same as the reasoning role, so the block parser never sees
			// two [[[end]]] markers per turn. Don't set streamingRole
			// either — reasoning was never opened with a header. Just
			// return silently.
			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"))
					streamingRole = ""
				} else {
					as.AppendLog([]byte(ev.Content))
				}
				return
			}
			if streamingRole != "" {
				if streamingFence != "" {
					as.AppendLog([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				as.AppendLog([]byte("\n[[[end]]]\n"))
				streamingRole = ""
			}
			if ev.Content == "" {
				as.AppendLog([]byte("[[[tool:" + ev.Name + "]]]\n```" + ev.OutputFormat + "\n"))
				streamingFence = "```"
				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[[[end]]]\n"))
				streamingRole = ""
			}
			as.AppendLog([]byte("[[[retry]]]\n"))
			streamingRole = "retry"
			as.AppendLog(FormatEvent(ev))

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

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