WebSocket Protocol

Quick Start

Start a TnsAI server with just a few lines of code. The server opens a WebSocket endpoint that clients can connect to for real-time agent interaction.

TnsServer server = TnsServer.builder()
    .port(7777)
    .llmClientSupplier(() -> LLMClientFactory.create("openai", "gpt-4o", 0.7f))
    .idleTimeout(Duration.ofMinutes(30))
    .build();
server.start();
// WebSocket endpoint: ws://localhost:7777/chat

The builder requires a Supplier<LLMClient> -- the factory is called once per agent creation. Default port is 7777. An idle shutdown manager auto-terminates the server after 30 minutes of inactivity.

Connection Lifecycle

Understanding the connection lifecycle helps you build reliable clients. Here is what happens from the moment a client connects to when it disconnects.

1. Client opens a WebSocket connection to /chat
2. WsHandler.onConnect fires, registers the connection with IdleShutdownManager
3. Client sends a chat event with a sessionId -- the server creates a SessionInfo (with a default assistant agent) on first access
4. The connection is registered to that session in WsConnectionManager -- multiple connections can share one session
5. On disconnect, onClose removes the connection and notifies the idle manager

Every message must carry the protocol version (v: 1) and a sessionId. The server rejects messages with mismatched versions (PROTOCOL_MISMATCH error).

Protocol v1 Events

Client to Server

These are the events that a client can send to the server. Every event must include v: 1 (protocol version) and a sessionId.

Server to Client

These are the events that the server sends to connected clients. Your client should handle at least token, done, and error for basic functionality.

Session Management

Sessions are the server's way of keeping track of ongoing conversations. Each session has its own set of agents, an audit trail, and an autonomy level. The SessionManager creates a session automatically the first time a client sends a message with a given sessionId.

// SessionInfo holds:
// - AgentGroup (one or more agents)
// - AuditService (per-session audit trail)
// - AutonomyLevel (default: SUPERVISED)
// - createdAt / lastActivity timestamps

SessionInfo session = sessionManager.getOrCreate("my-session-id");
session.getGroup().getMembers();  // List<Agent>
session.getAutonomyLevel();       // SUPERVISED
session.getAuditService();        // Per-session audit trail

Each session has its own RagService, lazily created on first access via sessionManager.getRag(sessionId).

Multi-Agent Support

A session can contain multiple agents with different roles and LLM configurations. You can add and remove agents dynamically at runtime using WebSocket events, which is useful for building team-based workflows where different agents handle different parts of a task.

// Add an agent with specific LLM configuration
{
  "v": 1,
  "type": "agent_add",
  "sessionId": "sess-1",
  "role": "reviewer",
  "provider": "anthropic",
  "model": "claude-sonnet-4-20250514",
  "name": "Code Reviewer",
  "goal": "Review code for security issues"
}

Agent IDs are generated as {name}-{8hexchars} (e.g., Code Reviewer-abcd1234). After each add/remove, the server broadcasts a team_update event with the full agent roster, including each agent's tools, roles, and LLM configuration.

Multi-role agents can be created by providing a roles array instead of a single goal/domain:

{
  "v": 1,
  "type": "agent_add",
  "sessionId": "sess-1",
  "role": "fullstack",
  "roles": [
    {"name": "frontend", "goal": "Build React UIs", "domain": "web"},
    {"name": "backend", "goal": "Build REST APIs", "domain": "server"}
  ]
}

MCP Tool Injection

If your frontend application connects to MCP servers (like a filesystem or database server), you can register those tools with the TnsAI server session. The server wraps each tool as a McpProxyTool and injects them into all agents, enabling a proxy pattern where the server orchestrates and the client executes.

// Client sends mcp_tools event
{
  "v": 1,
  "type": "mcp_tools",
  "sessionId": "sess-1",
  "tools": [
    {
      "server": "filesystem",
      "name": "read_file",
      "description": "Read a file",
      "inputSchema": {"type": "object", "properties": {"path": {"type": "string"}}}
    }
  ]
}

When the LLM invokes an MCP tool, the server sends an mcp_call event to the client. The client executes the tool via its local MCP connection and returns the result via mcp_result:

Server -> Client:  mcp_call {toolCallId: "abc", server: "filesystem", tool: "read_file", args: {path: "/src/App.tsx"}}
Client -> Server:  mcp_result {toolCallId: "abc", result: "import React from 'react';..."}

The McpProxyTool uses a CompletableFuture to block the agent's virtual thread until the client responds.

Knowledge Base Events

Each session has an optional knowledge base for Retrieval-Augmented Generation (RAG). The kb_* events let clients add, search, list, and remove documents in real time, which the agent can then use as context when answering questions.

// Add a document
{"v": 1, "type": "kb_add", "sessionId": "s1",
 "content": "React hooks must follow the rules of hooks...",
 "metadata": {"source": "docs", "topic": "react"}}

// Search
{"v": 1, "type": "kb_search", "sessionId": "s1",
 "query": "rules of hooks", "limit": 5, "threshold": 0.3}

// List all documents
{"v": 1, "type": "kb_list", "sessionId": "s1"}

// Remove a document
{"v": 1, "type": "kb_remove", "sessionId": "s1", "documentId": "doc-abc123"}

Search results include documentId, content, score, and metadata for each hit. The kb_list_result event returns document previews (first 100 chars) and content lengths.

WsHandler Internals

This section covers the internal workings of the WebSocket handler for contributors and advanced users. The WsHandler processes all WebSocket events on the Javalin thread, then dispatches chat work to a virtual-thread executor for non-blocking concurrency.

• Chat routing: Messages are routed to the first agent in the group. The agent runs streamChatWithTools, which handles the full LLM-tool-LLM loop.
• RAG injection: Before sending a message to the LLM, the handler checks if the session has indexed content. If so, it calls rag.buildContextPrompt(message, 5) to prepend relevant code context.
• Cancellation: cancel events interrupt the running virtual thread and cancel all pending approval futures.
• Keepalive: Plain "ping" text frames are silently ignored.

REST Endpoints

In addition to the WebSocket endpoint, the server exposes REST APIs for health checking, code search indexing, and audit trail retrieval. These are useful for monitoring, CI/CD integration, and administrative tasks.

Connection Example

Here is a complete JavaScript example showing how to connect to the server, send a chat message, handle streamed tokens, and respond to tool approval requests.

const ws = new WebSocket("ws://localhost:7777/chat");

ws.onopen = () => {
  // Send a chat message
  ws.send(JSON.stringify({
    v: 1,
    type: "chat",
    sessionId: "my-session",
    message: "Explain the WebSocket protocol"
  }));
};

ws.onmessage = (event) => {
  const data = JSON.parse(event.data);
  switch (data.type) {
    case "token":
      process.stdout.write(data.content);
      break;
    case "tool_approve_request":
      // Auto-approve for this example
      ws.send(JSON.stringify({
        v: 1,
        type: "tool_approve",
        sessionId: data.sessionId,
        toolCallId: data.toolCallId,
        decision: "approve"
      }));
      break;
    case "done":
      console.log("\n[Done]", data.usage);
      break;
    case "error":
      console.error("[Error]", data.message);
      break;
  }
};