backend redesign

backend.old/src/gateway/channels/websocket.py

@@ -0,0 +1,99 @@
+import asyncio
+import json
+import logging
+from typing import AsyncIterator, Optional
+from datetime import datetime
+from fastapi import WebSocket
+
+from gateway.channels.base import Channel
+from gateway.protocol import UserMessage, AgentMessage, WebSocketAgentUserMessage, WebSocketAgentChunk
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketChannel(Channel):
+    """WebSocket-based communication channel.
+
+    Integrates with the existing WebSocket endpoint to provide
+    bidirectional agent communication with streaming support.
+    """
+
+    def __init__(self, channel_id: str, websocket: WebSocket, session_id: str):
+        super().__init__(channel_id, "websocket")
+        self.websocket = websocket
+        self.session_id = session_id
+        self._connected = True
+        self._receive_queue: asyncio.Queue[UserMessage] = asyncio.Queue()
+        self._receive_task: Optional[asyncio.Task] = None
+
+    def supports_streaming(self) -> bool:
+        """WebSocket supports streaming responses."""
+        return True
+
+    def supports_attachments(self) -> bool:
+        """WebSocket can support attachments via URLs."""
+        return True
+
+    async def send(self, message: AgentMessage) -> None:
+        """Send agent message through WebSocket.
+
+        For streaming messages, sends chunks as they arrive.
+        For complete messages, sends as a single chunk.
+        """
+        if not self._connected:
+            logger.warning(f"Cannot send message, channel {self.channel_id} not connected")
+            return
+
+        try:
+            chunk = WebSocketAgentChunk(
+                session_id=message.session_id,
+                content=message.content,
+                done=message.done,
+                metadata=message.metadata
+            )
+            chunk_data = chunk.model_dump(mode="json")
+            logger.debug(f"Sending WebSocket message: done={message.done}, content_length={len(message.content)}")
+            await self.websocket.send_json(chunk_data)
+            logger.debug(f"WebSocket message sent successfully")
+        except Exception as e:
+            logger.error(f"WebSocket send error: {e}", exc_info=True)
+            self._connected = False
+
+    async def receive(self) -> AsyncIterator[UserMessage]:
+        """Receive messages from WebSocket.
+
+        Yields:
+            UserMessage objects as they arrive from the client
+        """
+        try:
+            while self._connected:
+                # Read from WebSocket
+                data = await self.websocket.receive_text()
+                message_json = json.loads(data)
+
+                # Only process agent_user_message types
+                if message_json.get("type") == "agent_user_message":
+                    msg = WebSocketAgentUserMessage(**message_json)
+
+                    user_msg = UserMessage(
+                        session_id=msg.session_id,
+                        channel_id=self.channel_id,
+                        content=msg.content,
+                        attachments=msg.attachments,
+                        timestamp=datetime.utcnow()
+                    )
+                    yield user_msg
+        except Exception as e:
+            print(f"WebSocket receive error: {e}")
+            self._connected = False
+
+    async def close(self) -> None:
+        """Close the WebSocket connection."""
+        self._connected = False
+        if self._receive_task:
+            self._receive_task.cancel()
+            try:
+                await self._receive_task
+            except asyncio.CancelledError:
+                pass
+        # Note: WebSocket close is handled by the main WebSocket endpoint