diff --git a/go/client.go b/go/client.go index fb02897f9..318db293c 100644 --- a/go/client.go +++ b/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,9 +1064,7 @@ 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) } @@ -1065,9 +1072,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses 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 { diff --git a/go/client_test.go b/go/client_test.go index f21442679..86f59c950 100644 --- a/go/client_test.go +++ b/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 { diff --git a/go/session.go b/go/session.go index 600a4bbeb..44bcf0f45 100644 --- a/go/session.go +++ b/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()