ollie-core/session/history.go

package session

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"slices"
	"strings"
	"time"

	"ollie/backend"
)

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
)

type Reaction struct {
	ID         string    `json:"id"`
	ResponseID string    `json:"responseId"`
	Emoji      string    `json:"emoji"`
	Category   string    `json:"category"`
	CreatedAt  time.Time `json:"createdAt"`
}

// PersistedAgent is the on-disk format for a saved session.
type PersistedAgent struct {
	ID        string            `json:"id"`
	Agent     string            `json:"agent,omitempty"`
	Backend   string            `json:"backend,omitempty"`
	Model     string            `json:"model,omitempty"`
	CWD       string            `json:"cwd,omitempty"`
	Remote    string            `json:"remote,omitempty"`
	Messages  []backend.Message `json:"messages"`
	TaskState *TaskState        `json:"taskState,omitempty"`
	// Usage and cost tracking, persisted across restarts.
	TotalInputTokens         int     `json:"totalInputTokens,omitempty"`
	TotalCachedInputTokens   int     `json:"totalCachedInputTokens,omitempty"`
	TotalCacheCreationTokens int     `json:"totalCacheCreationTokens,omitempty"`
	TotalOutputTokens        int     `json:"totalOutputTokens,omitempty"`
	TotalRequests            int     `json:"totalRequests,omitempty"`
	Estimated                bool    `json:"estimated,omitempty"`
	LastTurnCostUSD          float64 `json:"lastTurnCostUSD,omitempty"`
	SessionCostUSD           float64 `json:"sessionCostUSD,omitempty"`
	// Reaction tracking.
	PositiveReactions int        `json:"positiveReactions,omitempty"`
	NegativeReactions int        `json:"negativeReactions,omitempty"`
	Reactions         []Reaction `json:"reactions,omitempty"`
}

// TaskState is a compact structured overlay that summarizes the agent's
// current position in a task. Injected at the top of every turn so the
// model doesn't rely on full history recall.
type TaskState struct {
	Objective     string   `json:"objective"`
	PlanStep      string   `json:"plan_step"`
	Constraints   []string `json:"constraints,omitempty"`
	LastAction    string   `json:"last_action"`
	OpenQuestions []string `json:"open_questions,omitempty"`
	NextDecision  string   `json:"next_decision"`
}

// render produces a compact text representation for injection into the message history.
func (ts *TaskState) render() string {
	if ts.Objective == "" {
		return ""
	}
	var sb strings.Builder
	sb.WriteString("[task state]\n")
	sb.WriteString("objective: " + ts.Objective + "\n")
	if ts.PlanStep != "" {
		sb.WriteString("plan_step: " + ts.PlanStep + "\n")
	}
	if len(ts.Constraints) > 0 {
		sb.WriteString("constraints: " + strings.Join(ts.Constraints, "; ") + "\n")
	}
	if ts.LastAction != "" {
		sb.WriteString("last_action: " + ts.LastAction + "\n")
	}
	if len(ts.OpenQuestions) > 0 {
		sb.WriteString("open_questions: " + strings.Join(ts.OpenQuestions, "; ") + "\n")
	}
	if ts.NextDecision != "" {
		sb.WriteString("next_decision: " + ts.NextDecision + "\n")
	}
	sb.WriteString("[/task state]")
	return sb.String()
}

// saveTo writes the full message history to path as JSON.
func (s *History) saveTo(path, id, agentName, cwd string) error {
	return s.saveToFull(path, id, agentName, "", "", cwd, "")
}

// saveToFull writes session state including backend/model info.
func sanitizeMessages(msgs []backend.Message) []backend.Message {
	out := make([]backend.Message, 0, len(msgs))
	for _, m := range msgs {
		for j, tc := range m.ToolCalls {
			if len(tc.Arguments) > 0 {
				var dummy json.RawMessage
				if err := json.Unmarshal(tc.Arguments, &dummy); err != nil {
					// Corrupted arguments — replace with empty object so save succeeds.
					m.ToolCalls[j].Arguments = json.RawMessage("{}")
				}
			}
		}
		out = append(out, m)
	}
	return out
}

func (s *History) saveToFull(path, id, agentName, backendName, modelName, cwd, remote string) error {
	ps := PersistedAgent{
		ID:                       id,
		Agent:                    agentName,
		Backend:                  backendName,
		Model:                    modelName,
		CWD:                      cwd,
		Remote:                   remote,
		Messages:                 sanitizeMessages(s.messages),
		TaskState:                s.TaskState,
		TotalInputTokens:         s.TotalInputTokens,
		TotalCachedInputTokens:   s.TotalCachedInputTokens,
		TotalCacheCreationTokens: s.TotalCacheCreationTokens,
		TotalOutputTokens:        s.TotalOutputTokens,
		TotalRequests:            s.TotalRequests,
		Estimated:                s.Estimated,
		LastTurnCostUSD:          s.LastTurnCostUSD,
		SessionCostUSD:           s.SessionCostUSD,
		PositiveReactions:        s.PositiveReactions,
		NegativeReactions:        s.NegativeReactions,
		Reactions:                s.Reactions,
	}
	data, err := json.Marshal(ps)
	if err != nil {
		return fmt.Errorf("session save: %w", err)
	}
	// Atomic write: write to temp file, then rename to preserve last good snapshot.
	tmpPath := path + ".tmp"
	if err := os.WriteFile(tmpPath, data, 0600); err != nil {
		return fmt.Errorf("session save tmp: %w", err)
	}
	return os.Rename(tmpPath, path)
}

// LoadPersistedAgent reads a PersistedAgent from a JSON file.
func LoadPersistedAgent(path string) (*PersistedAgent, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var ps PersistedAgent
	if err := json.Unmarshal(data, &ps); err != nil {
		return nil, fmt.Errorf("parse session %s: %w", path, err)
	}
	return &ps, nil
}

// RestoreHistory reconstructs a Session from a persisted message list.
func RestoreHistory(ps *PersistedAgent) *History {
	s := &History{
		messages:                 ps.Messages,
		TaskState:                ps.TaskState,
		TotalInputTokens:         ps.TotalInputTokens,
		TotalCachedInputTokens:   ps.TotalCachedInputTokens,
		TotalCacheCreationTokens: ps.TotalCacheCreationTokens,
		TotalOutputTokens:        ps.TotalOutputTokens,
		TotalRequests:            ps.TotalRequests,
		Estimated:                ps.Estimated,
		LastTurnCostUSD:          ps.LastTurnCostUSD,
		SessionCostUSD:           ps.SessionCostUSD,
		PositiveReactions:        ps.PositiveReactions,
		NegativeReactions:        ps.NegativeReactions,
		Reactions:                ps.Reactions,
	}
	for i := range s.messages {
		if s.messages[i].Role == "assistant" && s.messages[i].ID == "" {
			s.messages[i].ID = NewResponseID()
		}
	}
	for _, m := range ps.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
	TaskState *TaskState
	// 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)
	// Reaction tracking.
	PositiveReactions int
	NegativeReactions int
	Reactions         []Reaction
}

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

// Checkpoint forks the session: returns a new Session that inherits the given
// TaskState but starts with a clean message history. This enables narrow-context
// sub-agents that know what to do without inheriting all parent message noise.
func (s *History) Checkpoint(ts TaskState) *History {
	child := &History{
		goal:      ts.Objective,
		TaskState: &ts,
	}
	// Seed with a user message so the child has a valid initial turn.
	child.messages = []backend.Message{{
		Role:    "user",
		Content: ts.render(),
	}}
	return child
}

func (s *History) history() []backend.Message {
	if len(s.Reactions) == 0 {
		return s.messages
	}
	// Append feedback after the conversation so assistant tool-call messages stay
	// adjacent to their tool results, as required by provider APIs.
	out := make([]backend.Message, 0, len(s.messages)+len(s.Reactions))
	out = append(out, s.messages...)
	for _, r := range s.Reactions {
		_, desc, _ := classifyReaction(r.Emoji)
		out = append(out, backend.Message{Role: "user", Content: "[reaction to assistant response " + r.ResponseID + ": " + r.Category + " (" + r.Emoji + ")] " + desc})
	}
	return out
}

func (s *History) taskState() *TaskState {
	return s.TaskState
}

func (s *History) updateTaskState(ts TaskState) {
	s.TaskState = &ts
}

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) recomputeReactionCounts() {
	s.PositiveReactions = 0
	s.NegativeReactions = 0
	for _, r := range s.Reactions {
		switch r.Category {
		case "positive", "excellent":
			s.PositiveReactions++
		case "negative", "terrible":
			s.NegativeReactions++
		}
	}
}

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 {
		s.messages = append(s.messages, backend.Message{
			Role:          "tool",
			Content:       r.Content,
			ContentBlocks: r.ContentBlocks,
			ToolCallID:    r.ToolCallID,
		})
	}
}

// llmSummarizeToolResult uses the LLM to produce a concise summary of a tool result.
func llmSummarizeToolResult(ctx context.Context, b backend.Backend, name, content string) string {
	prompt := fmt.Sprintf("Summarize this tool result from %q in 1-3 sentences, preserving key facts, file paths, values, and findings:\n\n%s", name, content)
	ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: prompt}}, nil, backend.GenerationParams{})
	if err != nil {
		return content
	}
	var sb strings.Builder
	for ev := range ch {
		if ev.Content != "" {
			sb.WriteString(ev.Content)
		}
		if ev.Done {
			break
		}
	}
	if s := strings.TrimSpace(sb.String()); s != "" {
		return fmt.Sprintf("[summary of %s result]: %s", name, s)
	}
	return content
}

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

// PreCompactionSnapshot returns a copy of the current messages for persistence
// before compaction. Call this before compact().
func (s *History) PreCompactionSnapshot() []backend.Message {
	return slices.Clone(s.messages)
}

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

	// Parse structured task state from the model's JSON response.
	var ts TaskState
	raw := extractJSON(summaryText)
	if err := json.Unmarshal([]byte(raw), &ts); err != nil {
		// Fallback: treat as unstructured summary if JSON parse fails.
		ts = TaskState{Objective: summaryText}
	}
	s.TaskState = &ts

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

// buildCompactedHistory constructs a three-zone post-compaction message list:
//   - Cold: structured task state summary (everything older than warm+hot)
//   - Warm: brief index of recent decisions/actions (warmIndexSize messages)
//   - Hot: last hotTailSize messages verbatim
func buildCompactedHistory(ts TaskState, 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: structured task state.
	stateJSON, _ := json.Marshal(ts)
	cold := backend.Message{
		Role:    "user",
		Content: "[compacted context — cold zone]\n```json\n" + string(stateJSON) + "\n```",
	}

	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 ""
}

// extractJSON finds the first JSON object in s (handling markdown fences).
func extractJSON(s string) string {
	// Strip markdown code fence if present.
	if i := strings.Index(s, "```json"); i >= 0 {
		s = s[i+7:]
		if j := strings.Index(s, "```"); j >= 0 {
			s = s[:j]
		}
	} else if i := strings.Index(s, "```"); i >= 0 {
		s = s[i+3:]
		if j := strings.Index(s, "```"); j >= 0 {
			s = s[:j]
		}
	}
	// Find first { ... }
	start := strings.Index(s, "{")
	if start < 0 {
		return s
	}
	end := strings.LastIndex(s, "}")
	if end < start {
		return s
	}
	return s[start : end+1]
}

// 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
}

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

// cloneMessages returns a deep copy of the message slice.
func cloneMessages(msgs []backend.Message) []backend.Message {
	out := make([]backend.Message, len(msgs))
	for i, m := range msgs {
		out[i] = backend.Message{
			Role:       m.Role,
			Content:    m.Content,
			ToolCallID: m.ToolCallID,
		}
		if len(m.ToolCalls) > 0 {
			out[i].ToolCalls = make([]backend.ToolCall, len(m.ToolCalls))
			for j, tc := range m.ToolCalls {
				out[i].ToolCalls[j] = backend.ToolCall{
					ID:        tc.ID,
					Name:      tc.Name,
					Arguments: append(json.RawMessage(nil), tc.Arguments...),
				}
			}
		}
	}
	return out
}

// 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
}

// stripColdResults summarizes large tool-result messages outside the hot tail
// using an LLM call. Messages in the last hotTailSize slots are left verbatim.
func (s *History) stripColdResults(ctx context.Context, b backend.Backend) {
	hot := len(s.messages) - hotTailSize
	for i := range s.messages {
		if i >= hot {
			break
		}
		m := &s.messages[i]
		if m.Role == "tool" && len(m.Content) > 200 {
			m.Content = llmSummarizeToolResult(ctx, b, m.ToolCallID, m.Content)
		}
	}
}

// stripCold implements the state interface.
func (s *History) stripCold(ctx context.Context, b backend.Backend) {
	s.stripColdResults(ctx, b)
}

// contextDebug returns a multi-line breakdown of the history.
func (s *History) contextDebug() string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("=== %d messages ===\n", len(s.messages)))
	for i, m := range s.messages {
		preview := m.Content
		if len(preview) > 80 {
			preview = preview[:80] + "..."
		}
		chars := len(m.Content)
		for _, tc := range m.ToolCalls {
			chars += len(tc.Name) + len(tc.Arguments)
		}
		sb.WriteString(fmt.Sprintf("  [%d] role=%-10s chars=%-6d %q\n", i, m.Role, chars, preview))
	}
	return sb.String()
}