Defer sessionId to server for cloud sessions (#1479)

When a caller creates a session with `cloud` set and does not supply a
`sessionId`, the SDK now omits `sessionId` from the `session.create`
request and lets the CLI/server assign one. The returned id is captured
synchronously when the response is parsed (via an inline response
callback) so that any subsequent `session.event` notification routes to
the registered session without race.

For all other cases (no `cloud`, or caller supplied `sessionId`), the
SDK continues to generate a UUID client-side (when none is supplied)
and pre-register the session BEFORE issuing the RPC. Pre-registration
is required because the CLI may issue session-scoped requests (e.g.
`sessionFs.writeFile` for workspace metadata) during `session.create`
processing, before it has sent the response — without pre-registration
those requests would have no handler. A caller-supplied `sessionId` is
always passed to the server so the server can validate it; mismatched
returned ids fail the call with a clear error.

In addition, the .NET implementation extracts the session-init local
function from CreateSessionAsync into a private InitializeSession
helper that is now shared by both CreateSessionAsync and
ResumeSessionAsync, eliminating duplicated handler-wiring code.

Per-SDK summary:

.NET (dotnet/src/Client.cs): New private InitializeSession helper
shared by CreateSessionAsync and ResumeSessionAsync. CreateSessionAsync
branches on (Cloud != null && SessionId == null): the cloud/no-id path
registers the session lazily from an `onResponseInline` callback; all
other paths pre-register before the RPC. Validates that the server-
returned id matches a caller-supplied id.

Python (python/copilot/client.py): create_session uses the same
conditional. Cloud+no-id installs an `on_response_inline` callback that
registers the session as soon as the response is parsed. Otherwise the
session is pre-registered via uuid.uuid4() (or the caller-supplied id).

Go (go/client.go): createSession uses the same conditional. Cloud+no-id
installs an inline callback that registers from the read loop the
instant the response arrives. Otherwise the session is pre-registered
via uuid.NewString() (or the caller-supplied id).

Rust (rust/src/session.rs): create_session uses the same conditional.
Cloud+no-id installs an InlineResponseCallback that validates the
returned id and registers in the router. Otherwise the session is
pre-registered via uuid::Uuid::new_v4() (or the caller-supplied id).
SessionConfig::into_wire now takes the local session id (or None when
the server is generating one).

Java (java/src/main/java/com/github/copilot/CopilotClient.java):
createSession uses the same conditional. Cloud+no-id defers
registration to the .thenCompose continuation (race-free because
CompletableFuture sync continuations run on the JSON-RPC reader
thread). Otherwise the session is pre-registered via UUID.randomUUID()
(or the caller-supplied id) before the RPC.

Node.js (nodejs/src/client.ts): createSession uses the same
conditional. Cloud+no-id lazy-initializes after the response (race-free
because vscode-jsonrpc dispatches each message via setImmediate, so
microtasks flush before the next message is read). Otherwise the
session is pre-registered via randomUUID() (or the caller-supplied id).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: stephentoub

dotnet/src/Client.cs

@@ -73,11 +73,22 @@ public void StartListening()
     /// <summary>
     /// Sends a JSON-RPC request and waits for the response.
     /// </summary>
-    public async Task<T> InvokeAsync<T>(string method, object?[]? args, CancellationToken cancellationToken)
+    /// <param name="method">The JSON-RPC method name.</param>
+    /// <param name="args">Positional arguments for the call.</param>
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <param name="onResponseInline">
+    /// Optional callback invoked synchronously from the read loop after the
+    /// response is parsed but before the awaiter resumes. Use this when you
+    /// need to mutate client-side state (for example, register a server-assigned
+    /// session id) before any subsequent notification on the same connection is
+    /// dispatched. The callback receives the raw JSON-RPC <c>result</c> element.
+    /// If the callback throws, the exception is propagated to the awaiter.
+    /// </param>
+    public async Task<T> InvokeAsync<T>(string method, object?[]? args, CancellationToken cancellationToken, Action<JsonElement>? onResponseInline = null)
     {
         var timingTimestamp = Stopwatch.GetTimestamp();
         var id = Interlocked.Increment(ref _nextId);
-        var pending = new PendingRequest();
+        var pending = new PendingRequest(onResponseInline);
         _pendingRequests[id] = pending;
 
         CancellationTokenRegistration cancelRegistration = default;
@@ -166,7 +177,7 @@ private void LogInvokeTiming(
     /// </summary>
     public void SetLocalRpcMethod(string methodName, Delegate handler, bool singleObjectParam = false)
     {
-        _methods[methodName] = new MethodRegistration(handler, singleObjectParam);
+        _methods[methodName] = new(handler, singleObjectParam);
     }
 
     /// <inheritdoc />
@@ -447,7 +458,24 @@ private void HandleResponse(JsonElement message, JsonElement idProp)
         }
         else if (message.TryGetProperty("result", out var resultProp))
         {
-            pending.TrySetResult(resultProp.Clone());
+            var cloned = resultProp.Clone();
+            if (pending.OnResultInline is { } inline)
+            {
+                // Run the inline callback synchronously in the read loop so any
+                // state it mutates (e.g. session registration) is visible before
+                // the read loop dispatches the next message.
+                try
+                {
+                    inline(cloned);
+                }
+                catch (Exception ex)
+                {
+                    _logger.LogWarning(ex, "Inline response callback for request {RequestId} threw", id);
+                    pending.TrySetException(ex);
+                    return;
+                }
+            }
+            pending.TrySetResult(cloned);
         }
         else
         {
@@ -765,7 +793,17 @@ await SendMessageAsync(new JsonRpcNotification
         }
     }
 
-    private sealed class PendingRequest() : TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously);
+    private sealed class PendingRequest(Action<JsonElement>? onResultInline = null) : TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously)
+    {
+        /// <summary>
+        /// Optional callback invoked synchronously from the read loop after the
+        /// response is parsed but before the awaiter resumes. Used to perform
+        /// state changes that must happen before any subsequent notification on
+        /// the same connection is dispatched (e.g. registering a session whose
+        /// id was assigned by the server in the response).
+        /// </summary>
+        public Action<JsonElement>? OnResultInline { get; } = onResultInline;
+    }
 
     private static readonly MethodInfo s_taskGetResult = typeof(Task<>).GetProperty(nameof(Task<int>.Result), BindingFlags.Instance | BindingFlags.Public)!.GetMethod!;
     private static readonly MethodInfo s_valueTaskAsTask = typeof(ValueTask<>).GetMethod(nameof(ValueTask<int>.AsTask), BindingFlags.Instance | BindingFlags.Public)!;

dotnet/src/JsonRpc.cs

@@ -307,10 +307,19 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel
 
         private Dictionary<string, object?> CreateSessionResult(JsonElement request)
         {
-            _lastSessionId = request
-                .GetProperty("params")
-                .GetProperty("sessionId")
-                .GetString();
+            string? sessionId = null;
+            if (request.TryGetProperty("params", out var paramsProp)
+                && paramsProp.ValueKind == JsonValueKind.Object
+                && paramsProp.TryGetProperty("sessionId", out var sidProp)
+                && sidProp.ValueKind == JsonValueKind.String)
+            {
+                sessionId = sidProp.GetString();
+            }
+            if (string.IsNullOrEmpty(sessionId))
+            {
+                sessionId = Guid.NewGuid().ToString();
+            }
+            _lastSessionId = sessionId;
 
             return new Dictionary<string, object?>
             {

dotnet/test/Unit/ClientSessionLifetimeTests.cs

@@ -190,7 +190,8 @@ public async Task<T> InvokeAsync<T>(string methodName, object?[]? args, Cancella
                 .GetMethod("InvokeAsync")!
                 .MakeGenericMethod(typeof(T));
 
-            var task = (Task<T>)method.Invoke(_instance, [methodName, args, cancellationToken])!;
+            // Pass null for the optional onResponseInline parameter.
+            var task = (Task<T>)method.Invoke(_instance, [methodName, args, cancellationToken, null])!;
             return await task.ConfigureAwait(false);
         }
 

dotnet/test/Unit/JsonRpcTests.cs

@@ -689,83 +689,161 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 	req.Traceparent = traceparent
 	req.Tracestate = tracestate
 
-	sessionID := config.SessionID
-	if sessionID == "" {
-		sessionID = uuid.New().String()
-	}
-	req.SessionID = sessionID
+	// For cloud sessions, let the CLI/server assign the session id and
+	// register the session lazily once the response arrives. For non-cloud
+	// sessions we generate the id client-side (when the caller didn't
+	// supply one) so the session can be registered BEFORE the RPC — the
+	// CLI may issue session-scoped requests (e.g. sessionFs.writeFile for
+	// workspace metadata) during session.create processing, before it has
+	// sent the response.
+	useServerGeneratedID := config.Cloud != nil && config.SessionID == ""
+	var localSessionID string
+	if useServerGeneratedID {
+		localSessionID = ""
+	} else if config.SessionID != "" {
+		localSessionID = config.SessionID
+	} else {
+		localSessionID = uuid.NewString()
+	}
+	req.SessionID = localSessionID
+
+	// 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
+	// message is dispatched) so notifications for the new session id are
+	// routed to a registered session.
+	initializeSession := func(sessionID string) (*Session, error) {
+		s := newSession(sessionID, c.client, "")
+
+		s.registerTools(config.Tools)
+		s.registerPermissionHandler(config.OnPermissionRequest)
+		if config.OnUserInputRequest != nil {
+			s.registerUserInputHandler(config.OnUserInputRequest)
+		}
+		if config.Hooks != nil {
+			s.registerHooks(config.Hooks)
+		}
+		if transformCallbacks != nil {
+			s.registerTransformCallbacks(transformCallbacks)
+		}
+		if config.OnEvent != nil {
+			s.On(config.OnEvent)
+		}
+		if len(config.Commands) > 0 {
+			s.registerCommands(config.Commands)
+		}
+		if config.OnElicitationRequest != nil {
+			s.registerElicitationHandler(config.OnElicitationRequest)
+		}
+		if config.OnExitPlanModeRequest != nil {
+			s.registerExitPlanModeHandler(config.OnExitPlanModeRequest)
+		}
+		if config.OnAutoModeSwitchRequest != nil {
+			s.registerAutoModeSwitchHandler(config.OnAutoModeSwitchRequest)
+		}
+		if config.CanvasHandler != nil {
+			s.registerCanvasHandler(config.CanvasHandler)
+		}
 
-	// Create and register the session before issuing the RPC so that
-	// events emitted by the CLI (e.g. session.start) are not dropped.
-	session := newSession(sessionID, c.client, "")
+		c.sessionsMux.Lock()
+		c.sessions[sessionID] = s
+		c.sessionsMux.Unlock()
 
-	session.registerTools(config.Tools)
-	session.registerPermissionHandler(config.OnPermissionRequest)
-	if config.OnUserInputRequest != nil {
-		session.registerUserInputHandler(config.OnUserInputRequest)
-	}
-	if config.Hooks != nil {
-		session.registerHooks(config.Hooks)
-	}
-	if transformCallbacks != nil {
-		session.registerTransformCallbacks(transformCallbacks)
-	}
-	if config.OnEvent != nil {
-		session.On(config.OnEvent)
-	}
-	if len(config.Commands) > 0 {
-		session.registerCommands(config.Commands)
-	}
-	if config.OnElicitationRequest != nil {
-		session.registerElicitationHandler(config.OnElicitationRequest)
-	}
-	if config.OnExitPlanModeRequest != nil {
-		session.registerExitPlanModeHandler(config.OnExitPlanModeRequest)
-	}
-	if config.OnAutoModeSwitchRequest != nil {
-		session.registerAutoModeSwitchHandler(config.OnAutoModeSwitchRequest)
-	}
-	if config.CanvasHandler != nil {
-		session.registerCanvasHandler(config.CanvasHandler)
+		if c.options.SessionFs != nil {
+			if config.CreateSessionFsProvider == nil {
+				c.sessionsMux.Lock()
+				delete(c.sessions, sessionID)
+				c.sessionsMux.Unlock()
+				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()
+					return nil, fmt.Errorf("SessionFs capabilities declare SQLite support but the provider does not implement SessionFsSqliteProvider")
+				}
+			}
+			s.clientSessionApis.SessionFs = newSessionFsAdapter(provider)
+		}
+		return s, nil
 	}
 
-	c.sessionsMux.Lock()
-	c.sessions[sessionID] = session
-	c.sessionsMux.Unlock()
+	var session *Session
+	var registeredSessionID string
 
-	if c.options.SessionFs != nil {
-		if config.CreateSessionFsProvider == nil {
-			c.sessionsMux.Lock()
-			delete(c.sessions, sessionID)
-			c.sessionsMux.Unlock()
-			return nil, fmt.Errorf("CreateSessionFsProvider is required in session config when SessionFs is enabled in client options")
+	// Pre-register non-cloud sessions BEFORE issuing the RPC so any
+	// session-scoped requests the CLI emits during session.create processing
+	// (e.g. sessionFs.writeFile for workspace metadata) can be routed to the
+	// correct handlers.
+	if localSessionID != "" {
+		s, err := initializeSession(localSessionID)
+		if err != nil {
+			return nil, err
 		}
-		provider := config.CreateSessionFsProvider(session)
-		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()
-				return nil, fmt.Errorf("SessionFs capabilities declare SQLite support but the provider does not implement SessionFsSqliteProvider")
+		session = s
+		registeredSessionID = localSessionID
+	}
+
+	// For the server-assigned (cloud) path, register the session
+	// synchronously from the read loop the instant the response arrives,
+	// before the read loop dispatches the next message. Without this hook
+	// the awaiter goroutine may not run until after the read loop has
+	// dispatched the first session.event notification, which would be
+	// silently dropped because the session id isn't yet in the lookup
+	// table. Non-cloud sessions are already registered above.
+	var inlineCb func(raw json.RawMessage) error
+	if session == nil {
+		inlineCb = func(raw json.RawMessage) error {
+			var early struct {
+				SessionID string `json:"sessionId"`
+			}
+			if err := json.Unmarshal(raw, &early); err != nil {
+				return fmt.Errorf("failed to parse sessionId from response: %w", err)
+			}
+			if early.SessionID == "" {
+				return fmt.Errorf("session.create response did not include a sessionId")
 			}
+			s, err := initializeSession(early.SessionID)
+			if err != nil {
+				return err
+			}
+			session = s
+			registeredSessionID = early.SessionID
+			return nil
 		}
-		session.clientSessionApis.SessionFs = newSessionFsAdapter(provider)
 	}
 
-	result, err := c.client.Request("session.create", req)
+	result, err := c.client.RequestWithInlineResponse("session.create", req, inlineCb)
 	if err != nil {
-		c.sessionsMux.Lock()
-		delete(c.sessions, sessionID)
-		c.sessionsMux.Unlock()
+		if registeredSessionID != "" {
+			c.sessionsMux.Lock()
+			delete(c.sessions, registeredSessionID)
+			c.sessionsMux.Unlock()
+		}
 		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()
+		}
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	if session == nil {
+		return nil, fmt.Errorf("session.create response did not include a sessionId")
+	}
+
+	if localSessionID != "" && response.SessionID != "" && response.SessionID != localSessionID {
 		c.sessionsMux.Lock()
-		delete(c.sessions, sessionID)
+		delete(c.sessions, registeredSessionID)
 		c.sessionsMux.Unlock()
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+		return nil, fmt.Errorf("session.create returned sessionId %s but the caller requested %s", response.SessionID, localSessionID)
 	}
 
 	session.workspacePath = response.WorkspacePath