@@ -232,7 +232,7 @@ func TestSession_CommandRouting(t *testing.T) {
232232 },
233233 })
234234
235- // Simulate the dispatch — executeCommandAndRespond will fail on RPC (nil client)
235+ // Simulate the dispatch ΓÇö executeCommandAndRespond will fail on RPC (nil client)
236236 // but the handler will still be invoked. We test routing only.
237237 _ , ok := session .getCommandHandler ("deploy" )
238238 if ! ok {
@@ -574,7 +574,7 @@ func TestSession_ElicitationRequestSchema(t *testing.T) {
574574 "type" : "object" ,
575575 "properties" : properties ,
576576 }
577- // Simulate: if len(schema.Required) > 0 { ... } — with empty required
577+ // Simulate: if len(schema.Required) > 0 { ... } ΓÇö with empty required
578578 var required []string
579579 if len (required ) > 0 {
580580 requestedSchema ["required" ] = required
@@ -585,3 +585,251 @@ func TestSession_ElicitationRequestSchema(t *testing.T) {
585585 }
586586 })
587587}
588+
589+ // newTestSessionWithFakeClient creates a Session backed by an in-process fake
590+ // JSON-RPC 2.0 server that immediately acknowledges any "session.send" call.
591+ // The returned requestReceived channel is signalled once per incoming request,
592+ // so tests can synchronize dispatch of follow-up events.
593+ // The cleanup function must be deferred by the caller.
594+ func newTestSessionWithFakeClient (t * testing.T ) (sess * Session , requestReceived <- chan struct {}, cleanup func ()) {
595+ t .Helper ()
596+
597+ stdinR , stdinW := io .Pipe ()
598+ stdoutR , stdoutW := io .Pipe ()
599+
600+ client := jsonrpc2 .NewClient (stdinW , stdoutR )
601+ client .Start ()
602+
603+ reqCh := make (chan struct {}, 4 )
604+
605+ go func () {
606+ reader := bufio .NewReader (stdinR )
607+ buf := make ([]byte , 8192 )
608+ for {
609+ var contentLength int
610+ for {
611+ line , err := reader .ReadString ('\n' )
612+ if err != nil {
613+ return
614+ }
615+ line = strings .TrimSpace (line )
616+ if line == "" {
617+ break
618+ }
619+ fmt .Sscanf (line , "Content-Length: %d" , & contentLength )
620+ }
621+ if contentLength == 0 {
622+ continue
623+ }
624+ if contentLength > len (buf ) {
625+ buf = make ([]byte , contentLength )
626+ }
627+ if _ , err := io .ReadFull (reader , buf [:contentLength ]); err != nil {
628+ return
629+ }
630+ var req struct {
631+ ID json.RawMessage `json:"id"`
632+ }
633+ json .Unmarshal (buf [:contentLength ], & req ) //nolint:errcheck
634+
635+ // Signal that a request was received before sending the response,
636+ // so tests can dispatch events once Send() is in-flight.
637+ select {
638+ case reqCh <- struct {}{}:
639+ default :
640+ }
641+
642+ result := `{"messageId":"test-msg"}`
643+ body := fmt .Sprintf (`{"jsonrpc":"2.0","id":%s,"result":%s}` , req .ID , result )
644+ header := fmt .Sprintf ("Content-Length: %d\r \n \r \n " , len (body ))
645+ stdoutW .Write ([]byte (header + body )) //nolint:errcheck
646+ }
647+ }()
648+
649+ s := & Session {
650+ SessionID : "test-session" ,
651+ handlers : make ([]sessionHandler , 0 ),
652+ eventCh : make (chan SessionEvent , 128 ),
653+ client : client ,
654+ }
655+ go s .processEvents ()
656+
657+ cleanupFn := func () {
658+ stdoutW .Close ()
659+ stdinW .Close ()
660+ client .Stop ()
661+ close (s .eventCh )
662+ }
663+
664+ return s , reqCh , cleanupFn
665+ }
666+
667+ // TestSendAndWait_BackgroundTasks verifies the fix for PolyPilot#299:
668+ // SendAndWait must NOT resolve when session.idle carries active background tasks.
669+ func TestSendAndWait_BackgroundTasks (t * testing.T ) {
670+ t .Run ("does not resolve when session.idle has active background agents" , func (t * testing.T ) {
671+ session , requestReceived , cleanup := newTestSessionWithFakeClient (t )
672+ defer cleanup ()
673+
674+ ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
675+ defer cancel ()
676+
677+ var resolved atomic.Bool
678+ done := make (chan struct {})
679+ go func () {
680+ defer close (done )
681+ session .SendAndWait (ctx , MessageOptions {Prompt : "test" }) //nolint:errcheck
682+ resolved .Store (true )
683+ }()
684+
685+ // Wait for the fake server to receive the session.send request.
686+ select {
687+ case <- requestReceived :
688+ case <- time .After (2 * time .Second ):
689+ t .Fatal ("timeout waiting for session.send request" )
690+ }
691+
692+ // Dispatch session.idle WITH active background agents — must NOT resolve.
693+ session .dispatchEvent (SessionEvent {
694+ Type : SessionEventTypeSessionIdle ,
695+ Data : Data {
696+ BackgroundTasks : & BackgroundTasks {
697+ Agents : []BackgroundTasksAgent {{AgentID : "bg-1" , AgentType : "worker" }},
698+ Shells : []Shell {},
699+ },
700+ },
701+ })
702+
703+ time .Sleep (100 * time .Millisecond )
704+ if resolved .Load () {
705+ t .Error ("BUG #299: SendAndWait resolved prematurely while background agents were active" )
706+ }
707+
708+ // Dispatch a clean idle (no background tasks) — now it SHOULD resolve.
709+ session .dispatchEvent (SessionEvent {
710+ Type : SessionEventTypeSessionIdle ,
711+ Data : Data {},
712+ })
713+
714+ select {
715+ case <- done :
716+ case <- time .After (2 * time .Second ):
717+ t .Fatal ("SendAndWait did not resolve after clean session.idle" )
718+ }
719+ if ! resolved .Load () {
720+ t .Error ("expected SendAndWait to resolve after clean session.idle" )
721+ }
722+ })
723+
724+ t .Run ("does not resolve when session.idle has active background shells" , func (t * testing.T ) {
725+ session , requestReceived , cleanup := newTestSessionWithFakeClient (t )
726+ defer cleanup ()
727+
728+ ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
729+ defer cancel ()
730+
731+ var resolved atomic.Bool
732+ done := make (chan struct {})
733+ go func () {
734+ defer close (done )
735+ session .SendAndWait (ctx , MessageOptions {Prompt : "test" }) //nolint:errcheck
736+ resolved .Store (true )
737+ }()
738+
739+ select {
740+ case <- requestReceived :
741+ case <- time .After (2 * time .Second ):
742+ t .Fatal ("timeout waiting for session.send request" )
743+ }
744+
745+ // Dispatch session.idle WITH active background shell — must NOT resolve.
746+ session .dispatchEvent (SessionEvent {
747+ Type : SessionEventTypeSessionIdle ,
748+ Data : Data {
749+ BackgroundTasks : & BackgroundTasks {
750+ Agents : []BackgroundTasksAgent {},
751+ Shells : []Shell {{ShellID : "sh-1" }},
752+ },
753+ },
754+ })
755+
756+ time .Sleep (100 * time .Millisecond )
757+ if resolved .Load () {
758+ t .Error ("BUG #299: SendAndWait resolved prematurely while background shells were active" )
759+ }
760+
761+ // Clean idle — should now resolve.
762+ session .dispatchEvent (SessionEvent {
763+ Type : SessionEventTypeSessionIdle ,
764+ Data : Data {BackgroundTasks : & BackgroundTasks {Agents : []BackgroundTasksAgent {}, Shells : []Shell {}}},
765+ })
766+
767+ select {
768+ case <- done :
769+ case <- time .After (2 * time .Second ):
770+ t .Fatal ("SendAndWait did not resolve after clean session.idle" )
771+ }
772+ })
773+
774+ t .Run ("resolves when session.idle has no backgroundTasks" , func (t * testing.T ) {
775+ session , requestReceived , cleanup := newTestSessionWithFakeClient (t )
776+ defer cleanup ()
777+
778+ ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
779+ defer cancel ()
780+
781+ done := make (chan struct {})
782+ go func () {
783+ defer close (done )
784+ session .SendAndWait (ctx , MessageOptions {Prompt : "test" }) //nolint:errcheck
785+ }()
786+
787+ select {
788+ case <- requestReceived :
789+ case <- time .After (2 * time .Second ):
790+ t .Fatal ("timeout waiting for session.send request" )
791+ }
792+
793+ session .dispatchEvent (SessionEvent {
794+ Type : SessionEventTypeSessionIdle ,
795+ Data : Data {},
796+ })
797+
798+ select {
799+ case <- done :
800+ case <- time .After (2 * time .Second ):
801+ t .Fatal ("SendAndWait did not resolve when session.idle had no backgroundTasks" )
802+ }
803+ })
804+
805+ t .Run ("resolves when session.idle has empty backgroundTasks arrays" , func (t * testing.T ) {
806+ session , requestReceived , cleanup := newTestSessionWithFakeClient (t )
807+ defer cleanup ()
808+
809+ ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
810+ defer cancel ()
811+
812+ done := make (chan struct {})
813+ go func () {
814+ defer close (done )
815+ session .SendAndWait (ctx , MessageOptions {Prompt : "test" }) //nolint:errcheck
816+ }()
817+
818+ select {
819+ case <- requestReceived :
820+ case <- time .After (2 * time .Second ):
821+ t .Fatal ("timeout waiting for session.send request" )
822+ }
823+
824+ session .dispatchEvent (SessionEvent {
825+ Type : SessionEventTypeSessionIdle ,
826+ Data : Data {BackgroundTasks : & BackgroundTasks {Agents : []BackgroundTasksAgent {}, Shells : []Shell {}}},
827+ })
828+
829+ select {
830+ case <- done :
831+ case <- time .After (2 * time .Second ):
832+ t .Fatal ("SendAndWait did not resolve when backgroundTasks arrays were empty" )
833+ }
834+ })
835+ }
0 commit comments