Skip to content

fix: async lifecycle hooks + locked display-callback swap #3174

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Leebana merged 2 commits into from
Jul 20, 2026
+87 −28
Merged

fix: async lifecycle hooks + locked display-callback swap #3174

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c804d93
fix: async lifecycle hooks + locked display-callback swap (fixes #3173)
praisonai-triage-agent[bot] Jul 19, 2026
175cf28
fix: guard callback restore under identity check + async-aware task h…
praisonai-triage-agent[bot] Jul 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions src/praisonai-agents/praisonaiagents/agent/execution_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -742,7 +742,7 @@ def start(self, prompt: Optional[str] = None, **kwargs: Any) -> Union[str, Gener

# Show animated status during LLM call if verbose
if self.verbose and is_tty:
from ..main import PRAISON_COLORS, sync_display_callbacks
from ..main import PRAISON_COLORS, sync_display_callbacks, _callbacks_lock
import threading
import time as time_module

Expand All @@ -763,9 +763,11 @@ def status_tool_callback(**kwargs):
tools_called.append(tool_name)
current_status[0] = f"Calling tool: {tool_name}..."

# Store original callback and register ours
original_tool_callback = sync_display_callbacks.get('tool_call')
sync_display_callbacks['tool_call'] = status_tool_callback
# Store original callback and register ours (use the module's
# own lock so this doesn't race concurrent verbose agents)
with _callbacks_lock:
original_tool_callback = sync_display_callbacks.get('tool_call')
sync_display_callbacks['tool_call'] = status_tool_callback

# Animation state
result_holder = [None]
Expand Down Expand Up @@ -803,11 +805,16 @@ def run_chat():
error_holder[0] = e
finally:
self.verbose = original_verbose_chat
# Restore original callback
if original_tool_callback:
sync_display_callbacks['tool_call'] = original_tool_callback
elif 'tool_call' in sync_display_callbacks:
del sync_display_callbacks['tool_call']
# Restore original callback under the lock. The identity
# check guards the entire restore, so we only mutate the
# entry while it's still ours and never clobber a callback
# a concurrent verbose agent has since installed.
with _callbacks_lock:
if sync_display_callbacks.get('tool_call') is status_tool_callback:
if original_tool_callback:
sync_display_callbacks['tool_call'] = original_tool_callback
else:
del sync_display_callbacks['tool_call']

# Start chat in background thread
chat_thread = threading.Thread(target=run_chat)
Expand Down
90 changes: 71 additions & 19 deletions src/praisonai-agents/praisonaiagents/agents/agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1115,6 +1115,63 @@ def _apply_task_guardrail(self, task, task_id, task_output):
logger.warning(f"Task {task_id}: Guardrail processing error (retry {task.retry_count}/{task.max_retries}): {e}")
return task_output, True # Signal retry needed

def _run_task_start_hook(self, task, task_id):
"""Run the on_task_start hook and propagate global variables to the task.

Shared by run_task (sync) and arun_task (async) so the lifecycle hooks
and variable propagation stay consistent across both paths.
"""
if self.on_task_start:
try:
self.on_task_start(task, task_id)
except Exception as e:
logger.error(f"Error in on_task_start callback: {e}")

# Apply global variables to task if not already set. Use a shallow copy so a
# task mutating its variables doesn't leak back into the shared AgentTeam state.
if self.variables and not getattr(task, 'variables', None):
task.variables = self.variables.copy()

def _run_task_complete_hook(self, task, task_output):
"""Run the on_task_complete hook. Shared by run_task and arun_task."""
if self.on_task_complete:
try:
self.on_task_complete(task, task_output)
except Exception as e:
logger.error(f"Error in on_task_complete callback: {e}")

Comment on lines +1118 to +1142

@coderabbitai coderabbitai Bot Jul 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Support asynchronous lifecycle callbacks and prevent event loop blocking.

By sharing synchronous helpers between run_task and arun_task, two issues arise for async execution:

  1. If a user provides an asynchronous callback (async def), it will return a coroutine that is never awaited, causing the hook to silently fail and emit a RuntimeWarning.
  2. If a user provides a synchronous callback that performs I/O, executing it directly inside arun_task will block the event loop.

Unlike Task.execute_callback() which is properly awaited natively, these hooks lack async-aware execution. Consider adding async variants of these helpers for arun_task that use asyncio.iscoroutinefunction() to await async hooks, and use loop.run_in_executor() to safely execute synchronous blocking hooks.

As per coding guidelines, "All I/O operations must provide async variants; never block the event loop with synchronous I/O".

🧰 Tools
🪛 Ruff (0.15.21)

[warning] 1127-1127: Do not catch blind exception: Exception

(BLE001)

[warning] 1139-1139: Do not catch blind exception: Exception

(BLE001)

🤖 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-agents/praisonaiagents/agents/agents.py` around lines 1118 -
1141, Add async lifecycle-hook helpers alongside _run_task_start_hook and
_run_task_complete_hook, and update arun_task to use them. Detect coroutine
callbacks with asyncio.iscoroutinefunction() and await them; run synchronous
callbacks through the event loop’s executor so blocking I/O does not block
arun_task. Preserve the existing exception logging and variable propagation
behavior.

Source: Coding guidelines

async def _arun_task_start_hook(self, task, task_id):
"""Async-aware on_task_start hook for arun_task.

Awaits coroutine callbacks and offloads synchronous ones to the executor so
a blocking hook can't stall the event loop. Falls back to the sync helper's
variable propagation.
"""
if self.on_task_start:
try:
if asyncio.iscoroutinefunction(self.on_task_start):
await self.on_task_start(task, task_id)
else:
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self.on_task_start, task, task_id)
except Exception as e:
logger.error(f"Error in on_task_start callback: {e}")

if self.variables and not getattr(task, 'variables', None):
task.variables = self.variables.copy()

async def _arun_task_complete_hook(self, task, task_output):
"""Async-aware on_task_complete hook for arun_task (see _arun_task_start_hook)."""
if self.on_task_complete:
try:
if asyncio.iscoroutinefunction(self.on_task_complete):
await self.on_task_complete(task, task_output)
else:
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self.on_task_complete, task, task_output)
except Exception as e:
logger.error(f"Error in on_task_complete callback: {e}")


async def arun_task(self, task_id):
"""Async version of run_task method"""
Expand All @@ -1126,6 +1183,9 @@ async def arun_task(self, task_id):
logger.info(f"Task with ID {task_id} is already completed")
return

# Call on_task_start callback and propagate variables (async-aware, mirrors run_task)
await self._arun_task_start_hook(task, task_id)

# Use per-task max_retries if available
task_max = getattr(task, "max_retries", self.max_retries)
retries = 0
Expand Down Expand Up @@ -1154,6 +1214,10 @@ async def arun_task(self, task_id):
raise

self.save_output_to_file(task, task_output)

# Call on_task_complete callback (async-aware, mirrors run_task)
await self._arun_task_complete_hook(task, task_output)

if self.verbose >= 1:
logger.info(f"Task {task_id} completed successfully.")
else:
Expand Down Expand Up @@ -1349,17 +1413,9 @@ def run_task(self, task_id):
logger.info(f"Task with ID {task_id} is already completed")
return

# Call on_task_start callback if provided
if self.on_task_start:
try:
self.on_task_start(task, task_id)
except Exception as e:
logger.error(f"Error in on_task_start callback: {e}")

# Apply global variables to task if not already set
if self.variables and not getattr(task, 'variables', None):
task.variables = self.variables

# Call on_task_start callback and propagate variables (shared with arun_task)
self._run_task_start_hook(task, task_id)

# Use per-task max_retries if available
task_max = getattr(task, "max_retries", self.max_retries)
retries = 0
Expand All @@ -1384,14 +1440,10 @@ def run_task(self, task_id):
logger.exception(e)

self.save_output_to_file(task, task_output)

# Call on_task_complete callback if provided
if self.on_task_complete:
try:
self.on_task_complete(task, task_output)
except Exception as e:
logger.error(f"Error in on_task_complete callback: {e}")


# Call on_task_complete callback (shared with arun_task)
self._run_task_complete_hook(task, task_output)

if self.verbose >= 1:
logger.info(f"Task {task_id} completed successfully.")
else:
Expand Down
Loading