mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 01:16:47 +08:00
fix(responses): preserve context across websocket transport changes
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user