omnimancer-cli 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (126) hide show
  1. omnimancer/__init__.py +42 -0
  2. omnimancer/__main__.py +9 -0
  3. omnimancer/cli/__init__.py +8 -0
  4. omnimancer/cli/agent_manager.py +826 -0
  5. omnimancer/cli/agent_wizard.py +944 -0
  6. omnimancer/cli/approval_formatter.py +734 -0
  7. omnimancer/cli/approval_integration.py +734 -0
  8. omnimancer/cli/approval_prompt.py +1286 -0
  9. omnimancer/cli/batch_approval_display.py +617 -0
  10. omnimancer/cli/batch_approval_filters.py +729 -0
  11. omnimancer/cli/batch_operation_monitor.py +775 -0
  12. omnimancer/cli/commands/performance_cmd.py +624 -0
  13. omnimancer/cli/commands.py +757 -0
  14. omnimancer/cli/completion_handler.py +281 -0
  15. omnimancer/cli/config_setup_wizard.py +783 -0
  16. omnimancer/cli/dynamic_commands.py +371 -0
  17. omnimancer/cli/handlers/__init__.py +18 -0
  18. omnimancer/cli/handlers/agent_handler.py +865 -0
  19. omnimancer/cli/handlers/agent_persona_handler.py +1048 -0
  20. omnimancer/cli/handlers/custom_agent_handler.py +352 -0
  21. omnimancer/cli/handlers/permissions_handler.py +1609 -0
  22. omnimancer/cli/interface.py +6695 -0
  23. omnimancer/cli/performance_dashboard.py +612 -0
  24. omnimancer/core/__init__.py +34 -0
  25. omnimancer/core/agent/__init__.py +256 -0
  26. omnimancer/core/agent/agent_config.py +715 -0
  27. omnimancer/core/agent/agent_switcher.py +668 -0
  28. omnimancer/core/agent/approval_context.py +460 -0
  29. omnimancer/core/agent/approval_dialog.py +902 -0
  30. omnimancer/core/agent/approval_interface.py +628 -0
  31. omnimancer/core/agent/approval_manager.py +760 -0
  32. omnimancer/core/agent/cli_integration.py +470 -0
  33. omnimancer/core/agent/config.py +845 -0
  34. omnimancer/core/agent/custom_agent_integration.py +654 -0
  35. omnimancer/core/agent/diff_renderer.py +706 -0
  36. omnimancer/core/agent/file_content_display.py +503 -0
  37. omnimancer/core/agent/file_modification_workflow.py +658 -0
  38. omnimancer/core/agent/file_system_manager.py +1790 -0
  39. omnimancer/core/agent/input_handler.py +614 -0
  40. omnimancer/core/agent/metrics_collector.py +771 -0
  41. omnimancer/core/agent/optimization_engine.py +729 -0
  42. omnimancer/core/agent/performance_monitor.py +717 -0
  43. omnimancer/core/agent/persona.py +870 -0
  44. omnimancer/core/agent/persona_validator.py +672 -0
  45. omnimancer/core/agent/program_executor.py +748 -0
  46. omnimancer/core/agent/proposed_changes_integration.py +663 -0
  47. omnimancer/core/agent/read_before_write_errors.py +492 -0
  48. omnimancer/core/agent/read_before_write_ui.py +647 -0
  49. omnimancer/core/agent/rich_renderer.py +763 -0
  50. omnimancer/core/agent/status_config.py +238 -0
  51. omnimancer/core/agent/status_core.py +446 -0
  52. omnimancer/core/agent/status_integrated_engine.py +540 -0
  53. omnimancer/core/agent/status_manager.py +878 -0
  54. omnimancer/core/agent/status_ui_components.py +535 -0
  55. omnimancer/core/agent/template_applicator.py +548 -0
  56. omnimancer/core/agent/terminal_status_listener.py +394 -0
  57. omnimancer/core/agent/token_tracker.py +679 -0
  58. omnimancer/core/agent/types.py +60 -0
  59. omnimancer/core/agent/web_client.py +745 -0
  60. omnimancer/core/agent/workflow_orchestrator.py +826 -0
  61. omnimancer/core/agent_engine.py +1625 -0
  62. omnimancer/core/agent_mode_manager.py +706 -0
  63. omnimancer/core/agent_progress_ui.py +545 -0
  64. omnimancer/core/chat_manager.py +141 -0
  65. omnimancer/core/config_manager.py +2199 -0
  66. omnimancer/core/config_migration.py +772 -0
  67. omnimancer/core/config_provider.py +1974 -0
  68. omnimancer/core/config_security_validator.py +393 -0
  69. omnimancer/core/config_validator.py +768 -0
  70. omnimancer/core/conversation_manager.py +352 -0
  71. omnimancer/core/engine.py +787 -0
  72. omnimancer/core/enhanced_input.py +523 -0
  73. omnimancer/core/env_loader.py +83 -0
  74. omnimancer/core/fallback_manager.py +568 -0
  75. omnimancer/core/health_monitor.py +462 -0
  76. omnimancer/core/history_manager.py +462 -0
  77. omnimancer/core/mcp_integration_layer.py +878 -0
  78. omnimancer/core/models.py +1889 -0
  79. omnimancer/core/provider_capabilities.py +403 -0
  80. omnimancer/core/provider_initializer.py +450 -0
  81. omnimancer/core/provider_registry.py +432 -0
  82. omnimancer/core/security/__init__.py +16 -0
  83. omnimancer/core/security/approval_workflow.py +406 -0
  84. omnimancer/core/security/audit_logger.py +511 -0
  85. omnimancer/core/security/permission_controller.py +511 -0
  86. omnimancer/core/security/sandbox_manager.py +379 -0
  87. omnimancer/core/security/security_manager.py +530 -0
  88. omnimancer/core/setup_wizard.py +345 -0
  89. omnimancer/core/setup_wizard_core.py +449 -0
  90. omnimancer/core/setup_wizard_guide_generator.py +376 -0
  91. omnimancer/core/setup_wizard_provider_setup.py +454 -0
  92. omnimancer/core/setup_wizard_ui.py +317 -0
  93. omnimancer/core/setup_wizard_validation.py +371 -0
  94. omnimancer/core/signal_handler.py +174 -0
  95. omnimancer/mcp/__init__.py +11 -0
  96. omnimancer/mcp/client.py +725 -0
  97. omnimancer/mcp/manager.py +649 -0
  98. omnimancer/providers/__init__.py +24 -0
  99. omnimancer/providers/azure.py +608 -0
  100. omnimancer/providers/base.py +262 -0
  101. omnimancer/providers/bedrock.py +711 -0
  102. omnimancer/providers/claude.py +426 -0
  103. omnimancer/providers/claude_code.py +406 -0
  104. omnimancer/providers/cohere.py +412 -0
  105. omnimancer/providers/factory.py +988 -0
  106. omnimancer/providers/gemini.py +708 -0
  107. omnimancer/providers/mistral.py +554 -0
  108. omnimancer/providers/ollama.py +510 -0
  109. omnimancer/providers/openai.py +362 -0
  110. omnimancer/providers/openrouter.py +709 -0
  111. omnimancer/providers/perplexity.py +476 -0
  112. omnimancer/providers/vertex.py +617 -0
  113. omnimancer/providers/xai.py +484 -0
  114. omnimancer/ui/__init__.py +6 -0
  115. omnimancer/ui/cancellation_handler.py +235 -0
  116. omnimancer/ui/progress_indicator.py +318 -0
  117. omnimancer/utils/__init__.py +22 -0
  118. omnimancer/utils/error_handler.py +589 -0
  119. omnimancer/utils/errors.py +667 -0
  120. omnimancer/utils/retry.py +238 -0
  121. omnimancer_cli-0.1.0.dist-info/METADATA +413 -0
  122. omnimancer_cli-0.1.0.dist-info/RECORD +126 -0
  123. omnimancer_cli-0.1.0.dist-info/WHEEL +5 -0
  124. omnimancer_cli-0.1.0.dist-info/entry_points.txt +4 -0
  125. omnimancer_cli-0.1.0.dist-info/licenses/LICENSE +21 -0
  126. omnimancer_cli-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,878 @@
1
+ """
2
+ Unified Status Management System for Omnimancer Agents.
3
+
4
+ This module provides a consolidated status manager that combines operation tracking,
5
+ event management, and real-time streaming capabilities into a single cohesive system.
6
+ """
7
+
8
+ import asyncio
9
+ import logging
10
+ import time
11
+ from collections import defaultdict, deque
12
+ from datetime import datetime
13
+ from typing import (
14
+ Any,
15
+ Dict,
16
+ List,
17
+ Optional,
18
+ )
19
+
20
+ from .status_core import (
21
+ AgentEvent,
22
+ AgentOperation,
23
+ AgentStatus,
24
+ EventListener,
25
+ EventType,
26
+ OperationStatus,
27
+ StatusDisplayConfig,
28
+ StatusStreamEvent,
29
+ StatusStreamListener,
30
+ StreamMetrics,
31
+ StreamPriority,
32
+ )
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+
37
+ class StatusUpdateCollector:
38
+ """Batches status updates for efficient rendering."""
39
+
40
+ def __init__(self, batch_size: int = 10, flush_interval: float = 0.1):
41
+ """
42
+ Initialize status update collector.
43
+
44
+ Args:
45
+ batch_size: Maximum number of updates to batch
46
+ flush_interval: Maximum time to wait before flushing batch (seconds)
47
+ """
48
+ self.batch_size = batch_size
49
+ self.flush_interval = flush_interval
50
+ self.updates: List[StatusStreamEvent] = []
51
+ self.last_flush = time.time()
52
+ self._lock = asyncio.Lock()
53
+
54
+ async def add_update(self, stream_event: StatusStreamEvent) -> bool:
55
+ """
56
+ Add update to batch.
57
+
58
+ Args:
59
+ stream_event: Status stream event to add
60
+
61
+ Returns:
62
+ True if batch should be flushed immediately
63
+ """
64
+ async with self._lock:
65
+ self.updates.append(stream_event)
66
+
67
+ # Check if we should flush
68
+ now = time.time()
69
+ should_flush = (
70
+ len(self.updates) >= self.batch_size
71
+ or (now - self.last_flush) >= self.flush_interval
72
+ or stream_event.priority
73
+ in [StreamPriority.HIGH, StreamPriority.CRITICAL]
74
+ )
75
+
76
+ return should_flush
77
+
78
+ async def flush(self) -> List[StatusStreamEvent]:
79
+ """
80
+ Flush and return batched updates.
81
+
82
+ Returns:
83
+ List of batched status stream events
84
+ """
85
+ async with self._lock:
86
+ updates = self.updates.copy()
87
+ self.updates.clear()
88
+ self.last_flush = time.time()
89
+ return updates
90
+
91
+
92
+ class UnifiedStatusManager:
93
+ """
94
+ Unified status manager that combines operation tracking, event management,
95
+ and real-time streaming capabilities.
96
+
97
+ This class consolidates the functionality previously split between
98
+ AgentStatusManager and AsyncStatusStream.
99
+ """
100
+
101
+ def __init__(self, config: Optional[StatusDisplayConfig] = None):
102
+ """
103
+ Initialize the unified status manager.
104
+
105
+ Args:
106
+ config: Configuration for the status system (uses defaults if None)
107
+ """
108
+ self.config = config or StatusDisplayConfig()
109
+
110
+ # Agent status tracking
111
+ self.agent_statuses: Dict[str, AgentStatus] = {}
112
+ self.agent_metadata: Dict[str, Dict[str, Any]] = defaultdict(dict)
113
+
114
+ # Operation tracking
115
+ self.active_operations: Dict[str, AgentOperation] = {}
116
+ self.operation_history: deque = deque(maxlen=self.config.max_operation_history)
117
+ self.operations_by_agent: Dict[str, List[str]] = defaultdict(list)
118
+
119
+ # Event system
120
+ self.event_listeners: List[EventListener] = []
121
+ self.event_history: deque = deque(maxlen=self.config.max_event_history)
122
+ self.event_queue: asyncio.Queue = asyncio.Queue()
123
+
124
+ # Streaming system
125
+ self.stream_listeners: Dict[str, StatusStreamListener] = {}
126
+ self.stream_event_queue: asyncio.Queue = asyncio.Queue(
127
+ maxsize=self.config.max_queue_size
128
+ )
129
+ self.priority_queue: asyncio.Queue = asyncio.Queue(
130
+ maxsize=self.config.max_queue_size // 2
131
+ )
132
+
133
+ # Throttling and batching
134
+ self.last_event_time = time.time()
135
+ self.event_count = 0
136
+ self.update_collector = StatusUpdateCollector(
137
+ batch_size=self.config.batch_size,
138
+ flush_interval=self.config.flush_interval,
139
+ )
140
+
141
+ # Metrics and monitoring
142
+ self.metrics = StreamMetrics()
143
+
144
+ # Statistics tracking
145
+ self.stats = {
146
+ "operations_started": 0,
147
+ "operations_completed": 0,
148
+ "operations_failed": 0,
149
+ "operations_cancelled": 0,
150
+ "events_emitted": 0,
151
+ }
152
+
153
+ # Control flags
154
+ self.running = False
155
+ self.shutdown_event = asyncio.Event()
156
+
157
+ # Background tasks
158
+ self._event_processor_task: Optional[asyncio.Task] = None
159
+ self._stream_processor_task: Optional[asyncio.Task] = None
160
+ self._distributor_task: Optional[asyncio.Task] = None
161
+
162
+ # Thread safety
163
+ self._lock = asyncio.Lock()
164
+ self._shutdown = False
165
+
166
+ async def initialize(self) -> None:
167
+ """Initialize the status manager and start all processing tasks."""
168
+ if self.running:
169
+ return
170
+
171
+ self.running = True
172
+ self.shutdown_event.clear()
173
+
174
+ # Start background processing tasks
175
+ self._event_processor_task = asyncio.create_task(self._process_events())
176
+ self._stream_processor_task = asyncio.create_task(self._process_stream_events())
177
+ self._distributor_task = asyncio.create_task(self._distribute_events())
178
+
179
+ logger.info("UnifiedStatusManager initialized and started")
180
+
181
+ async def shutdown(self) -> None:
182
+ """Shutdown the status manager and cleanup resources."""
183
+ if not self.running:
184
+ return
185
+
186
+ self._shutdown = True
187
+ self.running = False
188
+ self.shutdown_event.set()
189
+
190
+ # Cancel all active operations
191
+ async with self._lock:
192
+ for operation in self.active_operations.values():
193
+ if operation.is_active:
194
+ operation.cancel("System shutdown")
195
+
196
+ # Stop background tasks
197
+ tasks = [
198
+ self._event_processor_task,
199
+ self._stream_processor_task,
200
+ self._distributor_task,
201
+ ]
202
+
203
+ for task in tasks:
204
+ if task:
205
+ task.cancel()
206
+ try:
207
+ await task
208
+ except asyncio.CancelledError:
209
+ pass
210
+
211
+ # Stop all listeners
212
+ for listener in self.event_listeners:
213
+ listener.stop()
214
+
215
+ for listener in self.stream_listeners.values():
216
+ listener.stop()
217
+
218
+ logger.info("UnifiedStatusManager shutdown complete")
219
+
220
+ # Agent Status Management
221
+ async def set_agent_status(
222
+ self,
223
+ agent_id: str,
224
+ status: AgentStatus,
225
+ metadata: Optional[Dict[str, Any]] = None,
226
+ ) -> None:
227
+ """
228
+ Set agent status and emit state change event.
229
+
230
+ Args:
231
+ agent_id: ID of the agent
232
+ status: New status
233
+ metadata: Optional metadata about the status change
234
+ """
235
+ async with self._lock:
236
+ old_status = self.agent_statuses.get(agent_id)
237
+ self.agent_statuses[agent_id] = status
238
+
239
+ if metadata:
240
+ self.agent_metadata[agent_id].update(metadata)
241
+
242
+ # Emit state change event
243
+ await self._emit_event(
244
+ AgentEvent(
245
+ event_type=EventType.AGENT_STATE_CHANGED,
246
+ agent_id=agent_id,
247
+ data={
248
+ "old_status": old_status.value if old_status else None,
249
+ "new_status": status.value,
250
+ "metadata": metadata or {},
251
+ },
252
+ source="UnifiedStatusManager",
253
+ )
254
+ )
255
+
256
+ def get_agent_status(self, agent_id: str) -> Optional[AgentStatus]:
257
+ """Get current status of an agent."""
258
+ return self.agent_statuses.get(agent_id)
259
+
260
+ def get_agent_metadata(self, agent_id: str) -> Dict[str, Any]:
261
+ """Get metadata for an agent."""
262
+ return self.agent_metadata.get(agent_id, {}).copy()
263
+
264
+ # Operation Management
265
+ async def start_operation(self, operation: AgentOperation) -> str:
266
+ """
267
+ Start tracking a new operation.
268
+
269
+ Args:
270
+ operation: The operation to track
271
+
272
+ Returns:
273
+ The operation ID
274
+ """
275
+ async with self._lock:
276
+ operation.status = OperationStatus.RUNNING
277
+ operation.start_time = operation.start_time or time.time()
278
+
279
+ self.active_operations[operation.operation_id] = operation
280
+
281
+ if operation.agent_id:
282
+ self.operations_by_agent[operation.agent_id].append(
283
+ operation.operation_id
284
+ )
285
+
286
+ self.stats["operations_started"] += 1
287
+
288
+ # Emit operation started event
289
+ await self._emit_event(
290
+ AgentEvent(
291
+ event_type=EventType.OPERATION_STARTED,
292
+ agent_id=operation.agent_id,
293
+ operation_id=operation.operation_id,
294
+ data={
295
+ "operation_type": operation.operation_type.value,
296
+ "description": operation.description,
297
+ "metadata": operation.metadata,
298
+ },
299
+ source="UnifiedStatusManager",
300
+ )
301
+ )
302
+
303
+ return operation.operation_id
304
+
305
+ async def update_operation_progress(
306
+ self,
307
+ operation_id: str,
308
+ progress: float,
309
+ description: Optional[str] = None,
310
+ metadata: Optional[Dict[str, Any]] = None,
311
+ ) -> None:
312
+ """
313
+ Update operation progress.
314
+
315
+ Args:
316
+ operation_id: ID of the operation
317
+ progress: Progress percentage (0-100)
318
+ description: Optional progress description
319
+ metadata: Optional additional metadata
320
+ """
321
+ async with self._lock:
322
+ operation = self.active_operations.get(operation_id)
323
+ if not operation:
324
+ return
325
+
326
+ operation.update_progress(progress, description)
327
+ if metadata:
328
+ operation.metadata.update(metadata)
329
+
330
+ # Emit progress event
331
+ await self._emit_event(
332
+ AgentEvent(
333
+ event_type=EventType.OPERATION_PROGRESS,
334
+ agent_id=operation.agent_id,
335
+ operation_id=operation_id,
336
+ data={
337
+ "progress": progress,
338
+ "description": description or operation.description,
339
+ "metadata": metadata or {},
340
+ },
341
+ source="UnifiedStatusManager",
342
+ )
343
+ )
344
+
345
+ async def complete_operation(
346
+ self,
347
+ operation_id: str,
348
+ result_metadata: Optional[Dict[str, Any]] = None,
349
+ ) -> None:
350
+ """
351
+ Mark operation as completed.
352
+
353
+ Args:
354
+ operation_id: ID of the operation
355
+ result_metadata: Optional result metadata
356
+ """
357
+ async with self._lock:
358
+ operation = self.active_operations.get(operation_id)
359
+ if not operation:
360
+ return
361
+
362
+ operation.complete(result_metadata)
363
+ self._move_to_history(operation)
364
+
365
+ self.stats["operations_completed"] += 1
366
+
367
+ # Emit completion event
368
+ await self._emit_event(
369
+ AgentEvent(
370
+ event_type=EventType.OPERATION_COMPLETED,
371
+ agent_id=operation.agent_id,
372
+ operation_id=operation_id,
373
+ data={
374
+ "duration": (
375
+ operation.duration.total_seconds()
376
+ if operation.duration
377
+ else 0
378
+ ),
379
+ "result_metadata": result_metadata or {},
380
+ },
381
+ source="UnifiedStatusManager",
382
+ )
383
+ )
384
+
385
+ async def fail_operation(
386
+ self,
387
+ operation_id: str,
388
+ error: str,
389
+ error_metadata: Optional[Dict[str, Any]] = None,
390
+ ) -> None:
391
+ """
392
+ Mark operation as failed.
393
+
394
+ Args:
395
+ operation_id: ID of the operation
396
+ error: Error description
397
+ error_metadata: Optional error metadata
398
+ """
399
+ async with self._lock:
400
+ operation = self.active_operations.get(operation_id)
401
+ if not operation:
402
+ return
403
+
404
+ operation.fail(error, error_metadata)
405
+ self._move_to_history(operation)
406
+
407
+ self.stats["operations_failed"] += 1
408
+
409
+ # Emit failure event
410
+ await self._emit_event(
411
+ AgentEvent(
412
+ event_type=EventType.OPERATION_FAILED,
413
+ agent_id=operation.agent_id,
414
+ operation_id=operation_id,
415
+ data={
416
+ "error": error,
417
+ "duration": (
418
+ operation.duration.total_seconds()
419
+ if operation.duration
420
+ else 0
421
+ ),
422
+ "error_metadata": error_metadata or {},
423
+ },
424
+ source="UnifiedStatusManager",
425
+ )
426
+ )
427
+
428
+ async def cancel_operation(
429
+ self, operation_id: str, reason: Optional[str] = None
430
+ ) -> None:
431
+ """
432
+ Cancel an operation.
433
+
434
+ Args:
435
+ operation_id: ID of the operation
436
+ reason: Optional cancellation reason
437
+ """
438
+ async with self._lock:
439
+ operation = self.active_operations.get(operation_id)
440
+ if not operation:
441
+ return
442
+
443
+ operation.cancel(reason)
444
+ self._move_to_history(operation)
445
+
446
+ self.stats["operations_cancelled"] += 1
447
+
448
+ # Emit cancellation event
449
+ await self._emit_event(
450
+ AgentEvent(
451
+ event_type=EventType.OPERATION_CANCELLED,
452
+ agent_id=operation.agent_id,
453
+ operation_id=operation_id,
454
+ data={
455
+ "reason": reason or "No reason provided",
456
+ "duration": (
457
+ operation.duration.total_seconds()
458
+ if operation.duration
459
+ else 0
460
+ ),
461
+ },
462
+ source="UnifiedStatusManager",
463
+ )
464
+ )
465
+
466
+ def get_operation(self, operation_id: str) -> Optional[AgentOperation]:
467
+ """Get operation by ID."""
468
+ return self.active_operations.get(operation_id)
469
+
470
+ def get_active_operations(
471
+ self, agent_id: Optional[str] = None
472
+ ) -> List[AgentOperation]:
473
+ """Get list of active operations, optionally filtered by agent."""
474
+ operations = list(self.active_operations.values())
475
+ if agent_id:
476
+ operations = [op for op in operations if op.agent_id == agent_id]
477
+ return operations
478
+
479
+ def get_operation_history(
480
+ self, agent_id: Optional[str] = None, limit: Optional[int] = None
481
+ ) -> List[AgentOperation]:
482
+ """Get operation history, optionally filtered by agent."""
483
+ history = list(self.operation_history)
484
+ if agent_id:
485
+ history = [op for op in history if op.agent_id == agent_id]
486
+ if limit:
487
+ history = history[-limit:]
488
+ return history
489
+
490
+ # Event System
491
+ def add_event_listener(self, listener: EventListener) -> None:
492
+ """Add an event listener."""
493
+ self.event_listeners.append(listener)
494
+
495
+ def remove_event_listener(self, listener: EventListener) -> None:
496
+ """Remove an event listener."""
497
+ if listener in self.event_listeners:
498
+ self.event_listeners.remove(listener)
499
+ listener.stop()
500
+
501
+ # Streaming System
502
+ async def emit_stream_event(
503
+ self,
504
+ event: AgentEvent,
505
+ priority: StreamPriority = StreamPriority.NORMAL,
506
+ ) -> bool:
507
+ """
508
+ Emit an event to the status stream.
509
+
510
+ Args:
511
+ event: Agent event to emit
512
+ priority: Priority level for the event
513
+
514
+ Returns:
515
+ True if event was queued successfully
516
+ """
517
+ if not self.running:
518
+ return False
519
+
520
+ try:
521
+ stream_event = StatusStreamEvent(event=event, priority=priority)
522
+
523
+ # Use priority queue for high/critical events
524
+ if priority in [StreamPriority.HIGH, StreamPriority.CRITICAL]:
525
+ await self.priority_queue.put(stream_event)
526
+ else:
527
+ await self.stream_event_queue.put(stream_event)
528
+
529
+ return True
530
+
531
+ except asyncio.QueueFull:
532
+ self.metrics.events_dropped += 1
533
+ logger.warning(
534
+ f"Status stream queue full, dropping event: {event.event_type}"
535
+ )
536
+ return False
537
+ except Exception as e:
538
+ logger.error(f"Error emitting stream event: {e}")
539
+ return False
540
+
541
+ async def add_stream_listener(self, listener: StatusStreamListener) -> bool:
542
+ """
543
+ Add a listener to the status stream.
544
+
545
+ Args:
546
+ listener: Status stream listener to add
547
+
548
+ Returns:
549
+ True if listener was added successfully
550
+ """
551
+ if len(self.stream_listeners) >= self.config.max_listeners:
552
+ logger.warning(f"Maximum listeners ({self.config.max_listeners}) reached")
553
+ return False
554
+
555
+ async with self._lock:
556
+ if listener.listener_id in self.stream_listeners:
557
+ logger.warning(f"Listener {listener.listener_id} already exists")
558
+ return False
559
+
560
+ self.stream_listeners[listener.listener_id] = listener
561
+ self.metrics.listeners_count = len(self.stream_listeners)
562
+
563
+ logger.debug(f"Added stream listener: {listener.listener_id}")
564
+ return True
565
+
566
+ async def remove_stream_listener(self, listener_id: str) -> bool:
567
+ """
568
+ Remove a listener from the status stream.
569
+
570
+ Args:
571
+ listener_id: ID of listener to remove
572
+
573
+ Returns:
574
+ True if listener was removed successfully
575
+ """
576
+ async with self._lock:
577
+ if listener_id not in self.stream_listeners:
578
+ return False
579
+
580
+ # Stop the listener
581
+ self.stream_listeners[listener_id].stop()
582
+
583
+ # Clean up
584
+ del self.stream_listeners[listener_id]
585
+ self.metrics.listeners_count = len(self.stream_listeners)
586
+
587
+ logger.debug(f"Removed stream listener: {listener_id}")
588
+ return True
589
+
590
+ # Statistics and Health
591
+ def get_stats(self) -> Dict[str, Any]:
592
+ """Get current statistics."""
593
+ return {
594
+ **self.stats,
595
+ "active_operations": len(self.active_operations),
596
+ "total_agents": len(self.agent_statuses),
597
+ "enabled_agents": sum(
598
+ 1
599
+ for status in self.agent_statuses.values()
600
+ if status == AgentStatus.ENABLED
601
+ ),
602
+ "event_listeners": len(self.event_listeners),
603
+ "stream_listeners": len(self.stream_listeners),
604
+ }
605
+
606
+ def get_metrics(self) -> StreamMetrics:
607
+ """Get current stream metrics."""
608
+ self.metrics.queue_size = (
609
+ self.stream_event_queue.qsize() + self.priority_queue.qsize()
610
+ )
611
+ return self.metrics
612
+
613
+ async def health_check(self) -> Dict[str, Any]:
614
+ """Perform health check on the status system."""
615
+ metrics = self.get_metrics()
616
+
617
+ return {
618
+ "running": self.running,
619
+ "queue_size": metrics.queue_size,
620
+ "queue_full": metrics.queue_size >= self.config.max_queue_size * 0.9,
621
+ "listeners_count": metrics.listeners_count,
622
+ "events_processed": metrics.events_processed,
623
+ "events_dropped": metrics.events_dropped,
624
+ "avg_processing_time": metrics.avg_processing_time,
625
+ "last_update": (
626
+ metrics.last_update.isoformat() if metrics.last_update else None
627
+ ),
628
+ "healthy": (
629
+ self.running
630
+ and metrics.queue_size < self.config.max_queue_size * 0.9
631
+ and metrics.avg_processing_time < 0.1
632
+ ),
633
+ }
634
+
635
+ # Private Methods
636
+ async def _emit_event(
637
+ self,
638
+ event: AgentEvent,
639
+ stream_priority: Optional[StreamPriority] = None,
640
+ ) -> None:
641
+ """
642
+ Emit an event to both event listeners and stream listeners.
643
+
644
+ Args:
645
+ event: The event to emit
646
+ stream_priority: Priority level for status stream (if available)
647
+ """
648
+ self.stats["events_emitted"] += 1
649
+ self.event_history.append(event)
650
+
651
+ # Add to event processing queue
652
+ await self.event_queue.put(event)
653
+
654
+ # Emit to stream if priority is specified
655
+ if stream_priority is not None:
656
+ await self.emit_stream_event(event, stream_priority)
657
+ else:
658
+ # Determine priority based on event type
659
+ if event.event_type in [
660
+ EventType.ERROR_OCCURRED,
661
+ EventType.OPERATION_FAILED,
662
+ ]:
663
+ stream_priority = StreamPriority.HIGH
664
+ elif event.event_type in [
665
+ EventType.AGENT_STATE_CHANGED,
666
+ EventType.APPROVAL_REQUESTED,
667
+ ]:
668
+ stream_priority = StreamPriority.NORMAL
669
+ else:
670
+ stream_priority = StreamPriority.LOW
671
+
672
+ await self.emit_stream_event(event, stream_priority)
673
+
674
+ async def _process_events(self) -> None:
675
+ """Background task to process events for event listeners."""
676
+ logger.debug("Event processor started")
677
+
678
+ while self.running and not self.shutdown_event.is_set():
679
+ try:
680
+ # Wait for event with timeout to allow shutdown checking
681
+ event = await asyncio.wait_for(self.event_queue.get(), timeout=1.0)
682
+
683
+ # Process event for all listeners
684
+ for listener in self.event_listeners[
685
+ :
686
+ ]: # Copy list to avoid modification during iteration
687
+ try:
688
+ await listener.handle_event(event)
689
+ except Exception as e:
690
+ logger.error(
691
+ f"Error in event listener {type(listener).__name__}: {e}"
692
+ )
693
+
694
+ except asyncio.TimeoutError:
695
+ continue # Check shutdown flag
696
+ except Exception as e:
697
+ logger.error(f"Error processing event: {e}")
698
+
699
+ logger.debug("Event processor stopped")
700
+
701
+ async def _process_stream_events(self) -> None:
702
+ """Background task to process stream events."""
703
+ logger.debug("Stream event processor started")
704
+
705
+ while self.running and not self.shutdown_event.is_set():
706
+ try:
707
+ # Process priority events first
708
+ try:
709
+ stream_event = await asyncio.wait_for(
710
+ self.priority_queue.get(), timeout=0.01
711
+ )
712
+ await self._handle_stream_event(stream_event)
713
+ continue
714
+ except asyncio.TimeoutError:
715
+ pass
716
+
717
+ # Process normal events with throttling
718
+ try:
719
+ stream_event = await asyncio.wait_for(
720
+ self.stream_event_queue.get(), timeout=0.1
721
+ )
722
+
723
+ # Apply throttling for normal events
724
+ if self.config.throttle_rate > 0:
725
+ now = time.time()
726
+ time_since_last = now - self.last_event_time
727
+ min_interval = 1.0 / self.config.throttle_rate
728
+
729
+ if time_since_last < min_interval:
730
+ await asyncio.sleep(min_interval - time_since_last)
731
+
732
+ await self._handle_stream_event(stream_event)
733
+ self.last_event_time = time.time()
734
+
735
+ except asyncio.TimeoutError:
736
+ # Timeout is normal, just continue
737
+ continue
738
+
739
+ except Exception as e:
740
+ logger.error(f"Error in stream event processor: {e}")
741
+ await asyncio.sleep(0.1)
742
+
743
+ logger.debug("Stream event processor stopped")
744
+
745
+ async def _handle_stream_event(self, stream_event: StatusStreamEvent) -> None:
746
+ """Handle a single stream event."""
747
+ start_time = time.time()
748
+
749
+ try:
750
+ # Update metrics
751
+ self.metrics.queue_size = (
752
+ self.stream_event_queue.qsize() + self.priority_queue.qsize()
753
+ )
754
+
755
+ # Add to collector for batching
756
+ should_flush = await self.update_collector.add_update(stream_event)
757
+
758
+ # Distribute to listeners if we should flush
759
+ if should_flush:
760
+ updates = await self.update_collector.flush()
761
+ for update in updates:
762
+ await self._distribute_to_stream_listeners(update)
763
+
764
+ self.metrics.events_processed += 1
765
+
766
+ except Exception as e:
767
+ logger.error(f"Error handling stream event: {e}")
768
+ finally:
769
+ processing_time = time.time() - start_time
770
+ self.metrics.update_processing_time(processing_time)
771
+ self.metrics.last_update = datetime.now()
772
+
773
+ async def _distribute_events(self) -> None:
774
+ """Background task to distribute events to stream listeners."""
775
+ logger.debug("Event distributor started")
776
+
777
+ while self.running and not self.shutdown_event.is_set():
778
+ try:
779
+ # Periodically flush any remaining updates
780
+ await asyncio.sleep(0.05) # 20 times per second
781
+
782
+ updates = await self.update_collector.flush()
783
+ for update in updates:
784
+ await self._distribute_to_stream_listeners(update)
785
+
786
+ except Exception as e:
787
+ logger.error(f"Error in event distributor: {e}")
788
+ await asyncio.sleep(0.1)
789
+
790
+ logger.debug("Event distributor stopped")
791
+
792
+ async def _distribute_to_stream_listeners(
793
+ self, stream_event: StatusStreamEvent
794
+ ) -> None:
795
+ """Distribute event to all relevant stream listeners."""
796
+ if not self.stream_listeners:
797
+ return
798
+
799
+ # Create list of listeners to avoid modification during iteration
800
+ listeners_list = list(self.stream_listeners.values())
801
+
802
+ # Distribute to listeners concurrently
803
+ tasks = []
804
+ for listener in listeners_list:
805
+ if listener.active:
806
+ task = asyncio.create_task(listener.handle_event(stream_event))
807
+ tasks.append(task)
808
+
809
+ # Wait for all listeners to process (with timeout)
810
+ if tasks:
811
+ try:
812
+ await asyncio.wait_for(
813
+ asyncio.gather(*tasks, return_exceptions=True), timeout=1.0
814
+ )
815
+ except asyncio.TimeoutError:
816
+ logger.warning("Some stream listeners timed out processing event")
817
+
818
+ def _move_to_history(self, operation: AgentOperation) -> None:
819
+ """Move completed operation from active to history."""
820
+ if operation.operation_id in self.active_operations:
821
+ del self.active_operations[operation.operation_id]
822
+
823
+ self.operation_history.append(operation)
824
+
825
+ # Clean up agent operation list
826
+ if (
827
+ operation.agent_id
828
+ and operation.operation_id in self.operations_by_agent[operation.agent_id]
829
+ ):
830
+ self.operations_by_agent[operation.agent_id].remove(operation.operation_id)
831
+
832
+
833
+ # Global instance management
834
+ _global_status_manager: Optional[UnifiedStatusManager] = None
835
+
836
+
837
+ def get_status_manager() -> UnifiedStatusManager:
838
+ """Get the global status manager instance."""
839
+ global _global_status_manager
840
+ if _global_status_manager is None:
841
+ _global_status_manager = UnifiedStatusManager()
842
+ return _global_status_manager
843
+
844
+
845
+ def set_status_manager(manager: UnifiedStatusManager) -> None:
846
+ """Set the global status manager instance."""
847
+ global _global_status_manager
848
+ _global_status_manager = manager
849
+
850
+
851
+ async def initialize_status_system(
852
+ config: Optional[StatusDisplayConfig] = None,
853
+ ) -> UnifiedStatusManager:
854
+ """
855
+ Initialize the global status system.
856
+
857
+ Args:
858
+ config: Optional configuration for the status system
859
+
860
+ Returns:
861
+ The initialized status manager
862
+ """
863
+ if config:
864
+ manager = UnifiedStatusManager(config)
865
+ set_status_manager(manager)
866
+ else:
867
+ manager = get_status_manager()
868
+
869
+ await manager.initialize()
870
+ return manager
871
+
872
+
873
+ async def shutdown_status_system() -> None:
874
+ """Shutdown the global status system."""
875
+ global _global_status_manager
876
+ if _global_status_manager:
877
+ await _global_status_manager.shutdown()
878
+ _global_status_manager = None