From d13a4c9a48be66db84fd631cab549599f92cfaa3 Mon Sep 17 00:00:00 2001 From: liangxu Date: Wed, 8 Jul 2026 12:06:50 +0800 Subject: [PATCH] fix(agent): let stopped threads start a new task and scope thread state to the latest task Replace the stored cancel func with a handle so a stale deferred cleanup can't delete a newer task's cancellation, and clear a thread's cancelled flag when it submits fresh work instead of refusing to restart. Thread progress/finished state now derives from the latest user task in the thread rather than the whole history. Adds coverage for continuing a stopped thread and for sanitized generation-failure messages. Co-Authored-By: Claude Opus 4.8 (1M context) --- server/internal/application/design_service.go | 58 ++++++++++++++---- .../application/generation_failure_test.go | 61 ++++++++++++++++--- server/internal/logic/project_events_logic.go | 43 ++++++++++++- .../logic/project_events_logic_test.go | 12 ++++ 4 files changed, 151 insertions(+), 23 deletions(-) diff --git a/server/internal/application/design_service.go b/server/internal/application/design_service.go index b466ec8..0aa77ef 100644 --- a/server/internal/application/design_service.go +++ b/server/internal/application/design_service.go @@ -72,6 +72,10 @@ type DesignService struct { now func() time.Time } +type agentTaskCancelHandle struct { + cancel context.CancelFunc +} + type GenerateResult struct { Project design.Project Message design.Message @@ -537,14 +541,7 @@ func (s *DesignService) AgentChat(ctx context.Context, projectID string, req des if err != nil { return design.AgentThread{}, err } - if threadIsCancelled(project.Messages, threadID) || s.isAgentTaskCancelled(project.ID, threadID) { - return design.AgentThread{ - ThreadID: threadID, - Status: "cancelled", - ThreadIDType: threadIDType, - Project: project, - }, nil - } + s.clearAgentTaskCancellation(project.ID, threadID) now := s.now() actionID := newID() @@ -1700,7 +1697,7 @@ func (s *DesignService) isUserCancelledGeneration(ctx context.Context, projectID if getErr != nil { return false } - return threadIsCancelled(project.Messages, threadID) + return threadHasCancelledMessage(project.Messages, threadID) } func detachedUserContext(ctx context.Context) context.Context { @@ -1737,13 +1734,16 @@ func (s *DesignService) withAgentTaskCancellation(ctx context.Context, projectID return ctx, func() {} } taskCtx, cancel := context.WithCancel(ctx) - s.agentTaskCancels.Store(key, context.CancelFunc(cancel)) + handle := &agentTaskCancelHandle{cancel: cancel} + s.agentTaskCancels.Store(key, handle) if s.isAgentTaskCancelled(projectID, threadID) { cancel() } return taskCtx, func() { cancel() - s.agentTaskCancels.Delete(key) + if current, ok := s.agentTaskCancels.Load(key); ok && current == handle { + s.agentTaskCancels.Delete(key) + } } } @@ -1754,7 +1754,9 @@ func (s *DesignService) cancelAgentTask(projectID string, threadID string) { } s.cancelledAgentTasks.Store(key, s.now()) if cancel, ok := s.agentTaskCancels.Load(key); ok { - if cancelFunc, ok := cancel.(context.CancelFunc); ok { + if handle, ok := cancel.(*agentTaskCancelHandle); ok && handle != nil { + handle.cancel() + } else if cancelFunc, ok := cancel.(context.CancelFunc); ok { cancelFunc() } } @@ -1763,6 +1765,14 @@ func (s *DesignService) cancelAgentTask(projectID string, threadID string) { }) } +func (s *DesignService) clearAgentTaskCancellation(projectID string, threadID string) { + key := agentTaskKey(projectID, threadID) + if key == "" { + return + } + s.cancelledAgentTasks.Delete(key) +} + func (s *DesignService) isAgentTaskCancelled(projectID string, threadID string) bool { key := agentTaskKey(projectID, threadID) if key == "" { @@ -3122,6 +3132,30 @@ func (s *DesignService) applyIncrementalGeneratedNode(ctx context.Context, proje } func threadIsCancelled(messages []design.Message, threadID string) bool { + threadID = strings.TrimSpace(threadID) + if threadID == "" { + return false + } + for index := len(messages) - 1; index >= 0; index-- { + message := messages[index] + if strings.TrimSpace(message.ThreadID) != threadID { + continue + } + if strings.EqualFold(strings.TrimSpace(message.Status), "cancelled") { + return true + } + if strings.EqualFold(strings.TrimSpace(message.Role), "user") { + return false + } + switch strings.ToLower(strings.TrimSpace(message.Status)) { + case "running", "streaming", "submitted", "success", "failed": + return false + } + } + return false +} + +func threadHasCancelledMessage(messages []design.Message, threadID string) bool { threadID = strings.TrimSpace(threadID) if threadID == "" { return false diff --git a/server/internal/application/generation_failure_test.go b/server/internal/application/generation_failure_test.go index a5b03de..b95e4f4 100644 --- a/server/internal/application/generation_failure_test.go +++ b/server/internal/application/generation_failure_test.go @@ -64,6 +64,42 @@ func TestMarkGenerationFailedRemovesCanvasPlaceholders(t *testing.T) { } } +func TestMarkGenerationFailedDoesNotExposeInternalEndpoint(t *testing.T) { + ctx := context.Background() + repo := memory.NewProjectRepository() + service := NewDesignService(repo, nil, nil, nil, nil) + project := design.Project{ + ID: "project-safe-error", + Title: "Safe failure", + Brief: "prompt", + Status: design.StatusExploring, + UpdatedAt: time.Now(), + Nodes: []design.Node{ + {ID: "placeholder", Type: design.NodeTypeImage, Title: "生成中", Content: "prompt", Status: "generating"}, + }, + } + if err := repo.Save(ctx, project); err != nil { + t.Fatal(err) + } + + failure := errors.New(`gpt image 2 generation failed: Post "http://154.21.87.3:8080/v1/images/generations": context deadline exceeded (Client.Timeout exceeded while awaiting headers)`) + if err := service.markGenerationFailed(ctx, project.ID, failure); err != nil { + t.Fatal(err) + } + + saved, err := repo.Get(ctx, project.ID) + if err != nil { + t.Fatal(err) + } + content := saved.Messages[len(saved.Messages)-1].Content + if strings.Contains(content, "154.21.87.3") || strings.Contains(content, "/v1/images/generations") || strings.Contains(content, "Client.Timeout") { + t.Fatalf("expected sanitized user-facing failure, got %q", content) + } + if !strings.Contains(content, "图片生成耗时较久") { + t.Fatalf("expected friendly timeout message, got %q", content) + } +} + func TestMarkGenerationFailedKeepsRealImageNodeActionTarget(t *testing.T) { ctx := context.Background() repo := memory.NewProjectRepository() @@ -895,7 +931,7 @@ func TestCreateProjectGenerationPartialFailurePreservesSuccessfulImagesAndCloses if message.Name == "generate_media" && message.Status == "success" { foundArtifact = true } - if message.Role == "error" && message.Title == "图片生成部分失败" && strings.Contains(message.Content, "image 2") { + if message.Role == "error" && message.Title == "图片生成部分失败" && strings.Contains(message.Content, "已保留成功结果") { foundError = true } } @@ -1154,8 +1190,8 @@ func TestCreateProjectGenerationFailureUsesJobUserContext(t *testing.T) { t.Fatalf("expected failed project, got %s", saved.Status) } last := saved.Messages[len(saved.Messages)-1] - if last.Role != "error" || last.Content != "upstream image response timed out" { - t.Fatalf("expected original failure message for job user, got %#v", last) + if last.Role != "error" || strings.Contains(last.Content, "upstream") || !strings.Contains(last.Content, "画布内容已保留") { + t.Fatalf("expected safe failure message for job user, got %#v", last) } } @@ -1486,9 +1522,10 @@ func TestStopAgentTaskClearsAgentPlaceholdersWithoutStoppingImageGenerator(t *te } } -func TestAgentChatDoesNotRestartAlreadyCancelledThread(t *testing.T) { +func TestAgentChatContinuesStoppedThreadWithNewTask(t *testing.T) { ctx := context.Background() repo := memory.NewProjectRepository() + queue := &fakeJobQueue{} service := NewDesignService(repo, nil, nil, nil, nil, fakeCreativeChatAgent{ decision: design.AgentDecision{Intent: "image"}, taskPlan: design.AgentTaskPlan{ @@ -1498,6 +1535,7 @@ func TestAgentChatDoesNotRestartAlreadyCancelledThread(t *testing.T) { ImageTasks: []design.AgentImageTask{{Title: "视觉方案", Brief: "视觉方案"}}, }, }) + service.SetJobQueue(queue) project := design.Project{ ID: "project-cancelled-thread", Title: "Cancelled", @@ -1511,6 +1549,7 @@ func TestAgentChatDoesNotRestartAlreadyCancelledThread(t *testing.T) { if err := repo.Save(ctx, project); err != nil { t.Fatal(err) } + service.cancelAgentTask(project.ID, "thread-cancelled") thread, err := service.AgentChat(ctx, project.ID, design.AgentChatRequest{ ThreadID: "thread-cancelled", @@ -1526,14 +1565,20 @@ func TestAgentChatDoesNotRestartAlreadyCancelledThread(t *testing.T) { if err != nil { t.Fatal(err) } - if thread.Status != "cancelled" { - t.Fatalf("expected cancelled thread response, got %s", thread.Status) + if thread.Status != "running" { + t.Fatalf("expected stopped thread to accept a new running task, got %s", thread.Status) } saved, err := repo.Get(ctx, project.ID) if err != nil { t.Fatal(err) } - if len(saved.Nodes) != 0 || len(saved.Messages) != 1 { - t.Fatalf("expected cancelled thread not to create new work, got nodes=%#v messages=%#v", saved.Nodes, saved.Messages) + if len(saved.Nodes) != 1 || saved.Nodes[0].Status != "generating" { + t.Fatalf("expected new image task placeholder, got %#v", saved.Nodes) + } + if len(saved.Messages) <= 1 || saved.Messages[len(saved.Messages)-1].Status != "success" { + t.Fatalf("expected new thread messages after cancelled task, got %#v", saved.Messages) + } + if len(queue.jobs) != 1 || queue.jobs[0].Kind != design.JobFollowupGeneration { + t.Fatalf("expected new followup generation job, got %#v", queue.jobs) } } diff --git a/server/internal/logic/project_events_logic.go b/server/internal/logic/project_events_logic.go index dbcb6ca..8dda8ab 100644 --- a/server/internal/logic/project_events_logic.go +++ b/server/internal/logic/project_events_logic.go @@ -198,7 +198,12 @@ func projectThreadState(messages []design.Message, threadID string) (started boo if threadID == "" { return false, false, false } - for _, message := range messages { + startIndex := latestThreadTaskStartIndex(messages, threadID) + if startIndex < 0 { + return false, false, false + } + for index := startIndex; index < len(messages); index++ { + message := messages[index] if strings.TrimSpace(message.ThreadID) != threadID { continue } @@ -211,7 +216,7 @@ func projectThreadState(messages []design.Message, threadID string) (started boo failed = true } } - if hasUnfinishedThreadWork(messages, threadID) { + if hasUnfinishedThreadWorkFrom(messages, threadID, startIndex) { finished = false } return started, finished, failed @@ -222,8 +227,16 @@ func hasUnfinishedThreadWork(messages []design.Message, threadID string) bool { if threadID == "" { return false } + return hasUnfinishedThreadWorkFrom(messages, threadID, latestThreadTaskStartIndex(messages, threadID)) +} + +func hasUnfinishedThreadWorkFrom(messages []design.Message, threadID string, startIndex int) bool { + if startIndex < 0 { + return false + } running := make(map[string]bool) - for _, message := range messages { + for index := startIndex; index < len(messages); index++ { + message := messages[index] if strings.TrimSpace(message.ThreadID) != threadID { continue } @@ -242,6 +255,30 @@ func hasUnfinishedThreadWork(messages []design.Message, threadID string) bool { return len(running) > 0 } +func latestThreadTaskStartIndex(messages []design.Message, threadID string) int { + threadID = strings.TrimSpace(threadID) + if threadID == "" { + return -1 + } + firstIndex := -1 + latestUserIndex := -1 + for index, message := range messages { + if strings.TrimSpace(message.ThreadID) != threadID { + continue + } + if firstIndex < 0 { + firstIndex = index + } + if strings.EqualFold(strings.TrimSpace(message.Role), "user") { + latestUserIndex = index + } + } + if latestUserIndex >= 0 { + return latestUserIndex + } + return firstIndex +} + func threadWorkKey(message design.Message) string { role := strings.ToLower(strings.TrimSpace(message.Role)) if role == "" { diff --git a/server/internal/logic/project_events_logic_test.go b/server/internal/logic/project_events_logic_test.go index 7d3503a..92efd0e 100644 --- a/server/internal/logic/project_events_logic_test.go +++ b/server/internal/logic/project_events_logic_test.go @@ -100,6 +100,18 @@ func TestProjectThreadStateFinishesAfterGenerationToolSucceeds(t *testing.T) { } } +func TestProjectThreadStateKeepsNewTaskOpenAfterStoppedTask(t *testing.T) { + started, finished, failed := projectThreadState([]design.Message{ + {ID: "old-user", Role: "user", Type: "user", Content: "生成三张图", ThreadID: "thread-agent", CreatedAt: time.Now().Add(-3 * time.Second)}, + {ID: "old-stop", Role: "assistant", Type: "assistant", Content: "已停止", ThreadID: "thread-agent", Status: "cancelled", CreatedAt: time.Now().Add(-2 * time.Second)}, + {ID: "new-user", Role: "user", Type: "user", Content: "继续生成一张", ThreadID: "thread-agent", CreatedAt: time.Now().Add(-time.Second)}, + {ID: "planner", Role: "assistant", Type: "assistant", Name: "agent_planner", ThreadID: "thread-agent", Status: "success", CreatedAt: time.Now()}, + }, "thread-agent") + if !started || finished || failed { + t.Fatalf("expected new task to keep stopped thread open, got started=%v finished=%v failed=%v", started, finished, failed) + } +} + func TestProjectEventFingerprintIgnoresMessageOnlyChanges(t *testing.T) { now := time.Now() project := design.Project{