From ea2b9f121e1195ef18bce1acc3aef05e5a8d660e Mon Sep 17 00:00:00 2001
From: weichen <weichen@example.com>
Date: Wed, 16 Sep 2026 09:33:41 +0000
Subject: [PATCH] feat(taskrun): add WaitRunDone and WaitIdle for async workers

Allow tests and shutdown to drain in-memory workers before TempDir cleanup.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 internal/taskrun/manager.go        | 43 +++++++++++++++++++++
 internal/taskrun/wait_idle_test.go | 62 ++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 internal/taskrun/wait_idle_test.go

diff --git a/internal/taskrun/manager.go b/internal/taskrun/manager.go
index 21d86eee..4df35073 100644
--- a/internal/taskrun/manager.go
+++ b/internal/taskrun/manager.go
@@ -675,6 +675,49 @@ func (m *Manager) Cancel(runID string) error {
 	return fmt.Errorf("run not found")
 }
 
+// WaitRunDone 等待指定 run 的 worker 退出（runWorker 关闭 ar.done）。
+// 若 run 从未启动或已从内存表移除，立即返回 nil。
+func (m *Manager) WaitRunDone(ctx context.Context, runID string) error {
+	if m == nil || runID == "" {
+		return nil
+	}
+	m.mu.RLock()
+	ar, ok := m.runs[runID]
+	m.mu.RUnlock()
+	if !ok || ar == nil || ar.done == nil {
+		return nil
+	}
+	select {
+	case <-ar.done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// WaitIdle 等待当前内存中全部 active run 的 worker 退出。
+func (m *Manager) WaitIdle(ctx context.Context) error {
+	if m == nil {
+		return nil
+	}
+	m.mu.RLock()
+	dones := make([]<-chan struct{}, 0, len(m.runs))
+	for _, ar := range m.runs {
+		if ar != nil && ar.done != nil {
+			dones = append(dones, ar.done)
+		}
+	}
+	m.mu.RUnlock()
+	for _, done := range dones {
+		select {
+		case <-done:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+	return nil
+}
+
 func (m *Manager) heartbeatLoop(ctx context.Context, runID string) {
 	ticker := time.NewTicker(defaultHeartbeatPeriod)
 	defer ticker.Stop()
diff --git a/internal/taskrun/wait_idle_test.go b/internal/taskrun/wait_idle_test.go
new file mode 100644
index 00000000..8bfd4ae5
--- /dev/null
+++ b/internal/taskrun/wait_idle_test.go
@@ -0,0 +1,62 @@
+package taskrun
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"chat-agent-system/internal/bizdev"
+	"chat-agent-system/internal/stream"
+)
+
+func TestManager_WaitIdle_WaitsForWorker(t *testing.T) {
+	m := NewManager(nil, nil)
+	block := make(chan struct{})
+	m.SetExecutor(
+		func(ctx context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
+			select {
+			case <-ctx.Done():
+				return &bizdev.TaskRunResult{Success: false}, ctx.Err()
+			case <-block:
+				return &bizdev.TaskRunResult{Success: true}, nil
+			}
+		},
+		nil,
+		nil,
+	)
+
+	rec := &RunRecord{
+		RunID:     "run-wait-idle",
+		SessionID: "sess",
+		UserID:    1,
+		Status:    "pending",
+		CreatedAt: time.Now(),
+		UpdatedAt: time.Now(),
+	}
+	if _, err := m.CreateAndStart(context.Background(), rec.RunID, rec, bizdev.TaskRunRequest{}, bizdev.StageRunRequest{}, ""); err != nil {
+		t.Fatalf("CreateAndStart: %v", err)
+	}
+
+	idleDone := make(chan error, 1)
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		idleDone <- m.WaitIdle(ctx)
+	}()
+
+	select {
+	case err := <-idleDone:
+		t.Fatalf("WaitIdle returned before worker released: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(block)
+	select {
+	case err := <-idleDone:
+		if err != nil {
+			t.Fatalf("WaitIdle: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("WaitIdle timed out after releasing worker")
+	}
+}