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.
- omnimancer/__init__.py +42 -0
- omnimancer/__main__.py +9 -0
- omnimancer/cli/__init__.py +8 -0
- omnimancer/cli/agent_manager.py +826 -0
- omnimancer/cli/agent_wizard.py +944 -0
- omnimancer/cli/approval_formatter.py +734 -0
- omnimancer/cli/approval_integration.py +734 -0
- omnimancer/cli/approval_prompt.py +1286 -0
- omnimancer/cli/batch_approval_display.py +617 -0
- omnimancer/cli/batch_approval_filters.py +729 -0
- omnimancer/cli/batch_operation_monitor.py +775 -0
- omnimancer/cli/commands/performance_cmd.py +624 -0
- omnimancer/cli/commands.py +757 -0
- omnimancer/cli/completion_handler.py +281 -0
- omnimancer/cli/config_setup_wizard.py +783 -0
- omnimancer/cli/dynamic_commands.py +371 -0
- omnimancer/cli/handlers/__init__.py +18 -0
- omnimancer/cli/handlers/agent_handler.py +865 -0
- omnimancer/cli/handlers/agent_persona_handler.py +1048 -0
- omnimancer/cli/handlers/custom_agent_handler.py +352 -0
- omnimancer/cli/handlers/permissions_handler.py +1609 -0
- omnimancer/cli/interface.py +6695 -0
- omnimancer/cli/performance_dashboard.py +612 -0
- omnimancer/core/__init__.py +34 -0
- omnimancer/core/agent/__init__.py +256 -0
- omnimancer/core/agent/agent_config.py +715 -0
- omnimancer/core/agent/agent_switcher.py +668 -0
- omnimancer/core/agent/approval_context.py +460 -0
- omnimancer/core/agent/approval_dialog.py +902 -0
- omnimancer/core/agent/approval_interface.py +628 -0
- omnimancer/core/agent/approval_manager.py +760 -0
- omnimancer/core/agent/cli_integration.py +470 -0
- omnimancer/core/agent/config.py +845 -0
- omnimancer/core/agent/custom_agent_integration.py +654 -0
- omnimancer/core/agent/diff_renderer.py +706 -0
- omnimancer/core/agent/file_content_display.py +503 -0
- omnimancer/core/agent/file_modification_workflow.py +658 -0
- omnimancer/core/agent/file_system_manager.py +1790 -0
- omnimancer/core/agent/input_handler.py +614 -0
- omnimancer/core/agent/metrics_collector.py +771 -0
- omnimancer/core/agent/optimization_engine.py +729 -0
- omnimancer/core/agent/performance_monitor.py +717 -0
- omnimancer/core/agent/persona.py +870 -0
- omnimancer/core/agent/persona_validator.py +672 -0
- omnimancer/core/agent/program_executor.py +748 -0
- omnimancer/core/agent/proposed_changes_integration.py +663 -0
- omnimancer/core/agent/read_before_write_errors.py +492 -0
- omnimancer/core/agent/read_before_write_ui.py +647 -0
- omnimancer/core/agent/rich_renderer.py +763 -0
- omnimancer/core/agent/status_config.py +238 -0
- omnimancer/core/agent/status_core.py +446 -0
- omnimancer/core/agent/status_integrated_engine.py +540 -0
- omnimancer/core/agent/status_manager.py +878 -0
- omnimancer/core/agent/status_ui_components.py +535 -0
- omnimancer/core/agent/template_applicator.py +548 -0
- omnimancer/core/agent/terminal_status_listener.py +394 -0
- omnimancer/core/agent/token_tracker.py +679 -0
- omnimancer/core/agent/types.py +60 -0
- omnimancer/core/agent/web_client.py +745 -0
- omnimancer/core/agent/workflow_orchestrator.py +826 -0
- omnimancer/core/agent_engine.py +1625 -0
- omnimancer/core/agent_mode_manager.py +706 -0
- omnimancer/core/agent_progress_ui.py +545 -0
- omnimancer/core/chat_manager.py +141 -0
- omnimancer/core/config_manager.py +2199 -0
- omnimancer/core/config_migration.py +772 -0
- omnimancer/core/config_provider.py +1974 -0
- omnimancer/core/config_security_validator.py +393 -0
- omnimancer/core/config_validator.py +768 -0
- omnimancer/core/conversation_manager.py +352 -0
- omnimancer/core/engine.py +787 -0
- omnimancer/core/enhanced_input.py +523 -0
- omnimancer/core/env_loader.py +83 -0
- omnimancer/core/fallback_manager.py +568 -0
- omnimancer/core/health_monitor.py +462 -0
- omnimancer/core/history_manager.py +462 -0
- omnimancer/core/mcp_integration_layer.py +878 -0
- omnimancer/core/models.py +1889 -0
- omnimancer/core/provider_capabilities.py +403 -0
- omnimancer/core/provider_initializer.py +450 -0
- omnimancer/core/provider_registry.py +432 -0
- omnimancer/core/security/__init__.py +16 -0
- omnimancer/core/security/approval_workflow.py +406 -0
- omnimancer/core/security/audit_logger.py +511 -0
- omnimancer/core/security/permission_controller.py +511 -0
- omnimancer/core/security/sandbox_manager.py +379 -0
- omnimancer/core/security/security_manager.py +530 -0
- omnimancer/core/setup_wizard.py +345 -0
- omnimancer/core/setup_wizard_core.py +449 -0
- omnimancer/core/setup_wizard_guide_generator.py +376 -0
- omnimancer/core/setup_wizard_provider_setup.py +454 -0
- omnimancer/core/setup_wizard_ui.py +317 -0
- omnimancer/core/setup_wizard_validation.py +371 -0
- omnimancer/core/signal_handler.py +174 -0
- omnimancer/mcp/__init__.py +11 -0
- omnimancer/mcp/client.py +725 -0
- omnimancer/mcp/manager.py +649 -0
- omnimancer/providers/__init__.py +24 -0
- omnimancer/providers/azure.py +608 -0
- omnimancer/providers/base.py +262 -0
- omnimancer/providers/bedrock.py +711 -0
- omnimancer/providers/claude.py +426 -0
- omnimancer/providers/claude_code.py +406 -0
- omnimancer/providers/cohere.py +412 -0
- omnimancer/providers/factory.py +988 -0
- omnimancer/providers/gemini.py +708 -0
- omnimancer/providers/mistral.py +554 -0
- omnimancer/providers/ollama.py +510 -0
- omnimancer/providers/openai.py +362 -0
- omnimancer/providers/openrouter.py +709 -0
- omnimancer/providers/perplexity.py +476 -0
- omnimancer/providers/vertex.py +617 -0
- omnimancer/providers/xai.py +484 -0
- omnimancer/ui/__init__.py +6 -0
- omnimancer/ui/cancellation_handler.py +235 -0
- omnimancer/ui/progress_indicator.py +318 -0
- omnimancer/utils/__init__.py +22 -0
- omnimancer/utils/error_handler.py +589 -0
- omnimancer/utils/errors.py +667 -0
- omnimancer/utils/retry.py +238 -0
- omnimancer_cli-0.1.0.dist-info/METADATA +413 -0
- omnimancer_cli-0.1.0.dist-info/RECORD +126 -0
- omnimancer_cli-0.1.0.dist-info/WHEEL +5 -0
- omnimancer_cli-0.1.0.dist-info/entry_points.txt +4 -0
- omnimancer_cli-0.1.0.dist-info/licenses/LICENSE +21 -0
- 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
|