From 9c98e94a8bab6526ea57d5fb5429eb7c5bfaa490 Mon Sep 17 00:00:00 2001
From: Juan Antonio Osorio <ozz@stacklok.com>
Date: Mon, 27 Jul 2026 13:51:24 +0300
Subject: [PATCH] Audit stream opens on flush and pre-write panics

Stacked on #5874. Two fast-follows from rdimitrov's review:

- Flush() now logs the stream connection event: a handler that
  establishes the stream by flushing headers and then blocks never hits
  WriteHeader/Write, so the event was delayed until the first data
  write or lost entirely.
- The stream-open log moved into a defer registered before ServeHTTP,
  so a panic in an inner handler still produces the event during
  unwinding. On chains whose recovery middleware runs outside audit
  (the vMCP Serve path) the event previously vanished.
---
 pkg/audit/auditor.go      | 14 +++++++++++---
 pkg/audit/auditor_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/pkg/audit/auditor.go b/pkg/audit/auditor.go
index 96e3134879..42bffb53f8 100644
--- a/pkg/audit/auditor.go
+++ b/pkg/audit/auditor.go
@@ -231,10 +231,13 @@ func (a *Auditor) Middleware(next http.Handler) http.Handler {
 		// carries the authenticated identity (or records the 401/403 denial).
 		if a.isMCPStreamOpenRequest(r) {
 			sw := &streamOpenWriter{ResponseWriter: w, auditor: a, req: r}
+			// Deferred BEFORE ServeHTTP so a panic in an inner handler still
+			// produces the connection event during unwinding — some chains run
+			// the recovery middleware OUTSIDE audit (e.g. the vMCP Serve path),
+			// which would otherwise swallow the event entirely. If the stream
+			// already logged on first write this is a no-op.
+			defer sw.logOnce(http.StatusOK)
 			next.ServeHTTP(sw, r)
-			// Streams that end without a single write still get an event
-			// (net/http sends an implicit 200 in that case).
-			sw.logOnce(http.StatusOK)
 			return
 		}
 
@@ -703,7 +706,12 @@ func (sw *streamOpenWriter) Write(data []byte) (int, error) {
 }
 
 // Flush implements http.Flusher if the underlying ResponseWriter supports it.
+// A flush commits the response headers with an implicit 200, so it counts as
+// the stream being established — log the connection event here too, otherwise
+// a handler that flushes before its first write would delay (or, on a stream
+// that only ever flushes, lose) the event.
 func (sw *streamOpenWriter) Flush() {
+	sw.logOnce(http.StatusOK)
 	if flusher, ok := sw.ResponseWriter.(http.Flusher); ok {
 		flusher.Flush()
 	}
diff --git a/pkg/audit/auditor_test.go b/pkg/audit/auditor_test.go
index 633f3424d5..4c9005abcd 100644
--- a/pkg/audit/auditor_test.go
+++ b/pkg/audit/auditor_test.go
@@ -1224,4 +1224,43 @@ func TestStreamOpenAuditEvents(t *testing.T) {
 		assert.Equal(t, OutcomeSuccess, events[0]["outcome"],
 			"net/http sends an implicit 200 when the handler writes nothing")
 	})
+
+	t.Run("flush before first write logs the connection event", func(t *testing.T) {
+		t.Parallel()
+		auditor, logBuf := newBufferAuditor(t)
+
+		// A handler that establishes the stream by flushing headers and then
+		// blocks (waiting for events to send) never hits WriteHeader/Write.
+		flusher := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "text/event-stream")
+			if f, ok := w.(http.Flusher); ok {
+				f.Flush()
+			}
+		})
+		auditor.Middleware(flusher).ServeHTTP(httptest.NewRecorder(), newStreamRequest())
+
+		events := decodeAuditEvents(t, logBuf)
+		require.Len(t, events, 1, "the flush must not bypass the connection event")
+		assert.Equal(t, EventTypeSSEConnection, events[0]["type"])
+		assert.Equal(t, OutcomeSuccess, events[0]["outcome"])
+	})
+
+	t.Run("panic before first write still logs the connection event", func(t *testing.T) {
+		t.Parallel()
+		auditor, logBuf := newBufferAuditor(t)
+
+		panicker := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+			panic("boom before first write")
+		})
+		handler := auditor.Middleware(panicker)
+		require.Panics(t, func() {
+			handler.ServeHTTP(httptest.NewRecorder(), newStreamRequest())
+		}, "no recovery middleware on this chain: the panic must propagate")
+
+		events := decodeAuditEvents(t, logBuf)
+		require.Len(t, events, 1,
+			"chains whose recovery middleware runs OUTSIDE audit (e.g. the vMCP Serve path) "+
+				"must not lose the connection event to a panic")
+		assert.Equal(t, EventTypeSSEConnection, events[0]["type"])
+	})
 }