mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-29 06:00:52 +02:00
Fix finalization cleanup for pending tool executions
This commit is contained in:
+22
-16
@@ -698,18 +698,19 @@ func (h *AgentHandler) mergeAssistantMessagePartialOnCancel(messageID, partial s
|
||||
|
||||
// ChatResponse 聊天响应
|
||||
type ChatResponse struct {
|
||||
Response string `json:"response"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
||||
ConversationID string `json:"conversationId"` // 对话ID
|
||||
Time time.Time `json:"time"`
|
||||
Finalizable bool `json:"finalizable"`
|
||||
Finalized bool `json:"finalized"`
|
||||
Status string `json:"status,omitempty"`
|
||||
CompletionReason string `json:"completionReason,omitempty"`
|
||||
EvidenceVerified bool `json:"evidenceVerified"`
|
||||
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||
Response string `json:"response"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
||||
ConversationID string `json:"conversationId"` // 对话ID
|
||||
Time time.Time `json:"time"`
|
||||
Finalizable bool `json:"finalizable"`
|
||||
Finalized bool `json:"finalized"`
|
||||
Status string `json:"status,omitempty"`
|
||||
CompletionReason string `json:"completionReason,omitempty"`
|
||||
EvidenceVerified bool `json:"evidenceVerified"`
|
||||
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||
AutoCancelledPendingExecutionIDs []string `json:"autoCancelledPendingExecutionIds,omitempty"`
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) {
|
||||
@@ -724,8 +725,13 @@ func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMes
|
||||
return "", conversationID, errMA
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput), true)
|
||||
func (h *AgentHandler) finalizeRobotAgentSuccess(taskCtx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||
reasoningContent := multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput)
|
||||
decision := h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, true)
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, true)
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, "robot", resultMA.MCPExecutionIDs, reasoningContent, decision)
|
||||
responseText := decision.FinalText
|
||||
if !decision.Finalizable {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
@@ -758,7 +764,7 @@ func (h *AgentHandler) runRobotEinoSingleWithRetry(
|
||||
*taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
return h.finalizeRobotAgentSuccess(taskCtx, assistantMessageID, conversationID, resultMA)
|
||||
}
|
||||
|
||||
func (h *AgentHandler) runRobotMultiAgentWithRetry(
|
||||
@@ -779,7 +785,7 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry(
|
||||
*taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
return h.finalizeRobotAgentSuccess(taskCtx, assistantMessageID, conversationID, resultMA)
|
||||
}
|
||||
|
||||
// ProcessMessageForRobot 供机器人(企业微信/钉钉/飞书)调用:Eino 单/多代理执行路径(含 progressCallback、过程详情),仅不发送 SSE,最后返回完整回复
|
||||
|
||||
@@ -281,7 +281,12 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
if useBatchMulti {
|
||||
agentMode = "batch_eino_" + batchOrch
|
||||
}
|
||||
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, reasoningContent, true)
|
||||
decision := h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, true)
|
||||
autoCancelledPendingExecutionIDs := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback)
|
||||
if len(autoCancelledPendingExecutionIDs) > 0 {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, true)
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, mcpIDs, reasoningContent, decision)
|
||||
resText := decision.FinalText
|
||||
if !decision.Finalizable {
|
||||
resText = finalizationBlockedMessage(decision)
|
||||
@@ -289,14 +294,15 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
sendEvent("finalization_check", resText, decision)
|
||||
}
|
||||
sendEvent("response", resText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": mcpIDs,
|
||||
"batchQueueId": queueID,
|
||||
"batchTaskId": task.ID,
|
||||
"batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable],
|
||||
"candidatePreview": safeTruncateString(resultMA.Response, 500),
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": mcpIDs,
|
||||
"batchQueueId": queueID,
|
||||
"batchTaskId": task.ID,
|
||||
"batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable],
|
||||
"candidatePreview": safeTruncateString(resultMA.Response, 500),
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
|
||||
if assistantMessageID == "" {
|
||||
|
||||
@@ -76,6 +76,65 @@ func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testi
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsPageUsesPersistedExecutionStatusAfterBackgroundCancel(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-cancelled.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
conversation, err := db.CreateConversation("cancelled background", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
execID := "exec-cancelled-after-background"
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "exec", "toolCallId": "call-cancelled", "index": 1, "total": 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_result", "background", map[string]interface{}{
|
||||
"toolName": "exec", "toolCallId": "call-cancelled", "executionId": execID, "status": "background_running", "success": true,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: execID,
|
||||
ToolName: "exec",
|
||||
Status: mcp.ToolExecutionStatusCancelled,
|
||||
StartTime: now,
|
||||
EndTime: &now,
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveToolExecution: %v", err)
|
||||
}
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?limit=10&offset=0", nil)
|
||||
c.Params = gin.Params{{Key: "id", Value: message.ID}}
|
||||
NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("status = %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var response struct {
|
||||
ToolExecutions []database.ProcessDetailsToolExecution `json:"toolExecutions"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if len(response.ToolExecutions) != 1 {
|
||||
t.Fatalf("tool executions = %d, want 1", len(response.ToolExecutions))
|
||||
}
|
||||
if got := response.ToolExecutions[0].Status; got != mcp.ToolExecutionStatusCancelled {
|
||||
t.Fatalf("tool execution status = %q, want cancelled", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsFullBackfillsEmptyToolCallArgumentsFromExecution(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-args.db"), zap.NewNop())
|
||||
|
||||
@@ -192,6 +192,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
var emptyResponseContinueAttempt int
|
||||
var finalizationAutoContinueAttempt int
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
|
||||
for {
|
||||
segmentMainIterationMax := 0
|
||||
@@ -268,6 +269,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
@@ -384,6 +389,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
|
||||
if decision.CompletionReason == "" {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, "eino_single", cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||
|
||||
@@ -401,10 +410,11 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
}
|
||||
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
}
|
||||
@@ -464,6 +474,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
var emptyResponseContinueAttempt int
|
||||
var finalizationAutoContinueAttempt int
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
for {
|
||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||
taskCtx,
|
||||
@@ -493,6 +504,10 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, prep.ConversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||
continue
|
||||
}
|
||||
@@ -509,18 +524,19 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"response": responseText,
|
||||
"conversationId": prep.ConversationID,
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"assistantMessageId": prep.AssistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"finalized": decision.Finalized,
|
||||
"finalizable": decision.Finalizable,
|
||||
"status": decision.Status,
|
||||
"completionReason": decision.CompletionReason,
|
||||
"evidenceVerified": decision.EvidenceVerified,
|
||||
"evidenceRefs": decision.EvidenceRefs,
|
||||
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||
"missingChecks": decision.MissingChecks,
|
||||
"response": responseText,
|
||||
"conversationId": prep.ConversationID,
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"assistantMessageId": prep.AssistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"finalized": decision.Finalized,
|
||||
"finalizable": decision.Finalizable,
|
||||
"status": decision.Status,
|
||||
"completionReason": decision.CompletionReason,
|
||||
"evidenceVerified": decision.EvidenceVerified,
|
||||
"evidenceRefs": decision.EvidenceRefs,
|
||||
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||
"missingChecks": decision.MissingChecks,
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2,16 +2,21 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/agentfinalizer"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const finalizationAutoContinueMaxAttempts = 2
|
||||
const finalizationPendingToolCancelWait = 2 * time.Second
|
||||
const finalizationPendingToolCancelPoll = 50 * time.Millisecond
|
||||
const finalizationPendingToolCancelNote = "Agent 迭代已结束,最终回复前自动终止未完成的工具执行"
|
||||
|
||||
func shouldAutoContinueAfterFinalization(d agentfinalizer.Decision, attempt int) bool {
|
||||
if d.Finalizable || d.Finalized {
|
||||
@@ -75,3 +80,105 @@ func finalizationAutoContinueBackoff(attempt int) time.Duration {
|
||||
}
|
||||
return time.Duration(attempt) * time.Second
|
||||
}
|
||||
|
||||
func (h *AgentHandler) cleanupPendingToolExecutionsAfterIteration(
|
||||
taskCtx context.Context,
|
||||
conversationID string,
|
||||
decision agentfinalizer.Decision,
|
||||
progressCallback func(eventType, message string, data interface{}),
|
||||
) []string {
|
||||
if h == nil || h.agent == nil || decision.CompletionReason != agentfinalizer.ReasonPendingTools {
|
||||
return nil
|
||||
}
|
||||
pending := uniqueNonEmptyStrings(decision.PendingExecutionIDs)
|
||||
if len(pending) == 0 {
|
||||
return nil
|
||||
}
|
||||
cancelled := make([]string, 0, len(pending))
|
||||
for _, executionID := range pending {
|
||||
if h.agent.CancelMCPToolExecutionWithNote(executionID, finalizationPendingToolCancelNote) {
|
||||
cancelled = append(cancelled, executionID)
|
||||
} else if h.logger != nil {
|
||||
h.logger.Warn("finalization pending tool cleanup could not cancel execution",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("executionId", executionID))
|
||||
}
|
||||
}
|
||||
if len(cancelled) == 0 {
|
||||
return nil
|
||||
}
|
||||
if progressCallback != nil {
|
||||
progressCallback("finalization_pending_tools_cancelled", "迭代结束,已自动终止仍在运行的工具执行。", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "finalizer",
|
||||
"autoCancelledPendingExecutionIds": cancelled,
|
||||
"pendingExecutionIds": pending,
|
||||
"reason": agentfinalizer.ReasonPendingTools,
|
||||
})
|
||||
}
|
||||
h.waitForToolExecutionsToLeavePending(taskCtx, cancelled, finalizationPendingToolCancelWait)
|
||||
return cancelled
|
||||
}
|
||||
|
||||
func (h *AgentHandler) waitForToolExecutionsToLeavePending(ctx context.Context, executionIDs []string, wait time.Duration) {
|
||||
if h == nil || h.db == nil || len(executionIDs) == 0 || wait <= 0 {
|
||||
return
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
defer timer.Stop()
|
||||
ticker := time.NewTicker(finalizationPendingToolCancelPoll)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
if !h.hasPendingToolExecutions(executionIDs) {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-contextDone(ctx):
|
||||
return
|
||||
case <-timer.C:
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hasPendingToolExecutions(executionIDs []string) bool {
|
||||
if h == nil || h.db == nil {
|
||||
return false
|
||||
}
|
||||
for _, executionID := range uniqueNonEmptyStrings(executionIDs) {
|
||||
exec, err := h.db.GetToolExecution(executionID)
|
||||
if err != nil || exec == nil {
|
||||
continue
|
||||
}
|
||||
switch strings.TrimSpace(exec.Status) {
|
||||
case mcp.ToolExecutionStatusQueued, mcp.ToolExecutionStatusRunning:
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func uniqueNonEmptyStrings(values []string) []string {
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
out := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[value]; ok {
|
||||
continue
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
out = append(out, value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func contextDone(ctx context.Context) <-chan struct{} {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
return ctx.Done()
|
||||
}
|
||||
|
||||
@@ -1,9 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
agentpkg "cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/agentfinalizer"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestShouldAutoContinueAfterFinalization(t *testing.T) {
|
||||
@@ -57,3 +66,66 @@ func TestRequestRequiresExecutionEvidenceUsesExplicitPolicyOnly(t *testing.T) {
|
||||
t.Fatal("explicit false policy should not require execution evidence")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanupPendingToolExecutionsAfterIterationAllowsFinalization(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "cleanup-finalization.db"), logger)
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
server := mcp.NewServerWithStorage(logger, db)
|
||||
server.ConfigureToolWaitTimeoutSeconds(1)
|
||||
server.RegisterTool(mcp.Tool{Name: "block", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
ag := agentpkg.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 10)
|
||||
h := &AgentHandler{agent: ag, db: db, logger: logger}
|
||||
|
||||
callCtx := mcp.WithMCPConversationID(context.Background(), "conv-cleanup")
|
||||
result, execID, err := server.CallTool(callCtx, "block", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || execID == "" {
|
||||
t.Fatalf("expected background wait result, result=%#v execID=%q", result, execID)
|
||||
}
|
||||
|
||||
decision := agentfinalizer.Decide(db, agentfinalizer.Input{
|
||||
Response: "基于已完成信息的阶段性总结。",
|
||||
MCPExecutionIDs: []string{execID},
|
||||
})
|
||||
if decision.CompletionReason != agentfinalizer.ReasonPendingTools {
|
||||
t.Fatalf("decision reason = %s, want pending tools: %+v", decision.CompletionReason, decision)
|
||||
}
|
||||
|
||||
var eventType string
|
||||
cancelled := h.cleanupPendingToolExecutionsAfterIteration(context.Background(), "conv-cleanup", decision, func(et, _ string, _ interface{}) {
|
||||
eventType = et
|
||||
})
|
||||
if len(cancelled) != 1 || cancelled[0] != execID {
|
||||
t.Fatalf("cancelled = %#v, want [%s]", cancelled, execID)
|
||||
}
|
||||
if eventType != "finalization_pending_tools_cancelled" {
|
||||
t.Fatalf("event type = %q", eventType)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
exec, err := db.GetToolExecution(execID)
|
||||
if err == nil && exec != nil && exec.Status == mcp.ToolExecutionStatusCancelled {
|
||||
after := agentfinalizer.Decide(db, agentfinalizer.Input{
|
||||
Response: "基于已完成信息的阶段性总结。",
|
||||
MCPExecutionIDs: []string{execID},
|
||||
})
|
||||
if !after.Finalizable || !after.Finalized {
|
||||
t.Fatalf("decision should finalize after cleanup: %+v", after)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("execution did not become cancelled")
|
||||
}
|
||||
|
||||
@@ -205,6 +205,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
agentMode := "eino_" + effectiveOrch
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
|
||||
for {
|
||||
segmentMainIterationMax := 0
|
||||
@@ -282,6 +283,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
@@ -398,6 +403,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
if decision.CompletionReason == "" {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||
|
||||
@@ -415,10 +424,11 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
}
|
||||
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
}
|
||||
@@ -478,6 +488,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
}
|
||||
agentMode := "eino_" + effectiveOrch
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
for {
|
||||
result, runErr = multiagent.RunDeepAgent(
|
||||
taskCtx,
|
||||
@@ -514,6 +525,10 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, prep.ConversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||
continue
|
||||
}
|
||||
@@ -533,18 +548,19 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
}
|
||||
c.JSON(http.StatusOK, ChatResponse{
|
||||
Response: responseText,
|
||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||
ConversationID: prep.ConversationID,
|
||||
Time: time.Now(),
|
||||
Finalizable: decision.Finalizable,
|
||||
Finalized: decision.Finalized,
|
||||
Status: decision.Status,
|
||||
CompletionReason: decision.CompletionReason,
|
||||
EvidenceVerified: decision.EvidenceVerified,
|
||||
EvidenceRefs: decision.EvidenceRefs,
|
||||
PendingExecutionIDs: decision.PendingExecutionIDs,
|
||||
MissingChecks: decision.MissingChecks,
|
||||
Response: responseText,
|
||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||
ConversationID: prep.ConversationID,
|
||||
Time: time.Now(),
|
||||
Finalizable: decision.Finalizable,
|
||||
Finalized: decision.Finalized,
|
||||
Status: decision.Status,
|
||||
CompletionReason: decision.CompletionReason,
|
||||
EvidenceVerified: decision.EvidenceVerified,
|
||||
EvidenceRefs: decision.EvidenceRefs,
|
||||
PendingExecutionIDs: decision.PendingExecutionIDs,
|
||||
MissingChecks: decision.MissingChecks,
|
||||
AutoCancelledPendingExecutionIDs: autoCancelledPendingExecutionIDs,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user