test(mcp): drive interrupted-write reconnect test by cancellation instead of a timer

commit: ec2657b620f9e6b75e7753b4b3c0ff36e672265b
author: brice · 7d ago · 2026-10-15 15:06 UTC
parent: 75c7116
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

1 file · +29 -7

M packages/agent/mcp/round2_test.go +29 -7


@@ -3,12 +3,15 @@ package mcp
3	3		import (
4	4		"bytes"
5	5		"context"
	6	+	"encoding/json"
6	7		"errors"
7	8		"io"
8	9		"net/url"
9	10		"strings"
10	11		"testing"
11	12		"time"
	13	+
	14	+	"github.com/bnema/zut/packages/core"
12	15		)
13	16
14	17		// Regression tests for the second review round.

@@ -82,18 +85,37 @@ func TestStdioQueuedFrameIsWithdrawnButStartedWriteFailsTransport(t *testing.T)
82	85		}
83	86
84	87		func TestManagerReconnectsAfterInterruptedWrite(t *testing.T) {
85		-	m, events := startFake(t, nil, func(c *ServerConfig) { c.ToolTimeoutSec = 0.4 })
	88	+	// Cancellation is driven by the test, not by toolTimeoutSec: a timer could
	89	+	// fire before the 4 MiB frame is even queued on a slow (-race) runner, and
	90	+	// a withdrawn frame correctly leaves the transport up.
	91	+	m, events := startFake(t, nil, nil)
	92	+	st := m.servers["fake"].live().t.(*stdioTransport)
86	93		// Park the server: the call is delivered, then it stops reading stdin.
87		-	if res := execTool(t, findTool(t, m, "mcp__fake__stall"), `{}`); !res.IsError {
88		-	t.Fatalf("stall result = %q", resultText(res))
89		-	}
	94	+	// The stall call stays pending (no cancel notification is queued) and
	95	+	// ends when the transport fails.
	96	+	stall, echo := findTool(t, m, "mcp__fake__stall"), findTool(t, m, "mcp__fake__echo")
	97	+	stallDone := make(chan struct{})
	98	+	go func() {
	99	+	defer close(stallDone)
	100	+	_, _ = stall.Execute(context.Background(), json.RawMessage(`{}`), nil)
	101	+	}()
90	102		waitEvent(t, events, "stalled")
91		-	// A request too large for the pipe times out mid-write.
92		-	big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}`
93		-	if res := execTool(t, findTool(t, m, "mcp__fake__echo"), big); !res.IsError {
	103	+	// A request too large for the pipe is cancelled once its write started.
	104	+	started := st.writesStarted.Load()
	105	+	bigCtx, bigCancel := context.WithCancel(context.Background())
	106	+	bigDone := make(chan core.ToolResult, 1)
	107	+	go func() {
	108	+	big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}`
	109	+	res, _ := echo.Execute(bigCtx, json.RawMessage(big), nil)
	110	+	bigDone <- res
	111	+	}()
	112	+	eventually(t, "the big frame write to start", func() bool { return st.writesStarted.Load() > started })
	113	+	bigCancel()
	114	+	if res := <-bigDone; !res.IsError {
94	115		t.Fatalf("big call result = %q", resultText(res))
95	116		}
96	117		eventually(t, "status to show the dropped connection", func() bool { return !m.Status()[0].Connected })
	118	+	<-stallDone
97	119		// The next call starts a fresh server and works.
98	120		if got := resultText(execTool(t, findTool(t, m, "mcp__fake__echo"), `{"text":"again"}`)); got != "echo: again" {
99	121		t.Fatalf("after reconnect = %q", got)

 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
 	"net/url"
 	"strings"
 	"testing"
 	"time"
+	"github.com/bnema/zut/packages/core"
 )
 // Regression tests for the second review round.
 }
 func TestManagerReconnectsAfterInterruptedWrite(t *testing.T) {
-	m, events := startFake(t, nil, func(c *ServerConfig) { c.ToolTimeoutSec = 0.4 })
+	// Cancellation is driven by the test, not by toolTimeoutSec: a timer could
+	// fire before the 4 MiB frame is even queued on a slow (-race) runner, and
+	// a withdrawn frame correctly leaves the transport up.
+	m, events := startFake(t, nil, nil)
+	st := m.servers["fake"].live().t.(*stdioTransport)
 	// Park the server: the call is delivered, then it stops reading stdin.
-	if res := execTool(t, findTool(t, m, "mcp__fake__stall"), `{}`); !res.IsError {
+	// The stall call stays pending (no cancel notification is queued) and
-		t.Fatalf("stall result = %q", resultText(res))
+	// ends when the transport fails.
-	}
+	stall, echo := findTool(t, m, "mcp__fake__stall"), findTool(t, m, "mcp__fake__echo")
+	stallDone := make(chan struct{})
+	go func() {
+		defer close(stallDone)
+		_, _ = stall.Execute(context.Background(), json.RawMessage(`{}`), nil)
+	}()
 	waitEvent(t, events, "stalled")
-	// A request too large for the pipe times out mid-write.
+	// A request too large for the pipe is cancelled once its write started.
-	big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}`
+	started := st.writesStarted.Load()
-	if res := execTool(t, findTool(t, m, "mcp__fake__echo"), big); !res.IsError {
+	bigCtx, bigCancel := context.WithCancel(context.Background())
+	bigDone := make(chan core.ToolResult, 1)
+	go func() {
+		big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}`
+		res, _ := echo.Execute(bigCtx, json.RawMessage(big), nil)
+		bigDone <- res
+	}()
+	eventually(t, "the big frame write to start", func() bool { return st.writesStarted.Load() > started })
+	bigCancel()
+	if res := <-bigDone; !res.IsError {
 		t.Fatalf("big call result = %q", resultText(res))
 	}
 	eventually(t, "status to show the dropped connection", func() bool { return !m.Status()[0].Connected })
+	<-stallDone
 	// The next call starts a fresh server and works.
 	if got := resultText(execTool(t, findTool(t, m, "mcp__fake__echo"), `{"text":"again"}`)); got != "echo: again" {
 		t.Fatalf("after reconnect = %q", got)