Repository navigation
Conversation
- Treat cancel as soft interrupt; reset CANCELLED -> RUNNING on next execute - Add interrupt timeout and force process cleanup on timeout/error - Fail fast on error_during_execution to avoid stuck RUNNING - Add regression tests for cancellation/resume
Avoid anyio cancel-scope mismatch when disconnecting from a different task by falling back to process termination and cache cleanup; also fix AgentService ClaudeCode name checks.
When a task is cancelled and the Claude Code subprocess is terminated (exit 143), avoid reporting FAILED from response processing; mark workbench completed and return COMPLETED.
📝 WalkthroughWalkthroughAdds timeout-based shutdown and forceful termination to ClaudeCode client lifecycle, implements soft-cancel recovery and retry-aware response processing, adjusts service-level agent type checks, improves sensitive-data logging, and adds tests covering cancellation, error handling, and timeout-driven cleanup. Changes
Sequence Diagram(s)sequenceDiagram
participant Agent as ClaudeCodeAgent
participant Timer as Interrupt Timer
participant Proc as Client Process
participant Cleanup as cleanup_task_clients
Agent->>Timer: Start interrupt timeout (INTERRUPT_TIMEOUT_SECONDS)
Agent->>Proc: Send interrupt / soft-cancel
alt Interrupt completes before timeout
Proc-->>Agent: Acknowledge interrupt
Agent->>Agent: Continue or recreate client as needed
else Interrupt times out
rect rgba(255,100,100,0.5)
Timer-->>Agent: Interrupt timeout
end
Agent->>Cleanup: Invoke cleanup_task_clients (force termination)
Cleanup-->>Agent: Clients removed / cleared
Agent->>Proc: Force terminate process if still present
end
sequenceDiagram
participant Processor as process_response
participant Handler as _handle_system_message
participant Result as _process_result_message
participant Client as SDK Client
Processor->>Handler: Inspect system message
alt Handler returns TaskStatus (e.g., FAILED)
Handler-->>Processor: TaskStatus (terminal)
Processor-->>Processor: Propagate terminal status and stop
else None
Processor->>Result: Handle result message
alt Result detects retryable API error
Result->>Client: Attempt retry via client
Client-->>Result: Retry outcome
Result-->>Processor: Return "RETRY" (requeue)
else Success
Result-->>Processor: Return TaskStatus.COMPLETED
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
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 |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
executor/services/agent_service.py (1)
437-451:⚠️ Potential issue | 🟡 MinorInconsistent naming for Agno error detail key.
The ClaudeCode agent uses consistent naming (
"ClaudeCode"on lines 437 and 443), but the Agno agent uses"Agno"for the type check (line 445) but"AgnoAgent"for the error_detail key (line 451). This should be consistent.🔧 Proposed fix for consistency
if status == TaskStatus.SUCCESS: results.append("Agno") else: errors.append("Agno") - error_detail["AgnoAgent"] = msg or "Unknown error" + error_detail["Agno"] = msg or "Unknown error"
🤖 Fix all issues with AI agents
In `@executor/agents/claude_code/claude_code_agent.py`:
- Around line 1423-1467: The close_client method currently returns
TaskStatus.SUCCESS even when an exception is caught; update the exception
handler in claude_code_agent. close_client to return TaskStatus.FAILED instead
of TaskStatus.SUCCESS after performing the existing cleanup (cls._clients.pop
and cls._session_id_map removals); keep the logger.exception call as-is and
ensure cleanup still runs (the finally block and the except cleanup should
remain). This change affects the exception return only (use TaskStatus.FAILED)
so callers checking the returned TaskStatus (e.g.,
executor.services.agent_service) will see failures correctly.
🧹 Nitpick comments (2)
executor/tests/agents/test_claude_code_cancellation_resume.py (1)
116-124: Unreachable code after exception raise.Lines 122-123 (
if False: yield None) are technically unreachable. While this is a common pattern to make an async generator that only raises, consider using a cleaner approach.♻️ Cleaner async generator pattern
class FakeClient: async def receive_response(self): raise Exception( "Command failed with exit code 143 (exit code: 143)\n" "Error output: Check stderr output for details" ) - if False: - yield None + # Make this an async generator by yielding before raise would also work, + # but the current pattern is acceptable for test purposes + return + yield # type: ignore[misc] # Makes this an async generatorAlternatively, the existing pattern is acceptable for test code since it achieves the intended behavior.
executor/agents/claude_code/claude_code_agent.py (1)
1394-1421: Consider logging on successful termination path.The
_force_close_client_processmethod handles process termination well, but the broadexcept Exception(flagged by static analysis BLE001) could silently swallow unexpected errors. The current implementation is acceptable for cleanup code, but consider logging the specific exception type for debugging.♻️ Optional: More specific exception handling
except Exception as e: logger.warning( - f"Failed to force-terminate process for session_id={session_id}, PID={pid}: {e}" + f"Failed to force-terminate process for session_id={session_id}, PID={pid}: {type(e).__name__}: {e}" )
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
executor/agents/claude_code/session_manager.py (1)
376-399:⚠️ Potential issue | 🟠 Major
del cls._clients[session_id]can raiseKeyErrordue to async interleaving.Between the
if session_id in cls._clientscheck (Line 377) and thedel(Line 380), theawaiton Line 379 yields to the event loop. A concurrentclose_clientcall could.pop()the same key during this window, causingdelto raiseKeyError. The same pattern appears at Lines 396-398.Use
.pop(session_id, None)for safe removal:🐛 Proposed fix
for internal_key, session_id in internal_keys_to_cleanup: if session_id in cls._clients: client = cls._clients[session_id] await cls._terminate_client_process(client, session_id) - del cls._clients[session_id] + cls._clients.pop(session_id, None) cleaned_count += 1 else:if session_id.startswith(task_id_prefix) or session_id == task_id_str: client = cls._clients[session_id] await cls._terminate_client_process(client, session_id) - del cls._clients[session_id] + cls._clients.pop(session_id, None) cleaned_count += 1Note: Line 396 also has
client = cls._clients[session_id]after an iteration check, which could similarly fail. Consider:- if session_id.startswith(task_id_prefix) or session_id == task_id_str: - client = cls._clients[session_id] + if session_id.startswith(task_id_prefix) or session_id == task_id_str: + client = cls._clients.get(session_id) + if not client: + continue
🤖 Fix all issues with AI agents
In `@executor/agents/claude_code/claude_code_agent.py`:
- Around line 66-67: Replace the duplicated timeout constants by removing
INTERRUPT_TIMEOUT_SECONDS and DISCONNECT_TIMEOUT_SECONDS definitions from
claude_code_agent.py and import the single source of truth from session_manager
(use the existing DISCONNECT_TIMEOUT_SECONDS/INTERRUPT_TIMEOUT_SECONDS
definitions there); update claude_code_agent.py to import
INTERRUPT_TIMEOUT_SECONDS from session_manager and delete the unused local
DISCONNECT_TIMEOUT_SECONDS constant so each constant is defined only in
session_manager and used via import in claude_code_agent.py.
🧹 Nitpick comments (5)
executor/agents/claude_code/response_processor.py (2)
569-579: Consider replacing the"RETRY"string sentinel with a dedicated type.The
Union[TaskStatus, str, None]return type uses a magic string"RETRY"as a signal. This is fragile and could collide with futureTaskStatusstring values. A small enum or a@dataclasswrapper would make the contract explicit.♻️ Suggested approach
# At module level: class _ResultAction(str, enum.Enum): RETRY = "RETRY" # Return type becomes: Union[TaskStatus, _ResultAction, None]Also applies to: 756-756
786-792: Inconsistent null-guarding ofstate_manager.In
_handle_system_message(Line 248),state_manageris guarded withif state_manager:before use. Here and in the success path (Lines 684, 711), it is used without a guard. While callers currently always pass a validstate_manager, the inconsistency could cause anAttributeErrorif invoked differently in the future.executor/agents/claude_code/claude_code_agent.py (1)
1022-1044: Timeout-guarded interrupt with forced cleanup fallback looks solid.The three-level approach (interrupt → timeout → force-kill) provides a robust cancellation path. The
returnat Line 1031 correctly prevents the cleanup fallback from running on success.Nit from static analysis (Ruff TRY401):
logger.exceptionat Line 1038 already attaches the traceback, so interpolating{interrupt_error}in the message is redundant.♻️ Minor cleanup
except Exception as interrupt_error: logger.exception( - f"Interrupt failed for task_id={self.task_id}, session_id={self.session_id}: {interrupt_error}. " + f"Interrupt failed for task_id={self.task_id}, session_id={self.session_id}. " "Forcing process cleanup" )executor/agents/claude_code/session_manager.py (2)
232-274: Multi-level shutdown inclose_clientis well-structured.The graceful disconnect → timeout → force-terminate cascade with finally-based cleanup is a solid pattern. The
RuntimeErrorcatch for cancel-scope mismatches (Line 249) addresses a known real-world edge case with async task boundaries.Nit (Ruff RUF010/TRY401): at Line 268,
str(e)is redundant insidelogger.exceptionwhich already appends the exception info.♻️ Minor cleanup
except Exception as e: logger.exception( - f"Error closing client for session_id {session_id}: {str(e)}" + f"Error closing client for session_id {session_id}" )
321-327: Extract hardcoded timeouts to named constants.The
2.0and1.0second waits at Lines 321 and 327 are magic numbers. Extracting them to module-level constants (e.g.,TERMINATE_WAIT_SECONDSandKILL_WAIT_SECONDS) improves readability and makes them easier to tune.
| INTERRUPT_TIMEOUT_SECONDS = 5.0 | ||
| DISCONNECT_TIMEOUT_SECONDS = 2.0 |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Duplicated timeout constants — define them in one place.
INTERRUPT_TIMEOUT_SECONDS and DISCONNECT_TIMEOUT_SECONDS are declared identically in both this file and session_manager.py (Lines 24-25). Moreover, each file only uses one of the two constants it defines:
- This file uses
INTERRUPT_TIMEOUT_SECONDS(Line 1024) but neverDISCONNECT_TIMEOUT_SECONDS. session_manager.pyusesDISCONNECT_TIMEOUT_SECONDS(Line 239) but neverINTERRUPT_TIMEOUT_SECONDS.
Define them once in session_manager.py and import here, to avoid drift.
♻️ Suggested fix
In this file, replace the local constants with an import:
-INTERRUPT_TIMEOUT_SECONDS = 5.0
-DISCONNECT_TIMEOUT_SECONDS = 2.0
+from executor.agents.claude_code.session_manager import INTERRUPT_TIMEOUT_SECONDS🤖 Prompt for AI Agents
In `@executor/agents/claude_code/claude_code_agent.py` around lines 66 - 67,
Replace the duplicated timeout constants by removing INTERRUPT_TIMEOUT_SECONDS
and DISCONNECT_TIMEOUT_SECONDS definitions from claude_code_agent.py and import
the single source of truth from session_manager (use the existing
DISCONNECT_TIMEOUT_SECONDS/INTERRUPT_TIMEOUT_SECONDS definitions there); update
claude_code_agent.py to import INTERRUPT_TIMEOUT_SECONDS from session_manager
and delete the unused local DISCONNECT_TIMEOUT_SECONDS constant so each constant
is defined only in session_manager and used via import in claude_code_agent.py.
Summary by CodeRabbit
Bug Fixes
Logging / Security
Tests