AI Agent 工具调用中间件：Go 实现截断、超时与熔断

一、问题工具调用为什么需要中间件先看一个真实场景用户: 帮我分析这份 500 页 PDF 的内容 Agent: 调用 read_pdf 工具 Tool: 返回全部 500 页文本 (2MB) LLM: 上下文窗口溢出 → 报错 / 乱答这正是 LLM Agent 的典型故障模式工具返回值不受控撑爆上下文窗口。再看另一个用户: 把数据库里的用户数据导出来 Agent: 调用 db_query 工具 Tool: SQL 慢查询耗时 120 秒 Agent: 一直等待前端超时 → 用户以为系统挂了这些问题不能靠修改每个工具来解决 —— 每个工具开发者都自己写一遍超时和截断逻辑重复且易出错。中间件模式正是为此而生。二、中间件模式设计2.1 核心接口类似 HTTP 中间件我们定义工具调用的 Handler 和 Middleware// middleware/middleware.gopackagemiddlewareimport(contextgithub.com/example/mcp-server-go/internal/protocol)// ToolHandler 工具执行函数与上篇的 ToolHandler 一致typeToolHandlerfunc(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error)// Middleware 中间件 —— 包装一个 Handler返回新的 HandlertypeMiddlewarefunc(next ToolHandler)ToolHandler// Chain 将多个中间件串联// 执行顺序middlewares[0] → middlewares[1] → ... → handlerfuncChain(handler ToolHandler,middlewares...Middleware)ToolHandler{// 洋葱模型从外到内fori:len(middlewares)-1;i0;i--{handlermiddlewares[i](handler)}returnhandler}2.2 执行模型请求 → [截断中间件] → [超时中间件] → [熔断中间件] → [实际工具] ↓ 记录成功/失败 → 更新熔断状态每个中间件只做一件事组合起来形成完整的防御体系。三、中间件一输出截断TruncateMiddleware防止工具返回过大的结果撑爆 LLM 上下文窗口。// middleware/truncate.gopackagemiddlewareimport(contextfmtstringsunicode/utf8github.com/example/mcp-server-go/internal/protocol)// TruncateConfig 截断配置typeTruncateConfigstruct{MaxCharsint// 最大字符数0 不限制MaxTokensint// 最大 token 数粗略估计1 token ≈ 4 chars0 不限制TruncateMsgstring// 截断提示信息}funcDefaultTruncateConfig()TruncateConfig{returnTruncateConfig{MaxChars:32000,// 约 8K tokensMaxTokens:0,TruncateMsg:\n\n[...输出过长已截断请用更精确的查询...],}}// WithTruncate 创建截断中间件funcWithTruncate(cfg TruncateConfig)Middleware{returnfunc(next ToolHandler)ToolHandler{returnfunc(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error){result,err:next(ctx,args)iferr!nil||resultnil{returnresult,err}result.ContenttruncateContent(result.Content,cfg)returnresult,nil}}}functruncateContent(items[]protocol.ContentItem,cfg TruncateConfig)[]protocol.ContentItem{vartotalCharsinttruncated:falsefori:rangeitems{ifitems[i].Type!text{continue}itemChars:utf8.RuneCountInString(items[i].Text)ifcfg.MaxChars0totalCharsitemCharscfg.MaxChars{// 截断文本remaining:cfg.MaxChars-totalChars-utf8.RuneCountInString(cfg.TruncateMsg)ifremaining0{// 找到安全的截断点避免切断 UTF-8 多字节字符runes:[]rune(items[i].Text)ifremaininglen(runes){items[i].Textstring(runes[:remaining])cfg.TruncateMsg}}else{items[i].Textcfg.TruncateMsg}truncatedtruebreak}totalCharsitemChars// Token 估算截断ifcfg.MaxTokens0{estimatedTokens:totalChars/4ifestimatedTokenscfg.MaxTokens{items[i].Textitems[i].Text[:cfg.MaxTokens*4-len(items[i].Text)] [TRUNCATED]truncatedtruebreak}}}iftruncated{// 追加截断标记到最后一个 content itemlastText:items[len(items)-1]iflastText.Typetext{lastText.Textfmt.Sprintf(\n[截断统计: 原输出超过 %d 字符限制],cfg.MaxChars)}}returnitems}为什么不在工具内部截断工具开发者不需要关心 LLM 上下文窗口细节截断策略可以统一调整比如为不同模型设不同的 token 上限截断日志可以集中收集四、中间件二超时控制TimeoutMiddleware// middleware/timeout.gopackagemiddlewareimport(contextfmttimegithub.com/example/mcp-server-go/internal/protocol)// WithTimeout 创建超时中间件funcWithTimeout(timeout time.Duration)Middleware{returnfunc(next ToolHandler)ToolHandler{returnfunc(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error){// 创建带超时的子 contexttoolCtx,cancel:context.WithTimeout(ctx,timeout)defercancel()// 通过 channel 桥接同步调用typeresultstruct{res*protocol.CallToolResult errerror}done:make(chanresult,1)gofunc(){deferfunc(){ifr:recover();r!nil{done-result{err:fmt.Errorf(tool panic: %v,r)}}}()res,err:next(toolCtx,args)done-result{res:res,err:err}}()select{case-toolCtx.Done():// 超时 或 父 context 取消returnnil,fmt.Errorf(tool call timeout after %v: %w,timeout,toolCtx.Err())caser:-done:returnr.res,r.err}}}}超时的坑goroutine 泄漏上面的实现在超时后goroutine 中的next()仍在执行但结果被丢弃。如果工具内部有长时间运行的操作会造成 goroutine 泄漏。解决方案让工具本身尊重ctx.Done()// 工具实现必须检查 ctx.Done()funcslowDBQuery(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error){// 正确做法在每次阻塞操作前检查 contextselect{case-ctx.Done():returnnil,ctx.Err()default:}// 使用支持 context 的 APIrows,err:db.QueryContext(ctx,sql,args...)// ...}关键原则中间件的超时只是最后一层防线。真正可靠的超时需要工具内部配合ctx.Done()。五、中间件三熔断器CircuitBreakerMiddleware当某个工具连续失败时快速失败避免雪崩。// middleware/circuit_breaker.gopackagemiddlewareimport(contextfmtsyncsync/atomictimegithub.com/example/mcp-server-go/internal/protocol)// CircuitState 熔断器状态typeCircuitStateint32const(StateClosed CircuitStateiota// 正常StateOpen// 熔断中快速失败StateHalfOpen// 半开试探性放行)// CircuitBreakerConfig 熔断器配置typeCircuitBreakerConfigstruct{FailureThresholdint// 连续失败多少次后熔断SuccessThresholdint// 半开状态下多少次成功后恢复Timeout time.Duration// 熔断打开后多久进入半开HalfOpenMaxCallsint// 半开状态最多放行几个请求}funcDefaultCircuitBreakerConfig()CircuitBreakerConfig{returnCircuitBreakerConfig{FailureThreshold:5,SuccessThreshold:3,Timeout:30*time.Second,HalfOpenMaxCalls:1,}}// CircuitBreaker 熔断器实现typeCircuitBreakerstruct{cfg CircuitBreakerConfig state atomic.Int32// CircuitStateconsecutiveFail atomic.Int32// 连续失败计数consecutiveSuccess atomic.Int32// 连续成功计数半开状态用lastFailureTime time.Time halfOpenCalls atomic.Int32// 半开状态已放行数mu sync.Mutex}funcNewCircuitBreaker(cfg CircuitBreakerConfig)*CircuitBreaker{cb:CircuitBreaker{cfg:cfg}cb.state.Store(int32(StateClosed))returncb}// Allow 检查是否允许调用func(cb*CircuitBreaker)Allow()error{state:CircuitState(cb.state.Load())switchstate{caseStateClosed:returnnilcaseStateOpen:cb.mu.Lock()defercb.mu.Unlock()// 检查是否可以进入半开状态iftime.Since(cb.lastFailureTime)cb.cfg.Timeout{cb.state.Store(int32(StateHalfOpen))cb.consecutiveSuccess.Store(0)cb.halfOpenCalls.Store(0)returnnil}returnfmt.Errorf(circuit breaker is OPEN: too many failures, retry after %v,cb.cfg.Timeout-time.Since(cb.lastFailureTime).Round(time.Second))caseStateHalfOpen:// 限制半开状态的并发试探请求current:cb.halfOpenCalls.Add(1)ifcurrentint32(cb.cfg.HalfOpenMaxCalls){cb.halfOpenCalls.Add(-1)returnfmt.Errorf(circuit breaker is HALF-OPEN: max probe calls reached)}returnnildefault:returnnil}}// RecordSuccess 记录成功func(cb*CircuitBreaker)RecordSuccess(){state:CircuitState(cb.state.Load())switchstate{caseStateClosed:cb.consecutiveFail.Store(0)cb.consecutiveSuccess.Add(1)caseStateHalfOpen:success:cb.consecutiveSuccess.Add(1)ifint(success)cb.cfg.SuccessThreshold{cb.state.Store(int32(StateClosed))cb.consecutiveFail.Store(0)}cb.halfOpenCalls.Add(-1)}}// RecordFailure 记录失败func(cb*CircuitBreaker)RecordFailure(){cb.mu.Lock()defercb.mu.Unlock()state:CircuitState(cb.state.Load())cb.lastFailureTimetime.Now()switchstate{caseStateClosed:fails:cb.consecutiveFail.Add(1)ifint(fails)cb.cfg.FailureThreshold{cb.state.Store(int32(StateOpen))}caseStateHalfOpen:// 半开状态下的失败 → 立即重新熔断cb.state.Store(int32(StateOpen))cb.consecutiveFail.Store(int32(cb.cfg.FailureThreshold))cb.halfOpenCalls.Add(-1)}}// WithCircuitBreaker 创建熔断中间件funcWithCircuitBreaker(cb*CircuitBreaker)Middleware{returnfunc(next ToolHandler)ToolHandler{returnfunc(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error){iferr:cb.Allow();err!nil{returnnil,err}result,err:next(ctx,args)iferr!nil{cb.RecordFailure()}elseifresult!nil!result.IsError{cb.RecordSuccess()}else{// isErrortrue 也算失败cb.RecordFailure()}returnresult,err}}}六、组合使用洋葱模型// 实际用法packagemainimport(timegithub.com/example/mcp-server-go/internal/middlewaregithub.com/example/mcp-server-go/internal/protocol)funcmain(){// 原始工具处理器rawHandler:dbQueryHandler// 创建熔断器全局共享跨请求统计cb:middleware.NewCircuitBreaker(middleware.DefaultCircuitBreakerConfig())// 洋葱模型组合中间件// 执行顺序由外到内截断 → 超时 → 熔断 → 实际工具protectedHandler:middleware.Chain(rawHandler,middleware.WithTruncate(middleware.TruncateConfig{MaxChars:16000,// 约 4K tokens针对小模型TruncateMsg:\n\n[输出已截断],}),middleware.WithTimeout(15*time.Second),middleware.WithCircuitBreaker(cb),)// 将 protectedHandler 注册到 MCP Registry// registry.Register(tool, protectedHandler)}请求执行时序请求到达 │ ▼ [Truncate] ── 记录开始时间 │ ▼ [Timeout] ── ctx, cancel : context.WithTimeout(ctx, 15s) │ ▼ [CircuitBreaker] ── cb.Allow() → 状态检查 │ ├─ OPEN → 直接返回 circuit breaker open │ └─ CLOSED/HALF-OPEN → 继续 ▼ [实际工具] ── 执行业务逻辑 │ ▼ [CircuitBreaker] ◄── 记录成功/失败更新计数 │ ▼ [Timeout] ◄── cancel() 回收资源 │ ▼ [Truncate] ◄── 如果输出过长截断 │ ▼ 返回结果给 Client七、进阶可观测性中间件除了防御性中间件还可以加一层可观测性// middleware/observability.gopackagemiddlewareimport(contextlogtimegithub.com/example/mcp-server-go/internal/protocol)// WithMetrics 指标记录中间件funcWithMetrics(toolNamestring)Middleware{returnfunc(next ToolHandler)ToolHandler{returnfunc(ctx context.Context,argsmap[string]interface{})(*protocol.CallToolResult,error){start:time.Now()result,err:next(ctx,args)elapsed:time.Since(start)status:successiferr!nil{statuserror}// 结构化日志可接入 Prometheus / Lokilog.Printf([metrics] tool%s status%s duration%v args%v,toolName,status,elapsed,truncateArgs(args,200))// 生产环境emit metrics// toolCallDuration.WithLabelValues(toolName, status).Observe(elapsed.Seconds())// toolCallTotal.WithLabelValues(toolName, status).Inc()returnresult,err}}}functruncateArgs(argsmap[string]interface{},maxLenint)string{s:fmt.Sprintf(%v,args)iflen(s)maxLen{returns[:maxLen]...}returns}八、中间件对比总结中间件解决的问题适用场景性能开销Truncate输出过大撑爆上下文文件读取、数据库查询、API 调用低仅字符串操作Timeout工具卡死不返回网络调用、慢查询、外部 API低一个 goroutine channelCircuitBreaker连续失败雪崩外部依赖不可靠时极低原子操作 锁Metrics无感知问题发现滞后所有工具低日志 I/O 开销