""" Web Channel Adapter — WebSocket endpoint and message normalizer. The web channel lets portal users chat with AI employees directly from the Konstruct portal UI. Messages flow through the same agent pipeline as Slack and WhatsApp — the only difference is the transport layer. Message flow: 1. Browser opens WebSocket at /chat/ws/{conversation_id} 2. Client sends {"type": "auth", "userId": ..., "role": ..., "tenantId": ...} NOTE: Browsers cannot set custom HTTP headers on WebSocket connections, so auth credentials are sent as the first JSON message (Pitfall 1). 3. For each user message (type="message"): a. Server immediately sends {"type": "typing"} to client (CHAT-05) b. normalize_web_event() converts to KonstructMessage (channel=WEB) c. User message saved to web_conversation_messages d. handle_message.delay(msg | extras) dispatches to Celery pipeline e. Server subscribes to Redis pub-sub channel for the response f. When orchestrator publishes the response: - Save assistant message to web_conversation_messages - Send {"type": "response", "text": ..., "conversation_id": ...} to client 4. On disconnect: unsubscribe and close all Redis connections Design notes: - thread_id = conversation_id — scopes agent memory to one conversation (Pitfall 3) - Redis pub-sub connections closed in try/finally to prevent leaks (Pitfall 2) - DB access uses configure_rls_hook + current_tenant_id context var per project pattern - WebSocket is a long-lived connection; each message/response cycle is synchronous within the connection but non-blocking for other connections """ from __future__ import annotations import asyncio import json import logging import uuid from datetime import datetime, timezone from typing import Any import redis.asyncio as aioredis from fastapi import APIRouter, WebSocket, WebSocketDisconnect from sqlalchemy import select, text from orchestrator.agents.builder import build_messages_with_memory, build_system_prompt from orchestrator.agents.runner import run_agent_streaming from orchestrator.memory.short_term import get_recent_messages, append_message from orchestrator.tasks import handle_message, embed_and_store from shared.config import settings from shared.db import async_session_factory, engine from shared.models.chat import WebConversation, WebConversationMessage from shared.models.message import ChannelType, KonstructMessage, MessageContent, SenderInfo from shared.models.tenant import Agent from shared.redis_keys import webchat_response_key from shared.rls import configure_rls_hook, current_tenant_id logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Router — mounted in gateway/main.py # --------------------------------------------------------------------------- web_chat_router = APIRouter(tags=["web-chat"]) # Timeout for waiting for an agent response via Redis pub-sub (seconds) _RESPONSE_TIMEOUT_SECONDS = 180 def normalize_web_event(event: dict[str, Any]) -> KonstructMessage: """ Normalize a web channel event dict into a KonstructMessage. The web channel normalizer sets thread_id = conversation_id so that the agent memory pipeline scopes context to this conversation (Pitfall 3). Args: event: Dict with keys: text, tenant_id, agent_id, user_id, display_name, conversation_id. Returns: KonstructMessage with channel=WEB, thread_id=conversation_id. """ tenant_id: str = event.get("tenant_id", "") or "" user_id: str = event.get("user_id", "") or "" display_name: str = event.get("display_name", "Portal User") conversation_id: str = event.get("conversation_id", "") or "" text_content: str = event.get("text", "") or "" return KonstructMessage( id=str(uuid.uuid4()), tenant_id=tenant_id, channel=ChannelType.WEB, channel_metadata={ "portal_user_id": user_id, "tenant_id": tenant_id, "conversation_id": conversation_id, }, sender=SenderInfo( user_id=user_id, display_name=display_name, ), content=MessageContent( text=text_content, ), timestamp=datetime.now(timezone.utc), thread_id=conversation_id, reply_to=None, context={}, ) async def _handle_websocket_connection( websocket: WebSocket, conversation_id: str, ) -> None: """ Core WebSocket connection handler — separated for testability. Lifecycle: 1. Accept connection 2. Wait for auth message (browser cannot send custom headers) 3. Loop: receive messages → type indicator → Celery dispatch → Redis subscribe → response Args: websocket: The FastAPI WebSocket connection. conversation_id: The conversation UUID from the URL path. """ await websocket.accept() # ------------------------------------------------------------------------- # Step 1: Auth handshake # Browsers cannot send custom HTTP headers on WebSocket connections. # Auth credentials are sent as the first JSON message. # ------------------------------------------------------------------------- try: auth_msg = await websocket.receive_json() except WebSocketDisconnect: return if auth_msg.get("type") != "auth": await websocket.send_json({"type": "error", "message": "First message must be auth"}) await websocket.close(code=4001) return user_id_str: str = auth_msg.get("userId", "") or "" user_role: str = auth_msg.get("role", "") or "" tenant_id_str: str = auth_msg.get("tenantId", "") or "" if not user_id_str or not tenant_id_str: await websocket.send_json({"type": "error", "message": "Missing userId or tenantId in auth"}) await websocket.close(code=4001) return # Validate UUID format try: uuid.UUID(user_id_str) tenant_uuid = uuid.UUID(tenant_id_str) except (ValueError, AttributeError): await websocket.send_json({"type": "error", "message": "Invalid UUID format in auth"}) await websocket.close(code=4001) return logger.info( "WebSocket auth: user=%s role=%s tenant=%s conversation=%s", user_id_str, user_role, tenant_id_str, conversation_id, ) # ------------------------------------------------------------------------- # Step 2: Message loop # ------------------------------------------------------------------------- while True: try: msg_data = await websocket.receive_json() except (WebSocketDisconnect, Exception): break if msg_data.get("type") != "message": continue text_content: str = msg_data.get("text", "") or "" agent_id_str: str = msg_data.get("agentId", "") or "" msg_conversation_id: str = msg_data.get("conversationId", conversation_id) or conversation_id display_name: str = msg_data.get("displayName", "Portal User") # ------------------------------------------------------------------- # a. Send typing indicator IMMEDIATELY — before any DB or Celery work # ------------------------------------------------------------------- await websocket.send_json({"type": "typing"}) # ------------------------------------------------------------------- # b. Save user message to web_conversation_messages # ------------------------------------------------------------------- configure_rls_hook(engine) rls_token = current_tenant_id.set(tenant_uuid) saved_conversation_id = msg_conversation_id try: async with async_session_factory() as session: # Look up the conversation to get tenant-scoped context conv_stmt = select(WebConversation).where( WebConversation.id == uuid.UUID(msg_conversation_id) ) conv_result = await session.execute(conv_stmt) conversation = conv_result.scalar_one_or_none() if conversation is not None: # Save user message user_msg = WebConversationMessage( conversation_id=uuid.UUID(msg_conversation_id), tenant_id=tenant_uuid, role="user", content=text_content, ) session.add(user_msg) # Update conversation timestamp await session.execute( text( "UPDATE web_conversations SET updated_at = NOW() WHERE id = :conv_id" ), {"conv_id": str(msg_conversation_id)}, ) await session.commit() saved_conversation_id = msg_conversation_id except Exception: logger.exception( "Failed to save user message for conversation=%s", msg_conversation_id ) finally: current_tenant_id.reset(rls_token) # ------------------------------------------------------------------- # c. Build KonstructMessage and stream LLM response DIRECTLY # # Bypasses Celery entirely for web chat — calls the LLM pool's # streaming endpoint from the WebSocket handler. This eliminates # ~5-10s of Celery queue + Redis pub-sub round-trip overhead. # Slack/WhatsApp still use Celery (async webhook pattern). # ------------------------------------------------------------------- event = { "text": text_content, "tenant_id": tenant_id_str, "agent_id": agent_id_str, "user_id": user_id_str, "display_name": display_name, "conversation_id": saved_conversation_id, } normalized_msg = normalize_web_event(event) # Load agent for this tenant agent: Agent | None = None rls_token3 = current_tenant_id.set(tenant_uuid) try: async with async_session_factory() as session: from sqlalchemy import select as sa_select agent_stmt = sa_select(Agent).where( Agent.tenant_id == tenant_uuid, Agent.is_active == True, ).limit(1) agent_result = await session.execute(agent_stmt) agent = agent_result.scalar_one_or_none() finally: current_tenant_id.reset(rls_token3) if agent is None: await websocket.send_json({ "type": "done", "text": "No active AI employee is configured for this workspace.", "conversation_id": saved_conversation_id, }) continue # Build memory-enriched messages (Redis sliding window only — fast) redis_mem = aioredis.from_url(settings.redis_url) try: recent_messages = await get_recent_messages( redis_mem, tenant_id_str, str(agent.id), user_id_str ) finally: await redis_mem.aclose() enriched_messages = build_messages_with_memory( agent=agent, current_message=text_content, recent_messages=recent_messages, relevant_context=[], channel="web", ) # Stream LLM response directly to WebSocket — no Celery, no pub-sub response_text = "" try: async for token in run_agent_streaming( msg=normalized_msg, agent=agent, messages=enriched_messages, ): response_text += token try: await websocket.send_json({"type": "chunk", "text": token}) except Exception: break # Client disconnected except Exception: logger.exception("Direct streaming failed for conversation=%s", saved_conversation_id) if not response_text: response_text = "I encountered an error processing your message. Please try again." # Save to Redis sliding window (fire-and-forget, non-blocking) redis_mem2 = aioredis.from_url(settings.redis_url) try: await append_message(redis_mem2, tenant_id_str, str(agent.id), user_id_str, "user", text_content) if response_text: await append_message(redis_mem2, tenant_id_str, str(agent.id), user_id_str, "assistant", response_text) finally: await redis_mem2.aclose() # Fire-and-forget embedding for long-term memory try: embed_and_store.delay({ "tenant_id": tenant_id_str, "agent_id": str(agent.id), "user_id": user_id_str, "role": "user", "content": text_content, }) if response_text: embed_and_store.delay({ "tenant_id": tenant_id_str, "agent_id": str(agent.id), "user_id": user_id_str, "role": "assistant", "content": response_text, }) except Exception: pass # Non-fatal — memory will rebuild over time # ------------------------------------------------------------------- # e. Save assistant message and send final "done" to client # ------------------------------------------------------------------- if response_text: rls_token2 = current_tenant_id.set(tenant_uuid) try: async with async_session_factory() as session: assistant_msg = WebConversationMessage( conversation_id=uuid.UUID(saved_conversation_id), tenant_id=tenant_uuid, role="assistant", content=response_text, ) session.add(assistant_msg) await session.execute( text( "UPDATE web_conversations SET updated_at = NOW() WHERE id = :conv_id" ), {"conv_id": str(saved_conversation_id)}, ) await session.commit() except Exception: logger.exception( "Failed to save assistant message for conversation=%s", saved_conversation_id ) finally: current_tenant_id.reset(rls_token2) # Signal stream completion to the client await websocket.send_json({ "type": "done", "text": response_text, "conversation_id": saved_conversation_id, }) else: logger.warning( "No response received within %ds for conversation=%s", _RESPONSE_TIMEOUT_SECONDS, saved_conversation_id, ) await websocket.send_json({ "type": "error", "message": "Agent did not respond in time. Please try again.", }) @web_chat_router.websocket("/chat/ws/{conversation_id}") async def chat_websocket(websocket: WebSocket, conversation_id: str) -> None: """ WebSocket endpoint for web chat. URL: /chat/ws/{conversation_id} Protocol: 1. Connect 2. Send: {"type": "auth", "userId": "...", "role": "...", "tenantId": "..."} 3. Send: {"type": "message", "text": "...", "agentId": "...", "conversationId": "..."} 4. Receive: {"type": "typing"} 5. Receive: {"type": "response", "text": "...", "conversation_id": "..."} Closes with code 4001 on auth failure. """ try: await _handle_websocket_connection(websocket, conversation_id) except WebSocketDisconnect: logger.info("WebSocket disconnected for conversation=%s", conversation_id) except Exception: logger.exception("Unhandled error in WebSocket handler for conversation=%s", conversation_id)