vxcontrol/pentagi
1
0
Fork 0
mirror of https://github.com/vxcontrol/pentagi.git synced 2026-07-20 23:57:11 -04:00
Code Issues 45 Packages Projects Releases 9 Wiki Activity

fix(flow): task execution handling with cancellable contexts on subtasks generation stage

Browse Source 
Enhanced the flowWorker's task processing by introducing a cancellable context for task execution. This change ensures that tasks can be properly cancelled without reporting false success states. The `runTask` method has been refactored to utilize a new `execTask` method, which centralizes task execution logic and maintains context integrity. This update improves flow control and error handling during task creation and execution.
This commit is contained in:
Dmitry Ng
2026-05-29 00:08:01 +03:00
parent 83c263e98e
commit ce06a0d26a
1 changed files with 48 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/pkg/controller/flow.go
+48 -16
View File
@@ -998,18 +998,61 @@ func (fw *flowWorker) processInput(flin flowInput) (TaskWorker, error) {
// anyway there need to set flow status to Running to disable user input
_ = fw.SetStatus(fw.ctx, database.FlowStatusRunning)
if task, err := fw.tc.CreateTask(fw.ctx, flin.input, fw); err != nil {
// Pre-create the per-task cancellable context BEFORE calling CreateTask.
// GenerateSubtasks (an LLM call) runs synchronously inside CreateTask and may take
// many seconds. Without this, a concurrent Stop() would invoke a no-op taskST and
// find taskWG at zero—reporting success while the generator is still running.
fw.taskMX.Lock()
fw.taskST()
ctx, taskST := context.WithCancel(fw.ctx)
fw.taskST = taskST
fw.taskMX.Unlock()
defer taskST()
fw.taskWG.Add(1)
defer fw.taskWG.Done()
defer fw.signalTaskComplete()
task, err := fw.tc.CreateTask(ctx, flin.input, fw)
if err != nil {
if errors.Is(err, context.Canceled) && fw.ctx.Err() == nil {
// CreateTask was cancelled by Stop() — not a fatal flow error.
// Keep the worker alive and return the flow to Waiting state.
flin.done <- nil
_ = fw.SetStatus(fw.ctx, database.FlowStatusWaiting)
return nil, nil
}
err = fmt.Errorf("failed to create task for flow %d: %w", fw.flowCtx.FlowID, err)
flin.done <- err
return nil, err
} else {
flin.done <- nil
spanName := fmt.Sprintf("perform task %d: %s", task.GetTaskID(), task.GetTitle())
return task, fw.runTask(spanName, flin.input, task)
}
flin.done <- nil
spanName := fmt.Sprintf("perform task %d: %s", task.GetTaskID(), task.GetTitle())
return task, fw.execTask(ctx, spanName, flin.input, task)
}
// runTask creates a fresh per-task cancellable context and runs an already-created task.
// Use this for tasks that were previously created and are being resumed (e.g. after waiting).
func (fw *flowWorker) runTask(spanName, input string, task TaskWorker) error {
fw.taskMX.Lock()
fw.taskST()
ctx, taskST := context.WithCancel(fw.ctx)
fw.taskST = taskST
fw.taskMX.Unlock()
defer taskST()
fw.taskWG.Add(1)
defer fw.taskWG.Done()
defer fw.signalTaskComplete()
return fw.execTask(ctx, spanName, input, task)
}
// execTask executes a task using an already-prepared context and cancel function.
func (fw *flowWorker) execTask(ctx context.Context, spanName, input string, task TaskWorker) error {
_, observation := obs.Observer.NewObservation(fw.ctx)
span := observation.Span(
langfuse.WithSpanName(spanName),
@@ -1019,18 +1062,7 @@ func (fw *flowWorker) runTask(spanName, input string, task TaskWorker) error {
}),
)
fw.taskMX.Lock()
fw.taskST()
ctx, taskST := context.WithCancel(fw.ctx)
fw.taskST = taskST
fw.taskMX.Unlock()
ctx, _ = span.Observation(ctx)
defer taskST()
fw.taskWG.Add(1)
defer fw.taskWG.Done()
defer fw.signalTaskComplete()
if err := task.Run(ctx); err != nil {
// if task is stopped by user and it's not finished yet