fix(responses): preserve context across websocket transport changes

This commit is contained in:
sususu
2026-07-23 13:07:23 +08:00
parent e8ef4d7338
commit e05ae09425
12 changed files with 1499 additions and 211 deletions
@@ -1139,6 +1139,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
body = helps.SetStringIfDifferent(body, "model", baseModel)
body = helps.SetBoolIfDifferent(body, "stream", true)
body, _ = sjson.DeleteBytes(body, "previous_response_id")
body, _ = sjson.DeleteBytes(body, "generate")
body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
body, _ = sjson.DeleteBytes(body, "safety_identifier")
body, _ = sjson.DeleteBytes(body, "stream_options")
@@ -1408,6 +1409,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
requestPath := helps.PayloadRequestPath(opts)
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
body, _ = sjson.DeleteBytes(body, "previous_response_id")
body, _ = sjson.DeleteBytes(body, "generate")
body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
body, _ = sjson.DeleteBytes(body, "safety_identifier")
body, _ = sjson.DeleteBytes(body, "stream_options")
@@ -1580,6 +1582,7 @@ func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth
body = helps.SetStringIfDifferent(body, "model", baseModel)
body, _ = sjson.DeleteBytes(body, "previous_response_id")
body, _ = sjson.DeleteBytes(body, "generate")
body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
body, _ = sjson.DeleteBytes(body, "safety_identifier")
body, _ = sjson.DeleteBytes(body, "stream_options")
@@ -224,6 +224,22 @@ func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, w
return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL)
}
func existingWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) *websocket.Conn {
if sess == nil {
return nil
}
sess.connMu.Lock()
conn := sess.conn
matches := conn != nil &&
strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) &&
strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)
sess.connMu.Unlock()
if !matches || sess.upstreamDisconnectError(conn) != nil {
return nil
}
return conn
}
func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, string, string) {
if sess == nil {
return nil, "", ""
@@ -391,14 +407,24 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
}
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
var conn *websocket.Conn
var respHS *http.Response
var errDial error
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
conn = existingWebsocketSessionConn(sess, authID, wsURL)
if conn == nil {
return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
} else {
conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
}
if errDial != nil {
bodyErr := websocketHandshakeBody(respHS)
if respHS != nil {
helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
}
if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
return e.CodexExecutor.Execute(ctx, auth, req, opts)
return resp, statusErr{code: http.StatusUpgradeRequired, msg: string(bodyErr)}
}
if respHS != nil && respHS.StatusCode > 0 {
return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
@@ -433,6 +459,14 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {
errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
if sess != nil {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
return resp, errSend
}
return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
@@ -642,7 +676,20 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
}
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
var conn *websocket.Conn
var respHS *http.Response
var errDial error
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
conn = existingWebsocketSessionConn(sess, authID, wsURL)
if conn == nil {
if sess != nil {
sess.reqMu.Unlock()
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
} else {
conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
}
var upstreamHeaders http.Header
if respHS != nil {
upstreamHeaders = respHS.Header.Clone()
@@ -653,9 +700,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
}
if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)
if sess != nil {
sess.reqMu.Unlock()
}
return nil, statusErr{code: http.StatusUpgradeRequired, msg: string(bodyErr)}
}
if respHS != nil && respHS.StatusCode > 0 {
if sess != nil {
sess.reqMu.Unlock()
}
return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
}
helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial)
@@ -680,6 +733,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
if sess != nil {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
sess.clearActive(conn, readCh)
sess.reqMu.Unlock()
if !shouldRetryCodexWebsocketSend(errSend) {
return nil, errSend
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
sess.clearActive(conn, readCh)
@@ -1759,6 +1821,14 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession,
}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) {
e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true)
}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) {
e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false)
}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) {
if sess == nil || conn == nil {
return
}
@@ -1779,7 +1849,9 @@ func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSes
sess.connMu.Unlock()
logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err)
sess.notifyUpstreamDisconnect(err)
if notify {
sess.notifyUpstreamDisconnect(err)
}
if errClose := conn.Close(); errClose != nil {
log.Errorf("codex websockets executor: close websocket error: %v", errClose)
}
@@ -1984,6 +2056,9 @@ func (e *CodexAutoExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
return e.wsExec.Execute(ctx, auth, req, opts)
}
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return cliproxyexecutor.Response{}, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.httpExec.Execute(ctx, auth, req, opts)
}
@@ -1994,6 +2069,9 @@ func (e *CodexAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
return e.wsExec.ExecuteStream(ctx, auth, req, opts)
}
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.httpExec.ExecuteStream(ctx, auth, req, opts)
}
@@ -360,6 +360,174 @@ func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T)
}
}
func TestCodexWebsocketsExecuteStreamUpgradeRequiredReturnsWithoutLockingSession(t *testing.T) {
upgradeAttempts := make(chan struct{}, 2)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path)
w.WriteHeader(http.StatusInternalServerError)
return
}
upgradeAttempts <- struct{}{}
w.WriteHeader(http.StatusUpgradeRequired)
_, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`))
}))
defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
const executionSessionID = "ws-upgrade-required-session"
t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })
auth := &cliproxyauth.Auth{
ID: "codex-test",
Provider: "codex",
Attributes: map[string]string{
"api_key": "sk-test",
"base_url": server.URL,
},
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FromString("openai-response"),
ResponseFormat: sdktranslator.FromString("openai-response"),
Metadata: map[string]any{
cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,
},
}
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
execute := func(payload string) {
t.Helper()
done := make(chan error, 1)
go func() {
_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
Model: "gpt-5.4",
Payload: []byte(payload),
}, opts)
done <- errExecute
}()
select {
case errExecute := <-done:
if errExecute == nil {
t.Fatal("upgrade-required error = nil")
}
statusErr, ok := errExecute.(interface{ StatusCode() int })
if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {
t.Fatalf("upgrade-required error = %T %v, want status 426", errExecute, errExecute)
}
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for upgrade-required error; execution session may still be locked")
}
}
execute(`{"model":"gpt-5.4","generate":false,"input":[]}`)
execute(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`)
if got := len(upgradeAttempts); got != 2 {
t.Fatalf("websocket upgrade attempts = %d, want 2", got)
}
}
func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":{"message":"unauthorized"}}`))
}))
defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
const executionSessionID = "ws-handshake-error-session"
t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })
auth := &cliproxyauth.Auth{
ID: "codex-test",
Provider: "codex",
Attributes: map[string]string{
"api_key": "sk-test",
"base_url": server.URL,
},
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FromString("openai-response"),
Metadata: map[string]any{
cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,
},
}
for i := 0; i < 2; i++ {
done := make(chan error, 1)
go func() {
_, errExecute := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
Model: "gpt-5.4",
Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`),
}, opts)
done <- errExecute
}()
select {
case errExecute := <-done:
statusErr, ok := errExecute.(interface{ StatusCode() int })
if !ok || statusErr.StatusCode() != http.StatusUnauthorized {
t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute)
}
case <-time.After(5 * time.Second):
t.Fatalf("attempt %d timed out; execution session remained locked", i+1)
}
}
}
func TestExistingWebsocketSessionConnRequiresMatchingHealthyConnection(t *testing.T) {
conn := &websocket.Conn{}
sess := &codexWebsocketSession{
conn: conn,
authID: "auth-a",
wsURL: "ws://example.test/responses",
}
sess.resetUpstreamDisconnectError(conn)
if got := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != conn {
t.Fatal("matching healthy websocket session was not reusable")
}
if got := existingWebsocketSessionConn(sess, "auth-b", "ws://example.test/responses"); got != nil {
t.Fatal("websocket session matched a different auth")
}
if got := existingWebsocketSessionConn(sess, "auth-a", "ws://other.test/responses"); got != nil {
t.Fatal("websocket session matched a different URL")
}
sess.setUpstreamDisconnectError(conn, errors.New("upstream disconnected"))
if got := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != nil {
t.Fatal("disconnected websocket session remained reusable")
}
}
func TestCodexAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) {
exec := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
auth := &cliproxyauth.Auth{
ID: "codex-http-only",
Provider: "codex",
Attributes: map[string]string{
"api_key": "sk-test",
},
}
ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(
cliproxyexecutor.WithDownstreamWebsocket(context.Background()),
)
_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
Model: "gpt-5.4",
Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`),
}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")})
if errExecute == nil {
t.Fatal("ExecuteStream() error = nil, want replay-required error")
}
statusErr, ok := errExecute.(interface{ StatusCode() int })
if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {
t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute)
}
if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" {
t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got)
}
requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError)
if !ok || !requestScoped.IsRequestScoped() {
t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute)
}
}
func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) {
upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
capturedPayload := make(chan []byte, 1)
@@ -398,6 +398,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
}
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
if xaiInputHasItemType(req.Payload, "compaction_trigger") {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper)
}
@@ -463,7 +466,20 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody)
conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
var conn *websocket.Conn
var respHS *http.Response
var errDial error
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
conn = existingWebsocketSessionConn(sess, authID, wsURL)
if conn == nil {
if sess != nil {
sess.reqMu.Unlock()
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
} else {
conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
}
var upstreamHeaders http.Header
if respHS != nil {
upstreamHeaders = respHS.Header.Clone()
@@ -501,6 +517,15 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
errSend = mapXAIWebsocketWriteError(sess, conn, errSend)
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
if sess != nil {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
sess.clearActive(conn, readCh)
sess.reqMu.Unlock()
if !shouldRetryXAIWebsocketSend(errSend) {
return nil, errSend
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
if !shouldRetryXAIWebsocketSend(errSend) {
sess.clearActive(conn, readCh)
@@ -1431,6 +1456,9 @@ func (e *XAIAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.
if cliproxyexecutor.DownstreamWebsocket(ctx) && xaiWebsocketsEnabled(auth) {
return e.wsExec.ExecuteStream(ctx, auth, req, opts)
}
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.httpExec.ExecuteStream(ctx, auth, req, opts)
}
@@ -34,6 +34,50 @@ func TestXAIWebsocketsEnabledForConfigAPIKey(t *testing.T) {
}
}
func TestXAIAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) {
exec := NewXAIAutoExecutor(&config.Config{})
auth := &cliproxyauth.Auth{
ID: "xai-http-only",
Provider: "xai",
Attributes: map[string]string{
"api_key": "xai-key",
},
}
ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(
cliproxyexecutor.WithDownstreamWebsocket(context.Background()),
)
_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
Model: "grok-4",
Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`),
}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")})
if errExecute == nil {
t.Fatal("ExecuteStream() error = nil, want replay-required error")
}
statusErr, ok := errExecute.(interface{ StatusCode() int })
if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {
t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute)
}
if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" {
t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got)
}
requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError)
if !ok || !requestScoped.IsRequestScoped() {
t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute)
}
}
func TestXAIWebsocketsRequiredUpstreamRejectsCompactionHTTPFallback(t *testing.T) {
exec := NewXAIWebsocketsExecutor(&config.Config{})
ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background())
_, errExecute := exec.ExecuteStream(ctx, &cliproxyauth.Auth{}, cliproxyexecutor.Request{
Model: "grok-4",
Payload: []byte(`{"model":"grok-4","input":[{"type":"compaction_trigger"}]}`),
}, cliproxyexecutor.Options{})
if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) {
t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute)
}
}
func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) {
networkWriteErr := errors.New("write: broken pipe")
tests := []struct {
+25 -1
View File
@@ -64,6 +64,7 @@ const (
type pinnedAuthContextKey struct{}
type selectedAuthCallbackContextKey struct{}
type preparedModelRouteContextKey struct{}
type executionSessionContextKey struct{}
type disallowFreeAuthContextKey struct{}
@@ -142,6 +143,26 @@ func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) cont
return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback)
}
// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution.
// The boolean reports whether the route overrides normal model-to-provider resolution.
func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) {
if ctx == nil {
ctx = context.Background()
}
decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{})
ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision)
hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" || strings.TrimSpace(decision.Provider) != ""
return ctx, hasOverride
}
func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) {
if ctx == nil {
return modelRouteDecision{}, false
}
decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision)
return decision, ok
}
// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID.
func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context {
sessionID = strings.TrimSpace(sessionID)
@@ -1139,7 +1160,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl
func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) {
originalRequestedModel := modelName
routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)
routeDecision, preparedRoute := preparedModelRouteFromContext(ctx)
if !preparedRoute {
routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)
}
responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)
if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {
errChan := make(chan *interfaces.ErrorMessage, 1)
@@ -426,6 +426,38 @@ func TestHandlerModelRouterRoutesStreamBeforeRequestDetails(t *testing.T) {
}
}
func TestPrepareStreamModelRouteReusesDecisionDuringExecution(t *testing.T) {
const model = "prepared-router-model"
const targetPluginID = "prepared-stream-plugin"
routeCalls := 0
host := &handlerDirectExecutorRouteHost{}
host.hasRouters = true
host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
routeCalls++
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
}
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
handler.SetModelRouterHost(host)
body := []byte(`{"model":"prepared-router-model","stream":true}`)
ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai", model, body)
if !routedToPlugin {
t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route")
}
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, body, "")
for range dataChan {
}
if errMsg := <-errChan; errMsg != nil {
t.Fatalf("ExecuteStreamWithAuthManager() error = %+v", errMsg)
}
if routeCalls != 1 {
t.Fatalf("model router calls = %d, want 1", routeCalls)
}
if host.lastPluginID != targetPluginID {
t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
}
}
func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {
model := "model-execution-router-skip-model"
requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))
@@ -33,15 +33,19 @@ import (
)
const (
wsRequestTypeCreate = "response.create"
wsRequestTypeAppend = "response.append"
wsEventTypeError = "error"
wsEventTypeCompleted = "response.completed"
wsEventTypeDone = "response.done"
wsDoneMarker = "[DONE]"
wsTurnStateHeader = "x-codex-turn-state"
wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"
wsCloseReasonMaxBytes = 123
wsRequestTypeCreate = "response.create"
wsRequestTypeAppend = "response.append"
wsEventTypeError = "error"
wsEventTypeCompleted = "response.completed"
wsEventTypeDone = "response.done"
wsDoneMarker = "[DONE]"
wsTurnStateHeader = "x-codex-turn-state"
wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"
wsCloseReasonMaxBytes = 123
wsHTTPReplayRequiredCloseReason = "upstream requires HTTP replay"
responsesWebsocketUpstreamModeUnknown = ""
responsesWebsocketUpstreamModeWS = "websocket"
responsesWebsocketUpstreamModeHTTP = "http"
codexLocalCompactionSummaryPrefix = "Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work. Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:"
)
@@ -74,6 +78,14 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) {
return false, nil
}
errText := err.Error()
if cliproxyexecutor.IsUpstreamWebsocketReplayRequired(err) {
return true, websocket.FormatCloseMessage(
websocket.CloseServiceRestart,
truncateWebsocketCloseReason(wsHTTPReplayRequiredCloseReason, wsCloseReasonMaxBytes),
)
}
code := 0
reason := ""
var closeErr *websocket.CloseError
@@ -85,7 +97,6 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) {
StatusCode() int
}
var statusErr statusCoder
errText := err.Error()
if !errors.As(err, &statusErr) || statusErr.StatusCode() != http.StatusRequestEntityTooLarge ||
gjson.Get(errText, "error.code").String() != "message_too_big" {
return false, nil
@@ -419,20 +430,35 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
// Preserve independent upstream auth affinity when a downstream session switches providers.
pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState)
passthroughModelName := ""
upstreamMode := responsesWebsocketUpstreamModeUnknown
upstreamWebsocketAuthID := ""
sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) {
if h == nil || h.AuthManager == nil {
return nil, false, false
}
// Prefer the current manager view so hot-reloaded transport eligibility is
// observed even when the execution session still holds an older auth snapshot.
if auth, ok := h.AuthManager.GetByID(authID); ok {
return auth, false, true
}
if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {
return auth, true, true
}
auth, ok := h.AuthManager.GetByID(authID)
return auth, false, ok
return nil, false, false
}
sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {
auth, _, ok := sessionAuthByIDWithSource(authID)
return auth, ok
}
upstreamModeForAuth := func(auth *coreauth.Auth) string {
if auth != nil && websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) {
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
if provider == "codex" || provider == "xai" {
return responsesWebsocketUpstreamModeWS
}
}
return responsesWebsocketUpstreamModeHTTP
}
rememberPinnedAuth := func(authID string, modelName string) {
authID = strings.TrimSpace(authID)
auth, ok := sessionAuthByID(authID)
@@ -454,7 +480,6 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
}
pinnedAuthID = ""
}
forceTranscriptReplayNextRequest := false
for {
msgType, payload, errReadMessage := conn.ReadMessage()
@@ -488,6 +513,13 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
if requestModelName == "" {
requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
}
executionParent := context.WithValue(c.Request.Context(), "gin", c)
executionParent, routeOverridesModelResolution := h.PrepareStreamModelRoute(
executionParent,
h.HandlerType(),
requestModelName,
payload,
)
if pinnedAuthID != "" {
pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID)
providerKey := ""
@@ -514,15 +546,40 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
}
}
useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName)
if pinnedAuthID != "" {
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && responsesWebsocketAuthSupportsIncrementalInput(pinnedAuth) {
provider := strings.ToLower(strings.TrimSpace(pinnedAuth.Provider))
useUpstreamWebsocketPassthrough = provider == "codex" || provider == "xai"
}
}
nativeWebsocketPassthrough := !routeOverridesModelResolution && responsesWebsocketNativePassthroughAllowed(
upstreamMode,
useUpstreamWebsocketPassthrough,
pinnedAuthID,
upstreamWebsocketAuthID,
)
requestRequiresCurrentUpstreamWebsocket := responsesWebsocketRequestRequiresCurrentUpstream(payload)
if upstreamMode == responsesWebsocketUpstreamModeWS && !nativeWebsocketPassthrough {
if requestRequiresCurrentUpstreamWebsocket {
replayErr := responsesWebsocketHTTPReplayRequiredError()
wsTerminateErr = replayErr
matched, errClose := writer.closeForUpstreamError(replayErr)
if !matched {
_ = conn.Close()
} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
log.Debugf("responses websocket: replay close failed id=%s error=%v", passthroughSessionID, errClose)
}
return
}
// A full response.create is already a self-contained reset and can safely
// establish a new upstream transport without another replay.
}
if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough {
passthroughModelName = ""
}
allowIncrementalInputWithPreviousResponseID := false
allowCompactionReplayBypass := false
if !useUpstreamWebsocketPassthrough {
// Downstream websocket with CPA-mediated upstream (HTTP/SSE) always uses merged
// transcript replay. Incremental previous_response_id is reserved for end-to-end
// upstream websocket passthrough only.
if !nativeWebsocketPassthrough {
if pinnedAuthID != "" {
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
@@ -535,8 +592,10 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
var requestJSON []byte
var updatedLastRequest []byte
var errMsg *interfaces.ErrorMessage
if useUpstreamWebsocketPassthrough {
if nativeWebsocketPassthrough {
requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName)
} else if len(lastRequest) == 0 && strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" {
errMsg = responsesWebsocketPreviousResponseNotFoundError()
} else {
requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState(
payload,
@@ -544,7 +603,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
lastResponseOutput,
lastResponseID,
lastResponsePendingToolCallIDs,
allowIncrementalInputWithPreviousResponseID,
false,
allowCompactionReplayBypass,
)
}
@@ -570,7 +629,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
}
continue
}
if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, allowIncrementalInputWithPreviousResponseID) {
if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) {
if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {
requestJSON = updated
}
@@ -592,80 +651,112 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
previousLastResponseOutput := bytes.Clone(lastResponseOutput)
previousLastResponseID := lastResponseID
previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...)
forcedTranscriptReplay := forceTranscriptReplayNextRequest
if useUpstreamWebsocketPassthrough {
if nativeWebsocketPassthrough {
if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {
passthroughModelName = modelName
}
if forcedTranscriptReplay {
forceTranscriptReplayNextRequest = false
}
} else {
requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON)
requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)
updatedLastRequest = bytes.Clone(requestJSON)
lastRequest = updatedLastRequest
if forcedTranscriptReplay {
forceTranscriptReplayNextRequest = false
}
}
modelName := gjson.GetBytes(requestJSON, "model").String()
lastAttemptedAuthID := pinnedAuthID
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
attemptedUpstreamMode := responsesWebsocketUpstreamModeUnknown
selectedAuthObserved := false
pinnedAuthAttempted := false
cliCtx, cliCancel := h.GetContextWithCancel(h, c, executionParent)
cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx)
if nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket {
cliCtx = cliproxyexecutor.WithRequiredUpstreamWebsocket(cliCtx)
}
cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID)
if pinnedAuthID != "" {
cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
authID = strings.TrimSpace(authID)
if authID == "" || h == nil || h.AuthManager == nil {
return
}
lastAttemptedAuthID = authID
selectedAuthObserved = true
pinnedAuthAttempted = pinnedAuthAttempted || (pinnedAuthID != "" && authID == pinnedAuthID)
selectedAuth, ok := sessionAuthByID(authID)
if !ok || selectedAuth == nil {
return
}
attemptedUpstreamMode = upstreamModeForAuth(selectedAuth)
})
if pinnedAuthID != "" && !routeOverridesModelResolution {
cliCtx = handlers.WithPinnedAuthID(cliCtx, pinnedAuthID)
} else {
cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
authID = strings.TrimSpace(authID)
if authID == "" || h == nil || h.AuthManager == nil {
return
}
lastAttemptedAuthID = authID
selectedAuth, ok := sessionAuthByID(authID)
if !ok || selectedAuth == nil {
return
}
if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
rememberPinnedAuth(authID, modelName)
}
})
}
dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "")
if !selectedAuthObserved {
// Plugin/alternate routes bypass auth selection. Keep canonical HTTP-mode
// state instead of inheriting the previous pinned websocket mode.
attemptedUpstreamMode = responsesWebsocketUpstreamModeHTTP
}
// A connection-scoped continuation cannot rotate credentials in place. Suppress
// credential errors and make the client replay the full turn on a new socket.
replayPinnedAuthFailure := func(errMsg *interfaces.ErrorMessage) bool {
return nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket && pinnedAuthAttempted &&
shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg)
}
completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID)
completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(
c,
writer,
cliCancel,
dataChan,
errChan,
wsTimelineLog,
passthroughSessionID,
responsesWebsocketForwardOptions{
suppressError: replayPinnedAuthFailure,
},
)
if errForward != nil {
wsTerminateErr = errForward
log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
if !errors.Is(errForward, websocket.ErrCloseSent) {
log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
}
return
}
if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" {
if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil {
if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
rememberPinnedAuth(lastAttemptedAuthID, modelName)
} else if pinnedAuthID != "" {
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) {
rememberPinnedAuth(lastAttemptedAuthID, modelName)
}
}
if forwardErrMsg != nil {
lastRequest = previousLastRequest
lastResponseOutput = previousLastResponseOutput
lastResponseID = previousLastResponseID
lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
forgetPinnedAuth()
}
}
if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
forgetPinnedAuth()
forceTranscriptReplayNextRequest = true
if useUpstreamWebsocketPassthrough {
passthroughModelName = ""
} else {
lastRequest = previousLastRequest
lastResponseOutput = previousLastResponseOutput
lastResponseID = previousLastResponseID
lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
if replayPinnedAuthFailure(forwardErrMsg) {
replayErr := responsesWebsocketHTTPReplayRequiredError()
wsTerminateErr = replayErr
matched, errClose := writer.closeForUpstreamError(replayErr)
if !matched {
_ = conn.Close()
} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
log.Debugf("responses websocket: credential replay close failed id=%s error=%v", passthroughSessionID, errClose)
}
return
}
continue
}
if !useUpstreamWebsocketPassthrough {
upstreamMode = attemptedUpstreamMode
if upstreamMode == responsesWebsocketUpstreamModeWS {
upstreamWebsocketAuthID = lastAttemptedAuthID
if lastAttemptedAuthID != "" {
rememberPinnedAuth(lastAttemptedAuthID, modelName)
}
passthroughModelName = modelName
lastRequest = nil
lastResponseOutput = []byte("[]")
lastResponseID = ""
lastResponsePendingToolCallIDs = nil
} else {
upstreamWebsocketAuthID = ""
lastResponseOutput = completedOutput
lastResponseID = strings.TrimSpace(completedResponseID)
lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...)
@@ -673,6 +764,20 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
}
}
func responsesWebsocketHTTPReplayRequiredError() error {
return cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
func responsesWebsocketRequestRequiresCurrentUpstream(payload []byte) bool {
return strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" ||
strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == wsRequestTypeAppend
}
func responsesWebsocketNativePassthroughAllowed(upstreamMode string, useUpstreamWebsocket bool, pinnedAuthID string, upstreamAuthID string) bool {
return upstreamMode == responsesWebsocketUpstreamModeWS && useUpstreamWebsocket &&
strings.TrimSpace(pinnedAuthID) != "" && strings.TrimSpace(pinnedAuthID) == strings.TrimSpace(upstreamAuthID)
}
func websocketClientAddress(c *gin.Context) string {
if c == nil || c.Request == nil {
return ""
@@ -694,6 +799,15 @@ func websocketUpgradeHeaders(req *http.Request) http.Header {
return headers
}
func responsesWebsocketPreviousResponseNotFoundError() *interfaces.ErrorMessage {
return &interfaces.ErrorMessage{
StatusCode: http.StatusConflict,
Error: errors.New(
`{"error":{"message":"Previous response is not available on this websocket; resend the full conversation input without previous_response_id","type":"invalid_request_error","code":"previous_response_not_found","param":"previous_response_id"}}`,
),
}
}
func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)
}
@@ -1570,6 +1684,10 @@ func normalizeJSONArrayRaw(raw []byte) string {
return "[]"
}
type responsesWebsocketForwardOptions struct {
suppressError func(*interfaces.ErrorMessage) bool
}
func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
c *gin.Context,
writer *responsesWebsocketWriter,
@@ -1578,7 +1696,12 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
errs <-chan *interfaces.ErrorMessage,
wsTimelineLog websocketTimelineAppender,
sessionID string,
options ...responsesWebsocketForwardOptions,
) ([]byte, string, []string, *interfaces.ErrorMessage, error) {
var opts responsesWebsocketForwardOptions
if len(options) > 0 {
opts = options[0]
}
completed := false
completedOutput := []byte("[]")
completedResponseID := ""
@@ -1602,6 +1725,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
}
if errMsg != nil {
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
if opts.suppressError != nil && opts.suppressError(errMsg) {
cancel(errMsg.Error)
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil
}
markAPIResponseTimestamp(c)
if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched {
cancel(errMsg.Error)
@@ -1684,6 +1811,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
if h != nil {
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg)
}
if opts.suppressError != nil && opts.suppressError(payloadErrMsg) {
cancel(payloadErrMsg.Error)
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil
}
} else if isResponsesWebsocketCompletionEvent(eventType) {
completed = true
completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback)
@@ -1716,9 +1847,9 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
}
}
func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {
func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int {
if errMsg == nil {
return false
return 0
}
status := errMsg.StatusCode
if status <= 0 && errMsg.Error != nil {
@@ -1726,7 +1857,23 @@ func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage)
status = se.StatusCode()
}
}
switch status {
return status
}
func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool {
switch responsesWebsocketErrorStatus(errMsg) {
case http.StatusUnauthorized, http.StatusTooManyRequests:
return true
default:
return false
}
}
func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {
if errMsg == nil {
return false
}
switch responsesWebsocketErrorStatus(errMsg) {
case http.StatusUnauthorized,
http.StatusPaymentRequired,
http.StatusForbidden,
File diff suppressed because it is too large Load Diff
+19
View File
@@ -3,6 +3,7 @@ package executor
import "context"
type downstreamWebsocketContextKey struct{}
type requireUpstreamWebsocketContextKey struct{}
// WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection.
func WithDownstreamWebsocket(ctx context.Context) context.Context {
@@ -21,3 +22,21 @@ func DownstreamWebsocket(ctx context.Context) bool {
enabled, ok := raw.(bool)
return ok && enabled
}
// WithRequiredUpstreamWebsocket marks a request whose incremental context is valid only on the current upstream websocket.
func WithRequiredUpstreamWebsocket(ctx context.Context) context.Context {
if ctx == nil {
ctx = context.Background()
}
return context.WithValue(ctx, requireUpstreamWebsocketContextKey{}, true)
}
// RequiredUpstreamWebsocket reports whether falling back to an HTTP upstream would lose request context.
func RequiredUpstreamWebsocket(ctx context.Context) bool {
if ctx == nil {
return false
}
raw := ctx.Value(requireUpstreamWebsocketContextKey{})
enabled, ok := raw.(bool)
return ok && enabled
}
+29
View File
@@ -0,0 +1,29 @@
package executor
import (
"errors"
"net/http"
)
// UpstreamWebsocketReplayRequiredError indicates that an incremental request
// cannot safely continue because its upstream websocket is no longer reusable.
type UpstreamWebsocketReplayRequiredError struct{}
func (*UpstreamWebsocketReplayRequiredError) Error() string {
return `{"error":{"message":"upstream transport requires full HTTP replay","type":"server_error","code":"upstream_http_replay_required","status":426}}`
}
func (*UpstreamWebsocketReplayRequiredError) StatusCode() int { return http.StatusUpgradeRequired }
func (*UpstreamWebsocketReplayRequiredError) IsRequestScoped() bool { return true }
// NewUpstreamWebsocketReplayRequiredError creates a request-scoped replay signal.
func NewUpstreamWebsocketReplayRequiredError() error {
return &UpstreamWebsocketReplayRequiredError{}
}
// IsUpstreamWebsocketReplayRequired reports whether err is the internal replay signal.
func IsUpstreamWebsocketReplayRequired(err error) bool {
var replayErr *UpstreamWebsocketReplayRequiredError
return errors.As(err, &replayErr)
}
+25
View File
@@ -0,0 +1,25 @@
package executor
import (
"fmt"
"net/http"
"testing"
)
func TestUpstreamWebsocketReplayRequiredError(t *testing.T) {
err := NewUpstreamWebsocketReplayRequiredError()
if !IsUpstreamWebsocketReplayRequired(err) {
t.Fatal("replay error was not recognized")
}
if !IsUpstreamWebsocketReplayRequired(fmt.Errorf("wrapped: %w", err)) {
t.Fatal("wrapped replay error was not recognized")
}
statusErr, ok := err.(interface{ StatusCode() int })
if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {
t.Fatalf("replay error = %T %v, want status 426", err, err)
}
requestErr, ok := err.(RequestScopedError)
if !ok || !requestErr.IsRequestScoped() {
t.Fatalf("replay error = %T, want request scoped", err)
}
}