Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions pkg/audit/auditor.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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()
}
Expand Down
39 changes: 39 additions & 0 deletions pkg/audit/auditor_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
})
}
Loading