iasysadmin/internal/ctxmgr/manager.go

// Package ctxmgr administra la conversación enviada al modelo y estima
// cuánto contexto ocupa.
//
// Mantiene dos niveles (§20): el prompt de sistema (que incluye el contexto
// histórico comprimido) y el contexto reciente completo. La compresión la
// realiza Compressor (compressor.go) y la orquesta el agente.
//
// Se llama ctxmgr y no "context" para no colisionar con el paquete estándar.
package ctxmgr

import (
	"strings"
	"sync"

	"sysagent/internal/llm"
)

// perMessageOverhead aproxima los tokens de formato (rol, separadores) que
// cada mensaje añade en las plantillas de chat habituales.
const perMessageOverhead = 4

// Manager guarda el prompt de sistema y el historial reciente.
// Es seguro para uso concurrente.
type Manager struct {
	mu      sync.Mutex
	system  string
	history []llm.Message

	// contextSize es el máximo de tokens; 0 significa UNKNOWN.
	contextSize int
	// overhead son tokens fijos fuera de los mensajes (p.ej. el esquema de
	// herramientas en tool calling nativo).
	overhead int
	// factor calibra la estimación con los prompt_tokens reales reportados
	// por el servidor (0 = sin calibrar).
	factor float64
	// lastReported es el prompt_tokens de la última respuesta (dato real).
	lastReported int
}

// New crea un Manager con el prompt de sistema ya construido.
func New(system string, contextSize int) *Manager {
	return &Manager{system: system, contextSize: contextSize}
}

// Add añade mensajes al historial reciente.
func (m *Manager) Add(msgs ...llm.Message) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.history = append(m.history, msgs...)
}

// Messages devuelve una copia de la conversación completa (sistema + historial)
// lista para enviarse al modelo.
func (m *Manager) Messages() []llm.Message {
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make([]llm.Message, 0, len(m.history)+1)
	out = append(out, llm.Message{Role: llm.RoleSystem, Content: m.system})
	return append(out, m.history...)
}

// SetSystem reemplaza el prompt de sistema (p.ej. al cambiar de protocolo
// de herramientas).
func (m *Manager) SetSystem(system string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.system = system
}

// Len devuelve el número de mensajes del historial reciente.
func (m *Manager) Len() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.history)
}

// Truncate recorta el historial a n mensajes (para deshacer un turno fallido).
func (m *Manager) Truncate(n int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if n >= 0 && n < len(m.history) {
		m.history = m.history[:n]
	}
}

// ElideToolOutputs reemplaza el contenido de los resultados de herramientas
// antiguos (todos salvo los keep más recientes) por un marcador corto.
//
// Es el último recurso cuando la compresión no es posible o no basta: las
// salidas de comandos son lo más voluminoso y lo que antes pierde valor.
// Devuelve cuántos mensajes se recortaron.
func (m *Manager) ElideToolOutputs(keep int) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	var idx []int
	for i, msg := range m.history {
		if IsToolResult(msg) {
			idx = append(idx, i)
		}
	}
	n := 0
	for _, i := range idx[:max(len(idx)-keep, 0)] {
		if m.history[i].Content == elided {
			continue
		}
		m.history[i].Content = elided
		n++
	}
	return n
}

const elided = "[tool output elided to save context; run the tool again if you need it]"

// ToolResultPrefix marca los resultados de herramientas enviados como
// mensaje de usuario (protocolo JSON).
const ToolResultPrefix = "TOOL_RESULTS"

// IsToolResult informa si un mensaje contiene resultados de herramientas.
func IsToolResult(msg llm.Message) bool {
	return msg.Role == llm.RoleTool || msg.Role == llm.RoleUser && strings.HasPrefix(msg.Content, ToolResultPrefix)
}

// History devuelve una copia del historial reciente, sin el prompt de sistema.
func (m *Manager) History() []llm.Message {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]llm.Message(nil), m.history...)
}

// Clear vacía el historial reciente. El prompt de sistema se conserva.
func (m *Manager) Clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.history = nil
}

// Replace sustituye el historial reciente (tras una compactación).
func (m *Manager) Replace(history []llm.Message) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.history = append([]llm.Message(nil), history...)
}

// SetOverhead fija los tokens estimados fuera de los mensajes.
func (m *Manager) SetOverhead(n int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.overhead = n
}

// SetContextSize actualiza el máximo de tokens (p.ej. tras consultar al servidor).
func (m *Manager) SetContextSize(n int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.contextSize = n
}

// ContextSize devuelve el máximo configurado o descubierto (0 = UNKNOWN).
func (m *Manager) ContextSize() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.contextSize
}

// RecordUsage calibra la estimación con el consumo real del servidor.
// sentEstimate es la estimación local (sin calibrar) de la petición enviada.
func (m *Manager) RecordUsage(u llm.Usage, sentEstimate int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if u.PromptTokens <= 0 || sentEstimate <= 0 {
		return
	}
	m.lastReported = u.PromptTokens
	f := float64(u.PromptTokens) / float64(sentEstimate)
	// Límites de cordura: un servidor con usage erróneo no debe desactivar
	// la compactación ni dispararla continuamente.
	m.factor = min(max(f, 0.25), 2.0)
}

// RawEstimate estima (sin calibrar) los tokens de la conversación completa.
func (m *Manager) RawEstimate() int {
	msgs := m.Messages()
	m.mu.Lock()
	defer m.mu.Unlock()
	return EstimateMessages(msgs) + m.overhead
}

// Usage devuelve el uso estimado de contexto. calibrated indica si la
// estimación está ajustada con datos reales del servidor. Nunca es un
// conteo exacto (§84): el servidor solo informa tras cada petición.
func (m *Manager) Usage() (tokens int, calibrated bool) {
	raw := m.RawEstimate()
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.factor > 0 {
		return int(float64(raw) * m.factor), true
	}
	return raw, false
}

// Calibrate convierte una estimación cruda en calibrada.
func (m *Manager) Calibrate(raw int) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.factor > 0 {
		return int(float64(raw) * m.factor)
	}
	return raw
}

// LastReported devuelve prompt_tokens de la última respuesta (0 si no hay).
func (m *Manager) LastReported() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.lastReported
}

// Split divide el historial para compactarlo: old se resumirá y kept se
// conserva completo. kept cabe en keepTokens (estimación cruda) si es posible.
//
// El corte solo se hace en fronteras seguras: justo antes de un mensaje del
// asistente o de un mensaje del usuario que no sea resultado de
// herramientas. Así nunca se separa una llamada a herramienta de su
// resultado, que la API rechazaría.
// ok=false si no hay nada que compactar.
func (m *Manager) Split(keepTokens int) (old, kept []llm.Message, ok bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	h := m.history
	var cuts []int
	for i := 1; i < len(h); i++ {
		if h[i].Role == llm.RoleAssistant || h[i].Role == llm.RoleUser && !IsToolResult(h[i]) {
			cuts = append(cuts, i)
		}
	}
	if len(cuts) == 0 {
		return nil, nil, false
	}
	cut := cuts[len(cuts)-1] // por defecto, conservar el mínimo
	for _, c := range cuts {
		if EstimateMessages(h[c:]) <= keepTokens {
			cut = c
			break
		}
	}
	old = append([]llm.Message(nil), h[:cut]...)
	kept = append([]llm.Message(nil), h[cut:]...)
	return old, kept, true
}

// EstimateMessages estima los tokens de una lista de mensajes.
func EstimateMessages(msgs []llm.Message) int {
	total := 0
	for _, msg := range msgs {
		total += EstimateTokens(msg.Content) + perMessageOverhead
		for _, tc := range msg.ToolCalls {
			total += EstimateTokens(tc.Function.Name) + EstimateTokens(tc.Function.Arguments)
		}
	}
	return total
}

// EstimateTokens es una estimación conservadora (tiende a sobreestimar).
//
// Los tokenizadores BPE de modelos actuales promedian ~4 bytes por token en
// inglés y algo menos en español o código. Usar 3 bytes por token hace que la
// compactación se dispare antes de tiempo en lugar de tarde, que es el fallo
// menos grave. No es un conteo exacto y nunca se presenta como tal.
func EstimateTokens(s string) int {
	if s == "" {
		return 0
	}
	return (len(s) + 2) / 3
}