Improve context cancelation for agent mode #1707
Conversation
|
Caution Review failedThe pull request is closed. WalkthroughCentralizes agent session startup and OS signal handling via Agent.startSession, updates generator/toolmanager to clear previous tool state and return errors from tool calls, tightens logging message assembly with rune-aware truncation, adds Heroku HTTP helper and CLI token retrieval, and a minor CLI welcome punctuation fix. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant Agent
participant Generator
participant ToolManager
Note over Agent: StartWithUserPrompt / StartWithMessage -> startSession
User->>Agent: input / initial message
activate Agent
Agent->>Agent: create signal-aware context
Agent->>Generator: HandleMessage(ctx, message)
activate Generator
loop up to max turns
Generator->>ToolManager: ClearPrevious()
Generator->>ToolManager: HandleToolCalls(ctx, requests)
activate ToolManager
alt ctx canceled
ToolManager-->>Generator: (nil, context.Canceled)
else tool error
ToolManager-->>Generator: (nil, error)
else success
ToolManager-->>Generator: (*ai.Message, nil)
end
deactivate ToolManager
alt error returned
Generator-->>Agent: return/error (stop)
else append tool response and continue
end
end
deactivate Generator
alt context canceled during processing
Agent->>Agent: handle cancellation, prompt or exit
else other error
Agent->>Agent: print error and continue
end
deactivate Agent
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.5.0)level=warning msg="[linters_context] running gomodguard failed: unable to read module file go.mod: current working directory must have a go.mod file: if you are not using go modules it is suggested to disable this linter" Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
src/pkg/cli/client/byoc/aws/byoc.go (1)
783-789: Minor improvement: clarify the FIXME scope and reconsider the variable factoring.The comment change from an inline TODO to a FIXME is a reasonable documentation improvement, but consider:
FIXME clarity: The comment mentions storing task ID in
ProjectUpdate, but it's unclear whether this FIXME is a prerequisite for the current logic or a future enhancement. Consider adding context about why the conditional check on line 785 is necessary today and how storing the task ID would improve this flow.Variable factoring: While factoring
cdTailinto a local variable on line 783 improves readability, ensure the refactoring doesn't inadvertently change the intent. The code remains functionally identical, but the comment repositioning is the substantive change here.If the FIXME represents a future improvement, consider adding a brief explanation of the current workaround or linking to a tracking issue.
src/pkg/agent/agent.go (1)
147-156: Confirm defer cancel() in loop is intentional.The
defer cancel()at Line 148 is inside the loop, meaning it will be deferred on each iteration. Since the context created bysignal.NotifyContextis scoped to each message handling iteration, this pattern should work correctly, but it's unconventional.Consider explicitly calling
cancel()before continuing on Line 152 or at the end of the loop body to make the cleanup timing more explicit and avoid accumulating deferred calls.Apply this diff to make cleanup explicit:
// Handle Ctrl+C during message handling / tool calls ctx, cancel := signal.NotifyContext(ctx, os.Interrupt) - defer cancel() if err := a.handleUserMessage(ctx, input); err != nil { if errors.Is(err, context.Canceled) { + cancel() continue } a.printer.Println("Error handling message:", err) } + cancel() }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
src/cmd/cli/command/commands.go(1 hunks)src/pkg/agent/agent.go(3 hunks)src/pkg/agent/generator.go(3 hunks)src/pkg/agent/toolmanager.go(4 hunks)src/pkg/cli/client/byoc/aws/byoc.go(1 hunks)src/pkg/logs/slog.go(1 hunks)src/pkg/migrate/heroku.go(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
src/pkg/cli/client/byoc/aws/byoc.go (1)
src/pkg/clouds/aws/ecs/logs.go (1)
LogGroupInput(94-99)
src/pkg/migrate/heroku.go (1)
src/pkg/term/colorizer.go (1)
Errorf(322-324)
src/pkg/agent/toolmanager.go (1)
src/pkg/term/colorizer.go (1)
Println(282-284)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (13)
src/pkg/cli/client/byoc/aws/byoc.go (1)
783-789:⚠️ Mismatch between PR objectives and file changes.The PR objectives describe improvements to agent mode context cancellation and error propagation, but this file contains only a documentation/refactoring change in the
getLogGroupInputs()method. This change (repositioning a comment from an etag-filtering TODO to a task-ID storage FIXME) is unrelated to the stated PR objectives and appears to belong in a different PR or is incomplete.Please clarify:
- Is this file intended to be part of this PR?
- Are there other files in this PR that implement the actual agent mode context cancellation improvements described in the PR objectives?
src/cmd/cli/command/commands.go (1)
482-482: LGTM!The addition of the period improves the consistency of the welcome message.
src/pkg/migrate/heroku.go (1)
279-279: LGTM! Proper error wrapping.The changes from
%vto%wenable proper error chain inspection witherrors.Is()anderrors.As(), which is Go best practice.Also applies to: 304-304, 334-334
src/pkg/agent/agent.go (3)
99-107: LGTM! Clean session initialization refactor.The extraction of
startSessionproperly centralizes the initialization logic and separates user-facing prompts from agent initialization.
160-160: LGTM! Named constant improves readability.Defining
maxTurnsas a constant makes the code more maintainable.
109-118: Error handling at call sites is properly implemented.Verification confirms that both call sites correctly handle the wrapped error from
startSession:
src/cmd/cli/command/commands.go: The error fromStartWithUserPromptis returned directly to the Cobra command handler, which processes it appropriately.src/pkg/debug/debug.go: TheStartWithMessagecalls are wrapped bypromptAndTrackDebugSession, which captures the error withif err != nil { return err }before propagating it up the call stack.The error wrapping at line 116 is correct and necessary.
src/pkg/agent/generator.go (3)
65-66: LGTM! Proper state reset.Calling
ClearPrevious()at the start of each message ensures that infinite loop detection doesn't carry over between distinct user messages.
73-78: LGTM! Appropriate context cancellation handling.Early return on
context.Canceledensures that Ctrl+C interrupts are handled gracefully without treating them as retryable errors.
87-92: LGTM! Proper error propagation from tool calls.The updated signature and error handling correctly propagate fatal errors from tool execution while continuing to append successful tool responses.
src/pkg/agent/toolmanager.go (4)
61-68: LGTM! Signature change enables proper error propagation.Returning
(*ai.Message, error)allows the caller to distinguish between successful tool execution, context cancellation, and other errors. The early return for infinite loop detection correctly returnsnilerror.
74-89: LGTM! Context cancellation is properly handled as fatal.The distinction between
context.Canceled(returned immediately as fatal) and other errors (logged and included in the tool response for the agent to handle) correctly implements the PR's intent to not propagate cancellation errors back to the agent.
136-137: LGTM! Variable rename improves clarity.Renaming from
keytoprevJSONmakes the variable's purpose clearer.
148-150: LGTM! ClearPrevious enables proper session reset.The new method allows the generator to reset infinite-loop detection state between distinct user messages, as called in
generator.go:65.
|
Note Docstrings generation - SUCCESS |
Docstrings generation was requested by @lionello. * #1707 (comment) The following files were modified: * `src/pkg/migrate/heroku.go`
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/pkg/agent/agent.go (1)
160-160: Consider makingmaxTurnsconfigurable or documenting the constant.The magic number 8 is defined inline. Consider adding a brief comment explaining why 8 turns was chosen, or extracting it to a package-level constant if it might be adjusted later.
- const maxTurns = 8 + const maxTurns = 8 // limit agent turns to prevent runaway tool calls
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/pkg/agent/agent.go(3 hunks)src/pkg/logs/slog.go(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- src/pkg/logs/slog.go
🧰 Additional context used
🧬 Code graph analysis (1)
src/pkg/agent/agent.go (1)
src/pkg/term/colorizer.go (4)
Reset(388-390)Printf(286-288)Errorf(322-324)Println(282-284)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Analyze (go)
- GitHub Check: go-test
🔇 Additional comments (4)
src/pkg/agent/agent.go (4)
99-107: LGTM!Clean delegation to the new
startSessionmethod. The comment clearly explains thatuserPromptis displayed to the user but an empty message is sent to the agent, which aligns with the intended behavior.
146-151: Good handling of context cancellation in the input loop.Correctly continues the loop when
context.Canceledis detected, allowing the user to continue interacting after a Ctrl+C during message handling. This aligns with the PR objective.
156-158: Correct use ofsignal.NotifyContextfor interrupt handling.This pattern properly creates a cancellable context that responds to Ctrl+C during message handling. The
defer cancel()ensures proper cleanup of the signal notification.
110-110: This review comment is based on a misconception about Go's signal handling mechanisms.signal.Reset(os.Interrupt)only affects handlers registered viasignal.Notify(), not context-based cancellation fromsignal.NotifyContext(). The codebase usessignal.NotifyContext()for signal handling (main.go line 32, agent.go line 157), sosignal.Reset()has no functional effect on the signal flow. No restoration is needed because no handler-based signal setup is disrupted.Likely an incorrect or invalid review comment.
Docstrings generation was requested by @lionello. * #1707 (comment) The following files were modified: * `src/pkg/migrate/heroku.go` Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> [no ci]
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/pkg/migrate/heroku.go(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Analyze (go)
- GitHub Check: go-test
🔇 Additional comments (4)
src/pkg/migrate/heroku.go (4)
284-284: LGTM! Proper error wrapping enables cancellation detection.Changing from
%vto%wcorrectly enables error chain unwrapping, which is essential for detecting context cancellation errors as mentioned in the PR objectives.
309-309: LGTM! Consistent error wrapping.Proper use of
%wfor error chaining.
336-342: LGTM! Clear and comprehensive documentation.The documentation accurately describes the function's behavior, including its checks, execution flow, and error conditions.
346-346: LGTM! Proper error wrapping.Consistent use of
%wfor error chaining.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Description
This PR has two improvements for ctrl-c during agent mode:
Summary by CodeRabbit
Bug Fixes
Improvements
New Features
✏️ Tip: You can customize this high-level summary in your review settings.