@@ -91,6 +91,8 @@ type Client struct {
9191 lifecycleHandlers []SessionLifecycleHandler
9292 typedLifecycleHandlers map [SessionLifecycleEventType ][]SessionLifecycleHandler
9393 lifecycleHandlersMux sync.Mutex
94+ shellProcessMap map [string ]* Session
95+ shellProcessMapMux sync.Mutex
9496 startStopMux sync.RWMutex // protects process and state during start/[force]stop
9597 processDone chan struct {}
9698 processErrorPtr * error
@@ -130,6 +132,7 @@ func NewClient(options *ClientOptions) *Client {
130132 options : opts ,
131133 state : StateDisconnected ,
132134 sessions : make (map [string ]* Session ),
135+ shellProcessMap : make (map [string ]* Session ),
133136 actualHost : "localhost" ,
134137 isExternalServer : false ,
135138 useStdio : true ,
@@ -535,6 +538,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
535538 // Create and register the session before issuing the RPC so that
536539 // events emitted by the CLI (e.g. session.start) are not dropped.
537540 session := newSession (sessionID , c .client , "" )
541+ session .setShellProcessCallbacks (c .registerShellProcess , c .unregisterShellProcess )
538542
539543 session .registerTools (config .Tools )
540544 session .registerPermissionHandler (config .OnPermissionRequest )
@@ -648,6 +652,7 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string,
648652 // Create and register the session before issuing the RPC so that
649653 // events emitted by the CLI (e.g. session.start) are not dropped.
650654 session := newSession (sessionID , c .client , "" )
655+ session .setShellProcessCallbacks (c .registerShellProcess , c .unregisterShellProcess )
651656
652657 session .registerTools (config .Tools )
653658 session .registerPermissionHandler (config .OnPermissionRequest )
@@ -1379,6 +1384,8 @@ func (c *Client) setupNotificationHandler() {
13791384 c .client .SetRequestHandler ("permission.request" , jsonrpc2 .RequestHandlerFor (c .handlePermissionRequestV2 ))
13801385 c .client .SetRequestHandler ("userInput.request" , jsonrpc2 .RequestHandlerFor (c .handleUserInputRequest ))
13811386 c .client .SetRequestHandler ("hooks.invoke" , jsonrpc2 .RequestHandlerFor (c .handleHooksInvoke ))
1387+ c .client .SetRequestHandler ("shell.output" , jsonrpc2 .NotificationHandlerFor (c .handleShellOutput ))
1388+ c .client .SetRequestHandler ("shell.exit" , jsonrpc2 .NotificationHandlerFor (c .handleShellExit ))
13821389}
13831390
13841391func (c * Client ) handleSessionEvent (req sessionEventRequest ) {
@@ -1395,6 +1402,43 @@ func (c *Client) handleSessionEvent(req sessionEventRequest) {
13951402 }
13961403}
13971404
1405+ func (c * Client ) handleShellOutput (notification ShellOutputNotification ) {
1406+ c .shellProcessMapMux .Lock ()
1407+ session , ok := c .shellProcessMap [notification .ProcessID ]
1408+ c .shellProcessMapMux .Unlock ()
1409+
1410+ if ok {
1411+ session .dispatchShellOutput (notification )
1412+ }
1413+ }
1414+
1415+ func (c * Client ) handleShellExit (notification ShellExitNotification ) {
1416+ c .shellProcessMapMux .Lock ()
1417+ session , ok := c .shellProcessMap [notification .ProcessID ]
1418+ c .shellProcessMapMux .Unlock ()
1419+
1420+ if ok {
1421+ session .dispatchShellExit (notification )
1422+ // Clean up the mapping after exit
1423+ c .shellProcessMapMux .Lock ()
1424+ delete (c .shellProcessMap , notification .ProcessID )
1425+ c .shellProcessMapMux .Unlock ()
1426+ session .untrackShellProcess (notification .ProcessID )
1427+ }
1428+ }
1429+
1430+ func (c * Client ) registerShellProcess (processID string , session * Session ) {
1431+ c .shellProcessMapMux .Lock ()
1432+ c .shellProcessMap [processID ] = session
1433+ c .shellProcessMapMux .Unlock ()
1434+ }
1435+
1436+ func (c * Client ) unregisterShellProcess (processID string ) {
1437+ c .shellProcessMapMux .Lock ()
1438+ delete (c .shellProcessMap , processID )
1439+ c .shellProcessMapMux .Unlock ()
1440+ }
1441+
13981442// handleUserInputRequest handles a user input request from the CLI server.
13991443func (c * Client ) handleUserInputRequest (req userInputRequest ) (* userInputResponse , * jsonrpc2.Error ) {
14001444 if req .SessionID == "" || req .Question == "" {
0 commit comments