core/server/compaction/compaction.go

package compaction

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"time"

	"switchboard-core/models"
	"switchboard-core/providers"
	"switchboard-core/roles"
	"switchboard-core/store"
	"switchboard-core/treepath"
)

// ── Request / Result ────────────────────────

// CompactRequest specifies what to compact.
type CompactRequest struct {
	ChannelID string
	UserID    string  // channel owner
	TeamID    *string // for role resolution
	Trigger   string  // "manual" | "auto"
}

// CompactResult describes what compaction produced.
type CompactResult struct {
	SummaryID       string
	SummarizedCount int
	Model           string
	UsedFallback    bool
	Content         string
	InputTokens     int
	OutputTokens    int
}

// ── Errors ──────────────────────────────────

var (
	// ErrContextBudget is returned when the conversation to summarize
	// exceeds the utility model's context window.
	ErrContextBudget = fmt.Errorf("conversation exceeds utility model context window")
)

// ── Service ─────────────────────────────────

// Service provides conversation compaction (summarization).
// Used by both the HTTP handler (manual) and the background scanner (auto).
type Service struct {
	stores   store.Stores
	resolver *roles.Resolver
}

// NewService creates a compaction service.
func NewService(stores store.Stores, resolver *roles.Resolver) *Service {
	return &Service{stores: stores, resolver: resolver}
}

// Resolver exposes the underlying role resolver (for external checks like
// IsConfigured / IsPersonalOverride).
func (s *Service) Resolver() *roles.Resolver {
	return s.resolver
}

// ── Compact ─────────────────────────────────

// Compact summarizes the conversation in a channel, inserting a summary
// tree node. Returns the result or an error.
//
// The method:
//  1. Loads the active path via the message tree
//  2. Finds the most recent summary boundary
//  3. Collects post-boundary messages
//  4. Checks the content fits the utility model's context window
//  5. Calls the utility model role to generate a summary
//  6. Inserts the summary as a tree node with metadata
//  7. Updates the user's cursor
//  8. Logs usage
func (s *Service) Compact(ctx context.Context, req CompactRequest) (*CompactResult, error) {
	// ── Load active path ──
	path, err := treepath.GetActivePath(req.ChannelID, req.UserID)
	if err != nil {
		return nil, fmt.Errorf("load conversation: %w", err)
	}
	if len(path) < 4 {
		return nil, fmt.Errorf("conversation too short to summarize")
	}

	// Find existing summary boundary (if any) and only summarize after it
	startIdx := 0
	for i := range path {
		if treepath.IsSummaryMessage(&path[i]) {
			startIdx = i + 1
		}
	}

	// Build messages to summarize (skip system messages, skip previous summaries)
	var toSummarize []string
	messagesInScope := 0
	var lastMessageID string
	for _, m := range path[startIdx:] {
		if m.Role == "system" || treepath.IsSummaryMessage(&m) {
			continue
		}
		toSummarize = append(toSummarize, fmt.Sprintf("%s: %s", m.Role, m.Content))
		messagesInScope++
		lastMessageID = m.ID
	}

	if messagesInScope < 4 {
		return nil, fmt.Errorf("not enough new messages since last summary")
	}

	// ── Build the summarization prompt ──
	conversationText := strings.Join(toSummarize, "\n\n")
	systemPrompt := `You are a conversation summarizer. Create a concise but comprehensive summary of the following conversation that preserves:
- Key decisions and conclusions reached
- Important facts, names, numbers, and technical details mentioned
- Action items or commitments made
- The overall context and topic flow

Write the summary in a way that would allow the conversation to continue naturally. Use clear, factual language. Do not include meta-commentary about the summarization process.`

	userContent := "Summarize this conversation:\n\n" + conversationText

	// ── Context budget guard rail ──
	// Estimate whether the prompt fits the utility model's context window.
	// Prevents sending truncated input to small models (e.g. 32K utility
	// summarizing a 128K conversation).
	promptTokens := EstimateTokens(systemPrompt) + 4 + EstimateTokens(userContent) + 4
	utilityBudget := s.getUtilityContextBudget(ctx, req.UserID, req.TeamID)
	// Reserve 20% for output (the summary itself)
	inputCeiling := int(float64(utilityBudget) * 0.80)

	if promptTokens > inputCeiling {
		log.Printf("⚠ Compaction skipped for channel %s: prompt ~%dK tokens exceeds utility model budget %dK (ceiling %dK)",
			req.ChannelID, promptTokens/1000, utilityBudget/1000, inputCeiling/1000)
		return nil, fmt.Errorf("%w: ~%d tokens needed, utility model allows ~%d",
			ErrContextBudget, promptTokens, inputCeiling)
	}

	summaryPrompt := []providers.Message{
		{Role: "system", Content: systemPrompt},
		{Role: "user", Content: userContent},
	}

	// ── Call utility role ──
	log.Printf("📝 Compacting channel %s for user %s (%d messages, ~%dK tokens, trigger=%s)",
		req.ChannelID, req.UserID, messagesInScope, promptTokens/1000, req.Trigger)

	result, err := s.resolver.Complete(ctx, roles.RoleUtility, req.UserID, req.TeamID, summaryPrompt)
	if err != nil {
		log.Printf("⚠ Compaction failed for channel %s: %v", req.ChannelID, err)
		return nil, fmt.Errorf("summarization failed: %w", err)
	}

	// ── Log usage ──
	s.logUsage(ctx, req.ChannelID, req.UserID, req.Trigger, result)

	// ── Insert summary message as a tree node ──
	metadata := models.JSONMap{
		"type":                "summary",
		"summarized_until_id": lastMessageID,
		"summarized_count":    messagesInScope,
		"utility_model":       result.Model,
		"used_fallback":       result.UsedFallback,
		"trigger":             req.Trigger,
	}

	metaJSON, _ := json.Marshal(metadata)
	siblingIdx := treepath.NextSiblingIndex(req.ChannelID, &lastMessageID)

	summaryMsg := &models.Message{
		ChannelID:       req.ChannelID,
		ParentID:        &lastMessageID,
		Role:            "assistant",
		Content:         result.Content,
		Model:           result.Model,
		Metadata:        models.JSONMap{},
		SiblingIndex:    siblingIdx,
		ParticipantType: "system",
	}
	_ = json.Unmarshal(metaJSON, &summaryMsg.Metadata)
	if err := s.stores.Messages.Create(ctx, summaryMsg); err != nil {
		log.Printf("⚠ Failed to persist summary for channel %s: %v", req.ChannelID, err)
		return nil, fmt.Errorf("failed to save summary: %w", err)
	}
	summaryMsgID := summaryMsg.ID

	// Update cursor to point to the summary node
	if err := treepath.UpdateCursor(req.ChannelID, req.UserID, summaryMsgID); err != nil {
		log.Printf("⚠ Failed to update cursor after compaction: %v", err)
	}

	log.Printf("✅ Compaction done for channel %s: %s (%d messages → %d chars, trigger=%s)",
		req.ChannelID, summaryMsgID, messagesInScope, len(result.Content), req.Trigger)

	return &CompactResult{
		SummaryID:       summaryMsgID,
		SummarizedCount: messagesInScope,
		Model:           result.Model,
		UsedFallback:    result.UsedFallback,
		Content:         result.Content,
		InputTokens:     result.InputTokens,
		OutputTokens:    result.OutputTokens,
	}, nil
}

// ── Context Budget ──────────────────────────

// getUtilityContextBudget resolves the utility model's context window.
// Falls back to DefaultBudget if the model isn't in the catalog.
func (s *Service) getUtilityContextBudget(ctx context.Context, userID string, teamID *string) int {
	cfg, err := s.resolver.GetConfig(ctx, roles.RoleUtility, userID, teamID)
	if err != nil || cfg == nil || cfg.Primary == nil {
		return DefaultBudget
	}

	// Look up the primary model in catalog
	entry, err := s.stores.Catalog.GetByModelID(ctx, cfg.Primary.ProviderConfigID, cfg.Primary.ModelID)
	if err == nil && entry.Capabilities.MaxContext > 0 {
		return entry.Capabilities.MaxContext
	}

	// Try any-provider lookup (model may be synced under a different config)
	entry, err = s.stores.Catalog.GetByModelIDAny(ctx, cfg.Primary.ModelID)
	if err == nil && entry.Capabilities.MaxContext > 0 {
		return entry.Capabilities.MaxContext
	}

	return DefaultBudget
}

// ── Rate Limit Check ────────────────────────

// CheckRateLimit checks the utility rate limit for org-funded calls.
// Returns nil if within limits, an error with a descriptive message otherwise.
func (s *Service) CheckRateLimit(ctx context.Context, userID string) error {
	// Load rate limit from global settings (default: 20/hour, 0 = unlimited)
	limit := 20
	settings, err := s.stores.GlobalConfig.Get(ctx, "utility_rate_limit")
	if err == nil {
		if v, ok := settings["value"]; ok {
			switch n := v.(type) {
			case float64:
				limit = int(n)
			case int:
				limit = n
			}
		}
	}

	if limit <= 0 {
		return nil // unlimited
	}

	since := time.Now().Add(-1 * time.Hour)
	count, err := s.stores.Usage.CountRecentByRole(ctx, userID, roles.RoleUtility, since)
	if err != nil {
		log.Printf("⚠ Rate limit check failed: %v", err)
		return nil // fail open — don't block on DB errors
	}

	if count >= limit {
		return fmt.Errorf("rate limit exceeded: %d utility calls in the last hour (limit: %d)", count, limit)
	}
	return nil
}

// ── Helpers ─────────────────────────────────

// GetUserTeamID returns the user's first team ID (for role resolution).
func (s *Service) GetUserTeamID(ctx context.Context, userID string) *string {
	teamID, _ := s.stores.Teams.GetFirstTeamIDForUser(ctx, userID)
	if teamID == "" {
		return nil
	}
	return &teamID
}

func (s *Service) logUsage(ctx context.Context, channelID, userID, trigger string, result *roles.CompletionResult) {
	role := roles.RoleUtility
	entry := &models.UsageEntry{
		ChannelID:           &channelID,
		UserID:              userID,
		ProviderConfigID:    &result.ConfigID,
		ProviderScope:       result.ProviderScope,
		ModelID:             result.Model,
		Role:                &role,
		PromptTokens:        result.InputTokens,
		CompletionTokens:    result.OutputTokens,
		CacheCreationTokens: result.CacheCreationTokens,
		CacheReadTokens:     result.CacheReadTokens,
	}

	// Calculate cost from pricing
	pricing, err := s.stores.Pricing.GetForModel(ctx, result.ConfigID, result.Model)
	if err == nil && pricing != nil {
		if pricing.InputPerM != nil {
			costIn := float64(result.InputTokens) / 1_000_000 * *pricing.InputPerM
			entry.CostInput = &costIn
		}
		if pricing.OutputPerM != nil {
			costOut := float64(result.OutputTokens) / 1_000_000 * *pricing.OutputPerM
			entry.CostOutput = &costOut
		}
	}

	if err := s.stores.Usage.Log(ctx, entry); err != nil {
		log.Printf("⚠ Failed to log compaction usage: %v", err)
	}
}