feat(api): remove outdated planning document and implement streaming service for real-time task updates

planejamento_atualizado.md

@@ -1,191 +0,0 @@
-# Planejamento de Implementação - A2A Streaming (Atualizado)
-
-## 1. Visão Geral
-
-Implementar suporte a Server-Sent Events (SSE) para streaming de atualizações de tarefas em tempo real, seguindo a especificação oficial do A2A.
-
-## 2. Componentes Necessários
-
-### 2.2 Estrutura de Arquivos
-
-```
-src/
-├── api/
-│   └── agent_routes.py (modificação)
-├── schemas/
-│   └── streaming.py (novo)
-├── services/
-│   └── streaming_service.py (novo)
-└── utils/
-    └── streaming.py (novo)
-```
-
-## 3. Implementação
-
-### 3.1 Schemas (Pydantic)
-
-```python
-# schemas/streaming.py
-- TaskStatusUpdateEvent
-  - state: str (working, completed, failed)
-  - timestamp: datetime
-  - message: Optional[Message]
-  - error: Optional[Error]
-
-- TaskArtifactUpdateEvent
-  - type: str
-  - content: str
-  - metadata: Dict[str, Any]
-
-- JSONRPCRequest
-  - jsonrpc: str = "2.0"
-  - id: str
-  - method: str = "tasks/sendSubscribe"
-  - params: Dict[str, Any]
-
-- Message
-  - role: str
-  - parts: List[MessagePart]
-
-- MessagePart
-  - type: str
-  - text: str
-```
-
-### 3.2 Serviço de Streaming
-
-````python
-# services/streaming_service.py
-- send_task_streaming()
-  - Monta payload JSON-RPC conforme especificação:
-    ```json
-    {
-      "jsonrpc": "2.0",
-      "id": "<uuid>",
-      "method": "tasks/sendSubscribe",
-      "params": {
-        "id": "<uuid>",
-        "sessionId": "<opcional>",
-        "message": {
-          "role": "user",
-          "parts": [{"type": "text", "text": "<mensagem>"}]
-        }
-      }
-    }
-    ```
-  - Configura headers:
-    - Accept: text/event-stream
-    - Authorization: x-api-key <SUA_API_KEY>
-  - Gerencia conexão SSE
-  - Processa eventos em tempo real
-````
-
-### 3.3 Rota de Streaming
-
-```python
-# api/agent_routes.py
-- Nova rota POST /{agent_id}/tasks/sendSubscribe
-  - Validação de API key
-  - Gerenciamento de sessão
-  - Streaming de eventos SSE
-  - Tratamento de erros JSON-RPC
-```
-
-### 3.4 Utilitários
-
-```python
-# utils/streaming.py
-- Helpers para SSE
-  - Formatação de eventos
-  - Tratamento de reconexão
-  - Timeout e retry
-- Processamento de eventos
-  - Parsing de eventos SSE
-  - Validação de payloads
-- Formatação de respostas
-  - Conformidade com JSON-RPC 2.0
-```
-
-## 4. Fluxo de Dados
-
-1. Cliente envia requisição JSON-RPC para `/tasks/sendSubscribe`
-2. Servidor valida API key e configura sessão
-3. Inicia streaming de eventos SSE
-4. Envia atualizações em tempo real:
-   - TaskStatusUpdateEvent (estado da tarefa)
-   - TaskArtifactUpdateEvent (artefatos gerados)
-   - Mensagens do histórico
-
-## 5. Exemplo de Uso
-
-```python
-async def exemplo_uso():
-    agent_id = "uuid-do-agente"
-    api_key = "sua-api-key"
-    mensagem = "Olá, como posso ajudar?"
-
-    async with httpx.AsyncClient() as client:
-        # Configura headers
-        headers = {
-            "Accept": "text/event-stream",
-            "Authorization": f"x-api-key {api_key}"
-        }
-
-        # Monta payload JSON-RPC
-        payload = {
-            "jsonrpc": "2.0",
-            "id": str(uuid.uuid4()),
-            "method": "tasks/sendSubscribe",
-            "params": {
-                "id": str(uuid.uuid4()),
-                "message": {
-                    "role": "user",
-                    "parts": [{"type": "text", "text": mensagem}]
-                }
-            }
-        }
-
-        # Inicia streaming
-        async with connect_sse(client, "POST", f"/agents/{agent_id}/tasks/sendSubscribe",
-                             json=payload, headers=headers) as event_source:
-            async for event in event_source.aiter_sse():
-                if event.event == "message":
-                    data = json.loads(event.data)
-                    print(f"Evento recebido: {data}")
-```
-
-## 6. Considerações de Segurança
-
-- Validação rigorosa de API keys
-- Timeout de conexão SSE (30 segundos)
-- Tratamento de erros e reconexão automática
-- Limites de taxa (rate limiting)
-- Validação de payloads JSON-RPC
-- Sanitização de inputs
-
-## 7. Testes
-
-- Testes unitários para schemas
-- Testes de integração para streaming
-- Testes de carga e performance
-- Testes de reconexão e resiliência
-- Testes de conformidade JSON-RPC
-
-## 8. Documentação
-
-- Atualizar documentação da API
-- Adicionar exemplos de uso
-- Documentar formatos de eventos
-- Guia de troubleshooting
-- Referência à especificação A2A
-
-## 9. Próximos Passos
-
-1. Implementar schemas Pydantic conforme especificação
-2. Desenvolver serviço de streaming com suporte a JSON-RPC
-3. Adicionar rota SSE com validação de payloads
-4. Implementar utilitários de streaming
-5. Escrever testes de conformidade
-6. Atualizar documentação
-7. Revisão de código
-8. Deploy em ambiente de teste

src/api/agent_routes.py

@@ -24,6 +24,9 @@ from src.services.service_providers import (
     memory_service,
 )
 import logging
+from fastapi.responses import StreamingResponse
+from ..services.streaming_service import StreamingService
+from ..schemas.streaming import JSONRPCRequest
 
 from src.services.session_service import get_session_events
 
@@ -74,6 +77,8 @@ router = APIRouter(
     responses={404: {"description": "Not found"}},
 )
 
+streaming_service = StreamingService()
+
 
 @router.post("/", response_model=Agent, status_code=status.HTTP_201_CREATED)
 async def create_agent(
@@ -192,8 +197,8 @@ async def get_agent_json(
             },
             "version": os.getenv("API_VERSION", ""),
             "capabilities": {
-                "streaming": False,
+                "streaming": True,
-                "pushNotifications": False,
+                "pushNotifications": True,
                 "stateTransitionHistory": True,
             },
             "authentication": {
@@ -355,6 +360,11 @@ async def handle_task(
 
         except Exception as e:
             logger.error(f"Error processing history: {str(e)}")
+
+        # pushNotification = task_request.get("pushNotification", False)
+        # if pushNotification:
+        #     await send_push_notification(task_id, final_response_text)
+
         return response_task
 
     except HTTPException:
@@ -365,3 +375,51 @@ async def handle_task(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="Internal server error",
         )
+
+
+@router.post("/{agent_id}/tasks/sendSubscribe")
+async def subscribe_task_streaming(
+    agent_id: str,
+    request: JSONRPCRequest,
+    x_api_key: str = Header(None),
+    db: Session = Depends(get_db),
+):
+    """
+    Endpoint para streaming de eventos SSE de uma tarefa.
+
+    Args:
+        agent_id: ID do agente
+        request: Requisição JSON-RPC
+        x_api_key: Chave de API no header
+        db: Sessão do banco de dados
+
+    Returns:
+        StreamingResponse com eventos SSE
+    """
+    if not x_api_key:
+        raise HTTPException(status_code=401, detail="API key é obrigatória")
+
+    # Extrai mensagem do payload
+    message = request.params.get("message", {}).get("parts", [{}])[0].get("text", "")
+    session_id = request.params.get("sessionId")
+
+    # Configura streaming
+    async def event_generator():
+        async for event in streaming_service.send_task_streaming(
+            agent_id=agent_id,
+            api_key=x_api_key,
+            message=message,
+            session_id=session_id,
+            db=db,
+        ):
+            yield event
+
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )

src/schemas/streaming.py

@@ -0,0 +1,40 @@
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Literal
+from pydantic import BaseModel, Field
+
+
+class MessagePart(BaseModel):
+    type: str
+    text: str
+
+
+class Message(BaseModel):
+    role: str
+    parts: List[MessagePart]
+
+
+class TaskStatusUpdateEvent(BaseModel):
+    state: str = Field(..., description="Estado da tarefa (working, completed, failed)")
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    message: Optional[Message] = None
+    error: Optional[Dict[str, Any]] = None
+
+
+class TaskArtifactUpdateEvent(BaseModel):
+    type: str
+    content: str
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class JSONRPCRequest(BaseModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+    id: str
+    method: Literal["tasks/sendSubscribe"] = "tasks/sendSubscribe"
+    params: Dict[str, Any]
+
+
+class JSONRPCResponse(BaseModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+    id: str
+    result: Optional[Dict[str, Any]] = None
+    error: Optional[Dict[str, Any]] = None

src/services/streaming_service.py

@@ -0,0 +1,140 @@
+import uuid
+import json
+from typing import AsyncGenerator, Dict, Any
+from fastapi import HTTPException
+from datetime import datetime
+from ..schemas.streaming import (
+    JSONRPCRequest,
+    TaskStatusUpdateEvent,
+    TaskArtifactUpdateEvent,
+)
+from ..services.agent_runner import run_agent
+from ..services.service_providers import (
+    session_service,
+    artifacts_service,
+    memory_service,
+)
+from sqlalchemy.orm import Session
+
+
+class StreamingService:
+    def __init__(self):
+        self.active_connections: Dict[str, Any] = {}
+
+    async def send_task_streaming(
+        self,
+        agent_id: str,
+        api_key: str,
+        message: str,
+        session_id: str = None,
+        db: Session = None,
+    ) -> AsyncGenerator[str, None]:
+        """
+        Inicia o streaming de eventos SSE para uma tarefa.
+
+        Args:
+            agent_id: ID do agente
+            api_key: Chave de API para autenticação
+            message: Mensagem inicial
+            session_id: ID da sessão (opcional)
+            db: Sessão do banco de dados
+
+        Yields:
+            Eventos SSE formatados
+        """
+        # Validação básica da API key
+        if not api_key:
+            raise HTTPException(status_code=401, detail="API key é obrigatória")
+
+        # Gera IDs únicos
+        task_id = str(uuid.uuid4())
+        request_id = str(uuid.uuid4())
+
+        # Monta payload JSON-RPC
+        payload = JSONRPCRequest(
+            id=request_id,
+            params={
+                "id": task_id,
+                "sessionId": session_id,
+                "message": {
+                    "role": "user",
+                    "parts": [{"type": "text", "text": message}],
+                },
+            },
+        )
+
+        # Registra conexão
+        self.active_connections[task_id] = {
+            "agent_id": agent_id,
+            "api_key": api_key,
+            "session_id": session_id,
+        }
+
+        try:
+            # Envia evento de início
+            yield self._format_sse_event(
+                "status",
+                TaskStatusUpdateEvent(
+                    state="working",
+                    timestamp=datetime.now().isoformat(),
+                    message=payload.params["message"],
+                ).model_dump_json(),
+            )
+
+            # Executa o agente
+            result = await run_agent(
+                str(agent_id),
+                task_id,
+                message,
+                session_service,
+                artifacts_service,
+                memory_service,
+                db,
+                session_id,
+            )
+
+            # Envia a resposta do agente como um evento separado
+            yield self._format_sse_event(
+                "message",
+                json.dumps(
+                    {
+                        "role": "agent",
+                        "content": result,
+                        "timestamp": datetime.now().isoformat(),
+                    }
+                ),
+            )
+
+            # Evento de conclusão
+            yield self._format_sse_event(
+                "status",
+                TaskStatusUpdateEvent(
+                    state="completed",
+                    timestamp=datetime.now().isoformat(),
+                ).model_dump_json(),
+            )
+
+        except Exception as e:
+            # Evento de erro
+            yield self._format_sse_event(
+                "status",
+                TaskStatusUpdateEvent(
+                    state="failed",
+                    timestamp=datetime.now().isoformat(),
+                    error={"message": str(e)},
+                ).model_dump_json(),
+            )
+            raise
+
+        finally:
+            # Limpa conexão
+            self.active_connections.pop(task_id, None)
+
+    def _format_sse_event(self, event_type: str, data: str) -> str:
+        """Formata um evento SSE."""
+        return f"event: {event_type}\ndata: {data}\n\n"
+
+    async def close_connection(self, task_id: str):
+        """Fecha uma conexão de streaming."""
+        if task_id in self.active_connections:
+            self.active_connections.pop(task_id)

src/utils/streaming.py

@@ -0,0 +1,70 @@
+import asyncio
+from typing import AsyncGenerator, Optional
+from fastapi import HTTPException
+
+
+class SSEUtils:
+    @staticmethod
+    async def with_timeout(
+        generator: AsyncGenerator, timeout: int = 30, retry_attempts: int = 3
+    ) -> AsyncGenerator:
+        """
+        Adiciona timeout e retry a um gerador de eventos SSE.
+
+        Args:
+            generator: Gerador de eventos
+            timeout: Tempo máximo de espera em segundos
+            retry_attempts: Número de tentativas de reconexão
+
+        Yields:
+            Eventos do gerador
+        """
+        attempts = 0
+        while attempts < retry_attempts:
+            try:
+                async for event in asyncio.wait_for(generator, timeout):
+                    yield event
+                break
+            except asyncio.TimeoutError:
+                attempts += 1
+                if attempts >= retry_attempts:
+                    raise HTTPException(
+                        status_code=408, detail="Timeout após múltiplas tentativas"
+                    )
+                await asyncio.sleep(1)  # Espera antes de tentar novamente
+
+    @staticmethod
+    def format_error_event(error: Exception) -> str:
+        """
+        Formata um evento de erro SSE.
+
+        Args:
+            error: Exceção ocorrida
+
+        Returns:
+            String formatada do evento SSE
+        """
+        return f"event: error\ndata: {str(error)}\n\n"
+
+    @staticmethod
+    def validate_sse_headers(headers: dict) -> None:
+        """
+        Valida headers necessários para SSE.
+
+        Args:
+            headers: Dicionário de headers
+
+        Raises:
+            HTTPException se headers inválidos
+        """
+        required_headers = {
+            "Accept": "text/event-stream",
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        }
+
+        for header, value in required_headers.items():
+            if headers.get(header) != value:
+                raise HTTPException(
+                    status_code=400, detail=f"Header {header} inválido ou ausente"
+                )

static/test_a2a_stream.html

@@ -0,0 +1,295 @@
+<!DOCTYPE html>
+<html lang="pt-BR">
+
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Teste de Streaming A2A</title>
+    <link rel="icon" href="data:,">
+    <style>
+        :root {
+            --primary-color: #2563eb;
+            --secondary-color: #1e40af;
+            --background-color: #f8fafc;
+            --text-color: #1e293b;
+            --border-color: #e2e8f0;
+        }
+
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+
+        body {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            background-color: var(--background-color);
+            color: var(--text-color);
+            line-height: 1.6;
+            padding: 20px;
+        }
+
+        .container {
+            max-width: 800px;
+            margin: 0 auto;
+            background: white;
+            border-radius: 8px;
+            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
+            padding: 20px;
+        }
+
+        h1 {
+            color: var(--primary-color);
+            margin-bottom: 20px;
+            text-align: center;
+        }
+
+        .form-group {
+            margin-bottom: 20px;
+        }
+
+        label {
+            display: block;
+            margin-bottom: 5px;
+            font-weight: 600;
+        }
+
+        input[type="text"] {
+            width: 100%;
+            padding: 10px;
+            border: 1px solid var(--border-color);
+            border-radius: 4px;
+            font-size: 16px;
+        }
+
+        button {
+            background-color: var(--primary-color);
+            color: white;
+            border: none;
+            padding: 10px 20px;
+            border-radius: 4px;
+            cursor: pointer;
+            font-size: 16px;
+            transition: background-color 0.3s;
+        }
+
+        button:hover {
+            background-color: var(--secondary-color);
+        }
+
+        .chat-container {
+            margin-top: 20px;
+            border: 1px solid var(--border-color);
+            border-radius: 4px;
+            height: 400px;
+            overflow-y: auto;
+            padding: 10px;
+        }
+
+        .message {
+            margin-bottom: 10px;
+            padding: 10px;
+            border-radius: 4px;
+            max-width: 80%;
+        }
+
+        .user-message {
+            background-color: #e3f2fd;
+            margin-left: auto;
+        }
+
+        .agent-message {
+            background-color: #f3f4f6;
+        }
+
+        .status-message {
+            background-color: #fef3c7;
+            text-align: center;
+            font-style: italic;
+        }
+
+        .error-message {
+            background-color: #fee2e2;
+            color: #dc2626;
+        }
+
+        .controls {
+            display: flex;
+            gap: 10px;
+            margin-top: 20px;
+        }
+
+        .status {
+            margin-top: 10px;
+            padding: 10px;
+            border-radius: 4px;
+            background-color: #f3f4f6;
+            font-size: 14px;
+        }
+    </style>
+</head>
+
+<body>
+    <div class="container">
+        <h1>Teste de Streaming A2A</h1>
+
+        <div class="form-group">
+            <label for="agentId">Agent ID:</label>
+            <input type="text" id="agentId" placeholder="Digite o ID do agente">
+        </div>
+
+        <div class="form-group">
+            <label for="apiKey">API Key:</label>
+            <input type="text" id="apiKey" placeholder="Digite sua API Key">
+        </div>
+
+        <div class="form-group">
+            <label for="message">Mensagem:</label>
+            <input type="text" id="message" placeholder="Digite sua mensagem">
+        </div>
+
+        <div class="controls">
+            <button onclick="startStreaming()">Iniciar Streaming</button>
+            <button onclick="stopStreaming()">Parar Streaming</button>
+        </div>
+
+        <div class="status" id="connectionStatus">
+            Status: Não conectado
+        </div>
+
+        <div class="chat-container" id="chatContainer">
+            <!-- Mensagens serão adicionadas aqui -->
+        </div>
+    </div>
+
+    <script>
+        let controller = null;
+        const chatContainer = document.getElementById('chatContainer');
+        const statusElement = document.getElementById('connectionStatus');
+
+        function addMessage(content, type = 'agent') {
+            const messageDiv = document.createElement('div');
+            messageDiv.className = `message ${type}-message`;
+            messageDiv.textContent = content;
+            chatContainer.appendChild(messageDiv);
+            chatContainer.scrollTop = chatContainer.scrollHeight;
+        }
+
+        function updateStatus(status) {
+            statusElement.textContent = `Status: ${status}`;
+        }
+
+        async function startStreaming() {
+            const agentId = document.getElementById('agentId').value;
+            const apiKey = document.getElementById('apiKey').value;
+            const message = document.getElementById('message').value;
+
+            if (!agentId || !apiKey || !message) {
+                alert('Por favor, preencha todos os campos');
+                return;
+            }
+
+            // Limpa o chat
+            chatContainer.innerHTML = '';
+            addMessage('Iniciando conexão...', 'status');
+
+            // Configura o payload
+            const payload = {
+                jsonrpc: "2.0",
+                id: "1",
+                method: "tasks/sendSubscribe",
+                params: {
+                    id: "test-task",
+                    message: {
+                        role: "user",
+                        parts: [{
+                            type: "text",
+                            text: message
+                        }]
+                    }
+                }
+            };
+
+            try {
+                // Cria um novo AbortController para controlar o streaming
+                controller = new AbortController();
+                const signal = controller.signal;
+
+                // Faz a requisição POST com streaming
+                const response = await fetch(
+                    `http://${window.location.host}/api/v1/agents/${agentId}/tasks/sendSubscribe`, {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json',
+                            'Accept': 'text/event-stream',
+                            'X-API-Key': apiKey
+                        },
+                        body: JSON.stringify(payload),
+                        signal
+                    });
+
+                if (!response.ok) {
+                    throw new Error(`HTTP error! status: ${response.status}`);
+                }
+
+                updateStatus('Conectado');
+                addMessage('Conexão estabelecida', 'status');
+
+                // Lê o stream de dados
+                const reader = response.body.getReader();
+                const decoder = new TextDecoder();
+
+                while (true) {
+                    const {
+                        done,
+                        value
+                    } = await reader.read();
+                    if (done) break;
+
+                    // Decodifica o chunk de dados
+                    const chunk = decoder.decode(value);
+
+                    // Processa cada linha do chunk
+                    const lines = chunk.split('\n');
+                    for (const line of lines) {
+                        if (line.startsWith('data:')) {
+                            try {
+                                const data = JSON.parse(line.slice(5));
+
+                                if (data.state) {
+                                    // Evento de status
+                                    addMessage(`Estado: ${data.state}`, 'status');
+                                } else if (data.content) {
+                                    // Evento de conteúdo
+                                    addMessage(data.content, 'agent');
+                                }
+                            } catch (error) {
+                                addMessage('Erro ao processar mensagem: ' + error, 'error');
+                            }
+                        }
+                    }
+                }
+
+            } catch (error) {
+                if (error.name === 'AbortError') {
+                    updateStatus('Desconectado');
+                    addMessage('Conexão encerrada pelo usuário', 'status');
+                } else {
+                    updateStatus('Erro');
+                    addMessage('Erro ao iniciar streaming: ' + error.message, 'error');
+                }
+            }
+        }
+
+        function stopStreaming() {
+            if (controller) {
+                controller.abort();
+                controller = null;
+                updateStatus('Desconectado');
+                addMessage('Conexão encerrada', 'status');
+            }
+        }
+    </script>
+</body>
+
+</html>

static/test_chat_stream.html

@@ -3,40 +3,199 @@
 
 <head>
     <title>ADK Streaming Test</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <style>
+        :root {
+            --primary-color: #2563eb;
+            --secondary-color: #f3f4f6;
+            --text-color: #1f2937;
+            --border-color: #e5e7eb;
+            --success-color: #10b981;
+            --error-color: #ef4444;
+            --user-message-bg: #dbeafe;
+            --agent-message-bg: #f3f4f6;
+        }
+
+        * {
+            box-sizing: border-box;
+            margin: 0;
+            padding: 0;
+        }
+
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
+            line-height: 1.6;
+            color: var(--text-color);
+            background-color: #f9fafb;
+            padding: 20px;
+            max-width: 1200px;
+            margin: 0 auto;
+        }
+
+        h1 {
+            color: var(--primary-color);
+            margin-bottom: 20px;
+            text-align: center;
+            font-size: 2rem;
+        }
+
         #messages {
-            height: 400px;
+            height: 60vh;
             overflow-y: auto;
-            border: 1px solid #ccc;
+            border: 1px solid var(--border-color);
-            padding: 10px;
+            border-radius: 8px;
-            margin-bottom: 10px;
+            padding: 20px;
+            margin-bottom: 20px;
+            background-color: white;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
         }
 
         .message {
-            margin: 5px 0;
+            margin: 10px 0;
-            padding: 5px;
+            padding: 12px 16px;
-            border-radius: 5px;
+            border-radius: 12px;
+            max-width: 80%;
+            word-wrap: break-word;
+            animation: fadeIn 0.3s ease-in-out;
+        }
+
+        @keyframes fadeIn {
+            from {
+                opacity: 0;
+                transform: translateY(10px);
+            }
+            to {
+                opacity: 1;
+                transform: translateY(0);
+            }
         }
 
         .user-message {
-            background-color: #e3f2fd;
+            background-color: var(--user-message-bg);
-            margin-left: 20%;
+            margin-left: auto;
-            margin-right: 5px;
+            border-bottom-right-radius: 4px;
         }
 
         .agent-message {
-            background-color: #f5f5f5;
+            background-color: var(--agent-message-bg);
-            margin-right: 20%;
+            margin-right: auto;
-            margin-left: 5px;
+            border-bottom-left-radius: 4px;
+        }
+
+        #connection-status {
+            padding: 10px;
+            border-radius: 6px;
+            text-align: center;
+            font-weight: 500;
+            margin-bottom: 15px;
+        }
+
+        .connected {
+            background-color: rgba(16, 185, 129, 0.1);
+            color: var(--success-color);
+        }
+
+        .disconnected {
+            background-color: rgba(239, 68, 68, 0.1);
+            color: var(--error-color);
+        }
+
+        #messageForm {
+            background-color: white;
+            padding: 20px;
+            border-radius: 8px;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+        }
+
+        input[type="text"] {
+            width: 100%;
+            padding: 12px;
+            margin-bottom: 15px;
+            border: 1px solid var(--border-color);
+            border-radius: 6px;
+            font-size: 1rem;
+            transition: border-color 0.3s;
+        }
+
+        input[type="text"]:focus {
+            outline: none;
+            border-color: var(--primary-color);
+            box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1);
+        }
+
+        button {
+            padding: 12px 24px;
+            border: none;
+            border-radius: 6px;
+            font-size: 1rem;
+            font-weight: 500;
+            cursor: pointer;
+            transition: all 0.3s;
+        }
+
+        #connectButton {
+            background-color: var(--primary-color);
+            color: white;
+            width: 100%;
+        }
+
+        #connectButton:hover {
+            background-color: #1d4ed8;
+        }
+
+        #connectButton:disabled {
+            background-color: #93c5fd;
+            cursor: not-allowed;
+        }
+
+        #sendButton {
+            background-color: var(--primary-color);
+            color: white;
+            margin-left: 10px;
+        }
+
+        #sendButton:hover {
+            background-color: #1d4ed8;
+        }
+
+        #sendButton:disabled {
+            background-color: #93c5fd;
+            cursor: not-allowed;
         }
 
         #debug {
             margin-top: 20px;
+            padding: 15px;
+            background-color: #1f2937;
+            color: #e5e7eb;
+            border-radius: 8px;
+            font-family: 'Courier New', Courier, monospace;
+            font-size: 0.9rem;
+            line-height: 1.5;
+            max-height: 200px;
+            overflow-y: auto;
+        }
+
+        <blade media|%20(max-width%3A%20768px)%20%7B>body {
             padding: 10px;
-            background-color: #f8f9fa;
+        }
-            border: 1px solid #ddd;
+
-            font-family: monospace;
+        #messages {
-            white-space: pre-wrap;
+            height: 50vh;
+        }
+
+        .message {
+            max-width: 90%;
+        }
+
+        button {
+            width: 100%;
+            margin-bottom: 10px;
+        }
+
+        #sendButton {
+            margin-left: 0;
+        }
         }
     </style>
 </head>
@@ -44,14 +203,16 @@
 <body>
     <h1>ADK Streaming Test</h1>
     <div id="messages"></div>
-    <div id="connection-status" style="margin-bottom: 10px;"></div>
+    <div id="connection-status" class="disconnected">Desconectado</div>
     <form id="messageForm">
-        <input type="text" id="agentId" placeholder="Agent ID" style="margin-bottom: 10px; width: 300px;"><br>
+        <input type="text" id="agentId" placeholder="Agent ID">
-        <input type="text" id="contactId" placeholder="Contact ID" style="margin-bottom: 10px; width: 300px;"><br>
+        <input type="text" id="contactId" placeholder="Contact ID">
-        <input type="text" id="token" placeholder="JWT Token" style="margin-bottom: 10px; width: 300px;"><br>
+        <input type="text" id="token" placeholder="JWT Token">
-        <button type="button" id="connectButton">Conectar</button><br><br>
+        <button type="button" id="connectButton">Conectar</button>
-        <input type="text" id="message" placeholder="Digite sua mensagem..." style="width: 300px;">
+        <div style="display: flex; margin-top: 15px;">
-        <button type="submit" id="sendButton" disabled>Enviar</button>
+            <input type="text" id="message" placeholder="Digite sua mensagem...">
+            <button type="submit" id="sendButton" disabled>Enviar</button>
+        </div>
     </form>
 
     <div id="debug"></div>
@@ -86,12 +247,10 @@
                 return;
             }
 
-            // Fechar conexão existente se houver
             if (ws) {
                 ws.close();
             }
 
-            // Criar nova conexão WebSocket
             const ws_url = `ws://${window.location.host}/api/v1/chat/ws/${agentId}/${contactId}`;
             log('Connecting to WebSocket', {
                 url: ws_url
@@ -101,7 +260,6 @@
 
             ws.onopen = function () {
                 log('WebSocket connected, sending authentication');
-                // Adicionar token no header
                 const authMessage = {
                     type: 'authorization',
                     token: token
@@ -110,7 +268,7 @@
                 log('Authentication sent', authMessage);
 
                 statusDiv.textContent = 'Conectado';
-                statusDiv.style.color = 'green';
+                statusDiv.className = 'connected';
                 sendButton.disabled = false;
                 connectButton.disabled = true;
             };
@@ -144,7 +302,7 @@
                     wasClean: event.wasClean
                 });
                 statusDiv.textContent = 'Desconectado';
-                statusDiv.style.color = 'red';
+                statusDiv.className = 'disconnected';
                 sendButton.disabled = true;
                 connectButton.disabled = false;
                 ws = null;
@@ -153,7 +311,7 @@
             ws.onerror = function (error) {
                 log('WebSocket error', error);
                 statusDiv.textContent = 'Erro na conexão';
-                statusDiv.style.color = 'red';
+                statusDiv.className = 'disconnected';
             };
         };
 
@@ -163,13 +321,11 @@
             const message = messageInput.value;
 
             if (message && ws) {
-                // Criar div para mensagem do usuário
                 const userMessageDiv = document.createElement('div');
                 userMessageDiv.className = 'message user-message';
                 userMessageDiv.textContent = message;
                 messagesDiv.appendChild(userMessageDiv);
 
-                // Enviar mensagem via WebSocket
                 const messagePacket = {
                     message: message
                 };