test(teams): fix the data races in the teammate streaming test (#1556)

TestHandleTeammateMessageSchedulesStreamedRun fails CI intermittently under
-race. Two separate races, both in the test rather than in what it exercises:

  1. It shared `gotReq` between the scheduler's RunFunc and the assertions. The
     RunFunc runs on a scheduler goroutine, and the announce loop schedules a
     second run after the teammate one, so the write could land while the test
     was reading — and the second run also closed an already-closed channel.
     Requests now arrive over a buffered channel: no shared state, and a second
     run cannot clobber the first.

  2. The deferred sched.Stop() ran while handleTeammateMessage's background
     goroutine was still calling Schedule, so Lane.Submit's wg.Add raced
     Lane.Stop's wg.Wait. The gateway drains BgWg before stopping the scheduler
     (gateway_consumer.go waits on it; sched.Stop is an outer defer in
     gateway.go); the test skipped that step. It now drains too, which also makes
     the test match the shutdown order it is meant to represent.

Reproduced before the fix with `-race -count=60 -cpu=1,4` (fails within a few
iterations) and clean afterwards over `-count=200 -cpu=1,2,4`, plus the whole
cmd package under -race.

Nothing in production changes; `Stream: true` and the channel-manager assertion
are untouched.

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
yatulandClaude Opus 5 authored and GitHub committed 2026-09-09 09:01:31 +07:00
1 parent d98ea89e3b
commit 34cc6e011d
1 file changed
+19 -5
+19 -5
View File
@@ -22,8 +22,11 @@ import (
// channel manager, so HandleAgentEvent drops its chunks and nothing reaches a
// user incrementally — the task result keeps coming from the final RunResult.
func TestHandleTeammateMessageSchedulesStreamedRun(t *testing.T) {
var gotReq agent.RunRequest
ran := make(chan struct{})
// The scheduler runs its RunFunc on its own goroutine, and the announce loop
// schedules a second run after the teammate one, so a shared variable here is
// written concurrently with the assertions below. Hand the requests over a
// channel instead: no shared state, and a second run cannot clobber the first.
scheduled := make(chan agent.RunRequest, 4)
sched := scheduler.NewScheduler(
scheduler.DefaultLanes(),
@@ -34,8 +37,10 @@ func TestHandleTeammateMessageSchedulesStreamedRun(t *testing.T) {
MaxConcurrent: 1,
},
func(_ context.Context, req agent.RunRequest) (*agent.RunResult, error) {
gotReq = req
close(ran)
select {
case scheduled <- req:
default:
}
return &agent.RunResult{Content: "member deliverable"}, nil
},
)
@@ -47,6 +52,14 @@ func TestHandleTeammateMessageSchedulesStreamedRun(t *testing.T) {
Sched: sched,
ChannelMgr: channelMgr,
}
// handleTeammateMessage hands the announce loop to a background goroutine that
// keeps calling Schedule after this function returns. Drain it before the
// deferred Stop above runs, mirroring the shutdown order the gateway itself
// uses (gateway_consumer.go waits on BgWg; sched.Stop is an outer defer in
// gateway.go). Without this the teardown races the announce loop inside the
// lane's WaitGroup — Submit's Add against Stop's Wait — and -race fails the
// test intermittently.
defer deps.BgWg.Wait()
msg := bus.InboundMessage{
Channel: tools.ChannelSystem,
@@ -65,8 +78,9 @@ func TestHandleTeammateMessageSchedulesStreamedRun(t *testing.T) {
t.Fatal("handleTeammateMessage() = false, want true for a teammate: message on the system channel")
}
var gotReq agent.RunRequest
select {
case <-ran:
case gotReq = <-scheduled:
case <-time.After(5 * time.Second):
t.Fatal("teammate run was never scheduled")
}