fix(go-agent): support task-based run cancellation#17188
Draft
Hz-186 wants to merge 1 commit into
Draft
Conversation
dosubot
Bot
added
the
size:M
Hz-186
marked this pull request as draft
July 21, 2026 10:28
Contributor
📝 WalkthroughWalkthroughAdds task-scoped cancellation by tracking active task IDs in the runner, authorizing cancellation against the task’s canvas, and exposing an authenticated POST endpoint. ChangesTask cancellation flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant AgentHandler
participant AgentService
participant Runner
Client->>AgentHandler: POST /api/v1/tasks/:task_id/cancel
AgentHandler->>AgentService: CancelTask(userID, taskID)
AgentService->>Runner: TaskCanvas(taskID)
AgentService->>AgentService: loadCanvasForUser(userID, canvasID)
AgentService->>Runner: CancelTask(taskID)
Runner-->>AgentHandler: cancellation result
AgentHandler-->>Client: success response
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/service/agent.go`:
- Around line 1888-1901: Prevent the cancellation TOCTOU race by passing the
authorized canvasID from AgentService.CancelTask to runner.CancelTask in
internal/service/agent.go lines 1888-1901. Update runner.CancelTask in
internal/agent/canvas/runner.go lines 387-403 to accept expectedCanvasID and,
while holding the runner lock, verify the task still belongs to that canvas
before cancelling; otherwise leave it unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 82530c1f-1e47-4889-a9e5-313325a46324
📒 Files selected for processing (4)
internal/agent/canvas/runner.gointernal/handler/agent.gointernal/router/router.gointernal/service/agent.go
Comment on lines
+1888
to
+1901
| func (s *AgentService) CancelTask(ctx context.Context, userID, taskID string) error { | ||
| if taskID == "" || s.runner == nil { | ||
| return nil | ||
| } | ||
| canvasID, active := s.runner.TaskCanvas(taskID) | ||
| if !active { | ||
| return nil | ||
| } | ||
| if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil { | ||
| return err | ||
| } | ||
| s.runner.CancelTask(taskID) | ||
| return nil | ||
| } |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win
Prevent TOCTOU race condition during task cancellation.
There is a Time-Of-Check to Time-Of-Use (TOCTOU) race condition between s.runner.TaskCanvas(taskID) and s.runner.CancelTask(taskID). If an attacker quickly reassigns a taskID (e.g., via a colliding version_id) to a victim's canvas in the short window after the authorization check but before the actual cancellation, they could bypass authorization and cancel the victim's task. Fix this by passing the authorized canvas ID to the runner and atomically validating it inside the runner's lock.
internal/service/agent.go#L1888-L1901: pass the authorizedcanvasIDtos.runner.CancelTaskto ensure we only cancel the verified canvas.internal/agent/canvas/runner.go#L387-L403: updateCancelTaskto accept and enforceexpectedCanvasIDwhile holding the lock.
🔒️ Proposed fix to enforce canvas ownership during cancellation
In internal/service/agent.go:
if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil {
return err
}
- s.runner.CancelTask(taskID)
+ s.runner.CancelTask(taskID, canvasID)
return nilIn internal/agent/canvas/runner.go:
-// CancelTask signals an active run identified by the task_id emitted in its
-// events. It returns the owning canvas id when the task is still active.
-func (r *Runner) CancelTask(taskID string) (string, bool) {
+// CancelTask signals an active run by task_id, ensuring it belongs to the expected canvas.
+func (r *Runner) CancelTask(taskID, expectedCanvasID string) bool {
r.mu.Lock()
cancel, ok := r.taskCancels[taskID]
canvasID := r.taskCanvases[taskID]
r.mu.Unlock()
- if !ok {
- return "", false
+ if !ok || canvasID != expectedCanvasID {
+ return false
}
select {
case <-cancel:
default:
close(cancel)
}
- return canvasID, true
+ return true
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| func (s *AgentService) CancelTask(ctx context.Context, userID, taskID string) error { | |
| if taskID == "" || s.runner == nil { | |
| return nil | |
| } | |
| canvasID, active := s.runner.TaskCanvas(taskID) | |
| if !active { | |
| return nil | |
| } | |
| if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil { | |
| return err | |
| } | |
| s.runner.CancelTask(taskID) | |
| return nil | |
| } | |
| func (s *AgentService) CancelTask(ctx context.Context, userID, taskID string) error { | |
| if taskID == "" || s.runner == nil { | |
| return nil | |
| } | |
| canvasID, active := s.runner.TaskCanvas(taskID) | |
| if !active { | |
| return nil | |
| } | |
| if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil { | |
| return err | |
| } | |
| s.runner.CancelTask(taskID, canvasID) | |
| return nil | |
| } |
Suggested change
| func (s *AgentService) CancelTask(ctx context.Context, userID, taskID string) error { | |
| if taskID == "" || s.runner == nil { | |
| return nil | |
| } | |
| canvasID, active := s.runner.TaskCanvas(taskID) | |
| if !active { | |
| return nil | |
| } | |
| if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil { | |
| return err | |
| } | |
| s.runner.CancelTask(taskID) | |
| return nil | |
| } | |
| // CancelTask signals an active run by task_id, ensuring it belongs to the expected canvas. | |
| func (r *Runner) CancelTask(taskID, expectedCanvasID string) bool { | |
| r.mu.Lock() | |
| cancel, ok := r.taskCancels[taskID] | |
| canvasID := r.taskCanvases[taskID] | |
| r.mu.Unlock() | |
| if !ok || canvasID != expectedCanvasID { | |
| return false | |
| } | |
| select { | |
| case <-cancel: | |
| default: | |
| close(cancel) | |
| } | |
| return true | |
| } |
📍 Affects 2 files
internal/service/agent.go#L1888-L1901(this comment)internal/agent/canvas/runner.go#L387-L403
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/service/agent.go` around lines 1888 - 1901, Prevent the cancellation
TOCTOU race by passing the authorized canvasID from AgentService.CancelTask to
runner.CancelTask in internal/service/agent.go lines 1888-1901. Update
runner.CancelTask in internal/agent/canvas/runner.go lines 387-403 to accept
expectedCanvasID and, while holding the runner lock, verify the task still
belongs to that canvas before cancelling; otherwise leave it unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Testing
bash build.sh --test ./internal/agent/canvas ./internal/router