mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 01:16:47 +08:00
Merge pull request #4522 from sususu98/fix/responses-ws-continuity
This commit is contained in:
Luis Pater
@@ -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")
|
||||
|
||||
@@ -352,6 +352,23 @@ 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, *websocketConnectionCloser) {
|
||||
if sess == nil {
|
||||
return nil, nil
|
||||
}
|
||||
sess.connMu.Lock()
|
||||
conn := sess.conn
|
||||
closer := sess.connCloser
|
||||
matches := conn != nil && closer != 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, nil
|
||||
}
|
||||
return conn, closer
|
||||
}
|
||||
|
||||
func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) {
|
||||
if sess == nil {
|
||||
return nil, nil, "", "", nil
|
||||
@@ -532,14 +549,25 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
|
||||
}
|
||||
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
|
||||
|
||||
conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
||||
var conn *websocket.Conn
|
||||
var closer *websocketConnectionCloser
|
||||
var respHS *http.Response
|
||||
var errDial error
|
||||
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
||||
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
|
||||
if conn == nil {
|
||||
return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
||||
}
|
||||
} else {
|
||||
conn, closer, 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 {
|
||||
if opts.ExecutionLifecycle != nil {
|
||||
if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) {
|
||||
return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
|
||||
}
|
||||
return e.CodexExecutor.Execute(ctx, auth, req, opts)
|
||||
@@ -582,6 +610,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)
|
||||
@@ -810,7 +846,21 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
|
||||
}
|
||||
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
|
||||
|
||||
conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
||||
var conn *websocket.Conn
|
||||
var closer *websocketConnectionCloser
|
||||
var respHS *http.Response
|
||||
var errDial error
|
||||
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
||||
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
|
||||
if conn == nil {
|
||||
if sess != nil {
|
||||
sess.reqMu.Unlock()
|
||||
}
|
||||
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
||||
}
|
||||
} else {
|
||||
conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
||||
}
|
||||
var upstreamHeaders http.Header
|
||||
if respHS != nil {
|
||||
upstreamHeaders = respHS.Header.Clone()
|
||||
@@ -824,7 +874,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
|
||||
if sess != nil {
|
||||
sess.reqMu.Unlock()
|
||||
}
|
||||
if opts.ExecutionLifecycle != nil {
|
||||
if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) {
|
||||
return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
|
||||
}
|
||||
return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)
|
||||
@@ -864,6 +914,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)
|
||||
@@ -1964,6 +2023,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
|
||||
}
|
||||
@@ -1989,7 +2056,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 closer != nil {
|
||||
if errClose := closer.Close(); errClose != nil {
|
||||
log.Errorf("codex websockets executor: close websocket error: %v", errClose)
|
||||
@@ -2206,6 +2275,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)
|
||||
}
|
||||
|
||||
@@ -2216,6 +2288,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)
|
||||
}
|
||||
|
||||
|
||||
@@ -361,6 +361,176 @@ 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{}
|
||||
closer := newWebsocketConnectionCloser(conn)
|
||||
sess := &codexWebsocketSession{
|
||||
conn: conn,
|
||||
connCloser: closer,
|
||||
authID: "auth-a",
|
||||
wsURL: "ws://example.test/responses",
|
||||
}
|
||||
sess.resetUpstreamDisconnectError(conn)
|
||||
if gotConn, gotCloser := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); gotConn != conn || gotCloser != closer {
|
||||
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)
|
||||
|
||||
@@ -48,18 +48,21 @@ type xaiWebsocketIDStateStore struct {
|
||||
}
|
||||
|
||||
type xaiWebsocketIDState struct {
|
||||
mu sync.Mutex
|
||||
downstreamToUpstream map[string]string
|
||||
sequence int
|
||||
transcriptInput []json.RawMessage
|
||||
requestMu sync.Mutex
|
||||
mu sync.Mutex
|
||||
downstreamToUpstream map[string]string
|
||||
sequence int
|
||||
transcriptInput []json.RawMessage
|
||||
replayCompactedTranscriptOnReset bool
|
||||
}
|
||||
|
||||
type xaiWebsocketRequestIDMapper struct {
|
||||
state *xaiWebsocketIDState
|
||||
downstreamPreviousID string
|
||||
upstreamPreviousID string
|
||||
upstreamResponseID string
|
||||
downstreamResponseID string
|
||||
state *xaiWebsocketIDState
|
||||
downstreamPreviousID string
|
||||
upstreamPreviousID string
|
||||
upstreamResponseID string
|
||||
downstreamResponseID string
|
||||
replayedCompactedTranscript bool
|
||||
}
|
||||
|
||||
func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor {
|
||||
@@ -177,20 +180,21 @@ func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte {
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) {
|
||||
func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) {
|
||||
if s == nil || len(requestPayload) == 0 || len(completedPayload) == 0 {
|
||||
return
|
||||
}
|
||||
inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input"))
|
||||
outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output"))
|
||||
if len(inputItems) == 0 && len(outputItems) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if strings.TrimSpace(gjson.GetBytes(requestPayload, "previous_response_id").String()) == "" {
|
||||
if reset {
|
||||
s.transcriptInput = nil
|
||||
s.replayCompactedTranscriptOnReset = false
|
||||
}
|
||||
if len(inputItems) == 0 && len(outputItems) == 0 {
|
||||
return
|
||||
}
|
||||
s.transcriptInput = append(s.transcriptInput, inputItems...)
|
||||
s.transcriptInput = append(s.transcriptInput, outputItems...)
|
||||
@@ -210,9 +214,34 @@ func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) {
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.transcriptInput = next
|
||||
s.replayCompactedTranscriptOnReset = len(next) > 0
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) {
|
||||
if s == nil || len(payload) == 0 {
|
||||
return payload, false
|
||||
}
|
||||
s.mu.Lock()
|
||||
if !s.replayCompactedTranscriptOnReset || len(s.transcriptInput) == 0 {
|
||||
s.mu.Unlock()
|
||||
return payload, false
|
||||
}
|
||||
prefix := make([]json.RawMessage, 0, len(s.transcriptInput))
|
||||
for _, item := range s.transcriptInput {
|
||||
prefix = append(prefix, bytes.Clone(item))
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
current := xaiJSONRawMessages(gjson.GetBytes(payload, "input"))
|
||||
merged := append(prefix, current...)
|
||||
out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged))
|
||||
if errSet != nil {
|
||||
return payload, false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
func xaiJSONRawMessages(result gjson.Result) []json.RawMessage {
|
||||
if !result.Exists() || !result.IsArray() {
|
||||
return nil
|
||||
@@ -243,7 +272,16 @@ func xaiMarshalRawMessages(items []json.RawMessage) []byte {
|
||||
}
|
||||
|
||||
func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte {
|
||||
if m == nil || len(payload) == 0 || m.downstreamPreviousID == m.upstreamPreviousID {
|
||||
if m == nil || len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
if m.downstreamPreviousID == m.upstreamPreviousID {
|
||||
requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
|
||||
if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil {
|
||||
out, replayed := m.state.prependCompactedTranscriptOnReset(payload)
|
||||
m.replayedCompactedTranscript = replayed
|
||||
return out
|
||||
}
|
||||
return payload
|
||||
}
|
||||
if m.upstreamPreviousID == "" {
|
||||
@@ -251,6 +289,7 @@ func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []b
|
||||
if errDelete == nil {
|
||||
if m.downstreamPreviousID != "" && m.state != nil {
|
||||
out = m.state.prependTranscriptInput(out)
|
||||
m.replayedCompactedTranscript = true
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -396,8 +435,30 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
if stateSessionID == "" {
|
||||
stateSessionID = executionSessionID
|
||||
}
|
||||
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
|
||||
state := getXAIWebsocketIDState(e.idStore, stateSessionID)
|
||||
stateRequestLocked := false
|
||||
stateRequestLockTransferred := false
|
||||
if executionSessionID == "" && state != nil {
|
||||
state.requestMu.Lock()
|
||||
stateRequestLocked = true
|
||||
}
|
||||
defer func() {
|
||||
if stateRequestLocked && !stateRequestLockTransferred {
|
||||
state.requestMu.Unlock()
|
||||
}
|
||||
}()
|
||||
if xaiInputHasItemType(req.Payload, "compaction_trigger") {
|
||||
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
||||
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
||||
}
|
||||
if executionSessionID != "" {
|
||||
sess := e.getOrCreateSession(executionSessionID)
|
||||
if sess != nil {
|
||||
sess.reqMu.Lock()
|
||||
defer sess.reqMu.Unlock()
|
||||
}
|
||||
}
|
||||
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
|
||||
return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper)
|
||||
}
|
||||
|
||||
@@ -437,6 +498,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
sess.reqMu.Lock()
|
||||
}
|
||||
}
|
||||
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
|
||||
if idMapper != nil {
|
||||
if websocketSessionTargetChanged(sess, authID, wsURL) {
|
||||
idMapper.upstreamPreviousID = ""
|
||||
@@ -447,6 +509,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
|
||||
wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID)
|
||||
wsReqBody := buildXAIWebsocketRequestBody(prepared.body)
|
||||
requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String())
|
||||
transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" &&
|
||||
(requestType != "response.append" || (idMapper != nil && idMapper.replayedCompactedTranscript))
|
||||
warmupRequest := xaiWebsocketGenerateFalse(wsReqBody)
|
||||
|
||||
wsReqLog := helps.UpstreamRequestLog{
|
||||
@@ -463,7 +528,21 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
|
||||
logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody)
|
||||
|
||||
conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
||||
var conn *websocket.Conn
|
||||
var closer *websocketConnectionCloser
|
||||
var respHS *http.Response
|
||||
var errDial error
|
||||
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
||||
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
|
||||
if conn == nil {
|
||||
if sess != nil {
|
||||
sess.reqMu.Unlock()
|
||||
}
|
||||
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
||||
}
|
||||
} else {
|
||||
conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
||||
}
|
||||
var upstreamHeaders http.Header
|
||||
if respHS != nil {
|
||||
upstreamHeaders = respHS.Header.Clone()
|
||||
@@ -508,6 +587,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)
|
||||
@@ -570,7 +658,13 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
}
|
||||
|
||||
out := make(chan cliproxyexecutor.StreamChunk)
|
||||
if stateRequestLocked {
|
||||
stateRequestLockTransferred = true
|
||||
}
|
||||
go func() {
|
||||
if stateRequestLocked {
|
||||
defer state.requestMu.Unlock()
|
||||
}
|
||||
terminateReason := "completed"
|
||||
var terminateErr error
|
||||
|
||||
@@ -677,6 +771,10 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
case "response.created":
|
||||
if warmupRequest {
|
||||
warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload)
|
||||
if idMapper != nil && idMapper.state != nil && !recordedTranscript {
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset)
|
||||
recordedTranscript = true
|
||||
}
|
||||
logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload)
|
||||
}
|
||||
case "response.output_item.done":
|
||||
@@ -690,7 +788,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
payload = xaiNormalizeReasoningSummaryData(payload)
|
||||
cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload)
|
||||
if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload)
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
|
||||
recordedTranscript = true
|
||||
}
|
||||
case "response.done":
|
||||
@@ -699,7 +797,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
reporter.Publish(ctx, detail)
|
||||
}
|
||||
if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload)
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
|
||||
recordedTranscript = true
|
||||
}
|
||||
}
|
||||
@@ -793,8 +891,11 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx
|
||||
return nil, err
|
||||
}
|
||||
|
||||
responseID := xaiCompactionResponseID(data)
|
||||
idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID))
|
||||
responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data)
|
||||
if errValidate != nil {
|
||||
return nil, errValidate
|
||||
}
|
||||
idMapper.state.replaceTranscriptWithItems(compactionItem)
|
||||
idMapper.state.mapDownstreamToUpstream(responseID, "")
|
||||
|
||||
headers = headers.Clone()
|
||||
@@ -812,6 +913,30 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx
|
||||
return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil
|
||||
}
|
||||
|
||||
func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) {
|
||||
if len(data) == 0 || !json.Valid(data) {
|
||||
return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"}
|
||||
}
|
||||
responseIDResult := gjson.GetBytes(data, "id")
|
||||
output := gjson.GetBytes(data, "output")
|
||||
if responseIDResult.Type != gjson.String || strings.TrimSpace(responseIDResult.String()) == "" || !output.Exists() || !output.IsArray() {
|
||||
return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
|
||||
}
|
||||
items := output.Array()
|
||||
if len(items) == 0 {
|
||||
return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
|
||||
}
|
||||
item := items[0]
|
||||
itemType := item.Get("type")
|
||||
encryptedContent := item.Get("encrypted_content")
|
||||
if item.Type != gjson.JSON || itemType.Type != gjson.String || strings.TrimSpace(itemType.String()) != "compaction" ||
|
||||
encryptedContent.Type != gjson.String || strings.TrimSpace(encryptedContent.String()) == "" {
|
||||
return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
|
||||
}
|
||||
normalizedResponseID := xaiCompactionResponseID(data)
|
||||
return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil
|
||||
}
|
||||
|
||||
func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) {
|
||||
if len(payload) == 0 {
|
||||
payload = []byte(`{}`)
|
||||
@@ -1497,6 +1622,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 {
|
||||
@@ -1246,6 +1290,127 @@ func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedC
|
||||
}
|
||||
}
|
||||
|
||||
func TestXAIWebsocketPostCompactionAppendWithoutPreviousReplaysCompactedTranscript(t *testing.T) {
|
||||
store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
|
||||
state := getXAIWebsocketIDState(store, "post-compaction-append-session")
|
||||
state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`))
|
||||
state.mapDownstreamToUpstream("resp-compact", "")
|
||||
|
||||
fullReset := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-full"}]}`)
|
||||
fullMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", fullReset)
|
||||
if full := fullMapper.upstreamRequestPayload(fullReset); len(gjson.GetBytes(full, "input").Array()) != 1 {
|
||||
t.Fatalf("self-contained response.create unexpectedly replayed compacted transcript: %s", full)
|
||||
}
|
||||
|
||||
payload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`)
|
||||
mapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", payload)
|
||||
got := mapper.upstreamRequestPayload(payload)
|
||||
input := gjson.GetBytes(got, "input").Array()
|
||||
if len(input) != 2 {
|
||||
t.Fatalf("post-compaction append input len = %d, want 2: %s", len(input), got)
|
||||
}
|
||||
if gotType := input[0].Get("type").String(); gotType != "compaction" {
|
||||
t.Fatalf("post-compaction append input[0].type = %q, want compaction: %s", gotType, got)
|
||||
}
|
||||
if gotID := input[1].Get("id").String(); gotID != "msg-2" {
|
||||
t.Fatalf("post-compaction append input[1].id = %q, want msg-2: %s", gotID, got)
|
||||
}
|
||||
|
||||
state.recordTranscriptTurn(got, []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2"}]}}`), true)
|
||||
nextPayload := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-3"}]}`)
|
||||
nextMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", nextPayload)
|
||||
next := nextMapper.upstreamRequestPayload(nextPayload)
|
||||
if nextInput := gjson.GetBytes(next, "input").Array(); len(nextInput) != 1 || nextInput[0].Get("id").String() != "msg-3" {
|
||||
t.Fatalf("compacted transcript replay was not cleared after success: %s", next)
|
||||
}
|
||||
}
|
||||
|
||||
func TestXAIWebsocketPostCompactionWarmupPreservesTranscriptForLaterCompaction(t *testing.T) {
|
||||
store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
|
||||
state := getXAIWebsocketIDState(store, "warmup-reset-session")
|
||||
state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`))
|
||||
|
||||
warmupPayload := []byte(`{"type":"response.append","model":"grok-4.3","generate":false,"input":[{"type":"message","id":"warmup-context"}]}`)
|
||||
warmupMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", warmupPayload)
|
||||
warmupUpstream := warmupMapper.upstreamRequestPayload(warmupPayload)
|
||||
if !warmupMapper.replayedCompactedTranscript {
|
||||
t.Fatal("post-compaction warmup did not mark full transcript replay")
|
||||
}
|
||||
state.recordTranscriptTurn(
|
||||
warmupUpstream,
|
||||
[]byte(`{"type":"response.completed","response":{"id":"resp-warmup","output":[]}}`),
|
||||
true,
|
||||
)
|
||||
|
||||
appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-after-warmup"}]}`)
|
||||
appendMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", appendPayload)
|
||||
appendUpstream := appendMapper.upstreamRequestPayload(appendPayload)
|
||||
input := gjson.GetBytes(appendUpstream, "input").Array()
|
||||
if len(input) != 1 || input[0].Get("id").String() != "msg-after-warmup" {
|
||||
t.Fatalf("warmup retained pending replay instead of native append: %s", appendUpstream)
|
||||
}
|
||||
state.recordTranscriptTurn(
|
||||
appendUpstream,
|
||||
[]byte(`{"type":"response.completed","response":{"id":"resp-after-warmup","output":[{"type":"message","id":"out-after-warmup"}]}}`),
|
||||
false,
|
||||
)
|
||||
|
||||
transcript := gjson.ParseBytes(state.snapshotTranscriptInput()).Array()
|
||||
wantTypes := []string{"compaction", "message", "message", "message"}
|
||||
if len(transcript) != len(wantTypes) {
|
||||
t.Fatalf("post-warmup transcript len = %d, want %d: %s", len(transcript), len(wantTypes), state.snapshotTranscriptInput())
|
||||
}
|
||||
for i, wantType := range wantTypes {
|
||||
if gotType := transcript[i].Get("type").String(); gotType != wantType {
|
||||
t.Fatalf("post-warmup transcript[%d].type = %q, want %q: %s", i, gotType, wantType, state.snapshotTranscriptInput())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestXAIWebsocketEmptyFullResetClearsPendingCompactionReplay(t *testing.T) {
|
||||
store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
|
||||
state := getXAIWebsocketIDState(store, "empty-reset-session")
|
||||
state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"stale-compact-state"}`))
|
||||
state.recordTranscriptTurn(
|
||||
[]byte(`{"type":"response.create","model":"grok-4.3","input":[]}`),
|
||||
[]byte(`{"type":"response.completed","response":{"id":"resp-empty","output":[]}}`),
|
||||
true,
|
||||
)
|
||||
|
||||
appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-new"}]}`)
|
||||
mapper := newXAIWebsocketRequestIDMapper(store, "empty-reset-session", appendPayload)
|
||||
got := mapper.upstreamRequestPayload(appendPayload)
|
||||
input := gjson.GetBytes(got, "input").Array()
|
||||
if len(input) != 1 || input[0].Get("id").String() != "msg-new" {
|
||||
t.Fatalf("empty full reset retained stale compaction replay: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateXAIWebsocketCompactionResponse(t *testing.T) {
|
||||
valid := []byte(`{"id":"resp_compact","output":[{"type":"compaction","encrypted_content":"opaque-state"}]}`)
|
||||
responseID, item, err := validateXAIWebsocketCompactionResponse(valid)
|
||||
if err != nil {
|
||||
t.Fatalf("valid compaction response error: %v", err)
|
||||
}
|
||||
if responseID != "resp_compact" || gjson.GetBytes(item, "encrypted_content").String() != "opaque-state" {
|
||||
t.Fatalf("validated compaction response = id:%q item:%s", responseID, item)
|
||||
}
|
||||
|
||||
for _, payload := range [][]byte{
|
||||
nil,
|
||||
[]byte(`{}`),
|
||||
[]byte(`{"id":"resp_empty","output":[]}`),
|
||||
[]byte(`{"id":123,"output":[{"type":"compaction","encrypted_content":"opaque"}]}`),
|
||||
[]byte(`{"id":"resp_object","output":{"0":{"type":"compaction","encrypted_content":"opaque"}}}`),
|
||||
[]byte(`{"id":"resp_numeric_state","output":[{"type":"compaction","encrypted_content":123}]}`),
|
||||
[]byte(`{"id":"resp_missing_state","output":[{"type":"compaction"}]}`),
|
||||
} {
|
||||
if _, _, errInvalid := validateXAIWebsocketCompactionResponse(payload); errInvalid == nil {
|
||||
t.Fatalf("invalid compaction response accepted: %s", payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) {
|
||||
body := []byte(`{"model":"grok-4.3","stream":true,"stream_options":{"include_usage":true},"background":true,"prompt_cache_key":"cache-1","previous_response_id":"resp-prev","instructions":"system prompt","input":[{"type":"message","role":"user","content":"hello"}]}`)
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -1144,7 +1165,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)
|
||||
|
||||
@@ -464,6 +464,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
|
||||
}
|
||||
@@ -588,84 +647,120 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
|
||||
toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey)
|
||||
previousLastRequest := bytes.Clone(lastRequest)
|
||||
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)
|
||||
toolCacheTurn.recordRequest(requestJSON)
|
||||
requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(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{
|
||||
toolCacheTurn: toolCacheTurn,
|
||||
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 {
|
||||
|
||||
toolCacheTurn.commit()
|
||||
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 +768,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 +803,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 +1688,11 @@ func normalizeJSONArrayRaw(raw []byte) string {
|
||||
return "[]"
|
||||
}
|
||||
|
||||
type responsesWebsocketForwardOptions struct {
|
||||
toolCacheTurn *responsesWebsocketToolCacheTurn
|
||||
suppressError func(*interfaces.ErrorMessage) bool
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
|
||||
c *gin.Context,
|
||||
writer *responsesWebsocketWriter,
|
||||
@@ -1578,7 +1701,13 @@ 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]
|
||||
}
|
||||
toolCacheTurn := opts.toolCacheTurn
|
||||
completed := false
|
||||
completedOutput := []byte("[]")
|
||||
completedResponseID := ""
|
||||
@@ -1602,6 +1731,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)
|
||||
@@ -1676,7 +1809,11 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
|
||||
if isResponsesWebsocketCompletionEvent(eventType) {
|
||||
payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback)
|
||||
}
|
||||
recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
|
||||
if toolCacheTurn != nil {
|
||||
toolCacheTurn.recordResponse(payloads[i])
|
||||
} else {
|
||||
recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
|
||||
}
|
||||
recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])
|
||||
var payloadErrMsg *interfaces.ErrorMessage
|
||||
if eventType == wsEventTypeError {
|
||||
@@ -1684,6 +1821,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 +1857,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 +1867,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,6 +19,7 @@ const (
|
||||
var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
|
||||
var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
|
||||
var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()
|
||||
var defaultWebsocketToolCacheTransactionMu sync.RWMutex
|
||||
|
||||
type websocketToolOutputCache struct {
|
||||
mu sync.Mutex
|
||||
@@ -33,6 +34,14 @@ type websocketToolOutputSession struct {
|
||||
order []string
|
||||
}
|
||||
|
||||
type responsesWebsocketToolCacheTurn struct {
|
||||
sessionKey string
|
||||
outputs map[string]json.RawMessage
|
||||
outputOrder []string
|
||||
calls map[string]json.RawMessage
|
||||
callOrder []string
|
||||
}
|
||||
|
||||
func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache {
|
||||
if ttl < 0 {
|
||||
ttl = websocketToolOutputCacheTTL
|
||||
@@ -196,6 +205,8 @@ func (c *websocketToolSessionRefCounter) release(sessionKey string) bool {
|
||||
}
|
||||
|
||||
func retainResponsesWebsocketToolCaches(sessionKey string) {
|
||||
defaultWebsocketToolCacheTransactionMu.Lock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.Unlock()
|
||||
if defaultWebsocketToolSessionRefs == nil {
|
||||
return
|
||||
}
|
||||
@@ -203,13 +214,14 @@ func retainResponsesWebsocketToolCaches(sessionKey string) {
|
||||
}
|
||||
|
||||
func releaseResponsesWebsocketToolCaches(sessionKey string) {
|
||||
defaultWebsocketToolCacheTransactionMu.Lock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.Unlock()
|
||||
if defaultWebsocketToolSessionRefs == nil {
|
||||
return
|
||||
}
|
||||
if !defaultWebsocketToolSessionRefs.release(sessionKey) {
|
||||
return
|
||||
}
|
||||
|
||||
if defaultWebsocketToolOutputCache != nil {
|
||||
defaultWebsocketToolOutputCache.deleteSession(sessionKey)
|
||||
}
|
||||
@@ -218,15 +230,114 @@ func releaseResponsesWebsocketToolCaches(sessionKey string) {
|
||||
}
|
||||
}
|
||||
|
||||
func newResponsesWebsocketToolCacheTurn(sessionKey string) *responsesWebsocketToolCacheTurn {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return nil
|
||||
}
|
||||
return &responsesWebsocketToolCacheTurn{
|
||||
sessionKey: sessionKey,
|
||||
outputs: make(map[string]json.RawMessage),
|
||||
calls: make(map[string]json.RawMessage),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordRequest(payload []byte) {
|
||||
if t == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
input := gjson.GetBytes(payload, "input")
|
||||
if !input.Exists() || !input.IsArray() {
|
||||
return
|
||||
}
|
||||
for _, item := range input.Array() {
|
||||
t.recordItem(item)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordResponse(payload []byte) {
|
||||
if t == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) {
|
||||
case "response.completed":
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if !output.Exists() || !output.IsArray() {
|
||||
return
|
||||
}
|
||||
for _, item := range output.Array() {
|
||||
if isCompleteResponsesWebsocketToolCall(item) {
|
||||
t.recordItem(item)
|
||||
}
|
||||
}
|
||||
case "response.output_item.added", "response.output_item.done":
|
||||
item := gjson.GetBytes(payload, "item")
|
||||
if isCompleteResponsesWebsocketToolCall(item) {
|
||||
t.recordItem(item)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordItem(item gjson.Result) {
|
||||
if t == nil || !item.Exists() {
|
||||
return
|
||||
}
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
if callID == "" || strings.TrimSpace(item.Raw) == "" {
|
||||
return
|
||||
}
|
||||
raw := append(json.RawMessage(nil), item.Raw...)
|
||||
switch {
|
||||
case isResponsesToolCallOutputType(item.Get("type").String()):
|
||||
if _, exists := t.outputs[callID]; !exists {
|
||||
t.outputOrder = append(t.outputOrder, callID)
|
||||
}
|
||||
t.outputs[callID] = raw
|
||||
case isResponsesToolCallType(item.Get("type").String()):
|
||||
if _, exists := t.calls[callID]; !exists {
|
||||
t.callOrder = append(t.callOrder, callID)
|
||||
}
|
||||
t.calls[callID] = raw
|
||||
}
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) commit() {
|
||||
if t == nil || t.sessionKey == "" {
|
||||
return
|
||||
}
|
||||
defaultWebsocketToolCacheTransactionMu.Lock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.Unlock()
|
||||
if defaultWebsocketToolOutputCache != nil {
|
||||
for _, callID := range t.outputOrder {
|
||||
defaultWebsocketToolOutputCache.record(t.sessionKey, callID, t.outputs[callID])
|
||||
}
|
||||
}
|
||||
if defaultWebsocketToolCallCache != nil {
|
||||
for _, callID := range t.callOrder {
|
||||
defaultWebsocketToolCallCache.record(t.sessionKey, callID, t.calls[callID])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte {
|
||||
return repairResponsesWebsocketToolCallsWithCaches(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload)
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCallsWithoutRecording(sessionKey string, payload []byte) []byte {
|
||||
defaultWebsocketToolCacheTransactionMu.RLock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.RUnlock()
|
||||
return repairResponsesWebsocketToolCallsWithCachesMode(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload, false)
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
|
||||
return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload)
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
|
||||
return repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache, sessionKey, payload, true)
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte, record bool) []byte {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || outputCache == nil || len(payload) == 0 {
|
||||
return payload
|
||||
@@ -238,7 +349,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock
|
||||
}
|
||||
|
||||
allowOrphanOutputs := strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != ""
|
||||
updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs)
|
||||
updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs, record)
|
||||
if errRepair != nil || updatedRaw == "" || updatedRaw == input.Raw {
|
||||
return payload
|
||||
}
|
||||
@@ -250,7 +361,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock
|
||||
return updated
|
||||
}
|
||||
|
||||
func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool) (string, error) {
|
||||
func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool, record bool) (string, error) {
|
||||
rawArray = strings.TrimSpace(rawArray)
|
||||
if rawArray == "" {
|
||||
return "[]", nil
|
||||
@@ -276,14 +387,16 @@ func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCa
|
||||
continue
|
||||
}
|
||||
outputPresent[callID] = struct{}{}
|
||||
outputCache.record(sessionKey, callID, item)
|
||||
if record {
|
||||
outputCache.record(sessionKey, callID, item)
|
||||
}
|
||||
case isResponsesToolCallType(itemType):
|
||||
callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String())
|
||||
if callID == "" {
|
||||
continue
|
||||
}
|
||||
callPresent[callID] = struct{}{}
|
||||
if callCache != nil {
|
||||
if record && callCache != nil {
|
||||
callCache.record(sessionKey, callID, item)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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