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

package agent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"

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

const (
	compactionPrompt = `You are performing a CONTEXT CHECKPOINT COMPACTION. Produce a JSON object summarizing the current task state. Output ONLY valid JSON matching this schema:

{
  "objective": "the user's overall goal",
  "plan_step": "what step you are currently on",
  "constraints": ["any constraints or user preferences"],
  "last_action": "brief: last tool/action + outcome",
  "open_questions": ["unresolved questions"],
  "next_decision": "what to do next"
}

Be concise. Capture what another LLM needs to seamlessly continue.`

	// hotTailSize is the number of most recent messages kept verbatim (hot zone).
	hotTailSize = 8

	// warmIndexSize is the number of messages preceding the hot zone
	// summarized as a brief decision index (warm zone).
	warmIndexSize = 10
)

// RestoreHistoryFromMessages reconstructs a History from just the message list.
// Used for multi-agent session restore where each agent has separate messages.
func RestoreHistoryFromMessages(messages []backend.Message) *History {
	s := &History{
		messages:     messages,
		summaryCache: make(map[string]string),
	}
	for i := range s.messages {
		if s.messages[i].Role == "assistant" && s.messages[i].ID == "" {
			s.messages[i].ID = util.NewUUID()
		}
		if s.messages[i].Role == "tool" && s.messages[i].ID == "" {
			if s.messages[i].ToolCallID != "" {
				s.messages[i].ID = toolMessageID(s.messages[i])
			} else {
				s.messages[i].ID = util.NewUUID()
			}
		}
	}
	for _, m := range messages {
		if m.Role == "user" {
			s.goal = m.Content
			break
		}
	}
	return s
}

// History is an ephemeral in-memory state backend.
type History struct {
	goal     string
	messages []backend.Message
	// summaryCache reuses validated summaries for identical tool-result content
	// during the lifetime of this agent. SummaryApplied persists per-message
	// state across session restore; this cache avoids new LLM calls for repeats.
	summaryCache map[string]string
	// Cumulative usage tracking.
	TotalInputTokens         int
	TotalCachedInputTokens   int
	TotalCacheCreationTokens int
	TotalOutputTokens        int
	TotalRequests            int
	Estimated                bool    // true if any usage was estimated rather than reported by the backend
	LastTurnCostUSD          float64 // cost of the most recently completed turn in USD
	SessionCostUSD           float64 // cumulative cost of all turns in this session in USD
	// per-turn accumulators; reset at the start of each turn
	turnInputTokens    int
	turnCachedTokens   int
	turnCreationTokens int
	turnOutputTokens   int
	turnCostUSD        float64 // >0 when backend reported cost directly (e.g. OpenRouter)
}

// newHistory creates a new empty History. The caller is responsible for
// appending the initial user message via appendUserMessage.
func newHistory(goal string) *History {
	return &History{goal: goal, summaryCache: make(map[string]string)}
}

func toolMessageID(m backend.Message) string {
	if m.ToolCallID != "" {
		return "tool:" + m.ToolCallID
	}
	h := sha256.Sum256([]byte(m.Content))
	return "tool:content:" + hex.EncodeToString(h[:])
}

const coldSummaryCacheVersion = "v2"

func toolSummaryHash(m backend.Message) string {
	input := coldSummaryCacheVersion + "\x00" + m.ToolName + "\x00" + normalizeSummaryCacheContent(m.Content)
	h := sha256.Sum256([]byte(input))
	return "cold-summary:" + coldSummaryCacheVersion + ":" + hex.EncodeToString(h[:])
}

func normalizeSummaryCacheContent(content string) string {
	if looksStructured(content) {
		return content
	}
	lines := strings.Split(strings.ReplaceAll(content, "\r\n", "\n"), "\n")
	normalized := make([]string, 0, len(lines))
	blank := false
	for _, line := range lines {
		line = strings.TrimRight(line, " \t")
		if strings.TrimSpace(line) == "" {
			if blank {
				continue
			}
			blank = true
		} else {
			blank = false
		}
		normalized = append(normalized, line)
	}
	return strings.TrimSpace(strings.Join(normalized, "\n"))
}

func (s *History) history() []backend.Message {
	return s.messages
}

func (s *History) addUsage(u backend.Usage, estimated bool) {
	s.TotalInputTokens += u.InputTokens
	s.TotalCachedInputTokens += u.CachedInputTokens
	s.TotalCacheCreationTokens += u.CacheCreationTokens
	s.TotalOutputTokens += u.OutputTokens
	s.TotalRequests++
	s.turnInputTokens += u.InputTokens
	s.turnCachedTokens += u.CachedInputTokens
	s.turnCreationTokens += u.CacheCreationTokens
	s.turnOutputTokens += u.OutputTokens
	s.turnCostUSD += u.CostUSD
	if estimated {
		s.Estimated = true
	}
}

func (s *History) cacheHitRatio() float64 {
	total := s.TotalInputTokens + s.TotalCachedInputTokens
	if total == 0 {
		return 0
	}
	return float64(s.TotalCachedInputTokens) / float64(total)
}
func (s *History) resetTurnAccumulators() {
	s.turnInputTokens = 0
	s.turnCachedTokens = 0
	s.turnCreationTokens = 0
	s.turnOutputTokens = 0
	s.turnCostUSD = 0
}

// recordTurnCost computes and stores the last turn's cost, adding it to the
// session total. model is used for the pricing-table fallback when the backend
// did not report cost directly.
func (s *History) recordTurnCost(model string) {
	var cost float64
	if s.turnCostUSD > 0 {
		cost = s.turnCostUSD
	} else {
		cost = computeCostUSD(model, backend.Usage{
			InputTokens:         s.turnInputTokens,
			CachedInputTokens:   s.turnCachedTokens,
			CacheCreationTokens: s.turnCreationTokens,
			OutputTokens:        s.turnOutputTokens,
		})
	}
	s.LastTurnCostUSD = cost
	s.SessionCostUSD += cost
}

func (s *History) update(assistant backend.Message, results []toolResult) {
	s.messages = append(s.messages, assistant)
	for _, r := range results {
		m := backend.Message{
			Role:          "tool",
			Content:       r.Content,
			ContentBlocks: r.ContentBlocks,
			ToolCallID:    r.ToolCallID,
			ToolName:      r.Name,
		}
		if r.ToolCallID != "" {
			m.ID = toolMessageID(m)
		} else {
			m.ID = util.NewUUID()
		}
		s.messages = append(s.messages, m)
	}
}

// removeCancelledToolResults filters out tool results that were cancelled due
// to interrupt, keeping completed work. Also removes the corresponding tool
// calls from assistant messages to maintain a valid message sequence.
func (s *History) removeCancelledToolResults() {
	// First pass: collect cancelled tool call IDs
	cancelled := make(map[string]bool)
	for _, m := range s.messages {
		if m.Role == "tool" && isCancelledToolResult(m.Content) {
			cancelled[m.ToolCallID] = true
		}
	}
	if len(cancelled) == 0 {
		return
	}

	// Second pass: filter messages and prune tool calls from assistant messages
	filtered := s.messages[:0]
	for _, m := range s.messages {
		if m.Role == "tool" && cancelled[m.ToolCallID] {
			continue
		}
		if m.Role == "assistant" && len(m.ToolCalls) > 0 {
			// Remove cancelled tool calls from this assistant message
			kept := m.ToolCalls[:0]
			for _, tc := range m.ToolCalls {
				if !cancelled[tc.ID] {
					kept = append(kept, tc)
				}
			}
			if len(kept) == 0 && m.Content == "" {
				// All tool calls cancelled and no text content - skip message
				continue
			}
			m.ToolCalls = kept
		}
		filtered = append(filtered, m)
	}
	s.messages = filtered
}

func isCancelledToolResult(content string) bool {
	return strings.Contains(content, `"status":"cancelled"`) ||
		strings.Contains(content, "tool execution interrupted by user")
}

// Compact summarizes the conversation via an LLM call, replacing the history
// with system messages + preserved user messages + a structured summary.
// Returns (n compacted, summary text, error); n==0 means nothing to compact.
func (s *History) compact(ctx context.Context, b backend.Backend) (int, string, error) {
	if len(s.messages) <= hotTailSize+warmIndexSize {
		return 0, "", nil
	}

	// Flatten tool calls into plain text so the compaction request
	// doesn't need tool schemas.
	flattened := flattenToolMessages(s.messages)
	flattened = append(flattened, backend.Message{
		Role:    "user",
		Content: compactionPrompt,
	})

	ch, err := b.ChatStream(ctx, flattened, nil, backend.GenerationParams{})
	if err != nil {
		return 0, "", fmt.Errorf("compact: %w", err)
	}

	var summary strings.Builder
	for ev := range ch {
		if ev.Content != "" {
			summary.WriteString(ev.Content)
		}
		if ev.Done {
			break
		}
	}

	summaryText := strings.TrimSpace(summary.String())
	if summaryText == "" {
		return 0, "", fmt.Errorf("compact: empty summary")
	}

	beforeCount := len(s.messages)
	s.messages = buildCompactedHistory(summaryText, s.messages)
	return beforeCount - len(s.messages), summaryText, nil
}

// buildCompactedHistory constructs a three-zone post-compaction message list:
//   - Cold: compaction summary (everything older than warm+hot)
//   - Warm: brief index of recent decisions/actions (warmIndexSize messages)
//   - Hot: last hotTailSize messages verbatim
func buildCompactedHistory(summary string, allMessages []backend.Message) []backend.Message {
	total := len(allMessages)

	// Determine zone boundaries.
	hotStart := total - hotTailSize
	if hotStart < 0 {
		hotStart = 0
	}

	// If hotStart lands on a tool message, walk backward to include the
	// preceding assistant message with tool_calls. Without this, the hot
	// zone starts with an orphaned tool result which violates the OpenAI
	// API constraint that tool messages must follow an assistant message
	// containing the corresponding tool_calls.
	for hotStart > 0 && allMessages[hotStart].Role == "tool" {
		hotStart--
	}

	warmStart := hotStart - warmIndexSize
	if warmStart < 0 {
		warmStart = 0
	}

	// Cold zone: compaction summary.
	cold := backend.Message{
		Role:    "user",
		Content: "[compacted context]\n" + summary,
	}

	var out []backend.Message
	out = append(out, cold)

	// Warm zone: one-line summaries of decisions in the warm window.
	warmSlice := allMessages[warmStart:hotStart]
	if len(warmSlice) > 0 {
		var sb strings.Builder
		sb.WriteString("[warm zone — recent decisions]\n")
		for _, m := range warmSlice {
			line := summarizeMessage(m)
			if line != "" {
				sb.WriteString("- ")
				sb.WriteString(line)
				sb.WriteByte('\n')
			}
		}
		out = append(out, backend.Message{Role: "user", Content: sb.String()})
	}

	// Hot zone: verbatim recent messages.
	out = append(out, allMessages[hotStart:]...)
	return out
}

// summarizeMessage produces a one-line summary of a message for the warm index.
func summarizeMessage(m backend.Message) string {
	switch m.Role {
	case "assistant":
		if len(m.ToolCalls) > 0 {
			names := make([]string, len(m.ToolCalls))
			for i, tc := range m.ToolCalls {
				names[i] = tc.Name
			}
			return "called: " + strings.Join(names, ", ")
		}
		text := m.Content
		if len(text) > 120 {
			text = text[:120] + "…"
		}
		return "said: " + text
	case "user":
		text := m.Content
		if len(text) > 120 {
			text = text[:120] + "…"
		}
		return "user: " + text
	case "tool":
		text := m.Content
		if len(text) > 80 {
			text = text[:80] + "…"
		}
		return "result(" + m.ToolCallID + "): " + text
	}
	return ""
}

// flattenToolMessages converts tool call/result sequences into plain text
// so the compaction request doesn't include tool-specific structures that
// the API may reject when no tools are defined.
func flattenToolMessages(messages []backend.Message) []backend.Message {
	out := make([]backend.Message, 0, len(messages))
	for _, m := range messages {
		switch {
		case m.Role == "assistant" && len(m.ToolCalls) > 0:
			var sb strings.Builder
			if m.Content != "" {
				sb.WriteString(m.Content)
				sb.WriteString("\n\n")
			}
			for _, tc := range m.ToolCalls {
				fmt.Fprintf(&sb, "[Tool call: %s(%s)]\n", tc.Name, string(tc.Arguments))
			}
			out = append(out, backend.Message{Role: "assistant", Content: sb.String()})
		case m.Role == "tool":
			text := m.Content
			if len(text) > 4000 {
				text = text[:4000] + "..."
			}
			out = append(out, backend.Message{
				Role:    "user",
				Content: fmt.Sprintf("[Tool result for %s]:\n%s", m.ToolCallID, text),
			})
		default:
			out = append(out, m)
		}
	}
	return out
}

// injectPreamble ensures the preamble system message is at history position 0.
// If a system message already exists there, its Content is replaced; otherwise
// one is prepended.
func (s *History) injectPreamble(preamble string) {
	if preamble == "" {
		return
	}
	msg := backend.Message{Role: "system", Content: preamble}
	if len(s.messages) > 0 && s.messages[0].Role == "system" {
		if s.messages[0].Content == preamble {
			return
		}
		s.messages[0] = msg
	} else {
		s.messages = append([]backend.Message{msg}, s.messages...)
	}
}

func (s *History) appendUserMessage(content string) {
	s.messages = append(s.messages, backend.Message{Role: "user", Content: content})
}

// estimateTokens returns a rough token count (~4 chars per token).
func (s *History) estimateTokens() int {
	chars := 0
	for _, m := range s.messages {
		chars += len(m.Content)
		for _, tc := range m.ToolCalls {
			chars += len(tc.Name) + len(tc.Arguments)
		}
	}
	return chars / 4
}

const (
	maxColdSummariesPerCall = 8
	minColdSummaries        = 3
	minColdSummaryTokens    = 32_000
)

func (s *History) pendingColdSummaryStats() (count, tokens int) {
	hot := len(s.messages) - hotTailSize
	if hot <= 0 {
		return 0, 0
	}
	for i := 0; i < hot; i++ {
		m := &s.messages[i]
		if m.Role != "tool" || len(m.Content) <= 200 || m.SummaryApplied {
			continue
		}
		if s.summaryCache != nil {
			if _, ok := s.summaryCache[toolSummaryHash(*m)]; ok {
				continue
			}
		}
		count++
		tokens += len(m.Content) / 4
	}
	return count, tokens
}

// stripCold summarizes large tool-result messages outside the hot tail
// using a single batched LLM call. Messages in the last hotTailSize slots
// are left verbatim.
func (s *History) stripCold(ctx context.Context, b backend.Backend) (backend.Usage, bool) {
	if s.summaryCache == nil {
		s.summaryCache = make(map[string]string)
	}
	hot := len(s.messages) - hotTailSize
	if hot < 0 {
		return backend.Usage{}, false
	}

	// Collect cold-zone tool results exceeding 200 chars.
	type item struct {
		idx       int // index into s.messages
		orig      string
		messageID string
		hash      string
	}
	var items []item
	for i := range s.messages {
		if i >= hot {
			break
		}
		m := &s.messages[i]
		if m.Role == "tool" && (len(m.Content) > 200 || strings.TrimSpace(m.Content) == "") {
			id := m.ID
			if id == "" {
				if m.ToolCallID != "" {
					id = toolMessageID(*m)
				} else {
					id = util.NewUUID()
				}
				m.ID = id
			}
			if !m.SummaryApplied {
				if summary, ok := localColdSummary(m.Content); ok {
					hash := toolSummaryHash(*m)
					s.summaryCache[hash] = summary
					m.Content = summary
					m.SummaryApplied = true
					continue
				}
				if summary, ok := s.summaryCache[toolSummaryHash(*m)]; ok {
					m.Content = summary
					m.SummaryApplied = true
					continue
				}
				items = append(items, item{idx: i, orig: m.Content, messageID: id, hash: toolSummaryHash(*m)})
				if len(items) >= maxColdSummariesPerCall {
					break
				}
			}
		}
	}
	if len(items) == 0 {
		return backend.Usage{}, false
	}

	// Build a single prompt listing each item.
	var sb strings.Builder
	sb.WriteString("Compress each of the following tool results to 1-3 sentences preserving key facts, file paths, values, and findings.\n\nReturn ONLY a valid JSON object. Do not use Markdown fences, prose, or any text outside the JSON object. Each key must be exactly one of the supplied message IDs and each value must be a string summary, for example:\n{\"tool:call-123\": \"summarized result\", \"tool:call-456\": \"another summary\"}\n\nOnly include keys that were in the input.\n\n")
	for _, it := range items {
		fmt.Fprintf(&sb, "---\nid=%s:\n%s\n", it.messageID, it.orig)
	}

	ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: sb.String()}}, nil, backend.GenerationParams{})
	if err != nil {
		return backend.Usage{}, false
	}
	var resp strings.Builder
	var usage backend.Usage
	var reported bool
	for ev := range ch {
		if ev.Content != "" {
			resp.WriteString(ev.Content)
		}
		if ev.Done {
			usage = ev.Usage
			reported = true
			break
		}
	}
	if !reported {
		usage.InputTokens = len(sb.String()) / 4
		usage.OutputTokens = resp.Len() / 4
	}
	estimated := !reported

	summaries, err := parseColdSummaries(resp.String())
	if err != nil {
		return usage, estimated
	}

	// Apply summaries back into history.
	for _, it := range items {
		if summary, ok := summaries[it.messageID]; ok && summary != "" {
			s.summaryCache[it.hash] = summary
			s.messages[it.idx].Content = summary
			s.messages[it.idx].SummaryApplied = true
		}
	}
	return usage, estimated
}

// resolveCompactionModel returns the model to use for compaction.
// Priority: agent config > backends.conf compactionModel > backends.conf model > current model.
func resolveCompactionModel(cfgModel string, b backend.Backend) string {
	if cfgModel != "" {
		return cfgModel
	}
	if b != nil {
		if model := backend.CompactionModel(b.Name()); model != "" {
			return model
		}
		return b.Model()
	}
	return ""
}