Agent Chat using LangChain Part 2 - Token Streaming with WebSocketsThere's nothing worse than typing a message to an AI agent and then just staring at a spinning loader for ten seconds with zero feedback. It feels slow and lifeless. So in this follow-up post, I'm going to walk you through how I added proper token streaming over WebSockets, so users see the agent's response appear word by word, just like in ChatGPT.Why Token Streaming MattersWhen Claude (or any modern LLM) generates a response, it doesn't spit out the whole thing at once—it produces tokens one by one. If you wait for the complete response before sending anything to the client, the user is left hanging until it's all done. Streaming changes everything:It feels way faster, even if the total time is identical
People can start reading right away
They get immediate feedback that the agent is working

Why I Chose WebSocketsSure, you could use Server-Sent Events for one-way streaming, and that would work fine for just tokens. But WebSockets give me bidirectional communication, which turned out to be perfect for showing real-time progress when the agent calls tools. For example, when it decides to run search_contacts, I can push a quick "Using search_contacts..." update before the tool even finishes.The Backend ImplementationI built an AgentGateway in NestJS using the built-in WebSocket support with Socket.io. The main handler listens for incoming chat messages and streams everything back as it happens:@SubscribeMessage(AgentClientEvents.CHAT)
async handleChat(
  @ConnectedSocket
() client: AuthenticatedSocket,
  @MessageBody
() data: WsChatMessagePayload
) {
  client.emit(AgentSocketEvents.CHAT_START, {
    conversationId: data.conversationId,
    timestamp: new Date().toISOString(),
  });  const response = await agentService.chatWithStreaming(
    data.conversationId,
    data.message,
    (token: string) => {
      client.emit(AgentSocketEvents.CHAT_TOKEN, {
        conversationId: data.conversationId,
        token,
      });
    },
    (toolName: string, toolInput: any) => {
      client.emit(AgentSocketEvents.CHAT_TOOL_START, {
        conversationId: data.conversationId,
        toolName,
        toolInput,
      });
    },
    (toolName: string, toolResult: string) => {
      client.emit(AgentSocketEvents.CHAT_TOOL_END, {
        conversationId: data.conversationId,
        toolName,
        toolResult,
      });
    }
  );  client.emit(AgentSocketEvents.CHAT_COMPLETE, {
    ...response,
    timestamp: new Date().toISOString(),
  });
}The real work happens in chatWithStreaming, which takes callbacks for tokens and tool progress. On the Anthropic side, I switched to their streaming API instead of the regular synchronous call:const stream = await this.anthropic.messages.stream({
  model: 'claude-sonnet-4-20250514',
  max_tokens: 4096,
  system: SYSTEM_PROMPT,
  tools: this.anthropicTools,
  messages,
});for await (const event of stream) {
  if (event.type === 'content_block_delta') {
    const delta = event.delta as any;
    if (delta.type === 'text_delta' && delta.text) {
      onToken(delta.text);
    }
  }
}const finalMessage = await stream.finalMessage();I just forward each text delta to the callback, which emits it straight to the client. Once the stream ends, I grab the final message so I can process any tool calls that might still be pending.Handling AuthenticationWebSockets bypass the usual HTTP middleware, so I had to handle auth myself. When a client connects, they have 30 seconds to send an authenticate event with their JWT. If they do, I verify it and attach the user info to the socket. If not, I kick them off:handleConnection(client: AuthenticatedSocket) {
  setTimeout(() => {
    if (!this.authenticatedSockets.has(client.id)) {
      client.emit('error', { message: 'Authentication timeout' });
      client.disconnect();
    }
  }, 30000);
}The User Experience DifferenceHonestly, it's night and day. Now users see:An immediate "Agent is thinking..." indicator
Text appearing smoothly as Claude generates it
Live updates like "Using search_contacts..." when tools run
Confirmation when tools finish
The final polished response with any links or buttons

It stops feeling like you're submitting a form and waiting for processing. It actually feels like chatting with a helpful assistant.Keeping a REST FallbackI didn't want to break things for environments where WebSockets might be blocked, so I kept the original REST endpoint. It still works exactly like before—just returns the full response once everything is done. The WebSocket version is purely additive; the frontend picks the best option available.The Angular SideIn Angular, I wrapped everything in an AgentService that manages the socket connection and exposes RxJS observables the component can subscribe to. I used BehaviorSubjects so late subscribers still get the current state:@Injectable
({ providedIn: 'root' })
export class AgentService implements OnDestroy {
  private socket: Socket | null = null;  private streamingTextSubject = new BehaviorSubject<string>('');
  private isStreamingSubject = new BehaviorSubject<boolean>(false);
  private currentToolSubject = new BehaviorSubject<string | null>(null);  streamingText$ = this.streamingTextSubject.asObservable();
  isStreaming$ = this.isStreamingSubject.asObservable();
  currentTool$ = this.currentToolSubject.asObservable();
}On connection, I send the JWT for authentication and wire up all the event handlers. Tokens get appended as they arrive:this.socket.on(
  'chat:token',
  (data: { conversationId: string; token: string }) => {
    const current = this.streamingTextSubject.value;
    this.streamingTextSubject.next(current + data.token);
  }
);this.socket.on('chat:tool:start', (data: { toolName: string }) => {
  this.currentToolSubject.next(data.toolName);
});this.socket.on('chat:complete', (data) => {
  this.isStreamingSubject.next(false);
  this.chatCompleteSubject.next(data);
});In the template, while streaming I show a special bubble with a dashed border and a little typing animation:<div *ngIf="isStreaming" class="message assistant-message streaming">
  <p-tag
    *ngIf="currentTool"
    [value]="'Using: ' + currentTool"
    severity="warning"
    icon="pi pi-cog pi-spin"
  >
  </p-tag>
  <div [innerHTML]="formatMessage(streamingText)"></div>
  <span *ngIf="!streamingText" class="typing-indicator">
    <span></span><span></span><span></span>
  </span>
</div>

The sendMessage method checks if the socket is connected and uses WebSockets if possible, otherwise falls back to the REST call:sendMessage(): void {
  if (this.isConnected) {
    this.agentService.chatWs(this.conversationId, message);
  } else {
    this.agentService.chat(this.conversationId, message).subscribe({
      next: (response) => {
        this.messages.push({ role: 'assistant', content: response.message });
      }
    });
  }
}This way the feature works everywhere, and users only miss the fancy streaming if their network blocks WebSockets.Wrapping UpAdding proper token streaming completely changed how the agent feels—from "kinda slow but functional" to "this is actually fun to use." It took some extra work (auth on sockets, managing state, coordinating all the callbacks), but the improvement in user experience was absolutely worth it. If you're adding any kind of AI chat to your app, do yourself a favor and build streaming in from the beginning. You'll thank yourself later.