Fix webui streaming recovery regressions

audit/internal/webui/tasks.go

@@ -429,43 +429,8 @@ func (q *taskQueue) worker() {
 				wg.Add(1)
 				goRecoverOnce("task "+t.Target, func() {
 					defer wg.Done()
-					defer func() {
-						if rec := recover(); rec != nil {
-							msg := fmt.Sprintf("task panic: %v", rec)
-							slog.Error("task panic",
-								"task_id", t.ID,
-								"target", t.Target,
-								"panic", fmt.Sprint(rec),
-								"stack", string(debug.Stack()),
-							)
-							j.append("ERROR: " + msg)
-							j.finish(msg)
-						}
-					}()
-
-					if q.kmsgWatcher != nil && isSATTarget(t.Target) {
-						q.kmsgWatcher.NotifyTaskStarted(t.ID, t.Target)
-					}
-
-					q.runTask(t, j, taskCtx)
-
-					if q.kmsgWatcher != nil {
-						q.kmsgWatcher.NotifyTaskFinished(t.ID)
-					}
-
-					q.mu.Lock()
-					now2 := time.Now()
-					t.DoneAt = &now2
-					if t.Status == TaskRunning {
-						if j.err != "" {
-							t.Status = TaskFailed
-							t.ErrMsg = j.err
-						} else {
-							t.Status = TaskDone
-						}
-					}
-					q.persistLocked()
-					q.mu.Unlock()
+					defer taskCancel()
+					q.executeTask(t, j, taskCtx)
 				})
 			}
 			wg.Wait()
@@ -481,6 +446,54 @@ func (q *taskQueue) worker() {
 	}
 }
 
+func (q *taskQueue) executeTask(t *Task, j *jobState, ctx context.Context) {
+	startedKmsgWatch := false
+	defer q.finalizeTaskRun(t, j)
+	defer func() {
+		if startedKmsgWatch && q.kmsgWatcher != nil {
+			q.kmsgWatcher.NotifyTaskFinished(t.ID)
+		}
+	}()
+	defer func() {
+		if rec := recover(); rec != nil {
+			msg := fmt.Sprintf("task panic: %v", rec)
+			slog.Error("task panic",
+				"task_id", t.ID,
+				"target", t.Target,
+				"panic", fmt.Sprint(rec),
+				"stack", string(debug.Stack()),
+			)
+			j.append("ERROR: " + msg)
+			j.finish(msg)
+		}
+	}()
+
+	if q.kmsgWatcher != nil && isSATTarget(t.Target) {
+		q.kmsgWatcher.NotifyTaskStarted(t.ID, t.Target)
+		startedKmsgWatch = true
+	}
+
+	q.runTask(t, j, ctx)
+}
+
+func (q *taskQueue) finalizeTaskRun(t *Task, j *jobState) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	now := time.Now()
+	t.DoneAt = &now
+	if t.Status == TaskRunning {
+		if j.err != "" {
+			t.Status = TaskFailed
+			t.ErrMsg = j.err
+		} else {
+			t.Status = TaskDone
+			t.ErrMsg = ""
+		}
+	}
+	q.persistLocked()
+}
+
 // setCPUGovernor writes the given governor to all CPU scaling_governor sysfs files.
 // Silently ignores errors (e.g. when cpufreq is not available).
 func setCPUGovernor(governor string) {