#84 - APA passing the project via message_id-based dictionary instead of FIFO queue
#84 - APA passing the project via message_id-based dictionary instead of FIFO queue
deci committed
Dec 12, 2025 at 11:19 UTC
2c87e17f18ace3c67468ea793fef0b63895a891b
1 file changed
+60
-13
python/helpers/fasta2a_server.py
+60
-13
@@ -5,7 +5,6 @@ import atexit
5
from typing import Any, List
6
import contextlib
7
import threading
8
-from collections import deque
8
9
from python.helpers import settings, projects
10
from starlette.requests import Request
@@ -61,9 +60,9 @@ except ImportError: # pragma: no cover – library not installed
60
61
_PRINTER = PrintStyle(italic=True, font_color="purple", padding=False)
62
64
-# FIFO queue to pass project names from request context to worker context
65
-# Each request appends project (or None), worker pops in same order
66
-_a2a_project_queue: deque[str | None] = deque()
63
+# Map message_id to project name (thread-safe, no race conditions)
64
+# Each message has unique ID that correlates request to worker
65
+_a2a_message_projects: dict[str, str] = {}
66
_a2a_project_lock = threading.Lock()
67
68
@@ -90,12 +89,15 @@ class AgentZeroWorker(Worker): # type: ignore[misc]
89
cfg = initialize_agent()
90
context = AgentContext(cfg, type=AgentContextType.BACKGROUND)
91
93
- # Retrieve project from queue (FIFO matches task processing order)
92
+ # Retrieve project by message_id (direct lookup, no race conditions)
93
+ # Note: Request has messageId (camelCase), but FastA2A converts it to message_id (snake_case)
94
project_name = None
95
- with _a2a_project_lock:
96
- if _a2a_project_queue:
97
- project_name = _a2a_project_queue.popleft()
98
- _PRINTER.print(f"[A2A] Retrieved project from queue: {project_name}")
95
+ message_id = message.get('message_id') # FastA2A converts camelCase to snake_case
96
+ if message_id:
97
+ with _a2a_project_lock:
98
+ project_name = _a2a_message_projects.pop(message_id, None)
99
+ if project_name:
100
+ _PRINTER.print(f"[A2A] Retrieved project for message {message_id}: {project_name}")
101
102
# Activate project if specified
103
if project_name:
@@ -485,10 +487,55 @@ class DynamicA2AProxy:
487
})
488
return
489
488
- # Store project in queue for worker to retrieve (maintains FIFO order)
489
- with _a2a_project_lock:
490
- _a2a_project_queue.append(project_name) # None is valid (no project)
491
- _PRINTER.print(f"[A2A] Appended project to queue: {project_name}")
490
+ # If project specified, we need to extract message_id from request body
491
+ # to correlate project with the specific message (no race conditions)
492
+ if project_name:
493
+ # Buffer all messages for replay
494
+ received_messages = []
495
+ replay_index = 0
496
+ message_id_extracted = False
497
+ original_receive = receive
498
+
499
+ async def receive_wrapper():
500
+ nonlocal replay_index, message_id_extracted
501
+
502
+ # If replaying buffered messages, return them in order
503
+ if replay_index < len(received_messages):
504
+ msg = received_messages[replay_index]
505
+ replay_index += 1
506
+ return msg
507
+
508
+ # Otherwise, receive and buffer the next message
509
+ message = await original_receive()
510
+ received_messages.append(message)
511
+
512
+ # Parse message_id when we get the complete body
513
+ if message['type'] == 'http.request':
514
+ # If this is the last chunk, parse the full body
515
+ if not message.get('more_body', False) and not message_id_extracted:
516
+ message_id_extracted = True
517
+ try:
518
+ import json
519
+ # Reconstruct full body from all buffered messages
520
+ body_parts = [msg.get('body', b'') for msg in received_messages if msg['type'] == 'http.request']
521
+ full_body = b''.join(body_parts)
522
+ data = json.loads(full_body)
523
+ # Handle JSON-RPC format: params.message.messageId (camelCase!)
524
+ if 'params' in data and 'message' in data['params']:
525
+ msg = data['params']['message']
526
+ message_id = msg.get('messageId') # camelCase in raw JSON
527
+ if message_id:
528
+ with _a2a_project_lock:
529
+ _a2a_message_projects[message_id] = project_name
530
+ _PRINTER.print(f"[A2A] Stored project '{project_name}' for message {message_id}")
531
+ # Reset replay index so FastA2A can replay from start
532
+ replay_index = 0
533
+ except Exception as e:
534
+ _PRINTER.print(f"[A2A] Failed to parse message_id: {e}")
535
+
536
+ return message
537
+
538
+ receive = receive_wrapper
539
540
# Update scope with cleaned path
541
scope = dict(scope)