From 5cbd828cad0fa18408295a8fe2edac69dd77442c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:52:16 +0800 Subject: [PATCH] Add files via upload --- internal/database/conversation.go | 5 +- .../database/process_details_summary_test.go | 18 +++- internal/mcp/external_manager.go | 91 ++++++++++--------- internal/mcp/external_manager_test.go | 12 ++- internal/mcp/server.go | 34 +++---- internal/mcp/server_authorization_test.go | 13 ++- 6 files changed, 103 insertions(+), 70 deletions(-) diff --git a/internal/database/conversation.go b/internal/database/conversation.go index 958c1f29..d3e9ab0c 100644 --- a/internal/database/conversation.go +++ b/internal/database/conversation.go @@ -1424,7 +1424,8 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary seenExecIDs := make(map[string]bool) // A provider may reuse a fallback toolCallId across streaming rounds. Keep a // FIFO per ID instead of a single index so every persisted call gets at most - // one result. Results without an ID fall back to the oldest unmatched call. + // one result. Results without a stable ID are kept separate instead of being + // guessed by order; showing no link is safer than linking to the wrong tool. toolIndexesByCallID := make(map[string][]int) lastMatchedToolIndexByCallID := make(map[string]int) matchedToolIndexes := make([]bool, 0) @@ -1499,7 +1500,7 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary } } } - if idx < 0 { + if idx < 0 && toolCallID != "" { for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] { nextUnmatchedToolIdx++ } diff --git a/internal/database/process_details_summary_test.go b/internal/database/process_details_summary_test.go index 0ff8d9ab..39510da1 100644 --- a/internal/database/process_details_summary_test.go +++ b/internal/database/process_details_summary_test.go @@ -7,7 +7,7 @@ import ( "go.uber.org/zap" ) -func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) { +func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) { db, conversationID, messageID := setupProcessDetailsSummaryTest(t) for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} { if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{ @@ -32,14 +32,24 @@ func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) if err != nil { t.Fatalf("GetProcessDetailsSummary: %v", err) } - if len(summary.ToolExecutions) != 4 { - t.Fatalf("tool executions = %d, want 4", len(summary.ToolExecutions)) + if len(summary.ToolExecutions) != 6 { + t.Fatalf("tool executions = %d, want 6", len(summary.ToolExecutions)) } - for i, execution := range summary.ToolExecutions { + for i, execution := range summary.ToolExecutions[:2] { if execution.Status != "completed" { t.Fatalf("execution %d status = %q, want completed", i, execution.Status) } } + for i, execution := range summary.ToolExecutions[2:4] { + if execution.Status != "result_missing" { + t.Fatalf("unmatched call %d status = %q, want result_missing", i, execution.Status) + } + } + for i, execution := range summary.ToolExecutions[4:] { + if execution.Status != "completed" || execution.ToolCallID != "" { + t.Fatalf("idless result %d = %#v, want separate completed result without toolCallId", i, execution) + } + } } func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) { diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go index c6bd9b53..f51b0457 100644 --- a/internal/mcp/external_manager.go +++ b/internal/mcp/external_manager.go @@ -674,48 +674,6 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) { // CallTool 调用外部MCP工具（返回执行ID） func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { - _, authenticated := authctx.PrincipalFromContext(ctx) - m.mu.RLock() - authorizer := m.toolAuthorizer - m.mu.RUnlock() - if authorizer != nil { - if err := authorizer(ctx, toolName, args); err != nil { - return nil, "", fmt.Errorf("external tool authorization denied: %w", err) - } - } else if authenticated { - return nil, "", fmt.Errorf("external tool authorization policy is not configured") - } - // 解析工具名称：name::toolName - var mcpName, actualToolName string - if idx := findSubstring(toolName, "::"); idx > 0 { - mcpName = toolName[:idx] - actualToolName = toolName[idx+2:] - } else { - return nil, "", fmt.Errorf("无效的工具名称格式: %s", toolName) - } - - client, exists := m.GetClient(mcpName) - if !exists { - return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName) - } - if err := m.checkExternalMCPCircuit(mcpName); err != nil { - return nil, "", err - } - - // 检查连接状态，如果未连接或状态为error，不允许调用 - if !client.IsConnected() { - status := client.GetStatus() - if status == "error" { - // 获取错误信息（如果有） - errorMsg := m.GetError(mcpName) - if errorMsg != "" { - return nil, "", fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg) - } - return nil, "", fmt.Errorf("外部MCP连接失败: %s", mcpName) - } - return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status) - } - if m.executionService == nil { m.executionService = NewExecutionService(m.storage, m.logger) m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes) @@ -725,12 +683,57 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args if principal, ok := authctx.PrincipalFromContext(ctx); ok { ownerUserID = principal.UserID } + var mcpName, actualToolName string + var client ExternalMCPClient handle, err := m.executionService.Submit(ctx, ExecutionRequest{ ToolName: toolName, Arguments: args, ConversationID: MCPConversationIDFromContext(ctx), OwnerUserID: ownerUserID, PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + m.mu.RLock() + authorizer := m.toolAuthorizer + m.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("external tool authorization denied: %w", err) + } + } else if authenticated { + return nil, fmt.Errorf("external tool authorization policy is not configured") + } + + // 解析工具名称：name::toolName + if idx := findSubstring(toolName, "::"); idx > 0 { + mcpName = toolName[:idx] + actualToolName = toolName[idx+2:] + } else { + return nil, fmt.Errorf("无效的工具名称格式: %s", toolName) + } + + var exists bool + client, exists = m.GetClient(mcpName) + if !exists { + return nil, fmt.Errorf("外部MCP客户端不存在: %s", mcpName) + } + if err := m.checkExternalMCPCircuit(mcpName); err != nil { + return nil, err + } + + // 检查连接状态，如果未连接或状态为error，不允许调用 + if !client.IsConnected() { + status := client.GetStatus() + if status == "error" { + // 获取错误信息（如果有） + errorMsg := m.GetError(mcpName) + if errorMsg != "" { + return nil, fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg) + } + return nil, fmt.Errorf("外部MCP连接失败: %s", mcpName) + } + return nil, fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status) + } + release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName) if acquireErr != nil { return nil, acquireErr @@ -746,7 +749,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args }, OnDone: func(exec *ToolExecution) { failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled - m.recordExternalMCPResult(mcpName, failed) + if mcpName != "" { + m.recordExternalMCPResult(mcpName, failed) + } m.updateStats(toolName, failed) }, }) diff --git a/internal/mcp/external_manager_test.go b/internal/mcp/external_manager_test.go index 2ccf722b..3baff567 100644 --- a/internal/mcp/external_manager_test.go +++ b/internal/mcp/external_manager_test.go @@ -20,10 +20,20 @@ func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) { return errors.New("denied by policy") }) ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true})) - _, _, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{}) + _, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{}) if err == nil || !strings.Contains(err.Error(), "authorization denied") { t.Fatalf("external call bypassed authorizer: %v", err) } + if executionID == "" { + t.Fatal("denied external call should still return an execution id") + } + execution, ok := manager.GetExecution(executionID) + if !ok || execution == nil { + t.Fatalf("missing denied external execution %q", executionID) + } + if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") { + t.Fatalf("denied external execution = %#v, want failed with policy error", execution) + } } func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) { diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 0ed8abb3..3760183b 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -895,25 +895,6 @@ func (s *Server) GetAllTools() []Tool { // CallTool 直接调用工具（用于内部调用） func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { - _, authenticated := authctx.PrincipalFromContext(ctx) - s.mu.RLock() - authorizer := s.toolAuthorizer - s.mu.RUnlock() - if authorizer != nil { - if err := authorizer(ctx, toolName, args); err != nil { - return nil, "", fmt.Errorf("tool authorization denied: %w", err) - } - } else if authenticated { - return nil, "", errors.New("tool authorization policy is not configured") - } - s.mu.RLock() - handler, exists := s.tools[toolName] - s.mu.RUnlock() - - if !exists { - return nil, "", fmt.Errorf("工具 %s 未找到", toolName) - } - if s.executionService == nil { s.executionService = NewExecutionService(s.storage, s.logger) s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes) @@ -929,6 +910,21 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string] ConversationID: MCPConversationIDFromContext(ctx), OwnerUserID: ownerUserID, Run: func(runCtx context.Context) (*ToolResult, error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + handler, exists := s.tools[toolName] + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("tool authorization denied: %w", err) + } + } else if authenticated { + return nil, errors.New("tool authorization policy is not configured") + } + if !exists { + return nil, fmt.Errorf("工具 %s 未找到", toolName) + } return handler(runCtx, args) }, OnDone: func(exec *ToolExecution) { diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go index 90f0d1a9..bea489a3 100644 --- a/internal/mcp/server_authorization_test.go +++ b/internal/mcp/server_authorization_test.go @@ -23,9 +23,20 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) { } return nil }) - if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil { + _, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil) + if err == nil { t.Fatal("tool call without principal was allowed") } + if deniedExecutionID == "" { + t.Fatal("denied tool call should still return an execution id") + } + deniedExecution, ok := server.GetExecution(deniedExecutionID) + if !ok || deniedExecution == nil { + t.Fatalf("missing denied execution %q", deniedExecutionID) + } + if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") { + t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution) + } ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true})) _, executionID, err := server.CallTool(ctx, "echo", nil) if err != nil {