Honor context when acquiring init semaphore

Queued backend init must not block on the semaphore after the
session budget expires. Warn only on deadline exceeded; log
caller cancel at debug.

Co-authored-by: Ravi Tharuma <RaviTharuma@users.noreply.github.com>

Author: cursoragent

pkg/vmcp/session/default_session_test.go

@@ -1186,6 +1186,62 @@ func TestNewSessionFactory_SessionInitTimeoutBoundsWait(t *testing.T) {
 	require.NoError(t, sess.Close())
 }
 
+func TestNewSessionFactory_SessionInitTimeoutSkipsQueuedSemaphoreAcquire(t *testing.T) {
+	t.Parallel()
+
+	// One in-flight backend holds the only semaphore slot until the session
+	// budget expires, then delays release so queued backends would hang on
+	// `sem <-` if acquire ignored ctx. Those queued connectors ignore ctx
+	// and sleep long enough that MakeSession would miss the budget.
+	const queuedSleep = 3 * time.Second
+	started := make(chan struct{})
+	var holderStarted atomic.Bool
+	var queuedEntered atomic.Int64
+	connector := func(ctx context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) {
+		if holderStarted.CompareAndSwap(false, true) {
+			close(started)
+			select {
+			case <-ctx.Done():
+				// Keep the slot occupied after expiry so queued
+				// acquire must lose to ctx.Done() in the select.
+				time.Sleep(150 * time.Millisecond)
+				return nil, nil, ctx.Err()
+			case <-time.After(30 * time.Second):
+				return nil, nil, errors.New("holder did not observe session budget")
+			}
+		}
+		queuedEntered.Add(1)
+		time.Sleep(queuedSleep)
+		return nil, nil, errors.New("queued backend ran after session budget")
+	}
+
+	backends := []*vmcp.Backend{
+		{ID: "holder", Name: "holder", BaseURL: "http://x:1", TransportType: "streamable-http"},
+		{ID: "queued-a", Name: "queued-a", BaseURL: "http://x:2", TransportType: "streamable-http"},
+		{ID: "queued-b", Name: "queued-b", BaseURL: "http://x:3", TransportType: "streamable-http"},
+	}
+	factory := newSessionFactoryWithConnector(connector,
+		WithSessionInitTimeout(80*time.Millisecond),
+		WithBackendInitTimeout(30*time.Second),
+		WithMaxBackendInitConcurrency(1),
+	)
+
+	start := time.Now()
+	sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil)
+	elapsed := time.Since(start)
+
+	require.NoError(t, err)
+	require.NotNil(t, sess)
+	select {
+	case <-started:
+	default:
+		t.Fatal("expected the in-flight backend to start before MakeSession returned")
+	}
+	assert.Zero(t, queuedEntered.Load(), "queued backends must not acquire the semaphore after the session budget expires")
+	assert.Less(t, elapsed, queuedSleep, "MakeSession must return without waiting for queued connectors")
+	require.NoError(t, sess.Close())
+}
+
 func TestValidateSessionID(t *testing.T) {
 	t.Parallel()
 

pkg/vmcp/session/factory.go

@@ -7,6 +7,7 @@ package session
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"slices"
@@ -510,18 +511,30 @@ func (f *defaultMultiSessionFactory) makeBaseSession(
 	for i, b := range backends {
 		go func(i int, b *vmcp.Backend) {
 			defer wg.Done()
-			sem <- struct{}{}
-			defer func() { <-sem }()
+			// Acquire is context-aware so queued backends do not block on
+			// the semaphore after sessionInitTimeout (or cancel) has fired.
+			select {
+			case sem <- struct{}{}:
+				defer func() { <-sem }()
+			case <-ctx.Done():
+				return
+			}
 			rawResults[i], modernSkipped[i] = f.initOneBackend(ctx, b, identity, sessionHints[b.ID], sink)
 		}(i, b)
 	}
 	wg.Wait()
 
 	if err := ctx.Err(); err != nil {
-		slog.Warn("session initialize budget expired; returning with backends that connected in time",
-			"error", err,
-			"backendCount", len(backends),
-			"timeout", f.sessionInitTimeout)
+		if errors.Is(err, context.DeadlineExceeded) {
+			slog.Warn("session initialize budget expired; returning with backends that connected in time",
+				"error", err,
+				"backendCount", len(backends),
+				"timeout", f.sessionInitTimeout)
+		} else {
+			slog.Debug("session initialize cancelled; returning with backends that connected in time",
+				"error", err,
+				"backendCount", len(backends))
+		}
 	}
 
 	connections := make(map[string]backend.Session, len(backends))

pkg/vmcp/session/factory_revision_test.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"log/slog"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -212,3 +213,72 @@ func TestMakeSession_AllBackendsFailedWarning(t *testing.T) {
 		})
 	}
 }
+
+const (
+	sessionBudgetExpiredMsg = "session initialize budget expired"
+	sessionCancelledMsg     = "session initialize cancelled"
+)
+
+//nolint:paralleltest // setupLogRecorder swaps the global slog default logger.
+func TestMakeSession_SessionInitContextLogs(t *testing.T) {
+	t.Run("deadline exceeded logs a warning", func(t *testing.T) {
+		buf := setupLogRecorder(t)
+		released := make(chan struct{})
+		t.Cleanup(func() { close(released) })
+		connector := func(ctx context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) {
+			select {
+			case <-ctx.Done():
+				return nil, nil, ctx.Err()
+			case <-released:
+				return &mockConnectedBackend{}, &vmcp.CapabilityList{}, nil
+			}
+		}
+		factory := newSessionFactoryWithConnector(connector, WithSessionInitTimeout(50*time.Millisecond))
+		backend := &vmcp.Backend{ID: "slow", Name: "slow", BaseURL: "http://x:9", TransportType: "streamable-http"}
+
+		sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil)
+		require.NoError(t, err)
+		t.Cleanup(func() { _ = sess.Close() })
+
+		logs := buf.String()
+		assert.Contains(t, logs, sessionBudgetExpiredMsg)
+		assert.NotContains(t, logs, sessionCancelledMsg)
+	})
+
+	t.Run("caller cancel logs at debug", func(t *testing.T) {
+		buf := setupLogRecorder(t)
+		entered := make(chan struct{})
+		connector := func(ctx context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) {
+			close(entered)
+			<-ctx.Done()
+			return nil, nil, ctx.Err()
+		}
+		factory := newSessionFactoryWithConnector(connector, WithSessionInitTimeout(30*time.Second))
+		backend := &vmcp.Backend{ID: "slow", Name: "slow", BaseURL: "http://x:9", TransportType: "streamable-http"}
+
+		ctx, cancel := context.WithCancel(context.Background())
+		errCh := make(chan error, 1)
+		var sess MultiSession
+		go func() {
+			var err error
+			sess, err = factory.MakeSessionWithID(ctx, uuid.New().String(), nil, []*vmcp.Backend{backend}, nil)
+			errCh <- err
+		}()
+		select {
+		case <-entered:
+		case <-time.After(2 * time.Second):
+			t.Fatal("timeout waiting for connector to start")
+		}
+		cancel()
+		require.NoError(t, <-errCh)
+		t.Cleanup(func() {
+			if sess != nil {
+				_ = sess.Close()
+			}
+		})
+
+		logs := buf.String()
+		assert.Contains(t, logs, sessionCancelledMsg)
+		assert.NotContains(t, logs, sessionBudgetExpiredMsg)
+	})
+}