mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-09 01:46:48 +08:00
feat(websockets): add transcript state tracking and compaction trigger support
- Added methods for managing and tracking WebSocket transcript state, including recording, prepending, and replacing transcript inputs. - Implemented `executeCompactionTriggerFromWebsocketContext` to support compaction triggers using recorded transcript context. - Enhanced upstream-downstream ID mapping with additional utilities and state synchronization. - Expanded test coverage to validate transcript state management, compaction payload generation, and WebSocket response handling.
This commit is contained in:
Luis Pater
@@ -52,6 +52,7 @@ type xaiWebsocketIDState struct {
|
||||
mu sync.Mutex
|
||||
downstreamToUpstream map[string]string
|
||||
sequence int
|
||||
transcriptInput []json.RawMessage
|
||||
}
|
||||
|
||||
type xaiWebsocketRequestIDMapper struct {
|
||||
@@ -124,12 +125,124 @@ func (s *xaiWebsocketIDState) upstreamIDForDownstream(downstreamID string) strin
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if upstreamID := strings.TrimSpace(s.downstreamToUpstream[downstreamID]); upstreamID != "" {
|
||||
return upstreamID
|
||||
if upstreamID, ok := s.downstreamToUpstream[downstreamID]; ok {
|
||||
return strings.TrimSpace(upstreamID)
|
||||
}
|
||||
return downstreamID
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) mapDownstreamToUpstream(downstreamID string, upstreamID string) {
|
||||
downstreamID = strings.TrimSpace(downstreamID)
|
||||
if s == nil || downstreamID == "" {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.downstreamToUpstream == nil {
|
||||
s.downstreamToUpstream = make(map[string]string)
|
||||
}
|
||||
s.downstreamToUpstream[downstreamID] = strings.TrimSpace(upstreamID)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) snapshotTranscriptInput() []byte {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if len(s.transcriptInput) == 0 {
|
||||
return nil
|
||||
}
|
||||
return xaiMarshalRawMessages(s.transcriptInput)
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte {
|
||||
if s == nil || len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
s.mu.Lock()
|
||||
prefix := make([]json.RawMessage, 0, len(s.transcriptInput))
|
||||
for _, item := range s.transcriptInput {
|
||||
prefix = append(prefix, bytes.Clone(item))
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if len(prefix) == 0 {
|
||||
return payload
|
||||
}
|
||||
current := xaiJSONRawMessages(gjson.GetBytes(payload, "input"))
|
||||
merged := append(prefix, current...)
|
||||
out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged))
|
||||
if errSet != nil {
|
||||
return payload
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) {
|
||||
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()) == "" {
|
||||
s.transcriptInput = nil
|
||||
}
|
||||
s.transcriptInput = append(s.transcriptInput, inputItems...)
|
||||
s.transcriptInput = append(s.transcriptInput, outputItems...)
|
||||
}
|
||||
|
||||
func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
next := make([]json.RawMessage, 0, len(items))
|
||||
for _, item := range items {
|
||||
item = bytes.TrimSpace(item)
|
||||
if len(item) == 0 || !json.Valid(item) {
|
||||
continue
|
||||
}
|
||||
next = append(next, bytes.Clone(item))
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.transcriptInput = next
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func xaiJSONRawMessages(result gjson.Result) []json.RawMessage {
|
||||
if !result.Exists() || !result.IsArray() {
|
||||
return nil
|
||||
}
|
||||
items := result.Array()
|
||||
out := make([]json.RawMessage, 0, len(items))
|
||||
for _, item := range items {
|
||||
raw := bytes.TrimSpace([]byte(item.Raw))
|
||||
if len(raw) == 0 || !json.Valid(raw) {
|
||||
continue
|
||||
}
|
||||
out = append(out, bytes.Clone(raw))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func xaiMarshalRawMessages(items []json.RawMessage) []byte {
|
||||
var buf bytes.Buffer
|
||||
buf.WriteByte('[')
|
||||
for i, item := range items {
|
||||
if i > 0 {
|
||||
buf.WriteByte(',')
|
||||
}
|
||||
buf.Write(bytes.TrimSpace(item))
|
||||
}
|
||||
buf.WriteByte(']')
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte {
|
||||
if m == nil || len(payload) == 0 || m.downstreamPreviousID == m.upstreamPreviousID {
|
||||
return payload
|
||||
@@ -137,6 +250,9 @@ func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []b
|
||||
if m.upstreamPreviousID == "" {
|
||||
out, errDelete := sjson.DeleteBytes(payload, "previous_response_id")
|
||||
if errDelete == nil {
|
||||
if m.downstreamPreviousID != "" && m.state != nil {
|
||||
out = m.state.prependTranscriptInput(out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
return payload
|
||||
@@ -275,12 +391,16 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
if opts.Alt == "responses/compact" {
|
||||
return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"}
|
||||
}
|
||||
executionSessionID := executionSessionIDFromOptions(opts)
|
||||
stateSessionID := xaiExecutionSessionID(req, opts)
|
||||
if stateSessionID == "" {
|
||||
stateSessionID = executionSessionID
|
||||
}
|
||||
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
|
||||
if xaiInputHasItemType(req.Payload, "compaction_trigger") {
|
||||
return e.XAIExecutor.ExecuteStream(ctx, auth, req, opts)
|
||||
return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper)
|
||||
}
|
||||
|
||||
executionSessionID := executionSessionIDFromOptions(opts)
|
||||
idMapper := newXAIWebsocketRequestIDMapper(e.idStore, executionSessionID, req.Payload)
|
||||
token, baseURL := xaiCreds(auth)
|
||||
if baseURL == "" {
|
||||
baseURL = xaiauth.DefaultAPIBaseURL
|
||||
@@ -450,6 +570,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
var param any
|
||||
outputItemsByIndex := make(map[int64][]byte)
|
||||
var outputItemsFallback [][]byte
|
||||
recordedTranscript := false
|
||||
for {
|
||||
if ctx != nil && ctx.Err() != nil {
|
||||
terminateReason = "context_done"
|
||||
@@ -524,11 +645,19 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
reporter.Publish(ctx, detail)
|
||||
}
|
||||
payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback)
|
||||
if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload)
|
||||
recordedTranscript = true
|
||||
}
|
||||
case "response.done":
|
||||
logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload)
|
||||
if detail, ok := helps.ParseCodexUsage(payload); ok {
|
||||
reporter.Publish(ctx, detail)
|
||||
}
|
||||
if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
|
||||
idMapper.state.recordTranscriptTurn(wsReqBody, payload)
|
||||
recordedTranscript = true
|
||||
}
|
||||
}
|
||||
|
||||
if cliproxyexecutor.DownstreamWebsocket(ctx) {
|
||||
@@ -589,6 +718,72 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
|
||||
return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil
|
||||
}
|
||||
|
||||
func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, idMapper *xaiWebsocketRequestIDMapper) (*cliproxyexecutor.StreamResult, error) {
|
||||
if idMapper == nil || idMapper.state == nil {
|
||||
return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is unavailable"}
|
||||
}
|
||||
transcriptInput := idMapper.state.snapshotTranscriptInput()
|
||||
if len(transcriptInput) == 0 {
|
||||
return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is empty"}
|
||||
}
|
||||
authID := ""
|
||||
if auth != nil {
|
||||
authID = auth.ID
|
||||
}
|
||||
log.Infof(
|
||||
"xai websockets: compact fallback session=%s auth=%s input_items=%d",
|
||||
xaiExecutionSessionID(req, opts),
|
||||
strings.TrimSpace(authID),
|
||||
len(gjson.ParseBytes(transcriptInput).Array()),
|
||||
)
|
||||
compactPayload, err := buildXAIWebsocketCompactionPayload(req.Payload, transcriptInput)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
compactReq := req
|
||||
compactReq.Payload = compactPayload
|
||||
|
||||
prepared, data, headers, err := e.XAIExecutor.executeCompactRequest(ctx, auth, compactReq, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
responseID := xaiCompactionResponseID(data)
|
||||
idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID))
|
||||
idMapper.state.mapDownstreamToUpstream(responseID, "")
|
||||
|
||||
headers = headers.Clone()
|
||||
if headers == nil {
|
||||
headers = make(http.Header)
|
||||
}
|
||||
headers.Set("Content-Type", "text/event-stream")
|
||||
|
||||
chunks := xaiBuildCompactionTriggerStreamChunks(prepared, data)
|
||||
out := make(chan cliproxyexecutor.StreamChunk, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
out <- cliproxyexecutor.StreamChunk{Payload: chunk}
|
||||
}
|
||||
close(out)
|
||||
return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil
|
||||
}
|
||||
|
||||
func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) {
|
||||
if len(payload) == 0 {
|
||||
payload = []byte(`{}`)
|
||||
}
|
||||
if len(transcriptInput) == 0 {
|
||||
transcriptInput = []byte("[]")
|
||||
}
|
||||
out := bytes.Clone(payload)
|
||||
var err error
|
||||
out, err = sjson.SetRawBytes(out, "input", transcriptInput)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, _ = sjson.DeleteBytes(out, "previous_response_id")
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func xaiWebsocketGenerateFalse(payload []byte) bool {
|
||||
generate := gjson.GetBytes(payload, "generate")
|
||||
return generate.Exists() && !generate.Bool()
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -245,6 +246,170 @@ func TestXAIWebsocketsExecuteStreamRewritesRepeatedResponseIDForDownstream(t *te
|
||||
}
|
||||
}
|
||||
|
||||
func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedContext(t *testing.T) {
|
||||
upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
|
||||
capturedWebsocketPayload := make(chan []byte, 1)
|
||||
capturedCompactPayload := make(chan []byte, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/responses":
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
t.Errorf("upgrade websocket: %v", err)
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
_, payload, errRead := conn.ReadMessage()
|
||||
if errRead != nil {
|
||||
t.Errorf("read upstream websocket message: %v", errRead)
|
||||
return
|
||||
}
|
||||
capturedWebsocketPayload <- bytes.Clone(payload)
|
||||
completed := []byte(`{"type":"response.completed","response":{"id":"resp-real","output":[{"type":"message","id":"out-1","role":"assistant","content":"first answer"}],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`)
|
||||
if i == 1 {
|
||||
completed = []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2","role":"assistant","content":"second answer"}],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`)
|
||||
}
|
||||
if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {
|
||||
t.Errorf("write completed websocket message: %v", errWrite)
|
||||
return
|
||||
}
|
||||
}
|
||||
case "/responses/compact":
|
||||
body, errRead := io.ReadAll(r.Body)
|
||||
if errRead != nil {
|
||||
t.Errorf("read compact body: %v", errRead)
|
||||
return
|
||||
}
|
||||
capturedCompactPayload <- bytes.Clone(body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"resp_compact","model":"grok-4.3","output":[{"type":"compaction","encrypted_content":"opaque"}],"usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}`))
|
||||
default:
|
||||
t.Errorf("path = %q, want /responses", r.URL.Path)
|
||||
http.Error(w, "unexpected path", http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
exec := NewXAIWebsocketsExecutor(&config.Config{})
|
||||
exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)}
|
||||
exec.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
|
||||
auth := &cliproxyauth.Auth{
|
||||
ID: "xai-auth-compaction",
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
"websockets": "true",
|
||||
},
|
||||
Metadata: map[string]any{"access_token": "xai-token"},
|
||||
}
|
||||
opts := cliproxyexecutor.Options{
|
||||
SourceFormat: sdktranslator.FormatOpenAIResponse,
|
||||
ResponseFormat: sdktranslator.FormatOpenAIResponse,
|
||||
Stream: true,
|
||||
Metadata: map[string]any{
|
||||
cliproxyexecutor.ExecutionSessionMetadataKey: "xai-compaction-session",
|
||||
},
|
||||
}
|
||||
|
||||
result, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{
|
||||
Model: "grok-4.3",
|
||||
Payload: []byte(`{"model":"grok-4.3","stream":true,"input":[{"type":"message","id":"msg-1","role":"user","content":"first"}]}`),
|
||||
}, opts)
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteStream first turn error: %v", err)
|
||||
}
|
||||
for chunk := range result.Chunks {
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("stream chunk error = %v", chunk.Err)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case payload := <-capturedWebsocketPayload:
|
||||
if got := gjson.GetBytes(payload, "type").String(); got != "response.create" {
|
||||
t.Fatalf("type = %q, want response.create; payload=%s", got, payload)
|
||||
}
|
||||
input := gjson.GetBytes(payload, "input")
|
||||
if !input.IsArray() || len(input.Array()) != 1 {
|
||||
t.Fatalf("input = %s, want one first-turn item", input.Raw)
|
||||
}
|
||||
if gjson.GetBytes(payload, "stream").Exists() {
|
||||
t.Fatalf("stream must be omitted for xAI websocket payload: %s", payload)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for upstream websocket payload")
|
||||
}
|
||||
|
||||
compactResult, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{
|
||||
Model: "grok-4.3",
|
||||
Payload: []byte(`{"model":"grok-4.3","stream":true,"previous_response_id":"resp-real-xai-1","input":[{"type":"compaction_trigger"}]}`),
|
||||
}, opts)
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteStream compaction trigger error: %v", err)
|
||||
}
|
||||
for chunk := range compactResult.Chunks {
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("compact stream chunk error = %v", chunk.Err)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case payload := <-capturedCompactPayload:
|
||||
if xaiInputHasItemType(payload, "compaction_trigger") {
|
||||
t.Fatalf("compaction_trigger reached xai compact body: %s", payload)
|
||||
}
|
||||
input := gjson.GetBytes(payload, "input")
|
||||
if !input.IsArray() || len(input.Array()) != 2 {
|
||||
t.Fatalf("compact input = %s, want first request input plus response output", input.Raw)
|
||||
}
|
||||
if got := input.Array()[0].Get("id").String(); got != "msg-1" {
|
||||
t.Fatalf("compact input[0].id = %q, want msg-1; payload=%s", got, payload)
|
||||
}
|
||||
if got := input.Array()[1].Get("id").String(); got != "out-1" {
|
||||
t.Fatalf("compact input[1].id = %q, want out-1; payload=%s", got, payload)
|
||||
}
|
||||
if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "" {
|
||||
t.Fatalf("compact previous_response_id = %q, want empty; payload=%s", got, payload)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for compact HTTP payload")
|
||||
}
|
||||
|
||||
nextResult, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{
|
||||
Model: "grok-4.3",
|
||||
Payload: []byte(`{"model":"grok-4.3","stream":true,"previous_response_id":"resp_compact","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`),
|
||||
}, opts)
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteStream post-compaction turn error: %v", err)
|
||||
}
|
||||
for chunk := range nextResult.Chunks {
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("post-compaction stream chunk error = %v", chunk.Err)
|
||||
}
|
||||
}
|
||||
select {
|
||||
case payload := <-capturedWebsocketPayload:
|
||||
if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "" {
|
||||
t.Fatalf("post-compaction previous_response_id = %q, want empty; payload=%s", got, payload)
|
||||
}
|
||||
input := gjson.GetBytes(payload, "input")
|
||||
if !input.IsArray() || len(input.Array()) != 2 {
|
||||
t.Fatalf("post-compaction input = %s, want compaction item plus new message", input.Raw)
|
||||
}
|
||||
if got := input.Array()[0].Get("type").String(); got != "compaction" {
|
||||
t.Fatalf("post-compaction input[0].type = %q, want compaction; payload=%s", got, payload)
|
||||
}
|
||||
if got := input.Array()[1].Get("id").String(); got != "msg-2" {
|
||||
t.Fatalf("post-compaction input[1].id = %q, want msg-2; payload=%s", got, payload)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for post-compaction websocket 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"}]}`)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user