diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 455cecdba..0b4f7068b 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -72,6 +72,12 @@ public sealed partial class CopilotSession : IAsyncDisposable private SessionRpc? _sessionRpc; private int _isDisposed; + /// + /// Callback invoked when this session is disposed, so the owning client can + /// remove it from its session map. Set by . + /// + internal Action? OnDisposed { get; set; } + /// /// Channel that serializes event dispatch. enqueues; /// a single background consumer () dequeues and @@ -327,6 +333,10 @@ public IDisposable On(SessionEventHandler handler) /// internal void DispatchEvent(SessionEvent sessionEvent) { + // Guard: do not dispatch events to a disposed session. + if (Volatile.Read(ref _isDisposed) == 1) + return; + // Fire broadcast work concurrently (fire-and-forget with error logging). // This is done outside the channel so broadcast handlers don't block the // consumer loop — important when a secondary client's handler intentionally @@ -1186,6 +1196,10 @@ public async ValueTask DisposeAsync() return; } + // Remove from client's session map before the RPC so that events + // arriving during the round-trip are never dispatched (fail-closed). + OnDisposed?.Invoke(SessionId); + _eventChannel.Writer.TryComplete(); try diff --git a/dotnet/test/PermissionTests.cs b/dotnet/test/PermissionTests.cs index 3ab36dad1..e26e5a585 100644 --- a/dotnet/test/PermissionTests.cs +++ b/dotnet/test/PermissionTests.cs @@ -248,4 +248,45 @@ await session.SendAsync(new MessageOptions Assert.True(receivedToolCallId, "Should have received toolCallId in permission request"); } + + /// + /// Validates that disposing a session does not affect subsequent sessions on the same client. + /// After session A is disposed, session B should handle permissions and complete tool calls normally. + /// This exercises the OnDisposed cleanup path that removes disposed sessions from the client map. + /// + [Fact] + public async Task Should_Handle_Permission_After_Prior_Session_Disposed() + { + // Create and use session A + var session1 = await CreateSessionAsync(new SessionConfig + { + OnPermissionRequest = (request, invocation) => + Task.FromResult(new PermissionRequestResult { Kind = PermissionRequestResultKind.Approved }) + }); + + await session1.SendAndWaitAsync(new MessageOptions { Prompt = "What is 1+1?" }); + + // Dispose session A — exercises the OnDisposed callback that removes it from the client map + await session1.DisposeAsync(); + + // Create session B on the same client — should work normally + var session2PermissionReceived = false; + var session2 = await CreateSessionAsync(new SessionConfig + { + OnPermissionRequest = (request, invocation) => + { + session2PermissionReceived = true; + return Task.FromResult(new PermissionRequestResult { Kind = PermissionRequestResultKind.Approved }); + } + }); + + // Session B should handle permissions and complete normally + await session2.SendAndWaitAsync(new MessageOptions + { + Prompt = "Run 'echo hello' for me" + }, timeout: TimeSpan.FromSeconds(15)); + + Assert.True(session2PermissionReceived, + "Session B's permission handler should fire normally after session A was disposed."); + } } diff --git a/go/client.go b/go/client.go index 0c72e963f..5d5d72860 100644 --- a/go/client.go +++ b/go/client.go @@ -672,6 +672,11 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses c.sessionsMux.Lock() c.sessions[sessionID] = session c.sessionsMux.Unlock() + session.onDisposed = func(id string) { + c.sessionsMux.Lock() + delete(c.sessions, id) + c.sessionsMux.Unlock() + } if c.options.SessionFs != nil { if config.CreateSessionFsHandler == nil { @@ -831,6 +836,11 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, c.sessionsMux.Lock() c.sessions[sessionID] = session c.sessionsMux.Unlock() + session.onDisposed = func(id string) { + c.sessionsMux.Lock() + delete(c.sessions, id) + c.sessionsMux.Unlock() + } if c.options.SessionFs != nil { if config.CreateSessionFsHandler == nil { diff --git a/go/session.go b/go/session.go index 99256856d..b3ea7222f 100644 --- a/go/session.go +++ b/go/session.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "sync" + "sync/atomic" "time" "github.com/github/copilot-sdk/go/internal/jsonrpc2" @@ -74,6 +75,14 @@ type Session struct { capabilities SessionCapabilities capabilitiesMu sync.RWMutex + // disconnected is set atomically to 1 by Disconnect. Guards dispatchEvent + // so that events arriving after disconnect are silently dropped. + disconnected atomic.Int32 + + // onDisposed is called by Disconnect so the owning Client can remove the + // session from its sessions map before the RPC round-trip. Set by Client. + onDisposed func(sessionID string) + // eventCh serializes user event handler dispatch. dispatchEvent enqueues; // a single goroutine (processEvents) dequeues and invokes handlers in FIFO order. eventCh chan SessionEvent @@ -841,6 +850,11 @@ func fromRPCElicitationResult(r *rpc.UIElicitationResponse) *ElicitationResult { // are delivered by a single consumer goroutine (processEvents), guaranteeing // serial, FIFO dispatch without blocking the read loop. func (s *Session) dispatchEvent(event SessionEvent) { + // Guard: do not dispatch events to a disconnected session. + if s.disconnected.Load() != 0 { + return + } + go s.handleBroadcastEvent(event) // Send to the event channel in a closure with a recover guard. @@ -908,6 +922,8 @@ func (s *Session) handleBroadcastEvent(event SessionEvent) { } handler := s.getPermissionHandler() if handler == nil { + // No handler registered (e.g. session disconnected or single-client + // headless mode). Another client will handle it. return } s.executePermissionAndRespond(d.RequestID, d.PermissionRequest, handler) @@ -1120,33 +1136,47 @@ func (s *Session) GetMessages(ctx context.Context) ([]SessionEvent, error) { // log.Printf("Failed to disconnect session: %v", err) // } func (s *Session) Disconnect() error { - _, err := s.client.Request("session.destroy", sessionDestroyRequest{SessionID: s.SessionID}) - if err != nil { - return fmt.Errorf("failed to disconnect session: %w", err) + // Mark disconnected before the RPC round-trip so that events arriving + // during the round-trip are never dispatched (fail-closed). + s.disconnected.Store(1) + + // Remove from client's session map before the RPC so the client stops + // routing events to this session. + if s.onDisposed != nil { + s.onDisposed(s.SessionID) } - s.closeOnce.Do(func() { close(s.eventCh) }) + // Ensure cleanup always runs even if session.destroy fails + defer func() { + s.closeOnce.Do(func() { close(s.eventCh) }) - // Clear handlers - s.handlerMutex.Lock() - s.handlers = nil - s.handlerMutex.Unlock() + // Clear handlers + s.handlerMutex.Lock() + s.handlers = nil + s.handlerMutex.Unlock() - s.toolHandlersM.Lock() - s.toolHandlers = nil - s.toolHandlersM.Unlock() + s.toolHandlersM.Lock() + s.toolHandlers = nil + s.toolHandlersM.Unlock() - s.permissionMux.Lock() - s.permissionHandler = nil - s.permissionMux.Unlock() + s.permissionMux.Lock() + s.permissionHandler = nil + s.permissionMux.Unlock() - s.commandHandlersMu.Lock() - s.commandHandlers = nil - s.commandHandlersMu.Unlock() + s.commandHandlersMu.Lock() + s.commandHandlers = nil + s.commandHandlersMu.Unlock() - s.elicitationMu.Lock() - s.elicitationHandler = nil - s.elicitationMu.Unlock() + s.elicitationMu.Lock() + s.elicitationHandler = nil + s.elicitationMu.Unlock() + }() + + // Try to destroy session on server + _, err := s.client.Request("session.destroy", sessionDestroyRequest{SessionID: s.SessionID}) + if err != nil { + return fmt.Errorf("failed to disconnect session: %w", err) + } return nil } diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index 0ef19038f..0a9218df5 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -711,6 +711,7 @@ export class CopilotClient { session.on(config.onEvent); } this.sessions.set(sessionId, session); + session.onDisposed = (id) => this.sessions.delete(id); if (this.sessionFsConfig) { if (config.createSessionFsHandler) { session.clientSessionApis.sessionFs = createSessionFsAdapter( @@ -852,6 +853,7 @@ export class CopilotClient { session.on(config.onEvent); } this.sessions.set(sessionId, session); + session.onDisposed = (id) => this.sessions.delete(id); if (this.sessionFsConfig) { if (config.createSessionFsHandler) { session.clientSessionApis.sessionFs = createSessionFsAdapter( diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index eae4cab94..fe328cdc0 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -89,6 +89,14 @@ export class CopilotSession { private _rpc: ReturnType | null = null; private traceContextProvider?: TraceContextProvider; private _capabilities: SessionCapabilities = {}; + private _disposed = false; + + /** + * Callback invoked when this session is disposed, so the owning client can + * remove it from its session map. Set by {@link CopilotClient}. + * @internal + */ + onDisposed?: (sessionId: string) => void; /** @internal Client session API handlers, populated by CopilotClient during create/resume. */ clientSessionApis: ClientSessionApiHandlers = {}; @@ -353,6 +361,9 @@ export class CopilotSession { * @internal This method is for internal use by the SDK. */ _dispatchEvent(event: SessionEvent): void { + // Guard: do not dispatch events to a disposed session. + if (this._disposed) return; + // Handle broadcast request events internally (fire-and-forget) this._handleBroadcastEvent(event); @@ -421,6 +432,8 @@ export class CopilotSession { if (this.permissionHandler) { void this._executePermissionAndRespond(requestId, permissionRequest); } + // No handler registered (e.g. session disposed or single-client headless mode). + // Another client will handle it (multi-client scenario) or the request will timeout. } else if (event.type === "command.execute") { const { requestId, commandName, command, args } = event.data as { requestId: string; @@ -965,13 +978,29 @@ export class CopilotSession { * ``` */ async disconnect(): Promise { - await this.connection.sendRequest("session.destroy", { - sessionId: this.sessionId, - }); - this.eventHandlers.clear(); - this.typedEventHandlers.clear(); - this.toolHandlers.clear(); - this.permissionHandler = undefined; + if (this._disposed) return; + this._disposed = true; + + // Remove from client's session map before the RPC so that events + // arriving during the round-trip are never dispatched (fail-closed). + this.onDisposed?.(this.sessionId); + + try { + await this.connection.sendRequest("session.destroy", { + sessionId: this.sessionId, + }); + } catch { + // Connection already closed + } finally { + // Always clean up handlers + this.eventHandlers.clear(); + this.typedEventHandlers.clear(); + this.toolHandlers.clear(); + this.commandHandlers.clear(); + this.permissionHandler = undefined; + this.elicitationHandler = undefined; + this.userInputHandler = undefined; + } } /** diff --git a/test/snapshots/permissions/should_handle_permission_after_prior_session_disposed.yaml b/test/snapshots/permissions/should_handle_permission_after_prior_session_disposed.yaml new file mode 100644 index 000000000..b81eaec24 --- /dev/null +++ b/test/snapshots/permissions/should_handle_permission_after_prior_session_disposed.yaml @@ -0,0 +1,37 @@ +models: + - claude-sonnet-4.5 +conversations: + - messages: + - role: system + content: ${system} + - role: user + content: What is 1+1? + - role: assistant + content: 1+1 = 2 + - messages: + - role: system + content: ${system} + - role: user + content: Run 'echo hello' for me + - role: assistant + tool_calls: + - id: toolcall_0 + type: function + function: + name: report_intent + arguments: '{"intent":"Running echo command"}' + - id: toolcall_1 + type: function + function: + name: ${shell} + arguments: '{"description":"Run echo hello","command":"echo hello"}' + - role: tool + tool_call_id: toolcall_0 + content: Intent logged + - role: tool + tool_call_id: toolcall_1 + content: |- + hello + + - role: assistant + content: "The command ran successfully and output: **hello**"