# WebSocket Protocol

The Server module is a Javalin-based backend that bridges frontends (CLI, IDE, web) to the TnsAI agent framework via WebSocket. It provides multi-agent sessions, real-time streaming, risk-based tool approval, hybrid RAG search, and an audit trail.

## 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.

```java
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`.

| Type | Fields | Description |
|------|--------|-------------|
| `chat` | `message` | Send a user message to the session's active agent |
| `agent_add` | `role`, `provider?`, `model?`, `goal?`, `domain?`, `name?`, `roles?` | Add a new agent to the session |
| `agent_remove` | `agentId` | Remove an agent from the session |
| `tool_approve` | `toolCallId`, `decision` | Respond to a tool approval request (`approve`, `reject`, `always`) |
| `cancel` | -- | Cancel the running chat task and pending approvals |
| `mcp_tools` | `tools: McpToolDef[]` | Register MCP server tools for the session |
| `mcp_result` | `toolCallId`, `result?`, `error?` | Return the result of an MCP proxy tool call |
| `set_autonomy` | `level` | Change the autonomy level (`FULL_AUTO`, `SUPERVISED`, `CAUTIOUS`, `MANUAL`) |
| `kb_add` | `content`, `metadata?` | Add a document to the session knowledge base |
| `kb_search` | `query`, `limit?`, `threshold?` | Search the knowledge base |
| `kb_list` | -- | List all documents in the knowledge base |
| `kb_remove` | `documentId` | Remove a document from the knowledge base |

### 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.

| Type | Fields | Description |
|------|--------|-------------|
| `token` | `agentId`, `content` | A streamed content token from the LLM |
| `tool_start` | `agentId`, `toolCallId`, `tool`, `args`, `risk` | Tool execution has begun |
| `tool_approve_request` | `agentId`, `toolCallId`, `tool`, `risk`, `summary` | User approval required for a tool call |
| `tool_result` | `agentId`, `toolCallId`, `tool`, `duration`, `result` | Tool execution completed |
| `agent_state` | `agentId`, `status`, `role`, `provider?`, `model?` | Agent status change (`thinking`, `idle`) |
| `team_update` | `agents: AgentInfo[]`, `strategy` | Full roster update after agent add/remove |
| `done` | `usage: {tokens, cost}` | LLM response stream completed |
| `error` | `message`, `code?` | Error event with optional error code |
| `mcp_call` | `toolCallId`, `server`, `tool`, `args` | Request the client to execute an MCP tool |
| `index_start` | `path`, `totalFiles` | Directory indexing started |
| `index_progress` | `indexedFiles`, `totalFiles`, `currentFile` | Indexing progress update |
| `index_done` | `fileCount`, `chunkCount` | Indexing completed |
| `autonomy_changed` | `level` | Autonomy level was updated |
| `audit_entry` | `toolCallId`, `agentId`, `tool`, `risk`, `decision`, `timestamp`, `summary` | Audit trail entry |
| `kb_added` | `documentId` | Document added to knowledge base |
| `kb_search_result` | `results: KbSearchHit[]` | Knowledge base search results |
| `kb_list_result` | `documents: KbDocInfo[]`, `totalCount` | Knowledge base document listing |
| `kb_removed` | `documentId`, `found` | Document removal result |

## 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`.

```java
// 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.

```json
// Add an agent using the server's configured default LLM supplier
{
  "v": 1,
  "type": "agent_add",
  "sessionId": "sess-1",
  "role": "reviewer",
  "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.

### Provider Overrides and `agent_add` Results

The example omits `provider`, so the new agent uses the server's configured
default `LLMClient` supplier.

Maven Central `0.14.1` (`TnsAI@609a661d`) fails loudly:

- A null or blank `provider` selects the configured default supplier.
- `provider: "ollama"` is the only supported explicit override.
- Any other explicit provider produces an `error` event with code
  `AGENT_ADD_FAILED`. The attempted agent is not added and no `team_update`
  is emitted for that attempt.
- After a successful add, the idle `agent_state` reports `provider` and
  `model` from the agent's live `LLMClient`, not from the request fields.

To use a cloud client, configure it as the default supplier and omit
`provider`; the override field cannot select it by name.

Maven Central `0.13.0` was lenient: nonblank names other than `"ollama"`
logged a warning and fell back to the default supplier, and a successful
`agent_state` echoed the requested provider and model.

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

```json
{
  "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.

```json
// 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.

```json
// 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.



| Method | Path | Description |
|--------|------|-------------|
| GET | `/health` | Aggregated health status (200/503) |
| GET | `/health/live` | Liveness probe (always 200) |
| GET | `/health/ready` | Readiness probe (server + components) |
| GET | `/api/info` | Server info (version, protocol, active sessions) |
| POST | `/api/index` | Trigger directory indexing |
| GET | `/api/index/status` | Current index status |
| POST | `/api/search` | Search the index |
| DELETE | `/api/index` | Clear the index |
| GET | `/api/audit` | Retrieve audit trail for a session |

## 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.

```javascript
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;
  }
};
```