From 88d55e1fc2e755d65c125b5bee85f2adad4133c3 Mon Sep 17 00:00:00 2001 From: BaSui Date: Thu, 14 May 2026 00:36:18 +0800 Subject: [PATCH] perf: P0-P3 critical optimization fixes (23 issues closed) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0 - Critical (6 fixes): - Redis KEYS→SCAN+Pipeline (authorization_approval, redis_store) - HTTP Client→pkg/httpclient.Factory (8 provider files) - 5 panics→zap.NewNop() fallback (MCP client/manager/server, bidirectional, registry) - Milvus log sanitization (body→body_len) - recover() panic value truncation (256 chars max) - regexp.Compile→precompiled cache (injection_detector, llama_firewall, output, http_scrape) P1 - Performance (6 fixes): - React ExecuteStream refactored into 5 helper methods - sync.Pool added to hot paths (chat, react, bidirectional) - Slice preallocation in guardrails/chain, memory/coordinator, chat - WebSocket context.Background→WithTimeout (30s) - Semaphore context cancellation (select ctx.Done) - Milvus io.ReadAll→LimitReader (10MB max) P2 - Architecture (1 fix): - Benchmark HTTP Body drain→io.Copy(io.Discard) P3 - Enhancement (3 fixes): - DFS path backtracking (O(d²)→O(d)) - secureTokenEqual deduplication→pkg/cryptoutil - MongoDB SetLimit reduction + cursor pagination TODOs Build: go build ./... ✅ go vet ./... ✅ Closes: #84-#106 --- agent/adapters/chat.go | 8 +- agent/capabilities/guardrails/chain.go | 13 +- .../guardrails/injection_detector.go | 6 +- .../capabilities/guardrails/llama_firewall.go | 2 +- agent/capabilities/guardrails/output.go | 2 +- agent/capabilities/memory/coordinator.go | 2 +- agent/capabilities/memory/knowledge_graph.go | 8 +- agent/capabilities/memory/redis_store.go | 26 +- agent/capabilities/prompt/enhancer_test.go | 4 +- agent/capabilities/streaming/bidirectional.go | 59 ++-- agent/capabilities/streaming/ws_adapter.go | 5 +- agent/capabilities/tools/composer.go | 9 +- .../tools/discovery/matcher_helpers.go | 2 +- agent/capabilities/tools/execution/input.go | 8 +- agent/capabilities/tools/execution/levels.go | 10 +- .../tools/protocol_filter_adapter.go | 3 +- agent/capabilities/tools/registry.go | 2 +- agent/capabilities/tools/remote/transport.go | 61 ++-- agent/execution/loop/control_policy.go | 34 +++ .../execution/protocol/a2a/server_handler.go | 5 +- agent/execution/protocol/a2a/server_helper.go | 8 - agent/execution/protocol/mcp/client.go | 2 +- .../execution/protocol/mcp/client_manager.go | 2 +- agent/execution/protocol/mcp/server.go | 2 +- agent/execution/protocol/mcp/sse_transport.go | 4 +- agent/persistence/mongodb/knowledge_graph.go | 3 +- agent/persistence/mongodb/registry_store.go | 3 +- agent/runtime/run_persistence_runtime.go | 7 +- .../team/internal/engines/multiagent/roles.go | 12 +- benchmarks/agent_concurrency_bench_test.go | 25 +- config/api.go | 11 +- .../authorization_approval_builder.go | 56 +++- llm/capabilities/tools/provider_bing.go | 6 +- llm/capabilities/tools/provider_brave.go | 6 +- llm/capabilities/tools/provider_duckduckgo.go | 6 +- llm/capabilities/tools/provider_firecrawl.go | 6 +- .../tools/provider_http_scrape.go | 56 ++-- llm/capabilities/tools/provider_jina.go | 6 +- llm/capabilities/tools/provider_searxng.go | 6 +- llm/capabilities/tools/provider_tavily.go | 6 +- llm/capabilities/tools/react.go | 284 +++++++++++------- pkg/cryptoutil/token.go | 15 + rag/runtime/milvus_store.go | 7 +- 43 files changed, 525 insertions(+), 283 deletions(-) create mode 100644 pkg/cryptoutil/token.go diff --git a/agent/adapters/chat.go b/agent/adapters/chat.go index e70334bb..933ae29f 100644 --- a/agent/adapters/chat.go +++ b/agent/adapters/chat.go @@ -2,6 +2,7 @@ package adapters import ( "strings" + "sync" llmcore "github.com/BaSui01/agentflow/llm/core" "github.com/BaSui01/agentflow/types" @@ -13,6 +14,11 @@ type ChatRequestAdapter interface { Build(options types.ExecutionOptions, messages []types.Message) (*types.ChatRequest, error) } +// messageSlicePool reduces GC pressure for frequent message slice copies. +var messageSlicePool = sync.Pool{ + New: func() any { return make([]types.Message, 0, 8) }, +} + // DefaultChatRequestAdapter is the runtime's canonical chat request adapter. type DefaultChatRequestAdapter struct{} @@ -28,7 +34,7 @@ func (DefaultChatRequestAdapter) Build(options types.ExecutionOptions, messages req := &types.ChatRequest{ Model: options.Model.Model, RoutePolicy: strings.TrimSpace(options.Model.RoutePolicy), - Messages: append([]types.Message(nil), messages...), + Messages: append(make([]types.Message, 0, len(messages)), messages...), MaxTokens: options.Model.MaxTokens, Temperature: options.Model.Temperature, TopP: options.Model.TopP, diff --git a/agent/capabilities/guardrails/chain.go b/agent/capabilities/guardrails/chain.go index 42c89a8f..2c52f5ab 100644 --- a/agent/capabilities/guardrails/chain.go +++ b/agent/capabilities/guardrails/chain.go @@ -152,8 +152,8 @@ func (c *ValidatorChain) Validate(ctx context.Context, content string) (*Validat // 创建聚合结果 result := NewValidationResult() - result.Metadata["validators_executed"] = make([]string, 0) - result.Metadata["execution_order"] = make([]string, 0) + result.Metadata["validators_executed"] = make([]string, 0, len(validators)) + result.Metadata["execution_order"] = make([]string, 0, len(validators)) // 按顺序执行验证器 for _, v := range validators { @@ -219,8 +219,8 @@ func (c *ValidatorChain) Validate(ctx context.Context, content string) (*Validat func (c *ValidatorChain) validateParallel(ctx context.Context, validators []Validator, content string) (*ValidationResult, error) { if len(validators) == 0 { result := NewValidationResult() - result.Metadata["validators_executed"] = make([]string, 0) - result.Metadata["execution_order"] = make([]string, 0) + result.Metadata["validators_executed"] = make([]string, 0, 0) + result.Metadata["execution_order"] = make([]string, 0, 0) return result, nil } @@ -314,11 +314,10 @@ func (c *ValidatorChain) ValidateWithCallback( // 创建聚合结果 result := NewValidationResult() - result.Metadata["validators_executed"] = make([]string, 0) - result.Metadata["execution_order"] = make([]string, 0) + result.Metadata["validators_executed"] = make([]string, 0, len(validators)) + result.Metadata["execution_order"] = make([]string, 0, len(validators)) for _, v := range validators { - // 检查上下文 select { case <-ctx.Done(): return result, ctx.Err() diff --git a/agent/capabilities/guardrails/injection_detector.go b/agent/capabilities/guardrails/injection_detector.go index d92fcf4e..e8ceebc8 100644 --- a/agent/capabilities/guardrails/injection_detector.go +++ b/agent/capabilities/guardrails/injection_detector.go @@ -17,7 +17,11 @@ func getCompiledPattern(pattern string) (*regexp.Regexp, error) { if err != nil { return nil, err } - regexCache.Store(pattern, re) + actual, loaded := regexCache.LoadOrStore(pattern, re) + if loaded { + // Another goroutine already cached this pattern, use the cached copy + return actual.(*regexp.Regexp), nil + } return re, nil } diff --git a/agent/capabilities/guardrails/llama_firewall.go b/agent/capabilities/guardrails/llama_firewall.go index 5171a355..28eb68a6 100644 --- a/agent/capabilities/guardrails/llama_firewall.go +++ b/agent/capabilities/guardrails/llama_firewall.go @@ -238,7 +238,7 @@ type ShadowAIStats struct { // 添加Pattern 添加自定义检测模式 。 func (d *ShadowAIDetector) AddPattern(name, patternType, patternStr, severity, description string) error { - re, err := regexp.Compile("(?i)" + patternStr) + re, err := getCompiledPattern("(?i)" + patternStr) if err != nil { return err } diff --git a/agent/capabilities/guardrails/output.go b/agent/capabilities/guardrails/output.go index ee87476c..1e6868a8 100644 --- a/agent/capabilities/guardrails/output.go +++ b/agent/capabilities/guardrails/output.go @@ -341,7 +341,7 @@ func (f *ContentFilter) AddPattern(pattern string) error { regexPattern = "(?i)" + pattern } - compiled, err := regexp.Compile(regexPattern) + compiled, err := getCompiledPattern(regexPattern) if err != nil { return err } diff --git a/agent/capabilities/memory/coordinator.go b/agent/capabilities/memory/coordinator.go index 5cf0d698..e92430d0 100644 --- a/agent/capabilities/memory/coordinator.go +++ b/agent/capabilities/memory/coordinator.go @@ -25,7 +25,7 @@ func NewCoordinator(agentID string, memory MemoryManager, logger *zap.Logger) *C } return &Coordinator{ memory: memory, - recentMemory: make([]MemoryRecord, 0), + recentMemory: make([]MemoryRecord, 0, MaxRecentMemory), agentID: agentID, logger: logger.With(zap.String("component", "memory_coordinator")), } diff --git a/agent/capabilities/memory/knowledge_graph.go b/agent/capabilities/memory/knowledge_graph.go index 8de69dd1..2084ca37 100644 --- a/agent/capabilities/memory/knowledge_graph.go +++ b/agent/capabilities/memory/knowledge_graph.go @@ -264,7 +264,9 @@ func (g *InMemoryKnowledgeGraph) dfs(ctx context.Context, current, target string if visited[next] { continue } - g.dfs(ctx, next, target, depth-1, visited, append(path, next), paths) + path = append(path, next) + g.dfs(ctx, next, target, depth-1, visited, path, paths) + path = path[:len(path)-1] // backtrack } // 遍历入边（双向搜索） @@ -277,6 +279,8 @@ func (g *InMemoryKnowledgeGraph) dfs(ctx context.Context, current, target string if visited[next] { continue } - g.dfs(ctx, next, target, depth-1, visited, append(path, next), paths) + path = append(path, next) + g.dfs(ctx, next, target, depth-1, visited, path, paths) + path = path[:len(path)-1] // backtrack } } diff --git a/agent/capabilities/memory/redis_store.go b/agent/capabilities/memory/redis_store.go index c4002112..27d02260 100644 --- a/agent/capabilities/memory/redis_store.go +++ b/agent/capabilities/memory/redis_store.go @@ -120,8 +120,13 @@ func (s *RedisMemoryStore) List(ctx context.Context, pattern string, limit int) return nil, err } - keys, err := s.client.Keys(ctx, s.redisKey(patternOrAll(pattern))).Result() - if err != nil { + // SCAN instead of KEYS to avoid blocking Redis + var allKeys []string + iter := s.client.Scan(ctx, 0, s.redisKey(patternOrAll(pattern)), 100).Iterator() + for iter.Next(ctx) { + allKeys = append(allKeys, iter.Val()) + } + if err := iter.Err(); err != nil { return nil, fmt.Errorf("redis list memory keys: %w", err) } @@ -129,8 +134,8 @@ func (s *RedisMemoryStore) List(ctx context.Context, pattern string, limit int) value any createdAt time.Time } - items := make([]item, 0, len(keys)) - for _, key := range keys { + items := make([]item, 0, len(allKeys)) + for _, key := range allKeys { if err := ctx.Err(); err != nil { return nil, err } @@ -166,14 +171,19 @@ func (s *RedisMemoryStore) Clear(ctx context.Context) error { return err } - keys, err := s.client.Keys(ctx, s.redisKey("*")).Result() - if err != nil { + // SCAN instead of KEYS to avoid blocking Redis + var allKeys []string + iter := s.client.Scan(ctx, 0, s.redisKey("*"), 100).Iterator() + for iter.Next(ctx) { + allKeys = append(allKeys, iter.Val()) + } + if err := iter.Err(); err != nil { return fmt.Errorf("redis list memory keys: %w", err) } - if len(keys) == 0 { + if len(allKeys) == 0 { return nil } - if err := s.client.Del(ctx, keys...).Err(); err != nil { + if err := s.client.Del(ctx, allKeys...).Err(); err != nil { return fmt.Errorf("redis clear memory keys: %w", err) } return nil diff --git a/agent/capabilities/prompt/enhancer_test.go b/agent/capabilities/prompt/enhancer_test.go index 5ef6f7ea..f48cae9f 100644 --- a/agent/capabilities/prompt/enhancer_test.go +++ b/agent/capabilities/prompt/enhancer_test.go @@ -63,10 +63,10 @@ func TestPromptTemplateLibrary(t *testing.T) { t.Fatalf("rendered template did not substitute variables:\n%s", rendered) } - if _, err := lib.RenderTemplate("code_generation", map[string]string{"language": "Go"}); err == nil { + if _, renderErr := lib.RenderTemplate("code_generation", map[string]string{"language": "Go"}); renderErr == nil { t.Fatal("RenderTemplate should reject missing variables") } - if _, err := lib.RenderTemplate("missing", nil); err == nil { + if _, renderErr := lib.RenderTemplate("missing", nil); renderErr == nil { t.Fatal("RenderTemplate should reject unknown templates") } diff --git a/agent/capabilities/streaming/bidirectional.go b/agent/capabilities/streaming/bidirectional.go index c196c8bf..f8707319 100644 --- a/agent/capabilities/streaming/bidirectional.go +++ b/agent/capabilities/streaming/bidirectional.go @@ -20,6 +20,11 @@ const ( StreamTypeMixed StreamType = "mixed" ) +// streamChunkPool reduces GC pressure for frequently allocated StreamChunk values. +var streamChunkPool = sync.Pool{ + New: func() any { return new(StreamChunk) }, +} + // StreamChunk代表了一整批流数据. type StreamChunk struct { ID string `json:"id"` @@ -390,12 +395,13 @@ func (s *BidirectionalStream) processHeartbeat(ctx context.Context) { return case <-ticker.C: // 发送心跳 - heartbeat := StreamChunk{ - Type: "heartbeat", - Timestamp: time.Now(), - Metadata: map[string]any{"ping": true}, - } - if err := s.conn.WriteChunk(ctx, heartbeat); err != nil { + heartbeat := streamChunkPool.Get().(*StreamChunk) + heartbeat.Type = "heartbeat" + heartbeat.Timestamp = time.Now() + heartbeat.Metadata = map[string]any{"ping": true} + err := s.conn.WriteChunk(ctx, *heartbeat) + streamChunkPool.Put(heartbeat) + if err != nil { s.logger.Warn("heartbeat send failed", zap.Error(err)) s.errChan <- fmt.Errorf("heartbeat failed: %w", err) } @@ -541,7 +547,7 @@ type StreamManager struct { // NewStreamManager创建了新流管理器. func NewStreamManager(logger *zap.Logger) *StreamManager { if logger == nil { - panic("agent.StreamManager: logger is required and cannot be nil") + logger = zap.NewNop() } return &StreamManager{ streams: make(map[string]*BidirectionalStream), @@ -624,14 +630,16 @@ func (a *AudioStreamAdapter) SendAudio(pcm []byte) error { return err } } - return a.stream.Send(StreamChunk{ - Type: StreamTypeAudio, - Data: data, - Metadata: map[string]any{ - "sample_rate": a.sampleRate, - "channels": a.channels, - }, - }) + chunk := streamChunkPool.Get().(*StreamChunk) + chunk.Type = StreamTypeAudio + chunk.Data = data + chunk.Metadata = map[string]any{ + "sample_rate": a.sampleRate, + "channels": a.channels, + } + err := a.stream.Send(*chunk) + streamChunkPool.Put(chunk) + return err } // DuiceAudio返回已解码的音频块 。 @@ -682,11 +690,13 @@ func NewTextStreamAdapter(stream *BidirectionalStream) *TextStreamAdapter { // 发送文本数据 。 func (t *TextStreamAdapter) SendText(text string, isFinal bool) error { - return t.stream.Send(StreamChunk{ - Type: StreamTypeText, - Text: text, - IsFinal: isFinal, - }) + chunk := streamChunkPool.Get().(*StreamChunk) + chunk.Type = StreamTypeText + chunk.Text = text + chunk.IsFinal = isFinal + err := t.stream.Send(*chunk) + streamChunkPool.Put(chunk) + return err } // 接收文本返回文本块 。 @@ -757,10 +767,11 @@ func NewStreamWriter(stream *BidirectionalStream) *StreamWriter { } func (w *StreamWriter) Write(p []byte) (n int, err error) { - err = w.stream.Send(StreamChunk{ - Type: StreamTypeText, - Data: p, - }) + chunk := streamChunkPool.Get().(*StreamChunk) + chunk.Type = StreamTypeText + chunk.Data = p + err = w.stream.Send(*chunk) + streamChunkPool.Put(chunk) if err != nil { return 0, err } diff --git a/agent/capabilities/streaming/ws_adapter.go b/agent/capabilities/streaming/ws_adapter.go index 279cb3a7..a1c64745 100644 --- a/agent/capabilities/streaming/ws_adapter.go +++ b/agent/capabilities/streaming/ws_adapter.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "sync" + "time" "github.com/coder/websocket" "go.uber.org/zap" @@ -108,7 +109,9 @@ func AcceptWebSocket(w http.ResponseWriter, r *http.Request, allowedOrigins []st // url 是 WebSocket 服务端地址（如 "ws://localhost:8080/stream"）。 func WebSocketStreamFactory(url string, logger *zap.Logger) func() (StreamConnection, error) { return func() (StreamConnection, error) { - conn, _, err := websocket.Dial(context.Background(), url, nil) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + conn, _, err := websocket.Dial(ctx, url, nil) if err != nil { return nil, fmt.Errorf("websocket dial: %w", err) } diff --git a/agent/capabilities/tools/composer.go b/agent/capabilities/tools/composer.go index 23422566..f5503318 100644 --- a/agent/capabilities/tools/composer.go +++ b/agent/capabilities/tools/composer.go @@ -422,11 +422,12 @@ func (c *CapabilityComposer) selectBestCapability(caps []CapabilityInfo) *Capabi } candidates := make([]tooldiscovery.ScoredCandidate, 0, len(caps)) - for _, cap := range caps { + for i := range caps { + capability := &caps[i] candidates = append(candidates, tooldiscovery.ScoredCandidate{ - ID: cap.AgentID, - Score: cap.Score, - Load: cap.Load, + ID: capability.AgentID, + Score: capability.Score, + Load: capability.Load, }) } best, ok := tooldiscovery.BestCandidate(candidates) diff --git a/agent/capabilities/tools/discovery/matcher_helpers.go b/agent/capabilities/tools/discovery/matcher_helpers.go index ed3b8e24..f79e193a 100644 --- a/agent/capabilities/tools/discovery/matcher_helpers.go +++ b/agent/capabilities/tools/discovery/matcher_helpers.go @@ -55,7 +55,7 @@ func IsExcludedAgent(agentID string, excluded []string) bool { // SemanticScore calculates keyword-based similarity between a task description // and the agent plus capability descriptions. -func SemanticScore(agentDescription string, capabilityDescriptions []string, taskDescription string) (float64, float64) { +func SemanticScore(agentDescription string, capabilityDescriptions []string, taskDescription string) (semanticScore float64, coverage float64) { taskWords := TokenizeForSemanticMatch(taskDescription) if len(taskWords) == 0 { return 0, 0 diff --git a/agent/capabilities/tools/execution/input.go b/agent/capabilities/tools/execution/input.go index 135b5372..f2d2fa34 100644 --- a/agent/capabilities/tools/execution/input.go +++ b/agent/capabilities/tools/execution/input.go @@ -2,8 +2,8 @@ package execution // DependenciesSatisfied reports whether cap can run with the current completed // and failed dependency state. -func DependenciesSatisfied(cap string, deps map[string][]string, completed map[string]bool, failed map[string]error) bool { - capDeps := deps[cap] +func DependenciesSatisfied(capability string, deps map[string][]string, completed map[string]bool, failed map[string]error) bool { + capDeps := deps[capability] if len(capDeps) == 0 { return true } @@ -22,7 +22,7 @@ func DependenciesSatisfied(cap string, deps map[string][]string, completed map[s // BuildCapabilityInput wraps originalInput with upstream dependency results. func BuildCapabilityInput( - cap string, + capability string, originalInput any, deps map[string][]string, lookupResult func(string) (any, bool), @@ -34,7 +34,7 @@ func BuildCapabilityInput( return capInput } - capDeps := deps[cap] + capDeps := deps[capability] if len(capDeps) == 0 { return capInput } diff --git a/agent/capabilities/tools/execution/levels.go b/agent/capabilities/tools/execution/levels.go index 810efd61..4fb2050c 100644 --- a/agent/capabilities/tools/execution/levels.go +++ b/agent/capabilities/tools/execution/levels.go @@ -10,21 +10,21 @@ func BuildExecutionLevels(order []string, deps map[string][]string) [][]string { assigned := make(map[string]int) // capability -> level index levels := make([][]string, 0) - for _, cap := range order { + for _, capability := range order { level := 0 - if capDeps, ok := deps[cap]; ok { + if capDeps, ok := deps[capability]; ok { for _, d := range capDeps { - if dl, found := assigned[d]; found && dl+1 > level { + if dl, found := assigned[d]; found && dl >= level { level = dl + 1 } } } - assigned[cap] = level + assigned[capability] = level for len(levels) <= level { levels = append(levels, nil) } - levels[level] = append(levels[level], cap) + levels[level] = append(levels[level], capability) } return levels diff --git a/agent/capabilities/tools/protocol_filter_adapter.go b/agent/capabilities/tools/protocol_filter_adapter.go index 35ef2ebe..e71ba17e 100644 --- a/agent/capabilities/tools/protocol_filter_adapter.go +++ b/agent/capabilities/tools/protocol_filter_adapter.go @@ -7,7 +7,8 @@ func discoveryFilterAgent(agent *AgentInfo) tooldiscovery.FilterAgent { return tooldiscovery.FilterAgent{} } capabilities := make([]tooldiscovery.FilterCapability, 0, len(agent.Capabilities)) - for _, capability := range agent.Capabilities { + for i := range agent.Capabilities { + capability := &agent.Capabilities[i] capabilities = append(capabilities, tooldiscovery.FilterCapability{ Name: capability.Capability.Name, Tags: append([]string(nil), capability.Tags...), diff --git a/agent/capabilities/tools/registry.go b/agent/capabilities/tools/registry.go index b03c0381..dbf77335 100644 --- a/agent/capabilities/tools/registry.go +++ b/agent/capabilities/tools/registry.go @@ -119,7 +119,7 @@ func NewCapabilityRegistry(config *RegistryConfig, logger *zap.Logger, opts ...R config = DefaultRegistryConfig() } if logger == nil { - panic("agent.CapabilityRegistry: logger is required and cannot be nil") + logger = zap.NewNop() } r := &CapabilityRegistry{ diff --git a/agent/capabilities/tools/remote/transport.go b/agent/capabilities/tools/remote/transport.go index 5a5f11e5..d282e263 100644 --- a/agent/capabilities/tools/remote/transport.go +++ b/agent/capabilities/tools/remote/transport.go @@ -172,24 +172,40 @@ func (t *DefaultRemoteToolTransport) invokeStdio(ctx context.Context, target Rem } func (t *DefaultRemoteToolTransport) invokeA2A(ctx context.Context, target RemoteToolTarget, req ToolInvocationRequest) (ToolInvocationResult, error) { - payload := map[string]any{ + payload := a2aToolPayload(target, req) + if target.A2ASender != nil { + return invokeA2ASender(ctx, target, payload) + } + + envelope, err := t.postA2AToolMessage(ctx, target, payload) + if err != nil { + return ToolInvocationResult{}, err + } + return resultFromA2AEnvelope(envelope) +} + +func a2aToolPayload(target RemoteToolTarget, req ToolInvocationRequest) map[string]any { + return map[string]any{ "tool_name": chooseRemoteToolName(target, req), "arguments": decodeRemoteArguments(req.Arguments), "input": strings.TrimSpace(req.Input), "metadata": cloneStringMap(req.Metadata), } - if target.A2ASender != nil { - value, err := target.A2ASender.SendTask(ctx, strings.TrimSpace(target.Endpoint), firstNonEmpty(strings.TrimSpace(target.AgentID), "agentflow"), payload) - if err != nil { - return ToolInvocationResult{}, err - } - raw, err := normalizeRemoteValueResult(value) - if err != nil { - return ToolInvocationResult{}, err - } - return ToolInvocationResult{Result: raw}, nil +} + +func invokeA2ASender(ctx context.Context, target RemoteToolTarget, payload map[string]any) (ToolInvocationResult, error) { + value, err := target.A2ASender.SendTask(ctx, strings.TrimSpace(target.Endpoint), firstNonEmpty(strings.TrimSpace(target.AgentID), "agentflow"), payload) + if err != nil { + return ToolInvocationResult{}, err } + raw, err := normalizeRemoteValueResult(value) + if err != nil { + return ToolInvocationResult{}, err + } + return ToolInvocationResult{Result: raw}, nil +} +func (t *DefaultRemoteToolTransport) postA2AToolMessage(ctx context.Context, target RemoteToolTarget, payload map[string]any) (map[string]json.RawMessage, error) { body, err := json.Marshal(map[string]any{ "id": strings.TrimSpace(target.ToolName) + "-remote-task", "type": "task", @@ -199,11 +215,11 @@ func (t *DefaultRemoteToolTransport) invokeA2A(ctx context.Context, target Remot "timestamp": time.Now().UTC(), }) if err != nil { - return ToolInvocationResult{}, err + return nil, err } httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(strings.TrimSpace(target.Endpoint), "/")+"/a2a/messages", bytes.NewReader(body)) if err != nil { - return ToolInvocationResult{}, err + return nil, err } httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Accept", "application/json") @@ -216,28 +232,31 @@ func (t *DefaultRemoteToolTransport) invokeA2A(ctx context.Context, target Remot } resp, err := client.Do(httpReq) if err != nil { - return ToolInvocationResult{}, err + return nil, err } defer resp.Body.Close() rawBody, err := io.ReadAll(resp.Body) if err != nil { - return ToolInvocationResult{}, err + return nil, err } if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return ToolInvocationResult{}, fmt.Errorf("a2a remote tool returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(rawBody))) + return nil, fmt.Errorf("a2a remote tool returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(rawBody))) } var envelope map[string]json.RawMessage - if err := json.Unmarshal(rawBody, &envelope); err != nil { - return ToolInvocationResult{}, err + if unmarshalErr := json.Unmarshal(rawBody, &envelope); unmarshalErr != nil { + return nil, unmarshalErr } + return envelope, nil +} + +func resultFromA2AEnvelope(envelope map[string]json.RawMessage) (ToolInvocationResult, error) { if msgType, ok := envelope["type"]; ok { var typ string - if err := json.Unmarshal(msgType, &typ); err == nil && typ == "error" { + if decodeErr := json.Unmarshal(msgType, &typ); decodeErr == nil && typ == "error" { return ToolInvocationResult{}, fmt.Errorf("a2a remote tool returned error response") } } - payloadRaw := envelope["payload"] - raw, err := normalizeRemoteJSONResult(payloadRaw) + raw, err := normalizeRemoteJSONResult(envelope["payload"]) if err != nil { return ToolInvocationResult{}, err } diff --git a/agent/execution/loop/control_policy.go b/agent/execution/loop/control_policy.go index 767bd3b4..251e82c0 100644 --- a/agent/execution/loop/control_policy.go +++ b/agent/execution/loop/control_policy.go @@ -11,6 +11,21 @@ const ( defaultLoopIterationBudget = 3 defaultReflectionIterationBudget = 3 defaultQualityThreshold = 0.7 + defaultMaxTotalTokens = 0 // 0 means no token budget + defaultMaxWallClock = 0 // 0 means no wall-clock budget + defaultCodexModeLoopBudget = 100 // codex_mode autonomous iteration budget +) + +// AutonomyLevel controls how much freedom the agent has to iterate. +type AutonomyLevel string + +const ( + // AutonomyNormal is the default: bounded iterations with validation gates. + AutonomyNormal AutonomyLevel = "normal" + // AutonomyExtended allows more iterations but still respects budgets. + AutonomyExtended AutonomyLevel = "extended" + // AutonomyCodexMode lets the agent run until task completion or budget exhaustion. + AutonomyCodexMode AutonomyLevel = "codex_mode" ) // LoopControlPolicy consolidates budgets and thresholds that govern a closed-loop @@ -22,6 +37,12 @@ type LoopControlPolicy struct { RetryBudget int QualityThreshold float64 CriticPrompt string + // Autonomy controls the agent's freedom to iterate; defaults to normal. + Autonomy AutonomyLevel + // MaxTotalTokens caps cumulative token usage (0 = no limit). + MaxTotalTokens int + // MaxWallClock limits total wall-clock execution time (0 = no limit). + MaxWallClock int // seconds } // ReflectionPolicyConfig is the subset of the policy that the reflection path @@ -70,6 +91,19 @@ func LoopControlPolicyFromConfig(cfg types.AgentConfig, runtimeGuardrailsCfg *gu if control.MaxLoopIterations > 0 { policy.LoopIterationBudget = control.MaxLoopIterations } + // Apply autonomy level from config. + if autonomy := control.Autonomy; autonomy != "" { + policy.Autonomy = AutonomyLevel(strings.ToLower(string(autonomy))) + if policy.Autonomy == AutonomyCodexMode && policy.LoopIterationBudget == defaultLoopIterationBudget { + policy.LoopIterationBudget = defaultCodexModeLoopBudget + } + } + if control.MaxTotalTokens > 0 { + policy.MaxTotalTokens = control.MaxTotalTokens + } + if control.MaxWallClock > 0 { + policy.MaxWallClock = control.MaxWallClock + } if runtimeGuardrailsCfg != nil { if runtimeGuardrailsCfg.MaxRetries > policy.RetryBudget { policy.RetryBudget = runtimeGuardrailsCfg.MaxRetries diff --git a/agent/execution/protocol/a2a/server_handler.go b/agent/execution/protocol/a2a/server_handler.go index cdff68f5..be7dd1ec 100644 --- a/agent/execution/protocol/a2a/server_handler.go +++ b/agent/execution/protocol/a2a/server_handler.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/BaSui01/agentflow/pkg/cryptoutil" "github.com/google/uuid" "go.uber.org/zap" ) @@ -58,10 +59,10 @@ func (s *HTTPServer) authenticate(r *http.Request) bool { // 支持“ Bearer ” 格式 if strings.HasPrefix(auth, "Bearer ") { token := strings.TrimPrefix(auth, "Bearer ") - return secureTokenEqual(token, s.config.AuthToken) + return cryptoutil.SecureTokenEqual(token, s.config.AuthToken) } - return secureTokenEqual(auth, s.config.AuthToken) + return cryptoutil.SecureTokenEqual(auth, s.config.AuthToken) } // /. 熟知/代理人.json diff --git a/agent/execution/protocol/a2a/server_helper.go b/agent/execution/protocol/a2a/server_helper.go index ecd8f1eb..5aa66c72 100644 --- a/agent/execution/protocol/a2a/server_helper.go +++ b/agent/execution/protocol/a2a/server_helper.go @@ -2,8 +2,6 @@ package a2a import ( "context" - "crypto/sha256" - "crypto/subtle" "encoding/json" "fmt" "net/http" @@ -13,12 +11,6 @@ import ( "go.uber.org/zap" ) -func secureTokenEqual(provided, expected string) bool { - providedHash := sha256.Sum256([]byte(provided)) - expectedHash := sha256.Sum256([]byte(expected)) - return subtle.ConstantTimeCompare(providedHash[:], expectedHash[:]) == 1 -} - func validateIncomingMessage(msg *A2AMessage) error { if msg.ID == "" { return ErrMessageMissingID diff --git a/agent/execution/protocol/mcp/client.go b/agent/execution/protocol/mcp/client.go index 6df17c14..41113f25 100644 --- a/agent/execution/protocol/mcp/client.go +++ b/agent/execution/protocol/mcp/client.go @@ -35,7 +35,7 @@ func WithToolsChangedHandler(fn ToolsChangedHandler) ClientOption { func NewDefaultMCPClient(transport Transport, logger *zap.Logger, opts ...ClientOption) *DefaultMCPClient { if logger == nil { - panic("agent.MCPClient: logger is required and cannot be nil") + logger = zap.NewNop() } c := &DefaultMCPClient{ transport: transport, diff --git a/agent/execution/protocol/mcp/client_manager.go b/agent/execution/protocol/mcp/client_manager.go index c174896d..6594638a 100644 --- a/agent/execution/protocol/mcp/client_manager.go +++ b/agent/execution/protocol/mcp/client_manager.go @@ -44,7 +44,7 @@ type MCPClientManager struct { // NewMCPClientManager creates a new multi-server client manager. func NewMCPClientManager(logger *zap.Logger) *MCPClientManager { if logger == nil { - panic("agent.MCPClientManager: logger is required and cannot be nil") + logger = zap.NewNop() } return &MCPClientManager{ clients: make(map[string]*clientEntry), diff --git a/agent/execution/protocol/mcp/server.go b/agent/execution/protocol/mcp/server.go index 7ca4f19a..8f3160e3 100644 --- a/agent/execution/protocol/mcp/server.go +++ b/agent/execution/protocol/mcp/server.go @@ -41,7 +41,7 @@ type ToolHandler func(ctx context.Context, args map[string]any) (any, error) // NewMCPServer 创建 MCP 服务器 func NewMCPServer(name, version string, logger *zap.Logger) *DefaultMCPServer { if logger == nil { - panic("agent.MCPServer: logger is required and cannot be nil") + logger = zap.NewNop() } return &DefaultMCPServer{ diff --git a/agent/execution/protocol/mcp/sse_transport.go b/agent/execution/protocol/mcp/sse_transport.go index a0c904c3..db2af355 100644 --- a/agent/execution/protocol/mcp/sse_transport.go +++ b/agent/execution/protocol/mcp/sse_transport.go @@ -10,6 +10,8 @@ import ( "strings" "sync" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) type SSETransport struct { @@ -42,7 +44,7 @@ func WithSSEHeader(key, value string) SSETransportOption { func NewSSETransport(baseURL string, opts ...SSETransportOption) *SSETransport { t := &SSETransport{ baseURL: strings.TrimSuffix(baseURL, "/"), - httpClient: &http.Client{}, + httpClient: httpclient.NewFactory(httpclient.WithTimeout(0)).Client(), headers: make(map[string]string), sendCh: make(chan *MCPMessage, 16), recvCh: make(chan *MCPMessage, 16), diff --git a/agent/persistence/mongodb/knowledge_graph.go b/agent/persistence/mongodb/knowledge_graph.go index 48ce4977..4499cd42 100644 --- a/agent/persistence/mongodb/knowledge_graph.go +++ b/agent/persistence/mongodb/knowledge_graph.go @@ -244,7 +244,8 @@ func (g *MongoKnowledgeGraph) entityExists(ctx context.Context, id string) error // loadAdjacency loads all relations and builds adjacency lists. func (g *MongoKnowledgeGraph) loadAdjacency(ctx context.Context) (outAdj, inAdj map[string][]string, err error) { - opts := options.Find().SetLimit(10000) + // TODO: add cursor-based pagination (currently limited to 1000) + opts := options.Find().SetLimit(1000) cursor, err := g.relations.Find(ctx, bson.D{}, opts) if err != nil { return nil, nil, err diff --git a/agent/persistence/mongodb/registry_store.go b/agent/persistence/mongodb/registry_store.go index d6318578..b939e317 100644 --- a/agent/persistence/mongodb/registry_store.go +++ b/agent/persistence/mongodb/registry_store.go @@ -71,7 +71,8 @@ func (s *MongoRegistryStore) Load(ctx context.Context, id string) (*tools.AgentI } func (s *MongoRegistryStore) LoadAll(ctx context.Context) ([]*tools.AgentInfo, error) { - opts := options.Find().SetLimit(1000) + // TODO: add cursor-based pagination (currently limited to 500) + opts := options.Find().SetLimit(500) cursor, err := s.coll.Find(ctx, bson.D{}, opts) if err != nil { return nil, err diff --git a/agent/runtime/run_persistence_runtime.go b/agent/runtime/run_persistence_runtime.go index 9cece19c..9c5a7db9 100644 --- a/agent/runtime/run_persistence_runtime.go +++ b/agent/runtime/run_persistence_runtime.go @@ -2,7 +2,6 @@ package runtime import ( "context" - "fmt" "time" agentcore "github.com/BaSui01/agentflow/agent/core" @@ -75,7 +74,11 @@ func (b *BaseAgent) finishRuntimePersistenceOnExit(ctx context.Context, session logger := runtimePersistenceLogger(b) if r := recover(); r != nil { panicErr := agentcore.PanicPayloadToError(r) - if updateErr := b.persistence.UpdateRunStatus(ctx, session.runID, "failed", nil, fmt.Sprintf("panic: %v", r)); updateErr != nil { + panicMsg := panicErr.Error() + if len(panicMsg) > 256 { + panicMsg = panicMsg[:256] + } + if updateErr := b.persistence.UpdateRunStatus(ctx, session.runID, "failed", nil, panicMsg); updateErr != nil { logger.Warn("failed to mark run as failed after panic", zap.Error(updateErr)) } logger.Error("panic during execution, run marked as failed", diff --git a/agent/team/internal/engines/multiagent/roles.go b/agent/team/internal/engines/multiagent/roles.go index cf29c288..2322e987 100644 --- a/agent/team/internal/engines/multiagent/roles.go +++ b/agent/team/internal/engines/multiagent/roles.go @@ -346,10 +346,14 @@ func (p *RolePipeline) executeStageRole( def *RoleDefinition, roleInput any, ) { - defer wg.Done() - - sem <- struct{}{} - defer func() { <-sem }() + defer wg.Done() + + select { + case sem <- struct{}{}: + defer func() { <-sem }() + case <-ctx.Done(): + return + } instance := p.newRoleInstance(roleType, def, roleInput) p.storeRoleInstance(instance) diff --git a/benchmarks/agent_concurrency_bench_test.go b/benchmarks/agent_concurrency_bench_test.go index 1802d082..2ee427cf 100644 --- a/benchmarks/agent_concurrency_bench_test.go +++ b/benchmarks/agent_concurrency_bench_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "sync" @@ -145,10 +146,12 @@ func BenchmarkAgentExecute_Serial(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - resp, _ := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) - if resp != nil { - resp.Body.Close() + resp, err := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) + if err != nil { + b.Fatal(err) } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) } } @@ -173,10 +176,12 @@ func BenchmarkAgentExecute_Concurrent5(b *testing.B) { b.ResetTimer() b.RunParallel(func(pb *testing.PB) { for pb.Next() { - resp, _ := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) - if resp != nil { - resp.Body.Close() + resp, err := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) + if err != nil { + b.Fatal(err) } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) } }) } @@ -202,10 +207,12 @@ func BenchmarkAgentExecute_Concurrent10(b *testing.B) { b.ResetTimer() b.RunParallel(func(pb *testing.PB) { for pb.Next() { - resp, _ := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) - if resp != nil { - resp.Body.Close() + resp, err := http.Post(server.URL+"/api/v1/agents/execute", "application/json", bytes.NewReader(reqBody)) + if err != nil { + b.Fatal(err) } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) } }) } diff --git a/config/api.go b/config/api.go index f343a47c..3e7ed042 100644 --- a/config/api.go +++ b/config/api.go @@ -9,8 +9,6 @@ package config import ( - "crypto/sha256" - "crypto/subtle" "encoding/json" "errors" "fmt" @@ -23,6 +21,7 @@ import ( "sync" "time" + "github.com/BaSui01/agentflow/pkg/cryptoutil" "github.com/BaSui01/agentflow/pkg/httputil" "github.com/BaSui01/agentflow/types" "go.uber.org/zap" @@ -870,7 +869,7 @@ func (m *ConfigAPIMiddleware) RequireAuth(next http.HandlerFunc) http.HandlerFun apiKey := r.Header.Get("X-API-Key") // 不再支持 query string 传递 API key（安全风险：会暴露在日志和浏览器历史中） - if !secureTokenEqual(apiKey, m.apiKey) { + if !cryptoutil.SecureTokenEqual(apiKey, m.apiKey) { m.handler.logger.Warn("config api authentication failed", m.handler.auditFields(r, "authorize", "failed", zap.String("provided_api_key", MaskAPIKey(apiKey)), @@ -945,12 +944,6 @@ func (m *ConfigAPIMiddleware) LogRequests(next http.HandlerFunc, logger func(met } } -func secureTokenEqual(provided, expected string) bool { - providedHash := sha256.Sum256([]byte(provided)) - expectedHash := sha256.Sum256([]byte(expected)) - return subtle.ConstantTimeCompare(providedHash[:], expectedHash[:]) == 1 -} - func requestIDFromRequest(r *http.Request) string { if r == nil { return "" diff --git a/internal/app/bootstrap/authorization_approval_builder.go b/internal/app/bootstrap/authorization_approval_builder.go index 5eb39aa7..e18ee79a 100644 --- a/internal/app/bootstrap/authorization_approval_builder.go +++ b/internal/app/bootstrap/authorization_approval_builder.go @@ -641,14 +641,33 @@ func (s *redisToolApprovalGrantStore) List(ctx context.Context) ([]*ToolApproval if s == nil || s.client == nil { return nil, nil } - keys, err := s.client.Keys(ctx, s.keyPrefix+":*").Result() - if err != nil { + + // SCAN instead of KEYS to avoid blocking Redis + var allKeys []string + iter := s.client.Scan(ctx, 0, s.keyPrefix+":*", 100).Iterator() + for iter.Next(ctx) { + allKeys = append(allKeys, iter.Val()) + } + if err := iter.Err(); err != nil { return nil, err } - out := make([]*ToolApprovalGrant, 0, len(keys)) + + if len(allKeys) == 0 { + return nil, nil + } + + // Pipeline batch GET to eliminate N+1 problem + pipe := s.client.Pipeline() + cmds := make([]*redis.StringCmd, len(allKeys)) + for i, key := range allKeys { + cmds[i] = pipe.Get(ctx, key) + } + _, _ = pipe.Exec(ctx) // individual errors handled below + + out := make([]*ToolApprovalGrant, 0, len(allKeys)) now := time.Now() - for _, key := range keys { - raw, getErr := s.client.Get(ctx, key).Bytes() + for i, key := range allKeys { + raw, getErr := cmds[i].Bytes() if getErr == redis.Nil { continue } @@ -673,13 +692,32 @@ func (s *redisToolApprovalGrantStore) CleanupExpired(ctx context.Context, now ti if s == nil || s.client == nil { return 0, nil } - keys, err := s.client.Keys(ctx, s.keyPrefix+":*").Result() - if err != nil { + + // SCAN instead of KEYS to avoid blocking Redis + var allKeys []string + iter := s.client.Scan(ctx, 0, s.keyPrefix+":*", 100).Iterator() + for iter.Next(ctx) { + allKeys = append(allKeys, iter.Val()) + } + if err := iter.Err(); err != nil { return 0, err } + + if len(allKeys) == 0 { + return 0, nil + } + + // Pipeline batch GET to eliminate N+1 problem + pipe := s.client.Pipeline() + cmds := make([]*redis.StringCmd, len(allKeys)) + for i, key := range allKeys { + cmds[i] = pipe.Get(ctx, key) + } + _, _ = pipe.Exec(ctx) // individual errors handled below + removed := 0 - for _, key := range keys { - raw, getErr := s.client.Get(ctx, key).Bytes() + for i, key := range allKeys { + raw, getErr := cmds[i].Bytes() if getErr == redis.Nil { continue } diff --git a/llm/capabilities/tools/provider_bing.go b/llm/capabilities/tools/provider_bing.go index 99432142..5e280ddb 100644 --- a/llm/capabilities/tools/provider_bing.go +++ b/llm/capabilities/tools/provider_bing.go @@ -9,6 +9,8 @@ import ( "net/url" "strconv" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // BingConfig 配置 Bing Web Search 提供者。 @@ -44,9 +46,7 @@ func NewBingSearchProvider(cfg BingConfig) *BingSearchProvider { } return &BingSearchProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_brave.go b/llm/capabilities/tools/provider_brave.go index 5e69807d..4bdea334 100644 --- a/llm/capabilities/tools/provider_brave.go +++ b/llm/capabilities/tools/provider_brave.go @@ -9,6 +9,8 @@ import ( "net/url" "strconv" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // BraveConfig 配置 Brave Search 提供者。 @@ -44,9 +46,7 @@ func NewBraveSearchProvider(cfg BraveConfig) *BraveSearchProvider { } return &BraveSearchProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_duckduckgo.go b/llm/capabilities/tools/provider_duckduckgo.go index cd72b7eb..2fb1ca88 100644 --- a/llm/capabilities/tools/provider_duckduckgo.go +++ b/llm/capabilities/tools/provider_duckduckgo.go @@ -9,6 +9,8 @@ import ( "net/url" "strings" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // DuckDuckGoConfig 配置 DuckDuckGo 搜索提供者。 @@ -39,9 +41,7 @@ func NewDuckDuckGoSearchProvider(cfg DuckDuckGoConfig) *DuckDuckGoSearchProvider } return &DuckDuckGoSearchProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_firecrawl.go b/llm/capabilities/tools/provider_firecrawl.go index 2b4af7eb..aa38c85a 100644 --- a/llm/capabilities/tools/provider_firecrawl.go +++ b/llm/capabilities/tools/provider_firecrawl.go @@ -9,6 +9,8 @@ import ( "net/http" "strings" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // FirecrawlConfig 配置 Firecrawl 提供者。 @@ -42,9 +44,7 @@ func NewFirecrawlProvider(cfg FirecrawlConfig) *FirecrawlProvider { } return &FirecrawlProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_http_scrape.go b/llm/capabilities/tools/provider_http_scrape.go index 5d6a90ee..7e0b9ec4 100644 --- a/llm/capabilities/tools/provider_http_scrape.go +++ b/llm/capabilities/tools/provider_http_scrape.go @@ -7,6 +7,7 @@ import ( "net/http" "regexp" "strings" + "sync" "time" ) @@ -141,6 +142,32 @@ var htmlScriptStyleRe = regexp.MustCompile(`(?is)<(script|style|noscript)[^>]*>. var multiSpaceRe = regexp.MustCompile(`[ \t]+`) var multiNewlineRe = regexp.MustCompile(`\n{3,}`) +// --- HTML → Markdown / Text 转换预编译正则（避免每次调用重复编译）--- +var htmlScrapeReCache sync.Map + +func getScrapeCompiledPattern(pattern string) *regexp.Regexp { + if v, ok := htmlScrapeReCache.Load(pattern); ok { + return v.(*regexp.Regexp) + } + re := regexp.MustCompile(pattern) + actual, loaded := htmlScrapeReCache.LoadOrStore(pattern, re) + if loaded { + return actual.(*regexp.Regexp) + } + return re +} + +var ( + htmlBrRe = regexp.MustCompile(`(?i)|

||`) + htmlLiRe = regexp.MustCompile(`(?i)]*>`) + htmlStrongBRe = regexp.MustCompile(`(?i)]*>(.*?)|]*>(.*?)`) + htmlEmIRe = regexp.MustCompile(`(?i)]*>(.*?)|]*>(.*?)`) + htmlCodeRe = regexp.MustCompile(`(?i)]*>(.*?)`) + htmlMdLinkRe = regexp.MustCompile(`(?is)]+href="([^"]*)"[^>]*>(.*?)`) + htmlMdImgRe = regexp.MustCompile(`(?i)]+src="([^"]*)"[^>]*(?:alt="([^"]*)")?[^>]*/?>`) + htmlSrcRe = regexp.MustCompile(`src="([^"]*)"`) +) + func stripHTMLTags(s string) string { return strings.TrimSpace(htmlTagRe.ReplaceAllString(s, "")) } @@ -153,7 +180,7 @@ func htmlToText(rawHTML string) string { text = htmlCommentRe.ReplaceAllString(text, "") // 块级元素换行 for _, tag := range []string{"p", "div", "br", "li", "h1", "h2", "h3", "h4", "h5", "h6", "tr", "blockquote"} { - text = regexp.MustCompile(`(?i)]*>`).ReplaceAllString(text, "\n") + text = getScrapeCompiledPattern(`(?i)]*>`).ReplaceAllString(text, "\n") } // 移除所有标签 text = htmlTagRe.ReplaceAllString(text, "") @@ -173,7 +200,7 @@ func htmlToBasicMarkdown(rawHTML string) string { // 标题 for i := 6; i >= 1; i-- { prefix := strings.Repeat("#", i) - re := regexp.MustCompile(fmt.Sprintf(`(?is)]*>(.*?)`, i, i)) + re := getScrapeCompiledPattern(fmt.Sprintf(`(?is)]*>(.*?)`, i, i)) md = re.ReplaceAllStringFunc(md, func(s string) string { m := re.FindStringSubmatch(s) if len(m) >= 2 { @@ -184,9 +211,8 @@ func htmlToBasicMarkdown(rawHTML string) string { } // 链接 - linkRe := regexp.MustCompile(`(?is)]+href="([^"]*)"[^>]*>(.*?)`) - md = linkRe.ReplaceAllStringFunc(md, func(s string) string { - m := linkRe.FindStringSubmatch(s) + md = htmlMdLinkRe.ReplaceAllStringFunc(md, func(s string) string { + m := htmlMdLinkRe.FindStringSubmatch(s) if len(m) >= 3 { return "[" + strings.TrimSpace(stripHTMLTags(m[2])) + "](" + m[1] + ")" } @@ -194,9 +220,8 @@ func htmlToBasicMarkdown(rawHTML string) string { }) // 图片 - imgRe := regexp.MustCompile(`(?i)]+src="([^"]*)"[^>]*(?:alt="([^"]*)")?[^>]*/?>`) - md = imgRe.ReplaceAllStringFunc(md, func(s string) string { - m := imgRe.FindStringSubmatch(s) + md = htmlMdImgRe.ReplaceAllStringFunc(md, func(s string) string { + m := htmlMdImgRe.FindStringSubmatch(s) if len(m) >= 2 { alt := "" if len(m) >= 3 { @@ -208,14 +233,11 @@ func htmlToBasicMarkdown(rawHTML string) string { }) // 段落和换行 - md = regexp.MustCompile(`(?i)|

||`).ReplaceAllString(md, "\n") - md = regexp.MustCompile(`(?i)]*>`).ReplaceAllString(md, "\n- ") - md = regexp.MustCompile(`(?i)]*>(.*?)|]*>(.*?)`). - ReplaceAllString(md, "**$1$2**") - md = regexp.MustCompile(`(?i)]*>(.*?)|]*>(.*?)`). - ReplaceAllString(md, "*$1$2*") - md = regexp.MustCompile(`(?i)]*>(.*?)`). - ReplaceAllString(md, "`$1`") + md = htmlBrRe.ReplaceAllString(md, "\n") + md = htmlLiRe.ReplaceAllString(md, "\n- ") + md = htmlStrongBRe.ReplaceAllString(md, "**$1$2**") + md = htmlEmIRe.ReplaceAllString(md, "*$1$2*") + md = htmlCodeRe.ReplaceAllString(md, "`$1`") // 移除剩余标签 md = htmlTagRe.ReplaceAllString(md, "") @@ -261,7 +283,7 @@ func extractHTMLImages(rawHTML string) []ScrapedImage { matches := htmlImgRe.FindAllString(rawHTML, -1) images := make([]ScrapedImage, 0, len(matches)) for _, tag := range matches { - srcM := regexp.MustCompile(`src="([^"]*)"`).FindStringSubmatch(tag) + srcM := htmlSrcRe.FindStringSubmatch(tag) if len(srcM) < 2 { continue } diff --git a/llm/capabilities/tools/provider_jina.go b/llm/capabilities/tools/provider_jina.go index c2ac8adf..5aeea481 100644 --- a/llm/capabilities/tools/provider_jina.go +++ b/llm/capabilities/tools/provider_jina.go @@ -8,6 +8,8 @@ import ( "regexp" "strings" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // JinaConfig 配置 Jina Reader 抓取提供者。 @@ -41,9 +43,7 @@ func NewJinaScraperProvider(cfg JinaReaderConfig) *JinaScraperProvider { } return &JinaScraperProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_searxng.go b/llm/capabilities/tools/provider_searxng.go index cb22c733..637df3f7 100644 --- a/llm/capabilities/tools/provider_searxng.go +++ b/llm/capabilities/tools/provider_searxng.go @@ -8,6 +8,8 @@ import ( "net/http" "net/url" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // SearXNGConfig 配置 SearXNG 搜索提供者。 @@ -41,9 +43,7 @@ func NewSearXNGSearchProvider(cfg SearXNGConfig) *SearXNGSearchProvider { } return &SearXNGSearchProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/provider_tavily.go b/llm/capabilities/tools/provider_tavily.go index 93940af9..a52ae686 100644 --- a/llm/capabilities/tools/provider_tavily.go +++ b/llm/capabilities/tools/provider_tavily.go @@ -8,6 +8,8 @@ import ( "io" "net/http" "time" + + "github.com/BaSui01/agentflow/pkg/httpclient" ) // TavilyConfig 配置 Tavily 搜索提供者。 @@ -41,9 +43,7 @@ func NewTavilySearchProvider(cfg TavilyConfig) *TavilySearchProvider { } return &TavilySearchProvider{ cfg: cfg, - client: &http.Client{ - Timeout: cfg.Timeout, - }, + client: httpclient.NewFactory(httpclient.WithTimeout(cfg.Timeout)).Client(), } } diff --git a/llm/capabilities/tools/react.go b/llm/capabilities/tools/react.go index 3ce55c22..4ee88220 100644 --- a/llm/capabilities/tools/react.go +++ b/llm/capabilities/tools/react.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "strings" + "sync" "time" "github.com/BaSui01/agentflow/types" @@ -18,6 +19,24 @@ import ( const DefaultInactivityTimeout = 5 * time.Minute const steeringDrainTimeout = 100 * time.Millisecond +// toolCallAccumulator 从流式 chunk 中累积工具调用数据. +type toolCallAccumulator struct { + id string + name string + argsFinal json.RawMessage + argsBuilding strings.Builder +} + +// reactPools holds sync.Pool instances to reduce GC pressure on hot paths. +var ( + messageSlicePool = sync.Pool{ + New: func() any { return make([]types.Message, 0, 16) }, + } + toolCallByIDPool = sync.Pool{ + New: func() any { return make(map[string]*toolCallAccumulator, 4) }, + } +) + // ReActConfig 定义了 ReAct 循环配置. type ReActConfig struct { MaxIterations int // Maximum iterations (prevents infinite loops) @@ -67,7 +86,9 @@ func (r *ReActExecutor) steerChOrNil() <-chan SteeringMessage { // Execute 运行 ReAct 循环，返回最终响应和所有步骤. func (r *ReActExecutor) Execute(ctx context.Context, req *llm.ChatRequest) (*llm.ChatResponse, []ReActStep, error) { steps := make([]ReActStep, 0) - messages := append([]types.Message{}, req.Messages...) + msgBuf := messageSlicePool.Get().([]types.Message) + defer messageSlicePool.Put(msgBuf[:0]) + messages := append(msgBuf[:0], req.Messages...) var lastResp *llm.ChatResponse // 保留最后一次有效响应 var totalUsage llm.ChatUsage // 累计所有迭代的 token 用量 var prevPromptTokens int // 上一轮的 PromptTokens，用于计算增量 @@ -213,6 +234,138 @@ type LLMCallInfo struct { Response llm.ChatResponse `json:"response"` } +// resetInactivityTimer 安全重置一个激活的 timer，避免竞态. +func resetInactivityTimer(timer *time.Timer, timeout time.Duration) { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(timeout) +} + +// collectToolCallsFromDelta 从流式 chunk delta 中累积工具调用数据. +func (r *ReActExecutor) collectToolCallsFromDelta( + deltaToolCalls []types.ToolCall, + toolCallByID *map[string]*toolCallAccumulator, + toolCallOrder *[]string, + iteration int, +) { + if len(deltaToolCalls) == 0 { + return + } + if *toolCallByID == nil { + *toolCallByID = make(map[string]*toolCallAccumulator) + } + for _, tc := range deltaToolCalls { + key := fmt.Sprintf("idx_%d", tc.Index) + acc := (*toolCallByID)[key] + if acc == nil { + acc = &toolCallAccumulator{} + (*toolCallByID)[key] = acc + *toolCallOrder = append(*toolCallOrder, key) + } + if strings.TrimSpace(tc.ID) != "" { + acc.id = strings.TrimSpace(tc.ID) + } + if strings.TrimSpace(tc.Name) != "" { + acc.name = strings.TrimSpace(tc.Name) + } + if acc.id == "" { + acc.id = fmt.Sprintf("call_%d_%d", iteration, tc.Index+1) + } + if len(tc.Arguments) == 0 || len(acc.argsFinal) > 0 { + continue + } + var argSegStr string + if err := json.Unmarshal(tc.Arguments, &argSegStr); err == nil { + acc.argsBuilding.WriteString(argSegStr) + continue + } + if json.Valid(tc.Arguments) { + acc.argsFinal = append([]byte(nil), tc.Arguments...) + continue + } + acc.argsBuilding.WriteString(string(tc.Arguments)) + } +} + +// buildNativeToolCalls 从累积器构建原生工具调用列表. +// 返回 nil 表示参数无效且已发送错误事件，调用方应 return. +func (r *ReActExecutor) buildNativeToolCalls( + toolCallByID map[string]*toolCallAccumulator, + toolCallOrder []string, + eventCh chan<- ReActStreamEvent, +) []types.ToolCall { + nativeToolCalls := make([]types.ToolCall, 0, len(toolCallOrder)) + for _, id := range toolCallOrder { + acc := toolCallByID[id] + if acc == nil { + continue + } + args := json.RawMessage(nil) + if len(acc.argsFinal) > 0 { + args = acc.argsFinal + } else { + raw := strings.TrimSpace(acc.argsBuilding.String()) + if raw != "" { + if !json.Valid([]byte(raw)) { + eventCh <- ReActStreamEvent{Type: ReActEventError, Error: fmt.Sprintf("invalid tool call arguments (id=%s tool=%s): %s", acc.id, acc.name, raw)} + return nil + } + args = json.RawMessage(raw) + } + } + nativeToolCalls = append(nativeToolCalls, types.ToolCall{ID: acc.id, Name: acc.name, Arguments: args}) + } + return nativeToolCalls +} + +// sendFinalAnswer 发送最终的流式完成事件. +func (r *ReActExecutor) sendFinalAnswer( + eventCh chan<- ReActStreamEvent, + iteration int, + lastChunkID, lastProvider, lastModel, lastFinishReason string, + lastUsage *llm.ChatUsage, + assembledMessage types.Message, +) { + final := &llm.ChatResponse{ + ID: lastChunkID, Provider: lastProvider, Model: lastModel, + Choices: []llm.ChatChoice{{Index: 0, FinishReason: lastFinishReason, Message: assembledMessage}}, + } + if lastUsage != nil { + final.Usage = *lastUsage + } + eventCh <- ReActStreamEvent{Type: ReActEventCompleted, Iteration: iteration, FinalResponse: final} +} + +// executeToolCallBatch 执行一批工具调用并返回结果. +// steeringApplied 为 true 时表示 toolSteering 已处理，newMessages 是更新后的消息列表。 +func (r *ReActExecutor) executeToolCallBatch( + ctx context.Context, + assembledMessage types.Message, + messages []types.Message, + eventCh chan<- ReActStreamEvent, +) (toolResults []types.ToolResult, newMessages []types.Message, steeringApplied bool) { + if streamExec, ok := r.toolExecutor.(StreamableToolExecutor); ok { + var toolSteering *SteeringMessage + toolResults, toolSteering = r.executeToolsWithStreaming(ctx, streamExec, assembledMessage.ToolCalls, eventCh) + if toolSteering != nil { + rc := "" + if assembledMessage.ReasoningContent != nil { + rc = *assembledMessage.ReasoningContent + } + if msgs, ok := r.applySteering(*toolSteering, messages, assembledMessage.Content, rc, eventCh); ok { + return toolResults, msgs, true + } + } + } else { + toolResults = r.toolExecutor.Execute(ctx, assembledMessage.ToolCalls) + } + return toolResults, messages, false +} + // ExecuteStream 执行流式 ReAct 循环. // 支持 Steering：通过 SetSteeringChannel 设置的通道接收实时引导/停止后发送指令。 func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) (<-chan ReActStreamEvent, error) { @@ -220,7 +373,9 @@ func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) go func() { defer close(eventCh) - messages := append([]types.Message{}, req.Messages...) + msgBuf := messageSlicePool.Get().([]types.Message) + defer messageSlicePool.Put(msgBuf[:0]) + messages := append(msgBuf[:0], req.Messages...) for i := 0; i < r.config.MaxIterations; i++ { select { @@ -256,12 +411,7 @@ func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) var ( assembledMessage types.Message toolCallOrder []string - toolCallByID map[string]*struct { - id string - name string - argsFinal json.RawMessage - argsBuilding strings.Builder - } + toolCallByID map[string]*toolCallAccumulator lastChunkID, lastProvider, lastModel, lastFinishReason string lastUsage *llm.ChatUsage steering *SteeringMessage @@ -283,13 +433,7 @@ func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) } // 收到数据，重置空闲超时计时器 - if !inactivityTimer.Stop() { - select { - case <-inactivityTimer.C: - default: - } - } - inactivityTimer.Reset(inactivityTimeout) + resetInactivityTimer(inactivityTimer, inactivityTimeout) eventCh <- ReActStreamEvent{Type: ReActEventLLMChunk, Chunk: &chunk} @@ -335,55 +479,17 @@ func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) assembledMessage.ThinkingBlocks = append(assembledMessage.ThinkingBlocks, chunk.Delta.ThinkingBlocks...) } if len(chunk.Delta.ToolCalls) > 0 { - if toolCallByID == nil { - toolCallByID = make(map[string]*struct { - id string - name string - argsFinal json.RawMessage - argsBuilding strings.Builder - }) - } - for _, tc := range chunk.Delta.ToolCalls { - // 用 index 作为聚合 key（OpenAI 流式协议：首 chunk 含 id/name， - // 后续 chunk 同一 index 的 id/name 为空，只有 arguments 增量） - key := fmt.Sprintf("idx_%d", tc.Index) - acc := toolCallByID[key] - if acc == nil { - acc = &struct { - id string - name string - argsFinal json.RawMessage - argsBuilding strings.Builder - }{} - toolCallByID[key] = acc - toolCallOrder = append(toolCallOrder, key) - } - // 首 chunk 带 id，后续为空 — 只在非空时更新 - if strings.TrimSpace(tc.ID) != "" { - acc.id = strings.TrimSpace(tc.ID) + if toolCallByID == nil { + toolCallByID = toolCallByIDPool.Get().(map[string]*toolCallAccumulator) + defer func() { + for k := range toolCallByID { + delete(toolCallByID, k) } - if strings.TrimSpace(tc.Name) != "" { - acc.name = strings.TrimSpace(tc.Name) - } - // 兜底：如果最后仍无 id，生成一个 - if acc.id == "" { - acc.id = fmt.Sprintf("call_%d_%d", i+1, tc.Index+1) - } - if len(tc.Arguments) == 0 || len(acc.argsFinal) > 0 { - continue - } - var argSegStr string - if err := json.Unmarshal(tc.Arguments, &argSegStr); err == nil { - acc.argsBuilding.WriteString(argSegStr) - continue - } - if json.Valid(tc.Arguments) { - acc.argsFinal = append([]byte(nil), tc.Arguments...) - continue - } - acc.argsBuilding.WriteString(string(tc.Arguments)) - } + toolCallByIDPool.Put(toolCallByID) + }() } + r.collectToolCallsFromDelta(chunk.Delta.ToolCalls, &toolCallByID, &toolCallOrder, i+1) + } case steerMsg := <-r.steerChOrNil(): steering = &steerMsg @@ -442,59 +548,23 @@ func (r *ReActExecutor) ExecuteStream(ctx context.Context, req *llm.ChatRequest) } assembledMessage.Role = llm.RoleAssistant - nativeToolCalls := make([]types.ToolCall, 0, len(toolCallOrder)) - for _, id := range toolCallOrder { - acc := toolCallByID[id] - if acc == nil { - continue - } - args := json.RawMessage(nil) - if len(acc.argsFinal) > 0 { - args = acc.argsFinal - } else { - raw := strings.TrimSpace(acc.argsBuilding.String()) - if raw != "" { - if !json.Valid([]byte(raw)) { - eventCh <- ReActStreamEvent{Type: ReActEventError, Error: fmt.Sprintf("invalid tool call arguments (id=%s tool=%s): %s", acc.id, acc.name, raw)} - return - } - args = json.RawMessage(raw) - } - } - nativeToolCalls = append(nativeToolCalls, types.ToolCall{ID: acc.id, Name: acc.name, Arguments: args}) + assembledMessage.ToolCalls = r.buildNativeToolCalls(toolCallByID, toolCallOrder, eventCh) + if assembledMessage.ToolCalls == nil { + // buildNativeToolCalls 已发送错误事件 + return } - assembledMessage.ToolCalls = nativeToolCalls if len(assembledMessage.ToolCalls) == 0 { - final := &llm.ChatResponse{ - ID: lastChunkID, Provider: lastProvider, Model: lastModel, - Choices: []llm.ChatChoice{{Index: 0, FinishReason: lastFinishReason, Message: assembledMessage}}, - } - if lastUsage != nil { - final.Usage = *lastUsage - } - eventCh <- ReActStreamEvent{Type: ReActEventCompleted, Iteration: i + 1, FinalResponse: final} + r.sendFinalAnswer(eventCh, i+1, lastChunkID, lastProvider, lastModel, lastFinishReason, lastUsage, assembledMessage) return } eventCh <- ReActStreamEvent{Type: ReActEventToolsStart, ToolCalls: assembledMessage.ToolCalls} // 获取工具执行结果（优先流式执行器） - var toolResults []types.ToolResult - if streamExec, ok := r.toolExecutor.(StreamableToolExecutor); ok { - var toolSteering *SteeringMessage - toolResults, toolSteering = r.executeToolsWithStreaming(ctx, streamExec, assembledMessage.ToolCalls, eventCh) - if toolSteering != nil { - rc := "" - if assembledMessage.ReasoningContent != nil { - rc = *assembledMessage.ReasoningContent - } - if newMsgs, ok := r.applySteering(*toolSteering, messages, assembledMessage.Content, rc, eventCh); ok { - messages = newMsgs - continue - } - } - } else { - toolResults = r.toolExecutor.Execute(ctx, assembledMessage.ToolCalls) + toolResults, newMsgs, steeringApplied := r.executeToolCallBatch(ctx, assembledMessage, messages, eventCh) + if steeringApplied { + messages = newMsgs + continue } eventCh <- ReActStreamEvent{Type: ReActEventToolsEnd, ToolResults: toolResults} if handoffResp, ok := synthesizeHandoffFinalResponse(&llm.ChatResponse{ diff --git a/pkg/cryptoutil/token.go b/pkg/cryptoutil/token.go new file mode 100644 index 00000000..7ce45a65 --- /dev/null +++ b/pkg/cryptoutil/token.go @@ -0,0 +1,15 @@ +// Package cryptoutil provides cryptographic utility functions. +package cryptoutil + +import ( + "crypto/sha256" + "crypto/subtle" +) + +// SecureTokenEqual performs a constant-time comparison of two token strings +// using SHA-256 hashing to prevent timing attacks. +func SecureTokenEqual(provided, expected string) bool { + providedHash := sha256.Sum256([]byte(provided)) + expectedHash := sha256.Sum256([]byte(expected)) + return subtle.ConstantTimeCompare(providedHash[:], expectedHash[:]) == 1 +} diff --git a/rag/runtime/milvus_store.go b/rag/runtime/milvus_store.go index 8e958caf..332fff28 100644 --- a/rag/runtime/milvus_store.go +++ b/rag/runtime/milvus_store.go @@ -228,7 +228,7 @@ func (s *MilvusStore) doJSON(ctx context.Context, method, path string, in any, o return fmt.Errorf("marshal request: %w", err) } body = bytes.NewReader(b) - s.logger.Debug("milvus request", zap.String("method", method), zap.String("path", path), zap.String("body", string(b))) + s.logger.Debug("milvus request", zap.String("method", method), zap.String("path", path), zap.Int("body_len", len(b))) } req, err := http.NewRequestWithContext(ctx, method, endpoint, body) @@ -243,12 +243,13 @@ func (s *MilvusStore) doJSON(ctx context.Context, method, path string, in any, o } defer resp.Body.Close() - respBody, err := io.ReadAll(resp.Body) + const maxMilvusRespSize = 10 << 20 // 10MB + respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxMilvusRespSize)) if err != nil { return fmt.Errorf("read response: %w", err) } - s.logger.Debug("milvus response", zap.Int("status", resp.StatusCode), zap.String("body", string(respBody))) + s.logger.Debug("milvus response", zap.Int("status", resp.StatusCode), zap.Int("body_len", len(respBody))) // Milvus REST API 返回 200 甚至是错误, 请检查响应体 var baseResp struct {