ollie/cmd/olliesrv/internal/agent/chatlog.go

// chatlog.go — Chat output formatting and event handling.
//
// Converts agent events (streaming text, tool calls, tool results) into
// JSONL blocks written to the chat log. Handles streaming by emitting
// partial blocks that are replaced by final blocks.

package agent

import (
	"strings"

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

// initChatHandler sets up the agent's output handler to write to its own chat log.
func (ag *Agent) initChatHandler() {
	var streamingRole string
	var streamingID string
	var streamingContent strings.Builder
	var streamingResponseID string
	var streamingFormat string

	// flushPartial emits a partial block with current content.
	flushPartial := func() {
		if streamingID == "" {
			return
		}
		ag.SetPartial(format.Block{
			Role:    streamingRole,
			ID:      streamingID,
			Content: streamingContent.String(),
			Format:  streamingFormat,
			Partial: true,
		})
	}

	// closeBlock emits the final block and resets streaming state.
	closeBlock := func() {
		if streamingID == "" {
			return
		}
		ag.AppendBlock(format.Block{
			Role:    streamingRole,
			ID:      streamingID,
			Content: streamingContent.String(),
			Format:  streamingFormat,
			Partial: false,
		})
		streamingRole = ""
		streamingID = ""
		streamingContent.Reset()
		streamingFormat = ""
	}

	ag.output = func(ev Event) {
		switch ev.Role {
		case "assistant":
			// New response or different response ID = new block
			if streamingRole != ev.Role || ev.ResponseID != streamingResponseID {
				if streamingID != "" {
					closeBlock()
				}
				streamingRole = ev.Role
				streamingID = ag.NextBlockID()
				streamingResponseID = ev.ResponseID
				streamingContent.Reset()
			}
			streamingContent.WriteString(ev.Content)
			flushPartial()

		case "reasoning":
			if streamingRole != ev.Role {
				if streamingID != "" {
					closeBlock()
				}
				streamingRole = ev.Role
				streamingID = ag.NextBlockID()
				streamingContent.Reset()
			}
			streamingContent.WriteString(ev.Content)
			flushPartial()

		case "tool":
			if streamingRole == "tool" {
				if ev.Content == "" {
					// Empty content = end of streaming tool output
					closeBlock()
				} else {
					// Append to streaming tool content
					streamingContent.WriteString(ev.Content)
					flushPartial()
				}
				return
			}
			if streamingID != "" {
				closeBlock()
			}
			if ev.Content == "" {
				// Start streaming tool output
				streamingRole = "tool"
				streamingID = ag.NextBlockID()
				streamingFormat = ev.OutputFormat
				streamingContent.Reset()
			} else {
				// Complete tool output in one event
				ag.AppendBlock(format.Block{
					Role:    "tool",
					ID:      ag.NextBlockID(),
					Content: ev.Content,
					Format:  ev.OutputFormat,
				})
			}

		case "retry":
			if streamingRole == "retry" {
				streamingContent.WriteString("\n" + ev.Content)
				flushPartial()
				return
			}
			if streamingID != "" {
				closeBlock()
			}
			streamingRole = "retry"
			streamingID = ag.NextBlockID()
			streamingContent.Reset()
			streamingContent.WriteString(ev.Content)
			flushPartial()

		case "call":
			if streamingID != "" {
				closeBlock()
			}
			ag.AppendBlock(format.Block{
				Role:    "call",
				ID:      ag.NextBlockID(),
				Name:    ev.Name,
				Content: ev.Content,
			})

		case "state", "usage", "limitretry":
			// Internal events for status/stats, not chat display — skip
			return

		default:
			// error, info, stalled, etc.
			if streamingID != "" {
				closeBlock()
			}
			ag.AppendBlock(format.Block{
				Role:    ev.Role,
				ID:      ag.NextBlockID(),
				Content: ev.Content,
			})
		}
	}
}

// ReplayMessages renders persisted messages into the agent's chat log.
func (ag *Agent) ReplayMessages(messages []backend.Message) {
	const maxReplay = 20
	start := len(messages) - maxReplay
	if start < 0 {
		start = 0
	}
	for start < len(messages) && messages[start].Role == "system" {
		start++
	}

	for _, m := range messages[start:] {
		switch m.Role {
		case "system":
			// Skip
		case "user":
			ag.AppendBlock(format.Block{
				Role:    "user",
				ID:      ag.NextBlockID(),
				Content: m.Content,
			})
		case "assistant":
			if m.Reasoning != "" {
				ag.AppendBlock(format.Block{
					Role:    "reasoning",
					ID:      ag.NextBlockID(),
					Content: m.Reasoning,
				})
			}
			ag.AppendBlock(format.Block{
				Role:    "assistant",
				ID:      ag.NextBlockID(),
				Content: m.Content,
			})
			for _, tc := range m.ToolCalls {
				ag.AppendBlock(format.Block{
					Role:    "call",
					ID:      ag.NextBlockID(),
					Name:    tc.Name,
					Content: string(tc.Arguments),
				})
			}
		case "tool":
			ag.AppendBlock(format.Block{
				Role:    "tool",
				ID:      ag.NextBlockID(),
				Content: m.Content,
			})
		}
	}
}