初始化提交

This commit is contained in:
chenxin committed 2026-08-21 16:00:13 +08:00
commit 887c1f449e
963 files changed
+166869

No files matched your search

+36
View File
@@ -0,0 +1,36 @@
.PHONY: run-stdio run-sse run-dev build test clean lint
# 从项目根 .env 提取 API Key 并设为环境变量
#(@前缀隐藏配方回显,避免在终端暴露命令内容)
_ENV_KEYS = TAVILY_API_KEYS="$$(grep '^TAVILY_API_KEYS=' ../../.env | head -1 | cut -d= -f2-)" \
BRAVE_API_KEYS="$$(grep '^BRAVE_API_KEYS=' ../../.env | head -1 | cut -d= -f2-)"
# 本地 stdio 模式(用于 Claude Desktop / Cursor 等 MCP Client 子进程集成)
run-stdio:
@$(_ENV_KEYS) go run ./cmd/server --stdio --config ../../config/mcp/web-search.yaml
# 本地 SSE 模式(Streamable HTTP,灵台注册或远程 MCP Client 直连)
# 端口来自 config/mcp/web-search.yaml → server.port
# --port 未显式传时使用配置中的 port,传则覆盖
run-sse:
@$(_ENV_KEYS) go run ./cmd/server --config ../../config/mcp/web-search.yaml
# 本地开发:指定端口覆盖配置中的 port
run-dev:
@$(_ENV_KEYS) go run ./cmd/server --config ../../config/mcp/web-search.yaml --port $(PORT)
# 构建二进制
build:
go build -o bin/bowen-server ./cmd/server
# 测试
test:
go test ./... -v -count=1
# 清理
clean:
rm -rf bin/ tmp/
# Lint
lint:
golangci-lint run ./...
+426
View File
@@ -0,0 +1,426 @@
# 博闻 Bowen · MCP Web Search
博闻是魂镜的 MCP 联网搜索服务,支持多搜索引擎、多 API 密钥智能轮询、两级缓存和双传输模式。
> 文化命名出自《礼记·中庸》「博闻强识而让」,与辰正、一炁等同属魂镜文化命名体系。工程路径保留 `mcp/web-search`。
---
## 架构概览
```mermaid
flowchart TD
subgraph L1["① 传输层"]
ST["stdio<br/>stdin/stdout JSON-RPC"]
SSE["HTTP SSE<br/>/sse + /messages"]
HTTP["简化 HTTP<br/>/tools/list + /tools/call"]
end
subgraph L2["② 核心处理器"]
H["Handler<br/>传输层共享"]
C["Cache Manager<br/>L1 LRU → L2 Redis"]
end
subgraph L3["③ 负载均衡"]
LB["Load Balancer"]
KP["KeyPool<br/>EWMA 动态权重"]
CB["Circuit Breaker<br/>closed / open / half-open"]
end
subgraph L4["④ 搜索提供商"]
T[Tavily]
B[Bing]
V[Brave]
end
subgraph L5["⑤ 可观测性"]
M["Prometheus Metrics<br/>7 指标"]
S["/status + /health"]
A["管理 API<br/>enable / disable / clear"]
end
ST --> H
SSE --> H
HTTP --> H
H --> C
H --> LB
LB --> KP
KP --> CB
LB --> T
LB --> B
LB --> V
H --> M
H --> S
H --> A
```
### 核心特性
| 特性 | 说明 |
| ------ | ------ |
| **多 Provider** | Tavily、Bing、Brave 三家搜索,统一响应格式 |
| **多 Key 轮询** | 每 Provider 可配多个 API Key,按 EWMA 动态权重(延迟×成功率)自动排序 |
| **熔断器** | 每个 Key 独立熔断(closed → open → half-open),失败 key 自动隔离并恢复探测 |
| **故障转移** | 一个 Key 失败自动换 Key,一个 Provider 耗尽自动切其他 Provider |
| **双传输** | HTTP SSE(远程 Client 直连)+ stdio(子进程集成),由 `--port` / `--stdio` CLI 标志决定 |
| **两级缓存** | L1 内存 LRU(毫秒级)+ L2 Redis(可选,跨实例共享),L2 命中自动升温到 L1 |
| **配置热重载** | fsnotify 监听 YAML 变更,动态调整 provider 启用/禁用、缓存 TTL、重试次数,无需重启 |
| **管理 API** | 运行时启用/禁用 provider、清空缓存、查询密钥状态 |
| **SearchMode** | `first`(首个成功返回,默认低延迟)+ `merge`(并发搜所有引擎,URL 去重排序) |
| **可观测性** | 暗色风格 HTML 监控面板(`:9090`)、Prometheus 指标(`:9090/metrics`)、结构化彩色日志 |
---
## 快速开始
### 前置条件
- Go 1.22+
- 至少一个搜索 API Key(Tavily / Bing / Brave)
- (可选)Redis 用于二级缓存
#### 环境变量快速配置
```bash
# 必填——搜索 API Key(逗号分隔支持多 key 轮询)
export TAVILY_API_KEYS="tvly-key1,tvly-key2"
export BRAVE_API_KEYS="brave-key1"
```
完整环境变量见项目根目录 `.env.example` 中的「博闻」段。
### 启动
```bash
# HTTP SSE 模式(供 MCP Client 远程连接,或供灵台注册为工具)
go run ./cmd/server --port 8080
# stdio 模式(供 Claude Desktop / Cursor 等 MCP Client 子进程集成)
go run ./cmd/server --stdio
# 指定 YAML 配置文件(可选,不传则自动探测 config/mcp/web-search.yaml)
go run ./cmd/server --config ../../config/mcp/web-search.yaml
```
### 验证
```bash
# 存活探针
curl http://localhost:8080/health
# → {"status":"ok"}
# 系统状态
curl http://localhost:8080/status
# → {"version":"2.0.0","uptime":"...","providers":[...],"cache":{...}}
# 工具列表(MCP tools/list)
curl -X POST http://localhost:8080/tools/list
# → {"tools":[{"name":"web_search",...}]}
# 搜索调用
curl -X POST http://localhost:8080/tools/call \
-H "Content-Type: application/json" \
-d '{"tool":"web_search","input":{"query":"hello world"}}'
# 监控面板(浏览器打开)
open http://localhost:9090
```
---
## 端口规划
| 端口 | 用途 | 说明 |
| ------ | ------ | ------ |
| **8080** | MCP 服务端口 | tools/list、tools/call、SSE、管理 API、/health、/status |
| **9090** | 监控端口 | HTML 仪表盘(`/`)、Prometheus 指标(`/metrics`)、存活探针(`/health`) |
## MCP 服务端点(端口 8080)
### MCP 协议端点
| 路径 | 方法 | 说明 |
| ------ | ------ | ------ |
| `POST /tools/list` | POST | 列出可用工具(web_search),含 inputSchema |
| `POST /tools/call` | POST | 调用 web_search 工具执行搜索 |
| `GET /sse` | GET | SSE 事件流(MCP 标准握手) |
| `POST /messages` | POST | MCP JSON-RPC 消息(与 /sse 配对使用) |
| `GET /` | GET | 端点指引(纯文本) |
### web_search 工具
**inputSchema**:
| 参数 | 类型 | 默认 | 说明 |
| ------ | ------ | ------ | ------ |
| `query` | string | (必填) | 搜索关键词 |
| `max_results` | integer | 5 | 最大返回结果数(1–20) |
| `mode` | string | `"first"` | 搜索模式:`first`(首个成功返回) / `merge`(并发搜所有引擎后合并去重) |
| `providers` | string[] | 全部 | 优先使用的搜索引擎 |
**响应**(CallToolResult):
```json
{
"content": [{
"type": "json",
"text": "{\"provider\":\"tavily\",\"query\":\"...\",\"results\":[...],\"total_count\":5,\"cached\":false}"
}]
}
```
### 管理端点
| 路径 | 方法 | 说明 |
| ------ | ------ | ------ |
| `GET /health` | GET | 存活探针 → `{"status":"ok"}` |
| `GET /status` | GET | 系统状态:provider 列表、活跃密钥数、缓存命中率、uptime |
| `GET /admin/keys` | GET | 各 provider 启用状态 |
| `PUT /admin/provider/{name}/enable` | PUT | 启用指定 provider |
| `PUT /admin/provider/{name}/disable` | PUT | 禁用指定 provider |
| `POST /admin/cache/clear` | POST | 清空缓存 |
## 监控面板(端口 9090)
### HTML 仪表盘 `GET /`
暗色风格的可视化面板,展示:
- **四张概览卡片**:Provider 数、活跃密钥数、缓存命中率(绿/黄/红三档)、缓存条目总数
- **Provider 状态表格**:🟢/🔴 状态指示灯、名称、启用/禁用开关、活跃密钥计数
- **缓存统计表格**:L1(内存 LRU)命中、L2(Redis)命中、未命中、综合命中率
- **自动刷新**:每 30 秒自动刷新
- **导航链接**:一键跳转 /status、/admin/keys、/health
### Prometheus 原始指标 `GET /metrics`
| 指标名 | 类型 | 标签 | 说明 |
| -------- | ------ | ------ | ------ |
| `bowen_search_total` | Counter | `provider` | 搜索请求总数 |
| `bowen_search_duration_seconds` | Histogram | — | 搜索耗时分布 |
| `bowen_search_errors_total` | Counter | `reason` | 搜索错误数 |
| `bowen_cache_hits_total` | Counter | — | 缓存命中总数 |
| `bowen_cache_hits_l2_total` | Counter | — | Redis 二级缓存命中数 |
| `bowen_cache_misses_total` | Counter | — | 缓存未命中数 |
| `bowen_active_keys` | Gauge | `provider` | 各 provider 活跃密钥数 |
所有监控端点也可在端口 9090 访问,与 MCP 服务端口独立。
---
## 配置文件
非敏感配置放在项目根目录 `config/mcp/web-search.yaml`(从 `config/mcp/web-search.example.yaml` 复制)。
不传 `--config` 时自动按项目惯例探测:开发模式探测 `../../config/mcp/web-search.yaml`,容器内探测 `/etc/bowen/config.yaml`。
```yaml
server:
port: 8080 # SSE 模式监听端口
read_timeout: 30s # HTTP 读取超时
write_timeout: 30s # HTTP 写入超时
providers:
tavily:
enabled: true # 是否启用(热重载:支持动态变更)
timeout: 10s
max_results: 10
health_check:
interval: 60s
failure_threshold: 3
bing:
enabled: true
timeout: 10s
max_results: 20
brave:
enabled: true
timeout: 10s
max_results: 10
balancer:
strategy: weighted_round_robin
retry_count: 2 # 热重载:支持动态变更
circuit_breaker:
enabled: true
failure_threshold: 5
cooldown: 30s
cache:
enabled: true
ttl: 5m # 热重载:支持动态变更
max_items: 1000
redis:
enabled: false # 需要自行部署 Redis
addr: 127.0.0.1:6379
password: ""
db: 0
metrics:
enabled: true
port: 9090
path: /metrics
```
**配置优先级**:环境变量 > YAML 文件 > 内置默认值
**热重载**:指定 `--config` 后,修改 YAML 文件自动生效(500ms 防抖),支持动态变更:
| 配置项 | 变更后自动生效 |
| -------- | --------------- |
| `providers.{name}.enabled` | 运行时启用/禁用搜索源 |
| `cache.ttl` | 更新新缓存条目的过期时间 |
| `balancer.retry_count` | 更新重试次数 |
---
## 传输模式
博闻支持两种 MCP 传输模式,**不在配置中写死**,由启动 CLI 参数决定:
| 模式 | 启动方式 | 适用场景 |
| ------ | --------- | --------- |
| **SSE** | `--port 8080` | 灵台后端注册(简化 HTTP 端点 `/tools/list` + `/tools/call`)、远程 MCP Client 直连(`/sse` + `/messages`) |
| **Stdio** | `--stdio` | Claude Desktop / Cursor 等 MCP Client 以子进程方式启动,通过 stdin/stdout 走 JSON-RPC 通信 |
两种模式下核心业务逻辑完全共享(`internal/mcp/handler.go`),只有传输层不同。
---
## 环境变量
所有敏感信息通过环境变量注入(参见 `.env.example` 博闻段):
| 变量 | 必填 | 说明 |
| ------ | ------ | ------ |
| `TAVILY_API_KEYS` | 是 | Tavily 搜索 Key,多 key 逗号分隔 |
| `BRAVE_API_KEYS` | 是 | Brave Search Key |
| `MCP_WEB_SEARCH_REDIS_ENABLED` | 否 | Redis 二级缓存开关 |
| `MCP_WEB_SEARCH_REDIS_ADDR` | 否 | Redis 地址(默认 `127.0.0.1:6379`) |
| `MCP_WEB_SEARCH_REDIS_PASSWORD` | 否 | Redis 密码 |
| `MCP_WEB_SEARCH_REDIS_DB` | 否 | Redis 数据库编号(默认 0) |
---
## 开发
```bash
# 运行所有测试
make test
# 编译
make build # → bin/bowen-server
# 本地 SSE 模式
make run-sse
# 本地 stdio 模式
make run-stdio
# 代码检查
make lint
```
开发环境日志自动使用彩色输出(ANSI 颜色代码):
- `INFO` → 绿色
- `WARN` → 黄色
- `ERROR` → 红色
- `DEBUG` → 灰色
生产环境(`ENV=production`)自动切换为 JSON 格式,便于日志采集系统采集。
### 项目结构
```text
mcp/web-search/
├── cmd/server/main.go # 入口(--port / --stdio)
├── internal/
│ ├── config/
│ │ ├── config.go # YAML + 环境变量配置加载
│ │ ├── yaml.go # YAML 解析
│ │ └── reloader.go # 配置热重载(fsnotify)
│ ├── logger/logger.go # slog 双模式(dev 彩色 / prod JSON)
│ ├── models/types.go # 核心数据类型
│ ├── providers/ # 搜索提供商适配
│ │ ├── provider.go # SearchProvider 接口
│ │ ├── tavily.go / bing.go / brave.go
│ │ └── registry.go # 注册中心
│ ├── balancer/ # 负载均衡
│ │ ├── keypool.go # 密钥池(EWMA)
│ │ ├── circuitbreaker.go # 熔断器
│ │ └── balancer.go # LoadBalancer
│ ├── cache/cache.go # 两级缓存
│ ├── mcp/ # MCP 协议实现
│ │ ├── handler.go # 核心处理逻辑
│ │ ├── stdio.go # stdin/stdout JSON-RPC
│ │ ├── sse.go # HTTP SSE 传输
│ │ ├── http.go # 简化 HTTP 端点
│ │ ├── status.go # /health + /status
│ │ ├── admin.go # 管理 API
│ │ └── monitor.go # HTML 监控面板
│ └── metrics/metrics.go # Prometheus 指标
├── Makefile
├── go.mod / go.sum
```
---
## Docker
```bash
# 构建镜像
docker build -t bowen-web-search -f ../../docker/Dockerfile.web-search .
# 运行
docker run \
-e TAVILY_API_KEYS=tvly-xxx \
-e BRAVE_API_KEYS=brave-xxx \
-p 8080:8080 -p 9090:9090 \
bowen-web-search
```
Docker 镜像可通过魂镜的 `docker/docker-compose.app.yml` 编排启动。
---
## 集成到灵台
后端通过 MCP 注册 API 自动接入博闻:
```bash
# 注册博闻为 MCP Server
curl -X POST http://localhost:8000/api/v1/tools/mcp-servers \
-H "Content-Type: application/json" \
-d '{
"name": "bowen-web-search",
"endpoint_url": "http://localhost:8080",
"status": "enabled"
}'
```
注册后灵台会拉取工具清单,Agent 即可在对话中直接调用 `web_search` 工具。
---
## 搜索模式对比
| 模式 | 延迟 | 召回 | 使用场景 |
| ------ | ------ | ------ | --------- |
| `first`(默认) | 低(首个成功即返回) | 单 provider | 快速查天气、新闻、百科等简单查询 |
| `merge` | 中(等最慢 provider) | 高(多源合并去重) | 研究、对比多个来源、需要全面 coverage |
---
## 技术栈
| 类别 | 选型 |
| ------ | ------ |
| 语言 | Go 1.22+ |
| Web | 标准库 `net/http` |
| 日志 | 标准库 `log/slog` + 自定义彩色 Handler |
| 配置 | `gopkg.in/yaml.v3` + 环境变量 |
| 内存缓存 | `hashicorp/golang-lru/v2` |
| Redis | `redis/go-redis/v9`(可选) |
| 监控 | `prometheus/client_golang` |
| 热重载 | `fsnotify` |
| 测试 | 标准库 `testing` |
| 容器化 | Docker 多阶段构建 |
+160
View File
@@ -0,0 +1,160 @@
package main
import (
"flag"
"fmt"
"log/slog"
"net/http"
"os"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/chenxin/hunjing/mcp/web-search/internal/balancer"
"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/logger"
"github.com/chenxin/hunjing/mcp/web-search/internal/mcp"
"github.com/chenxin/hunjing/mcp/web-search/internal/metrics"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
"github.com/chenxin/hunjing/mcp/web-search/internal/providers"
)
func main() {
configPath := flag.String("config", "", "YAML 配置文件路径(不传则按项目惯例自动探测)")
port := flag.Int("port", 0, "HTTP SSE 模式监听端口(缺省时使用配置文件的 port)")
stdioMode := flag.Bool("stdio", false, "以 stdio 模式运行(通过 stdin/stdout 走 JSON-RPC,用于子进程集成)")
flag.Parse()
// 未指定 --config 时按项目惯例自动探测
if *configPath == "" {
for _, p := range []string{
"../../config/mcp/web-search.yaml", // 开发模式(mcp/web-search/ 下 go run)
"config/mcp/web-search.yaml", // 项目根目录下 go run mcp/... 或 make
"/etc/bowen/config.yaml", // Docker 容器内默认路径
} {
if _, err := os.Stat(p); err == nil {
*configPath = p
break
}
}
}
// 加载配置
cfg, err := config.Load(*configPath)
if err != nil {
fmt.Fprintf(os.Stderr, "配置加载失败: %v\n", err)
os.Exit(1)
}
// 设置日志
logCfg := logger.Config{
Level: slog.LevelInfo,
AddSource: false,
JSON: os.Getenv("ENV") == "production",
}
slog.SetDefault(logger.New(logCfg))
// 端口覆盖
if *port > 0 {
cfg.Server.Port = *port
}
// 初始化 Provider 注册中心
reg, err := providers.SetupFromConfig(cfg)
if err != nil {
slog.Error("provider 初始化失败", "error", err)
os.Exit(1)
}
// 初始化负载均衡器
lb := balancer.NewLoadBalancer(&cfg.Balancer)
// 为每个启用的 provider 及其密钥注册到负载均衡器
type providerEntry struct {
kind models.ProviderKind
keys []string
}
entries := []providerEntry{
{models.ProviderTavily, cfg.Providers.Tavily.APIKeys},
{models.ProviderBrave, cfg.Providers.Brave.APIKeys},
}
var providerKinds []models.ProviderKind
for _, e := range entries {
if len(e.keys) == 0 {
continue
}
prov, ok := reg.Get(e.kind)
if !ok {
continue
}
lb.AddProvider(e.kind, e.keys, 1, prov)
providerKinds = append(providerKinds, e.kind)
metrics.ActiveKeys.WithLabelValues(string(e.kind)).Set(float64(len(e.keys)))
slog.Info("provider 已注册",
"provider", e.kind,
"active_keys", len(e.keys),
)
}
// 初始化缓存
cacheMgr, err := cache.New(cfg.Cache)
if err != nil {
slog.Error("缓存初始化失败", "error", err)
os.Exit(1)
}
// 创建 handler
handler := mcp.NewHandler(lb, cacheMgr, providerKinds)
// 启动后台健康检查(使用第一个启用的 provider 的配置参数)
if len(providerKinds) > 0 && cfg.Providers.Tavily.HealthCheck.Interval > 0 {
checker := balancer.NewHealthChecker(lb, cfg.Providers.Tavily.HealthCheck)
checker.Start()
defer checker.Stop()
}
// 配置热重载(仅在指定了配置文件时启用)
if *configPath != "" {
onChange := func(old, new *config.Config) error {
handler.ReloadConfig(old, new)
return nil
}
reloader, err := config.NewReloader(*configPath, cfg, onChange)
if err != nil {
slog.Warn("配置热重载初始化失败(不影响启动)", "error", err)
} else {
defer reloader.Stop()
}
}
// 启动 Metrics 服务
if cfg.Metrics.Enabled {
go func() {
metricsMux := http.NewServeMux()
metricsMux.Handle(cfg.Metrics.Path, promhttp.Handler())
metricsMux.HandleFunc("/", handler.HandleMonitor)
metricsMux.HandleFunc("/health", handler.HandleHealth)
metricsMuxWithLog := mcp.LoggingMiddleware(metricsMux)
metricsAddr := fmt.Sprintf(":%d", cfg.Metrics.Port)
slog.Info("metrics 服务启动", "addr", metricsAddr)
if err := http.ListenAndServe(metricsAddr, metricsMuxWithLog); err != nil {
slog.Error("metrics 服务异常", "error", err)
}
}()
}
// 根据 CLI 标志选择传输模式:
// --stdio :stdin/stdout JSON-RPC(用于 Claude Desktop / Cursor 等 MCP Client 子进程集成)
// 默认 :Streamable HTTP 模式(用于灵台注册 / 远程 MCP Client 直连)
if *stdioMode {
handler.ServeStdio()
} else {
addr := fmt.Sprintf(":%d", cfg.Server.Port)
slog.Info("博闻 HTTP 服务启动", "addr", addr)
if err := handler.ServeHTTP(addr); err != nil {
slog.Error("博闻异常退出", "error", err)
os.Exit(1)
}
}
}
+25
View File
@@ -0,0 +1,25 @@
module github.com/chenxin/hunjing/mcp/web-search
go 1.25.5
require (
github.com/fsnotify/fsnotify v1.7.0
github.com/hashicorp/golang-lru/v2 v2.0.7
github.com/prometheus/client_golang v1.23.2
github.com/redis/go-redis/v9 v9.21.0
gopkg.in/yaml.v3 v3.0.1
)
require (
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/kr/text v0.2.0 // indirect
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
github.com/prometheus/client_model v0.6.2 // indirect
github.com/prometheus/common v0.66.1 // indirect
github.com/prometheus/procfs v0.16.1 // indirect
go.uber.org/atomic v1.11.0 // indirect
go.yaml.in/yaml/v2 v2.4.2 // indirect
golang.org/x/sys v0.35.0 // indirect
google.golang.org/protobuf v1.36.8 // indirect
)
+62
View File
@@ -0,0 +1,62 @@
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o=
github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg=
github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE=
github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9ZoGs=
github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA=
github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg=
github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is=
github.com/redis/go-redis/v9 v9.21.0 h1:FPBE4hhbAke+TLmcY3WkpbDffJEomdqPn3HYiqAtL9E=
github.com/redis/go-redis/v9 v9.21.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA=
github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs=
github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s=
go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI=
go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU=
golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
@@ -0,0 +1,264 @@
package balancer
import (
"context"
"fmt"
"sync"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// KeySearcher 按密钥执行搜索的接口。
// 与 SearchProvider 解耦,方便测试时注入 mock。
type KeySearcher interface {
Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error)
HealthCheck(ctx context.Context) error
Name() models.ProviderKind
}
// LoadBalancer 负载均衡器(goroutine-safe)。
type LoadBalancer struct {
pools map[models.ProviderKind]*KeyPool
searchers map[models.ProviderKind]KeySearcher
disabled map[models.ProviderKind]bool // 运行时动态禁用
retryCount int
mu sync.RWMutex
}
func NewLoadBalancer(cfg *config.BalancerConfig) *LoadBalancer {
return &LoadBalancer{
pools: make(map[models.ProviderKind]*KeyPool),
searchers: make(map[models.ProviderKind]KeySearcher),
disabled: make(map[models.ProviderKind]bool),
retryCount: cfg.RetryCount,
}
}
// AddProvider 注册 provider 及其密钥池。
func (lb *LoadBalancer) AddProvider(
provider models.ProviderKind,
keys []string,
weight int,
searcher KeySearcher,
) {
lb.mu.Lock()
defer lb.mu.Unlock()
cbThreshold := 5
cbCooldown := 30 * time.Second
lb.pools[provider] = NewKeyPool(provider, keys, weight, cbThreshold, cbCooldown)
lb.searchers[provider] = searcher
// 添加时默认启用
delete(lb.disabled, provider)
}
// RemoveProvider 移除 provider 及其密钥池。
func (lb *LoadBalancer) RemoveProvider(provider models.ProviderKind) {
lb.mu.Lock()
defer lb.mu.Unlock()
delete(lb.pools, provider)
delete(lb.searchers, provider)
delete(lb.disabled, provider)
}
// EnableProvider 启用 provider。
func (lb *LoadBalancer) EnableProvider(provider models.ProviderKind) {
lb.mu.Lock()
defer lb.mu.Unlock()
delete(lb.disabled, provider)
}
// DisableProvider 禁用 provider(搜索时跳过,不删除配置)。
func (lb *LoadBalancer) DisableProvider(provider models.ProviderKind) {
lb.mu.Lock()
defer lb.mu.Unlock()
lb.disabled[provider] = true
}
// IsProviderEnabled 检查 provider 是否启用。
func (lb *LoadBalancer) IsProviderEnabled(provider models.ProviderKind) bool {
lb.mu.RLock()
defer lb.mu.RUnlock()
return !lb.disabled[provider]
}
// SetRetryCount 更新重试次数。
func (lb *LoadBalancer) SetRetryCount(n int) {
lb.mu.Lock()
defer lb.mu.Unlock()
lb.retryCount = n
}
// HasProvider 检查 provider 是否已注册。
func (lb *LoadBalancer) HasProvider(provider models.ProviderKind) bool {
lb.mu.RLock()
defer lb.mu.RUnlock()
_, ok := lb.pools[provider]
return ok
}
// Search 执行搜索(自动选密钥、跨 provider 重试)。
func (lb *LoadBalancer) Search(ctx context.Context, req *models.SearchRequest) (*models.SearchResult, error) {
providers := lb.resolveProviders(req.Providers)
var lastErr error
for attempt := 0; attempt <= lb.getRetryCount(); attempt++ {
for _, prov := range providers {
pool, ok := lb.pools[prov]
if !ok {
continue
}
ks := pool.Select()
if ks == nil {
continue
}
searcher := lb.searchers[prov]
if searcher == nil {
continue
}
start := time.Now()
result, err := searcher.Search(ctx, ks.Key, req)
duration := time.Since(start)
if err != nil {
ks.RecordFailure()
lastErr = err
continue
}
ks.RecordSuccess(duration)
result.Provider = prov
return result, nil
}
}
if lastErr != nil {
return nil, fmt.Errorf("balancer: all providers exhausted: %w", lastErr)
}
return nil, fmt.Errorf("balancer: no available provider or key for query: %s", req.Query)
}
// SearchAll 并发搜索所有启用 provider(用于 Merge 模式)。
func (lb *LoadBalancer) SearchAll(ctx context.Context, req *models.SearchRequest) []*models.SearchResult {
providers := lb.resolveProviders(nil)
type result struct {
prov models.ProviderKind
res *models.SearchResult
err error
}
ch := make(chan result, len(providers))
for _, prov := range providers {
p := prov
go func() {
pool, ok := lb.pools[p]
if !ok {
ch <- result{prov: p, err: fmt.Errorf("no pool")}
return
}
ks := pool.Select()
if ks == nil {
ch <- result{prov: p, err: fmt.Errorf("no available key")}
return
}
searcher := lb.searchers[p]
if searcher == nil {
ch <- result{prov: p, err: fmt.Errorf("no searcher")}
return
}
start := time.Now()
res, err := searcher.Search(ctx, ks.Key, req)
duration := time.Since(start)
if err != nil {
ks.RecordFailure()
ch <- result{prov: p, err: err}
return
}
ks.RecordSuccess(duration)
res.Provider = p
ch <- result{prov: p, res: res}
}()
}
var results []*models.SearchResult
for range providers {
r := <-ch
if r.res != nil {
results = append(results, r.res)
}
}
return results
}
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
func (lb *LoadBalancer) ActiveKeyCount() map[models.ProviderKind]int {
lb.mu.RLock()
defer lb.mu.RUnlock()
out := make(map[models.ProviderKind]int)
for p, pool := range lb.pools {
if lb.disabled[p] {
continue
}
out[p] = pool.ActiveCount()
}
return out
}
// Providers 返回所有已注册的 provider 名称(无论启用与否)。
func (lb *LoadBalancer) Providers() []models.ProviderKind {
lb.mu.RLock()
defer lb.mu.RUnlock()
out := make([]models.ProviderKind, 0, len(lb.pools))
for p := range lb.pools {
out = append(out, p)
}
return out
}
// HealthCheck 对指定 provider 执行健康检查(委托给 KeySearcher)。
func (lb *LoadBalancer) HealthCheck(ctx context.Context, provider models.ProviderKind) error {
lb.mu.RLock()
searcher, ok := lb.searchers[provider]
lb.mu.RUnlock()
if !ok {
return fmt.Errorf("healthcheck: provider not found: %s", provider)
}
return searcher.HealthCheck(ctx)
}
func (lb *LoadBalancer) getRetryCount() int {
lb.mu.RLock()
defer lb.mu.RUnlock()
return lb.retryCount
}
// resolveProviders 按用户偏好和可用性确定 provider 顺序(跳过已禁用的)。
func (lb *LoadBalancer) resolveProviders(preferred []string) []models.ProviderKind {
lb.mu.RLock()
defer lb.mu.RUnlock()
if len(preferred) == 0 {
out := make([]models.ProviderKind, 0, len(lb.pools))
for p := range lb.pools {
if !lb.disabled[p] {
out = append(out, p)
}
}
return out
}
out := make([]models.ProviderKind, 0, len(preferred))
seen := make(map[models.ProviderKind]bool)
for _, name := range preferred {
p := models.ProviderKind(name)
if _, ok := lb.pools[p]; ok && !seen[p] && !lb.disabled[p] {
out = append(out, p)
seen[p] = true
}
}
return out
}
@@ -0,0 +1,90 @@
package balancer
import (
"context"
"errors"
"testing"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
type mockSearcher struct {
name models.ProviderKind
failOn int // 第几次调用失败(0 = 永不)
calls int
}
func (m *mockSearcher) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
m.calls++
if m.failOn > 0 && m.calls >= m.failOn {
return nil, errors.New("mock failure")
}
return &models.SearchResult{Provider: m.name, Query: req.Query, Results: []models.ResultItem{{Title: "mock"}}}, nil
}
func (m *mockSearcher) HealthCheck(ctx context.Context) error { return nil }
func (m *mockSearcher) Name() models.ProviderKind { return m.name }
func TestLoadBalancer_Search_Basic(t *testing.T) {
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if result.Provider != models.ProviderTavily {
t.Errorf("expected tavily, got %s", result.Provider)
}
}
func TestLoadBalancer_Search_FallbackProvider(t *testing.T) {
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if string(result.Provider) != "brave" {
t.Errorf("expected fallback to brave, got %s", result.Provider)
}
}
func TestLoadBalancer_Search_AllExhausted(t *testing.T) {
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
_, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
if err == nil {
t.Fatal("expected error when all keys exhausted")
}
}
func TestLoadBalancer_Search_PreferredProviders(t *testing.T) {
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
result, err := lb.Search(context.Background(), &models.SearchRequest{
Query: "test",
Providers: []string{"brave"},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if result.Provider != models.ProviderBrave {
t.Errorf("expected bing (preferred), got %s", result.Provider)
}
}
func TestLoadBalancer_ActiveKeyCount(t *testing.T) {
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
lb.AddProvider(models.ProviderTavily, []string{"key1", "key2"}, 1, &mockSearcher{name: models.ProviderTavily})
counts := lb.ActiveKeyCount()
if counts[models.ProviderTavily] != 2 {
t.Errorf("expected 2 active keys, got %d", counts[models.ProviderTavily])
}
}
+117
View File
@@ -0,0 +1,117 @@
package balancer
import (
"context"
"log/slog"
"sync"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// HealthChecker 定期对已注册 provider 执行健康检查。
// 连续失败超过阈值时自动禁用 provider,恢复成功后自动启用。
type HealthChecker struct {
lb *LoadBalancer
interval time.Duration
timeout time.Duration
threshold int
mu sync.Mutex
failures map[models.ProviderKind]int // 连续失败计数
done chan struct{}
}
// NewHealthChecker 创建健康检查器。cfg 为对应 provider 的健康检查配置。
func NewHealthChecker(lb *LoadBalancer, cfg config.HealthCheckConf) *HealthChecker {
return &HealthChecker{
lb: lb,
interval: cfg.Interval,
timeout: cfg.Timeout,
threshold: cfg.FailureThreshold,
failures: make(map[models.ProviderKind]int),
done: make(chan struct{}),
}
}
// Start 在后台 goroutine 中启动健康检查循环。
// 每隔 interval 时间对所有已注册 provider 执行一次健康检查。
func (hc *HealthChecker) Start() {
if hc.interval <= 0 {
return
}
slog.Info("健康检查已启动", "interval", hc.interval, "timeout", hc.timeout, "failure_threshold", hc.threshold)
go hc.loop()
}
// Stop 停止健康检查循环。
func (hc *HealthChecker) Stop() {
close(hc.done)
}
func (hc *HealthChecker) loop() {
ticker := time.NewTicker(hc.interval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
hc.checkAll()
case <-hc.done:
return
}
}
}
func (hc *HealthChecker) checkAll() {
// 获取当前所有已注册 provider 的名称
// 我们需要从 LoadBalancer 获取提供者列表
// 由于 LoadBalancer 没有直接暴露 provider 列表的 API,我们遍历 ActiveKeyCount
// 来获取所有已知的 provider
for _, prov := range hc.getAllProviders() {
hc.checkProvider(prov)
}
}
func (hc *HealthChecker) checkProvider(prov models.ProviderKind) {
ctx, cancel := context.WithTimeout(context.Background(), hc.timeout)
defer cancel()
err := hc.lb.HealthCheck(ctx, prov)
hc.mu.Lock()
defer hc.mu.Unlock()
if err != nil {
hc.failures[prov]++
fails := hc.failures[prov]
slog.Warn("健康检查失败",
"provider", prov,
"consecutive_failures", fails,
"error", err,
)
if fails >= hc.threshold && hc.lb.IsProviderEnabled(prov) {
hc.lb.DisableProvider(prov)
slog.Warn("健康检查: provider 已自动禁用",
"provider", prov,
"failures", fails,
)
}
} else {
// 恢复成功
if hc.failures[prov] > 0 {
slog.Info("健康检查恢复", "provider", prov)
}
hc.failures[prov] = 0
if !hc.lb.IsProviderEnabled(prov) {
hc.lb.EnableProvider(prov)
slog.Info("健康检查: provider 已自动启用", "provider", prov)
}
}
}
// getAllProviders 返回当前所有已注册的 provider 名称(包括已禁用的)。
func (hc *HealthChecker) getAllProviders() []models.ProviderKind {
return hc.lb.Providers()
}
@@ -0,0 +1,103 @@
// Package balancer 实现密钥池管理和负载均衡。
package balancer
import (
"sync"
"time"
)
// State 熔断器状态。
type State int
const (
StateClosed State = iota // 正常
StateOpen // 熔断
StateHalfOpen // 半开(探测)
)
func (s State) String() string {
switch s {
case StateClosed:
return "closed"
case StateOpen:
return "open"
case StateHalfOpen:
return "half-open"
default:
return "unknown"
}
}
// CircuitBreaker 熔断器。
type CircuitBreaker struct {
mu sync.Mutex
state State
failures int
threshold int
cooldown time.Duration
lastFailureTime time.Time
}
func NewCircuitBreaker(threshold int, cooldown time.Duration) *CircuitBreaker {
return &CircuitBreaker{
state: StateClosed,
threshold: threshold,
cooldown: cooldown,
}
}
// Allow 判断是否允许请求通过。
func (cb *CircuitBreaker) Allow() bool {
cb.mu.Lock()
defer cb.mu.Unlock()
switch cb.state {
case StateClosed:
return true
case StateOpen:
if time.Since(cb.lastFailureTime) > cb.cooldown {
cb.state = StateHalfOpen
return true
}
return false
case StateHalfOpen:
return true
default:
return true
}
}
// Success 记录成功。
func (cb *CircuitBreaker) Success() {
cb.mu.Lock()
defer cb.mu.Unlock()
cb.failures = 0
if cb.state == StateHalfOpen {
cb.state = StateClosed
}
}
// Failure 记录失败。
func (cb *CircuitBreaker) Failure() {
cb.mu.Lock()
defer cb.mu.Unlock()
cb.failures++
cb.lastFailureTime = time.Now()
if cb.state == StateHalfOpen || cb.failures >= cb.threshold {
cb.state = StateOpen
}
}
// State 返回当前状态(goroutine-safe 快照)。
func (cb *CircuitBreaker) State() State {
cb.mu.Lock()
defer cb.mu.Unlock()
return cb.state
}
// Failures 返回连续失败次数。
func (cb *CircuitBreaker) Failures() int {
cb.mu.Lock()
defer cb.mu.Unlock()
return cb.failures
}
+169
View File
@@ -0,0 +1,169 @@
package balancer
import (
"math"
"sync"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// KeyStatus 密钥状态。
type KeyStatus int
const (
KeyActive KeyStatus = iota // 正常
KeyCooldown // 冷却中(熔断后)
KeyDead // 死亡(多次失败不再恢复)
)
// KeyState 单个密钥状态。
type KeyState struct {
Provider models.ProviderKind
Key string
Weight int
mu sync.Mutex
status KeyStatus
latencyEWMA time.Duration // 指数加权滑动平均延迟
successCount int64
failCount int64
cb *CircuitBreaker
}
func newKeyState(provider models.ProviderKind, key string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyState {
return &KeyState{
Provider: provider,
Key: key,
Weight: weight,
status: KeyActive,
cb: NewCircuitBreaker(cbThreshold, cbCooldown),
}
}
// Allow 检查是否允许使用此密钥。
func (ks *KeyState) Allow() bool {
ks.mu.Lock()
defer ks.mu.Unlock()
if ks.status != KeyActive {
return false
}
return ks.cb.Allow()
}
// RecordSuccess 记录成功调用,更新 EWMA 延迟。
func (ks *KeyState) RecordSuccess(latency time.Duration) {
ks.mu.Lock()
defer ks.mu.Unlock()
ks.successCount++
ks.cb.Success()
// EWMA 更新:α = 0.3
alpha := 0.3
if ks.latencyEWMA == 0 {
ks.latencyEWMA = latency
} else {
ks.latencyEWMA = time.Duration(
(1-alpha)*float64(ks.latencyEWMA) + alpha*float64(latency),
)
}
ks.status = KeyActive
}
// RecordFailure 记录失败调用。
func (ks *KeyState) RecordFailure() {
ks.mu.Lock()
defer ks.mu.Unlock()
ks.failCount++
ks.cb.Failure()
if ks.cb.State() == StateOpen {
ks.status = KeyCooldown
}
}
// CombinedScore 计算综合分数(越高优先选用)。
// 综合 = 基础权重 × 延迟惩罚 × 可靠性因子
func (ks *KeyState) CombinedScore() float64 {
ks.mu.Lock()
defer ks.mu.Unlock()
if ks.status != KeyActive {
return 0
}
if ks.cb.State() != StateClosed {
return 0
}
base := float64(ks.Weight)
latencyPenalty := 1.0
if ks.latencyEWMA > 0 {
latencyPenalty = 1.0 / (1.0 + ks.latencyEWMA.Seconds())
}
reliability := 1.0
total := ks.successCount + ks.failCount
if total > 0 && ks.failCount > 0 {
reliability = 1.0 - float64(ks.failCount)/float64(total)
}
return base * latencyPenalty * math.Max(0.1, reliability)
}
// KeyPool 管理一个 provider 下的多密钥。
type KeyPool struct {
provider models.ProviderKind
keys []*KeyState
mu sync.RWMutex
cbThreshold int
cbCooldown time.Duration
}
func NewKeyPool(provider models.ProviderKind, keys []string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyPool {
kp := &KeyPool{
provider: provider,
keys: make([]*KeyState, 0, len(keys)),
cbThreshold: cbThreshold,
cbCooldown: cbCooldown,
}
for _, k := range keys {
kp.keys = append(kp.keys, newKeyState(provider, k, weight, cbThreshold, cbCooldown))
}
return kp
}
// Select 按综合分数选最佳密钥。
func (kp *KeyPool) Select() *KeyState {
kp.mu.RLock()
defer kp.mu.RUnlock()
var best *KeyState
var bestScore float64
for _, ks := range kp.keys {
if !ks.Allow() {
continue
}
score := ks.CombinedScore()
if score > bestScore {
bestScore = score
best = ks
}
}
return best
}
// ActiveCount 返回活跃密钥数。
func (kp *KeyPool) ActiveCount() int {
kp.mu.RLock()
defer kp.mu.RUnlock()
count := 0
for _, ks := range kp.keys {
if ks.status == KeyActive {
count++
}
}
return count
}
// TotalCount 返回密钥总数。
func (kp *KeyPool) TotalCount() int {
kp.mu.RLock()
defer kp.mu.RUnlock()
return len(kp.keys)
}
@@ -0,0 +1,123 @@
package balancer
import (
"testing"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
func TestKeyPool_Select_Basic(t *testing.T) {
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 5, 30*time.Second)
ks := kp.Select()
if ks == nil {
t.Fatal("expected to select a key")
}
if ks.Key != "key1" && ks.Key != "key2" {
t.Errorf("unexpected key: %s", ks.Key)
}
}
func TestKeyPool_Select_SkipsDead(t *testing.T) {
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 2, 1*time.Hour)
ks1 := kp.Select()
// Fail twice to trigger circuit breaker
ks1.RecordFailure()
ks1.RecordFailure()
selected := kp.Select()
if selected == nil {
t.Fatal("expected fallback to second key")
}
if selected.Key == ks1.Key {
t.Error("should not select the failed key")
}
}
func TestKeyState_CombinedScore(t *testing.T) {
ks := newKeyState(models.ProviderTavily, "test-key", 2, 5, 30*time.Second)
score := ks.CombinedScore()
if score <= 0 {
t.Errorf("expected positive score, got %f", score)
}
// After recording a slow latency, score should drop
ks.RecordSuccess(2 * time.Second)
slowScore := ks.CombinedScore()
if slowScore >= score {
t.Error("expected slow latency to reduce score")
}
}
func TestCircuitBreaker_OpenAndRecover(t *testing.T) {
cb := NewCircuitBreaker(3, 50*time.Millisecond)
// Should be closed initially
if !cb.Allow() {
t.Error("expected allow initially")
}
// Fail 3 times to open
cb.Failure()
cb.Failure()
cb.Failure()
if cb.Allow() {
t.Error("expected circuit open after 3 failures")
}
// Wait for cooldown
time.Sleep(60 * time.Millisecond)
if !cb.Allow() {
t.Error("expected allow after cooldown")
}
// Half-open: success should close
cb.Success()
if cb.State() != StateClosed {
t.Error("expected circuit closed after half-open success")
}
}
func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
cb := NewCircuitBreaker(3, 50*time.Millisecond)
cb.Failure()
cb.Failure()
cb.Failure()
if cb.Allow() {
t.Error("expected open")
}
time.Sleep(60 * time.Millisecond)
if !cb.Allow() {
t.Error("expected half-open allow")
}
// Half-open: failure should re-open immediately
cb.Failure()
if cb.State() != StateOpen {
t.Error("expected open after half-open failure")
}
}
func TestKeyPool_ActiveCount(t *testing.T) {
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2", "key3"}, 1, 3, 30*time.Second)
if kp.TotalCount() != 3 {
t.Errorf("expected 3 total, got %d", kp.TotalCount())
}
if kp.ActiveCount() != 3 {
t.Errorf("expected 3 active, got %d", kp.ActiveCount())
}
}
func TestKeyPool_AllKeysDead(t *testing.T) {
kp := NewKeyPool(models.ProviderTavily, []string{"key1"}, 1, 1, 1*time.Hour)
ks := kp.Select()
ks.RecordFailure()
selected := kp.Select()
if selected != nil {
t.Error("expected nil when all keys dead")
}
}
+222
View File
@@ -0,0 +1,222 @@
// Package cache 提供搜索缓存(两级:内存 LRU + 可选 Redis)。
package cache
import (
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"sync"
"time"
lru "github.com/hashicorp/golang-lru/v2"
"github.com/redis/go-redis/v9"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// entry 缓存条目。
type entry struct {
result *models.SearchResult
expiresAt time.Time
}
// Manager 两级缓存管理器。
// L1: 内存 LRU(本地毫秒级,单进程共享)
// L2: Redis(跨实例共享,可按开关关闭)
type Manager struct {
enabled bool
ttl time.Duration
lru *lru.Cache[string, *entry]
rdb *redis.Client
redisDB int
mu sync.RWMutex
l1Hits int64
l2Hits int64
misses int64
}
// New 创建缓存管理器。若 Redis 启用则初始化 L2 连接。
func New(cfg config.CacheConfig) (*Manager, error) {
if !cfg.Enabled {
return &Manager{enabled: false}, nil
}
c, err := lru.New[string, *entry](cfg.MaxItems)
if err != nil {
return nil, fmt.Errorf("cache: create lru: %w", err)
}
m := &Manager{
enabled: true,
ttl: cfg.TTL,
lru: c,
}
// 初始化 Redis 连接(仅 enabled 时)
if cfg.Redis.Enabled {
m.rdb = redis.NewClient(&redis.Options{
Addr: cfg.Redis.Addr,
Password: cfg.Redis.Password,
DB: cfg.Redis.DB,
})
m.redisDB = cfg.Redis.DB
// 探测 Redis 连通性
if err := m.rdb.Ping(context.Background()).Err(); err != nil {
return nil, fmt.Errorf("cache: redis ping failed: %w", err)
}
}
return m, nil
}
// cacheKey 生成缓存键。格式:bowen:cache:{provider}:{hash[:8]}:{max_results}
func cacheKey(req *models.SearchRequest, provider models.ProviderKind) string {
h := sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%d", provider, req.Query, req.MaxResults)))
return fmt.Sprintf("bowen:cache:%x:%d", h[:8], req.MaxResults)
}
// Get 获取缓存结果。
// 查询顺序:L1 (LRU) → L2 (Redis) → nil
// 若 L2 命中,自动升温到 L1。
func (m *Manager) Get(req *models.SearchRequest, provider models.ProviderKind) *models.SearchResult {
if !m.enabled {
return nil
}
key := cacheKey(req, provider)
// L1: 内存 LRU
if result := m.getL1(key); result != nil {
m.mu.Lock()
m.l1Hits++
m.mu.Unlock()
return result
}
// L2: Redis(如启用)
if m.rdb != nil {
if result := m.getL2(key); result != nil {
// 升温到 L1
m.setL1(key, result)
m.mu.Lock()
m.l2Hits++
m.mu.Unlock()
return result
}
}
m.mu.Lock()
m.misses++
m.mu.Unlock()
return nil
}
func (m *Manager) getL1(key string) *models.SearchResult {
m.mu.RLock()
e, ok := m.lru.Get(key)
m.mu.RUnlock()
if !ok {
return nil
}
if time.Now().After(e.expiresAt) {
m.mu.Lock()
m.lru.Remove(key)
m.mu.Unlock()
return nil
}
result := *e.result
result.Cached = true
return &result
}
func (m *Manager) getL2(key string) *models.SearchResult {
data, err := m.rdb.Get(context.Background(), key).Bytes()
if err != nil {
return nil
}
var result models.SearchResult
if err := json.Unmarshal(data, &result); err != nil {
return nil
}
result.Cached = true
return &result
}
// Set 写入缓存(L1 + L2,若 Redis 启用则同步写 L2)。
func (m *Manager) Set(req *models.SearchRequest, provider models.ProviderKind, result *models.SearchResult) {
if !m.enabled {
return
}
key := cacheKey(req, provider)
// 写 L1
m.setL1(key, result)
// 写 L2(异步 fire-and-forget,不阻塞)
if m.rdb != nil {
go func() {
data, err := json.Marshal(result)
if err != nil {
return
}
m.rdb.Set(context.Background(), key, data, m.ttl)
}()
}
}
func (m *Manager) setL1(key string, result *models.SearchResult) {
m.mu.Lock()
defer m.mu.Unlock()
r := *result
m.lru.Add(key, &entry{
result: &r,
expiresAt: time.Now().Add(m.ttl),
})
}
// Stats 缓存统计(L1 命中、L2 命中、L1+L2 总未命中)。
func (m *Manager) Stats() (l1Hits, l2Hits, misses int64) {
m.mu.RLock()
defer m.mu.RUnlock()
return m.l1Hits, m.l2Hits, m.misses
}
// HitRate 综合缓存命中率(L1+L2 命中 / 总请求)。
func (m *Manager) HitRate() float64 {
m.mu.RLock()
defer m.mu.RUnlock()
total := m.l1Hits + m.l2Hits + m.misses
if total == 0 {
return 0
}
return float64(m.l1Hits+m.l2Hits) / float64(total)
}
// SetTTL 运行时更新 TTL(仅影响后续写入的条目)。
func (m *Manager) SetTTL(d time.Duration) {
m.mu.Lock()
defer m.mu.Unlock()
m.ttl = d
}
// Clear 清空缓存。
func (m *Manager) Clear() {
if !m.enabled {
return
}
m.mu.Lock()
defer m.mu.Unlock()
m.lru.Purge()
m.l1Hits = 0
m.l2Hits = 0
m.misses = 0
}
// Close 关闭 Redis 连接。
func (m *Manager) Close() error {
if m.rdb != nil {
return m.rdb.Close()
}
return nil
}
+122
View File
@@ -0,0 +1,122 @@
package cache
import (
"testing"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
func TestCache_SetAndGet(t *testing.T) {
m, err := New(config.CacheConfig{
Enabled: true,
TTL: 5 * time.Minute,
MaxItems: 100,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
req := &models.SearchRequest{Query: "test", MaxResults: 5}
result := &models.SearchResult{Query: "test", TotalCount: 1}
m.Set(req, models.ProviderTavily, result)
got := m.Get(req, models.ProviderTavily)
if got == nil {
t.Fatal("expected cached result")
}
if !got.Cached {
t.Error("expected cached=true")
}
}
func TestCache_Miss(t *testing.T) {
m, err := New(config.CacheConfig{
Enabled: true,
TTL: 5 * time.Minute,
MaxItems: 100,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
got := m.Get(&models.SearchRequest{Query: "nosuch"}, models.ProviderTavily)
if got != nil {
t.Fatal("expected nil for cache miss")
}
}
func TestCache_Disabled(t *testing.T) {
m, err := New(config.CacheConfig{Enabled: false})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
m.Set(&models.SearchRequest{Query: "test"}, models.ProviderTavily, &models.SearchResult{})
got := m.Get(&models.SearchRequest{Query: "test"}, models.ProviderTavily)
if got != nil {
t.Fatal("expected nil when cache disabled")
}
}
func TestCache_Expired(t *testing.T) {
m, err := New(config.CacheConfig{
Enabled: true,
TTL: -1 * time.Second, // 立即过期
MaxItems: 100,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
m.Set(&models.SearchRequest{Query: "test"}, models.ProviderTavily, &models.SearchResult{})
got := m.Get(&models.SearchRequest{Query: "test"}, models.ProviderTavily)
if got != nil {
t.Fatal("expected nil for expired entry")
}
}
func TestCache_Stats(t *testing.T) {
m, err := New(config.CacheConfig{
Enabled: true,
TTL: 5 * time.Minute,
MaxItems: 100,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Miss then hit
m.Get(&models.SearchRequest{Query: "miss"}, models.ProviderTavily)
m.Set(&models.SearchRequest{Query: "hit"}, models.ProviderTavily, &models.SearchResult{})
m.Get(&models.SearchRequest{Query: "hit"}, models.ProviderTavily)
l1Hits, l2Hits, misses := m.Stats()
if l1Hits != 1 {
t.Errorf("expected 1 L1 hit, got %d", l1Hits)
}
if l2Hits != 0 {
t.Errorf("expected 0 L2 hits, got %d", l2Hits)
}
if misses != 1 {
t.Errorf("expected 1 miss, got %d", misses)
}
}
func TestCache_ProviderIsolation(t *testing.T) {
m, err := New(config.CacheConfig{
Enabled: true,
TTL: 5 * time.Minute,
MaxItems: 100,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
req := &models.SearchRequest{Query: "test", MaxResults: 5}
m.Set(req, models.ProviderTavily, &models.SearchResult{Query: "test"})
// Same query, different provider should miss
got := m.Get(req, models.ProviderBrave)
if got != nil {
t.Fatal("expected miss for different provider")
}
}
+179
View File
@@ -0,0 +1,179 @@
// Package config 管理博闻配置。
package config
import (
"fmt"
"os"
"strconv"
"strings"
"time"
)
// Config 顶层配置。
type Config struct {
Server ServerConfig `yaml:"server"`
Providers ProvidersConfig `yaml:"providers"`
Balancer BalancerConfig `yaml:"balancer"`
Cache CacheConfig `yaml:"cache"`
Metrics MetricsConfig `yaml:"metrics"`
}
// ServerConfig 服务端配置。
type ServerConfig struct {
Port int `yaml:"port"`
ReadTimeout time.Duration `yaml:"read_timeout"`
WriteTimeout time.Duration `yaml:"write_timeout"`
}
// ProvidersConfig 聚合所有搜索提供商配置。
type ProvidersConfig struct {
Tavily ProviderConfig `yaml:"tavily"`
Brave ProviderConfig `yaml:"brave"`
}
// ProviderConfig 单个搜索提供商配置。
type ProviderConfig struct {
Enabled bool `yaml:"enabled"`
Timeout time.Duration `yaml:"timeout"`
MaxResults int `yaml:"max_results"`
HealthCheck HealthCheckConf `yaml:"health_check"`
// APIKeys 从环境变量注入,不在 YAML 中
APIKeys []string `yaml:"-"`
}
// HealthCheckConf 健康检查配置。
type HealthCheckConf struct {
Interval time.Duration `yaml:"interval"`
Timeout time.Duration `yaml:"timeout"`
FailureThreshold int `yaml:"failure_threshold"`
}
// BalancerConfig 负载均衡配置。
type BalancerConfig struct {
Strategy string `yaml:"strategy"`
RetryCount int `yaml:"retry_count"`
CircuitBreaker CircuitBreakerCfg `yaml:"circuit_breaker"`
}
// CircuitBreakerCfg 熔断器配置。
type CircuitBreakerCfg struct {
Enabled bool `yaml:"enabled"`
FailureThreshold int `yaml:"failure_threshold"`
Cooldown time.Duration `yaml:"cooldown"`
}
// CacheConfig 缓存配置。
type CacheConfig struct {
Enabled bool `yaml:"enabled"`
TTL time.Duration `yaml:"ttl"`
MaxItems int `yaml:"max_items"`
Redis RedisConfig `yaml:"redis"`
}
// RedisConfig Redis 缓存配置(二级缓存层,可开关)。
type RedisConfig struct {
Enabled bool `yaml:"enabled"`
Addr string `yaml:"addr"`
Password string `yaml:"password"`
DB int `yaml:"db"`
}
// MetricsConfig 监控指标配置。
type MetricsConfig struct {
Enabled bool `yaml:"enabled"`
Port int `yaml:"port"`
Path string `yaml:"path"`
}
// envVarKeys 定义各 provider 的 API Key 环境变量名。
var envVarKeys = map[string]string{
"tavily": "TAVILY_API_KEYS",
"brave": "BRAVE_API_KEYS",
}
// Load 加载配置,API Keys 从环境变量读取。
// configPath 为可选 YAML 配置文件路径。
func Load(configPath string) (*Config, error) {
cfg := &Config{}
setDefaults(cfg)
// 从 YAML 文件加载(如存在)
if configPath != "" {
if err := loadFromFile(configPath, cfg); err != nil {
return nil, err
}
}
// 从环境变量加载 API Keys
cfg.Providers.Tavily.APIKeys = parseEnvKeys(envVarKeys["tavily"])
cfg.Providers.Brave.APIKeys = parseEnvKeys(envVarKeys["brave"])
// 从环境变量加载 Redis 配置
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ENABLED"); env != "" {
cfg.Cache.Redis.Enabled, _ = strconv.ParseBool(env)
}
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ADDR"); env != "" {
cfg.Cache.Redis.Addr = env
}
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_PASSWORD"); env != "" {
cfg.Cache.Redis.Password = env
}
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_DB"); env != "" {
if db, err := strconv.Atoi(env); err == nil {
cfg.Cache.Redis.DB = db
}
}
return cfg, nil
}
func setDefaults(cfg *Config) {
cfg.Server.Port = 8080
cfg.Server.ReadTimeout = 30 * time.Second
cfg.Server.WriteTimeout = 30 * time.Second
// Provider 默认启用(可通过 YAML 中 enabled: false 关闭)
cfg.Providers.Tavily.Enabled = true
cfg.Providers.Tavily.Timeout = 10 * time.Second
cfg.Providers.Tavily.MaxResults = 10
cfg.Providers.Brave.Enabled = true
cfg.Providers.Brave.Timeout = 10 * time.Second
cfg.Providers.Brave.MaxResults = 10
cfg.Balancer.Strategy = "weighted_round_robin"
cfg.Balancer.RetryCount = 2
cfg.Balancer.CircuitBreaker.Enabled = true
cfg.Balancer.CircuitBreaker.FailureThreshold = 5
cfg.Balancer.CircuitBreaker.Cooldown = 30 * time.Second
cfg.Cache.Enabled = true
cfg.Cache.TTL = 5 * time.Minute
cfg.Cache.MaxItems = 1000
cfg.Cache.Redis.Enabled = false
cfg.Cache.Redis.Addr = "127.0.0.1:6379"
cfg.Cache.Redis.Password = ""
cfg.Cache.Redis.DB = 0
cfg.Metrics.Enabled = true
cfg.Metrics.Port = 9090
cfg.Metrics.Path = "/metrics"
}
func loadFromFile(path string, cfg *Config) error {
data, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("read config: %w", err)
}
return yamlUnmarshal(data, cfg)
}
func parseEnvKeys(envName string) []string {
val := os.Getenv(envName)
if val == "" {
return nil
}
parts := strings.Split(val, ",")
keys := make([]string, 0, len(parts))
for _, p := range parts {
if trimmed := strings.TrimSpace(p); trimmed != "" {
keys = append(keys, trimmed)
}
}
return keys
}
+109
View File
@@ -0,0 +1,109 @@
package config
import (
"log/slog"
"sync"
"time"
"github.com/fsnotify/fsnotify"
)
// Reloader 配置文件热重载器。
// 监听 YAML 文件变更,通过 onChange 回调通知调用方。
type Reloader struct {
path string
cfg *Config
watcher *fsnotify.Watcher
onChange func(old, new *Config) error
mu sync.Mutex
debounce *time.Timer
done chan struct{}
}
// NewReloader 创建并启动配置热重载。
// onChange 在配置变更时被调用,接收 old/new 两份配置。
func NewReloader(path string, cfg *Config, onChange func(old, new *Config) error) (*Reloader, error) {
w, err := fsnotify.NewWatcher()
if err != nil {
return nil, err
}
if err := w.Add(path); err != nil {
w.Close()
return nil, err
}
r := &Reloader{
path: path,
cfg: cfg,
watcher: w,
onChange: onChange,
debounce: time.NewTimer(0),
done: make(chan struct{}),
}
// 停止初始化的空 timer
if !r.debounce.Stop() {
<-r.debounce.C
}
go r.loop()
slog.Info("配置热重载已启动", "path", path)
return r, nil
}
func (r *Reloader) loop() {
for {
select {
case event := <-r.watcher.Events:
if event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) == 0 {
continue
}
slog.Debug("配置文件变更", "event", event.Name)
r.debounce.Reset(500 * time.Millisecond)
case <-r.debounce.C:
r.reload()
case err := <-r.watcher.Errors:
slog.Error("配置热重载监听异常", "error", err)
case <-r.done:
return
}
}
}
func (r *Reloader) reload() {
r.mu.Lock()
defer r.mu.Unlock()
old := r.cfg
// 只重新加载 YAML 部分(API Keys 仍由环境变量管理)
updated := *old // 浅拷贝
if err := loadFromFile(r.path, &updated); err != nil {
slog.Error("配置热重载读取失败", "error", err)
return
}
// 保留环境变量注入的值
updated.Providers.Tavily.APIKeys = old.Providers.Tavily.APIKeys
updated.Providers.Brave.APIKeys = old.Providers.Brave.APIKeys
updated.Cache.Redis.Enabled = old.Cache.Redis.Enabled
updated.Cache.Redis.Addr = old.Cache.Redis.Addr
updated.Cache.Redis.Password = old.Cache.Redis.Password
updated.Cache.Redis.DB = old.Cache.Redis.DB
if err := r.onChange(r.cfg, &updated); err != nil {
slog.Error("配置热重载应用失败", "error", err)
return
}
r.cfg = &updated
slog.Info("配置已热重载")
}
// Stop 停止配置热重载。
func (r *Reloader) Stop() {
close(r.done)
r.watcher.Close()
}
+7
View File
@@ -0,0 +1,7 @@
package config
import "gopkg.in/yaml.v3"
func yamlUnmarshal(data []byte, cfg *Config) error {
return yaml.Unmarshal(data, cfg)
}
+136
View File
@@ -0,0 +1,136 @@
// Package logger 提供结构化日志配置。
package logger
import (
"context"
"fmt"
"io"
"log/slog"
"os"
"strings"
"sync"
)
// Config 日志配置。
type Config struct {
Level slog.Level
AddSource bool
JSON bool // true = JSONHandler,false = 彩色 TextHandler
}
// New 创建配置好的 Logger。
// 开发环境(JSON=false)输出彩色终端文本,生产环境输出 JSON 供采集。
func New(cfg Config) *slog.Logger {
var h slog.Handler
opts := &slog.HandlerOptions{
Level: cfg.Level,
AddSource: cfg.AddSource,
}
if cfg.JSON {
h = slog.NewJSONHandler(os.Stdout, opts)
} else {
h = newColorHandler(os.Stdout, opts)
}
return slog.New(h)
}
// ANSI 颜色代码
const (
colorReset = "\033[0m"
colorRed = "\033[31m"
colorGreen = "\033[32m"
colorYellow = "\033[33m"
colorBlue = "\033[34m"
colorPurple = "\033[35m"
colorCyan = "\033[36m"
colorGray = "\033[37m"
)
// levelColor 日志级别 → ANSI 颜色映射
var levelColor = map[slog.Level]string{
slog.LevelDebug: colorGray,
slog.LevelInfo: colorGreen,
slog.LevelWarn: colorYellow,
slog.LevelError: colorRed,
}
// levelLabel 日志级别 → 短标签映射
var levelLabel = map[slog.Level]string{
slog.LevelDebug: "DEBUG",
slog.LevelInfo: "INFO",
slog.LevelWarn: "WARN",
slog.LevelError: "ERROR",
}
// colorHandler 带颜色的 slog.Handler 包装。
type colorHandler struct {
inner slog.Handler
mu sync.Mutex
writer io.Writer
}
func newColorHandler(w io.Writer, opts *slog.HandlerOptions) *colorHandler {
return &colorHandler{
inner: slog.NewTextHandler(w, opts),
writer: w,
}
}
func (h *colorHandler) Enabled(ctx context.Context, level slog.Level) bool {
return h.inner.Enabled(ctx, level)
}
func (h *colorHandler) Handle(ctx context.Context, r slog.Record) error {
h.mu.Lock()
defer h.mu.Unlock()
// 时间(使用 -07:00 让 Go 自动填充时区偏移)
timeStr := r.Time.Format("2006-01-02T15:04:05.000-07:00")
// Level 带颜色
lvl := r.Level
color := levelColor[lvl]
label := levelLabel[lvl]
// 消息
msg := r.Message
// 组装属性
var attrs []string
r.Attrs(func(a slog.Attr) bool {
attrs = append(attrs, fmt.Sprintf("%s=%v", a.Key, a.Value.Any()))
return true
})
// 输出格式:[时间] Level颜色(LEVEL) 消息 attr1=val1 attr2=val2
var buf strings.Builder
buf.WriteString(timeStr)
buf.WriteString(" ")
buf.WriteString(color)
buf.WriteString(fmt.Sprintf("%-5s", label))
buf.WriteString(colorReset)
buf.WriteString(" ")
buf.WriteString(msg)
for _, a := range attrs {
buf.WriteString(" ")
buf.WriteString(a)
}
buf.WriteString("\n")
_, err := fmt.Fprint(h.writer, buf.String())
return err
}
func (h *colorHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
return &colorHandler{
inner: h.inner.WithAttrs(attrs),
writer: h.writer,
}
}
func (h *colorHandler) WithGroup(name string) slog.Handler {
return &colorHandler{
inner: h.inner.WithGroup(name),
writer: h.writer,
}
}
+116
View File
@@ -0,0 +1,116 @@
package mcp
import (
"encoding/json"
"log/slog"
"net/http"
"strings"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// registerAdminRoutes 注册管理端点到 mux。
func (h *Handler) registerAdminRoutes(mux *http.ServeMux) {
mux.HandleFunc("/admin/provider/", h.handleProviderAdmin)
mux.HandleFunc("/admin/cache/clear", h.handleCacheClear)
mux.HandleFunc("/admin/keys", h.handleKeys)
}
// handleProviderAdmin 管理 provider 启用/禁用。
// PUT /admin/provider/{name}/enable
// PUT /admin/provider/{name}/disable
func (h *Handler) handleProviderAdmin(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPut {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
// 解析路径:/admin/provider/{name}/{action}
parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/admin/provider/"), "/")
if len(parts) != 2 {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "路径格式: /admin/provider/{name}/{action}"})
return
}
name := models.ProviderKind(parts[0])
action := parts[1]
clientIP := requestIP(r)
if !h.balancer.HasProvider(name) {
writeJSON(w, http.StatusNotFound, map[string]string{"error": "provider 不存在: " + string(name)})
return
}
switch action {
case "enable":
h.balancer.EnableProvider(name)
slog.Info("管理: 启用 provider", "provider", name, "from", clientIP)
writeJSON(w, http.StatusOK, map[string]string{"status": "enabled", "provider": string(name)})
case "disable":
h.balancer.DisableProvider(name)
slog.Info("管理: 禁用 provider", "provider", name, "from", clientIP)
writeJSON(w, http.StatusOK, map[string]string{"status": "disabled", "provider": string(name)})
default:
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "未知操作: " + action + "(可用: enable, disable)"})
}
}
// handleCacheClear 清空缓存。
// POST /admin/cache/clear
func (h *Handler) handleCacheClear(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
h.cache.Clear()
slog.Info("管理: 缓存已清空", "from", requestIP(r))
writeJSON(w, http.StatusOK, map[string]string{"status": "cleared"})
}
// handleKeys 列出所有密钥状态。
// GET /admin/keys
func (h *Handler) handleKeys(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
type keyInfo struct {
Provider string `json:"provider"`
Enabled bool `json:"enabled"`
}
var keys []keyInfo
for _, p := range h.providerKinds {
keys = append(keys, keyInfo{
Provider: string(p),
Enabled: h.balancer.IsProviderEnabled(p),
})
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"keys": keys,
})
}
func writeJSON(w http.ResponseWriter, status int, data interface{}) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
json.NewEncoder(w).Encode(data)
}
// requestIP 从请求中提取客户端 IP。
func requestIP(r *http.Request) string {
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
return fwd
}
// RemoteAddr 格式 "ip:port",去掉端口
if host := r.RemoteAddr; host != "" {
for i := len(host) - 1; i >= 0; i-- {
if host[i] == ':' {
return host[:i]
}
}
return host
}
return "unknown"
}
+240
View File
@@ -0,0 +1,240 @@
// Package mcp 实现 MCP 协议的工具注册、调用处理(传输层无关)。
package mcp
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"sort"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/balancer"
"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/metrics"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// ToolDefinition MCP 工具定义。
type ToolDefinition struct {
Name string `json:"name"`
Description string `json:"description"`
InputSchema interface{} `json:"inputSchema"`
}
// CallToolRequest MCP tools/call JSON-RPC 参数。
type CallToolRequest struct {
Name string `json:"name"`
Input json.RawMessage `json:"input,omitempty"`
Arguments json.RawMessage `json:"arguments,omitempty"` // MCP 标准字段名
}
// CallToolResult MCP tools/call 响应体。
type CallToolResult struct {
Content []ContentItem `json:"content"`
IsError bool `json:"isError,omitempty"`
}
// ContentItem MCP 内容项。
type ContentItem struct {
Type string `json:"type"`
Text string `json:"text"`
}
// Handler 搜索工具核心处理逻辑(传输层无关)。
type Handler struct {
balancer *balancer.LoadBalancer
cache *cache.Manager
providerKinds []models.ProviderKind
}
func NewHandler(b *balancer.LoadBalancer, c *cache.Manager, providers []models.ProviderKind) *Handler {
return &Handler{
balancer: b,
cache: c,
providerKinds: providers,
}
}
// ReloadConfig 应用配置热重载变更。
// 由 config.Reloader 在文件变更后调用,只处理 YAML 中可动态变更的字段。
func (h *Handler) ReloadConfig(oldCfg, newCfg *config.Config) {
// Provider 启用/禁用
type pc struct {
name models.ProviderKind
old config.ProviderConfig
new config.ProviderConfig
}
changes := []pc{
{models.ProviderTavily, oldCfg.Providers.Tavily, newCfg.Providers.Tavily},
{models.ProviderBrave, oldCfg.Providers.Brave, newCfg.Providers.Brave},
}
for _, c := range changes {
if !c.old.Enabled && c.new.Enabled {
h.balancer.EnableProvider(c.name)
slog.Info("配置热重载: 启用 provider", "name", c.name)
} else if c.old.Enabled && !c.new.Enabled {
h.balancer.DisableProvider(c.name)
slog.Info("配置热重载: 禁用 provider", "name", c.name)
}
}
// 缓存 TTL
if oldCfg.Cache.TTL != newCfg.Cache.TTL {
h.cache.SetTTL(newCfg.Cache.TTL)
slog.Info("配置热重载: 缓存 TTL 更新", "old", oldCfg.Cache.TTL, "new", newCfg.Cache.TTL)
}
// 负载均衡重试次数
if oldCfg.Balancer.RetryCount != newCfg.Balancer.RetryCount {
h.balancer.SetRetryCount(newCfg.Balancer.RetryCount)
slog.Info("配置热重载: 重试次数更新", "old", oldCfg.Balancer.RetryCount, "new", newCfg.Balancer.RetryCount)
}
}
// ToolDefinition 返回 web_search 工具定义(MCP tools/list 用)。
func (h *Handler) ToolDefinition() ToolDefinition {
return ToolDefinition{
Name: "web_search",
Description: "联网搜索,获取最新的网络信息。支持多搜索引擎。",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "搜索关键词",
},
"max_results": map[string]interface{}{
"type": "integer",
"default": 5,
"minimum": 1,
"maximum": 20,
"description": "最大返回结果数",
},
"mode": map[string]interface{}{
"type": "string",
"enum": []string{"first", "merge"},
"default": "first",
"description": "搜索模式:first(首个成功返回,默认)、merge(并发搜索所有引擎,去重合并)",
},
"providers": map[string]interface{}{
"type": "array",
"items": map[string]interface{}{
"type": "string",
},
"description": "优先使用的搜索引擎(默认全部可用引擎)",
},
},
"required": []string{"query"},
},
}
}
// HandleSearch 执行搜索(被 stdio、SSE、HTTP 三种传输层共享)。
func (h *Handler) HandleSearch(ctx context.Context, req *models.SearchRequest) *CallToolResult {
start := time.Now()
// 校验
if req.Query == "" {
metrics.SearchErrors.WithLabelValues("empty_query").Inc()
return errorResult("搜索关键词不能为空")
}
if req.MaxResults <= 0 {
req.MaxResults = 5
}
if req.MaxResults > 20 {
req.MaxResults = 20
}
if req.Mode == "" {
req.Mode = models.ModeFirst
}
// 查缓存(用空 provider 查跨 provider 缓存)
if cached := h.cache.Get(req, ""); cached != nil {
metrics.CacheHits.Inc()
metrics.SearchDuration.Observe(time.Since(start).Seconds())
return successResult(cached)
}
metrics.CacheMisses.Inc()
// 执行搜索
if req.Mode == models.ModeMerge {
return h.mergeSearch(ctx, req, start)
}
result, err := h.balancer.Search(ctx, req)
if err != nil {
metrics.SearchErrors.WithLabelValues("search_failed").Inc()
return errorResult(fmt.Sprintf("搜索失败: %v", err))
}
// 写入缓存
h.cache.Set(req, result.Provider, result)
metrics.SearchTotal.WithLabelValues(string(result.Provider)).Inc()
return successResult(result)
}
func successResult(result *models.SearchResult) *CallToolResult {
jsonBytes, _ := json.Marshal(result)
return &CallToolResult{
Content: []ContentItem{
{Type: "json", Text: string(jsonBytes)},
},
}
}
func errorResult(msg string) *CallToolResult {
errJSON, _ := json.Marshal(map[string]string{"error": msg})
return &CallToolResult{
IsError: true,
Content: []ContentItem{
{Type: "text", Text: string(errJSON)},
},
}
}
// mergeSearch 并发搜索所有 provider,合并去重后返回。
func (h *Handler) mergeSearch(ctx context.Context, req *models.SearchRequest, start time.Time) *CallToolResult {
results := h.balancer.SearchAll(ctx, req)
if len(results) == 0 {
metrics.SearchErrors.WithLabelValues("merge_no_results").Inc()
return errorResult("所有搜索引擎均未返回结果")
}
// 合并 + URL 去重
seen := make(map[string]bool)
var merged []models.ResultItem
for _, res := range results {
for _, item := range res.Results {
if seen[item.URL] {
continue
}
seen[item.URL] = true
merged = append(merged, item)
}
}
// 按 score 降序排序(无 score 的排最后)
sort.Slice(merged, func(i, j int) bool {
return merged[i].Score > merged[j].Score
})
// 截断到 max_results
if len(merged) > req.MaxResults {
merged = merged[:req.MaxResults]
}
final := &models.SearchResult{
Provider: models.ProviderKind("merge"),
Query: req.Query,
Results: merged,
TotalCount: len(merged),
DurationMs: time.Since(start).Milliseconds(),
Cached: false,
}
h.cache.Set(req, "", final)
return successResult(final)
}
@@ -0,0 +1,43 @@
package mcp
import (
"context"
"testing"
"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
func TestHandler_ToolDefinition(t *testing.T) {
h := NewHandler(nil, nil, nil)
def := h.ToolDefinition()
if def.Name != "web_search" {
t.Errorf("expected 'web_search', got '%s'", def.Name)
}
if def.Description == "" {
t.Error("expected non-empty description")
}
}
func TestHandler_HandleSearch_EmptyQuery(t *testing.T) {
c, _ := cache.New(config.CacheConfig{Enabled: false})
h := NewHandler(nil, c, nil)
result := h.HandleSearch(context.Background(), &models.SearchRequest{Query: ""})
if !result.IsError {
t.Fatal("expected error for empty query")
}
}
func TestHandler_HandleSearch_DefaultMaxResults(t *testing.T) {
req := &models.SearchRequest{Query: "test"}
if req.MaxResults <= 0 {
req.MaxResults = 5
}
if req.MaxResults > 20 {
req.MaxResults = 20
}
if req.MaxResults != 5 {
t.Errorf("expected default 5, got %d", req.MaxResults)
}
}
+242
View File
@@ -0,0 +1,242 @@
package mcp
import (
"fmt"
"net/http"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// HandleMonitor 返回 HTML 监控面板(GET /)。
func (h *Handler) HandleMonitor(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html; charset=utf-8")
uptime := time.Since(startTime).Round(time.Second).String()
// Provider 表格行
var providerRows string
for _, p := range h.providerKinds {
counts := h.balancer.ActiveKeyCount()
activeKeys := counts[p]
enabled := h.balancer.IsProviderEnabled(p)
statusIcon := "🟢"
statusText := "启用"
if !enabled {
statusIcon = "🔴"
statusText = "禁用"
}
providerRows += fmt.Sprintf(`<tr>
<td>%s</td>
<td><span class="badge badge-provider">%s</span></td>
<td>%s <span class="status-text">%s</span></td>
<td class="num">%d</td>
</tr>`, statusIcon, p, statusIcon, statusText, activeKeys)
}
// 缓存统计
l1Hits, l2Hits, misses := h.cache.Stats()
hitRate := h.cache.HitRate()
html := fmt.Sprintf(`<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>博闻 · 监控面板</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body {
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Noto Sans SC", sans-serif;
background: #0f172a;
color: #e2e8f0;
padding: 24px;
line-height: 1.6;
}
.container { max-width: 960px; margin: 0 auto; }
h1 {
font-size: 24px;
font-weight: 600;
margin-bottom: 4px;
display: flex;
align-items: center;
gap: 12px;
}
h1 small {
font-size: 14px;
font-weight: 400;
color: #64748b;
}
.subtitle {
color: #64748b;
font-size: 14px;
margin-bottom: 24px;
}
.grid {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
gap: 16px;
margin-bottom: 24px;
}
.card {
background: #1e293b;
border-radius: 12px;
padding: 16px 20px;
border: 1px solid #334155;
}
.card .label {
font-size: 12px;
color: #64748b;
text-transform: uppercase;
letter-spacing: 0.05em;
margin-bottom: 4px;
}
.card .value {
font-size: 28px;
font-weight: 700;
color: #38bdf8;
}
.card .value.green { color: #4ade80; }
.card .value.yellow { color: #fbbf24; }
.card .value.red { color: #f87171; }
.section-title {
font-size: 16px;
font-weight: 600;
margin-bottom: 12px;
color: #94a3b8;
}
table {
width: 100%%;
border-collapse: collapse;
background: #1e293b;
border-radius: 12px;
overflow: hidden;
border: 1px solid #334155;
margin-bottom: 24px;
}
th {
text-align: left;
padding: 12px 16px;
font-size: 12px;
color: #64748b;
text-transform: uppercase;
letter-spacing: 0.05em;
background: #1a2332;
border-bottom: 1px solid #334155;
}
td {
padding: 10px 16px;
border-bottom: 1px solid #1e293b;
font-size: 14px;
}
tr:last-child td { border-bottom: none; }
tr:hover { background: #243348; }
.num { text-align: right; font-variant-numeric: tabular-nums; }
.badge {
display: inline-block;
padding: 2px 8px;
border-radius: 4px;
font-size: 12px;
font-weight: 500;
}
.badge-provider {
background: #1e3a5f;
color: #93c5fd;
}
.status-text {
font-size: 13px;
color: #94a3b8;
margin-left: 4px;
}
a { color: #38bdf8; text-decoration: none; }
a:hover { text-decoration: underline; }
.links { margin-top: 16px; display: flex; gap: 16px; font-size: 14px; }
.footer { margin-top: 32px; padding-top: 16px; border-top: 1px solid #334155; color: #475569; font-size: 12px; }
</style>
</head>
<body>
<div class="container">
<h1>🔍 博闻 <small>Bowen Web Search</small></h1>
<div class="subtitle">运行时间 %s</div>
<div class="grid">
<div class="card">
<div class="label">Provider</div>
<div class="value">%d</div>
</div>
<div class="card">
<div class="label">活跃密钥</div>
<div class="value">%d</div>
</div>
<div class="card">
<div class="label">缓存命中率</div>
<div class="value %s">%d%%%%</div>
</div>
<div class="card">
<div class="label">缓存条目</div>
<div class="value">%d</div>
</div>
</div>
<div class="section-title">📡 Provider 状态</div>
<table>
<thead>
<tr><th>状态</th><th>名称</th><th>开关</th><th class="num">活跃密钥</th></tr>
</thead>
<tbody>%s</tbody>
</table>
<div class="section-title">💾 缓存统计</div>
<table>
<thead>
<tr><th>指标</th><th class="num">值</th></tr>
</thead>
<tbody>
<tr><td>L1 命中(内存 LRU)</td><td class="num">%d</td></tr>
<tr><td>L2 命中(Redis)</td><td class="num">%d</td></tr>
<tr><td>未命中</td><td class="num">%d</td></tr>
<tr><td>总命中率</td><td class="num">%.1f%%%%</td></tr>
</tbody>
</table>
<div class="links">
<span>🔗</span>
<a href="/status">/status (JSON)</a>
<a href="/admin/keys">/admin/keys</a>
<a href="/health">/health</a>
</div>
<div class="footer">博闻 v2.0.0 · 数据实时 · 页面每 30 秒刷新</div>
</div>
<script>setTimeout(function(){ location.reload(); }, 30000);</script>
</body>
</html>`,
uptime,
len(h.providerKinds),
sumActiveKeys(h.balancer.ActiveKeyCount()),
hitRateColor(hitRate),
int(hitRate*100),
l1Hits+l2Hits+misses,
providerRows,
l1Hits, l2Hits, misses, hitRate*100,
)
fmt.Fprint(w, html)
}
func sumActiveKeys(m map[models.ProviderKind]int) int {
s := 0
for _, v := range m {
s += v
}
return s
}
func hitRateColor(rate float64) string {
if rate >= 0.6 {
return "green"
} else if rate >= 0.3 {
return "yellow"
}
return "red"
}
+101
View File
@@ -0,0 +1,101 @@
package mcp
import (
"encoding/json"
"fmt"
"log/slog"
"net/http"
"time"
)
// LoggingMiddleware 记录 HTTP 请求的源 IP、方法、路径和处理耗时。
func LoggingMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
ip := r.RemoteAddr
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
ip = fwd
}
method := r.Method
path := r.URL.Path
next.ServeHTTP(w, r)
slog.Info("HTTP",
"method", method,
"path", path,
"from", ip,
"ms", time.Since(start).Milliseconds(),
)
})
}
// ServeHTTP 启动 Streamable HTTP 传输模式的 MCP 服务。
//
// 传输协议:MCP Streamable HTTP(JSON-RPC over HTTP POST /)
// POST / → MCP 处理,GET / → 帮助页
func (h *Handler) ServeHTTP(addr string) error {
mux := http.NewServeMux()
mux.HandleFunc("/", h.handleRoot)
// 管理端点
mux.HandleFunc("/health", h.HandleHealth)
mux.HandleFunc("/status", h.handleStatus)
h.registerAdminRoutes(mux)
slog.Info("博闻 HTTP 服务启动", "addr", addr, "providers", h.providerKinds)
return http.ListenAndServe(addr, LoggingMiddleware(mux))
}
// handleRoot MCP 服务根路径。
// POST → MCP Streamable HTTP(JSON-RPC 2.0)
// GET → 帮助页
func (h *Handler) handleRoot(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodPost {
h.handleStreamableHTTP(w, r)
return
}
if r.URL.Path != "/" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
fmt.Fprint(w, `博闻 Bowen · MCP Web Search
传输协议:
POST / MCP Streamable HTTP(JSON-RPC 2.0,单端点)
管理端点:
GET /health 存活探针
GET /status 系统状态
GET /admin/keys 密钥状态
PUT /admin/provider/:name/enable 启用 provider
PUT /admin/provider/:name/disable 禁用 provider
POST /admin/cache/clear 清空缓存
监控面板(metrics 端口):
GET / HTML 仪表盘
GET /metrics Prometheus 原始指标
Stdio 模式:go run ./cmd/server --stdio
`)
}
// handleStreamableHTTP MCP Streamable HTTP 传输端点。
// 接收 JSON-RPC 2.0 POST 请求,返回 JSON-RPC 响应。
// 现阶段返回完整响应(非流式),后续可扩展为 text/event-stream 流式传输。
func (h *Handler) handleStreamableHTTP(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req JSONRPCRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
json.NewEncoder(w).Encode(jsonRPCError(nil, -32700, "Parse error"))
return
}
resp := h.handleJSONRPC(req)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(resp)
}
+75
View File
@@ -0,0 +1,75 @@
package mcp
import (
"encoding/json"
"net/http"
"time"
)
// SystemStatus 博闻系统状态(通过 /status 或管理 API 查询)。
type SystemStatus struct {
Version string `json:"version"`
Uptime string `json:"uptime"`
Providers []ProviderStatus `json:"providers"`
Cache CacheStats `json:"cache"`
}
// ProviderStatus provider 状态摘要。
type ProviderStatus struct {
Name string `json:"name"`
ActiveKeys int `json:"active_keys"`
Enabled bool `json:"enabled"`
}
// CacheStats 缓存状态摘要。
type CacheStats struct {
Enabled bool `json:"enabled"`
L1Hits int64 `json:"l1_hits"`
L2Hits int64 `json:"l2_hits"`
Misses int64 `json:"misses"`
HitRate float64 `json:"hit_rate"`
}
// startTime 记录博闻启动时间,用于 /status 的 uptime 计算。
var startTime = time.Now()
// HandleHealth 存活探针(GET /health → 200)。
func (h *Handler) HandleHealth(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}
// handleStatus 系统状态(GET /status → SystemStatus)。
func (h *Handler) handleStatus(w http.ResponseWriter, r *http.Request) {
// Provider 状态
var providers []ProviderStatus
for _, p := range h.providerKinds {
counts := h.balancer.ActiveKeyCount()
activeKeys := counts[p]
providers = append(providers, ProviderStatus{
Name: string(p),
ActiveKeys: activeKeys,
Enabled: activeKeys > 0,
})
}
// 缓存状态
l1Hits, l2Hits, misses := h.cache.Stats()
cacheStats := CacheStats{
Enabled: true,
L1Hits: l1Hits,
L2Hits: l2Hits,
Misses: misses,
HitRate: h.cache.HitRate(),
}
status := SystemStatus{
Version: "2.0.0",
Uptime: time.Since(startTime).Round(time.Second).String(),
Providers: providers,
Cache: cacheStats,
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(status)
}
+111
View File
@@ -0,0 +1,111 @@
package mcp
import (
"bufio"
"context"
"encoding/json"
"fmt"
"log/slog"
"os"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// JSONRPCRequest JSON-RPC 2.0 请求。
type JSONRPCRequest struct {
JSONRPC string `json:"jsonrpc"`
ID *int `json:"id"`
Method string `json:"method"`
Params json.RawMessage `json:"params,omitempty"`
}
// JSONRPCResponse JSON-RPC 2.0 响应。
type JSONRPCResponse struct {
JSONRPC string `json:"jsonrpc"`
ID *int `json:"id"`
Result interface{} `json:"result,omitempty"`
Error *RPCError `json:"error,omitempty"`
}
// RPCError JSON-RPC 错误。
type RPCError struct {
Code int `json:"code"`
Message string `json:"message"`
}
// ServeStdio 通过 stdin/stdout 提供 MCP JSON-RPC 服务。
func (h *Handler) ServeStdio() {
slog.Info("博闻 stdio 模式启动", "providers", h.providerKinds)
reader := bufio.NewReader(os.Stdin)
encoder := json.NewEncoder(os.Stdout)
for {
line, err := reader.ReadString('\n')
if err != nil {
slog.Error("stdio read error", "error", err)
return
}
var req JSONRPCRequest
if err := json.Unmarshal([]byte(line), &req); err != nil {
slog.Warn("stdio parse error", "error", err)
continue
}
resp := h.handleJSONRPC(req)
if err := encoder.Encode(resp); err != nil {
slog.Error("stdio write error", "error", err)
}
}
}
func (h *Handler) handleJSONRPC(req JSONRPCRequest) JSONRPCResponse {
switch req.Method {
case "tools/list":
return JSONRPCResponse{
JSONRPC: "2.0",
ID: req.ID,
Result: map[string]interface{}{
"tools": []ToolDefinition{h.ToolDefinition()},
},
}
case "tools/call":
var callReq CallToolRequest
if err := json.Unmarshal(req.Params, &callReq); err != nil {
return jsonRPCError(req.ID, -32602, "Invalid params")
}
if callReq.Name != "web_search" {
return jsonRPCError(req.ID, -32601, fmt.Sprintf("Tool not found: %s", callReq.Name))
}
// 优先使用 MCP 标准 arguments 字段,兼容旧 input 字段
rawInput := callReq.Arguments
if len(rawInput) == 0 {
rawInput = callReq.Input
}
var searchReq models.SearchRequest
if err := json.Unmarshal(rawInput, &searchReq); err != nil {
return jsonRPCError(req.ID, -32602, "Invalid search request parameters: "+err.Error())
}
result := h.HandleSearch(context.Background(), &searchReq)
return JSONRPCResponse{
JSONRPC: "2.0",
ID: req.ID,
Result: result,
}
default:
return jsonRPCError(req.ID, -32601, fmt.Sprintf("Method not found: %s", req.Method))
}
}
func jsonRPCError(id *int, code int, msg string) JSONRPCResponse {
return JSONRPCResponse{
JSONRPC: "2.0",
ID: id,
Error: &RPCError{Code: code, Message: msg},
}
}
@@ -0,0 +1,69 @@
// Package metrics 定义 Prometheus 监控指标。
package metrics
import (
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
)
var (
// SearchTotal 按 provider 统计搜索请求总数。
SearchTotal = promauto.NewCounterVec(
prometheus.CounterOpts{
Name: "bowen_search_total",
Help: "Total search requests by provider.",
},
[]string{"provider"},
)
// SearchDuration 搜索请求耗时直方图。
SearchDuration = promauto.NewHistogram(
prometheus.HistogramOpts{
Name: "bowen_search_duration_seconds",
Help: "Search request duration in seconds.",
Buckets: prometheus.DefBuckets,
},
)
// SearchErrors 按原因统计搜索错误。
SearchErrors = promauto.NewCounterVec(
prometheus.CounterOpts{
Name: "bowen_search_errors_total",
Help: "Search errors by reason.",
},
[]string{"reason"},
)
// CacheHits 缓存命中总数。
CacheHits = promauto.NewCounter(
prometheus.CounterOpts{
Name: "bowen_cache_hits_total",
Help: "Total cache hits.",
},
)
// CacheMisses 缓存未命中总数。
CacheMisses = promauto.NewCounter(
prometheus.CounterOpts{
Name: "bowen_cache_misses_total",
Help: "Total cache misses.",
},
)
// CacheHitsL2 Redis 二级缓存命中数。
CacheHitsL2 = promauto.NewCounter(
prometheus.CounterOpts{
Name: "bowen_cache_hits_l2_total",
Help: "L2 (Redis) cache hits.",
},
)
// ActiveKeys 各 provider 的活跃密钥数。
ActiveKeys = promauto.NewGaugeVec(
prometheus.GaugeOpts{
Name: "bowen_active_keys",
Help: "Number of active API keys by provider.",
},
[]string{"provider"},
)
)
+45
View File
@@ -0,0 +1,45 @@
// Package models 定义博闻的数据类型。
package models
// ProviderKind 搜索提供商枚举。
type ProviderKind string
const (
ProviderTavily ProviderKind = "tavily"
ProviderBrave ProviderKind = "brave"
)
// SearchMode 搜索结果聚合模式。
type SearchMode string
const (
ModeFirst SearchMode = "first" // 首个成功返回(低延迟,默认)
ModeMerge SearchMode = "merge" // 合并多 provider 结果(高召回)
)
// SearchRequest MCP 工具入参。
type SearchRequest struct {
Query string `json:"query"`
MaxResults int `json:"max_results"`
Providers []string `json:"providers,omitempty"`
Mode SearchMode `json:"mode,omitempty"`
}
// SearchResult 统一搜索响应。
type SearchResult struct {
Provider ProviderKind `json:"provider"`
Query string `json:"query"`
Results []ResultItem `json:"results"`
TotalCount int `json:"total_count"`
DurationMs int64 `json:"duration_ms"`
Cached bool `json:"cached"`
}
// ResultItem 单条搜索结果。
type ResultItem struct {
Title string `json:"title"`
URL string `json:"url"`
Snippet string `json:"snippet"`
PublishedAt string `json:"published_at,omitempty"`
Score float64 `json:"score,omitempty"`
}
+125
View File
@@ -0,0 +1,125 @@
package providers
import (
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/url"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
const braveAPIBase = "https://api.search.brave.com"
// BraveProvider 实现 SearchProvider。
type BraveProvider struct {
client *http.Client
apiBase string
maxResults int
}
func NewBraveProvider(timeout time.Duration, maxResults int) *BraveProvider {
return &BraveProvider{
client: &http.Client{
Timeout: timeout,
},
apiBase: braveAPIBase,
maxResults: maxResults,
}
}
func (b *BraveProvider) Name() models.ProviderKind {
return models.ProviderBrave
}
func (b *BraveProvider) MaxResults() int {
return b.maxResults
}
type braveWebResult struct {
Title string `json:"title"`
URL string `json:"url"`
Description string `json:"description"`
Age string `json:"age"`
}
type braveResponse struct {
Web struct {
Results []braveWebResult `json:"results"`
} `json:"web"`
}
func (b *BraveProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
maxRes := req.MaxResults
if maxRes <= 0 || maxRes > b.maxResults {
maxRes = b.maxResults
}
u, _ := url.Parse(b.apiBase + "/res/v1/web/search")
q := u.Query()
q.Set("q", req.Query)
q.Set("count", fmt.Sprintf("%d", maxRes))
u.RawQuery = q.Encode()
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
if err != nil {
return nil, fmt.Errorf("brave: create request: %w", err)
}
httpReq.Header.Set("Accept", "application/json")
httpReq.Header.Set("X-Subscription-Token", apiKey)
start := time.Now()
resp, err := b.client.Do(httpReq)
duration := time.Since(start)
if err != nil {
return nil, fmt.Errorf("brave: do request: %w", err)
}
defer resp.Body.Close()
raw, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("brave: read body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("brave: status %d: %s", resp.StatusCode, string(raw))
}
var br braveResponse
if err := json.Unmarshal(raw, &br); err != nil {
return nil, fmt.Errorf("brave: decode: %w", err)
}
items := make([]models.ResultItem, 0, len(br.Web.Results))
for _, r := range br.Web.Results {
items = append(items, models.ResultItem{
Title: r.Title,
URL: r.URL,
Snippet: r.Description,
PublishedAt: r.Age,
})
}
return &models.SearchResult{
Provider: models.ProviderBrave,
Query: req.Query,
Results: items,
TotalCount: len(items),
DurationMs: duration.Milliseconds(),
Cached: false,
}, nil
}
func (b *BraveProvider) HealthCheck(ctx context.Context) error {
var d net.Dialer
conn, err := d.DialContext(ctx, "tcp", "api.search.brave.com:443")
if err != nil {
return err
}
conn.Close()
return nil
}
@@ -0,0 +1,53 @@
package providers
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
func TestBraveProvider_Search_Success(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("X-Subscription-Token") == "" {
t.Error("missing subscription token header")
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.Write([]byte(`{
"web": {
"results": [
{"title": "Brave Result", "url": "https://brave.com/1", "description": "Brave snippet"}
]
}
}`))
}))
defer ts.Close()
p := &BraveProvider{
client: ts.Client(),
apiBase: ts.URL,
maxResults: 10,
}
r, err := p.Search(context.Background(), "test-key", &models.SearchRequest{Query: "test"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if r.TotalCount != 1 {
t.Fatalf("expected 1 result, got %d", r.TotalCount)
}
if r.Results[0].Title != "Brave Result" {
t.Errorf("expected 'Brave Result', got '%s'", r.Results[0].Title)
}
}
func TestBraveProvider_Name(t *testing.T) {
p := NewBraveProvider(5*time.Second, 10)
if p.Name() != models.ProviderBrave {
t.Errorf("expected brave, got %s", p.Name())
}
}
@@ -0,0 +1,20 @@
// Package providers 封装搜索提供商适配。
package providers
import (
"context"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// SearchProvider 搜索提供商接口。
type SearchProvider interface {
// Name 返回提供商标识。
Name() models.ProviderKind
// Search 执行搜索。apiKey 由负载均衡器传入,provider 本身不管理密钥。
Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error)
// HealthCheck 健康检查。
HealthCheck(ctx context.Context) error
// MaxResults 返回单次搜索最大数量。
MaxResults() int
}
@@ -0,0 +1,123 @@
package providers
import (
"fmt"
"log/slog"
"sync"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
// Registry 管理多 provider 注册和查询。
type Registry struct {
mu sync.RWMutex
providers map[models.ProviderKind]SearchProvider
}
func NewRegistry() *Registry {
return &Registry{
providers: make(map[models.ProviderKind]SearchProvider),
}
}
// Register 注册 provider。
func (r *Registry) Register(p SearchProvider) {
r.mu.Lock()
defer r.mu.Unlock()
r.providers[p.Name()] = p
}
// Get 按名称获取 provider。
func (r *Registry) Get(name models.ProviderKind) (SearchProvider, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
p, ok := r.providers[name]
return p, ok
}
// All 返回所有已注册 provider。
func (r *Registry) All() []SearchProvider {
r.mu.RLock()
defer r.mu.RUnlock()
out := make([]SearchProvider, 0, len(r.providers))
for _, p := range r.providers {
out = append(out, p)
}
return out
}
// Names 返回所有已注册的 provider 名称。
func (r *Registry) Names() []models.ProviderKind {
r.mu.RLock()
defer r.mu.RUnlock()
out := make([]models.ProviderKind, 0, len(r.providers))
for name := range r.providers {
out = append(out, name)
}
return out
}
// SetupFromConfig 根据配置创建并注册所有启用的 provider。
func SetupFromConfig(cfg *config.Config) (*Registry, error) {
reg := NewRegistry()
if cfg.Providers.Tavily.Enabled {
if len(cfg.Providers.Tavily.APIKeys) == 0 {
slog.Warn("tavily: enabled but no API keys configured, skipping")
} else {
reg.Register(NewTavilyProvider(cfg.Providers.Tavily.Timeout, cfg.Providers.Tavily.MaxResults))
}
}
if cfg.Providers.Brave.Enabled {
if len(cfg.Providers.Brave.APIKeys) == 0 {
slog.Warn("brave: enabled but no API keys configured, skipping")
} else {
reg.Register(NewBraveProvider(cfg.Providers.Brave.Timeout, cfg.Providers.Brave.MaxResults))
}
}
if len(reg.Names()) == 0 {
return nil, fmt.Errorf("no providers enabled or configured")
}
return reg, nil
}
// SetupFromConfigWithOverrides 从配置创建注册中心,并允许使用显式 provider 配置覆盖配置文件。
// 本函数方便测试和动态配置场景。
func SetupFromConfigWithOverrides(
cfg *config.Config,
overrideTimeout time.Duration,
overrideMaxResults int,
) (*Registry, error) {
reg := NewRegistry()
if cfg.Providers.Tavily.Enabled && len(cfg.Providers.Tavily.APIKeys) > 0 {
t := cfg.Providers.Tavily.Timeout
m := cfg.Providers.Tavily.MaxResults
if overrideTimeout > 0 {
t = overrideTimeout
}
if overrideMaxResults > 0 {
m = overrideMaxResults
}
reg.Register(NewTavilyProvider(t, m))
}
if cfg.Providers.Brave.Enabled && len(cfg.Providers.Brave.APIKeys) > 0 {
t := cfg.Providers.Brave.Timeout
m := cfg.Providers.Brave.MaxResults
if overrideTimeout > 0 {
t = overrideTimeout
}
if overrideMaxResults > 0 {
m = overrideMaxResults
}
reg.Register(NewBraveProvider(t, m))
}
if len(reg.Names()) == 0 {
return nil, fmt.Errorf("no providers enabled or configured")
}
return reg, nil
}
@@ -0,0 +1,63 @@
package providers
import (
"context"
"testing"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
type mockProvider struct {
name models.ProviderKind
}
func (m *mockProvider) Name() models.ProviderKind { return m.name }
func (m *mockProvider) Search(ctx context.Context, s string, req *models.SearchRequest) (*models.SearchResult, error) {
return &models.SearchResult{Provider: m.name, Query: req.Query}, nil
}
func (m *mockProvider) HealthCheck(ctx context.Context) error { return nil }
func (m *mockProvider) MaxResults() int { return 10 }
func TestRegistry_RegisterAndGet(t *testing.T) {
r := NewRegistry()
m := &mockProvider{name: models.ProviderTavily}
r.Register(m)
got, ok := r.Get(models.ProviderTavily)
if !ok {
t.Fatal("expected to find provider")
}
if got.Name() != models.ProviderTavily {
t.Errorf("expected tavily, got %s", got.Name())
}
}
func TestRegistry_All(t *testing.T) {
r := NewRegistry()
r.Register(&mockProvider{name: models.ProviderTavily})
r.Register(&mockProvider{name: models.ProviderBrave})
all := r.All()
if len(all) != 2 {
t.Errorf("expected 2 providers, got %d", len(all))
}
}
func TestRegistry_GetMissing(t *testing.T) {
r := NewRegistry()
_, ok := r.Get(models.ProviderBrave)
if ok {
t.Fatal("expected missing provider")
}
}
func TestRegistry_Names(t *testing.T) {
r := NewRegistry()
r.Register(&mockProvider{name: models.ProviderTavily})
r.Register(&mockProvider{name: models.ProviderBrave})
names := r.Names()
if len(names) != 2 {
t.Errorf("expected 2 names, got %d", len(names))
}
}
+134
View File
@@ -0,0 +1,134 @@
package providers
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
const tavilyAPIBase = "https://api.tavily.com"
type tavilyRequest struct {
APIKey string `json:"api_key"`
Query string `json:"query"`
MaxResults int `json:"max_results"`
}
type tavilyResponse struct {
Answer string `json:"answer"`
Query string `json:"query"`
ResponseTime float64 `json:"response_time"`
Results []tavilyResult `json:"results"`
}
type tavilyResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Score float64 `json:"score"`
PublishedAt string `json:"published_date"`
}
// TavilyProvider 实现 SearchProvider。
type TavilyProvider struct {
client *http.Client
apiBase string
maxResults int
}
func NewTavilyProvider(timeout time.Duration, maxResults int) *TavilyProvider {
return &TavilyProvider{
client: &http.Client{
Timeout: timeout,
},
apiBase: tavilyAPIBase,
maxResults: maxResults,
}
}
func (t *TavilyProvider) Name() models.ProviderKind {
return models.ProviderTavily
}
func (t *TavilyProvider) MaxResults() int {
return t.maxResults
}
func (t *TavilyProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
maxRes := req.MaxResults
if maxRes <= 0 || maxRes > t.maxResults {
maxRes = t.maxResults
}
body := tavilyRequest{
APIKey: apiKey,
Query: req.Query,
MaxResults: maxRes,
}
payload, _ := json.Marshal(body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, t.apiBase+"/search", bytes.NewReader(payload))
if err != nil {
return nil, fmt.Errorf("tavily: create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
start := time.Now()
resp, err := t.client.Do(httpReq)
duration := time.Since(start)
if err != nil {
return nil, fmt.Errorf("tavily: do request: %w", err)
}
defer resp.Body.Close()
raw, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("tavily: read body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("tavily: status %d: %s", resp.StatusCode, string(raw))
}
var tr tavilyResponse
if err := json.Unmarshal(raw, &tr); err != nil {
return nil, fmt.Errorf("tavily: decode response: %w", err)
}
items := make([]models.ResultItem, 0, len(tr.Results))
for _, r := range tr.Results {
items = append(items, models.ResultItem{
Title: r.Title,
URL: r.URL,
Snippet: r.Content,
Score: r.Score,
PublishedAt: r.PublishedAt,
})
}
return &models.SearchResult{
Provider: models.ProviderTavily,
Query: req.Query,
Results: items,
TotalCount: len(items),
DurationMs: duration.Milliseconds(),
Cached: false,
}, nil
}
func (t *TavilyProvider) HealthCheck(ctx context.Context) error {
var d net.Dialer
conn, err := d.DialContext(ctx, "tcp", "api.tavily.com:443")
if err != nil {
return err
}
conn.Close()
return nil
}
@@ -0,0 +1,71 @@
package providers
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)
func TestTavilyProvider_Search_Success(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/search" {
t.Errorf("unexpected path: %s", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.Write([]byte(`{
"query": "test",
"results": [
{"title": "Result 1", "url": "https://example.com/1", "content": "Content 1", "score": 0.95}
]
}`))
}))
defer ts.Close()
p := &TavilyProvider{
client: ts.Client(),
apiBase: ts.URL,
maxResults: 10,
}
result, err := p.Search(context.Background(), "test-key", &models.SearchRequest{Query: "test"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if result.TotalCount != 1 {
t.Fatalf("expected 1 result, got %d", result.TotalCount)
}
if result.Results[0].Title != "Result 1" {
t.Errorf("expected 'Result 1', got '%s'", result.Results[0].Title)
}
}
func TestTavilyProvider_Search_HTTPError(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
w.Write([]byte(`{"error": "invalid api key"}`))
}))
defer ts.Close()
p := &TavilyProvider{
client: ts.Client(),
apiBase: ts.URL,
maxResults: 10,
}
_, err := p.Search(context.Background(), "bad-key", &models.SearchRequest{Query: "test"})
if err == nil {
t.Fatal("expected error, got nil")
}
}
func TestTavilyProvider_Name(t *testing.T) {
p := NewTavilyProvider(5*time.Second, 10)
if p.Name() != models.ProviderTavily {
t.Errorf("expected tavily, got %s", p.Name())
}
}