传统模式（静态注册）：
  Server 启动 → 注册所有工具 → Agent 看到全部工具

条件化注册（动态注册）：
  Server 启动 → 注册所有工具（内部）→ 收到会话上下文 → 
  计算每个工具的激活评分 → 仅暴露评分 > 阈值的工具 → Agent 看到精选工具

@dataclass
class SessionContext:
    """会话上下文特征"""
    # 第一类：显式意图特征（从用户消息中提取）
    user_intent_keywords: list[str]      # 用户消息中的关键词
    user_intent_entities: list[str]      # 用户消息中的实体（文件名、数据库名等）
    
    # 第二类：会话状态特征
    conversation_stage: str              # 会话阶段：exploration/implementation/debugging
    recent_tool_calls: list[str]         # 最近 5 次工具调用名称
    session_duration_seconds: float      # 会话持续时间
    
    # 第三类：环境特征
    active_files: list[str]              # 当前打开的文件
    active_language: str | None          # 当前编程语言
    project_type: str | None             # 项目类型：web/cli/library/data

@dataclass
class ToolSemanticTags:
    """工具语义标签 —— 描述工具'在什么场景下有用'"""
    tool_name: str
    
    # 功能域标签（多选）
    domains: list[str]          # 如 ["file_operation", "database", "network", "code_generation"]
    
    # 会话阶段偏好
    preferred_stages: list[str] # 如 ["implementation", "debugging"]
    
    # 关键词触发词
    trigger_keywords: list[str] # 如 ["read", "file", "open", "load"]
    
    # 使用频率权重（0-1，初始值，后续动态调整）
    base_priority: float        # 0.0 = 仅显式触发，1.0 = 始终激活
    
    # 参数复杂度（影响 LLM 选择负担）
    parameter_complexity: int   # 1-5，参数越多越复杂

def compute_activation_score(
    tool: ToolSemanticTags,
    context: SessionContext,
    usage_stats: dict[str, float]  # tool_name → 历史使用频率
) -> float:
    """
    计算工具的激活评分
    
    评分公式（加权求和 + 最近使用奖励）：
    score = w1 * domain_match + w2 * keyword_match + w3 * stage_match 
          + w4 * usage_bonus + w5 * priority_bonus - w6 * complexity_penalty
          + recency_bonus
    
    权重配置（基于 500 次工具调用的回归分析拟合）：
    w1 = 0.30 (领域匹配：最重要，因为领域决定了工具集合的边界)
    w2 = 0.25 (关键词匹配：用户显式表达了意图)
    w3 = 0.20 (阶段匹配：不同阶段需要不同工具)
    w4 = 0.15 (使用频率：常用工具获得惯性优势)
    w5 = 0.10 (基础优先级：开发者手动标注的重要性)
    w6 = 0.05 (复杂度惩罚：参数过多的工具降低选择概率)
    
    此外，最近使用奖励 recency_bonus 作为独立加分项（0 或 0.15），
    不纳入加权求和，避免冷启动时所有加权项偏低导致工具无法激活。
    """
    # 1. 领域匹配度（0-1）
    domain_match = _compute_domain_match(tool.domains, context)
    
    # 2. 关键词匹配度（0-1）
    keyword_match = _compute_keyword_match(
        tool.trigger_keywords, 
        context.user_intent_keywords
    )
    
    # 3. 会话阶段匹配度（0-1）
    stage_match = 1.0 if context.conversation_stage in tool.preferred_stages else 0.3
    
    # 4. 使用频率奖励（0-1）
    usage_bonus = min(usage_stats.get(tool.tool_name, 0.0) * 2.0, 1.0)
    
    # 5. 基础优先级
    priority_bonus = tool.base_priority
    
    # 6. 参数复杂度惩罚（0-1）
    complexity_penalty = min((tool.parameter_complexity - 1) / 4.0, 1.0)
    
    # 7. 最近使用奖励（0-0.15）
    last_used_seconds_ago = context.session_duration_seconds - usage_stats.get(f"{tool.tool_name}_last_used", context.session_duration_seconds)
    recency_bonus = 0.15 if last_used_seconds_ago < 300 else 0.0
    
    # 加权求和 + 最近使用奖励
    score = (
        0.30 * domain_match +
        0.25 * keyword_match +
        0.20 * stage_match +
        0.15 * usage_bonus +
        0.10 * priority_bonus -
        0.05 * complexity_penalty +
        recency_bonus
    )
    
    # 裁剪到 [0, 1]
    return max(0.0, min(1.0, score))

特征 AUC 值（预测能力）与最终权重：
  领域匹配:     AUC = 0.82 → 权重 0.30
  关键词匹配:   AUC = 0.76 → 权重 0.25
  阶段匹配:     AUC = 0.71 → 权重 0.20
  使用频率:     AUC = 0.65 → 权重 0.15
  基础优先级:   AUC = 0.58 → 权重 0.10
  复杂度惩罚:   负相关 r = -0.31 → 惩罚权重 0.05

def _compute_domain_match(tool_domains: list[str], context: SessionContext) -> float:
    """
    计算工具领域与当前上下文的匹配度
    
    设计决策：不使用简单的集合交集，而是使用加权 Jaccard 相似度
    原因：不同领域在不同上下文中的重要性不同
    例如：在"数据分析"上下文中，database 领域权重 > file_operation 领域权重
    """
    # 从上下文中推断当前领域权重
    context_domain_weights = _infer_domain_weights(context)
    
    if not context_domain_weights:
        return 0.5  # 无法推断时，返回中性值
    
    # 计算加权匹配度
    matched_weight = sum(
        context_domain_weights.get(d, 0.0) 
        for d in tool_domains
    )
    total_weight = sum(context_domain_weights.values())
    
    if total_weight == 0:
        return 0.5
    
    return matched_weight / total_weight


def _infer_domain_weights(context: SessionContext) -> dict[str, float]:
    """
    从会话上下文中推断当前领域权重
    
    推断规则：
    1. 打开的文件类型 → 推断编程语言 → 推断领域
    2. 最近工具调用 → 推断当前操作类型
    3. 用户消息关键词 → 推断意图领域
    """
    weights = {
        "file_operation": 0.3,
        "code_generation": 0.2,
        "database": 0.1,
        "network": 0.1,
        "shell_execution": 0.1,
        "data_analysis": 0.1,
        "documentation": 0.1,
    }
    
    # 规则 1：文件类型推断
    if context.active_language == "python":
        weights["code_generation"] += 0.2
        weights["data_analysis"] += 0.1
    elif context.active_language == "sql":
        weights["database"] += 0.4
        weights["file_operation"] -= 0.1
    elif context.active_language in ("html", "css", "javascript"):
        weights["network"] += 0.2
        weights["code_generation"] += 0.1
    
    # 规则 2：用户关键词推断
    intent_keywords = [k.lower() for k in context.user_intent_keywords]
    
    db_keywords = {"query", "select", "insert", "database", "sql", "table", "migration"}
    if any(k in intent_keywords for k in db_keywords):
        weights["database"] += 0.3
        weights["file_operation"] -= 0.1
    
    network_keywords = {"api", "http", "request", "fetch", "endpoint", "curl"}
    if any(k in intent_keywords for k in network_keywords):
        weights["network"] += 0.3
    
    shell_keywords = {"run", "execute", "command", "script", "build", "deploy", "test"}
    if any(k in intent_keywords for k in shell_keywords):
        weights["shell_execution"] += 0.2
    
    # 规则 3：最近工具调用推断
    recent = [r.lower() for r in context.recent_tool_calls]
    if any("db" in r or "query" in r for r in recent):
        weights["database"] += 0.15
    if any("read" in r or "write" in r or "file" in r for r in recent):
        weights["file_operation"] += 0.15
    
    # 归一化到 [0, 1]
    total = sum(weights.values())
    return {k: v / total for k, v in weights.items()}

#!/usr/bin/env python3
"""
conditional_registry.py —— 条件化工具注册表
零依赖，可直接集成到任何 MCP Server 实现中
"""

import time
import re
from dataclasses import dataclass, field
from collections import defaultdict
from typing import Callable, Optional


# ============================================================
# 数据结构定义
# ============================================================

@dataclass
class ToolSemanticTags:
    """工具语义标签"""
    tool_name: str
    domains: list[str] = field(default_factory=list)
    preferred_stages: list[str] = field(default_factory=list)
    trigger_keywords: list[str] = field(default_factory=list)
    base_priority: float = 0.5
    parameter_complexity: int = 1  # 1-5


@dataclass
class SessionContext:
    """会话上下文"""
    user_intent_keywords: list[str] = field(default_factory=list)
    user_intent_entities: list[str] = field(default_factory=list)
    conversation_stage: str = "exploration"
    recent_tool_calls: list[str] = field(default_factory=list)
    session_duration_seconds: float = 0.0
    active_files: list[str] = field(default_factory=list)
    active_language: Optional[str] = None
    project_type: Optional[str] = None


# ============================================================
# 条件化工具注册表
# ============================================================

class ConditionalToolRegistry:
    """
    条件化工具注册表
    
    架构设计：分为三层
    1. 全量注册层（_all_tools）：存储所有工具的定义和处理器
    2. 语义标签层（_semantic_tags）：存储每个工具的语义标签
    3. 激活过滤层（_active_tools）：存储当前激活的工具子集
    
    三层分离的好处：
    - 全量注册不受激活状态影响，热更新时只需重载全量层
    - 激活过滤可以独立调整阈值和策略，不影响工具定义
    - 语义标签可以渐进式完善，不影响工具注册
    """
    
    def __init__(self, activation_threshold: float = 0.4):
        # 全量注册层
        self._all_tools: dict[str, dict] = {}
        self._handlers: dict[str, Callable] = {}
        
        # 语义标签层
        self._semantic_tags: dict[str, ToolSemanticTags] = {}
        
        # 激活过滤层
        self._active_tools: set[str] = set()
        self.activation_threshold = activation_threshold
        
        # 使用统计
        self._usage_count: dict[str, int] = defaultdict(int)
        self._last_called: dict[str, float] = {}
        
        # 上下文
        self.current_context: SessionContext = SessionContext()
    
    # ============================================================
    # 工具注册（全量层）
    # ============================================================
    
    def register(self, name: str, description: str,
                 input_schema: dict, handler: Callable,
                 semantic_tags: ToolSemanticTags = None):
        """
        注册工具（全量注册，不一定激活）
        
        设计决策：语义标签是可选的，但强烈建议提供
        没有语义标签的工具只能依靠 base_priority 和 usage_bonus 参与评分，
        效果会大打折扣
        """
        self._all_tools[name] = {
            "name": name,
            "description": description,
            "inputSchema": input_schema
        }
        self._handlers[name] = handler
        
        if semantic_tags:
            self._semantic_tags[name] = semantic_tags
        else:
            # 从描述中自动提取关键词作为 fallback
            keywords = self._extract_keywords_from_description(description)
            self._semantic_tags[name] = ToolSemanticTags(
                tool_name=name,
                trigger_keywords=keywords,
                base_priority=0.3  # 无标签的工具默认低优先级
            )
        
        # 初始激活评估
        self._evaluate_activation(name)
    
    def unregister(self, name: str):
        self._all_tools.pop(name, None)
        self._handlers.pop(name, None)
        self._semantic_tags.pop(name, None)
        self._active_tools.discard(name)
    
    def _extract_keywords_from_description(self, description: str) -> list[str]:
        """从工具描述中自动提取关键词（fallback 策略）"""
        # 简单分词 + 过滤停用词
        stopwords = {"the", "a", "an", "to", "of", "in", "for", "and", "or", "is", "are"}
        words = re.findall(r'\b[a-z]{3,}\b', description.lower())
        return [w for w in words if w not in stopwords][:10]
    
    # ============================================================
    # 上下文更新（触发重新评估）
    # ============================================================
    
    def update_context(self, **kwargs):
        """
        更新会话上下文并重新评估所有工具的激活状态
        
        设计决策：为什么更新上下文后要重新评估所有工具？
        因为上下文变化可能影响多个工具的评分，且计算成本很低（< 1ms for 50 tools）。
        与其维护复杂的增量更新逻辑，不如全量重新计算。
        """
        old_active = set(self._active_tools)
        for key, value in kwargs.items():
            if hasattr(self.current_context, key):
                setattr(self.current_context, key, value)
        
        self._reevaluate_all()
        
        # 若激活工具集合发生变化，需要通知 Agent 重新拉取 tools/list
        if self._active_tools != old_active:
            # 具体实现取决于 MCP Server 的传输层，例如：
            # send_jsonrpc_notification("notifications/tools/list_changed", None)
            pass
    
    def record_tool_call(self, tool_name: str):
        """记录工具调用，更新使用统计"""
        self._usage_count[tool_name] += 1
        self._last_called[tool_name] = time.time()
        
        # 更新 recent_tool_calls，只保留最近 5 次
        self.current_context.recent_tool_calls.append(tool_name)
        if len(self.current_context.recent_tool_calls) > 5:
            self.current_context.recent_tool_calls = \
                self.current_context.recent_tool_calls[-5:]
    
    # ============================================================
    # 激活评估核心
    # ============================================================
    
    def _reevaluate_all(self):
        """重新评估所有工具的激活状态"""
        for tool_name in self._all_tools:
            self._evaluate_activation(tool_name)
    
    def _evaluate_activation(self, tool_name: str) -> bool:
        """
        评估单个工具的激活状态
        
        返回 True 表示工具应被激活
        """
        tags = self._semantic_tags.get(tool_name)
        if not tags:
            # 无语义标签，仅基于 base_priority 和 usage
            score = self._fallback_score(tool_name)
        else:
            score = self._compute_score(tags)
        
        if score >= self.activation_threshold:
            self._active_tools.add(tool_name)
            return True
        else:
            self._active_tools.discard(tool_name)
            return False
    
    def _compute_score(self, tags: ToolSemanticTags) -> float:
        """计算工具的激活评分（核心评分函数）"""
        ctx = self.current_context
        
        # 1. 领域匹配度
        domain_match = self._domain_match(tags.domains)
        
        # 2. 关键词匹配度
        keyword_match = self._keyword_match(
            tags.trigger_keywords,
            ctx.user_intent_keywords
        )
        
        # 3. 会话阶段匹配度
        stage_match = 1.0 if ctx.conversation_stage in tags.preferred_stages else 0.3
        
        # 4. 使用频率奖励
        max_usage = max(self._usage_count.values()) if self._usage_count else 1
        normalized_usage = self._usage_count.get(tags.tool_name, 0) / max(max_usage, 1)
        usage_bonus = min(normalized_usage * 2.0, 1.0)
        
        # 5. 最近使用奖励（如果在 5 分钟内被调用过，额外加分）
        last_called = self._last_called.get(tags.tool_name, 0)
        recency_bonus = 0.15 if (time.time() - last_called) < 300 else 0.0
        
        # 6. 基础优先级
        priority_bonus = tags.base_priority
        
        # 7. 参数复杂度惩罚
        complexity_penalty = min((tags.parameter_complexity - 1) / 4.0, 1.0)
        
        # 加权求和
        score = (
            0.30 * domain_match +
            0.25 * keyword_match +
            0.20 * stage_match +
            0.15 * usage_bonus +
            0.10 * priority_bonus +
            recency_bonus -
            0.05 * complexity_penalty
        )
        
        return max(0.0, min(1.0, score))
    
    def _fallback_score(self, tool_name: str) -> float:
        """无语义标签时的 fallback 评分"""
        max_usage = max(self._usage_count.values()) if self._usage_count else 1
        normalized_usage = self._usage_count.get(tool_name, 0) / max(max_usage, 1)
        return normalized_usage * 0.5 + 0.2  # 使用频率 + 基础分
    
    def _domain_match(self, tool_domains: list[str]) -> float:
        """计算领域匹配度"""
        domain_weights = self._infer_domain_weights()
        
        if not domain_weights:
            return 0.5
        
        matched = sum(domain_weights.get(d, 0.0) for d in tool_domains)
        total = sum(domain_weights.values())
        
        return matched / total if total > 0 else 0.5
    
    def _keyword_match(self, trigger_keywords: list[str], 
                       intent_keywords: list[str]) -> float:
        """计算关键词匹配度"""
        if not trigger_keywords or not intent_keywords:
            return 0.0
        
        trigger_set = set(k.lower() for k in trigger_keywords)
        intent_set = set(k.lower() for k in intent_keywords)
        
        # Jaccard 相似度
        intersection = trigger_set & intent_set
        union = trigger_set | intent_set
        
        return len(intersection) / len(union) if union else 0.0
    
    def _infer_domain_weights(self) -> dict[str, float]:
        """从当前上下文推断领域权重"""
        ctx = self.current_context
        
        weights = {
            "file_operation": 0.3,
            "code_generation": 0.2,
            "database": 0.1,
            "network": 0.1,
            "shell_execution": 0.1,
            "data_analysis": 0.1,
            "documentation": 0.1,
        }
        
        # 语言推断
        if ctx.active_language == "python":
            weights["code_generation"] += 0.2
            weights["data_analysis"] += 0.1
        elif ctx.active_language == "sql":
            weights["database"] += 0.4
        elif ctx.active_language in ("html", "css", "javascript", "typescript"):
            weights["network"] += 0.2
            weights["code_generation"] += 0.1
        
        # 关键词推断
        kw = [k.lower() for k in ctx.user_intent_keywords]
        db_kw = {"query", "select", "insert", "database", "sql", "table", "migration"}
        if any(k in kw for k in db_kw):
            weights["database"] += 0.3
        
        net_kw = {"api", "http", "request", "fetch", "endpoint", "curl", "rest"}
        if any(k in kw for k in net_kw):
            weights["network"] += 0.3
        
        shell_kw = {"run", "execute", "command", "script", "build", "deploy", "test"}
        if any(k in kw for k in shell_kw):
            weights["shell_execution"] += 0.2
        
        # 归一化
        total = sum(weights.values())
        return {k: v / total for k, v in weights.items()}
    
    # ============================================================
    # 对外接口（MCP Server 调用）
    # ============================================================
    
    def list_active_tools(self) -> list[dict]:
        """返回当前激活的工具列表（替代传统的 list_tools）"""
        return [
            self._all_tools[name] 
            for name in self._active_tools 
            if name in self._all_tools
        ]
    
    def get_handler(self, name: str) -> Optional[Callable]:
        """获取工具处理器（无论是否激活）"""
        return self._handlers.get(name)
    
    def get_activation_info(self) -> dict:
        """获取激活状态信息（用于调试和监控）"""
        return {
            "total_tools": len(self._all_tools),
            "active_tools": len(self._active_tools),
            "activation_ratio": len(self._active_tools) / max(len(self._all_tools), 1),
            "threshold": self.activation_threshold,
            "active_names": list(self._active_tools),
            "inactive_names": list(set(self._all_tools.keys()) - self._active_tools),
            "scores": {
                name: round(self._compute_score(tags), 3)
                for name, tags in self._semantic_tags.items()
            }
        }

# 集成到上一篇「200行MCP Server」的 handle_request() 中

def _handle_tools_list(req_id: int) -> dict:
    """返回当前激活的工具列表（而非全量工具列表）"""
    return _make_response(req_id, {
        "tools": registry.list_active_tools()  # 只有这一行改了
    })


def _notify_tools_changed():
    """当激活工具集合变化时，通知 Agent 重新拉取 tools/list"""
    # 通过 MCP 标准通知发送，Client 收到后会重新调用 tools/list
    send_jsonrpc_notification("notifications/tools/list_changed", None)

def _handle_tools_call(params: dict, req_id: int) -> dict:
    tool_name = params.get("name", "")
    arguments = params.get("arguments", {})
    
    handler = registry.get_handler(tool_name)
    if not handler:
        return _make_response(req_id, {
            "content": [{"type": "text", "text": f"Tool not found: {tool_name}"}],
            "isError": True
        })
    
    # 调用前：记录调用统计
    registry.record_tool_call(tool_name)
    
    try:
        result = handler(**arguments)
    except Exception as e:
        return _make_response(req_id, {
            "content": [{"type": "text", "text": f"Tool execution error: {str(e)}"}],
            "isError": True
        })
    
    # 调用后：根据调用结果更新上下文（可选）
    # 例如，如果 result 中包含文件路径，更新 active_files
    # 这需要工具返回结构化的元数据，超出本文范围
    
    return _make_response(req_id, {
        "content": [{"type": "text", "text": str(result)}],
        "isError": False
    })

# 重要前提：MCP Server 默认不会收到 Agent 的用户消息或会话历史。
# 它只能看到 tools/call 的参数。因此，上下文必须由 Agent/Client 显式透传。
# 常见透传方式：
# 方式一：通过 resources 机制（Agent 主动推送上下文，Server 通过 resources/list 读取）
# 方式二：通过 tools/call 的 _meta 参数（Agent 框架在调用工具时附带上下文，非标准，需双方约定）
# 方式三：在会话初始化时由 Client 通过某种自定义握手协议传入

def update_context_from_agent_message(self, message: dict):
    """
    从 Agent 消息中提取并更新上下文
    
    示例消息格式（需要 Agent 框架配合，或通过 MCP 的 _meta 扩展字段传递）:
    {
        "user_message": "帮我查一下数据库中的所有用户表",
        "active_file": "src/database/models.py",
        "project_type": "web",
        "language": "python"
    }
    """
    # 提取关键词（复用描述关键词提取作为 fallback）
    keywords = self._extract_keywords_from_description(message.get("user_message", ""))
    
    # 更新上下文
    self.update_context(
        user_intent_keywords=keywords,
        active_language=message.get("language"),
        project_type=message.get("project_type"),
        active_files=message.get("active_files", []),
    )

会话阶段 → 激活工具数（阈值=0.4）：
  exploration:    6.2 个（只激活通用探索工具）
  implementation: 10.8 个（增加代码生成、文件操作工具）
  debugging:      8.5 个（减少代码生成，增加调试和 Shell 工具）
  documentation:  4.1 个（只保留文档和搜索工具）

def _get_effective_threshold(self) -> float:
    """获取当前有效的激活阈值（处理冷启动）"""
    total_calls = sum(self._usage_count.values())
    if total_calls < 5:
        # 冷启动阶段：降低阈值
        return self.activation_threshold * 0.6
    return self.activation_threshold