AI Agent API Design: REST vs WebSocket vs Streaming

Understanding AI Agent API Design Patterns

When building AI agents that communicate over the network, the choice of API transport protocol fundamentally shapes the user experience, system architecture, and agent responsiveness. This tutorial covers the three dominant patterns—REST, WebSocket, and Streaming (SSE)—with practical code examples, trade-off analysis, and best practices tailored specifically to AI agent workloads.

What is AI Agent API Design?

AI Agent API Design refers to the architectural decisions around how an AI agent—whether it's a chatbot, a coding assistant, an autonomous task executor, or a multi-step reasoning system—exposes its capabilities and communicates with clients. Unlike traditional CRUD APIs, AI agents often involve:

• Long-running operations that may take seconds or minutes to complete
• Incremental output where partial results are valuable before the full response is ready
• Bidirectional communication where the client may need to interrupt, steer, or provide additional context mid-execution
• Tool calls and reasoning chains that benefit from real-time visibility

These characteristics make the choice between REST, WebSocket, and Streaming far more consequential than in a typical microservice.

Why API Design Matters for AI Agents

The transport layer you choose directly impacts:

• Perceived responsiveness: Users expect AI to "think" in real-time. Streaming tokens can make a 10-second LLM response feel instantaneous, while a REST endpoint that returns all at once feels sluggish.
• Operational visibility: When an agent executes multi-step tool calls, streaming each step builds trust and lets users catch errors early.
• Cancellation and control: Long agent runs need interrupt mechanisms. WebSocket and streaming protocols support this natively; REST requires workarounds.
• Infrastructure complexity: WebSocket connections are stateful and require sticky sessions; REST is stateless and simpler to scale horizontally.
• Client compatibility: Not every environment supports WebSocket (e.g., some corporate proxies block it), while SSE works nearly everywhere HTTP does.

REST APIs for AI Agents

How REST Works for AI Agents

In a RESTful AI agent API, the client sends a complete request payload (prompt, context, configuration) and receives a complete response once the agent finishes processing. This is the simplest model: a single HTTP POST that blocks until completion, then returns JSON with the full result, tool call history, and any artifacts.

For agents that complete quickly (sub-second reasoning, simple retrievals), REST is perfectly adequate. For longer agents, you can combine REST with asynchronous patterns—submit a job, poll for status, then fetch results—but this adds complexity.

Example: REST Agent API with FastAPI

# server_rest.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import uuid
import asyncio

app = FastAPI()

class AgentRequest(BaseModel):
    prompt: str
    context: Optional[Dict[str, Any]] = None
    max_steps: int = 10

class ToolCall(BaseModel):
    tool_name: str
    arguments: Dict[str, Any]
    result: Optional[str] = None

class AgentResponse(BaseModel):
    run_id: str
    final_answer: str
    tool_calls: List[ToolCall]
    total_tokens: int
    execution_time_ms: int

# Simulated agent reasoning loop
async def run_agent(prompt: str, context: dict, max_steps: int) -> AgentResponse:
    # In a real app, this would call an LLM, execute tools, etc.
    await asyncio.sleep(2)  # Simulate processing time
    return AgentResponse(
        run_id=str(uuid.uuid4()),
        final_answer=f"Here is the analysis for: {prompt}",
        tool_calls=[
            ToolCall(tool_name="search", arguments={"query": prompt}, result="Relevant data found"),
            ToolCall(tool_name="calculator", arguments={"expression": "2+2"}, result="4")
        ],
        total_tokens=450,
        execution_time_ms=2100
    )

@app.post("/agent/run", response_model=AgentResponse)
async def run_agent_endpoint(request: AgentRequest):
    if not request.prompt.strip():
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
    response = await run_agent(request.prompt, request.context, request.max_steps)
    return response

@app.get("/agent/health")
async def health():
    return {"status": "ok"}

Client-Side REST Call

// client_rest.js
async function runAgentREST(prompt) {
  const response = await fetch("http://localhost:8000/agent/run", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      prompt: prompt,
      max_steps: 10
    })
  });
  
  if (!response.ok) {
    const error = await response.json();
    console.error("Agent error:", error);
    return;
  }
  
  const data = await response.json();
  console.log("Run ID:", data.run_id);
  console.log("Final Answer:", data.final_answer);
  console.log("Tool calls:", data.tool_calls.length);
  console.log("Time:", data.execution_time_ms, "ms");
  
  return data;
}

// Usage
runAgentREST("Analyze quarterly sales data");

When to Use REST

• Short-running agents: When agent execution completes in under 2-3 seconds, the simplicity of REST outweighs streaming complexity.
• Batch processing: Fire-and-forget jobs where the client doesn't need intermediate feedback.
• Simple integrations: Webhook consumers, cron jobs, or backend-to-backend communication where HTTP is already ubiquitous.
• Stateless scaling: REST endpoints behind a load balancer require no sticky sessions, making horizontal scaling trivial.

WebSocket APIs for AI Agents

How WebSocket Works for AI Agents

WebSocket provides a persistent, bidirectional TCP connection over which both client and server can send messages at any time. For AI agents, this enables:

• Real-time streaming of thoughts, tool calls, and partial responses
• Client-initiated cancellation: Send an "interrupt" message mid-execution
• Progressive context updates: The client can inject new information without starting over
• Long-lived sessions: Maintain conversation history and agent state across multiple turns

Example: WebSocket Agent Server with FastAPI

# server_websocket.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Dict, Any
import json
import asyncio
import uuid

app = FastAPI()

class AgentWebSocketSession:
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
        self.session_id = str(uuid.uuid4())
        self.cancelled = False
    
    async def send_event(self, event_type: str, data: Dict[str, Any]):
        message = json.dumps({"type": event_type, **data})
        await self.websocket.send_text(message)
    
    async def run_agent(self, prompt: str, max_steps: int):
        self.cancelled = False
        
        # Send initial acknowledgement
        await self.send_event("agent_start", {
            "session_id": self.session_id,
            "prompt": prompt
        })
        
        # Simulate step-by-step reasoning
        for step in range(1, max_steps + 1):
            if self.cancelled:
                await self.send_event("agent_cancelled", {
                    "message": f"Agent halted at step {step}",
                    "completed_steps": step - 1
                })
                return
            
            await asyncio.sleep(0.5)  # Simulate thinking
            
            # Send thought event
            await self.send_event("agent_thought", {
                "step": step,
                "thought": f"Analyzing aspect {step} of the problem...",
                "confidence": 0.7 + (step * 0.05)
            })
            
            # Simulate tool call
            if step == 3:
                await self.send_event("tool_call", {
                    "step": step,
                    "tool": "search_database",
                    "arguments": {"query": prompt},
                    "status": "executing"
                })
                await asyncio.sleep(1)
                await self.send_event("tool_result", {
                    "step": step,
                    "tool": "search_database",
                    "result": "Found 42 relevant records",
                    "status": "completed"
                })
        
        # Send final response
        await self.send_event("agent_complete", {
            "final_answer": f"Comprehensive analysis complete for: {prompt}",
            "total_steps": max_steps,
            "tokens_used": max_steps * 150
        })

@app.websocket("/agent/ws")
async def websocket_agent_endpoint(websocket: WebSocket):
    await websocket.accept()
    session = AgentWebSocketSession(websocket)
    
    try:
        while True:
            raw_message = await websocket.receive_text()
            message = json.loads(raw_message)
            msg_type = message.get("type")
            
            if msg_type == "run":
                # Start agent execution as a background task so we can
                # continue listening for cancellation messages
                asyncio.create_task(
                    session.run_agent(
                        message["prompt"],
                        message.get("max_steps", 10)
                    )
                )
            elif msg_type == "cancel":
                session.cancelled = True
                await session.send_event("cancel_ack", {
                    "message": "Cancellation requested"
                })
            elif msg_type == "ping":
                await session.send_event("pong", {"timestamp": asyncio.get_event_loop().time()})
    
    except WebSocketDisconnect:
        print(f"Session {session.session_id} disconnected")

Client-Side WebSocket Usage

// client_websocket.js
class AgentWebSocketClient {
  constructor(url) {
    this.ws = new WebSocket(url);
    this.sessionId = null;
    this.listeners = {};
    
    this.ws.onopen = () => {
      console.log("Connected to agent WebSocket");
    };
    
    this.ws.onmessage = (event) => {
      const data = JSON.parse(event.data);
      const handler = this.listeners[data.type];
      if (handler) handler(data);
      
      // Also notify generic message listener
      if (this.listeners["*"]) this.listeners["*"](data);
    };
    
    this.ws.onerror = (error) => {
      console.error("WebSocket error:", error);
    };
    
    this.ws.onclose = (event) => {
      console.log("WebSocket closed:", event.code, event.reason);
    };
  }
  
  on(eventType, callback) {
    this.listeners[eventType] = callback;
    return this; // For chaining
  }
  
  runAgent(prompt, maxSteps = 10) {
    this.send({ type: "run", prompt, max_steps: maxSteps });
  }
  
  cancel() {
    this.send({ type: "cancel" });
  }
  
  send(message) {
    if (this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
    } else {
      console.error("WebSocket is not open");
    }
  }
  
  close() {
    this.ws.close();
  }
}

// Usage example
const client = new AgentWebSocketClient("ws://localhost:8000/agent/ws");

client
  .on("agent_start", (data) => {
    console.log(`Agent started! Session: ${data.session_id}`);
  })
  .on("agent_thought", (data) => {
    console.log(`Step ${data.step}: ${data.thought} (confidence: ${data.confidence})`);
  })
  .on("tool_call", (data) => {
    console.log(`🔧 Calling ${data.tool} with args:`, data.arguments);
  })
  .on("tool_result", (data) => {
    console.log(`✅ Tool result: ${data.result}`);
  })
  .on("agent_complete", (data) => {
    console.log(`🎯 Final answer: ${data.final_answer}`);
    client.close();
  })
  .on("agent_cancelled", (data) => {
    console.log(`⛔ Cancelled: ${data.message}`);
  });

// Start agent after a short delay to ensure connection
setTimeout(() => {
  client.runAgent("Analyze market trends for Q4 2024");
}, 500);

// Cancel after 3 seconds for demonstration
// setTimeout(() => client.cancel(), 3000);

When to Use WebSocket

• Interactive agent sessions: When users need to see step-by-step reasoning, interrupt, or guide the agent in real-time.
• Bidirectional control: The client needs to push new context or change parameters mid-execution without restarting.
• Persistent conversations: Multi-turn agent interactions where maintaining state on the server reduces overhead.
• Collaborative environments: Multiple clients observing the same agent session (e.g., a shared workspace).

Streaming APIs (SSE) for AI Agents

How SSE Streaming Works

Server-Sent Events (SSE) use a standard HTTP connection where the server keeps the response stream open and pushes events to the client as they occur. Unlike WebSocket, SSE is unidirectional (server to client only) but works over standard HTTP, passes through proxies cleanly, and automatically reconnects on connection loss.

For AI agents, SSE is the sweet spot for most use cases: you get real-time streaming of tokens, thoughts, and tool calls without the complexity of bidirectional state management. Client cancellation can be achieved by simply closing the HTTP connection (the server detects the disconnected client).

Example: SSE Agent Server with FastAPI

# server_sse.py
from fastapi import FastAPI, Request
from pydantic import BaseModel
from typing import Optional
import json
import asyncio
import uuid
from starlette.responses import StreamingResponse

app = FastAPI()

class AgentSSERequest(BaseModel):
    prompt: str
    max_steps: int = 10
    stream_tokens: bool = True

async def event_generator(prompt: str, max_steps: int, stream_tokens: bool):
    run_id = str(uuid.uuid4())
    
    # Send initial metadata event
    yield f"event: agent_start\ndata: {json.dumps({'run_id': run_id, 'prompt': prompt})}\n\n"
    
    # Simulate token-by-token streaming of the final answer
    answer_tokens = [
        "Let", " me", " analyze", " the", " market", " trends", " for", " Q4.",
        " The", " data", " shows", " significant", " growth", " in", " tech", " sector.",
        " Revenue", " projections", " indicate", " a", " 15%", " increase."
    ]
    
    if stream_tokens:
        for i, token in enumerate(answer_tokens):
            await asyncio.sleep(0.08)  # Simulate token generation latency
            yield f"event: token\ndata: {json.dumps({'index': i, 'token': token, 'run_id': run_id})}\n\n"
    
    # Simulate agent reasoning steps
    for step in range(1, max_steps + 1):
        await asyncio.sleep(0.3)
        
        # Check for client disconnect (cancellation)
        # In FastAPI SSE, we rely on the generator's async context;
        # if the client disconnects, further yields may raise an error
        
        thought_data = {
            "step": step,
            "thought": f"Evaluating data source {step}: correlation analysis in progress...",
            "confidence": min(0.95, 0.5 + (step * 0.1)),
            "run_id": run_id
        }
        yield f"event: thought\ndata: {json.dumps(thought_data)}\n\n"
        
        # Simulate a tool call at step 4
        if step == 4:
            tool_call_data = {
                "step": step,
                "tool": "financial_database",
                "query": "SELECT revenue FROM quarterly WHERE year=2024",
                "status": "executing",
                "run_id": run_id
            }
            yield f"event: tool_call\ndata: {json.dumps(tool_call_data)}\n\n"
            await asyncio.sleep(0.7)
            tool_result_data = {
                "step": step,
                "tool": "financial_database",
                "result": {"rows": 120, "total_revenue": "$4.2M"},
                "status": "completed",
                "run_id": run_id
            }
            yield f"event: tool_result\ndata: {json.dumps(tool_result_data)}\n\n"
    
    # Send completion event with full assembled answer
    final_data = {
        "run_id": run_id,
        "final_answer": "".join(answer_tokens),
        "total_steps": max_steps,
        "tool_calls_count": 1,
        "tokens_generated": len(answer_tokens)
    }
    yield f"event: agent_complete\ndata: {json.dumps(final_data)}\n\n"
    
    # Send stream end marker
    yield f"event: stream_end\ndata: {json.dumps({'run_id': run_id})}\n\n"

@app.post("/agent/stream")
async def sse_agent_endpoint(request: AgentSSERequest):
    async def stream_with_disconnect_detection():
        try:
            async for event_chunk in event_generator(
                request.prompt, 
                request.max_steps, 
                request.stream_tokens
            ):
                yield event_chunk
        except (asyncio.CancelledError, ConnectionError):
            # Client disconnected - clean up resources here
            print(f"Client disconnected for prompt: {request.prompt}")
            # Re-raise to properly terminate the stream
            raise
    
    return StreamingResponse(
        stream_with_disconnect_detection(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Disable proxy buffering
        }
    )

Client-Side SSE Consumption

// client_sse.js
class AgentSSEClient {
  constructor(baseUrl) {
    this.baseUrl = baseUrl;
    this.abortController = null;
    this.listeners = {};
  }
  
  on(eventType, callback) {
    this.listeners[eventType] = callback;
    return this;
  }
  
  async runAgent(prompt, maxSteps = 10) {
    this.abortController = new AbortController();
    
    const response = await fetch(`${this.baseUrl}/agent/stream`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ prompt, max_steps: maxSteps, stream_tokens: true }),
      signal: this.abortController.signal
    });
    
    if (!response.ok) {
      console.error("Stream request failed:", response.status);
      return;
    }
    
    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let buffer = "";
    
    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        
        buffer += decoder.decode(value, { stream: true });
        
        // Parse SSE events from buffer
        const events = buffer.split("\n\n");
        // Keep the last potentially incomplete event in buffer
        buffer = events.pop() || "";
        
        for (const eventBlock of events) {
          if (!eventBlock.trim()) continue;
          
          const lines = eventBlock.split("\n");
          let eventType = "message";
          let dataStr = "";
          
          for (const line of lines) {
            if (line.startsWith("event: ")) {
              eventType = line.substring(7).trim();
            } else if (line.startsWith("data: ")) {
              dataStr = line.substring(6);
            }
          }
          
          if (dataStr) {
            try {
              const data = JSON.parse(dataStr);
              const handler = this.listeners[eventType];
              if (handler) handler(data);
              if (this.listeners["*"]) this.listeners["*"](eventType, data);
            } catch (parseError) {
              console.warn("Failed to parse SSE data:", dataStr);
            }
          }
        }
      }
    } catch (error) {
      if (error.name === "AbortError") {
        console.log("Stream cancelled by client");
        const cancelHandler = this.listeners["cancelled"];
        if (cancelHandler) cancelHandler({ reason: "user_abort" });
      } else {
        console.error("Stream error:", error);
      }
    }
  }
  
  cancel() {
    if (this.abortController) {
      this.abortController.abort();
      this.abortController = null;
    }
  }
}

// Usage example
const sseClient = new AgentSSEClient("http://localhost:8000");

sseClient
  .on("agent_start", (data) => {
    console.log(`🚀 Agent started! Run ID: ${data.run_id}`);
    console.log(`   Prompt: ${data.prompt}`);
  })
  .on("token", (data) => {
    // Stream each token as it arrives
    process.stdout.write(data.token);
  })
  .on("thought", (data) => {
    console.log(`\n💭 Step ${data.step}: ${data.thought} (${Math.round(data.confidence * 100)}% confidence)`);
  })
  .on("tool_call", (data) => {
    console.log(`🔧 Calling ${data.tool}...`);
  })
  .on("tool_result", (data) => {
    console.log(`✅ ${data.tool} returned ${data.result.total_revenue}`);
  })
  .on("agent_complete", (data) => {
    console.log(`\n\n🎯 Complete! ${data.tokens_generated} tokens, ${data.total_steps} steps`);
  });

sseClient.runAgent("Analyze market trends for Q4 2024");

// Cancel after 2 seconds to demonstrate cancellation
// setTimeout(() => sseClient.cancel(), 2000);

When to Use SSE Streaming

• Token-by-token LLM output: The gold standard for AI chat interfaces—users see responses appear in real-time.
• Observable agent pipelines: When you want clients to watch agent reasoning unfold without needing to send messages back.
• Proxy-friendly deployments: SSE passes through HTTP proxies, CDNs, and corporate firewalls that would block WebSocket upgrades.
• Simple cancellation: Clients abort by closing the HTTP connection; no special cancel protocol needed.
• Auto-reconnection: SSE clients can automatically resume from the last received event ID if the connection drops.

Comparison: REST vs WebSocket vs Streaming

The table below summarizes the key differences to help you choose the right pattern for your AI agent:

Best Practices for AI Agent API Design

1. Design for Observability from Day One

Even if you start with REST, structure your response objects to include step-by-step traces. This makes migrating to streaming later much easier and helps with debugging. Include fields like reasoning_trace, tool_calls, and intermediate_results in every response model.

2. Use Structured Event Types

Whether using WebSocket or SSE, define a consistent event taxonomy:

// Standard AI Agent event types
const AGENT_EVENTS = {
  START: "agent_start",           // Agent begins execution
  THOUGHT: "agent_thought",       // Reasoning step
  TOOL_CALL: "tool_call",         // Tool invocation started
  TOOL_RESULT: "tool_result",     // Tool invocation completed
  TOKEN: "token",                 // Streaming content token
  ERROR: "agent_error",           // Recoverable error
  COMPLETE: "agent_complete",     // Successful completion
  CANCELLED: "agent_cancelled",   // Interrupted by client
  STREAM_END: "stream_end"        // No more events
};

3. Implement Graceful Cancellation

Always support cancellation. For SSE, detect client disconnect. For WebSocket, handle a dedicated cancel message. For REST, provide a POST /agent/cancel/{run_id} endpoint and poll-based status checks.

# REST cancellation pattern
from fastapi import BackgroundTasks

active_runs = {}  # run_id -> asyncio.Task

@app.post("/agent/run_async")
async def run_async(request: AgentRequest, background_tasks: BackgroundTasks):
    run_id = str(uuid.uuid4())
    # Create a cancellable task
    task = asyncio.create_task(run_agent_cancellable(run_id, request))
    active_runs[run_id] = task
    background_tasks.add_task(cleanup_on_completion, run_id, task)
    return {"run_id": run_id, "status": "started"}

@app.post("/agent/cancel/{run_id}")
async def cancel_run(run_id: str):
    task = active_runs.get(run_id)
    if not task:
        raise HTTPException(404, "Run not found or already completed")
    task.cancel()
    return {"run_id": run_id, "status": "cancelling"}

@app.get("/agent/status/{run_id}")
async def get_status(run_id: str):
    task = active_runs.get(run_id)
    if not task:
        return {"run_id": run_id, "status": "completed_or_not_found"}
    return {
        "run_id": run_id,
        "status": "running" if not task.done() else "completed",
        "cancelled": task.cancelled()
    }

4. Version Your Event Schemas

Include a version field in every event payload. AI agent APIs evolve rapidly—adding new event types, deprecating old ones, changing payload shapes. A version field lets clients handle changes gracefully.

yield f"event: thought\ndata: {json.dumps({'version': '2.1', 'step': 1, 'thought': '...'})}\n\n"

5. Choose the Right Protocol Per Endpoint

A single AI agent service can expose multiple endpoints with different protocols:

• POST /agent/run — REST for quick, synchronous completions
• POST /agent/stream — SSE for real-time token streaming
• /agent/ws — WebSocket for interactive, cancellable sessions
• POST /agent/run_async + GET /agent/status/{id} — REST for long-running batch jobs

6. Handle Partial Results Gracefully

When an agent is cancelled or encounters an error, return what it accomplished so far. Partial results are valuable—they show the user the agent wasn't "stuck" and provide useful context for retry.

# In your SSE or WebSocket error handler
await send_event("agent_error", {
    "message": "Tool execution timed out after 30s",
    "partial_results": {
        "completed_steps": 3,
        "last_successful_tool": "search_database",
        "intermediate_findings": ["Revenue data for Q1-Q3 collected"]
    },
    "recoverable": True,
    "suggested_action": "retry_with_longer_timeout"
})

7. Monitor Stream Health

Implement heartbeat events for long-running streams. If no events are generated for a period (e.g., 30 seconds), send a keepalive to prevent proxy timeouts and help clients detect stale connections.

# Heartbeat pattern in SSE generator
async def with_heartbeat(main_generator, interval_seconds=15):
    heartbeat_task = None
    try:
        async for event in main_generator:
            yield event
            # Reset heartbeat timer after each real event
    except asyncio.CancelledError:
        raise
    
# Alternative: interleave heartbeat in the main loop
# yield f": heartbeat\n\n"  # SSE comment lines act as keepalives

8. Secure Agent Endpoints Thoroughly

AI agents often have access to tools, databases, and APIs—making them high-value targets. Apply these security measures:

• Authentication: Require bearer tokens even for WebSocket connections (pass token in query string or initial message)
• Input validation: Sanitize all user-provided prompts and context before they reach the agent's reasoning loop
• Tool sandboxing: Never pass raw user input directly to SQL queries or shell commands
• Rate limiting: Apply per-user quotas on agent runs, tokens generated, and tool calls
• Stream timeouts: Set maximum durations for SSE and WebSocket connections to prevent resource exhaustion

Conclusion

The choice between REST, WebSocket, and SSE streaming for AI agent APIs is not about finding a single "best" protocol—it's about matching the communication pattern to your agent's execution characteristics and your users' expectations. REST shines for short, atomic agent tasks where simplicity and stateless scaling matter most. WebSocket is the