feat: implement graceful shutdown with cooperative timeout strategy (#130)#154
feat: implement graceful shutdown with cooperative timeout strategy (#130)#154
Conversation
…130) Add INTERRUPTED task status, SHUTDOWN termination reason, ShutdownStrategy protocol, CooperativeTimeoutStrategy, and ShutdownManager with cross-platform signal handling. The execution loop checks shutdown at turn boundaries and before tool invocations. The engine transitions tasks to INTERRUPTED on shutdown. Includes GracefulShutdownConfig in RootConfig. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…130) Pre-reviewed by 10 agents, 30+ findings addressed: - CooperativeTimeoutStrategy + ShutdownManager with signal handling - ShutdownResult model, INTERRUPTED TaskStatus transitions - GracefulShutdownConfig with validation bounds - Turn-boundary shutdown checks in ReactLoop - Windows signal safety (deferred logging via call_soon_threadsafe) - Cleanup callback exception isolation - Extracted build_error_prompt to prompt.py, make_budget_checker to loop_protocol.py - DESIGN_SPEC §6.1/§6.7/§15.3/§15.5 documentation updates Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Dependency Review✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.Scanned FilesNone |
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (10)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds a pluggable graceful shutdown framework (strategy + manager), introduces TaskStatus.INTERRUPTED and related transitions, threads shutdown checks through execution loops and AgentEngine, exposes shutdown config and public API, adds observability events, and includes unit and integration tests for shutdown behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant OS as Operating System
participant Manager as ShutdownManager
participant Strategy as CooperativeTimeoutStrategy
participant Tasks as Running Tasks
participant Cleanup as Cleanup Callbacks
OS->>Manager: Signal (SIGTERM/SIGINT)
activate Manager
Manager->>Strategy: request_shutdown()
activate Strategy
Manager->>Manager: initiate_shutdown()
Strategy->>Strategy: emit SHUTDOWN_SIGNAL event
rect rgba(100,150,200,0.5)
Note over Strategy,Tasks: Grace Period (grace_seconds)
Strategy->>Tasks: wait for cooperative exit
Tasks-->>Strategy: tasks complete cooperatively
end
rect rgba(200,100,100,0.5)
Note over Strategy,Tasks: Force-cancel remaining
Strategy->>Tasks: task.cancel() for remaining
Tasks-->>Strategy: cancellations acknowledged
end
rect rgba(150,200,100,0.5)
Note over Strategy,Cleanup: Cleanup Phase (cleanup_seconds)
Strategy->>Cleanup: run cleanup callbacks
Cleanup-->>Strategy: callbacks complete / timeouts
Strategy->>Manager: return ShutdownResult
end
deactivate Strategy
Manager-->>OS: shutdown complete (result)
deactivate Manager
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
✨ Simplify code
Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the system's resilience by introducing a robust and configurable graceful shutdown mechanism. It ensures that when a termination signal is received, active agent tasks are cooperatively stopped, critical cleanup operations are performed, and task states are updated appropriately to allow for seamless recovery or reassignment. This prevents data loss and ensures a controlled system exit, improving overall stability and operational reliability. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Pull request overview
Implements a graceful shutdown capability in the engine, introducing a pluggable shutdown strategy and wiring shutdown checks into the ReAct loop and task lifecycle so in-flight work can exit cooperatively (or be cancelled) and tasks can be resumed after restart.
Changes:
- Added shutdown framework (
ShutdownStrategy,CooperativeTimeoutStrategy,ShutdownManager) plus execution-loop termination reasonSHUTDOWN. - Extended task lifecycle with
TaskStatus.INTERRUPTEDand corresponding transitions; engine maps shutdown termination → task interruption. - Added configuration (
GracefulShutdownConfig), observability events, and unit/integration tests; extracted helper functions to reduce engine module size.
Reviewed changes
Copilot reviewed 22 out of 22 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/engine/test_shutdown.py | Adds unit tests for shutdown strategy/manager and config validation. |
| tests/unit/engine/test_run_result.py | Updates imports/tests to use make_budget_checker from loop_protocol. |
| tests/unit/engine/test_react_loop.py | Adds unit tests covering shutdown-triggered termination paths. |
| tests/unit/engine/test_loop_protocol.py | Updates TerminationReason expectations for new SHUTDOWN member. |
| tests/unit/engine/test_agent_engine_lifecycle.py | Verifies engine transitions to INTERRUPTED on shutdown termination. |
| tests/unit/core/test_task_transitions.py | Expands transition tests to include INTERRUPTED transitions. |
| tests/unit/core/test_enums.py | Updates enum membership/value assertions for TaskStatus.INTERRUPTED. |
| tests/integration/engine/test_graceful_shutdown.py | Adds end-to-end shutdown flow tests using a shutdown checker/provider. |
| src/ai_company/observability/events/execution.py | Adds shutdown-related observability event constants. |
| src/ai_company/engine/shutdown.py | Introduces shutdown strategy implementation and manager w/ signal handling + cleanup. |
| src/ai_company/engine/react_loop.py | Adds shutdown checks at loop top and before tool execution; logs turn start at info. |
| src/ai_company/engine/prompt.py | Adds build_error_prompt() helper used when prompt construction fails. |
| src/ai_company/engine/loop_protocol.py | Adds ShutdownChecker type + TerminationReason.SHUTDOWN; defines make_budget_checker(). |
| src/ai_company/engine/agent_engine.py | Wires shutdown checker into loop execution and maps SHUTDOWN → INTERRUPTED. |
| src/ai_company/engine/init.py | Re-exports shutdown types and ShutdownChecker. |
| src/ai_company/core/task_transitions.py | Adds valid transitions for INTERRUPTED. |
| src/ai_company/core/enums.py | Adds TaskStatus.INTERRUPTED and updates lifecycle documentation. |
| src/ai_company/config/schema.py | Adds GracefulShutdownConfig and includes it in RootConfig. |
| src/ai_company/config/defaults.py | Adds graceful_shutdown default config block. |
| src/ai_company/config/init.py | Exposes GracefulShutdownConfig from config package. |
| DESIGN_SPEC.md | Updates lifecycle diagram and shutdown design sections to reflect implementation. |
| CLAUDE.md | Updates engine package description to include recovery/shutdown. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
tests/unit/engine/test_shutdown.py
Outdated
| cleanup_completed=True, | ||
| duration_seconds=0.1, | ||
| ) | ||
| with pytest.raises(Exception, match="frozen"): |
There was a problem hiding this comment.
ShutdownResult is a frozen Pydantic model, but this test asserts immutability via pytest.raises(Exception, ...). Elsewhere in the repo, frozen model reassignment is consistently asserted with ValidationError; tightening this makes the test more specific and avoids masking unrelated exceptions.
| with pytest.raises(Exception, match="frozen"): | |
| with pytest.raises(ValidationError, match="frozen"): |
| """Provider that triggers shutdown on the second call. | ||
|
|
||
| First call returns a tool-use response, second call triggers | ||
| shutdown and returns a stop response. |
There was a problem hiding this comment.
The class docstring says this provider "triggers shutdown on the second call", but the implementation calls request_shutdown() when _call_count == 1. Update the docstring or the condition so they match.
| """Provider that triggers shutdown on the second call. | |
| First call returns a tool-use response, second call triggers | |
| shutdown and returns a stop response. | |
| """Provider that triggers shutdown on the first call. | |
| First call triggers shutdown and returns a stop response. |
| await engine.run( | ||
| identity=identity, | ||
| task=task, | ||
| ) | ||
|
|
||
| # The first LLM call triggers shutdown, but since it returns | ||
| # STOP with no tool calls, the loop completes normally. | ||
| # Verify the manager state reflects the shutdown signal. |
There was a problem hiding this comment.
test_shutdown_interrupts_task_and_runs_cleanup (and its docstring/comments) claim to verify that the task is INTERRUPTED and that cleanup ran, but the test currently discards the engine.run() result and only asserts manager.is_shutting_down(). Consider asserting the returned TaskExecution.status and that the cleanup callback was invoked (or rename/re-scope the test to what it actually checks).
| await engine.run( | |
| identity=identity, | |
| task=task, | |
| ) | |
| # The first LLM call triggers shutdown, but since it returns | |
| # STOP with no tool calls, the loop completes normally. | |
| # Verify the manager state reflects the shutdown signal. | |
| execution = await engine.run( | |
| identity=identity, | |
| task=task, | |
| ) | |
| # The first LLM call triggers shutdown, but since it returns | |
| # STOP with no tool calls, the loop completes normally. | |
| # Verify the task is marked as INTERRUPTED, cleanup ran, and the | |
| # manager state reflects the shutdown signal. | |
| assert execution.status is TaskStatus.INTERRUPTED | |
| assert cleanup_ran == ["done"] |
| logger.warning( | ||
| EXECUTION_SHUTDOWN_FORCE_CANCEL, | ||
| error=( | ||
| f"Task raised during shutdown: " | ||
| f"{type(task.exception()).__name__}" | ||
| ), | ||
| ) |
There was a problem hiding this comment.
EXECUTION_SHUTDOWN_FORCE_CANCEL is logged even when a task simply raised while already in done (no force-cancel occurred). This makes the event name misleading for metrics/alerts; consider using a different shutdown event (or add an action field and reserve FORCE_CANCEL for the actual cancellation path).
Greptile SummaryThis PR implements the graceful shutdown system described in DESIGN_SPEC §6.7, introducing a pluggable Key changes:
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant OS as OS Signal
participant SM as ShutdownManager
participant ST as CooperativeTimeoutStrategy
participant RL as ReactLoop
participant AE as AgentEngine
OS->>SM: SIGINT / SIGTERM
SM->>ST: request_shutdown()
ST->>ST: _shutdown_event.set()
Note over RL: Next turn boundary check
RL->>RL: _check_shutdown(shutdown_checker)
RL-->>AE: ExecutionResult(SHUTDOWN)
Note over RL: OR before tool invocation
RL->>RL: _check_shutdown(shutdown_checker)
RL->>RL: clear tool_calls_made in last TurnRecord
RL-->>AE: ExecutionResult(SHUTDOWN)
AE->>AE: _apply_post_execution_transitions()
AE->>AE: _transition_to_interrupted()
Note over AE: task → INTERRUPTED
Note over SM: initiate_shutdown() called by caller
SM->>ST: execute_shutdown(running_tasks, cleanup_callbacks)
ST->>ST: asyncio.wait(task_set, timeout=grace_seconds)
alt tasks exited cooperatively
ST-->>ST: tasks_completed += 1
else timeout — tasks still pending
ST->>ST: task.cancel() for each pending
ST->>ST: asyncio.wait(pending, timeout=5s)
ST-->>ST: tasks_interrupted = len(pending)
end
ST->>ST: _run_cleanup(callbacks, timeout=cleanup_seconds)
ST-->>SM: ShutdownResult
Last reviewed commit: d9366f7 |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/ai_company/engine/react_loop.py (1)
115-123:⚠️ Potential issue | 🟠 MajorRecord tool calls only after they actually run.
With the new shutdown check before
invoke_all(), this path can now returnSHUTDOWNafterTurnRecordhas already been appended withresponse.tool_calls. That makestool_calls_madeandtotal_tool_callsreport tools as executed even when shutdown skipped them.Also applies to: 183-197
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ai_company/engine/react_loop.py` around lines 115 - 123, The TurnRecord for a turn is being appended to turns before tools are guaranteed to have executed, causing tool_calls_made/total_tool_calls to include tools skipped by shutdown; change the flow in react_loop so that you only call _make_turn_record and append it to turns after invoke_all completes successfully and after shutdown_checker confirms tools ran (i.e., move creation/appending of the TurnRecord out of the pre-invoke path and into the post-invoke path inside self._process_turn_response handling), and apply the same change to the duplicate block around the other invoke_all usage (the second occurrence in this method) so both code paths record tool_calls only after tools actually ran.DESIGN_SPEC.md (1)
1028-1058:⚠️ Potential issue | 🟡 MinorThe shutdown spec still disagrees with the implemented shutdown outcome.
These sections treat
INTERRUPTEDas something reserved for force-cancelled tasks, but the engine in this PR now maps anyTerminationReason.SHUTDOWNrun toINTERRUPTEDas well. §6.5 also still documents the pre-shutdownExecutionLoopAPI and termination reasons, so the spec is internally inconsistent right after marking graceful shutdown as adopted.Also applies to: 2485-2485
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@DESIGN_SPEC.md` around lines 1028 - 1058, The spec conflicts with the implementation: currently the code maps any TerminationReason.SHUTDOWN to TaskStatus.INTERRUPTED but §6.7/§6.5 describe INTERRUPTED as only for force-cancelled tasks; reconcile by updating the spec to match the implementation — explicitly state that TerminationReason.SHUTDOWN results in TaskStatus.INTERRUPTED (including cooperative and forced shutdown outcomes), adjust the lifecycle transition table to show ASSIGNED/IN_PROGRESS → INTERRUPTED on shutdown, and remove or revise the pre-shutdown wording in the ExecutionLoop section so it no longer contradicts this mapping (refer to TerminationReason.SHUTDOWN, TaskStatus.INTERRUPTED, and ExecutionLoop when making edits).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ai_company/engine/prompt.py`:
- Around line 628-660: The metadata dict in build_error_prompt duplicates logic
from _build_metadata; change build_error_prompt to call
_build_metadata(identity) and then set/override the "agent_id" key with the
provided agent_id so the rest of the fields come from the single source of truth
(_build_metadata) while preserving the ability to supply a different agent_id
for error paths; update references to SystemPrompt(metadata=...) in
build_error_prompt to use the merged metadata.
In `@src/ai_company/engine/shutdown.py`:
- Around line 366-378: Prevent new tasks from being added once shutdown has been
requested: in register_task check the shutdown flag used by
request_shutdown()/initiate_shutdown() (e.g., self._shutdown_requested or
equivalent) and, if shutdown is active, do not add to self._running_tasks;
instead log an appropriate message (EXECUTION_SHUTDOWN_GRACE_START with
action="task_rejected" and task_id), cancel or reject the passed asyncio_task
(so it won’t escape the grace window), and return/raise as appropriate. Apply
the same guard behavior to the related registration code referenced around the
other block (lines ~404-409) so no tasks can be enrolled after shutdown is
requested. Ensure initiate_shutdown still snapshots/operates on the tasks that
were registered prior to the shutdown flag being set.
- Around line 209-215: The code is reusing shutdown lifecycle event constants
(e.g., EXECUTION_SHUTDOWN_FORCE_CANCEL) for unrelated bookkeeping; define and
use distinct log/event constants (for example TASK_EXCEPTION_DURING_SHUTDOWN,
MANAGER_CONSTRUCTION_WARNING, TASK_REGISTRATION_NOTICE) and replace the
logger.warning calls that currently reference EXECUTION_SHUTDOWN_FORCE_CANCEL
(and the similar uses around the manager construction and task
registration/unregistration sites) so shutdown metrics/alerts remain accurate;
update the warning calls that include task.exception() and other messages to use
the new constants and keep the same error context in the log payload.
- Around line 205-228: The tasks_completed value is currently set to len(done)
which counts tasks that finished with exceptions; update the logic in shutdown
handling (in shutdown.py where variables done and pending are processed, e.g.,
the block that logs exceptions and cancels pending tasks) to compute
tasks_completed as the number of done tasks that exited cooperatively (i.e., not
cancelled and task.exception() is None) so that ShutdownResult.tasks_completed
reflects only successful cooperative exits; then return that filtered count
along with len(pending).
In `@tests/integration/engine/test_graceful_shutdown.py`:
- Around line 29-61: The provider currently requests shutdown on the first call
but still returns FinishReason.STOP so engine.run can complete normally; update
_ShutdownTriggeringProvider.complete so it only calls
self._strategy.request_shutdown() on the second call and, when shutdown is
requested, return an interrupted finish (e.g., FinishReason.INTERRUPTED) so
engine.run yields a shutdown/interruption result; also ensure the test
test_shutdown_interrupts_task_and_runs_cleanup explicitly calls
manager.initiate_shutdown() (or the equivalent shutdown trigger used in the
test) so cleanup_cb is executed. Include references to
_ShutdownTriggeringProvider.complete,
CooperativeTimeoutStrategy.request_shutdown, FinishReason,
manager.initiate_shutdown, and cleanup_cb when making the changes.
In `@tests/unit/engine/test_shutdown.py`:
- Around line 298-307: Update the test_handle_signal_threadsafe_with_loop to
execute the captured callback returned from mock_loop.call_soon_threadsafe and
then assert that the CooperativeTimeoutStrategy instance has entered shutdown
mode; specifically, after retrieving callback =
mock_loop.call_soon_threadsafe.call_args[0][0], call callback() (the _on_loop
closure) and assert the strategy reports shutdown (e.g., strategy.is_shutdown()
or strategy.shutdown_requested) so the test fails if _on_loop stops calling
request_shutdown() on ShutdownManager.
---
Outside diff comments:
In `@DESIGN_SPEC.md`:
- Around line 1028-1058: The spec conflicts with the implementation: currently
the code maps any TerminationReason.SHUTDOWN to TaskStatus.INTERRUPTED but
§6.7/§6.5 describe INTERRUPTED as only for force-cancelled tasks; reconcile by
updating the spec to match the implementation — explicitly state that
TerminationReason.SHUTDOWN results in TaskStatus.INTERRUPTED (including
cooperative and forced shutdown outcomes), adjust the lifecycle transition table
to show ASSIGNED/IN_PROGRESS → INTERRUPTED on shutdown, and remove or revise the
pre-shutdown wording in the ExecutionLoop section so it no longer contradicts
this mapping (refer to TerminationReason.SHUTDOWN, TaskStatus.INTERRUPTED, and
ExecutionLoop when making edits).
In `@src/ai_company/engine/react_loop.py`:
- Around line 115-123: The TurnRecord for a turn is being appended to turns
before tools are guaranteed to have executed, causing
tool_calls_made/total_tool_calls to include tools skipped by shutdown; change
the flow in react_loop so that you only call _make_turn_record and append it to
turns after invoke_all completes successfully and after shutdown_checker
confirms tools ran (i.e., move creation/appending of the TurnRecord out of the
pre-invoke path and into the post-invoke path inside self._process_turn_response
handling), and apply the same change to the duplicate block around the other
invoke_all usage (the second occurrence in this method) so both code paths
record tool_calls only after tools actually ran.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1a976a81-c8f7-4c12-961b-b96ad9d02076
📒 Files selected for processing (22)
CLAUDE.mdDESIGN_SPEC.mdsrc/ai_company/config/__init__.pysrc/ai_company/config/defaults.pysrc/ai_company/config/schema.pysrc/ai_company/core/enums.pysrc/ai_company/core/task_transitions.pysrc/ai_company/engine/__init__.pysrc/ai_company/engine/agent_engine.pysrc/ai_company/engine/loop_protocol.pysrc/ai_company/engine/prompt.pysrc/ai_company/engine/react_loop.pysrc/ai_company/engine/shutdown.pysrc/ai_company/observability/events/execution.pytests/integration/engine/test_graceful_shutdown.pytests/unit/core/test_enums.pytests/unit/core/test_task_transitions.pytests/unit/engine/test_agent_engine_lifecycle.pytests/unit/engine/test_loop_protocol.pytests/unit/engine/test_react_loop.pytests/unit/engine/test_run_result.pytests/unit/engine/test_shutdown.py
📜 Review details
⏰ 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: Agent
- GitHub Check: Greptile Review
🧰 Additional context used
📓 Path-based instructions (4)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Do NOT usefrom __future__ import annotationsin Python code — Python 3.14 has PEP 649 native lazy annotations.
Useexcept A, B:syntax (no parentheses) — ruff enforces PEP 758 except syntax on Python 3.14.
All public functions and classes must have type hints. Use mypy strict mode.
Use Google style docstrings, required on all public classes and functions (enforced by ruff D rules).
Create new objects instead of mutating existing ones. For non-Pydantic internal collections (registries,BaseTool), usecopy.deepcopy()at construction andMappingProxyTypewrapping for read-only enforcement.
Fordict/listfields in frozen Pydantic models, rely onfrozen=Truefor field reassignment prevention and usecopy.deepcopy()at system boundaries (tool execution, LLM provider serialization, inter-agent delegation, serializing for persistence).
Use frozen Pydantic models for config/identity; use separate mutable-via-copy models (usingmodel_copy(update=...)) for runtime state that evolves. Never mix static config fields with mutable runtime fields in one model.
Use Pydantic v2 (BaseModel,model_validator,computed_field,ConfigDict). Use@computed_fieldfor derived values instead of storing redundant fields. UseNotBlankStr(fromcore.types) for all identifier/name fields — including optional (NotBlankStr | None) and tuple (tuple[NotBlankStr, ...]) variants — instead of manual whitespace validators.
Preferasyncio.TaskGroupfor fan-out/fan-in parallel operations in new code (e.g. multiple tool invocations, parallel agent calls). Prefer structured concurrency over barecreate_task. Existing code is being migrated incrementally.
Keep functions under 50 lines and files under 800 lines.
Handle errors explicitly, never silently swallow exceptions.
Validate at system boundaries (user input, external APIs, config files).
Set line length to 88 characters (ruff).
Files:
tests/unit/core/test_enums.pytests/unit/core/test_task_transitions.pytests/unit/engine/test_loop_protocol.pysrc/ai_company/config/__init__.pytests/unit/engine/test_react_loop.pytests/unit/engine/test_agent_engine_lifecycle.pytests/unit/engine/test_run_result.pysrc/ai_company/engine/prompt.pysrc/ai_company/config/defaults.pysrc/ai_company/engine/__init__.pytests/unit/engine/test_shutdown.pysrc/ai_company/observability/events/execution.pysrc/ai_company/engine/shutdown.pytests/integration/engine/test_graceful_shutdown.pysrc/ai_company/engine/loop_protocol.pysrc/ai_company/config/schema.pysrc/ai_company/core/enums.pysrc/ai_company/engine/react_loop.pysrc/ai_company/core/task_transitions.pysrc/ai_company/engine/agent_engine.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Use test markers:@pytest.mark.unit,@pytest.mark.integration,@pytest.mark.e2e,@pytest.mark.slow.
Maintain 80% minimum code coverage (enforced in CI).
Prefer@pytest.mark.parametrizefor testing similar cases.
Files:
tests/unit/core/test_enums.pytests/unit/core/test_task_transitions.pytests/unit/engine/test_loop_protocol.pytests/unit/engine/test_react_loop.pytests/unit/engine/test_agent_engine_lifecycle.pytests/unit/engine/test_run_result.pytests/unit/engine/test_shutdown.pytests/integration/engine/test_graceful_shutdown.py
src/ai_company/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/ai_company/**/*.py: Every module with business logic MUST have:from ai_company.observability import get_loggerthenlogger = get_logger(__name__). Never useimport logging/logging.getLogger()/print()in application code.
Always useloggeras the variable name for the logger (not_logger, notlog).
Use event name constants from the domain-specific module underai_company.observability.events(e.g.PROVIDER_CALL_STARTfromevents.provider,BUDGET_RECORD_ADDEDfromevents.budget). Import directly:from ai_company.observability.events.<domain> import EVENT_CONSTANT.
Always use structured logging kwargs:logger.info(EVENT, key=value)— neverlogger.info("msg %s", val).
All error paths must log at WARNING or ERROR with context before raising.
All state transitions must log at INFO level.
DEBUG logging should be used for object creation, internal flow, and entry/exit of key functions.
Pure data models, enums, and re-exports do NOT need logging.
NEVER use real vendor names (Anthropic, OpenAI, Claude, GPT, etc.) in project-owned code, docstrings, comments, tests, or config examples. Use generic names:example-provider,example-large-001,example-medium-001,example-small-001,large/medium/smallas aliases. Vendor names may only appear in: (1) DESIGN_SPEC.md provider list, (2).claude/skill/agent files, (3) third-party import paths/module names. Tests must usetest-provider,test-small-001, etc.
Files:
src/ai_company/config/__init__.pysrc/ai_company/engine/prompt.pysrc/ai_company/config/defaults.pysrc/ai_company/engine/__init__.pysrc/ai_company/observability/events/execution.pysrc/ai_company/engine/shutdown.pysrc/ai_company/engine/loop_protocol.pysrc/ai_company/config/schema.pysrc/ai_company/core/enums.pysrc/ai_company/engine/react_loop.pysrc/ai_company/core/task_transitions.pysrc/ai_company/engine/agent_engine.py
src/ai_company/engine/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
RetryExhaustedErrorsignals that all retries failed — the engine layer catches this to trigger fallback chains.
Files:
src/ai_company/engine/prompt.pysrc/ai_company/engine/__init__.pysrc/ai_company/engine/shutdown.pysrc/ai_company/engine/loop_protocol.pysrc/ai_company/engine/react_loop.pysrc/ai_company/engine/agent_engine.py
🧠 Learnings (5)
📚 Learning: 2026-03-07T10:57:52.980Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-07T10:57:52.980Z
Learning: Applies to src/ai_company/providers/**/*.py : RetryConfig and RateLimiterConfig are set per-provider in `ProviderConfig`.
Applied to files:
src/ai_company/config/__init__.py
📚 Learning: 2026-03-07T10:57:52.980Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-07T10:57:52.980Z
Learning: Applies to src/ai_company/**/*.py : Use event name constants from the domain-specific module under `ai_company.observability.events` (e.g. `PROVIDER_CALL_START` from `events.provider`, `BUDGET_RECORD_ADDED` from `events.budget`). Import directly: `from ai_company.observability.events.<domain> import EVENT_CONSTANT`.
Applied to files:
src/ai_company/observability/events/execution.pysrc/ai_company/engine/react_loop.py
📚 Learning: 2026-03-07T10:57:52.980Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-07T10:57:52.980Z
Learning: Applies to **/*.py : Use Pydantic v2 (`BaseModel`, `model_validator`, `computed_field`, `ConfigDict`). Use `computed_field` for derived values instead of storing redundant fields. Use `NotBlankStr` (from `core.types`) for all identifier/name fields — including optional (`NotBlankStr | None`) and tuple (`tuple[NotBlankStr, ...]`) variants — instead of manual whitespace validators.
Applied to files:
src/ai_company/engine/loop_protocol.py
📚 Learning: 2026-03-07T10:57:52.980Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-07T10:57:52.980Z
Learning: Applies to **/*.py : Prefer `asyncio.TaskGroup` for fan-out/fan-in parallel operations in new code (e.g. multiple tool invocations, parallel agent calls). Prefer structured concurrency over bare `create_task`. Existing code is being migrated incrementally.
Applied to files:
DESIGN_SPEC.md
📚 Learning: 2026-03-07T10:57:52.980Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-07T10:57:52.980Z
Learning: Applies to src/ai_company/**/*.py : All state transitions must log at INFO level.
Applied to files:
src/ai_company/core/task_transitions.py
🧬 Code graph analysis (16)
tests/unit/core/test_enums.py (1)
src/ai_company/core/enums.py (1)
TaskStatus(122-148)
tests/unit/core/test_task_transitions.py (1)
src/ai_company/core/enums.py (1)
TaskStatus(122-148)
tests/unit/engine/test_loop_protocol.py (1)
src/ai_company/engine/loop_protocol.py (1)
TerminationReason(26-33)
src/ai_company/config/__init__.py (1)
src/ai_company/config/schema.py (1)
GracefulShutdownConfig(324-352)
tests/unit/engine/test_agent_engine_lifecycle.py (3)
src/ai_company/core/task.py (1)
Task(38-212)src/ai_company/engine/context.py (3)
AgentContext(87-307)from_identity(140-171)with_task_transition(235-278)src/ai_company/core/enums.py (1)
TaskStatus(122-148)
tests/unit/engine/test_run_result.py (1)
src/ai_company/engine/loop_protocol.py (1)
make_budget_checker(173-188)
src/ai_company/engine/prompt.py (1)
src/ai_company/core/agent.py (1)
AgentIdentity(177-235)
src/ai_company/engine/__init__.py (1)
src/ai_company/engine/shutdown.py (4)
CooperativeTimeoutStrategy(110-271)ShutdownManager(274-409)ShutdownResult(37-71)ShutdownStrategy(75-107)
tests/unit/engine/test_shutdown.py (2)
src/ai_company/config/schema.py (1)
GracefulShutdownConfig(324-352)src/ai_company/engine/shutdown.py (19)
ShutdownResult(37-71)ShutdownStrategy(75-107)strategy(299-301)is_shutting_down(82-84)is_shutting_down(139-141)is_shutting_down(400-402)request_shutdown(78-80)request_shutdown(135-137)get_strategy_type(105-107)get_strategy_type(143-145)execute_shutdown(86-103)execute_shutdown(147-185)register_task(366-378)unregister_task(380-388)register_cleanup(390-398)install_signal_handlers(303-320)initiate_shutdown(404-409)_handle_signal(322-335)_handle_signal_threadsafe(337-364)
src/ai_company/engine/shutdown.py (1)
src/ai_company/observability/_logger.py (1)
get_logger(8-28)
tests/integration/engine/test_graceful_shutdown.py (8)
src/ai_company/core/enums.py (2)
SeniorityLevel(6-21)ToolCategory(218-232)src/ai_company/engine/shutdown.py (9)
CooperativeTimeoutStrategy(110-271)ShutdownManager(274-409)strategy(299-301)request_shutdown(78-80)request_shutdown(135-137)register_cleanup(390-398)is_shutting_down(82-84)is_shutting_down(139-141)is_shutting_down(400-402)src/ai_company/providers/enums.py (1)
FinishReason(15-22)src/ai_company/providers/capabilities.py (1)
ModelCapabilities(10-83)src/ai_company/core/agent.py (2)
AgentIdentity(177-235)ModelConfig(76-105)src/ai_company/tools/base.py (4)
description(115-117)BaseTool(56-161)ToolExecutionResult(24-53)category(110-112)src/ai_company/tools/invoker.py (1)
registry(100-102)src/ai_company/engine/run_result.py (1)
termination_reason(64-66)
src/ai_company/engine/loop_protocol.py (3)
src/ai_company/core/task.py (1)
Task(38-212)src/ai_company/engine/react_loop.py (1)
execute(60-135)src/ai_company/engine/context.py (1)
AgentContext(87-307)
src/ai_company/config/schema.py (2)
src/ai_company/engine/shutdown.py (1)
strategy(299-301)src/ai_company/tools/base.py (1)
description(115-117)
src/ai_company/engine/react_loop.py (2)
src/ai_company/engine/loop_protocol.py (1)
execute(140-166)src/ai_company/engine/context.py (1)
AgentContext(87-307)
src/ai_company/core/task_transitions.py (1)
src/ai_company/core/enums.py (1)
TaskStatus(122-148)
src/ai_company/engine/agent_engine.py (3)
src/ai_company/engine/loop_protocol.py (2)
make_budget_checker(173-188)TerminationReason(26-33)src/ai_company/engine/prompt.py (2)
SystemPrompt(55-81)build_error_prompt(628-660)src/ai_company/engine/errors.py (1)
ExecutionStateError(12-13)
🔇 Additional comments (10)
src/ai_company/engine/__init__.py (1)
22-50: LGTM!The shutdown-related re-exports are correctly imported from their source modules and properly added to
__all__in alphabetical order. The public API surface for the graceful shutdown framework is now cleanly exposed.Also applies to: 63-64, 78-81
src/ai_company/config/defaults.py (1)
28-28: LGTM!The empty dict default for
graceful_shutdownis consistent with other config sections (e.g.,budget,communication) and correctly allowsGracefulShutdownConfigto use its field defaults when no user overrides are provided.tests/unit/engine/test_loop_protocol.py (1)
21-29: LGTM!The tests correctly validate the new
SHUTDOWNtermination reason and update the member count expectation. The test ordering follows the enum definition order.tests/unit/core/test_enums.py (1)
61-62: LGTM!The tests correctly validate the new
INTERRUPTEDtask status enum member and update the member count. The test placement is consistent with the enum definition order.Also applies to: 114-114
src/ai_company/config/__init__.py (1)
13-13: LGTM!The
GracefulShutdownConfigexport is correctly added to the docstring autosummary, imported from the schema module, and included in__all__in alphabetical order.Also applies to: 41-41, 56-56
tests/unit/engine/test_react_loop.py (2)
861-968: LGTM!The
TestReactLoopShutdowntest suite comprehensively covers the shutdown checker integration:
- Immediate shutdown before any LLM calls
- Shutdown after a completed turn with tool execution
- Shutdown detected between LLM response and tool invocation (turn-boundary check)
- Normal completion path when no shutdown checker is provided
The stateful checkers with
call_counttracking clearly validate the expected check points in the loop.
970-988: LGTM!The cost accounting test correctly verifies that error responses (like
content_filter) still propagate their token costs into the resulting context, ensuring cost tracking integrity even on failure paths.src/ai_company/config/schema.py (2)
324-352: LGTM!The
GracefulShutdownConfigmodel is well-designed:
- Frozen with
allow_inf_nan=Falseconsistent with other numeric config models (e.g.,RetryConfig,RateLimiterConfig)- Uses
NotBlankStrfor the strategy field per coding guidelines- Constraints are sensible (
grace_seconds ≤ 300,cleanup_seconds ≤ 60)- Defaults align with the PR objectives (30s grace, 5s cleanup)
373-373: LGTM!The
graceful_shutdownfield is correctly integrated intoRootConfig:
- Uses
default_factory=GracefulShutdownConfigto provide defaults- Added to the class docstring attributes
- Follows the same pattern as other nested config fields (
budget,communication, etc.)Also applies to: 421-424
src/ai_company/engine/agent_engine.py (1)
540-543: The SHUTDOWN-to-INTERRUPTED transition assumes ShutdownManager does not manage engine.run() tasks.This transition only triggers when the loop returns
TerminationReason.SHUTDOWN(cooperative shutdown detected inside the loop viashutdown_checkercallback). However, if an external orchestration layer registers engine tasks withShutdownManagerand initiates shutdown, force-cancelled tasks receiveasyncio.CancelledErrorbefore anyExecutionResultis created, bypassing the INTERRUPTED transition entirely.The current codebase shows no integration between ShutdownManager and engine task registration—ShutdownManager exposes
register_task()/unregister_task()for external callers to manage, but no code insrc/ai_companyregisters engine.run() tasks with it. If such integration is planned, an outer wrapper must catchCancelledErrorfrom forced cancellations and update task status to INTERRUPTED before the task is discarded.
| def build_error_prompt( | ||
| identity: AgentIdentity, | ||
| agent_id: str, | ||
| system_prompt: SystemPrompt | None, | ||
| ) -> SystemPrompt: | ||
| """Return the existing system prompt or a minimal error placeholder. | ||
|
|
||
| Used by the engine when the execution pipeline fails and a | ||
| ``SystemPrompt`` was never built (or was partially built). | ||
|
|
||
| Args: | ||
| identity: Agent identity for metadata. | ||
| agent_id: String agent identifier. | ||
| system_prompt: Previously built prompt, or ``None``. | ||
|
|
||
| Returns: | ||
| The existing prompt if available, else a minimal placeholder. | ||
| """ | ||
| if system_prompt is not None: | ||
| return system_prompt | ||
| return SystemPrompt( | ||
| content="", | ||
| template_version="error", | ||
| estimated_tokens=0, | ||
| sections=(), | ||
| metadata={ | ||
| "agent_id": agent_id, | ||
| "name": identity.name, | ||
| "role": identity.role, | ||
| "department": identity.department, | ||
| "level": identity.level.value, | ||
| }, | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider reusing _build_metadata to reduce duplication.
The metadata dictionary construction duplicates the logic in _build_metadata (lines 421-436). While the agent_id parameter allows flexibility in error paths, the other fields come directly from identity.
♻️ Optional refactor to reduce duplication
def build_error_prompt(
identity: AgentIdentity,
agent_id: str,
system_prompt: SystemPrompt | None,
) -> SystemPrompt:
"""Return the existing system prompt or a minimal error placeholder.
Used by the engine when the execution pipeline fails and a
``SystemPrompt`` was never built (or was partially built).
Args:
identity: Agent identity for metadata.
agent_id: String agent identifier.
system_prompt: Previously built prompt, or ``None``.
Returns:
The existing prompt if available, else a minimal placeholder.
"""
if system_prompt is not None:
return system_prompt
+ metadata = _build_metadata(identity)
+ metadata["agent_id"] = agent_id # Override with provided agent_id
return SystemPrompt(
content="",
template_version="error",
estimated_tokens=0,
sections=(),
- metadata={
- "agent_id": agent_id,
- "name": identity.name,
- "role": identity.role,
- "department": identity.department,
- "level": identity.level.value,
- },
+ metadata=metadata,
)Note: If the agent_id parameter is intentionally separate from identity.id for error recovery scenarios, the current explicit approach is acceptable.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/engine/prompt.py` around lines 628 - 660, The metadata dict in
build_error_prompt duplicates logic from _build_metadata; change
build_error_prompt to call _build_metadata(identity) and then set/override the
"agent_id" key with the provided agent_id so the rest of the fields come from
the single source of truth (_build_metadata) while preserving the ability to
supply a different agent_id for error paths; update references to
SystemPrompt(metadata=...) in build_error_prompt to use the merged metadata.
| def register_task( | ||
| self, | ||
| task_id: str, | ||
| asyncio_task: asyncio.Task[Any], | ||
| ) -> None: | ||
| """Track a running agent task.""" | ||
| self._running_tasks[task_id] = asyncio_task | ||
| logger.debug( | ||
| EXECUTION_SHUTDOWN_GRACE_START, | ||
| action="task_registered", | ||
| task_id=task_id, | ||
| running_tasks=len(self._running_tasks), | ||
| ) |
There was a problem hiding this comment.
Stop accepting new tasks once shutdown starts.
register_task() still enrolls new executions after request_shutdown(), and initiate_shutdown() snapshots only the tasks already present at call time. Anything registered during the grace window can miss cancellation entirely, which breaks the drain-gate behavior described in §6.7.
Also applies to: 404-409
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/engine/shutdown.py` around lines 366 - 378, Prevent new tasks
from being added once shutdown has been requested: in register_task check the
shutdown flag used by request_shutdown()/initiate_shutdown() (e.g.,
self._shutdown_requested or equivalent) and, if shutdown is active, do not add
to self._running_tasks; instead log an appropriate message
(EXECUTION_SHUTDOWN_GRACE_START with action="task_rejected" and task_id), cancel
or reject the passed asyncio_task (so it won’t escape the grace window), and
return/raise as appropriate. Apply the same guard behavior to the related
registration code referenced around the other block (lines ~404-409) so no tasks
can be enrolled after shutdown is requested. Ensure initiate_shutdown still
snapshots/operates on the tasks that were registered prior to the shutdown flag
being set.
| class _ShutdownTriggeringProvider: | ||
| """Provider that triggers shutdown on the second call. | ||
|
|
||
| First call returns a tool-use response, second call triggers | ||
| shutdown and returns a stop response. | ||
| """ | ||
|
|
||
| def __init__(self, strategy: CooperativeTimeoutStrategy) -> None: | ||
| self._strategy = strategy | ||
| self._call_count = 0 | ||
|
|
||
| async def complete( | ||
| self, | ||
| messages: list[ChatMessage], | ||
| model: str, | ||
| *, | ||
| tools: list[ToolDefinition] | None = None, | ||
| config: CompletionConfig | None = None, | ||
| ) -> CompletionResponse: | ||
| self._call_count += 1 | ||
| if self._call_count == 1: | ||
| # Trigger shutdown after first LLM call | ||
| self._strategy.request_shutdown() | ||
| return CompletionResponse( | ||
| content="Working on it.", | ||
| finish_reason=FinishReason.STOP, | ||
| usage=TokenUsage( | ||
| input_tokens=50, | ||
| output_tokens=25, | ||
| cost_usd=0.005, | ||
| ), | ||
| model="test-model-001", | ||
| ) |
There was a problem hiding this comment.
This integration test never exercises interruption or cleanup.
Line 49 sets the shutdown flag on the first provider call but still returns FinishReason.STOP, so engine.run() can complete without ever producing a shutdown result. The test also never calls manager.initiate_shutdown(), so cleanup_cb is never run. As written, test_shutdown_interrupts_task_and_runs_cleanup can pass while both behaviors in its name are broken.
Also applies to: 93-162
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/integration/engine/test_graceful_shutdown.py` around lines 29 - 61, The
provider currently requests shutdown on the first call but still returns
FinishReason.STOP so engine.run can complete normally; update
_ShutdownTriggeringProvider.complete so it only calls
self._strategy.request_shutdown() on the second call and, when shutdown is
requested, return an interrupted finish (e.g., FinishReason.INTERRUPTED) so
engine.run yields a shutdown/interruption result; also ensure the test
test_shutdown_interrupts_task_and_runs_cleanup explicitly calls
manager.initiate_shutdown() (or the equivalent shutdown trigger used in the
test) so cleanup_cb is executed. Include references to
_ShutdownTriggeringProvider.complete,
CooperativeTimeoutStrategy.request_shutdown, FinishReason,
manager.initiate_shutdown, and cleanup_cb when making the changes.
There was a problem hiding this comment.
Code Review
This pull request introduces a robust and well-designed graceful shutdown mechanism. The implementation is comprehensive, covering everything from signal handling and task lifecycle updates to configuration and extensive testing. The separation of concerns into a ShutdownStrategy protocol and a ShutdownManager is excellent, allowing for future extensibility. The addition of the INTERRUPTED task status and its integration into the agent engine and execution loop are well-handled.
I've identified one critical syntax issue that will prevent the code from running and one medium-severity suggestion to improve maintainability. Once these are addressed, this will be a solid contribution to the project.
Note: Security Review is unavailable for this PR.
| except MemoryError, RecursionError: | ||
| raise |
There was a problem hiding this comment.
This except syntax is from Python 2 and is invalid in Python 3, which will cause a SyntaxError. To catch multiple exceptions, they should be enclosed in a tuple. Other parts of the codebase correctly use except (ExceptionA, ExceptionB):, so this seems to be an inconsistency.
except (MemoryError, RecursionError):
raise
src/ai_company/engine/shutdown.py
Outdated
| for task in pending: | ||
| task.cancel() | ||
| # Wait for cancellation to propagate (bounded) | ||
| await asyncio.wait(pending, timeout=5.0) |
There was a problem hiding this comment.
The timeout value 5.0 for waiting for task cancellation is hardcoded. It's better to define this as a named constant at the module level to improve readability and make it easier to change if needed.
For example, you could add at the top of the file:
_CANCELLATION_WAIT_SECONDS = 5.0And then use it here.
| await asyncio.wait(pending, timeout=5.0) | |
| await asyncio.wait(pending, timeout=_CANCELLATION_WAIT_SECONDS) |
- Add drain gate to ShutdownManager.register_task (rejects after shutdown) - Add input_token_estimate to turn-start log (char_count // 4 heuristic) - Fix semantic misuse of event constants in shutdown module - Fix tasks_completed counting (exclude errored tasks) - Add exception guards to Windows signal handler callback - Add exception guard to Windows signal handler fallback path - Fix raise chain: raise exc from build_exc in agent_engine - Fix TurnRecord tool_calls_made when shutdown preempts tool execution - Rework integration test_shutdown_signal_propagates_through_manager - Add is_success=False test for SHUTDOWN termination reason - Add shutdown_checker exception test (returns ERROR) - Retrieve post-cancellation exceptions to prevent warnings - Add CANCELLED transition arrows to §6.1 lifecycle diagram - Add duplicate task_id warning in register_task - Fix module docstring: shutdown module doesn't mark tasks INTERRUPTED - Add Returns section to _run_cleanup docstring - Clarify tasks_interrupted field description - Fix _frame type annotation to types.FrameType | None - Clarify §6.7: ALL shutdown-terminated tasks become INTERRUPTED - Add MemoryError propagation test for shutdown_checker - Add _transition_to_interrupted failure path test - Fix test_handle_signal_threadsafe to execute callback - Fix provider docstring (triggers on first call, not second) - Use Mapping/Sequence in protocol execute_shutdown params - Use ValidationError instead of Exception in frozen test - Use named constant for cancel propagation timeout - Fix get_strategy_type return type alignment with NotBlankStr Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
| except MemoryError, RecursionError: | ||
| raise |
There was a problem hiding this comment.
Missing parentheses for multiple exception types
except MemoryError, RecursionError: is the Python 2 syntax for binding an exception to a name (except ExcType, varname:). In Python 3 this requires explicit parentheses to catch multiple types. Without them, the comma may be parsed as a binding expression rather than a tuple on older Python 3 versions, which would produce a SyntaxError at import time and prevent the module from loading.
The exact same pattern exists in _check_budget at line 253 — both should be fixed together.
| except MemoryError, RecursionError: | |
| raise | |
| except (MemoryError, RecursionError): |
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/ai_company/engine/react_loop.py
Line: 217-218
Comment:
**Missing parentheses for multiple exception types**
`except MemoryError, RecursionError:` is the Python 2 syntax for binding an exception to a name (`except ExcType, varname:`). In Python 3 this requires explicit parentheses to catch multiple types. Without them, the comma may be parsed as a binding expression rather than a tuple on older Python 3 versions, which would produce a `SyntaxError` at import time and prevent the module from loading.
The exact same pattern exists in `_check_budget` at line 253 — both should be fixed together.
```suggestion
except (MemoryError, RecursionError):
```
How can I resolve this? If you propose a fix, please make it concise.| logger.info( | ||
| EXECUTION_SHUTDOWN_CLEANUP, | ||
| callback_count=len(callbacks), | ||
| cleanup_seconds=self._cleanup_seconds, | ||
| ) | ||
|
|
||
| all_succeeded = True | ||
|
|
||
| async def _run_all() -> None: | ||
| nonlocal all_succeeded | ||
| for i, callback in enumerate(callbacks): | ||
| try: | ||
| await callback() | ||
| except Exception: | ||
| all_succeeded = False | ||
| logger.exception( | ||
| EXECUTION_SHUTDOWN_CLEANUP, | ||
| callback_index=i, | ||
| callback_count=len(callbacks), | ||
| error="Cleanup callback failed", | ||
| ) |
There was a problem hiding this comment.
EXECUTION_SHUTDOWN_CLEANUP reused for semantically distinct events
EXECUTION_SHUTDOWN_CLEANUP is logged at two different semantic points in _run_cleanup:
- Line 267 — info: "cleanup phase is starting" (callback count + time budget)
- Line 281 — exception: "a cleanup callback raised"
Using the same event string for both makes structured log queries ambiguous — filtering on "execution.shutdown.cleanup" will return both successful-start entries and failure entries, making it hard to build reliable alerts or dashboards for cleanup failures.
A dedicated constant (e.g. EXECUTION_SHUTDOWN_CLEANUP_FAILED) would follow the existing pattern where EXECUTION_SHUTDOWN_CLEANUP_TIMEOUT is already a separate constant for the timeout path.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/ai_company/engine/shutdown.py
Line: 266-286
Comment:
**`EXECUTION_SHUTDOWN_CLEANUP` reused for semantically distinct events**
`EXECUTION_SHUTDOWN_CLEANUP` is logged at two different semantic points in `_run_cleanup`:
1. **Line 267** — info: "cleanup phase is starting" (callback count + time budget)
2. **Line 281** — exception: "a cleanup callback raised"
Using the same event string for both makes structured log queries ambiguous — filtering on `"execution.shutdown.cleanup"` will return both successful-start entries and failure entries, making it hard to build reliable alerts or dashboards for cleanup failures.
A dedicated constant (e.g. `EXECUTION_SHUTDOWN_CLEANUP_FAILED`) would follow the existing pattern where `EXECUTION_SHUTDOWN_CLEANUP_TIMEOUT` is already a separate constant for the timeout path.
How can I resolve this? If you propose a fix, please make it concise.CI fixes: - Remove unused type: ignore on Linux (use dual suppression for cross-platform compat) in subprocess_sandbox.py - Fix test assertion for raise-from-build_exc chaining in test_agent_engine_errors.py Review feedback (Greptile): - Rename _check_ref → _check_git_arg with updated docstring (was misleadingly named when used for author/date validation) - Fix missing space in clone URL error message - Add EXECUTION_SHUTDOWN_CLEANUP_FAILED event constant (was reusing EXECUTION_SHUTDOWN_CLEANUP for both start and failure) - Capture partial stderr on timeout in direct subprocess path (mirrors sandbox behavior for better debuggability) - Make Windows Git PATH prefixes configurable via SubprocessSandboxConfig.extra_safe_path_prefixes - Add clarifying comment for total_entries semantics when truncated - Fix ResourceWarning on Windows by explicitly closing subprocess transports after communicate() Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Summary
ShutdownStrategyprotocol,CooperativeTimeoutStrategy,ShutdownManagerwith cross-platform signal handling (Unixadd_signal_handler+ Windowssignal.signalfallback)INTERRUPTEDstatus with transitions fromASSIGNED/IN_PROGRESS, reassignable viaINTERRUPTED → ASSIGNEDon restartReactLoopchecks shutdown at loop top and before tool invocations;_check_shutdownhas full exception handling matching_check_budgetpatterncleanup_secondstimeoutGracefulShutdownConfigwith validated bounds (grace_seconds≤ 300,cleanup_seconds≤ 60,NotBlankStrstrategy)Pre-PR review findings addressed (10 agents, 30+ findings)
call_soon_threadsafeto avoid structlog deadlocksCooperativeTimeoutStrategyrejects non-positivegrace_seconds/cleanup_secondsShutdownResultmodel — addedallow_inf_nan=Falsefor float fieldsbuild_error_prompt→engine/prompt.pyandmake_budget_checker→engine/loop_protocol.py(agent_engine.py now under 800 lines)_wait_and_cancelto prevent "exception never retrieved" warningsasyncio.waitwithtimeout=5.0shutdown.py), §15.5 Adopted status; CLAUDE.md engine descriptionFiles changed (14)
engine/shutdown.pyengine/react_loop.pyengine/agent_engine.pyengine/loop_protocol.pymake_budget_checker,Taskimportengine/prompt.pybuild_error_promptextracted hereconfig/schema.pyGracefulShutdownConfigvalidation boundscore/task_transitions.pyobservability/events/execution.pyEXECUTION_SHUTDOWN_CLEANUP_TIMEOUTeventDESIGN_SPEC.mdCLAUDE.mdTest plan
uv run ruff check src/ tests/— all checks passeduv run ruff format src/ tests/— formatteduv run mypy src/ tests/— no issues in 260 filesuv run pytest tests/ -n auto --cov=ai_company --cov-fail-under=80— 2370 passed, 96.33% coverageCloses #130
🤖 Generated with Claude Code