Add WebSocket transport for OpenAI Responses API streaming

Introduce an optional WebSocket transport as an alternative to SSE for
the OpenAI Responses API. Users can enable it via provider_opts:

  provider_opts:
    transport: websocket

Key changes:
- Add responseEventStream interface to abstract SSE and WebSocket transports
- Refactor ResponseStreamAdapter to accept any responseEventStream
- Implement wsStream (WebSocket transport) and wsPool (connection pool
  with 55-min TTL, auto-reconnect, and lastResponseID tracking)
- Integrate WebSocket path in CreateResponseStream with automatic SSE
  fallback on connection failure
- No new dependencies (reuses existing gorilla/websocket)

The existing ResponseStreamAdapter.Recv() logic is fully reused since
WebSocket events use the same JSON schema as SSE events.

Assisted-By: docker-agent

Author: dgageot

agent-schema.json

@@ -547,7 +547,7 @@
         },
         "provider_opts": {
           "type": "object",
-          "description": "Provider-specific options. dmr: runtime_flags. anthropic/amazon-bedrock (Claude): interleaved_thinking (boolean, default true). openai/anthropic/google: rerank_prompt (string) to fully override the system prompt used for RAG reranking (advanced - prefer using results.reranking.criteria for domain-specific guidance).",
+          "description": "Provider-specific options. dmr: runtime_flags. anthropic/amazon-bedrock (Claude): interleaved_thinking (boolean, default true). openai: transport ('sse' or 'websocket') to choose between SSE and WebSocket streaming for the Responses API. openai/anthropic/google: rerank_prompt (string) to fully override the system prompt used for RAG reranking (advanced - prefer using results.reranking.criteria for domain-specific guidance).",
           "additionalProperties": true
         },
         "track_usage": {

docs/providers/openai/index.md

@@ -77,3 +77,33 @@ models:
     model: gpt-4o
     base_url: https://your-proxy.example.com/v1
 ```
+
+## WebSocket Transport
+
+For OpenAI Responses API models (gpt-4.1+, o-series, gpt-5), you can use WebSocket streaming instead of the default SSE (Server-Sent Events):
+
+```yaml
+models:
+  fast-gpt:
+    provider: openai
+    model: gpt-4.1
+    provider_opts:
+      transport: websocket  # Use WebSocket instead of SSE
+```
+
+### Benefits
+
+- **~40% faster** for workflows with 20+ tool calls
+- **Persistent connection** reduces per-turn overhead
+- **Server-side caching** of connection state
+- **Automatic fallback** to SSE if WebSocket fails
+
+### Requirements
+
+- Only works with Responses API models: `gpt-4.1+`, `o1`, `o3`, `o4`, `gpt-5`
+- NOT compatible with `--gateway` flag (automatically falls back to SSE)
+- Requires `OPENAI_API_KEY` environment variable
+
+### Example
+
+See [`examples/websocket_transport.yaml`]({{ '/examples/websocket_transport/' | relative_url }}) for a complete example.

examples/websocket_transport.yaml

@@ -0,0 +1,42 @@
+#!/usr/bin/env docker agent run
+
+# Example: WebSocket Transport for OpenAI Responses API
+#
+# This example demonstrates how to use WebSocket streaming instead of
+# Server-Sent Events (SSE) for the OpenAI Responses API.
+#
+# WebSocket transport maintains a persistent connection across tool-call
+# rounds, reducing per-turn overhead and improving end-to-end latency
+# for agentic workflows with many tool calls.
+#
+# Benefits of WebSocket over SSE:
+#   - ~40% faster end-to-end execution for workflows with 20+ tool calls
+#   - Persistent connection reduces per-turn continuation overhead
+#   - Connection-local state caching on the server
+#   - Falls back to SSE automatically if WebSocket connection fails
+#
+# Requirements:
+#   - Works only with OpenAI Responses API models (gpt-4.1+, o-series, gpt-5)
+#   - Requires OPENAI_API_KEY environment variable (or use token_key)
+#   - NOT compatible with --gateway flag (automatically falls back to SSE)
+#
+# Run with:
+#   docker agent run websocket_transport.yaml
+
+models:
+  gpt-ws:
+    provider: openai
+    model: gpt-4.1
+    provider_opts:
+      transport: websocket  # Use WebSocket instead of SSE
+
+agents:
+  root:
+    model: gpt-ws
+    description: Assistant using WebSocket streaming
+    instruction: |
+      You are a helpful assistant. Answer questions concisely.
+    toolsets:
+      - type: shell  # Real toolset for demonstrating multi-turn tool calls
+    commands:
+      demo: "List the files in the current directory, then count how many are YAML files"

pkg/model/provider/openai/client.go

@@ -7,7 +7,9 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"net/http"
 	"net/url"
+	"os"
 	"strings"
 
 	"github.com/openai/openai-go/v3"
@@ -29,12 +31,16 @@ import (
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
-// Client represents an OpenAI client wrapper
-// It implements the provider.Provider interface
+// Client represents an OpenAI client wrapper.
+// It implements the provider.Provider interface.
 type Client struct {
 	base.Config
 
 	clientFn func(context.Context) (*openai.Client, error)
+
+	// wsPool is lazily initialized when transport=websocket is configured.
+	// It maintains a persistent WebSocket connection across requests.
+	wsPool *wsPool
 }
 
 // NewClient creates a new OpenAI client from the provided configuration
@@ -307,12 +313,6 @@ func (c *Client) CreateResponseStream(
 		return nil, errors.New("at least one message is required")
 	}
 
-	client, err := c.clientFn(ctx)
-	if err != nil {
-		slog.Error("Failed to create OpenAI client", "error", err)
-		return nil, err
-	}
-
 	input := convertMessagesToResponseInput(messages)
 
 	params := responses.ResponseNewParams{
@@ -398,10 +398,88 @@ func (c *Client) CreateResponseStream(
 		slog.Error("Failed to marshal OpenAI responses request to JSON", "error", err)
 	}
 
+	// Choose transport: WebSocket or SSE (default).
+	// WebSocket is disabled when using a Gateway since most gateways don't support it.
+	transport := getTransport(&c.ModelConfig)
+	trackUsage := c.ModelConfig.TrackUsage == nil || *c.ModelConfig.TrackUsage
+
+	if transport == "websocket" && c.ModelOptions.Gateway() == "" {
+		stream, err := c.createWebSocketStream(ctx, params)
+		if err != nil {
+			slog.Error("WebSocket stream failed, falling back to SSE", "error", err)
+			// Fall through to SSE below.
+		} else {
+			slog.Debug("OpenAI responses WebSocket stream created successfully", "model", c.ModelConfig.Model)
+			return newResponseStreamAdapter(stream, trackUsage), nil
+		}
+	} else if transport == "websocket" {
+		slog.Debug("WebSocket transport requested but Gateway is configured, using SSE",
+			"model", c.ModelConfig.Model,
+			"gateway", c.ModelOptions.Gateway())
+	}
+
+	client, err := c.clientFn(ctx)
+	if err != nil {
+		slog.Error("Failed to create OpenAI client", "error", err)
+		return nil, err
+	}
 	stream := client.Responses.NewStreaming(ctx, params)
 
 	slog.Debug("OpenAI responses stream created successfully", "model", c.ModelConfig.Model)
-	return newResponseStreamAdapter(stream, c.ModelConfig.TrackUsage == nil || *c.ModelConfig.TrackUsage), nil
+	return newResponseStreamAdapter(stream, trackUsage), nil
+}
+
+// createWebSocketStream initializes (or reuses) a WebSocket connection and
+// sends the response.create message, returning a responseEventStream.
+func (c *Client) createWebSocketStream(
+	ctx context.Context,
+	params responses.ResponseNewParams,
+) (responseEventStream, error) {
+	if c.wsPool == nil {
+		// Lazy-init the pool on first WebSocket call.
+		baseURL := cmp.Or(c.ModelConfig.BaseURL, "https://api.openai.com/v1")
+		wsURL := httpToWSURL(baseURL)
+
+		headerFn := c.buildWSHeaderFn()
+		c.wsPool = newWSPool(wsURL, headerFn)
+	}
+
+	return c.wsPool.Stream(ctx, params)
+}
+
+// buildWSHeaderFn returns a function that produces the HTTP headers needed
+// for the WebSocket handshake, including the Authorization header.
+func (c *Client) buildWSHeaderFn() func(ctx context.Context) (http.Header, error) {
+	return func(ctx context.Context) (http.Header, error) {
+		h := http.Header{}
+
+		// Resolve the API key using the same logic as the HTTP client.
+		var apiKey string
+		if c.ModelConfig.TokenKey != "" {
+			apiKey, _ = c.Env.Get(ctx, c.ModelConfig.TokenKey)
+		}
+		if apiKey == "" {
+			// Fall back to the standard OPENAI_API_KEY env var.
+			apiKey = os.Getenv("OPENAI_API_KEY")
+		}
+		if apiKey != "" {
+			h.Set("Authorization", "Bearer "+apiKey)
+		}
+
+		return h, nil
+	}
+}
+
+// getTransport returns the streaming transport preference from ProviderOpts.
+// Valid values are "sse" (default) and "websocket".
+func getTransport(cfg *latest.ModelConfig) string {
+	if cfg == nil || cfg.ProviderOpts == nil {
+		return "sse"
+	}
+	if t, ok := cfg.ProviderOpts["transport"].(string); ok {
+		return strings.ToLower(t)
+	}
+	return "sse"
 }
 
 func convertMessagesToResponseInput(messages []chat.Message) []responses.ResponseInputItemUnionParam {

pkg/model/provider/openai/event_stream.go

@@ -0,0 +1,23 @@
+package openai
+
+import "github.com/openai/openai-go/v3/responses"
+
+// responseEventStream abstracts over SSE and WebSocket transports for
+// streaming Responses API events.
+//
+// The ssestream.Stream[responses.ResponseStreamEventUnion] type already
+// satisfies this interface, so it can be used directly.
+type responseEventStream interface {
+	// Next advances the stream to the next event.
+	// Returns false when the stream is exhausted or an error occurred.
+	Next() bool
+
+	// Current returns the most recently decoded event.
+	Current() responses.ResponseStreamEventUnion
+
+	// Err returns the first non-EOF error encountered by the stream.
+	Err() error
+
+	// Close releases resources held by the stream.
+	Close() error
+}

pkg/model/provider/openai/response_stream.go

@@ -13,15 +13,19 @@ import (
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
-// ResponseStreamAdapter adapts the OpenAI responses stream to our interface
+// Compile-time check: ssestream.Stream satisfies responseEventStream.
+var _ responseEventStream = (*ssestream.Stream[responses.ResponseStreamEventUnion])(nil)
+
+// ResponseStreamAdapter adapts the OpenAI responses stream to our interface.
+// It works with any responseEventStream implementation (SSE or WebSocket).
 type ResponseStreamAdapter struct {
-	stream         *ssestream.Stream[responses.ResponseStreamEventUnion]
+	stream         responseEventStream
 	trackUsage     bool
 	itemCallIDMap  map[string]string
 	itemHasContent map[string]bool
 }
 
-func newResponseStreamAdapter(stream *ssestream.Stream[responses.ResponseStreamEventUnion], trackUsage bool) *ResponseStreamAdapter {
+func newResponseStreamAdapter(stream responseEventStream, trackUsage bool) *ResponseStreamAdapter {
 	return &ResponseStreamAdapter{
 		stream:         stream,
 		trackUsage:     trackUsage,
@@ -254,5 +258,5 @@ func (a *ResponseStreamAdapter) Recv() (chat.MessageStreamResponse, error) {
 
 // Close closes the stream
 func (a *ResponseStreamAdapter) Close() {
-	a.stream.Close()
+	_ = a.stream.Close()
 }