fix(providers): send stable session headers to OpenCode

Fixes #1023

CHANGELOG.md

@@ -35,6 +35,7 @@ tagged. Until then, source builds report the version `dev`.
 * **mcp:** reject a server-initiated request as a streamable-HTTP response ([#1054](https://github.com/Gitlawb/zero/issues/1054)) ([2479d94](https://github.com/Gitlawb/zero/commit/2479d946c15dbf74ee3c04a637e97181d1786f08))
 * **oauth:** force-close loopback callback connections past the shutdown deadline ([#1051](https://github.com/Gitlawb/zero/issues/1051)) ([176f5a2](https://github.com/Gitlawb/zero/commit/176f5a242b754a86ea4c7121a8dcc1c0d72cf2e4))
 * **oauth:** Split an oversized keyring token blob across entries ([#1007](https://github.com/Gitlawb/zero/issues/1007)) ([275e248](https://github.com/Gitlawb/zero/commit/275e2486695338ce38193a6e3319f8b6a9a5c33c))
+* **providerio:** close idle connections periodically to prevent HTTP/2 PROTOCOL_ERROR ([#1058](https://github.com/Gitlawb/zero/issues/1058)) ([b76c764](https://github.com/Gitlawb/zero/commit/b76c764b77926891aa67163e6605a2d8540aacaa))
 * **providers:** register Atomic Chat as local OpenAI-compatible runtime ([#550](https://github.com/Gitlawb/zero/issues/550)) ([79fa142](https://github.com/Gitlawb/zero/commit/79fa1422ef1a3691b93342788453ebc284d57f9c))
 * **redaction:** split camelCase keys in normalizeKey ([#994](https://github.com/Gitlawb/zero/issues/994)) ([08fbea7](https://github.com/Gitlawb/zero/commit/08fbea7a3221d2061d88038924f90a4cba9af1ac))
 * **sandbox:** keep a derived ACL target inside the write root it came from ([#1040](https://github.com/Gitlawb/zero/issues/1040)) ([3b43c43](https://github.com/Gitlawb/zero/commit/3b43c43eca968c8619adadc81fcc1a17e220216a))

internal/agent/async_diagnostics_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Gitlawb/zero/internal/testutil"
 	"github.com/Gitlawb/zero/internal/tools"
 	"github.com/Gitlawb/zero/internal/zeroruntime"
 )
@@ -111,7 +112,11 @@ func TestAsyncDiagnosticsReEditReplacesResult(t *testing.T) {
 	diagnostics := newAsyncDiagnostics(check, "/ws")
 
 	diagnostics.enqueue(context.Background(), []string{"a.go"})
-	waitForIdle(t, diagnostics)
+	testutil.WaitFor(t, "worker idle", func() bool {
+		diagnostics.mu.Lock()
+		defer diagnostics.mu.Unlock()
+		return diagnostics.working == nil
+	})
 	mu.Lock()
 	response = "ERR new"
 	mu.Unlock()
@@ -123,23 +128,6 @@ func TestAsyncDiagnosticsReEditReplacesResult(t *testing.T) {
 	}
 }
 
-func waitForIdle(t *testing.T, diagnostics *asyncDiagnostics) {
-	t.Helper()
-	deadline := time.Now().Add(2 * time.Second)
-	for {
-		diagnostics.mu.Lock()
-		busy := diagnostics.working
-		diagnostics.mu.Unlock()
-		if busy == nil {
-			return
-		}
-		if time.Now().After(deadline) {
-			t.Fatal("worker never went idle")
-		}
-		time.Sleep(2 * time.Millisecond)
-	}
-}
-
 // changedFilesTool is a fake mutating tool reporting a changed file without
 // touching the filesystem.
 type changedFilesTool struct{}

internal/daemon/pool_test.go

@@ -7,6 +7,8 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
+
+	"github.com/Gitlawb/zero/internal/testutil"
 )
 
 // --- test doubles ---------------------------------------------------------
@@ -191,7 +193,7 @@ func TestPoolQueuesWhenFull(t *testing.T) {
 	}()
 	<-started
 	// Wait until the first run holds the only slot.
-	waitFor(t, func() bool { return pool.QueueDepth() == 1 })
+	testutil.WaitFor(t, "queue depth 1", func() bool { return pool.QueueDepth() == 1 })
 
 	secondDone := make(chan struct{})
 	go func() {
@@ -224,7 +226,7 @@ func TestPoolDrainKillsStraggler(t *testing.T) {
 		_, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{})
 		close(runDone)
 	}()
-	waitFor(t, func() bool { return pool.QueueDepth() == 1 })
+	testutil.WaitFor(t, "queue depth 1", func() bool { return pool.QueueDepth() == 1 })
 
 	pool.Drain() // KillTimeout elapses, straggler is force-killed
 	if atomic.LoadInt32(&straggler.killed) != 1 {
@@ -242,18 +244,6 @@ func TestPoolDrainKillsStraggler(t *testing.T) {
 	}
 }
 
-func waitFor(t *testing.T, cond func() bool) {
-	t.Helper()
-	deadline := time.Now().Add(2 * time.Second)
-	for time.Now().Before(deadline) {
-		if cond() {
-			return
-		}
-		time.Sleep(time.Millisecond)
-	}
-	t.Fatal("condition not met within timeout")
-}
-
 func TestPoolRunSurfacesStdoutReadError(t *testing.T) {
 	// A worker stdout read error must surface as a failure (not be swallowed and
 	// reported as clean success), and the worker is killed so it can't block on an

internal/daemon/server_test.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/Gitlawb/zero/internal/observability"
+	"github.com/Gitlawb/zero/internal/testutil"
 )
 
 func TestServeSupportsReadOnlyCustomRuntimeDirectory(t *testing.T) {
@@ -34,7 +35,7 @@ func TestServeSupportsReadOnlyCustomRuntimeDirectory(t *testing.T) {
 	srv := newTestServerWithPaths(t, launcher, paths)
 	serveErr := make(chan error, 1)
 	go func() { serveErr <- srv.Serve() }()
-	waitForFile(t, paths.Status)
+	testutil.WaitFor(t, "file "+paths.Status, func() bool { _, err := os.Stat(paths.Status); return err == nil })
 	srv.Shutdown()
 	if err := <-serveErr; err != nil {
 		t.Fatalf("Serve with 0755 custom runtime directory: %v", err)
@@ -212,18 +213,6 @@ func newTestServerWithPaths(t *testing.T, launcher Launcher, paths Paths) *Serve
 	return srv
 }
 
-func waitForFile(t *testing.T, path string) {
-	t.Helper()
-	deadline := time.Now().Add(2 * time.Second)
-	for time.Now().Before(deadline) {
-		if _, err := os.Stat(path); err == nil {
-			return
-		}
-		time.Sleep(2 * time.Millisecond)
-	}
-	t.Fatalf("file %s did not appear within timeout", path)
-}
-
 func TestServerEndToEnd(t *testing.T) {
 	out := []string{`{"type":"event","seq":1}`, `{"type":"event","seq":2}`}
 	launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: out, exitCode: 0})
@@ -231,7 +220,7 @@ func TestServerEndToEnd(t *testing.T) {
 
 	serveErr := make(chan error, 1)
 	go func() { serveErr <- srv.Serve() }()
-	waitForFile(t, paths.Status)
+	testutil.WaitFor(t, "file "+paths.Status, func() bool { _, err := os.Stat(paths.Status); return err == nil })
 
 	// --- run a session and collect its stream-json output ---
 	runClient, err := Dial(paths.Socket)
@@ -384,7 +373,8 @@ func TestServeCleansStatusThroughBoundDirectoryAfterSwap(t *testing.T) {
 
 	serveErr := make(chan error, 1)
 	go func() { serveErr <- srv.Serve() }()
-	waitForFile(t, filepath.Join(movedDir, "d.status"))
+	movedStatus := filepath.Join(movedDir, "d.status")
+	testutil.WaitFor(t, "file "+movedStatus, func() bool { _, err := os.Stat(movedStatus); return err == nil })
 	srv.Shutdown()
 	select {
 	case err := <-serveErr:
@@ -472,7 +462,7 @@ func TestServerSecondInstanceFails(t *testing.T) {
 
 	serveErr := make(chan error, 1)
 	go func() { serveErr <- srv1.Serve() }()
-	waitForFile(t, paths.Status)
+	testutil.WaitFor(t, "file "+paths.Status, func() bool { _, err := os.Stat(paths.Status); return err == nil })
 
 	// A second server on the same paths must refuse to start (single instance).
 	pool2, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
@@ -493,7 +483,7 @@ func TestServerRejectsUnknownCommand(t *testing.T) {
 	srv, paths := newTestServer(t, launcher)
 	go func() { _ = srv.Serve() }()
 	defer srv.Shutdown()
-	waitForFile(t, paths.Status)
+	testutil.WaitFor(t, "file "+paths.Status, func() bool { _, err := os.Stat(paths.Status); return err == nil })
 
 	client, err := Dial(paths.Socket)
 	if err != nil {

internal/daemon/session_test.go

@@ -6,6 +6,8 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
+
+	"github.com/Gitlawb/zero/internal/testutil"
 )
 
 func drain(t *testing.T, buffered []string, live <-chan string) []string {
@@ -111,7 +113,7 @@ func TestSessionLeaseQueuesWhenPoolFull(t *testing.T) {
 	mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool})
 
 	sa, _ := mgr.Start(context.Background(), WorkerSpec{Session: "a"})
-	waitFor(t, func() bool { return sa.State() == SessionRunning })
+	testutil.WaitFor(t, "session running", func() bool { return sa.State() == SessionRunning })
 
 	sb, _ := mgr.Start(context.Background(), WorkerSpec{Session: "b"})
 	// b must stay queued (its worker not yet launched) while a holds the slot.
@@ -197,7 +199,7 @@ func TestSessionManagerKeepsRunningOverCap(t *testing.T) {
 
 	s1, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r1"})
 	s2, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r2"})
-	waitFor(t, func() bool { return s1.State() == SessionRunning && s2.State() == SessionRunning })
+	testutil.WaitFor(t, "both sessions running", func() bool { return s1.State() == SessionRunning && s2.State() == SessionRunning })
 	if _, ok := mgr.Get("r1"); !ok {
 		t.Fatal("a running session must never be evicted, even past the cap")
 	}

internal/providers/anthropic/provider.go

@@ -90,6 +90,7 @@ type Options struct {
 	// When unset, Zero uses providerio.ResolveStreamIdleTimeout — the
 	// ZERO_STREAM_IDLE_TIMEOUT override or providerio.DefaultStreamIdleTimeout.
 	StreamIdleTimeout time.Duration
+	SessionHeader     string
 }
 
 // Provider streams completions from Anthropic's Messages API.
@@ -108,6 +109,7 @@ type Provider struct {
 	httpClient        *http.Client
 	userAgent         string
 	streamIdleTimeout time.Duration
+	sessionHeader     providerio.SessionHeader
 }
 
 // New creates an Anthropic provider.
@@ -143,6 +145,7 @@ func New(options Options) (*Provider, error) {
 		httpClient:        providerio.HTTPClient(options.HTTPClient),
 		userAgent:         options.UserAgent,
 		streamIdleTimeout: providerio.ResolveStreamIdleTimeout(options.StreamIdleTimeout),
+		sessionHeader:     providerio.NewSessionHeader(options.SessionHeader),
 	}, nil
 }
 
@@ -163,12 +166,12 @@ func (provider *Provider) StreamCompletion(
 	events := make(chan zeroruntime.StreamEvent, 16)
 	go func() {
 		defer close(events)
-		provider.stream(ctx, body, events)
+		provider.stream(ctx, body, request.PromptCacheKey, events)
 	}()
 	return events, nil
 }
 
-func (provider *Provider) stream(ctx context.Context, body []byte, events chan<- zeroruntime.StreamEvent) {
+func (provider *Provider) stream(ctx context.Context, body []byte, sessionID string, events chan<- zeroruntime.StreamEvent) {
 	// streamCtx lets the idle watchdog abort an in-flight body read by cancelling
 	// the request, which unblocks the SSE reader goroutine.
 	streamCtx, cancelStream := context.WithCancel(ctx)
@@ -185,6 +188,7 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
 		},
 		provider.oauthResolver,
 		func(request *http.Request) {
+			provider.sessionHeader.Apply(request, sessionID)
 			request.Header.Set("Content-Type", "application/json")
 			request.Header.Set("anthropic-version", provider.version)
 			if provider.beta != "" {

internal/providers/factory.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 	"strings"
 
 	"github.com/Gitlawb/zero/internal/config"
@@ -53,6 +54,11 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
 		return newCodexProvider(profile, resolved, options)
 	}
 
+	sessionHeader := sessionHeaderForProfile(profile, resolved)
+	if sessionHeader != "" && strings.TrimSpace(options.UserAgent) == "" {
+		options.UserAgent = "zero"
+	}
+
 	switch resolved.providerKind {
 	case config.ProviderKindOpenAI, config.ProviderKindOpenAICompatible:
 		// prompt_cache_key is an OpenAI-only chat-completions field. Strict
@@ -74,6 +80,7 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
 			UserAgent:             options.UserAgent,
 			ParseThinkTags:        parseThinkTagsForProfile(profile, resolved),
 			DisablePromptCacheKey: resolved.providerKind == config.ProviderKindOpenAICompatible,
+			SessionHeader:         sessionHeader,
 		})
 	case config.ProviderKindAnthropic, config.ProviderKindAnthropicCompat:
 		return anthropic.New(anthropic.Options{
@@ -88,6 +95,7 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
 			MaxTokens:       resolved.maxOutputTokens,
 			HTTPClient:      options.HTTPClient,
 			UserAgent:       options.UserAgent,
+			SessionHeader:   sessionHeader,
 		})
 	case config.ProviderKindGoogle:
 		return gemini.New(gemini.Options{
@@ -345,6 +353,18 @@ func isCodexCatalog(profile config.ProviderProfile, _ resolvedProfile) bool {
 	return providercatalog.NormalizeID(profile.CatalogID) == "chatgpt"
 }
 
+func sessionHeaderForProfile(profile config.ProviderProfile, resolved resolvedProfile) string {
+	switch providercatalog.NormalizeID(profile.CatalogID) {
+	case "opencode", "opencode-go", "opencode-go-anthropic-compatible":
+		return "x-opencode-session"
+	}
+	endpoint, err := url.Parse(resolved.baseURL)
+	if err == nil && strings.EqualFold(endpoint.Hostname(), "opencode.ai") && (endpoint.Path == "/zen" || strings.HasPrefix(endpoint.Path, "/zen/")) {
+		return "x-opencode-session"
+	}
+	return ""
+}
+
 // newCodexProvider builds a Codex-flavored openai provider for the chatgpt
 // catalog. The Codex headers (`originator`, `chatgpt-account-id`) are
 // injected by the CodexProvider wrapper. The `chatgpt-account-id` is read from

internal/providers/factory_test.go

@@ -3,13 +3,17 @@ package providers
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/Gitlawb/zero/internal/config"
 	"github.com/Gitlawb/zero/internal/oauth"
+	"github.com/Gitlawb/zero/internal/providers/providerio"
 	"github.com/Gitlawb/zero/internal/zeroruntime"
 )
 
@@ -794,3 +798,566 @@ func TestResolveRuntimeMetadataRejectsModelOutsideProviderAllowlist(t *testing.T
 		t.Fatalf("error = %q, want it to name the rejected model", err.Error())
 	}
 }
+
+func TestOpenCodeSessionHeader(t *testing.T) {
+	const sessionHeader = "x-opencode-session"
+	tests := []struct {
+		name       string
+		profile    config.ProviderProfile
+		response   string
+		wantHeader bool
+		wantAuth   map[string]string
+	}{
+		{
+			name: "catalog opencode-go default URL",
+			profile: config.ProviderProfile{
+				Name:         "opencode-go",
+				CatalogID:    "opencode-go",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				APIKey:       "sk-oc-go",
+				Model:        "deepseek-v4-pro",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-oc-go"},
+		},
+		{
+			name: "catalog opencode-go-anthropic-compatible",
+			profile: config.ProviderProfile{
+				Name:         "opencode-go-anthropic",
+				CatalogID:    "opencode-go-anthropic-compatible",
+				ProviderKind: config.ProviderKindAnthropicCompat,
+				APIKey:       "sk-oc-ant",
+				Model:        "minimax-m3",
+			},
+			response:   "data: {\"type\":\"message_stop\"}\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"x-api-key": "sk-oc-ant"},
+		},
+		{
+			name: "catalog opencode Zen",
+			profile: config.ProviderProfile{
+				Name:         "opencode-zen",
+				CatalogID:    "opencode",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				APIKey:       "sk-zen",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-zen"},
+		},
+		{
+			name: "renamed catalog profile with custom proxy URL",
+			profile: config.ProviderProfile{
+				Name:         "my-opencode",
+				CatalogID:    "opencode-go",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://proxy.example/v1",
+				APIKey:       "sk-proxy",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-proxy"},
+		},
+		{
+			name: "catalog-less OpenAI-compatible OpenCode Go URL",
+			profile: config.ProviderProfile{
+				Name:         "manual-go",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://opencode.ai/zen/go/v1",
+				APIKey:       "sk-manual",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-manual"},
+		},
+		{
+			name: "catalog-less Anthropic-compatible OpenCode Go URL",
+			profile: config.ProviderProfile{
+				Name:         "manual-go-anthropic",
+				ProviderKind: config.ProviderKindAnthropicCompat,
+				BaseURL:      "https://opencode.ai/zen/go",
+				APIKey:       "sk-manual-ant",
+				Model:        "test-model",
+			},
+			response:   "data: {\"type\":\"message_stop\"}\n\n",
+			wantHeader: true,
+			wantAuth:   map[string]string{"x-api-key": "sk-manual-ant"},
+		},
+		{
+			name: "official OpenAI",
+			profile: config.ProviderProfile{
+				Name:         "openai",
+				ProviderKind: config.ProviderKindOpenAI,
+				APIKey:       "sk-openai",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-openai"},
+		},
+		{
+			name: "generic OpenAI-compatible",
+			profile: config.ProviderProfile{
+				Name:         "generic",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://provider.example/v1",
+				APIKey:       "sk-generic",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-generic"},
+		},
+		{
+			name: "official Anthropic",
+			profile: config.ProviderProfile{
+				Name:         "anthropic",
+				ProviderKind: config.ProviderKindAnthropic,
+				APIKey:       "sk-ant",
+				Model:        "test-model",
+			},
+			response:   "data: {\"type\":\"message_stop\"}\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"x-api-key": "sk-ant"},
+		},
+		{
+			name: "lookalike subdomain is not OpenCode",
+			profile: config.ProviderProfile{
+				Name:         "deceptive-host",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://opencode.ai.example/zen/go/v1",
+				APIKey:       "sk-fake",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-fake"},
+		},
+		{
+			name: "lookalike path is not OpenCode",
+			profile: config.ProviderProfile{
+				Name:         "deceptive-path",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://example.com/opencode.ai/zen/go",
+				APIKey:       "sk-fake",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-fake"},
+		},
+		{
+			name: "zenith path is not Zen",
+			profile: config.ProviderProfile{
+				Name:         "zenith",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				BaseURL:      "https://opencode.ai/zenith",
+				APIKey:       "sk-zenith",
+				Model:        "test-model",
+			},
+			response:   "data: [DONE]\n\n",
+			wantHeader: false,
+			wantAuth:   map[string]string{"Authorization": "Bearer sk-zenith"},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			transport := &captureTransport{responseBody: test.response}
+			provider, err := New(test.profile, Options{
+				HTTPClient: &http.Client{Transport: transport},
+				UserAgent:  "zero-test",
+			})
+			if err != nil {
+				t.Fatalf("New() error = %v", err)
+			}
+			stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+				Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+				PromptCacheKey: "conversation-a",
+			})
+			if err != nil {
+				t.Fatalf("StreamCompletion() error = %v", err)
+			}
+			drainStream(t, stream)
+			if transport.request == nil {
+				t.Fatal("HTTP client was not used")
+			}
+			got := transport.request.Header.Get(sessionHeader)
+			if test.wantHeader && got != "conversation-a" {
+				t.Fatalf("%s = %q, want %q", sessionHeader, got, "conversation-a")
+			}
+			if !test.wantHeader && got != "" {
+				t.Fatalf("%s = %q, want no session header", sessionHeader, got)
+			}
+			for header, want := range test.wantAuth {
+				if got := transport.request.Header.Get(header); got != want {
+					t.Fatalf("%s = %q, want %q", header, got, want)
+				}
+			}
+			if got := transport.request.Header.Get("User-Agent"); got != "zero-test" {
+				t.Fatalf("User-Agent = %q, want zero-test", got)
+			}
+			if test.profile.ProviderKind == config.ProviderKindOpenAICompatible {
+				var body map[string]any
+				if err := json.NewDecoder(transport.body()).Decode(&body); err != nil {
+					t.Fatalf("decode request body: %v", err)
+				}
+				if _, ok := body["prompt_cache_key"]; ok {
+					t.Fatalf("openai-compatible body carried prompt_cache_key: %#v", body["prompt_cache_key"])
+				}
+			}
+		})
+	}
+}
+
+func TestOpenCodeSessionHeaderFollowsConversationKey(t *testing.T) {
+	transport := &recordingTransport{status: http.StatusOK, body: "data: [DONE]\n\n"}
+	profile := config.ProviderProfile{
+		Name:         "opencode-go",
+		CatalogID:    "opencode-go",
+		ProviderKind: config.ProviderKindOpenAICompatible,
+		APIKey:       "sk-oc",
+		Model:        "test-model",
+	}
+	options := Options{HTTPClient: &http.Client{Transport: transport}, UserAgent: "zero-test"}
+	provider, err := New(profile, options)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	for _, key := range []string{"conversation-a", "conversation-a", "conversation-b"} {
+		stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+			Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+			PromptCacheKey: key,
+		})
+		if err != nil {
+			t.Fatalf("StreamCompletion() error = %v", err)
+		}
+		drainStream(t, stream)
+	}
+	got := transport.sessionHeaders()
+	want := []string{"conversation-a", "conversation-a", "conversation-b"}
+	if len(got) != len(want) {
+		t.Fatalf("captured %d requests, want %d", len(got), len(want))
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("request %d x-opencode-session = %q, want %q", i, got[i], want[i])
+		}
+	}
+
+	reconstructed, err := New(profile, options)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	stream, err := reconstructed.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+		Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+		PromptCacheKey: "conversation-a",
+	})
+	if err != nil {
+		t.Fatalf("StreamCompletion() error = %v", err)
+	}
+	drainStream(t, stream)
+	got = transport.sessionHeaders()
+	if last := got[len(got)-1]; last != "conversation-a" {
+		t.Fatalf("reconstructed provider sent %s = %q, want conversation-a", "x-opencode-session", last)
+	}
+}
+
+func TestOpenCodeSessionHeaderCustomOverride(t *testing.T) {
+	custom := map[string]string{"X-Opencode-Session": "explicit-session"}
+	profile := config.ProviderProfile{
+		Name:          "opencode-go",
+		CatalogID:     "opencode-go",
+		ProviderKind:  config.ProviderKindOpenAICompatible,
+		APIKey:        "sk-oc",
+		Model:         "test-model",
+		CustomHeaders: custom,
+	}
+	transport := &captureTransport{responseBody: "data: [DONE]\n\n"}
+	provider, err := New(profile, Options{HTTPClient: &http.Client{Transport: transport}, UserAgent: "zero-test"})
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+		Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+		PromptCacheKey: "conversation-a",
+	})
+	if err != nil {
+		t.Fatalf("StreamCompletion() error = %v", err)
+	}
+	drainStream(t, stream)
+	if got := transport.request.Header.Get("x-opencode-session"); got != "explicit-session" {
+		t.Fatalf("x-opencode-session = %q, want explicit custom value", got)
+	}
+	if len(custom) != 1 || custom["X-Opencode-Session"] != "explicit-session" {
+		t.Fatalf("profile CustomHeaders mutated: %#v", custom)
+	}
+}
+
+func TestOpenCodeSessionHeaderFallback(t *testing.T) {
+	transport := &recordingTransport{status: http.StatusOK, body: "data: [DONE]\n\n"}
+	profile := config.ProviderProfile{
+		Name:         "opencode-go",
+		CatalogID:    "opencode-go",
+		ProviderKind: config.ProviderKindOpenAICompatible,
+		APIKey:       "sk-oc",
+		Model:        "test-model",
+	}
+	options := Options{HTTPClient: &http.Client{Transport: transport}}
+	first, err := New(profile, options)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	second, err := New(profile, options)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	request := zeroruntime.CompletionRequest{
+		Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+	}
+	for _, provider := range []zeroruntime.Provider{first, first, second} {
+		stream, err := provider.StreamCompletion(context.Background(), request)
+		if err != nil {
+			t.Fatalf("StreamCompletion() error = %v", err)
+		}
+		drainStream(t, stream)
+	}
+	headers := transport.sessionHeaders()
+	if len(headers) != 3 {
+		t.Fatalf("captured %d requests, want 3", len(headers))
+	}
+	if headers[0] == "" || headers[1] != headers[0] {
+		t.Fatalf("fallback headers = %q, %q — want one stable nonempty value per provider", headers[0], headers[1])
+	}
+	if headers[2] == headers[0] {
+		t.Fatalf("distinct provider instances shared fallback %q", headers[0])
+	}
+}
+
+func TestOpenCodeSessionHeaderDefaultsUserAgent(t *testing.T) {
+	transport := &captureTransport{responseBody: "data: [DONE]\n\n"}
+	provider, err := New(config.ProviderProfile{
+		Name:         "opencode-go",
+		CatalogID:    "opencode-go",
+		ProviderKind: config.ProviderKindOpenAICompatible,
+		APIKey:       "sk-oc",
+		Model:        "test-model",
+	}, Options{HTTPClient: &http.Client{Transport: transport}})
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+	stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+		Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+	})
+	if err != nil {
+		t.Fatalf("StreamCompletion() error = %v", err)
+	}
+	drainStream(t, stream)
+	if got := transport.request.Header.Get("User-Agent"); got != "zero" {
+		t.Fatalf("User-Agent = %q, want zero", got)
+	}
+}
+
+func TestOpenCodeSessionHeaderConcurrent(t *testing.T) {
+	transport := &recordingTransport{status: http.StatusOK, body: "data: [DONE]\n\n"}
+	provider, err := New(config.ProviderProfile{
+		Name:         "opencode-go",
+		CatalogID:    "opencode-go",
+		ProviderKind: config.ProviderKindOpenAICompatible,
+		APIKey:       "sk-oc",
+		Model:        "test-model",
+	}, Options{HTTPClient: &http.Client{Transport: transport}})
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+
+	var wg sync.WaitGroup
+	var errsMu sync.Mutex
+	var errs []string
+	for _, session := range []string{"session-a", "session-b"} {
+		wg.Add(1)
+		go func(session string) {
+			defer wg.Done()
+			for i := 0; i < 5; i++ {
+				stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+					Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: fmt.Sprintf("req-%s-%d", session, i)}},
+					PromptCacheKey: session,
+				})
+				if err != nil {
+					errsMu.Lock()
+					errs = append(errs, err.Error())
+					errsMu.Unlock()
+					return
+				}
+				for event := range stream {
+					if event.Type == zeroruntime.StreamEventError {
+						errsMu.Lock()
+						errs = append(errs, event.Error)
+						errsMu.Unlock()
+					}
+				}
+			}
+		}(session)
+	}
+	wg.Wait()
+	for _, errText := range errs {
+		t.Errorf("concurrent stream error: %s", errText)
+	}
+
+	requests := transport.recorded()
+	if len(requests) != 10 {
+		t.Fatalf("captured %d requests, want 10", len(requests))
+	}
+	for _, request := range requests {
+		var payload struct {
+			Messages []struct {
+				Content string `json:"content"`
+			} `json:"messages"`
+		}
+		if err := json.Unmarshal([]byte(request.body), &payload); err != nil {
+			t.Fatalf("decode request body: %v", err)
+		}
+		if len(payload.Messages) == 0 {
+			t.Fatal("request carried no user message")
+		}
+		content := payload.Messages[0].Content
+		want := "session-a"
+		if strings.Contains(content, "session-b") {
+			want = "session-b"
+		}
+		if got := request.header.Get("x-opencode-session"); got != want {
+			t.Fatalf("request for %q carried x-opencode-session %q, want %q", content, got, want)
+		}
+	}
+}
+
+func TestOpenCodeSessionHeaderAuthRetry(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		profile  config.ProviderProfile
+		response string
+	}{
+		{
+			name: "openai-compatible",
+			profile: config.ProviderProfile{
+				Name:         "opencode-go",
+				CatalogID:    "opencode-go",
+				ProviderKind: config.ProviderKindOpenAICompatible,
+				APIKey:       "sk-oc",
+				Model:        "test-model",
+			},
+			response: "data: [DONE]\n\n",
+		},
+		{
+			name: "anthropic-compatible",
+			profile: config.ProviderProfile{
+				Name:         "opencode-go-anthropic",
+				CatalogID:    "opencode-go-anthropic-compatible",
+				ProviderKind: config.ProviderKindAnthropicCompat,
+				APIKey:       "sk-oc-ant",
+				Model:        "minimax-m3",
+			},
+			response: "data: {\"type\":\"message_stop\"}\n\n",
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			transport := &recordingTransport{status: http.StatusOK, body: test.response, unauthorized: 1}
+			var calls atomic.Int32
+			resolver := providerio.TokenResolver(func(context.Context, bool) (string, string, bool, error) {
+				n := calls.Add(1)
+				return "Authorization", fmt.Sprintf("Bearer tok-%d", n), true, nil
+			})
+			provider, err := New(test.profile, Options{
+				HTTPClient:    &http.Client{Transport: transport},
+				OAuthResolver: resolver,
+			})
+			if err != nil {
+				t.Fatalf("New() error = %v", err)
+			}
+			stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
+				Messages:       []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}},
+				PromptCacheKey: "conversation-a",
+			})
+			if err != nil {
+				t.Fatalf("StreamCompletion() error = %v", err)
+			}
+			drainStream(t, stream)
+			requests := transport.recorded()
+			if len(requests) != 2 {
+				t.Fatalf("captured %d requests, want 2 (401 + retry)", len(requests))
+			}
+			for i, request := range requests {
+				if got := request.header.Get("x-opencode-session"); got != "conversation-a" {
+					t.Fatalf("attempt %d x-opencode-session = %q, want conversation-a", i, got)
+				}
+			}
+			if first, second := requests[0].header.Get("Authorization"), requests[1].header.Get("Authorization"); first == second {
+				t.Fatalf("Authorization not refreshed across retry: %q", first)
+			}
+		})
+	}
+}
+
+func drainStream(t *testing.T, stream <-chan zeroruntime.StreamEvent) {
+	t.Helper()
+	for event := range stream {
+		if event.Type == zeroruntime.StreamEventError {
+			t.Fatalf("stream error: %s", event.Error)
+		}
+	}
+}
+
+type recordedRequest struct {
+	header http.Header
+	body   string
+}
+
+type recordingTransport struct {
+	mu           sync.Mutex
+	requests     []recordedRequest
+	status       int
+	body         string
+	unauthorized int
+}
+
+func (transport *recordingTransport) RoundTrip(request *http.Request) (*http.Response, error) {
+	var body string
+	if request.Body != nil {
+		data, _ := io.ReadAll(request.Body)
+		body = string(data)
+	}
+	transport.mu.Lock()
+	transport.requests = append(transport.requests, recordedRequest{header: request.Header.Clone(), body: body})
+	attempt := len(transport.requests)
+	transport.mu.Unlock()
+
+	status := transport.status
+	if attempt <= transport.unauthorized {
+		status = http.StatusUnauthorized
+	}
+	return &http.Response{
+		StatusCode: status,
+		Status:     fmt.Sprintf("%d", status),
+		Header:     make(http.Header),
+		Body:       io.NopCloser(strings.NewReader(transport.body)),
+		Request:    request,
+	}, nil
+}
+
+func (transport *recordingTransport) recorded() []recordedRequest {
+	transport.mu.Lock()
+	defer transport.mu.Unlock()
+	return append([]recordedRequest(nil), transport.requests...)
+}
+
+func (transport *recordingTransport) sessionHeaders() []string {
+	requests := transport.recorded()
+	headers := make([]string, 0, len(requests))
+	for _, request := range requests {
+		headers = append(headers, request.header.Get("x-opencode-session"))
+	}
+	return headers
+}

internal/providers/openai/provider.go

@@ -65,6 +65,7 @@ type Options struct {
 	// enabled; the ZERO_DISABLE_PROMPT_CACHE_KEY env kill switch still applies
 	// on top for any endpoint.
 	DisablePromptCacheKey bool
+	SessionHeader         string
 }
 
 // Provider streams completions from an OpenAI-compatible chat completions API.
@@ -85,6 +86,7 @@ type Provider struct {
 	parseThinkTags        bool
 	setRequestExtra       func(*http.Request)
 	disablePromptCacheKey bool
+	sessionHeader         providerio.SessionHeader
 }
 
 // New creates an OpenAI-compatible provider.
@@ -140,6 +142,7 @@ func New(options Options) (*Provider, error) {
 		parseThinkTags:        options.ParseThinkTags,
 		setRequestExtra:       options.SetRequestExtra,
 		disablePromptCacheKey: options.DisablePromptCacheKey,
+		sessionHeader:         providerio.NewSessionHeader(options.SessionHeader),
 	}, nil
 }
 
@@ -156,13 +159,13 @@ func (provider *Provider) StreamCompletion(
 	events := make(chan zeroruntime.StreamEvent, 16)
 	go func() {
 		defer close(events)
-		provider.stream(ctx, body, events)
+		provider.stream(ctx, body, request.PromptCacheKey, events)
 	}()
 
 	return events, nil
 }
 
-func (provider *Provider) stream(ctx context.Context, body []byte, events chan<- zeroruntime.StreamEvent) {
+func (provider *Provider) stream(ctx context.Context, body []byte, sessionID string, events chan<- zeroruntime.StreamEvent) {
 	endpoint := provider.endpoint
 
 	// streamCtx lets the idle watchdog abort an in-flight body read by cancelling
@@ -186,6 +189,7 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
 		},
 		provider.oauthResolver,
 		func(request *http.Request) {
+			provider.sessionHeader.Apply(request, sessionID)
 			request.Header.Set("Content-Type", "application/json")
 			if provider.userAgent != "" {
 				request.Header.Set("User-Agent", provider.userAgent)

internal/providers/providerio/headers.go

@@ -1,6 +1,7 @@
 package providerio
 
 import (
+	"crypto/rand"
 	"net/http"
 	"strings"
 )
@@ -52,6 +53,28 @@ func ApplyAuthHeaders(request *http.Request, options AuthHeaders) {
 	request.Header.Set(header, value)
 }
 
+type SessionHeader struct {
+	name     string
+	fallback string
+}
+
+func NewSessionHeader(name string) SessionHeader {
+	if name == "" {
+		return SessionHeader{}
+	}
+	return SessionHeader{name: name, fallback: rand.Text()}
+}
+
+func (header SessionHeader) Apply(request *http.Request, sessionID string) {
+	if header.name == "" {
+		return
+	}
+	if sessionID = strings.TrimSpace(sessionID); sessionID == "" {
+		sessionID = header.fallback
+	}
+	request.Header.Set(header.name, sessionID)
+}
+
 func CopyHeaders(headers map[string]string) map[string]string {
 	if headers == nil {
 		return nil

internal/swarm/lifecycle_test.go

@@ -8,6 +8,8 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
+
+	"github.com/Gitlawb/zero/internal/testutil"
 )
 
 // controllableLauncher records every launched spec and lets a test control each
@@ -85,18 +87,6 @@ func newSwarmForWithSize(t *testing.T, l MemberLauncher, maxTeamSize int) *Swarm
 	return sw
 }
 
-func waitFor(t *testing.T, what string, cond func() bool) {
-	t.Helper()
-	deadline := time.Now().Add(3 * time.Second)
-	for time.Now().Before(deadline) {
-		if cond() {
-			return
-		}
-		time.Sleep(5 * time.Millisecond)
-	}
-	t.Fatalf("timed out waiting for %s", what)
-}
-
 func okFor(spec MemberSpec, _ int) (MemberResult, error) {
 	return MemberResult{Result: "ok:" + spec.Task, SessionID: "sess-" + spec.ID}, nil
 }
@@ -108,7 +98,7 @@ func TestSpawnCompletes(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Spawn: %v", err)
 	}
-	waitFor(t, "task done", func() bool {
+	testutil.WaitFor(t, "task done", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusDone
 	})
@@ -125,7 +115,7 @@ func TestSpawnInheritsPolicy(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Spawn: %v", err)
 	}
-	waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
+	testutil.WaitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
 	spec := l.recorded()[0]
 	if spec.Model != "orch-model" {
 		t.Fatalf("member model = %q, want inherited orch-model", spec.Model)
@@ -158,7 +148,7 @@ func TestConcurrencyCapAndQueueDrains(t *testing.T) {
 	}
 	// Release everyone; the queue should drain one-per-slot until all are done.
 	close(gate)
-	waitFor(t, "all tasks done", func() bool { return sw.Coordinator().Summarize().Done == 5 })
+	testutil.WaitFor(t, "all tasks done", func() bool { return sw.Coordinator().Summarize().Done == 5 })
 	if team.Running() != 0 || team.QueueDepth() != 0 {
 		t.Fatalf("after drain running=%d queue=%d, want 0/0", team.Running(), team.QueueDepth())
 	}
@@ -176,7 +166,7 @@ func TestRetryOnTemporaryError(t *testing.T) {
 	})
 	sw := newSwarmFor(t, l)
 	id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
-	waitFor(t, "task recovered", func() bool {
+	testutil.WaitFor(t, "task recovered", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusDone
 	})
@@ -195,7 +185,7 @@ func TestRetryExhaustionFails(t *testing.T) {
 	})
 	sw := newSwarmFor(t, l)
 	id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
-	waitFor(t, "task failed", func() bool {
+	testutil.WaitFor(t, "task failed", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusFailed
 	})
@@ -210,7 +200,7 @@ func TestPermanentErrorNoRetry(t *testing.T) {
 	})
 	sw := newSwarmFor(t, l)
 	id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
-	waitFor(t, "task failed", func() bool {
+	testutil.WaitFor(t, "task failed", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusFailed
 	})
@@ -296,7 +286,7 @@ func TestClosePreventsMemberRetry(t *testing.T) {
 		sw.Close()
 		close(closed)
 	}()
-	waitFor(t, "swarm closed state", func() bool {
+	testutil.WaitFor(t, "swarm closed state", func() bool {
 		sw.lifecycleMu.RLock()
 		defer sw.lifecycleMu.RUnlock()
 		return sw.closed
@@ -325,7 +315,7 @@ func TestCloseWaitsForMemberWatchers(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Spawn: %v", err)
 	}
-	waitFor(t, "task running", func() bool {
+	testutil.WaitFor(t, "task running", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusRunning
 	})
@@ -365,7 +355,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
 	sw := newSwarmFor(t, l)
 	pol := Policy{Model: "m"}
 	origID, _ := sw.Spawn(pol, "team", "teammate", "original task", "/w")
-	waitFor(t, "original running", func() bool {
+	testutil.WaitFor(t, "original running", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusRunning
 	})
@@ -389,7 +379,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
 	}
 	// The new member carries the handoff note in its task and preserves cwd.
 	close(gate)
-	waitFor(t, "spec for new member", func() bool {
+	testutil.WaitFor(t, "spec for new member", func() bool {
 		for _, s := range l.recorded() {
 			if s.ID == newID {
 				return true
@@ -405,7 +395,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
 		}
 	}
 	// A handoff of an already-terminal task is rejected.
-	waitFor(t, "new task done", func() bool {
+	testutil.WaitFor(t, "new task done", func() bool {
 		task, ok := sw.Coordinator().Get(newID)
 		return ok && task.Status == StatusDone
 	})
@@ -470,7 +460,7 @@ func TestHandoffMailboxFailureRestoresSourceCompletion(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Spawn: %v", err)
 	}
-	waitFor(t, "original running", func() bool {
+	testutil.WaitFor(t, "original running", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusRunning
 	})
@@ -485,7 +475,7 @@ func TestHandoffMailboxFailureRestoresSourceCompletion(t *testing.T) {
 	}
 
 	close(gate)
-	waitFor(t, "source completion after aborted handoff", func() bool {
+	testutil.WaitFor(t, "source completion after aborted handoff", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusDone
 	})
@@ -532,7 +522,7 @@ func TestHandoffSuccessorCollisionDoesNotStopSource(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Spawn: %v", err)
 	}
-	waitFor(t, "original running", func() bool {
+	testutil.WaitFor(t, "original running", func() bool {
 		task, ok := coord.Get(origID)
 		return ok && task.Status == StatusRunning
 	})
@@ -557,7 +547,7 @@ func TestHandoffSuccessorCollisionDoesNotStopSource(t *testing.T) {
 		t.Fatalf("collision misdelivered %d handoff notes to the existing task", len(msgs))
 	}
 	close(gate)
-	waitFor(t, "source completion after rejected collision", func() bool {
+	testutil.WaitFor(t, "source completion after rejected collision", func() bool {
 		task, ok := coord.Get(origID)
 		return ok && task.Status == StatusDone
 	})
@@ -643,7 +633,7 @@ func TestCloseReleasesBlockedHandoffBeforeMemberExit(t *testing.T) {
 	case <-time.After(3 * time.Second):
 		t.Fatal("Close did not return after the source member exited")
 	}
-	waitFor(t, "source completion after handoff cancellation", func() bool {
+	testutil.WaitFor(t, "source completion after handoff cancellation", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusDone
 	})
@@ -704,7 +694,7 @@ func TestHandoffTimesOutWithoutStartingSuccessor(t *testing.T) {
 	}
 
 	close(releaseOriginal)
-	waitFor(t, "source completion after handoff timeout", func() bool {
+	testutil.WaitFor(t, "source completion after handoff timeout", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusDone
 	})
@@ -876,7 +866,7 @@ func TestHandoffRemovesQueuedOriginalBeforeDispatchingSuccessor(t *testing.T) {
 	if _, err := sw.Spawn(pol, "team", "teammate", "slot blocker", "/w"); err != nil {
 		t.Fatalf("Spawn blocker: %v", err)
 	}
-	waitFor(t, "blocker running", func() bool { return len(l.recorded()) == 1 })
+	testutil.WaitFor(t, "blocker running", func() bool { return len(l.recorded()) == 1 })
 	origID, err := sw.Spawn(pol, "team", "teammate", "queued original", "/w")
 	if err != nil {
 		t.Fatalf("Spawn queued original: %v", err)
@@ -890,7 +880,7 @@ func TestHandoffRemovesQueuedOriginalBeforeDispatchingSuccessor(t *testing.T) {
 		t.Fatalf("Handoff: %v", err)
 	}
 	close(gate)
-	waitFor(t, "successor launch", func() bool {
+	testutil.WaitFor(t, "successor launch", func() bool {
 		for _, spec := range l.recorded() {
 			if spec.ID == newID {
 				return true
@@ -923,7 +913,7 @@ func TestAdoptOrphans(t *testing.T) {
 		t.Fatalf("adopted = %v, want [orphan-1]", adopted)
 	}
 	// The orphan is relaunched under a fresh agent and completes.
-	waitFor(t, "orphan completed", func() bool {
+	testutil.WaitFor(t, "orphan completed", func() bool {
 		task, ok := sw.Coordinator().Get("orphan-1")
 		return ok && task.Status == StatusDone
 	})
@@ -964,7 +954,7 @@ func TestCollectScopesToTeam(t *testing.T) {
 	sw := newSwarmFor(t, l)
 	a, _ := sw.Spawn(Policy{}, "alpha", "teammate", "ta", "")
 	_, _ = sw.Spawn(Policy{}, "beta", "teammate", "tb", "")
-	waitFor(t, "alpha done", func() bool {
+	testutil.WaitFor(t, "alpha done", func() bool {
 		task, ok := sw.Coordinator().Get(a)
 		return ok && task.Status == StatusDone
 	})
@@ -1110,7 +1100,7 @@ func TestAdmittedDispatchDoesNotLaunchAfterClose(t *testing.T) {
 		sw.Close()
 		close(closed)
 	}()
-	waitFor(t, "swarm closed state", func() bool {
+	testutil.WaitFor(t, "swarm closed state", func() bool {
 		sw.lifecycleMu.RLock()
 		defer sw.lifecycleMu.RUnlock()
 		return sw.closed
@@ -1170,7 +1160,7 @@ func TestCloseBetweenLaunchPrecheckAndReturn(t *testing.T) {
 		sw.Close()
 		close(closed)
 	}()
-	waitFor(t, "swarm closed state", func() bool {
+	testutil.WaitFor(t, "swarm closed state", func() bool {
 		sw.lifecycleMu.RLock()
 		defer sw.lifecycleMu.RUnlock()
 		return sw.closed
@@ -1268,7 +1258,7 @@ func TestCloseBetweenRetryLaunchPrecheckAndReturn(t *testing.T) {
 		sw.Close()
 		close(closed)
 	}()
-	waitFor(t, "swarm closed state", func() bool {
+	testutil.WaitFor(t, "swarm closed state", func() bool {
 		sw.lifecycleMu.RLock()
 		defer sw.lifecycleMu.RUnlock()
 		return sw.closed
@@ -1460,7 +1450,7 @@ func TestCloseFailsQueuedSpecOnLateCreatedTeam(t *testing.T) {
 		sw.Close()
 		close(closed)
 	}()
-	waitFor(t, "swarm closed state", func() bool {
+	testutil.WaitFor(t, "swarm closed state", func() bool {
 		sw.lifecycleMu.RLock()
 		defer sw.lifecycleMu.RUnlock()
 		return sw.closed

internal/swarm/scheduler_test.go

@@ -4,6 +4,8 @@ import (
 	"math"
 	"testing"
 	"time"
+
+	"github.com/Gitlawb/zero/internal/testutil"
 )
 
 // testTicker returns a ticker factory backed by a single unbounded-handshake
@@ -39,11 +41,11 @@ func TestSchedulerFiresAndCountsRuns(t *testing.T) {
 	for i := 0; i < 3; i++ {
 		ticks <- time.Time{}
 		want := i + 1
-		waitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == want })
+		testutil.WaitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == want })
 	}
 
 	// After MaxRuns the job retires itself.
-	waitFor(t, "job retired", func() bool { _, ok := findJob(sched.List(), id); return !ok })
+	testutil.WaitFor(t, "job retired", func() bool { _, ok := findJob(sched.List(), id); return !ok })
 	if got := len(l.recorded()); got != 3 {
 		t.Fatalf("spawned %d members, want 3", got)
 	}
@@ -68,11 +70,11 @@ func TestSchedulerSkipsWhilePreviousRuns(t *testing.T) {
 
 	// Fire 1: spawns and the member stays running (gated).
 	ticks <- time.Time{}
-	waitFor(t, "first spawn", func() bool { return len(l.recorded()) == 1 })
+	testutil.WaitFor(t, "first spawn", func() bool { return len(l.recorded()) == 1 })
 
 	// Fire 2: previous still running => skipped, no new spawn.
 	ticks <- time.Time{}
-	waitFor(t, "skip recorded", func() bool {
+	testutil.WaitFor(t, "skip recorded", func() bool {
 		j, ok := findJob(sched.List(), id)
 		return ok && j.Skipped == 1
 	})
@@ -82,14 +84,14 @@ func TestSchedulerSkipsWhilePreviousRuns(t *testing.T) {
 
 	// Release the first member, then fire 3: previous terminal => spawns again.
 	close(gate)
-	waitFor(t, "first done", func() bool { return sw.Coordinator().Summarize().Done == 1 })
+	testutil.WaitFor(t, "first done", func() bool { return sw.Coordinator().Summarize().Done == 1 })
 	ticks <- time.Time{}
-	waitFor(t, "second spawn", func() bool { return len(l.recorded()) == 2 })
+	testutil.WaitFor(t, "second spawn", func() bool { return len(l.recorded()) == 2 })
 	// fireIfIdle's spawn (what "second spawn" observes via the launcher) and
 	// run's subsequent job.incRuns() are sequential but distinct steps in the
 	// scheduler's goroutine; wait for Runs itself rather than assuming the
 	// launcher recording it means the job's counter is updated too.
-	waitFor(t, "second run recorded", func() bool {
+	testutil.WaitFor(t, "second run recorded", func() bool {
 		j, ok := findJob(sched.List(), id)
 		return ok && j.Runs == 2
 	})
@@ -215,7 +217,7 @@ func TestSchedulerDailyRecomputesNextDelay(t *testing.T) {
 	}
 	for i := 0; i < 2; i++ {
 		ticks <- time.Time{}
-		waitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == i+1 })
+		testutil.WaitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == i+1 })
 	}
 	if len(delays) < 2 {
 		t.Fatalf("expected at least two requested delays, got %v", delays)

internal/swarm/tools_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 	"unicode/utf8"
 
+	"github.com/Gitlawb/zero/internal/testutil"
 	"github.com/Gitlawb/zero/internal/tools"
 )
 
@@ -94,7 +95,7 @@ func TestSpawnToolThroughRegistry(t *testing.T) {
 	if id == "" {
 		t.Fatal("spawn must return a task_id in Meta")
 	}
-	waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
+	testutil.WaitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
 	spec := l.recorded()[0]
 	if spec.Model != "m1" || spec.Cwd != "/work" {
 		t.Fatalf("policy/cwd not threaded into member: model=%q cwd=%q", spec.Model, spec.Cwd)
@@ -165,7 +166,7 @@ func TestStatusAndCollectTools(t *testing.T) {
 		"agent_type": "teammate", "task": "compute", "team": "alpha",
 	}, tools.RunOptions{PermissionGranted: true, Model: "m"})
 	id := spawn.Meta["task_id"]
-	waitFor(t, "task done", func() bool {
+	testutil.WaitFor(t, "task done", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusDone
 	})
@@ -213,7 +214,7 @@ func TestCollectBlocksUntilMembersFinish(t *testing.T) {
 		"agent_type": "teammate", "task": "compute", "team": "alpha",
 	}, tools.RunOptions{PermissionGranted: true, Model: "m"})
 	id := spawn.Meta["task_id"]
-	waitFor(t, "member running", func() bool {
+	testutil.WaitFor(t, "member running", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusRunning
 	})
@@ -258,7 +259,7 @@ func TestCollectReturnsPartialOnTimeout(t *testing.T) {
 		"agent_type": "teammate", "task": "compute", "team": "alpha",
 	}, tools.RunOptions{PermissionGranted: true, Model: "m"})
 	id := spawn.Meta["task_id"]
-	waitFor(t, "member running", func() bool {
+	testutil.WaitFor(t, "member running", func() bool {
 		task, ok := sw.Coordinator().Get(id)
 		return ok && task.Status == StatusRunning
 	})
@@ -324,7 +325,7 @@ func TestNewWithUserDefinitions(t *testing.T) {
 	if _, err := sw.Spawn(Policy{Model: "m"}, "team", "researcher", "find things", ""); err != nil {
 		t.Fatalf("Spawn user-defined agent: %v", err)
 	}
-	waitFor(t, "researcher used", func() bool { return used })
+	testutil.WaitFor(t, "researcher used", func() bool { return used })
 
 	// A bad definition makes New fail closed.
 	if _, err := New(Options{BaseDir: t.TempDir(), Launcher: newLauncher(okFor), Definitions: []Definition{{AgentType: "  "}}}); err == nil {
@@ -341,7 +342,7 @@ func TestHandoffToolThroughRegistry(t *testing.T) {
 		"agent_type": "teammate", "task": "long", "team": "alpha",
 	}, tools.RunOptions{PermissionGranted: true, Model: "m"})
 	origID := spawn.Meta["task_id"]
-	waitFor(t, "running", func() bool {
+	testutil.WaitFor(t, "running", func() bool {
 		task, ok := sw.Coordinator().Get(origID)
 		return ok && task.Status == StatusRunning
 	})

internal/testutil/waitfor.go

@@ -0,0 +1,23 @@
+// Package testutil provides small, shared helpers for tests across the
+// codebase.
+package testutil
+
+import (
+	"testing"
+	"time"
+)
+
+// WaitFor polls cond until it returns true, or fails the test if it does not
+// become true within a generous deadline. what names the condition being
+// waited for so a timeout failure message is actionable.
+func WaitFor(t *testing.T, what string, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatalf("timed out waiting for %s", what)
+}