Compare commits

...
Author SHA1 Message Date
Vasanthdev2004 1d034d8027 fix(daemon): stop worker process groups by their launch-time identity
Kill and the Cancel hook now go through background.TerminateCommandGroup,
which signals the group the worker was launched into (ConfigureChild-
ProcessGroup's Setpgid) instead of rediscovering it from the PID, so a
worker that has exited, or on Darwin one that is still a zombie, no longer
leaves its children running. It is TerminateCommand without the reap:
the pool waits on its workers itself and exec.Cmd calls Cancel from
inside Wait, so reaping there would be a second Wait.

On Windows it is the same tree kill, with the process handle pinned while
it runs. A command that has already been waited cannot be pinned; that
now reports os.ErrProcessDone, which exec.Cmd's Cancel contract treats
as nothing to cancel, instead of the unexported "already released" error.

Fixes #861
2026-09-26 10:07:00 +05:30
Vasanthdev2004 5eff187a79 test(daemon): stop what a worker left running after it was reaped
A daemon worker is launched into its own process group, but Kill and the
Cancel hook pass its PID to background.TerminateProcess, which looks the
group up again with Getpgid when it signals. Once the worker has exited
and been waited, that lookup fails, only the dead PID is signalled, and a
child the worker left running survives. This pins the case with a shell
worker that backgrounds a sleep and exits: after the wait, Kill and the
Cancel hook must still stop the sleep.

Committed ahead of the fix so CI shows it failing on the current
behaviour on Linux and macOS.
2026-09-26 10:06:59 +05:30
Anandan 99721c762f fix(providers): send stable session headers to OpenCode (#1060) 2026-09-16 07:58:40 +08:00
10 changed files with 822 additions and 14 deletions

No files matched your search

+8 -1
View File
@@ -4,6 +4,7 @@ package background
import (
"fmt"
"os"
"os/exec"
"github.com/Gitlawb/zero/internal/execution"
@@ -46,7 +47,13 @@ func terminateOwnedProcess(cmd *exec.Cmd) (bool, error) {
terminateErr = terminateProcessForTest(cmd.Process.Pid)
})
if err != nil {
return false, fmt.Errorf("pin process %d for termination: %w", cmd.Process.Pid, err)
// WithHandle refuses only a process that has already been waited or
// released, so it has finished. Say so in the form exec.Cmd understands:
// a Cancel error wrapping os.ErrProcessDone is "nothing to cancel", where
// the unexported "already released" error it would otherwise carry made
// Wait report a spurious cancel failure. A dead root has no tree left to
// find either way.
return true, fmt.Errorf("pin process %d for termination: %w: %w", cmd.Process.Pid, os.ErrProcessDone, err)
}
return alreadyExited, terminateErr
}
+23
View File
@@ -73,6 +73,29 @@ func TerminateCommand(cmd *exec.Cmd) error {
return nil
}
// TerminateCommandGroup stops a started command's process group on POSIX, or its
// process tree on Windows, WITHOUT reaping it.
//
// It is TerminateCommand without the Wait, for a caller whose Wait already runs
// somewhere else: a worker pool waiting on the command in its own goroutine, or
// the command's own Cancel hook, which exec.Cmd invokes while its Wait is in
// progress. TerminateCommand would be a second Wait there. Both used to pass
// cmd.Process.Pid to TerminateProcess, which rediscovers the group at signal time
// with Getpgid, and once the leader has exited that lookup fails, so only the dead
// leader was signalled and its descendants kept running. When cmd was configured
// by ConfigureChildProcessGroup the group now comes from that launch-time
// configuration instead. On Windows it is the tree kill TerminateProcess performs,
// with the process handle pinned while it runs so the PID cannot be recycled
// underneath it; once the command has been waited that pin fails with an error
// wrapping os.ErrProcessDone.
func TerminateCommandGroup(cmd *exec.Cmd) error {
if cmd == nil || cmd.Process == nil {
return nil
}
_, err := terminateOwnedProcess(cmd)
return err
}
// classifyWaitError treats a non-zero exit status as success: a process being
// terminated is expected to report one (or a signal), so the only interesting
// failure is Wait itself not working.
@@ -0,0 +1,71 @@
package background
import (
"errors"
"os"
"os/exec"
"runtime"
"testing"
"time"
)
// TerminateCommandGroup LEAVES THE WAIT TO ITS CALLER. Its callers already have a
// Wait running or about to run: the daemon pool waits on its workers, and
// exec.Cmd calls the Cancel hook from inside Wait. Reaping here, the way
// TerminateCommand does, would make their Wait fail with "Wait was already
// called" instead of reporting how the command ended.
func TestTerminateCommandGroupLeavesTheWaitToTheCaller(t *testing.T) {
cmd := exec.Command("sleep", "30")
if runtime.GOOS == "windows" {
cmd = exec.Command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command", "Start-Sleep -Seconds 30")
}
ConfigureChildProcessGroup(cmd)
if err := cmd.Start(); err != nil {
t.Fatalf("start: %v", err)
}
if err := TerminateCommandGroup(cmd); err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("TerminateCommandGroup: %v", err)
}
waited := make(chan error, 1)
go func() { waited <- cmd.Wait() }()
select {
case err := <-waited:
var exitErr *exec.ExitError
if err != nil && !errors.As(err, &exitErr) {
t.Fatalf("the caller's own Wait failed after TerminateCommandGroup: %v", err)
}
case <-time.After(10 * time.Second):
_ = cmd.Process.Kill()
t.Fatal("the command was still running 10s after TerminateCommandGroup")
}
}
// A command that was never started, or a nil one, is nothing to stop.
func TestTerminateCommandGroupWithNoProcessIsANoOp(t *testing.T) {
if err := TerminateCommandGroup(nil); err != nil {
t.Fatalf("nil command: %v", err)
}
if err := TerminateCommandGroup(exec.Command("sleep", "1")); err != nil {
t.Fatalf("unstarted command: %v", err)
}
}
// After the caller has waited, Windows cannot pin the process any more. The error
// must still read as "already finished", which is what exec.Cmd's Cancel contract
// and the pool both treat as nothing wrong.
func TestTerminateCommandGroupAfterWaitReportsProcessDone(t *testing.T) {
if runtime.GOOS != "windows" {
t.Skip("POSIX signals the group, which is simply gone after the wait")
}
cmd := exec.Command("cmd.exe", "/c", "exit", "0")
ConfigureChildProcessGroup(cmd)
if err := cmd.Run(); err != nil {
t.Fatalf("run: %v", err)
}
if err := TerminateCommandGroup(cmd); err != nil && !errors.Is(err, os.ErrProcessDone) {
t.Fatalf("after Wait the error is %v, want nil or one wrapping os.ErrProcessDone", err)
}
}
+8 -9
View File
@@ -72,9 +72,10 @@ func (w *execWorker) Kill() error {
if w.cmd.Process == nil {
return nil
}
// background.TerminateProcess is the cross-platform terminate (kills the
// process group on POSIX, taskkill /T on Windows).
return background.TerminateProcess(w.cmd.Process.Pid)
// The group the worker was launched into, not one rediscovered from its PID:
// the pool may call this after the worker has exited and been waited, and its
// children can still be running. See background.TerminateCommandGroup.
return background.TerminateCommandGroup(w.cmd)
}
// readerLines adapts a bufio.Reader to the Lines interface. Unlike a capped
@@ -154,13 +155,11 @@ func NewExecLauncher(cfg ExecLauncherConfig) (Launcher, error) {
background.ConfigureChildProcessGroup(cmd)
// CommandContext's default cancel sends os.Process.Kill to the LEADER only,
// orphaning the process group we just configured (a stuck worker's children
// would survive ctx cancellation). Terminate the whole group instead — the
// same cross-platform group terminate Kill() uses (D11).
// would survive ctx cancellation). Terminate the whole group instead, the
// same launch-time group terminate Kill() uses (D11). Cancel runs while Wait
// is in progress, so it must not reap.
cmd.Cancel = func() error {
if cmd.Process == nil {
return nil
}
return background.TerminateProcess(cmd.Process.Pid)
return background.TerminateCommandGroup(cmd)
}
stdout, err := cmd.StdoutPipe()
@@ -0,0 +1,90 @@
//go:build !windows
package daemon
import (
"context"
"errors"
"os/exec"
"strconv"
"strings"
"syscall"
"testing"
"time"
)
// A WORKER'S CHILDREN ARE STOPPED THROUGH THE GROUP IT WAS LAUNCHED INTO.
//
// Kill and the Cancel hook used to hand the worker's PID to
// background.TerminateProcess, which looks the group up again with Getpgid when it
// signals. Once the worker has exited and been waited, that lookup fails, only the
// dead PID is signalled, and whatever the worker left running survives; on Darwin
// the lookup already fails while the exited worker is still a zombie. The worker
// here is a shell that backgrounds a sleep, prints its PID and exits, so after the
// wait the sleep is the only process left in the worker's group. Reported in #861.
func TestExecWorkerStopsWhatItLeftRunningAfterItExits(t *testing.T) {
for _, stop := range []struct {
name string
run func(*execWorker) error
}{
{"Kill", func(worker *execWorker) error { return worker.Kill() }},
{"Cancel hook", func(worker *execWorker) error { return worker.cmd.Cancel() }},
} {
t.Run(stop.name, func(t *testing.T) {
launcher, err := NewExecLauncher(ExecLauncherConfig{
Executable: "/bin/sh",
BaseArgs: []string{"-c", "sleep 300 >/dev/null 2>&1 & echo $!"},
})
if err != nil {
t.Fatalf("NewExecLauncher: %v", err)
}
handle, err := launcher(context.Background(), WorkerSpec{})
if err != nil {
t.Fatalf("launch: %v", err)
}
worker, ok := handle.(*execWorker)
if !ok {
t.Fatalf("SETUP INVALID: the launcher returned %T, not the exec worker under test", handle)
}
line, ok, err := worker.Stdout().Next()
if err != nil || !ok {
t.Fatalf("read the child's pid: ok=%v err=%v", ok, err)
}
childPID, err := strconv.Atoi(strings.TrimSpace(line))
if err != nil {
t.Fatalf("parse the child's pid %q: %v", line, err)
}
t.Cleanup(func() { _ = syscall.Kill(childPID, syscall.SIGKILL) })
if _, err := worker.Wait(); err != nil {
t.Fatalf("wait for the worker: %v", err)
}
if childStopped(childPID) {
t.Fatalf("SETUP INVALID: the worker's child %d is not running, so there is nothing left to stop", childPID)
}
// The pool ignores what Kill returns; the outcome is what matters.
_ = stop.run(worker)
deadline := time.Now().Add(5 * time.Second)
for !childStopped(childPID) {
if time.Now().After(deadline) {
t.Fatalf("the worker's child %d survived %s after the worker was reaped", childPID, stop.name)
}
time.Sleep(20 * time.Millisecond)
}
})
}
}
// childStopped reports whether pid is gone or only waiting to be reaped.
func childStopped(pid int) bool {
if errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH) {
return true
}
state, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
if err != nil {
return errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH)
}
return strings.HasPrefix(strings.TrimSpace(string(state)), "Z")
}
+6 -2
View File
@@ -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 != "" {
+20
View File
@@ -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
+567
View File
@@ -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
}
+6 -2
View File
@@ -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)
+23
View File
@@ -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