diff --git a/go.mod b/go.mod index 1c63dd1..347c350 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.25.0 require ( github.com/fsnotify/fsnotify v1.10.1 github.com/getlantern/systray v1.2.2 + github.com/google/uuid v1.6.0 github.com/pkoukk/tiktoken-go v0.1.8 github.com/spf13/cobra v1.10.2 golang.org/x/mod v0.40.0 @@ -22,7 +23,6 @@ require ( github.com/getlantern/hidden v0.0.0-20190325191715-f02dbb02be55 // indirect github.com/getlantern/ops v0.0.0-20190325191751-d70cb0d6f85f // indirect github.com/go-stack/stack v1.8.0 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/mattn/go-isatty v0.0.24 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect diff --git a/internal/client/opencode.go b/internal/client/opencode.go index 38d0afe..9f9e6b6 100644 --- a/internal/client/opencode.go +++ b/internal/client/opencode.go @@ -351,6 +351,16 @@ func setOpenRouterHeaders(req *http.Request) { req.Header.Set("X-OpenRouter-Categories", openRouterCategories) } +// setOpenCodeSessionHeader forwards the OpenCode session ID from ctx to the +// upstream request, but only for OpenCode Go models. Zen, Bedrock, and +// OpenRouter do not receive it. +func setOpenCodeSessionHeader(h http.Header, ctx context.Context, modelConfig config.ModelConfig) { + if config.NormalizeProvider(modelConfig.Provider) != config.ProviderOpenCodeGo { + return + } + core.SetOpenCodeSessionHeader(h, ctx) +} + // EndpointType determines which Zen endpoint format to use. type EndpointType int @@ -473,6 +483,7 @@ func (c *OpenCodeClient) ChatCompletion( if IsOpenRouter(modelConfig) { setOpenRouterHeaders(httpReq) } + setOpenCodeSessionHeader(httpReq.Header, ctx, modelConfig) if req.Stream != nil && *req.Stream { httpReq.Header.Set("Accept", "text/event-stream") @@ -575,6 +586,7 @@ func (c *OpenCodeClient) SendAnthropicRequest( httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Authorization", "Bearer "+apiKey) httpReq.Header.Set("x-api-key", apiKey) + setOpenCodeSessionHeader(httpReq.Header, ctx, modelConfig) if stream { httpReq.Header.Set("Accept", "text/event-stream") @@ -623,6 +635,7 @@ func (c *OpenCodeClient) ResponsesCompletion( if IsOpenRouter(modelConfig) { setOpenRouterHeaders(httpReq) } + setOpenCodeSessionHeader(httpReq.Header, ctx, modelConfig) resp, err := c.httpClient.Do(httpReq) if err != nil { @@ -723,6 +736,7 @@ func (c *OpenCodeClient) GeminiCompletion( if IsOpenRouter(modelConfig) { setOpenRouterHeaders(httpReq) } + setOpenCodeSessionHeader(httpReq.Header, ctx, modelConfig) resp, err := c.httpClient.Do(httpReq) if err != nil { diff --git a/internal/client/opencode_test.go b/internal/client/opencode_test.go index 57e94d0..7cd0cb3 100644 --- a/internal/client/opencode_test.go +++ b/internal/client/opencode_test.go @@ -9,6 +9,7 @@ import ( "time" "github.com/routatic/proxy/internal/config" + "github.com/routatic/proxy/internal/core" "github.com/routatic/proxy/pkg/types" ) @@ -1000,6 +1001,7 @@ func TestOpenRouterChatCompletion_UsesAttributionHeaders(t *testing.T) { {name: "referer", header: "HTTP-Referer", want: "https://github.com/routatic/proxy"}, {name: "title", header: "X-OpenRouter-Title", want: "routatic-proxy"}, {name: "category", header: "X-OpenRouter-Categories", want: "cli-agent"}, + {name: "no opencode session", header: "x-opencode-session", want: ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -1082,3 +1084,69 @@ func TestGetProviderAPIKeys_EmptyReturnsGlobal(t *testing.T) { t.Errorf("getProviderAPIKeys() = %v, want %v (should fallback to global)", got, want) } } + +func TestOpenCodeClient_SetsOpenCodeSessionHeader(t *testing.T) { + var gotSession string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession = r.Header.Get(core.OpenCodeSessionHeader) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"resp-1","object":"chat.completion","created":1,"model":"deepseek-v4-pro","choices":[],"usage":{}}`)) + })) + defer ts.Close() + + cfg := &config.Config{ + APIKey: "test-key", + OpenCodeGo: config.OpenCodeGoConfig{ + BaseURL: ts.URL, + }, + } + atomicCfg := config.NewAtomicConfig(cfg, "") + c := NewOpenCodeClient(atomicCfg, nil) + + model := config.ModelConfig{Provider: ProviderOpenCodeGo, ModelID: "deepseek-v4-pro"} + req := &types.ChatCompletionRequest{ + Model: "deepseek-v4-pro", + Messages: []types.ChatMessage{{Role: "user", Content: json.RawMessage(`"hello"`)}}, + } + ctx := core.WithSessionID(context.Background(), "client-session-1") + if _, err := c.ChatCompletionNonStreaming(ctx, "deepseek-v4-pro", req, model); err != nil { + t.Fatalf("ChatCompletionNonStreaming() error = %v", err) + } + + if gotSession != "client-session-1" { + t.Errorf("%s = %q, want %q", core.OpenCodeSessionHeader, gotSession, "client-session-1") + } +} + +func TestOpenCodeClient_ZenOmitsOpenCodeSessionHeader(t *testing.T) { + var gotSession string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession = r.Header.Get(core.OpenCodeSessionHeader) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"resp-1","object":"chat.completion","created":1,"model":"deepseek-v4-flash-free","choices":[],"usage":{}}`)) + })) + defer ts.Close() + + cfg := &config.Config{ + APIKey: "test-key", + OpenCodeZen: config.OpenCodeZenConfig{ + BaseURL: ts.URL, + }, + } + atomicCfg := config.NewAtomicConfig(cfg, "") + c := NewOpenCodeClient(atomicCfg, nil) + + model := config.ModelConfig{Provider: ProviderOpenCodeZen, ModelID: "deepseek-v4-flash-free"} + req := &types.ChatCompletionRequest{ + Model: "deepseek-v4-flash-free", + Messages: []types.ChatMessage{{Role: "user", Content: json.RawMessage(`"hello"`)}}, + } + ctx := core.WithSessionID(context.Background(), "client-session-1") + if _, err := c.ChatCompletionNonStreaming(ctx, "deepseek-v4-flash-free", req, model); err != nil { + t.Fatalf("ChatCompletionNonStreaming() error = %v", err) + } + + if gotSession != "" { + t.Errorf("OpenCode Zen must not receive %s, got %q", core.OpenCodeSessionHeader, gotSession) + } +} diff --git a/internal/core/session.go b/internal/core/session.go new file mode 100644 index 0000000..bc653da --- /dev/null +++ b/internal/core/session.go @@ -0,0 +1,33 @@ +package core + +import ( + "context" + "net/http" +) + +// OpenCodeSessionHeader is the header sent to OpenCode Go upstreams identifying +// the conversation. Its value is the Claude Code conversation UUID verbatim. +const OpenCodeSessionHeader = "x-opencode-session" + +type sessionIDKey struct{} + +// WithSessionID returns a context carrying the OpenCode session ID. +func WithSessionID(ctx context.Context, id string) context.Context { + return context.WithValue(ctx, sessionIDKey{}, id) +} + +// SessionIDFromContext returns the OpenCode session ID carried by ctx, or "" +// when the context has none. +func SessionIDFromContext(ctx context.Context) string { + id, _ := ctx.Value(sessionIDKey{}).(string) + return id +} + +// SetOpenCodeSessionHeader sets the OpenCode session header on h when ctx +// carries a session ID. Uses Set (never Add) so a reused header map cannot +// accumulate values. +func SetOpenCodeSessionHeader(h http.Header, ctx context.Context) { + if id := SessionIDFromContext(ctx); id != "" { + h.Set(OpenCodeSessionHeader, id) + } +} diff --git a/internal/handlers/messages.go b/internal/handlers/messages.go index 88ac9ed..6df86bd 100644 --- a/internal/handlers/messages.go +++ b/internal/handlers/messages.go @@ -16,6 +16,8 @@ import ( "sync/atomic" "time" + "github.com/google/uuid" + "github.com/routatic/proxy/internal/client" "github.com/routatic/proxy/internal/config" "github.com/routatic/proxy/internal/core" @@ -265,6 +267,10 @@ func (w *responseWriter) Flush() { const ( defaultKeepaliveInterval = 3 * time.Second keepaliveWriteTimeout = 5 * time.Second + + // claudeCodeSessionHeader carries the Claude Code conversation UUID. Its + // value is forwarded verbatim to OpenCode Go as x-opencode-session. + claudeCodeSessionHeader = "x-claude-code-session-id" ) // WriteKeepalive writes a keepalive comment frame (":keepalive\n\n") to the @@ -390,6 +396,18 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) } w.Header().Set("X-Request-ID", requestID) + // Resolve the OpenCode session ID from the Claude Code conversation header + // before any dispatch: this context reaches every provider call and every + // fallback attempt. Header.Get returns the first value when duplicates are + // present; the first value wins. Clients that do not send the header (curl, + // older Claude Code) get a per-request UUID so the header is always present + // on OpenCode Go. + sessionID := r.Header.Get(claudeCodeSessionHeader) + if sessionID == "" { + sessionID = uuid.NewString() + } + r = r.WithContext(core.WithSessionID(r.Context(), sessionID)) + // Rate limiting clientIP := middleware.GetClientIP(r) if !h.rateLimiter.Allow(clientIP) { diff --git a/internal/handlers/messages_test.go b/internal/handlers/messages_test.go index 89eb10f..272adb4 100644 --- a/internal/handlers/messages_test.go +++ b/internal/handlers/messages_test.go @@ -14,6 +14,8 @@ import ( "testing" "time" + "github.com/google/uuid" + "github.com/routatic/proxy/internal/client" "github.com/routatic/proxy/internal/config" "github.com/routatic/proxy/internal/core" @@ -1890,3 +1892,334 @@ func TestHandleStreaming_AnthropicRaw_NoKeepaliveInjection(t *testing.T) { t.Errorf("keepalive comment leaked into Anthropic raw stream output (concurrent write bug):\n%s", body) } } + +// newSessionPropagationHandler builds a fully-wired MessagesHandler (legacy +// client path, nil provider registry) whose OpenCode Go models hit the +// anthropic endpoint of upstreamURL. +func newSessionPropagationHandler(t *testing.T, upstreamURL string) *MessagesHandler { + t.Helper() + cfg := &config.Config{ + APIKey: "test-key", + Models: map[string]config.ModelConfig{ + "default": {Provider: "opencode-go", ModelID: "kimi-k2.6"}, + }, + Fallbacks: map[string][]config.ModelConfig{ + "default": {{Provider: "opencode-go", ModelID: "glm-5"}}, + }, + ModelOverrides: map[string]config.ModelConfig{ + "minimax-m3": {Provider: "opencode-go", ModelID: "minimax-m3"}, + }, + OpenCodeGo: config.OpenCodeGoConfig{ + AnthropicBaseURL: upstreamURL, + BaseURL: upstreamURL, + TimeoutMs: 5000, + }, + } + atomicCfg := config.NewAtomicConfig(cfg, "/tmp/test-config.json") + tokenCounter, err := token.NewCounter() + if err != nil { + t.Fatalf("NewCounter: %v", err) + } + handler := NewMessagesHandler( + client.NewOpenCodeClient(atomicCfg, nil), + nil, // providerRegistry + router.NewModelRouter(atomicCfg), + router.NewFallbackHandler(slog.Default(), 3, 30*time.Second), + tokenCounter, + metrics.New(), + nil, // captureLogger + nil, // hist + nil, // storage + ) + handler.logger = slog.Default() + return handler +} + +// TestHandleMessages_PropagatesClaudeCodeSessionID pins that the inbound +// x-claude-code-session-id is forwarded verbatim as x-opencode-session on both +// the streaming and non-streaming paths. +func TestHandleMessages_PropagatesClaudeCodeSessionID(t *testing.T) { + const sessionID = "test-session-0001" + + for _, tc := range []struct { + name string + stream bool + }{ + {name: "streaming", stream: true}, + {name: "non-streaming"}, + } { + t.Run(tc.name, func(t *testing.T) { + var gotSession string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession = r.Header.Get(core.OpenCodeSessionHeader) + if tc.stream { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprintf(w, "event: message_start\ndata: {}\n\n") + _, _ = fmt.Fprintf(w, "event: message_stop\ndata: {}\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"minimax-m3","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`)) + } + })) + defer upstream.Close() + + handler := newSessionPropagationHandler(t, upstream.URL) + + streamField := "false" + if tc.stream { + streamField = "true" + } + requestBody := fmt.Sprintf(`{"model":"minimax-m3","stream":%s,"max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}`, streamField) + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(requestBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-claude-code-session-id", sessionID) + + handler.HandleMessages(recorder, req) + + if gotSession != sessionID { + t.Errorf("upstream %s = %q, want %q", core.OpenCodeSessionHeader, gotSession, sessionID) + } + }) + } +} + +// TestHandleMessages_SessionIDStableAcrossRequests pins the invariant: the same +// inbound session maps to the same upstream value, and different inbound +// sessions map to different values. +func TestHandleMessages_SessionIDStableAcrossRequests(t *testing.T) { + var sessions []string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sessions = append(sessions, r.Header.Get(core.OpenCodeSessionHeader)) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"minimax-m3","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`)) + })) + defer upstream.Close() + + handler := newSessionPropagationHandler(t, upstream.URL) + + requestBody := `{"model":"minimax-m3","max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}` + send := func(sid string) { + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(requestBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-claude-code-session-id", sid) + handler.HandleMessages(recorder, req) + } + + send("session-A") + send("session-A") + send("session-B") + + if len(sessions) != 3 { + t.Fatalf("expected 3 upstream calls, got %d", len(sessions)) + } + if sessions[0] != "session-A" || sessions[1] != "session-A" { + t.Errorf("same inbound session must map to the same upstream value, got %v", sessions[:2]) + } + if sessions[2] != "session-B" { + t.Errorf("different inbound session must map to a different upstream value, got %q", sessions[2]) + } + if sessions[0] == sessions[2] { + t.Error("different inbound sessions collided to the same upstream value") + } +} + +// TestHandleMessages_MissingSessionID_UsesUUIDFallback pins that clients that +// do not send x-claude-code-session-id get a per-request UUID fallback so the +// upstream header is always present. +func TestHandleMessages_MissingSessionID_UsesUUIDFallback(t *testing.T) { + var gotSession string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession = r.Header.Get(core.OpenCodeSessionHeader) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"minimax-m3","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`)) + })) + defer upstream.Close() + + handler := newSessionPropagationHandler(t, upstream.URL) + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"minimax-m3","max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}`)) + req.Header.Set("Content-Type", "application/json") + + handler.HandleMessages(recorder, req) + + if _, err := uuid.Parse(gotSession); err != nil { + t.Errorf("fallback session %q is not a valid UUID: %v", gotSession, err) + } +} + +// TestHandleStreaming_FallbackAttemptsShareSessionID pins that every fallback +// attempt for one inbound request carries the same session value: the first +// model fails, the second succeeds, and both upstream requests share the +// inbound session. +func TestHandleStreaming_FallbackAttemptsShareSessionID(t *testing.T) { + var sessions []string + callCount := int32(0) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sessions = append(sessions, r.Header.Get(core.OpenCodeSessionHeader)) + if atomic.AddInt32(&callCount, 1) == 1 { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprintf(w, "event: message_start\ndata: {}\n\n") + _, _ = fmt.Fprintf(w, "event: message_stop\ndata: {}\n\n") + })) + defer upstream.Close() + + handler := newSessionPropagationHandler(t, upstream.URL) + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"minimax-m3","stream":true,"max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-claude-code-session-id", "session-fallback-1") + + handler.HandleMessages(recorder, req) + + if len(sessions) != 2 { + t.Fatalf("expected 2 upstream calls (1 fail + 1 fallback), got %d", len(sessions)) + } + if sessions[0] != "session-fallback-1" || sessions[1] != "session-fallback-1" { + t.Errorf("all fallback attempts must share the session, got %v", sessions) + } +} + +// sessionRecordingProvider is a fake OpenCode Go provider that records the +// session ID it receives via context and returns a minimal Anthropic SSE body. +type sessionRecordingProvider struct { + session *string +} + +func (p *sessionRecordingProvider) Name() string { return "opencode-go" } +func (p *sessionRecordingProvider) Capabilities() core.ProviderCapabilities { + return core.ProviderCapabilities{SupportsStreaming: true, SupportsTools: true} +} +func (p *sessionRecordingProvider) ModelCapabilities(string) (core.ProviderCapabilities, bool) { + return p.Capabilities(), true +} +func (p *sessionRecordingProvider) WireFormat(config.ModelConfig) core.WireFormat { + return core.WireFormatAnthropic +} +func (p *sessionRecordingProvider) Execute(context.Context, *core.NormalizedRequest, config.ModelConfig) (*core.ExecuteResult, error) { + return nil, nil +} +func (p *sessionRecordingProvider) Stream(ctx context.Context, _ *core.NormalizedRequest, _ config.ModelConfig) (io.ReadCloser, error) { + *p.session = core.SessionIDFromContext(ctx) + return io.NopCloser(strings.NewReader("event: message_start\ndata: {}\n\nevent: message_stop\ndata: {}\n\n")), nil +} +func (p *sessionRecordingProvider) RoundTripName(model config.ModelConfig) string { + return model.ModelID +} +func (p *sessionRecordingProvider) StreamIdleTimeout(config.ModelConfig) time.Duration { + return time.Minute +} + +// TestHandleStreaming_SessionIDReachesProvider pins that the session ID reaches +// the provider interface via context on the registry (provider) dispatch path. +func TestHandleStreaming_SessionIDReachesProvider(t *testing.T) { + var gotSession string + registry := core.NewProviderRegistry() + _ = registry.Register(&sessionRecordingProvider{session: &gotSession}) + + cfg := &config.Config{ + APIKey: "test-key", + Models: map[string]config.ModelConfig{ + "default": {Provider: "opencode-go", ModelID: "kimi-k2.6"}, + }, + Fallbacks: map[string][]config.ModelConfig{ + "default": {{Provider: "opencode-go", ModelID: "glm-5"}}, + }, + ModelOverrides: map[string]config.ModelConfig{ + "minimax-m3": {Provider: "opencode-go", ModelID: "minimax-m3"}, + }, + OpenCodeGo: config.OpenCodeGoConfig{BaseURL: "http://127.0.0.1:1", TimeoutMs: 5000}, + } + atomicCfg := config.NewAtomicConfig(cfg, "/tmp/test-config.json") + tokenCounter, err := token.NewCounter() + if err != nil { + t.Fatalf("NewCounter: %v", err) + } + handler := NewMessagesHandler( + client.NewOpenCodeClient(atomicCfg, nil), + registry, + router.NewModelRouter(atomicCfg), + router.NewFallbackHandler(slog.Default(), 3, 30*time.Second), + tokenCounter, + metrics.New(), + nil, // captureLogger + nil, // hist + nil, // storage + ) + handler.logger = slog.Default() + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"minimax-m3","stream":true,"max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-claude-code-session-id", "session-provider-1") + + handler.HandleMessages(recorder, req) + + if gotSession != "session-provider-1" { + t.Errorf("provider received session %q via context, want %q", gotSession, "session-provider-1") + } +} + +// TestHandleMessages_OpenCodeZen_OmitsSessionHeader pins that the session ID +// does not leak onto OpenCode Zen requests even when the inbound header is +// present. +func TestHandleMessages_OpenCodeZen_OmitsSessionHeader(t *testing.T) { + var gotSession string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession = r.Header.Get(core.OpenCodeSessionHeader) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"cmpl-1","object":"chat.completion","created":1,"model":"deepseek-v4-flash-free","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`)) + })) + defer upstream.Close() + + cfg := &config.Config{ + APIKey: "test-key", + Models: map[string]config.ModelConfig{ + "default": {Provider: "opencode-zen", ModelID: "deepseek-v4-flash-free"}, + }, + Fallbacks: map[string][]config.ModelConfig{ + "default": {{Provider: "opencode-zen", ModelID: "mimo-v2.5-free"}}, + }, + ModelOverrides: map[string]config.ModelConfig{ + "deepseek-v4-flash-free": {Provider: "opencode-zen", ModelID: "deepseek-v4-flash-free"}, + }, + OpenCodeZen: config.OpenCodeZenConfig{ + BaseURL: upstream.URL, + TimeoutMs: 5000, + }, + } + atomicCfg := config.NewAtomicConfig(cfg, "/tmp/test-config.json") + tokenCounter, err := token.NewCounter() + if err != nil { + t.Fatalf("NewCounter: %v", err) + } + handler := NewMessagesHandler( + client.NewOpenCodeClient(atomicCfg, nil), + nil, // providerRegistry + router.NewModelRouter(atomicCfg), + router.NewFallbackHandler(slog.Default(), 3, 30*time.Second), + tokenCounter, + metrics.New(), + nil, // captureLogger + nil, // hist + nil, // storage + ) + handler.logger = slog.Default() + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"deepseek-v4-flash-free","max_tokens":256,"messages":[{"role":"user","content":"Say hello"}]}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-claude-code-session-id", "session-zen-1") + + handler.HandleMessages(recorder, req) + + if gotSession != "" { + t.Errorf("OpenCode Zen must not receive %s, got %q", core.OpenCodeSessionHeader, gotSession) + } +} diff --git a/internal/provider/opencode_go.go b/internal/provider/opencode_go.go index da6d8f4..ae48e7e 100644 --- a/internal/provider/opencode_go.go +++ b/internal/provider/opencode_go.go @@ -257,6 +257,7 @@ func (p *OpenCodeGoProvider) executeAnthropic(ctx context.Context, req *core.Nor httpReq.Header.Set("Authorization", "Bearer "+apiKey) httpReq.Header.Set("User-Agent", upstreamUserAgent) httpReq.Header.Set("x-api-key", apiKey) + core.SetOpenCodeSessionHeader(httpReq.Header, ctx) start := time.Now() resp, err := p.httpClient.Do(httpReq) @@ -302,6 +303,7 @@ func (p *OpenCodeGoProvider) streamAnthropic(ctx context.Context, req *core.Norm httpReq.Header.Set("User-Agent", upstreamUserAgent) httpReq.Header.Set("x-api-key", apiKey) httpReq.Header.Set("Accept", "text/event-stream") + core.SetOpenCodeSessionHeader(httpReq.Header, ctx) resp, err := p.httpClient.Do(httpReq) if err != nil { @@ -332,6 +334,7 @@ func (p *OpenCodeGoProvider) doRequest(ctx context.Context, endpoint, apiKey str httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Authorization", "Bearer "+apiKey) httpReq.Header.Set("User-Agent", upstreamUserAgent) + core.SetOpenCodeSessionHeader(httpReq.Header, ctx) if stream { httpReq.Header.Set("Accept", "text/event-stream") } diff --git a/internal/provider/opencode_session_test.go b/internal/provider/opencode_session_test.go new file mode 100644 index 0000000..d300bac --- /dev/null +++ b/internal/provider/opencode_session_test.go @@ -0,0 +1,246 @@ +package provider + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/routatic/proxy/internal/config" + "github.com/routatic/proxy/internal/core" + "github.com/routatic/proxy/pkg/types" +) + +const testSessionID = "session-123e4567-e89b-12d3-a456-426614174000" + +// sessionAssertServer returns a server that asserts the x-opencode-session +// header on every upstream request, then delegates to handler. +func sessionAssertServer(t *testing.T, want string, handler func(w http.ResponseWriter, r *http.Request)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get(core.OpenCodeSessionHeader); got != want { + t.Errorf("%s = %q, want %q", core.OpenCodeSessionHeader, got, want) + } + handler(w, r) + })) +} + +// sessionAbsentServer returns a server that asserts the header is NOT present. +func sessionAbsentServer(t *testing.T, handler func(w http.ResponseWriter, r *http.Request)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get(core.OpenCodeSessionHeader); got != "" { + t.Errorf("%s = %q, want absent", core.OpenCodeSessionHeader, got) + } + handler(w, r) + })) +} + +func chatCompletionJSON(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(types.ChatCompletionResponse{ + ID: "cmpl-test", + Model: "test-model", + Choices: []types.Choice{ + {Index: 0, Message: types.ChatMessage{Role: "assistant", Content: json.RawMessage(`"hi"`)}, FinishReason: "stop"}, + }, + Usage: types.UsageInfo{PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2}, + }) +} + +func anthropicJSON(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg-test","content":[{"type":"text","text":"hi"}]}`)) +} + +func responsesJSON(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(types.ResponsesResponse{ + ID: "resp-test", Object: "response", Created: 1, Model: "muse-spark-1.2-contributor", + Output: []types.ResponsesOutput{{ + Type: "message", Role: "assistant", + Content: []types.ResponsesContent{{Type: "output_text", Text: "hi"}}, + }}, + Usage: types.ResponsesUsage{InputTokens: 1, OutputTokens: 1}, + }) +} + +func chatCompletionSSE(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) +} + +func anthropicSSE(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: message_start\n")) + _, _ = w.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\"}}\n\n")) +} + +func responsesSSE(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: [DONE]\n\n")) +} + +// TestOpenCodeGoProvider_SetsOpenCodeSessionHeader pins that every OpenCode Go +// request builder (chat, anthropic, and responses; execute and stream) forwards +// the session ID from the context verbatim as x-opencode-session. +func TestOpenCodeGoProvider_SetsOpenCodeSessionHeader(t *testing.T) { + testCases := []struct { + name string + model config.ModelConfig + stream bool + cfg func(serverURL string) config.Config + respond func(w http.ResponseWriter, r *http.Request) + }{ + { + name: "chat completions execute", + model: config.ModelConfig{ModelID: "deepseek-v4-pro"}, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: u}} + }, + respond: chatCompletionJSON, + }, + { + name: "chat completions stream", + model: config.ModelConfig{ModelID: "deepseek-v4-pro"}, + stream: true, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: u}} + }, + respond: chatCompletionSSE, + }, + { + name: "anthropic execute", + model: config.ModelConfig{ModelID: "qwen3.5-plus"}, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: u, AnthropicBaseURL: u}} + }, + respond: anthropicJSON, + }, + { + name: "anthropic stream", + model: config.ModelConfig{ModelID: "qwen3.5-plus"}, + stream: true, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: u, AnthropicBaseURL: u}} + }, + respond: anthropicSSE, + }, + { + name: "responses execute", + model: config.ModelConfig{ModelID: "muse-spark-1.2-contributor", WireFormat: "responses"}, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: "http://127.0.0.1:1", ResponsesBaseURL: u}} + }, + respond: responsesJSON, + }, + { + name: "responses stream", + model: config.ModelConfig{ModelID: "muse-spark-1.2-contributor", WireFormat: "responses"}, + stream: true, + cfg: func(u string) config.Config { + return config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: "http://127.0.0.1:1", ResponsesBaseURL: u}} + }, + respond: responsesSSE, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + server := sessionAssertServer(t, testSessionID, tc.respond) + defer server.Close() + + cfg := tc.cfg(server.URL) + p := NewOpenCodeGoProvider(config.NewAtomicConfig(&cfg, "")) + + req := &core.NormalizedRequest{ + Model: tc.model.ModelID, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, + Stream: tc.stream, + } + ctx := core.WithSessionID(context.Background(), testSessionID) + + if tc.stream { + body, err := p.Stream(ctx, req, tc.model) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + defer func() { _ = body.Close() }() + buf := make([]byte, 1024) + if n, _ := body.Read(buf); n == 0 { + t.Error("Stream() returned empty body") + } + } else { + if _, err := p.Execute(ctx, req, tc.model); err != nil { + t.Fatalf("Execute() error = %v", err) + } + } + }) + } +} + +// TestOpenCodeGoProvider_NoSessionID_OmitsHeader pins that a context without a +// session ID produces no header (the handler fills the fallback UUID before it +// reaches the provider, so this only happens on direct provider use). +func TestOpenCodeGoProvider_NoSessionID_OmitsHeader(t *testing.T) { + server := sessionAbsentServer(t, chatCompletionJSON) + defer server.Close() + + cfg := &config.Config{APIKey: "test-key", OpenCodeGo: config.OpenCodeGoConfig{BaseURL: server.URL}} + p := NewOpenCodeGoProvider(config.NewAtomicConfig(cfg, "")) + + req := &core.NormalizedRequest{ + Model: "deepseek-v4-pro", + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, + } + model := config.ModelConfig{ModelID: "deepseek-v4-pro"} + if _, err := p.Execute(context.Background(), req, model); err != nil { + t.Fatalf("Execute() error = %v", err) + } +} + +// TestOpenCodeZenProvider_DoesNotReceiveSessionHeader pins the scope boundary: +// x-opencode-session is for OpenCode Go only, so a session ID in the context +// must not leak onto OpenCode Zen requests. +func TestOpenCodeZenProvider_DoesNotReceiveSessionHeader(t *testing.T) { + server := sessionAbsentServer(t, chatCompletionJSON) + defer server.Close() + + cfg := &config.Config{APIKey: "test-key", OpenCodeZen: config.OpenCodeZenConfig{BaseURL: server.URL}} + p := NewOpenCodeZenProvider(config.NewAtomicConfig(cfg, "")) + + req := &core.NormalizedRequest{ + Model: "deepseek-v4-flash-free", + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, + } + model := config.ModelConfig{ModelID: "deepseek-v4-flash-free"} + if _, err := p.Execute(core.WithSessionID(context.Background(), testSessionID), req, model); err != nil { + t.Fatalf("Execute() error = %v", err) + } +} + +// TestAWSBedrockProvider_DoesNotReceiveSessionHeader pins that the session ID +// does not leak onto Bedrock requests. +func TestAWSBedrockProvider_DoesNotReceiveSessionHeader(t *testing.T) { + server := sessionAbsentServer(t, chatCompletionJSON) + defer server.Close() + + cfg := &config.Config{ + AWSBedrock: config.AWSBedrockConfig{ + BaseURL: server.URL, + APIKey: "test-key", + }, + } + p := NewAWSBedrockProvider(config.NewAtomicConfig(cfg, "")) + + req := &core.NormalizedRequest{ + Model: "moonshotai.kimi-k2.5", + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, + } + model := config.ModelConfig{ModelID: "moonshotai.kimi-k2.5"} + if _, err := p.Execute(core.WithSessionID(context.Background(), testSessionID), req, model); err != nil { + t.Fatalf("Execute() error = %v", err) + } +}