fix(go): close failed session event loops

Author: green3sf

go/client.go

@@ -935,6 +935,19 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 	}
 	req.SessionID = localSessionID
 
+	// unregisterSession removes only the session created by this call and stops
+	// its event consumer. The latter is essential on CreateSession error paths:
+	// newSession starts processEvents eagerly, and no caller receives the failed
+	// session to disconnect it.
+	unregisterSession := func(sessionID string, s *Session) {
+		c.sessionsMux.Lock()
+		if c.sessions[sessionID] == s {
+			delete(c.sessions, sessionID)
+		}
+		c.sessionsMux.Unlock()
+		s.closeEventChannel()
+	}
+
 	// initializeSession creates the session, wires up handlers, and registers
 	// it in the sessions map. Invoked from the read loop the instant the
 	// session.create response arrives (synchronously, before the next
@@ -988,17 +1001,13 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 
 		if c.options.SessionFS != nil {
 			if config.CreateSessionFSProvider == nil {
-				c.sessionsMux.Lock()
-				delete(c.sessions, sessionID)
-				c.sessionsMux.Unlock()
+				unregisterSession(sessionID, s)
 				return nil, fmt.Errorf("CreateSessionFSProvider is required in session config when SessionFS is enabled in client options")
 			}
 			provider := config.CreateSessionFSProvider(s)
 			if c.options.SessionFS.Capabilities != nil && c.options.SessionFS.Capabilities.Sqlite {
 				if _, ok := provider.(SessionFSSqliteProvider); !ok {
-					c.sessionsMux.Lock()
-					delete(c.sessions, sessionID)
-					c.sessionsMux.Unlock()
+					unregisterSession(sessionID, s)
 					return nil, fmt.Errorf("SessionFS capabilities declare SQLite support but the provider does not implement SessionFSSqliteProvider")
 				}
 			}
@@ -1055,19 +1064,15 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 	result, err := c.client.RequestWithInlineResponse(ctx, "session.create", req, inlineCb)
 	if err != nil {
 		if registeredSessionID != "" {
-			c.sessionsMux.Lock()
-			delete(c.sessions, registeredSessionID)
-			c.sessionsMux.Unlock()
+			unregisterSession(registeredSessionID, session)
 		}
 		return nil, fmt.Errorf("failed to create session: %w", err)
 	}
 
 	var response createSessionResponse
 	if err := json.Unmarshal(result, &response); err != nil {
 		if registeredSessionID != "" {
-			c.sessionsMux.Lock()
-			delete(c.sessions, registeredSessionID)
-			c.sessionsMux.Unlock()
+			unregisterSession(registeredSessionID, session)
 		}
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
 	}
@@ -1077,9 +1082,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 	}
 
 	if localSessionID != "" && response.SessionID != "" && response.SessionID != localSessionID {
-		c.sessionsMux.Lock()
-		delete(c.sessions, registeredSessionID)
-		c.sessionsMux.Unlock()
+		unregisterSession(registeredSessionID, session)
 		return nil, fmt.Errorf("session.create returned sessionId %s but the caller requested %s", response.SessionID, localSessionID)
 	}
 	if config.OnMCPAuthRequest != nil {

go/client_test.go

@@ -689,6 +689,139 @@ func sessionIDFromParams(t *testing.T, params json.RawMessage) string {
 	return decoded.SessionID
 }
 
+func TestClient_CreateSessionFailureClosesRegisteredSession(t *testing.T) {
+	tests := []struct {
+		name       string
+		response   func(string) (json.RawMessage, *jsonrpc2.Error)
+		wantErrSub string
+	}{
+		{
+			name: "RPC failure",
+			response: func(string) (json.RawMessage, *jsonrpc2.Error) {
+				return nil, &jsonrpc2.Error{Code: -32000, Message: "session creation failed"}
+			},
+			wantErrSub: "failed to create session",
+		},
+		{
+			name: "invalid response",
+			response: func(string) (json.RawMessage, *jsonrpc2.Error) {
+				return json.RawMessage(`"invalid"`), nil
+			},
+			wantErrSub: "failed to unmarshal response",
+		},
+		{
+			name: "session ID mismatch",
+			response: func(string) (json.RawMessage, *jsonrpc2.Error) {
+				return json.RawMessage(`{"sessionId":"different-session"}`), nil
+			},
+			wantErrSub: "but the caller requested failed-session",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rpcClient, server, _ := newRuntimeShutdownRpcPair(t)
+			t.Cleanup(server.Stop)
+			client := &Client{
+				client:   rpcClient,
+				RPC:      rpc.NewServerRPC(rpcClient),
+				sessions: make(map[string]*Session),
+			}
+
+			captured := make(chan *Session, 1)
+			server.SetRequestHandler("session.create", func(params json.RawMessage) (json.RawMessage, *jsonrpc2.Error) {
+				sessionID := sessionIDFromParams(t, params)
+				client.sessionsMux.Lock()
+				session := client.sessions[sessionID]
+				client.sessionsMux.Unlock()
+				captured <- session
+				return tt.response(sessionID)
+			})
+
+			_, err := client.CreateSession(t.Context(), &SessionConfig{SessionID: "failed-session"})
+			if err == nil || !strings.Contains(err.Error(), tt.wantErrSub) {
+				t.Fatalf("CreateSession error = %v, want substring %q", err, tt.wantErrSub)
+			}
+
+			session := <-captured
+			if session == nil {
+				t.Fatal("session was not registered before session.create")
+			}
+			assertSessionEventChannelClosed(t, session)
+			assertSessionNotRegistered(t, client, "failed-session")
+		})
+	}
+}
+
+func TestClient_CreateSessionInitializationFailureClosesRegisteredSession(t *testing.T) {
+	rpcClient, server, _ := newRuntimeShutdownRpcPair(t)
+	t.Cleanup(server.Stop)
+	client := &Client{
+		client:   rpcClient,
+		RPC:      rpc.NewServerRPC(rpcClient),
+		sessions: make(map[string]*Session),
+		options: ClientOptions{SessionFS: &SessionFSConfig{
+			InitialWorkingDirectory: "/",
+			SessionStatePath:        "/session-state",
+			Conventions:             rpc.SessionFSSetProviderConventionsPosix,
+			Capabilities:            &SessionFSCapabilities{Sqlite: true},
+		}},
+	}
+
+	var captured *Session
+	_, err := client.CreateSession(t.Context(), &SessionConfig{
+		SessionID: "failed-session-fs",
+		CreateSessionFSProvider: func(session *Session) SessionFSProvider {
+			captured = session
+			return noSQLiteSessionFSProvider{}
+		},
+	})
+	if err == nil || !strings.Contains(err.Error(), "does not implement SessionFSSqliteProvider") {
+		t.Fatalf("CreateSession error = %v, want SQLite provider validation error", err)
+	}
+	if captured == nil {
+		t.Fatal("CreateSessionFSProvider did not receive the registered session")
+	}
+	assertSessionEventChannelClosed(t, captured)
+	assertSessionNotRegistered(t, client, "failed-session-fs")
+}
+
+func assertSessionEventChannelClosed(t *testing.T, session *Session) {
+	t.Helper()
+	select {
+	case _, ok := <-session.eventCh:
+		if ok {
+			t.Fatal("session event channel is still open")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for session event channel to close")
+	}
+}
+
+func assertSessionNotRegistered(t *testing.T, client *Client, sessionID string) {
+	t.Helper()
+	client.sessionsMux.Lock()
+	defer client.sessionsMux.Unlock()
+	if _, ok := client.sessions[sessionID]; ok {
+		t.Fatalf("session %q is still registered", sessionID)
+	}
+}
+
+type noSQLiteSessionFSProvider struct{}
+
+func (noSQLiteSessionFSProvider) ReadFile(string) (string, error)         { return "", nil }
+func (noSQLiteSessionFSProvider) WriteFile(string, string, *int) error    { return nil }
+func (noSQLiteSessionFSProvider) AppendFile(string, string, *int) error   { return nil }
+func (noSQLiteSessionFSProvider) Exists(string) (bool, error)             { return false, nil }
+func (noSQLiteSessionFSProvider) Stat(string) (*SessionFSFileInfo, error) { return nil, nil }
+func (noSQLiteSessionFSProvider) MakeDirectory(string, bool, *int) error  { return nil }
+func (noSQLiteSessionFSProvider) ReadDirectory(string) ([]string, error)  { return nil, nil }
+func (noSQLiteSessionFSProvider) ReadDirectoryWithTypes(string) ([]rpc.SessionFSReaddirWithTypesEntry, error) {
+	return nil, nil
+}
+func (noSQLiteSessionFSProvider) Remove(string, bool, bool) error { return nil }
+func (noSQLiteSessionFSProvider) Rename(string, string) error     { return nil }
+
 func assertRuntimeShutdownNotCalled(t *testing.T, shutdownCalled <-chan struct{}) {
 	t.Helper()
 	select {

go/session.go

@@ -96,7 +96,7 @@ type Session struct {
 	// eventCh serializes user event handler dispatch. dispatchEvent enqueues;
 	// a single goroutine (processEvents) dequeues and invokes handlers in FIFO order.
 	eventCh   chan SessionEvent
-	closeOnce sync.Once // guards eventCh close so Disconnect is safe to call more than once
+	closeOnce sync.Once // guards eventCh close across disconnect and failed session creation
 
 	// RPC provides typed session-scoped RPC methods.
 	RPC *rpc.SessionRPC
@@ -1421,6 +1421,13 @@ func (s *Session) processEvents() {
 	}
 }
 
+// closeEventChannel stops the session event consumer without making an RPC.
+// CreateSession uses this when a locally registered session fails before it can
+// be returned to the caller.
+func (s *Session) closeEventChannel() {
+	s.closeOnce.Do(func() { close(s.eventCh) })
+}
+
 // handleBroadcastEvent handles broadcast request events by executing local handlers
 // and responding via RPC. This implements the protocol v3 broadcast model where tool
 // calls and permission requests are broadcast as session events to all clients.
@@ -1726,7 +1733,7 @@ func (s *Session) Disconnect() error {
 		return fmt.Errorf("failed to disconnect session: %w", err)
 	}
 
-	s.closeOnce.Do(func() { close(s.eventCh) })
+	s.closeEventChannel()
 
 	// Clear handlers
 	s.handlerMutex.Lock()