🐛 Forward Proxy blocks response with SSE base MCP server #642

Describe the bug

Hi team, I'm evaluating the authbridge-proxy as a standalone locally with podman (not using kagenti), and I found a problem handling the SSE events from an MCP server.

Both my agent and MCP are Quarkus-based (langchain4j extension), and the agent is connecting using the SSE transport.

I added some log messages to the authbridge code (not a GoLang developer, not proud of debugging code with prints 😢), and I discovered that the code blocks around here:

https://github.com/kagenti/kagenti-extensions/blob/e9a7dd42db87843760e7bccb97ba5d0b15436822/authbridge/authlib/listener/forwardproxy/server.go#L450-L452

I added a log message after line 452, and it is only printed after the agent times out and closes the connection.

Steps To Reproduce

My configuration is the following, although I suspect the issue should be reproducible with any pluging involved.

config.yaml

mode: proxy-sidecar

listener:
  reverse_proxy_addr: ":8091"
  reverse_proxy_backend: "http://chatbot-backend:8090"
  forward_proxy_addr: ":8092"
  session_api_addr: ":8093"

pipeline:
  inbound:
    plugins:
      - name: jwt-validation
        config:
          jwks_url: "http://keycloak:8180/realms/demo/protocol/openid-connect/certs"
          issuer: "http://localhost:8180/realms/demo"
          audience: ****
  outbound:
    plugins:
      - name: token-exchange
        config:
          token_url: "http://keycloak:8180/realms/demo/protocol/openid-connect/token"
          identity:
            type: client-secret
            client_id: ****
            client_secret: ****
          routes:
            file: "/etc/authproxy/routes.yaml"

routes.yaml

- host: "mcp-server"
  target_audience: "***"
  token_scopes: ""

Let me know if you want a reproducer; my experiments has agents, authbridge, mcp server, keyloak, web-ui and so on and are difficult to share. I can investigate if it is reproducible with cURL if asked.

Expected Behavior

I expect the SSE event to reach the agent code.

Kagenti Version

main

Additional Context

I tried to use the io.ReadAll, but it blocks as well.

https://github.com/kagenti/kagenti-extensions/blob/e9a7dd42db87843760e7bccb97ba5d0b15436822/authbridge/authlib/listener/forwardproxy/server.go#L395

I asked AI for a possible replacement for the io.Copy, and it generated the following snippet. It fixes the problem, but I don't have the knowledge of whether it is a valid fix. And there are other places to fix as well.

slog.Info("Starting real-time response streaming...")

// 1. Assert that the ResponseWriter supports the Flusher interface
flusher, ok := w.(http.Flusher)
if !ok {
	slog.Error("Streaming unsupported by the current response writer")
	http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
	return
}

// 2. Read from the MCP server response in small chunks
buf := make([]byte, 4096)
for {
	n, readErr := resp.Body.Read(buf)
	if n > 0 {
		// Write the chunk to the client
		_, writeErr := w.Write(buf[:n])
		if writeErr != nil {
			slog.Debug("Client disconnected while writing stream", "error", writeErr)
			break
		}
		
		// CRITICAL: Force Go to push the data out over the network immediately!
		flusher.Flush() 
	}

	if readErr != nil {
		if readErr != io.EOF {
			slog.Debug("Error reading from upstream MCP server", "error", readErr)
		}
		break
	}
}

slog.Info("Stream connection closed!")