Fix CI after shell notification merge

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

Author: jmoseley

dotnet/src/Generated/Rpc.cs

@@ -1325,11 +1325,13 @@ public class SessionRpc
 {
     private readonly JsonRpc _rpc;
     private readonly string _sessionId;
+    private readonly Action<string>? _onShellExec;
 
-    internal SessionRpc(JsonRpc rpc, string sessionId)
+    internal SessionRpc(JsonRpc rpc, string sessionId, Action<string>? onShellExec = null)
     {
         _rpc = rpc;
         _sessionId = sessionId;
+        _onShellExec = onShellExec;
         Model = new ModelApi(rpc, sessionId);
         Mode = new ModeApi(rpc, sessionId);
         Plan = new PlanApi(rpc, sessionId);
@@ -1345,7 +1347,7 @@ internal SessionRpc(JsonRpc rpc, string sessionId)
         Commands = new CommandsApi(rpc, sessionId);
         Ui = new UiApi(rpc, sessionId);
         Permissions = new PermissionsApi(rpc, sessionId);
-        Shell = new ShellApi(rpc, sessionId);
+        Shell = new ShellApi(rpc, sessionId, _onShellExec);
     }
 
     /// <summary>Model APIs.</summary>
@@ -1849,18 +1851,22 @@ public class ShellApi
 {
     private readonly JsonRpc _rpc;
     private readonly string _sessionId;
+    private readonly Action<string>? _onExec;
 
-    internal ShellApi(JsonRpc rpc, string sessionId)
+    internal ShellApi(JsonRpc rpc, string sessionId, Action<string>? onExec = null)
     {
         _rpc = rpc;
         _sessionId = sessionId;
+        _onExec = onExec;
     }
 
     /// <summary>Calls "session.shell.exec".</summary>
     public async Task<SessionShellExecResult> ExecAsync(string command, string? cwd = null, double? timeout = null, CancellationToken cancellationToken = default)
     {
         var request = new SessionShellExecRequest { SessionId = _sessionId, Command = command, Cwd = cwd, Timeout = timeout };
-        return await CopilotClient.InvokeRpcAsync<SessionShellExecResult>(_rpc, "session.shell.exec", [request], cancellationToken);
+        var result = await CopilotClient.InvokeRpcAsync<SessionShellExecResult>(_rpc, "session.shell.exec", [request], cancellationToken);
+        _onExec?.Invoke(result.ProcessId);
+        return result;
     }
 
     /// <summary>Calls "session.shell.kill".</summary>
@@ -1971,4 +1977,4 @@ public async Task<SessionShellKillResult> KillAsync(string processId, SessionShe
 [JsonSerializable(typeof(Tool))]
 [JsonSerializable(typeof(ToolsListRequest))]
 [JsonSerializable(typeof(ToolsListResult))]
-internal partial class RpcJsonContext : JsonSerializerContext;
+internal partial class RpcJsonContext : JsonSerializerContext;

dotnet/src/Session.cs

@@ -93,7 +93,7 @@ public sealed partial class CopilotSession : IAsyncDisposable
     /// <summary>
     /// Gets the typed RPC client for session-scoped methods.
     /// </summary>
-    public SessionRpc Rpc => _sessionRpc ??= new SessionRpc(_rpc, SessionId);
+    public SessionRpc Rpc => _sessionRpc ??= new SessionRpc(_rpc, SessionId, TrackShellProcess);
 
     /// <summary>
     /// Gets the path to the session workspace directory when infinite sessions are enabled.

dotnet/src/Types.cs

@@ -2099,6 +2099,12 @@ public class ShellOutputNotification
     [JsonPropertyName("processId")]
     public string ProcessId { get; set; } = string.Empty;
 
+    /// <summary>
+    /// Identifier of the session that produced this notification, when provided by the runtime.
+    /// </summary>
+    [JsonPropertyName("sessionId")]
+    public string? SessionId { get; set; }
+
     /// <summary>
     /// Which output stream produced this chunk ("stdout" or "stderr").
     /// </summary>
@@ -2124,6 +2130,12 @@ public class ShellExitNotification
     [JsonPropertyName("processId")]
     public string ProcessId { get; set; } = string.Empty;
 
+    /// <summary>
+    /// Identifier of the session that produced this notification, when provided by the runtime.
+    /// </summary>
+    [JsonPropertyName("sessionId")]
+    public string? SessionId { get; set; }
+
     /// <summary>
     /// Process exit code (0 = success).
     /// </summary>

go/client.go

@@ -1466,15 +1466,23 @@ func (c *Client) handleSessionEvent(req sessionEventRequest) {
 }
 
 func (c *Client) handleShellOutput(notification ShellOutputNotification) {
-	session, ok := c.getShellNotificationSession(notification.SessionID, notification.ProcessID)
+	var sessionID string
+	if notification.SessionID != nil {
+		sessionID = *notification.SessionID
+	}
+	session, ok := c.getShellNotificationSession(sessionID, notification.ProcessID)
 
 	if ok {
 		session.dispatchShellOutput(notification)
 	}
 }
 
 func (c *Client) handleShellExit(notification ShellExitNotification) {
-	session, ok := c.getShellNotificationSession(notification.SessionID, notification.ProcessID)
+	var sessionID string
+	if notification.SessionID != nil {
+		sessionID = *notification.SessionID
+	}
+	session, ok := c.getShellNotificationSession(sessionID, notification.ProcessID)
 
 	if ok {
 		session.dispatchShellExit(notification)

go/rpc/generated_rpc.go

@@ -1383,6 +1383,7 @@ func (a *PermissionsRpcApi) HandlePendingPermissionRequest(ctx context.Context,
 type ShellRpcApi struct {
 	client    *jsonrpc2.Client
 	sessionID string
+	onExec    func(string)
 }
 
 func (a *ShellRpcApi) Exec(ctx context.Context, params *SessionShellExecParams) (*SessionShellExecResult, error) {
@@ -1404,6 +1405,9 @@ func (a *ShellRpcApi) Exec(ctx context.Context, params *SessionShellExecParams)
 	if err := json.Unmarshal(raw, &result); err != nil {
 		return nil, err
 	}
+	if a.onExec != nil {
+		a.onExec(result.ProcessID)
+	}
 	return &result, nil
 }
 
@@ -1473,7 +1477,11 @@ func (a *SessionRpc) Log(ctx context.Context, params *SessionLogParams) (*Sessio
 	return &result, nil
 }
 
-func NewSessionRpc(client *jsonrpc2.Client, sessionID string) *SessionRpc {
+func NewSessionRpc(client *jsonrpc2.Client, sessionID string, onShellExec ...func(string)) *SessionRpc {
+	var shellExecHandler func(string)
+	if len(onShellExec) > 0 {
+		shellExecHandler = onShellExec[0]
+	}
 	return &SessionRpc{client: client, sessionID: sessionID,
 		Model:       &ModelRpcApi{client: client, sessionID: sessionID},
 		Mode:        &ModeRpcApi{client: client, sessionID: sessionID},
@@ -1490,6 +1498,6 @@ func NewSessionRpc(client *jsonrpc2.Client, sessionID string) *SessionRpc {
 		Commands:    &CommandsRpcApi{client: client, sessionID: sessionID},
 		Ui:          &UiRpcApi{client: client, sessionID: sessionID},
 		Permissions: &PermissionsRpcApi{client: client, sessionID: sessionID},
-		Shell:       &ShellRpcApi{client: client, sessionID: sessionID},
+		Shell:       &ShellRpcApi{client: client, sessionID: sessionID, onExec: shellExecHandler},
 	}
 }

go/session.go

@@ -112,8 +112,8 @@ func newSession(sessionID string, client *jsonrpc2.Client, workspacePath string)
 		toolHandlers:      make(map[string]ToolHandler),
 		eventCh:           make(chan SessionEvent, 128),
 		trackedProcessIDs: make(map[string]struct{}),
-		RPC:               rpc.NewSessionRpc(client, sessionID),
 	}
+	s.RPC = rpc.NewSessionRpc(client, sessionID, s.trackShellProcess)
 	go s.processEvents()
 	return s
 }
@@ -674,6 +674,16 @@ func (s *Session) dispatchShellExit(notification ShellExitNotification) {
 	}
 }
 
+// trackShellProcess starts tracking a shell process ID.
+func (s *Session) trackShellProcess(processID string) {
+	s.trackedProcessMux.Lock()
+	s.trackedProcessIDs[processID] = struct{}{}
+	s.trackedProcessMux.Unlock()
+	if s.registerShellProc != nil {
+		s.registerShellProc(processID, s)
+	}
+}
+
 // untrackShellProcess stops tracking a shell process ID.
 func (s *Session) untrackShellProcess(processID string) {
 	s.trackedProcessMux.Lock()

go/types.go

@@ -755,6 +755,8 @@ const (
 type ShellOutputNotification struct {
 	// ProcessID is the process identifier returned by shell.exec.
 	ProcessID string `json:"processId"`
+	// SessionID is the optional session identifier for direct routing when provided by the runtime.
+	SessionID *string `json:"sessionId,omitempty"`
 	// Stream indicates which output stream produced this chunk.
 	Stream ShellOutputStream `json:"stream"`
 	// Data is the output data (UTF-8 string).
@@ -766,6 +768,8 @@ type ShellOutputNotification struct {
 type ShellExitNotification struct {
 	// ProcessID is the process identifier returned by shell.exec.
 	ProcessID string `json:"processId"`
+	// SessionID is the optional session identifier for direct routing when provided by the runtime.
+	SessionID *string `json:"sessionId,omitempty"`
 	// ExitCode is the process exit code (0 = success).
 	ExitCode int `json:"exitCode"`
 }

python/copilot/generated/rpc.py

@@ -2836,14 +2836,25 @@ async def handle_pending_permission_request(self, params: SessionPermissionsHand
 
 
 class ShellApi:
-    def __init__(self, client: "JsonRpcClient", session_id: str):
+    def __init__(
+        self,
+        client: "JsonRpcClient",
+        session_id: str,
+        on_exec: Callable[[str], None] | None = None,
+    ):
         self._client = client
         self._session_id = session_id
+        self._on_exec = on_exec
 
     async def exec(self, params: SessionShellExecParams, *, timeout: float | None = None) -> SessionShellExecResult:
         params_dict = {k: v for k, v in params.to_dict().items() if v is not None}
         params_dict["sessionId"] = self._session_id
-        return SessionShellExecResult.from_dict(await self._client.request("session.shell.exec", params_dict, **_timeout_kwargs(timeout)))
+        result = SessionShellExecResult.from_dict(
+            await self._client.request("session.shell.exec", params_dict, **_timeout_kwargs(timeout))
+        )
+        if self._on_exec is not None:
+            self._on_exec(result.process_id)
+        return result
 
     async def kill(self, params: SessionShellKillParams, *, timeout: float | None = None) -> SessionShellKillResult:
         params_dict = {k: v for k, v in params.to_dict().items() if v is not None}
@@ -2853,7 +2864,12 @@ async def kill(self, params: SessionShellKillParams, *, timeout: float | None =
 
 class SessionRpc:
     """Typed session-scoped RPC methods."""
-    def __init__(self, client: "JsonRpcClient", session_id: str):
+    def __init__(
+        self,
+        client: "JsonRpcClient",
+        session_id: str,
+        on_shell_exec: Callable[[str], None] | None = None,
+    ):
         self._client = client
         self._session_id = session_id
         self.model = ModelApi(client, session_id)
@@ -2871,10 +2887,9 @@ def __init__(self, client: "JsonRpcClient", session_id: str):
         self.commands = CommandsApi(client, session_id)
         self.ui = UiApi(client, session_id)
         self.permissions = PermissionsApi(client, session_id)
-        self.shell = ShellApi(client, session_id)
+        self.shell = ShellApi(client, session_id, on_shell_exec)
 
     async def log(self, params: SessionLogParams, *, timeout: float | None = None) -> SessionLogResult:
         params_dict = {k: v for k, v in params.to_dict().items() if v is not None}
         params_dict["sessionId"] = self._session_id
         return SessionLogResult.from_dict(await self._client.request("session.log", params_dict, **_timeout_kwargs(timeout)))
-

python/copilot/session.py

@@ -120,15 +120,7 @@ def __init__(self, session_id: str, client: Any, workspace_path: str | None = No
     def rpc(self) -> SessionRpc:
         """Typed session-scoped RPC methods."""
         if self._rpc is None:
-            self._rpc = SessionRpc(self._client, self.session_id)
-            original_exec = self._rpc.shell.exec
-
-            async def exec_with_tracking(params, *, timeout=None):
-                result = await original_exec(params, timeout=timeout)
-                self._track_shell_process(result.process_id)
-                return result
-
-            self._rpc.shell.exec = exec_with_tracking
+            self._rpc = SessionRpc(self._client, self.session_id, self._track_shell_process)
         return self._rpc
 
     @property

python/pyproject.toml

@@ -39,6 +39,7 @@ dev = [
     "pytest-asyncio>=0.21.0",
     "pytest-timeout>=2.0.0",
     "httpx>=0.24.0",
+    "opentelemetry-api>=1.0.0",
 ]
 telemetry = [
     "opentelemetry-api>=1.0.0",