WebSocket Protocol

Real-time communication protocol for Aurora chat interface

Overview

Aurora uses WebSockets for real-time bidirectional communication between the frontend and the chatbot service. This enables streaming responses, live status updates, and interactive tool execution.

Connection Details

Endpoint

ws://localhost:5006  # Development
wss://your-domain.com:5006  # Production

Connection Lifecycle

1. Client connects to WebSocket endpoint
2. Server sends initial START status
3. Client sends init message with user context
4. Bidirectional messaging begins
5. Server sends END status when workflow completes

Message Format

All messages are JSON objects with a consistent structure:

{
  "type": "message_type",
  "data": { /* type-specific data */ },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Client-to-Server Messages

Init Message

Initialize the WebSocket connection with user context.

Type: init

{
  "type": "init",
  "user_id": "user_abc123"
}

Purpose:

• Establishes user identity
• Triggers MCP (Model Context Protocol) preloading
• Starts API cost cache warming
• Initializes deployment listeners

Chat Query

Send a chat message to the agent.

Type: query

{
  "query": "Deploy my application to GKE",
  "user_id": "user_abc123",
  "session_id": "550e8400-e29b-41d4-a716-446655440000",
  "mode": "agent",
  "model": "anthropic/claude-3.5-sonnet",
  "provider_preference": ["gcp", "aws"],
  "selected_project_id": "my-gcp-project",
  "attachments": [
    {
      "filename": "config.yaml",
      "file_type": "text/yaml",
      "file_data": "base64EncodedContent==",
      "is_server_path": false
    }
  ],
  "ui_state": {
    "selectedModel": "anthropic/claude-3.5-sonnet",
    "selectedMode": "agent",
    "selectedProviders": ["gcp", "aws"]
  }
}

Fields:

• query: User’s message text
• user_id: Authenticated user ID
• session_id: Chat session ID (created by REST API)
• mode: “agent” (full execution) or “ask” (read-only)
• model: LLM model to use
• provider_preference: List of preferred cloud providers
• selected_project_id: (Optional) Specific cloud project
• attachments: (Optional) File attachments
• ui_state: UI configuration to persist

Control Messages

Control workflow execution.

Type: control

{
  "type": "control",
  "action": "cancel",
  "session_id": "550e8400-e29b-41d4-a716-446655440000",
  "user_id": "user_abc123"
}

Actions:

• cancel: Stop the current workflow execution

Cancellation Behavior:

• Cancels pending infrastructure confirmations
• Terminates running Celery tasks
• Consolidates message chunks
• Saves context for session resumption
• Sends END status to frontend

Confirmation Response

Respond to infrastructure change confirmations.

Type: confirmation_response

{
  "type": "confirmation_response",
  "confirmation_id": "conf_abc123",
  "approved": true,
  "user_id": "user_abc123",
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Purpose:

• Approve or reject infrastructure changes
• Refreshes WebSocket connection if user reconnected
• Allows workflow to proceed or abort

Direct Tool Call

Execute a specific tool without AI routing.

{
  "query": "Commit changes",
  "user_id": "user_abc123",
  "session_id": "550e8400-e29b-41d4-a716-446655440000",
  "direct_tool_call": {
    "tool_name": "github_commit",
    "parameters": {
      "repo": "my-org/my-repo",
      "commit_message": "Fix: Update deployment config",
      "branch": "main",
      "push": true
    }
  }
}

Purpose:

• Bypass AI decision-making for specific operations
• Direct execution of known actions (e.g., Git commits)
• Faster response for deterministic operations

Server-to-Client Messages

Status Messages

Indicate workflow state changes.

Type: status

{
  "type": "status",
  "data": {
    "status": "START"
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Status Values:

• START: Workflow beginning
• END: Workflow completed
• ERROR: Workflow failed

Message Chunks

Streaming text responses from the LLM.

Type: message

{
  "type": "message",
  "data": {
    "text": "I'll help you deploy to GKE. Let me check your clusters...",
    "is_chunk": true,
    "is_complete": false,
    "streaming": true
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Streaming Behavior:

• Tokens arrive as soon as generated by LLM
• Multiple chunks per response
• is_chunk: true indicates partial content
• Final chunk may have is_complete: true

Tool Call Messages

Indicate tool execution start.

Type: tool_call

{
  "type": "tool_call",
  "data": {
    "tool_name": "kubectl_get",
    "input": {
      "resource": "pods",
      "namespace": "default"
    },
    "status": "running",
    "timestamp": "2024-03-01T10:00:00Z",
    "tool_call_id": "call_abc123"
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Tool Output Messages

Streaming output from tool execution.

Type: tool_output

{
  "type": "tool_output",
  "data": {
    "tool_name": "kubectl_get",
    "output": "NAME                    READY   STATUS\napi-deployment-abc123   1/1     Running\n",
    "is_chunk": true,
    "tool_call_id": "call_abc123"
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Streaming Output:

• Tool output streams in real-time
• Large outputs split across multiple chunks
• Final chunk indicates completion

Tool Result Messages

Final result of tool execution.

Type: tool_result

{
  "type": "tool_result",
  "data": {
    "tool_name": "kubectl_get",
    "result": {
      "success": true,
      "pods": [
        {
          "name": "api-deployment-abc123",
          "status": "Running"
        }
      ]
    },
    "session_id": "550e8400-e29b-41d4-a716-446655440000"
  }
}

Confirmation Request

Request user approval for infrastructure changes.

Type: confirmation_request

{
  "type": "confirmation_request",
  "data": {
    "confirmation_id": "conf_abc123",
    "action": "create_gke_cluster",
    "details": {
      "cluster_name": "my-cluster",
      "region": "us-central1",
      "node_count": 3,
      "machine_type": "e2-medium"
    },
    "risk_level": "medium",
    "estimated_cost": "$150/month"
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Risk Levels:

• low: Safe operations (read-only, reversible)
• medium: Infrastructure changes (creation, updates)
• high: Destructive operations (deletion, data loss)

Usage Info

API cost tracking information.

Type: usage_info

{
  "type": "usage_info",
  "data": {
    "total_cost": 2.45
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Error Messages

Error information.

Type: error

{
  "type": "error",
  "data": {
    "text": "Failed to execute kubectl command: cluster not found",
    "code": "TOOL_EXECUTION_FAILED"
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Connection Management

Rate Limiting

WebSocket connections are rate-limited to prevent abuse:

Rate: 5 messages per 60 seconds per client

Exceeded Response:

{
  "type": "error",
  "data": {
    "text": "Rate limit exceeded. Please wait and try again."
  }
}

Connection Recovery

If the WebSocket connection drops:

1. Client reconnects to the same endpoint
2. Client sends init message
3. Client sends confirmation_response to refresh connection
4. Workflow continues in the background
5. Messages resume streaming to the new connection

Connection Refresh:

{
  "type": "confirmation_response",
  "user_id": "user_abc123",
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Background Execution

Workflows continue running even if the WebSocket disconnects:

• Messages saved to database
• Tool execution continues
• Context preserved for reconnection
• Results available via REST API

Timeout Handling

Workflows have a 30-minute timeout:

Timeout Message:

{
  "type": "message",
  "data": {
    "text": "Workflow timeout - the operation may have completed but the response took too long. Please check your resources manually."
  },
  "session_id": "550e8400-e29b-41d4-a716-446655440000"
}

Authentication

kubectl Agent WebSocket

Special authentication for kubectl agent connections.

Header:

Authorization: Bearer <token>

Purpose:

• Authenticate kubectl agent running in user clusters
• Enable bi-directional command execution
• Secured with long-lived bearer tokens

Code Examples

JavaScript/TypeScript Client

const ws = new WebSocket('ws://localhost:5006');

// Handle connection open
ws.onopen = () => {
  console.log('Connected to Aurora chatbot');
  
  // Send init message
  ws.send(JSON.stringify({
    type: 'init',
    user_id: 'user_abc123'
  }));
  
  // Send chat query
  ws.send(JSON.stringify({
    query: 'Deploy to GKE',
    user_id: 'user_abc123',
    session_id: 'session123',
    mode: 'agent'
  }));
};

// Handle incoming messages
ws.onmessage = (event) => {
  const message = JSON.parse(event.data);
  
  switch (message.type) {
    case 'status':
      console.log('Status:', message.data.status);
      break;
      
    case 'message':
      console.log('Message chunk:', message.data.text);
      break;
      
    case 'tool_call':
      console.log('Tool executing:', message.data.tool_name);
      break;
      
    case 'confirmation_request':
      // Show confirmation UI
      const approved = confirm(`Approve: ${message.data.action}?`);
      ws.send(JSON.stringify({
        type: 'confirmation_response',
        confirmation_id: message.data.confirmation_id,
        approved: approved,
        user_id: 'user_abc123',
        session_id: 'session123'
      }));
      break;
  }
};

// Handle errors
ws.onerror = (error) => {
  console.error('WebSocket error:', error);
};

// Handle connection close
ws.onclose = () => {
  console.log('Disconnected from Aurora chatbot');
  // Implement reconnection logic here
};

Python Client

import asyncio
import json
import websockets

async def chat_with_aurora():
    uri = "ws://localhost:5006"
    
    async with websockets.connect(uri) as websocket:
        # Send init message
        await websocket.send(json.dumps({
            "type": "init",
            "user_id": "user_abc123"
        }))
        
        # Send query
        await websocket.send(json.dumps({
            "query": "List my GKE clusters",
            "user_id": "user_abc123",
            "session_id": "session123",
            "mode": "agent"
        }))
        
        # Receive messages
        async for message in websocket:
            data = json.loads(message)
            
            if data["type"] == "status":
                print(f"Status: {data['data']['status']}")
                if data['data']['status'] == 'END':
                    break
                    
            elif data["type"] == "message":
                print(data['data']['text'], end='', flush=True)
                
            elif data["type"] == "tool_call":
                print(f"\nExecuting: {data['data']['tool_name']}")

asyncio.run(chat_with_aurora())

Best Practices

Connection Handling

1. Always send init message after connecting
2. Handle reconnections gracefully with exponential backoff
3. Store session_id to resume conversations
4. Monitor connection health with periodic pings

Message Processing

1. Buffer streaming chunks for smooth UI updates
2. Queue messages if UI can’t keep up
3. Handle out-of-order messages using timestamps
4. Validate message structure before processing

Error Recovery

1. Retry failed connections with backoff
2. Cache unsent messages for retry after reconnection
3. Fetch missed messages via REST API after reconnection
4. Display clear error messages to users

Performance

1. Use WebSocket compression for large messages
2. Batch multiple small messages when possible
3. Implement client-side throttling for rapid updates
4. Monitor memory usage with long-lived connections