Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 18 additions & 15 deletions go/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
}
}
Expand Down Expand Up @@ -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)
}
Expand All @@ -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 {
Expand Down
133 changes: 133 additions & 0 deletions go/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
11 changes: 9 additions & 2 deletions go/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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()
Expand Down