diff --git a/README.md b/README.md index f617d3a..3d948cb 100644 --- a/README.md +++ b/README.md @@ -128,6 +128,7 @@ ocswitch opencode sync 这个命令会做一件事：把当前可路由的 alias 列表同步进 OpenCode 的 `provider.ocswitch.models`。 +注意：如果目标文件原本是 JSONC，`sync` 写回时会规范化成普通 JSON，因此注释和尾逗号不会保留。 默认行为： - 优先复用全局 OpenCode 配置文件：`opencode.jsonc` > `opencode.json` > `config.json` @@ -229,8 +230,7 @@ ocswitch provider import-opencode --from ./examples/opencode.jsonc - `npm: @ai-sdk/openai` - 有 `options.baseURL` -- 有 `options.apiKey` - +- `options.apiKey` 可以为空；导入后由 `ocswitch doctor` / `serve` 前校验帮助你发现风险 注意： - 这不是完整迁移工具 diff --git a/README_EN.md b/README_EN.md index 9a203bc..9bc5c24 100644 --- a/README_EN.md +++ b/README_EN.md @@ -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 different file. -Only `@ai-sdk/openai` custom providers with a `baseURL` and `apiKey` are -imported. Everything else is out of MVP scope. - +Only `@ai-sdk/openai` custom providers with a `baseURL` are imported. An empty +`apiKey` is allowed and kept as-is so you can complete credentials later. ### Doctor (static) ```bash diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index 868df99..997cbc2 100644 --- a/internal/cli/cli_test.go +++ b/internal/cli/cli_test.go @@ -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/ form" { + t.Fatalf("error = %q", got) + } +} + func TestHelpTextIncludesOperationalGuidance(t *testing.T) { tests := []struct { name string diff --git a/internal/cli/opencode.go b/internal/cli/opencode.go index 8d55e87..31c4411 100644 --- a/internal/cli/opencode.go +++ b/internal/cli/opencode.go @@ -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 "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 OPENCODE_CONFIG_DIR unless you pass --target yourself. The command writes alias exposure into provider.ocswitch.models using only aliases that are currently routable. - 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 needed.`, @@ -77,20 +81,26 @@ needed.`, return err } 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 raw["model"] != setModel { - raw["model"] = setModel - changed = true + if err := validateSyncedModelSelection(setModel, aliasNames, "--set-model"); err != nil { + return err } } if setSmallModel != "" { - if raw["small_model"] != setSmallModel { - raw["small_model"] = setSmallModel - changed = true + if err := validateSyncedModelSelection(setSmallModel, aliasNames, "--set-small-model"); err != nil { + return err } } + 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 { fmt.Fprintf(cmd.OutOrStdout(), "✓ no changes required at %s\n", path) return nil diff --git a/internal/cli/opencode_validate.go b/internal/cli/opencode_validate.go new file mode 100644 index 0000000..3c8be3d --- /dev/null +++ b/internal/cli/opencode_validate.go @@ -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/ form", flagName) + } + alias := strings.TrimPrefix(value, prefix) + if alias == "" { + return fmt.Errorf("%s must use the ocswitch/ 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, ", ")) +} diff --git a/internal/cli/provider.go b/internal/cli/provider.go index 4761284..1b70a2f 100644 --- a/internal/cli/provider.go +++ b/internal/cli/provider.go @@ -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 --from when you want a different file. -Only config-defined @ai-sdk/openai custom providers with both baseURL and apiKey -are imported. Unsupported provider shapes are skipped by design. Existing ocswitch -providers are skipped unless --overwrite is given. - +Only config-defined @ai-sdk/openai custom providers with baseURL are imported. +Providers with an empty apiKey are allowed and kept as-is. Unsupported provider +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.`, Example: ` ocswitch provider import-opencode ocswitch provider import-opencode --from /path/to/opencode.jsonc diff --git a/internal/config/config.go b/internal/config/config.go index b823bc3..c151d3c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -5,11 +5,14 @@ import ( "encoding/json" "errors" "fmt" + "net" "os" "path/filepath" "sort" "strings" "sync" + + "github.com/Apale7/opencode-provider-switch/internal/fileutil" ) const ( @@ -146,14 +149,10 @@ func (c *Config) Path() string { return c.path } // Save writes config atomically. func (c *Config) Save() error { c.mu.RLock() - defer c.mu.RUnlock() if c.path == "" { c.path = DefaultPath() } - if err := os.MkdirAll(filepath.Dir(c.path), 0o755); err != nil { - return fmt.Errorf("mkdir: %w", err) - } - // sort for stability + path := c.path providers := append([]Provider(nil), c.Providers...) sort.Slice(providers, func(i, j int) bool { return providers[i].ID < providers[j].ID }) aliases := append([]Alias(nil), c.Aliases...) @@ -163,16 +162,18 @@ func (c *Config) Save() error { Providers []Provider `json:"providers"` Aliases []Alias `json:"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, "", " ") if err != nil { return fmt.Errorf("marshal: %w", err) } data = append(data, '\n') - tmp := c.path + ".tmp" - if err := os.WriteFile(tmp, data, 0o600); err != nil { - return fmt.Errorf("write tmp: %w", err) - } - return os.Rename(tmp, c.path) + return fileutil.WithLockedFile(path, func() error { + return fileutil.AtomicWriteFile(path, data, 0o600) + }) } // 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 { 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 } +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() +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 678f1b9..5a9824c 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,6 +1,7 @@ package config import ( + "strings" "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) { t.Parallel() diff --git a/internal/fileutil/fileutil.go b/internal/fileutil/fileutil.go new file mode 100644 index 0000000..2279468 --- /dev/null +++ b/internal/fileutil/fileutil.go @@ -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 +} diff --git a/internal/opencode/opencode.go b/internal/opencode/opencode.go index 9dc4100..8c77a5a 100644 --- a/internal/opencode/opencode.go +++ b/internal/opencode/opencode.go @@ -13,6 +13,7 @@ import ( "reflect" "sort" + "github.com/Apale7/opencode-provider-switch/internal/fileutil" "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 { return fmt.Errorf("mkdir: %w", err) } - data, err := renderSaveData(path, raw) - if err != nil { - return err - } - tmp := path + ".tmp" - if err := os.WriteFile(tmp, data, 0o600); err != nil { - return fmt.Errorf("write tmp: %w", err) - } - return os.Rename(tmp, path) + return fileutil.WithLockedFile(path, func() error { + data, err := renderSaveData(path, raw) + if err != nil { + return err + } + if err := fileutil.AtomicWriteFile(path, data, 0o600); err != nil { + return err + } + return nil + }) } - func renderSaveData(path string, raw Raw) ([]byte, error) { original, err := os.ReadFile(path) if err != nil { diff --git a/internal/opencode/opencode_test.go b/internal/opencode/opencode_test.go index 59ffd7f..a41dd92 100644 --- a/internal/opencode/opencode_test.go +++ b/internal/opencode/opencode_test.go @@ -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) { raw := Raw{ "provider": map[string]any{ diff --git a/internal/proxy/server.go b/internal/proxy/server.go index f4d192a..9064c70 100644 --- a/internal/proxy/server.go +++ b/internal/proxy/server.go @@ -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 import ( @@ -21,11 +18,19 @@ import ( ) var firstByteTimeout = 15 * time.Second +var requestReadTimeout = 30 * time.Second +var streamIdleTimeout = 60 * time.Second type openAIErrorEnvelope struct { Error openAIError `json:"error"` } +type upstreamFailure struct { + status int + header http.Header + body []byte +} + type openAIError struct { Message string `json:"message"` Type string `json:"type,omitempty"` @@ -52,15 +57,14 @@ func New(cfg *config.Config) *Server { TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: 1 * time.Second, ResponseHeaderTimeout: firstByteTimeout, - // no response buffering so streams flow through immediately - DisableCompression: false, - ForceAttemptHTTP2: true, + DisableCompression: false, + ForceAttemptHTTP2: true, } return &Server{ cfg: cfg, client: &http.Client{ Transport: transport, - Timeout: 0, // streaming, no overall timeout + Timeout: 0, }, logger: log.New(log.Writer(), "[ocswitch] ", log.LstdFlags|log.Lmicroseconds), } @@ -80,10 +84,10 @@ func (s *Server) ListenAndServe(ctx context.Context) error { Addr: addr, Handler: mux, ReadHeaderTimeout: 10 * time.Second, + ReadTimeout: requestReadTimeout, } errCh := make(chan error, 1) go func() { errCh <- srv.ListenAndServe() }() - s.logger.Printf("listening on http://%s", addr) select { case <-ctx.Done(): 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)) 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 } var payload map[string]any @@ -182,6 +187,7 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { } failoverCount := 0 + var lastRetryable *upstreamFailure for attempt, t := range targets { p := s.cfg.FindProvider(t.Provider) if p == nil || !p.IsEnabled() { @@ -190,7 +196,6 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { 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) - // rewrite payload.model for this upstream cloned := cloneMap(payload) cloned["model"] = t.Model newBody, err := json.Marshal(cloned) @@ -200,24 +205,34 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { 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 { return } 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) return } + if failure != nil { + lastRetryable = failure + } s.logger.Printf("req=%d alias=%s attempt=%d retryable: %v", reqID, aliasName, attempt+1, upstreamErr) 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)) } -// 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. // retryable=true means failure happened before any bytes flushed downstream. func (s *Server) tryOnce( @@ -230,11 +245,14 @@ func (s *Server) tryOnce( aliasName string, attempt 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" upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) 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) upReq.Header.Set("Content-Type", "application/json") @@ -250,46 +268,41 @@ func (s *Server) tryOnce( startedAt := time.Now() resp, err := s.client.Do(upReq) 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() - if resp.StatusCode >= 500 || resp.StatusCode == 429 { - // drain up to small cap for logging context - peek, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) - return false, true, fmt.Errorf("upstream %d: %s", resp.StatusCode, truncate(string(peek), 200)) + if resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests { + failure = captureRetryableFailure(resp) + return false, true, fmt.Errorf("upstream %d: %s", resp.StatusCode, truncate(string(failure.body), 200)), failure } 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.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount) copyResponseHeaders(w.Header(), resp.Header) w.WriteHeader(resp.StatusCode) _, _ = 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) if remaining <= 0 { - _ = resp.Body.Close() - return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout) + return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout), nil } firstChunk, firstErr := readFirstChunk(resp.Body, remaining) if firstErr != nil { if errors.Is(firstErr, errFirstByteTimeout) { - _ = resp.Body.Close() - return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout) + return false, true, fmt.Errorf("upstream first byte timeout after %s", firstByteTimeout), nil } if errors.Is(firstErr, io.EOF) { if len(firstChunk) == 0 { firstChunk = nil } } 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.writeDebugHeaders(w, aliasName, provider.ID, target.Model, attempt, failoverCount) copyResponseHeaders(w.Header(), resp.Header) @@ -297,7 +310,7 @@ func (s *Server) tryOnce( flusher, _ := w.(http.Flusher) if len(firstChunk) > 0 { if _, werr := w.Write(firstChunk); werr != nil { - return true, false, werr + return true, false, werr, nil } if flusher != nil { flusher.Flush() @@ -305,10 +318,10 @@ func (s *Server) tryOnce( } buf := make([]byte, 16<<10) for { - n, rerr := resp.Body.Read(buf) + n, rerr := readChunkWithTimeout(resp.Body, buf, streamIdleTimeout) if n > 0 { if _, werr := w.Write(buf[:n]); werr != nil { - return true, false, werr + return true, false, werr, nil } if flusher != nil { flusher.Flush() @@ -316,15 +329,15 @@ func (s *Server) tryOnce( } if rerr != nil { 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 + return true, false, rerr, nil } } } var errFirstByteTimeout = errors.New("first byte timeout") +var errStreamIdleTimeout = errors.New("stream idle timeout") func readFirstChunk(r io.Reader, timeout time.Duration) ([]byte, error) { 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) { h := w.Header() h.Set("Content-Type", "application/json") @@ -370,6 +421,18 @@ func errorTypeForStatus(status int) string { 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 { prefix := config.AppName + "/" 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)) } -// hopByHopHeaders lists headers that must not be forwarded per RFC 7230. var hopByHopHeaders = map[string]bool{ "Connection": true, "Proxy-Connection": true, @@ -404,16 +466,15 @@ var hopByHopHeaders = map[string]bool{ "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) { + connectionHeaders := connectionDeclaredHeaders(src) for k, vs := range src { ck := http.CanonicalHeaderKey(k) - if hopByHopHeaders[ck] { + if hopByHopHeaders[ck] || connectionHeaders[ck] { continue } 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 } 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) { for k, vs := range src { ck := http.CanonicalHeaderKey(k) @@ -430,7 +503,6 @@ func copyResponseHeaders(dst, src http.Header) { continue } if ck == "Content-Length" { - // We may transform nothing, but streaming responses often omit this. continue } 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 { out := make(map[string]any, len(m)) for k, v := range m { @@ -448,10 +519,9 @@ func cloneMap(m map[string]any) map[string]any { return out } -// truncate returns s truncated to at most n bytes (best-effort). func truncate(s string, n int) string { if len(s) <= n { return s } - return s[:n] + "…" + return s[:n] } diff --git a/internal/proxy/server_test.go b/internal/proxy/server_test.go index df51c50..97ea718 100644 --- a/internal/proxy/server_test.go +++ b/internal/proxy/server_test.go @@ -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) { 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) { t.Parallel() @@ -309,8 +426,12 @@ func TestReadFirstChunk(t *testing.T) { } type blockingReader 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) { time.Sleep(200 * time.Millisecond)