mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 17:56:14 -10:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1d034d8027 | ||
|
|
5eff187a79 | ||
|
|
99721c762f |
No files matched your search
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user