@@ -89,6 +89,8 @@ type Client struct {
8989 lifecycleHandlers []SessionLifecycleHandler
9090 typedLifecycleHandlers map [SessionLifecycleEventType ][]SessionLifecycleHandler
9191 lifecycleHandlersMux sync.Mutex
92+ shellProcessMap map [string ]* Session
93+ shellProcessMapMux sync.Mutex
9294 startStopMux sync.RWMutex // protects process and state during start/[force]stop
9395 processDone chan struct {}
9496 processErrorPtr * error
@@ -128,6 +130,7 @@ func NewClient(options *ClientOptions) *Client {
128130 options : opts ,
129131 state : StateDisconnected ,
130132 sessions : make (map [string ]* Session ),
133+ shellProcessMap : make (map [string ]* Session ),
131134 actualHost : "localhost" ,
132135 isExternalServer : false ,
133136 useStdio : true ,
@@ -537,6 +540,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
537540 // Create and register the session before issuing the RPC so that
538541 // events emitted by the CLI (e.g. session.start) are not dropped.
539542 session := newSession (sessionID , c .client , "" )
543+ session .setShellProcessCallbacks (c .registerShellProcess , c .unregisterShellProcess )
540544
541545 session .registerTools (config .Tools )
542546 session .registerPermissionHandler (config .OnPermissionRequest )
@@ -650,6 +654,7 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string,
650654 // Create and register the session before issuing the RPC so that
651655 // events emitted by the CLI (e.g. session.start) are not dropped.
652656 session := newSession (sessionID , c .client , "" )
657+ session .setShellProcessCallbacks (c .registerShellProcess , c .unregisterShellProcess )
653658
654659 session .registerTools (config .Tools )
655660 session .registerPermissionHandler (config .OnPermissionRequest )
@@ -1365,6 +1370,8 @@ func (c *Client) setupNotificationHandler() {
13651370 c .client .SetRequestHandler ("permission.request" , jsonrpc2 .RequestHandlerFor (c .handlePermissionRequestV2 ))
13661371 c .client .SetRequestHandler ("userInput.request" , jsonrpc2 .RequestHandlerFor (c .handleUserInputRequest ))
13671372 c .client .SetRequestHandler ("hooks.invoke" , jsonrpc2 .RequestHandlerFor (c .handleHooksInvoke ))
1373+ c .client .SetRequestHandler ("shell.output" , jsonrpc2 .NotificationHandlerFor (c .handleShellOutput ))
1374+ c .client .SetRequestHandler ("shell.exit" , jsonrpc2 .NotificationHandlerFor (c .handleShellExit ))
13681375}
13691376
13701377func (c * Client ) handleSessionEvent (req sessionEventRequest ) {
@@ -1381,6 +1388,43 @@ func (c *Client) handleSessionEvent(req sessionEventRequest) {
13811388 }
13821389}
13831390
1391+ func (c * Client ) handleShellOutput (notification ShellOutputNotification ) {
1392+ c .shellProcessMapMux .Lock ()
1393+ session , ok := c .shellProcessMap [notification .ProcessID ]
1394+ c .shellProcessMapMux .Unlock ()
1395+
1396+ if ok {
1397+ session .dispatchShellOutput (notification )
1398+ }
1399+ }
1400+
1401+ func (c * Client ) handleShellExit (notification ShellExitNotification ) {
1402+ c .shellProcessMapMux .Lock ()
1403+ session , ok := c .shellProcessMap [notification .ProcessID ]
1404+ c .shellProcessMapMux .Unlock ()
1405+
1406+ if ok {
1407+ session .dispatchShellExit (notification )
1408+ // Clean up the mapping after exit
1409+ c .shellProcessMapMux .Lock ()
1410+ delete (c .shellProcessMap , notification .ProcessID )
1411+ c .shellProcessMapMux .Unlock ()
1412+ session .untrackShellProcess (notification .ProcessID )
1413+ }
1414+ }
1415+
1416+ func (c * Client ) registerShellProcess (processID string , session * Session ) {
1417+ c .shellProcessMapMux .Lock ()
1418+ c .shellProcessMap [processID ] = session
1419+ c .shellProcessMapMux .Unlock ()
1420+ }
1421+
1422+ func (c * Client ) unregisterShellProcess (processID string ) {
1423+ c .shellProcessMapMux .Lock ()
1424+ delete (c .shellProcessMap , processID )
1425+ c .shellProcessMapMux .Unlock ()
1426+ }
1427+
13841428// handleUserInputRequest handles a user input request from the CLI server.
13851429func (c * Client ) handleUserInputRequest (req userInputRequest ) (* userInputResponse , * jsonrpc2.Error ) {
13861430 if req .SessionID == "" || req .Question == "" {
0 commit comments