@@ -717,10 +717,27 @@ func RunOutpostMCPSubprocess(t *testing.T, projectRoot, configPath string, args
717717 return RunMCPSubprocess (t , projectRoot , configPath , append ([]string {"outpost" , "mcp" }, args ... ), extraEnv , stdin , runDuration )
718718}
719719
720- // RunMCPSubprocess builds the CLI binary, runs the given MCP command with optional stdin,
721- // lets it run for runDuration, then kills the process. Returns stdout, stderr, and the
722- // error from Wait (often non-nil because the process was killed). configPath, when non-empty,
723- // is passed as HOOKDECK_CONFIG_FILE. extraEnv entries override the process environment.
720+ // mcpShutdownGrace is how long a server gets to exit after its stdin is closed
721+ // before it is killed.
722+ const mcpShutdownGrace = 5 * time .Second
723+
724+ // RunMCPSubprocess builds the CLI binary and drives the given MCP command over
725+ // stdio the way an MCP client does. stdin holds newline-delimited JSON-RPC
726+ // messages. An initialize request is sent on its own, and the messages after it
727+ // only once the server has answered it. Stdout is read as it arrives, and as soon
728+ // as every request in stdin has a response, stdin is closed (how a client ends a
729+ // stdio session) and the process is left to exit.
730+ //
731+ // runDuration is an upper bound, not a fixed wait: if the responses have not all
732+ // arrived by then, the process is killed and the caller sees whatever it did
733+ // write. With no requests in stdin (none at all, or only notifications) there is
734+ // nothing to wait for, so the process runs until it exits or runDuration passes.
735+ // An empty stdin reads as EOF at once, for commands expected to fail before they
736+ // serve.
737+ //
738+ // Returns stdout, stderr, and the error from Wait (non-nil when the process was
739+ // killed or exited non-zero). configPath, when non-empty, is passed as
740+ // HOOKDECK_CONFIG_FILE. extraEnv entries override the process environment.
724741func RunMCPSubprocess (t * testing.T , projectRoot , configPath string , args []string , extraEnv map [string ]string , stdin string , runDuration time.Duration ) (stdout , stderr string , err error ) {
725742 t .Helper ()
726743 tmpBinary := filepath .Join (projectRoot , "hookdeck-mcp-test-" + generateTimestamp ())
@@ -743,44 +760,213 @@ func RunMCPSubprocess(t *testing.T, projectRoot, configPath string, args []strin
743760 }
744761 cmd .Env = env
745762
746- var stdoutBuf , stderrBuf bytes.Buffer
747- cmd .Stdout = & stdoutBuf
763+ out := newMCPStdout ()
764+ var stderrBuf bytes.Buffer
765+ cmd .Stdout = out
748766 cmd .Stderr = & stderrBuf
749767
750- closeStdin := func () {}
768+ messages := splitMCPMessages (stdin )
769+ var stdinPipe io.WriteCloser
751770 if stdin != "" {
752- pr , pw , pipeErr := os . Pipe ()
771+ pipe , pipeErr := cmd . StdinPipe ()
753772 if pipeErr != nil {
754773 return "" , "" , pipeErr
755774 }
756- cmd .Stdin = pr
757- if _ , werr := io .WriteString (pw , stdin ); werr != nil {
758- _ = pr .Close ()
759- _ = pw .Close ()
760- return "" , "" , werr
761- }
762- if ! strings .HasSuffix (stdin , "\n " ) {
763- _ , _ = pw .Write ([]byte ("\n " ))
764- }
765- closeStdin = func () {
766- _ = pw .Close ()
767- _ = pr .Close ()
768- }
775+ stdinPipe = pipe
769776 }
770777
771778 if err := cmd .Start (); err != nil {
772- closeStdin ()
773779 return "" , "" , err
774780 }
781+ exited := make (chan struct {})
782+ var waitErr error
783+ go func () {
784+ waitErr = cmd .Wait ()
785+ close (exited )
786+ }()
787+ // Killing here, not only signalling, is what bounds a write to a server
788+ // that has stopped reading.
789+ expired := make (chan struct {})
775790 timer := time .AfterFunc (runDuration , func () {
776- if cmd .Process != nil {
777- _ = cmd .Process .Kill ()
778- }
791+ close (expired )
792+ _ = cmd .Process .Kill ()
779793 })
780794 defer timer .Stop ()
781- waitErr := cmd .Wait ()
782- closeStdin ()
783- return stdoutBuf .String (), stderrBuf .String (), waitErr
795+
796+ // awaitResponses reports whether every id has a response, waiting until
797+ // they do, the process exits, or runDuration passes.
798+ awaitResponses := func (ids []string ) bool {
799+ for {
800+ if out .hasResponses (ids ) {
801+ return true
802+ }
803+ select {
804+ case <- out .changed :
805+ case <- exited :
806+ return out .hasResponses (ids )
807+ case <- expired :
808+ return false
809+ }
810+ }
811+ }
812+
813+ requests := mcpRequestIDs (messages )
814+ answered := false
815+ if stdinPipe != nil {
816+ head , tail , initID := splitAtMCPInitialize (messages )
817+ writeMCPMessages (stdinPipe , head )
818+ if initID == "" || awaitResponses ([]string {initID }) {
819+ writeMCPMessages (stdinPipe , tail )
820+ answered = len (requests ) > 0 && awaitResponses (requests )
821+ }
822+ }
823+
824+ if answered {
825+ _ = stdinPipe .Close ()
826+ select {
827+ case <- exited :
828+ case <- expired :
829+ case <- time .After (mcpShutdownGrace ):
830+ }
831+ } else {
832+ select {
833+ case <- exited :
834+ case <- expired :
835+ }
836+ }
837+ _ = cmd .Process .Kill ()
838+ <- exited
839+ return out .String (), stderrBuf .String (), waitErr
840+ }
841+
842+ // mcpStdout collects a server's stdout and tracks which JSON-RPC responses have
843+ // arrived, so the harness can stop waiting as soon as the one it needs is there.
844+ type mcpStdout struct {
845+ mu sync.Mutex
846+ buf bytes.Buffer
847+ partial []byte
848+ responses map [string ]bool
849+ // changed is signalled after every write.
850+ changed chan struct {}
851+ }
852+
853+ func newMCPStdout () * mcpStdout {
854+ return & mcpStdout {responses : map [string ]bool {}, changed : make (chan struct {}, 1 )}
855+ }
856+
857+ func (o * mcpStdout ) Write (p []byte ) (int , error ) {
858+ o .mu .Lock ()
859+ o .buf .Write (p )
860+ o .partial = append (o .partial , p ... )
861+ for {
862+ i := bytes .IndexByte (o .partial , '\n' )
863+ if i < 0 {
864+ break
865+ }
866+ if id , ok := mcpResponseID (o .partial [:i ]); ok {
867+ o .responses [id ] = true
868+ }
869+ o .partial = o .partial [i + 1 :]
870+ }
871+ o .mu .Unlock ()
872+ select {
873+ case o .changed <- struct {}{}:
874+ default :
875+ }
876+ return len (p ), nil
877+ }
878+
879+ func (o * mcpStdout ) hasResponses (ids []string ) bool {
880+ o .mu .Lock ()
881+ defer o .mu .Unlock ()
882+ for _ , id := range ids {
883+ if ! o .responses [id ] {
884+ return false
885+ }
886+ }
887+ return true
888+ }
889+
890+ func (o * mcpStdout ) String () string {
891+ o .mu .Lock ()
892+ defer o .mu .Unlock ()
893+ return o .buf .String ()
894+ }
895+
896+ // jsonRPCEnvelope is the part of a JSON-RPC message the harness routes on.
897+ type jsonRPCEnvelope struct {
898+ ID json.RawMessage `json:"id"`
899+ Method string `json:"method"`
900+ }
901+
902+ // mcpMessageID returns a JSON-RPC message's id as a comparable key, and whether
903+ // it has one. The raw JSON is the key, so 2 and "2" stay distinct.
904+ func mcpMessageID (env jsonRPCEnvelope ) (string , bool ) {
905+ id := strings .TrimSpace (string (env .ID ))
906+ if id == "" || id == "null" {
907+ return "" , false
908+ }
909+ return id , true
910+ }
911+
912+ // mcpResponseID returns the id of a response line. Requests the server sends to
913+ // the client carry an id too, so a method rules a line out.
914+ func mcpResponseID (line []byte ) (string , bool ) {
915+ var env jsonRPCEnvelope
916+ if json .Unmarshal (line , & env ) != nil || env .Method != "" {
917+ return "" , false
918+ }
919+ return mcpMessageID (env )
920+ }
921+
922+ // splitMCPMessages returns the non-blank lines of stdin, one message each.
923+ func splitMCPMessages (stdin string ) []string {
924+ var messages []string
925+ for _ , line := range strings .Split (stdin , "\n " ) {
926+ if strings .TrimSpace (line ) != "" {
927+ messages = append (messages , line )
928+ }
929+ }
930+ return messages
931+ }
932+
933+ // mcpRequestIDs returns the ids of the requests in messages: everything that
934+ // expects a response. Notifications and lines that are not JSON are skipped.
935+ func mcpRequestIDs (messages []string ) []string {
936+ var ids []string
937+ for _ , message := range messages {
938+ var env jsonRPCEnvelope
939+ if json .Unmarshal ([]byte (message ), & env ) != nil || env .Method == "" {
940+ continue
941+ }
942+ if id , ok := mcpMessageID (env ); ok {
943+ ids = append (ids , id )
944+ }
945+ }
946+ return ids
947+ }
948+
949+ // splitAtMCPInitialize splits messages after the initialize request, returning
950+ // its id. Without one, everything is head and the id is empty.
951+ func splitAtMCPInitialize (messages []string ) (head , tail []string , initID string ) {
952+ for i , message := range messages {
953+ var env jsonRPCEnvelope
954+ if json .Unmarshal ([]byte (message ), & env ) != nil || env .Method != "initialize" {
955+ continue
956+ }
957+ if id , ok := mcpMessageID (env ); ok {
958+ return messages [:i + 1 ], messages [i + 1 :], id
959+ }
960+ }
961+ return messages , nil , ""
962+ }
963+
964+ // writeMCPMessages writes each message as one line. Errors are ignored: a server
965+ // that has exited cannot read them, and its output says why.
966+ func writeMCPMessages (w io.Writer , messages []string ) {
967+ for _ , message := range messages {
968+ _ , _ = io .WriteString (w , message + "\n " )
969+ }
784970}
785971
786972// MCPToolCallResult is the parsed JSON-RPC response for an MCP tools/call request.
0 commit comments