Honor configured vMCP backend timeouts

pkg/transport/middleware/write_timeout.go

@@ -7,29 +7,82 @@ import (
 	"log/slog"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 )
 
-// WriteTimeout clears the write deadline for qualifying SSE connections
-// (GET + Accept: text/event-stream + matching path) so http.Server.WriteTimeout
-// does not kill long-lived streams (golang/go#16100). All other requests are
-// left untouched.
-func WriteTimeout(endpointPath string) func(http.Handler) http.Handler {
+const defaultResponseWriteTimeout = 30 * time.Second
+
+type postResponseWriter struct {
+	http.ResponseWriter
+	writeTimeout time.Duration
+	armOnce      sync.Once
+}
+
+func (w *postResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
+
+func (w *postResponseWriter) armWriteDeadline() {
+	w.armOnce.Do(func() {
+		// A streamable-HTTP POST may legitimately negotiate an SSE response.
+		// Like a qualifying GET stream, its lifetime is governed by request
+		// cancellation rather than a socket write deadline.
+		if strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") {
+			return
+		}
+		if err := http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(w.writeTimeout)); err != nil {
+			slog.Warn("failed to arm MCP response write deadline", "error", err)
+		}
+	})
+}
+
+func (w *postResponseWriter) WriteHeader(statusCode int) {
+	w.armWriteDeadline()
+	w.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (w *postResponseWriter) Write(p []byte) (int, error) {
+	w.armWriteDeadline()
+	return w.ResponseWriter.Write(p)
+}
+
+func (w *postResponseWriter) Flush() {
+	w.armWriteDeadline()
+	if err := http.NewResponseController(w.ResponseWriter).Flush(); err != nil {
+		slog.Debug("failed to flush MCP response", "error", err)
+	}
+}
+
+// WriteTimeout clears the server-level write deadline while an MCP POST is
+// computing, then arms a fresh deadline when a non-streaming response begins.
+// Qualifying SSE requests remain unbounded and are canceled through request
+// contexts. Other requests are left untouched. The optional duration controls
+// the response-write allowance and defaults to 30 seconds.
+func WriteTimeout(endpointPath string, responseWriteTimeout ...time.Duration) func(http.Handler) http.Handler {
+	writeTimeout := defaultResponseWriteTimeout
+	if len(responseWriteTimeout) > 0 && responseWriteTimeout[0] > 0 {
+		writeTimeout = responseWriteTimeout[0]
+	}
+
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if r.Method == http.MethodGet &&
+			isMCPPost := r.Method == http.MethodPost && r.URL.Path == endpointPath
+			isMCPStream := r.Method == http.MethodGet &&
 				strings.Contains(r.Header.Get("Accept"), "text/event-stream") &&
-				r.URL.Path == endpointPath {
+				r.URL.Path == endpointPath
+			if isMCPPost || isMCPStream {
 				rc := http.NewResponseController(w)
 				if err := rc.SetWriteDeadline(time.Time{}); err != nil {
-					slog.Warn("failed to clear write deadline for SSE connection; stream may be killed by server WriteTimeout",
+					slog.Warn("failed to clear write deadline for MCP request; request may be killed by server WriteTimeout",
 						"error", err,
 						"method", r.Method,
 						"path", r.URL.Path,
 						"remote", r.RemoteAddr,
 					)
 				}
 			}
+			if isMCPPost {
+				w = &postResponseWriter{ResponseWriter: w, writeTimeout: writeTimeout}
+			}
 			next.ServeHTTP(w, r)
 		})
 	}

pkg/transport/middleware/write_timeout_test.go

@@ -27,11 +27,13 @@ type deadlineTrackingResponseWriter struct {
 	*httptest.ResponseRecorder
 	deadlineSet bool
 	deadline    time.Time
+	deadlines   []time.Time
 }
 
 func (d *deadlineTrackingResponseWriter) SetWriteDeadline(t time.Time) error {
 	d.deadlineSet = true
 	d.deadline = t
+	d.deadlines = append(d.deadlines, t)
 	return nil
 }
 
@@ -97,17 +99,52 @@ func TestWriteTimeout_GETOnWrongPathLeavesDeadlineUntouched(t *testing.T) {
 	assert.Equal(t, http.StatusOK, w.Code)
 }
 
-// TestWriteTimeout_POSTLeavesDeadlineUntouched verifies that POST requests are not
-// touched by the middleware — their deadline comes from http.Server.WriteTimeout.
-func TestWriteTimeout_POSTLeavesDeadlineUntouched(t *testing.T) {
+// TestWriteTimeout_MCPPOSTDefersDeadlineUntilResponse verifies that MCP POST
+// requests are unbounded while computing and receive a fresh write deadline
+// when their non-streaming response starts.
+func TestWriteTimeout_MCPPOSTDefersDeadlineUntilResponse(t *testing.T) {
 	t.Parallel()
 
 	w := newDeadlineTracker()
 	r := httptest.NewRequest(http.MethodPost, testEndpointPath, nil)
 
 	mw(noopHandler).ServeHTTP(w, r)
 
-	assert.False(t, w.deadlineSet, "POST deadline is managed by http.Server.WriteTimeout, not the middleware")
+	require.Len(t, w.deadlines, 2)
+	assert.True(t, w.deadlines[0].IsZero(), "the handler computation deadline must be cleared")
+	assert.False(t, w.deadlines[1].IsZero(), "response writing must receive a bounded deadline")
+	assert.True(t, w.deadlines[1].After(time.Now()), "response write deadline must be in the future")
+	assert.Equal(t, http.StatusOK, w.Code)
+}
+
+func TestWriteTimeout_MCPPostSSEResponseRemainsUnbounded(t *testing.T) {
+	t.Parallel()
+
+	w := newDeadlineTracker()
+	r := httptest.NewRequest(http.MethodPost, testEndpointPath, nil)
+	handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+	})
+
+	mw(handler).ServeHTTP(w, r)
+
+	require.Len(t, w.deadlines, 1)
+	assert.True(t, w.deadlines[0].IsZero(), "SSE response must retain the cleared deadline")
+	assert.Equal(t, http.StatusOK, w.Code)
+}
+
+// TestWriteTimeout_POSTOnWrongPathLeavesDeadlineUntouched verifies that only
+// the configured MCP endpoint gets the relaxed POST deadline.
+func TestWriteTimeout_POSTOnWrongPathLeavesDeadlineUntouched(t *testing.T) {
+	t.Parallel()
+
+	w := newDeadlineTracker()
+	r := httptest.NewRequest(http.MethodPost, "/health", nil)
+
+	mw(noopHandler).ServeHTTP(w, r)
+
+	assert.False(t, w.deadlineSet, "POST on non-MCP path must retain the server WriteTimeout")
 	assert.Equal(t, http.StatusOK, w.Code)
 }
 

pkg/vmcp/cli/serve.go

@@ -354,6 +354,10 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 	if revisions, ok := backendClient.(vmcp.RevisionReporter); ok {
 		sessionFactoryOpts = append(sessionFactoryOpts, vmcpsession.WithRevisionLookup(revisions.CachedRevision))
 	}
+	sessionFactoryOpts = append(
+		sessionFactoryOpts,
+		vmcpsession.WithRequestTimeoutResolver(backendRequestTimeoutResolver(vmcpCfg)),
+	)
 	sessionFactory := vmcpsession.NewSessionFactory(outgoingRegistry, sessionFactoryOpts...)
 
 	// When the optimizer is enabled, its meta-tools are pass-through tools.
@@ -520,6 +524,24 @@ func getStatusReportingInterval(cfg *config.Config) time.Duration {
 	return 0
 }
 
+// backendRequestTimeoutResolver resolves the documented operational timeout
+// for a backend workload. Configuration loaded from YAML has defaults applied
+// and is immutable after startup; the defensive fallback also supports quick
+// mode and direct embedders that omit Operational.
+func backendRequestTimeoutResolver(cfg *config.Config) func(workloadID string) time.Duration {
+	timeouts := config.DefaultOperationalConfig().Timeouts
+	if cfg != nil && cfg.Operational != nil && cfg.Operational.Timeouts != nil {
+		timeouts = cfg.Operational.Timeouts
+	}
+
+	return func(workloadID string) time.Duration {
+		if timeout, ok := timeouts.PerWorkload[workloadID]; ok && timeout > 0 {
+			return time.Duration(timeout)
+		}
+		return time.Duration(timeouts.Default)
+	}
+}
+
 // loadAndValidateConfig loads and validates the vMCP configuration file.
 func loadAndValidateConfig(configPath string) (*config.Config, error) {
 	slog.Info(fmt.Sprintf("Loading configuration from: %s", configPath))
@@ -624,7 +646,10 @@ func discoverBackends(
 		return nil, nil, nil, fmt.Errorf("failed to create outgoing authentication registry: %w", err)
 	}
 
-	backendClient, err := vmcpclient.NewHTTPBackendClient(outgoingRegistry)
+	backendClient, err := vmcpclient.NewHTTPBackendClient(
+		outgoingRegistry,
+		vmcpclient.WithRequestTimeoutResolver(backendRequestTimeoutResolver(cfg)),
+	)
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("failed to create backend client: %w", err)
 	}

pkg/vmcp/cli/serve_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -23,6 +24,52 @@ import (
 	vmcpmocks "github.com/stacklok/toolhive/pkg/vmcp/mocks"
 )
 
+func TestBackendRequestTimeoutResolver(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		cfg        *config.Config
+		workloadID string
+		want       time.Duration
+	}{
+		{
+			name:       "missing operational config uses default",
+			cfg:        &config.Config{},
+			workloadID: "elastic",
+			want:       30 * time.Second,
+		},
+		{
+			name: "configured default applies to unmatched workload",
+			cfg: &config.Config{Operational: &config.OperationalConfig{
+				Timeouts: &config.TimeoutConfig{Default: config.Duration(90 * time.Second)},
+			}},
+			workloadID: "other",
+			want:       90 * time.Second,
+		},
+		{
+			name: "per workload timeout overrides configured default",
+			cfg: &config.Config{Operational: &config.OperationalConfig{
+				Timeouts: &config.TimeoutConfig{
+					Default: config.Duration(90 * time.Second),
+					PerWorkload: map[string]config.Duration{
+						"elastic": config.Duration(240 * time.Second),
+					},
+				},
+			}},
+			workloadID: "elastic",
+			want:       240 * time.Second,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			assert.Equal(t, tt.want, backendRequestTimeoutResolver(tt.cfg)(tt.workloadID))
+		})
+	}
+}
+
 // TestLoadAndValidateConfig covers all config-loading paths.
 func TestLoadAndValidateConfig(t *testing.T) {
 	t.Parallel()