diff --git a/internal/launcher/getorlaunchforsession_test.go b/internal/launcher/getorlaunchforsession_test.go index ad5994481..54e1aa0e2 100644 --- a/internal/launcher/getorlaunchforsession_test.go +++ b/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) + })) +}