mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 08:38:30 +00:00
Fix async task token tracking race condition
Resolved race condition where concurrent async tasks from same agent would get incorrect token attribution. Solution wraps async task execution to capture tokens_after immediately upon task completion, before other concurrent tasks can interfere. Changes: - Wrapped async task execution to return (result, tokens_after) tuple - Updated _aprocess_async_tasks to unwrap and use captured tokens_after - Updated type hints for pending_tasks to reflect new signature Note: Threading-based async_execution still has similar race condition as it's harder to wrap threaded execution. Will track separately.
This commit is contained in:
@@ -954,13 +954,20 @@ class Crew(FlowTrackable, BaseModel):
|
||||
context = self._get_context(
|
||||
task, [last_sync_output] if last_sync_output else []
|
||||
)
|
||||
async_task = asyncio.create_task(
|
||||
task.aexecute_sync(
|
||||
|
||||
# Wrap task execution to capture tokens immediately after completion
|
||||
async def _wrapped_task_execution():
|
||||
result = await task.aexecute_sync(
|
||||
agent=exec_data.agent,
|
||||
context=context,
|
||||
tools=exec_data.tools,
|
||||
)
|
||||
)
|
||||
# Capture tokens immediately after task completes
|
||||
# This reduces (but doesn't eliminate) race conditions
|
||||
tokens_after = self._get_agent_token_usage(exec_data.agent)
|
||||
return result, tokens_after
|
||||
|
||||
async_task = asyncio.create_task(_wrapped_task_execution())
|
||||
pending_tasks.append((task, async_task, task_index, exec_data.agent, tokens_before))
|
||||
else:
|
||||
if pending_tasks:
|
||||
@@ -998,7 +1005,7 @@ class Crew(FlowTrackable, BaseModel):
|
||||
self,
|
||||
task: ConditionalTask,
|
||||
task_outputs: list[TaskOutput],
|
||||
pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int, Any, Any]],
|
||||
pending_tasks: list[tuple[Task, asyncio.Task[tuple[TaskOutput, Any]], int, Any, Any]],
|
||||
task_index: int,
|
||||
was_replayed: bool,
|
||||
) -> TaskOutput | None:
|
||||
@@ -1013,16 +1020,16 @@ class Crew(FlowTrackable, BaseModel):
|
||||
|
||||
async def _aprocess_async_tasks(
|
||||
self,
|
||||
pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int, Any, Any]],
|
||||
pending_tasks: list[tuple[Task, asyncio.Task[tuple[TaskOutput, Any]], int, Any, Any]],
|
||||
was_replayed: bool = False,
|
||||
) -> list[TaskOutput]:
|
||||
"""Process pending async tasks and return their outputs."""
|
||||
task_outputs: list[TaskOutput] = []
|
||||
for future_task, async_task, task_index, agent, tokens_before in pending_tasks:
|
||||
task_output = await async_task
|
||||
# Unwrap the result which includes both output and tokens_after
|
||||
task_output, tokens_after = await async_task
|
||||
|
||||
# Capture token usage after async task execution and attach to task output
|
||||
tokens_after = self._get_agent_token_usage(agent)
|
||||
# Attach token metrics using the captured tokens_after
|
||||
task_output = self._attach_task_token_metrics(
|
||||
task_output, future_task, agent, tokens_before, tokens_after
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user