fix: harden proxy and sync flows

This commit is contained in:
apale7 2026-04-18 00:06:38 +08:00
parent 51f030fb8c
commit 8e5b247ca1
13 changed files with 511 additions and 84 deletions

View File

@ -128,6 +128,7 @@ ocswitch opencode sync
这个命令会做一件事:把当前可路由的 alias 列表同步进 OpenCode 的 `provider.ocswitch.models` 这个命令会做一件事:把当前可路由的 alias 列表同步进 OpenCode 的 `provider.ocswitch.models`
注意:如果目标文件原本是 JSONC`sync` 写回时会规范化成普通 JSON因此注释和尾逗号不会保留。
默认行为: 默认行为:
- 优先复用全局 OpenCode 配置文件:`opencode.jsonc` > `opencode.json` > `config.json` - 优先复用全局 OpenCode 配置文件:`opencode.jsonc` > `opencode.json` > `config.json`
@ -229,8 +230,7 @@ ocswitch provider import-opencode --from ./examples/opencode.jsonc
- `npm: @ai-sdk/openai` - `npm: @ai-sdk/openai`
- 有 `options.baseURL` - 有 `options.baseURL`
- 有 `options.apiKey` - `options.apiKey` 可以为空;导入后由 `ocswitch doctor` / `serve` 前校验帮助你发现风险
注意: 注意:
- 这不是完整迁移工具 - 这不是完整迁移工具

View File

@ -56,9 +56,8 @@ The default import/sync target is the global user config only. It does not
follow `OPENCODE_CONFIG_DIR`; use `--from` or `--target` when you want a follow `OPENCODE_CONFIG_DIR`; use `--from` or `--target` when you want a
different file. different file.
Only `@ai-sdk/openai` custom providers with a `baseURL` and `apiKey` are Only `@ai-sdk/openai` custom providers with a `baseURL` are imported. An empty
imported. Everything else is out of MVP scope. `apiKey` is allowed and kept as-is so you can complete credentials later.
### Doctor (static) ### Doctor (static)
```bash ```bash

View File

@ -224,6 +224,50 @@ func TestOpencodeSyncDoesNotPanicOnSliceModelMetadata(t *testing.T) {
} }
} }
func TestOpencodeSyncRejectsInvalidSelectedModel(t *testing.T) {
t.Setenv(config.ConfigEnvVar, filepath.Join(t.TempDir(), "ocswitch.json"))
configPath = ""
cfg, err := loadCfg()
if err != nil {
t.Fatalf("loadCfg: %v", err)
}
cfg.UpsertProvider(config.Provider{ID: "p1", BaseURL: "https://example.com/v1"})
cfg.UpsertAlias(config.Alias{
Alias: "gpt-5.4",
Enabled: true,
Targets: []config.Target{{Provider: "p1", Model: "up-1", Enabled: true}},
})
if err := cfg.Save(); err != nil {
t.Fatalf("save config: %v", err)
}
target := filepath.Join(t.TempDir(), "opencode.jsonc")
cmd := newOpencodeSyncCmd()
cmd.SetArgs([]string{"--target", target, "--set-model", "ocswitch/missing"})
err = cmd.Execute()
if err == nil {
t.Fatal("expected invalid --set-model error")
}
if got := err.Error(); got != `--set-model "ocswitch/missing" is not a routable alias; available: ocswitch/gpt-5.4` {
t.Fatalf("error = %q", got)
}
if _, statErr := os.Stat(target); !os.IsNotExist(statErr) {
t.Fatalf("expected no target file write, stat err = %v", statErr)
}
}
func TestOpencodeSyncRejectsNonPrefixedSelectedModel(t *testing.T) {
err := validateSyncedModelSelection("gpt-5.4", []string{"gpt-5.4"}, "--set-model")
if err == nil {
t.Fatal("expected prefix validation error")
}
if got := err.Error(); got != "--set-model must use the ocswitch/<alias> form" {
t.Fatalf("error = %q", got)
}
}
func TestHelpTextIncludesOperationalGuidance(t *testing.T) { func TestHelpTextIncludesOperationalGuidance(t *testing.T) {
tests := []struct { tests := []struct {
name string name string

View File

@ -43,11 +43,15 @@ existing file in precedence order opencode.jsonc > opencode.json > config.json,
or creating opencode.jsonc if none exists. It does NOT touch the top-level or creating opencode.jsonc if none exists. It does NOT touch the top-level
"model" or "small_model" unless --set-model / --set-small-model are given. "model" or "small_model" unless --set-model / --set-small-model are given.
If the target file is JSONC, sync rewrites it as normalized plain JSON and any
comments/trailing commas are lost. Existing provider.ocswitch model metadata is
preserved when alias names stay the same, but this command is still a writing
operation rather than a comment-preserving patcher.
The default target scope is only the global user config path; it does not follow The default target scope is only the global user config path; it does not follow
OPENCODE_CONFIG_DIR unless you pass --target yourself. The command writes alias OPENCODE_CONFIG_DIR unless you pass --target yourself. The command writes alias
exposure into provider.ocswitch.models using only aliases that are currently exposure into provider.ocswitch.models using only aliases that are currently
routable. routable.
Use --dry-run to preview the resolved target file without writing it. Typical Use --dry-run to preview the resolved target file without writing it. Typical
workflow: run ocswitch doctor first, then sync, then start or restart ocswitch serve if workflow: run ocswitch doctor first, then sync, then start or restart ocswitch serve if
needed.`, needed.`,
@ -77,20 +81,26 @@ needed.`,
return err return err
} }
aliasNames := cfg.AvailableAliasNames() aliasNames := cfg.AvailableAliasNames()
baseURL := fmt.Sprintf("http://%s:%d/v1", cfg.Server.Host, cfg.Server.Port)
changed := opencode.EnsureOcswitchProvider(raw, baseURL, cfg.Server.APIKey, aliasNames)
if setModel != "" { if setModel != "" {
if raw["model"] != setModel { if err := validateSyncedModelSelection(setModel, aliasNames, "--set-model"); err != nil {
raw["model"] = setModel return err
changed = true
} }
} }
if setSmallModel != "" { if setSmallModel != "" {
if raw["small_model"] != setSmallModel { if err := validateSyncedModelSelection(setSmallModel, aliasNames, "--set-small-model"); err != nil {
raw["small_model"] = setSmallModel return err
changed = true
} }
} }
baseURL := fmt.Sprintf("http://%s:%d/v1", cfg.Server.Host, cfg.Server.Port)
changed := opencode.EnsureOcswitchProvider(raw, baseURL, cfg.Server.APIKey, aliasNames)
if setModel != "" && raw["model"] != setModel {
raw["model"] = setModel
changed = true
}
if setSmallModel != "" && raw["small_model"] != setSmallModel {
raw["small_model"] = setSmallModel
changed = true
}
if !changed { if !changed {
fmt.Fprintf(cmd.OutOrStdout(), "✓ no changes required at %s\n", path) fmt.Fprintf(cmd.OutOrStdout(), "✓ no changes required at %s\n", path)
return nil return nil

View File

@ -0,0 +1,33 @@
package cli
import (
"fmt"
"sort"
"strings"
)
func validateSyncedModelSelection(value string, aliases []string, flagName string) error {
const prefix = "ocswitch/"
if !strings.HasPrefix(value, prefix) {
return fmt.Errorf("%s must use the ocswitch/<alias> form", flagName)
}
alias := strings.TrimPrefix(value, prefix)
if alias == "" {
return fmt.Errorf("%s must use the ocswitch/<alias> form", flagName)
}
for _, name := range aliases {
if name == alias {
return nil
}
}
sorted := append([]string(nil), aliases...)
sort.Strings(sorted)
if len(sorted) == 0 {
return fmt.Errorf("%s requires at least one routable alias; run ocswitch alias list or doctor first", flagName)
}
choices := make([]string, 0, len(sorted))
for _, name := range sorted {
choices = append(choices, prefix+name)
}
return fmt.Errorf("%s %q is not a routable alias; available: %s", flagName, value, strings.Join(choices, ", "))
}

View File

@ -270,10 +270,10 @@ opencode.jsonc > opencode.json > config.json under ~/.config/opencode (XDG
aware). It does not follow OPENCODE_CONFIG_DIR for this default source; use aware). It does not follow OPENCODE_CONFIG_DIR for this default source; use
--from when you want a different file. --from when you want a different file.
Only config-defined @ai-sdk/openai custom providers with both baseURL and apiKey Only config-defined @ai-sdk/openai custom providers with baseURL are imported.
are imported. Unsupported provider shapes are skipped by design. Existing ocswitch Providers with an empty apiKey are allowed and kept as-is. Unsupported provider
providers are skipped unless --overwrite is given. shapes are skipped by design. Existing ocswitch providers are skipped unless
--overwrite is given.
Typical next step: run ocswitch provider list, then create aliases and bindings.`, Typical next step: run ocswitch provider list, then create aliases and bindings.`,
Example: ` ocswitch provider import-opencode Example: ` ocswitch provider import-opencode
ocswitch provider import-opencode --from /path/to/opencode.jsonc ocswitch provider import-opencode --from /path/to/opencode.jsonc

View File

@ -5,11 +5,14 @@ import (
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"net"
"os" "os"
"path/filepath" "path/filepath"
"sort" "sort"
"strings" "strings"
"sync" "sync"
"github.com/Apale7/opencode-provider-switch/internal/fileutil"
) )
const ( const (
@ -146,14 +149,10 @@ func (c *Config) Path() string { return c.path }
// Save writes config atomically. // Save writes config atomically.
func (c *Config) Save() error { func (c *Config) Save() error {
c.mu.RLock() c.mu.RLock()
defer c.mu.RUnlock()
if c.path == "" { if c.path == "" {
c.path = DefaultPath() c.path = DefaultPath()
} }
if err := os.MkdirAll(filepath.Dir(c.path), 0o755); err != nil { path := c.path
return fmt.Errorf("mkdir: %w", err)
}
// sort for stability
providers := append([]Provider(nil), c.Providers...) providers := append([]Provider(nil), c.Providers...)
sort.Slice(providers, func(i, j int) bool { return providers[i].ID < providers[j].ID }) sort.Slice(providers, func(i, j int) bool { return providers[i].ID < providers[j].ID })
aliases := append([]Alias(nil), c.Aliases...) aliases := append([]Alias(nil), c.Aliases...)
@ -163,16 +162,18 @@ func (c *Config) Save() error {
Providers []Provider `json:"providers"` Providers []Provider `json:"providers"`
Aliases []Alias `json:"aliases"` Aliases []Alias `json:"aliases"`
}{c.Server, providers, aliases} }{c.Server, providers, aliases}
c.mu.RUnlock()
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return fmt.Errorf("mkdir: %w", err)
}
data, err := json.MarshalIndent(snap, "", " ") data, err := json.MarshalIndent(snap, "", " ")
if err != nil { if err != nil {
return fmt.Errorf("marshal: %w", err) return fmt.Errorf("marshal: %w", err)
} }
data = append(data, '\n') data = append(data, '\n')
tmp := c.path + ".tmp" return fileutil.WithLockedFile(path, func() error {
if err := os.WriteFile(tmp, data, 0o600); err != nil { return fileutil.AtomicWriteFile(path, data, 0o600)
return fmt.Errorf("write tmp: %w", err) })
}
return os.Rename(tmp, c.path)
} }
// FindProvider returns the provider with matching id or nil. // FindProvider returns the provider with matching id or nil.
@ -392,5 +393,16 @@ func (c *Config) Validate() []error {
if c.Server.Port <= 0 || c.Server.Port > 65535 { if c.Server.Port <= 0 || c.Server.Port > 65535 {
errs = append(errs, fmt.Errorf("invalid server port %d", c.Server.Port)) errs = append(errs, fmt.Errorf("invalid server port %d", c.Server.Port))
} }
if c.Server.APIKey == DefaultLocalAPIKey && !isLoopbackHost(c.Server.Host) {
errs = append(errs, fmt.Errorf("server.api_key must not use the default value when listening on non-loopback host %q", c.Server.Host))
}
return errs return errs
} }
func isLoopbackHost(host string) bool {
host = strings.TrimSpace(host)
if host == "" || strings.EqualFold(host, "localhost") {
return true
}
ip := net.ParseIP(host)
return ip != nil && ip.IsLoopback()
}

View File

@ -1,6 +1,7 @@
package config package config
import ( import (
"strings"
"testing" "testing"
) )
@ -89,6 +90,22 @@ func TestAvailableAliasNamesOnlyReturnsRoutableAliases(t *testing.T) {
} }
} }
func TestValidateRejectsDefaultKeyOnNonLoopbackHost(t *testing.T) {
t.Parallel()
cfg := &Config{
Server: Server{Host: "0.0.0.0", Port: 9982, APIKey: DefaultLocalAPIKey},
}
errs := cfg.Validate()
if len(errs) != 1 {
t.Fatalf("Validate() errors = %v, want 1 error", errs)
}
if !strings.Contains(errs[0].Error(), "must not use the default value") {
t.Fatalf("Validate() error = %q", errs[0].Error())
}
}
func TestValidateReportsAliasWithoutAvailableTargets(t *testing.T) { func TestValidateReportsAliasWithoutAvailableTargets(t *testing.T) {
t.Parallel() t.Parallel()

View File

@ -0,0 +1,96 @@
package fileutil
import (
"fmt"
"os"
"path/filepath"
"syscall"
)
type LockedFile struct {
file *os.File
}
func WithLockedFile(path string, fn func() error) error {
lock, err := acquireLock(path)
if err != nil {
return err
}
defer lock.Close()
return fn()
}
func AtomicWriteFile(path string, data []byte, perm os.FileMode) error {
dir := filepath.Dir(path)
base := filepath.Base(path)
tmp, err := os.CreateTemp(dir, base+".*.tmp")
if err != nil {
return fmt.Errorf("create tmp: %w", err)
}
tmpPath := tmp.Name()
defer func() {
_ = os.Remove(tmpPath)
}()
if err := tmp.Chmod(perm); err != nil {
_ = tmp.Close()
return fmt.Errorf("chmod tmp: %w", err)
}
if _, err := tmp.Write(data); err != nil {
_ = tmp.Close()
return fmt.Errorf("write tmp: %w", err)
}
if err := tmp.Sync(); err != nil {
_ = tmp.Close()
return fmt.Errorf("sync tmp: %w", err)
}
if err := tmp.Close(); err != nil {
return fmt.Errorf("close tmp: %w", err)
}
if err := os.Rename(tmpPath, path); err != nil {
return fmt.Errorf("rename tmp: %w", err)
}
if err := syncDir(dir); err != nil {
return err
}
return nil
}
func acquireLock(path string) (*LockedFile, error) {
lockPath := path + ".lock"
file, err := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0o600)
if err != nil {
return nil, fmt.Errorf("open lock: %w", err)
}
if err := syscall.Flock(int(file.Fd()), syscall.LOCK_EX); err != nil {
_ = file.Close()
return nil, fmt.Errorf("lock file: %w", err)
}
return &LockedFile{file: file}, nil
}
func (l *LockedFile) Close() error {
if l == nil || l.file == nil {
return nil
}
unlockErr := syscall.Flock(int(l.file.Fd()), syscall.LOCK_UN)
closeErr := l.file.Close()
if unlockErr != nil {
return fmt.Errorf("unlock file: %w", unlockErr)
}
if closeErr != nil {
return fmt.Errorf("close lock: %w", closeErr)
}
return nil
}
func syncDir(dir string) error {
f, err := os.Open(dir)
if err != nil {
return fmt.Errorf("open dir: %w", err)
}
defer f.Close()
if err := f.Sync(); err != nil {
return fmt.Errorf("sync dir: %w", err)
}
return nil
}

View File

@ -13,6 +13,7 @@ import (
"reflect" "reflect"
"sort" "sort"
"github.com/Apale7/opencode-provider-switch/internal/fileutil"
"github.com/tidwall/jsonc" "github.com/tidwall/jsonc"
) )
@ -85,17 +86,17 @@ func Save(path string, raw Raw) error {
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return fmt.Errorf("mkdir: %w", err) return fmt.Errorf("mkdir: %w", err)
} }
data, err := renderSaveData(path, raw) return fileutil.WithLockedFile(path, func() error {
if err != nil { data, err := renderSaveData(path, raw)
return err if err != nil {
} return err
tmp := path + ".tmp" }
if err := os.WriteFile(tmp, data, 0o600); err != nil { if err := fileutil.AtomicWriteFile(path, data, 0o600); err != nil {
return fmt.Errorf("write tmp: %w", err) return err
} }
return os.Rename(tmp, path) return nil
})
} }
func renderSaveData(path string, raw Raw) ([]byte, error) { func renderSaveData(path string, raw Raw) ([]byte, error) {
original, err := os.ReadFile(path) original, err := os.ReadFile(path)
if err != nil { if err != nil {

View File

@ -327,6 +327,30 @@ func TestRenderSaveDataAcceptsJSONCAndProducesValidJSON(t *testing.T) {
} }
} }
func TestImportCustomProvidersAllowsEmptyAPIKey(t *testing.T) {
raw := Raw{
"provider": map[string]any{
"openai-empty": map[string]any{
"npm": "@ai-sdk/openai",
"name": "Empty Key",
"options": map[string]any{
"baseURL": "https://example.com/v1",
},
},
},
}
imports := ImportCustomProviders(raw)
if len(imports) != 1 {
t.Fatalf("len(imports) = %d, want 1", len(imports))
}
if imports[0].ID != "openai-empty" {
t.Fatalf("id = %q, want openai-empty", imports[0].ID)
}
if imports[0].APIKey != "" {
t.Fatalf("api key = %q, want empty", imports[0].APIKey)
}
}
func TestRenderSaveDataRejectsInvalidJSONC(t *testing.T) { func TestRenderSaveDataRejectsInvalidJSONC(t *testing.T) {
raw := Raw{ raw := Raw{
"provider": map[string]any{ "provider": map[string]any{

View File

@ -1,6 +1,3 @@
// Package proxy implements the local `/v1/responses` HTTP server that resolves
// ocswitch aliases and forwards requests to upstream providers with deterministic
// pre-first-byte failover.
package proxy package proxy
import ( import (
@ -21,11 +18,19 @@ import (
) )
var firstByteTimeout = 15 * time.Second var firstByteTimeout = 15 * time.Second
var requestReadTimeout = 30 * time.Second
var streamIdleTimeout = 60 * time.Second
type openAIErrorEnvelope struct { type openAIErrorEnvelope struct {
Error openAIError `json:"error"` Error openAIError `json:"error"`
} }
type upstreamFailure struct {
status int
header http.Header
body []byte
}
type openAIError struct { type openAIError struct {
Message string `json:"message"` Message string `json:"message"`
Type string `json:"type,omitempty"` Type string `json:"type,omitempty"`
@ -52,15 +57,14 @@ func New(cfg *config.Config) *Server {
TLSHandshakeTimeout: 10 * time.Second, TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second, ExpectContinueTimeout: 1 * time.Second,
ResponseHeaderTimeout: firstByteTimeout, ResponseHeaderTimeout: firstByteTimeout,
// no response buffering so streams flow through immediately DisableCompression: false,
DisableCompression: false, ForceAttemptHTTP2: true,
ForceAttemptHTTP2: true,
} }
return &Server{ return &Server{
cfg: cfg, cfg: cfg,
client: &http.Client{ client: &http.Client{
Transport: transport, Transport: transport,
Timeout: 0, // streaming, no overall timeout Timeout: 0,
}, },
logger: log.New(log.Writer(), "[ocswitch] ", log.LstdFlags|log.Lmicroseconds), logger: log.New(log.Writer(), "[ocswitch] ", log.LstdFlags|log.Lmicroseconds),
} }
@ -80,10 +84,10 @@ func (s *Server) ListenAndServe(ctx context.Context) error {
Addr: addr, Addr: addr,
Handler: mux, Handler: mux,
ReadHeaderTimeout: 10 * time.Second, ReadHeaderTimeout: 10 * time.Second,
ReadTimeout: requestReadTimeout,
} }
errCh := make(chan error, 1) errCh := make(chan error, 1)
go func() { errCh <- srv.ListenAndServe() }() go func() { errCh <- srv.ListenAndServe() }()
s.logger.Printf("listening on http://%s", addr)
select { select {
case <-ctx.Done(): case <-ctx.Done():
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@ -147,7 +151,8 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
} }
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 50<<20)) body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 50<<20))
if err != nil { if err != nil {
writeOpenAIError(w, http.StatusBadRequest, "invalid_request_error", "read body: "+err.Error()) status, msg := requestReadError(err)
writeOpenAIError(w, status, "invalid_request_error", msg)
return return
} }
var payload map[string]any var payload map[string]any
@ -182,6 +187,7 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
} }
failoverCount := 0 failoverCount := 0
var lastRetryable *upstreamFailure
for attempt, t := range targets { for attempt, t := range targets {
p := s.cfg.FindProvider(t.Provider) p := s.cfg.FindProvider(t.Provider)
if p == nil || !p.IsEnabled() { if p == nil || !p.IsEnabled() {
@ -190,7 +196,6 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
continue continue
} }
s.logger.Printf("req=%d alias=%s attempt=%d provider=%s remote_model=%s failovers=%d", reqID, aliasName, attempt+1, p.ID, t.Model, failoverCount) s.logger.Printf("req=%d alias=%s attempt=%d provider=%s remote_model=%s failovers=%d", reqID, aliasName, attempt+1, p.ID, t.Model, failoverCount)
// rewrite payload.model for this upstream
cloned := cloneMap(payload) cloned := cloneMap(payload)
cloned["model"] = t.Model cloned["model"] = t.Model
newBody, err := json.Marshal(cloned) newBody, err := json.Marshal(cloned)
@ -200,24 +205,34 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
return return
} }
ok, retryable, upstreamErr := s.tryOnce(r.Context(), w, r, p, t, newBody, aliasName, attempt+1, failoverCount) ok, retryable, upstreamErr, failure := s.tryOnce(r.Context(), w, r, p, t, newBody, aliasName, attempt+1, failoverCount)
if ok { if ok {
return return
} }
if !retryable { if !retryable {
// final — response was either already committed or unrecoverable
s.logger.Printf("req=%d alias=%s attempt=%d final failure: %v", reqID, aliasName, attempt+1, upstreamErr) s.logger.Printf("req=%d alias=%s attempt=%d final failure: %v", reqID, aliasName, attempt+1, upstreamErr)
return return
} }
if failure != nil {
lastRetryable = failure
}
s.logger.Printf("req=%d alias=%s attempt=%d retryable: %v", reqID, aliasName, attempt+1, upstreamErr) s.logger.Printf("req=%d alias=%s attempt=%d retryable: %v", reqID, aliasName, attempt+1, upstreamErr)
failoverCount++ failoverCount++
} }
// exhausted all targets with retryable failures if lastRetryable != nil {
copyResponseHeaders(w.Header(), lastRetryable.header)
w.WriteHeader(lastRetryable.status)
if len(lastRetryable.body) > 0 {
_, _ = w.Write(lastRetryable.body)
}
return
}
writeOpenAIError(w, http.StatusBadGateway, "server_error", fmt.Sprintf("all upstream targets failed for alias %q", aliasName)) writeOpenAIError(w, http.StatusBadGateway, "server_error", fmt.Sprintf("all upstream targets failed for alias %q", aliasName))
} }
// tryOnce proxies one attempt. Returns (ok, retryable, err). // tryOnce proxies one attempt. Returns (ok, retryable, err, failure).
// ok=true means successful response fully/partially written to client. // ok=true means successful response fully/partially written to client.
// retryable=true means failure happened before any bytes flushed downstream. // retryable=true means failure happened before any bytes flushed downstream.
func (s *Server) tryOnce( func (s *Server) tryOnce(
@ -230,11 +245,14 @@ func (s *Server) tryOnce(
aliasName string, aliasName string,
attempt int, attempt int,
failoverCount int, failoverCount int,
) (ok bool, retryable bool, err error) { ) (ok bool, retryable bool, err error, failure *upstreamFailure) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
upstreamURL := strings.TrimRight(provider.BaseURL, "/") + "/responses" upstreamURL := strings.TrimRight(provider.BaseURL, "/") + "/responses"
upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body))
if err != nil { if err != nil {
return false, false, fmt.Errorf("build request: %w", err) return false, false, fmt.Errorf("build request: %w", err), nil
} }
copyForwardHeaders(upReq.Header, clientReq.Header) copyForwardHeaders(upReq.Header, clientReq.Header)
upReq.Header.Set("Content-Type", "application/json") upReq.Header.Set("Content-Type", "application/json")
@ -250,46 +268,41 @@ func (s *Server) tryOnce(
startedAt := time.Now() startedAt := time.Now()
resp, err := s.client.Do(upReq) resp, err := s.client.Do(upReq)
if err != nil { if err != nil {
return false, true, fmt.Errorf("upstream dial/transport: %w", err) return false, true, fmt.Errorf("upstream dial/transport: %w", err), nil
} }
defer resp.Body.Close() defer resp.Body.Close()
if resp.StatusCode >= 500 || resp.StatusCode == 429 { if resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests {
// drain up to small cap for logging context failure = captureRetryableFailure(resp)
peek, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) return false, true, fmt.Errorf("upstream %d: %s", resp.StatusCode, truncate(string(failure.body), 200)), failure
return false, true, fmt.Errorf("upstream %d: %s", resp.StatusCode, truncate(string(peek), 200))
} }
if resp.StatusCode >= 400 { if resp.StatusCode >= 400 {
// non-retryable: forward status + body to client
s.logger.Printf("alias=%s attempt=%d provider=%s remote_model=%s upstream_status=%d", aliasName, attempt, provider.ID, target.Model, resp.StatusCode) s.logger.Printf("alias=%s attempt=%d provider=%s remote_model=%s upstream_status=%d", aliasName, attempt, provider.ID, target.Model, resp.StatusCode)
s.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount) s.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount)
copyResponseHeaders(w.Header(), resp.Header) copyResponseHeaders(w.Header(), resp.Header)
w.WriteHeader(resp.StatusCode) w.WriteHeader(resp.StatusCode)
_, _ = io.Copy(w, resp.Body) _, _ = io.Copy(w, resp.Body)
return true, false, fmt.Errorf("upstream %d", resp.StatusCode) return true, false, fmt.Errorf("upstream %d", resp.StatusCode), nil
} }
remaining := firstByteTimeout - time.Since(startedAt) remaining := firstByteTimeout - time.Since(startedAt)
if remaining <= 0 { if remaining <= 0 {
_ = resp.Body.Close() return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout), nil
return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout)
} }
firstChunk, firstErr := readFirstChunk(resp.Body, remaining) firstChunk, firstErr := readFirstChunk(resp.Body, remaining)
if firstErr != nil { if firstErr != nil {
if errors.Is(firstErr, errFirstByteTimeout) { if errors.Is(firstErr, errFirstByteTimeout) {
_ = resp.Body.Close() return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout), nil
return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout)
} }
if errors.Is(firstErr, io.EOF) { if errors.Is(firstErr, io.EOF) {
if len(firstChunk) == 0 { if len(firstChunk) == 0 {
firstChunk = nil firstChunk = nil
} }
} else { } else {
return false, true, fmt.Errorf("upstream first read: %w", firstErr) return false, true, fmt.Errorf("upstream first read: %w", firstErr), nil
} }
} }
// 2xx: start streaming pass-through. From this point no failover is allowed.
s.logger.Printf("alias=%s attempt=%d provider=%s remote_model=%s upstream_status=%d", aliasName, attempt, provider.ID, target.Model, resp.StatusCode) s.logger.Printf("alias=%s attempt=%d provider=%s remote_model=%s upstream_status=%d", aliasName, attempt, provider.ID, target.Model, resp.StatusCode)
s.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount) s.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount)
copyResponseHeaders(w.Header(), resp.Header) copyResponseHeaders(w.Header(), resp.Header)
@ -297,7 +310,7 @@ func (s *Server) tryOnce(
flusher, _ := w.(http.Flusher) flusher, _ := w.(http.Flusher)
if len(firstChunk) > 0 { if len(firstChunk) > 0 {
if _, werr := w.Write(firstChunk); werr != nil { if _, werr := w.Write(firstChunk); werr != nil {
return true, false, werr return true, false, werr, nil
} }
if flusher != nil { if flusher != nil {
flusher.Flush() flusher.Flush()
@ -305,10 +318,10 @@ func (s *Server) tryOnce(
} }
buf := make([]byte, 16<<10) buf := make([]byte, 16<<10)
for { for {
n, rerr := resp.Body.Read(buf) n, rerr := readChunkWithTimeout(resp.Body, buf, streamIdleTimeout)
if n > 0 { if n > 0 {
if _, werr := w.Write(buf[:n]); werr != nil { if _, werr := w.Write(buf[:n]); werr != nil {
return true, false, werr return true, false, werr, nil
} }
if flusher != nil { if flusher != nil {
flusher.Flush() flusher.Flush()
@ -316,15 +329,15 @@ func (s *Server) tryOnce(
} }
if rerr != nil { if rerr != nil {
if errors.Is(rerr, io.EOF) { if errors.Is(rerr, io.EOF) {
return true, false, nil return true, false, nil, nil
} }
// mid-stream error — already committed, cannot failover return true, false, rerr, nil
return true, false, rerr
} }
} }
} }
var errFirstByteTimeout = errors.New("first byte timeout") var errFirstByteTimeout = errors.New("first byte timeout")
var errStreamIdleTimeout = errors.New("stream idle timeout")
func readFirstChunk(r io.Reader, timeout time.Duration) ([]byte, error) { func readFirstChunk(r io.Reader, timeout time.Duration) ([]byte, error) {
type result struct { type result struct {
@ -350,6 +363,44 @@ func readFirstChunk(r io.Reader, timeout time.Duration) ([]byte, error) {
} }
} }
func readChunkWithTimeout(r io.Reader, buf []byte, timeout time.Duration) (int, error) {
type result struct {
n int
err error
}
ch := make(chan result, 1)
go func() {
n, err := r.Read(buf)
ch <- result{n: n, err: err}
}()
select {
case res := <-ch:
return res.n, res.err
case <-time.After(timeout):
return 0, errStreamIdleTimeout
}
}
func captureRetryableFailure(resp *http.Response) *upstreamFailure {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 32<<10))
return &upstreamFailure{
status: resp.StatusCode,
header: cloneHeaderSubset(resp.Header, "Content-Type", "Retry-After"),
body: body,
}
}
func cloneHeaderSubset(src http.Header, names ...string) http.Header {
dst := make(http.Header)
for _, name := range names {
ck := http.CanonicalHeaderKey(name)
for _, v := range src.Values(ck) {
dst.Add(ck, v)
}
}
return dst
}
func writeOpenAIError(w http.ResponseWriter, status int, code, message string) { func writeOpenAIError(w http.ResponseWriter, status int, code, message string) {
h := w.Header() h := w.Header()
h.Set("Content-Type", "application/json") h.Set("Content-Type", "application/json")
@ -370,6 +421,18 @@ func errorTypeForStatus(status int) string {
return "invalid_request_error" return "invalid_request_error"
} }
func requestReadError(err error) (int, string) {
var netErr net.Error
switch {
case errors.As(err, &netErr) && netErr.Timeout():
return http.StatusRequestTimeout, "request body read timeout"
case strings.Contains(strings.ToLower(err.Error()), "timeout"):
return http.StatusRequestTimeout, "request body read timeout"
default:
return http.StatusBadRequest, "read body: " + err.Error()
}
}
func normalizeAliasName(model string) string { func normalizeAliasName(model string) string {
prefix := config.AppName + "/" prefix := config.AppName + "/"
if strings.HasPrefix(model, prefix) { if strings.HasPrefix(model, prefix) {
@ -391,7 +454,6 @@ func (s *Server) writeDebugHeaders(w http.ResponseWriter, alias, provider, remot
h.Set("X-OCSWITCH-Failover-Count", fmt.Sprintf("%d", failoverCount)) h.Set("X-OCSWITCH-Failover-Count", fmt.Sprintf("%d", failoverCount))
} }
// hopByHopHeaders lists headers that must not be forwarded per RFC 7230.
var hopByHopHeaders = map[string]bool{ var hopByHopHeaders = map[string]bool{
"Connection": true, "Connection": true,
"Proxy-Connection": true, "Proxy-Connection": true,
@ -404,16 +466,15 @@ var hopByHopHeaders = map[string]bool{
"Upgrade": true, "Upgrade": true,
} }
// copyForwardHeaders copies safe request headers from client to upstream,
// dropping Authorization (the upstream key replaces it) and hop-by-hop headers.
func copyForwardHeaders(dst, src http.Header) { func copyForwardHeaders(dst, src http.Header) {
connectionHeaders := connectionDeclaredHeaders(src)
for k, vs := range src { for k, vs := range src {
ck := http.CanonicalHeaderKey(k) ck := http.CanonicalHeaderKey(k)
if hopByHopHeaders[ck] { if hopByHopHeaders[ck] || connectionHeaders[ck] {
continue continue
} }
switch ck { switch ck {
case "Authorization", "X-Api-Key", "Host", "Content-Length": case "Authorization", "X-Api-Key", "Host", "Content-Length", "Transfer-Encoding", "Forwarded", "X-Forwarded-For", "X-Forwarded-Host", "X-Forwarded-Proto", "Via":
continue continue
} }
for _, v := range vs { for _, v := range vs {
@ -422,7 +483,19 @@ func copyForwardHeaders(dst, src http.Header) {
} }
} }
// copyResponseHeaders copies upstream response headers into client response. func connectionDeclaredHeaders(src http.Header) map[string]bool {
declared := map[string]bool{}
for _, raw := range src.Values("Connection") {
for _, part := range strings.Split(raw, ",") {
name := http.CanonicalHeaderKey(strings.TrimSpace(part))
if name != "" {
declared[name] = true
}
}
}
return declared
}
func copyResponseHeaders(dst, src http.Header) { func copyResponseHeaders(dst, src http.Header) {
for k, vs := range src { for k, vs := range src {
ck := http.CanonicalHeaderKey(k) ck := http.CanonicalHeaderKey(k)
@ -430,7 +503,6 @@ func copyResponseHeaders(dst, src http.Header) {
continue continue
} }
if ck == "Content-Length" { if ck == "Content-Length" {
// We may transform nothing, but streaming responses often omit this.
continue continue
} }
for _, v := range vs { for _, v := range vs {
@ -439,7 +511,6 @@ func copyResponseHeaders(dst, src http.Header) {
} }
} }
// cloneMap performs a shallow copy of a top-level map.
func cloneMap(m map[string]any) map[string]any { func cloneMap(m map[string]any) map[string]any {
out := make(map[string]any, len(m)) out := make(map[string]any, len(m))
for k, v := range m { for k, v := range m {
@ -448,10 +519,9 @@ func cloneMap(m map[string]any) map[string]any {
return out return out
} }
// truncate returns s truncated to at most n bytes (best-effort).
func truncate(s string, n int) string { func truncate(s string, n int) string {
if len(s) <= n { if len(s) <= n {
return s return s
} }
return s[:n] + "…" return s[:n]
} }

View File

@ -157,6 +157,111 @@ func TestHandleResponsesDoesNotFailOverOn400(t *testing.T) {
} }
} }
func TestHandleResponsesReturnsLastRetryableFailure(t *testing.T) {
t.Parallel()
first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.Header().Set("Retry-After", "7")
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"error":{"message":"rate limited"}}`))
}))
defer first.Close()
second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadGateway)
_, _ = w.Write([]byte(`{"error":{"message":"upstream unavailable"}}`))
}))
defer second.Close()
srv := New(&config.Config{
Server: config.Server{APIKey: config.DefaultLocalAPIKey},
Providers: []config.Provider{
{ID: "p1", BaseURL: first.URL + "/v1"},
{ID: "p2", BaseURL: second.URL + "/v1"},
},
Aliases: []config.Alias{{
Alias: "gpt-5.4",
Enabled: true,
Targets: []config.Target{{Provider: "p1", Model: "up-1", Enabled: true}, {Provider: "p2", Model: "up-2", Enabled: true}},
}},
})
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","stream":true}`))
req.Header.Set("Authorization", "Bearer "+config.DefaultLocalAPIKey)
req.Header.Set("Content-Type", "application/json")
rr := httptest.NewRecorder()
srv.handleResponses(rr, req)
if rr.Code != http.StatusBadGateway {
t.Fatalf("status = %d, want %d", rr.Code, http.StatusBadGateway)
}
if got := rr.Header().Get("Content-Type"); got != "application/json" {
t.Fatalf("Content-Type = %q, want application/json", got)
}
if got := rr.Header().Get("Retry-After"); got != "" {
t.Fatalf("Retry-After = %q, want empty from last failure", got)
}
if body := rr.Body.String(); body != `{"error":{"message":"upstream unavailable"}}` {
t.Fatalf("body = %q", body)
}
}
func TestCopyForwardHeadersDropsDynamicConnectionHeaders(t *testing.T) {
t.Parallel()
src := http.Header{}
src.Set("Connection", "X-Trace-Id, Keep-Alive")
src.Set("X-Trace-Id", "abc")
src.Set("Keep-Alive", "timeout=5")
src.Set("OpenAI-Beta", "assistants=v2")
src.Set("X-Forwarded-For", "1.2.3.4")
dst := http.Header{}
copyForwardHeaders(dst, src)
if got := dst.Get("X-Trace-Id"); got != "" {
t.Fatalf("X-Trace-Id = %q, want empty", got)
}
if got := dst.Get("Keep-Alive"); got != "" {
t.Fatalf("Keep-Alive = %q, want empty", got)
}
if got := dst.Get("X-Forwarded-For"); got != "" {
t.Fatalf("X-Forwarded-For = %q, want empty", got)
}
if got := dst.Get("OpenAI-Beta"); got != "assistants=v2" {
t.Fatalf("OpenAI-Beta = %q, want assistants=v2", got)
}
}
func TestReadChunkWithTimeout(t *testing.T) {
t.Parallel()
t.Run("returns data", func(t *testing.T) {
t.Parallel()
buf := make([]byte, 8)
n, err := readChunkWithTimeout(bytes.NewBufferString("abc"), buf, 50*time.Millisecond)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if n != 3 || string(buf[:n]) != "abc" {
t.Fatalf("read = %d %q, want 3 abc", n, string(buf[:n]))
}
})
t.Run("times out", func(t *testing.T) {
t.Parallel()
buf := make([]byte, 8)
n, err := readChunkWithTimeout(blockingReader{}, buf, 20*time.Millisecond)
if !errors.Is(err, errStreamIdleTimeout) {
t.Fatalf("err = %v, want errStreamIdleTimeout", err)
}
if n != 0 {
t.Fatalf("n = %d, want 0", n)
}
})
}
func TestHandleResponsesSkipsDisabledProviders(t *testing.T) { func TestHandleResponsesSkipsDisabledProviders(t *testing.T) {
t.Parallel() t.Parallel()
@ -260,6 +365,18 @@ func TestHandleModelsSkipsAliasesWithoutAvailableTargets(t *testing.T) {
} }
} }
func TestRequestReadErrorTimeout(t *testing.T) {
t.Parallel()
status, message := requestReadError(timeoutErr{})
if status != http.StatusRequestTimeout {
t.Fatalf("status = %d, want %d", status, http.StatusRequestTimeout)
}
if message != "request body read timeout" {
t.Fatalf("message = %q", message)
}
}
func TestReadFirstChunk(t *testing.T) { func TestReadFirstChunk(t *testing.T) {
t.Parallel() t.Parallel()
@ -309,8 +426,12 @@ func TestReadFirstChunk(t *testing.T) {
} }
type blockingReader struct{} type blockingReader struct{}
type dataEOFReader struct{} type dataEOFReader struct{}
type timeoutErr struct{}
func (timeoutErr) Error() string { return "i/o timeout" }
func (timeoutErr) Timeout() bool { return true }
func (timeoutErr) Temporary() bool { return false }
func (blockingReader) Read(p []byte) (int, error) { func (blockingReader) Read(p []byte) (int, error) {
time.Sleep(200 * time.Millisecond) time.Sleep(200 * time.Millisecond)