[test] Add tests for launcher.GetOrLaunchForSession

internal/launcher/getorlaunchforsession_test.go

@@ -2,10 +2,14 @@ package launcher
 
 import (
 	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/github/gh-aw-mcpg/internal/config"
 )
@@ -133,3 +137,170 @@ func TestGetOrLaunchForSession_StartupTimeoutConfig(t *testing.T) {
 	// verifies that the startupTimeout field can be configured correctly.
 	// The actual timeout behavior is tested in integration tests.
 }
+
+// TestGetOrLaunchForSession_StdioSessionPoolHit tests that a cached stdio session connection
+// is returned from the session pool without launching a new backend.
+func TestGetOrLaunchForSession_StdioSessionPoolHit(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+
+	// Create a mock HTTP server to get a real *mcp.Connection to put in the pool
+	mockServer := newMockHTTPMCPServer(t)
+	defer mockServer.Close()
+
+	ctx := context.Background()
+	cfg := newTestConfig(map[string]*config.ServerConfig{
+		"http-backend": {
+			Type: "http",
+			URL:  mockServer.URL,
+		},
+		"stdio-backend": {
+			Type:    "stdio",
+			Command: "docker",
+			Args:    []string{"run", "--rm", "-i", "nonexistent:latest"},
+		},
+	})
+
+	l := New(ctx, cfg)
+	defer l.Close()
+
+	// Get a real HTTP connection to use as a stand-in in the session pool
+	httpConn, err := GetOrLaunch(l, "http-backend")
+	require.NoError(err)
+	require.NotNil(httpConn)
+
+	// Pre-populate the session pool with the connection for the stdio backend
+	sessionID := "test-session-123"
+	l.sessionPool.Set("stdio-backend", sessionID, httpConn)
+
+	// GetOrLaunchForSession should return the cached connection without launching a new process
+	result, err := GetOrLaunchForSession(l, "stdio-backend", sessionID)
+	require.NoError(err)
+	require.NotNil(result)
+
+	// Verify we got back the same cached connection
+	assert.Equal(httpConn, result, "Should return the pre-cached connection from session pool")
+}
+
+// TestGetOrLaunchForSession_StdioLaunchFailure tests that a failed stdio launch returns an error.
+func TestGetOrLaunchForSession_StdioLaunchFailure(t *testing.T) {
+	require := require.New(t)
+
+	ctx := context.Background()
+	cfg := newTestConfig(map[string]*config.ServerConfig{
+		"stdio-backend": {
+			Type:    "stdio",
+			Command: "nonexistent-command-xyz-99999",
+			Args:    []string{"--flag"},
+		},
+	})
+
+	l := New(ctx, cfg)
+	defer l.Close()
+
+	// GetOrLaunchForSession should fail for an invalid command
+	conn, err := GetOrLaunchForSession(l, "stdio-backend", "session-abc")
+	require.Error(err, "Should return error for invalid command")
+	require.Nil(conn)
+	assert.Contains(t, err.Error(), "failed to create connection")
+}
+
+// TestGetOrLaunchForSession_DirectCommandWarningInContainer tests that a security warning
+// is logged when a direct (non-docker) command is used inside a container.
+func TestGetOrLaunchForSession_DirectCommandWarningInContainer(t *testing.T) {
+	require := require.New(t)
+
+	ctx := context.Background()
+	cfg := newTestConfig(map[string]*config.ServerConfig{
+		"stdio-backend": {
+			Type:    "stdio",
+			Command: "echo", // direct command, not docker
+			Args:    []string{"hello"},
+		},
+	})
+
+	l := New(ctx, cfg)
+	defer l.Close()
+
+	// Simulate running inside a container
+	l.runningInContainer = true
+
+	// The launch will fail (echo doesn't implement MCP), but the security
+	// warning path (lines 222-226 of launcher.go) will be exercised.
+	conn, err := GetOrLaunchForSession(l, "stdio-backend", "session-warn")
+	require.Error(err, "Should fail since echo doesn't implement MCP protocol")
+	require.Nil(conn)
+}
+
+// TestGetOrLaunchForSession_StdioSessionPoolHit_DifferentSessions tests that different
+// session IDs for the same backend return different cached connections.
+func TestGetOrLaunchForSession_StdioSessionPoolHit_DifferentSessions(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+
+	mockServer := newMockHTTPMCPServer(t)
+	defer mockServer.Close()
+
+	ctx := context.Background()
+	cfg := newTestConfig(map[string]*config.ServerConfig{
+		"http-helper": {
+			Type: "http",
+			URL:  mockServer.URL,
+		},
+		"stdio-backend": {
+			Type:    "stdio",
+			Command: "docker",
+			Args:    []string{"run", "--rm", "-i", "nonexistent:latest"},
+		},
+	})
+
+	l := New(ctx, cfg)
+	defer l.Close()
+
+	// Get a shared HTTP connection to use as two distinct pool entries
+	httpConn, err := GetOrLaunch(l, "http-helper")
+	require.NoError(err)
+
+	// Pre-populate session pool with the same connection pointer for two different sessions
+	l.sessionPool.Set("stdio-backend", "session-A", httpConn)
+	l.sessionPool.Set("stdio-backend", "session-B", httpConn)
+
+	// Both sessions should return from cache
+	connA, err := GetOrLaunchForSession(l, "stdio-backend", "session-A")
+	require.NoError(err)
+	require.NotNil(connA)
+
+	connB, err := GetOrLaunchForSession(l, "stdio-backend", "session-B")
+	require.NoError(err)
+	require.NotNil(connB)
+
+	// Both return the same underlying connection (we used the same pointer)
+	assert.Equal(httpConn, connA)
+	assert.Equal(httpConn, connB)
+}
+
+// newMockHTTPMCPServer creates a test HTTP server that responds to MCP initialize requests.
+func newMockHTTPMCPServer(t *testing.T) *httptest.Server {
+	t.Helper()
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var req map[string]interface{}
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+		response := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      req["id"],
+			"result": map[string]interface{}{
+				"protocolVersion": "2024-11-05",
+				"capabilities":    map[string]interface{}{},
+				"serverInfo": map[string]interface{}{
+					"name":    "mock-server",
+					"version": "1.0.0",
+				},
+			},
+		}
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(response)
+	}))
+}