Add connector message queue protocol
Advertise message queue support from the Agent Zero connector backend and add WebSocket handlers for queue add, remove, and send operations. Include queue snapshots in context subscriptions and emit queue updates as the backend state changes so the CLI can stay in sync.
Alessandro committed
May 15, 2026 at 18:13 UTC
38bbff3d9a45a417ca5a0b2b9b955994d6154bca
2 files changed
+261
plugins/_a0_connector/api/v1/capabilities.py
+1
@@ -17,6 +17,7 @@ _BASE_FEATURES = [
17
"pause",
18
"nudge",
19
"message_send",
20
+ "message_queue",
21
"log_tail",
22
"projects",
23
"text_editor_remote",
plugins/_a0_connector/api/ws_connector.py
+260
@@ -49,6 +49,10 @@ PROTOCOL_VERSION = "a0-connector.v1"
49
WS_FEATURES = [
50
"connector_subscribe_context",
51
"connector_send_message",
52
+ "message_queue",
53
+ "connector_message_queue_add",
54
+ "connector_message_queue_remove",
55
+ "connector_message_queue_send",
56
"text_editor_remote",
57
"remote_file_tree",
58
"code_execution_remote",
@@ -129,6 +133,15 @@ class WsConnector(WsHandler):
133
if event == "connector_send_message":
134
return await self._handle_send_message(data, sid)
135
136
+ if event == "connector_message_queue_add":
137
+ return await self._handle_message_queue_add(data, sid)
138
+
139
+ if event == "connector_message_queue_remove":
140
+ return await self._handle_message_queue_remove(data, sid)
141
+
142
+ if event == "connector_message_queue_send":
143
+ return await self._handle_message_queue_send(data, sid)
144
+
145
if event == "connector_file_op_result":
146
return self._handle_file_op_result(data, sid)
147
@@ -246,6 +259,7 @@ class WsConnector(WsHandler):
259
"context_id": context_id,
260
"events": events,
261
"last_sequence": last_sequence,
262
+ "message_queue": self._queue_items_for_context(context),
263
},
264
correlation_id=data.get("correlationId"),
265
)
@@ -341,6 +355,7 @@ class WsConnector(WsHandler):
355
"context_id": context_id,
356
"events": events,
357
"last_sequence": last_sequence,
358
+ "message_queue": self._queue_items_for_context(context),
359
},
360
correlation_id=data.get("correlationId"),
361
)
@@ -370,6 +385,162 @@ class WsConnector(WsHandler):
385
"client_message_id": client_message_id or None,
386
}
387
388
+ async def _handle_message_queue_add(
389
+ self,
390
+ data: dict[str, Any],
391
+ sid: str,
392
+ ) -> dict[str, Any] | WsResult:
393
+ from agent import AgentContext
394
+ from helpers import message_queue as mq
395
+ from helpers.state_monitor_integration import mark_dirty_for_context
396
+
397
+ context_id = str(data.get("context_id", data.get("context", ""))).strip()
398
+ message = str(data.get("message", data.get("text", ""))).strip()
399
+ client_message_id = str(data.get("client_message_id", data.get("item_id", ""))).strip()
400
+ raw_attachments = list(data.get("attachments", [])) if isinstance(data.get("attachments"), list) else []
401
+ attachments, attachment_error = self._normalize_attachment_refs(raw_attachments)
402
+ if attachment_error:
403
+ return WsResult.error(
404
+ code="INVALID_ATTACHMENTS",
405
+ message=attachment_error,
406
+ correlation_id=data.get("correlationId"),
407
+ )
408
+ if not context_id:
409
+ return WsResult.error(
410
+ code="MISSING_CONTEXT_ID",
411
+ message="context_id is required",
412
+ correlation_id=data.get("correlationId"),
413
+ )
414
+ if not message and not attachments:
415
+ return WsResult.error(
416
+ code="MISSING_MESSAGE",
417
+ message="message or attachments are required",
418
+ correlation_id=data.get("correlationId"),
419
+ )
420
+
421
+ context = AgentContext.get(context_id)
422
+ if context is None:
423
+ return WsResult.error(
424
+ code="CONTEXT_NOT_FOUND",
425
+ message=f"Context '{context_id}' not found",
426
+ correlation_id=data.get("correlationId"),
427
+ )
428
+
429
+ item = mq.add(
430
+ context,
431
+ message,
432
+ attachments,
433
+ item_id=client_message_id or data.get("correlationId") or None,
434
+ )
435
+ mark_dirty_for_context(context_id, reason="connector_message_queue_add")
436
+ await self._emit_message_queue_updated(context_id=context_id, context=context)
437
+
438
+ return {
439
+ "context_id": context_id,
440
+ "status": "queued",
441
+ "item": self._queue_item_payload(item),
442
+ "message_queue": self._queue_items_for_context(context),
443
+ }
444
+
445
+ async def _handle_message_queue_remove(
446
+ self,
447
+ data: dict[str, Any],
448
+ sid: str,
449
+ ) -> dict[str, Any] | WsResult:
450
+ from agent import AgentContext
451
+ from helpers import message_queue as mq
452
+ from helpers.state_monitor_integration import mark_dirty_for_context
453
+
454
+ context_id = str(data.get("context_id", data.get("context", ""))).strip()
455
+ item_id = str(data.get("item_id", "") or "").strip() or None
456
+ if not context_id:
457
+ return WsResult.error(
458
+ code="MISSING_CONTEXT_ID",
459
+ message="context_id is required",
460
+ correlation_id=data.get("correlationId"),
461
+ )
462
+
463
+ context = AgentContext.get(context_id)
464
+ if context is None:
465
+ return WsResult.error(
466
+ code="CONTEXT_NOT_FOUND",
467
+ message=f"Context '{context_id}' not found",
468
+ correlation_id=data.get("correlationId"),
469
+ )
470
+
471
+ remaining = mq.remove(context, item_id)
472
+ mark_dirty_for_context(context_id, reason="connector_message_queue_remove")
473
+ await self._emit_message_queue_updated(context_id=context_id, context=context)
474
+
475
+ return {
476
+ "context_id": context_id,
477
+ "status": "removed",
478
+ "remaining": remaining,
479
+ "message_queue": self._queue_items_for_context(context),
480
+ }
481
+
482
+ async def _handle_message_queue_send(
483
+ self,
484
+ data: dict[str, Any],
485
+ sid: str,
486
+ ) -> dict[str, Any] | WsResult:
487
+ from agent import AgentContext
488
+ from helpers import message_queue as mq
489
+ from helpers.state_monitor_integration import mark_dirty_for_context
490
+
491
+ context_id = str(data.get("context_id", data.get("context", ""))).strip()
492
+ item_id = str(data.get("item_id", "") or "").strip() or None
493
+ send_all = bool(data.get("send_all", False))
494
+ if not context_id:
495
+ return WsResult.error(
496
+ code="MISSING_CONTEXT_ID",
497
+ message="context_id is required",
498
+ correlation_id=data.get("correlationId"),
499
+ )
500
+
501
+ context = AgentContext.get(context_id)
502
+ if context is None:
503
+ return WsResult.error(
504
+ code="CONTEXT_NOT_FOUND",
505
+ message=f"Context '{context_id}' not found",
506
+ correlation_id=data.get("correlationId"),
507
+ )
508
+
509
+ if not mq.has_queue(context):
510
+ await self._emit_message_queue_updated(context_id=context_id, context=context)
511
+ return {
512
+ "context_id": context_id,
513
+ "status": "empty",
514
+ "sent_count": 0,
515
+ "message_queue": [],
516
+ }
517
+
518
+ if send_all:
519
+ sent_count = mq.send_all_aggregated(context)
520
+ sent_item_id = None
521
+ else:
522
+ item = mq.pop_item(context, item_id) if item_id else mq.pop_first(context)
523
+ if not item:
524
+ return WsResult.error(
525
+ code="QUEUE_ITEM_NOT_FOUND",
526
+ message="Queued message was not found",
527
+ correlation_id=data.get("correlationId"),
528
+ )
529
+ sent_item_id = item.get("id")
530
+ mq.send_message(context, item)
531
+ sent_count = 1
532
+
533
+ mark_dirty_for_context(context_id, reason="connector_message_queue_send")
534
+ await self._emit_message_queue_updated(context_id=context_id, context=context)
535
+
536
+ return {
537
+ "context_id": context_id,
538
+ "status": "sent",
539
+ "sent_count": sent_count,
540
+ "sent_item_id": sent_item_id,
541
+ "message_queue": self._queue_items_for_context(context),
542
+ }
543
+
544
def _normalize_attachment_refs(self, attachments: list[Any]) -> tuple[list[str], str]:
545
refs: list[str] = []
546
for attachment in attachments:
@@ -398,6 +569,71 @@ class WsConnector(WsHandler):
569
570
return refs, ""
571
572
+ def _queue_item_payload(self, item: dict[str, Any]) -> dict[str, Any]:
573
+ text = str(item.get("text", "") or "")
574
+ attachments = [
575
+ str(attachment).split("/")[-1]
576
+ for attachment in item.get("attachments", [])
577
+ if str(attachment or "").strip()
578
+ ]
579
+ return {
580
+ "id": str(item.get("id", "") or ""),
581
+ "seq": int(item.get("seq", 0) or 0),
582
+ "text": text[:100] + "..." if len(text) > 100 else text,
583
+ "attachments": attachments,
584
+ "attachment_count": len(item.get("attachments", []) or []),
585
+ }
586
+
587
+ def _queue_items_for_context(self, context: AgentContext | None) -> list[dict[str, Any]]:
588
+ if context is None:
589
+ return []
590
+ try:
591
+ from helpers import message_queue as mq
592
+
593
+ return [self._queue_item_payload(item) for item in mq.get_queue(context)]
594
+ except Exception:
595
+ return []
596
+
597
+ def _queue_state_for_context_id(self, context_id: str) -> tuple[str, list[dict[str, Any]]]:
598
+ try:
599
+ from agent import AgentContext
600
+
601
+ context = AgentContext.get(context_id)
602
+ except Exception:
603
+ context = None
604
+
605
+ items = self._queue_items_for_context(context)
606
+ signature = repr(items)
607
+ return signature, items
608
+
609
+ def _context_is_running(self, context_id: str) -> bool:
610
+ try:
611
+ from agent import AgentContext
612
+
613
+ context = AgentContext.get(context_id)
614
+ return bool(context is not None and context.is_running())
615
+ except Exception:
616
+ return False
617
+
618
+ async def _emit_message_queue_updated(
619
+ self,
620
+ *,
621
+ context_id: str,
622
+ context: AgentContext | None = None,
623
+ ) -> None:
624
+ payload = {
625
+ "context_id": context_id,
626
+ "message_queue": self._queue_items_for_context(context),
627
+ }
628
+ for target_sid in subscribed_sids_for_context(context_id):
629
+ try:
630
+ await self.emit_to(target_sid, "connector_message_queue_updated", payload)
631
+ except Exception as exc:
632
+ PrintStyle.error(
633
+ f"[a0-connector] failed to emit connector_message_queue_updated "
634
+ f"to {target_sid}: {exc}"
635
+ )
636
+
637
def _handle_file_op_result(
638
self,
639
data: dict[str, Any],
@@ -652,12 +888,36 @@ class WsConnector(WsHandler):
888
) -> None:
889
# `from_sequence` is a log-output cursor (not an event sequence number).
890
cursor = max(int(from_sequence or 0), 0)
891
+ last_queue_signature, _ = self._queue_state_for_context_id(context_id)
892
+ was_running = self._context_is_running(context_id)
893
try:
894
while context_id in subscribed_contexts_for_sid(sid):
895
events, next_cursor = get_context_log_entries(context_id, after=cursor)
896
for event in events:
897
await self.emit_to(sid, "connector_context_event", event)
898
cursor = max(cursor, int(next_cursor or cursor))
899
+ queue_signature, queue_items = self._queue_state_for_context_id(context_id)
900
+ if queue_signature != last_queue_signature:
901
+ last_queue_signature = queue_signature
902
+ await self.emit_to(
903
+ sid,
904
+ "connector_message_queue_updated",
905
+ {
906
+ "context_id": context_id,
907
+ "message_queue": queue_items,
908
+ },
909
+ )
910
+ is_running = self._context_is_running(context_id)
911
+ if was_running and not is_running:
912
+ await self.emit_to(
913
+ sid,
914
+ "connector_context_complete",
915
+ {
916
+ "context_id": context_id,
917
+ "status": "completed",
918
+ },
919
+ )
920
+ was_running = is_running
921
await asyncio.sleep(0.5)
922
except asyncio.CancelledError:
923
raise