fix: architectural gaps in wrapper package - thread safety, async safety, feature parity#1771
fix: architectural gaps in wrapper package - thread safety, async safety, feature parity#1771praisonai-triage-agent[bot] wants to merge 2 commits into
Conversation
…ety, feature parity - Fix TOCTOU race in jobs server singleton initialization with thread-safe locking - Replace deprecated asyncio.get_event_loop() calls with asyncio.run() and asyncio.to_thread() - Restore AsyncAgentScheduler feature parity: add missing timeout and max_cost parameters - Ensure async/sync scheduler classes have identical safety features Fixes #1770 Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>
|
@coderabbitai review |
|
/review |
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
✅ Actions performedReview triggered.
|
|
Important Review skippedBot user detected. To trigger a single review, invoke the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR modernizes asyncio API usage, fixes thread-safety in global initialization, and adds timeout/budget support to the async agent scheduler. Deprecated ChangesAsyncIO modernization and scheduler budget feature
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
@copilot Do a thorough review of this PR. Read ALL existing reviewer comments above from Qodo, Coderabbit, and Gemini first — incorporate their findings. Review areas:
|
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/praisonai/praisonai/scheduler/async_agent_scheduler.py (1)
96-118:⚠️ Potential issue | 🟠 Major | ⚡ Quick winExpose
timeoutandmax_costthrough the factory API too.
AsyncAgentScheduler.__init__now supports the new options, butcreate_async_agent_scheduler()still only forwards(agent, task, config). Factory-based callers still cannot configure the restored features.Suggested fix
def create_async_agent_scheduler( agent, task: str, - config: Optional[Dict[str, Any]] = None + config: Optional[Dict[str, Any]] = None, + on_success: Optional[Callable] = None, + on_failure: Optional[Callable] = None, + timeout: Optional[int] = None, + max_cost: Optional[float] = 1.00, ) -> AsyncAgentScheduler: @@ - return AsyncAgentScheduler(agent, task, config) + return AsyncAgentScheduler( + agent, + task, + config=config, + on_success=on_success, + on_failure=on_failure, + timeout=timeout, + max_cost=max_cost, + )🤖 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 `@src/praisonai/praisonai/scheduler/async_agent_scheduler.py` around lines 96 - 118, The factory function create_async_agent_scheduler currently only forwards (agent, task, config) so callers cannot set the new AsyncAgentScheduler.__init__ options; update create_async_agent_scheduler to accept optional timeout: Optional[int] = None and max_cost: Optional[float] = 1.00 (or accept **kwargs that include them) and forward those values into AsyncAgentScheduler(...) when constructing the scheduler, and update any internal call sites to pass through or surface these parameters so factory-based callers can configure timeout and max_cost.src/praisonai/praisonai/cli/commands/standardise.py (1)
371-382:⚠️ Potential issue | 🟠 Major | ⚡ Quick winKeep the runtime lifecycle on one event loop.
_run_aiinsrc/praisonai/praisonai/cli/commands/standardise.pystarts the runtime withasyncio.run(runtime.start())(~381) and later stops it withasyncio.run(runtime.stop())(~491).InteractiveRuntime.start()/stop()both await asyncio-based subsystem work (LSP_lsp_client.start()/stop()viaasyncio.gather), so splitting across event loops can break shutdown/cleanup and leave loop-bound tasks/handles in an inconsistent state. Wrap start → generation → stop in a single async helper and callasyncio.run(...)once around the whole lifecycle.🤖 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 `@src/praisonai/praisonai/cli/commands/standardise.py` around lines 371 - 382, The code in _run_ai currently calls asyncio.run(runtime.start()) and later asyncio.run(runtime.stop()), which creates and tears down separate event loops and can leave loop-bound tasks in a bad state; instead create a single async helper (e.g., async def _run_runtime_and_generate(...)) that awaits runtime.start(), performs the generation steps, then awaits runtime.stop(), and replace the two asyncio.run calls with a single asyncio.run(_run_runtime_and_generate(...)); reference InteractiveRuntime.start, InteractiveRuntime.stop, runtime.start(), runtime.stop(), and _run_ai to locate where to wrap the lifecycle.src/praisonai/praisonai/sandbox/e2b.py (1)
215-233:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFinish the asyncio migration for
envexports in E2B sandbox
src/praisonai/praisonai/sandbox/e2b.pystill usesawait asyncio.get_event_loop().run_in_executor(...)in theenvexport loop inside_execute_bash_command; switch that loop toasyncio.to_thread(...).src/praisonai/tests/unit/sandbox/test_e2b_sandbox.pystill patchesasyncio.get_event_loop().run_in_executorin multiple execution tests; update those mocks to align with theasyncio.to_thread-based code paths (and only mock env-export behavior appropriately).🤖 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 `@src/praisonai/praisonai/sandbox/e2b.py` around lines 215 - 233, The env export loop in _execute_bash_command still uses asyncio.get_event_loop().run_in_executor; replace that per-item run_in_executor call with asyncio.to_thread calls that call self._sandbox.commands.run (mirroring the working_dir and command calls), i.e., for each key,value use await asyncio.to_thread(self._sandbox.commands.run, f"export {shlex.quote(key)}={shlex.quote(value)}", timeout=5). Then update unit tests in src/praisonai/tests/unit/sandbox/test_e2b_sandbox.py to stop patching asyncio.get_event_loop().run_in_executor and instead mock asyncio.to_thread (or directly patch self._sandbox.commands.run for env-export behavior) so tests only intercept the env export calls while the rest of _execute_bash_command continues to use asyncio.to_thread as implemented.
🤖 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 `@src/praisonai/praisonai/scheduler/async_agent_scheduler.py`:
- Around line 361-390: The budget-check branch in _execute_with_retry currently
only returns, so _run_schedule keeps calling it and failures don't increment
_total_cost; change it to set the scheduler stop event (e.g.,
self._stop_event.set()) when budget is exhausted and return so the scheduler
halts; also deduct/charge an estimated_cost (or per-attempt charge) before each
attempt (inside _execute_with_retry, before retries loop or at start of each
attempt) while holding self._stats_lock so _total_cost is updated even for
failed/timeouts and prevent starting an attempt that would exceed self.max_cost
by checking self._total_cost + estimated_cost >= self.max_cost and setting
self._stop_event if exceeded.
---
Outside diff comments:
In `@src/praisonai/praisonai/cli/commands/standardise.py`:
- Around line 371-382: The code in _run_ai currently calls
asyncio.run(runtime.start()) and later asyncio.run(runtime.stop()), which
creates and tears down separate event loops and can leave loop-bound tasks in a
bad state; instead create a single async helper (e.g., async def
_run_runtime_and_generate(...)) that awaits runtime.start(), performs the
generation steps, then awaits runtime.stop(), and replace the two asyncio.run
calls with a single asyncio.run(_run_runtime_and_generate(...)); reference
InteractiveRuntime.start, InteractiveRuntime.stop, runtime.start(),
runtime.stop(), and _run_ai to locate where to wrap the lifecycle.
In `@src/praisonai/praisonai/sandbox/e2b.py`:
- Around line 215-233: The env export loop in _execute_bash_command still uses
asyncio.get_event_loop().run_in_executor; replace that per-item run_in_executor
call with asyncio.to_thread calls that call self._sandbox.commands.run
(mirroring the working_dir and command calls), i.e., for each key,value use
await asyncio.to_thread(self._sandbox.commands.run, f"export
{shlex.quote(key)}={shlex.quote(value)}", timeout=5). Then update unit tests in
src/praisonai/tests/unit/sandbox/test_e2b_sandbox.py to stop patching
asyncio.get_event_loop().run_in_executor and instead mock asyncio.to_thread (or
directly patch self._sandbox.commands.run for env-export behavior) so tests only
intercept the env export calls while the rest of _execute_bash_command continues
to use asyncio.to_thread as implemented.
In `@src/praisonai/praisonai/scheduler/async_agent_scheduler.py`:
- Around line 96-118: The factory function create_async_agent_scheduler
currently only forwards (agent, task, config) so callers cannot set the new
AsyncAgentScheduler.__init__ options; update create_async_agent_scheduler to
accept optional timeout: Optional[int] = None and max_cost: Optional[float] =
1.00 (or accept **kwargs that include them) and forward those values into
AsyncAgentScheduler(...) when constructing the scheduler, and update any
internal call sites to pass through or surface these parameters so factory-based
callers can configure timeout and max_cost.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: d998db42-af96-4169-8aee-f9b3f3b21591
📒 Files selected for processing (5)
src/praisonai/praisonai/cli/commands/standardise.pysrc/praisonai/praisonai/jobs/executor.pysrc/praisonai/praisonai/jobs/server.pysrc/praisonai/praisonai/sandbox/e2b.pysrc/praisonai/praisonai/scheduler/async_agent_scheduler.py
| # Check budget limit | ||
| if self.max_cost and self._total_cost >= self.max_cost: | ||
| logger.warning(f"Budget limit reached: ${self._total_cost:.4f} >= ${self.max_cost}") | ||
| logger.warning("Stopping scheduler to prevent additional costs") | ||
| return | ||
|
|
||
| last_exc: Optional[Exception] = None | ||
| for attempt in range(max_retries): | ||
| try: | ||
| logger.info(f"Async attempt {attempt + 1}/{max_retries}") | ||
|
|
||
| # TODO: Add timeout support from sync version: | ||
| # if self.timeout: | ||
| # result = await asyncio.wait_for( | ||
| # self._executor.execute(self.task), | ||
| # timeout=self.timeout | ||
| # ) | ||
| # else: | ||
| result = await self._executor.execute(self.task) | ||
| # Execute with timeout if specified | ||
| if self.timeout: | ||
| result = await asyncio.wait_for( | ||
| self._executor.execute(self.task), | ||
| timeout=self.timeout | ||
| ) | ||
| else: | ||
| result = await self._executor.execute(self.task) | ||
|
|
||
| logger.info(f"Async agent execution successful on attempt {attempt + 1}") | ||
| logger.info(f"Result: {result}") | ||
|
|
||
| # Estimate cost (rough: ~$0.0001 per execution for gpt-4o-mini) | ||
| estimated_cost = 0.0001 # Base cost estimate | ||
|
|
||
| async with self._stats_lock: | ||
| self._success_count += 1 | ||
| # TODO: Add cost tracking from sync version: | ||
| # estimated_cost = self._estimate_cost(result) | ||
| # self._total_cost += estimated_cost | ||
| self._total_cost += estimated_cost | ||
|
|
There was a problem hiding this comment.
Budget exhaustion does not actually stop the scheduler.
This branch logs that it is stopping, but it only returns from _execute_with_retry(). _run_schedule() will keep waking up and calling it forever. On top of that, _total_cost is only updated after successful runs, so repeated failures/timeouts are effectively free and can overshoot max_cost. Set the stop event when the budget is exhausted and charge each attempt before deciding whether another retry is allowed.
🤖 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 `@src/praisonai/praisonai/scheduler/async_agent_scheduler.py` around lines 361
- 390, The budget-check branch in _execute_with_retry currently only returns, so
_run_schedule keeps calling it and failures don't increment _total_cost; change
it to set the scheduler stop event (e.g., self._stop_event.set()) when budget is
exhausted and return so the scheduler halts; also deduct/charge an
estimated_cost (or per-attempt charge) before each attempt (inside
_execute_with_retry, before retries loop or at start of each attempt) while
holding self._stats_lock so _total_cost is updated even for failed/timeouts and
prevent starting an attempt that would exceed self.max_cost by checking
self._total_cost + estimated_cost >= self.max_cost and setting self._stop_event
if exceeded.
Greptile SummaryThis PR closes three architectural gaps in the wrapper package: thread-safe singleton initialization for the jobs server, modernised async patterns across executor/sandbox files, and missing
Confidence Score: 5/5Safe to merge; all three targeted fixes are mechanically correct and the previously reported loop-lifecycle and budget-guard bugs are resolved. The double-checked locking, asyncio.to_thread() migration, and single-event-loop refactor in standardise.py are all straightforward and well-scoped. The only remaining nits are unused imports and a falsy timeout guard, neither of which affects runtime correctness in practice. async_agent_scheduler.py — the timeout guard and the cost-tracking TODO stubs warrant a second look before the budget feature is considered production-ready. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["_run_ai (sync)"] -->|"asyncio.run()"| B["_run_ai_with_runtime (async)"]
B --> C["runtime.start()"]
C --> D["_run_generation_logic()"]
D --> E["return 0"]
E --> F["runtime.stop() [finally]"]
G["_run_schedule loop"] --> H{"budget exhausted?"}
H -->|yes| I["_stop_event.set() exit loop"]
H -->|no| J{"timeout set?"}
J -->|yes| K["asyncio.wait_for(execute, timeout)"]
J -->|no| L["await execute()"]
K --> M["_total_cost += 0.0001"]
L --> M
M --> N["wait interval or stop event"]
N --> G
O["get_store() / get_executor()"] --> P{"instance is None?"}
P -->|yes| Q["acquire Lock"]
Q --> R{"still None?"}
R -->|yes| S["create instance"]
R -->|no| T["return existing"]
S --> T
P -->|no| T
Reviews (2): Last reviewed commit: "fix: address critical architectural issu..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Focus ONLY on Python packages (praisonaiagents, praisonai). Do NOT modify praisonai-rust or praisonai-ts. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding. Phase 1: Review per AGENTS.md
Phase 2: FIX Valid Issues Phase 3: Final Verdict |
|
Claude finished @MervinPraison's task —— View job Final Architecture Review - Complete ✅Acting as final architecture reviewer for PR #1771. I have reviewed all reviewer feedback and implemented critical fixes. Todo List:
Architecture Compliance Review:✅ Protocol-Driven: Changes in wrapper package only (praisonai), not core SDK Critical Fixes Implemented:🔧 Fixed: Event Loop Splitting (standardise.py)
🔧 Fixed: Budget Scheduler Bug (async_agent_scheduler.py)
🔧 Fixed: Asyncio Migration (e2b.py)
🔧 Fixed: Factory Function (async_agent_scheduler.py)
Original PR Assessment:✅ Thread Safety: Double-checked locking pattern correctly implemented Final Verdict: APPROVED ✅All critical architectural issues have been resolved. The PR now:
Recommendation: Merge after CI passes The fixes ensure production-ready safety and restore the promised 3-way feature surface (CLI/YAML/Python) for scheduled agents as stated in the PR objectives. Files Modified in Additional Fixes:
|
- Fix event loop splitting in CLI standardise.py by using single asyncio.run() lifecycle - Fix budget scheduler stopping issue by setting stop event when budget exceeded - Complete asyncio migration in e2b.py by replacing remaining get_event_loop() calls - Update factory function to expose timeout and max_cost parameters Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>
1 participant
Fixes three validated architectural gaps in the wrapper package that violated PraisonAI's stated pillars:
Summary
Changes Made
1. Jobs Server Race Condition Fix (praisonai/jobs/server.py)
2. Deprecated asyncio.get_event_loop() Fixes
3. AsyncAgentScheduler Feature Parity (scheduler/async_agent_scheduler.py)
Testing
Impact
These fixes eliminate production-ready safety issues and restore the promised 3-way feature surface (CLI/YAML/Python) for scheduled agents.
Fixes #1770
🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes
New Features
timeoutandmax_costconfiguration parameters to agent scheduler for per-execution control and budget enforcement.Performance & Stability