ollie/cmd/olliesrv/internal/fs/agent_log.go

package fs

import (
	"sync"

	"ollie/cmd/olliesrv/internal/agent"
	"ollie/cmd/olliesrv/internal/backend"
	"ollie/format"
)

// AgentLog holds per-agent 9P streaming state.
type AgentLog struct {
	mu         sync.RWMutex
	log        []byte
	logVers    uint32
	ChatOffset int
	plan       []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) Log() []byte { return as.log }

// 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, ""))
		}
	}
}