WebSocket Protocol

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

Client connects to WebSocket endpoint
Server sends initial START status
Client sends init message with user context
Bidirectional messaging begins
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:

Client reconnects to the same endpoint
Client sends init message
Client sends confirmation_response to refresh connection
Workflow continues in the background
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

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

Message Processing

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

Error Recovery

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

Performance

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

Get Started

Core Features

Architecture

Deployment

Configuration

Integrations

Cloud Providers

Observability

Development

Guides

Reference

Help

Documentation Index

​Overview

​Connection Details

​Endpoint

​Connection Lifecycle

​Message Format

​Client-to-Server Messages

​Init Message

​Chat Query

​Control Messages

​Confirmation Response

​Direct Tool Call

​Server-to-Client Messages

​Status Messages

​Message Chunks

​Tool Call Messages

​Tool Output Messages

​Tool Result Messages

​Confirmation Request

​Usage Info

​Error Messages

​Connection Management

​Rate Limiting

​Connection Recovery

​Background Execution

​Timeout Handling

​Authentication

​kubectl Agent WebSocket

​Code Examples

​JavaScript/TypeScript Client

​Python Client

​Best Practices

​Connection Handling

​Message Processing

​Error Recovery

​Performance

Build docs developers (and LLMs) love

Overview

Connection Details

Endpoint

Connection Lifecycle

Message Format

Client-to-Server Messages

Init Message

Chat Query

Control Messages

Confirmation Response

Direct Tool Call

Server-to-Client Messages

Status Messages

Message Chunks

Tool Call Messages

Tool Output Messages

Tool Result Messages

Confirmation Request

Usage Info

Error Messages

Connection Management

Rate Limiting

Connection Recovery

Background Execution

Timeout Handling

Authentication

kubectl Agent WebSocket

Code Examples

JavaScript/TypeScript Client

Python Client

Best Practices

Connection Handling

Message Processing

Error Recovery

Performance