mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-09 05:19:33 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
45cec9c6cf | ||
|
|
347db49978 | ||
|
|
99721c762f |
No files matched your search
@@ -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
|
||||
|
||||
@@ -246,20 +246,62 @@ func RedactResult(result Result) Result {
|
||||
return redacted
|
||||
}
|
||||
|
||||
// redactMetadata redacts EVERY text field of the metadata, not the subset a
|
||||
// reader happens to think is user-facing. `zero search --json` prints this
|
||||
// straight to stdout, and the fields carrying the most free text are the ones
|
||||
// least likely to be on a hand-kept list: SpecUserComment and SpecRejectReason
|
||||
// are review text the user typed, SpecDraftReasoning is the model's own prose,
|
||||
// and Goal.Objective and Goal.StatusReason are whatever the goal was phrased
|
||||
// as. A field added to sessions.Metadata later is covered by
|
||||
// TestRedactMetadataCoversEveryTextField, which walks the struct rather than
|
||||
// repeating this list, so the two cannot drift apart silently.
|
||||
func redactMetadata(session sessions.Metadata, options redaction.Options) sessions.Metadata {
|
||||
session.SessionID = redaction.RedactString(session.SessionID, options)
|
||||
session.SessionKind = sessions.SessionKind(redaction.RedactString(string(session.SessionKind), options))
|
||||
session.Title = redaction.RedactString(session.Title, options)
|
||||
session.Cwd = redaction.RedactString(session.Cwd, options)
|
||||
session.ModelID = redaction.RedactString(session.ModelID, options)
|
||||
session.Provider = redaction.RedactString(session.Provider, options)
|
||||
session.Tag = redaction.RedactString(session.Tag, options)
|
||||
session.ParentSessionID = redaction.RedactString(session.ParentSessionID, options)
|
||||
session.RootSessionID = redaction.RedactString(session.RootSessionID, options)
|
||||
session.AgentName = redaction.RedactString(session.AgentName, options)
|
||||
session.TaskID = redaction.RedactString(session.TaskID, options)
|
||||
session.ForkedFromEventID = redaction.RedactString(session.ForkedFromEventID, options)
|
||||
session.SpawnedFromEventID = redaction.RedactString(session.SpawnedFromEventID, options)
|
||||
session.SpecID = redaction.RedactString(session.SpecID, options)
|
||||
session.SpecFilePath = redaction.RedactString(session.SpecFilePath, options)
|
||||
session.SpecStatus = sessions.SpecStatus(redaction.RedactString(string(session.SpecStatus), options))
|
||||
session.SpecDraftModelID = redaction.RedactString(session.SpecDraftModelID, options)
|
||||
session.SpecDraftReasoning = redaction.RedactString(session.SpecDraftReasoning, options)
|
||||
session.SpecUserComment = redaction.RedactString(session.SpecUserComment, options)
|
||||
session.SpecRejectReason = redaction.RedactString(session.SpecRejectReason, options)
|
||||
session.SpecSourceSessionID = redaction.RedactString(session.SpecSourceSessionID, options)
|
||||
session.SpecImplSessionID = redaction.RedactString(session.SpecImplSessionID, options)
|
||||
session.CreatedAt = redaction.RedactString(session.CreatedAt, options)
|
||||
session.UpdatedAt = redaction.RedactString(session.UpdatedAt, options)
|
||||
session.LastEventType = sessions.EventType(redaction.RedactString(string(session.LastEventType), options))
|
||||
session.Goal = redactGoal(session.Goal, options)
|
||||
return session
|
||||
}
|
||||
|
||||
// redactGoal returns a redacted COPY. Metadata is passed by value, but Goal is
|
||||
// a pointer inside it, so redacting through it would rewrite the goal the
|
||||
// caller still holds: a search would quietly edit the live session record it
|
||||
// was only supposed to report on.
|
||||
func redactGoal(goal *sessions.Goal, options redaction.Options) *sessions.Goal {
|
||||
if goal == nil {
|
||||
return nil
|
||||
}
|
||||
redacted := *goal
|
||||
redacted.Objective = redaction.RedactString(redacted.Objective, options)
|
||||
redacted.Status = sessions.GoalStatus(redaction.RedactString(string(redacted.Status), options))
|
||||
redacted.StatusReason = redaction.RedactString(redacted.StatusReason, options)
|
||||
redacted.CreatedAt = redaction.RedactString(redacted.CreatedAt, options)
|
||||
redacted.UpdatedAt = redaction.RedactString(redacted.UpdatedAt, options)
|
||||
return &redacted
|
||||
}
|
||||
|
||||
func redactEventSummary(event EventSummary, options redaction.Options) EventSummary {
|
||||
event.ID = redaction.RedactString(event.ID, options)
|
||||
event.Type = sessions.EventType(redaction.RedactString(string(event.Type), options))
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
)
|
||||
|
||||
@@ -107,3 +110,156 @@ func fixedSearchClock(value string) func() time.Time {
|
||||
}
|
||||
return func() time.Time { return parsed }
|
||||
}
|
||||
|
||||
// EVERY TEXT FIELD, NOT A HAND-KEPT SUBSET.
|
||||
//
|
||||
// `zero search --json` prints Hit.Session straight to stdout, so a field
|
||||
// redactMetadata forgets is a field that reaches the terminal with whatever was
|
||||
// in it. The list in redactMetadata is written out by hand; this walks the
|
||||
// struct instead, so a string field added to sessions.Metadata (or to the Goal
|
||||
// hanging off it) fails here on the day it lands rather than the day someone
|
||||
// notices it in output. It named twenty uncovered fields when it was written,
|
||||
// including the review text the user types and the model's own draft reasoning.
|
||||
func TestRedactMetadataCoversEveryTextField(t *testing.T) {
|
||||
const secret = "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaa0123456789ABCD"
|
||||
|
||||
// The premise: the secret is one RedactString already recognizes, so a field
|
||||
// that still carries it was not passed through the redactor at all.
|
||||
if redaction.RedactString(secret, redaction.Options{}) == secret {
|
||||
t.Fatalf("SETUP INVALID: %q is not a shape RedactString redacts, so every leg below passes vacuously", secret)
|
||||
}
|
||||
|
||||
var session sessions.Metadata
|
||||
filled := fillMetadataText(t, reflect.ValueOf(&session).Elem(), secret, 0)
|
||||
|
||||
var leaked []string
|
||||
var checked int
|
||||
walkMetadataText(t, reflect.ValueOf(redactMetadata(session, redaction.Options{})), "", 0, func(name, value string) {
|
||||
checked++
|
||||
if strings.Contains(value, secret) {
|
||||
leaked = append(leaked, name)
|
||||
}
|
||||
})
|
||||
|
||||
// Filling and walking are separate traversals of the same shape, so a field
|
||||
// one of them understands and the other does not would otherwise show up as
|
||||
// a quiet pass rather than a failure.
|
||||
if checked != filled {
|
||||
t.Fatalf("SETUP INVALID: filled %d text fields and walked %d, so the two traversals disagree about the shape of Metadata", filled, checked)
|
||||
}
|
||||
if filled == 0 {
|
||||
t.Fatal("SETUP INVALID: no text fields were populated, so this proves nothing")
|
||||
}
|
||||
if len(leaked) > 0 {
|
||||
t.Errorf("redactMetadata left the secret in %d of %d text fields: %s",
|
||||
len(leaked), checked, strings.Join(leaked, ", "))
|
||||
}
|
||||
}
|
||||
|
||||
// fillMetadataText writes text into every text-bearing field reachable from
|
||||
// value, allocating nil pointers and one-element slices on the way, and returns
|
||||
// how many strings it set.
|
||||
//
|
||||
// Recursive on purpose. The first version of this walked only the direct string
|
||||
// fields of one struct and hand-allocated the Goal next to it, which left the
|
||||
// guard with the same blind spot it exists to close: a nested struct added to
|
||||
// sessions.Metadata later would arrive empty, the walk would find nothing in it
|
||||
// and the test would pass while the new field went unredacted. Raised by
|
||||
// CodeRabbit on the PR.
|
||||
//
|
||||
// It fails rather than skips on a kind it does not understand, because a field
|
||||
// this cannot populate is a field the guard is not guarding.
|
||||
func fillMetadataText(t *testing.T, value reflect.Value, text string, depth int) int {
|
||||
t.Helper()
|
||||
if depth > 6 {
|
||||
t.Fatalf("fillMetadataText recursed past depth %d, which suggests a cycle in the type", depth)
|
||||
}
|
||||
if !value.CanSet() {
|
||||
return 0
|
||||
}
|
||||
switch value.Kind() {
|
||||
case reflect.String:
|
||||
value.SetString(text)
|
||||
return 1
|
||||
case reflect.Pointer:
|
||||
if value.IsNil() {
|
||||
value.Set(reflect.New(value.Type().Elem()))
|
||||
}
|
||||
return fillMetadataText(t, value.Elem(), text, depth+1)
|
||||
case reflect.Struct:
|
||||
filled := 0
|
||||
for index := 0; index < value.NumField(); index++ {
|
||||
if !value.Type().Field(index).IsExported() {
|
||||
continue
|
||||
}
|
||||
filled += fillMetadataText(t, value.Field(index), text, depth+1)
|
||||
}
|
||||
return filled
|
||||
case reflect.Slice:
|
||||
value.Set(reflect.MakeSlice(value.Type(), 1, 1))
|
||||
return fillMetadataText(t, value.Index(0), text, depth+1)
|
||||
case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
|
||||
reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
|
||||
reflect.Float32, reflect.Float64:
|
||||
return 0
|
||||
default:
|
||||
t.Fatalf("fillMetadataText does not know how to populate %s (kind %s); teach it and walkMetadataText together, or this guard stops covering that field",
|
||||
value.Type(), value.Kind())
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// walkMetadataText visits every text-bearing field reachable from value, in the
|
||||
// same shape fillMetadataText populates. The two are deliberately symmetric and
|
||||
// the caller compares their counts.
|
||||
func walkMetadataText(t *testing.T, value reflect.Value, prefix string, depth int, visit func(name, value string)) {
|
||||
t.Helper()
|
||||
if depth > 6 {
|
||||
t.Fatalf("walkMetadataText recursed past depth %d, which suggests a cycle in the type", depth)
|
||||
}
|
||||
switch value.Kind() {
|
||||
case reflect.String:
|
||||
visit(prefix, value.String())
|
||||
case reflect.Pointer:
|
||||
if !value.IsNil() {
|
||||
walkMetadataText(t, value.Elem(), prefix, depth+1, visit)
|
||||
}
|
||||
case reflect.Struct:
|
||||
for index := 0; index < value.NumField(); index++ {
|
||||
field := value.Type().Field(index)
|
||||
if !field.IsExported() {
|
||||
continue
|
||||
}
|
||||
name := field.Name
|
||||
if prefix != "" {
|
||||
name = prefix + "." + name
|
||||
}
|
||||
walkMetadataText(t, value.Field(index), name, depth+1, visit)
|
||||
}
|
||||
case reflect.Slice:
|
||||
for index := 0; index < value.Len(); index++ {
|
||||
walkMetadataText(t, value.Index(index), fmt.Sprintf("%s[%d]", prefix, index), depth+1, visit)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Metadata is passed by value, but Goal is a pointer inside it. Redacting
|
||||
// through that pointer would edit the session the caller still holds, so a
|
||||
// search would quietly rewrite the live record it was asked to report on.
|
||||
func TestRedactMetadataDoesNotRewriteTheCallersGoal(t *testing.T) {
|
||||
const secret = "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaa0123456789ABCD"
|
||||
goal := &sessions.Goal{Objective: "ship " + secret, StatusReason: "blocked on " + secret}
|
||||
session := sessions.Metadata{SessionID: "s", Goal: goal}
|
||||
|
||||
redacted := redactMetadata(session, redaction.Options{})
|
||||
|
||||
if goal.Objective != "ship "+secret || goal.StatusReason != "blocked on "+secret {
|
||||
t.Errorf("redactMetadata rewrote the caller's goal: %+v", *goal)
|
||||
}
|
||||
if redacted.Goal == goal {
|
||||
t.Error("redacted metadata still points at the caller's goal")
|
||||
}
|
||||
if strings.Contains(redacted.Goal.Objective, secret) || strings.Contains(redacted.Goal.StatusReason, secret) {
|
||||
t.Errorf("goal text was not redacted: %+v", *redacted.Goal)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user