xbot 上下文管理架构抽象与开关机制设计

package agent

import (
    "context"
    "xbot/llm"
    "xbot/session"
)

// ContextMode 上下文管理模式
type ContextMode string

const (
    // ContextModePhase1 Phase 1 双视图架构（当前默认）
    ContextModePhase1 ContextMode = "phase1"
    // ContextModePhase2 Phase 2 三层渐进压缩（未来实现）
    ContextModePhase2 ContextMode = "phase2"
    // ContextModeNone 禁用自动上下文压缩
    ContextModeNone ContextMode = "none"
)

// ValidContextModes 所有可能的上下文模式
var ValidContextModes = []ContextMode{ContextModePhase1, ContextModePhase2, ContextModeNone}

// IsValidContextMode 检查是否为有效的上下文模式
func IsValidContextMode(mode ContextMode) bool {
    for _, m := range ValidContextModes {
        if m == mode {
            return true
        }
    }
    return false
}

// ContextManager 上下文管理器统一接口。
// 所有压缩策略实现此接口，通过策略模式实现新旧架构可切换。
type ContextManager interface {
    // Mode 返回当前管理模式标识。
    Mode() ContextMode

    // ShouldCompress 判断是否需要触发自动压缩。
    // 参数：
    //   - messages: 当前上下文消息
    //   - model: LLM 模型名（用于 token 计数）
    //   - toolTokens: 工具定义占用的 token 数
    ShouldCompress(messages []llm.ChatMessage, model string, toolTokens int) bool

    // Compress 执行上下文压缩。
    Compress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error)

    // ManualCompress 手动压缩（/compress 命令使用）。
    // 关键契约：无论当前模式如何，ManualCompress 都应尽力执行压缩。
    // 即使 auto=false 的 noopManager，ManualCompress 也降级到 Phase1 执行，
    // 保留 /compress 手动命令始终可用的现有语义。
    ManualCompress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error)

    // ContextInfo 返回上下文统计信息（/context info 命令使用）。
    ContextInfo(messages []llm.ChatMessage, model string, toolTokens int) *ContextStats

    // SessionHook 返回压缩后的 session 持久化钩子（可选，返回 nil 表示无特殊处理）。
    // Phase 2 可能需要在此钩子中做额外操作（如更新话题分区索引）。
    SessionHook() SessionCompressHook

    // SetMemoryTools 注册 memory 相关工具（recall_memory / store_memory）。
    // ContextManager 需要感知这些工具以在压缩时正确处理记忆消息。
    SetMemoryTools(tools []Tool)
}

// ContextStats 上下文统计信息
type ContextStats struct {
    SystemTokens    int
    UserTokens      int
    AssistantTokens int
    ToolMsgTokens   int
    ToolDefTokens   int
    TotalTokens     int
    MaxTokens       int
    Threshold       int
    Mode            ContextMode
    IsRuntimeOverride bool // 是否为运行时覆盖
    DefaultMode     ContextMode
}

// SessionCompressHook 压缩后的 session 处理钩子
type SessionCompressHook interface {
    // AfterPersist 在 session 持久化压缩结果后调用
    AfterPersist(ctx context.Context, tenantSession *session.TenantSession, result *CompressResult)
}

import "sync"

// ContextManagerConfig 上下文管理器配置。
// 包含全局配置（环境变量/Agent.Config）和运行时开关（命令行切换）。
// 所有读写操作通过 sync.RWMutex 保护，确保并发安全。
type ContextManagerConfig struct {
    mu sync.RWMutex

    // MaxContextTokens 最大上下文 token 数（默认 100000）
    MaxContextTokens int
    // CompressionThreshold 触发压缩的 token 比例阈值（默认 0.7）
    CompressionThreshold float64

    // DefaultMode 默认压缩模式（启动时决定，来自环境变量或 Agent.Config）
    DefaultMode ContextMode

    // runtimeMode 运行时模式覆盖（通过 /context mode 命令切换）
    // 空值表示使用 DefaultMode，非空值覆盖 DefaultMode
    runtimeMode ContextMode
}

// EffectiveMode 返回当前生效的模式（RuntimeMode 优先）。
// 读锁保护。
func (c *ContextManagerConfig) EffectiveMode() ContextMode {
    c.mu.RLock()
    defer c.mu.RUnlock()
    if c.runtimeMode != "" {
        return c.runtimeMode
    }
    return c.DefaultMode
}

// RuntimeMode 返回当前运行时覆盖模式（无覆盖时返回空字符串）。
// 读锁保护。
func (c *ContextManagerConfig) RuntimeMode() ContextMode {
    c.mu.RLock()
    defer c.mu.RUnlock()
    return c.runtimeMode
}

// SetRuntimeMode 设置运行时模式覆盖。
// 写锁保护。
func (c *ContextManagerConfig) SetRuntimeMode(mode ContextMode) {
    c.mu.Lock()
    defer c.mu.Unlock()
    c.runtimeMode = mode
}

// ResetRuntimeMode 清除运行时覆盖，恢复默认模式。
// 写锁保护。
func (c *ContextManagerConfig) ResetRuntimeMode() {
    c.mu.Lock()
    defer c.mu.Unlock()
    c.runtimeMode = ""
}

type Agent struct {
    // ... 现有字段 ...

    // 上下文压缩配置（新增，替换旧的三个字段）
    contextManagerConfig *ContextManagerConfig
    contextManagerMu     sync.RWMutex            // 保护 contextManager 的并发读写
    contextManager       ContextManager           // 当前激活的管理器实例

    // 以下旧字段将在步骤 7（清理）中删除，步骤 1-6 期间保留用于向后兼容：
    // maxContextTokens     int
    // compressionThreshold float64
    // enableAutoCompress   bool
}

// GetContextManager 获取当前上下文管理器（读锁保护）。
// 用于 buildMainRunConfig / buildSubAgentRunConfig / handleCompress 等场景。
func (a *Agent) GetContextManager() ContextManager {
    a.contextManagerMu.RLock()
    defer a.contextManagerMu.RUnlock()
    return a.contextManager
}

// SetContextManager 替换当前上下文管理器（写锁保护）。
// 用于 /context mode 命令运行时切换。
func (a *Agent) SetContextManager(cm ContextManager) {
    a.contextManagerMu.Lock()
    defer a.contextManagerMu.Unlock()
    a.contextManager = cm
}

type Config struct {
    // ... 现有字段 ...

    // 上下文管理配置（新增）
    // 优先级：ContextMode > EnableAutoCompress 旧字段
    // 默认 "phase1"（保持现有行为）
    ContextMode ContextMode

    // 旧字段保留向后兼容（步骤 1-6 期间保留）：
    // MaxContextTokens     int
    // CompressionThreshold float64
    // EnableAutoCompress   bool
}

package agent

import (
    "context"
    "fmt"
    "strings"

    "xbot/llm"
    log "xbot/logger"
)

// phase1Manager Phase 1 双视图压缩管理器。
// 封装现有 compress.go 中的逻辑，行为与现有完全一致。
// 不持有 *Agent 引用，仅依赖配置和独立函数。
type phase1Manager struct {
    config *ContextManagerConfig
}

func newPhase1Manager(cfg *ContextManagerConfig) *phase1Manager {
    return &phase1Manager{config: cfg}
}

func (m *phase1Manager) Mode() ContextMode { return ContextModePhase1 }

func (m *phase1Manager) ShouldCompress(messages []llm.ChatMessage, model string, toolTokens int) bool {
    if len(messages) <= 3 {
        return false
    }
    msgTokens, err := llm.CountMessagesTokens(messages, model)
    if err != nil {
        return false
    }
    tokenCount := msgTokens + toolTokens
    threshold := int(float64(m.config.MaxContextTokens) * m.config.CompressionThreshold)
    return tokenCount >= threshold
}

func (m *phase1Manager) Compress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    return compressMessages(ctx, messages, client, model)
}

func (m *phase1Manager) ManualCompress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    return compressMessages(ctx, messages, client, model)
}

func (m *phase1Manager) ContextInfo(messages []llm.ChatMessage, model string, toolTokens int) *ContextStats {
    cfg := m.config
    var systemTokens, userTokens, assistantTokens, toolMsgTokens int

    for _, msg := range messages {
        tokens, err := llm.CountMessagesTokens([]llm.ChatMessage{msg}, model)
        if err != nil {
            continue
        }
        switch msg.Role {
        case "system":
            systemTokens += tokens
        case "user":
            userTokens += tokens
        case "assistant":
            assistantTokens += tokens
        case "tool":
            toolMsgTokens += tokens
        }
    }

    total := systemTokens + userTokens + assistantTokens + toolMsgTokens + toolTokens
    threshold := int(float64(cfg.MaxContextTokens) * cfg.CompressionThreshold)

    return &ContextStats{
        SystemTokens:    systemTokens,
        UserTokens:      userTokens,
        AssistantTokens: assistantTokens,
        ToolMsgTokens:   toolMsgTokens,
        ToolDefTokens:   toolTokens,
        TotalTokens:     total,
        MaxTokens:       cfg.MaxContextTokens,
        Threshold:       threshold,
        Mode:            cfg.EffectiveMode(),
        IsRuntimeOverride: cfg.RuntimeMode() != "",
        DefaultMode:     cfg.DefaultMode,
    }
}

func (m *phase1Manager) SessionHook() SessionCompressHook { return nil }

// compressMessages 使用 LLM 压缩对话历史（独立函数，不依赖 Agent receiver）。
// 逻辑与现有 compressContext() 完全一致，仅为消除 phase1Manager 对 *Agent 的引用。
// 现有 compressContext() 改为调用此函数：
//
//   func (a *Agent) compressContext(...) { return compressMessages(...) }
func compressMessages(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    // 第一步：找到尾部安全切割点
    tailStart := len(messages)
    for i := len(messages) - 1; i >= 1; i-- {
        msg := messages[i]
        if msg.Role == "user" {
            tailStart = i
            break
        }
        if msg.Role == "assistant" && len(msg.ToolCalls) == 0 {
            tailStart = i
            break
        }
        if i == 1 {
            tailStart = 1
        }
    }
    // ... 后续逻辑完全复用现有 compressContext 的代码 ...
    // （extractDialogueFromTail、thinTail、LLM 压缩调用等，不动）
}

// phase2Manager Phase 2 三层渐进压缩管理器。
// 目前为空壳实现，Phase 2 实现时填充。
// NewContextManager() 中会自动降级到 Phase 1。
type phase2Manager struct {
    config *ContextManagerConfig
    // Phase 2 专用字段（未来实现）：
    // topicPartitioner  *TopicPartitioner
    // qualityChecker    *QualityChecker
    // progressiveLevels [3]CompressionLevel
}

func newPhase2Manager(cfg *ContextManagerConfig) *phase2Manager {
    return &phase2Manager{config: cfg}
}

func (m *phase2Manager) Mode() ContextMode { return ContextModePhase2 }

func (m *phase2Manager) ShouldCompress(messages []llm.ChatMessage, model string, toolTokens int) bool {
    // Phase 2 智能触发逻辑（未来实现）
    // 临时 fallback：与 Phase 1 相同的阈值判断
    if len(messages) <= 3 {
        return false
    }
    msgTokens, err := llm.CountMessagesTokens(messages, model)
    if err != nil {
        return false
    }
    tokenCount := msgTokens + toolTokens
    threshold := int(float64(m.config.MaxContextTokens) * m.config.CompressionThreshold)
    return tokenCount >= threshold
}

func (m *phase2Manager) Compress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    // TODO: Phase 2 三层渐进压缩实现
    return nil, fmt.Errorf("phase 2 compression not yet implemented")
}

func (m *phase2Manager) ManualCompress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    // ManualCompress 契约：无论模式如何，都尽力执行。
    // Phase 2 未实现时，降级到 compressMessages（Phase 1 逻辑）。
    return compressMessages(ctx, messages, client, model)
}

func (m *phase2Manager) ContextInfo(messages []llm.ChatMessage, model string, toolTokens int) *ContextStats {
    // 复用 Phase 1 的统计逻辑（统计方式相同）
    return newPhase1Manager(m.config).ContextInfo(messages, model, toolTokens)
}

func (m *phase2Manager) SessionHook() SessionCompressHook { return nil }

// noopManager 禁用自动压缩的管理器。
// ShouldCompress 始终返回 false，但 ManualCompress 仍可执行（降级到 Phase 1）。
type noopManager struct {
    config  *ContextManagerConfig
    phase1  *phase1Manager // 内嵌 Phase1 用于 ManualCompress 和 ContextInfo
}

func newNoopManager(cfg *ContextManagerConfig) *noopManager {
    return &noopManager{
        config: cfg,
        phase1: newPhase1Manager(cfg),
    }
}

func (m *noopManager) Mode() ContextMode { return ContextModeNone }

func (m *noopManager) ShouldCompress([]llm.ChatMessage, string, int) bool {
    return false // 自动压缩始终禁用
}

func (m *noopManager) Compress(context.Context, []llm.ChatMessage, llm.LLM, string) (*CompressResult, error) {
    // 自动路径不应到达这里（ShouldCompress 返回 false）
    return nil, fmt.Errorf("auto compression is disabled (mode=none)")
}

func (m *noopManager) ManualCompress(ctx context.Context, messages []llm.ChatMessage, client llm.LLM, model string) (*CompressResult, error) {
    // /compress 手动命令：降级到 Phase 1 执行，保持命令始终可用
    return m.phase1.ManualCompress(ctx, messages, client, model)
}

func (m *noopManager) ContextInfo(messages []llm.ChatMessage, model string, toolTokens int) *ContextStats {
    // 仍返回完整统计信息（复用 Phase1 的实现）
    stats := m.phase1.ContextInfo(messages, model, toolTokens)
    stats.Mode = ContextModeNone
    return stats
}

func (m *noopManager) SessionHook() SessionCompressHook { return nil }

// NewContextManager 根据配置创建对应的 ContextManager 实例。
func NewContextManager(cfg *ContextManagerConfig) ContextManager {
    mode := cfg.EffectiveMode()
    switch mode {
    case ContextModePhase2:
        // Phase 2: 即使未实现也创建 phase2Manager，
        // Compress 会返回错误并自动降级，ManualCompress 降级到 Phase 1
        log.WithField("mode", mode).Warn("Phase 2 not yet implemented, will fallback to Phase 1 on actual compression")
        return newPhase2Manager(cfg)
    case ContextModeNone:
        return newNoopManager(cfg)
    case ContextModePhase1, "":
        return newPhase1Manager(cfg)
    default:
        log.WithField("mode", mode).Warnf("Unknown context mode %q, falling back to Phase 1", mode)
        return newPhase1Manager(cfg)
    }
}

// config/config.go Load() 函数中 Agent 配置段新增：
func Load() Config {
    // ... 现有代码 ...
    agent := AgentConfig{
        // ... 现有字段 ...
        ContextMode:         getEnvString("AGENT_CONTEXT_MODE", ""),      // 新增
        EnableAutoCompress:  getEnvBoolOrDefault("AGENT_ENABLE_AUTO_COMPRESS", true),
        MaxContextTokens:    getEnvInt("AGENT_MAX_CONTEXT_TOKENS", 100000),
        CompressionThreshold: getEnvFloat("AGENT_COMPRESSION_THRESHOLD", 0.7),
    }
    // ...
}

// resolveContextMode 根据新旧配置确定上下文管理模式。
// 优先级：ContextMode > EnableAutoCompress > 默认 phase1
func resolveContextMode(cfg Config) ContextMode {
    // 1. 优先使用新配置
    if cfg.ContextMode != "" {
        if IsValidContextMode(cfg.ContextMode) {
            return cfg.ContextMode
        }
        log.WithField("mode", cfg.ContextMode).Warn("Invalid AGENT_CONTEXT_MODE, ignoring")
    }
    // 2. 向后兼容：旧字段
    if !cfg.EnableAutoCompress {
        return ContextModeNone
    }
    // 3. 默认 phase1
    return ContextModePhase1
}

// command_builtin.go:168-175
type contextCmd struct{}
func (c *contextCmd) Name() string      { return "/context" }
func (c *contextCmd) Aliases() []string { return nil }
func (c *contextCmd) Match(s string) bool {
    trimmed := strings.TrimSpace(s)
    return strings.ToLower(trimmed) == "/context"
}
func (c *contextCmd) Execute(ctx context.Context, a *Agent, msg bus.InboundMessage) (*bus.OutboundMessage, error) {
    return a.handleContext(ctx, msg, ...)
}

// command_builtin.go — 修改现有 contextCmd
type contextCmd struct{}

func (c *contextCmd) Name() string      { return "/context" }
func (c *contextCmd) Aliases() []string { return nil }

// Match 支持子命令匹配：/context, /context info, /context mode ...
func (c *contextCmd) Match(s string) bool {
    trimmed := strings.TrimSpace(strings.ToLower(s))
    return trimmed == "/context" || strings.HasPrefix(trimmed, "/context ")
}

func (c *contextCmd) Concurrent() bool  { return true }

func (c *contextCmd) Execute(ctx context.Context, a *Agent, msg bus.InboundMessage) (*bus.OutboundMessage, error) {
    content := strings.TrimSpace(msg.Content)
    subCmd := strings.TrimSpace(strings.TrimPrefix(strings.ToLower(content), "/context"))

    switch {
    case subCmd == "" || subCmd == "info":
        return a.handleContextInfo(ctx, msg)
    case strings.HasPrefix(subCmd, "mode"):
        mode := strings.TrimSpace(strings.TrimPrefix(subCmd, "mode"))
        return a.handleContextMode(ctx, msg, mode)
    default:
        return &bus.OutboundMessage{
            Channel: msg.Channel,
            ChatID:  msg.ChatID,
            Content: "未知子命令。用法: /context [info|mode [phase1|phase2|none|default]]",
        }, nil
    }
}

// handleContextMode 处理 /context mode 子命令
func (a *Agent) handleContextMode(ctx context.Context, msg bus.InboundMessage, modeStr string) (*bus.OutboundMessage, error) {
    cfg := a.contextManagerConfig

    if modeStr == "" {
        // 仅查询当前模式
        stats := a.GetContextManager().ContextInfo(nil, "", 0)
        overrideInfo := ""
        if stats.IsRuntimeOverride {
            overrideInfo = fmt.Sprintf("（运行时覆盖，默认为 %s）", stats.DefaultMode)
        }
        return &bus.OutboundMessage{
            Channel: msg.Channel,
            ChatID:  msg.ChatID,
            Content: fmt.Sprintf("当前上下文模式: %s %s", cfg.EffectiveMode(), overrideInfo),
        }, nil
    }

    target := ContextMode(modeStr)
    if target == "default" {
        cfg.ResetRuntimeMode()
        a.SetContextManager(NewContextManager(cfg))
        return &bus.OutboundMessage{
            Channel: msg.Channel,
            ChatID:  msg.ChatID,
            Content: fmt.Sprintf("已恢复默认上下文模式: %s", cfg.DefaultMode),
        }, nil
    }

    if !IsValidContextMode(target) {
        return &bus.OutboundMessage{
            Channel: msg.Channel,
            ChatID:  msg.ChatID,
            Content: "无效模式。可选: phase1, phase2, none, default",
        }, nil
    }

    // Phase 2 未实现时的额外提示
    extraMsg := ""
    if target == ContextModePhase2 {
        extraMsg = "（Phase 2 尚未实现，压缩时将自动降级到 Phase 1）"
    }

    // 原子操作：先设置配置，再替换 manager
    cfg.SetRuntimeMode(target)
    a.SetContextManager(NewContextManager(cfg))

    return &bus.OutboundMessage{
        Channel: msg.Channel,
        ChatID:  msg.ChatID,
        Content: fmt.Sprintf("已切换上下文模式: %s %s", target, extraMsg),
    }, nil
}

type RunConfig struct {
    // ... 现有字段 ...

    // AutoCompress 自动压缩配置（旧字段，向后兼容，步骤 7 删除）
    AutoCompress *CompressConfig

    // ContextManager 上下文管理器（新字段，优先级高于 AutoCompress）
    // 如果设置了此字段，AutoCompress 被忽略
    ContextManager ContextManager
}

// engine.go Run() 中修改 maybeCompress 闭包：
maybeCompress := func() {
    // 优先使用 ContextManager（新路径）
    if cm := cfg.ContextManager; cm != nil {
        toolDefs := cfg.Tools.AsDefinitionsForSession(sessionKey)
        toolTokens, _ := llm.CountToolsTokens(toolDefs, cfg.Model)

        if !cm.ShouldCompress(messages, cfg.Model, toolTokens) {
            return
        }

        // ... 通知逻辑不变 ...

        result, compressErr := cm.Compress(ctx, messages, cfg.LLMClient, cfg.Model)
        // ... 后续持久化逻辑不变 ...

        // 调用 SessionHook（如果有的话）
        if hook := cm.SessionHook(); hook != nil && cfg.Session != nil {
            hook.AfterPersist(ctx, cfg.Session, result)
        }
        return
    }

    // 回退路径：旧的 CompressFunc（向后兼容）
    cc := cfg.AutoCompress
    if cc == nil || len(messages) <= 3 {
        return
    }
    // ... 现有 maybeCompress 逻辑完全不变 ...
}

// 在输入超限强制压缩处，同样优先使用 ContextManager：
if cc := cfg.AutoCompress; cc != nil {
    // 旧路径（保持不变）
    result, compressErr := cc.CompressFunc(ctx, messages, cfg.LLMClient, cfg.Model)
    // ...
} else if cm := cfg.ContextManager; cm != nil {
    // 新路径：手动压缩（不检查阈值）
    result, compressErr := cm.ManualCompress(ctx, messages, cfg.LLMClient, cfg.Model)
    // ... 后续重试逻辑不变 ...
}

// handleCompress 处理 /compress 命令：手动触发上下文压缩
func (a *Agent) handleCompress(ctx context.Context, msg bus.InboundMessage, tenantSession *session.TenantSession) (*bus.OutboundMessage, error) {
    // 手动 /compress 始终可用，不受模式开关限制

    llmClient, model, _, _ := a.llmFactory.GetLLM(msg.SenderID)

    // buildPrompt、token 计数等外壳逻辑完全保留 ...
    messages, err := a.buildPrompt(ctx, msg, tenantSession)
    // ... 阈值检查、进度消息等保留不变 ...

    // 核心改动：通过 ContextManager.ManualCompress 执行
    cm := a.GetContextManager()
    result, err := cm.ManualCompress(ctx, messages, llmClient, model)
    // ... 后续 session 持久化逻辑完全保留不变 ...
}

func (a *Agent) buildMainRunConfig(...) RunConfig {
    // ... 现有代码 ...
    cfg := RunConfig{ /* ... */ }

    // 新路径：注入 ContextManager
    cfg.ContextManager = a.GetContextManager()

    // 旧路径兼容（步骤 7 删除）：
    // if cfg.ContextManager == nil && a.enableAutoCompress { ... }

    return cfg
}

func (a *Agent) buildSubAgentRunConfig(...) RunConfig {
    // ... 现有代码 ...
    cfg := RunConfig{ /* ... */ }

    // SubAgent 共享主 Agent 的 ContextManager（caps.Memory=true 时）
    if caps.Memory && a.enableAutoCompress {
        cfg.ContextManager = a.GetContextManager()
    }
    // SubAgent 不需要独立的 ContextManager——它与主 Agent 共用同一套压缩配置
    // 未来如果需要 SubAgent 独立模式，可通过 caps.ContextMode 字段扩展

    return cfg
}