Fix: update message type to current implementation
This commit is contained in:
ziesorx
parent
b940790e4a
commit
96ecc321ec
3 changed files with 549 additions and 230 deletions
|
|
@ -83,16 +83,17 @@ class WebSocketHandler:
|
|||
|
||||
# Message handlers
|
||||
self.message_handlers: Dict[str, MessageHandler] = {
|
||||
"subscribe": self._handle_subscribe,
|
||||
"unsubscribe": self._handle_unsubscribe,
|
||||
"setSubscriptionList": self._handle_set_subscription_list,
|
||||
"requestState": self._handle_request_state,
|
||||
"setSessionId": self._handle_set_session,
|
||||
"patchSession": self._handle_patch_session,
|
||||
"setProgressionStage": self._handle_set_progression_stage
|
||||
"setProgressionStage": self._handle_set_progression_stage,
|
||||
"patchSessionResult": self._handle_patch_session_result
|
||||
}
|
||||
|
||||
# Session and display management
|
||||
self.session_ids: Dict[str, str] = {} # display_identifier -> session_id
|
||||
self.progression_stages: Dict[str, str] = {} # display_identifier -> progression_stage
|
||||
self.display_identifiers: Set[str] = set()
|
||||
|
||||
# Camera monitor
|
||||
|
|
@ -171,22 +172,40 @@ class WebSocketHandler:
|
|||
# Get system metrics
|
||||
metrics = get_system_metrics()
|
||||
|
||||
# Get active streams info
|
||||
active_streams = self.stream_manager.get_active_streams()
|
||||
active_models = self.model_manager.get_loaded_models()
|
||||
# Build cameraConnections array as required by protocol
|
||||
camera_connections = []
|
||||
with self.stream_manager.streams_lock:
|
||||
for camera_id, stream_info in self.stream_manager.streams.items():
|
||||
# Check if camera is online
|
||||
is_online = self.stream_manager.is_stream_active(camera_id)
|
||||
|
||||
connection_info = {
|
||||
"subscriptionIdentifier": stream_info.get("subscriptionIdentifier", camera_id),
|
||||
"modelId": stream_info.get("modelId", 0),
|
||||
"modelName": stream_info.get("modelName", "Unknown Model"),
|
||||
"online": is_online
|
||||
}
|
||||
|
||||
# Add crop coordinates if available
|
||||
if "cropX1" in stream_info:
|
||||
connection_info["cropX1"] = stream_info["cropX1"]
|
||||
if "cropY1" in stream_info:
|
||||
connection_info["cropY1"] = stream_info["cropY1"]
|
||||
if "cropX2" in stream_info:
|
||||
connection_info["cropX2"] = stream_info["cropX2"]
|
||||
if "cropY2" in stream_info:
|
||||
connection_info["cropY2"] = stream_info["cropY2"]
|
||||
|
||||
camera_connections.append(connection_info)
|
||||
|
||||
# Protocol-compliant stateReport format (worker.md lines 169-189)
|
||||
state_data = {
|
||||
"type": "stateReport",
|
||||
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
"data": {
|
||||
"activeStreams": len(active_streams),
|
||||
"loadedModels": len(active_models),
|
||||
"cpuUsage": metrics.get("cpu_percent", 0),
|
||||
"memoryUsage": metrics.get("memory_percent", 0),
|
||||
"gpuUsage": metrics.get("gpu_percent", 0),
|
||||
"gpuMemory": metrics.get("gpu_memory_percent", 0),
|
||||
"uptime": time.time() - metrics.get("start_time", time.time())
|
||||
}
|
||||
"cpuUsage": metrics.get("cpu_percent", 0),
|
||||
"memoryUsage": metrics.get("memory_percent", 0),
|
||||
"gpuUsage": metrics.get("gpu_percent", 0),
|
||||
"gpuMemoryUsage": metrics.get("gpu_memory_percent", 0), # Fixed field name
|
||||
"cameraConnections": camera_connections
|
||||
}
|
||||
|
||||
# Compact JSON for RX/TX logging
|
||||
|
|
@ -351,75 +370,105 @@ class WebSocketHandler:
|
|||
traceback.print_exc()
|
||||
return persistent_data
|
||||
|
||||
async def _handle_subscribe(self, data: Dict[str, Any]) -> None:
|
||||
"""Handle stream subscription request."""
|
||||
payload = data.get("payload", {})
|
||||
subscription_id = payload.get("subscriptionIdentifier")
|
||||
async def _handle_set_subscription_list(self, data: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Handle setSubscriptionList command - declarative subscription management.
|
||||
|
||||
This is the primary subscription command per worker.md protocol.
|
||||
Workers must reconcile the new subscription list with current state.
|
||||
"""
|
||||
subscriptions = data.get("subscriptions", [])
|
||||
|
||||
if not subscription_id:
|
||||
logger.error("Missing subscriptionIdentifier in subscribe payload")
|
||||
return
|
||||
|
||||
try:
|
||||
# Extract display and camera IDs
|
||||
parts = subscription_id.split(";")
|
||||
if len(parts) >= 2:
|
||||
display_id = parts[0]
|
||||
camera_id = parts[1]
|
||||
self.display_identifiers.add(display_id)
|
||||
else:
|
||||
camera_id = subscription_id
|
||||
# Get current subscription identifiers
|
||||
current_subscriptions = set(subscription_to_camera.keys())
|
||||
|
||||
# Get desired subscription identifiers
|
||||
desired_subscriptions = set()
|
||||
subscription_configs = {}
|
||||
|
||||
for sub_config in subscriptions:
|
||||
sub_id = sub_config.get("subscriptionIdentifier")
|
||||
if sub_id:
|
||||
desired_subscriptions.add(sub_id)
|
||||
subscription_configs[sub_id] = sub_config
|
||||
|
||||
# Extract display ID for session management
|
||||
parts = sub_id.split(";")
|
||||
if len(parts) >= 2:
|
||||
display_id = parts[0]
|
||||
self.display_identifiers.add(display_id)
|
||||
|
||||
# Calculate changes needed
|
||||
to_add = desired_subscriptions - current_subscriptions
|
||||
to_remove = current_subscriptions - desired_subscriptions
|
||||
to_update = desired_subscriptions & current_subscriptions
|
||||
|
||||
logger.info(f"Subscription reconciliation: add={len(to_add)}, remove={len(to_remove)}, update={len(to_update)}")
|
||||
|
||||
# Remove obsolete subscriptions
|
||||
for sub_id in to_remove:
|
||||
camera_id = subscription_to_camera.get(sub_id)
|
||||
if camera_id:
|
||||
await self.stream_manager.stop_stream(camera_id)
|
||||
self.model_manager.unload_models(camera_id)
|
||||
subscription_to_camera.pop(sub_id, None)
|
||||
self.session_cache.clear_session(camera_id)
|
||||
logger.info(f"Removed subscription: {sub_id}")
|
||||
|
||||
# Add new subscriptions
|
||||
for sub_id in to_add:
|
||||
await self._start_subscription(sub_id, subscription_configs[sub_id])
|
||||
logger.info(f"Added subscription: {sub_id}")
|
||||
|
||||
# Update existing subscriptions if needed
|
||||
for sub_id in to_update:
|
||||
# Check if configuration changed (model URL, crop coordinates, etc.)
|
||||
current_config = subscription_to_camera.get(sub_id)
|
||||
new_config = subscription_configs[sub_id]
|
||||
|
||||
# For now, restart subscription if model URL changed (handles S3 expiration)
|
||||
current_model_url = getattr(current_config, 'model_url', None) if current_config else None
|
||||
new_model_url = new_config.get("modelUrl")
|
||||
|
||||
if current_model_url != new_model_url:
|
||||
# Restart with new configuration
|
||||
camera_id = subscription_to_camera.get(sub_id)
|
||||
if camera_id:
|
||||
await self.stream_manager.stop_stream(camera_id)
|
||||
self.model_manager.unload_models(camera_id)
|
||||
|
||||
await self._start_subscription(sub_id, new_config)
|
||||
logger.info(f"Updated subscription: {sub_id}")
|
||||
|
||||
logger.info(f"Subscription list reconciliation completed. Active: {len(desired_subscriptions)}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling setSubscriptionList: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
async def _start_subscription(self, subscription_id: str, config: Dict[str, Any]) -> None:
|
||||
"""Start a single subscription with given configuration."""
|
||||
try:
|
||||
# Extract camera ID from subscription identifier
|
||||
parts = subscription_id.split(";")
|
||||
camera_id = parts[1] if len(parts) >= 2 else subscription_id
|
||||
|
||||
# Store subscription mapping
|
||||
subscription_to_camera[subscription_id] = camera_id
|
||||
|
||||
# Start camera stream
|
||||
await self.stream_manager.start_stream(camera_id, payload)
|
||||
await self.stream_manager.start_stream(camera_id, config)
|
||||
|
||||
# Load model
|
||||
model_id = payload.get("modelId")
|
||||
model_url = payload.get("modelUrl")
|
||||
model_id = config.get("modelId")
|
||||
model_url = config.get("modelUrl")
|
||||
if model_id and model_url:
|
||||
await self.model_manager.load_model(camera_id, model_id, model_url)
|
||||
|
||||
logger.info(f"Subscribed to stream: {subscription_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling subscription: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
async def _handle_unsubscribe(self, data: Dict[str, Any]) -> None:
|
||||
"""Handle stream unsubscription request."""
|
||||
payload = data.get("payload", {})
|
||||
subscription_id = payload.get("subscriptionIdentifier")
|
||||
|
||||
if not subscription_id:
|
||||
logger.error("Missing subscriptionIdentifier in unsubscribe payload")
|
||||
return
|
||||
|
||||
try:
|
||||
# Get camera ID from subscription
|
||||
camera_id = subscription_to_camera.get(subscription_id)
|
||||
if not camera_id:
|
||||
logger.warning(f"No camera found for subscription: {subscription_id}")
|
||||
return
|
||||
|
||||
# Stop stream
|
||||
await self.stream_manager.stop_stream(camera_id)
|
||||
|
||||
# Unload model
|
||||
self.model_manager.unload_models(camera_id)
|
||||
|
||||
# Clean up mappings
|
||||
subscription_to_camera.pop(subscription_id, None)
|
||||
|
||||
# Clean up session state
|
||||
self.session_cache.clear_session(camera_id)
|
||||
|
||||
logger.info(f"Unsubscribed from stream: {subscription_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling unsubscription: {e}")
|
||||
logger.error(f"Error starting subscription {subscription_id}: {e}")
|
||||
raise
|
||||
traceback.print_exc()
|
||||
|
||||
async def _handle_request_state(self, data: Dict[str, Any]) -> None:
|
||||
|
|
@ -534,6 +583,26 @@ class WebSocketHandler:
|
|||
elif progression_stage in ["car_wait_staff"]:
|
||||
pipeline_state["progression_stage"] = progression_stage
|
||||
logger.info(f"📋 Camera {camera_id}: Progression stage set to {progression_stage}")
|
||||
|
||||
# Store progression stage for this display
|
||||
if display_id and progression_stage is not None:
|
||||
if progression_stage:
|
||||
self.progression_stages[display_id] = progression_stage
|
||||
else:
|
||||
# Clear progression stage if null
|
||||
self.progression_stages.pop(display_id, None)
|
||||
|
||||
async def _handle_patch_session_result(self, data: Dict[str, Any]) -> None:
|
||||
"""Handle patchSessionResult message from backend."""
|
||||
payload = data.get("payload", {})
|
||||
session_id = payload.get("sessionId")
|
||||
success = payload.get("success", False)
|
||||
message = payload.get("message", "")
|
||||
|
||||
if success:
|
||||
logger.info(f"Patch session {session_id} successful: {message}")
|
||||
else:
|
||||
logger.warning(f"Patch session {session_id} failed: {message}")
|
||||
|
||||
async def _send_detection_result(
|
||||
self,
|
||||
|
|
@ -542,10 +611,16 @@ class WebSocketHandler:
|
|||
detection_result: DetectionResult
|
||||
) -> None:
|
||||
"""Send detection result over WebSocket."""
|
||||
# Get session ID for this display
|
||||
subscription_id = stream_info["subscriptionIdentifier"]
|
||||
display_id = subscription_id.split(";")[0] if ";" in subscription_id else subscription_id
|
||||
session_id = self.session_ids.get(display_id)
|
||||
|
||||
detection_data = {
|
||||
"type": "imageDetection",
|
||||
"subscriptionIdentifier": stream_info["subscriptionIdentifier"],
|
||||
"subscriptionIdentifier": subscription_id,
|
||||
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
"sessionId": session_id, # Required by protocol
|
||||
"data": {
|
||||
"detection": detection_result.to_dict(),
|
||||
"modelId": stream_info["modelId"],
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue