Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
488
backend/sdk/api/handlers/claude/code_handlers.go
Normal file
488
backend/sdk/api/handlers/claude/code_handlers.go
Normal file
|
|
@ -0,0 +1,488 @@
|
|||
// Package claude provides HTTP handlers for Claude API code-related functionality.
|
||||
// This package implements Claude-compatible streaming chat completions with sophisticated
|
||||
// client rotation and quota management systems to ensure high availability and optimal
|
||||
// resource utilization across multiple backend clients. It handles request translation
|
||||
// between Claude API format and the underlying Gemini backend, providing seamless
|
||||
// API compatibility while maintaining robust error handling and connection management.
|
||||
package claude
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
claudemodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/claude/models"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// ClaudeCodeAPIHandler contains the handlers for Claude API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service.
|
||||
type ClaudeCodeAPIHandler struct {
|
||||
*handlers.BaseAPIHandler
|
||||
}
|
||||
|
||||
// NewClaudeCodeAPIHandler creates a new Claude API handlers instance.
|
||||
// It takes an BaseAPIHandler instance as input and returns a ClaudeCodeAPIHandler.
|
||||
//
|
||||
// Parameters:
|
||||
// - apiHandlers: The base API handler instance.
|
||||
//
|
||||
// Returns:
|
||||
// - *ClaudeCodeAPIHandler: A new Claude code API handler instance.
|
||||
func NewClaudeCodeAPIHandler(apiHandlers *handlers.BaseAPIHandler) *ClaudeCodeAPIHandler {
|
||||
return &ClaudeCodeAPIHandler{
|
||||
BaseAPIHandler: apiHandlers,
|
||||
}
|
||||
}
|
||||
|
||||
// HandlerType returns the identifier for this handler implementation.
|
||||
func (h *ClaudeCodeAPIHandler) HandlerType() string {
|
||||
return Claude
|
||||
}
|
||||
|
||||
// Models returns a list of models supported by this handler.
|
||||
func (h *ClaudeCodeAPIHandler) Models() []map[string]any {
|
||||
// Get dynamic models from the global registry
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
return modelRegistry.GetAvailableModels("claude")
|
||||
}
|
||||
|
||||
// ClaudeMessages handles Claude-compatible streaming chat completions.
|
||||
// This function implements a sophisticated client rotation and quota management system
|
||||
// to ensure high availability and optimal resource utilization across multiple backend clients.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request.
|
||||
func (h *ClaudeCodeAPIHandler) ClaudeMessages(c *gin.Context) {
|
||||
// Extract raw JSON data from the incoming request
|
||||
rawJSON, err := c.GetRawData()
|
||||
// If data retrieval fails, return a 400 Bad Request error.
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Decode claude-fable-5-dd-<reversed> model IDs back to the real model name for routing.
|
||||
rawJSON = rewriteClaudeDDModelInBody(rawJSON)
|
||||
|
||||
// Check if the client requested a streaming response.
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
if !streamResult.Exists() || streamResult.Type == gjson.False {
|
||||
h.handleNonStreamingResponse(c, rawJSON)
|
||||
} else {
|
||||
h.handleStreamingResponse(c, rawJSON)
|
||||
}
|
||||
}
|
||||
|
||||
// ClaudeMessages handles Claude-compatible streaming chat completions.
|
||||
// This function implements a sophisticated client rotation and quota management system
|
||||
// to ensure high availability and optimal resource utilization across multiple backend clients.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request.
|
||||
func (h *ClaudeCodeAPIHandler) ClaudeCountTokens(c *gin.Context) {
|
||||
// Extract raw JSON data from the incoming request
|
||||
rawJSON, err := c.GetRawData()
|
||||
// If data retrieval fails, return a 400 Bad Request error.
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Decode claude-fable-5-dd-<reversed> model IDs back to the real model name for routing.
|
||||
rawJSON = rewriteClaudeDDModelInBody(rawJSON)
|
||||
|
||||
c.Header("Content-Type", "application/json")
|
||||
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// rewriteClaudeDDModelInBody decodes model IDs of the form claude-fable-5-dd-<reversed>
|
||||
// back into the original model name used for routing and upstream requests.
|
||||
func rewriteClaudeDDModelInBody(rawJSON []byte) []byte {
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
resolved := claudemodels.ResolveClaudeModelIDPrefix(modelName)
|
||||
if resolved == modelName {
|
||||
return rawJSON
|
||||
}
|
||||
updated, errSet := sjson.SetBytes(rawJSON, "model", resolved)
|
||||
if errSet != nil {
|
||||
return rawJSON
|
||||
}
|
||||
return updated
|
||||
}
|
||||
|
||||
// ClaudeModels handles the Claude models listing endpoint.
|
||||
// It returns a JSON response containing available Claude models and their specifications.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request.
|
||||
func (h *ClaudeCodeAPIHandler) ClaudeModels(c *gin.Context) {
|
||||
disableCloaking := h.Cfg != nil && h.Cfg.ClaudeCode.DisableCloakingModelList
|
||||
c.JSON(http.StatusOK, claudemodels.BuildResponse(h.Models(), disableCloaking))
|
||||
}
|
||||
|
||||
// handleNonStreamingResponse handles non-streaming content generation requests for Claude models.
|
||||
// This function processes the request synchronously and returns the complete generated
|
||||
// response in a single API call. It supports various generation parameters and
|
||||
// response formats.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for content generation
|
||||
// - rawJSON: The raw JSON request body containing generation parameters and content
|
||||
func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
|
||||
// Decompress gzipped responses - Claude API sometimes returns gzip without Content-Encoding header
|
||||
// This fixes title generation and other non-streaming responses that arrive compressed
|
||||
if len(resp) >= 2 && resp[0] == 0x1f && resp[1] == 0x8b {
|
||||
gzReader, errGzip := gzip.NewReader(bytes.NewReader(resp))
|
||||
if errGzip != nil {
|
||||
log.Warnf("failed to decompress gzipped Claude response: %v", errGzip)
|
||||
} else {
|
||||
defer func() {
|
||||
if errClose := gzReader.Close(); errClose != nil {
|
||||
log.Warnf("failed to close Claude gzip reader: %v", errClose)
|
||||
}
|
||||
}()
|
||||
decompressed, errRead := io.ReadAll(gzReader)
|
||||
if errRead != nil {
|
||||
log.Warnf("failed to read decompressed Claude response: %v", errRead)
|
||||
} else {
|
||||
resp = decompressed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleStreamingResponse streams Claude-compatible responses backed by Gemini.
|
||||
// It sets up SSE, selects a backend client with rotation/quota logic,
|
||||
// forwards chunks, and translates them to Claude CLI format.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request.
|
||||
// - rawJSON: The raw JSON request body.
|
||||
func (h *ClaudeCodeAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
// This is crucial for streaming as it allows immediate sending of data chunks
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
|
||||
// Create a cancellable context for the backend client request
|
||||
// This allows proper cleanup and cancellation of ongoing requests
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk to determine success or failure before setting headers
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg, hasPendingError := handlers.PendingStreamError(errChan); hasPendingError {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Stream closed without data? Send DONE or just headers.
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers now.
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
|
||||
// Write the first chunk
|
||||
if len(chunk) > 0 {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
// Continue streaming the rest
|
||||
h.forwardClaudeStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ClaudeCodeAPIHandler) forwardClaudeStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if len(chunk) == 0 {
|
||||
return
|
||||
}
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
c.Status(status)
|
||||
|
||||
errorBytes, _ := json.Marshal(h.toClaudeError(errMsg))
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", errorBytes)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
type claudeErrorDetail struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type claudeErrorResponse struct {
|
||||
Type string `json:"type"`
|
||||
Error claudeErrorDetail `json:"error"`
|
||||
}
|
||||
|
||||
func (h *ClaudeCodeAPIHandler) toClaudeError(msg *interfaces.ErrorMessage) claudeErrorResponse {
|
||||
status := http.StatusInternalServerError
|
||||
errText := http.StatusText(status)
|
||||
if msg != nil {
|
||||
if msg.StatusCode > 0 {
|
||||
status = msg.StatusCode
|
||||
errText = http.StatusText(status)
|
||||
}
|
||||
if msg.Error != nil {
|
||||
if v := strings.TrimSpace(msg.Error.Error()); v != "" {
|
||||
errText = v
|
||||
}
|
||||
}
|
||||
}
|
||||
errType, message := claudeErrorDetailFromText(status, errText)
|
||||
return claudeErrorResponse{
|
||||
Type: "error",
|
||||
Error: claudeErrorDetail{
|
||||
Type: errType,
|
||||
Message: message,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ClaudeCodeAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.ErrorMessage) {
|
||||
status := http.StatusInternalServerError
|
||||
if msg != nil && msg.StatusCode > 0 {
|
||||
status = msg.StatusCode
|
||||
}
|
||||
if msg != nil && msg.DirectResponse {
|
||||
for key, values := range handlers.FilterUpstreamHeaders(msg.Headers) {
|
||||
if len(values) == 0 || handlers.IsCPAReservedResponseHeader(key) {
|
||||
continue
|
||||
}
|
||||
c.Writer.Header().Del(key)
|
||||
for _, value := range values {
|
||||
c.Writer.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
body := bytes.Clone(msg.Body)
|
||||
appendClaudeAPIResponse(c, body)
|
||||
if !c.Writer.Written() && c.Writer.Header().Get("Content-Type") == "" {
|
||||
c.Writer.Header().Set("Content-Type", "application/json")
|
||||
}
|
||||
c.Status(status)
|
||||
_, _ = c.Writer.Write(body)
|
||||
return
|
||||
}
|
||||
if msg != nil && msg.Addon != nil && handlers.PassthroughHeadersEnabled(h.Cfg) {
|
||||
for key, values := range msg.Addon {
|
||||
if len(values) == 0 || handlers.IsCPAReservedResponseHeader(key) {
|
||||
continue
|
||||
}
|
||||
c.Writer.Header().Del(key)
|
||||
for _, value := range values {
|
||||
c.Writer.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
body, err := json.Marshal(h.toClaudeError(msg))
|
||||
if err != nil {
|
||||
body = []byte(`{"type":"error","error":{"type":"api_error","message":"Internal Server Error"}}`)
|
||||
}
|
||||
appendClaudeAPIResponse(c, body)
|
||||
if !c.Writer.Written() {
|
||||
c.Writer.Header().Set("Content-Type", "application/json")
|
||||
}
|
||||
c.Status(status)
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
|
||||
func claudeErrorDetailFromText(status int, errText string) (string, string) {
|
||||
message := strings.TrimSpace(errText)
|
||||
if message == "" {
|
||||
message = http.StatusText(status)
|
||||
}
|
||||
errType := claudeErrorTypeFromStatus(status)
|
||||
|
||||
var payload map[string]any
|
||||
if json.Valid([]byte(message)) {
|
||||
if err := json.Unmarshal([]byte(message), &payload); err == nil {
|
||||
if e, ok := payload["error"].(map[string]any); ok {
|
||||
if t, ok := e["type"].(string); ok && strings.TrimSpace(t) != "" {
|
||||
errType = strings.TrimSpace(t)
|
||||
}
|
||||
if m, ok := e["message"].(string); ok && strings.TrimSpace(m) != "" {
|
||||
message = strings.TrimSpace(m)
|
||||
} else if c, ok := e["code"].(string); ok && strings.TrimSpace(c) != "" {
|
||||
message = strings.TrimSpace(c)
|
||||
}
|
||||
} else {
|
||||
if t, ok := payload["type"].(string); ok && strings.TrimSpace(t) != "" && strings.TrimSpace(t) != "error" {
|
||||
errType = strings.TrimSpace(t)
|
||||
}
|
||||
if m, ok := payload["message"].(string); ok && strings.TrimSpace(m) != "" {
|
||||
message = strings.TrimSpace(m)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return errType, message
|
||||
}
|
||||
|
||||
func claudeErrorTypeFromStatus(status int) string {
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
return "authentication_error"
|
||||
case http.StatusPaymentRequired:
|
||||
return "billing_error"
|
||||
case http.StatusForbidden:
|
||||
return "permission_error"
|
||||
case http.StatusNotFound:
|
||||
return "not_found_error"
|
||||
case http.StatusRequestEntityTooLarge:
|
||||
return "request_too_large"
|
||||
case http.StatusTooManyRequests:
|
||||
return "rate_limit_error"
|
||||
case http.StatusGatewayTimeout:
|
||||
return "timeout_error"
|
||||
case 529:
|
||||
return "overloaded_error"
|
||||
default:
|
||||
if status >= http.StatusInternalServerError {
|
||||
return "api_error"
|
||||
}
|
||||
return "invalid_request_error"
|
||||
}
|
||||
}
|
||||
|
||||
func appendClaudeAPIResponse(c *gin.Context, data []byte) {
|
||||
if c == nil || len(data) == 0 {
|
||||
return
|
||||
}
|
||||
if _, exists := c.Get("API_RESPONSE_TIMESTAMP"); !exists {
|
||||
c.Set("API_RESPONSE_TIMESTAMP", time.Now())
|
||||
}
|
||||
if existing, exists := c.Get("API_RESPONSE"); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
|
||||
combined := make([]byte, 0, len(existingBytes)+len(data)+1)
|
||||
combined = append(combined, existingBytes...)
|
||||
if existingBytes[len(existingBytes)-1] != '\n' {
|
||||
combined = append(combined, '\n')
|
||||
}
|
||||
combined = append(combined, data...)
|
||||
c.Set("API_RESPONSE", combined)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Set("API_RESPONSE", bytes.Clone(data))
|
||||
}
|
||||
95
backend/sdk/api/handlers/claude/code_handlers_error_test.go
Normal file
95
backend/sdk/api/handlers/claude/code_handlers_error_test.go
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
package claude
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestClaudeErrorExtractsOpenAIStyleUpstreamJSON(t *testing.T) {
|
||||
handler := &ClaudeCodeAPIHandler{}
|
||||
msg := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errors.New(`{"error":{"message":"Your input exceeds the context window of this model. Please adjust your input and try again.","type":"invalid_request_error","code":"context_too_large"}}`),
|
||||
}
|
||||
|
||||
got := handler.toClaudeError(msg)
|
||||
|
||||
if got.Type != "error" {
|
||||
t.Fatalf("type = %q, want error", got.Type)
|
||||
}
|
||||
if got.Error.Type != "invalid_request_error" {
|
||||
t.Fatalf("error.type = %q, want invalid_request_error", got.Error.Type)
|
||||
}
|
||||
if got.Error.Message != "Your input exceeds the context window of this model. Please adjust your input and try again." {
|
||||
t.Fatalf("error.message = %q", got.Error.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaudeErrorExtractsClaudeStyleUpstreamJSON(t *testing.T) {
|
||||
handler := &ClaudeCodeAPIHandler{}
|
||||
msg := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: errors.New(`{"type":"error","error":{"type":"rate_limit_error","message":"This request would exceed your account's rate limit. Please try again later."},"request_id":"req_123"}`),
|
||||
}
|
||||
|
||||
got := handler.toClaudeError(msg)
|
||||
|
||||
if got.Error.Type != "rate_limit_error" {
|
||||
t.Fatalf("error.type = %q, want rate_limit_error", got.Error.Type)
|
||||
}
|
||||
if got.Error.Message != "This request would exceed your account's rate limit. Please try again later." {
|
||||
t.Fatalf("error.message = %q", got.Error.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteClaudeErrorResponseUsesClaudeEnvelope(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
handler := &ClaudeCodeAPIHandler{}
|
||||
msg := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errors.New(`{"error":{"message":"Your input exceeds the context window of this model. Please adjust your input and try again.","type":"invalid_request_error","code":"context_too_large"}}`),
|
||||
}
|
||||
|
||||
handler.WriteErrorResponse(c, msg)
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest)
|
||||
}
|
||||
body := recorder.Body.Bytes()
|
||||
if got := gjson.GetBytes(body, "type").String(); got != "error" {
|
||||
t.Fatalf("type = %q, want error; body=%s", got, body)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "error.type").String(); got != "invalid_request_error" {
|
||||
t.Fatalf("error.type = %q, want invalid_request_error; body=%s", got, body)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "error.message").String(); got != "Your input exceeds the context window of this model. Please adjust your input and try again." {
|
||||
t.Fatalf("error.message = %q; body=%s", got, body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPendingClaudeStreamErrorUsesBufferedError(t *testing.T) {
|
||||
wantErr := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errors.New(`{"error":{"message":"Your input exceeds the context window of this model. Please adjust your input and try again.","type":"invalid_request_error","code":"context_too_large"}}`),
|
||||
}
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- wantErr
|
||||
close(errs)
|
||||
|
||||
gotErr, ok := handlers.PendingStreamError(errs)
|
||||
if !ok {
|
||||
t.Fatal("expected pending stream error")
|
||||
}
|
||||
if gotErr != wantErr {
|
||||
t.Fatalf("pending error = %p, want %p", gotErr, wantErr)
|
||||
}
|
||||
}
|
||||
120
backend/sdk/api/handlers/claude/code_handlers_model_test.go
Normal file
120
backend/sdk/api/handlers/claude/code_handlers_model_test.go
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
package claude
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestClaudeModelsResponseUsesConfiguredDisplayName(t *testing.T) {
|
||||
const clientID = "claude-display-name-catalog-test"
|
||||
const modelID = "claude-display-name-catalog-test"
|
||||
registryRef := registry.GetGlobalRegistry()
|
||||
registryRef.RegisterClient(clientID, "claude", []*registry.ModelInfo{{
|
||||
ID: modelID, Object: "model", OwnedBy: "test", DisplayName: "Configured Claude Name",
|
||||
}})
|
||||
t.Cleanup(func() {
|
||||
registryRef.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
NewClaudeCodeAPIHandler(&handlers.BaseAPIHandler{}).ClaudeModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
DisplayName string `json:"display_name"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(recorder.Body.Bytes(), &response); errUnmarshal != nil {
|
||||
t.Fatalf("decode response: %v", errUnmarshal)
|
||||
}
|
||||
for _, model := range response.Data {
|
||||
if model.ID == modelID {
|
||||
if model.DisplayName != "Configured Claude Name" {
|
||||
t.Fatalf("display_name = %q, want Configured Claude Name", model.DisplayName)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("model %q not found in response", modelID)
|
||||
}
|
||||
|
||||
func TestClaudeModelsResponseDisablesModelListCloaking(t *testing.T) {
|
||||
const clientID = "claude-disable-model-list-cloaking-test"
|
||||
const modelID = "gpt-disable-model-list-cloaking-test"
|
||||
registryRef := registry.GetGlobalRegistry()
|
||||
registryRef.RegisterClient(clientID, "claude", []*registry.ModelInfo{{
|
||||
ID: modelID, Object: "model", OwnedBy: "test",
|
||||
}})
|
||||
t.Cleanup(func() {
|
||||
registryRef.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
baseHandler := &handlers.BaseAPIHandler{Cfg: &sdkconfig.SDKConfig{
|
||||
ClaudeCode: sdkconfig.ClaudeCodeConfig{DisableCloakingModelList: true},
|
||||
}}
|
||||
NewClaudeCodeAPIHandler(baseHandler).ClaudeModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(recorder.Body.Bytes(), &response); errUnmarshal != nil {
|
||||
t.Fatalf("decode response: %v", errUnmarshal)
|
||||
}
|
||||
for _, model := range response.Data {
|
||||
if model.ID == modelID {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("uncloaked model %q not found in response", modelID)
|
||||
}
|
||||
|
||||
func TestRewriteClaudeDDModelInBody(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantModel string
|
||||
}{
|
||||
{
|
||||
name: "encoded model is decoded",
|
||||
body: `{"model":"claude-fable-5-dd-o4-tpg","messages":[]}`,
|
||||
wantModel: "gpt-4o",
|
||||
},
|
||||
{
|
||||
name: "plain claude model unchanged",
|
||||
body: `{"model":"claude-sonnet-4-6","messages":[]}`,
|
||||
wantModel: "claude-sonnet-4-6",
|
||||
},
|
||||
{
|
||||
name: "encoded model with thinking suffix",
|
||||
body: `{"model":"claude-fable-5-dd-o4-tpg(high)","stream":true}`,
|
||||
wantModel: "gpt-4o(high)",
|
||||
},
|
||||
{
|
||||
name: "missing model field unchanged",
|
||||
body: `{"messages":[]}`,
|
||||
wantModel: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := rewriteClaudeDDModelInBody([]byte(tt.body))
|
||||
if model := gjson.GetBytes(got, "model").String(); model != tt.wantModel {
|
||||
t.Fatalf("model = %q, want %q; body=%s", model, tt.wantModel, string(got))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
350
backend/sdk/api/handlers/gemini/gemini_handlers.go
Normal file
350
backend/sdk/api/handlers/gemini/gemini_handlers.go
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
// Package gemini provides HTTP handlers for Gemini API endpoints.
|
||||
// This package implements handlers for managing Gemini model operations including
|
||||
// model listing, content generation, streaming content generation, and token counting.
|
||||
// It serves as a proxy layer between clients and the Gemini backend service,
|
||||
// handling request translation, client management, and response processing.
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
)
|
||||
|
||||
// GeminiAPIHandler contains the handlers for Gemini API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service.
|
||||
type GeminiAPIHandler struct {
|
||||
*handlers.BaseAPIHandler
|
||||
}
|
||||
|
||||
// NewGeminiAPIHandler creates a new Gemini API handlers instance.
|
||||
// It takes an BaseAPIHandler instance as input and returns a GeminiAPIHandler.
|
||||
func NewGeminiAPIHandler(apiHandlers *handlers.BaseAPIHandler) *GeminiAPIHandler {
|
||||
return &GeminiAPIHandler{
|
||||
BaseAPIHandler: apiHandlers,
|
||||
}
|
||||
}
|
||||
|
||||
// HandlerType returns the identifier for this handler implementation.
|
||||
func (h *GeminiAPIHandler) HandlerType() string {
|
||||
return Gemini
|
||||
}
|
||||
|
||||
// Models returns the Gemini-compatible model metadata supported by this handler.
|
||||
func (h *GeminiAPIHandler) Models() []map[string]any {
|
||||
// Get dynamic models from the global registry
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
return modelRegistry.GetAvailableModels("gemini")
|
||||
}
|
||||
|
||||
// GeminiModels handles the Gemini models listing endpoint.
|
||||
// It returns a JSON response containing available Gemini models and their specifications.
|
||||
func (h *GeminiAPIHandler) GeminiModels(c *gin.Context) {
|
||||
rawModels := h.Models()
|
||||
normalizedModels := make([]map[string]any, 0, len(rawModels))
|
||||
defaultMethods := []string{"generateContent"}
|
||||
for _, model := range rawModels {
|
||||
normalizedModel := make(map[string]any, len(model))
|
||||
for k, v := range model {
|
||||
normalizedModel[k] = v
|
||||
}
|
||||
if name, ok := normalizedModel["name"].(string); ok && name != "" {
|
||||
if !strings.HasPrefix(name, "models/") {
|
||||
normalizedModel["name"] = "models/" + name
|
||||
}
|
||||
if displayName, _ := normalizedModel["displayName"].(string); displayName == "" {
|
||||
normalizedModel["displayName"] = name
|
||||
}
|
||||
if description, _ := normalizedModel["description"].(string); description == "" {
|
||||
normalizedModel["description"] = name
|
||||
}
|
||||
}
|
||||
if _, ok := normalizedModel["supportedGenerationMethods"]; !ok {
|
||||
normalizedModel["supportedGenerationMethods"] = defaultMethods
|
||||
}
|
||||
normalizedModels = append(normalizedModels, normalizedModel)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"models": normalizedModels,
|
||||
})
|
||||
}
|
||||
|
||||
// GeminiGetHandler handles GET requests for specific Gemini model information.
|
||||
// It returns detailed information about a specific Gemini model based on the action parameter.
|
||||
func (h *GeminiAPIHandler) GeminiGetHandler(c *gin.Context) {
|
||||
var request struct {
|
||||
Action string `uri:"action" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindUri(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
action := strings.TrimPrefix(request.Action, "/")
|
||||
|
||||
// Get dynamic models from the global registry and find the matching one
|
||||
availableModels := h.Models()
|
||||
var targetModel map[string]any
|
||||
|
||||
for _, model := range availableModels {
|
||||
name, _ := model["name"].(string)
|
||||
// Match name with or without 'models/' prefix
|
||||
if name == action || name == "models/"+action {
|
||||
targetModel = model
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if targetModel != nil {
|
||||
// Ensure the name has 'models/' prefix in the output if it's a Gemini model
|
||||
if name, ok := targetModel["name"].(string); ok && name != "" && !strings.HasPrefix(name, "models/") {
|
||||
targetModel["name"] = "models/" + name
|
||||
}
|
||||
c.JSON(http.StatusOK, targetModel)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusNotFound, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Not Found",
|
||||
Type: "not_found",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// GeminiHandler handles POST requests for Gemini API operations.
|
||||
// It routes requests to appropriate handlers based on the action parameter (model:method format).
|
||||
func (h *GeminiAPIHandler) GeminiHandler(c *gin.Context) {
|
||||
var request struct {
|
||||
Action string `uri:"action" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindUri(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
action := strings.Split(strings.TrimPrefix(request.Action, "/"), ":")
|
||||
if len(action) != 2 {
|
||||
c.JSON(http.StatusNotFound, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("%s not found.", c.Request.URL.Path),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
method := action[1]
|
||||
rawJSON, _ := c.GetRawData()
|
||||
|
||||
switch method {
|
||||
case "generateContent":
|
||||
h.handleGenerateContent(c, action[0], rawJSON)
|
||||
case "streamGenerateContent":
|
||||
h.handleStreamGenerateContent(c, action[0], rawJSON)
|
||||
case "countTokens":
|
||||
h.handleCountTokens(c, action[0], rawJSON)
|
||||
}
|
||||
}
|
||||
|
||||
// handleStreamGenerateContent handles streaming content generation requests for Gemini models.
|
||||
// This function establishes a Server-Sent Events connection and streams the generated content
|
||||
// back to the client in real-time. It supports both SSE format and direct streaming based
|
||||
// on the 'alt' query parameter.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for content generation
|
||||
// - rawJSON: The raw JSON request body containing generation parameters
|
||||
func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
alt := h.GetAlt(c)
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg, hasPendingError := handlers.PendingStreamError(errChan); hasPendingError {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Closed without data
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
|
||||
// Write first chunk
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
flusher.Flush()
|
||||
|
||||
// Continue
|
||||
h.forwardGeminiStream(c, flusher, alt, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleCountTokens handles token counting requests for Gemini models.
|
||||
// This function counts the number of tokens in the provided content without
|
||||
// generating a response. It's useful for quota management and content validation.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for token counting
|
||||
// - rawJSON: The raw JSON request body containing the content to count
|
||||
func (h *GeminiAPIHandler) handleCountTokens(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleGenerateContent handles non-streaming content generation requests for Gemini models.
|
||||
// This function processes the request synchronously and returns the complete generated
|
||||
// response in a single API call. It supports various generation parameters and
|
||||
// response formats.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for content generation
|
||||
// - rawJSON: The raw JSON request body containing generation parameters and content
|
||||
func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) forwardGeminiStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
var keepAliveInterval *time.Duration
|
||||
if alt != "" {
|
||||
keepAliveInterval = new(time.Duration(0))
|
||||
}
|
||||
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
KeepAliveInterval: keepAliveInterval,
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
if alt == "" {
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,93 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
const (
|
||||
initialFailureGeminiModel = "initial-failure-gemini-model"
|
||||
)
|
||||
|
||||
type initialFailureGeminiStreamExecutor struct{}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Identifier() string {
|
||||
return "initial-failure-gemini-stream-executor"
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed before first payload")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func TestGeminiStreamGenerateContentDoesNotLoseErrorBeforeFirstPayload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
executor := &initialFailureGeminiStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
authID := fmt.Sprintf("initial-failure-gemini-auth-%d", idx)
|
||||
auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Errorf("register auth %d: %v", idx, errRegister)
|
||||
return
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: initialFailureGeminiModel}})
|
||||
defer registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewGeminiAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1beta/models/*action", h.GeminiHandler)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1beta/models/initial-failure-gemini-model:streamGenerateContent", strings.NewReader(`{"contents":[{"parts":[{"text":"hi"}]}]}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Errorf("request %d lost the buffered initial error and returned HTTP 200: %q", idx, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "upstream failed before first payload") {
|
||||
t.Errorf("request %d lost the initial upstream error: status=%d body=%q", idx, recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
)
|
||||
|
||||
func TestGeminiModelsResponseUsesConfiguredDisplayName(t *testing.T) {
|
||||
const clientID = "gemini-display-name-catalog-test"
|
||||
const modelID = "gemini-display-name-catalog-test"
|
||||
registryRef := registry.GetGlobalRegistry()
|
||||
registryRef.RegisterClient(clientID, "gemini", []*registry.ModelInfo{{
|
||||
ID: modelID, Name: modelID, DisplayName: "Configured Gemini Name",
|
||||
}})
|
||||
t.Cleanup(func() {
|
||||
registryRef.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
NewGeminiAPIHandler(&handlers.BaseAPIHandler{}).GeminiModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
DisplayName string `json:"displayName"`
|
||||
} `json:"models"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(recorder.Body.Bytes(), &response); errUnmarshal != nil {
|
||||
t.Fatalf("decode response: %v", errUnmarshal)
|
||||
}
|
||||
for _, model := range response.Models {
|
||||
if model.Name == "models/"+modelID {
|
||||
if model.DisplayName != "Configured Gemini Name" {
|
||||
t.Fatalf("displayName = %q, want Configured Gemini Name", model.DisplayName)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("model %q not found in response", modelID)
|
||||
}
|
||||
202
backend/sdk/api/handlers/gemini/interactions_handlers.go
Normal file
202
backend/sdk/api/handlers/gemini/interactions_handlers.go
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
const interactionsAgentAuthSelectionModel = "gemini-2.5-flash"
|
||||
|
||||
type interactionsRequestTarget struct {
|
||||
Model string
|
||||
Agent string
|
||||
Stream bool
|
||||
}
|
||||
|
||||
func parseInteractionsRequestTarget(rawJSON []byte) (interactionsRequestTarget, error) {
|
||||
if !gjson.ValidBytes(rawJSON) {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("invalid JSON body")
|
||||
}
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
model := strings.TrimSpace(root.Get("model").String())
|
||||
agent := strings.TrimSpace(root.Get("agent").String())
|
||||
if model == "" && agent == "" {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("request requires exactly one of model or agent")
|
||||
}
|
||||
if model != "" && agent != "" {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("request requires exactly one of model or agent")
|
||||
}
|
||||
streamNode := root.Get("stream")
|
||||
stream := false
|
||||
if streamNode.Exists() {
|
||||
if !streamNode.IsBool() {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("stream must be a boolean")
|
||||
}
|
||||
stream = streamNode.Bool()
|
||||
}
|
||||
return interactionsRequestTarget{Model: model, Agent: agent, Stream: stream}, nil
|
||||
}
|
||||
|
||||
func prepareInteractionsExecutionTarget(rawJSON []byte, target interactionsRequestTarget) (string, []byte) {
|
||||
if target.Agent != "" {
|
||||
return target.Agent, rawJSON
|
||||
}
|
||||
model := normalizeGeminiModelResourceName(target.Model)
|
||||
if model == target.Model {
|
||||
return model, rawJSON
|
||||
}
|
||||
updatedRawJSON, errSet := sjson.SetBytes(rawJSON, "model", model)
|
||||
if errSet != nil {
|
||||
return model, rawJSON
|
||||
}
|
||||
return model, updatedRawJSON
|
||||
}
|
||||
|
||||
func normalizeGeminiModelResourceName(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if strings.HasPrefix(model, "models/") && len(model) > len("models/") {
|
||||
return strings.TrimPrefix(model, "models/")
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
func buildInteractionsExecutionRequest(target interactionsRequestTarget, modelName string, rawJSON []byte, alt string) handlers.ProtocolExecutionRequest {
|
||||
forcedProvider := ""
|
||||
authSelectionModel := ""
|
||||
if target.Agent != "" {
|
||||
forcedProvider = GeminiInteractions
|
||||
authSelectionModel = interactionsAgentAuthSelectionModel
|
||||
}
|
||||
return handlers.ProtocolExecutionRequest{
|
||||
EntryProtocol: Interactions,
|
||||
ExitProtocol: Interactions,
|
||||
ForcedProvider: forcedProvider,
|
||||
AuthSelectionModel: authSelectionModel,
|
||||
Model: modelName,
|
||||
Stream: target.Stream,
|
||||
Body: rawJSON,
|
||||
Alt: alt,
|
||||
}
|
||||
}
|
||||
|
||||
// Interactions handles POST /v1beta/interactions.
|
||||
func (h *GeminiAPIHandler) Interactions(c *gin.Context) {
|
||||
rawJSON, errRead := c.GetRawData()
|
||||
if errRead != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: errRead.Error(), Type: "invalid_request_error"}})
|
||||
return
|
||||
}
|
||||
target, errParse := parseInteractionsRequestTarget(rawJSON)
|
||||
if errParse != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: errParse.Error(), Type: "invalid_request_error"}})
|
||||
return
|
||||
}
|
||||
|
||||
modelName, resolvedRawJSON := prepareInteractionsExecutionTarget(rawJSON, target)
|
||||
rawJSON = resolvedRawJSON
|
||||
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
defer cliCancel(nil)
|
||||
|
||||
req := buildInteractionsExecutionRequest(target, modelName, rawJSON, alt)
|
||||
if target.Stream {
|
||||
h.handleInteractionsStream(c, cliCtx, cliCancel, req)
|
||||
return
|
||||
}
|
||||
h.handleInteractionsNonStream(c, cliCtx, cliCancel, req)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) handleInteractionsNonStream(c *gin.Context, cliCtx context.Context, cliCancel handlers.APIHandlerCancelFunc, req handlers.ProtocolExecutionRequest) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, errMsg := h.ExecuteProtocolWithAuthManager(cliCtx, req)
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), resp.Headers)
|
||||
_, _ = c.Writer.Write(resp.Body)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) handleInteractionsStream(c *gin.Context, cliCtx context.Context, cliCancel handlers.APIHandlerCancelFunc, req handlers.ProtocolExecutionRequest) {
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: "Streaming not supported", Type: "server_error"}})
|
||||
return
|
||||
}
|
||||
stream, errMsg := h.ExecuteProtocolStreamWithAuthManager(cliCtx, req)
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), stream.Headers)
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
go func() {
|
||||
defer close(data)
|
||||
defer close(errs)
|
||||
for chunk := range stream.Chunks {
|
||||
if chunk.Err != nil {
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: chunk.Err.StatusCode, Error: chunk.Err}
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) > 0 {
|
||||
data <- chunk.Payload
|
||||
}
|
||||
}
|
||||
}()
|
||||
h.forwardInteractionsStream(c, flusher, func(err error) { cliCancel(err) }, data, errs)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) forwardInteractionsStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if len(chunk) == 0 {
|
||||
return
|
||||
}
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
if bytes.HasPrefix(trimmed, []byte("event:")) || bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
} else {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
if !bytes.HasSuffix(chunk, []byte("\n\n")) {
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
}
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
},
|
||||
})
|
||||
}
|
||||
320
backend/sdk/api/handlers/gemini/interactions_handlers_test.go
Normal file
320
backend/sdk/api/handlers/gemini/interactions_handlers_test.go
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestParseInteractionsRequestTarget(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantModel string
|
||||
wantAgent string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "model", body: `{"model":"gemini-3.5-flash","input":"hi"}`, wantModel: "gemini-3.5-flash"},
|
||||
{name: "model resource name", body: `{"model":"models/gemini-3.5-flash","input":"hi"}`, wantModel: "models/gemini-3.5-flash"},
|
||||
{name: "agent", body: `{"agent":"agents/test-agent","input":"hi"}`, wantAgent: "agents/test-agent"},
|
||||
{name: "missing", body: `{"input":"hi"}`, wantErr: true},
|
||||
{name: "both", body: `{"model":"gemini-3.5-flash","agent":"agents/test-agent","input":"hi"}`, wantErr: true},
|
||||
{name: "stream string", body: `{"model":"gemini-3.5-flash","stream":"true","input":"hi"}`, wantErr: true},
|
||||
{name: "stream true", body: `{"model":"gemini-3.5-flash","stream":true,"input":"hi"}`, wantModel: "gemini-3.5-flash"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(tt.body))
|
||||
if tt.wantErr {
|
||||
if errParse == nil {
|
||||
t.Fatal("parseInteractionsRequestTarget() error = nil, want error")
|
||||
}
|
||||
return
|
||||
}
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
if target.Model != tt.wantModel || target.Agent != tt.wantAgent {
|
||||
t.Fatalf("target = %#v, want model %q agent %q", target, tt.wantModel, tt.wantAgent)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareInteractionsExecutionTargetNormalizesModelResourceName(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"model":"models/gemini-3.5-flash","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
model, body := prepareInteractionsExecutionTarget([]byte(`{"model":"models/gemini-3.5-flash","input":"hi"}`), target)
|
||||
if model != "gemini-3.5-flash" {
|
||||
t.Fatalf("model = %q, want gemini-3.5-flash", model)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "model").String(); got != "gemini-3.5-flash" {
|
||||
t.Fatalf("body model = %q, want gemini-3.5-flash. Body: %s", got, string(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareInteractionsExecutionTargetPreservesBareModel(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"model":"gemini-3.5-flash","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
model, body := prepareInteractionsExecutionTarget([]byte(`{"model":"gemini-3.5-flash","input":"hi"}`), target)
|
||||
if model != "gemini-3.5-flash" {
|
||||
t.Fatalf("model = %q, want gemini-3.5-flash", model)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "model").String(); got != "gemini-3.5-flash" {
|
||||
t.Fatalf("body model = %q, want gemini-3.5-flash. Body: %s", got, string(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildInteractionsExecutionRequestUsesAgentAuthSelectionModel(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"agent":"agents/test-agent","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
req := buildInteractionsExecutionRequest(target, "agents/test-agent", []byte(`{"agent":"agents/test-agent","input":"hi"}`), "")
|
||||
if req.ForcedProvider != "gemini-interactions" {
|
||||
t.Fatalf("ForcedProvider = %q, want gemini-interactions", req.ForcedProvider)
|
||||
}
|
||||
if req.AuthSelectionModel != interactionsAgentAuthSelectionModel {
|
||||
t.Fatalf("AuthSelectionModel = %q, want %q", req.AuthSelectionModel, interactionsAgentAuthSelectionModel)
|
||||
}
|
||||
if req.Model != "agents/test-agent" {
|
||||
t.Fatalf("Model = %q, want agents/test-agent", req.Model)
|
||||
}
|
||||
if got := gjson.GetBytes(req.Body, "agent").String(); got != "agents/test-agent" {
|
||||
t.Fatalf("body agent = %q, want agents/test-agent", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsInvalidJSON(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "invalid_request_error") {
|
||||
t.Fatalf("body = %s, want invalid_request_error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsMissingModelAndAgent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "exactly one of model or agent") {
|
||||
t.Fatalf("body = %s, want model/agent validation error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsBothModelAndAgent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"gemini-3.5-flash","agent":"agents/test-agent","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "exactly one of model or agent") {
|
||||
t.Fatalf("body = %s, want model/agent validation error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsNonBooleanStream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"gemini-3.5-flash","stream":"true","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "invalid_request_error") {
|
||||
t.Fatalf("body = %s, want invalid_request_error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsAgentUsesNativeInteractionsEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
var gotPath string
|
||||
var upstreamBody []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
body, errRead := io.ReadAll(r.Body)
|
||||
if errRead != nil {
|
||||
http.Error(w, errRead.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
upstreamBody = body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"text":"ok"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor.NewGeminiInteractionsExecutor(&config.Config{RequestRetry: 1}))
|
||||
auth := &coreauth.Auth{
|
||||
ID: "interactions-agent-native-auth",
|
||||
Provider: "gemini-interactions",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
"base_url": server.URL,
|
||||
},
|
||||
Metadata: map[string]any{"email": "interactions-agent@example.com"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: interactionsAgentAuthSelectionModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"agent":"agents/test-agent","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager))
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if gotPath != "/v1beta/interactions" {
|
||||
t.Fatalf("path = %q, want /v1beta/interactions", gotPath)
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "agent").String(); got != "agents/test-agent" {
|
||||
t.Fatalf("upstream agent = %q, want agents/test-agent. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(rec.Body.Bytes(), "id").String(); got != "interaction_1" {
|
||||
t.Fatalf("response id = %q, want interaction_1. Body: %s", got, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsAntigravityModelUsesTranslatorBridge(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
model := "interactions-antigravity-bridge-model"
|
||||
var upstreamBody []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1internal:generateContent" {
|
||||
http.Error(w, "unexpected path: "+r.URL.Path, http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
body, errRead := io.ReadAll(r.Body)
|
||||
if errRead != nil {
|
||||
http.Error(w, errRead.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
upstreamBody = body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"response":{"responseId":"resp_1","candidates":[{"content":{"role":"model","parts":[{"text":"translated-ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":2,"totalTokenCount":3}}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor.NewAntigravityExecutor(&config.Config{RequestRetry: 1}))
|
||||
auth := &coreauth.Auth{
|
||||
ID: "interactions-antigravity-bridge-auth",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"access_token": "token",
|
||||
"project_id": "project-1",
|
||||
"expired": time.Now().Add(time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"`+model+`","input":"hi","generation_config":{"top_p":0.8}}`))
|
||||
h := NewGeminiAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager))
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if gjson.GetBytes(upstreamBody, "input").Exists() {
|
||||
t.Fatalf("upstream body still contains raw interactions input: %s", string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "request.contents.0.parts.0.text").String(); got != "hi" {
|
||||
t.Fatalf("upstream request text = %q, want hi. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "request.generationConfig.topP").Float(); got != 0.8 {
|
||||
t.Fatalf("upstream topP = %v, want 0.8. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(rec.Body.Bytes(), "steps.0.content.0.text").String(); got != "translated-ok" {
|
||||
t.Fatalf("response text = %q, want translated-ok. Body: %s", got, rec.Body.String())
|
||||
}
|
||||
if gjson.GetBytes(rec.Body.Bytes(), "response").Exists() {
|
||||
t.Fatalf("response still contains raw antigravity response wrapper: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardInteractionsStreamWrapsBareJSONAsSSEData(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{}`))
|
||||
data := make(chan []byte, 1)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte(`{"type":"interaction.completed"}`)
|
||||
close(data)
|
||||
close(errs)
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.forwardInteractionsStream(ctx, rec, func(error) {}, data, errs)
|
||||
|
||||
if got := rec.Body.String(); got != "data: {\"type\":\"interaction.completed\"}\n\n" {
|
||||
t.Fatalf("body = %q, want SSE data frame", got)
|
||||
}
|
||||
}
|
||||
581
backend/sdk/api/handlers/handlers.go
Normal file
581
backend/sdk/api/handlers/handlers.go
Normal file
|
|
@ -0,0 +1,581 @@
|
|||
// Package handlers provides core API handler functionality for the CLI Proxy API server.
|
||||
// It includes common types, client management, load balancing, and error handling
|
||||
// shared across all API endpoint handlers (OpenAI, Claude, Gemini).
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coresession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// ErrorResponse represents a standard error response format for the API.
|
||||
// It contains a single ErrorDetail field.
|
||||
type ErrorResponse struct {
|
||||
// Error contains detailed information about the error that occurred.
|
||||
Error ErrorDetail `json:"error"`
|
||||
}
|
||||
|
||||
// ErrorDetail provides specific information about an error that occurred.
|
||||
// It includes a human-readable message, an error type, and an optional error code.
|
||||
type ErrorDetail struct {
|
||||
// Message is a human-readable message providing more details about the error.
|
||||
Message string `json:"message"`
|
||||
|
||||
// Type is the category of error that occurred (e.g., "invalid_request_error").
|
||||
Type string `json:"type"`
|
||||
|
||||
// Code is a short code identifying the error, if applicable.
|
||||
Code string `json:"code,omitempty"`
|
||||
}
|
||||
|
||||
const idempotencyKeyMetadataKey = "idempotency_key"
|
||||
|
||||
const (
|
||||
defaultStreamingKeepAliveSeconds = 0
|
||||
defaultStreamingBootstrapRetries = 0
|
||||
// Stream interceptor history is intentionally bounded and not configurable in the first SDK surface.
|
||||
maxStreamInterceptorHistoryChunks = 64
|
||||
maxStreamInterceptorHistoryBytes = 1 << 20
|
||||
)
|
||||
|
||||
// BuildErrorResponseBody builds an OpenAI-compatible JSON error response body.
|
||||
// If errText is already valid JSON, it is returned as-is to preserve upstream error payloads.
|
||||
func BuildErrorResponseBody(status int, errText string) []byte {
|
||||
if status <= 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
if strings.TrimSpace(errText) == "" {
|
||||
errText = http.StatusText(status)
|
||||
}
|
||||
|
||||
trimmed := strings.TrimSpace(errText)
|
||||
if trimmed != "" && json.Valid([]byte(trimmed)) {
|
||||
return []byte(trimmed)
|
||||
}
|
||||
|
||||
errType := "invalid_request_error"
|
||||
var code string
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
errType = "authentication_error"
|
||||
code = "invalid_api_key"
|
||||
case http.StatusForbidden:
|
||||
errType = "permission_error"
|
||||
code = "insufficient_quota"
|
||||
case http.StatusTooManyRequests:
|
||||
errType = "rate_limit_error"
|
||||
code = "rate_limit_exceeded"
|
||||
case http.StatusNotFound:
|
||||
errType = "invalid_request_error"
|
||||
code = "model_not_found"
|
||||
default:
|
||||
if status >= http.StatusInternalServerError {
|
||||
errType = "server_error"
|
||||
code = "internal_server_error"
|
||||
}
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(ErrorResponse{
|
||||
Error: ErrorDetail{
|
||||
Message: errText,
|
||||
Type: errType,
|
||||
Code: code,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return []byte(fmt.Sprintf(`{"error":{"message":%q,"type":"server_error","code":"internal_server_error"}}`, errText))
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
// StreamingKeepAliveInterval returns the SSE keep-alive interval for this server.
|
||||
// Returning 0 disables keep-alives (default when unset).
|
||||
func StreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
|
||||
seconds := defaultStreamingKeepAliveSeconds
|
||||
if cfg != nil {
|
||||
seconds = cfg.Streaming.KeepAliveSeconds
|
||||
}
|
||||
if seconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
// NonStreamingKeepAliveInterval returns the keep-alive interval for non-streaming responses.
|
||||
// Returning 0 disables keep-alives (default when unset).
|
||||
func NonStreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
|
||||
seconds := 0
|
||||
if cfg != nil {
|
||||
seconds = cfg.NonStreamKeepAliveInterval
|
||||
}
|
||||
if seconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
// StreamingBootstrapRetries returns how many times a streaming request may be retried before any bytes are sent.
|
||||
func StreamingBootstrapRetries(cfg *config.SDKConfig) int {
|
||||
retries := defaultStreamingBootstrapRetries
|
||||
if cfg != nil {
|
||||
retries = cfg.Streaming.BootstrapRetries
|
||||
}
|
||||
if retries < 0 {
|
||||
retries = 0
|
||||
}
|
||||
return retries
|
||||
}
|
||||
|
||||
// PassthroughHeadersEnabled returns whether upstream response headers should be forwarded to clients.
|
||||
// Default is false.
|
||||
func PassthroughHeadersEnabled(cfg *config.SDKConfig) bool {
|
||||
return cfg != nil && cfg.PassthroughHeaders
|
||||
}
|
||||
|
||||
func requestExecutionMetadata(ctx context.Context) map[string]any {
|
||||
// Idempotency-Key is an optional client-supplied header used to correlate retries.
|
||||
// Only include it if the client explicitly provides it.
|
||||
key := ""
|
||||
requestPath := ""
|
||||
var ginCtx *gin.Context
|
||||
if ctx != nil {
|
||||
if requestGinCtx, ok := ctx.Value("gin").(*gin.Context); ok && requestGinCtx != nil && requestGinCtx.Request != nil {
|
||||
ginCtx = requestGinCtx
|
||||
key = strings.TrimSpace(ginCtx.GetHeader("Idempotency-Key"))
|
||||
requestPath = strings.TrimSpace(ginCtx.FullPath())
|
||||
if requestPath == "" && ginCtx.Request.URL != nil {
|
||||
requestPath = strings.TrimSpace(ginCtx.Request.URL.Path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
meta := make(map[string]any)
|
||||
if key != "" {
|
||||
meta[idempotencyKeyMetadataKey] = key
|
||||
}
|
||||
if requestPath != "" {
|
||||
meta[coreexecutor.RequestPathMetadataKey] = requestPath
|
||||
}
|
||||
if pinnedAuthID := pinnedAuthIDFromContext(ctx); pinnedAuthID != "" {
|
||||
meta[coreexecutor.PinnedAuthMetadataKey] = pinnedAuthID
|
||||
}
|
||||
if selectedCallback := selectedAuthIDCallbackFromContext(ctx); selectedCallback != nil {
|
||||
meta[coreexecutor.SelectedAuthCallbackMetadataKey] = selectedCallback
|
||||
}
|
||||
if ginCtx != nil && !websocket.IsWebSocketUpgrade(ginCtx.Request) {
|
||||
if traceCallback := logging.GinCPATraceIDCallback(ginCtx); traceCallback != nil {
|
||||
meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey] = traceCallback
|
||||
}
|
||||
}
|
||||
if executionSessionID := executionSessionIDFromContext(ctx); executionSessionID != "" {
|
||||
meta[coreexecutor.ExecutionSessionMetadataKey] = executionSessionID
|
||||
}
|
||||
if callerScope := requestCallerScope(ginCtx); callerScope != "" {
|
||||
meta[coreexecutor.CallerScopeMetadataKey] = callerScope
|
||||
}
|
||||
if disallowFreeAuthFromContext(ctx) {
|
||||
meta[coreexecutor.DisallowFreeAuthMetadataKey] = true
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
func requestClientIP(request *http.Request) string {
|
||||
if request == nil {
|
||||
return ""
|
||||
}
|
||||
remoteAddr := strings.TrimSpace(request.RemoteAddr)
|
||||
if host, _, errSplit := net.SplitHostPort(remoteAddr); errSplit == nil {
|
||||
return strings.TrimSpace(host)
|
||||
}
|
||||
return remoteAddr
|
||||
}
|
||||
|
||||
func requestCallerScope(ginCtx *gin.Context) string {
|
||||
if ginCtx == nil {
|
||||
return ""
|
||||
}
|
||||
value, exists := ginCtx.Get("userApiKey")
|
||||
if !exists || value == nil {
|
||||
return ""
|
||||
}
|
||||
return coresession.CallerScope(fmt.Sprint(value))
|
||||
}
|
||||
|
||||
func addAuthSelectionModelMetadata(meta map[string]any, model string) {
|
||||
if meta == nil {
|
||||
return
|
||||
}
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return
|
||||
}
|
||||
meta[coreexecutor.AuthSelectionModelMetadataKey] = model
|
||||
}
|
||||
|
||||
func setReasoningEffortMetadata(meta map[string]any, handlerType, model string, rawJSON []byte) {
|
||||
if meta == nil {
|
||||
return
|
||||
}
|
||||
effort := thinking.ExtractReasoningEffort(rawJSON, handlerType, model)
|
||||
if effort == "" {
|
||||
return
|
||||
}
|
||||
meta[coreexecutor.ReasoningEffortMetadataKey] = effort
|
||||
}
|
||||
|
||||
func setServiceTierMetadata(meta map[string]any, rawJSON []byte) {
|
||||
if meta == nil {
|
||||
return
|
||||
}
|
||||
serviceTier := coreusage.AutoServiceTier
|
||||
node := gjson.GetBytes(rawJSON, "service_tier")
|
||||
if node.Exists() {
|
||||
value := strings.TrimSpace(node.String())
|
||||
if value != "" {
|
||||
serviceTier = value
|
||||
}
|
||||
}
|
||||
meta[coreexecutor.ServiceTierMetadataKey] = serviceTier
|
||||
}
|
||||
|
||||
func setGenerateMetadata(meta map[string]any, rawJSON []byte) {
|
||||
if meta == nil {
|
||||
return
|
||||
}
|
||||
// Missing or true means generation is enabled; only an explicit false disables generation.
|
||||
generate := true
|
||||
node := gjson.GetBytes(rawJSON, "generate")
|
||||
if node.Exists() && node.IsBool() && !node.Bool() {
|
||||
generate = false
|
||||
}
|
||||
meta[coreexecutor.GenerateMetadataKey] = generate
|
||||
}
|
||||
|
||||
// BaseAPIHandler contains the handlers for API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service and manages
|
||||
// load balancing, client selection, and configuration.
|
||||
type BaseAPIHandler struct {
|
||||
// AuthManager manages auth lifecycle and execution in the new architecture.
|
||||
AuthManager *coreauth.Manager
|
||||
|
||||
// Cfg holds the current application configuration.
|
||||
Cfg *config.SDKConfig
|
||||
|
||||
// PluginHost optionally applies plugin interceptors around upstream execution.
|
||||
PluginHost PluginInterceptorHost
|
||||
|
||||
// ModelRouterHost optionally routes matching requests to a plugin executor, the router's own
|
||||
// executor, or a built-in provider before model-to-provider resolution and auth selection.
|
||||
ModelRouterHost PluginModelRouterHost
|
||||
}
|
||||
|
||||
// NewBaseAPIHandlers creates a new API handlers instance.
|
||||
// It takes a slice of clients and configuration as input.
|
||||
//
|
||||
// Parameters:
|
||||
// - cliClients: A slice of AI service clients
|
||||
// - cfg: The application configuration
|
||||
//
|
||||
// Returns:
|
||||
// - *BaseAPIHandler: A new API handlers instance
|
||||
func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager) *BaseAPIHandler {
|
||||
return &BaseAPIHandler{
|
||||
Cfg: cfg,
|
||||
AuthManager: authManager,
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateClients updates the handlers' client list and configuration.
|
||||
// This method is called when the configuration or authentication tokens change.
|
||||
//
|
||||
// Parameters:
|
||||
// - clients: The new slice of AI service clients
|
||||
// - cfg: The new application configuration
|
||||
func (h *BaseAPIHandler) UpdateClients(cfg *config.SDKConfig) { h.Cfg = cfg }
|
||||
|
||||
// SetPluginHost configures the optional plugin interceptor host.
|
||||
func (h *BaseAPIHandler) SetPluginHost(host PluginInterceptorHost) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
if isNilPluginInterceptorHost(host) {
|
||||
h.PluginHost = nil
|
||||
return
|
||||
}
|
||||
h.PluginHost = host
|
||||
}
|
||||
|
||||
// SetModelRouterHost configures the optional plugin model router host.
|
||||
func (h *BaseAPIHandler) SetModelRouterHost(host PluginModelRouterHost) {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
if isNilPluginModelRouterHost(host) {
|
||||
h.ModelRouterHost = nil
|
||||
return
|
||||
}
|
||||
h.ModelRouterHost = host
|
||||
}
|
||||
|
||||
func isNilPluginInterceptorHost(host PluginInterceptorHost) bool {
|
||||
return isNilInterface(host)
|
||||
}
|
||||
|
||||
func isNilPluginModelRouterHost(host PluginModelRouterHost) bool {
|
||||
return isNilInterface(host)
|
||||
}
|
||||
|
||||
func isNilInterface(value any) bool {
|
||||
if value == nil {
|
||||
return true
|
||||
}
|
||||
// A typed nil pointer stored in an interface is not equal to nil.
|
||||
reflected := reflect.ValueOf(value)
|
||||
switch reflected.Kind() {
|
||||
case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
|
||||
return reflected.IsNil()
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// GetAlt extracts the 'alt' parameter from the request query string.
|
||||
// It checks both 'alt' and '$alt' parameters and returns the appropriate value.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request
|
||||
//
|
||||
// Returns:
|
||||
// - string: The alt parameter value, or empty string if it's "sse"
|
||||
func (h *BaseAPIHandler) GetAlt(c *gin.Context) string {
|
||||
var alt string
|
||||
var hasAlt bool
|
||||
alt, hasAlt = c.GetQuery("alt")
|
||||
if !hasAlt {
|
||||
alt, _ = c.GetQuery("$alt")
|
||||
}
|
||||
if alt == "sse" {
|
||||
return ""
|
||||
}
|
||||
return alt
|
||||
}
|
||||
|
||||
// GetContextWithCancel creates a new context with cancellation capabilities.
|
||||
// It embeds the Gin context and the API handler into the new context for later use.
|
||||
// The returned cancel function also handles logging the API response if request logging is enabled.
|
||||
//
|
||||
// Parameters:
|
||||
// - handler: The API handler associated with the request.
|
||||
// - c: The Gin context of the current request.
|
||||
// - ctx: The parent context (caller values/deadlines are preserved; request context adds cancellation and request ID).
|
||||
//
|
||||
// Returns:
|
||||
// - context.Context: The new context with cancellation and embedded values.
|
||||
// - APIHandlerCancelFunc: A function to cancel the context and log the response.
|
||||
func (h *BaseAPIHandler) GetContextWithCancel(handler interfaces.APIHandler, c *gin.Context, ctx context.Context) (context.Context, APIHandlerCancelFunc) {
|
||||
parentCtx := ctx
|
||||
if parentCtx == nil {
|
||||
parentCtx = context.Background()
|
||||
}
|
||||
|
||||
var requestCtx context.Context
|
||||
if c != nil && c.Request != nil {
|
||||
requestCtx = c.Request.Context()
|
||||
}
|
||||
|
||||
if requestCtx != nil && logging.GetRequestID(parentCtx) == "" {
|
||||
if requestID := logging.GetRequestID(requestCtx); requestID != "" {
|
||||
parentCtx = logging.WithRequestID(parentCtx, requestID)
|
||||
} else if requestID = logging.GetGinRequestID(c); requestID != "" {
|
||||
parentCtx = logging.WithRequestID(parentCtx, requestID)
|
||||
}
|
||||
}
|
||||
newCtx, cancel := context.WithCancel(parentCtx)
|
||||
|
||||
endpoint := ""
|
||||
if c != nil && c.Request != nil {
|
||||
path := strings.TrimSpace(c.FullPath())
|
||||
if path == "" && c.Request.URL != nil {
|
||||
path = strings.TrimSpace(c.Request.URL.Path)
|
||||
}
|
||||
if path != "" {
|
||||
method := strings.TrimSpace(c.Request.Method)
|
||||
if method != "" {
|
||||
endpoint = method + " " + path
|
||||
} else {
|
||||
endpoint = path
|
||||
}
|
||||
}
|
||||
}
|
||||
if endpoint != "" {
|
||||
newCtx = logging.WithEndpoint(newCtx, endpoint)
|
||||
}
|
||||
if c != nil && c.Request != nil {
|
||||
newCtx = logging.WithClientRequestMetadata(newCtx, logging.ClientRequestMetadata{
|
||||
ClientIP: requestClientIP(c.Request),
|
||||
XForwardedFor: strings.TrimSpace(strings.Join(c.Request.Header.Values("X-Forwarded-For"), ", ")),
|
||||
UserAgent: strings.TrimSpace(c.Request.UserAgent()),
|
||||
})
|
||||
}
|
||||
newCtx = logging.WithResponseStatusHolder(newCtx)
|
||||
newCtx = logging.WithResponseHeadersHolder(newCtx)
|
||||
|
||||
cancelCtx := newCtx
|
||||
if requestCtx != nil && requestCtx != parentCtx {
|
||||
go func() {
|
||||
select {
|
||||
case <-requestCtx.Done():
|
||||
cancel()
|
||||
case <-cancelCtx.Done():
|
||||
}
|
||||
}()
|
||||
}
|
||||
newCtx = context.WithValue(newCtx, "gin", c)
|
||||
newCtx = context.WithValue(newCtx, "handler", handler)
|
||||
return newCtx, func(params ...interface{}) {
|
||||
if c != nil {
|
||||
logging.SetResponseStatus(cancelCtx, c.Writer.Status())
|
||||
}
|
||||
if h.Cfg.RequestLog && len(params) == 1 {
|
||||
if captured, exists := c.Get(logging.APIResponseCapturedContextKey); exists {
|
||||
if capturedBool, ok := captured.(bool); ok && capturedBool {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
if existing, exists := c.Get("API_RESPONSE"); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(bytes.TrimSpace(existingBytes)) > 0 {
|
||||
switch params[0].(type) {
|
||||
case error, string:
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var payload []byte
|
||||
switch data := params[0].(type) {
|
||||
case []byte:
|
||||
payload = data
|
||||
case error:
|
||||
if data != nil {
|
||||
payload = []byte(data.Error())
|
||||
}
|
||||
case string:
|
||||
payload = []byte(data)
|
||||
}
|
||||
if len(payload) > 0 {
|
||||
if existing, exists := c.Get("API_RESPONSE"); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
|
||||
trimmedPayload := bytes.TrimSpace(payload)
|
||||
if len(trimmedPayload) > 0 && bytes.Contains(existingBytes, trimmedPayload) {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
appendAPIResponse(c, payload)
|
||||
}
|
||||
}
|
||||
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// StartNonStreamingKeepAlive emits blank lines every 5 seconds while waiting for a non-streaming response.
|
||||
// It returns a stop function that must be called before writing the final response.
|
||||
func (h *BaseAPIHandler) StartNonStreamingKeepAlive(c *gin.Context, ctx context.Context) func() {
|
||||
if h == nil || c == nil {
|
||||
return func() {}
|
||||
}
|
||||
interval := NonStreamingKeepAliveInterval(h.Cfg)
|
||||
if interval <= 0 {
|
||||
return func() {}
|
||||
}
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
return func() {}
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
stopChan := make(chan struct{})
|
||||
var stopOnce sync.Once
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return func() {
|
||||
stopOnce.Do(func() {
|
||||
close(stopChan)
|
||||
})
|
||||
wg.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
// appendAPIResponse preserves any previously captured API response and appends new data.
|
||||
func appendAPIResponse(c *gin.Context, data []byte) {
|
||||
if c == nil || len(data) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Capture timestamp on first API response
|
||||
if _, exists := c.Get("API_RESPONSE_TIMESTAMP"); !exists {
|
||||
c.Set("API_RESPONSE_TIMESTAMP", time.Now())
|
||||
}
|
||||
|
||||
if existing, exists := c.Get("API_RESPONSE"); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
|
||||
combined := make([]byte, 0, len(existingBytes)+len(data)+1)
|
||||
combined = append(combined, existingBytes...)
|
||||
if existingBytes[len(existingBytes)-1] != '\n' {
|
||||
combined = append(combined, '\n')
|
||||
}
|
||||
combined = append(combined, data...)
|
||||
c.Set("API_RESPONSE", combined)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
c.Set("API_RESPONSE", bytes.Clone(data))
|
||||
}
|
||||
|
||||
// APIHandlerCancelFunc is a function type for canceling an API handler's context.
|
||||
// It can optionally accept parameters, which are used for logging the response.
|
||||
type APIHandlerCancelFunc func(params ...interface{})
|
||||
208
backend/sdk/api/handlers/handlers_context.go
Normal file
208
backend/sdk/api/handlers/handlers_context.go
Normal file
|
|
@ -0,0 +1,208 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
type pinnedAuthContextKey struct{}
|
||||
|
||||
type selectedAuthCallbackContextKey struct{}
|
||||
|
||||
type preparedModelRouteContextKey struct{}
|
||||
|
||||
type executionSessionContextKey struct{}
|
||||
|
||||
type disallowFreeAuthContextKey struct{}
|
||||
|
||||
type nestedExecutionTrackerKey struct{}
|
||||
|
||||
type nestedExecutionTracker struct {
|
||||
mu sync.Mutex
|
||||
called bool
|
||||
}
|
||||
|
||||
func (t *nestedExecutionTracker) mark() {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
t.mu.Lock()
|
||||
t.called = true
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
func (t *nestedExecutionTracker) hasNestedExecution() bool {
|
||||
if t == nil {
|
||||
return false
|
||||
}
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.called
|
||||
}
|
||||
|
||||
func withNestedExecutionTracker(ctx context.Context) (context.Context, *nestedExecutionTracker) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if existing, ok := ctx.Value(nestedExecutionTrackerKey{}).(*nestedExecutionTracker); ok && existing != nil {
|
||||
return ctx, existing
|
||||
}
|
||||
tracker := &nestedExecutionTracker{}
|
||||
return context.WithValue(ctx, nestedExecutionTrackerKey{}, tracker), tracker
|
||||
}
|
||||
|
||||
func markNestedExecution(ctx context.Context) {
|
||||
if ctx == nil {
|
||||
return
|
||||
}
|
||||
if tracker, ok := ctx.Value(nestedExecutionTrackerKey{}).(*nestedExecutionTracker); ok && tracker != nil {
|
||||
tracker.mark()
|
||||
}
|
||||
}
|
||||
|
||||
// WithPinnedAuthID returns a child context that requests execution on a specific auth ID.
|
||||
func WithPinnedAuthID(ctx context.Context, authID string) context.Context {
|
||||
authID = strings.TrimSpace(authID)
|
||||
if authID == "" {
|
||||
return ctx
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, pinnedAuthContextKey{}, authID)
|
||||
}
|
||||
|
||||
// WithSelectedAuthIDCallback returns a child context that receives the selected auth ID.
|
||||
func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) context.Context {
|
||||
if callback == nil {
|
||||
return ctx
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
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, skipRouterPluginID string) (modelRouteDecision, bool) {
|
||||
// A host.model.execute_stream callback is a nested execution. Its caller is
|
||||
// excluded from model routing, so an outer prepared route cannot be reused:
|
||||
// it may point straight back at that caller.
|
||||
if ctx == nil || strings.TrimSpace(skipRouterPluginID) != "" {
|
||||
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)
|
||||
if sessionID == "" {
|
||||
return ctx
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, executionSessionContextKey{}, sessionID)
|
||||
}
|
||||
|
||||
// WithDisallowFreeAuth returns a child context that requests skipping known free-tier credentials.
|
||||
func WithDisallowFreeAuth(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, disallowFreeAuthContextKey{}, true)
|
||||
}
|
||||
|
||||
// headersFromContext extracts the original HTTP request headers from the gin context
|
||||
// embedded in the provided context. This allows session affinity selectors to read
|
||||
// client-provided session headers.
|
||||
func headersFromContext(ctx context.Context) http.Header {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
|
||||
return ginCtx.Request.Header.Clone()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// queryFromContext extracts the original HTTP request query parameters from the
|
||||
// gin context embedded in the provided context. Mirrors headersFromContext so
|
||||
// model routers can observe inbound query parameters for plain HTTP requests,
|
||||
// where execOptions.Query is not populated by callers.
|
||||
func queryFromContext(ctx context.Context) url.Values {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil && ginCtx.Request.URL != nil {
|
||||
return ginCtx.Request.URL.Query()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func pinnedAuthIDFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
raw := ctx.Value(pinnedAuthContextKey{})
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
case []byte:
|
||||
return strings.TrimSpace(string(v))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func selectedAuthIDCallbackFromContext(ctx context.Context) func(string) {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
raw := ctx.Value(selectedAuthCallbackContextKey{})
|
||||
if callback, ok := raw.(func(string)); ok && callback != nil {
|
||||
return callback
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func executionSessionIDFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
raw := ctx.Value(executionSessionContextKey{})
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
case []byte:
|
||||
return strings.TrimSpace(string(v))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func disallowFreeAuthFromContext(ctx context.Context) bool {
|
||||
if ctx == nil {
|
||||
return false
|
||||
}
|
||||
raw, ok := ctx.Value(disallowFreeAuthContextKey{}).(bool)
|
||||
return ok && raw
|
||||
}
|
||||
280
backend/sdk/api/handlers/handlers_error_response_test.go
Normal file
280
backend/sdk/api/handlers/handlers_error_response_test.go
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestWriteErrorResponse_AddonHeadersDisabledByDefault(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
handler := NewBaseAPIHandlers(nil, nil)
|
||||
handler.WriteErrorResponse(c, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: errors.New("rate limit"),
|
||||
Addon: http.Header{
|
||||
"Retry-After": {"30"},
|
||||
"X-Request-Id": {"req-1"},
|
||||
},
|
||||
})
|
||||
|
||||
if recorder.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
if got := recorder.Header().Get("Retry-After"); got != "" {
|
||||
t.Fatalf("Retry-After should be empty when passthrough is disabled, got %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("X-Request-Id"); got != "" {
|
||||
t.Fatalf("X-Request-Id should be empty when passthrough is disabled, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteErrorResponseDirectResponse(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
c.Writer.Header().Set("X-Cpa-Trace-Id", "local-trace")
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", "https://trusted.example")
|
||||
|
||||
handler := NewBaseAPIHandlers(nil, nil)
|
||||
handler.WriteErrorResponse(c, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusForbidden,
|
||||
DirectResponse: true,
|
||||
Body: []byte(`{"error":"blocked"}`),
|
||||
Headers: http.Header{
|
||||
"Content-Type": {"application/problem+json"},
|
||||
"X-Plugin-Policy": {"blocked"},
|
||||
"X-Cpa-Trace-Id": {"plugin-trace"},
|
||||
"Access-Control-Allow-Origin": {"https://untrusted.example"},
|
||||
},
|
||||
})
|
||||
|
||||
if recorder.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusForbidden)
|
||||
}
|
||||
if got := recorder.Body.String(); got != `{"error":"blocked"}` {
|
||||
t.Fatalf("body = %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Content-Type"); got != "application/problem+json" {
|
||||
t.Fatalf("Content-Type = %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("X-Plugin-Policy"); got != "blocked" {
|
||||
t.Fatalf("X-Plugin-Policy = %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("X-Cpa-Trace-Id"); got != "local-trace" {
|
||||
t.Fatalf("X-Cpa-Trace-Id = %q, want local value", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Access-Control-Allow-Origin"); got != "https://trusted.example" {
|
||||
t.Fatalf("Access-Control-Allow-Origin = %q, want trusted origin", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInternalConcurrencyBusyWritesRetryAfterWithoutPassthrough(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
handler := NewBaseAPIHandlers(nil, nil)
|
||||
handler.WriteErrorResponse(c, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: coreauth.NewHomeConcurrencyBusyError("busy", 750*time.Millisecond),
|
||||
})
|
||||
|
||||
if recorder.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
if got := recorder.Header().Get("Retry-After"); got != "1" {
|
||||
t.Fatalf("Retry-After = %q, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteErrorResponseHomeBusyNormalAndStreamHeaders(t *testing.T) {
|
||||
for _, stream := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "normal", true: "stream"}[stream], func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
if stream {
|
||||
c.Request.Header.Set("Accept", "text/event-stream")
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(nil, nil)
|
||||
handler.WriteErrorResponse(c, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: coreauth.NewHomeConcurrencyBusyError("busy", 750*time.Millisecond),
|
||||
})
|
||||
if recorder.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
if got := recorder.Header().Get("Retry-After"); got != "1" {
|
||||
t.Fatalf("Retry-After = %q, want 1", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteErrorResponse_AddonHeadersEnabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Writer.Header().Set("X-Request-Id", "old-value")
|
||||
c.Writer.Header().Set("x-cpa-trace-id", "local-trace")
|
||||
c.Writer.Header().Set("Access-Control-Expose-Headers", "x-cpa-trace-id")
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true}, nil)
|
||||
handler.WriteErrorResponse(c, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: errors.New("rate limit"),
|
||||
Addon: http.Header{
|
||||
"Retry-After": {"30"},
|
||||
"X-Request-Id": {"new-1", "new-2"},
|
||||
"x-cpa-trace-id": {"upstream-trace"},
|
||||
"Access-Control-Expose-Headers": {"upstream-header"},
|
||||
},
|
||||
})
|
||||
|
||||
if recorder.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
if got := recorder.Header().Get("Retry-After"); got != "30" {
|
||||
t.Fatalf("Retry-After = %q, want %q", got, "30")
|
||||
}
|
||||
if got := recorder.Header().Values("X-Request-Id"); !reflect.DeepEqual(got, []string{"new-1", "new-2"}) {
|
||||
t.Fatalf("X-Request-Id = %#v, want %#v", got, []string{"new-1", "new-2"})
|
||||
}
|
||||
if got := recorder.Header().Get("x-cpa-trace-id"); got != "local-trace" {
|
||||
t.Fatalf("x-cpa-trace-id = %q, want local trace", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Access-Control-Expose-Headers"); got != "x-cpa-trace-id" {
|
||||
t.Fatalf("Access-Control-Expose-Headers = %q, want CPA value", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnrichAuthSelectionError_DefaultsTo503WithContext(t *testing.T) {
|
||||
in := &coreauth.Error{Code: "auth_not_found", Message: "no auth available"}
|
||||
out := enrichAuthSelectionError(in, []string{"claude"}, "claude-sonnet-4-6")
|
||||
|
||||
var got *coreauth.Error
|
||||
if !errors.As(out, &got) || got == nil {
|
||||
t.Fatalf("expected coreauth.Error, got %T", out)
|
||||
}
|
||||
if got.StatusCode() != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status = %d, want %d", got.StatusCode(), http.StatusServiceUnavailable)
|
||||
}
|
||||
if !strings.Contains(got.Message, "providers=claude") {
|
||||
t.Fatalf("message missing provider context: %q", got.Message)
|
||||
}
|
||||
if !strings.Contains(got.Message, "model=claude-sonnet-4-6") {
|
||||
t.Fatalf("message missing model context: %q", got.Message)
|
||||
}
|
||||
if !strings.Contains(got.Message, "/v0/management/auth-files") {
|
||||
t.Fatalf("message missing management hint: %q", got.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnrichAuthSelectionError_PreservesExplicitStatus(t *testing.T) {
|
||||
in := &coreauth.Error{Code: "auth_unavailable", Message: "no auth available", HTTPStatus: http.StatusTooManyRequests}
|
||||
out := enrichAuthSelectionError(in, []string{"gemini"}, "gemini-2.5-pro")
|
||||
|
||||
var got *coreauth.Error
|
||||
if !errors.As(out, &got) || got == nil {
|
||||
t.Fatalf("expected coreauth.Error, got %T", out)
|
||||
}
|
||||
if got.StatusCode() != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", got.StatusCode(), http.StatusTooManyRequests)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnrichAuthSelectionError_IgnoresOtherErrors(t *testing.T) {
|
||||
in := errors.New("boom")
|
||||
out := enrichAuthSelectionError(in, []string{"claude"}, "claude-sonnet-4-6")
|
||||
if out != in {
|
||||
t.Fatalf("expected original error to be returned unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionErrorMessageMapsContextStatuses(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
err error
|
||||
want int
|
||||
}{
|
||||
{name: "canceled", err: context.Canceled, want: clienterror.StatusClientClosedRequest},
|
||||
{name: "deadline", err: context.DeadlineExceeded, want: http.StatusGatewayTimeout},
|
||||
{
|
||||
name: "url error wraps canceled",
|
||||
err: &url.Error{Op: "Post", URL: "https://example.com", Err: context.Canceled},
|
||||
want: clienterror.StatusClientClosedRequest,
|
||||
},
|
||||
{name: "plain error defaults to 500", err: errors.New("boom"), want: http.StatusInternalServerError},
|
||||
{
|
||||
name: "explicit status wins",
|
||||
err: &coreauth.Error{Code: "rate_limited", Message: "slow down", HTTPStatus: http.StatusTooManyRequests},
|
||||
want: http.StatusTooManyRequests,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
msg := executionErrorMessage(tc.err)
|
||||
if msg == nil {
|
||||
t.Fatalf("executionErrorMessage() returned nil")
|
||||
}
|
||||
if msg.StatusCode != tc.want {
|
||||
t.Fatalf("StatusCode = %d, want %d", msg.StatusCode, tc.want)
|
||||
}
|
||||
if msg.Error != tc.err {
|
||||
t.Fatalf("Error = %v, want original %v", msg.Error, tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusFromErrorMapsContextStatuses(t *testing.T) {
|
||||
if got := statusFromError(context.Canceled); got != clienterror.StatusClientClosedRequest {
|
||||
t.Fatalf("statusFromError(canceled) = %d, want %d", got, clienterror.StatusClientClosedRequest)
|
||||
}
|
||||
if got := statusFromError(context.DeadlineExceeded); got != http.StatusGatewayTimeout {
|
||||
t.Fatalf("statusFromError(deadline) = %d, want %d", got, http.StatusGatewayTimeout)
|
||||
}
|
||||
if got := statusFromError(&url.Error{Op: "Post", URL: "https://example.com", Err: context.Canceled}); got != clienterror.StatusClientClosedRequest {
|
||||
t.Fatalf("statusFromError(url canceled) = %d, want %d", got, clienterror.StatusClientClosedRequest)
|
||||
}
|
||||
if got := statusFromError(errors.New("boom")); got != 0 {
|
||||
t.Fatalf("statusFromError(plain) = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteErrorResponse_ContextCanceledUses499(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
handler := NewBaseAPIHandlers(nil, nil)
|
||||
handler.WriteErrorResponse(c, executionErrorMessage(context.Canceled))
|
||||
|
||||
if recorder.Code != clienterror.StatusClientClosedRequest {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, clienterror.StatusClientClosedRequest)
|
||||
}
|
||||
}
|
||||
170
backend/sdk/api/handlers/handlers_errors.go
Normal file
170
backend/sdk/api/handlers/handlers_errors.go
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
func statusFromError(err error) int {
|
||||
return clienterror.HTTPStatusFromError(err)
|
||||
}
|
||||
|
||||
func isAuthSelectionUnavailable(err error) bool {
|
||||
var authErr *coreauth.Error
|
||||
if !errors.As(err, &authErr) || authErr == nil {
|
||||
return false
|
||||
}
|
||||
code := strings.TrimSpace(authErr.Code)
|
||||
return code == "auth_not_found" || code == "auth_unavailable"
|
||||
}
|
||||
|
||||
func enrichAuthSelectionError(err error, providers []string, model string) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var authErr *coreauth.Error
|
||||
if !errors.As(err, &authErr) || authErr == nil {
|
||||
return err
|
||||
}
|
||||
|
||||
code := strings.TrimSpace(authErr.Code)
|
||||
if code != "auth_not_found" && code != "auth_unavailable" {
|
||||
return err
|
||||
}
|
||||
|
||||
providerText := strings.Join(providers, ",")
|
||||
if providerText == "" {
|
||||
providerText = "unknown"
|
||||
}
|
||||
modelText := strings.TrimSpace(model)
|
||||
if modelText == "" {
|
||||
modelText = "unknown"
|
||||
}
|
||||
|
||||
baseMessage := strings.TrimSpace(authErr.Message)
|
||||
if baseMessage == "" {
|
||||
baseMessage = "no auth available"
|
||||
}
|
||||
detail := fmt.Sprintf("%s (providers=%s, model=%s)", baseMessage, providerText, modelText)
|
||||
|
||||
// Clarify the most common alias confusion between Anthropic route names and internal provider keys.
|
||||
if strings.Contains(","+providerText+",", ",claude,") {
|
||||
detail += "; check Claude auth/key session and cooldown state via /v0/management/auth-files"
|
||||
}
|
||||
|
||||
status := authErr.HTTPStatus
|
||||
if status <= 0 {
|
||||
status = http.StatusServiceUnavailable
|
||||
}
|
||||
|
||||
return &coreauth.Error{
|
||||
Code: authErr.Code,
|
||||
Message: detail,
|
||||
Retryable: authErr.Retryable,
|
||||
HTTPStatus: status,
|
||||
}
|
||||
}
|
||||
|
||||
// WriteErrorResponse writes an error message to the response writer using the HTTP status embedded in the message.
|
||||
func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.ErrorMessage) {
|
||||
status := http.StatusInternalServerError
|
||||
if msg != nil && msg.StatusCode > 0 {
|
||||
status = msg.StatusCode
|
||||
}
|
||||
if msg != nil && msg.DirectResponse {
|
||||
writeDirectErrorResponse(c, status, msg)
|
||||
return
|
||||
}
|
||||
if msg != nil && msg.Error != nil {
|
||||
for _, value := range coreauth.SafeResponseHeaders(msg.Error).Values("Retry-After") {
|
||||
c.Writer.Header().Add("Retry-After", value)
|
||||
}
|
||||
}
|
||||
if msg != nil && msg.Addon != nil && PassthroughHeadersEnabled(h.Cfg) {
|
||||
for key, values := range msg.Addon {
|
||||
if len(values) == 0 || IsCPAReservedResponseHeader(key) {
|
||||
continue
|
||||
}
|
||||
c.Writer.Header().Del(key)
|
||||
for _, value := range values {
|
||||
c.Writer.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
errText := http.StatusText(status)
|
||||
if msg != nil && msg.Error != nil {
|
||||
if v := strings.TrimSpace(msg.Error.Error()); v != "" {
|
||||
errText = v
|
||||
}
|
||||
}
|
||||
|
||||
body := BuildErrorResponseBody(status, errText)
|
||||
// Append first to preserve upstream response logs, then drop duplicate payloads if already recorded.
|
||||
var previous []byte
|
||||
if existing, exists := c.Get("API_RESPONSE"); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
|
||||
previous = existingBytes
|
||||
}
|
||||
}
|
||||
appendAPIResponse(c, body)
|
||||
trimmedErrText := strings.TrimSpace(errText)
|
||||
trimmedBody := bytes.TrimSpace(body)
|
||||
if len(previous) > 0 {
|
||||
if (trimmedErrText != "" && bytes.Contains(previous, []byte(trimmedErrText))) ||
|
||||
(len(trimmedBody) > 0 && bytes.Contains(previous, trimmedBody)) {
|
||||
c.Set("API_RESPONSE", previous)
|
||||
}
|
||||
}
|
||||
|
||||
if !c.Writer.Written() {
|
||||
c.Writer.Header().Set("Content-Type", "application/json")
|
||||
}
|
||||
c.Status(status)
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
|
||||
func writeDirectErrorResponse(c *gin.Context, status int, msg *interfaces.ErrorMessage) {
|
||||
for key, values := range FilterUpstreamHeaders(msg.Headers) {
|
||||
if len(values) == 0 || IsCPAReservedResponseHeader(key) {
|
||||
continue
|
||||
}
|
||||
c.Writer.Header().Del(key)
|
||||
for _, value := range values {
|
||||
c.Writer.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
body := bytes.Clone(msg.Body)
|
||||
appendAPIResponse(c, body)
|
||||
if !c.Writer.Written() && c.Writer.Header().Get("Content-Type") == "" {
|
||||
c.Writer.Header().Set("Content-Type", "application/json")
|
||||
}
|
||||
c.Status(status)
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) LoggingAPIResponseError(ctx context.Context, err *interfaces.ErrorMessage) {
|
||||
if h.Cfg.RequestLog {
|
||||
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
|
||||
if apiResponseErrors, isExist := ginContext.Get("API_RESPONSE_ERROR"); isExist {
|
||||
if slicesAPIResponseError, isOk := apiResponseErrors.([]*interfaces.ErrorMessage); isOk {
|
||||
slicesAPIResponseError = append(slicesAPIResponseError, err)
|
||||
ginContext.Set("API_RESPONSE_ERROR", slicesAPIResponseError)
|
||||
}
|
||||
} else {
|
||||
// Create new response data entry
|
||||
ginContext.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{err})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
349
backend/sdk/api/handlers/handlers_execution.go
Normal file
349
backend/sdk/api/handlers/handlers_execution.go
Normal file
|
|
@ -0,0 +1,349 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// PluginExecutorHost executes a routed request with a specific plugin executor.
|
||||
type PluginExecutorHost interface {
|
||||
ExecutePluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error)
|
||||
ExecutePluginExecutorStream(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error)
|
||||
CountPluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error)
|
||||
}
|
||||
|
||||
type pluginExecutorFormatResolver interface {
|
||||
PluginExecutorRequestToFormat(string, coreexecutor.Request, coreexecutor.Options) sdktranslator.Format
|
||||
}
|
||||
|
||||
// ExecuteWithAuthManager executes a non-streaming request via the core auth manager.
|
||||
// This path is the only supported execution route.
|
||||
func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false)
|
||||
}
|
||||
|
||||
// ExecuteImageWithAuthManager executes an OpenAI-compatible image endpoint request.
|
||||
func (h *BaseAPIHandler) ExecuteImageWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true)
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
return h.executeWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{})
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
originalRequestedModel := modelName
|
||||
routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, false, execOptions)
|
||||
responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)
|
||||
if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
if routeDecision.ExecutorPluginID != "" {
|
||||
return h.executeWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions)
|
||||
}
|
||||
providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions)
|
||||
if errMsg != nil {
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers)
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel
|
||||
addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel)
|
||||
addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource)
|
||||
setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
setGenerateMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
}
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: payload,
|
||||
}
|
||||
afterAuthCapture := &requestAfterAuthCapture{}
|
||||
lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID)
|
||||
opts := coreexecutor.Options{
|
||||
Stream: false,
|
||||
Alt: alt,
|
||||
OriginalRequest: rawJSON,
|
||||
SourceFormat: sdktranslator.FromString(entryProtocol),
|
||||
ResponseFormat: sdktranslator.FromString(responseProtocol),
|
||||
Headers: modelExecutionHeaders(ctx, execOptions.Headers),
|
||||
Query: modelExecutionQuery(ctx, execOptions.Query),
|
||||
RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID),
|
||||
}
|
||||
opts.Metadata = reqMeta
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(ctx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
resp, err := h.AuthManager.Execute(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
err = enrichAuthSelectionError(err, providers, normalizedModel)
|
||||
errMsg := executionErrorMessage(err)
|
||||
lifecycle.completeError(ctx, errMsg)
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
executedReq, executedOpts := afterAuthCapture.apply(req, opts)
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), responseProtocol, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID)
|
||||
lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
// ExecuteCountWithAuthManager executes a non-streaming request via the core auth manager.
|
||||
// This path is the only supported execution route.
|
||||
func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
return h.executeCountWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, modelExecutionOptions{})
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
originalRequestedModel := modelName
|
||||
routeDecision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, false, execOptions)
|
||||
if routeDecision.ExecutorPluginID != "" {
|
||||
return h.countWithPluginExecutor(ctx, handlerType, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions)
|
||||
}
|
||||
providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, false, routeDecision, execOptions)
|
||||
if errMsg != nil {
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
providers = adjustExecutionProvidersForEntryProtocol(handlerType, providers)
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel
|
||||
addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel)
|
||||
setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
setGenerateMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
}
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: payload,
|
||||
}
|
||||
afterAuthCapture := &requestAfterAuthCapture{}
|
||||
lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID)
|
||||
opts := coreexecutor.Options{
|
||||
Stream: false,
|
||||
Alt: alt,
|
||||
OriginalRequest: rawJSON,
|
||||
SourceFormat: sdktranslator.FromString(handlerType),
|
||||
Headers: modelExecutionHeaders(ctx, execOptions.Headers),
|
||||
Query: modelExecutionQuery(ctx, execOptions.Query),
|
||||
RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID),
|
||||
}
|
||||
opts.Metadata = reqMeta
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(ctx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
err = enrichAuthSelectionError(err, providers, normalizedModel)
|
||||
errMsg := executionErrorMessage(err)
|
||||
lifecycle.completeError(ctx, errMsg)
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
executedReq, executedOpts := afterAuthCapture.apply(req, opts)
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), handlerType, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID)
|
||||
lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
if h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")}
|
||||
}
|
||||
host := h.pluginExecutorHost()
|
||||
if host == nil {
|
||||
return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")}
|
||||
}
|
||||
execCtx, nestedTracker := withNestedExecutionTracker(ctx)
|
||||
req, opts := h.pluginExecutorRequest(execCtx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions)
|
||||
lifecycle := h.newRequestLifecycleTracker(execCtx, entryProtocol, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID)
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(execCtx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(execCtx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(execCtx, host, executorPluginID, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(execCtx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
var reporter *helps.UsageReporter
|
||||
if !execOptions.InternalSource {
|
||||
reporter = helps.NewUsageReporter(execCtx, executorPluginID, modelName, nil)
|
||||
reporter.SetTranslatedReasoningEffort(req.Payload, entryProtocol)
|
||||
}
|
||||
resp, errExecute := host.ExecutePluginExecutor(execCtx, executorPluginID, req, opts)
|
||||
if errExecute != nil {
|
||||
if reporter != nil && !nestedTracker.hasNestedExecution() {
|
||||
reporter.PublishFailure(execCtx, errExecute)
|
||||
}
|
||||
errMsg := executionErrorMessage(errExecute)
|
||||
lifecycle.completeError(execCtx, errMsg)
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
if reporter != nil && !nestedTracker.hasNestedExecution() {
|
||||
detail := parsePluginExecutorResponseUsage(responseProtocol, resp.Payload)
|
||||
reporter.Publish(execCtx, detail)
|
||||
reporter.EnsurePublished(execCtx)
|
||||
}
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(execCtx, lifecycle.requestID(), responseProtocol, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID)
|
||||
lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerType, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) {
|
||||
if h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")}
|
||||
}
|
||||
host := h.pluginExecutorHost()
|
||||
if host == nil {
|
||||
return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")}
|
||||
}
|
||||
req, opts := h.pluginExecutorRequest(ctx, handlerType, handlerType, modelName, originalRequestedModel, rawJSON, alt, false, execOptions)
|
||||
lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID)
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(ctx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(ctx, interceptErr)
|
||||
return nil, nil, interceptErr
|
||||
}
|
||||
resp, errCount := host.CountPluginExecutor(ctx, executorPluginID, req, opts)
|
||||
if errCount != nil {
|
||||
errMsg := executionErrorMessage(errCount)
|
||||
lifecycle.completeError(ctx, errMsg)
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), handlerType, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID)
|
||||
lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt string, stream bool, execOptions modelExecutionOptions) (coreexecutor.Request, coreexecutor.Options) {
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel
|
||||
addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel)
|
||||
addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource)
|
||||
setReasoningEffortMetadata(reqMeta, entryProtocol, modelName, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
setGenerateMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
}
|
||||
req := coreexecutor.Request{Model: modelName, Payload: payload}
|
||||
opts := coreexecutor.Options{
|
||||
Stream: stream,
|
||||
Alt: alt,
|
||||
OriginalRequest: rawJSON,
|
||||
SourceFormat: sdktranslator.FromString(entryProtocol),
|
||||
ResponseFormat: sdktranslator.FromString(responseProtocol),
|
||||
Headers: modelExecutionHeaders(ctx, execOptions.Headers),
|
||||
Query: modelExecutionQuery(ctx, execOptions.Query),
|
||||
Metadata: reqMeta,
|
||||
}
|
||||
return req, opts
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx context.Context, host PluginExecutorHost, executorPluginID, entryProtocol, originalRequestedModel, requestID string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options, *interfaces.ErrorMessage) {
|
||||
if !requestInterceptorsEnabled(h.interceptorHost()) {
|
||||
return req, opts, nil
|
||||
}
|
||||
toFormat := sdktranslator.FromString(entryProtocol)
|
||||
if resolver, ok := host.(pluginExecutorFormatResolver); ok && resolver != nil {
|
||||
if resolved := resolver.PluginExecutorRequestToFormat(executorPluginID, req, opts); resolved != "" {
|
||||
toFormat = resolved
|
||||
}
|
||||
}
|
||||
resp := h.applyRequestInterceptorsAfterAuth(ctx, coreexecutor.RequestAfterAuthInterceptRequest{
|
||||
SourceFormat: opts.SourceFormat,
|
||||
ToFormat: toFormat,
|
||||
Model: req.Model,
|
||||
RequestedModel: originalRequestedModel,
|
||||
Stream: opts.Stream,
|
||||
Headers: cloneHeader(opts.Headers),
|
||||
Body: cloneBytes(req.Payload),
|
||||
Metadata: opts.Metadata,
|
||||
}, requestID, skipPluginID)
|
||||
opts.Headers = mergeRequestInterceptorHeaders(opts.Headers, resp.Headers, resp.ClearHeaders)
|
||||
if len(resp.Body) > 0 {
|
||||
req.Payload = cloneBytes(resp.Body)
|
||||
opts.OriginalRequest = cloneBytes(resp.Body)
|
||||
}
|
||||
if resp.Terminate {
|
||||
return req, opts, directTerminationError(resp.StatusCode, resp.ResponseHeaders, resp.ResponseBody)
|
||||
}
|
||||
return req, opts, nil
|
||||
}
|
||||
|
||||
func ExecutionErrorMessage(err error) *interfaces.ErrorMessage {
|
||||
return executionErrorMessage(err)
|
||||
}
|
||||
|
||||
func executionErrorMessage(err error) *interfaces.ErrorMessage {
|
||||
var terminated *coreexecutor.RequestTerminatedError
|
||||
if errors.As(err, &terminated) && terminated != nil {
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: normalizedTerminationStatus(terminated.StatusCode()),
|
||||
Error: err,
|
||||
DirectResponse: true,
|
||||
Body: terminated.ResponseBody(),
|
||||
Headers: terminated.ResponseHeaders(),
|
||||
}
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if code := clienterror.HTTPStatusFromError(err); code > 0 {
|
||||
status = code
|
||||
}
|
||||
var addon http.Header
|
||||
if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil {
|
||||
if hdr := he.Headers(); hdr != nil {
|
||||
addon = hdr.Clone()
|
||||
}
|
||||
}
|
||||
return &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon}
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) pluginExecutorHost() PluginExecutorHost {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
if executorHost, ok := h.ModelRouterHost.(PluginExecutorHost); ok && executorHost != nil {
|
||||
return executorHost
|
||||
}
|
||||
if executorHost, ok := h.PluginHost.(PluginExecutorHost); ok && executorHost != nil {
|
||||
return executorHost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
518
backend/sdk/api/handlers/handlers_interceptors.go
Normal file
518
backend/sdk/api/handlers/handlers_interceptors.go
Normal file
|
|
@ -0,0 +1,518 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// PluginInterceptorHost applies plugin interceptors around handler execution.
|
||||
type PluginInterceptorHost interface {
|
||||
InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse
|
||||
InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse
|
||||
InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse
|
||||
InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse
|
||||
}
|
||||
|
||||
type pluginInterceptorSkipHost interface {
|
||||
InterceptRequestBeforeAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse
|
||||
InterceptRequestAfterAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse
|
||||
InterceptResponseExcept(context.Context, pluginapi.ResponseInterceptRequest, string) pluginapi.ResponseInterceptResponse
|
||||
InterceptStreamChunkExcept(context.Context, pluginapi.StreamChunkInterceptRequest, string) pluginapi.StreamChunkInterceptResponse
|
||||
}
|
||||
|
||||
type streamInterceptorDetector interface {
|
||||
HasStreamInterceptors() bool
|
||||
}
|
||||
|
||||
// streamChunkRequestBodyPolicy reports whether payload stream-chunk interceptors
|
||||
// still require OriginalRequest/RequestBody (legacy schema_version < 3).
|
||||
type streamChunkRequestBodyPolicy interface {
|
||||
StreamChunkPayloadIncludesRequestBody() bool
|
||||
}
|
||||
|
||||
// streamChunkPayloadIncludesRequestBody returns true when at least one active
|
||||
// stream interceptor needs per-chunk request bodies. Evaluated per call so
|
||||
// mid-stream plugin reloads stay correct. Unknown hosts default to true.
|
||||
func streamChunkPayloadIncludesRequestBody(host PluginInterceptorHost) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
if policy, ok := host.(streamChunkRequestBodyPolicy); ok {
|
||||
return policy.StreamChunkPayloadIncludesRequestBody()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
type requestInterceptorDetector interface {
|
||||
HasRequestInterceptors() bool
|
||||
}
|
||||
|
||||
type requestLifecycleHost interface {
|
||||
CompleteRequest(context.Context, pluginapi.RequestCompletion)
|
||||
}
|
||||
|
||||
type requestLifecycleSkipHost interface {
|
||||
CompleteRequestExcept(context.Context, pluginapi.RequestCompletion, string)
|
||||
}
|
||||
|
||||
type requestLifecycleTracker struct {
|
||||
once sync.Once
|
||||
ctx context.Context
|
||||
host PluginInterceptorHost
|
||||
skipPluginID string
|
||||
completion pluginapi.RequestCompletion
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceFormat, model, requestedModel string, stream bool, metadata map[string]any, skipPluginID string) *requestLifecycleTracker {
|
||||
requestID := uuid.NewString()
|
||||
traceID := logging.GetRequestID(ctx)
|
||||
return &requestLifecycleTracker{
|
||||
ctx: ctx,
|
||||
host: h.interceptorHost(),
|
||||
skipPluginID: skipPluginID,
|
||||
completion: pluginapi.RequestCompletion{
|
||||
RequestID: requestID,
|
||||
TraceID: traceID,
|
||||
SourceFormat: sourceFormat,
|
||||
Model: model,
|
||||
RequestedModel: requestedModel,
|
||||
Stream: stream,
|
||||
StartedAt: time.Now(),
|
||||
Metadata: metadata,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *requestLifecycleTracker) requestID() string {
|
||||
if t == nil {
|
||||
return ""
|
||||
}
|
||||
return t.completion.RequestID
|
||||
}
|
||||
|
||||
func (t *requestLifecycleTracker) complete(outcome pluginapi.RequestCompletionOutcome, statusCode int, err error) {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
t.once.Do(func() {
|
||||
completion := t.completion
|
||||
completion.Outcome = outcome
|
||||
completion.StatusCode = statusCode
|
||||
completion.CompletedAt = time.Now()
|
||||
if err != nil {
|
||||
completion.Error = err.Error()
|
||||
}
|
||||
if t.skipPluginID != "" {
|
||||
if host, ok := t.host.(requestLifecycleSkipHost); ok {
|
||||
host.CompleteRequestExcept(t.ctx, completion, t.skipPluginID)
|
||||
return
|
||||
}
|
||||
}
|
||||
if host, ok := t.host.(requestLifecycleHost); ok {
|
||||
host.CompleteRequest(t.ctx, completion)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (t *requestLifecycleTracker) completeError(ctx context.Context, msg *interfaces.ErrorMessage) {
|
||||
outcome := pluginapi.RequestCompletionFailed
|
||||
if msg != nil && msg.DirectResponse {
|
||||
outcome = pluginapi.RequestCompletionRejected
|
||||
} else if ctx != nil && ctx.Err() != nil {
|
||||
outcome = pluginapi.RequestCompletionCanceled
|
||||
}
|
||||
statusCode := 0
|
||||
var err error
|
||||
if msg != nil {
|
||||
statusCode = msg.StatusCode
|
||||
err = msg.Error
|
||||
}
|
||||
if outcome == pluginapi.RequestCompletionCanceled {
|
||||
statusCode = 0
|
||||
}
|
||||
t.complete(outcome, statusCode, err)
|
||||
}
|
||||
|
||||
func normalizedTerminationStatus(statusCode int) int {
|
||||
if statusCode < http.StatusOK || statusCode > 599 {
|
||||
return http.StatusForbidden
|
||||
}
|
||||
return statusCode
|
||||
}
|
||||
|
||||
func requestTerminationError(resp pluginapi.RequestInterceptResponse) *interfaces.ErrorMessage {
|
||||
return directTerminationError(resp.StatusCode, resp.ResponseHeaders, resp.ResponseBody)
|
||||
}
|
||||
|
||||
func directTerminationError(statusCode int, headers http.Header, body []byte) *interfaces.ErrorMessage {
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: normalizedTerminationStatus(statusCode),
|
||||
DirectResponse: true,
|
||||
Body: cloneBytes(body),
|
||||
Headers: cloneHeader(headers),
|
||||
}
|
||||
}
|
||||
|
||||
func cloneHeader(src http.Header) http.Header {
|
||||
if src == nil {
|
||||
return nil
|
||||
}
|
||||
dst := make(http.Header, len(src))
|
||||
for key, values := range src {
|
||||
dst[key] = append([]string(nil), values...)
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func cloneByteSlices(src [][]byte) [][]byte {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
dst := make([][]byte, 0, len(src))
|
||||
for _, item := range src {
|
||||
dst = append(dst, cloneBytes(item))
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func nextStreamChunk(ctx context.Context, pending *[]coreexecutor.StreamChunk, closed *bool, chunks <-chan coreexecutor.StreamChunk) (coreexecutor.StreamChunk, bool, bool) {
|
||||
if pending != nil && len(*pending) > 0 {
|
||||
chunk := (*pending)[0]
|
||||
(*pending)[0] = coreexecutor.StreamChunk{}
|
||||
*pending = (*pending)[1:]
|
||||
return chunk, true, false
|
||||
}
|
||||
if closed != nil && *closed {
|
||||
return coreexecutor.StreamChunk{}, false, false
|
||||
}
|
||||
var chunk coreexecutor.StreamChunk
|
||||
var ok bool
|
||||
if ctx != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return coreexecutor.StreamChunk{}, false, true
|
||||
case chunk, ok = <-chunks:
|
||||
}
|
||||
} else {
|
||||
chunk, ok = <-chunks
|
||||
}
|
||||
if !ok && closed != nil {
|
||||
*closed = true
|
||||
}
|
||||
return chunk, ok, false
|
||||
}
|
||||
|
||||
func appendStreamInterceptorHistory(history [][]byte, chunk []byte) [][]byte {
|
||||
if len(chunk) == 0 {
|
||||
return history
|
||||
}
|
||||
history = append(history, cloneBytes(chunk))
|
||||
for len(history) > maxStreamInterceptorHistoryChunks || byteSlicesSize(history) > maxStreamInterceptorHistoryBytes {
|
||||
history[0] = nil
|
||||
history = history[1:]
|
||||
}
|
||||
if len(history) == 0 {
|
||||
return nil
|
||||
}
|
||||
return history
|
||||
}
|
||||
|
||||
func byteSlicesSize(items [][]byte) int {
|
||||
total := 0
|
||||
for _, item := range items {
|
||||
total += len(item)
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func finalInterceptorHeaders(current, intercepted http.Header) http.Header {
|
||||
if intercepted == nil {
|
||||
return current
|
||||
}
|
||||
if len(intercepted) == 0 {
|
||||
return nil
|
||||
}
|
||||
return cloneHeader(intercepted)
|
||||
}
|
||||
|
||||
func downstreamHeadersFromExecutor(headers http.Header, passthrough bool) http.Header {
|
||||
if !passthrough {
|
||||
return nil
|
||||
}
|
||||
return FilterUpstreamHeaders(headers)
|
||||
}
|
||||
|
||||
func downstreamHeadersAfterInterceptors(baseRaw, finalRaw http.Header, passthrough bool) http.Header {
|
||||
if passthrough {
|
||||
return FilterUpstreamHeaders(finalRaw)
|
||||
}
|
||||
return FilterUpstreamHeaders(diffHeaders(baseRaw, finalRaw))
|
||||
}
|
||||
|
||||
func diffHeaders(base, next http.Header) http.Header {
|
||||
if len(next) == 0 {
|
||||
return nil
|
||||
}
|
||||
baseValues := make(map[string][]string, len(base))
|
||||
for key, values := range base {
|
||||
baseValues[http.CanonicalHeaderKey(key)] = values
|
||||
}
|
||||
out := make(http.Header)
|
||||
for key, values := range next {
|
||||
canonicalKey := http.CanonicalHeaderKey(key)
|
||||
if stringSlicesEqual(baseValues[canonicalKey], values) {
|
||||
continue
|
||||
}
|
||||
out[canonicalKey] = append([]string(nil), values...)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stringSlicesEqual(left, right []string) bool {
|
||||
if len(left) != len(right) {
|
||||
return false
|
||||
}
|
||||
for i := range left {
|
||||
if left[i] != right[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) interceptorHost() PluginInterceptorHost {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
return h.PluginHost
|
||||
}
|
||||
|
||||
func streamInterceptorsEnabled(host PluginInterceptorHost) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
if detector, ok := host.(streamInterceptorDetector); ok {
|
||||
return detector.HasStreamInterceptors()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func requestInterceptorsEnabled(host PluginInterceptorHost) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
if detector, ok := host.(requestInterceptorDetector); ok {
|
||||
return detector.HasRequestInterceptors()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
type requestAfterAuthCapture struct {
|
||||
mu sync.Mutex
|
||||
set bool
|
||||
headers http.Header
|
||||
body []byte
|
||||
originalRequest []byte
|
||||
originalRequestReplaced bool
|
||||
}
|
||||
|
||||
func (c *requestAfterAuthCapture) record(req coreexecutor.RequestAfterAuthInterceptRequest, resp coreexecutor.RequestAfterAuthInterceptResponse) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
headers := mergeRequestInterceptorHeaders(req.Headers, resp.Headers, resp.ClearHeaders)
|
||||
body := cloneBytes(req.Body)
|
||||
var originalRequest []byte
|
||||
originalRequestReplaced := false
|
||||
if len(resp.Body) > 0 {
|
||||
body = cloneBytes(resp.Body)
|
||||
originalRequest = cloneBytes(resp.Body)
|
||||
originalRequestReplaced = true
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.set = true
|
||||
c.headers = headers
|
||||
c.body = body
|
||||
c.originalRequest = originalRequest
|
||||
c.originalRequestReplaced = originalRequestReplaced
|
||||
}
|
||||
|
||||
func (c *requestAfterAuthCapture) apply(req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) {
|
||||
if c == nil {
|
||||
return req, opts
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if !c.set {
|
||||
return req, opts
|
||||
}
|
||||
req.Payload = cloneBytes(c.body)
|
||||
opts.Headers = cloneHeader(c.headers)
|
||||
if c.originalRequestReplaced {
|
||||
opts.OriginalRequest = cloneBytes(c.originalRequest)
|
||||
}
|
||||
return req, opts
|
||||
}
|
||||
|
||||
func mergeRequestInterceptorHeaders(current, updates http.Header, clear []string) http.Header {
|
||||
if updates == nil && len(clear) == 0 {
|
||||
return cloneHeader(current)
|
||||
}
|
||||
out := cloneHeader(current)
|
||||
if out == nil && (len(updates) > 0 || len(clear) > 0) {
|
||||
out = make(http.Header)
|
||||
}
|
||||
for _, key := range clear {
|
||||
out.Del(key)
|
||||
}
|
||||
for key, values := range updates {
|
||||
out.Del(key)
|
||||
for _, value := range values {
|
||||
out.Add(key, value)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func interceptRequestBeforeAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse {
|
||||
if skipPluginID != "" {
|
||||
if skipper, ok := host.(pluginInterceptorSkipHost); ok {
|
||||
return skipper.InterceptRequestBeforeAuthExcept(ctx, req, skipPluginID)
|
||||
}
|
||||
}
|
||||
return host.InterceptRequestBeforeAuth(ctx, req)
|
||||
}
|
||||
|
||||
func interceptRequestAfterAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse {
|
||||
if skipPluginID != "" {
|
||||
if skipper, ok := host.(pluginInterceptorSkipHost); ok {
|
||||
return skipper.InterceptRequestAfterAuthExcept(ctx, req, skipPluginID)
|
||||
}
|
||||
}
|
||||
return host.InterceptRequestAfterAuth(ctx, req)
|
||||
}
|
||||
|
||||
func interceptResponse(ctx context.Context, host PluginInterceptorHost, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse {
|
||||
if skipPluginID != "" {
|
||||
if skipper, ok := host.(pluginInterceptorSkipHost); ok {
|
||||
return skipper.InterceptResponseExcept(ctx, req, skipPluginID)
|
||||
}
|
||||
}
|
||||
return host.InterceptResponse(ctx, req)
|
||||
}
|
||||
|
||||
func interceptStreamChunk(ctx context.Context, host PluginInterceptorHost, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse {
|
||||
if skipPluginID != "" {
|
||||
if skipper, ok := host.(pluginInterceptorSkipHost); ok {
|
||||
return skipper.InterceptStreamChunkExcept(ctx, req, skipPluginID)
|
||||
}
|
||||
}
|
||||
return host.InterceptStreamChunk(ctx, req)
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel, requestID string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options, *interfaces.ErrorMessage) {
|
||||
host := h.interceptorHost()
|
||||
if !requestInterceptorsEnabled(host) {
|
||||
return req, opts, nil
|
||||
}
|
||||
resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{
|
||||
RequestID: requestID,
|
||||
TraceID: logging.GetRequestID(ctx),
|
||||
SourceFormat: handlerType,
|
||||
Model: req.Model,
|
||||
RequestedModel: requestedModel,
|
||||
Stream: opts.Stream,
|
||||
Headers: cloneHeader(opts.Headers),
|
||||
Body: cloneBytes(req.Payload),
|
||||
Metadata: opts.Metadata,
|
||||
}, skipPluginID)
|
||||
opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers)
|
||||
if len(resp.Body) > 0 {
|
||||
req.Payload = cloneBytes(resp.Body)
|
||||
opts.OriginalRequest = cloneBytes(resp.Body)
|
||||
}
|
||||
if resp.Terminate {
|
||||
return req, opts, requestTerminationError(resp)
|
||||
}
|
||||
return req, opts, nil
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor {
|
||||
if !requestInterceptorsEnabled(h.interceptorHost()) {
|
||||
return nil
|
||||
}
|
||||
return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse {
|
||||
resp := h.applyRequestInterceptorsAfterAuth(ctx, req, requestID, skipPluginID)
|
||||
if capture != nil {
|
||||
capture.record(req, resp)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse {
|
||||
host := h.interceptorHost()
|
||||
if !requestInterceptorsEnabled(host) {
|
||||
return coreexecutor.RequestAfterAuthInterceptResponse{}
|
||||
}
|
||||
resp := interceptRequestAfterAuth(ctx, host, pluginapi.RequestInterceptRequest{
|
||||
RequestID: requestID,
|
||||
TraceID: logging.GetRequestID(ctx),
|
||||
SourceFormat: req.SourceFormat.String(),
|
||||
ToFormat: req.ToFormat.String(),
|
||||
Model: req.Model,
|
||||
RequestedModel: req.RequestedModel,
|
||||
Stream: req.Stream,
|
||||
Headers: cloneHeader(req.Headers),
|
||||
Body: cloneBytes(req.Body),
|
||||
Metadata: req.Metadata,
|
||||
}, skipPluginID)
|
||||
return coreexecutor.RequestAfterAuthInterceptResponse{
|
||||
Headers: resp.Headers,
|
||||
Body: resp.Body,
|
||||
ClearHeaders: resp.ClearHeaders,
|
||||
Terminate: resp.Terminate,
|
||||
StatusCode: normalizedTerminationStatus(resp.StatusCode),
|
||||
ResponseHeaders: resp.ResponseHeaders,
|
||||
ResponseBody: resp.ResponseBody,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, requestID, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int, skipPluginID string) ([]byte, http.Header) {
|
||||
host := h.interceptorHost()
|
||||
if host == nil {
|
||||
return body, responseHeaders
|
||||
}
|
||||
resp := interceptResponse(ctx, host, pluginapi.ResponseInterceptRequest{
|
||||
RequestID: requestID,
|
||||
SourceFormat: handlerType,
|
||||
Model: normalizedModel,
|
||||
RequestedModel: requestedModel,
|
||||
Stream: false,
|
||||
RequestHeaders: cloneHeader(opts.Headers),
|
||||
ResponseHeaders: cloneHeader(rawResponseHeaders),
|
||||
OriginalRequest: cloneBytes(originalRequest),
|
||||
RequestBody: cloneBytes(requestBody),
|
||||
Body: cloneBytes(body),
|
||||
StatusCode: statusCode,
|
||||
Metadata: opts.Metadata,
|
||||
}, skipPluginID)
|
||||
responseHeaders = downstreamHeadersAfterInterceptors(rawResponseHeaders, finalInterceptorHeaders(rawResponseHeaders, resp.Headers), PassthroughHeadersEnabled(h.Cfg))
|
||||
if len(resp.Body) > 0 {
|
||||
body = cloneBytes(resp.Body)
|
||||
}
|
||||
return body, responseHeaders
|
||||
}
|
||||
1483
backend/sdk/api/handlers/handlers_interceptors_test.go
Normal file
1483
backend/sdk/api/handlers/handlers_interceptors_test.go
Normal file
File diff suppressed because it is too large
Load diff
184
backend/sdk/api/handlers/handlers_metadata_test.go
Normal file
184
backend/sdk/api/handlers/handlers_metadata_test.go
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coresession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
func TestGetContextWithCancelCapturesClientRequestMetadata(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ginCtx.Request.RemoteAddr = "192.0.2.10:43123"
|
||||
ginCtx.Request.Header.Add("X-Forwarded-For", "203.0.113.5")
|
||||
ginCtx.Request.Header.Add("X-Forwarded-For", "198.51.100.8")
|
||||
ginCtx.Request.Header.Set("User-Agent", "test-client/1.0")
|
||||
|
||||
handler := &BaseAPIHandler{Cfg: &config.SDKConfig{}}
|
||||
ctx, cancel := handler.GetContextWithCancel(nil, ginCtx, context.Background())
|
||||
defer cancel()
|
||||
|
||||
metadata := logging.GetClientRequestMetadata(ctx)
|
||||
if metadata.ClientIP != "192.0.2.10" {
|
||||
t.Fatalf("ClientIP = %q, want direct peer IP", metadata.ClientIP)
|
||||
}
|
||||
if metadata.XForwardedFor != "203.0.113.5, 198.51.100.8" {
|
||||
t.Fatalf("XForwardedFor = %q", metadata.XForwardedFor)
|
||||
}
|
||||
if metadata.UserAgent != "test-client/1.0" {
|
||||
t.Fatalf("UserAgent = %q", metadata.UserAgent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestExecutionMetadataIncludesExecutionSessionWithoutIdempotencyKey(t *testing.T) {
|
||||
ctx := WithExecutionSessionID(context.Background(), "session-1")
|
||||
|
||||
meta := requestExecutionMetadata(ctx)
|
||||
if got := meta[coreexecutor.ExecutionSessionMetadataKey]; got != "session-1" {
|
||||
t.Fatalf("ExecutionSessionMetadataKey = %v, want %q", got, "session-1")
|
||||
}
|
||||
if _, ok := meta[idempotencyKeyMetadataKey]; ok {
|
||||
t.Fatalf("unexpected idempotency key in metadata: %v", meta[idempotencyKeyMetadataKey])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestExecutionMetadataIncludesHashedCallerScope(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ginCtx.Set("userApiKey", "downstream-secret")
|
||||
ctx := context.WithValue(context.Background(), "gin", ginCtx)
|
||||
|
||||
meta := requestExecutionMetadata(ctx)
|
||||
got, _ := meta[coreexecutor.CallerScopeMetadataKey].(string)
|
||||
want := coresession.CallerScope("downstream-secret")
|
||||
if got != want {
|
||||
t.Fatalf("CallerScopeMetadataKey = %q, want %q", got, want)
|
||||
}
|
||||
if got == "downstream-secret" {
|
||||
t.Fatal("caller scope contains the raw downstream credential")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestExecutionMetadataTraceCallbackWebsocketDetection(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("skips websocket upgrade", func(t *testing.T) {
|
||||
ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
ginCtx.Request.Header.Set("Connection", "Upgrade")
|
||||
ginCtx.Request.Header.Set("Upgrade", "websocket")
|
||||
logging.SetGinRequestID(ginCtx, "1234abcd")
|
||||
ctx := context.WithValue(context.Background(), "gin", ginCtx)
|
||||
|
||||
meta := requestExecutionMetadata(ctx)
|
||||
|
||||
if _, exists := meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey]; exists {
|
||||
t.Fatal("unexpected selected auth index callback for websocket upgrade")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("keeps callback for incomplete upgrade headers", func(t *testing.T) {
|
||||
ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
ginCtx.Request.Header.Set("Upgrade", "websocket")
|
||||
logging.SetGinRequestID(ginCtx, "1234abcd")
|
||||
ctx := context.WithValue(context.Background(), "gin", ginCtx)
|
||||
|
||||
meta := requestExecutionMetadata(ctx)
|
||||
|
||||
if _, exists := meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey]; !exists {
|
||||
t.Fatal("missing selected auth index callback for ordinary HTTP request")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSetReasoningEffortMetadataUsesSuffixOverBody(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setReasoningEffortMetadata(meta, "openai", "gpt-5.4(high)", []byte(`{"reasoning_effort":"low"}`))
|
||||
|
||||
if got := meta[coreexecutor.ReasoningEffortMetadataKey]; got != "high" {
|
||||
t.Fatalf("ReasoningEffortMetadataKey = %v, want %q", got, "high")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetReasoningEffortMetadataSupportsOpenAIResponses(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setReasoningEffortMetadata(meta, "openai-response", "gpt-5.4", []byte(`{"reasoning":{"effort":"medium"}}`))
|
||||
|
||||
if got := meta[coreexecutor.ReasoningEffortMetadataKey]; got != "medium" {
|
||||
t.Fatalf("ReasoningEffortMetadataKey = %v, want %q", got, "medium")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetServiceTierMetadataExtractsValue(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setServiceTierMetadata(meta, []byte(`{"service_tier":"priority"}`))
|
||||
|
||||
gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey]
|
||||
if gotServiceTier != "priority" {
|
||||
t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "priority")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetServiceTierMetadataDefaultsWhenMissing(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setServiceTierMetadata(meta, []byte(`{"model":"gpt-5.4"}`))
|
||||
|
||||
gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey]
|
||||
if gotServiceTier != "auto" {
|
||||
t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "auto")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetServiceTierMetadataPreservesExplicitDefault(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setServiceTierMetadata(meta, []byte(`{"service_tier":"default"}`))
|
||||
|
||||
if gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey]; gotServiceTier != "default" {
|
||||
t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetGenerateMetadataDefaultsWhenMissing(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setGenerateMetadata(meta, []byte(`{"model":"gpt-5.4"}`))
|
||||
|
||||
if got := meta[coreexecutor.GenerateMetadataKey]; got != true {
|
||||
t.Fatalf("GenerateMetadataKey = %v, want true", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetGenerateMetadataPreservesTrue(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setGenerateMetadata(meta, []byte(`{"generate":true}`))
|
||||
|
||||
if got := meta[coreexecutor.GenerateMetadataKey]; got != true {
|
||||
t.Fatalf("GenerateMetadataKey = %v, want true", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetGenerateMetadataHonorsExplicitFalse(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setGenerateMetadata(meta, []byte(`{"generate":false}`))
|
||||
|
||||
if got := meta[coreexecutor.GenerateMetadataKey]; got != false {
|
||||
t.Fatalf("GenerateMetadataKey = %v, want false", got)
|
||||
}
|
||||
}
|
||||
832
backend/sdk/api/handlers/handlers_model_router_test.go
Normal file
832
backend/sdk/api/handlers/handlers_model_router_test.go
Normal file
|
|
@ -0,0 +1,832 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
)
|
||||
|
||||
type handlerModelRouterTestHost struct {
|
||||
hasRouters bool
|
||||
route func(context.Context, pluginapi.ModelRouteRequest, string) (pluginapi.ModelRouteResponse, bool)
|
||||
routeSkip string
|
||||
lastReq *pluginapi.ModelRouteRequest
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) RouteModel(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return h.RouteModelExcept(ctx, req, "")
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) RouteModelExcept(ctx context.Context, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) {
|
||||
h.routeSkip = skipPluginID
|
||||
reqCopy := req
|
||||
h.lastReq = &reqCopy
|
||||
if h != nil && h.route != nil {
|
||||
return h.route(ctx, req, skipPluginID)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) HasModelRouters() bool { return h != nil && h.hasRouters }
|
||||
|
||||
func (h *handlerModelRouterTestHost) HasModelRoutersExcept(skipPluginID string) bool {
|
||||
return h != nil && h.hasRouters
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) HasRequestInterceptors() bool { return false }
|
||||
|
||||
func (h *handlerModelRouterTestHost) HasStreamInterceptors() bool { return false }
|
||||
|
||||
func (h *handlerModelRouterTestHost) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
return pluginapi.RequestInterceptResponse{Headers: cloneHeader(req.Headers), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
return pluginapi.RequestInterceptResponse{Headers: cloneHeader(req.Headers), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse {
|
||||
return pluginapi.ResponseInterceptResponse{Headers: cloneHeader(req.ResponseHeaders), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerModelRouterTestHost) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse {
|
||||
return pluginapi.StreamChunkInterceptResponse{Headers: cloneHeader(req.ResponseHeaders), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
type handlerRouterOnlyTestHost struct {
|
||||
route func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool)
|
||||
hasRouters bool
|
||||
called bool
|
||||
}
|
||||
|
||||
func (h *handlerRouterOnlyTestHost) RouteModel(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if h != nil {
|
||||
h.called = true
|
||||
}
|
||||
if h != nil && h.route != nil {
|
||||
return h.route(ctx, req)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
|
||||
func (h *handlerRouterOnlyTestHost) HasModelRouters() bool {
|
||||
return h != nil && h.hasRouters
|
||||
}
|
||||
|
||||
type handlerDirectExecutorRouteHost struct {
|
||||
handlerRouterOnlyTestHost
|
||||
lastPluginID string
|
||||
lastRequest coreexecutor.Request
|
||||
lastOptions coreexecutor.Options
|
||||
stream func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error)
|
||||
}
|
||||
|
||||
type handlerSkipAwareDirectExecutorRouteHost struct {
|
||||
handlerDirectExecutorRouteHost
|
||||
routeSkip string
|
||||
}
|
||||
|
||||
func (h *handlerSkipAwareDirectExecutorRouteHost) RouteModelExcept(ctx context.Context, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) {
|
||||
h.routeSkip = skipPluginID
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
|
||||
func (h *handlerSkipAwareDirectExecutorRouteHost) HasModelRoutersExcept(string) bool {
|
||||
return h != nil && h.hasRouters
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorRouteHost) ExecutePluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
h.lastPluginID = pluginID
|
||||
h.lastRequest = req
|
||||
h.lastOptions = opts
|
||||
return coreexecutor.Response{Payload: []byte("direct-ok")}, nil
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorRouteHost) ExecutePluginExecutorStream(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
h.lastPluginID = pluginID
|
||||
h.lastRequest = req
|
||||
h.lastOptions = opts
|
||||
if h.stream != nil {
|
||||
return h.stream(ctx, pluginID, req, opts)
|
||||
}
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("direct-stream")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorRouteHost) CountPluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
h.lastPluginID = pluginID
|
||||
h.lastRequest = req
|
||||
h.lastOptions = opts
|
||||
return coreexecutor.Response{Payload: []byte("7")}, nil
|
||||
}
|
||||
|
||||
type handlerDirectExecutorInterceptorHost struct {
|
||||
handlerDirectExecutorRouteHost
|
||||
afterAuthCalled bool
|
||||
afterAuthReq pluginapi.RequestInterceptRequest
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) HasRequestInterceptors() bool { return true }
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) HasStreamInterceptors() bool { return false }
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
return pluginapi.RequestInterceptResponse{Headers: cloneHeader(req.Headers), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
h.afterAuthCalled = true
|
||||
h.afterAuthReq = req
|
||||
headers := cloneHeader(req.Headers)
|
||||
if headers == nil {
|
||||
headers = make(http.Header)
|
||||
}
|
||||
headers.Set("X-After-Auth", "yes")
|
||||
return pluginapi.RequestInterceptResponse{Headers: headers, Body: []byte(`{"after":true}`)}
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse {
|
||||
return pluginapi.ResponseInterceptResponse{Headers: cloneHeader(req.ResponseHeaders), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse {
|
||||
return pluginapi.StreamChunkInterceptResponse{Headers: cloneHeader(req.ResponseHeaders), Body: cloneBytes(req.Body)}
|
||||
}
|
||||
|
||||
func (h *handlerDirectExecutorInterceptorHost) PluginExecutorRequestToFormat(pluginID string, req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format {
|
||||
return sdktranslator.FormatCodex
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterRoutesBeforeRequestDetails(t *testing.T) {
|
||||
originalModel := "handler-router-original-model"
|
||||
targetPluginID := "websearch-plugin"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.SourceFormat != "openai" || req.RequestedModel != originalModel || req.Stream {
|
||||
t.Fatalf("unexpected route request = %#v", req)
|
||||
}
|
||||
if req.Headers.Get("X-Original") != "client" {
|
||||
t.Fatalf("route headers = %#v, want client header", req.Headers)
|
||||
}
|
||||
if string(req.Body) != fmt.Sprintf(`{"model":%q}`, originalModel) {
|
||||
t.Fatalf("route body = %q, want original body", req.Body)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID, Reason: "test"}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
ctx := contextWithHeaders(http.Header{"X-Original": []string{"client"}})
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(body) != "direct-ok" {
|
||||
t.Fatalf("body = %q, want direct plugin executor response", body)
|
||||
}
|
||||
if host.lastPluginID != targetPluginID {
|
||||
t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
|
||||
}
|
||||
if host.lastRequest.Model != originalModel {
|
||||
t.Fatalf("executor model = %q, want original model", host.lastRequest.Model)
|
||||
}
|
||||
if host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey] != originalModel {
|
||||
t.Fatalf("requested model metadata = %#v, want original model", host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey])
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterDirectExecutorRunsAfterAuthInterceptor(t *testing.T) {
|
||||
originalModel := "handler-router-after-auth-original-model"
|
||||
targetPluginID := "websearch-plugin"
|
||||
host := &handlerDirectExecutorInterceptorHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetPluginHost(host)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(body) != "direct-ok" {
|
||||
t.Fatalf("body = %q, want direct plugin executor response", body)
|
||||
}
|
||||
if !host.afterAuthCalled {
|
||||
t.Fatal("after-auth interceptor was not called")
|
||||
}
|
||||
if host.afterAuthReq.SourceFormat != "openai" || host.afterAuthReq.ToFormat != "codex" {
|
||||
t.Fatalf("after-auth formats = %q -> %q, want openai -> codex", host.afterAuthReq.SourceFormat, host.afterAuthReq.ToFormat)
|
||||
}
|
||||
if host.afterAuthReq.Model != originalModel || host.afterAuthReq.RequestedModel != originalModel {
|
||||
t.Fatalf("after-auth models = %q/%q, want original model", host.afterAuthReq.Model, host.afterAuthReq.RequestedModel)
|
||||
}
|
||||
if string(host.lastRequest.Payload) != `{"after":true}` {
|
||||
t.Fatalf("executor payload = %q, want after-auth body", host.lastRequest.Payload)
|
||||
}
|
||||
if host.lastOptions.Headers.Get("X-After-Auth") != "yes" {
|
||||
t.Fatalf("executor headers = %#v, want after-auth header", host.lastOptions.Headers)
|
||||
}
|
||||
if string(host.lastOptions.OriginalRequest) != `{"after":true}` {
|
||||
t.Fatalf("original request = %q, want after-auth body", host.lastOptions.OriginalRequest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterPluginExecutorFailsClosedWhenHomeEnabled(t *testing.T) {
|
||||
originalModel := "home-plugin-route"
|
||||
targetPluginID := "plugin-executor"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route"}`), "")
|
||||
if body != nil || errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("ExecuteWithAuthManager() = %q, %#v; want 503", body, errMsg)
|
||||
}
|
||||
body, _, errMsg = handler.ExecuteCountWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route"}`), "")
|
||||
if body != nil || errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("ExecuteCountWithAuthManager() = %q, %#v; want 503", body, errMsg)
|
||||
}
|
||||
data, _, errors := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route","stream":true}`), "")
|
||||
if data != nil {
|
||||
t.Fatalf("ExecuteStreamWithAuthManager() data = %v, want nil", data)
|
||||
}
|
||||
if errMsg = <-errors; errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("ExecuteStreamWithAuthManager() error = %#v, want 503", errMsg)
|
||||
}
|
||||
if host.lastPluginID != "" {
|
||||
t.Fatalf("plugin executor was invoked with %q while Home was enabled", host.lastPluginID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterRequiresPluginExecutorHost(t *testing.T) {
|
||||
originalModel := "handler-router-only-original-model"
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(&handlerRouterOnlyTestHost{
|
||||
hasRouters: true,
|
||||
route: func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.RequestedModel != originalModel {
|
||||
t.Fatalf("requested model = %q, want %q", req.RequestedModel, originalModel)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: "websearch-plugin"}, true
|
||||
},
|
||||
})
|
||||
|
||||
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg == nil || errMsg.StatusCode != http.StatusBadGateway {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v, want BadGateway", errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterCanTargetPluginExecutorWithoutChangingModel(t *testing.T) {
|
||||
originalModel := "handler-router-direct-original-model"
|
||||
targetPluginID := "websearch-plugin"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.RequestedModel != originalModel {
|
||||
t.Fatalf("requested model = %q, want %q", req.RequestedModel, originalModel)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(body) != "direct-ok" {
|
||||
t.Fatalf("body = %q, want direct plugin executor response", body)
|
||||
}
|
||||
if host.lastPluginID != targetPluginID {
|
||||
t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
|
||||
}
|
||||
if host.lastRequest.Model != originalModel {
|
||||
t.Fatalf("executor model = %q, want original model", host.lastRequest.Model)
|
||||
}
|
||||
if host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey] != originalModel {
|
||||
t.Fatalf("requested model metadata = %#v, want original model", host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey])
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterRoutesCountBeforeRequestDetails(t *testing.T) {
|
||||
originalModel := "handler-router-count-original-model"
|
||||
targetPluginID := "count-plugin"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.SourceFormat != "claude" || req.RequestedModel != originalModel || req.Stream {
|
||||
t.Fatalf("unexpected count route request = %#v", req)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
body, _, errMsg := handler.ExecuteCountWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteCountWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(body) != "7" {
|
||||
t.Fatalf("body = %q, want count response", body)
|
||||
}
|
||||
if host.lastPluginID != targetPluginID {
|
||||
t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
|
||||
}
|
||||
if host.lastRequest.Model != originalModel {
|
||||
t.Fatalf("executor model = %q, want original model", host.lastRequest.Model)
|
||||
}
|
||||
if host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey] != originalModel {
|
||||
t.Fatalf("requested model metadata = %#v, want original model", host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouteModelDoesNotFallbackWhenSkipUnsupported(t *testing.T) {
|
||||
host := &handlerRouterOnlyTestHost{hasRouters: true}
|
||||
resp, ok := routeModel(context.Background(), host, pluginapi.ModelRouteRequest{RequestedModel: "model"}, "origin-plugin")
|
||||
if ok || resp.Handled {
|
||||
t.Fatalf("routeModel() = %#v, %v; want unhandled when skip is unsupported", resp, ok)
|
||||
}
|
||||
if host.called {
|
||||
t.Fatal("RouteModel was called despite unsupported skip")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyModelRouterSkipsHostsWithoutRouters(t *testing.T) {
|
||||
host := &handlerRouterOnlyTestHost{hasRouters: false}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
got := handler.applyModelRouter(context.Background(), "openai", "model", []byte(`{"model":"model"}`), false, modelExecutionOptions{})
|
||||
if got.ExecutorPluginID != "" {
|
||||
t.Fatalf("applyModelRouter() = %#v, want no routing decision", got)
|
||||
}
|
||||
if host.called {
|
||||
t.Fatal("RouteModel was called even though detector reported no routers")
|
||||
}
|
||||
}
|
||||
|
||||
// routeModelOnlyHost implements PluginModelRouterHost without HasModelRouters (conservative default).
|
||||
type routeModelOnlyHost struct {
|
||||
called bool
|
||||
}
|
||||
|
||||
func (h *routeModelOnlyHost) RouteModel(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if h != nil {
|
||||
h.called = true
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
|
||||
func TestModelRoutersEnabledFalseWithoutDetector(t *testing.T) {
|
||||
host := &routeModelOnlyHost{}
|
||||
if modelRoutersEnabled(host, "") {
|
||||
t.Fatal("modelRoutersEnabled() = true, want false when host has no HasModelRouters")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyModelRouterSkipsHostWithoutDetector(t *testing.T) {
|
||||
host := &routeModelOnlyHost{}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
got := handler.applyModelRouter(context.Background(), "openai", "model", []byte(`{"model":"model"}`), false, modelExecutionOptions{})
|
||||
if got.ExecutorPluginID != "" || got.Provider != "" {
|
||||
t.Fatalf("applyModelRouter() = %#v, want no routing decision", got)
|
||||
}
|
||||
if host.called {
|
||||
t.Fatal("RouteModel was called on host without HasModelRouters")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyModelRouterRestoresQueryFromContext(t *testing.T) {
|
||||
var gotQuery url.Values
|
||||
host := &handlerRouterOnlyTestHost{hasRouters: true}
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
gotQuery = cloneURLValues(req.Query)
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
// execOptions.Query is intentionally empty; the inbound query must be recovered
|
||||
// from the embedded gin context, mirroring plain HTTP requests.
|
||||
ctx := contextWithQuery(url.Values{"session": []string{"abc"}})
|
||||
handler.applyModelRouter(ctx, "openai", "model", []byte(`{"model":"model"}`), false, modelExecutionOptions{})
|
||||
|
||||
if gotQuery.Get("session") != "abc" {
|
||||
t.Fatalf("route query = %#v, want session=abc recovered from gin context", gotQuery)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterRoutesStreamBeforeRequestDetails(t *testing.T) {
|
||||
originalModel := "handler-router-stream-original-model"
|
||||
targetPluginID := "stream-plugin"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.SourceFormat != "openai" || req.RequestedModel != originalModel || !req.Stream {
|
||||
t.Fatalf("unexpected stream route request = %#v", req)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
var gotPayload bool
|
||||
for range dataChan {
|
||||
gotPayload = true
|
||||
}
|
||||
if !gotPayload {
|
||||
t.Fatal("stream produced no payload")
|
||||
}
|
||||
if errMsg := <-errChan; errMsg != nil {
|
||||
t.Fatalf("ExecuteStreamWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if host.lastPluginID != targetPluginID {
|
||||
t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
|
||||
}
|
||||
if host.lastRequest.Model != originalModel {
|
||||
t.Fatalf("executor model = %q, want original model", host.lastRequest.Model)
|
||||
}
|
||||
if host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey] != originalModel {
|
||||
t.Fatalf("requested model metadata = %#v, want original model", host.lastOptions.Metadata[coreexecutor.RequestedModelMetadataKey])
|
||||
}
|
||||
}
|
||||
|
||||
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 TestExecuteModelStreamDoesNotReusePreparedRouteWhenRouterPluginSkipped(t *testing.T) {
|
||||
const originalModel = "prepared-router-model"
|
||||
const mappedModel = "mapped-upstream-model"
|
||||
const originPluginID = "origin-plugin"
|
||||
host := &handlerSkipAwareDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: originPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
body := []byte(`{"model":"prepared-router-model","stream":true}`)
|
||||
ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai-response", originalModel, body)
|
||||
if !routedToPlugin {
|
||||
t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route")
|
||||
}
|
||||
|
||||
_, errMsg := handler.ExecuteModelStream(ctx, ModelExecutionRequest{
|
||||
EntryProtocol: "openai-response",
|
||||
ExitProtocol: "openai-response",
|
||||
Model: mappedModel,
|
||||
Stream: true,
|
||||
Body: []byte(`{"model":"mapped-upstream-model","stream":true}`),
|
||||
SkipRouterPluginID: originPluginID,
|
||||
})
|
||||
if host.routeSkip != originPluginID {
|
||||
t.Fatalf("router skip id = %q, want %q", host.routeSkip, originPluginID)
|
||||
}
|
||||
if host.lastPluginID == originPluginID {
|
||||
t.Fatalf("plugin executor %q was re-entered despite SkipRouterPluginID", host.lastPluginID)
|
||||
}
|
||||
if errMsg == nil {
|
||||
t.Fatal("ExecuteModelStream() error = nil, want normal provider resolution failure with empty auth manager")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {
|
||||
model := "model-execution-router-skip-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
routerHost := &handlerModelRouterTestHost{hasRouters: true}
|
||||
handler.SetPluginHost(routerHost)
|
||||
|
||||
resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: model,
|
||||
Body: requestBody,
|
||||
SkipRouterPluginID: "origin-plugin",
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModel() error = %+v", errMsg)
|
||||
}
|
||||
if string(resp.Body) != "model-execution-ok" {
|
||||
t.Fatalf("body = %q, want executor response", resp.Body)
|
||||
}
|
||||
if routerHost.routeSkip != "origin-plugin" {
|
||||
t.Fatalf("router skip id = %q, want origin-plugin", routerHost.routeSkip)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerProvidersForExecutionUsesRouterProvider(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
decision := modelRouteDecision{Provider: "claude", Model: "claude-sonnet-4"}
|
||||
providers, normalizedModel, errMsg := handler.providersForExecution("ignored-by-router", "original-model", false, decision, modelExecutionOptions{})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("providersForExecution() error = %+v", errMsg)
|
||||
}
|
||||
if fmt.Sprint(providers) != "[claude]" {
|
||||
t.Fatalf("providers = %v, want [claude]", providers)
|
||||
}
|
||||
if normalizedModel != "claude-sonnet-4" {
|
||||
t.Fatalf("normalizedModel = %q, want claude-sonnet-4", normalizedModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerProvidersForExecutionFallsBackToOriginalModel(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
decision := modelRouteDecision{Provider: "claude"}
|
||||
providers, normalizedModel, errMsg := handler.providersForExecution("ignored-by-router", "original-model", false, decision, modelExecutionOptions{})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("providersForExecution() error = %+v", errMsg)
|
||||
}
|
||||
if fmt.Sprint(providers) != "[claude]" {
|
||||
t.Fatalf("providers = %v, want [claude]", providers)
|
||||
}
|
||||
if normalizedModel != "original-model" {
|
||||
t.Fatalf("normalizedModel = %q, want original-model", normalizedModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterProviderRouteUsesAuthManager(t *testing.T) {
|
||||
originalModel := "provider-route-original-model"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetProvider, Target: "claude"}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
handler.AuthManager = coreauth.NewManager(nil, nil, nil)
|
||||
|
||||
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
// The empty AuthManager has no claude auth, so execution surfaces an auth selection error
|
||||
// rather than succeeding. The point is that the request reached the AuthManager path.
|
||||
if errMsg == nil {
|
||||
t.Fatal("ExecuteWithAuthManager() error = nil, want auth selection error for routed provider")
|
||||
}
|
||||
if !host.called {
|
||||
t.Fatal("model router was not consulted")
|
||||
}
|
||||
if host.lastPluginID != "" {
|
||||
t.Fatalf("plugin executor path was used (plugin id = %q); want provider path via AuthManager", host.lastPluginID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerProvidersForExecutionRejectsImageOnlyModelOnProviderRoute(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
cases := []struct {
|
||||
name string
|
||||
originalModel string
|
||||
decision modelRouteDecision
|
||||
}{
|
||||
{
|
||||
name: "target-model",
|
||||
originalModel: "original-model",
|
||||
decision: modelRouteDecision{Provider: "claude", Model: "gpt-image-2"},
|
||||
},
|
||||
{
|
||||
name: "target-model-thinking-suffix",
|
||||
originalModel: "original-model",
|
||||
decision: modelRouteDecision{Provider: "claude", Model: "gpt-image-2(auto)"},
|
||||
},
|
||||
{
|
||||
name: "original-model-thinking-suffix",
|
||||
originalModel: "gpt-image-2(auto)",
|
||||
decision: modelRouteDecision{Provider: "claude"},
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, _, errMsg := handler.providersForExecution("ignored", tc.originalModel, false, tc.decision, modelExecutionOptions{})
|
||||
if errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("providersForExecution() error = %+v, want image-only service unavailable", errMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteCountWithAuthManagerPropagatesRouterSkipAndQuery(t *testing.T) {
|
||||
model := "model-execution-count-router-context-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
routerHost := &handlerModelRouterTestHost{hasRouters: true}
|
||||
handler.SetPluginHost(routerHost)
|
||||
ctx := contextWithQuery(url.Values{"session": []string{"abc"}})
|
||||
|
||||
_, _, errMsg := handler.executeCountWithAuthManager(ctx, "openai", model, requestBody, "", modelExecutionOptions{
|
||||
SkipRouterPluginID: "origin-plugin",
|
||||
})
|
||||
if errMsg == nil {
|
||||
t.Fatal("executeCountWithAuthManager() error = nil, want auth selection error on empty manager")
|
||||
}
|
||||
if routerHost.routeSkip != "origin-plugin" {
|
||||
t.Fatalf("router skip id = %q, want origin-plugin", routerHost.routeSkip)
|
||||
}
|
||||
if routerHost.lastReq == nil || routerHost.lastReq.Query.Get("session") != "abc" {
|
||||
t.Fatalf("route query = %#v, want session=abc", routerHost.lastReq)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerModelRouterDirectExecutorPropagatesQueryFromContext(t *testing.T) {
|
||||
originalModel := "handler-router-query-model"
|
||||
targetPluginID := "query-plugin"
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
ctx := contextWithQuery(url.Values{"session": []string{"abc"}})
|
||||
|
||||
_, _, errMsg := handler.ExecuteWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if host.lastOptions.Query == nil || host.lastOptions.Query.Get("session") != "abc" {
|
||||
t.Fatalf("executor query = %#v, want session=abc from gin context", host.lastOptions.Query)
|
||||
}
|
||||
}
|
||||
|
||||
type handlerStuckPluginStreamHost struct {
|
||||
handlerDirectExecutorRouteHost
|
||||
}
|
||||
|
||||
func (h *handlerStuckPluginStreamHost) ExecutePluginExecutorStream(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func TestStreamWithPluginExecutorExitsOnContextCancel(t *testing.T) {
|
||||
originalModel := "handler-router-stream-cancel-model"
|
||||
targetPluginID := "stuck-stream-plugin"
|
||||
host := &handlerStuckPluginStreamHost{}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
deadline := time.After(2 * time.Second)
|
||||
for {
|
||||
select {
|
||||
case _, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg := <-errChan; errMsg != nil {
|
||||
t.Fatalf("unexpected stream error: %+v", errMsg)
|
||||
}
|
||||
return
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatal("plugin executor stream goroutine did not exit after context cancel")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamWithPluginExecutorReturnedHeadersImmutableAfterReturn(t *testing.T) {
|
||||
originalModel := "handler-router-plugin-immutable-headers-model"
|
||||
targetPluginID := "immutable-headers-plugin"
|
||||
releaseSecond := make(chan struct{})
|
||||
bodyStarted := make(chan struct{})
|
||||
releaseBody := make(chan struct{})
|
||||
host := &handlerDirectExecutorRouteHost{}
|
||||
host.stream = func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk)
|
||||
go func() {
|
||||
defer close(chunks)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("first")}
|
||||
<-releaseSecond
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("second")}
|
||||
}()
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
host.hasRouters = true
|
||||
host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true}, nil)
|
||||
handler.SetModelRouterHost(host)
|
||||
handler.SetPluginHost(&handlerInterceptorTestHost{
|
||||
interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse {
|
||||
headers := cloneHeader(req.ResponseHeaders)
|
||||
if headers == nil {
|
||||
headers = make(http.Header)
|
||||
}
|
||||
switch req.ChunkIndex {
|
||||
case pluginapi.StreamChunkHeaderInitIndex:
|
||||
headers.Set("X-Init", "plugin")
|
||||
case 1:
|
||||
close(bodyStarted)
|
||||
<-releaseBody
|
||||
headers.Set("X-Body", "plugin")
|
||||
}
|
||||
return pluginapi.StreamChunkInterceptResponse{Headers: headers, Body: cloneBytes(req.Body)}
|
||||
},
|
||||
})
|
||||
|
||||
dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
dataDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(dataDone)
|
||||
for range dataChan {
|
||||
}
|
||||
}()
|
||||
stopReading := make(chan struct{})
|
||||
readerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(readerDone)
|
||||
for {
|
||||
select {
|
||||
case <-stopReading:
|
||||
return
|
||||
default:
|
||||
_ = upstreamHeaders.Get("X-Init")
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
close(releaseSecond)
|
||||
<-bodyStarted
|
||||
close(releaseBody)
|
||||
<-dataDone
|
||||
for msg := range errChan {
|
||||
if msg != nil {
|
||||
t.Fatalf("unexpected stream error: %+v", msg)
|
||||
}
|
||||
}
|
||||
close(stopReading)
|
||||
<-readerDone
|
||||
if upstreamHeaders.Get("X-Init") != "plugin" || upstreamHeaders.Get("X-Body") != "" {
|
||||
t.Fatalf("returned headers mutated after return: %#v", upstreamHeaders)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryFromContextNilURLDoesNotPanic(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = &http.Request{Header: make(http.Header)}
|
||||
ctx := context.WithValue(context.Background(), "gin", c)
|
||||
if got := queryFromContext(ctx); got != nil {
|
||||
t.Fatalf("queryFromContext() = %#v, want nil when URL is nil", got)
|
||||
}
|
||||
}
|
||||
197
backend/sdk/api/handlers/handlers_plugin_executor_usage.go
Normal file
197
backend/sdk/api/handlers/handlers_plugin_executor_usage.go
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func parsePluginExecutorResponseUsage(protocol string, payload []byte) usage.Detail {
|
||||
if len(payload) == 0 {
|
||||
return usage.Detail{}
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(protocol)) {
|
||||
case "claude":
|
||||
return parseClaudePayloadUsage(payload)
|
||||
case "gemini":
|
||||
return helps.ParseGeminiUsage(payload)
|
||||
case "interactions", "interactions-response":
|
||||
return helps.ParseInteractionsUsage(payload)
|
||||
case "antigravity":
|
||||
return helps.ParseAntigravityUsage(payload)
|
||||
case "codex", "openai-response":
|
||||
if detail, ok := helps.ParseCodexUsage(payload); ok {
|
||||
return detail
|
||||
}
|
||||
return helps.ParseOpenAIUsage(payload)
|
||||
default:
|
||||
return helps.ParseOpenAIUsage(payload)
|
||||
}
|
||||
}
|
||||
|
||||
func observePluginExecutorStreamUsage(protocol string, payload []byte, buffer *helps.StreamUsageBuffer) {
|
||||
if buffer == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(protocol)) {
|
||||
case "claude":
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := parseClaudeStreamLine(line); ok {
|
||||
observeMergedStreamUsage(buffer, detail)
|
||||
}
|
||||
})
|
||||
case "gemini":
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := helps.ParseGeminiStreamUsage(line); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
}
|
||||
})
|
||||
case "interactions", "interactions-response":
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := helps.ParseInteractionsStreamUsage(line); ok {
|
||||
observeMergedStreamUsage(buffer, detail)
|
||||
}
|
||||
})
|
||||
case "antigravity":
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
if detail, ok := helps.ParseAntigravityStreamUsage(line); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
}
|
||||
})
|
||||
case "codex", "openai-response":
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
if jsonBytes := extractStreamJSONPayload(line); len(jsonBytes) > 0 {
|
||||
if detail, ok := helps.ParseCodexUsage(jsonBytes); ok {
|
||||
buffer.Observe(detail, ok)
|
||||
return
|
||||
}
|
||||
}
|
||||
buffer.ObserveOpenAIStream(line)
|
||||
})
|
||||
default:
|
||||
iterateStreamLines(payload, func(line []byte) {
|
||||
buffer.ObserveOpenAIStream(line)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func parseClaudePayloadUsage(payload []byte) usage.Detail {
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
usageNode = gjson.GetBytes(payload, "message.usage")
|
||||
}
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
return helps.ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`))
|
||||
}
|
||||
|
||||
func parseClaudeStreamLine(line []byte) (usage.Detail, bool) {
|
||||
payload := extractStreamJSONPayload(line)
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
usageNode := gjson.GetBytes(payload, "usage")
|
||||
if !usageNode.Exists() {
|
||||
usageNode = gjson.GetBytes(payload, "message.usage")
|
||||
}
|
||||
if !usageNode.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := helps.ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`))
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func observeMergedStreamUsage(buffer *helps.StreamUsageBuffer, update usage.Detail) {
|
||||
if buffer == nil {
|
||||
return
|
||||
}
|
||||
if existing, ok := buffer.Detail(); ok {
|
||||
merged := mergeStreamUsageDetail(existing, update)
|
||||
buffer.Observe(merged, true)
|
||||
return
|
||||
}
|
||||
buffer.Observe(update, true)
|
||||
}
|
||||
|
||||
func mergeStreamUsageDetail(existing, update usage.Detail) usage.Detail {
|
||||
merged := update
|
||||
if merged.InputTokens == 0 && existing.InputTokens > 0 {
|
||||
merged.InputTokens = existing.InputTokens
|
||||
}
|
||||
if merged.CachedTokens == 0 && existing.CachedTokens > 0 {
|
||||
merged.CachedTokens = existing.CachedTokens
|
||||
}
|
||||
if merged.CacheReadTokens == 0 && existing.CacheReadTokens > 0 {
|
||||
merged.CacheReadTokens = existing.CacheReadTokens
|
||||
}
|
||||
if merged.CacheCreationTokens == 0 && existing.CacheCreationTokens > 0 {
|
||||
merged.CacheCreationTokens = existing.CacheCreationTokens
|
||||
}
|
||||
if merged.OutputTokens == 0 && existing.OutputTokens > 0 {
|
||||
merged.OutputTokens = existing.OutputTokens
|
||||
}
|
||||
if merged.ReasoningTokens == 0 && existing.ReasoningTokens > 0 {
|
||||
merged.ReasoningTokens = existing.ReasoningTokens
|
||||
}
|
||||
if merged.ResponseServiceTier == "" {
|
||||
merged.ResponseServiceTier = existing.ResponseServiceTier
|
||||
}
|
||||
cached := merged.CacheReadTokens + merged.CacheCreationTokens
|
||||
if cached == 0 {
|
||||
cached = merged.CachedTokens
|
||||
}
|
||||
calculatedTotal := merged.InputTokens + merged.OutputTokens + cached
|
||||
if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal {
|
||||
merged.TotalTokens = calculatedTotal
|
||||
}
|
||||
nonReasoningOutput := merged.OutputTokens - merged.ReasoningTokens
|
||||
if nonReasoningOutput < 0 {
|
||||
nonReasoningOutput = 0
|
||||
}
|
||||
merged.TokenBreakdown = usage.NewIndependentTokenBreakdown(
|
||||
merged.InputTokens,
|
||||
merged.CacheReadTokens,
|
||||
merged.CacheCreationTokens,
|
||||
nonReasoningOutput,
|
||||
merged.ReasoningTokens,
|
||||
merged.TotalTokens,
|
||||
)
|
||||
return merged
|
||||
}
|
||||
|
||||
func iterateStreamLines(payload []byte, fn func(line []byte)) {
|
||||
for _, line := range bytes.Split(payload, []byte("\n")) {
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if len(trimmed) == 0 {
|
||||
continue
|
||||
}
|
||||
fn(trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
func extractStreamJSONPayload(line []byte) []byte {
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if len(trimmed) == 0 {
|
||||
return nil
|
||||
}
|
||||
if bytes.Equal(trimmed, []byte("[DONE]")) {
|
||||
return nil
|
||||
}
|
||||
if bytes.HasPrefix(trimmed, []byte("event:")) {
|
||||
return nil
|
||||
}
|
||||
if bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
trimmed = bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:")))
|
||||
}
|
||||
if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("[DONE]")) {
|
||||
return nil
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
718
backend/sdk/api/handlers/handlers_plugin_executor_usage_test.go
Normal file
718
backend/sdk/api/handlers/handlers_plugin_executor_usage_test.go
Normal file
|
|
@ -0,0 +1,718 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
type noopUsagePlugin struct{}
|
||||
|
||||
func (noopUsagePlugin) HandleUsage(context.Context, usage.Record) {}
|
||||
|
||||
type capturePluginExecutorUsagePlugin struct {
|
||||
targetProvider string
|
||||
records chan usage.Record
|
||||
}
|
||||
|
||||
func newCapturePluginExecutorUsagePlugin(targetProvider string) *capturePluginExecutorUsagePlugin {
|
||||
return &capturePluginExecutorUsagePlugin{
|
||||
targetProvider: targetProvider,
|
||||
records: make(chan usage.Record, 50),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *capturePluginExecutorUsagePlugin) HandleUsage(_ context.Context, record usage.Record) {
|
||||
if p.targetProvider != "" && record.Provider != p.targetProvider {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case p.records <- record:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (p *capturePluginExecutorUsagePlugin) waitRecord(t *testing.T) usage.Record {
|
||||
t.Helper()
|
||||
select {
|
||||
case rec := <-p.records:
|
||||
return rec
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for usage record")
|
||||
return usage.Record{}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *capturePluginExecutorUsagePlugin) assertNoRecord(t *testing.T) {
|
||||
t.Helper()
|
||||
select {
|
||||
case rec := <-p.records:
|
||||
t.Fatalf("expected no usage record for %q, got %+v", p.targetProvider, rec)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
|
||||
func registerUsagePluginForTest(t *testing.T, name string, plugin usage.Plugin) {
|
||||
t.Helper()
|
||||
usage.RegisterNamedPlugin(name, plugin)
|
||||
t.Cleanup(func() {
|
||||
usage.RegisterNamedPlugin(name, noopUsagePlugin{})
|
||||
})
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageNonStreamOpenAI(t *testing.T) {
|
||||
targetPluginID := "custom-openai-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-nonstream-openai", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
openAIResponseBody := []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":12,"completion_tokens":34,"total_tokens":46}}`)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
execResp: coreexecutor.Response{
|
||||
Payload: openAIResponseBody,
|
||||
},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if len(body) == 0 {
|
||||
t.Fatal("empty response body")
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 12 || record.Detail.OutputTokens != 34 || record.Detail.TotalTokens != 46 {
|
||||
t.Errorf("record.Detail = %+v, want prompt=12 completion=34 total=46", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageStreamOpenAI(t *testing.T) {
|
||||
targetPluginID := "custom-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-openai", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
chunks := make(chan coreexecutor.StreamChunk, 3)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":25,\"total_tokens\":40}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error = %+v", err)
|
||||
}
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 15 || record.Detail.OutputTokens != 25 || record.Detail.TotalTokens != 40 {
|
||||
t.Errorf("record.Detail = %+v, want prompt=15 completion=25 total=40", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageStreamCodex(t *testing.T) {
|
||||
targetPluginID := "custom-codex-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-codex", plugin)
|
||||
|
||||
originalModel := "codex-5.2"
|
||||
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":18,\"output_tokens\":22,\"total_tokens\":40}}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai-response", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error = %+v", err)
|
||||
}
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 18 || record.Detail.OutputTokens != 22 || record.Detail.TotalTokens != 40 {
|
||||
t.Errorf("record.Detail = %+v, want input=18 output=22 total=40", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageNonStreamClaude(t *testing.T) {
|
||||
targetPluginID := "custom-claude-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-nonstream-claude", plugin)
|
||||
|
||||
originalModel := "claude-3-5-sonnet"
|
||||
|
||||
claudeResponseBody := []byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hello"}],"usage":{"input_tokens":50,"output_tokens":30,"output_tokens_details":{"thinking_tokens":10}}}`)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
execResp: coreexecutor.Response{
|
||||
Payload: claudeResponseBody,
|
||||
},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hi"}]}`, originalModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if len(body) == 0 {
|
||||
t.Fatal("empty response body")
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 50 || record.Detail.OutputTokens != 30 || record.Detail.ReasoningTokens != 10 || record.Detail.TotalTokens != 80 {
|
||||
t.Errorf("record.Detail = %+v, want input=50 output=30 reasoning=10 total=80", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageStreamClaude(t *testing.T) {
|
||||
targetPluginID := "custom-claude-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-claude", plugin)
|
||||
|
||||
originalModel := "claude-3-5-sonnet"
|
||||
|
||||
// Claude streams split usage between message_start (input, cache) and message_delta (output, thinking)
|
||||
chunks := make(chan coreexecutor.StreamChunk, 4)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"usage\":{\"input_tokens\":100,\"cache_read_input_tokens\":50,\"cache_creation_input_tokens\":20,\"output_tokens\":1}}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":25,\"output_tokens_details\":{\"thinking_tokens\":5}}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error = %+v", err)
|
||||
}
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 100 || record.Detail.CacheReadTokens != 50 || record.Detail.CacheCreationTokens != 20 || record.Detail.OutputTokens != 25 || record.Detail.ReasoningTokens != 5 || record.Detail.TotalTokens != 195 {
|
||||
t.Errorf("record.Detail = %+v, want input=100 cache_read=50 cache_creation=20 output=25 reasoning=5 total=195", record.Detail)
|
||||
}
|
||||
tb := record.Detail.TokenBreakdown
|
||||
if !tb.Valid() || tb.TotalTokens != 195 || tb.Input.TotalTokens != 170 || tb.Input.UncachedTokens != 100 || tb.Input.CacheReadTokens != 50 || tb.Input.CacheWriteTokens != 20 || tb.Output.TotalTokens != 25 || tb.Output.NonReasoningTokens != 20 || tb.Output.ReasoningTokens != 5 {
|
||||
t.Errorf("record.Detail.TokenBreakdown = %+v, want valid independent breakdown with total=195 input=170 output=25", tb)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageGemini(t *testing.T) {
|
||||
targetPluginID := "custom-gemini-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-gemini", plugin)
|
||||
|
||||
originalModel := "gemini-2.5-flash"
|
||||
|
||||
geminiResponseBody := []byte(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":40,"candidatesTokenCount":60,"totalTokenCount":100}}`)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
execResp: coreexecutor.Response{
|
||||
Payload: geminiResponseBody,
|
||||
},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "gemini", originalModel, []byte(fmt.Sprintf(`{"contents":[{"parts":[{"text":"hi"}]}]}`)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if len(body) == 0 {
|
||||
t.Fatal("empty response body")
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 40 || record.Detail.OutputTokens != 60 || record.Detail.TotalTokens != 100 {
|
||||
t.Errorf("record.Detail = %+v, want prompt=40 candidates=60 total=100", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageStreamInteractions(t *testing.T) {
|
||||
targetPluginID := "custom-interactions-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-interactions", plugin)
|
||||
|
||||
originalModel := "gemini-2.5-flash"
|
||||
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"event_type\":\"finish\",\"metadata\":{\"total_usage\":{\"total_input_tokens\":30,\"total_output_tokens\":70,\"total_tokens\":100}}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "interactions", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error = %+v", err)
|
||||
}
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 30 || record.Detail.OutputTokens != 70 || record.Detail.TotalTokens != 100 {
|
||||
t.Errorf("record.Detail = %+v, want input=30 output=70 total=100", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesUsageStreamAntigravity(t *testing.T) {
|
||||
targetPluginID := "custom-antigravity-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-antigravity", plugin)
|
||||
|
||||
originalModel := "claude-3-5-sonnet"
|
||||
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"response\":{\"usageMetadata\":{\"promptTokenCount\":33,\"candidatesTokenCount\":67,\"totalTokenCount\":100}}}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "antigravity", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
t.Fatalf("stream error = %+v", err)
|
||||
}
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if record.Detail.InputTokens != 33 || record.Detail.OutputTokens != 67 || record.Detail.TotalTokens != 100 {
|
||||
t.Errorf("record.Detail = %+v, want prompt=33 candidates=67 total=100", record.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesFailure(t *testing.T) {
|
||||
targetPluginID := "failing-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-failure", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
execErr: errors.New("upstream plugin failure"),
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "")
|
||||
if errMsg == nil {
|
||||
t.Fatal("expected ExecuteWithAuthManager() to fail")
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if !record.Failed {
|
||||
t.Error("record.Failed = false, want true")
|
||||
}
|
||||
if record.Fail.Body != "upstream plugin failure" {
|
||||
t.Errorf("record.Fail.Body = %q, want %q", record.Fail.Body, "upstream plugin failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesStreamFailure(t *testing.T) {
|
||||
targetPluginID := "failing-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-failure", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream stream broke")}
|
||||
close(chunks)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
for range dataChan {
|
||||
}
|
||||
seenErr := false
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
seenErr = true
|
||||
}
|
||||
}
|
||||
if !seenErr {
|
||||
t.Fatal("expected stream error")
|
||||
}
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if !record.Failed {
|
||||
t.Error("record.Failed = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorPublishesStreamCancellation(t *testing.T) {
|
||||
targetPluginID := "canceling-stream-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-cancel", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
chunks := make(chan coreexecutor.StreamChunk, 5)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")}
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
streamResult: &coreexecutor.StreamResult{Chunks: chunks},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
dataChan, _, _ := handler.ExecuteStreamWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
|
||||
// Receive first chunk then cancel context
|
||||
<-dataChan
|
||||
cancel()
|
||||
close(chunks)
|
||||
|
||||
record := plugin.waitRecord(t)
|
||||
if record.Provider != targetPluginID {
|
||||
t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID)
|
||||
}
|
||||
if !record.Failed {
|
||||
t.Error("record.Failed = false, want true for canceled stream")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorSkipsUsageForNestedExecution(t *testing.T) {
|
||||
targetPluginID := "nested-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-usage-nested", plugin)
|
||||
|
||||
originalModel := "gpt-4o"
|
||||
|
||||
openAIResponseBody := []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":12,"completion_tokens":34,"total_tokens":46}}`)
|
||||
|
||||
mockHost := &mockPluginUsageHost{
|
||||
execResp: coreexecutor.Response{
|
||||
Payload: openAIResponseBody,
|
||||
},
|
||||
}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
// ExecuteModel triggers execution with InternalSource = true (host.model.execute callback)
|
||||
resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: originalModel,
|
||||
Body: []byte(fmt.Sprintf(`{"model":%q}`, originalModel)),
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModel() error = %+v", errMsg)
|
||||
}
|
||||
if len(resp.Body) == 0 {
|
||||
t.Fatal("empty response body")
|
||||
}
|
||||
|
||||
plugin.assertNoRecord(t)
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorSkipsOuterUsageWhenPluginCallsHostModelExecute(t *testing.T) {
|
||||
targetPluginID := "agent-wrapper-plugin"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-nested-callback", plugin)
|
||||
|
||||
outerModel := "agent-wrapper-model"
|
||||
innerModel := "inner-model"
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
innerExecutor := &modelExecutionCaptureExecutor{
|
||||
provider: "openai",
|
||||
execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{
|
||||
Payload: []byte(`{"id":"chatcmpl-inner","choices":[{"message":{"role":"assistant","content":"inner"}}],"usage":{"prompt_tokens":5,"completion_tokens":5,"total_tokens":10}}`),
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
manager.RegisterExecutor(innerExecutor)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "auth-" + innerModel,
|
||||
Provider: innerExecutor.Identifier(),
|
||||
Status: coreauth.StatusActive,
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: innerModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
|
||||
mockHost := &mockPluginUsageHost{}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.RequestedModel == outerModel {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
// When the plugin executor executes, it simulates calling back into the host via ExecuteModel
|
||||
mockHost.execFunc = func(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
// Plugin calls back into host.model.execute using the provided ctx
|
||||
innerResp, errInner := handler.ExecuteModel(ctx, ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: innerModel,
|
||||
Body: []byte(fmt.Sprintf(`{"model":%q}`, innerModel)),
|
||||
})
|
||||
if errInner != nil {
|
||||
return coreexecutor.Response{}, errInner.Error
|
||||
}
|
||||
// Return response payload back to outer caller
|
||||
return coreexecutor.Response{Payload: innerResp.Body}, nil
|
||||
}
|
||||
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", outerModel, []byte(fmt.Sprintf(`{"model":%q}`, outerModel)), "")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if len(body) == 0 {
|
||||
t.Fatal("empty response body")
|
||||
}
|
||||
|
||||
// Since inner ExecuteModel was executed, the outer plugin executor must NOT publish a duplicate record for targetPluginID
|
||||
plugin.assertNoRecord(t)
|
||||
}
|
||||
|
||||
func TestHandlerPluginExecutorSkipsOuterFailureWhenPluginCallsHostModelExecute(t *testing.T) {
|
||||
targetPluginID := "agent-wrapper-plugin-failure"
|
||||
plugin := newCapturePluginExecutorUsagePlugin(targetPluginID)
|
||||
registerUsagePluginForTest(t, "test-plugin-executor-nested-callback-failure", plugin)
|
||||
|
||||
outerModel := "agent-wrapper-model-fail"
|
||||
innerModel := "inner-model-fail"
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
innerExecutor := &modelExecutionCaptureExecutor{
|
||||
provider: "openai",
|
||||
execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("inner failure")
|
||||
},
|
||||
}
|
||||
manager.RegisterExecutor(innerExecutor)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "auth-" + innerModel,
|
||||
Provider: innerExecutor.Identifier(),
|
||||
Status: coreauth.StatusActive,
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: innerModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
|
||||
mockHost := &mockPluginUsageHost{}
|
||||
mockHost.hasRouters = true
|
||||
mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
|
||||
if req.RequestedModel == outerModel {
|
||||
return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
mockHost.execFunc = func(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
_, errInner := handler.ExecuteModel(ctx, ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: innerModel,
|
||||
Body: []byte(fmt.Sprintf(`{"model":%q}`, innerModel)),
|
||||
})
|
||||
if errInner != nil {
|
||||
return coreexecutor.Response{}, errInner.Error
|
||||
}
|
||||
return coreexecutor.Response{Payload: []byte("ok")}, nil
|
||||
}
|
||||
|
||||
handler.SetModelRouterHost(mockHost)
|
||||
|
||||
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", outerModel, []byte(fmt.Sprintf(`{"model":%q}`, outerModel)), "")
|
||||
if errMsg == nil {
|
||||
t.Fatal("expected failure")
|
||||
}
|
||||
|
||||
// Since inner ExecuteModel was executed, the outer plugin executor must NOT publish a duplicate failure record
|
||||
plugin.assertNoRecord(t)
|
||||
}
|
||||
|
||||
type mockPluginUsageHost struct {
|
||||
handlerDirectExecutorRouteHost
|
||||
execResp coreexecutor.Response
|
||||
execErr error
|
||||
execFunc func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error)
|
||||
streamResult *coreexecutor.StreamResult
|
||||
}
|
||||
|
||||
func (h *mockPluginUsageHost) ExecutePluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
h.lastPluginID = pluginID
|
||||
h.lastRequest = req
|
||||
h.lastOptions = opts
|
||||
if h.execFunc != nil {
|
||||
return h.execFunc(ctx, pluginID, req, opts)
|
||||
}
|
||||
if h.execErr != nil {
|
||||
return coreexecutor.Response{}, h.execErr
|
||||
}
|
||||
return h.execResp, nil
|
||||
}
|
||||
|
||||
func (h *mockPluginUsageHost) ExecutePluginExecutorStream(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
h.lastPluginID = pluginID
|
||||
h.lastRequest = req
|
||||
h.lastOptions = opts
|
||||
return h.streamResult, nil
|
||||
}
|
||||
288
backend/sdk/api/handlers/handlers_request_details_test.go
Normal file
288
backend/sdk/api/handlers/handlers_request_details_test.go
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestGetRequestDetails_PreservesSuffix(t *testing.T) {
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
now := time.Now().Unix()
|
||||
|
||||
modelRegistry.RegisterClient("test-request-details-gemini", "gemini", []*registry.ModelInfo{
|
||||
{ID: "gemini-2.5-pro", Created: now + 30},
|
||||
{ID: "gemini-2.5-flash", Created: now + 25},
|
||||
})
|
||||
modelRegistry.RegisterClient("test-request-details-openai", "openai", []*registry.ModelInfo{
|
||||
{ID: "gpt-5.2", Created: now + 20},
|
||||
})
|
||||
modelRegistry.RegisterClient("test-request-details-claude", "claude", []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4-5", Created: now + 5},
|
||||
})
|
||||
|
||||
// Ensure cleanup of all test registrations.
|
||||
clientIDs := []string{
|
||||
"test-request-details-gemini",
|
||||
"test-request-details-openai",
|
||||
"test-request-details-claude",
|
||||
}
|
||||
for _, clientID := range clientIDs {
|
||||
id := clientID
|
||||
t.Cleanup(func() {
|
||||
modelRegistry.UnregisterClient(id)
|
||||
})
|
||||
}
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
inputModel string
|
||||
wantProviders []string
|
||||
wantModel string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "numeric suffix preserved",
|
||||
inputModel: "gemini-2.5-pro(8192)",
|
||||
wantProviders: []string{"gemini"},
|
||||
wantModel: "gemini-2.5-pro(8192)",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "level suffix preserved",
|
||||
inputModel: "gpt-5.2(high)",
|
||||
wantProviders: []string{"openai"},
|
||||
wantModel: "gpt-5.2(high)",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "no suffix unchanged",
|
||||
inputModel: "claude-sonnet-4-5",
|
||||
wantProviders: []string{"claude"},
|
||||
wantModel: "claude-sonnet-4-5",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "unknown model with suffix",
|
||||
inputModel: "unknown-model(8192)",
|
||||
wantProviders: nil,
|
||||
wantModel: "",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "auto suffix resolved",
|
||||
inputModel: "auto(high)",
|
||||
wantProviders: []string{"gemini"},
|
||||
wantModel: "gemini-2.5-pro(high)",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "special suffix none preserved",
|
||||
inputModel: "gemini-2.5-flash(none)",
|
||||
wantProviders: []string{"gemini"},
|
||||
wantModel: "gemini-2.5-flash(none)",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "special suffix auto preserved",
|
||||
inputModel: "claude-sonnet-4-5(auto)",
|
||||
wantProviders: []string{"claude"},
|
||||
wantModel: "claude-sonnet-4-5(auto)",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
providers, model, errMsg := handler.getRequestDetails(tt.inputModel)
|
||||
if (errMsg != nil) != tt.wantErr {
|
||||
t.Fatalf("getRequestDetails() error = %v, wantErr %v", errMsg, tt.wantErr)
|
||||
}
|
||||
if errMsg != nil {
|
||||
return
|
||||
}
|
||||
if !reflect.DeepEqual(providers, tt.wantProviders) {
|
||||
t.Fatalf("getRequestDetails() providers = %v, want %v", providers, tt.wantProviders)
|
||||
}
|
||||
if model != tt.wantModel {
|
||||
t.Fatalf("getRequestDetails() model = %v, want %v", model, tt.wantModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetRequestDetails_UnknownModelErrorResistsJSONInjection pins the unroutable
|
||||
// model error body against client-controlled model names. The name is echoed into
|
||||
// the body, so formatting it into a JSON literal would let a caller corrupt the
|
||||
// payload or overwrite the error code that clients branch on.
|
||||
func TestGetRequestDetails_UnknownModelErrorResistsJSONInjection(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
for _, model := range []string{
|
||||
"unroutable-model",
|
||||
`foo"bar`,
|
||||
`x","code":"insufficient_quota","x":"`,
|
||||
`x"}}`,
|
||||
`foo\bar`,
|
||||
"foo\nbar",
|
||||
} {
|
||||
t.Run(model, func(t *testing.T) {
|
||||
_, _, errMsg := handler.getRequestDetails(model)
|
||||
if errMsg == nil || errMsg.Error == nil {
|
||||
t.Fatal("expected an error for an unroutable model")
|
||||
}
|
||||
if errMsg.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", errMsg.StatusCode, http.StatusBadRequest)
|
||||
}
|
||||
body := errMsg.Error.Error()
|
||||
if !json.Valid([]byte(body)) {
|
||||
t.Fatalf("error body is not valid JSON: %s", body)
|
||||
}
|
||||
if got := gjson.Get(body, "error.code").String(); got != "model_not_found" {
|
||||
t.Fatalf("error code = %q, want model_not_found; the caller controlled the body: %s", got, body)
|
||||
}
|
||||
if got, want := gjson.Get(body, "error.message").String(), "unknown provider for model "+model; got != want {
|
||||
t.Fatalf("error message = %q, want %q", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetRequestDetails_ImageModelReturns503(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
imageOnlyModels := []string{
|
||||
"gpt-image-1.5",
|
||||
"gpt-image-2",
|
||||
"codex/gpt-image-2",
|
||||
"grok-imagine-image",
|
||||
"xai/grok-imagine-image",
|
||||
"grok-imagine-image-quality",
|
||||
"xai/grok-imagine-image-quality",
|
||||
"grok-imagine-image-2.0",
|
||||
"xai/grok-imagine-image-2.0",
|
||||
}
|
||||
for _, model := range imageOnlyModels {
|
||||
t.Run(model, func(t *testing.T) {
|
||||
_, _, errMsg := handler.getRequestDetails(model)
|
||||
if errMsg == nil {
|
||||
t.Fatalf("expected error for %s, got nil", model)
|
||||
}
|
||||
if errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("unexpected status code: got %d want %d", errMsg.StatusCode, http.StatusServiceUnavailable)
|
||||
}
|
||||
if errMsg.Error == nil {
|
||||
t.Fatalf("expected error message, got nil")
|
||||
}
|
||||
msg := errMsg.Error.Error()
|
||||
if !strings.Contains(msg, "/v1/images/generations") || !strings.Contains(msg, "/v1/images/edits") {
|
||||
t.Fatalf("unexpected error message: %q", msg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateImageOnlyModel_AllowsImageEndpoints(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
imageOnlyModels := []string{
|
||||
"gpt-image-1.5",
|
||||
"gpt-image-2",
|
||||
"codex/gpt-image-2",
|
||||
"grok-imagine-image",
|
||||
"xai/grok-imagine-image",
|
||||
"grok-imagine-image-quality",
|
||||
"xai/grok-imagine-image-quality",
|
||||
"grok-imagine-image-2.0",
|
||||
"xai/grok-imagine-image-2.0",
|
||||
}
|
||||
for _, model := range imageOnlyModels {
|
||||
t.Run(model, func(t *testing.T) {
|
||||
if errMsg := handler.validateImageOnlyModel(model, true); errMsg != nil {
|
||||
t.Fatalf("validateImageOnlyModel(%q, true) = %+v, want nil", model, errMsg)
|
||||
}
|
||||
if errMsg := handler.validateImageOnlyModel(model, false); errMsg == nil {
|
||||
t.Fatalf("validateImageOnlyModel(%q, false) = nil, want image-only error", model)
|
||||
} else if errMsg.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("unexpected status code: got %d want %d", errMsg.StatusCode, http.StatusServiceUnavailable)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsOpenAIImageOnlyModel(t *testing.T) {
|
||||
tests := []struct {
|
||||
model string
|
||||
want bool
|
||||
}{
|
||||
{model: "gpt-image-1.5", want: true},
|
||||
{model: "gpt-image-2", want: true},
|
||||
{model: "codex/gpt-image-1.5", want: true},
|
||||
{model: "grok-imagine-image", want: true},
|
||||
{model: "xai/grok-imagine-image", want: true},
|
||||
{model: "XAI/Grok-Imagine-Image-Quality", want: true},
|
||||
{model: "grok-imagine-image-quality", want: true},
|
||||
{model: "grok-imagine-image-2.0", want: true},
|
||||
{model: "xai/grok-imagine-image-2.0", want: true},
|
||||
{model: "grok-3", want: false},
|
||||
{model: "gpt-5.2", want: false},
|
||||
{model: "grok-imagine-video", want: false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.model, func(t *testing.T) {
|
||||
if got := isOpenAIImageOnlyModel(tt.model); got != tt.want {
|
||||
t.Fatalf("isOpenAIImageOnlyModel(%q) = %v, want %v", tt.model, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteImageWithAuthManager_AllowsImageOnlyModels(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
imageOnlyModels := []string{
|
||||
"gpt-image-1.5",
|
||||
"gpt-image-2",
|
||||
"grok-imagine-image",
|
||||
"grok-imagine-image-quality",
|
||||
"xai/grok-imagine-image-quality",
|
||||
"grok-imagine-image-2.0",
|
||||
"xai/grok-imagine-image-2.0",
|
||||
}
|
||||
for _, model := range imageOnlyModels {
|
||||
t.Run(model, func(t *testing.T) {
|
||||
body := []byte(`{"model":"` + model + `","prompt":"draw"}`)
|
||||
_, _, errMsg := handler.ExecuteImageWithAuthManager(context.Background(), "openai-image", model, body, "")
|
||||
if errMsg == nil {
|
||||
t.Fatal("expected auth selection error, got nil")
|
||||
}
|
||||
if errMsg.Error == nil {
|
||||
t.Fatal("expected error message, got nil")
|
||||
}
|
||||
msg := errMsg.Error.Error()
|
||||
if strings.Contains(msg, "only supported on /v1/images/generations") {
|
||||
t.Fatalf("ExecuteImageWithAuthManager rejected image-only model: %q", msg)
|
||||
}
|
||||
|
||||
_, _, errMsg = handler.ExecuteWithAuthManager(context.Background(), "openai-image", model, body, "")
|
||||
if errMsg == nil {
|
||||
t.Fatal("expected image-only rejection for non-image execution path, got nil")
|
||||
}
|
||||
if errMsg.Error == nil || !strings.Contains(errMsg.Error.Error(), "only supported on /v1/images/generations") {
|
||||
t.Fatalf("unexpected non-image execution error: %+v", errMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
354
backend/sdk/api/handlers/handlers_routing.go
Normal file
354
backend/sdk/api/handlers/handlers_routing.go
Normal file
|
|
@ -0,0 +1,354 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/tidwall/sjson"
|
||||
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// PluginModelRouterHost routes matching requests to a plugin executor, the router's own executor,
|
||||
// or a built-in provider before model-to-provider resolution and auth selection.
|
||||
type PluginModelRouterHost interface {
|
||||
RouteModel(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool)
|
||||
}
|
||||
|
||||
type pluginModelRouterSkipHost interface {
|
||||
RouteModelExcept(context.Context, pluginapi.ModelRouteRequest, string) (pluginapi.ModelRouteResponse, bool)
|
||||
}
|
||||
|
||||
type modelRouterDetector interface {
|
||||
HasModelRouters() bool
|
||||
}
|
||||
|
||||
type modelRouterSkipDetector interface {
|
||||
HasModelRoutersExcept(string) bool
|
||||
}
|
||||
|
||||
func preferExecutionProvider(providers []string, preferred string) []string {
|
||||
preferred = strings.ToLower(strings.TrimSpace(preferred))
|
||||
if preferred == "" || len(providers) < 2 {
|
||||
return providers
|
||||
}
|
||||
preferredIndex := -1
|
||||
for i := range providers {
|
||||
if strings.ToLower(strings.TrimSpace(providers[i])) == preferred {
|
||||
preferredIndex = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if preferredIndex <= 0 {
|
||||
return providers
|
||||
}
|
||||
out := make([]string, 0, len(providers))
|
||||
out = append(out, providers[preferredIndex])
|
||||
out = append(out, providers[:preferredIndex]...)
|
||||
out = append(out, providers[preferredIndex+1:]...)
|
||||
return out
|
||||
}
|
||||
|
||||
func adjustExecutionProvidersForEntryProtocol(entryProtocol string, providers []string) []string {
|
||||
if entryProtocol == Interactions {
|
||||
return preferExecutionProvider(providers, GeminiInteractions)
|
||||
}
|
||||
if supportsNativeInteractionsEntryProtocol(entryProtocol) {
|
||||
return providers
|
||||
}
|
||||
return excludeExecutionProvider(providers, GeminiInteractions)
|
||||
}
|
||||
|
||||
func supportsNativeInteractionsEntryProtocol(entryProtocol string) bool {
|
||||
switch entryProtocol {
|
||||
case Interactions, OpenAI, OpenaiResponse, Claude, Gemini:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func excludeExecutionProvider(providers []string, excluded string) []string {
|
||||
excluded = strings.ToLower(strings.TrimSpace(excluded))
|
||||
if excluded == "" || len(providers) == 0 {
|
||||
return providers
|
||||
}
|
||||
excludedIndex := -1
|
||||
for i := range providers {
|
||||
if strings.ToLower(strings.TrimSpace(providers[i])) == excluded {
|
||||
excludedIndex = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if excludedIndex == -1 {
|
||||
return providers
|
||||
}
|
||||
out := make([]string, 0, len(providers)-1)
|
||||
out = append(out, providers[:excludedIndex]...)
|
||||
out = append(out, providers[excludedIndex+1:]...)
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) {
|
||||
return h.getRequestDetailsWithOptions(modelName, false)
|
||||
}
|
||||
|
||||
func validateNativeInteractionsExecution(entryProtocol string, execOptions modelExecutionOptions, routeDecision modelRouteDecision) *interfaces.ErrorMessage {
|
||||
forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider))
|
||||
if forcedProvider == "" || entryProtocol != Interactions {
|
||||
return nil
|
||||
}
|
||||
if routeDecision.ExecutorPluginID != "" {
|
||||
return nativeInteractionsExecutionError()
|
||||
}
|
||||
if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider {
|
||||
return nativeInteractionsExecutionError()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func nativeInteractionsExecutionError() *interfaces.ErrorMessage {
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("agent is only supported for native interactions execution"),
|
||||
}
|
||||
}
|
||||
|
||||
// providersForExecution resolves the providers and normalized model for a request. When a model
|
||||
// router selected a built-in provider, it skips model->provider resolution and uses the router's
|
||||
// provider (with an optional target model); otherwise it falls back to the registry-based path.
|
||||
func (h *BaseAPIHandler) providersForExecution(modelName, originalRequestedModel string, allowImageModel bool, routeDecision modelRouteDecision, execOptions modelExecutionOptions) ([]string, string, *interfaces.ErrorMessage) {
|
||||
forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider))
|
||||
if forcedProvider != "" {
|
||||
if routeDecision.ExecutorPluginID != "" {
|
||||
return nil, "", nativeInteractionsExecutionError()
|
||||
}
|
||||
if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider {
|
||||
return nil, "", nativeInteractionsExecutionError()
|
||||
}
|
||||
normalizedModel := strings.TrimSpace(modelName)
|
||||
if normalizedModel == "" {
|
||||
normalizedModel = strings.TrimSpace(originalRequestedModel)
|
||||
}
|
||||
if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil {
|
||||
return nil, "", errMsg
|
||||
}
|
||||
return []string{forcedProvider}, normalizedModel, nil
|
||||
}
|
||||
if routeDecision.Provider != "" {
|
||||
normalizedModel := originalRequestedModel
|
||||
if routeDecision.Model != "" {
|
||||
normalizedModel = routeDecision.Model
|
||||
}
|
||||
if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil {
|
||||
return nil, "", errMsg
|
||||
}
|
||||
return []string{routeDecision.Provider}, normalizedModel, nil
|
||||
}
|
||||
return h.getRequestDetailsWithOptions(modelName, allowImageModel)
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) getRequestDetailsWithOptions(modelName string, allowImageModel bool) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) {
|
||||
resolvedModelName := modelName
|
||||
initialSuffix := thinking.ParseSuffix(modelName)
|
||||
if initialSuffix.ModelName == "auto" {
|
||||
if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
resolvedModelName = modelName
|
||||
} else {
|
||||
resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName)
|
||||
if initialSuffix.HasSuffix {
|
||||
resolvedModelName = fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix)
|
||||
} else {
|
||||
resolvedModelName = resolvedBase
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
resolvedModelName = modelName
|
||||
} else {
|
||||
resolvedModelName = util.ResolveAutoModel(modelName)
|
||||
}
|
||||
}
|
||||
|
||||
parsed := thinking.ParseSuffix(resolvedModelName)
|
||||
baseModel := strings.TrimSpace(parsed.ModelName)
|
||||
|
||||
if errMsg := h.validateImageOnlyModel(baseModel, allowImageModel); errMsg != nil {
|
||||
return nil, "", errMsg
|
||||
}
|
||||
|
||||
if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
return []string{"home"}, resolvedModelName, nil
|
||||
}
|
||||
|
||||
providers = util.GetProviderName(baseModel)
|
||||
// Fallback: if baseModel has no provider but differs from resolvedModelName,
|
||||
// try using the full model name. This handles edge cases where custom models
|
||||
// may be registered with their full suffixed name (e.g., "my-model(8192)").
|
||||
// Evaluated in Story 11.8: This fallback is intentionally preserved to support
|
||||
// custom model registrations that include thinking suffixes.
|
||||
if len(providers) == 0 && baseModel != resolvedModelName {
|
||||
providers = util.GetProviderName(resolvedModelName)
|
||||
}
|
||||
|
||||
if len(providers) == 0 {
|
||||
// The client asked for a model this proxy cannot route. Report it as a request
|
||||
// error so streaming clients receive an actionable message instead of a
|
||||
// gateway failure they would keep retrying. 400 is used rather than 404 to keep
|
||||
// it distinguishable from an unregistered HTTP route.
|
||||
// The model name is client supplied, so it is inserted through sjson rather
|
||||
// than formatted into the JSON literal: an unescaped quote would otherwise
|
||||
// corrupt the body or let the caller overwrite the error code.
|
||||
body := `{"error":{"message":"","type":"invalid_request_error","code":"model_not_found","param":"model"}}`
|
||||
body, errSet := sjson.Set(body, "error.message", "unknown provider for model "+modelName)
|
||||
if errSet != nil {
|
||||
body = `{"error":{"message":"unknown provider for model","type":"invalid_request_error","code":"model_not_found","param":"model"}}`
|
||||
}
|
||||
return nil, "", &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errors.New(body),
|
||||
}
|
||||
}
|
||||
|
||||
// The thinking suffix is preserved in the model name itself, so no
|
||||
// metadata-based configuration passing is needed.
|
||||
return providers, resolvedModelName, nil
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) validateImageOnlyModel(modelName string, allowImageModel bool) *interfaces.ErrorMessage {
|
||||
baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName)
|
||||
if baseModel == "" {
|
||||
baseModel = strings.TrimSpace(modelName)
|
||||
}
|
||||
if isOpenAIImageOnlyModel(baseModel) && !allowImageModel {
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
Error: fmt.Errorf("model %s is only supported on /v1/images/generations and /v1/images/edits", routeModelBaseName(baseModel)),
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isOpenAIImageOnlyModel(model string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(routeModelBaseName(model))) {
|
||||
case "gpt-image-1.5", "gpt-image-2", "grok-imagine-image", "grok-imagine-image-quality", "grok-imagine-image-2.0":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func routeModelBaseName(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if idx := strings.LastIndex(model, "/"); idx >= 0 && idx < len(model)-1 {
|
||||
return strings.TrimSpace(model[idx+1:])
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
func cloneBytes(src []byte) []byte {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
dst := make([]byte, len(src))
|
||||
copy(dst, src)
|
||||
return dst
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) modelRouterHost() PluginModelRouterHost {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
if !isNilPluginModelRouterHost(h.ModelRouterHost) {
|
||||
return h.ModelRouterHost
|
||||
}
|
||||
host := h.interceptorHost()
|
||||
if host == nil {
|
||||
return nil
|
||||
}
|
||||
router, ok := host.(PluginModelRouterHost)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return router
|
||||
}
|
||||
|
||||
type modelRouteDecision struct {
|
||||
ExecutorPluginID string
|
||||
Provider string
|
||||
Model string
|
||||
}
|
||||
|
||||
func routeModel(ctx context.Context, host PluginModelRouterHost, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) {
|
||||
if host == nil {
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
skipPluginID = strings.TrimSpace(skipPluginID)
|
||||
if skipPluginID != "" {
|
||||
if skipper, ok := host.(pluginModelRouterSkipHost); ok {
|
||||
return skipper.RouteModelExcept(ctx, req, skipPluginID)
|
||||
}
|
||||
return pluginapi.ModelRouteResponse{}, false
|
||||
}
|
||||
return host.RouteModel(ctx, req)
|
||||
}
|
||||
|
||||
func modelRoutersEnabled(host PluginModelRouterHost, skipPluginID string) bool {
|
||||
if host == nil {
|
||||
return false
|
||||
}
|
||||
skipPluginID = strings.TrimSpace(skipPluginID)
|
||||
if skipPluginID != "" {
|
||||
if _, ok := host.(pluginModelRouterSkipHost); !ok {
|
||||
return false
|
||||
}
|
||||
if detector, ok := host.(modelRouterSkipDetector); ok {
|
||||
return detector.HasModelRoutersExcept(skipPluginID)
|
||||
}
|
||||
}
|
||||
if detector, ok := host.(modelRouterDetector); ok {
|
||||
return detector.HasModelRouters()
|
||||
}
|
||||
// No detector: treat routing as disabled (same conservative default as before any
|
||||
// ModelRouter existed). Hosts that route must implement HasModelRouters (pluginhost.Host does).
|
||||
return false
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) applyModelRouter(ctx context.Context, handlerType, modelName string, rawJSON []byte, stream bool, execOptions modelExecutionOptions) modelRouteDecision {
|
||||
var decision modelRouteDecision
|
||||
host := h.modelRouterHost()
|
||||
if host == nil || !modelRoutersEnabled(host, execOptions.SkipRouterPluginID) {
|
||||
return decision
|
||||
}
|
||||
meta := requestExecutionMetadata(ctx)
|
||||
meta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
addModelExecutionSourceMetadata(meta, execOptions.InternalSource)
|
||||
resp, ok := routeModel(ctx, host, pluginapi.ModelRouteRequest{
|
||||
SourceFormat: handlerType,
|
||||
RequestedModel: modelName,
|
||||
Stream: stream,
|
||||
Headers: modelExecutionHeaders(ctx, execOptions.Headers),
|
||||
Query: modelExecutionQuery(ctx, execOptions.Query),
|
||||
Body: cloneBytes(rawJSON),
|
||||
Metadata: meta,
|
||||
}, execOptions.SkipRouterPluginID)
|
||||
if !ok || !resp.Handled {
|
||||
return decision
|
||||
}
|
||||
switch resp.TargetKind {
|
||||
case pluginapi.ModelRouteTargetSelf, pluginapi.ModelRouteTargetExecutor:
|
||||
decision.ExecutorPluginID = strings.TrimSpace(resp.Target)
|
||||
case pluginapi.ModelRouteTargetProvider:
|
||||
decision.Provider = strings.ToLower(strings.TrimSpace(resp.Target))
|
||||
decision.Model = strings.TrimSpace(resp.TargetModel)
|
||||
}
|
||||
return decision
|
||||
}
|
||||
842
backend/sdk/api/handlers/handlers_stream.go
Normal file
842
backend/sdk/api/handlers/handlers_stream.go
Normal file
|
|
@ -0,0 +1,842 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// ExecuteStreamWithAuthManager executes a streaming request via the core auth manager.
|
||||
// This path is the only supported execution route.
|
||||
// The returned http.Header carries upstream response headers captured before streaming begins.
|
||||
func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) {
|
||||
return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false)
|
||||
}
|
||||
|
||||
// ExecuteImageStreamWithAuthManager executes a streaming OpenAI-compatible image endpoint request.
|
||||
func (h *BaseAPIHandler) ExecuteImageStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) {
|
||||
return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true)
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) {
|
||||
if h.AuthManager != nil && h.AuthManager.HomeEnabled() {
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")}
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
host := h.pluginExecutorHost()
|
||||
if host == nil {
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")}
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
execCtx, nestedTracker := withNestedExecutionTracker(ctx)
|
||||
req, opts := h.pluginExecutorRequest(execCtx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions)
|
||||
lifecycle := h.newRequestLifecycleTracker(execCtx, entryProtocol, modelName, originalRequestedModel, true, opts.Metadata, execOptions.SkipInterceptorPluginID)
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(execCtx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(execCtx, interceptErr)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- interceptErr
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(execCtx, host, executorPluginID, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(execCtx, interceptErr)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- interceptErr
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
var reporter *helps.UsageReporter
|
||||
if !execOptions.InternalSource {
|
||||
reporter = helps.NewUsageReporter(execCtx, executorPluginID, modelName, nil)
|
||||
reporter.SetTranslatedReasoningEffort(req.Payload, entryProtocol)
|
||||
}
|
||||
streamResult, errStream := host.ExecutePluginExecutorStream(execCtx, executorPluginID, req, opts)
|
||||
if errStream != nil {
|
||||
if reporter != nil && !nestedTracker.hasNestedExecution() {
|
||||
reporter.PublishFailure(execCtx, errStream)
|
||||
}
|
||||
errMsg := executionErrorMessage(errStream)
|
||||
lifecycle.completeError(execCtx, errMsg)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
if streamResult == nil {
|
||||
errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor returned nil stream")}
|
||||
if reporter != nil && !nestedTracker.hasNestedExecution() {
|
||||
reporter.PublishFailure(execCtx, errMsg.Error)
|
||||
}
|
||||
lifecycle.completeError(execCtx, errMsg)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
|
||||
passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg)
|
||||
interceptorHost := h.interceptorHost()
|
||||
streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost)
|
||||
rawStreamHeaders := cloneHeader(streamResult.Headers)
|
||||
baseStreamHeaders := cloneHeader(streamResult.Headers)
|
||||
// Request headers and request bodies are stream-invariant. Keep a private snapshot
|
||||
// and clone into each interceptor call so plugins cannot mutate shared storage.
|
||||
// Schema v3+ payload chunks omit these bodies (host also strips per plugin).
|
||||
var streamRequestHeaders http.Header
|
||||
var streamOriginalRequest []byte
|
||||
var streamRequestBody []byte
|
||||
applyStreamHeaders := func(headers http.Header) {
|
||||
rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers)
|
||||
}
|
||||
if streamInterceptorsActive {
|
||||
streamRequestHeaders = cloneHeader(opts.Headers)
|
||||
streamOriginalRequest = cloneBytes(opts.OriginalRequest)
|
||||
streamRequestBody = cloneBytes(req.Payload)
|
||||
intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{
|
||||
RequestID: lifecycle.requestID(),
|
||||
SourceFormat: responseProtocol,
|
||||
Model: modelName,
|
||||
RequestedModel: originalRequestedModel,
|
||||
RequestHeaders: cloneHeader(streamRequestHeaders),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
OriginalRequest: cloneBytes(streamOriginalRequest),
|
||||
RequestBody: cloneBytes(streamRequestBody),
|
||||
ChunkIndex: pluginapi.StreamChunkHeaderInitIndex,
|
||||
Metadata: opts.Metadata,
|
||||
}, execOptions.SkipInterceptorPluginID)
|
||||
applyStreamHeaders(intercepted.Headers)
|
||||
}
|
||||
upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled)
|
||||
if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) {
|
||||
upstreamHeaders = make(http.Header)
|
||||
}
|
||||
|
||||
dataChan := make(chan []byte)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
var done <-chan struct{}
|
||||
if ctx != nil {
|
||||
done = ctx.Done()
|
||||
}
|
||||
chunks := streamResult.Chunks
|
||||
if chunks == nil {
|
||||
closed := make(chan coreexecutor.StreamChunk)
|
||||
close(closed)
|
||||
chunks = closed
|
||||
}
|
||||
var responseSSEValidator *sseJSONValidationState
|
||||
if responseProtocol == "openai-response" {
|
||||
responseSSEValidator = &sseJSONValidationState{}
|
||||
}
|
||||
go func() {
|
||||
completionOutcome := pluginapi.RequestCompletionSucceeded
|
||||
completionStatus := http.StatusOK
|
||||
var completionErr error
|
||||
var streamUsage helps.StreamUsageBuffer
|
||||
defer func() {
|
||||
lifecycle.complete(completionOutcome, completionStatus, completionErr)
|
||||
if reporter != nil && !nestedTracker.hasNestedExecution() {
|
||||
if completionOutcome != pluginapi.RequestCompletionSucceeded && completionErr != nil {
|
||||
if !streamUsage.PublishFailure(execCtx, reporter, completionErr) {
|
||||
reporter.PublishFailure(execCtx, completionErr)
|
||||
}
|
||||
} else {
|
||||
streamUsage.Publish(execCtx, reporter)
|
||||
reporter.EnsurePublished(execCtx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
defer close(dataChan)
|
||||
defer close(errChan)
|
||||
chunkIndex := 0
|
||||
var historyChunks [][]byte
|
||||
for {
|
||||
chunk, ok, canceled := nextStreamChunk(ctx, nil, nil, chunks)
|
||||
if canceled {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if !ok {
|
||||
if responseSSEValidator != nil {
|
||||
if errValidate := responseSSEValidator.Finish(); errValidate != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = http.StatusBadGateway
|
||||
completionErr = errValidate
|
||||
select {
|
||||
case errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}:
|
||||
case <-done:
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
errMsg := executionErrorMessage(chunk.Err)
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = errMsg.StatusCode
|
||||
completionErr = chunk.Err
|
||||
select {
|
||||
case errChan <- errMsg:
|
||||
case <-done:
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) == 0 {
|
||||
continue
|
||||
}
|
||||
observePluginExecutorStreamUsage(responseProtocol, chunk.Payload, &streamUsage)
|
||||
payload := cloneBytes(chunk.Payload)
|
||||
if streamInterceptorsActive {
|
||||
chunkReq := pluginapi.StreamChunkInterceptRequest{
|
||||
RequestID: lifecycle.requestID(),
|
||||
SourceFormat: responseProtocol,
|
||||
Model: modelName,
|
||||
RequestedModel: originalRequestedModel,
|
||||
RequestHeaders: cloneHeader(streamRequestHeaders),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
Body: payload,
|
||||
HistoryChunks: cloneByteSlices(historyChunks),
|
||||
ChunkIndex: chunkIndex,
|
||||
Metadata: opts.Metadata,
|
||||
}
|
||||
// Re-evaluate each chunk so mid-stream plugin reloads stay correct.
|
||||
// Schema v3+ omits bodies here (one header-init clone only).
|
||||
if streamChunkPayloadIncludesRequestBody(interceptorHost) {
|
||||
chunkReq.OriginalRequest = cloneBytes(streamOriginalRequest)
|
||||
chunkReq.RequestBody = cloneBytes(streamRequestBody)
|
||||
}
|
||||
intercepted := interceptStreamChunk(ctx, interceptorHost, chunkReq, execOptions.SkipInterceptorPluginID)
|
||||
applyStreamHeaders(intercepted.Headers)
|
||||
if len(intercepted.Body) > 0 {
|
||||
payload = cloneBytes(intercepted.Body)
|
||||
}
|
||||
chunkIndex++
|
||||
if intercepted.DropChunk {
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
chunkIndex++
|
||||
}
|
||||
if responseSSEValidator != nil {
|
||||
validatedPayload, errValidate := responseSSEValidator.AddChunk(payload)
|
||||
if errValidate != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = http.StatusBadGateway
|
||||
completionErr = errValidate
|
||||
select {
|
||||
case errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}:
|
||||
case <-done:
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
payload = validatedPayload
|
||||
if len(payload) == 0 {
|
||||
continue
|
||||
}
|
||||
}
|
||||
select {
|
||||
case dataChan <- payload:
|
||||
if streamInterceptorsActive {
|
||||
historyChunks = appendStreamInterceptorHistory(historyChunks, payload)
|
||||
}
|
||||
case <-done:
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
return dataChan, upstreamHeaders, errChan
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) {
|
||||
return h.executeStreamWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{})
|
||||
}
|
||||
|
||||
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, preparedRoute := preparedModelRouteFromContext(ctx, execOptions.SkipRouterPluginID)
|
||||
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)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
if routeDecision.ExecutorPluginID != "" {
|
||||
return h.streamWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions)
|
||||
}
|
||||
providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions)
|
||||
if errMsg != nil {
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers)
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel
|
||||
addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel)
|
||||
addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource)
|
||||
setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
setGenerateMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
}
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: payload,
|
||||
}
|
||||
afterAuthCapture := &requestAfterAuthCapture{}
|
||||
lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, true, reqMeta, execOptions.SkipInterceptorPluginID)
|
||||
opts := coreexecutor.Options{
|
||||
Stream: true,
|
||||
Alt: alt,
|
||||
OriginalRequest: rawJSON,
|
||||
SourceFormat: sdktranslator.FromString(entryProtocol),
|
||||
ResponseFormat: sdktranslator.FromString(responseProtocol),
|
||||
Headers: modelExecutionHeaders(ctx, execOptions.Headers),
|
||||
Query: modelExecutionQuery(ctx, execOptions.Query),
|
||||
RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID),
|
||||
}
|
||||
opts.Metadata = reqMeta
|
||||
var interceptErr *interfaces.ErrorMessage
|
||||
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID)
|
||||
if interceptErr != nil {
|
||||
lifecycle.completeError(ctx, interceptErr)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- interceptErr
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
err = enrichAuthSelectionError(err, providers, normalizedModel)
|
||||
errMsg := executionErrorMessage(err)
|
||||
lifecycle.completeError(ctx, errMsg)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
if streamResult == nil {
|
||||
errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")}
|
||||
lifecycle.completeError(ctx, errMsg)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
errChan <- errMsg
|
||||
close(errChan)
|
||||
return nil, nil, errChan
|
||||
}
|
||||
executedRequest := func() (coreexecutor.Request, coreexecutor.Options) {
|
||||
return afterAuthCapture.apply(req, opts)
|
||||
}
|
||||
passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg)
|
||||
interceptorHost := h.interceptorHost()
|
||||
streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost)
|
||||
// Resolve bootstrap retries and header initialization before returning so the
|
||||
// returned header snapshot is never modified by the stream goroutine.
|
||||
rawStreamHeaders := cloneHeader(streamResult.Headers)
|
||||
baseStreamHeaders := cloneHeader(streamResult.Headers)
|
||||
chunks := streamResult.Chunks
|
||||
if chunks == nil {
|
||||
closed := make(chan coreexecutor.StreamChunk)
|
||||
close(closed)
|
||||
chunks = closed
|
||||
}
|
||||
streamClosedBeforeRead := false
|
||||
streamCanceledBeforeRead := false
|
||||
streamHeaderInitialized := false
|
||||
// Request headers/bodies are stream-invariant after after-auth capture. Keep a private
|
||||
// snapshot and clone into each interceptor call so plugins cannot mutate shared storage.
|
||||
// Schema v3+ payload chunks omit these bodies (host also strips per plugin).
|
||||
var streamRequestHeaders http.Header
|
||||
var streamOriginalRequest []byte
|
||||
var streamRequestBody []byte
|
||||
|
||||
applyStreamHeaders := func(headers http.Header) {
|
||||
rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers)
|
||||
}
|
||||
|
||||
applyStreamHeaderInit := func() {
|
||||
if !streamInterceptorsActive || streamHeaderInitialized {
|
||||
return
|
||||
}
|
||||
executedReq, executedOpts := executedRequest()
|
||||
streamRequestHeaders = cloneHeader(executedOpts.Headers)
|
||||
streamOriginalRequest = cloneBytes(executedOpts.OriginalRequest)
|
||||
streamRequestBody = cloneBytes(executedReq.Payload)
|
||||
intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{
|
||||
RequestID: lifecycle.requestID(),
|
||||
SourceFormat: responseProtocol,
|
||||
Model: normalizedModel,
|
||||
RequestedModel: originalRequestedModel,
|
||||
RequestHeaders: cloneHeader(streamRequestHeaders),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
OriginalRequest: cloneBytes(streamOriginalRequest),
|
||||
RequestBody: cloneBytes(streamRequestBody),
|
||||
ChunkIndex: pluginapi.StreamChunkHeaderInitIndex,
|
||||
Metadata: executedOpts.Metadata,
|
||||
}, execOptions.SkipInterceptorPluginID)
|
||||
applyStreamHeaders(intercepted.Headers)
|
||||
streamHeaderInitialized = true
|
||||
}
|
||||
|
||||
var responseSSEValidator *sseJSONValidationState
|
||||
if responseProtocol == "openai-response" {
|
||||
responseSSEValidator = &sseJSONValidationState{}
|
||||
}
|
||||
|
||||
transformStreamPayload := func(payload []byte, chunkIndex *int, historyChunks [][]byte) ([]byte, bool, *interfaces.ErrorMessage) {
|
||||
applyStreamHeaderInit()
|
||||
payload = cloneBytes(payload)
|
||||
if streamInterceptorsActive {
|
||||
chunkReq := pluginapi.StreamChunkInterceptRequest{
|
||||
RequestID: lifecycle.requestID(),
|
||||
SourceFormat: responseProtocol,
|
||||
Model: normalizedModel,
|
||||
RequestedModel: originalRequestedModel,
|
||||
RequestHeaders: cloneHeader(streamRequestHeaders),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
Body: payload,
|
||||
HistoryChunks: cloneByteSlices(historyChunks),
|
||||
ChunkIndex: *chunkIndex,
|
||||
Metadata: opts.Metadata,
|
||||
}
|
||||
// Re-evaluate each chunk so mid-stream plugin reloads stay correct.
|
||||
// Schema v3+ omits bodies here (one header-init clone only).
|
||||
if streamChunkPayloadIncludesRequestBody(interceptorHost) {
|
||||
chunkReq.OriginalRequest = cloneBytes(streamOriginalRequest)
|
||||
chunkReq.RequestBody = cloneBytes(streamRequestBody)
|
||||
}
|
||||
intercepted := interceptStreamChunk(ctx, interceptorHost, chunkReq, execOptions.SkipInterceptorPluginID)
|
||||
applyStreamHeaders(intercepted.Headers)
|
||||
if len(intercepted.Body) > 0 {
|
||||
payload = cloneBytes(intercepted.Body)
|
||||
}
|
||||
(*chunkIndex)++
|
||||
if intercepted.DropChunk {
|
||||
return nil, false, nil
|
||||
}
|
||||
} else {
|
||||
(*chunkIndex)++
|
||||
}
|
||||
if responseSSEValidator != nil {
|
||||
validatedPayload, errValidate := responseSSEValidator.AddChunk(payload)
|
||||
if errValidate != nil {
|
||||
return nil, false, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}
|
||||
}
|
||||
payload = validatedPayload
|
||||
if len(payload) == 0 {
|
||||
return nil, false, nil
|
||||
}
|
||||
}
|
||||
return payload, true, nil
|
||||
}
|
||||
|
||||
var bootstrapPayload []byte
|
||||
bootstrapChunkIndex := 0
|
||||
var bootstrapHistoryChunks [][]byte
|
||||
var bootstrapStreamErr error
|
||||
var bootstrapErr *interfaces.ErrorMessage
|
||||
readInitialStreamChunks := func() {
|
||||
for {
|
||||
var chunk coreexecutor.StreamChunk
|
||||
var ok bool
|
||||
if ctx != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
streamCanceledBeforeRead = true
|
||||
return
|
||||
case chunk, ok = <-chunks:
|
||||
}
|
||||
} else {
|
||||
chunk, ok = <-chunks
|
||||
}
|
||||
if !ok {
|
||||
streamClosedBeforeRead = true
|
||||
applyStreamHeaderInit()
|
||||
return
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
bootstrapStreamErr = chunk.Err
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) == 0 {
|
||||
continue
|
||||
}
|
||||
payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &bootstrapChunkIndex, bootstrapHistoryChunks)
|
||||
if errMsg != nil {
|
||||
bootstrapErr = errMsg
|
||||
return
|
||||
}
|
||||
if !deliverable {
|
||||
continue
|
||||
}
|
||||
bootstrapPayload = payload
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
bootstrapEligible := func(err error) bool {
|
||||
status := statusFromError(err)
|
||||
if status == 0 {
|
||||
return true
|
||||
}
|
||||
switch status {
|
||||
case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired,
|
||||
http.StatusRequestTimeout, http.StatusTooManyRequests:
|
||||
return true
|
||||
default:
|
||||
return status >= http.StatusInternalServerError
|
||||
}
|
||||
}
|
||||
|
||||
maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg)
|
||||
if h.AuthManager.HomeEnabled() {
|
||||
maxBootstrapRetries = 0
|
||||
}
|
||||
for bootstrapRetries := 0; !streamCanceledBeforeRead; {
|
||||
readInitialStreamChunks()
|
||||
if streamCanceledBeforeRead || bootstrapErr != nil || bootstrapStreamErr == nil {
|
||||
break
|
||||
}
|
||||
if bootstrapRetries >= maxBootstrapRetries || !bootstrapEligible(bootstrapStreamErr) {
|
||||
bootstrapErr = executionErrorMessage(bootstrapStreamErr)
|
||||
break
|
||||
}
|
||||
bootstrapRetries++
|
||||
retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
|
||||
if retryErr != nil {
|
||||
originalBootstrapErr := executionErrorMessage(bootstrapStreamErr)
|
||||
if isAuthSelectionUnavailable(retryErr) && originalBootstrapErr.StatusCode >= http.StatusInternalServerError {
|
||||
bootstrapErr = originalBootstrapErr
|
||||
} else {
|
||||
bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel))
|
||||
}
|
||||
break
|
||||
}
|
||||
if retryResult == nil {
|
||||
bootstrapErr = executionErrorMessage(fmt.Errorf("auth manager returned nil stream"))
|
||||
break
|
||||
}
|
||||
rawStreamHeaders = cloneHeader(retryResult.Headers)
|
||||
baseStreamHeaders = cloneHeader(retryResult.Headers)
|
||||
streamHeaderInitialized = false
|
||||
streamClosedBeforeRead = false
|
||||
bootstrapStreamErr = nil
|
||||
bootstrapPayload = nil
|
||||
bootstrapChunkIndex = 0
|
||||
bootstrapHistoryChunks = nil
|
||||
if responseSSEValidator != nil {
|
||||
responseSSEValidator = &sseJSONValidationState{}
|
||||
}
|
||||
chunks = retryResult.Chunks
|
||||
if chunks == nil {
|
||||
closed := make(chan coreexecutor.StreamChunk)
|
||||
close(closed)
|
||||
chunks = closed
|
||||
}
|
||||
}
|
||||
|
||||
upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled)
|
||||
if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) {
|
||||
upstreamHeaders = make(http.Header)
|
||||
}
|
||||
dataChan := make(chan []byte)
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
|
||||
go func() {
|
||||
completionOutcome := pluginapi.RequestCompletionSucceeded
|
||||
completionStatus := http.StatusOK
|
||||
var completionErr error
|
||||
defer func() {
|
||||
lifecycle.complete(completionOutcome, completionStatus, completionErr)
|
||||
}()
|
||||
defer close(dataChan)
|
||||
defer close(errChan)
|
||||
if streamCanceledBeforeRead {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
sendErr := func(msg *interfaces.ErrorMessage) bool {
|
||||
if ctx == nil {
|
||||
errChan <- msg
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case errChan <- msg:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
sendData := func(chunk []byte) bool {
|
||||
if ctx == nil {
|
||||
dataChan <- chunk
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case dataChan <- chunk:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
if bootstrapErr != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
if bootstrapErr.DirectResponse {
|
||||
completionOutcome = pluginapi.RequestCompletionRejected
|
||||
}
|
||||
completionStatus = bootstrapErr.StatusCode
|
||||
completionErr = bootstrapErr.Error
|
||||
if !sendErr(bootstrapErr) && ctx != nil && ctx.Err() != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
chunkIndex := bootstrapChunkIndex
|
||||
historyChunks := bootstrapHistoryChunks
|
||||
if bootstrapPayload != nil {
|
||||
if okSendData := sendData(bootstrapPayload); !okSendData {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if streamInterceptorsActive {
|
||||
historyChunks = appendStreamInterceptorHistory(historyChunks, bootstrapPayload)
|
||||
}
|
||||
}
|
||||
for {
|
||||
chunk, ok, canceled := nextStreamChunk(ctx, nil, &streamClosedBeforeRead, chunks)
|
||||
if canceled {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if !ok {
|
||||
if responseSSEValidator != nil {
|
||||
if errValidate := responseSSEValidator.Finish(); errValidate != nil {
|
||||
errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = errMsg.StatusCode
|
||||
completionErr = errMsg.Error
|
||||
_ = sendErr(errMsg)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
errMsg := executionErrorMessage(chunk.Err)
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = errMsg.StatusCode
|
||||
completionErr = chunk.Err
|
||||
if !sendErr(errMsg) && ctx != nil && ctx.Err() != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) == 0 {
|
||||
continue
|
||||
}
|
||||
payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &chunkIndex, historyChunks)
|
||||
if errMsg != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionFailed
|
||||
completionStatus = errMsg.StatusCode
|
||||
completionErr = errMsg.Error
|
||||
if !sendErr(errMsg) && ctx != nil && ctx.Err() != nil {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if !deliverable {
|
||||
continue
|
||||
}
|
||||
if okSendData := sendData(payload); !okSendData {
|
||||
completionOutcome = pluginapi.RequestCompletionCanceled
|
||||
completionStatus = 0
|
||||
if ctx != nil {
|
||||
completionErr = ctx.Err()
|
||||
}
|
||||
return
|
||||
}
|
||||
if streamInterceptorsActive {
|
||||
historyChunks = appendStreamInterceptorHistory(historyChunks, payload)
|
||||
}
|
||||
}
|
||||
}()
|
||||
return dataChan, upstreamHeaders, errChan
|
||||
}
|
||||
|
||||
type sseJSONValidationState struct {
|
||||
pending []byte
|
||||
pendingErr error
|
||||
}
|
||||
|
||||
func (s *sseJSONValidationState) AddChunk(chunk []byte) ([]byte, error) {
|
||||
if s.pendingErr != nil {
|
||||
errPending := s.pendingErr
|
||||
s.pendingErr = nil
|
||||
return nil, errPending
|
||||
}
|
||||
if len(chunk) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
chunk = bytes.ReplaceAll(chunk, []byte("\r\n"), []byte("\n"))
|
||||
chunk = bytes.ReplaceAll(chunk, []byte("\r"), []byte("\n"))
|
||||
if len(s.pending) > 0 && !bytes.HasSuffix(s.pending, []byte("\n")) && !bytes.HasPrefix(chunk, []byte("\n")) {
|
||||
first := bytes.TrimSpace(bytes.SplitN(chunk, []byte("\n"), 2)[0])
|
||||
if bytes.HasPrefix(first, []byte("data:")) || bytes.HasPrefix(first, []byte("event:")) {
|
||||
s.pending = append(s.pending, '\n')
|
||||
}
|
||||
}
|
||||
s.pending = append(s.pending, chunk...)
|
||||
|
||||
var output []byte
|
||||
for {
|
||||
frameEnd := bytes.Index(s.pending, []byte("\n\n"))
|
||||
if frameEnd < 0 {
|
||||
break
|
||||
}
|
||||
frameEnd += 2
|
||||
frame := s.pending[:frameEnd]
|
||||
if errValidate := validateSSEFrameDataJSON(frame); errValidate != nil {
|
||||
if len(output) > 0 {
|
||||
s.pending = s.pending[:0]
|
||||
s.pendingErr = errValidate
|
||||
return output, nil
|
||||
}
|
||||
return nil, errValidate
|
||||
}
|
||||
output = append(output, frame...)
|
||||
copy(s.pending, s.pending[frameEnd:])
|
||||
s.pending = s.pending[:len(s.pending)-frameEnd]
|
||||
}
|
||||
|
||||
if len(bytes.TrimSpace(s.pending)) == 0 {
|
||||
s.pending = s.pending[:0]
|
||||
return output, nil
|
||||
}
|
||||
payload, found := sseJSONValidationDataPayload(s.pending)
|
||||
payload = bytes.TrimSpace(payload)
|
||||
if !found || len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) || json.Valid(payload) {
|
||||
output = append(output, s.pending...)
|
||||
s.pending = s.pending[:0]
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
func (s *sseJSONValidationState) Finish() error {
|
||||
if s.pendingErr != nil {
|
||||
errPending := s.pendingErr
|
||||
s.pendingErr = nil
|
||||
s.pending = nil
|
||||
return errPending
|
||||
}
|
||||
if len(bytes.TrimSpace(s.pending)) == 0 {
|
||||
s.pending = nil
|
||||
return nil
|
||||
}
|
||||
errValidate := validateSSEFrameDataJSON(s.pending)
|
||||
s.pending = nil
|
||||
return errValidate
|
||||
}
|
||||
|
||||
func sseJSONValidationDataPayload(frame []byte) ([]byte, bool) {
|
||||
var payload []byte
|
||||
found := false
|
||||
for _, line := range bytes.Split(frame, []byte("\n")) {
|
||||
line = bytes.TrimSpace(line)
|
||||
if !bytes.HasPrefix(line, []byte("data:")) {
|
||||
continue
|
||||
}
|
||||
if found {
|
||||
payload = append(payload, '\n')
|
||||
}
|
||||
payload = append(payload, bytes.TrimSpace(line[len("data:"):])...)
|
||||
found = true
|
||||
}
|
||||
return payload, found
|
||||
}
|
||||
|
||||
func validateSSEFrameDataJSON(frame []byte) error {
|
||||
payload, found := sseJSONValidationDataPayload(frame)
|
||||
payload = bytes.TrimSpace(payload)
|
||||
if !found || len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) || json.Valid(payload) {
|
||||
return nil
|
||||
}
|
||||
const max = 512
|
||||
preview := payload
|
||||
if len(preview) > max {
|
||||
preview = preview[:max]
|
||||
}
|
||||
return fmt.Errorf("invalid SSE data JSON (len=%d): %q", len(payload), preview)
|
||||
}
|
||||
|
||||
func validateSSEDataJSON(chunk []byte) error {
|
||||
state := &sseJSONValidationState{}
|
||||
if _, errAdd := state.AddChunk(chunk); errAdd != nil {
|
||||
return errAdd
|
||||
}
|
||||
return state.Finish()
|
||||
}
|
||||
1188
backend/sdk/api/handlers/handlers_stream_bootstrap_test.go
Normal file
1188
backend/sdk/api/handlers/handlers_stream_bootstrap_test.go
Normal file
File diff suppressed because it is too large
Load diff
124
backend/sdk/api/handlers/header_filter.go
Normal file
124
backend/sdk/api/handlers/header_filter.go
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// gatewayHeaderPrefixes lists header name prefixes injected by known AI gateway
|
||||
// proxies. Claude Code's client-side telemetry detects these and reports the
|
||||
// gateway type, so we strip them from upstream responses to avoid detection.
|
||||
var gatewayHeaderPrefixes = []string{
|
||||
"x-litellm-",
|
||||
"helicone-",
|
||||
"x-portkey-",
|
||||
"cf-aig-",
|
||||
"x-kong-",
|
||||
"x-bt-",
|
||||
}
|
||||
|
||||
// hopByHopHeaders lists RFC 7230 Section 6.1 hop-by-hop headers that MUST NOT
|
||||
// be forwarded by proxies, plus security-sensitive headers that should not leak.
|
||||
var hopByHopHeaders = map[string]struct{}{
|
||||
// RFC 7230 hop-by-hop
|
||||
"Connection": {},
|
||||
"Keep-Alive": {},
|
||||
"Proxy-Authenticate": {},
|
||||
"Proxy-Authorization": {},
|
||||
"Te": {},
|
||||
"Trailer": {},
|
||||
"Transfer-Encoding": {},
|
||||
"Upgrade": {},
|
||||
// Security-sensitive
|
||||
"Set-Cookie": {},
|
||||
// CPA-managed (set by handlers, not upstream)
|
||||
"Content-Length": {},
|
||||
"Content-Encoding": {},
|
||||
}
|
||||
|
||||
var cpaReservedResponseHeaders = map[string]struct{}{
|
||||
"Access-Control-Allow-Credentials": {},
|
||||
"Access-Control-Allow-Headers": {},
|
||||
"Access-Control-Allow-Methods": {},
|
||||
"Access-Control-Allow-Origin": {},
|
||||
"Access-Control-Expose-Headers": {},
|
||||
"Access-Control-Max-Age": {},
|
||||
"X-Cpa-Trace-Id": {},
|
||||
}
|
||||
|
||||
// IsCPAReservedResponseHeader reports whether a downstream response header is managed by CPA.
|
||||
func IsCPAReservedResponseHeader(name string) bool {
|
||||
_, reserved := cpaReservedResponseHeaders[http.CanonicalHeaderKey(name)]
|
||||
return reserved
|
||||
}
|
||||
|
||||
// FilterUpstreamHeaders returns a copy of src with hop-by-hop and security-sensitive
|
||||
// headers removed. Returns nil if src is nil or empty after filtering.
|
||||
func FilterUpstreamHeaders(src http.Header) http.Header {
|
||||
if src == nil {
|
||||
return nil
|
||||
}
|
||||
connectionScoped := connectionScopedHeaders(src)
|
||||
dst := make(http.Header)
|
||||
for key, values := range src {
|
||||
canonicalKey := http.CanonicalHeaderKey(key)
|
||||
if _, blocked := hopByHopHeaders[canonicalKey]; blocked {
|
||||
continue
|
||||
}
|
||||
if _, reserved := cpaReservedResponseHeaders[canonicalKey]; reserved {
|
||||
continue
|
||||
}
|
||||
if _, scoped := connectionScoped[canonicalKey]; scoped {
|
||||
continue
|
||||
}
|
||||
// Strip headers injected by known AI gateway proxies to avoid
|
||||
// Claude Code client-side gateway detection.
|
||||
lowerKey := strings.ToLower(key)
|
||||
gatewayMatch := false
|
||||
for _, prefix := range gatewayHeaderPrefixes {
|
||||
if strings.HasPrefix(lowerKey, prefix) {
|
||||
gatewayMatch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if gatewayMatch {
|
||||
continue
|
||||
}
|
||||
dst[key] = values
|
||||
}
|
||||
if len(dst) == 0 {
|
||||
return nil
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func connectionScopedHeaders(src http.Header) map[string]struct{} {
|
||||
scoped := make(map[string]struct{})
|
||||
for _, rawValue := range src.Values("Connection") {
|
||||
for _, token := range strings.Split(rawValue, ",") {
|
||||
headerName := strings.TrimSpace(token)
|
||||
if headerName == "" {
|
||||
continue
|
||||
}
|
||||
scoped[http.CanonicalHeaderKey(headerName)] = struct{}{}
|
||||
}
|
||||
}
|
||||
return scoped
|
||||
}
|
||||
|
||||
// WriteUpstreamHeaders writes filtered upstream headers to the gin response writer.
|
||||
// Headers already set by CPA (e.g., Content-Type) are NOT overwritten.
|
||||
func WriteUpstreamHeaders(dst http.Header, src http.Header) {
|
||||
if src == nil {
|
||||
return
|
||||
}
|
||||
for key, values := range src {
|
||||
// Don't overwrite headers already set by CPA handlers
|
||||
if dst.Get(key) != "" {
|
||||
continue
|
||||
}
|
||||
for _, v := range values {
|
||||
dst.Add(key, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
59
backend/sdk/api/handlers/header_filter_test.go
Normal file
59
backend/sdk/api/handlers/header_filter_test.go
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFilterUpstreamHeaders_RemovesConnectionScopedHeaders(t *testing.T) {
|
||||
src := http.Header{}
|
||||
src.Add("Connection", "keep-alive, x-hop-a, x-hop-b")
|
||||
src.Add("Connection", "x-hop-c")
|
||||
src.Set("Keep-Alive", "timeout=5")
|
||||
src.Set("X-Hop-A", "a")
|
||||
src.Set("X-Hop-B", "b")
|
||||
src.Set("X-Hop-C", "c")
|
||||
src.Set("X-Request-Id", "req-1")
|
||||
src.Set("Set-Cookie", "session=secret")
|
||||
src.Set("x-cpa-trace-id", "upstream-trace")
|
||||
src.Set("Access-Control-Expose-Headers", "upstream-header")
|
||||
|
||||
filtered := FilterUpstreamHeaders(src)
|
||||
if filtered == nil {
|
||||
t.Fatalf("expected filtered headers, got nil")
|
||||
}
|
||||
|
||||
requestID := filtered.Get("X-Request-Id")
|
||||
if requestID != "req-1" {
|
||||
t.Fatalf("expected X-Request-Id to be preserved, got %q", requestID)
|
||||
}
|
||||
|
||||
blockedHeaderKeys := []string{
|
||||
"Connection",
|
||||
"Keep-Alive",
|
||||
"X-Hop-A",
|
||||
"X-Hop-B",
|
||||
"X-Hop-C",
|
||||
"Set-Cookie",
|
||||
"x-cpa-trace-id",
|
||||
"Access-Control-Expose-Headers",
|
||||
}
|
||||
for _, key := range blockedHeaderKeys {
|
||||
value := filtered.Get(key)
|
||||
if value != "" {
|
||||
t.Fatalf("expected %s to be removed, got %q", key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFilterUpstreamHeaders_ReturnsNilWhenAllHeadersBlocked(t *testing.T) {
|
||||
src := http.Header{}
|
||||
src.Add("Connection", "x-hop-a")
|
||||
src.Set("X-Hop-A", "a")
|
||||
src.Set("Set-Cookie", "session=secret")
|
||||
|
||||
filtered := FilterUpstreamHeaders(src)
|
||||
if filtered != nil {
|
||||
t.Fatalf("expected nil when all headers are filtered, got %#v", filtered)
|
||||
}
|
||||
}
|
||||
338
backend/sdk/api/handlers/model_execution.go
Normal file
338
backend/sdk/api/handlers/model_execution.go
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
const (
|
||||
modelExecutionMetadataSourceKey = "source"
|
||||
modelExecutionInternalSource = "plugin_host_model_callback"
|
||||
)
|
||||
|
||||
type modelExecutionOptions struct {
|
||||
Headers http.Header
|
||||
Query url.Values
|
||||
InternalSource bool
|
||||
SkipInterceptorPluginID string
|
||||
SkipRouterPluginID string
|
||||
ForcedProvider string
|
||||
AuthSelectionModel string
|
||||
}
|
||||
|
||||
// ProtocolExecutionRequest describes a route-level model execution request with explicit protocols.
|
||||
type ProtocolExecutionRequest struct {
|
||||
EntryProtocol string
|
||||
ExitProtocol string
|
||||
ForcedProvider string
|
||||
AuthSelectionModel string
|
||||
Model string
|
||||
Stream bool
|
||||
Body []byte
|
||||
Headers http.Header
|
||||
Query url.Values
|
||||
Alt string
|
||||
}
|
||||
|
||||
// ModelExecutionRequest describes an internal model execution request.
|
||||
type ModelExecutionRequest struct {
|
||||
EntryProtocol string
|
||||
ExitProtocol string
|
||||
Model string
|
||||
Stream bool
|
||||
Body []byte
|
||||
Headers http.Header
|
||||
Query url.Values
|
||||
Alt string
|
||||
SkipInterceptorPluginID string
|
||||
SkipRouterPluginID string
|
||||
}
|
||||
|
||||
// ModelExecutionResponse describes a non-streaming internal model execution response.
|
||||
type ModelExecutionResponse struct {
|
||||
StatusCode int
|
||||
Headers http.Header
|
||||
Body []byte
|
||||
}
|
||||
|
||||
// ModelExecutionStream describes a streaming internal model execution response.
|
||||
type ModelExecutionStream struct {
|
||||
StatusCode int
|
||||
Headers http.Header
|
||||
Chunks <-chan ModelExecutionChunk
|
||||
}
|
||||
|
||||
// ModelExecutionChunk carries either a streaming payload or a terminal stream error.
|
||||
type ModelExecutionChunk struct {
|
||||
Payload []byte
|
||||
Err *ModelExecutionStreamError
|
||||
}
|
||||
|
||||
// ModelExecutionStreamError carries a JSON-friendly terminal stream error.
|
||||
type ModelExecutionStreamError struct {
|
||||
StatusCode int `json:"status_code"`
|
||||
Message string `json:"message"`
|
||||
Headers http.Header `json:"headers"`
|
||||
}
|
||||
|
||||
// Error returns the stream error message or the HTTP status text.
|
||||
func (e *ModelExecutionStreamError) Error() string {
|
||||
if e == nil {
|
||||
return ""
|
||||
}
|
||||
if e.Message != "" {
|
||||
return e.Message
|
||||
}
|
||||
return http.StatusText(e.StatusCode)
|
||||
}
|
||||
|
||||
// ExecuteModel executes an internal non-streaming model request.
|
||||
// Host model callbacks are non-recursive for their caller: when
|
||||
// skip plugin IDs are set, that plugin's interceptors and router are skipped
|
||||
// for the nested model execution while other plugins may still run.
|
||||
func (h *BaseAPIHandler) ExecuteModel(ctx context.Context, req ModelExecutionRequest) (ModelExecutionResponse, *interfaces.ErrorMessage) {
|
||||
markNestedExecution(ctx)
|
||||
if req.Stream {
|
||||
return ModelExecutionResponse{}, modelExecutionModeError("ExecuteModel requires Stream=false")
|
||||
}
|
||||
body, headers, errMsg := h.executeWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{
|
||||
Headers: req.Headers,
|
||||
Query: req.Query,
|
||||
InternalSource: true,
|
||||
SkipInterceptorPluginID: req.SkipInterceptorPluginID,
|
||||
SkipRouterPluginID: req.SkipRouterPluginID,
|
||||
})
|
||||
if errMsg != nil {
|
||||
return ModelExecutionResponse{}, errMsg
|
||||
}
|
||||
return ModelExecutionResponse{
|
||||
StatusCode: http.StatusOK,
|
||||
Headers: cloneHeader(headers),
|
||||
Body: cloneBytes(body),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExecuteModelStream executes an internal streaming model request.
|
||||
// Host model callbacks are non-recursive for their caller: when
|
||||
// skip plugin IDs are set, that plugin's interceptors and router are skipped
|
||||
// for the nested model execution while other plugins may still run.
|
||||
func (h *BaseAPIHandler) ExecuteModelStream(ctx context.Context, req ModelExecutionRequest) (ModelExecutionStream, *interfaces.ErrorMessage) {
|
||||
markNestedExecution(ctx)
|
||||
if !req.Stream {
|
||||
return ModelExecutionStream{}, modelExecutionModeError("ExecuteModelStream requires Stream=true")
|
||||
}
|
||||
dataChan, headers, errChan := h.executeStreamWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{
|
||||
Headers: req.Headers,
|
||||
Query: req.Query,
|
||||
InternalSource: true,
|
||||
SkipInterceptorPluginID: req.SkipInterceptorPluginID,
|
||||
SkipRouterPluginID: req.SkipRouterPluginID,
|
||||
})
|
||||
chunks, errMsg := prepareModelExecutionStream(ctx, dataChan, errChan)
|
||||
if errMsg != nil {
|
||||
return ModelExecutionStream{}, errMsg
|
||||
}
|
||||
return ModelExecutionStream{
|
||||
StatusCode: http.StatusOK,
|
||||
Headers: cloneHeader(headers),
|
||||
Chunks: chunks,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExecuteProtocolWithAuthManager executes a route-level non-streaming request with explicit protocols.
|
||||
func (h *BaseAPIHandler) ExecuteProtocolWithAuthManager(ctx context.Context, req ProtocolExecutionRequest) (ModelExecutionResponse, *interfaces.ErrorMessage) {
|
||||
if req.Stream {
|
||||
return ModelExecutionResponse{}, modelExecutionModeError("ExecuteProtocolWithAuthManager requires Stream=false")
|
||||
}
|
||||
body, headers, errMsg := h.executeWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{
|
||||
Headers: req.Headers,
|
||||
Query: req.Query,
|
||||
ForcedProvider: req.ForcedProvider,
|
||||
AuthSelectionModel: req.AuthSelectionModel,
|
||||
})
|
||||
if errMsg != nil {
|
||||
return ModelExecutionResponse{}, errMsg
|
||||
}
|
||||
return ModelExecutionResponse{
|
||||
StatusCode: http.StatusOK,
|
||||
Headers: cloneHeader(headers),
|
||||
Body: cloneBytes(body),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExecuteProtocolStreamWithAuthManager executes a route-level streaming request with explicit protocols.
|
||||
func (h *BaseAPIHandler) ExecuteProtocolStreamWithAuthManager(ctx context.Context, req ProtocolExecutionRequest) (ModelExecutionStream, *interfaces.ErrorMessage) {
|
||||
if !req.Stream {
|
||||
return ModelExecutionStream{}, modelExecutionModeError("ExecuteProtocolStreamWithAuthManager requires Stream=true")
|
||||
}
|
||||
dataChan, headers, errChan := h.executeStreamWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{
|
||||
Headers: req.Headers,
|
||||
Query: req.Query,
|
||||
ForcedProvider: req.ForcedProvider,
|
||||
AuthSelectionModel: req.AuthSelectionModel,
|
||||
})
|
||||
chunks, errMsg := prepareModelExecutionStream(ctx, dataChan, errChan)
|
||||
if errMsg != nil {
|
||||
return ModelExecutionStream{}, errMsg
|
||||
}
|
||||
return ModelExecutionStream{
|
||||
StatusCode: http.StatusOK,
|
||||
Headers: cloneHeader(headers),
|
||||
Chunks: chunks,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func modelExecutionModeError(message string) *interfaces.ErrorMessage {
|
||||
return &interfaces.ErrorMessage{StatusCode: http.StatusBadRequest, Error: errors.New(message)}
|
||||
}
|
||||
|
||||
func modelExecutionResponseProtocol(entryProtocol, exitProtocol string) string {
|
||||
if exitProtocol == "" {
|
||||
return entryProtocol
|
||||
}
|
||||
return exitProtocol
|
||||
}
|
||||
|
||||
func modelExecutionHeaders(ctx context.Context, headers http.Header) http.Header {
|
||||
if len(headers) > 0 {
|
||||
return cloneHeader(headers)
|
||||
}
|
||||
return headersFromContext(ctx)
|
||||
}
|
||||
|
||||
// modelExecutionQuery prefers an explicitly provided query and otherwise falls
|
||||
// back to the inbound query embedded in the request context. This lets model
|
||||
// routers observe query parameters for plain HTTP requests even when callers
|
||||
// do not populate execOptions.Query (mirrors modelExecutionHeaders).
|
||||
func modelExecutionQuery(ctx context.Context, query url.Values) url.Values {
|
||||
if len(query) > 0 {
|
||||
return cloneURLValues(query)
|
||||
}
|
||||
return queryFromContext(ctx)
|
||||
}
|
||||
|
||||
func cloneURLValues(src url.Values) url.Values {
|
||||
if src == nil {
|
||||
return nil
|
||||
}
|
||||
dst := make(url.Values, len(src))
|
||||
for key, values := range src {
|
||||
dst[key] = append([]string(nil), values...)
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func addModelExecutionSourceMetadata(meta map[string]any, internalSource bool) {
|
||||
if !internalSource || meta == nil {
|
||||
return
|
||||
}
|
||||
meta[modelExecutionMetadataSourceKey] = modelExecutionInternalSource
|
||||
}
|
||||
|
||||
func prepareModelExecutionStream(ctx context.Context, dataChan <-chan []byte, errChan <-chan *interfaces.ErrorMessage) (<-chan ModelExecutionChunk, *interfaces.ErrorMessage) {
|
||||
pending, nextDataChan, nextErrChan, errMsg := receiveInitialModelExecutionChunk(ctx, dataChan, errChan)
|
||||
if errMsg != nil {
|
||||
return nil, errMsg
|
||||
}
|
||||
return wrapModelExecutionChunks(ctx, nextDataChan, nextErrChan, pending), nil
|
||||
}
|
||||
|
||||
func receiveInitialModelExecutionChunk(ctx context.Context, dataChan <-chan []byte, errChan <-chan *interfaces.ErrorMessage) ([]ModelExecutionChunk, <-chan []byte, <-chan *interfaces.ErrorMessage, *interfaces.ErrorMessage) {
|
||||
var done <-chan struct{}
|
||||
if ctx != nil {
|
||||
done = ctx.Done()
|
||||
}
|
||||
for dataChan != nil || errChan != nil {
|
||||
select {
|
||||
case payload, ok := <-dataChan:
|
||||
if !ok {
|
||||
dataChan = nil
|
||||
continue
|
||||
}
|
||||
return []ModelExecutionChunk{{Payload: cloneBytes(payload)}}, dataChan, errChan, nil
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
return nil, dataChan, errChan, errMsg
|
||||
}
|
||||
case <-done:
|
||||
return nil, dataChan, errChan, nil
|
||||
}
|
||||
}
|
||||
return nil, dataChan, errChan, nil
|
||||
}
|
||||
|
||||
func wrapModelExecutionChunks(ctx context.Context, dataChan <-chan []byte, errChan <-chan *interfaces.ErrorMessage, pending []ModelExecutionChunk) <-chan ModelExecutionChunk {
|
||||
chunks := make(chan ModelExecutionChunk)
|
||||
go func() {
|
||||
defer close(chunks)
|
||||
var done <-chan struct{}
|
||||
if ctx != nil {
|
||||
done = ctx.Done()
|
||||
}
|
||||
for _, chunk := range pending {
|
||||
if !sendModelExecutionChunk(ctx, chunks, chunk) {
|
||||
return
|
||||
}
|
||||
}
|
||||
for dataChan != nil || errChan != nil {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case payload, ok := <-dataChan:
|
||||
if !ok {
|
||||
dataChan = nil
|
||||
continue
|
||||
}
|
||||
if !sendModelExecutionChunk(ctx, chunks, ModelExecutionChunk{Payload: cloneBytes(payload)}) {
|
||||
return
|
||||
}
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
_ = sendModelExecutionChunk(ctx, chunks, ModelExecutionChunk{Err: modelExecutionStreamErrorFromMessage(errMsg)})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
return chunks
|
||||
}
|
||||
|
||||
func modelExecutionStreamErrorFromMessage(errMsg *interfaces.ErrorMessage) *ModelExecutionStreamError {
|
||||
if errMsg == nil {
|
||||
return nil
|
||||
}
|
||||
message := ""
|
||||
if errMsg.Error != nil {
|
||||
message = errMsg.Error.Error()
|
||||
}
|
||||
return &ModelExecutionStreamError{
|
||||
StatusCode: errMsg.StatusCode,
|
||||
Message: message,
|
||||
Headers: cloneHeader(errMsg.Addon),
|
||||
}
|
||||
}
|
||||
|
||||
func sendModelExecutionChunk(ctx context.Context, chunks chan<- ModelExecutionChunk, chunk ModelExecutionChunk) bool {
|
||||
if ctx == nil {
|
||||
chunks <- chunk
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case chunks <- chunk:
|
||||
return true
|
||||
}
|
||||
}
|
||||
788
backend/sdk/api/handlers/model_execution_test.go
Normal file
788
backend/sdk/api/handlers/model_execution_test.go
Normal file
|
|
@ -0,0 +1,788 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
)
|
||||
|
||||
type modelExecutionCaptureExecutor struct {
|
||||
provider string
|
||||
|
||||
mu sync.Mutex
|
||||
lastRequest coreexecutor.Request
|
||||
lastOptions coreexecutor.Options
|
||||
execute func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error)
|
||||
stream func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error)
|
||||
}
|
||||
|
||||
type modelExecutionStatusHeaderError struct {
|
||||
statusCode int
|
||||
message string
|
||||
headers http.Header
|
||||
}
|
||||
|
||||
type modelExecutionSkipHost struct {
|
||||
beforeSkip string
|
||||
afterSkip string
|
||||
respSkip string
|
||||
streamSkip []string
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
panic("InterceptRequestBeforeAuth called without skip")
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse {
|
||||
panic("InterceptRequestAfterAuth called without skip")
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse {
|
||||
panic("InterceptResponse called without skip")
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse {
|
||||
panic("InterceptStreamChunk called without skip")
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptRequestBeforeAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse {
|
||||
h.beforeSkip = skipPluginID
|
||||
return pluginapi.RequestInterceptResponse{
|
||||
Headers: cloneHeader(req.Headers),
|
||||
Body: cloneBytes(req.Body),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptRequestAfterAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse {
|
||||
h.afterSkip = skipPluginID
|
||||
return pluginapi.RequestInterceptResponse{
|
||||
Headers: cloneHeader(req.Headers),
|
||||
Body: cloneBytes(req.Body),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptResponseExcept(ctx context.Context, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse {
|
||||
h.respSkip = skipPluginID
|
||||
return pluginapi.ResponseInterceptResponse{
|
||||
Headers: cloneHeader(req.ResponseHeaders),
|
||||
Body: cloneBytes(req.Body),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *modelExecutionSkipHost) InterceptStreamChunkExcept(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse {
|
||||
h.streamSkip = append(h.streamSkip, skipPluginID)
|
||||
return pluginapi.StreamChunkInterceptResponse{
|
||||
Headers: cloneHeader(req.ResponseHeaders),
|
||||
Body: cloneBytes(req.Body),
|
||||
}
|
||||
}
|
||||
|
||||
func (e modelExecutionStatusHeaderError) Error() string {
|
||||
return e.message
|
||||
}
|
||||
|
||||
func (e modelExecutionStatusHeaderError) StatusCode() int {
|
||||
return e.statusCode
|
||||
}
|
||||
|
||||
func (e modelExecutionStatusHeaderError) Headers() http.Header {
|
||||
return e.headers
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) Identifier() string {
|
||||
if e.provider != "" {
|
||||
return e.provider
|
||||
}
|
||||
return "codex"
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) Execute(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
e.capture(req, opts)
|
||||
if e.execute != nil {
|
||||
return e.execute(ctx, auth, req, opts)
|
||||
}
|
||||
return coreexecutor.Response{Payload: []byte("model-execution-ok")}, nil
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
e.capture(req, opts)
|
||||
if e.stream != nil {
|
||||
return e.stream(ctx, auth, req, opts)
|
||||
}
|
||||
chunks := make(chan coreexecutor.StreamChunk)
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{Payload: []byte("0")}, nil
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented}
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) capture(req coreexecutor.Request, opts coreexecutor.Options) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
e.lastRequest = coreexecutor.Request{
|
||||
Model: req.Model,
|
||||
Payload: cloneBytes(req.Payload),
|
||||
Format: req.Format,
|
||||
Metadata: req.Metadata,
|
||||
}
|
||||
e.lastOptions = coreexecutor.Options{
|
||||
Stream: opts.Stream,
|
||||
Alt: opts.Alt,
|
||||
Headers: cloneHeader(opts.Headers),
|
||||
Query: cloneURLValues(opts.Query),
|
||||
OriginalRequest: cloneBytes(opts.OriginalRequest),
|
||||
SourceFormat: opts.SourceFormat,
|
||||
ResponseFormat: opts.ResponseFormat,
|
||||
Metadata: opts.Metadata,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *modelExecutionCaptureExecutor) captured() (coreexecutor.Request, coreexecutor.Options) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
return e.lastRequest, e.lastOptions
|
||||
}
|
||||
|
||||
func newModelExecutionHandler(t *testing.T, model string, executor *modelExecutionCaptureExecutor, cfg *sdkconfig.SDKConfig) *BaseAPIHandler {
|
||||
t.Helper()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "model-execution-" + model,
|
||||
Provider: executor.Identifier(),
|
||||
Status: coreauth.StatusActive,
|
||||
Metadata: map[string]any{"email": model + "@example.com"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
return NewBaseAPIHandlers(cfg, manager)
|
||||
}
|
||||
|
||||
func TestExecuteModelCarriesEntryAndExitProtocols(t *testing.T) {
|
||||
model := "model-execution-nonstream-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{
|
||||
Payload: []byte(`{"ok":true}`),
|
||||
Headers: http.Header{
|
||||
"X-Upstream": []string{"nonstream"},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{PassthroughHeaders: true})
|
||||
|
||||
resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "claude",
|
||||
Model: model,
|
||||
Body: requestBody,
|
||||
Headers: http.Header{"X-Callback": []string{"nonstream"}},
|
||||
Query: url.Values{"q": []string{"callback"}},
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModel() error = %+v", errMsg)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusOK)
|
||||
}
|
||||
if string(resp.Body) != `{"ok":true}` {
|
||||
t.Fatalf("body = %q, want executor response", resp.Body)
|
||||
}
|
||||
if resp.Headers.Get("X-Upstream") != "nonstream" {
|
||||
t.Fatalf("headers = %#v, want upstream header", resp.Headers)
|
||||
}
|
||||
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != model {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, model)
|
||||
}
|
||||
if string(gotReq.Payload) != string(requestBody) {
|
||||
t.Fatalf("executor payload = %q, want %q", gotReq.Payload, requestBody)
|
||||
}
|
||||
if gotOpts.Stream {
|
||||
t.Fatal("executor stream option = true, want false")
|
||||
}
|
||||
if gotOpts.SourceFormat != sdktranslator.FormatOpenAI {
|
||||
t.Fatalf("SourceFormat = %q, want %q", gotOpts.SourceFormat, sdktranslator.FormatOpenAI)
|
||||
}
|
||||
if gotOpts.ResponseFormat != sdktranslator.FormatClaude {
|
||||
t.Fatalf("ResponseFormat = %q, want %q", gotOpts.ResponseFormat, sdktranslator.FormatClaude)
|
||||
}
|
||||
if gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey] != model {
|
||||
t.Fatalf("requested model metadata = %#v, want %q", gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey], model)
|
||||
}
|
||||
if gotOpts.Metadata[modelExecutionMetadataSourceKey] != modelExecutionInternalSource {
|
||||
t.Fatalf("source metadata = %#v, want %q", gotOpts.Metadata[modelExecutionMetadataSourceKey], modelExecutionInternalSource)
|
||||
}
|
||||
if gotOpts.Headers.Get("X-Callback") != "nonstream" {
|
||||
t.Fatalf("executor headers = %#v, want callback header", gotOpts.Headers)
|
||||
}
|
||||
if gotOpts.Query.Get("q") != "callback" {
|
||||
t.Fatalf("executor query = %#v, want callback query", gotOpts.Query)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelSkipsOriginatingPluginInterceptors(t *testing.T) {
|
||||
model := "model-execution-skip-origin-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
skipHost := &modelExecutionSkipHost{}
|
||||
handler.SetPluginHost(skipHost)
|
||||
|
||||
resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: model,
|
||||
Body: requestBody,
|
||||
SkipInterceptorPluginID: "origin-plugin",
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModel() error = %+v", errMsg)
|
||||
}
|
||||
if string(resp.Body) != "model-execution-ok" {
|
||||
t.Fatalf("body = %q, want executor response", resp.Body)
|
||||
}
|
||||
if skipHost.beforeSkip != "origin-plugin" || skipHost.afterSkip != "origin-plugin" || skipHost.respSkip != "origin-plugin" {
|
||||
t.Fatalf("skip ids = before:%q after:%q response:%q, want origin-plugin", skipHost.beforeSkip, skipHost.afterSkip, skipHost.respSkip)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStream(t *testing.T) {
|
||||
model := "model-execution-stream-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("stream-one")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{
|
||||
Headers: http.Header{"X-Upstream": []string{"stream"}},
|
||||
Chunks: chunks,
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{PassthroughHeaders: true})
|
||||
|
||||
stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "claude",
|
||||
Model: model,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
Headers: http.Header{"X-Callback": []string{"stream"}},
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModelStream() error = %+v", errMsg)
|
||||
}
|
||||
if stream.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d", stream.StatusCode, http.StatusOK)
|
||||
}
|
||||
if stream.Headers.Get("X-Upstream") != "stream" {
|
||||
t.Fatalf("headers = %#v, want upstream header", stream.Headers)
|
||||
}
|
||||
chunk, ok := <-stream.Chunks
|
||||
if !ok {
|
||||
t.Fatal("stream chunks closed before payload")
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("stream chunk error = %+v", chunk.Err)
|
||||
}
|
||||
if string(chunk.Payload) != "stream-one" {
|
||||
t.Fatalf("stream chunk payload = %q, want stream-one", chunk.Payload)
|
||||
}
|
||||
if chunk, ok = <-stream.Chunks; ok {
|
||||
t.Fatalf("unexpected extra stream chunk: %+v", chunk)
|
||||
}
|
||||
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != model {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, model)
|
||||
}
|
||||
if string(gotReq.Payload) != string(requestBody) {
|
||||
t.Fatalf("executor payload = %q, want %q", gotReq.Payload, requestBody)
|
||||
}
|
||||
if !gotOpts.Stream {
|
||||
t.Fatal("executor stream option = false, want true")
|
||||
}
|
||||
if gotOpts.SourceFormat != sdktranslator.FormatOpenAI {
|
||||
t.Fatalf("SourceFormat = %q, want %q", gotOpts.SourceFormat, sdktranslator.FormatOpenAI)
|
||||
}
|
||||
if gotOpts.ResponseFormat != sdktranslator.FormatClaude {
|
||||
t.Fatalf("ResponseFormat = %q, want %q", gotOpts.ResponseFormat, sdktranslator.FormatClaude)
|
||||
}
|
||||
if gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey] != model {
|
||||
t.Fatalf("requested model metadata = %#v, want %q", gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey], model)
|
||||
}
|
||||
if gotOpts.Metadata[modelExecutionMetadataSourceKey] != modelExecutionInternalSource {
|
||||
t.Fatalf("source metadata = %#v, want %q", gotOpts.Metadata[modelExecutionMetadataSourceKey], modelExecutionInternalSource)
|
||||
}
|
||||
if gotOpts.Headers.Get("X-Callback") != "stream" {
|
||||
t.Fatalf("executor headers = %#v, want callback header", gotOpts.Headers)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStreamSkipsOriginatingPluginInterceptors(t *testing.T) {
|
||||
model := "model-execution-stream-skip-origin-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("stream-one")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
skipHost := &modelExecutionSkipHost{}
|
||||
handler.SetPluginHost(skipHost)
|
||||
|
||||
stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "openai",
|
||||
Model: model,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
SkipInterceptorPluginID: "origin-plugin",
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModelStream() error = %+v", errMsg)
|
||||
}
|
||||
chunk, ok := <-stream.Chunks
|
||||
if !ok {
|
||||
t.Fatal("stream chunks closed before payload")
|
||||
}
|
||||
if string(chunk.Payload) != "stream-one" {
|
||||
t.Fatalf("stream chunk payload = %q, want stream-one", chunk.Payload)
|
||||
}
|
||||
if skipHost.beforeSkip != "origin-plugin" || skipHost.afterSkip != "origin-plugin" {
|
||||
t.Fatalf("request skip ids = before:%q after:%q, want origin-plugin", skipHost.beforeSkip, skipHost.afterSkip)
|
||||
}
|
||||
if len(skipHost.streamSkip) == 0 {
|
||||
t.Fatal("stream interceptor was not called with skip")
|
||||
}
|
||||
for _, skipID := range skipHost.streamSkip {
|
||||
if skipID != "origin-plugin" {
|
||||
t.Fatalf("stream skip id = %q, want origin-plugin", skipID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStreamStartupError(t *testing.T) {
|
||||
model := "model-execution-stream-startup-error-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model))
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Err: fmt.Errorf("startup failed")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
|
||||
stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "claude",
|
||||
Model: model,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg == nil {
|
||||
t.Fatal("ExecuteModelStream() error = nil, want startup error")
|
||||
}
|
||||
if errMsg.StatusCode != http.StatusInternalServerError {
|
||||
t.Fatalf("status = %d, want %d", errMsg.StatusCode, http.StatusInternalServerError)
|
||||
}
|
||||
if errMsg.Error == nil || errMsg.Error.Error() != "startup failed" {
|
||||
t.Fatalf("error = %v, want startup failed", errMsg.Error)
|
||||
}
|
||||
if stream.Chunks != nil {
|
||||
t.Fatal("stream chunks created for startup error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStreamTerminalError(t *testing.T) {
|
||||
model := "model-execution-stream-terminal-error-model"
|
||||
requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model))
|
||||
errorHeaders := http.Header{"X-Stream-Error": []string{"terminal"}}
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("stream-before-error")}
|
||||
chunks <- coreexecutor.StreamChunk{Err: modelExecutionStatusHeaderError{
|
||||
statusCode: http.StatusTooManyRequests,
|
||||
message: "rate limited",
|
||||
headers: errorHeaders,
|
||||
}}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
|
||||
stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: "openai",
|
||||
ExitProtocol: "claude",
|
||||
Model: model,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModelStream() error = %+v", errMsg)
|
||||
}
|
||||
|
||||
chunk, ok := <-stream.Chunks
|
||||
if !ok {
|
||||
t.Fatal("stream chunks closed before payload")
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("first stream chunk error = %+v", chunk.Err)
|
||||
}
|
||||
if string(chunk.Payload) != "stream-before-error" {
|
||||
t.Fatalf("first stream chunk payload = %q, want stream-before-error", chunk.Payload)
|
||||
}
|
||||
|
||||
chunk, ok = <-stream.Chunks
|
||||
if !ok {
|
||||
t.Fatal("stream chunks closed before terminal error")
|
||||
}
|
||||
if len(chunk.Payload) != 0 {
|
||||
t.Fatalf("terminal stream chunk payload = %q, want empty", chunk.Payload)
|
||||
}
|
||||
if chunk.Err == nil {
|
||||
t.Fatal("terminal stream chunk error = nil")
|
||||
}
|
||||
if chunk.Err.StatusCode != http.StatusTooManyRequests {
|
||||
t.Fatalf("terminal status = %d, want %d", chunk.Err.StatusCode, http.StatusTooManyRequests)
|
||||
}
|
||||
if chunk.Err.Message != "rate limited" {
|
||||
t.Fatalf("terminal message = %q, want rate limited", chunk.Err.Message)
|
||||
}
|
||||
if chunk.Err.Error() != "rate limited" {
|
||||
t.Fatalf("terminal Error() = %q, want rate limited", chunk.Err.Error())
|
||||
}
|
||||
if chunk.Err.Headers.Get("X-Stream-Error") != "terminal" {
|
||||
t.Fatalf("terminal headers = %#v, want stream error header", chunk.Err.Headers)
|
||||
}
|
||||
if chunk, ok = <-stream.Chunks; ok {
|
||||
t.Fatalf("unexpected extra stream chunk: %+v", chunk)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStreamContextCancel(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
dataChan := make(chan []byte)
|
||||
errChan := make(chan *interfaces.ErrorMessage)
|
||||
chunks := wrapModelExecutionChunks(ctx, dataChan, errChan, nil)
|
||||
|
||||
cancel()
|
||||
|
||||
timeout := time.NewTimer(time.Second)
|
||||
defer timeout.Stop()
|
||||
select {
|
||||
case chunk, ok := <-chunks:
|
||||
if ok {
|
||||
t.Fatalf("stream chunks yielded after cancel: %+v", chunk)
|
||||
}
|
||||
case <-timeout.C:
|
||||
t.Fatal("stream chunks did not close after context cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteProtocolWithAuthManagerUsesForcedProvider(t *testing.T) {
|
||||
model := "interactions-agent-target"
|
||||
requestBody := []byte(`{"agent":"agents/test-agent","input":"hi"}`)
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
provider: "gemini",
|
||||
execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{Payload: []byte(`{"id":"interaction_1"}`)}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
|
||||
resp, errMsg := handler.ExecuteProtocolWithAuthManager(context.Background(), ProtocolExecutionRequest{
|
||||
EntryProtocol: "interactions",
|
||||
ExitProtocol: "interactions",
|
||||
ForcedProvider: "gemini",
|
||||
Model: model,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteProtocolWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(resp.Body) != `{"id":"interaction_1"}` {
|
||||
t.Fatalf("body = %q, want native interactions response", resp.Body)
|
||||
}
|
||||
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != model {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, model)
|
||||
}
|
||||
if gotOpts.SourceFormat != sdktranslator.FormatInteractions {
|
||||
t.Fatalf("SourceFormat = %q, want %q", gotOpts.SourceFormat, sdktranslator.FormatInteractions)
|
||||
}
|
||||
if gotOpts.ResponseFormat != sdktranslator.FormatInteractions {
|
||||
t.Fatalf("ResponseFormat = %q, want %q", gotOpts.ResponseFormat, sdktranslator.FormatInteractions)
|
||||
}
|
||||
if gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey] != model {
|
||||
t.Fatalf("requested model metadata = %#v, want %q", gotOpts.Metadata[coreexecutor.RequestedModelMetadataKey], model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreferExecutionProviderMovesPreferredFirst(t *testing.T) {
|
||||
providers := preferExecutionProvider([]string{"gemini", "gemini-interactions", "claude"}, "gemini-interactions")
|
||||
want := []string{"gemini-interactions", "gemini", "claude"}
|
||||
if len(providers) != len(want) {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
for i := range want {
|
||||
if providers[i] != want[i] {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdjustExecutionProvidersExcludesInteractionsProviderForUnsupportedEntry(t *testing.T) {
|
||||
providers := adjustExecutionProvidersForEntryProtocol("codex", []string{"gemini-interactions", "codex"})
|
||||
want := []string{"codex"}
|
||||
if len(providers) != len(want) {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
for i := range want {
|
||||
if providers[i] != want[i] {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdjustExecutionProvidersKeepsInteractionsProviderForSupportedNativeInteractionsEntries(t *testing.T) {
|
||||
for _, entryProtocol := range []string{constant.OpenAI, constant.OpenaiResponse, constant.Claude, constant.Gemini} {
|
||||
t.Run(entryProtocol, func(t *testing.T) {
|
||||
providers := adjustExecutionProvidersForEntryProtocol(entryProtocol, []string{"gemini-interactions"})
|
||||
want := []string{"gemini-interactions"}
|
||||
if len(providers) != len(want) {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
for i := range want {
|
||||
if providers[i] != want[i] {
|
||||
t.Fatalf("providers = %#v, want %#v", providers, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteModelStreamKeepsInteractionsProviderForOpenAIEntry(t *testing.T) {
|
||||
model := "gemini-3.1-flash-lite"
|
||||
requestBody := []byte(`{"model":"gemini-3.1-flash-lite","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
provider: constant.GeminiInteractions,
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"id":"chunk_1","object":"chat.completion.chunk","choices":[]}`)}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
},
|
||||
}
|
||||
handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{})
|
||||
|
||||
stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{
|
||||
EntryProtocol: constant.OpenAI,
|
||||
ExitProtocol: constant.OpenAI,
|
||||
Model: model,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteModelStream() error = %+v", errMsg)
|
||||
}
|
||||
for range stream.Chunks {
|
||||
}
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != model {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, model)
|
||||
}
|
||||
if gotOpts.SourceFormat != sdktranslator.FormatOpenAI {
|
||||
t.Fatalf("SourceFormat = %q, want %q", gotOpts.SourceFormat, sdktranslator.FormatOpenAI)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteProtocolWithAuthManagerAgentUsesSelectionModelForAuth(t *testing.T) {
|
||||
selectionModel := "gemini-2.5-flash"
|
||||
agentModel := "agents/test-agent"
|
||||
requestBody := []byte(`{"agent":"agents/test-agent","input":"hi"}`)
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
provider: constant.GeminiInteractions,
|
||||
execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{Payload: []byte(`{"id":"interaction_1"}`)}, nil
|
||||
},
|
||||
}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "model-execution-agent-selection",
|
||||
Provider: constant.GeminiInteractions,
|
||||
Status: coreauth.StatusActive,
|
||||
Metadata: map[string]any{"email": "agent-selection@example.com"},
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: selectionModel}, {ID: agentModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
manager.RefreshSchedulerEntry(auth.ID)
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
|
||||
resp, errMsg := handler.ExecuteProtocolWithAuthManager(context.Background(), ProtocolExecutionRequest{
|
||||
EntryProtocol: "interactions",
|
||||
ExitProtocol: "interactions",
|
||||
ForcedProvider: constant.GeminiInteractions,
|
||||
AuthSelectionModel: selectionModel,
|
||||
Model: agentModel,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteProtocolWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
if string(resp.Body) != `{"id":"interaction_1"}` {
|
||||
t.Fatalf("body = %q, want native interactions response", resp.Body)
|
||||
}
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != agentModel {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, agentModel)
|
||||
}
|
||||
if string(gotReq.Payload) != string(requestBody) {
|
||||
t.Fatalf("executor payload = %q, want %q", gotReq.Payload, requestBody)
|
||||
}
|
||||
if gotOpts.Metadata[coreexecutor.AuthSelectionModelMetadataKey] != selectionModel {
|
||||
t.Fatalf("auth selection metadata = %#v, want %q", gotOpts.Metadata[coreexecutor.AuthSelectionModelMetadataKey], selectionModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteProtocolStreamWithAuthManagerAgentUsesSelectionModelForAuth(t *testing.T) {
|
||||
selectionModel := "gemini-2.5-flash"
|
||||
agentModel := "agents/test-agent"
|
||||
requestBody := []byte(`{"agent":"agents/test-agent","input":"hi","stream":true}`)
|
||||
executor := &modelExecutionCaptureExecutor{
|
||||
provider: constant.GeminiInteractions,
|
||||
stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"id":"interaction_1"}`)}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
},
|
||||
}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{
|
||||
ID: "model-execution-agent-stream-selection",
|
||||
Provider: constant.GeminiInteractions,
|
||||
Status: coreauth.StatusActive,
|
||||
Metadata: map[string]any{"email": "agent-stream-selection@example.com"},
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: selectionModel}, {ID: agentModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
manager.RefreshSchedulerEntry(auth.ID)
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
|
||||
stream, errMsg := handler.ExecuteProtocolStreamWithAuthManager(context.Background(), ProtocolExecutionRequest{
|
||||
EntryProtocol: "interactions",
|
||||
ExitProtocol: "interactions",
|
||||
ForcedProvider: constant.GeminiInteractions,
|
||||
AuthSelectionModel: selectionModel,
|
||||
Model: agentModel,
|
||||
Stream: true,
|
||||
Body: requestBody,
|
||||
})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("ExecuteProtocolStreamWithAuthManager() error = %+v", errMsg)
|
||||
}
|
||||
chunk, ok := <-stream.Chunks
|
||||
if !ok {
|
||||
t.Fatal("stream chunks closed before payload")
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("stream chunk error = %+v", chunk.Err)
|
||||
}
|
||||
if string(chunk.Payload) != `{"id":"interaction_1"}` {
|
||||
t.Fatalf("stream chunk payload = %q, want native interactions response", chunk.Payload)
|
||||
}
|
||||
gotReq, gotOpts := executor.captured()
|
||||
if gotReq.Model != agentModel {
|
||||
t.Fatalf("executor model = %q, want %q", gotReq.Model, agentModel)
|
||||
}
|
||||
if string(gotReq.Payload) != string(requestBody) {
|
||||
t.Fatalf("executor payload = %q, want %q", gotReq.Payload, requestBody)
|
||||
}
|
||||
if gotOpts.Metadata[coreexecutor.AuthSelectionModelMetadataKey] != selectionModel {
|
||||
t.Fatalf("auth selection metadata = %#v, want %q", gotOpts.Metadata[coreexecutor.AuthSelectionModelMetadataKey], selectionModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersForExecutionForcedGeminiRejectsRouterProvider(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
decision := modelRouteDecision{Provider: "claude", Model: "claude-sonnet-4"}
|
||||
_, _, errMsg := handler.providersForExecution("agents/test-agent", "agents/test-agent", false, decision, modelExecutionOptions{ForcedProvider: "gemini"})
|
||||
if errMsg == nil {
|
||||
t.Fatal("providersForExecution() error = nil, want native interactions error")
|
||||
}
|
||||
if errMsg.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", errMsg.StatusCode, http.StatusBadRequest)
|
||||
}
|
||||
if errMsg.Error == nil || !strings.Contains(errMsg.Error.Error(), "native interactions") {
|
||||
t.Fatalf("error = %v, want native interactions message", errMsg.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersForExecutionForcedGeminiUsesGeminiProvider(t *testing.T) {
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
providers, model, errMsg := handler.providersForExecution("agents/test-agent", "agents/test-agent", false, modelRouteDecision{}, modelExecutionOptions{ForcedProvider: "gemini"})
|
||||
if errMsg != nil {
|
||||
t.Fatalf("providersForExecution() error = %+v", errMsg)
|
||||
}
|
||||
if len(providers) != 1 || providers[0] != "gemini" {
|
||||
t.Fatalf("providers = %#v, want [gemini]", providers)
|
||||
}
|
||||
if model != "agents/test-agent" {
|
||||
t.Fatalf("model = %q, want agents/test-agent", model)
|
||||
}
|
||||
}
|
||||
22
backend/sdk/api/handlers/openai/codex_client_models.go
Normal file
22
backend/sdk/api/handlers/openai/codex_client_models.go
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
codexmodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/models"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
)
|
||||
|
||||
func (h *OpenAIAPIHandler) codexClientModelsResponse() map[string]any {
|
||||
optimizeMultiAgentV2 := h != nil && h.Cfg != nil && h.Cfg.CodexOptimizeMultiAgentV2
|
||||
return codexmodels.BuildResponse(h.Models(), registry.GetGlobalRegistry().GetModelProviders, optimizeMultiAgentV2)
|
||||
}
|
||||
|
||||
// CodexClientModelsResponse builds a Codex client model response.
|
||||
func CodexClientModelsResponse(models []map[string]any) map[string]any {
|
||||
return codexmodels.BuildResponse(models, nil, false)
|
||||
}
|
||||
|
||||
// CodexClientModelsResponseWithMultiAgentV2 builds a Codex client model response
|
||||
// and advertises multi-agent v2 for synthesized models when enabled.
|
||||
func CodexClientModelsResponseWithMultiAgentV2(models []map[string]any, enabled bool) map[string]any {
|
||||
return codexmodels.BuildResponse(models, nil, enabled)
|
||||
}
|
||||
59
backend/sdk/api/handlers/openai/codex_client_models_test.go
Normal file
59
backend/sdk/api/handlers/openai/codex_client_models_test.go
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
)
|
||||
|
||||
func TestCodexClientModelsResponseMultiAgentV2FollowsConfig(t *testing.T) {
|
||||
modelID := "codex-client-multi-agent-v2-test"
|
||||
clientID := "codex-client-multi-agent-v2-test-client"
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
modelRegistry.RegisterClient(clientID, "openai-compatibility", []*registry.ModelInfo{{ID: modelID}})
|
||||
t.Cleanup(func() {
|
||||
modelRegistry.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&config.SDKConfig{}, nil)
|
||||
handler := NewOpenAIAPIHandler(base)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
enabled bool
|
||||
}{
|
||||
{name: "disabled", enabled: false},
|
||||
{name: "enabled", enabled: true},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
base.Cfg.CodexOptimizeMultiAgentV2 = tt.enabled
|
||||
response := handler.codexClientModelsResponse()
|
||||
models, ok := response["models"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("models type = %T, want []map[string]any", response["models"])
|
||||
}
|
||||
var entry map[string]any
|
||||
for _, model := range models {
|
||||
slug, _ := model["slug"].(string)
|
||||
if slug == modelID {
|
||||
entry = model
|
||||
break
|
||||
}
|
||||
}
|
||||
if entry == nil {
|
||||
t.Fatalf("missing synthesized model %q", modelID)
|
||||
}
|
||||
value, exists := entry["multi_agent_version"]
|
||||
if tt.enabled {
|
||||
if !exists || value != "v2" {
|
||||
t.Fatalf("multi_agent_version = %#v, want v2", value)
|
||||
}
|
||||
return
|
||||
}
|
||||
if !exists || value != nil {
|
||||
t.Fatalf("multi_agent_version = %#v, want preserved null", value)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
709
backend/sdk/api/handlers/openai/openai_handlers.go
Normal file
709
backend/sdk/api/handlers/openai/openai_handlers.go
Normal file
|
|
@ -0,0 +1,709 @@
|
|||
// Package openai provides HTTP handlers for OpenAI API endpoints.
|
||||
// This package implements the OpenAI-compatible API interface, including model listing
|
||||
// and chat completion functionality. It supports both streaming and non-streaming responses,
|
||||
// and manages a pool of clients to interact with backend services.
|
||||
// The handlers translate OpenAI API requests to the appropriate backend format and
|
||||
// convert responses back to OpenAI-compatible format.
|
||||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
responsesconverter "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/openai/openai/responses"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// OpenAIAPIHandler contains the handlers for OpenAI API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service.
|
||||
type OpenAIAPIHandler struct {
|
||||
*handlers.BaseAPIHandler
|
||||
}
|
||||
|
||||
// NewOpenAIAPIHandler creates a new OpenAI API handlers instance.
|
||||
// It takes an BaseAPIHandler instance as input and returns an OpenAIAPIHandler.
|
||||
//
|
||||
// Parameters:
|
||||
// - apiHandlers: The base API handlers instance
|
||||
//
|
||||
// Returns:
|
||||
// - *OpenAIAPIHandler: A new OpenAI API handlers instance
|
||||
func NewOpenAIAPIHandler(apiHandlers *handlers.BaseAPIHandler) *OpenAIAPIHandler {
|
||||
return &OpenAIAPIHandler{
|
||||
BaseAPIHandler: apiHandlers,
|
||||
}
|
||||
}
|
||||
|
||||
// HandlerType returns the identifier for this handler implementation.
|
||||
func (h *OpenAIAPIHandler) HandlerType() string {
|
||||
return OpenAI
|
||||
}
|
||||
|
||||
// Models returns the OpenAI-compatible model metadata supported by this handler.
|
||||
func (h *OpenAIAPIHandler) Models() []map[string]any {
|
||||
// Get dynamic models from the global registry
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
return modelRegistry.GetAvailableModels("openai")
|
||||
}
|
||||
|
||||
// OpenAIModels handles the /v1/models endpoint.
|
||||
// It returns a list of available AI models with their capabilities
|
||||
// and specifications in OpenAI-compatible format.
|
||||
func (h *OpenAIAPIHandler) OpenAIModels(c *gin.Context) {
|
||||
if _, ok := c.Request.URL.Query()["client_version"]; ok {
|
||||
c.JSON(http.StatusOK, h.codexClientModelsResponse())
|
||||
return
|
||||
}
|
||||
|
||||
// Get all available models
|
||||
allModels := h.Models()
|
||||
|
||||
// Filter to only include the 4 required fields: id, object, created, owned_by
|
||||
filteredModels := make([]map[string]any, len(allModels))
|
||||
for i, model := range allModels {
|
||||
filteredModel := map[string]any{
|
||||
"id": model["id"],
|
||||
"object": model["object"],
|
||||
}
|
||||
|
||||
// Add created field if it exists
|
||||
if created, exists := model["created"]; exists {
|
||||
filteredModel["created"] = created
|
||||
}
|
||||
|
||||
// Add owned_by field if it exists
|
||||
if ownedBy, exists := model["owned_by"]; exists {
|
||||
filteredModel["owned_by"] = ownedBy
|
||||
}
|
||||
|
||||
filteredModels[i] = filteredModel
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": filteredModels,
|
||||
})
|
||||
}
|
||||
|
||||
// ChatCompletions handles the /v1/chat/completions endpoint.
|
||||
// It determines whether the request is for a streaming or non-streaming response
|
||||
// and calls the appropriate handler based on the model provider.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
func (h *OpenAIAPIHandler) ChatCompletions(c *gin.Context) {
|
||||
rawJSON, err := handlers.ReadRequestBody(c)
|
||||
// If data retrieval fails, return a 400 Bad Request error.
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the client requested a streaming response.
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
stream := streamResult.Type == gjson.True
|
||||
|
||||
// Some clients send OpenAI Responses-format payloads to /v1/chat/completions.
|
||||
// Convert them to Chat Completions so downstream translators preserve tool metadata.
|
||||
if shouldTreatAsResponsesFormat(rawJSON) {
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
rawJSON = responsesconverter.ConvertOpenAIResponsesRequestToOpenAIChatCompletions(modelName, rawJSON, stream)
|
||||
stream = gjson.GetBytes(rawJSON, "stream").Bool()
|
||||
}
|
||||
|
||||
if stream {
|
||||
h.handleStreamingResponse(c, rawJSON)
|
||||
} else {
|
||||
h.handleNonStreamingResponse(c, rawJSON)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// shouldTreatAsResponsesFormat detects OpenAI Responses-style payloads that are
|
||||
// accidentally sent to the Chat Completions endpoint.
|
||||
func shouldTreatAsResponsesFormat(rawJSON []byte) bool {
|
||||
if gjson.GetBytes(rawJSON, "messages").Exists() {
|
||||
return false
|
||||
}
|
||||
if gjson.GetBytes(rawJSON, "input").Exists() {
|
||||
return true
|
||||
}
|
||||
if gjson.GetBytes(rawJSON, "instructions").Exists() {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Completions handles the /v1/completions endpoint.
|
||||
// It determines whether the request is for a streaming or non-streaming response
|
||||
// and calls the appropriate handler based on the model provider.
|
||||
// This endpoint follows the OpenAI completions API specification.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
func (h *OpenAIAPIHandler) Completions(c *gin.Context) {
|
||||
rawJSON, err := handlers.ReadRequestBody(c)
|
||||
// If data retrieval fails, return a 400 Bad Request error.
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the client requested a streaming response.
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
if streamResult.Type == gjson.True {
|
||||
h.handleCompletionsStreamingResponse(c, rawJSON)
|
||||
} else {
|
||||
h.handleCompletionsNonStreamingResponse(c, rawJSON)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// convertCompletionsRequestToChatCompletions converts OpenAI completions API request to chat completions format.
|
||||
// This allows the completions endpoint to use the existing chat completions infrastructure.
|
||||
//
|
||||
// Parameters:
|
||||
// - rawJSON: The raw JSON bytes of the completions request
|
||||
//
|
||||
// Returns:
|
||||
// - []byte: The converted chat completions request
|
||||
func convertCompletionsRequestToChatCompletions(rawJSON []byte) []byte {
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
// Extract prompt from completions request
|
||||
prompt := root.Get("prompt").String()
|
||||
if prompt == "" {
|
||||
prompt = "Complete this:"
|
||||
}
|
||||
|
||||
// Create chat completions structure
|
||||
out := []byte(`{"model":"","messages":[{"role":"user","content":""}]}`)
|
||||
|
||||
// Set model
|
||||
if model := root.Get("model"); model.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "model", model.String())
|
||||
}
|
||||
|
||||
// Set the prompt as user message content
|
||||
out, _ = sjson.SetBytes(out, "messages.0.content", prompt)
|
||||
|
||||
// Copy other parameters from completions to chat completions
|
||||
if maxTokens := root.Get("max_tokens"); maxTokens.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "max_tokens", maxTokens.Int())
|
||||
}
|
||||
|
||||
if temperature := root.Get("temperature"); temperature.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "temperature", temperature.Float())
|
||||
}
|
||||
|
||||
if topP := root.Get("top_p"); topP.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "top_p", topP.Float())
|
||||
}
|
||||
|
||||
if frequencyPenalty := root.Get("frequency_penalty"); frequencyPenalty.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "frequency_penalty", frequencyPenalty.Float())
|
||||
}
|
||||
|
||||
if presencePenalty := root.Get("presence_penalty"); presencePenalty.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "presence_penalty", presencePenalty.Float())
|
||||
}
|
||||
|
||||
if stop := root.Get("stop"); stop.Exists() {
|
||||
out, _ = sjson.SetRawBytes(out, "stop", []byte(stop.Raw))
|
||||
}
|
||||
|
||||
if stream := root.Get("stream"); stream.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "stream", stream.Bool())
|
||||
}
|
||||
|
||||
if logprobs := root.Get("logprobs"); logprobs.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "logprobs", logprobs.Bool())
|
||||
}
|
||||
|
||||
if topLogprobs := root.Get("top_logprobs"); topLogprobs.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "top_logprobs", topLogprobs.Int())
|
||||
}
|
||||
|
||||
if echo := root.Get("echo"); echo.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "echo", echo.Bool())
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// convertChatCompletionsResponseToCompletions converts chat completions API response back to completions format.
|
||||
// This ensures the completions endpoint returns data in the expected format.
|
||||
//
|
||||
// Parameters:
|
||||
// - rawJSON: The raw JSON bytes of the chat completions response
|
||||
//
|
||||
// Returns:
|
||||
// - []byte: The converted completions response
|
||||
func convertChatCompletionsResponseToCompletions(rawJSON []byte) []byte {
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
// Base completions response structure
|
||||
out := []byte(`{"id":"","object":"text_completion","created":0,"model":"","choices":[]}`)
|
||||
|
||||
// Copy basic fields
|
||||
if id := root.Get("id"); id.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "id", id.String())
|
||||
}
|
||||
|
||||
if created := root.Get("created"); created.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "created", created.Int())
|
||||
}
|
||||
|
||||
if model := root.Get("model"); model.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "model", model.String())
|
||||
}
|
||||
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
out, _ = sjson.SetRawBytes(out, "usage", []byte(usage.Raw))
|
||||
}
|
||||
|
||||
// Convert choices from chat completions to completions format
|
||||
var choices []interface{}
|
||||
if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
|
||||
chatChoices.ForEach(func(_, choice gjson.Result) bool {
|
||||
completionsChoice := map[string]interface{}{
|
||||
"index": choice.Get("index").Int(),
|
||||
}
|
||||
|
||||
// Extract text content from message.content
|
||||
if message := choice.Get("message"); message.Exists() {
|
||||
if content := message.Get("content"); content.Exists() {
|
||||
completionsChoice["text"] = content.String()
|
||||
}
|
||||
} else if delta := choice.Get("delta"); delta.Exists() {
|
||||
// For streaming responses, use delta.content
|
||||
if content := delta.Get("content"); content.Exists() {
|
||||
completionsChoice["text"] = content.String()
|
||||
}
|
||||
}
|
||||
|
||||
// Copy finish_reason
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() {
|
||||
completionsChoice["finish_reason"] = finishReason.String()
|
||||
}
|
||||
|
||||
// Copy logprobs if present
|
||||
if logprobs := choice.Get("logprobs"); logprobs.Exists() {
|
||||
completionsChoice["logprobs"] = logprobs.Value()
|
||||
}
|
||||
|
||||
choices = append(choices, completionsChoice)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
if len(choices) > 0 {
|
||||
choicesJSON, _ := json.Marshal(choices)
|
||||
out, _ = sjson.SetRawBytes(out, "choices", choicesJSON)
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// convertChatCompletionsStreamChunkToCompletions converts a streaming chat completions chunk to completions format.
|
||||
// This handles the real-time conversion of streaming response chunks and filters out empty text responses.
|
||||
//
|
||||
// Parameters:
|
||||
// - chunkData: The raw JSON bytes of a single chat completions stream chunk
|
||||
//
|
||||
// Returns:
|
||||
// - []byte: The converted completions stream chunk, or nil if should be filtered out
|
||||
func convertChatCompletionsStreamChunkToCompletions(chunkData []byte) []byte {
|
||||
root := gjson.ParseBytes(chunkData)
|
||||
|
||||
// Check if this chunk has any meaningful content
|
||||
hasContent := false
|
||||
hasUsage := root.Get("usage").Exists()
|
||||
if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
|
||||
chatChoices.ForEach(func(_, choice gjson.Result) bool {
|
||||
// Check if delta has content or finish_reason
|
||||
if delta := choice.Get("delta"); delta.Exists() {
|
||||
if content := delta.Get("content"); content.Exists() && content.String() != "" {
|
||||
hasContent = true
|
||||
return false // Break out of forEach
|
||||
}
|
||||
}
|
||||
// Also check for finish_reason to ensure we don't skip final chunks
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() && finishReason.String() != "" && finishReason.String() != "null" {
|
||||
hasContent = true
|
||||
return false // Break out of forEach
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// If no meaningful content and no usage, return nil to indicate this chunk should be skipped
|
||||
if !hasContent && !hasUsage {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Base completions stream response structure
|
||||
out := []byte(`{"id":"","object":"text_completion","created":0,"model":"","choices":[]}`)
|
||||
|
||||
// Copy basic fields
|
||||
if id := root.Get("id"); id.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "id", id.String())
|
||||
}
|
||||
|
||||
if created := root.Get("created"); created.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "created", created.Int())
|
||||
}
|
||||
|
||||
if model := root.Get("model"); model.Exists() {
|
||||
out, _ = sjson.SetBytes(out, "model", model.String())
|
||||
}
|
||||
|
||||
// Convert choices from chat completions delta to completions format
|
||||
var choices []interface{}
|
||||
if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
|
||||
chatChoices.ForEach(func(_, choice gjson.Result) bool {
|
||||
completionsChoice := map[string]interface{}{
|
||||
"index": choice.Get("index").Int(),
|
||||
}
|
||||
|
||||
// Extract text content from delta.content
|
||||
if delta := choice.Get("delta"); delta.Exists() {
|
||||
if content := delta.Get("content"); content.Exists() && content.String() != "" {
|
||||
completionsChoice["text"] = content.String()
|
||||
} else {
|
||||
completionsChoice["text"] = ""
|
||||
}
|
||||
} else {
|
||||
completionsChoice["text"] = ""
|
||||
}
|
||||
|
||||
// Copy finish_reason
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() && finishReason.String() != "null" {
|
||||
completionsChoice["finish_reason"] = finishReason.String()
|
||||
}
|
||||
|
||||
// Copy logprobs if present
|
||||
if logprobs := choice.Get("logprobs"); logprobs.Exists() {
|
||||
completionsChoice["logprobs"] = logprobs.Value()
|
||||
}
|
||||
|
||||
choices = append(choices, completionsChoice)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
if len(choices) > 0 {
|
||||
choicesJSON, _ := json.Marshal(choices)
|
||||
out, _ = sjson.SetRawBytes(out, "choices", choicesJSON)
|
||||
}
|
||||
|
||||
// Copy usage if present
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
out, _ = sjson.SetRawBytes(out, "usage", []byte(usage.Raw))
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// handleNonStreamingResponse handles non-streaming chat completion responses
|
||||
// for Gemini models. It selects a client from the pool, sends the request, and
|
||||
// aggregates the response before sending it back to the client in OpenAI format.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible request
|
||||
func (h *OpenAIAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c))
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleStreamingResponse handles streaming responses for Gemini models.
|
||||
// It establishes a streaming connection with the backend service and forwards
|
||||
// the response chunks to the client in real-time using Server-Sent Events.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible request
|
||||
func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c))
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk to determine success or failure before setting headers
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg, hasPendingError := handlers.PendingStreamError(errChan); hasPendingError {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Stream closed without data? Send DONE or just headers.
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Commit to streaming headers.
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
|
||||
flusher.Flush()
|
||||
|
||||
// Continue streaming the rest
|
||||
h.handleStreamResult(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleCompletionsNonStreamingResponse handles non-streaming completions responses.
|
||||
// It converts completions request to chat completions format, sends to backend,
|
||||
// then converts the response back to completions format before sending to client.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible completions request
|
||||
func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
|
||||
// Convert completions request to chat completions format
|
||||
chatCompletionsJSON := convertCompletionsRequestToChatCompletions(rawJSON)
|
||||
|
||||
modelName := gjson.GetBytes(chatCompletionsJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
completionsResp := convertChatCompletionsResponseToCompletions(resp)
|
||||
_, _ = c.Writer.Write(completionsResp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleCompletionsStreamingResponse handles streaming completions responses.
|
||||
// It converts completions request to chat completions format, streams from backend,
|
||||
// then converts each response chunk back to completions format before sending to client.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible completions request
|
||||
func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Convert completions request to chat completions format
|
||||
chatCompletionsJSON := convertCompletionsRequestToChatCompletions(rawJSON)
|
||||
|
||||
modelName := gjson.GetBytes(chatCompletionsJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg, hasPendingError := handlers.PendingStreamError(errChan); hasPendingError {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
|
||||
// Write the first chunk
|
||||
converted := convertChatCompletionsStreamChunkToCompletions(chunk)
|
||||
if converted != nil {
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(converted))
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
var doneOnce sync.Once
|
||||
stop := func() { doneOnce.Do(func() { close(done) }) }
|
||||
|
||||
convertedChan := make(chan []byte)
|
||||
go func() {
|
||||
defer close(convertedChan)
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
converted := convertChatCompletionsStreamChunkToCompletions(chunk)
|
||||
if converted == nil {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case convertedChan <- converted:
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
h.handleStreamResult(c, flusher, func(err error) {
|
||||
stop()
|
||||
cliCancel(err)
|
||||
}, convertedChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
func (h *OpenAIAPIHandler) handleStreamResult(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(body))
|
||||
},
|
||||
WriteDone: func() {
|
||||
_, _ = fmt.Fprint(c.Writer, "data: [DONE]\n\n")
|
||||
},
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,103 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
const (
|
||||
initialFailureChatModel = "initial-failure-chat-model"
|
||||
)
|
||||
|
||||
type initialFailureStreamExecutor struct{}
|
||||
|
||||
func (*initialFailureStreamExecutor) Identifier() string { return "initial-failure-stream-executor" }
|
||||
|
||||
func (*initialFailureStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed before first payload")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (*initialFailureStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (*initialFailureStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func runOpenAIStreamErrorTest(t *testing.T, endpoint string, body string) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
executor := &initialFailureStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
authID := fmt.Sprintf("initial-failure-auth-%s-%d", strings.ReplaceAll(endpoint, "/", "-"), idx)
|
||||
auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Errorf("register auth %d: %v", idx, errRegister)
|
||||
return
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: initialFailureChatModel}})
|
||||
defer registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIAPIHandler(base)
|
||||
router := gin.New()
|
||||
if endpoint == "/v1/chat/completions" {
|
||||
router.POST(endpoint, h.ChatCompletions)
|
||||
} else {
|
||||
router.POST(endpoint, h.Completions)
|
||||
}
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, endpoint, strings.NewReader(body))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Errorf("[%s] request %d lost the buffered initial error and returned HTTP 200: %q", endpoint, idx, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "upstream failed before first payload") {
|
||||
t.Errorf("[%s] request %d lost the initial upstream error: status=%d body=%q", endpoint, idx, recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestChatCompletionsHandlerDoesNotLoseErrorBeforeFirstPayload(t *testing.T) {
|
||||
runOpenAIStreamErrorTest(t, "/v1/chat/completions", `{"model":"initial-failure-chat-model","messages":[{"role":"user","content":"hi"}],"stream":true}`)
|
||||
}
|
||||
|
||||
func TestCompletionsHandlerDoesNotLoseErrorBeforeFirstPayload(t *testing.T) {
|
||||
runOpenAIStreamErrorTest(t, "/v1/completions", `{"model":"initial-failure-chat-model","prompt":"hi","stream":true}`)
|
||||
}
|
||||
2015
backend/sdk/api/handlers/openai/openai_images_handlers.go
Normal file
2015
backend/sdk/api/handlers/openai/openai_images_handlers.go
Normal file
File diff suppressed because it is too large
Load diff
500
backend/sdk/api/handlers/openai/openai_images_handlers_test.go
Normal file
500
backend/sdk/api/handlers/openai/openai_images_handlers_test.go
Normal file
|
|
@ -0,0 +1,500 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/textproto"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func performImagesEndpointRequest(t *testing.T, endpointPath string, contentType string, body io.Reader, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.POST(endpointPath, handler)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, endpointPath, body)
|
||||
if contentType != "" {
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
}
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
return resp
|
||||
}
|
||||
|
||||
func assertUnsupportedImagesModelResponse(t *testing.T, resp *httptest.ResponseRecorder, model string) {
|
||||
t.Helper()
|
||||
|
||||
if resp.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
|
||||
}
|
||||
|
||||
message := gjson.GetBytes(resp.Body.Bytes(), "error.message").String()
|
||||
expectedMessage := "Model " + model + " is not supported on " + imagesGenerationsPath + " or " + imagesEditsPath + ". Use " + gptImage15Model + ", " + defaultImagesToolModel + ", " + defaultXAIImagesModel + ", " + xaiImagesQualityModel + ", " + xaiImages20Model + ", or a configured openai-compatibility image model."
|
||||
if message != expectedMessage {
|
||||
t.Fatalf("error message = %q, want %q", message, expectedMessage)
|
||||
}
|
||||
if errorType := gjson.GetBytes(resp.Body.Bytes(), "error.type").String(); errorType != "invalid_request_error" {
|
||||
t.Fatalf("error type = %q, want invalid_request_error", errorType)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesModelValidationAllowsGPTImageAndXAIModels(t *testing.T) {
|
||||
for _, model := range []string{"gpt-image-1.5", "codex/gpt-image-1.5", "gpt-image-2", "codex/gpt-image-2", "grok-imagine-image", "xai/grok-imagine-image", "grok-imagine-image-quality", "xai/grok-imagine-image-quality", "grok-imagine-image-2.0", "xai/grok-imagine-image-2.0"} {
|
||||
if !isSupportedImagesModel(model) {
|
||||
t.Fatalf("expected %s to be supported", model)
|
||||
}
|
||||
}
|
||||
if isSupportedImagesModel("gpt-5.4-mini") {
|
||||
t.Fatal("expected gpt-5.4-mini to be rejected")
|
||||
}
|
||||
if isSupportedImagesModel("codex/grok-imagine-image") {
|
||||
t.Fatal("expected codex/grok-imagine-image to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesModelValidationAllowsOpenAICompatImageModels(t *testing.T) {
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
clientID := "test-openai-compat-image-model-validation"
|
||||
modelRegistry.RegisterClient(clientID, "openai-compatibility", []*registry.ModelInfo{
|
||||
{ID: "compat-image-model", Object: "model", OwnedBy: "compat", Type: registry.OpenAIImageModelType},
|
||||
{ID: "compat-chat-model", Object: "model", OwnedBy: "compat", Type: "openai-compatibility"},
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
modelRegistry.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
if !isSupportedImagesModel("compat-image-model") {
|
||||
t.Fatal("expected configured openai-compatibility image model to be supported")
|
||||
}
|
||||
if isSupportedImagesModel("compat-chat-model") {
|
||||
t.Fatal("expected non-image openai-compatibility model to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanonicalXAIImagesModelPreservesImage20(t *testing.T) {
|
||||
for _, model := range []string{"grok-imagine-image-2.0", "xai/grok-imagine-image-2.0", "XAI/Grok-Imagine-Image-2.0"} {
|
||||
if got := canonicalXAIImagesModel(model); got != xaiImages20Model {
|
||||
t.Fatalf("canonicalXAIImagesModel(%q) = %q, want %s", model, got, xaiImages20Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildXAIImagesGenerationsRequest(t *testing.T) {
|
||||
rawJSON := []byte(`{"model":"xai/grok-imagine-image-quality","prompt":"abstract art","aspect_ratio":"landscape","resolution":"2k","n":2,"response_format":"url"}`)
|
||||
|
||||
req := buildXAIImagesGenerationsRequest(rawJSON, "xai/grok-imagine-image-quality", "url")
|
||||
|
||||
if got := gjson.GetBytes(req, "model").String(); got != "grok-imagine-image-quality" {
|
||||
t.Fatalf("model = %q, want grok-imagine-image-quality", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "prompt").String(); got != "abstract art" {
|
||||
t.Fatalf("prompt = %q, want abstract art", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "aspect_ratio").String(); got != "16:9" {
|
||||
t.Fatalf("aspect_ratio = %q, want 16:9", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "resolution").String(); got != "2k" {
|
||||
t.Fatalf("resolution = %q, want 2k", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "response_format").String(); got != "url" {
|
||||
t.Fatalf("response_format = %q, want url", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "n").Int(); got != 2 {
|
||||
t.Fatalf("n = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildXAIImagesEditRequest(t *testing.T) {
|
||||
req := buildXAIImagesEditRequest("grok-imagine-image", "edit it", []string{"data:image/png;base64,AA==", "https://example.com/image.png"}, "b64_json", "3:2", "1k", 0)
|
||||
|
||||
if got := gjson.GetBytes(req, "model").String(); got != "grok-imagine-image" {
|
||||
t.Fatalf("model = %q, want grok-imagine-image", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "images.0.type").String(); got != "image_url" {
|
||||
t.Fatalf("images.0.type = %q, want image_url", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "images.0.url").String(); got != "data:image/png;base64,AA==" {
|
||||
t.Fatalf("images.0.url = %q", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "images.1.url").String(); got != "https://example.com/image.png" {
|
||||
t.Fatalf("images.1.url = %q", got)
|
||||
}
|
||||
if gjson.GetBytes(req, "image").Exists() {
|
||||
t.Fatalf("multiple image edits must use images array: %s", string(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildXAIImagesEditRequestSingleImage(t *testing.T) {
|
||||
req := buildXAIImagesEditRequest("grok-imagine-image", "edit it", []string{"https://example.com/image.png"}, "url", "", "", 0)
|
||||
|
||||
if got := gjson.GetBytes(req, "image.type").String(); got != "image_url" {
|
||||
t.Fatalf("image.type = %q, want image_url", got)
|
||||
}
|
||||
if got := gjson.GetBytes(req, "image.url").String(); got != "https://example.com/image.png" {
|
||||
t.Fatalf("image.url = %q", got)
|
||||
}
|
||||
if gjson.GetBytes(req, "images").Exists() {
|
||||
t.Fatalf("single image edit must use image object: %s", string(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAICompatImagesJSONRequestPreservesStreamForStreaming(t *testing.T) {
|
||||
req := buildOpenAICompatImagesJSONRequest([]byte(`{"model":"compat-image","prompt":"draw","stream":false}`), "upstream-image", true)
|
||||
|
||||
if got := gjson.GetBytes(req, "model").String(); got != "upstream-image" {
|
||||
t.Fatalf("model = %q, want upstream-image; body=%s", got, string(req))
|
||||
}
|
||||
if !gjson.GetBytes(req, "stream").Bool() {
|
||||
t.Fatalf("stream flag missing: %s", string(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAICompatImagesJSONRequestDropsStreamForNonStreaming(t *testing.T) {
|
||||
req := buildOpenAICompatImagesJSONRequest([]byte(`{"model":"compat-image","prompt":"draw","stream":true}`), "upstream-image", false)
|
||||
|
||||
if got := gjson.GetBytes(req, "model").String(); got != "upstream-image" {
|
||||
t.Fatalf("model = %q, want upstream-image; body=%s", got, string(req))
|
||||
}
|
||||
if gjson.GetBytes(req, "stream").Exists() {
|
||||
t.Fatalf("stream flag should be removed from non-streaming request: %s", string(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAICompatImagesMultipartRequestPreservesStreamAndFileContentType(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
if errWrite := writer.WriteField("model", "compat-image"); errWrite != nil {
|
||||
t.Fatalf("write model field: %v", errWrite)
|
||||
}
|
||||
if errWrite := writer.WriteField("stream", "false"); errWrite != nil {
|
||||
t.Fatalf("write stream field: %v", errWrite)
|
||||
}
|
||||
if errWrite := writer.WriteField("prompt", "edit"); errWrite != nil {
|
||||
t.Fatalf("write prompt field: %v", errWrite)
|
||||
}
|
||||
header := make(textproto.MIMEHeader)
|
||||
header.Set("Content-Disposition", multipart.FileContentDisposition("image", "image.png"))
|
||||
header.Set("Content-Type", "image/png")
|
||||
part, errCreate := writer.CreatePart(header)
|
||||
if errCreate != nil {
|
||||
t.Fatalf("create image field: %v", errCreate)
|
||||
}
|
||||
if _, errWrite := part.Write([]byte("png-data")); errWrite != nil {
|
||||
t.Fatalf("write image field: %v", errWrite)
|
||||
}
|
||||
if errClose := writer.Close(); errClose != nil {
|
||||
t.Fatalf("close multipart writer: %v", errClose)
|
||||
}
|
||||
|
||||
reader := multipart.NewReader(bytes.NewReader(body.Bytes()), writer.Boundary())
|
||||
form, errRead := reader.ReadForm(32 << 20)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read source form: %v", errRead)
|
||||
}
|
||||
defer func() {
|
||||
if errRemove := form.RemoveAll(); errRemove != nil {
|
||||
t.Fatalf("remove source form files: %v", errRemove)
|
||||
}
|
||||
}()
|
||||
|
||||
out, contentType, errBuild := buildOpenAICompatImagesMultipartRequest(form, "upstream-image", true)
|
||||
if errBuild != nil {
|
||||
t.Fatalf("buildOpenAICompatImagesMultipartRequest error: %v", errBuild)
|
||||
}
|
||||
mediaType, params, errParse := mime.ParseMediaType(contentType)
|
||||
if errParse != nil {
|
||||
t.Fatalf("parse content type: %v", errParse)
|
||||
}
|
||||
if mediaType != "multipart/form-data" {
|
||||
t.Fatalf("media type = %q, want multipart/form-data", mediaType)
|
||||
}
|
||||
rewrittenReader := multipart.NewReader(bytes.NewReader(out), params["boundary"])
|
||||
rewrittenForm, errRead := rewrittenReader.ReadForm(32 << 20)
|
||||
if errRead != nil {
|
||||
t.Fatalf("read rewritten form: %v", errRead)
|
||||
}
|
||||
defer func() {
|
||||
if errRemove := rewrittenForm.RemoveAll(); errRemove != nil {
|
||||
t.Fatalf("remove rewritten form files: %v", errRemove)
|
||||
}
|
||||
}()
|
||||
if got := rewrittenForm.Value["model"]; len(got) != 1 || got[0] != "upstream-image" {
|
||||
t.Fatalf("model values = %#v, want upstream-image", got)
|
||||
}
|
||||
if got := rewrittenForm.Value["stream"]; len(got) != 1 || got[0] != "true" {
|
||||
t.Fatalf("stream values = %#v, want true", got)
|
||||
}
|
||||
if got := rewrittenForm.Value["prompt"]; len(got) != 1 || got[0] != "edit" {
|
||||
t.Fatalf("prompt values = %#v, want edit", got)
|
||||
}
|
||||
if got := rewrittenForm.File["image"]; len(got) != 1 || got[0].Header.Get("Content-Type") != "image/png" {
|
||||
t.Fatalf("image headers = %#v, want image/png", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildImagesAPIResponseFromXAI(t *testing.T) {
|
||||
payload := []byte(`{"created":123,"data":[{"b64_json":"AA==","revised_prompt":"refined","mime_type":"image/png"}],"usage":{"total_tokens":0}}`)
|
||||
|
||||
out, err := buildImagesAPIResponseFromXAI(payload, "b64_json")
|
||||
if err != nil {
|
||||
t.Fatalf("buildImagesAPIResponseFromXAI() error = %v", err)
|
||||
}
|
||||
|
||||
if got := gjson.GetBytes(out, "created").Int(); got != 123 {
|
||||
t.Fatalf("created = %d, want 123", got)
|
||||
}
|
||||
if got := gjson.GetBytes(out, "data.0.b64_json").String(); got != "AA==" {
|
||||
t.Fatalf("data.0.b64_json = %q, want AA==", got)
|
||||
}
|
||||
if got := gjson.GetBytes(out, "data.0.revised_prompt").String(); got != "refined" {
|
||||
t.Fatalf("data.0.revised_prompt = %q, want refined", got)
|
||||
}
|
||||
if !gjson.GetBytes(out, "usage").Exists() {
|
||||
t.Fatalf("usage missing: %s", string(out))
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesGenerationsRejectsUnsupportedModel(t *testing.T) {
|
||||
handler := &OpenAIAPIHandler{}
|
||||
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"draw a square"}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
|
||||
|
||||
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
||||
}
|
||||
|
||||
func TestImagesEditsJSONRejectsUnsupportedModel(t *testing.T) {
|
||||
handler := &OpenAIAPIHandler{}
|
||||
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
|
||||
|
||||
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
||||
}
|
||||
|
||||
func TestImagesEditsMultipartRejectsUnsupportedModel(t *testing.T) {
|
||||
handler := &OpenAIAPIHandler{}
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
if err := writer.WriteField("model", "gpt-5.4-mini"); err != nil {
|
||||
t.Fatalf("write model field: %v", err)
|
||||
}
|
||||
if err := writer.WriteField("prompt", "edit this"); err != nil {
|
||||
t.Fatalf("write prompt field: %v", err)
|
||||
}
|
||||
if errClose := writer.Close(); errClose != nil {
|
||||
t.Fatalf("close multipart writer: %v", errClose)
|
||||
}
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesEditsPath, writer.FormDataContentType(), &body, handler.ImagesEdits)
|
||||
|
||||
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
||||
}
|
||||
|
||||
func TestImagesGenerations_DisableImageGeneration_Returns404(t *testing.T) {
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationAll}, nil)
|
||||
handler := NewOpenAIAPIHandler(base)
|
||||
body := strings.NewReader(`{"prompt":"draw a square"}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
|
||||
|
||||
if resp.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusNotFound, resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesEdits_DisableImageGeneration_Returns404(t *testing.T) {
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationAll}, nil)
|
||||
handler := NewOpenAIAPIHandler(base)
|
||||
body := strings.NewReader(`{"prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
|
||||
|
||||
if resp.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusNotFound, resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesGenerations_DisableImageGenerationChat_DoesNotReturn404(t *testing.T) {
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationChat}, nil)
|
||||
handler := NewOpenAIAPIHandler(base)
|
||||
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"draw a square"}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
|
||||
|
||||
if resp.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestImagesEdits_DisableImageGenerationChat_DoesNotReturn404(t *testing.T) {
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationChat}, nil)
|
||||
handler := NewOpenAIAPIHandler(base)
|
||||
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
|
||||
|
||||
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
|
||||
|
||||
if resp.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEFrameAccumulatorFlushesDataOnlyFrame(t *testing.T) {
|
||||
accumulator := &sseFrameAccumulator{}
|
||||
chunk := []byte(`data: {"type":"image_generation.partial","partial_image_index":0}`)
|
||||
|
||||
if frames := accumulator.AddChunk(chunk); len(frames) != 0 {
|
||||
t.Fatalf("AddChunk() emitted an unterminated data-only frame: %q", frames)
|
||||
}
|
||||
frames := accumulator.Flush()
|
||||
if len(frames) != 1 || string(frames[0]) != string(chunk) {
|
||||
t.Fatalf("Flush() frames = %q, want [%q]", frames, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteImagesStreamErrorEventSanitizesPayload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
raw := `{"error":{"code":"upstream_failed","message":"token=image-secret"},"debug":"` + strings.Repeat("x", 8192) + `"}`
|
||||
writeImagesStreamErrorEvent(c, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New(raw)})
|
||||
|
||||
body := recorder.Body.String()
|
||||
if strings.Contains(body, "image-secret") || len(body) > 4096 || !strings.Contains(body, "[REDACTED]") {
|
||||
t.Fatalf("image stream error was not safely bounded: len=%d body=%q", len(body), body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectImagesRejectsPayloadErrorBeforeCompleted(t *testing.T) {
|
||||
data := make(chan []byte, 1)
|
||||
data <- []byte("event: error\ndata: {\"type\":\"provider.error\",\"error\":{\"code\":\"failed\",\"message\":\"token=image-secret\"}}\n\n" +
|
||||
"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2U=\"}]}}\n\n")
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
close(errs)
|
||||
|
||||
out, errMsg := collectImagesFromResponsesStream(context.Background(), data, errs, "b64_json")
|
||||
if len(out) != 0 || errMsg == nil || errMsg.Error == nil {
|
||||
t.Fatalf("payload error result out=%q err=%#v", out, errMsg)
|
||||
}
|
||||
if strings.Contains(errMsg.Error.Error(), "image-secret") || !strings.Contains(errMsg.Error.Error(), "[REDACTED]") {
|
||||
t.Fatalf("payload error was not sanitized: %q", errMsg.Error.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardImagesStreamCancelsWithPayloadError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := NewOpenAIAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil))
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
data := make(chan []byte)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
close(errs)
|
||||
var canceled error
|
||||
firstChunk := []byte("event: error\ndata: {\"error\":{\"message\":\"token=image-secret\"}}\n\n")
|
||||
|
||||
h.forwardImagesStream(context.Background(), c, flusher, func(err error) { canceled = err }, data, errs, firstChunk, "b64_json", "image_generation", func(string, []byte) {})
|
||||
if canceled == nil || strings.Contains(canceled.Error(), "image-secret") || !strings.Contains(canceled.Error(), "[REDACTED]") {
|
||||
t.Fatalf("payload error cancel = %v body=%q", canceled, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "event: error") {
|
||||
t.Fatalf("payload error event missing: %q", recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardRawImageStreamPrefersPendingErrorOnClose(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := NewOpenAIAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil))
|
||||
for i := 0; i < 100; i++ {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||||
data := make(chan []byte)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusTooManyRequests, Error: errors.New("image upstream busy")}
|
||||
close(errs)
|
||||
var canceled error
|
||||
|
||||
h.forwardRawImageStream(context.Background(), c, func(err error) { canceled = err }, data, errs)
|
||||
if canceled == nil || !strings.Contains(canceled.Error(), "image upstream busy") {
|
||||
t.Fatalf("iteration %d: cancel=%v body=%q", i, canceled, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectImagesPrefersPendingErrorWhenDataChannelCloses(t *testing.T) {
|
||||
for i := 0; i < 100; i++ {
|
||||
data := make(chan []byte)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
want := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Error: errors.New("image upstream busy"),
|
||||
DirectResponse: true,
|
||||
Headers: http.Header{"Retry-After": []string{"9"}},
|
||||
}
|
||||
errs <- want
|
||||
close(errs)
|
||||
|
||||
_, got := collectImagesFromResponsesStream(context.Background(), data, errs, "b64_json")
|
||||
if got != want {
|
||||
t.Fatalf("iteration %d: pending error = %#v, want original %#v", i, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectImagesAllowsMultilineSSEData(t *testing.T) {
|
||||
data := make(chan []byte, 1)
|
||||
data <- []byte("event: response.completed\n" +
|
||||
"data: {\"type\":\"response.completed\",\n" +
|
||||
"data: \"response\":{\"created_at\":1,\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2U=\"}]}}\n\n")
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
close(errs)
|
||||
|
||||
out, errMsg := collectImagesFromResponsesStream(context.Background(), data, errs, "b64_json")
|
||||
if errMsg != nil {
|
||||
t.Fatalf("collectImagesFromResponsesStream() error = %v", errMsg.Error)
|
||||
}
|
||||
if !strings.Contains(string(out), `"b64_json":"aW1hZ2U="`) {
|
||||
t.Fatalf("multiline image response = %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEFrameAccumulatorKeepsMultipleFramesDistinct(t *testing.T) {
|
||||
accumulator := &sseFrameAccumulator{}
|
||||
first := "event: first\ndata: {\"type\":\"first\"}\n\n"
|
||||
second := "event: second\ndata: {\"type\":\"second\"}\n\n"
|
||||
|
||||
frames := accumulator.AddChunk([]byte(first + second))
|
||||
if len(frames) != 2 {
|
||||
t.Fatalf("AddChunk() returned %d frames, want 2: %q", len(frames), frames)
|
||||
}
|
||||
if string(frames[0]) != first || string(frames[1]) != second {
|
||||
t.Fatalf("frames were overwritten during buffer compaction: %q", frames)
|
||||
}
|
||||
}
|
||||
375
backend/sdk/api/handlers/openai/openai_responses_compact_test.go
Normal file
375
backend/sdk/api/handlers/openai/openai_responses_compact_test.go
Normal file
|
|
@ -0,0 +1,375 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
type compactCaptureExecutor struct {
|
||||
alt string
|
||||
sourceFormat string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (e *compactCaptureExecutor) Identifier() string { return "test-provider" }
|
||||
|
||||
func (e *compactCaptureExecutor) Execute(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
e.calls++
|
||||
e.alt = opts.Alt
|
||||
e.sourceFormat = opts.SourceFormat.String()
|
||||
return coreexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil
|
||||
}
|
||||
|
||||
func (e *compactCaptureExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (e *compactCaptureExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (e *compactCaptureExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (e *compactCaptureExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactRejectsStream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactCaptureExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth := &coreauth.Auth{ID: "auth1", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("Register auth: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", strings.NewReader(`{"model":"test-model","stream":true}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", resp.Code, http.StatusBadRequest)
|
||||
}
|
||||
if executor.calls != 0 {
|
||||
t.Fatalf("executor calls = %d, want 0", executor.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactExecute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactCaptureExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth := &coreauth.Auth{ID: "auth2", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("Register auth: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", strings.NewReader(`{"model":"test-model","input":"hello"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d", resp.Code, http.StatusOK)
|
||||
}
|
||||
if executor.alt != "responses/compact" {
|
||||
t.Fatalf("alt = %q, want %q", executor.alt, "responses/compact")
|
||||
}
|
||||
if executor.sourceFormat != "openai-response" {
|
||||
t.Fatalf("source format = %q, want %q", executor.sourceFormat, "openai-response")
|
||||
}
|
||||
if strings.TrimSpace(resp.Body.String()) != `{"ok":true}` {
|
||||
t.Fatalf("body = %s", resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactDecodesZstdRequestBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactCaptureExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth := &coreauth.Auth{ID: "auth3", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("Register auth: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
|
||||
var compressed bytes.Buffer
|
||||
encoder, err := zstd.NewWriter(&compressed)
|
||||
if err != nil {
|
||||
t.Fatalf("zstd.NewWriter: %v", err)
|
||||
}
|
||||
if _, errWrite := encoder.Write([]byte(`{"model":"test-model","input":"hello"}`)); errWrite != nil {
|
||||
t.Fatalf("zstd write: %v", errWrite)
|
||||
}
|
||||
if errClose := encoder.Close(); errClose != nil {
|
||||
t.Fatalf("zstd close: %v", errClose)
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader(compressed.Bytes()))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Content-Encoding", "zstd")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", resp.Code, http.StatusOK, resp.Body.String())
|
||||
}
|
||||
if executor.calls != 1 {
|
||||
t.Fatalf("executor calls = %d, want 1", executor.calls)
|
||||
}
|
||||
if executor.alt != "responses/compact" {
|
||||
t.Fatalf("alt = %q, want %q", executor.alt, "responses/compact")
|
||||
}
|
||||
if strings.TrimSpace(resp.Body.String()) != `{"ok":true}` {
|
||||
t.Fatalf("body = %s", resp.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
type compactMockStatusError struct {
|
||||
code int
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e compactMockStatusError) Error() string { return e.msg }
|
||||
func (e compactMockStatusError) StatusCode() int { return e.code }
|
||||
|
||||
type compactFailureMockExecutor struct {
|
||||
compactErr error
|
||||
normalResp []byte
|
||||
calls int
|
||||
lastAlt string
|
||||
lastAuthID string
|
||||
}
|
||||
|
||||
func (e *compactFailureMockExecutor) Identifier() string { return "test-compact-provider" }
|
||||
|
||||
func (e *compactFailureMockExecutor) Execute(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
e.calls++
|
||||
e.lastAlt = opts.Alt
|
||||
if auth != nil {
|
||||
e.lastAuthID = auth.ID
|
||||
}
|
||||
if opts.Alt == "responses/compact" {
|
||||
if e.compactErr != nil {
|
||||
return coreexecutor.Response{}, e.compactErr
|
||||
}
|
||||
}
|
||||
respPayload := e.normalResp
|
||||
if len(respPayload) == 0 {
|
||||
respPayload = []byte(`{"id":"resp_123","object":"response","status":"completed"}`)
|
||||
}
|
||||
return coreexecutor.Response{Payload: respPayload}, nil
|
||||
}
|
||||
|
||||
func (e *compactFailureMockExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (e *compactFailureMockExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (e *compactFailureMockExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (e *compactFailureMockExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactTransientFailureDoesNotCooldownAuthAndPreservesError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactFailureMockExecutor{
|
||||
compactErr: compactMockStatusError{
|
||||
code: http.StatusInternalServerError,
|
||||
msg: `{"error":{"message":"compact upstream temporary error","type":"api_error"}}`,
|
||||
},
|
||||
}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth1 := &coreauth.Auth{ID: "auth1", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
auth2 := &coreauth.Auth{ID: "auth2", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth1); err != nil {
|
||||
t.Fatalf("Register auth1: %v", err)
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth2); err != nil {
|
||||
t.Fatalf("Register auth2: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth1.ID)
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth2.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
// Send compact request which fails upstream on all auths with 500
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", strings.NewReader(`{"model":"test-model","input":"hello"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
// 1. Should return upstream status 500 and upstream error message (not generic 503 Service temporarily unavailable)
|
||||
if resp.Code != http.StatusInternalServerError {
|
||||
t.Fatalf("compact status = %d, want %d; body = %s", resp.Code, http.StatusInternalServerError, resp.Body.String())
|
||||
}
|
||||
if !strings.Contains(resp.Body.String(), "compact upstream temporary error") {
|
||||
t.Fatalf("compact body = %s, want containing 'compact upstream temporary error'", resp.Body.String())
|
||||
}
|
||||
|
||||
// 2. Auth model states should NOT be marked unavailable for normal traffic
|
||||
for _, authID := range []string{"auth1", "auth2"} {
|
||||
a, ok := manager.GetByID(authID)
|
||||
if !ok {
|
||||
t.Fatalf("auth %s not found", authID)
|
||||
}
|
||||
if state, exists := a.ModelStates["test-model"]; exists && state != nil {
|
||||
if state.Unavailable {
|
||||
t.Fatalf("auth %s model state marked Unavailable after compact failure", authID)
|
||||
}
|
||||
if !state.NextRetryAfter.IsZero() && state.NextRetryAfter.After(time.Now()) {
|
||||
t.Fatalf("auth %s model state has NextRetryAfter %v in future", authID, state.NextRetryAfter)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Normal /v1/responses request should succeed immediately without auth cooldown errors
|
||||
reqNormal := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"test-model","input":"hello"}`))
|
||||
reqNormal.Header.Set("Content-Type", "application/json")
|
||||
respNormal := httptest.NewRecorder()
|
||||
router.ServeHTTP(respNormal, reqNormal)
|
||||
|
||||
if respNormal.Code != http.StatusOK {
|
||||
t.Fatalf("normal responses status = %d, want %d; body = %s", respNormal.Code, http.StatusOK, respNormal.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactRequestFaultStopsFallbackAndPreservesError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactFailureMockExecutor{
|
||||
compactErr: compactMockStatusError{
|
||||
code: http.StatusNotFound,
|
||||
msg: `404 page not found`,
|
||||
},
|
||||
}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth1 := &coreauth.Auth{ID: "auth1", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
auth2 := &coreauth.Auth{ID: "auth2", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth1); err != nil {
|
||||
t.Fatalf("Register auth1: %v", err)
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth2); err != nil {
|
||||
t.Fatalf("Register auth2: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth1.ID)
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth2.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
// Send compact request which fails upstream with 404 (endpoint not supported / invalid)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", strings.NewReader(`{"model":"test-model","input":"hello"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
// 1. Should return upstream status 404 and upstream error message
|
||||
if resp.Code != http.StatusNotFound {
|
||||
t.Fatalf("compact status = %d, want %d; body = %s", resp.Code, http.StatusNotFound, resp.Body.String())
|
||||
}
|
||||
if !strings.Contains(resp.Body.String(), "404 page not found") {
|
||||
t.Fatalf("compact body = %s, want containing '404 page not found'", resp.Body.String())
|
||||
}
|
||||
|
||||
// 2. Should stop fallback on request/capability fault (calls == 1)
|
||||
if executor.calls != 1 {
|
||||
t.Fatalf("executor calls = %d, want 1 (fallback should stop)", executor.calls)
|
||||
}
|
||||
|
||||
// 3. Auth model states should NOT be marked unavailable for normal traffic
|
||||
for _, authID := range []string{"auth1", "auth2"} {
|
||||
a, ok := manager.GetByID(authID)
|
||||
if !ok {
|
||||
t.Fatalf("auth %s not found", authID)
|
||||
}
|
||||
if state, exists := a.ModelStates["test-model"]; exists && state != nil {
|
||||
if state.Unavailable {
|
||||
t.Fatalf("auth %s model state marked Unavailable after compact failure", authID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Normal /v1/responses request should succeed immediately
|
||||
reqNormal := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"test-model","input":"hello"}`))
|
||||
reqNormal.Header.Set("Content-Type", "application/json")
|
||||
respNormal := httptest.NewRecorder()
|
||||
router.ServeHTTP(respNormal, reqNormal)
|
||||
|
||||
if respNormal.Code != http.StatusOK {
|
||||
t.Fatalf("normal responses status = %d, want %d; body = %s", respNormal.Code, http.StatusOK, respNormal.Body.String())
|
||||
}
|
||||
}
|
||||
973
backend/sdk/api/handlers/openai/openai_responses_handlers.go
Normal file
973
backend/sdk/api/handlers/openai/openai_responses_handlers.go
Normal file
|
|
@ -0,0 +1,973 @@
|
|||
// Package openai provides HTTP handlers for OpenAIResponses API endpoints.
|
||||
// This package implements the OpenAIResponses-compatible API interface, including model listing
|
||||
// and chat completion functionality. It supports both streaming and non-streaming responses,
|
||||
// and manages a pool of clients to interact with backend services.
|
||||
// The handlers translate OpenAIResponses API requests to the appropriate backend format and
|
||||
// convert responses back to OpenAIResponses-compatible format.
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/optimize-multi-agent-v2"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func writeResponsesSSEChunk(w io.Writer, chunk []byte) {
|
||||
if w == nil || len(chunk) == 0 {
|
||||
return
|
||||
}
|
||||
if _, err := w.Write(chunk); err != nil {
|
||||
return
|
||||
}
|
||||
if bytes.HasSuffix(chunk, []byte("\n\n")) || bytes.HasSuffix(chunk, []byte("\r\n\r\n")) {
|
||||
return
|
||||
}
|
||||
suffix := []byte("\n\n")
|
||||
if bytes.HasSuffix(chunk, []byte("\r\n")) {
|
||||
suffix = []byte("\r\n")
|
||||
} else if bytes.HasSuffix(chunk, []byte("\n")) {
|
||||
suffix = []byte("\n")
|
||||
}
|
||||
if _, err := w.Write(suffix); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
type responsesSSEFramer struct {
|
||||
pending []byte
|
||||
outputItems map[int][]byte
|
||||
outputOrder []int
|
||||
unindexedOutputItems [][]byte
|
||||
lastEvent string
|
||||
terminalEvent string
|
||||
terminalError *interfaces.ErrorMessage
|
||||
failureEvent string
|
||||
dataFrames int
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) {
|
||||
if len(chunk) == 0 || f.terminalEvent != "" {
|
||||
return
|
||||
}
|
||||
if responsesSSEStartsNewDataFrame(f.pending, chunk) {
|
||||
f.writeFrame(w, f.pending)
|
||||
f.pending = f.pending[:0]
|
||||
if f.terminalEvent != "" {
|
||||
return
|
||||
}
|
||||
}
|
||||
if responsesSSENeedsLineBreak(f.pending, chunk) {
|
||||
f.pending = append(f.pending, '\n')
|
||||
}
|
||||
f.pending = append(f.pending, chunk...)
|
||||
for {
|
||||
frameLen := responsesSSEFrameLen(f.pending)
|
||||
if frameLen == 0 {
|
||||
break
|
||||
}
|
||||
f.writeFrame(w, f.pending[:frameLen])
|
||||
copy(f.pending, f.pending[frameLen:])
|
||||
f.pending = f.pending[:len(f.pending)-frameLen]
|
||||
if f.terminalEvent != "" {
|
||||
f.pending = f.pending[:0]
|
||||
return
|
||||
}
|
||||
}
|
||||
if len(bytes.TrimSpace(f.pending)) == 0 {
|
||||
f.pending = f.pending[:0]
|
||||
return
|
||||
}
|
||||
if len(f.pending) == 0 || !responsesSSECanEmitWithoutDelimiter(f.pending) {
|
||||
return
|
||||
}
|
||||
f.writeFrame(w, f.pending)
|
||||
f.pending = f.pending[:0]
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) Flush(w io.Writer) {
|
||||
if len(f.pending) == 0 || f.terminalEvent != "" {
|
||||
return
|
||||
}
|
||||
if len(bytes.TrimSpace(f.pending)) == 0 {
|
||||
f.pending = f.pending[:0]
|
||||
return
|
||||
}
|
||||
if !responsesSSECanFlushWithoutDelimiter(f.pending) {
|
||||
f.pending = f.pending[:0]
|
||||
return
|
||||
}
|
||||
f.writeFrame(w, f.pending)
|
||||
f.pending = f.pending[:0]
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) writeFrame(w io.Writer, frame []byte) {
|
||||
writeResponsesSSEChunk(w, f.repairFrame(frame))
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) repairFrame(frame []byte) []byte {
|
||||
payload, ok := responsesSSEDataPayload(frame)
|
||||
if !ok || len(payload) == 0 {
|
||||
return frame
|
||||
}
|
||||
if bytes.Equal(payload, []byte("[DONE]")) {
|
||||
f.dataFrames++
|
||||
return frame
|
||||
}
|
||||
if !json.Valid(payload) {
|
||||
return frame
|
||||
}
|
||||
f.dataFrames++
|
||||
|
||||
payloadType := gjson.GetBytes(payload, "type").String()
|
||||
if responsesSSEErrorEvent(payloadType) || responsesSSEPayloadHasError(payload) {
|
||||
if payloadType != "" {
|
||||
f.lastEvent = sanitizeResponsesStreamEventName(payloadType)
|
||||
}
|
||||
return f.repairErrorPayload(payload)
|
||||
}
|
||||
streamEvent := responsesSSEEventName(frame)
|
||||
eventType := payloadType
|
||||
if responsesSSETerminalEvent(streamEvent) {
|
||||
eventType = streamEvent
|
||||
} else if eventType == "" {
|
||||
eventType = streamEvent
|
||||
}
|
||||
if eventType != "" {
|
||||
f.lastEvent = sanitizeResponsesStreamEventName(eventType)
|
||||
}
|
||||
if responsesSSEErrorEvent(eventType) {
|
||||
return f.repairErrorPayload(payload)
|
||||
}
|
||||
if responsesSSETerminalEvent(eventType) {
|
||||
f.terminalEvent = eventType
|
||||
}
|
||||
|
||||
switch eventType {
|
||||
case "response.output_item.done":
|
||||
f.recordOutputItem(payload)
|
||||
case "response.completed":
|
||||
repaired := f.repairCompletedPayload(payload)
|
||||
if !bytes.Equal(repaired, payload) {
|
||||
return responsesSSEFrameWithData(frame, repaired)
|
||||
}
|
||||
}
|
||||
return frame
|
||||
}
|
||||
|
||||
func responsesSSEPayloadErrorMessage(payload []byte) *interfaces.ErrorMessage {
|
||||
status := http.StatusBadGateway
|
||||
for _, path := range []string{"status", "status_code", "error.status", "error.status_code", "response.error.status", "response.error.status_code"} {
|
||||
candidate := int(gjson.GetBytes(payload, path).Int())
|
||||
if candidate >= http.StatusBadRequest && candidate <= 599 {
|
||||
status = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
return sanitizeResponsesStreamErrorMessage(&interfaces.ErrorMessage{StatusCode: status, Error: fmt.Errorf("%s", payload)})
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) repairErrorPayload(payload []byte) []byte {
|
||||
errMsg := responsesSSEPayloadErrorMessage(payload)
|
||||
status := errMsg.StatusCode
|
||||
f.terminalError = errMsg
|
||||
failureEvent := f.failureEvent
|
||||
if failureEvent != "response.failed" {
|
||||
failureEvent = "error"
|
||||
}
|
||||
f.terminalEvent = failureEvent
|
||||
errText := responsesStreamErrorText(errMsg, status)
|
||||
if failureEvent == "response.failed" {
|
||||
chunk := handlers.BuildOpenAIResponsesStreamFailedChunk(status, errText, 0)
|
||||
return []byte(fmt.Sprintf("event: response.failed\ndata: %s\n\n", chunk))
|
||||
}
|
||||
chunk := handlers.BuildOpenAIResponsesStreamErrorChunk(status, errText, 0)
|
||||
return []byte(fmt.Sprintf("event: error\ndata: %s\n\n", chunk))
|
||||
}
|
||||
|
||||
func responsesSSEErrorEvent(eventType string) bool {
|
||||
switch eventType {
|
||||
case "response.failed", "response.error", "error":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func responsesSSETerminalEvent(eventType string) bool {
|
||||
switch eventType {
|
||||
case "response.completed", "response.incomplete", "response.failed", "response.done", "response.error", "error":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func responsesSSEPayloadHasError(payload []byte) bool {
|
||||
for _, path := range []string{"error", "response.error"} {
|
||||
result := gjson.GetBytes(payload, path)
|
||||
if result.Exists() && result.Type != gjson.Null {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return gjson.GetBytes(payload, "code").Exists() && gjson.GetBytes(payload, "message").Exists()
|
||||
}
|
||||
|
||||
func responsesSSEDataPayload(frame []byte) ([]byte, bool) {
|
||||
var payload []byte
|
||||
found := false
|
||||
for _, line := range bytes.Split(frame, []byte("\n")) {
|
||||
line = bytes.TrimRight(line, "\r")
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if !bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
continue
|
||||
}
|
||||
data := bytes.TrimSpace(trimmed[len("data:"):])
|
||||
if found {
|
||||
payload = append(payload, '\n')
|
||||
}
|
||||
payload = append(payload, data...)
|
||||
found = true
|
||||
}
|
||||
return payload, found
|
||||
}
|
||||
|
||||
func responsesSSEFrameWithData(frame, payload []byte) []byte {
|
||||
var out bytes.Buffer
|
||||
for _, line := range bytes.Split(frame, []byte("\n")) {
|
||||
line = bytes.TrimRight(line, "\r")
|
||||
trimmed := bytes.TrimSpace(line)
|
||||
if len(trimmed) == 0 || bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
continue
|
||||
}
|
||||
out.Write(line)
|
||||
out.WriteByte('\n')
|
||||
}
|
||||
for _, line := range bytes.Split(payload, []byte("\n")) {
|
||||
out.WriteString("data: ")
|
||||
out.Write(line)
|
||||
out.WriteByte('\n')
|
||||
}
|
||||
out.WriteByte('\n')
|
||||
return out.Bytes()
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) recordOutputItem(payload []byte) {
|
||||
item := gjson.GetBytes(payload, "item")
|
||||
if !item.Exists() || !item.IsObject() || item.Get("type").String() == "" {
|
||||
return
|
||||
}
|
||||
|
||||
if outputIndex := gjson.GetBytes(payload, "output_index"); outputIndex.Exists() {
|
||||
index := int(outputIndex.Int())
|
||||
if f.outputItems == nil {
|
||||
f.outputItems = make(map[int][]byte)
|
||||
}
|
||||
if _, exists := f.outputItems[index]; !exists {
|
||||
f.outputOrder = append(f.outputOrder, index)
|
||||
}
|
||||
f.outputItems[index] = append([]byte(nil), item.Raw...)
|
||||
return
|
||||
}
|
||||
|
||||
f.unindexedOutputItems = append(f.unindexedOutputItems, append([]byte(nil), item.Raw...))
|
||||
}
|
||||
|
||||
func (f *responsesSSEFramer) repairCompletedPayload(payload []byte) []byte {
|
||||
if len(f.outputOrder) == 0 && len(f.unindexedOutputItems) == 0 {
|
||||
return payload
|
||||
}
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if output.Exists() && (!output.IsArray() || len(output.Array()) > 0) {
|
||||
return payload
|
||||
}
|
||||
|
||||
var outputJSON bytes.Buffer
|
||||
outputJSON.WriteByte('[')
|
||||
indexes := append([]int(nil), f.outputOrder...)
|
||||
sort.Ints(indexes)
|
||||
written := 0
|
||||
for _, index := range indexes {
|
||||
item, ok := f.outputItems[index]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if written > 0 {
|
||||
outputJSON.WriteByte(',')
|
||||
}
|
||||
outputJSON.Write(item)
|
||||
written++
|
||||
}
|
||||
for _, item := range f.unindexedOutputItems {
|
||||
if written > 0 {
|
||||
outputJSON.WriteByte(',')
|
||||
}
|
||||
outputJSON.Write(item)
|
||||
written++
|
||||
}
|
||||
outputJSON.WriteByte(']')
|
||||
|
||||
repaired, err := sjson.SetRawBytes(payload, "response.output", outputJSON.Bytes())
|
||||
if err != nil {
|
||||
return payload
|
||||
}
|
||||
return repaired
|
||||
}
|
||||
|
||||
func responsesSSEFrameLen(chunk []byte) int {
|
||||
if len(chunk) == 0 {
|
||||
return 0
|
||||
}
|
||||
lf := bytes.Index(chunk, []byte("\n\n"))
|
||||
crlf := bytes.Index(chunk, []byte("\r\n\r\n"))
|
||||
switch {
|
||||
case lf < 0:
|
||||
if crlf < 0 {
|
||||
return 0
|
||||
}
|
||||
return crlf + 4
|
||||
case crlf < 0:
|
||||
return lf + 2
|
||||
case lf < crlf:
|
||||
return lf + 2
|
||||
default:
|
||||
return crlf + 4
|
||||
}
|
||||
}
|
||||
|
||||
func responsesSSENeedsMoreData(chunk []byte) bool {
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
if len(trimmed) == 0 {
|
||||
return false
|
||||
}
|
||||
return responsesSSEHasField(trimmed, []byte("event:")) && !responsesSSEHasField(trimmed, []byte("data:"))
|
||||
}
|
||||
|
||||
func responsesSSEHasField(chunk []byte, prefix []byte) bool {
|
||||
s := chunk
|
||||
for len(s) > 0 {
|
||||
line := s
|
||||
if i := bytes.IndexByte(s, '\n'); i >= 0 {
|
||||
line = s[:i]
|
||||
s = s[i+1:]
|
||||
} else {
|
||||
s = nil
|
||||
}
|
||||
line = bytes.TrimSpace(line)
|
||||
if bytes.HasPrefix(line, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func responsesSSECanEmitWithoutDelimiter(chunk []byte) bool {
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
if len(trimmed) == 0 || responsesSSENeedsMoreData(trimmed) ||
|
||||
!responsesSSEHasField(trimmed, []byte("event:")) || !responsesSSEHasField(trimmed, []byte("data:")) {
|
||||
return false
|
||||
}
|
||||
return responsesSSEDataLinesValid(trimmed)
|
||||
}
|
||||
|
||||
func responsesSSECanFlushWithoutDelimiter(chunk []byte) bool {
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
return len(trimmed) > 0 && responsesSSEHasField(trimmed, []byte("data:")) && responsesSSEDataLinesValid(trimmed)
|
||||
}
|
||||
|
||||
func responsesSSEStartsNewDataFrame(pending, chunk []byte) bool {
|
||||
trimmedPending := bytes.TrimSpace(pending)
|
||||
if len(trimmedPending) == 0 || responsesSSEHasField(trimmedPending, []byte("event:")) ||
|
||||
!responsesSSEHasField(trimmedPending, []byte("data:")) || !responsesSSEDataLinesValid(trimmedPending) {
|
||||
return false
|
||||
}
|
||||
trimmedChunk := bytes.TrimLeft(chunk, " \t\r\n")
|
||||
return bytes.HasPrefix(trimmedChunk, []byte("data:"))
|
||||
}
|
||||
|
||||
func responsesSSEEventName(frame []byte) string {
|
||||
for _, line := range bytes.Split(frame, []byte("\n")) {
|
||||
trimmed := bytes.TrimSpace(bytes.TrimRight(line, "\r"))
|
||||
if bytes.HasPrefix(trimmed, []byte("event:")) {
|
||||
return strings.TrimSpace(string(trimmed[len("event:"):]))
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func responsesSSEDataLinesValid(chunk []byte) bool {
|
||||
payload, found := responsesSSEDataPayload(chunk)
|
||||
if !found {
|
||||
return true
|
||||
}
|
||||
payload = bytes.TrimSpace(payload)
|
||||
return len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) || json.Valid(payload)
|
||||
}
|
||||
|
||||
func responsesSSENeedsLineBreak(pending, chunk []byte) bool {
|
||||
if len(pending) == 0 || len(chunk) == 0 {
|
||||
return false
|
||||
}
|
||||
if bytes.HasSuffix(pending, []byte("\n")) || bytes.HasSuffix(pending, []byte("\r")) {
|
||||
return false
|
||||
}
|
||||
if chunk[0] == '\n' || chunk[0] == '\r' {
|
||||
return false
|
||||
}
|
||||
trimmed := bytes.TrimLeft(chunk, " \t")
|
||||
if len(trimmed) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, prefix := range [][]byte{[]byte("data:"), []byte("event:"), []byte("id:"), []byte("retry:"), []byte(":")} {
|
||||
if bytes.HasPrefix(trimmed, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// OpenAIResponsesAPIHandler contains the handlers for OpenAIResponses API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service.
|
||||
type OpenAIResponsesAPIHandler struct {
|
||||
*handlers.BaseAPIHandler
|
||||
}
|
||||
|
||||
// NewOpenAIResponsesAPIHandler creates a new OpenAIResponses API handlers instance.
|
||||
// It takes an BaseAPIHandler instance as input and returns an OpenAIResponsesAPIHandler.
|
||||
//
|
||||
// Parameters:
|
||||
// - apiHandlers: The base API handlers instance
|
||||
//
|
||||
// Returns:
|
||||
// - *OpenAIResponsesAPIHandler: A new OpenAIResponses API handlers instance
|
||||
func NewOpenAIResponsesAPIHandler(apiHandlers *handlers.BaseAPIHandler) *OpenAIResponsesAPIHandler {
|
||||
return &OpenAIResponsesAPIHandler{
|
||||
BaseAPIHandler: apiHandlers,
|
||||
}
|
||||
}
|
||||
|
||||
// HandlerType returns the identifier for this handler implementation.
|
||||
func (h *OpenAIResponsesAPIHandler) HandlerType() string {
|
||||
return OpenaiResponse
|
||||
}
|
||||
|
||||
// Models returns the OpenAIResponses-compatible model metadata supported by this handler.
|
||||
func (h *OpenAIResponsesAPIHandler) Models() []map[string]any {
|
||||
// Get dynamic models from the global registry
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
return modelRegistry.GetAvailableModels("openai")
|
||||
}
|
||||
|
||||
// OpenAIResponsesModels handles the /v1/models endpoint.
|
||||
// It returns a list of available AI models with their capabilities
|
||||
// and specifications in OpenAIResponses-compatible format.
|
||||
func (h *OpenAIResponsesAPIHandler) OpenAIResponsesModels(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": h.Models(),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) prepareCodexMultiAgentV2Tools(c *gin.Context, payload []byte) []byte {
|
||||
if h == nil || h.Cfg == nil {
|
||||
return payload
|
||||
}
|
||||
|
||||
requestCtx := context.Background()
|
||||
if c != nil && c.Request != nil {
|
||||
requestCtx = c.Request.Context()
|
||||
}
|
||||
requestCtx = context.WithValue(requestCtx, "gin", c)
|
||||
|
||||
var requestHeaders http.Header
|
||||
if c != nil && c.Request != nil {
|
||||
requestHeaders = c.Request.Header
|
||||
}
|
||||
homeEnabled := h.AuthManager != nil && h.AuthManager.HomeEnabled()
|
||||
updated, prepared := multiagentv2.PrepareCodexMultiAgentV2Tools(
|
||||
requestCtx,
|
||||
requestHeaders,
|
||||
payload,
|
||||
h.Cfg.CodexOptimizeMultiAgentV2,
|
||||
homeEnabled,
|
||||
)
|
||||
if prepared && c != nil {
|
||||
c.Set(multiagentv2.CodexMultiAgentV2ToolsPreparedContextKey, true)
|
||||
}
|
||||
return updated
|
||||
}
|
||||
|
||||
// Responses handles the /v1/responses endpoint.
|
||||
// It determines whether the request is for a streaming or non-streaming response
|
||||
// and calls the appropriate handler based on the model provider.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
func (h *OpenAIResponsesAPIHandler) Responses(c *gin.Context) {
|
||||
rawJSON, err := handlers.ReadRequestBody(c)
|
||||
// If data retrieval fails, return a 400 Bad Request error.
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
rawJSON = h.prepareCodexMultiAgentV2Tools(c, rawJSON)
|
||||
|
||||
// Check if the client requested a streaming response.
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
if streamResult.Type == gjson.True {
|
||||
h.handleStreamingResponse(c, rawJSON)
|
||||
} else {
|
||||
h.handleNonStreamingResponse(c, rawJSON)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) Compact(c *gin.Context) {
|
||||
rawJSON, err := handlers.ReadRequestBody(c)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
if streamResult.Type == gjson.True {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported for compact responses",
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
if streamResult.Exists() {
|
||||
if updated, err := sjson.DeleteBytes(rawJSON, "stream"); err == nil {
|
||||
rawJSON = updated
|
||||
}
|
||||
}
|
||||
|
||||
c.Header("Content-Type", "application/json")
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "responses/compact")
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleNonStreamingResponse handles non-streaming chat completion responses
|
||||
// for Gemini models. It selects a client from the pool, sends the request, and
|
||||
// aggregates the response before sending it back to the client in OpenAIResponses format.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAIResponses-compatible request
|
||||
func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleStreamingResponse handles streaming responses for Gemini models.
|
||||
// It establishes a streaming connection with the backend service and forwards
|
||||
// the response chunks to the client in real-time using Server-Sent Events.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAIResponses-compatible request
|
||||
func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// New core execution path
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
failureEvent := "error"
|
||||
if isCodexResponsesClientRequest(c) {
|
||||
failureEvent = "response.failed"
|
||||
}
|
||||
framer := &responsesSSEFramer{failureEvent: failureEvent}
|
||||
var initialOutput bytes.Buffer
|
||||
|
||||
// Peek at the first complete SSE data frame.
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
framer.Flush(&initialOutput)
|
||||
safeErrMsg := sanitizeResponsesStreamErrorMessage(errMsg)
|
||||
if framer.dataFrames == 0 {
|
||||
safeErrMsg = sanitizeResponsesInitialErrorMessage(errMsg)
|
||||
}
|
||||
if safeErrMsg != nil && framer.dataFrames > 0 {
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(initialOutput.Bytes())
|
||||
flusher.Flush()
|
||||
pendingErrors := make(chan *interfaces.ErrorMessage, 1)
|
||||
pendingErrors <- safeErrMsg
|
||||
close(pendingErrors)
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, make(chan []byte), pendingErrors, framer)
|
||||
return
|
||||
}
|
||||
// Upstream failed before a complete SSE data frame. Return JSON.
|
||||
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), safeErrMsg)
|
||||
h.WriteErrorResponse(c, safeErrMsg)
|
||||
if safeErrMsg != nil {
|
||||
cliCancel(safeErrMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
framer.Flush(&initialOutput)
|
||||
errMsg, hasPendingError := handlers.PendingStreamError(errChan)
|
||||
if !hasPendingError && framer.terminalEvent == "" {
|
||||
message := "upstream stream closed before first payload"
|
||||
if framer.dataFrames > 0 {
|
||||
message = "upstream stream closed before a terminal event"
|
||||
}
|
||||
errMsg = &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("%s", message)}
|
||||
}
|
||||
if framer.dataFrames > 0 {
|
||||
errMsg = sanitizeResponsesStreamErrorMessage(errMsg)
|
||||
} else {
|
||||
errMsg = sanitizeResponsesInitialErrorMessage(errMsg)
|
||||
}
|
||||
|
||||
if framer.dataFrames > 0 {
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(initialOutput.Bytes())
|
||||
flusher.Flush()
|
||||
if framer.terminalError != nil {
|
||||
h.logResponsesStreamError(c, framer, framer.terminalError)
|
||||
cliCancel(framer.terminalError.Error)
|
||||
return
|
||||
}
|
||||
if errMsg == nil {
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
pendingErrors := make(chan *interfaces.ErrorMessage, 1)
|
||||
pendingErrors <- errMsg
|
||||
close(pendingErrors)
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, make(chan []byte), pendingErrors, framer)
|
||||
return
|
||||
}
|
||||
|
||||
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
framer.WriteChunk(&initialOutput, chunk)
|
||||
if framer.dataFrames == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
setSSEHeaders()
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(initialOutput.Bytes())
|
||||
flusher.Flush()
|
||||
if framer.terminalError != nil {
|
||||
h.logResponsesStreamError(c, framer, framer.terminalError)
|
||||
cliCancel(framer.terminalError.Error)
|
||||
return
|
||||
}
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan, framer)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// isCodexResponsesClientRequest limits the alternate terminal event to official Codex clients.
|
||||
func isCodexResponsesClientRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil {
|
||||
return false
|
||||
}
|
||||
if multiagentv2.IsCodexClientUserAgent(c.GetHeader("User-Agent")) {
|
||||
return true
|
||||
}
|
||||
|
||||
switch originator := strings.ToLower(strings.TrimSpace(c.GetHeader("Originator"))); originator {
|
||||
case "codex desktop", "codex-tui", "codex_cli_rs":
|
||||
return true
|
||||
default:
|
||||
return strings.HasPrefix(originator, "codex desktop/") || strings.HasPrefix(originator, "codex-tui/") || strings.HasPrefix(originator, "codex_cli_rs/")
|
||||
}
|
||||
}
|
||||
|
||||
const (
|
||||
responsesStreamErrorMessageLimit = 2048
|
||||
responsesStreamErrorFieldLimit = 256
|
||||
)
|
||||
|
||||
var (
|
||||
responsesStreamSensitiveValuePattern = regexp.MustCompile(`(?i)((?:"?(?:api[_-]?key|access[_-]?token|token|authorization|secret)"?)\s*[=:]\s*"?)([^\s"&,;}]+)`)
|
||||
responsesStreamBearerPattern = regexp.MustCompile(`(?i)\bBearer\s+[A-Za-z0-9._~+/=-]+`)
|
||||
)
|
||||
|
||||
func truncateResponsesStreamErrorText(text string, limit int) string {
|
||||
runes := []rune(text)
|
||||
if len(runes) <= limit {
|
||||
return text
|
||||
}
|
||||
return string(runes[:limit]) + "…"
|
||||
}
|
||||
|
||||
func redactResponsesStreamErrorText(text string) string {
|
||||
text = responsesStreamSensitiveValuePattern.ReplaceAllString(text, `${1}[REDACTED]`)
|
||||
return responsesStreamBearerPattern.ReplaceAllString(text, "Bearer [REDACTED]")
|
||||
}
|
||||
|
||||
func sanitizeResponsesStreamEventName(eventName string) string {
|
||||
return truncateResponsesStreamErrorText(redactResponsesStreamErrorText(strings.TrimSpace(eventName)), responsesStreamErrorFieldLimit)
|
||||
}
|
||||
|
||||
func responsesStreamErrorText(errMsg *interfaces.ErrorMessage, status int) string {
|
||||
text := http.StatusText(status)
|
||||
if errMsg != nil && errMsg.Error != nil && strings.TrimSpace(errMsg.Error.Error()) != "" {
|
||||
text = strings.TrimSpace(errMsg.Error.Error())
|
||||
}
|
||||
if !json.Valid([]byte(text)) {
|
||||
return truncateResponsesStreamErrorText(redactResponsesStreamErrorText(text), responsesStreamErrorMessageLimit)
|
||||
}
|
||||
|
||||
root := gjson.Parse(text)
|
||||
errorNode := root.Get("error")
|
||||
if !errorNode.Exists() || !errorNode.IsObject() {
|
||||
errorNode = root.Get("response.error")
|
||||
}
|
||||
if errorNode.Exists() && errorNode.IsObject() {
|
||||
safe := []byte(`{"error":{}}`)
|
||||
copied := false
|
||||
for _, field := range []string{"type", "code", "message", "param"} {
|
||||
value := errorNode.Get(field)
|
||||
if !value.Exists() || value.Type == gjson.Null {
|
||||
continue
|
||||
}
|
||||
limit := responsesStreamErrorFieldLimit
|
||||
if field == "message" {
|
||||
limit = responsesStreamErrorMessageLimit
|
||||
}
|
||||
safe, _ = sjson.SetBytes(safe, "error."+field, truncateResponsesStreamErrorText(redactResponsesStreamErrorText(value.String()), limit))
|
||||
copied = true
|
||||
}
|
||||
if copied {
|
||||
return string(safe)
|
||||
}
|
||||
}
|
||||
|
||||
safe := []byte(`{"type":"error"}`)
|
||||
copied := false
|
||||
for _, field := range []string{"code", "message", "param"} {
|
||||
value := root.Get(field)
|
||||
if !value.Exists() || value.Type == gjson.Null {
|
||||
continue
|
||||
}
|
||||
limit := responsesStreamErrorFieldLimit
|
||||
if field == "message" {
|
||||
limit = responsesStreamErrorMessageLimit
|
||||
}
|
||||
safe, _ = sjson.SetBytes(safe, field, truncateResponsesStreamErrorText(redactResponsesStreamErrorText(value.String()), limit))
|
||||
copied = true
|
||||
}
|
||||
if copied {
|
||||
return string(safe)
|
||||
}
|
||||
return http.StatusText(status)
|
||||
}
|
||||
|
||||
type responsesStreamSanitizedError struct {
|
||||
message string
|
||||
cause error
|
||||
}
|
||||
|
||||
func (e *responsesStreamSanitizedError) Error() string { return e.message }
|
||||
func (e *responsesStreamSanitizedError) Unwrap() error { return e.cause }
|
||||
|
||||
func sanitizeResponsesInitialErrorMessage(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage {
|
||||
if errMsg != nil && errMsg.DirectResponse {
|
||||
return errMsg
|
||||
}
|
||||
return sanitizeResponsesStreamErrorMessage(errMsg)
|
||||
}
|
||||
|
||||
func sanitizeResponsesStreamErrorMessage(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage {
|
||||
if errMsg == nil {
|
||||
return nil
|
||||
}
|
||||
status := errMsg.StatusCode
|
||||
if status < http.StatusBadRequest || status > 599 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
safe := *errMsg
|
||||
safe.StatusCode = status
|
||||
safe.Error = &responsesStreamSanitizedError{message: responsesStreamErrorText(errMsg, status), cause: errMsg.Error}
|
||||
safe.DirectResponse = false
|
||||
safe.Body = nil
|
||||
return &safe
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) logResponsesStreamError(c *gin.Context, framer *responsesSSEFramer, errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := errMsg.StatusCode
|
||||
if status < http.StatusBadRequest || status > 599 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
lastEvent := "none"
|
||||
if framer != nil && framer.lastEvent != "" {
|
||||
lastEvent = framer.lastEvent
|
||||
}
|
||||
errText := responsesStreamErrorText(errMsg, status)
|
||||
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), &interfaces.ErrorMessage{
|
||||
StatusCode: status,
|
||||
Error: fmt.Errorf("responses stream terminated after %s: %s", lastEvent, errText),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) forwardResponsesStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage, framer *responsesSSEFramer) {
|
||||
if framer == nil {
|
||||
framer = &responsesSSEFramer{}
|
||||
}
|
||||
if isCodexResponsesClientRequest(c) {
|
||||
framer.failureEvent = "response.failed"
|
||||
} else {
|
||||
framer.failureEvent = "error"
|
||||
}
|
||||
writeTerminalError := func(errMsg *interfaces.ErrorMessage) {
|
||||
framer.Flush(c.Writer)
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := responsesStreamErrorText(errMsg, status)
|
||||
h.logResponsesStreamError(c, framer, errMsg)
|
||||
if framer.terminalEvent != "" {
|
||||
return
|
||||
}
|
||||
if isCodexResponsesClientRequest(c) {
|
||||
chunk := handlers.BuildOpenAIResponsesStreamFailedChunk(status, errText, 0)
|
||||
_, _ = fmt.Fprintf(c.Writer, "\nevent: response.failed\ndata: %s\n\n", string(chunk))
|
||||
return
|
||||
}
|
||||
chunk := handlers.BuildOpenAIResponsesStreamErrorChunk(status, errText, 0)
|
||||
_, _ = fmt.Fprintf(c.Writer, "\nevent: error\ndata: %s\n\n", string(chunk))
|
||||
}
|
||||
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
NormalizeTerminalError: sanitizeResponsesStreamErrorMessage,
|
||||
WriteChunk: func(chunk []byte) {
|
||||
framer.WriteChunk(c.Writer, chunk)
|
||||
},
|
||||
ChunkError: func() *interfaces.ErrorMessage {
|
||||
if framer.terminalError != nil {
|
||||
h.logResponsesStreamError(c, framer, framer.terminalError)
|
||||
}
|
||||
return framer.terminalError
|
||||
},
|
||||
WriteTerminalError: writeTerminalError,
|
||||
CloseError: func() *interfaces.ErrorMessage {
|
||||
framer.Flush(c.Writer)
|
||||
if framer.terminalError != nil {
|
||||
return framer.terminalError
|
||||
}
|
||||
if framer.terminalEvent != "" {
|
||||
return nil
|
||||
}
|
||||
lastEvent := framer.lastEvent
|
||||
if lastEvent == "" {
|
||||
lastEvent = "none"
|
||||
}
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
Error: fmt.Errorf("upstream stream closed before a terminal event (last event: %s)", lastEvent),
|
||||
}
|
||||
},
|
||||
WriteDone: func() {
|
||||
framer.Flush(c.Writer)
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
},
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,893 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
const (
|
||||
prematureResponsesStreamModel = "premature-responses-stream-model"
|
||||
initialFailureResponsesModel = "initial-failure-responses-stream-model"
|
||||
emptyResponsesStreamModel = "empty-responses-stream-model"
|
||||
incompleteFirstFrameResponsesModel = "incomplete-first-frame-responses-model"
|
||||
dataOnlyFirstFrameResponsesModel = "data-only-first-frame-responses-model"
|
||||
dataOnlyCleanCloseResponsesModel = "data-only-clean-close-responses-model"
|
||||
sensitiveInitialErrorResponsesModel = "sensitive-initial-error-responses-model"
|
||||
directInitialErrorResponsesModel = "direct-initial-error-responses-model"
|
||||
crossChunkMultilineResponsesModel = "cross-chunk-multiline-responses-model"
|
||||
validThenMalformedResponsesModel = "valid-then-malformed-responses-model"
|
||||
)
|
||||
|
||||
type prematureResponsesStreamExecutor struct{}
|
||||
|
||||
func (*prematureResponsesStreamExecutor) Identifier() string { return "premature-responses-stream" }
|
||||
|
||||
func (*prematureResponsesStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*prematureResponsesStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
if req.Model == directInitialErrorResponsesModel {
|
||||
return nil, &coreexecutor.RequestTerminatedError{
|
||||
HTTPStatus: http.StatusTooManyRequests,
|
||||
Header: http.Header{"Retry-After": []string{"17"}, "X-Plugin-Response": []string{"true"}},
|
||||
Body: []byte(`{"error":{"message":"plugin direct response"}}`),
|
||||
}
|
||||
}
|
||||
chunks := make(chan coreexecutor.StreamChunk, 2)
|
||||
if req.Model == validThenMalformedResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n" +
|
||||
"event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\"\n\n")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == crossChunkMultilineResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.completed\ndata: {\"type\":\"response.completed\",")}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: \"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n\n")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == sensitiveInitialErrorResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New(`{"error":{"type":"server_error","code":"upstream_failed","message":"initial upstream failure: {\"api_key\":\"initial-message-secret\"}"},"debug":{"token":"initial-debug-secret","trace":"` + strings.Repeat("x", 8192) + `"}}`)}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == dataOnlyFirstFrameResponsesModel || req.Model == dataOnlyCleanCloseResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.output_text.delta","delta":"partial"}`)}
|
||||
if req.Model == dataOnlyFirstFrameResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed after data-only frame")}
|
||||
}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == incompleteFirstFrameResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.created")}
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed before first complete frame")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == emptyResponsesStreamModel {
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
if req.Model == initialFailureResponsesModel {
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed before first payload")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n")}
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("unexpected EOF")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (*prematureResponsesStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (*prematureResponsesStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*prematureResponsesStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func TestResponsesHandlerEmitsFailureWhenExecutorStopsAfterPartialOutput(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "premature-responses-stream-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: prematureResponsesStreamModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"premature-responses-stream-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 after stream start; body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "response.output_text.delta") || !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("handler did not preserve partial output and terminal failure: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, "unexpected EOF") {
|
||||
t.Fatalf("handler terminal failure lost executor error: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeResponsesStreamErrorMessageNormalizesSuccessStatus(t *testing.T) {
|
||||
got := sanitizeResponsesStreamErrorMessage(&interfaces.ErrorMessage{StatusCode: http.StatusOK, Error: errors.New("upstream failed")})
|
||||
if got == nil || got.StatusCode != http.StatusInternalServerError {
|
||||
t.Fatalf("sanitized status = %#v, want %d", got, http.StatusInternalServerError)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerCommitsValidFrameBeforeMalformedFrameInSameChunk(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "valid-then-malformed-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: validThenMalformedResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"valid-then-malformed-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), "response.output_text.delta") || !strings.Contains(recorder.Body.String(), "event: response.failed") {
|
||||
t.Fatalf("valid then malformed response status=%d body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerAcceptsMultilineDataAcrossExecutorChunks(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "cross-chunk-multiline-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: crossChunkMultilineResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"cross-chunk-multiline-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), "event: response.completed") {
|
||||
t.Fatalf("cross-chunk multiline response status=%d body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerPreservesDirectResponseBeforeFirstFrame(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "direct-initial-error-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: directInitialErrorResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"direct-initial-error-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusTooManyRequests || recorder.Header().Get("Retry-After") != "17" || recorder.Header().Get("X-Plugin-Response") != "true" {
|
||||
t.Fatalf("direct response status=%d headers=%v body=%q", recorder.Code, recorder.Header(), recorder.Body.String())
|
||||
}
|
||||
if recorder.Body.String() != `{"error":{"message":"plugin direct response"}}` {
|
||||
t.Fatalf("direct response body = %q", recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerSanitizesErrorBeforeFirstFrame(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "sensitive-initial-error-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: sensitiveInitialErrorResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"sensitive-initial-error-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
body := recorder.Body.String()
|
||||
if recorder.Code == http.StatusOK || !strings.Contains(body, "upstream_failed") || !strings.Contains(body, "initial upstream failure") {
|
||||
t.Fatalf("initial error response = status %d body %q", recorder.Code, body)
|
||||
}
|
||||
for _, secret := range []string{"initial-message-secret", "initial-debug-secret"} {
|
||||
if strings.Contains(body, secret) {
|
||||
t.Fatalf("initial error leaked %q: %q", secret, body)
|
||||
}
|
||||
}
|
||||
if len(body) > 4096 {
|
||||
t.Fatalf("initial error response remained unbounded: len=%d", len(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerFlushesDataOnlyFrameBeforeStreamingError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "data-only-first-frame-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: dataOnlyFirstFrameResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"data-only-first-frame-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 after complete data frame; body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "response.output_text.delta") || !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("data-only frame or terminal failure was lost: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerEmitsFailureWhenDataOnlyStreamClosesCleanly(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "data-only-clean-close-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: dataOnlyCleanCloseResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"data-only-clean-close-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 after complete data frame; body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "response.output_text.delta") || !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("clean close did not retain data and emit terminal failure: %q", body)
|
||||
}
|
||||
if strings.Contains(body, "event: response.completed") {
|
||||
t.Fatalf("clean close synthesized completion: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerDoesNotCommitHeadersForIncompleteFirstFrame(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "incomplete-first-frame-responses-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: incompleteFirstFrameResponsesModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"incomplete-first-frame-responses-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Fatalf("incomplete first SSE frame committed HTTP 200: %q", recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "upstream failed before first complete frame") {
|
||||
t.Fatalf("initial frame error was lost: status=%d body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerRejectsStreamClosedBeforeFirstPayload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "empty-responses-stream-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: emptyResponsesStreamModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"empty-responses-stream-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Fatalf("empty upstream stream returned HTTP 200: %q", recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "closed before first payload") {
|
||||
t.Fatalf("empty upstream stream error is unclear: status=%d body=%q", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesHandlerDoesNotLoseErrorBeforeFirstPayload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
executor := &prematureResponsesStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: fmt.Sprintf("initial-failure-responses-stream-auth-%d", i), Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth %d: %v", i, errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: initialFailureResponsesModel}})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"initial-failure-responses-stream-model","input":"hi","stream":true}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Fatalf("request %d lost the buffered initial error and returned HTTP 200: %q", i, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "upstream failed before first payload") {
|
||||
t.Fatalf("request %d lost the initial upstream error: status=%d body=%q", i, recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestForwardResponsesStreamExposesTerminalErrors pins the SSE side: once a
|
||||
// Responses stream has started, every terminal upstream error reaches the client.
|
||||
func TestForwardResponsesStreamExposesTerminalErrors(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
message string
|
||||
wantExposed bool
|
||||
}{
|
||||
{
|
||||
name: "bad request",
|
||||
status: http.StatusBadRequest,
|
||||
message: `{"error":{"type":"invalid_request","code":"cyber_policy","message":"blocked"}}`,
|
||||
wantExposed: true,
|
||||
},
|
||||
{
|
||||
// Observed in production: the same cyber_policy rejection arrives with 502
|
||||
// when it is surfaced through the websocket disconnect channel.
|
||||
name: "cyber policy behind bad gateway status",
|
||||
status: http.StatusBadGateway,
|
||||
message: `{"error":{"type":"invalid_request","code":"cyber_policy","message":"This content was flagged for possible cybersecurity risk.","param":null}}`,
|
||||
wantExposed: true,
|
||||
},
|
||||
{
|
||||
name: "context length exceeded behind bad gateway status",
|
||||
status: http.StatusBadGateway,
|
||||
message: `{"error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"Your input exceeds the context window."}}`,
|
||||
wantExposed: true,
|
||||
},
|
||||
{name: "conflict", status: http.StatusConflict, message: "conflict", wantExposed: true},
|
||||
{name: "message too big", status: http.StatusRequestEntityTooLarge, message: "too large", wantExposed: true},
|
||||
{name: "unprocessable entity", status: http.StatusUnprocessableEntity, message: "invalid input", wantExposed: true},
|
||||
{name: "authentication", status: http.StatusUnauthorized, message: "invalid credential", wantExposed: true},
|
||||
{name: "payment required", status: http.StatusPaymentRequired, message: "insufficient credits", wantExposed: true},
|
||||
{name: "quota error", status: http.StatusTooManyRequests, message: "usage limit reached", wantExposed: true},
|
||||
{name: "request timeout", status: http.StatusRequestTimeout, message: "upstream timeout", wantExposed: true},
|
||||
{name: "transport error", status: http.StatusInternalServerError, message: "unexpected EOF", wantExposed: true},
|
||||
{name: "upstream websocket drop", status: http.StatusInternalServerError,
|
||||
message: `{"error":{"message":"websocket: close 1006 (abnormal closure): unexpected EOF","type":"server_error","code":"internal_server_error"}}`, wantExposed: true},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: tc.status, Error: errors.New(tc.message)}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
body := recorder.Body.String()
|
||||
exposed := strings.Contains(body, `"type":"error"`)
|
||||
if exposed != tc.wantExposed {
|
||||
t.Fatalf("error exposed = %t, want %t: %q", exposed, tc.wantExposed, body)
|
||||
}
|
||||
if exposed && strings.Contains(body, `"error":{`) {
|
||||
t.Fatalf("expected streaming error chunk, got HTTP error body: %q", body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamUsesResponseFailedForCodex(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errors.New(`{"error":{"type":"invalid_request","code":"cyber_policy","message":"blocked"}}`),
|
||||
}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("missing response.failed event: %q", body)
|
||||
}
|
||||
if strings.Contains(body, "event: error") {
|
||||
t.Fatalf("unexpected legacy error event for Codex: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, `"type":"invalid_request"`) || !strings.Contains(body, `"code":"cyber_policy"`) {
|
||||
t.Fatalf("missing nested Codex error detail: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamExposesTransportErrorAfterOutputForCodex(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"))
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("unexpected EOF")}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("transport failure ended without response.failed: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, "unexpected EOF") {
|
||||
t.Fatalf("response.failed lost the upstream error: %q", body)
|
||||
}
|
||||
|
||||
loggedValue, ok := c.Get("API_RESPONSE_ERROR")
|
||||
if !ok {
|
||||
t.Fatal("request log did not retain the stream error")
|
||||
}
|
||||
loggedErrors, ok := loggedValue.([]*interfaces.ErrorMessage)
|
||||
if !ok || len(loggedErrors) != 1 || loggedErrors[0] == nil || loggedErrors[0].Error == nil {
|
||||
t.Fatalf("unexpected request-log errors: %#v", loggedValue)
|
||||
}
|
||||
diagnostic := loggedErrors[0].Error.Error()
|
||||
if !strings.Contains(diagnostic, "response.output_text.delta") || !strings.Contains(diagnostic, "unexpected EOF") {
|
||||
t.Fatalf("request-log diagnostic lacks last event or upstream error: %q", diagnostic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamSanitizesDiagnosticErrorDetails(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
debugSecret := "super-secret-provider-debug-value"
|
||||
messageSecret := "super-secret-provider-message-value"
|
||||
rawError := `{"error":{"type":"server_error","code":"upstream_failed","message":"upstream failed: {\"api_key\":\"` + messageSecret + `\"}"},"debug":{"api_key":"` + debugSecret + `","trace":"` + strings.Repeat("x", 8192) + `"}}`
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"))
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New(rawError)}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "upstream failed") || !strings.Contains(body, "upstream_failed") {
|
||||
t.Fatalf("client error lost safe structured fields: %q", body)
|
||||
}
|
||||
if strings.Contains(body, debugSecret) || strings.Contains(body, messageSecret) {
|
||||
t.Fatalf("client error leaked provider secret: %q", body)
|
||||
}
|
||||
|
||||
loggedValue, ok := c.Get("API_RESPONSE_ERROR")
|
||||
if !ok {
|
||||
t.Fatal("request log did not retain the sanitized stream error")
|
||||
}
|
||||
loggedErrors, ok := loggedValue.([]*interfaces.ErrorMessage)
|
||||
if !ok || len(loggedErrors) != 1 || loggedErrors[0] == nil || loggedErrors[0].Error == nil {
|
||||
t.Fatalf("unexpected request-log errors: %#v", loggedValue)
|
||||
}
|
||||
diagnostic := loggedErrors[0].Error.Error()
|
||||
if strings.Contains(diagnostic, debugSecret) || strings.Contains(diagnostic, messageSecret) || len(diagnostic) > 4096 {
|
||||
t.Fatalf("request-log diagnostic leaked or retained an unbounded upstream body: len=%d diagnostic=%q", len(diagnostic), diagnostic)
|
||||
}
|
||||
if !strings.Contains(diagnostic, "upstream failed") {
|
||||
t.Fatalf("sanitized request-log diagnostic lost upstream message: %q", diagnostic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamPreservesNestedResponseError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"))
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New(`{"type":"response.failed","response":{"error":{"type":"server_error","code":"upstream_failed","message":"nested response failure","param":"input"}}}`)}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
body := recorder.Body.String()
|
||||
for _, want := range []string{"nested response failure", "upstream_failed", "server_error"} {
|
||||
if !strings.Contains(body, want) {
|
||||
t.Fatalf("response.failed lost nested response error field %q: %q", want, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamSanitizesLastEventDiagnostic(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
eventSecret := "event-secret-value"
|
||||
eventName := "custom-event-Bearer " + eventSecret + strings.Repeat("x", 1024)
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: "+eventName+"\ndata: {\"message\":\"partial\"}\n\n"))
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("unexpected EOF")}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
loggedValue, ok := c.Get("API_RESPONSE_ERROR")
|
||||
if !ok {
|
||||
t.Fatal("request log did not retain the stream error")
|
||||
}
|
||||
loggedErrors, ok := loggedValue.([]*interfaces.ErrorMessage)
|
||||
if !ok || len(loggedErrors) != 1 || loggedErrors[0] == nil || loggedErrors[0].Error == nil {
|
||||
t.Fatalf("unexpected request-log errors: %#v", loggedValue)
|
||||
}
|
||||
diagnostic := loggedErrors[0].Error.Error()
|
||||
if strings.Contains(diagnostic, eventSecret) || len(diagnostic) > 1024 {
|
||||
t.Fatalf("last-event diagnostic leaked or remained unbounded: len=%d diagnostic=%q", len(diagnostic), diagnostic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamSanitizesPayloadErrorsAndStopsAtFailure(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
frame string
|
||||
}{
|
||||
{
|
||||
name: "event error with payload type",
|
||||
frame: "event: error\ndata: {\"type\":\"provider.error\",\"error\":{\"code\":\"failed\",\"message\":\"token=payload-secret\"}}\n\n",
|
||||
},
|
||||
{
|
||||
name: "typed nested error",
|
||||
frame: "data: {\"type\":\"provider.error\",\"error\":{\"code\":\"failed\",\"message\":\"token=payload-secret\"}}\n\n",
|
||||
},
|
||||
{
|
||||
name: "top level error fields",
|
||||
frame: "data: {\"code\":\"failed\",\"message\":\"token=payload-secret\"}\n\n",
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
data := make(chan []byte, 1)
|
||||
data <- []byte(tc.frame + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\"}}\n\n")
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
close(errs)
|
||||
var canceled error
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { canceled = err }, data, errs, &responsesSSEFramer{})
|
||||
body := recorder.Body.String()
|
||||
if canceled == nil {
|
||||
t.Fatalf("payload error canceled with nil: %q", body)
|
||||
}
|
||||
if strings.Contains(body, "payload-secret") || strings.Contains(body, "event: response.completed") {
|
||||
t.Fatalf("payload error leaked or accepted later completion: %q", body)
|
||||
}
|
||||
if strings.Count(body, "event: response.failed") != 1 || !strings.Contains(body, "[REDACTED]") {
|
||||
t.Fatalf("payload error was not converted to one sanitized response.failed: %q", body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamReportsDataOnlyErrorFlushedAtEOF(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
data := make(chan []byte, 1)
|
||||
data <- []byte(`data: {"type":"error","error":{"message":"failed at EOF"}}`)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
close(errs)
|
||||
var canceled error
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { canceled = err }, data, errs, &responsesSSEFramer{})
|
||||
|
||||
if canceled == nil || !strings.Contains(canceled.Error(), "failed at EOF") {
|
||||
t.Fatalf("EOF error cancel = %v, body=%q", canceled, recorder.Body.String())
|
||||
}
|
||||
if strings.Count(recorder.Body.String(), "event: response.failed") != 1 {
|
||||
t.Fatalf("EOF error terminal output = %q", recorder.Body.String())
|
||||
}
|
||||
if _, okLog := c.Get("API_RESPONSE_ERROR"); !okLog {
|
||||
t.Fatal("EOF error was not retained in request diagnostics")
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamDoesNotAppendFailureAfterTerminalEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{RequestLog: true}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n\n"))
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("unexpected EOF after completion")}
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
body := recorder.Body.String()
|
||||
if strings.Contains(body, "event: response.failed") || strings.Contains(body, "event: error") {
|
||||
t.Fatalf("stream appended a second terminal event after response.completed: %q", body)
|
||||
}
|
||||
|
||||
loggedValue, ok := c.Get("API_RESPONSE_ERROR")
|
||||
if !ok {
|
||||
t.Fatal("request log did not retain the post-terminal upstream error")
|
||||
}
|
||||
loggedErrors, ok := loggedValue.([]*interfaces.ErrorMessage)
|
||||
if !ok || len(loggedErrors) != 1 || loggedErrors[0] == nil || loggedErrors[0].Error == nil {
|
||||
t.Fatalf("unexpected request-log errors: %#v", loggedValue)
|
||||
}
|
||||
diagnostic := loggedErrors[0].Error.Error()
|
||||
if !strings.Contains(diagnostic, "response.completed") || !strings.Contains(diagnostic, "unexpected EOF after completion") {
|
||||
t.Fatalf("request-log diagnostic lacks terminal event or upstream error: %q", diagnostic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamFailsWhenUpstreamClosesWithoutTerminalEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", "Codex Desktop/26.803.41515")
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(c.Writer, []byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"))
|
||||
data := make(chan []byte)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, framer)
|
||||
body := recorder.Body.String()
|
||||
if !strings.Contains(body, "event: response.failed") {
|
||||
t.Fatalf("unterminated stream ended without response.failed: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, "closed before a terminal event") {
|
||||
t.Fatalf("response.failed does not explain the premature close: %q", body)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,314 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func newResponsesStreamTestHandler(t *testing.T) (*OpenAIResponsesAPIHandler, *httptest.ResponseRecorder, *gin.Context, http.Flusher) {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatalf("expected gin writer to implement http.Flusher")
|
||||
}
|
||||
|
||||
return h, recorder, c, flusher
|
||||
}
|
||||
|
||||
func TestResponsesSSEFramerWaitsForEventFieldAfterData(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
framer := &responsesSSEFramer{}
|
||||
|
||||
framer.WriteChunk(&output, []byte(`data: {"response":{"id":"resp-1","status":"completed"}}`))
|
||||
if output.Len() != 0 {
|
||||
t.Fatalf("framer emitted data before a following event field arrived: %q", output.String())
|
||||
}
|
||||
|
||||
framer.WriteChunk(&output, []byte("event: response.completed"))
|
||||
if framer.terminalEvent != "response.completed" {
|
||||
t.Fatalf("terminal event = %q, want response.completed", framer.terminalEvent)
|
||||
}
|
||||
got := output.String()
|
||||
if !strings.Contains(got, "data: ") || !strings.Contains(got, "event: response.completed") {
|
||||
t.Fatalf("framer did not preserve data-before-event fields in one frame: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesSSEFramerFlushesMultilineDataWithoutDelimiter(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
framer := &responsesSSEFramer{}
|
||||
chunk := []byte("event: response.completed\n" +
|
||||
"data: {\"type\":\"response.completed\",\n" +
|
||||
"data: \"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}")
|
||||
framer.WriteChunk(&output, chunk)
|
||||
framer.Flush(&output)
|
||||
|
||||
if framer.terminalEvent != "response.completed" || !strings.Contains(output.String(), "response.completed") {
|
||||
t.Fatalf("multiline data-only terminal frame was dropped: terminal=%q output=%q", framer.terminalEvent, output.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesSSEFramerUsesPayloadErrorOverCompletedEvent(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
framer := &responsesSSEFramer{failureEvent: "response.failed"}
|
||||
framer.WriteChunk(&output, []byte("data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\nevent: response.completed\n\n"))
|
||||
|
||||
if framer.terminalEvent != "response.failed" || strings.Contains(output.String(), "event: response.completed") {
|
||||
t.Fatalf("payload error was overridden by completed event: terminal=%q output=%q", framer.terminalEvent, output.String())
|
||||
}
|
||||
if strings.Count(output.String(), "event: response.failed") != 1 {
|
||||
t.Fatalf("payload error output = %q, want one response.failed", output.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesSSEFramerUsesErrorEventOverPayloadType(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
framer := &responsesSSEFramer{}
|
||||
framer.WriteChunk(&output, []byte("event: error\ndata: {\"type\":\"provider.error\",\"message\":\"failed\"}\n\n"))
|
||||
if framer.terminalEvent != "error" {
|
||||
t.Fatalf("terminal event = %q, want error", framer.terminalEvent)
|
||||
}
|
||||
|
||||
framer = &responsesSSEFramer{}
|
||||
framer.WriteChunk(&output, []byte("data: {\"response\":{\"error\":{\"message\":\"failed\"}}}\n\n"))
|
||||
if framer.terminalEvent != "error" {
|
||||
t.Fatalf("nested response error terminal event = %q, want error", framer.terminalEvent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamSeparatesDataOnlySSEChunks(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 2)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte("data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"function_call\",\"arguments\":\"{}\"}}")
|
||||
data <- []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}")
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
body := recorder.Body.String()
|
||||
parts := strings.Split(strings.TrimSpace(body), "\n\n")
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 SSE events, got %d. Body: %q", len(parts), body)
|
||||
}
|
||||
|
||||
expectedPart1 := "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"function_call\",\"arguments\":\"{}\"}}"
|
||||
if parts[0] != expectedPart1 {
|
||||
t.Errorf("unexpected first event.\nGot: %q\nWant: %q", parts[0], expectedPart1)
|
||||
}
|
||||
|
||||
expectedPart2 := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[{\"type\":\"function_call\",\"arguments\":\"{}\"}]}}"
|
||||
if parts[1] != expectedPart2 {
|
||||
t.Errorf("unexpected second event.\nGot: %q\nWant: %q", parts[1], expectedPart2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamRepairsEmptyCompletedOutputFromDoneItems(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 3)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte(`data: {"type":"response.output_item.done","output_index":0,"item":{"type":"reasoning","id":"rs-1","summary":[]}}`)
|
||||
data <- []byte(`data: {"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"shell","arguments":"{\"cmd\":\"pwd\"}","status":"completed"}}`)
|
||||
data <- []byte(`data: {"type":"response.completed","response":{"id":"resp-1","output":[]}}`)
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
|
||||
if len(parts) != 3 {
|
||||
t.Fatalf("expected 3 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
|
||||
}
|
||||
|
||||
payload := strings.TrimPrefix(parts[2], "data: ")
|
||||
output := gjson.Get(payload, "response.output")
|
||||
if !output.IsArray() || len(output.Array()) != 2 {
|
||||
t.Fatalf("expected repaired completed output with 2 items, got %s", output.Raw)
|
||||
}
|
||||
if got := gjson.Get(payload, "response.output.1.name").String(); got != "shell" {
|
||||
t.Fatalf("expected function_call name to be preserved, got %q in %s", got, payload)
|
||||
}
|
||||
if got := gjson.Get(payload, "response.output.1.arguments").String(); got != `{"cmd":"pwd"}` {
|
||||
t.Fatalf("expected function_call arguments to be preserved, got %q in %s", got, payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamRepairsMixedIndexedAndUnindexedDoneItems(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 3)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte(`data: {"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"shell","arguments":"{}","status":"completed"}}`)
|
||||
data <- []byte(`data: {"type":"response.output_item.done","item":{"type":"message","id":"msg-1","role":"assistant","content":[{"type":"output_text","text":"done"}]}}`)
|
||||
data <- []byte(`data: {"type":"response.completed","response":{"id":"resp-1","output":[]}}`)
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
|
||||
if len(parts) != 3 {
|
||||
t.Fatalf("expected 3 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
|
||||
}
|
||||
|
||||
payload := strings.TrimPrefix(parts[2], "data: ")
|
||||
output := gjson.Get(payload, "response.output")
|
||||
if !output.IsArray() || len(output.Array()) != 2 {
|
||||
t.Fatalf("expected repaired completed output with 2 items, got %s", output.Raw)
|
||||
}
|
||||
if got := gjson.Get(payload, "response.output.0.name").String(); got != "shell" {
|
||||
t.Fatalf("expected indexed function_call to be preserved first, got %q in %s", got, payload)
|
||||
}
|
||||
if got := gjson.Get(payload, "response.output.1.id").String(); got != "msg-1" {
|
||||
t.Fatalf("expected unindexed message to be appended, got %q in %s", got, payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamRepairsMultilineCompletedOutputAsSSEDataLines(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 2)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte(`data: {"type":"response.output_item.done","item":{"type":"function_call","arguments":"{}"}}`)
|
||||
data <- []byte("data: {\"type\":\"response.completed\",\ndata: \"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n")
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
|
||||
}
|
||||
|
||||
completedFrame := []byte(parts[1])
|
||||
for _, line := range strings.Split(parts[1], "\n") {
|
||||
if line != "" && !strings.HasPrefix(line, "data: ") {
|
||||
t.Fatalf("expected every completed payload line to be an SSE data line, got %q in %q", line, parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
payload, ok := responsesSSEDataPayload(completedFrame)
|
||||
if !ok {
|
||||
t.Fatalf("expected completed frame to contain data payload: %q", parts[1])
|
||||
}
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if !output.IsArray() || len(output.Array()) != 1 {
|
||||
t.Fatalf("expected repaired completed output with 1 item, got %s from %q", output.Raw, payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamReassemblesSplitSSEEventChunks(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 3)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte("event: response.created")
|
||||
data <- []byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}")
|
||||
data <- []byte("\n")
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
got := recorder.Body.String()
|
||||
wantPrefix := "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}\n\n"
|
||||
if !strings.HasPrefix(got, wantPrefix) {
|
||||
t.Fatalf("unexpected split-event framing.\nGot: %q\nWant prefix: %q", got, wantPrefix)
|
||||
}
|
||||
if !strings.Contains(got, "event: error") {
|
||||
t.Fatalf("unterminated framing test stream did not end with an error: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamPreservesValidFullSSEEventChunks(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 1)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
chunk := []byte("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}\n\n")
|
||||
data <- chunk
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
got := recorder.Body.String()
|
||||
if !strings.HasPrefix(got, string(chunk)) {
|
||||
t.Fatalf("unexpected full-event framing.\nGot: %q\nWant prefix: %q", got, string(chunk))
|
||||
}
|
||||
if !strings.Contains(got, "event: error") {
|
||||
t.Fatalf("unterminated framing test stream did not end with an error: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamBuffersSplitDataPayloadChunks(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 2)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte("data: {\"type\":\"response.created\"")
|
||||
data <- []byte(",\"response\":{\"id\":\"resp-1\"}}")
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
got := recorder.Body.String()
|
||||
wantPrefix := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}\n\n"
|
||||
if !strings.HasPrefix(got, wantPrefix) {
|
||||
t.Fatalf("unexpected split-data framing.\nGot: %q\nWant prefix: %q", got, wantPrefix)
|
||||
}
|
||||
if !strings.Contains(got, "event: error") {
|
||||
t.Fatalf("unterminated framing test stream did not end with an error: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesSSENeedsLineBreakSkipsChunksThatAlreadyStartWithNewline(t *testing.T) {
|
||||
if responsesSSENeedsLineBreak([]byte("event: response.created"), []byte("\n")) {
|
||||
t.Fatal("expected no injected newline before newline-only chunk")
|
||||
}
|
||||
if responsesSSENeedsLineBreak([]byte("event: response.created"), []byte("\r\n")) {
|
||||
t.Fatal("expected no injected newline before CRLF chunk")
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardResponsesStreamDropsIncompleteTrailingDataChunkOnFlush(t *testing.T) {
|
||||
h, recorder, c, flusher := newResponsesStreamTestHandler(t)
|
||||
|
||||
data := make(chan []byte, 1)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte("data: {\"type\":\"response.created\"")
|
||||
close(data)
|
||||
close(errs)
|
||||
|
||||
h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
|
||||
|
||||
got := recorder.Body.String()
|
||||
if strings.Contains(got, `data: {"type":"response.created"`) {
|
||||
t.Fatalf("incomplete trailing data was not dropped on flush: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "event: error") {
|
||||
t.Fatalf("unterminated framing test stream did not end with an error: %q", got)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,199 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
multiagentv2 "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/optimize-multi-agent-v2"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestPrepareCodexMultiAgentV2ToolsAtResponsesBoundary(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, nil)
|
||||
handler := NewOpenAIResponsesAPIHandler(base)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
|
||||
ginContext, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginContext.Request = request
|
||||
|
||||
payload := []byte(`{
|
||||
"tools":[{"type":"namespace","name":"collaboration","tools":[
|
||||
{"type":"function","name":"spawn_agent","description":"Spawns an agent.","parameters":{"properties":{"message":{"encrypted":true}}}},
|
||||
{"type":"function","name":"send_message","parameters":{"properties":{"message":{"encrypted":true}}}}
|
||||
]}]
|
||||
}`)
|
||||
got := handler.prepareCodexMultiAgentV2Tools(ginContext, payload)
|
||||
|
||||
if namespace := gjson.GetBytes(got, "tools.0.name").String(); namespace != "collaboration" {
|
||||
t.Fatalf("namespace = %q, want collaboration", namespace)
|
||||
}
|
||||
for _, path := range []string{"tools.0.tools.0", "tools.0.tools.1"} {
|
||||
if encrypted := gjson.GetBytes(got, path+".parameters.properties.message.encrypted"); encrypted.Exists() {
|
||||
t.Fatalf("%s message.encrypted was not removed: %s", path, encrypted.Raw)
|
||||
}
|
||||
}
|
||||
prepared, exists := ginContext.Get(multiagentv2.CodexMultiAgentV2ToolsPreparedContextKey)
|
||||
if !exists || prepared != true {
|
||||
t.Fatalf("prepared marker = %#v, want true", prepared)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesPreparesCodexMultiAgentV2ToolsForHTTPAndSSE(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, stream := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
|
||||
executor := &responsesMultiAgentCaptureExecutor{}
|
||||
handler, modelID := newResponsesMultiAgentTestHandler(t, executor)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", handler.Responses)
|
||||
|
||||
payload := fmt.Sprintf(`{"model":%q,"stream":%t,"tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent","description":"Spawns an agent.","parameters":{"properties":{"message":{"encrypted":true}}}}]}]}`, modelID, stream)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewBufferString(payload))
|
||||
request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
|
||||
payloads := executor.Payloads()
|
||||
if len(payloads) != 1 {
|
||||
t.Fatalf("captured payload count = %d, want 1", len(payloads))
|
||||
}
|
||||
captured := payloads[0]
|
||||
if encrypted := gjson.GetBytes(captured, "tools.0.tools.0.parameters.properties.message.encrypted"); encrypted.Exists() {
|
||||
t.Fatalf("message.encrypted was not removed: %s", captured)
|
||||
}
|
||||
if namespace := gjson.GetBytes(captured, "tools.0.name").String(); namespace != "collaboration" {
|
||||
t.Fatalf("namespace = %q, want collaboration", namespace)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type responsesMultiAgentCaptureExecutor struct {
|
||||
websocketDirectCaptureExecutor
|
||||
}
|
||||
|
||||
func (e *responsesMultiAgentCaptureExecutor) Execute(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
e.mu.Lock()
|
||||
e.payloads = append(e.payloads, bytes.Clone(req.Payload))
|
||||
e.mu.Unlock()
|
||||
return coreexecutor.Response{Payload: []byte(`{"id":"resp-1","output":[]}`)}, nil
|
||||
}
|
||||
|
||||
func (e *responsesMultiAgentCaptureExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
e.mu.Lock()
|
||||
e.payloads = append(e.payloads, bytes.Clone(req.Payload))
|
||||
e.mu.Unlock()
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func newResponsesMultiAgentTestHandler(t *testing.T, executor *responsesMultiAgentCaptureExecutor) (*OpenAIResponsesAPIHandler, string) {
|
||||
t.Helper()
|
||||
|
||||
modelID := "responses-multi-agent-test-model"
|
||||
authID := "responses-multi-agent-test-auth"
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: authID, Provider: "codex", Status: coreauth.StatusActive, ProxyURL: "direct"}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("Register auth: %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(authID, auth.Provider, []*registry.ModelInfo{{ID: modelID}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(authID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, manager)
|
||||
return NewOpenAIResponsesAPIHandler(base), modelID
|
||||
}
|
||||
|
||||
func TestResponsesWebsocketPreparesCodexMultiAgentV2Tools(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &websocketDirectCaptureExecutor{provider: "codex"}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
auth := &coreauth.Auth{ID: "responses-multi-agent-ws-auth", Provider: "codex", Status: coreauth.StatusActive, ProxyURL: "direct"}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("Register auth: %v", errRegister)
|
||||
}
|
||||
modelID := "responses-multi-agent-ws-model"
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelID}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, manager)
|
||||
handler := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.GET("/v1/responses", handler.ResponsesWebsocket)
|
||||
server := httptest.NewServer(router)
|
||||
defer server.Close()
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses"
|
||||
conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, http.Header{"User-Agent": []string{"codex_cli_rs/0.144.1"}})
|
||||
if errDial != nil {
|
||||
t.Fatalf("dial websocket: %v", errDial)
|
||||
}
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
request := fmt.Sprintf(`{"type":"response.create","model":%q,"input":[],"tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent","description":"Spawns an agent.","parameters":{"properties":{"message":{"encrypted":true}}}}]}]}`, modelID)
|
||||
if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil {
|
||||
t.Fatalf("write websocket request: %v", errWrite)
|
||||
}
|
||||
if _, _, errRead := conn.ReadMessage(); errRead != nil {
|
||||
t.Fatalf("read websocket response: %v", errRead)
|
||||
}
|
||||
|
||||
payloads := executor.Payloads()
|
||||
if len(payloads) != 1 {
|
||||
t.Fatalf("captured payload count = %d, want 1", len(payloads))
|
||||
}
|
||||
captured := payloads[0]
|
||||
if encrypted := gjson.GetBytes(captured, "tools.0.tools.0.parameters.properties.message.encrypted"); encrypted.Exists() {
|
||||
t.Fatalf("message.encrypted was not removed: %s", captured)
|
||||
}
|
||||
if namespace := gjson.GetBytes(captured, "tools.0.name").String(); namespace != "collaboration" {
|
||||
t.Fatalf("namespace = %q, want collaboration", namespace)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareCodexMultiAgentV2ToolsAtResponsesBoundarySkipsOtherClients(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{CodexOptimizeMultiAgentV2: true}, nil)
|
||||
handler := NewOpenAIResponsesAPIHandler(base)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
request.Header.Set("User-Agent", "curl/8.7.1")
|
||||
ginContext, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ginContext.Request = request
|
||||
|
||||
payload := []byte(`{"tools":[{"type":"function","name":"send_message","parameters":{"properties":{"message":{"encrypted":true}}}}]}`)
|
||||
got := handler.prepareCodexMultiAgentV2Tools(ginContext, payload)
|
||||
|
||||
if string(got) != string(payload) {
|
||||
t.Fatalf("other client payload changed: %s", got)
|
||||
}
|
||||
if _, exists := ginContext.Get(multiagentv2.CodexMultiAgentV2ToolsPreparedContextKey); exists {
|
||||
t.Fatal("other client unexpectedly received prepared marker")
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,86 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestOpenAIResponsesForwardsInvalidReasoningEncryptedContentToExecutor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactCaptureExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth := &coreauth.Auth{ID: "signature-auth-responses", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("Register auth: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "test-signature-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", h.Responses)
|
||||
|
||||
body := `{"model":"test-signature-model","stream":false,"input":[{"id":"rs_bad","type":"reasoning","encrypted_content":"gAAAAABqFTIa\u2026abc","summary":[]}]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", resp.Code, http.StatusOK, resp.Body.String())
|
||||
}
|
||||
if executor.calls != 1 {
|
||||
t.Fatalf("executor calls = %d, want 1", executor.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactForwardsInvalidReasoningEncryptedContentToExecutor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
executor := &compactCaptureExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth := &coreauth.Auth{ID: "signature-auth-compact", Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("Register auth: %v", err)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "test-signature-compact-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewOpenAIResponsesAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses/compact", h.Compact)
|
||||
|
||||
body := `{"model":"test-signature-compact-model","input":[{"id":"rs_bad","type":"reasoning","encrypted_content":"bad","summary":[]}]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses/compact", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp := httptest.NewRecorder()
|
||||
router.ServeHTTP(resp, req)
|
||||
|
||||
if resp.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", resp.Code, http.StatusOK, resp.Body.String())
|
||||
}
|
||||
if executor.calls != 1 {
|
||||
t.Fatalf("executor calls = %d, want 1", executor.calls)
|
||||
}
|
||||
if executor.alt != "responses/compact" {
|
||||
t.Fatalf("alt = %q, want responses/compact", executor.alt)
|
||||
}
|
||||
}
|
||||
704
backend/sdk/api/handlers/openai/openai_responses_websocket.go
Normal file
704
backend/sdk/api/handlers/openai/openai_responses_websocket.go
Normal file
|
|
@ -0,0 +1,704 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
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
|
||||
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:"
|
||||
)
|
||||
|
||||
var responsesWebsocketUpgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 4096,
|
||||
WriteBufferSize: 4096,
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
// writeWebsocketCloseForUpstreamError mirrors transport-level upstream close
|
||||
// codes to the downstream WebSocket client before the connection is torn down.
|
||||
// Without this the client only observes an abnormal closure (1006) and cannot
|
||||
// apply its own close-code based handling (e.g. falling back to SSE on 1009).
|
||||
func writeWebsocketCloseForUpstreamError(conn *websocket.Conn, err error) (bool, error) {
|
||||
if conn == nil {
|
||||
return false, nil
|
||||
}
|
||||
matched, payload := websocketClosePayloadForUpstreamError(err)
|
||||
if !matched {
|
||||
return false, nil
|
||||
}
|
||||
return true, conn.WriteControl(websocket.CloseMessage, payload, time.Time{})
|
||||
}
|
||||
|
||||
func websocketClosePayloadForUpstreamError(err error) (bool, []byte) {
|
||||
if err == nil {
|
||||
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
|
||||
if errors.As(err, &closeErr) && closeErr.Code == websocket.CloseMessageTooBig {
|
||||
code = closeErr.Code
|
||||
reason = closeErr.Text
|
||||
} else {
|
||||
type statusCoder interface {
|
||||
StatusCode() int
|
||||
}
|
||||
var statusErr statusCoder
|
||||
if !errors.As(err, &statusErr) || statusErr.StatusCode() != http.StatusRequestEntityTooLarge ||
|
||||
gjson.Get(errText, "error.code").String() != "message_too_big" {
|
||||
return false, nil
|
||||
}
|
||||
code = websocket.CloseMessageTooBig
|
||||
reason = strings.TrimSpace(gjson.Get(errText, "error.message").String())
|
||||
}
|
||||
if reason == "" {
|
||||
reason = "message too big"
|
||||
}
|
||||
reason = truncateWebsocketCloseReason(reason, wsCloseReasonMaxBytes)
|
||||
return true, websocket.FormatCloseMessage(code, reason)
|
||||
}
|
||||
|
||||
type responsesWebsocketWriter struct {
|
||||
conn *websocket.Conn
|
||||
writeMu sync.Mutex
|
||||
closing atomic.Bool
|
||||
}
|
||||
|
||||
func newResponsesWebsocketWriter(conn *websocket.Conn) *responsesWebsocketWriter {
|
||||
return &responsesWebsocketWriter{conn: conn}
|
||||
}
|
||||
|
||||
// closeForUpstreamError sends a best-effort close frame without waiting behind
|
||||
// an active downstream data writer. If a data write already owns writeMu, the
|
||||
// connection is closed immediately so the blocked writer and session can exit.
|
||||
func (w *responsesWebsocketWriter) closeForUpstreamError(err error) (bool, error) {
|
||||
if w == nil || w.conn == nil {
|
||||
return false, nil
|
||||
}
|
||||
matched, payload := websocketClosePayloadForUpstreamError(err)
|
||||
if !matched {
|
||||
return false, nil
|
||||
}
|
||||
if !w.closing.CompareAndSwap(false, true) {
|
||||
return true, nil
|
||||
}
|
||||
if !w.writeMu.TryLock() {
|
||||
return true, w.conn.Close()
|
||||
}
|
||||
defer w.writeMu.Unlock()
|
||||
|
||||
errWrite := w.conn.WriteControl(websocket.CloseMessage, payload, time.Time{})
|
||||
errClose := w.conn.Close()
|
||||
if errWrite != nil {
|
||||
return true, errWrite
|
||||
}
|
||||
return true, errClose
|
||||
}
|
||||
|
||||
func (w *responsesWebsocketWriter) closeWithoutError() (bool, error) {
|
||||
if w == nil || w.conn == nil {
|
||||
return false, nil
|
||||
}
|
||||
if !w.closing.CompareAndSwap(false, true) {
|
||||
return false, nil
|
||||
}
|
||||
return true, w.conn.Close()
|
||||
}
|
||||
|
||||
func (w *responsesWebsocketWriter) closeWithPayload(payload []byte) (bool, error) {
|
||||
if w == nil || w.conn == nil {
|
||||
return false, nil
|
||||
}
|
||||
if !w.closing.CompareAndSwap(false, true) {
|
||||
return false, nil
|
||||
}
|
||||
if !w.writeMu.TryLock() {
|
||||
return false, w.conn.Close()
|
||||
}
|
||||
defer w.writeMu.Unlock()
|
||||
|
||||
errWrite := w.conn.WriteMessage(websocket.TextMessage, payload)
|
||||
errClose := w.conn.Close()
|
||||
if errWrite != nil {
|
||||
return false, errWrite
|
||||
}
|
||||
return true, errClose
|
||||
}
|
||||
|
||||
func (w *responsesWebsocketWriter) closeForUpstreamDisconnect(err error) {
|
||||
if w == nil || w.conn == nil {
|
||||
return
|
||||
}
|
||||
if matched, _ := w.closeForUpstreamError(err); matched {
|
||||
return
|
||||
}
|
||||
|
||||
errMsg := handlers.ExecutionErrorMessage(err)
|
||||
if !shouldExposeResponsesUpstreamError(errMsg) {
|
||||
_, _ = w.closeWithoutError()
|
||||
return
|
||||
}
|
||||
payload, errBuild := buildResponsesWebsocketErrorPayload(errMsg)
|
||||
if errBuild != nil {
|
||||
_, _ = w.closeWithoutError()
|
||||
return
|
||||
}
|
||||
wrote, errClose := w.closeWithPayload(payload)
|
||||
if wrote {
|
||||
log.Infof(
|
||||
"responses websocket: downstream_out disconnect_error event=%s payload=%s",
|
||||
websocketPayloadEventType(payload),
|
||||
websocketPayloadPreview(payload),
|
||||
)
|
||||
}
|
||||
if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
|
||||
log.Debugf("responses websocket: upstream disconnect close failed: %v", errClose)
|
||||
}
|
||||
}
|
||||
|
||||
// isWebsocketConnectionClosedError reports whether the error only means the
|
||||
// connection was already torn down. These are expected during shutdown races
|
||||
// (the proxy closes after sending a terminal frame, or the client hangs up mid
|
||||
// write) and must not be logged as proxy failures.
|
||||
func isWebsocketConnectionClosedError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, net.ErrClosed) || errors.Is(err, websocket.ErrCloseSent) {
|
||||
return true
|
||||
}
|
||||
return strings.Contains(err.Error(), "use of closed network connection")
|
||||
}
|
||||
|
||||
func truncateWebsocketCloseReason(reason string, maxBytes int) string {
|
||||
if maxBytes <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(reason) <= maxBytes && utf8.ValidString(reason) {
|
||||
return reason
|
||||
}
|
||||
|
||||
// Decode from the front so work and output stay bounded by maxBytes.
|
||||
var truncated strings.Builder
|
||||
truncated.Grow(min(len(reason), maxBytes))
|
||||
remaining := maxBytes
|
||||
runeErrorSize := utf8.RuneLen(utf8.RuneError)
|
||||
for len(reason) > 0 && remaining > 0 {
|
||||
r, size := utf8.DecodeRuneInString(reason)
|
||||
if r == utf8.RuneError && size == 1 {
|
||||
if runeErrorSize > remaining {
|
||||
break
|
||||
}
|
||||
truncated.WriteRune(utf8.RuneError)
|
||||
reason = reason[1:]
|
||||
remaining -= runeErrorSize
|
||||
continue
|
||||
}
|
||||
if size > remaining {
|
||||
break
|
||||
}
|
||||
truncated.WriteString(reason[:size])
|
||||
reason = reason[size:]
|
||||
remaining -= size
|
||||
}
|
||||
return truncated.String()
|
||||
}
|
||||
|
||||
// ResponsesWebsocket handles websocket requests for /v1/responses.
|
||||
// It accepts `response.create` and `response.append` requests and streams
|
||||
// response events back as JSON websocket text messages.
|
||||
func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
|
||||
conn, err := responsesWebsocketUpgrader.Upgrade(c.Writer, c.Request, websocketUpgradeHeaders(c.Request))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
writer := newResponsesWebsocketWriter(conn)
|
||||
passthroughSessionID := uuid.NewString()
|
||||
downstreamSessionKey := websocketDownstreamSessionKey(c.Request)
|
||||
retainResponsesWebsocketToolCaches(downstreamSessionKey)
|
||||
clientIP := websocketClientAddress(c)
|
||||
log.Infof("responses websocket: client connected id=%s remote=%s", passthroughSessionID, clientIP)
|
||||
|
||||
requestLogEnabled := h != nil && h.Cfg != nil && h.Cfg.RequestLog
|
||||
wsTimelineLog := newWebsocketTimelineLog(requestLogEnabled, websocketTimelineSourceFromContext(c))
|
||||
|
||||
wsDone := make(chan struct{})
|
||||
defer close(wsDone)
|
||||
|
||||
if h != nil && h.AuthManager != nil {
|
||||
type upstreamDisconnectSubscriber interface {
|
||||
UpstreamDisconnectChan(sessionID string) <-chan error
|
||||
}
|
||||
for _, provider := range []string{"codex", "xai"} {
|
||||
exec, ok := h.AuthManager.Executor(provider)
|
||||
if !ok || exec == nil {
|
||||
continue
|
||||
}
|
||||
if subscriber, ok := exec.(upstreamDisconnectSubscriber); ok && subscriber != nil {
|
||||
disconnectCh := subscriber.UpstreamDisconnectChan(passthroughSessionID)
|
||||
if disconnectCh != nil {
|
||||
go func() {
|
||||
select {
|
||||
case <-wsDone:
|
||||
return
|
||||
case disconnectErr := <-disconnectCh:
|
||||
writer.closeForUpstreamDisconnect(disconnectErr)
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var wsTerminateErr error
|
||||
defer func() {
|
||||
releaseResponsesWebsocketToolCaches(downstreamSessionKey)
|
||||
if wsTerminateErr != nil {
|
||||
appendWebsocketTimelineDisconnect(wsTimelineLog, wsTerminateErr, time.Now())
|
||||
// log.Infof("responses websocket: session closing id=%s reason=%v", passthroughSessionID, wsTerminateErr)
|
||||
} else {
|
||||
log.Infof("responses websocket: session closing id=%s", passthroughSessionID)
|
||||
}
|
||||
if h != nil && h.AuthManager != nil {
|
||||
h.AuthManager.CloseExecutionSession(passthroughSessionID)
|
||||
log.Infof("responses websocket: upstream execution session closed id=%s", passthroughSessionID)
|
||||
}
|
||||
wsTimelineLog.SetContext(c)
|
||||
if errClose := conn.Close(); errClose != nil && !isWebsocketConnectionClosedError(errClose) {
|
||||
log.Warnf("responses websocket: close connection error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
var lastRequest []byte
|
||||
lastResponseOutput := []byte("[]")
|
||||
lastResponseID := ""
|
||||
var lastResponsePendingToolCallIDs []string
|
||||
pinnedAuthID := ""
|
||||
// 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
|
||||
}
|
||||
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)
|
||||
if authID == "" || !ok || auth == nil {
|
||||
return
|
||||
}
|
||||
pinnedAuthID = authID
|
||||
providerKey := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
_, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName))
|
||||
if providerKey != "" {
|
||||
pinnedAuthByProvider[providerKey] = responsesWebsocketPinnedAuthState{authID: authID, modelKey: modelKey}
|
||||
}
|
||||
}
|
||||
forgetPinnedAuth := func() {
|
||||
for providerKey, state := range pinnedAuthByProvider {
|
||||
if state.authID == pinnedAuthID {
|
||||
delete(pinnedAuthByProvider, providerKey)
|
||||
}
|
||||
}
|
||||
pinnedAuthID = ""
|
||||
}
|
||||
|
||||
for {
|
||||
msgType, payload, errReadMessage := conn.ReadMessage()
|
||||
if errReadMessage != nil {
|
||||
wsTerminateErr = errReadMessage
|
||||
if websocket.IsCloseError(errReadMessage, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
|
||||
log.Infof("responses websocket: client disconnected id=%s error=%v", passthroughSessionID, errReadMessage)
|
||||
} else {
|
||||
// log.Warnf("responses websocket: read message failed id=%s error=%v", passthroughSessionID, errReadMessage)
|
||||
}
|
||||
return
|
||||
}
|
||||
if msgType != websocket.TextMessage && msgType != websocket.BinaryMessage {
|
||||
continue
|
||||
}
|
||||
// log.Infof(
|
||||
// "responses websocket: downstream_in id=%s type=%d event=%s payload=%s",
|
||||
// passthroughSessionID,
|
||||
// msgType,
|
||||
// websocketPayloadEventType(payload),
|
||||
// websocketPayloadPreview(payload),
|
||||
// )
|
||||
wsTimelineLog.BeginRequest()
|
||||
wsTimelineLog.Append("request", payload, time.Now())
|
||||
|
||||
explicitRequestModelName := strings.TrimSpace(gjson.GetBytes(payload, "model").String())
|
||||
requestModelName := explicitRequestModelName
|
||||
if requestModelName == "" {
|
||||
requestModelName = passthroughModelName
|
||||
}
|
||||
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 := ""
|
||||
if pinnedAuth != nil {
|
||||
providerKey = strings.ToLower(strings.TrimSpace(pinnedAuth.Provider))
|
||||
}
|
||||
state, hasState := pinnedAuthByProvider[providerKey]
|
||||
if !ok || !hasState || state.authID != pinnedAuthID || !responsesWebsocketPinnedAuthMatchesModel(pinnedAuth, requestModelName, state.modelKey, homeRuntime) {
|
||||
pinnedAuthID = ""
|
||||
}
|
||||
}
|
||||
if pinnedAuthID == "" {
|
||||
providerSet, _ := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(requestModelName))
|
||||
if len(providerSet) == 1 {
|
||||
for providerKey := range providerSet {
|
||||
state, ok := pinnedAuthByProvider[providerKey]
|
||||
candidateAuth, homeRuntime, okAuth := sessionAuthByIDWithSource(state.authID)
|
||||
if ok && okAuth && responsesWebsocketPinnedAuthMatchesModel(candidateAuth, requestModelName, state.modelKey, homeRuntime) {
|
||||
pinnedAuthID = state.authID
|
||||
} else {
|
||||
delete(pinnedAuthByProvider, providerKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
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 = ""
|
||||
}
|
||||
|
||||
allowCompactionReplayBypass := false
|
||||
if !nativeWebsocketPassthrough {
|
||||
if pinnedAuthID != "" {
|
||||
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
|
||||
allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
|
||||
}
|
||||
} else {
|
||||
allowCompactionReplayBypass = h.websocketUpstreamSupportsCompactionReplayForModel(requestModelName)
|
||||
}
|
||||
}
|
||||
|
||||
var requestJSON []byte
|
||||
var updatedLastRequest []byte
|
||||
var errMsg *interfaces.ErrorMessage
|
||||
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,
|
||||
lastRequest,
|
||||
lastResponseOutput,
|
||||
lastResponseID,
|
||||
lastResponsePendingToolCallIDs,
|
||||
false,
|
||||
allowCompactionReplayBypass,
|
||||
)
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
|
||||
markAPIResponseTimestamp(c)
|
||||
errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg)
|
||||
log.Infof(
|
||||
"responses websocket: downstream_out id=%s type=%d event=%s payload=%s",
|
||||
passthroughSessionID,
|
||||
websocket.TextMessage,
|
||||
websocketPayloadEventType(errorPayload),
|
||||
websocketPayloadPreview(errorPayload),
|
||||
)
|
||||
if errWrite != nil {
|
||||
log.Warnf(
|
||||
"responses websocket: downstream_out write failed id=%s event=%s error=%v",
|
||||
passthroughSessionID,
|
||||
websocketPayloadEventType(errorPayload),
|
||||
errWrite,
|
||||
)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
requestJSON = h.prepareCodexMultiAgentV2Tools(c, requestJSON)
|
||||
|
||||
if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) {
|
||||
if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {
|
||||
requestJSON = updated
|
||||
}
|
||||
if updated, errDelete := sjson.DeleteBytes(updatedLastRequest, "generate"); errDelete == nil {
|
||||
updatedLastRequest = updated
|
||||
}
|
||||
lastRequest = updatedLastRequest
|
||||
lastResponseOutput = []byte("[]")
|
||||
lastResponseID = ""
|
||||
lastResponsePendingToolCallIDs = nil
|
||||
if errWrite := writeResponsesWebsocketSyntheticPrewarm(c, writer, requestJSON, wsTimelineLog, passthroughSessionID); errWrite != nil {
|
||||
wsTerminateErr = errWrite
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
var toolCacheTurn *responsesWebsocketToolCacheTurn
|
||||
nextLastRequest := lastRequest
|
||||
if nativeWebsocketPassthrough {
|
||||
if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {
|
||||
passthroughModelName = modelName
|
||||
}
|
||||
} else {
|
||||
requestJSON, toolCacheTurn = prepareResponsesWebsocketFallbackTurn(downstreamSessionKey, requestJSON)
|
||||
nextLastRequest = requestJSON
|
||||
}
|
||||
|
||||
modelName := gjson.GetBytes(requestJSON, "model").String()
|
||||
lastAttemptedAuthID := pinnedAuthID
|
||||
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)
|
||||
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)
|
||||
}
|
||||
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,
|
||||
responsesWebsocketForwardOptions{
|
||||
toolCacheTurn: toolCacheTurn,
|
||||
suppressError: replayPinnedAuthFailure,
|
||||
},
|
||||
)
|
||||
if errForward != nil {
|
||||
wsTerminateErr = errForward
|
||||
switch {
|
||||
case errors.Is(errForward, websocket.ErrCloseSent):
|
||||
case isWebsocketConnectionClosedError(errForward):
|
||||
// The client hung up while a downstream write was in flight. This is a
|
||||
// normal shutdown race, not a proxy failure.
|
||||
log.Debugf("responses websocket: client closed during forward id=%s error=%v", passthroughSessionID, errForward)
|
||||
default:
|
||||
log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
|
||||
}
|
||||
return
|
||||
}
|
||||
if forwardErrMsg != nil {
|
||||
if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
|
||||
forgetPinnedAuth()
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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 = ""
|
||||
lastRequest = nextLastRequest
|
||||
lastResponseOutput = completedOutput
|
||||
lastResponseID = strings.TrimSpace(completedResponseID)
|
||||
lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 ""
|
||||
}
|
||||
return strings.TrimSpace(c.ClientIP())
|
||||
}
|
||||
|
||||
func websocketUpgradeHeaders(req *http.Request) http.Header {
|
||||
headers := http.Header{}
|
||||
if req == nil {
|
||||
return headers
|
||||
}
|
||||
|
||||
// Keep the same sticky turn-state across reconnects when provided by the client.
|
||||
turnState := strings.TrimSpace(req.Header.Get(wsTurnStateHeader))
|
||||
if turnState != "" {
|
||||
headers.Set(wsTurnStateHeader, turnState)
|
||||
}
|
||||
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"}}`,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,607 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
type responsesWebsocketForwardOptions struct {
|
||||
toolCacheTurn *responsesWebsocketToolCacheTurn
|
||||
suppressError func(*interfaces.ErrorMessage) bool
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
|
||||
c *gin.Context,
|
||||
writer *responsesWebsocketWriter,
|
||||
cancel handlers.APIHandlerCancelFunc,
|
||||
data <-chan []byte,
|
||||
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 := ""
|
||||
outputItemsByIndex := make(map[int64][]byte)
|
||||
var outputItemsFallback [][]byte
|
||||
pendingToolCallIDs := make(map[string]struct{})
|
||||
downstreamSessionKey := ""
|
||||
if c != nil && c.Request != nil {
|
||||
downstreamSessionKey = websocketDownstreamSessionKey(c.Request)
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, c.Request.Context().Err()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
errs = nil
|
||||
continue
|
||||
}
|
||||
if errMsg == nil {
|
||||
cancel(nil)
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, 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)
|
||||
if errClose != nil {
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errClose
|
||||
}
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent
|
||||
}
|
||||
|
||||
errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, errMsg, nil)
|
||||
if wrote {
|
||||
log.Infof(
|
||||
"responses websocket: downstream_out id=%s type=%d event=%s payload=%s",
|
||||
sessionID,
|
||||
websocket.TextMessage,
|
||||
websocketPayloadEventType(errorPayload),
|
||||
websocketPayloadPreview(errorPayload),
|
||||
)
|
||||
}
|
||||
cancel(errMsg.Error)
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errTerminate
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
if !completed {
|
||||
errMsg := &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusRequestTimeout,
|
||||
Error: fmt.Errorf("stream closed before response.completed"),
|
||||
}
|
||||
h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
|
||||
markAPIResponseTimestamp(c)
|
||||
_, errClose := writer.closeWithoutError()
|
||||
cancel(errMsg.Error)
|
||||
if errClose != nil {
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errClose
|
||||
}
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent
|
||||
}
|
||||
cancel(nil)
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, nil
|
||||
}
|
||||
|
||||
payloads := websocketJSONPayloadsFromChunk(chunk)
|
||||
for i := range payloads {
|
||||
collectResponsesWebsocketOutputItem(payloads[i], outputItemsByIndex, &outputItemsFallback)
|
||||
eventType := gjson.GetBytes(payloads[i], "type").String()
|
||||
if isResponsesWebsocketCompletionEvent(eventType) {
|
||||
payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback)
|
||||
}
|
||||
if toolCacheTurn != nil {
|
||||
toolCacheTurn.recordResponse(payloads[i])
|
||||
} else {
|
||||
recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
|
||||
}
|
||||
recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])
|
||||
var payloadErrMsg *interfaces.ErrorMessage
|
||||
if eventType == wsEventTypeError {
|
||||
payloadErrMsg = responsesWebsocketErrorMessageFromPayload(payloads[i])
|
||||
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)
|
||||
completedResponseID = responseCompletedIDFromPayload(payloads[i])
|
||||
}
|
||||
markAPIResponseTimestamp(c)
|
||||
if payloadErrMsg != nil {
|
||||
if matched, errClose := writer.closeForUpstreamError(payloadErrMsg.Error); matched {
|
||||
cancel(payloadErrMsg.Error)
|
||||
if errClose != nil {
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, errClose
|
||||
}
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, websocket.ErrCloseSent
|
||||
}
|
||||
errorPayload, wrote, errTerminate := writeResponsesWebsocketTerminalError(writer, wsTimelineLog, payloadErrMsg, payloads[i])
|
||||
if wrote {
|
||||
log.Infof(
|
||||
"responses websocket: downstream_out id=%s type=%d event=%s payload=%s",
|
||||
sessionID,
|
||||
websocket.TextMessage,
|
||||
websocketPayloadEventType(errorPayload),
|
||||
websocketPayloadPreview(errorPayload),
|
||||
)
|
||||
}
|
||||
cancel(payloadErrMsg.Error)
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, errTerminate
|
||||
}
|
||||
// log.Infof(
|
||||
// "responses websocket: downstream_out id=%s type=%d event=%s payload=%s",
|
||||
// sessionID,
|
||||
// websocket.TextMessage,
|
||||
// websocketPayloadEventType(payloads[i]),
|
||||
// websocketPayloadPreview(payloads[i]),
|
||||
// )
|
||||
if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil {
|
||||
log.Warnf(
|
||||
"responses websocket: downstream_out write failed id=%s event=%s error=%v",
|
||||
sessionID,
|
||||
websocketPayloadEventType(payloads[i]),
|
||||
errWrite,
|
||||
)
|
||||
cancel(errWrite)
|
||||
return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int {
|
||||
if errMsg == nil {
|
||||
return 0
|
||||
}
|
||||
if errMsg.StatusCode > 0 {
|
||||
return errMsg.StatusCode
|
||||
}
|
||||
return clienterror.HTTPStatusFromError(errMsg.Error)
|
||||
}
|
||||
|
||||
// shouldExposeResponsesUpstreamError reports whether a terminal upstream error
|
||||
// must reach the downstream client.
|
||||
//
|
||||
// Only request-shape failures are exposed: the client can act on them and no
|
||||
// credential rotation or retry can make the request succeed. Credential, quota
|
||||
// and transport failures stay silent so the client simply reconnects and retries;
|
||||
// a fresh connection carries no server-side transcript, so reconnecting already
|
||||
// implies a full context resend.
|
||||
func shouldExposeResponsesUpstreamError(errMsg *interfaces.ErrorMessage) bool {
|
||||
if errMsg == nil {
|
||||
return false
|
||||
}
|
||||
return clienterror.IsRequestFault(responsesWebsocketErrorStatus(errMsg), errMsg.Error)
|
||||
}
|
||||
|
||||
func writeResponsesWebsocketTerminalError(
|
||||
writer *responsesWebsocketWriter,
|
||||
wsTimelineLog websocketTimelineAppender,
|
||||
errMsg *interfaces.ErrorMessage,
|
||||
payload []byte,
|
||||
) ([]byte, bool, error) {
|
||||
if !shouldExposeResponsesUpstreamError(errMsg) {
|
||||
// Keep the upstream reason in the request-log timeline even though the client
|
||||
// only observes a closed connection, otherwise silent failures are
|
||||
// undiagnosable after the fact.
|
||||
if wsTimelineLog != nil && errMsg != nil {
|
||||
appendWebsocketTimelineDisconnect(wsTimelineLog, errMsg.Error, time.Now())
|
||||
}
|
||||
_, errClose := writer.closeWithoutError()
|
||||
if errClose != nil {
|
||||
return nil, false, errClose
|
||||
}
|
||||
return nil, false, websocket.ErrCloseSent
|
||||
}
|
||||
|
||||
if len(payload) == 0 {
|
||||
var errBuild error
|
||||
payload, errBuild = buildResponsesWebsocketErrorPayload(errMsg)
|
||||
if errBuild != nil {
|
||||
_, _ = writer.closeWithoutError()
|
||||
return nil, false, errBuild
|
||||
}
|
||||
}
|
||||
|
||||
wrote, errClose := writer.closeWithPayload(payload)
|
||||
if wrote && wsTimelineLog != nil {
|
||||
wsTimelineLog.Append("response", payload, time.Now())
|
||||
}
|
||||
if errClose != nil {
|
||||
return payload, wrote, errClose
|
||||
}
|
||||
return payload, wrote, websocket.ErrCloseSent
|
||||
}
|
||||
|
||||
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,
|
||||
http.StatusTooManyRequests,
|
||||
http.StatusRequestTimeout,
|
||||
http.StatusBadGateway,
|
||||
http.StatusServiceUnavailable,
|
||||
http.StatusGatewayTimeout:
|
||||
return true
|
||||
default:
|
||||
}
|
||||
if errMsg.Error != nil {
|
||||
msg := strings.ToLower(errMsg.Error.Error())
|
||||
switch {
|
||||
case strings.Contains(msg, "stream closed before response.completed"),
|
||||
strings.Contains(msg, "previous_response_not_found"),
|
||||
strings.Contains(msg, "ws_failed"),
|
||||
strings.Contains(msg, "upstream stream closed before first payload"),
|
||||
strings.Contains(msg, "empty_stream"):
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func collectResponsesWebsocketOutputItem(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) {
|
||||
if gjson.GetBytes(payload, "type").String() != "response.output_item.done" {
|
||||
return
|
||||
}
|
||||
item := gjson.GetBytes(payload, "item")
|
||||
if !item.Exists() || !item.IsObject() {
|
||||
return
|
||||
}
|
||||
outputIndex := gjson.GetBytes(payload, "output_index")
|
||||
if outputIndex.Exists() {
|
||||
outputItemsByIndex[outputIndex.Int()] = bytes.Clone([]byte(item.Raw))
|
||||
return
|
||||
}
|
||||
*outputItemsFallback = append(*outputItemsFallback, bytes.Clone([]byte(item.Raw)))
|
||||
}
|
||||
|
||||
func restoreResponsesWebsocketCompletionOutput(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte {
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if output.Exists() && output.IsArray() && len(output.Array()) > 0 {
|
||||
reconciledOutput, changed := reconcileResponsesWebsocketCompletionToolCalls(output, outputItemsByIndex, outputItemsFallback)
|
||||
if !changed {
|
||||
return payload
|
||||
}
|
||||
restored, errSet := sjson.SetRawBytes(payload, "response.output", reconciledOutput)
|
||||
if errSet != nil {
|
||||
return payload
|
||||
}
|
||||
return restored
|
||||
}
|
||||
if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 {
|
||||
return payload
|
||||
}
|
||||
|
||||
restored, errSet := sjson.SetRawBytes(payload, "response.output", responseCompletedOutputFromPayload(payload, outputItemsByIndex, outputItemsFallback))
|
||||
if errSet != nil {
|
||||
return payload
|
||||
}
|
||||
return restored
|
||||
}
|
||||
|
||||
func reconcileResponsesWebsocketCompletionToolCalls(output gjson.Result, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) ([]byte, bool) {
|
||||
collectedToolCalls := make(map[string]json.RawMessage)
|
||||
recordCollectedToolCall := func(raw []byte) {
|
||||
item := gjson.ParseBytes(raw)
|
||||
if !isCompleteResponsesWebsocketToolCall(item) {
|
||||
return
|
||||
}
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
collectedToolCalls[callID] = append(json.RawMessage(nil), raw...)
|
||||
}
|
||||
|
||||
indexes := make([]int64, 0, len(outputItemsByIndex))
|
||||
for index := range outputItemsByIndex {
|
||||
indexes = append(indexes, index)
|
||||
}
|
||||
sort.Slice(indexes, func(i, j int) bool {
|
||||
return indexes[i] < indexes[j]
|
||||
})
|
||||
for _, index := range indexes {
|
||||
recordCollectedToolCall(outputItemsByIndex[index])
|
||||
}
|
||||
for _, item := range outputItemsFallback {
|
||||
recordCollectedToolCall(item)
|
||||
}
|
||||
if len(collectedToolCalls) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
items := output.Array()
|
||||
reconciled := make([]json.RawMessage, 0, len(items))
|
||||
changed := false
|
||||
for _, item := range items {
|
||||
raw := json.RawMessage(item.Raw)
|
||||
if isResponsesToolCallType(item.Get("type").String()) {
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
if collected, ok := collectedToolCalls[callID]; ok && !bytes.Equal(raw, collected) {
|
||||
raw = collected
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
reconciled = append(reconciled, raw)
|
||||
}
|
||||
if !changed {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
marshaledOutput, errMarshal := json.Marshal(reconciled)
|
||||
if errMarshal != nil {
|
||||
return nil, false
|
||||
}
|
||||
return marshaledOutput, true
|
||||
}
|
||||
|
||||
func isCompleteResponsesWebsocketToolCall(item gjson.Result) bool {
|
||||
if !item.Exists() || !item.IsObject() {
|
||||
return false
|
||||
}
|
||||
callID := item.Get("call_id")
|
||||
name := item.Get("name")
|
||||
if callID.Type != gjson.String || strings.TrimSpace(callID.String()) == "" || name.Type != gjson.String || strings.TrimSpace(name.String()) == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
switch strings.TrimSpace(item.Get("type").String()) {
|
||||
case "function_call":
|
||||
arguments := item.Get("arguments")
|
||||
return arguments.Exists() && arguments.Type == gjson.String
|
||||
case "custom_tool_call":
|
||||
input := item.Get("input")
|
||||
return input.Exists() && input.Type == gjson.String
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func responseCompletedOutputFromPayload(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte {
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if output.Exists() && output.IsArray() && len(output.Array()) > 0 {
|
||||
return bytes.Clone([]byte(output.Raw))
|
||||
}
|
||||
if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 {
|
||||
return []byte("[]")
|
||||
}
|
||||
|
||||
indexes := make([]int64, 0, len(outputItemsByIndex))
|
||||
for index := range outputItemsByIndex {
|
||||
indexes = append(indexes, index)
|
||||
}
|
||||
sort.Slice(indexes, func(i, j int) bool {
|
||||
return indexes[i] < indexes[j]
|
||||
})
|
||||
|
||||
items := make([]json.RawMessage, 0, len(outputItemsByIndex)+len(outputItemsFallback))
|
||||
appendCollectedItem := func(raw []byte) {
|
||||
item := gjson.ParseBytes(raw)
|
||||
if isResponsesToolCallType(item.Get("type").String()) && !isCompleteResponsesWebsocketToolCall(item) {
|
||||
return
|
||||
}
|
||||
items = append(items, append(json.RawMessage(nil), raw...))
|
||||
}
|
||||
for _, index := range indexes {
|
||||
appendCollectedItem(outputItemsByIndex[index])
|
||||
}
|
||||
for _, item := range outputItemsFallback {
|
||||
appendCollectedItem(item)
|
||||
}
|
||||
|
||||
marshaledOutput, errMarshal := json.Marshal(items)
|
||||
if errMarshal != nil {
|
||||
return []byte("[]")
|
||||
}
|
||||
return marshaledOutput
|
||||
}
|
||||
|
||||
func responseCompletedIDFromPayload(payload []byte) string {
|
||||
return strings.TrimSpace(gjson.GetBytes(payload, "response.id").String())
|
||||
}
|
||||
|
||||
func recordPendingToolCallIDsFromPayload(pending map[string]struct{}, payload []byte) {
|
||||
if pending == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
updatePendingToolCallIDsFromItem(pending, gjson.GetBytes(payload, "item"))
|
||||
output := gjson.GetBytes(payload, "response.output")
|
||||
if output.IsArray() {
|
||||
for _, item := range output.Array() {
|
||||
updatePendingToolCallIDsFromItem(pending, item)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func updatePendingToolCallIDsFromItem(pending map[string]struct{}, item gjson.Result) {
|
||||
if pending == nil || !item.Exists() {
|
||||
return
|
||||
}
|
||||
switch strings.TrimSpace(item.Get("type").String()) {
|
||||
case "function_call", "custom_tool_call":
|
||||
if !isCompleteResponsesWebsocketToolCall(item) {
|
||||
return
|
||||
}
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
pending[callID] = struct{}{}
|
||||
case "function_call_output", "custom_tool_call_output":
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
if callID != "" {
|
||||
delete(pending, callID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func sortedStringSet(values map[string]struct{}) []string {
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(values))
|
||||
for value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" {
|
||||
out = append(out, value)
|
||||
}
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
func websocketJSONPayloadsFromChunk(chunk []byte) [][]byte {
|
||||
payloads := make([][]byte, 0, 2)
|
||||
lines := bytes.Split(chunk, []byte("\n"))
|
||||
for i := range lines {
|
||||
line := bytes.TrimSpace(lines[i])
|
||||
if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) {
|
||||
continue
|
||||
}
|
||||
if bytes.HasPrefix(line, []byte("data:")) {
|
||||
line = bytes.TrimSpace(line[len("data:"):])
|
||||
}
|
||||
if len(line) == 0 || bytes.Equal(line, []byte(wsDoneMarker)) {
|
||||
continue
|
||||
}
|
||||
if json.Valid(line) {
|
||||
payloads = append(payloads, bytes.Clone(line))
|
||||
}
|
||||
}
|
||||
|
||||
if len(payloads) > 0 {
|
||||
return payloads
|
||||
}
|
||||
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
if bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
trimmed = bytes.TrimSpace(trimmed[len("data:"):])
|
||||
}
|
||||
if len(trimmed) > 0 && !bytes.Equal(trimmed, []byte(wsDoneMarker)) && json.Valid(trimmed) {
|
||||
payloads = append(payloads, bytes.Clone(trimmed))
|
||||
}
|
||||
return payloads
|
||||
}
|
||||
|
||||
func buildResponsesWebsocketErrorPayload(errMsg *interfaces.ErrorMessage) ([]byte, error) {
|
||||
status := http.StatusInternalServerError
|
||||
errText := http.StatusText(status)
|
||||
if errMsg != nil {
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
errText = http.StatusText(status)
|
||||
}
|
||||
if errMsg.Error != nil && strings.TrimSpace(errMsg.Error.Error()) != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
}
|
||||
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
payload := []byte(`{}`)
|
||||
var errSet error
|
||||
payload, errSet = sjson.SetBytes(payload, "type", wsEventTypeError)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
payload, errSet = sjson.SetBytes(payload, "status", status)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
|
||||
if errMsg != nil && errMsg.Addon != nil {
|
||||
headers := []byte(`{}`)
|
||||
hasHeaders := false
|
||||
for key, values := range errMsg.Addon {
|
||||
if len(values) == 0 {
|
||||
continue
|
||||
}
|
||||
headerPath := strings.ReplaceAll(strings.ReplaceAll(key, `\\`, `\\\\`), ".", `\\.`)
|
||||
headers, errSet = sjson.SetBytes(headers, headerPath, values[0])
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
hasHeaders = true
|
||||
}
|
||||
if hasHeaders {
|
||||
payload, errSet = sjson.SetRawBytes(payload, "headers", headers)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(body) > 0 && json.Valid(body) {
|
||||
errorNode := gjson.GetBytes(body, "error")
|
||||
if errorNode.Exists() {
|
||||
payload, errSet = sjson.SetRawBytes(payload, "error", []byte(errorNode.Raw))
|
||||
} else {
|
||||
payload, errSet = sjson.SetRawBytes(payload, "error", body)
|
||||
}
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
}
|
||||
|
||||
if !gjson.GetBytes(payload, "error").Exists() {
|
||||
payload, errSet = sjson.SetBytes(payload, "error.type", "server_error")
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
payload, errSet = sjson.SetBytes(payload, "error.message", errText)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
}
|
||||
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func writeResponsesWebsocketError(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage) ([]byte, error) {
|
||||
payload, errBuild := buildResponsesWebsocketErrorPayload(errMsg)
|
||||
if errBuild != nil {
|
||||
return nil, errBuild
|
||||
}
|
||||
return payload, writeResponsesWebsocketPayload(writer, wsTimelineLog, payload, time.Now())
|
||||
}
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func shouldHandleResponsesWebsocketPrewarmLocally(rawJSON []byte, lastRequest []byte, allowIncrementalInputWithPreviousResponseID bool) bool {
|
||||
if allowIncrementalInputWithPreviousResponseID || len(lastRequest) != 0 {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) != wsRequestTypeCreate {
|
||||
return false
|
||||
}
|
||||
generateResult := gjson.GetBytes(rawJSON, "generate")
|
||||
return generateResult.Exists() && !generateResult.Bool()
|
||||
}
|
||||
|
||||
func writeResponsesWebsocketSyntheticPrewarm(
|
||||
c *gin.Context,
|
||||
writer *responsesWebsocketWriter,
|
||||
requestJSON []byte,
|
||||
wsTimelineLog websocketTimelineAppender,
|
||||
sessionID string,
|
||||
) error {
|
||||
payloads, errPayloads := syntheticResponsesWebsocketPrewarmPayloads(requestJSON)
|
||||
if errPayloads != nil {
|
||||
return errPayloads
|
||||
}
|
||||
for i := 0; i < len(payloads); i++ {
|
||||
markAPIResponseTimestamp(c)
|
||||
// log.Infof(
|
||||
// "responses websocket: downstream_out id=%s type=%d event=%s payload=%s",
|
||||
// sessionID,
|
||||
// websocket.TextMessage,
|
||||
// websocketPayloadEventType(payloads[i]),
|
||||
// websocketPayloadPreview(payloads[i]),
|
||||
// )
|
||||
if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil {
|
||||
log.Warnf(
|
||||
"responses websocket: downstream_out write failed id=%s event=%s error=%v",
|
||||
sessionID,
|
||||
websocketPayloadEventType(payloads[i]),
|
||||
errWrite,
|
||||
)
|
||||
return errWrite
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func syntheticResponsesWebsocketPrewarmPayloads(requestJSON []byte) ([][]byte, error) {
|
||||
responseID := "resp_prewarm_" + uuid.NewString()
|
||||
createdAt := time.Now().Unix()
|
||||
modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String())
|
||||
|
||||
createdPayload := []byte(`{"type":"response.created","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","background":false,"error":null,"output":[]}}`)
|
||||
var errSet error
|
||||
createdPayload, errSet = sjson.SetBytes(createdPayload, "response.id", responseID)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
createdPayload, errSet = sjson.SetBytes(createdPayload, "response.created_at", createdAt)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
if modelName != "" {
|
||||
createdPayload, errSet = sjson.SetBytes(createdPayload, "response.model", modelName)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
completedPayload := []byte(`{"type":"response.completed","sequence_number":1,"response":{"id":"","object":"response","created_at":0,"status":"completed","background":false,"error":null,"output":[],"usage":{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0}}}`)
|
||||
completedPayload, errSet = sjson.SetBytes(completedPayload, "response.id", responseID)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
completedPayload, errSet = sjson.SetBytes(completedPayload, "response.created_at", createdAt)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
if modelName != "" {
|
||||
completedPayload, errSet = sjson.SetBytes(completedPayload, "response.model", modelName)
|
||||
if errSet != nil {
|
||||
return nil, errSet
|
||||
}
|
||||
}
|
||||
|
||||
return [][]byte{createdPayload, completedPayload}, nil
|
||||
}
|
||||
|
||||
// inputContainsFullTranscript returns true when the input array carries compact
|
||||
// replay markers that indicate the client already sent the full conversation
|
||||
// transcript. Merging that input with stale lastRequest/lastResponseOutput
|
||||
// would duplicate or break function_call/function_call_output pairings, so the
|
||||
// caller should use the input as-is.
|
||||
//
|
||||
// Assistant messages alone are not enough to classify the payload as a replay:
|
||||
// incremental websocket requests may legitimately append assistant items.
|
||||
func inputContainsFullTranscript(input gjson.Result) bool {
|
||||
if !input.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, item := range input.Array() {
|
||||
t := item.Get("type").String()
|
||||
if t == "compaction" || t == "compaction_summary" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func inputWithoutCompactionItems(input gjson.Result) string {
|
||||
if !input.IsArray() {
|
||||
return normalizeJSONArrayRaw([]byte(input.Raw))
|
||||
}
|
||||
filtered := make([]string, 0, len(input.Array()))
|
||||
for _, item := range input.Array() {
|
||||
t := item.Get("type").String()
|
||||
if t == "compaction" || t == "compaction_summary" {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item.Raw)
|
||||
}
|
||||
return "[" + strings.Join(filtered, ",") + "]"
|
||||
}
|
||||
|
||||
func normalizeJSONArrayRaw(raw []byte) string {
|
||||
trimmed := strings.TrimSpace(string(raw))
|
||||
if trimmed == "" {
|
||||
return "[]"
|
||||
}
|
||||
result := gjson.Parse(trimmed)
|
||||
if result.Type == gjson.JSON && result.IsArray() {
|
||||
return trimmed
|
||||
}
|
||||
return "[]"
|
||||
}
|
||||
|
|
@ -0,0 +1,737 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)
|
||||
}
|
||||
|
||||
func normalizeResponsesWebsocketRequestWithMode(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
return normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON, lastRequest, lastResponseOutput, "", allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass)
|
||||
}
|
||||
|
||||
func normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
return normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON, lastRequest, lastResponseOutput, lastResponseID, nil, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass)
|
||||
}
|
||||
|
||||
func normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String())
|
||||
switch requestType {
|
||||
case wsRequestTypeCreate:
|
||||
// log.Infof("responses websocket: response.create request")
|
||||
if len(lastRequest) == 0 {
|
||||
return normalizeResponseCreateRequest(rawJSON)
|
||||
}
|
||||
return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass)
|
||||
case wsRequestTypeAppend:
|
||||
// log.Infof("responses websocket: response.append request")
|
||||
return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass)
|
||||
default:
|
||||
return nil, lastRequest, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("unsupported websocket request type: %s", requestType),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeResponseCreateRequest(rawJSON []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
normalized, errDelete := sjson.DeleteBytes(rawJSON, "type")
|
||||
if errDelete != nil {
|
||||
normalized = bytes.Clone(rawJSON)
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "stream", true)
|
||||
if !gjson.GetBytes(normalized, "input").Exists() {
|
||||
normalized, _ = sjson.SetRawBytes(normalized, "input", []byte("[]"))
|
||||
}
|
||||
|
||||
modelName := strings.TrimSpace(gjson.GetBytes(normalized, "model").String())
|
||||
if modelName == "" {
|
||||
return nil, nil, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("missing model in response.create request"),
|
||||
}
|
||||
}
|
||||
return normalized, bytes.Clone(normalized), nil
|
||||
}
|
||||
|
||||
func normalizeResponseSubsequentRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) {
|
||||
if len(lastRequest) == 0 {
|
||||
return nil, lastRequest, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("websocket request received before response.create"),
|
||||
}
|
||||
}
|
||||
|
||||
nextInput := gjson.GetBytes(rawJSON, "input")
|
||||
if !nextInput.Exists() || !nextInput.IsArray() {
|
||||
return nil, lastRequest, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("websocket request requires array field: input"),
|
||||
}
|
||||
}
|
||||
|
||||
// Compaction can cause clients to replace local websocket history with a new
|
||||
// compact transcript on the next `response.create`. When the input already
|
||||
// contains historical model output items, treating it as an incremental append
|
||||
// duplicates stale turn-state and can leave late orphaned function_call items.
|
||||
if shouldReplaceWebsocketTranscript(rawJSON, nextInput) {
|
||||
normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest)
|
||||
return normalized, bytes.Clone(normalized), nil
|
||||
}
|
||||
|
||||
// Websocket v2 mode uses response.create with previous_response_id + incremental input.
|
||||
// Do not expand it into a full input transcript; upstream expects the incremental payload.
|
||||
if allowIncrementalInputWithPreviousResponseID {
|
||||
prev := strings.TrimSpace(gjson.GetBytes(rawJSON, "previous_response_id").String())
|
||||
if prev == "" {
|
||||
if !inputSatisfiesPendingToolCalls(nextInput, lastResponsePendingToolCallIDs) {
|
||||
normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest)
|
||||
return normalized, bytes.Clone(normalized), nil
|
||||
}
|
||||
prev = strings.TrimSpace(lastResponseID)
|
||||
}
|
||||
if prev != "" {
|
||||
normalized, errDelete := sjson.DeleteBytes(rawJSON, "type")
|
||||
if errDelete != nil {
|
||||
normalized = bytes.Clone(rawJSON)
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "previous_response_id", prev)
|
||||
if !gjson.GetBytes(normalized, "model").Exists() {
|
||||
modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
|
||||
if modelName != "" {
|
||||
normalized, _ = sjson.SetBytes(normalized, "model", modelName)
|
||||
}
|
||||
}
|
||||
if !gjson.GetBytes(normalized, "instructions").Exists() {
|
||||
instructions := gjson.GetBytes(lastRequest, "instructions")
|
||||
if instructions.Exists() {
|
||||
normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw))
|
||||
}
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "stream", true)
|
||||
return normalized, bytes.Clone(normalized), nil
|
||||
}
|
||||
}
|
||||
|
||||
// When the client sends a compact replay for a downstream that can consume it
|
||||
// directly, the input already carries the canonical history. In that case,
|
||||
// skip merging with stale lastRequest/lastResponseOutput to avoid breaking
|
||||
// function_call / function_call_output pairings.
|
||||
// See: https://github.com/router-for-me/CLIProxyAPI/issues/2207
|
||||
var mergedInput []byte
|
||||
if allowCompactionReplayBypass && inputContainsFullTranscript(nextInput) {
|
||||
log.Infof("responses websocket: full transcript detected, skipping stale merge (input items=%d)", len(nextInput.Array()))
|
||||
mergedInput = []byte(nextInput.Raw)
|
||||
} else {
|
||||
appendInputRaw := nextInput.Raw
|
||||
if inputContainsFullTranscript(nextInput) {
|
||||
appendInputRaw = inputWithoutCompactionItems(nextInput)
|
||||
}
|
||||
|
||||
var errMerge error
|
||||
mergedInput, errMerge = mergeResponsesWebsocketInput(lastRequest, lastResponseOutput, appendInputRaw)
|
||||
if errMerge != nil {
|
||||
return nil, lastRequest, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: errMerge,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
normalized, errDelete := sjson.DeleteBytes(rawJSON, "type")
|
||||
if errDelete != nil {
|
||||
normalized = bytes.Clone(rawJSON)
|
||||
}
|
||||
normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id")
|
||||
if !gjson.GetBytes(normalized, "model").Exists() {
|
||||
modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
|
||||
if modelName != "" {
|
||||
normalized, _ = sjson.SetBytes(normalized, "model", modelName)
|
||||
}
|
||||
}
|
||||
if !gjson.GetBytes(normalized, "instructions").Exists() {
|
||||
instructions := gjson.GetBytes(lastRequest, "instructions")
|
||||
if instructions.Exists() {
|
||||
normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw))
|
||||
}
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "stream", true)
|
||||
var errSet error
|
||||
normalized, errSet = sjson.SetRawBytes(normalized, "input", mergedInput)
|
||||
if errSet != nil {
|
||||
return nil, lastRequest, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("failed to merge websocket input: %w", errSet),
|
||||
}
|
||||
}
|
||||
return normalized, normalized, nil
|
||||
}
|
||||
|
||||
func shouldReplaceWebsocketTranscript(rawJSON []byte, nextInput gjson.Result) bool {
|
||||
requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String())
|
||||
if requestType != wsRequestTypeCreate && requestType != wsRequestTypeAppend {
|
||||
return false
|
||||
}
|
||||
previousResponseID := gjson.GetBytes(rawJSON, "previous_response_id")
|
||||
if strings.TrimSpace(previousResponseID.String()) != "" {
|
||||
return false
|
||||
}
|
||||
if !nextInput.Exists() || !nextInput.IsArray() {
|
||||
return false
|
||||
}
|
||||
if requestType == wsRequestTypeCreate && !previousResponseID.Exists() && inputHasCodexLocalCompactionSummary(nextInput) {
|
||||
return true
|
||||
}
|
||||
|
||||
for _, item := range nextInput.Array() {
|
||||
switch strings.TrimSpace(item.Get("type").String()) {
|
||||
case "function_call", "custom_tool_call":
|
||||
return true
|
||||
case "message":
|
||||
if strings.TrimSpace(item.Get("role").String()) == "assistant" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func inputHasCodexLocalCompactionSummary(input gjson.Result) bool {
|
||||
if !input.IsArray() {
|
||||
return false
|
||||
}
|
||||
|
||||
hasSummary := false
|
||||
for index, item := range input.Array() {
|
||||
itemType := strings.TrimSpace(item.Get("type").String())
|
||||
if itemType == "additional_tools" {
|
||||
tools := item.Get("tools")
|
||||
if index != 0 || strings.TrimSpace(item.Get("role").String()) != "developer" || !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) == "" {
|
||||
return false
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if itemType != "" && itemType != "message" {
|
||||
return false
|
||||
}
|
||||
|
||||
role := strings.TrimSpace(item.Get("role").String())
|
||||
if role != "user" && role != "developer" {
|
||||
return false
|
||||
}
|
||||
if role == "user" && strings.HasPrefix(codexLocalCompactionMessageText(item), codexLocalCompactionSummaryPrefix+"\n") {
|
||||
hasSummary = true
|
||||
}
|
||||
}
|
||||
return hasSummary
|
||||
}
|
||||
|
||||
func codexLocalCompactionMessageText(message gjson.Result) string {
|
||||
content := message.Get("content")
|
||||
if content.Type == gjson.String {
|
||||
return content.String()
|
||||
}
|
||||
if !content.IsArray() {
|
||||
return ""
|
||||
}
|
||||
|
||||
var text strings.Builder
|
||||
for _, part := range content.Array() {
|
||||
if strings.TrimSpace(part.Get("type").String()) == "input_text" {
|
||||
text.WriteString(part.Get("text").String())
|
||||
}
|
||||
}
|
||||
return text.String()
|
||||
}
|
||||
|
||||
func inputSatisfiesPendingToolCalls(input gjson.Result, pendingCallIDs []string) bool {
|
||||
if len(pendingCallIDs) == 0 {
|
||||
return true
|
||||
}
|
||||
if !input.IsArray() {
|
||||
return false
|
||||
}
|
||||
outputs := make(map[string]struct{}, len(pendingCallIDs))
|
||||
for _, item := range input.Array() {
|
||||
switch strings.TrimSpace(item.Get("type").String()) {
|
||||
case "function_call_output", "custom_tool_call_output":
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
if callID != "" {
|
||||
outputs[callID] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, callID := range pendingCallIDs {
|
||||
callID = strings.TrimSpace(callID)
|
||||
if callID == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := outputs[callID]; !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func normalizeResponseTranscriptReplacement(rawJSON []byte, lastRequest []byte) []byte {
|
||||
normalized, errDelete := sjson.DeleteBytes(rawJSON, "type")
|
||||
if errDelete != nil {
|
||||
normalized = bytes.Clone(rawJSON)
|
||||
}
|
||||
normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id")
|
||||
if !gjson.GetBytes(normalized, "model").Exists() {
|
||||
modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
|
||||
if modelName != "" {
|
||||
normalized, _ = sjson.SetBytes(normalized, "model", modelName)
|
||||
}
|
||||
}
|
||||
if !gjson.GetBytes(normalized, "instructions").Exists() {
|
||||
instructions := gjson.GetBytes(lastRequest, "instructions")
|
||||
if instructions.Exists() {
|
||||
normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw))
|
||||
}
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "stream", true)
|
||||
return bytes.Clone(normalized)
|
||||
}
|
||||
|
||||
type responsesWebsocketInputItem struct {
|
||||
raw json.RawMessage
|
||||
itemType string
|
||||
id string
|
||||
callID string
|
||||
}
|
||||
|
||||
type responsesWebsocketMergeInputItem struct {
|
||||
// raw may reference a caller-owned request buffer. Merge items must remain
|
||||
// local to mergeResponsesWebsocketInput, which copies every item into the
|
||||
// owned output buffer before returning.
|
||||
raw string
|
||||
itemType string
|
||||
id string
|
||||
callID string
|
||||
}
|
||||
|
||||
func mergeResponsesWebsocketInput(lastRequest []byte, lastResponseOutput []byte, appendRaw string) ([]byte, error) {
|
||||
previousInput, errPrevious := responsesWebsocketPreviousInputNoCopy(lastRequest)
|
||||
if errPrevious != nil {
|
||||
return nil, fmt.Errorf("invalid previous request input: %w", errPrevious)
|
||||
}
|
||||
items, errExisting := appendResponsesWebsocketMergeInputResult(nil, previousInput)
|
||||
if errExisting != nil {
|
||||
return nil, fmt.Errorf("invalid previous request input: %w", errExisting)
|
||||
}
|
||||
|
||||
trimmedResponse := bytes.TrimSpace(lastResponseOutput)
|
||||
if len(trimmedResponse) > 0 && trimmedResponse[0] == '[' && json.Valid(trimmedResponse) {
|
||||
responseInput := util.ParseGJSONBytesNoCopy(trimmedResponse)
|
||||
if inputContainsFullTranscript(responseInput) {
|
||||
items = slices.DeleteFunc(items, func(item responsesWebsocketMergeInputItem) bool {
|
||||
return item.itemType == "compaction_trigger"
|
||||
})
|
||||
}
|
||||
var errResponse error
|
||||
items, errResponse = appendResponsesWebsocketMergeInputResult(items, responseInput)
|
||||
if errResponse != nil {
|
||||
return nil, fmt.Errorf("invalid previous response output: %w", errResponse)
|
||||
}
|
||||
}
|
||||
|
||||
items, errAppend := appendResponsesWebsocketMergeInputItems(items, appendRaw)
|
||||
if errAppend != nil {
|
||||
return nil, fmt.Errorf("invalid request input: %w", errAppend)
|
||||
}
|
||||
|
||||
items = dedupeResponsesWebsocketMergeFunctionCalls(items)
|
||||
items = dedupeResponsesWebsocketMergeInputItems(items)
|
||||
return marshalResponsesWebsocketMergeInputItems(items), nil
|
||||
}
|
||||
|
||||
func responsesWebsocketPreviousInputNoCopy(lastRequest []byte) (gjson.Result, error) {
|
||||
if !json.Valid(lastRequest) {
|
||||
return gjson.Result{}, responsesWebsocketPreviousInputDecodeError(lastRequest)
|
||||
}
|
||||
|
||||
root := util.ParseGJSONBytesNoCopy(lastRequest)
|
||||
if root.Type == gjson.Null {
|
||||
return gjson.Parse("[]"), nil
|
||||
}
|
||||
if !root.IsObject() {
|
||||
return gjson.Result{}, responsesWebsocketPreviousInputDecodeError(lastRequest)
|
||||
}
|
||||
|
||||
var input gjson.Result
|
||||
inputFound := false
|
||||
invalidInput := false
|
||||
root.ForEach(func(key, value gjson.Result) bool {
|
||||
if !strings.EqualFold(key.String(), "input") {
|
||||
return true
|
||||
}
|
||||
// encoding/json processes matching duplicate fields in source order,
|
||||
// retains the last value, and still reports a type error from any
|
||||
// incompatible duplicate. Preserve those semantics without copying the
|
||||
// selected array out of the caller-owned request buffer.
|
||||
inputFound = true
|
||||
input = value
|
||||
if value.Type != gjson.Null && !value.IsArray() {
|
||||
invalidInput = true
|
||||
}
|
||||
return true
|
||||
})
|
||||
if invalidInput {
|
||||
return gjson.Result{}, responsesWebsocketPreviousInputDecodeError(lastRequest)
|
||||
}
|
||||
if !inputFound || input.Type == gjson.Null {
|
||||
return gjson.Parse("[]"), nil
|
||||
}
|
||||
return input, nil
|
||||
}
|
||||
|
||||
func responsesWebsocketPreviousInputDecodeError(lastRequest []byte) error {
|
||||
var previousRequest struct {
|
||||
Input []json.RawMessage `json:"input"`
|
||||
}
|
||||
return json.Unmarshal(lastRequest, &previousRequest)
|
||||
}
|
||||
|
||||
func appendResponsesWebsocketMergeInputItems(items []responsesWebsocketMergeInputItem, rawArray string) ([]responsesWebsocketMergeInputItem, error) {
|
||||
rawArray = strings.TrimSpace(rawArray)
|
||||
if rawArray == "" {
|
||||
rawArray = "[]"
|
||||
}
|
||||
parsed := gjson.Parse(rawArray)
|
||||
if gjson.Valid(rawArray) {
|
||||
return appendResponsesWebsocketMergeInputResult(items, parsed)
|
||||
}
|
||||
|
||||
var rawItems []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal([]byte(rawArray), &rawItems); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func appendResponsesWebsocketMergeInputResult(items []responsesWebsocketMergeInputItem, input gjson.Result) ([]responsesWebsocketMergeInputItem, error) {
|
||||
if input.Type == gjson.Null {
|
||||
return items, nil
|
||||
}
|
||||
if !input.IsArray() {
|
||||
var rawItems []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal([]byte(input.Raw), &rawItems); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
rawItems := input.Array()
|
||||
items = slices.Grow(items, len(rawItems))
|
||||
for _, rawItem := range rawItems {
|
||||
item := responsesWebsocketMergeInputItem{raw: rawItem.Raw}
|
||||
if rawItem.IsObject() {
|
||||
rawItem.ForEach(func(key, value gjson.Result) bool {
|
||||
metadataKey := key.String()
|
||||
switch {
|
||||
case strings.EqualFold(metadataKey, "type"):
|
||||
item.itemType = strings.TrimSpace(value.String())
|
||||
case strings.EqualFold(metadataKey, "id"):
|
||||
item.id = strings.TrimSpace(value.String())
|
||||
case strings.EqualFold(metadataKey, "call_id"):
|
||||
item.callID = strings.TrimSpace(value.String())
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func dedupeResponsesWebsocketMergeFunctionCalls(items []responsesWebsocketMergeInputItem) []responsesWebsocketMergeInputItem {
|
||||
seenCallIDs := make(map[string]struct{}, len(items))
|
||||
filtered := items[:0]
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallType(item.itemType) && item.callID != "" {
|
||||
if _, ok := seenCallIDs[item.callID]; ok {
|
||||
continue
|
||||
}
|
||||
seenCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
clear(items[len(filtered):])
|
||||
return filtered
|
||||
}
|
||||
|
||||
func dedupeResponsesWebsocketMergeInputItems(items []responsesWebsocketMergeInputItem) []responsesWebsocketMergeInputItem {
|
||||
referencedCallIDs := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallOutputType(item.itemType) && item.callID != "" {
|
||||
referencedCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
keepIndexByID := make(map[string]int, len(items))
|
||||
keepReferencedByID := make(map[string]bool, len(items))
|
||||
for index, item := range items {
|
||||
if item.id == "" {
|
||||
continue
|
||||
}
|
||||
_, referenced := referencedCallIDs[item.callID]
|
||||
referenced = referenced && item.callID != ""
|
||||
if _, seen := keepIndexByID[item.id]; !seen {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
continue
|
||||
}
|
||||
if referenced || !keepReferencedByID[item.id] {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
}
|
||||
}
|
||||
|
||||
filtered := items[:0]
|
||||
for index, item := range items {
|
||||
if item.id != "" && keepIndexByID[item.id] != index {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
clear(items[len(filtered):])
|
||||
return filtered
|
||||
}
|
||||
|
||||
func marshalResponsesWebsocketMergeInputItems(items []responsesWebsocketMergeInputItem) []byte {
|
||||
outputLength := 2
|
||||
if len(items) > 1 {
|
||||
outputLength += len(items) - 1
|
||||
}
|
||||
for _, item := range items {
|
||||
outputLength += len(item.raw)
|
||||
}
|
||||
|
||||
// This allocation establishes ownership of the merged transcript and is the
|
||||
// only large allocation the merge path retains after it returns.
|
||||
out := make([]byte, 0, outputLength)
|
||||
out = append(out, '[')
|
||||
for index, item := range items {
|
||||
if index > 0 {
|
||||
out = append(out, ',')
|
||||
}
|
||||
out = append(out, item.raw...)
|
||||
}
|
||||
out = append(out, ']')
|
||||
return out
|
||||
}
|
||||
|
||||
func parseResponsesWebsocketInputItems(rawArray string) ([]responsesWebsocketInputItem, error) {
|
||||
return appendResponsesWebsocketInputItems(nil, rawArray)
|
||||
}
|
||||
|
||||
func appendResponsesWebsocketInputItems(items []responsesWebsocketInputItem, rawArray string) ([]responsesWebsocketInputItem, error) {
|
||||
rawArray = strings.TrimSpace(rawArray)
|
||||
if rawArray == "" {
|
||||
rawArray = "[]"
|
||||
}
|
||||
var rawItems []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal([]byte(rawArray), &rawItems); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
return appendResponsesWebsocketRawInputItems(items, rawItems)
|
||||
}
|
||||
|
||||
func appendResponsesWebsocketRawInputItems(items []responsesWebsocketInputItem, rawItems []json.RawMessage) ([]responsesWebsocketInputItem, error) {
|
||||
for _, rawItem := range rawItems {
|
||||
item, errItem := parseResponsesWebsocketInputItem(rawItem)
|
||||
if errItem != nil {
|
||||
return nil, errItem
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func parseResponsesWebsocketInputItem(rawItem json.RawMessage) (responsesWebsocketInputItem, error) {
|
||||
item := responsesWebsocketInputItem{raw: rawItem}
|
||||
trimmed := bytes.TrimSpace(rawItem)
|
||||
if len(trimmed) == 0 || trimmed[0] != '{' {
|
||||
return item, nil
|
||||
}
|
||||
var metadata struct {
|
||||
Type json.RawMessage `json:"type"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
CallID json.RawMessage `json:"call_id"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(trimmed, &metadata); errUnmarshal != nil {
|
||||
return responsesWebsocketInputItem{}, errUnmarshal
|
||||
}
|
||||
item.itemType = responsesWebsocketMetadataString(metadata.Type)
|
||||
item.id = responsesWebsocketMetadataString(metadata.ID)
|
||||
item.callID = responsesWebsocketMetadataString(metadata.CallID)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func responsesWebsocketMetadataString(raw json.RawMessage) string {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
|
||||
return ""
|
||||
}
|
||||
if raw[0] == '"' {
|
||||
var value string
|
||||
if errUnmarshal := json.Unmarshal(raw, &value); errUnmarshal == nil {
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(string(raw))
|
||||
}
|
||||
|
||||
func marshalResponsesWebsocketInputItems(items []responsesWebsocketInputItem) (string, error) {
|
||||
rawItems := make([]json.RawMessage, len(items))
|
||||
for index := range items {
|
||||
rawItems[index] = items[index].raw
|
||||
}
|
||||
out, errMarshal := json.Marshal(rawItems)
|
||||
if errMarshal != nil {
|
||||
return "", errMarshal
|
||||
}
|
||||
return string(out), nil
|
||||
}
|
||||
|
||||
func dedupeResponsesWebsocketFunctionCalls(items []responsesWebsocketInputItem) []responsesWebsocketInputItem {
|
||||
seenCallIDs := make(map[string]struct{}, len(items))
|
||||
filtered := items[:0]
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallType(item.itemType) && item.callID != "" {
|
||||
if _, ok := seenCallIDs[item.callID]; ok {
|
||||
continue
|
||||
}
|
||||
seenCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
clear(items[len(filtered):])
|
||||
return filtered
|
||||
}
|
||||
|
||||
func dedupeResponsesWebsocketInputItems(items []responsesWebsocketInputItem) []responsesWebsocketInputItem {
|
||||
// Collect the call_ids that are still referenced by tool-call output
|
||||
// items. When several input items share the same id, the one we keep must
|
||||
// preserve any call_id that has a matching output; otherwise the upstream
|
||||
// rejects the request with "No tool call found for function call output".
|
||||
referencedCallIDs := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
switch item.itemType {
|
||||
case "function_call_output", "custom_tool_call_output":
|
||||
if item.callID != "" {
|
||||
referencedCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For each id, choose the index to keep. The default is the last
|
||||
// occurrence (matching the original dedupe behavior), but we never replace
|
||||
// an item whose call_id still has a matching output with one that does not.
|
||||
keepIndexByID := make(map[string]int, len(items))
|
||||
keepReferencedByID := make(map[string]bool, len(items))
|
||||
for index, item := range items {
|
||||
if item.id == "" {
|
||||
continue
|
||||
}
|
||||
_, referenced := referencedCallIDs[item.callID]
|
||||
referenced = referenced && item.callID != ""
|
||||
if _, seen := keepIndexByID[item.id]; !seen {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
continue
|
||||
}
|
||||
if referenced || !keepReferencedByID[item.id] {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
}
|
||||
}
|
||||
|
||||
filtered := items[:0]
|
||||
for index, item := range items {
|
||||
if item.id != "" && keepIndexByID[item.id] != index {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
clear(items[len(filtered):])
|
||||
return filtered
|
||||
}
|
||||
|
||||
func dedupeResponsesWebsocketInputItemsByID(payload []byte) []byte {
|
||||
input := gjson.GetBytes(payload, "input")
|
||||
if !input.Exists() || !input.IsArray() {
|
||||
return payload
|
||||
}
|
||||
dedupedInput, errDedupe := dedupeInputItemsByID(input.Raw)
|
||||
if errDedupe != nil || dedupedInput == input.Raw {
|
||||
return payload
|
||||
}
|
||||
updated, errSet := sjson.SetRawBytes(payload, "input", []byte(dedupedInput))
|
||||
if errSet != nil {
|
||||
return payload
|
||||
}
|
||||
return updated
|
||||
}
|
||||
|
||||
func dedupeInputItemsByID(rawArray string) (string, error) {
|
||||
items, errParse := parseResponsesWebsocketInputItems(rawArray)
|
||||
if errParse != nil {
|
||||
return "", errParse
|
||||
}
|
||||
return marshalResponsesWebsocketInputItems(dedupeResponsesWebsocketInputItems(items))
|
||||
}
|
||||
|
||||
func normalizeResponsesWebsocketPassthroughRequest(rawJSON []byte, modelName string) ([]byte, *interfaces.ErrorMessage) {
|
||||
if !json.Valid(rawJSON) {
|
||||
return nil, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("invalid websocket request JSON"),
|
||||
}
|
||||
}
|
||||
|
||||
requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String())
|
||||
switch requestType {
|
||||
case wsRequestTypeCreate, wsRequestTypeAppend:
|
||||
default:
|
||||
return nil, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("unsupported websocket request type: %s", requestType),
|
||||
}
|
||||
}
|
||||
|
||||
normalized := bytes.Clone(rawJSON)
|
||||
if strings.TrimSpace(gjson.GetBytes(normalized, "model").String()) == "" {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if modelName == "" {
|
||||
return nil, &interfaces.ErrorMessage{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Error: fmt.Errorf("missing model in response.create request"),
|
||||
}
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "model", modelName)
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "stream", true)
|
||||
return normalized, nil
|
||||
}
|
||||
|
|
@ -0,0 +1,575 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const (
|
||||
responsesWebsocketLargeTranscriptSize = 1 << 20
|
||||
responsesWebsocketBenchmarkTranscriptSize = 8 << 20
|
||||
)
|
||||
|
||||
var (
|
||||
responsesWebsocketMergedInputSink any
|
||||
responsesWebsocketNormalizedRequestSink []byte
|
||||
)
|
||||
|
||||
func TestMergeResponsesWebsocketInputMatchesCompatibilityScenarios(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
lastRequest string
|
||||
lastResponseOutput string
|
||||
appendInput string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "messages and paired tool call",
|
||||
lastRequest: `{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1","role":"user","content":"hello"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
want: `[{"type":"message","id":"msg-1","role":"user","content":"hello"},{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"},{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "duplicate function call keeps first",
|
||||
lastRequest: `{"input":[{"type":"function_call","id":"fc-first","call_id":"call-1","name":"first","arguments":"{}"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-second","call_id":"call-1","name":"second","arguments":"{}"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
want: `[{"type":"function_call","id":"fc-first","call_id":"call-1","name":"first","arguments":"{}"},{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "duplicate id keeps item referenced by output",
|
||||
lastRequest: `{"input":[{"type":"function_call","id":"fc-1","call_id":"call-kept","name":"first","arguments":"{}"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-1","call_id":"call-other","name":"second","arguments":"{}"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
want: `[{"type":"function_call","id":"fc-1","call_id":"call-kept","name":"first","arguments":"{}"},{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "raw JSON values and escaping",
|
||||
lastRequest: `{"input":[ {"type":"message","id":"msg-1","content":"<tag> & \\u263a"}, true ]}`,
|
||||
lastResponseOutput: `[null, 42, "line\\nvalue"]`,
|
||||
appendInput: `[{"id":"last","nested":{"value":[1,2,3]}}]`,
|
||||
want: `[{"type":"message","id":"msg-1","content":"<tag> & \\u263a"},true,null,42,"line\\nvalue",{"id":"last","nested":{"value":[1,2,3]}}]`,
|
||||
},
|
||||
{
|
||||
name: "large numbers retain exact JSON values",
|
||||
lastRequest: `{"input":[9007199254740993,{"id":"n","value":9223372036854775807}]}`,
|
||||
lastResponseOutput: `[18446744073709551615]`,
|
||||
appendInput: `[{"id":"decimal","value":1.0000000000000000001}]`,
|
||||
want: `[9007199254740993,{"id":"n","value":9223372036854775807},18446744073709551615,{"id":"decimal","value":1.0000000000000000001}]`,
|
||||
},
|
||||
{
|
||||
name: "invalid response output remains ignored",
|
||||
lastRequest: `{"input":[{"id":"first"}]}`,
|
||||
lastResponseOutput: `[{"id":`,
|
||||
appendInput: `[{"id":"last"}]`,
|
||||
want: `[{"id":"first"},{"id":"last"}]`,
|
||||
},
|
||||
{
|
||||
name: "missing previous input and null append",
|
||||
lastRequest: `{"model":"gpt-5.4"}`,
|
||||
lastResponseOutput: `[{"id":"response"}]`,
|
||||
appendInput: `null`,
|
||||
want: `[{"id":"response"}]`,
|
||||
},
|
||||
{
|
||||
name: "null previous request",
|
||||
lastRequest: `null`,
|
||||
lastResponseOutput: `[]`,
|
||||
appendInput: `[{"id":"last"}]`,
|
||||
want: `[{"id":"last"}]`,
|
||||
},
|
||||
{
|
||||
name: "duplicate metadata keys follow encoding json",
|
||||
lastRequest: `{"input":[{"type":"message","type":"function_call","id":"first","id":"fc-1","call_id":"call-other","call_id":"call-kept"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-2","call_id":"call-kept"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
want: `[{"type":"function_call","id":"fc-1","call_id":"call-kept"},{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "case insensitive metadata dedupes function calls",
|
||||
lastRequest: `{"input":[{"Type":"function_call","ID":"fc-old","CALL_ID":"call-1","name":"first"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-new","call_id":"call-1","name":"second"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
want: `[{"Type":"function_call","ID":"fc-old","CALL_ID":"call-1","name":"first"},{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "case insensitive metadata keeps referenced duplicate id",
|
||||
lastRequest: `{"input":[{"Type":"function_call","Id":"fc-1","Call_Id":"call-kept","name":"first"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-1","call_id":"call-other","name":"second"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
want: `[{"Type":"function_call","Id":"fc-1","Call_Id":"call-kept","name":"first"},{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "mixed case duplicate metadata keeps last values",
|
||||
lastRequest: `{"input":[{"type":"message","TYPE":"function_call","id":"first","ID":"fc-1","call_id":"call-other","CALL_ID":"call-kept"}]}`,
|
||||
lastResponseOutput: `[{"type":"function_call","id":"fc-2","call_id":"call-kept"}]`,
|
||||
appendInput: `[{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
want: `[{"type":"message","TYPE":"function_call","id":"first","ID":"fc-1","call_id":"call-other","CALL_ID":"call-kept"},{"type":"function_call_output","id":"fco-1","call_id":"call-kept","output":"done"}]`,
|
||||
},
|
||||
{
|
||||
name: "duplicate previous input keeps last array",
|
||||
lastRequest: `{"input":[{"id":"old"}],"input":[{"id":"new"}]}`,
|
||||
lastResponseOutput: `[]`,
|
||||
appendInput: `[]`,
|
||||
want: `[{"id":"new"}]`,
|
||||
},
|
||||
{
|
||||
name: "previous input field matching is case insensitive",
|
||||
lastRequest: `{"Input":[{"id":"old"}],"INPUT":[{"id":"new"}]}`,
|
||||
lastResponseOutput: `[]`,
|
||||
appendInput: `[]`,
|
||||
want: `[{"id":"new"}]`,
|
||||
},
|
||||
{
|
||||
name: "last duplicate null clears previous input",
|
||||
lastRequest: `{"input":[{"id":"old"}],"input":null}`,
|
||||
lastResponseOutput: `[{"id":"response"}]`,
|
||||
appendInput: `[]`,
|
||||
want: `[{"id":"response"}]`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
legacy, errLegacy := mergeResponsesWebsocketInputReference([]byte(test.lastRequest), []byte(test.lastResponseOutput), test.appendInput)
|
||||
if errLegacy != nil {
|
||||
t.Fatalf("legacy merge failed: %v", errLegacy)
|
||||
}
|
||||
assertJSONSemanticallyEqual(t, []byte(legacy), test.want)
|
||||
|
||||
got, errGot := mergeResponsesWebsocketInput([]byte(test.lastRequest), []byte(test.lastResponseOutput), test.appendInput)
|
||||
if errGot != nil {
|
||||
t.Fatalf("mergeResponsesWebsocketInput() error = %v", errGot)
|
||||
}
|
||||
assertJSONSemanticallyEqual(t, []byte(got), test.want)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeResponsesWebsocketInputReturnsCompatibleErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
lastRequest string
|
||||
lastResponseOutput string
|
||||
appendInput string
|
||||
wantPrefix string
|
||||
wantSyntaxError bool
|
||||
}{
|
||||
{name: "invalid previous request", lastRequest: `{"input":`, appendInput: `[]`, wantPrefix: "invalid previous request input", wantSyntaxError: true},
|
||||
{name: "non-array previous input", lastRequest: `{"input":{"id":"item"}}`, appendInput: `[]`, wantPrefix: "invalid previous request input"},
|
||||
{name: "invalid appended input", lastRequest: `{"input":[]}`, appendInput: `[{"id":`, wantPrefix: "invalid request input", wantSyntaxError: true},
|
||||
{name: "non-array appended input", lastRequest: `{"input":[]}`, appendInput: `{"id":"item"}`, wantPrefix: "invalid request input"},
|
||||
{name: "array previous request", lastRequest: `[]`, appendInput: `[]`, wantPrefix: "invalid previous request input"},
|
||||
{name: "last duplicate previous input is non-array", lastRequest: `{"input":[],"input":{"id":"item"}}`, appendInput: `[]`, wantPrefix: "invalid previous request input"},
|
||||
{name: "earlier non-array previous input remains invalid", lastRequest: `{"input":{"id":"item"},"input":[]}`, appendInput: `[]`, wantPrefix: "invalid previous request input"},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
_, errGot := mergeResponsesWebsocketInput([]byte(test.lastRequest), []byte(test.lastResponseOutput), test.appendInput)
|
||||
if errGot == nil {
|
||||
t.Fatal("expected merge error")
|
||||
}
|
||||
if !strings.HasPrefix(errGot.Error(), test.wantPrefix+": ") {
|
||||
t.Fatalf("error = %q, want prefix %q", errGot, test.wantPrefix+": ")
|
||||
}
|
||||
if test.wantSyntaxError {
|
||||
var syntaxError *json.SyntaxError
|
||||
if !errors.As(errGot, &syntaxError) {
|
||||
t.Fatalf("error cause = %T, want *json.SyntaxError", errGot)
|
||||
}
|
||||
return
|
||||
}
|
||||
var typeError *json.UnmarshalTypeError
|
||||
if !errors.As(errGot, &typeError) {
|
||||
t.Fatalf("error cause = %T, want *json.UnmarshalTypeError", errGot)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeResponsesWebsocketInputMatchesReferenceAcrossGeneratedTranscripts(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
random := rand.New(rand.NewSource(0xC0DE))
|
||||
for iteration := 0; iteration < 250; iteration++ {
|
||||
previous := generatedResponsesWebsocketInput(t, random, random.Intn(12))
|
||||
response := generatedResponsesWebsocketInput(t, random, random.Intn(12))
|
||||
appendInput := generatedResponsesWebsocketInput(t, random, random.Intn(12))
|
||||
lastRequest := append(append([]byte(`{"model":"gpt-5.4","input":`), previous...), '}')
|
||||
|
||||
want, errWant := mergeResponsesWebsocketInputReference(lastRequest, response, string(appendInput))
|
||||
if errWant != nil {
|
||||
t.Fatalf("iteration %d reference merge failed: %v", iteration, errWant)
|
||||
}
|
||||
got, errGot := mergeResponsesWebsocketInput(lastRequest, response, string(appendInput))
|
||||
if errGot != nil {
|
||||
t.Fatalf("iteration %d merge failed: %v", iteration, errGot)
|
||||
}
|
||||
assertJSONSemanticallyEqual(t, []byte(got), want)
|
||||
}
|
||||
}
|
||||
|
||||
func generatedResponsesWebsocketInput(t *testing.T, random *rand.Rand, count int) []byte {
|
||||
t.Helper()
|
||||
|
||||
items := make([]any, 0, count)
|
||||
itemTypes := []string{"message", "function_call", "function_call_output", "custom_tool_call", "custom_tool_call_output", "reasoning"}
|
||||
for index := 0; index < count; index++ {
|
||||
if random.Intn(10) == 0 {
|
||||
items = append(items, []any{true, float64(index), nil}[random.Intn(3)])
|
||||
continue
|
||||
}
|
||||
item := map[string]any{
|
||||
"type": itemTypes[random.Intn(len(itemTypes))],
|
||||
"id": fmt.Sprintf("item-%d", random.Intn(8)),
|
||||
"call_id": fmt.Sprintf("call-%d", random.Intn(6)),
|
||||
"content": fmt.Sprintf("iteration-%d <tag> & \\u263a", index),
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
out, errMarshal := json.Marshal(items)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal generated input: %v", errMarshal)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestNormalizeResponseSubsequentRequestDetachesSourceBuffers(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
lastRequest := []byte(`{"model":"gpt-5.4","instructions":"keep me","input":[{"type":"message","id":"msg-1","role":"user","content":"history sentinel"}]}`)
|
||||
lastResponseOutput := []byte(`[{"type":"message","id":"msg-2","role":"assistant","content":[{"type":"output_text","text":"response sentinel"}]}]`)
|
||||
raw := []byte(`{"type":"response.create","input":[{"type":"message","id":"msg-3","role":"user","content":"append sentinel"}]}`)
|
||||
|
||||
normalized, next, errMessage := normalizeResponseSubsequentRequest(raw, lastRequest, lastResponseOutput, "", nil, false, false)
|
||||
if errMessage != nil {
|
||||
t.Fatalf("normalizeResponseSubsequentRequest() error = %v", errMessage.Error)
|
||||
}
|
||||
wantNormalized := bytes.Clone(normalized)
|
||||
wantNext := bytes.Clone(next)
|
||||
|
||||
for _, source := range [][]byte{lastRequest, lastResponseOutput, raw} {
|
||||
for index := range source {
|
||||
source[index] = 'x'
|
||||
}
|
||||
}
|
||||
runtime.KeepAlive(lastRequest)
|
||||
runtime.KeepAlive(lastResponseOutput)
|
||||
runtime.KeepAlive(raw)
|
||||
|
||||
if !bytes.Equal(normalized, wantNormalized) {
|
||||
t.Fatal("normalized request aliases a source buffer")
|
||||
}
|
||||
if !bytes.Equal(next, wantNext) {
|
||||
t.Fatal("stored next request aliases a source buffer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeResponsesWebsocketInputBoundsLargeTranscriptAllocations(t *testing.T) {
|
||||
if raceDetectorEnabled {
|
||||
t.Skip("allocation budgets are not meaningful with race detector instrumentation")
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
makeFixture func() ([]byte, []byte, string)
|
||||
}{
|
||||
{name: "single_large_item", makeFixture: responsesWebsocketLargeTranscriptFixture},
|
||||
{name: "many_messages_and_tool_pairs", makeFixture: func() ([]byte, []byte, string) {
|
||||
return responsesWebsocketManyItemsTranscriptFixture(responsesWebsocketLargeTranscriptSize)
|
||||
}},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
lastRequest, lastResponseOutput, appendInput := test.makeFixture()
|
||||
inputBytes := len(lastRequest) + len(lastResponseOutput) + len(appendInput)
|
||||
|
||||
result := testing.Benchmark(func(b *testing.B) {
|
||||
b.SetBytes(int64(inputBytes))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
merged, errMerge := mergeResponsesWebsocketInput(lastRequest, lastResponseOutput, appendInput)
|
||||
if errMerge != nil {
|
||||
b.Fatalf("mergeResponsesWebsocketInput() error = %v", errMerge)
|
||||
}
|
||||
responsesWebsocketMergedInputSink = merged
|
||||
}
|
||||
responsesWebsocketMergedInputSink = nil
|
||||
})
|
||||
|
||||
const (
|
||||
maxAllocationNumerator = 3
|
||||
maxAllocationDenominator = 2
|
||||
)
|
||||
maxAllocatedBytes := int64(inputBytes) * maxAllocationNumerator / maxAllocationDenominator
|
||||
t.Logf("merge allocated %d bytes per operation for %d input bytes", result.AllocedBytesPerOp(), inputBytes)
|
||||
if allocatedBytes := result.AllocedBytesPerOp(); allocatedBytes > maxAllocatedBytes {
|
||||
t.Fatalf("merging %d input bytes allocated %d bytes per operation, want at most %d", inputBytes, allocatedBytes, maxAllocatedBytes)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkNormalizeResponseSubsequentRequestTranscripts(b *testing.B) {
|
||||
tests := []struct {
|
||||
name string
|
||||
makeFixture func() ([]byte, []byte, string)
|
||||
}{
|
||||
{name: "single_large_item", makeFixture: func() ([]byte, []byte, string) {
|
||||
return responsesWebsocketTranscriptFixture(responsesWebsocketBenchmarkTranscriptSize)
|
||||
}},
|
||||
{name: "many_messages_and_tool_pairs", makeFixture: func() ([]byte, []byte, string) {
|
||||
return responsesWebsocketManyItemsTranscriptFixture(responsesWebsocketBenchmarkTranscriptSize)
|
||||
}},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
b.Run(test.name, func(b *testing.B) {
|
||||
lastRequest, lastResponseOutput, appendInput := test.makeFixture()
|
||||
raw := []byte(`{"type":"response.create","input":` + appendInput + `}`)
|
||||
|
||||
b.SetBytes(int64(len(lastRequest) + len(lastResponseOutput) + len(raw)))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
normalized, _, errMessage := normalizeResponseSubsequentRequest(raw, lastRequest, lastResponseOutput, "", nil, false, false)
|
||||
if errMessage != nil {
|
||||
b.Fatalf("normalizeResponseSubsequentRequest() error = %v", errMessage.Error)
|
||||
}
|
||||
responsesWebsocketNormalizedRequestSink = normalized
|
||||
}
|
||||
responsesWebsocketNormalizedRequestSink = nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func responsesWebsocketLargeTranscriptFixture() ([]byte, []byte, string) {
|
||||
return responsesWebsocketTranscriptFixture(responsesWebsocketLargeTranscriptSize)
|
||||
}
|
||||
|
||||
func responsesWebsocketTranscriptFixture(transcriptSize int) ([]byte, []byte, string) {
|
||||
lastRequest := []byte(`{"model":"gpt-5.4","instructions":"coding","stream":true,"input":[{"type":"message","id":"msg-large","role":"user","content":"` + strings.Repeat("x", transcriptSize) + `"}]}`)
|
||||
lastResponseOutput := []byte(`[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"}]`)
|
||||
appendInput := `[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"done"}]`
|
||||
return lastRequest, lastResponseOutput, appendInput
|
||||
}
|
||||
|
||||
func responsesWebsocketManyItemsTranscriptFixture(transcriptSize int) ([]byte, []byte, string) {
|
||||
const (
|
||||
messageCount = 512
|
||||
toolPairCount = 128
|
||||
)
|
||||
contentSize := max(transcriptSize/messageCount, 1)
|
||||
content := strings.Repeat("x", contentSize)
|
||||
|
||||
var lastRequest strings.Builder
|
||||
lastRequest.Grow(transcriptSize + messageCount*96)
|
||||
lastRequest.WriteString(`{"model":"gpt-5.4","instructions":"coding","stream":true,"input":[`)
|
||||
for index := range messageCount {
|
||||
if index > 0 {
|
||||
lastRequest.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&lastRequest, `{"type":"message","id":"msg-%d","role":"user","content":"%s"}`, index, content)
|
||||
}
|
||||
lastRequest.WriteString(`]}`)
|
||||
|
||||
var lastResponseOutput strings.Builder
|
||||
lastResponseOutput.Grow(toolPairCount * 112)
|
||||
lastResponseOutput.WriteByte('[')
|
||||
for index := range toolPairCount {
|
||||
if index > 0 {
|
||||
lastResponseOutput.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&lastResponseOutput, `{"type":"function_call","id":"fc-%d","call_id":"call-%d","name":"lookup","arguments":"{}"}`, index, index)
|
||||
}
|
||||
lastResponseOutput.WriteByte(']')
|
||||
|
||||
var appendInput strings.Builder
|
||||
appendInput.Grow(toolPairCount * 104)
|
||||
appendInput.WriteByte('[')
|
||||
for index := range toolPairCount {
|
||||
if index > 0 {
|
||||
appendInput.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&appendInput, `{"type":"function_call_output","id":"fco-%d","call_id":"call-%d","output":"done"}`, index, index)
|
||||
}
|
||||
appendInput.WriteByte(']')
|
||||
|
||||
return []byte(lastRequest.String()), []byte(lastResponseOutput.String()), appendInput.String()
|
||||
}
|
||||
|
||||
// The legacy oracle detects broad behavior drift from the implementation that
|
||||
// preceded the allocation optimization. Explicit compatibility scenarios above
|
||||
// remain the independent specification for important merge behavior.
|
||||
type referenceResponsesWebsocketInputItem struct {
|
||||
raw json.RawMessage
|
||||
itemType string
|
||||
id string
|
||||
callID string
|
||||
}
|
||||
|
||||
func mergeResponsesWebsocketInputReference(lastRequest []byte, lastResponseOutput []byte, appendRaw string) (string, error) {
|
||||
var previousRequest struct {
|
||||
Input []json.RawMessage `json:"input"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(lastRequest, &previousRequest); errUnmarshal != nil {
|
||||
return "", fmt.Errorf("invalid previous request input: %w", errUnmarshal)
|
||||
}
|
||||
items, errExisting := appendReferenceResponsesWebsocketRawInputItems(nil, previousRequest.Input)
|
||||
if errExisting != nil {
|
||||
return "", fmt.Errorf("invalid previous request input: %w", errExisting)
|
||||
}
|
||||
|
||||
var responseItems []json.RawMessage
|
||||
trimmedResponse := bytes.TrimSpace(lastResponseOutput)
|
||||
if len(trimmedResponse) > 0 && trimmedResponse[0] == '[' && json.Valid(trimmedResponse) {
|
||||
if errUnmarshal := json.Unmarshal(trimmedResponse, &responseItems); errUnmarshal != nil {
|
||||
return "", fmt.Errorf("invalid previous response output: %w", errUnmarshal)
|
||||
}
|
||||
}
|
||||
items, errResponse := appendReferenceResponsesWebsocketRawInputItems(items, responseItems)
|
||||
if errResponse != nil {
|
||||
return "", fmt.Errorf("invalid previous response output: %w", errResponse)
|
||||
}
|
||||
|
||||
appendRaw = strings.TrimSpace(appendRaw)
|
||||
if appendRaw == "" {
|
||||
appendRaw = "[]"
|
||||
}
|
||||
var appendItems []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal([]byte(appendRaw), &appendItems); errUnmarshal != nil {
|
||||
return "", fmt.Errorf("invalid request input: %w", errUnmarshal)
|
||||
}
|
||||
items, errAppend := appendReferenceResponsesWebsocketRawInputItems(items, appendItems)
|
||||
if errAppend != nil {
|
||||
return "", fmt.Errorf("invalid request input: %w", errAppend)
|
||||
}
|
||||
|
||||
items = dedupeReferenceResponsesWebsocketFunctionCalls(items)
|
||||
items = dedupeReferenceResponsesWebsocketInputItems(items)
|
||||
rawItems := make([]json.RawMessage, len(items))
|
||||
for index := range items {
|
||||
rawItems[index] = items[index].raw
|
||||
}
|
||||
out, errMarshal := json.Marshal(rawItems)
|
||||
if errMarshal != nil {
|
||||
return "", errMarshal
|
||||
}
|
||||
return string(out), nil
|
||||
}
|
||||
|
||||
func appendReferenceResponsesWebsocketRawInputItems(items []referenceResponsesWebsocketInputItem, rawItems []json.RawMessage) ([]referenceResponsesWebsocketInputItem, error) {
|
||||
for _, rawItem := range rawItems {
|
||||
item := referenceResponsesWebsocketInputItem{raw: rawItem}
|
||||
trimmed := bytes.TrimSpace(rawItem)
|
||||
if len(trimmed) > 0 && trimmed[0] == '{' {
|
||||
var metadata struct {
|
||||
Type json.RawMessage `json:"type"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
CallID json.RawMessage `json:"call_id"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(trimmed, &metadata); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
item.itemType = responsesWebsocketMetadataString(metadata.Type)
|
||||
item.id = responsesWebsocketMetadataString(metadata.ID)
|
||||
item.callID = responsesWebsocketMetadataString(metadata.CallID)
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func dedupeReferenceResponsesWebsocketFunctionCalls(items []referenceResponsesWebsocketInputItem) []referenceResponsesWebsocketInputItem {
|
||||
seenCallIDs := make(map[string]struct{}, len(items))
|
||||
filtered := items[:0]
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallType(item.itemType) && item.callID != "" {
|
||||
if _, ok := seenCallIDs[item.callID]; ok {
|
||||
continue
|
||||
}
|
||||
seenCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func dedupeReferenceResponsesWebsocketInputItems(items []referenceResponsesWebsocketInputItem) []referenceResponsesWebsocketInputItem {
|
||||
referencedCallIDs := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallOutputType(item.itemType) && item.callID != "" {
|
||||
referencedCallIDs[item.callID] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
keepIndexByID := make(map[string]int, len(items))
|
||||
keepReferencedByID := make(map[string]bool, len(items))
|
||||
for index, item := range items {
|
||||
if item.id == "" {
|
||||
continue
|
||||
}
|
||||
_, referenced := referencedCallIDs[item.callID]
|
||||
referenced = referenced && item.callID != ""
|
||||
if _, seen := keepIndexByID[item.id]; !seen {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
continue
|
||||
}
|
||||
if referenced || !keepReferencedByID[item.id] {
|
||||
keepIndexByID[item.id] = index
|
||||
keepReferencedByID[item.id] = referenced
|
||||
}
|
||||
}
|
||||
|
||||
filtered := items[:0]
|
||||
for index, item := range items {
|
||||
if item.id != "" && keepIndexByID[item.id] != index {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func assertJSONSemanticallyEqual(t *testing.T, got []byte, want string) {
|
||||
t.Helper()
|
||||
if !json.Valid(got) {
|
||||
t.Fatalf("invalid actual JSON:\n%s", got)
|
||||
}
|
||||
var gotValue any
|
||||
gotDecoder := json.NewDecoder(bytes.NewReader(got))
|
||||
gotDecoder.UseNumber()
|
||||
if errUnmarshal := gotDecoder.Decode(&gotValue); errUnmarshal != nil {
|
||||
t.Fatalf("invalid actual JSON: %v\n%s", errUnmarshal, got)
|
||||
}
|
||||
if !json.Valid([]byte(want)) {
|
||||
t.Fatalf("invalid expected JSON:\n%s", want)
|
||||
}
|
||||
var wantValue any
|
||||
wantDecoder := json.NewDecoder(strings.NewReader(want))
|
||||
wantDecoder.UseNumber()
|
||||
if errUnmarshal := wantDecoder.Decode(&wantValue); errUnmarshal != nil {
|
||||
t.Fatalf("invalid reference JSON: %v\n%s", errUnmarshal, want)
|
||||
}
|
||||
if !reflect.DeepEqual(gotValue, wantValue) {
|
||||
t.Fatalf("JSON values differ:\n got: %s\nwant: %s", got, want)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,237 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func websocketUpstreamSupportsIncrementalInput(attributes map[string]string, metadata map[string]any) bool {
|
||||
if len(attributes) > 0 {
|
||||
if raw := strings.TrimSpace(attributes["websockets"]); raw != "" {
|
||||
parsed, errParse := strconv.ParseBool(raw)
|
||||
if errParse == nil {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(metadata) == 0 {
|
||||
return false
|
||||
}
|
||||
raw, ok := metadata["websockets"]
|
||||
if !ok || raw == nil {
|
||||
return false
|
||||
}
|
||||
switch value := raw.(type) {
|
||||
case bool:
|
||||
return value
|
||||
case string:
|
||||
parsed, errParse := strconv.ParseBool(strings.TrimSpace(value))
|
||||
if errParse == nil {
|
||||
return parsed
|
||||
}
|
||||
default:
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsIncrementalInputForModel(modelName string) bool {
|
||||
auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName)
|
||||
for _, auth := range auths {
|
||||
if responsesWebsocketAuthSupportsIncrementalInput(auth) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsCompactionReplayForModel(modelName string) bool {
|
||||
auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName)
|
||||
if len(auths) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, auth := range auths {
|
||||
if !responsesWebsocketAuthSupportsCompactionReplay(auth) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) responsesWebsocketAvailableAuthsForModel(modelName string) ([]*coreauth.Auth, string) {
|
||||
if h == nil || h.AuthManager == nil {
|
||||
return nil, ""
|
||||
}
|
||||
resolvedModelName := responsesWebsocketResolvedModelName(modelName)
|
||||
providerSet, modelKey := responsesWebsocketProviderSetForModel(resolvedModelName)
|
||||
if len(providerSet) == 0 {
|
||||
return nil, modelKey
|
||||
}
|
||||
|
||||
registryRef := registry.GetGlobalRegistry()
|
||||
now := time.Now()
|
||||
auths := h.AuthManager.List()
|
||||
available := make([]*coreauth.Auth, 0, len(auths))
|
||||
for _, auth := range auths {
|
||||
if !responsesWebsocketAuthMatchesModel(auth, providerSet, modelKey, registryRef, now) {
|
||||
continue
|
||||
}
|
||||
available = append(available, auth)
|
||||
}
|
||||
return available, modelKey
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesCodexWebsocketPassthrough(modelName string) bool {
|
||||
return h.responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName)
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName string) bool {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if h == nil || h.AuthManager == nil || modelName == "" {
|
||||
return false
|
||||
}
|
||||
auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName)
|
||||
if len(auths) == 0 {
|
||||
return false
|
||||
}
|
||||
provider := ""
|
||||
for _, auth := range auths {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
authProvider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
if authProvider != "codex" && authProvider != "xai" {
|
||||
return false
|
||||
}
|
||||
if provider == "" {
|
||||
provider = authProvider
|
||||
if _, ok := h.AuthManager.Executor(provider); !ok {
|
||||
return false
|
||||
}
|
||||
} else if authProvider != provider {
|
||||
return false
|
||||
}
|
||||
if !websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return provider != ""
|
||||
}
|
||||
|
||||
func responsesWebsocketAuthSupportsIncrementalInput(auth *coreauth.Auth) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
return websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata)
|
||||
}
|
||||
|
||||
func responsesWebsocketPinnedAuthMatchesModel(auth *coreauth.Auth, modelName string, pinnedModelKey string, homeRuntime bool) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
providerSet, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName))
|
||||
providerKey := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
if _, ok := providerSet[providerKey]; !ok {
|
||||
return false
|
||||
}
|
||||
if !responsesWebsocketAuthAvailableForModel(auth, modelKey, time.Now()) {
|
||||
return false
|
||||
}
|
||||
|
||||
if homeRuntime {
|
||||
return strings.EqualFold(strings.TrimSpace(pinnedModelKey), strings.TrimSpace(modelKey))
|
||||
}
|
||||
return registry.GetGlobalRegistry().ClientSupportsModel(auth.ID, modelKey)
|
||||
}
|
||||
|
||||
func responsesWebsocketResolvedModelName(modelName string) string {
|
||||
initialSuffix := thinking.ParseSuffix(modelName)
|
||||
if initialSuffix.ModelName == "auto" {
|
||||
resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName)
|
||||
if initialSuffix.HasSuffix {
|
||||
return fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix)
|
||||
}
|
||||
return resolvedBase
|
||||
}
|
||||
return util.ResolveAutoModel(modelName)
|
||||
}
|
||||
|
||||
func responsesWebsocketProviderSetForModel(resolvedModelName string) (map[string]struct{}, string) {
|
||||
parsed := thinking.ParseSuffix(resolvedModelName)
|
||||
baseModel := strings.TrimSpace(parsed.ModelName)
|
||||
providers := util.GetProviderName(baseModel)
|
||||
if len(providers) == 0 && baseModel != resolvedModelName {
|
||||
providers = util.GetProviderName(resolvedModelName)
|
||||
}
|
||||
providerSet := make(map[string]struct{}, len(providers))
|
||||
for _, provider := range providers {
|
||||
providerKey := strings.TrimSpace(strings.ToLower(provider))
|
||||
if providerKey == "" {
|
||||
continue
|
||||
}
|
||||
providerSet[providerKey] = struct{}{}
|
||||
}
|
||||
modelKey := baseModel
|
||||
if modelKey == "" {
|
||||
modelKey = strings.TrimSpace(resolvedModelName)
|
||||
}
|
||||
return providerSet, modelKey
|
||||
}
|
||||
|
||||
func responsesWebsocketAuthMatchesModel(auth *coreauth.Auth, providerSet map[string]struct{}, modelKey string, registryRef *registry.ModelRegistry, now time.Time) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
providerKey := strings.TrimSpace(strings.ToLower(auth.Provider))
|
||||
if _, ok := providerSet[providerKey]; !ok {
|
||||
return false
|
||||
}
|
||||
if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(auth.ID, modelKey) {
|
||||
return false
|
||||
}
|
||||
return responsesWebsocketAuthAvailableForModel(auth, modelKey, now)
|
||||
}
|
||||
|
||||
func responsesWebsocketAuthSupportsCompactionReplay(auth *coreauth.Auth) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(strings.TrimSpace(auth.Provider), "codex")
|
||||
}
|
||||
|
||||
func responsesWebsocketAuthAvailableForModel(auth *coreauth.Auth, modelName string, now time.Time) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
if auth.Disabled || auth.Status == coreauth.StatusDisabled {
|
||||
return false
|
||||
}
|
||||
if modelName != "" && len(auth.ModelStates) > 0 {
|
||||
state, ok := auth.ModelStates[modelName]
|
||||
if (!ok || state == nil) && modelName != "" {
|
||||
baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName)
|
||||
if baseModel != "" && baseModel != modelName {
|
||||
state, ok = auth.ModelStates[baseModel]
|
||||
}
|
||||
}
|
||||
if ok && state != nil {
|
||||
if state.Status == coreauth.StatusDisabled {
|
||||
return false
|
||||
}
|
||||
if state.Unavailable && !state.NextRetryAfter.IsZero() && state.NextRetryAfter.After(now) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
}
|
||||
if auth.Unavailable && !auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
5803
backend/sdk/api/handlers/openai/openai_responses_websocket_test.go
Normal file
5803
backend/sdk/api/handlers/openai/openai_responses_websocket_test.go
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,336 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
requestlogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type websocketTimelineAppender interface {
|
||||
Append(eventType string, payload []byte, timestamp time.Time)
|
||||
}
|
||||
|
||||
type responsesWebsocketPinnedAuthState struct {
|
||||
authID string
|
||||
modelKey string
|
||||
}
|
||||
|
||||
type websocketTimelineLog struct {
|
||||
enabled bool
|
||||
source *requestlogging.FileBodySource
|
||||
builder *strings.Builder
|
||||
|
||||
currentPart io.WriteCloser
|
||||
currentPartHasLog bool
|
||||
}
|
||||
|
||||
func newWebsocketTimelineLog(enabled bool, source *requestlogging.FileBodySource) *websocketTimelineLog {
|
||||
if !enabled {
|
||||
return &websocketTimelineLog{}
|
||||
}
|
||||
if source == nil {
|
||||
return newInMemoryWebsocketTimelineLog()
|
||||
}
|
||||
return &websocketTimelineLog{
|
||||
enabled: true,
|
||||
source: source,
|
||||
}
|
||||
}
|
||||
|
||||
func newInMemoryWebsocketTimelineLog() *websocketTimelineLog {
|
||||
return &websocketTimelineLog{
|
||||
enabled: true,
|
||||
builder: &strings.Builder{},
|
||||
}
|
||||
}
|
||||
|
||||
func websocketTimelineSourceFromContext(c *gin.Context) *requestlogging.FileBodySource {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
value, exists := c.Get(requestlogging.WebsocketTimelineSourceContextKey)
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
source, ok := value.(*requestlogging.FileBodySource)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return source
|
||||
}
|
||||
|
||||
func (l *websocketTimelineLog) BeginRequest() {
|
||||
if l == nil || !l.enabled || l.source == nil {
|
||||
return
|
||||
}
|
||||
l.closeCurrentPart()
|
||||
part, errCreate := l.source.CreatePart("request")
|
||||
if errCreate != nil {
|
||||
log.WithError(errCreate).Warn("failed to create websocket request detail log")
|
||||
return
|
||||
}
|
||||
l.currentPart = part
|
||||
l.currentPartHasLog = false
|
||||
}
|
||||
|
||||
func (l *websocketTimelineLog) Append(eventType string, payload []byte, timestamp time.Time) {
|
||||
if l == nil || !l.enabled {
|
||||
return
|
||||
}
|
||||
data := formatWebsocketTimelineEvent(eventType, payload, timestamp)
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
if l.source != nil {
|
||||
if l.currentPart == nil {
|
||||
l.BeginRequest()
|
||||
}
|
||||
if l.currentPart == nil {
|
||||
return
|
||||
}
|
||||
if errWrite := writeWebsocketTimelinePart(l.currentPart, data, l.currentPartHasLog); errWrite != nil {
|
||||
log.WithError(errWrite).Warn("failed to write websocket request detail log")
|
||||
return
|
||||
}
|
||||
l.currentPartHasLog = true
|
||||
return
|
||||
}
|
||||
if l.builder != nil {
|
||||
writeWebsocketTimelineBuilder(l.builder, data)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *websocketTimelineLog) SetContext(c *gin.Context) {
|
||||
if l == nil || !l.enabled {
|
||||
return
|
||||
}
|
||||
l.closeCurrentPart()
|
||||
if l.source != nil {
|
||||
if l.source.HasPayload() {
|
||||
c.Set(requestlogging.WebsocketTimelineSourceContextKey, l.source)
|
||||
return
|
||||
}
|
||||
if errCleanup := l.source.Cleanup(); errCleanup != nil {
|
||||
log.WithError(errCleanup).Warn("failed to clean up empty websocket timeline log parts")
|
||||
}
|
||||
}
|
||||
if l.builder != nil {
|
||||
setWebsocketTimelineBody(c, l.builder.String())
|
||||
}
|
||||
}
|
||||
|
||||
func (l *websocketTimelineLog) String() string {
|
||||
if l == nil || !l.enabled {
|
||||
return ""
|
||||
}
|
||||
l.closeCurrentPart()
|
||||
if l.source != nil {
|
||||
data, errRead := l.source.Bytes()
|
||||
if errRead != nil {
|
||||
return ""
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
if l.builder == nil {
|
||||
return ""
|
||||
}
|
||||
return l.builder.String()
|
||||
}
|
||||
|
||||
func (l *websocketTimelineLog) closeCurrentPart() {
|
||||
if l == nil || l.currentPart == nil {
|
||||
return
|
||||
}
|
||||
if errClose := l.currentPart.Close(); errClose != nil {
|
||||
log.WithError(errClose).Warn("failed to close websocket request detail log")
|
||||
}
|
||||
l.currentPart = nil
|
||||
l.currentPartHasLog = false
|
||||
}
|
||||
|
||||
func writeWebsocketTimelinePart(w io.Writer, data []byte, prependNewline bool) error {
|
||||
if w == nil || len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
if prependNewline {
|
||||
if _, errWrite := io.WriteString(w, "\n"); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
}
|
||||
_, errWrite := w.Write(data)
|
||||
return errWrite
|
||||
}
|
||||
|
||||
func writeWebsocketTimelineBuilder(builder *strings.Builder, data []byte) {
|
||||
if builder == nil || len(data) == 0 {
|
||||
return
|
||||
}
|
||||
if builder.Len() > 0 {
|
||||
builder.WriteString("\n")
|
||||
}
|
||||
builder.Write(data)
|
||||
}
|
||||
|
||||
func appendWebsocketEvent(builder *strings.Builder, eventType string, payload []byte) {
|
||||
if builder == nil {
|
||||
return
|
||||
}
|
||||
trimmedPayload := bytes.TrimSpace(payload)
|
||||
if len(trimmedPayload) == 0 {
|
||||
return
|
||||
}
|
||||
if builder.Len() > 0 {
|
||||
builder.WriteString("\n")
|
||||
}
|
||||
builder.WriteString("websocket.")
|
||||
builder.WriteString(eventType)
|
||||
builder.WriteString("\n")
|
||||
builder.Write(trimmedPayload)
|
||||
builder.WriteString("\n")
|
||||
}
|
||||
|
||||
func websocketPayloadEventType(payload []byte) string {
|
||||
eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
|
||||
if eventType == "" {
|
||||
return "-"
|
||||
}
|
||||
return eventType
|
||||
}
|
||||
|
||||
func websocketPayloadPreview(payload []byte) string {
|
||||
trimmedPayload := bytes.TrimSpace(payload)
|
||||
if len(trimmedPayload) == 0 {
|
||||
return "<empty>"
|
||||
}
|
||||
previewText := strings.ReplaceAll(string(trimmedPayload), "\n", "\\n")
|
||||
previewText = strings.ReplaceAll(previewText, "\r", "\\r")
|
||||
return previewText
|
||||
}
|
||||
|
||||
func isResponsesWebsocketCompletionEvent(eventType string) bool {
|
||||
return eventType == wsEventTypeCompleted || eventType == wsEventTypeDone
|
||||
}
|
||||
|
||||
type responsesWebsocketPayloadError struct {
|
||||
status int
|
||||
payload []byte
|
||||
}
|
||||
|
||||
func (e *responsesWebsocketPayloadError) Error() string {
|
||||
if e == nil {
|
||||
return ""
|
||||
}
|
||||
return string(e.payload)
|
||||
}
|
||||
|
||||
func (e *responsesWebsocketPayloadError) StatusCode() int {
|
||||
if e == nil {
|
||||
return 0
|
||||
}
|
||||
return e.status
|
||||
}
|
||||
|
||||
func responsesWebsocketErrorMessageFromPayload(payload []byte) *interfaces.ErrorMessage {
|
||||
status := int(gjson.GetBytes(payload, "status").Int())
|
||||
if status <= 0 {
|
||||
status = int(gjson.GetBytes(payload, "status_code").Int())
|
||||
}
|
||||
if status <= 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
|
||||
trimmedPayload := bytes.TrimSpace(payload)
|
||||
if len(trimmedPayload) > 0 {
|
||||
return &interfaces.ErrorMessage{
|
||||
StatusCode: status,
|
||||
Error: &responsesWebsocketPayloadError{
|
||||
status: status,
|
||||
payload: bytes.Clone(trimmedPayload),
|
||||
},
|
||||
}
|
||||
}
|
||||
return &interfaces.ErrorMessage{StatusCode: status, Error: fmt.Errorf("%s", http.StatusText(status))}
|
||||
}
|
||||
|
||||
func setWebsocketTimelineBody(c *gin.Context, body string) {
|
||||
setWebsocketBody(c, wsTimelineBodyKey, body)
|
||||
}
|
||||
|
||||
func setWebsocketBody(c *gin.Context, key string, body string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
trimmedBody := strings.TrimSpace(body)
|
||||
if trimmedBody == "" {
|
||||
return
|
||||
}
|
||||
c.Set(key, []byte(trimmedBody))
|
||||
}
|
||||
|
||||
func writeResponsesWebsocketPayload(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, payload []byte, timestamp time.Time) error {
|
||||
if wsTimelineLog != nil {
|
||||
wsTimelineLog.Append("response", payload, timestamp)
|
||||
}
|
||||
if writer == nil || writer.conn == nil {
|
||||
return fmt.Errorf("responses websocket: writer is nil")
|
||||
}
|
||||
writer.writeMu.Lock()
|
||||
defer writer.writeMu.Unlock()
|
||||
if writer.closing.Load() {
|
||||
return websocket.ErrCloseSent
|
||||
}
|
||||
return writer.conn.WriteMessage(websocket.TextMessage, payload)
|
||||
}
|
||||
|
||||
func appendWebsocketTimelineDisconnect(timeline websocketTimelineAppender, err error, timestamp time.Time) {
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
if timeline != nil {
|
||||
timeline.Append("disconnect", []byte(err.Error()), timestamp)
|
||||
}
|
||||
}
|
||||
|
||||
func appendWebsocketTimelineEvent(builder *strings.Builder, eventType string, payload []byte, timestamp time.Time) {
|
||||
if builder == nil {
|
||||
return
|
||||
}
|
||||
writeWebsocketTimelineBuilder(builder, formatWebsocketTimelineEvent(eventType, payload, timestamp))
|
||||
}
|
||||
|
||||
func formatWebsocketTimelineEvent(eventType string, payload []byte, timestamp time.Time) []byte {
|
||||
trimmedPayload := bytes.TrimSpace(payload)
|
||||
if len(trimmedPayload) == 0 {
|
||||
return nil
|
||||
}
|
||||
var builder strings.Builder
|
||||
builder.WriteString("Timestamp: ")
|
||||
builder.WriteString(timestamp.Format(time.RFC3339Nano))
|
||||
builder.WriteString("\n")
|
||||
builder.WriteString("Event: websocket.")
|
||||
builder.WriteString(eventType)
|
||||
builder.WriteString("\n")
|
||||
builder.Write(trimmedPayload)
|
||||
builder.WriteString("\n")
|
||||
return []byte(builder.String())
|
||||
}
|
||||
|
||||
func markAPIResponseTimestamp(c *gin.Context) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if _, exists := c.Get("API_RESPONSE_TIMESTAMP"); exists {
|
||||
return
|
||||
}
|
||||
c.Set("API_RESPONSE_TIMESTAMP", time.Now())
|
||||
}
|
||||
|
|
@ -0,0 +1,675 @@
|
|||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
const (
|
||||
websocketToolOutputCacheMaxPerSession = 256
|
||||
websocketToolOutputCacheTTL = 30 * time.Minute
|
||||
)
|
||||
|
||||
var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
|
||||
var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
|
||||
var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()
|
||||
var defaultWebsocketToolCacheTransactionMu sync.RWMutex
|
||||
|
||||
type websocketToolOutputCache struct {
|
||||
mu sync.Mutex
|
||||
ttl time.Duration
|
||||
maxPerSession int
|
||||
sessions map[string]*websocketToolOutputSession
|
||||
}
|
||||
|
||||
type websocketToolOutputSession struct {
|
||||
lastSeen time.Time
|
||||
outputs map[string]json.RawMessage
|
||||
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
|
||||
}
|
||||
if maxPerSession <= 0 {
|
||||
maxPerSession = websocketToolOutputCacheMaxPerSession
|
||||
}
|
||||
return &websocketToolOutputCache{
|
||||
ttl: ttl,
|
||||
maxPerSession: maxPerSession,
|
||||
sessions: make(map[string]*websocketToolOutputSession),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketToolOutputCache) record(sessionKey string, callID string, item json.RawMessage) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
callID = strings.Clone(strings.TrimSpace(callID))
|
||||
if sessionKey == "" || callID == "" || c == nil {
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.cleanupLocked(now)
|
||||
|
||||
session, ok := c.sessions[sessionKey]
|
||||
if !ok || session == nil {
|
||||
session = &websocketToolOutputSession{
|
||||
lastSeen: now,
|
||||
outputs: make(map[string]json.RawMessage),
|
||||
}
|
||||
c.sessions[sessionKey] = session
|
||||
}
|
||||
session.lastSeen = now
|
||||
|
||||
if _, exists := session.outputs[callID]; !exists {
|
||||
session.order = append(session.order, callID)
|
||||
}
|
||||
session.outputs[callID] = append(json.RawMessage(nil), item...)
|
||||
|
||||
for len(session.order) > c.maxPerSession {
|
||||
evict := session.order[0]
|
||||
session.order[0] = ""
|
||||
session.order = session.order[1:]
|
||||
delete(session.outputs, evict)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketToolOutputCache) get(sessionKey string, callID string) (json.RawMessage, bool) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
callID = strings.TrimSpace(callID)
|
||||
if sessionKey == "" || callID == "" || c == nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.cleanupLocked(now)
|
||||
|
||||
session, ok := c.sessions[sessionKey]
|
||||
if !ok || session == nil {
|
||||
return nil, false
|
||||
}
|
||||
session.lastSeen = now
|
||||
item, ok := session.outputs[callID]
|
||||
if !ok || len(item) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return append(json.RawMessage(nil), item...), true
|
||||
}
|
||||
|
||||
func (c *websocketToolOutputCache) cleanupLocked(now time.Time) {
|
||||
if c == nil || c.ttl <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
for key, session := range c.sessions {
|
||||
if session == nil {
|
||||
delete(c.sessions, key)
|
||||
continue
|
||||
}
|
||||
if now.Sub(session.lastSeen) > c.ttl {
|
||||
delete(c.sessions, key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketToolOutputCache) deleteSession(sessionKey string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || c == nil {
|
||||
return
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
delete(c.sessions, sessionKey)
|
||||
}
|
||||
|
||||
func websocketDownstreamSessionKey(req *http.Request) string {
|
||||
if req == nil {
|
||||
return ""
|
||||
}
|
||||
if requestID := strings.TrimSpace(req.Header.Get("X-Client-Request-Id")); requestID != "" {
|
||||
return requestID
|
||||
}
|
||||
if raw := strings.TrimSpace(req.Header.Get("X-Codex-Turn-Metadata")); raw != "" {
|
||||
if sessionID := strings.TrimSpace(gjson.Get(raw, "session_id").String()); sessionID != "" {
|
||||
return sessionID
|
||||
}
|
||||
}
|
||||
if sessionID := strings.TrimSpace(req.Header.Get("Session-Id")); sessionID != "" {
|
||||
return sessionID
|
||||
}
|
||||
if sessionID := strings.TrimSpace(req.Header.Get("Session_id")); sessionID != "" {
|
||||
return sessionID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type websocketToolSessionRefCounter struct {
|
||||
mu sync.Mutex
|
||||
counts map[string]int
|
||||
}
|
||||
|
||||
func newWebsocketToolSessionRefCounter() *websocketToolSessionRefCounter {
|
||||
return &websocketToolSessionRefCounter{counts: make(map[string]int)}
|
||||
}
|
||||
|
||||
func (c *websocketToolSessionRefCounter) acquire(sessionKey string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || c == nil {
|
||||
return
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.counts[sessionKey]++
|
||||
}
|
||||
|
||||
func (c *websocketToolSessionRefCounter) release(sessionKey string) bool {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || c == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
count := c.counts[sessionKey]
|
||||
if count <= 1 {
|
||||
delete(c.counts, sessionKey)
|
||||
return true
|
||||
}
|
||||
c.counts[sessionKey] = count - 1
|
||||
return false
|
||||
}
|
||||
|
||||
func retainResponsesWebsocketToolCaches(sessionKey string) {
|
||||
defaultWebsocketToolCacheTransactionMu.Lock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.Unlock()
|
||||
if defaultWebsocketToolSessionRefs == nil {
|
||||
return
|
||||
}
|
||||
defaultWebsocketToolSessionRefs.acquire(sessionKey)
|
||||
}
|
||||
|
||||
func releaseResponsesWebsocketToolCaches(sessionKey string) {
|
||||
defaultWebsocketToolCacheTransactionMu.Lock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.Unlock()
|
||||
if defaultWebsocketToolSessionRefs == nil {
|
||||
return
|
||||
}
|
||||
if !defaultWebsocketToolSessionRefs.release(sessionKey) {
|
||||
return
|
||||
}
|
||||
if defaultWebsocketToolOutputCache != nil {
|
||||
defaultWebsocketToolOutputCache.deleteSession(sessionKey)
|
||||
}
|
||||
if defaultWebsocketToolCallCache != nil {
|
||||
defaultWebsocketToolCallCache.deleteSession(sessionKey)
|
||||
}
|
||||
}
|
||||
|
||||
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) recordResponse(payload []byte) {
|
||||
if t == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
switch strings.TrimSpace(util.GetGJSONBytesNoCopy(payload, "type").String()) {
|
||||
case "response.completed":
|
||||
output := util.GetGJSONBytesNoCopy(payload, "response.output")
|
||||
if !output.Exists() || !output.IsArray() {
|
||||
return
|
||||
}
|
||||
output.ForEach(func(_, item gjson.Result) bool {
|
||||
if isCompleteResponsesWebsocketToolCall(item) {
|
||||
t.recordItem(payload, item)
|
||||
}
|
||||
return true
|
||||
})
|
||||
case "response.output_item.added", "response.output_item.done":
|
||||
item := util.GetGJSONBytesNoCopy(payload, "item")
|
||||
if isCompleteResponsesWebsocketToolCall(item) {
|
||||
t.recordItem(payload, item)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordItem(payload []byte, item gjson.Result) {
|
||||
if t == nil || !item.Exists() {
|
||||
return
|
||||
}
|
||||
rawItem, ok := responsesWebsocketRawMessageForResult(payload, item)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
t.recordRawItem(item.Get("type").String(), item.Get("call_id").String(), rawItem)
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordInputItem(item responsesWebsocketInputItem) {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
t.recordRawItem(item.itemType, item.callID, item.raw)
|
||||
}
|
||||
|
||||
func (t *responsesWebsocketToolCacheTurn) recordRawItem(itemType string, callID string, rawItem []byte) {
|
||||
if t == nil || (!isResponsesToolCallOutputType(itemType) && !isResponsesToolCallType(itemType)) {
|
||||
return
|
||||
}
|
||||
callID = strings.Clone(strings.TrimSpace(callID))
|
||||
if callID == "" || len(bytes.TrimSpace(rawItem)) == 0 {
|
||||
return
|
||||
}
|
||||
raw := append(json.RawMessage(nil), rawItem...)
|
||||
if isResponsesToolCallOutputType(itemType) {
|
||||
if _, exists := t.outputs[callID]; !exists {
|
||||
t.outputOrder = append(t.outputOrder, callID)
|
||||
}
|
||||
t.outputs[callID] = raw
|
||||
return
|
||||
}
|
||||
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, nil)
|
||||
}
|
||||
|
||||
func prepareResponsesWebsocketFallbackTurn(sessionKey string, payload []byte) ([]byte, *responsesWebsocketToolCacheTurn) {
|
||||
turn := newResponsesWebsocketToolCacheTurn(sessionKey)
|
||||
defaultWebsocketToolCacheTransactionMu.RLock()
|
||||
defer defaultWebsocketToolCacheTransactionMu.RUnlock()
|
||||
payload = repairResponsesWebsocketToolCallsWithCachesMode(
|
||||
defaultWebsocketToolOutputCache,
|
||||
defaultWebsocketToolCallCache,
|
||||
sessionKey,
|
||||
payload,
|
||||
false,
|
||||
turn,
|
||||
)
|
||||
return payload, turn
|
||||
}
|
||||
|
||||
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, nil)
|
||||
}
|
||||
|
||||
func repairResponsesWebsocketToolCallsWithCachesMode(
|
||||
outputCache, callCache *websocketToolOutputCache,
|
||||
sessionKey string,
|
||||
payload []byte,
|
||||
record bool,
|
||||
turn *responsesWebsocketToolCacheTurn,
|
||||
) []byte {
|
||||
if len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
|
||||
input, previousResponseID, ok := parseResponsesWebsocketRepairRequest(payload)
|
||||
if !ok {
|
||||
return payload
|
||||
}
|
||||
items, rawItems, ok := parseResponsesWebsocketInputItemsNoCopy(payload, input)
|
||||
if !ok {
|
||||
return payload
|
||||
}
|
||||
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
repairEnabled := sessionKey != "" && outputCache != nil
|
||||
updatedItems, errRepair := repairResponsesToolCallItems(
|
||||
outputCache,
|
||||
callCache,
|
||||
sessionKey,
|
||||
items,
|
||||
repairEnabled && responsesWebsocketMetadataString(previousResponseID) != "",
|
||||
record && repairEnabled,
|
||||
turn,
|
||||
repairEnabled,
|
||||
)
|
||||
if errRepair != nil || responsesWebsocketInputItemsEqualRaw(updatedItems, rawItems) {
|
||||
return payload
|
||||
}
|
||||
|
||||
updatedRaw, errMarshal := marshalResponsesWebsocketInputItems(updatedItems)
|
||||
if errMarshal != nil {
|
||||
return payload
|
||||
}
|
||||
updated, ok := replaceResponsesWebsocketRawResult(payload, input, []byte(updatedRaw))
|
||||
if !ok {
|
||||
return payload
|
||||
}
|
||||
return updated
|
||||
}
|
||||
|
||||
func parseResponsesWebsocketRepairRequest(payload []byte) (gjson.Result, json.RawMessage, bool) {
|
||||
if !json.Valid(payload) {
|
||||
return gjson.Result{}, nil, false
|
||||
}
|
||||
root := util.ParseGJSONBytesNoCopy(payload)
|
||||
if !root.IsObject() {
|
||||
return gjson.Result{}, nil, false
|
||||
}
|
||||
|
||||
var input gjson.Result
|
||||
var previousResponseID json.RawMessage
|
||||
inputFound := false
|
||||
valid := true
|
||||
root.ForEach(func(key, value gjson.Result) bool {
|
||||
switch {
|
||||
case strings.EqualFold(key.String(), "input"):
|
||||
if !value.IsArray() && strings.TrimSpace(value.Raw) != "null" {
|
||||
valid = false
|
||||
return false
|
||||
}
|
||||
input = value
|
||||
inputFound = true
|
||||
case strings.EqualFold(key.String(), "previous_response_id"):
|
||||
var ok bool
|
||||
previousResponseID, ok = responsesWebsocketRawMessageForResult(payload, value)
|
||||
if !ok {
|
||||
valid = false
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if !valid || !inputFound || !input.IsArray() {
|
||||
return gjson.Result{}, nil, false
|
||||
}
|
||||
return input, previousResponseID, true
|
||||
}
|
||||
|
||||
func replaceResponsesWebsocketRawResult(payload []byte, result gjson.Result, replacement []byte) ([]byte, bool) {
|
||||
if result.Index < 0 || result.Index > len(payload) || len(result.Raw) > len(payload)-result.Index {
|
||||
return nil, false
|
||||
}
|
||||
updated := make([]byte, 0, len(payload)-len(result.Raw)+len(replacement))
|
||||
updated = append(updated, payload[:result.Index]...)
|
||||
updated = append(updated, replacement...)
|
||||
updated = append(updated, payload[result.Index+len(result.Raw):]...)
|
||||
return updated, true
|
||||
}
|
||||
|
||||
func parseResponsesWebsocketInputItemsNoCopy(payload []byte, input gjson.Result) ([]responsesWebsocketInputItem, []json.RawMessage, bool) {
|
||||
var items []responsesWebsocketInputItem
|
||||
var rawItems []json.RawMessage
|
||||
valid := true
|
||||
input.ForEach(func(_, itemResult gjson.Result) bool {
|
||||
rawItem, ok := responsesWebsocketRawMessageForResult(payload, itemResult)
|
||||
if !ok {
|
||||
valid = false
|
||||
return false
|
||||
}
|
||||
item, errItem := parseResponsesWebsocketInputItem(rawItem)
|
||||
if errItem != nil {
|
||||
valid = false
|
||||
return false
|
||||
}
|
||||
items = append(items, item)
|
||||
rawItems = append(rawItems, rawItem)
|
||||
return true
|
||||
})
|
||||
if !valid {
|
||||
return nil, nil, false
|
||||
}
|
||||
return items, rawItems, true
|
||||
}
|
||||
|
||||
func responsesWebsocketRawMessageForResult(payload []byte, result gjson.Result) (json.RawMessage, bool) {
|
||||
if result.Index < 0 || result.Index > len(payload) || len(result.Raw) > len(payload)-result.Index {
|
||||
return nil, false
|
||||
}
|
||||
return payload[result.Index : result.Index+len(result.Raw)], true
|
||||
}
|
||||
|
||||
func repairResponsesToolCallItems(
|
||||
outputCache, callCache *websocketToolOutputCache,
|
||||
sessionKey string,
|
||||
items []responsesWebsocketInputItem,
|
||||
allowOrphanOutputs bool,
|
||||
record bool,
|
||||
turn *responsesWebsocketToolCacheTurn,
|
||||
repairEnabled bool,
|
||||
) ([]responsesWebsocketInputItem, error) {
|
||||
if !repairEnabled {
|
||||
return dedupeResponsesWebsocketInputItems(items), nil
|
||||
}
|
||||
|
||||
// First pass: record tool outputs and remember which call_ids have outputs in this payload.
|
||||
outputPresent := make(map[string]struct{}, len(items))
|
||||
callPresent := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
if turn != nil {
|
||||
turn.recordInputItem(item)
|
||||
}
|
||||
switch {
|
||||
case isResponsesToolCallOutputType(item.itemType):
|
||||
if item.callID == "" {
|
||||
continue
|
||||
}
|
||||
outputPresent[item.callID] = struct{}{}
|
||||
if record {
|
||||
outputCache.record(sessionKey, item.callID, item.raw)
|
||||
}
|
||||
case isResponsesToolCallType(item.itemType):
|
||||
if item.callID == "" {
|
||||
continue
|
||||
}
|
||||
callPresent[item.callID] = struct{}{}
|
||||
if record && callCache != nil {
|
||||
callCache.record(sessionKey, item.callID, item.raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
filtered := make([]responsesWebsocketInputItem, 0, len(items))
|
||||
insertedCalls := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
if isResponsesToolCallOutputType(item.itemType) {
|
||||
if item.callID == "" {
|
||||
// Upstream rejects tool outputs without a call_id; drop it.
|
||||
continue
|
||||
}
|
||||
|
||||
if _, ok := callPresent[item.callID]; ok {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
|
||||
if allowOrphanOutputs {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
|
||||
if callCache != nil {
|
||||
if cached, ok := callCache.get(sessionKey, item.callID); ok {
|
||||
if _, already := insertedCalls[item.callID]; !already {
|
||||
cachedItem, errCached := parseResponsesWebsocketInputItem(cached)
|
||||
if errCached != nil {
|
||||
return nil, errCached
|
||||
}
|
||||
filtered = append(filtered, cachedItem)
|
||||
insertedCalls[item.callID] = struct{}{}
|
||||
callPresent[item.callID] = struct{}{}
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Drop orphaned function_call_output items; upstream rejects transcripts with missing calls.
|
||||
continue
|
||||
}
|
||||
if !isResponsesToolCallType(item.itemType) {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
|
||||
if item.callID == "" {
|
||||
// Upstream rejects tool calls without a call_id; drop it.
|
||||
continue
|
||||
}
|
||||
|
||||
if _, ok := outputPresent[item.callID]; ok {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
|
||||
if allowOrphanOutputs {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
|
||||
if cached, ok := outputCache.get(sessionKey, item.callID); ok {
|
||||
cachedItem, errCached := parseResponsesWebsocketInputItem(cached)
|
||||
if errCached != nil {
|
||||
return nil, errCached
|
||||
}
|
||||
filtered = append(filtered, item, cachedItem)
|
||||
outputPresent[item.callID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
|
||||
// Drop orphaned function_call items; upstream rejects transcripts with missing outputs.
|
||||
}
|
||||
|
||||
return dedupeResponsesWebsocketInputItems(filtered), nil
|
||||
}
|
||||
|
||||
func responsesWebsocketInputItemsEqualRaw(items []responsesWebsocketInputItem, rawItems []json.RawMessage) bool {
|
||||
if len(items) != len(rawItems) {
|
||||
return false
|
||||
}
|
||||
for index := range items {
|
||||
if !bytes.Equal(items[index].raw, rawItems[index]) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func recordResponsesWebsocketToolCallsFromPayload(sessionKey string, payload []byte) {
|
||||
recordResponsesWebsocketToolCallsFromPayloadWithCache(defaultWebsocketToolCallCache, sessionKey, payload)
|
||||
}
|
||||
|
||||
func recordResponsesWebsocketToolCallsFromPayloadWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || cache == nil || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
eventType := strings.TrimSpace(util.GetGJSONBytesNoCopy(payload, "type").String())
|
||||
switch eventType {
|
||||
case "response.completed":
|
||||
output := util.GetGJSONBytesNoCopy(payload, "response.output")
|
||||
if !output.Exists() || !output.IsArray() {
|
||||
return
|
||||
}
|
||||
output.ForEach(func(_, item gjson.Result) bool {
|
||||
if !isCompleteResponsesWebsocketToolCall(item) {
|
||||
return true
|
||||
}
|
||||
rawItem, ok := responsesWebsocketRawMessageForResult(payload, item)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
cache.record(sessionKey, callID, rawItem)
|
||||
return true
|
||||
})
|
||||
case "response.output_item.added", "response.output_item.done":
|
||||
item := util.GetGJSONBytesNoCopy(payload, "item")
|
||||
if !isCompleteResponsesWebsocketToolCall(item) {
|
||||
return
|
||||
}
|
||||
rawItem, ok := responsesWebsocketRawMessageForResult(payload, item)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
callID := strings.TrimSpace(item.Get("call_id").String())
|
||||
cache.record(sessionKey, callID, rawItem)
|
||||
}
|
||||
}
|
||||
|
||||
func isResponsesToolCallType(itemType string) bool {
|
||||
switch strings.TrimSpace(itemType) {
|
||||
case "function_call", "custom_tool_call":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func isResponsesToolCallOutputType(itemType string) bool {
|
||||
switch strings.TrimSpace(itemType) {
|
||||
case "function_call_output", "custom_tool_call_output":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
1052
backend/sdk/api/handlers/openai/openai_videos_handlers.go
Normal file
1052
backend/sdk/api/handlers/openai/openai_videos_handlers.go
Normal file
File diff suppressed because it is too large
Load diff
1020
backend/sdk/api/handlers/openai/openai_videos_handlers_test.go
Normal file
1020
backend/sdk/api/handlers/openai/openai_videos_handlers_test.go
Normal file
File diff suppressed because it is too large
Load diff
5
backend/sdk/api/handlers/openai/race_disabled_test.go
Normal file
5
backend/sdk/api/handlers/openai/race_disabled_test.go
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
//go:build !race
|
||||
|
||||
package openai
|
||||
|
||||
const raceDetectorEnabled = false
|
||||
5
backend/sdk/api/handlers/openai/race_enabled_test.go
Normal file
5
backend/sdk/api/handlers/openai/race_enabled_test.go
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
//go:build race
|
||||
|
||||
package openai
|
||||
|
||||
const raceDetectorEnabled = true
|
||||
190
backend/sdk/api/handlers/openai_responses_stream_error.go
Normal file
190
backend/sdk/api/handlers/openai_responses_stream_error.go
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type openAIResponsesStreamErrorChunk struct {
|
||||
Type string `json:"type"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
SequenceNumber int `json:"sequence_number"`
|
||||
}
|
||||
|
||||
type openAIResponsesStreamFailedChunk struct {
|
||||
Type string `json:"type"`
|
||||
SequenceNumber int `json:"sequence_number"`
|
||||
Response openAIResponsesStreamFailedResponse `json:"response"`
|
||||
}
|
||||
|
||||
type openAIResponsesStreamFailedResponse struct {
|
||||
Status string `json:"status"`
|
||||
Error map[string]any `json:"error"`
|
||||
}
|
||||
|
||||
func openAIResponsesStreamErrorCode(status int) string {
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
return "invalid_api_key"
|
||||
case http.StatusForbidden:
|
||||
return "insufficient_quota"
|
||||
case http.StatusTooManyRequests:
|
||||
return "rate_limit_exceeded"
|
||||
case http.StatusNotFound:
|
||||
return "model_not_found"
|
||||
case http.StatusRequestTimeout:
|
||||
return "request_timeout"
|
||||
default:
|
||||
if status >= http.StatusInternalServerError {
|
||||
return "internal_server_error"
|
||||
}
|
||||
if status >= http.StatusBadRequest {
|
||||
return "invalid_request_error"
|
||||
}
|
||||
return "unknown_error"
|
||||
}
|
||||
}
|
||||
|
||||
// BuildOpenAIResponsesStreamErrorChunk builds an OpenAI Responses streaming error chunk.
|
||||
//
|
||||
// Important: OpenAI's HTTP error bodies are shaped like {"error":{...}}; those are valid for
|
||||
// non-streaming responses, but streaming clients validate SSE `data:` payloads against a union
|
||||
// of chunks that requires a top-level `type` field.
|
||||
func BuildOpenAIResponsesStreamErrorChunk(status int, errText string, sequenceNumber int) []byte {
|
||||
if status <= 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
if sequenceNumber < 0 {
|
||||
sequenceNumber = 0
|
||||
}
|
||||
|
||||
message := strings.TrimSpace(errText)
|
||||
if message == "" {
|
||||
message = http.StatusText(status)
|
||||
}
|
||||
|
||||
code := openAIResponsesStreamErrorCode(status)
|
||||
|
||||
trimmed := strings.TrimSpace(errText)
|
||||
if trimmed != "" && json.Valid([]byte(trimmed)) {
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(trimmed), &payload); err == nil {
|
||||
if t, ok := payload["type"].(string); ok && strings.TrimSpace(t) == "error" {
|
||||
if m, ok := payload["message"].(string); ok && strings.TrimSpace(m) != "" {
|
||||
message = strings.TrimSpace(m)
|
||||
}
|
||||
if v, ok := payload["code"]; ok && v != nil {
|
||||
if c, ok := v.(string); ok && strings.TrimSpace(c) != "" {
|
||||
code = strings.TrimSpace(c)
|
||||
} else {
|
||||
code = strings.TrimSpace(fmt.Sprint(v))
|
||||
}
|
||||
}
|
||||
if v, ok := payload["sequence_number"].(float64); ok && sequenceNumber == 0 {
|
||||
sequenceNumber = int(v)
|
||||
}
|
||||
}
|
||||
if e, ok := payload["error"].(map[string]any); ok {
|
||||
if m, ok := e["message"].(string); ok && strings.TrimSpace(m) != "" {
|
||||
message = strings.TrimSpace(m)
|
||||
}
|
||||
if v, ok := e["code"]; ok && v != nil {
|
||||
if c, ok := v.(string); ok && strings.TrimSpace(c) != "" {
|
||||
code = strings.TrimSpace(c)
|
||||
} else {
|
||||
code = strings.TrimSpace(fmt.Sprint(v))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if strings.TrimSpace(code) == "" {
|
||||
code = "unknown_error"
|
||||
}
|
||||
|
||||
data, err := json.Marshal(openAIResponsesStreamErrorChunk{
|
||||
Type: "error",
|
||||
Code: code,
|
||||
Message: message,
|
||||
SequenceNumber: sequenceNumber,
|
||||
})
|
||||
if err == nil {
|
||||
return data
|
||||
}
|
||||
|
||||
// Extremely defensive fallback.
|
||||
data, _ = json.Marshal(openAIResponsesStreamErrorChunk{
|
||||
Type: "error",
|
||||
Code: "internal_server_error",
|
||||
Message: message,
|
||||
SequenceNumber: sequenceNumber,
|
||||
})
|
||||
if len(data) > 0 {
|
||||
return data
|
||||
}
|
||||
return []byte(`{"type":"error","code":"internal_server_error","message":"internal error","sequence_number":0}`)
|
||||
}
|
||||
|
||||
func openAIResponsesStreamFailedErrorDetail(status int, errText, code, message string) map[string]any {
|
||||
var payload map[string]any
|
||||
if errUnmarshal := json.Unmarshal([]byte(strings.TrimSpace(errText)), &payload); errUnmarshal == nil {
|
||||
if errorDetail, ok := payload["error"].(map[string]any); ok {
|
||||
return errorDetail
|
||||
}
|
||||
if response, ok := payload["response"].(map[string]any); ok {
|
||||
if errorDetail, ok := response["error"].(map[string]any); ok {
|
||||
return errorDetail
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
errorType := "invalid_request_error"
|
||||
if status >= http.StatusInternalServerError {
|
||||
errorType = "server_error"
|
||||
}
|
||||
return map[string]any{
|
||||
"type": errorType,
|
||||
"code": code,
|
||||
"message": message,
|
||||
}
|
||||
}
|
||||
|
||||
// BuildOpenAIResponsesStreamFailedChunk builds the terminal Responses event used by official Codex clients.
|
||||
// It is intentionally separate from BuildOpenAIResponsesStreamErrorChunk so existing clients keep the legacy shape.
|
||||
func BuildOpenAIResponsesStreamFailedChunk(status int, errText string, sequenceNumber int) []byte {
|
||||
if status <= 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
if sequenceNumber < 0 {
|
||||
sequenceNumber = 0
|
||||
}
|
||||
|
||||
legacyChunk := BuildOpenAIResponsesStreamErrorChunk(status, errText, sequenceNumber)
|
||||
var legacyPayload openAIResponsesStreamErrorChunk
|
||||
if errUnmarshal := json.Unmarshal(legacyChunk, &legacyPayload); errUnmarshal != nil {
|
||||
legacyPayload.Code = openAIResponsesStreamErrorCode(status)
|
||||
legacyPayload.Message = http.StatusText(status)
|
||||
legacyPayload.SequenceNumber = sequenceNumber
|
||||
}
|
||||
if sequenceNumber == 0 && legacyPayload.SequenceNumber > 0 {
|
||||
sequenceNumber = legacyPayload.SequenceNumber
|
||||
}
|
||||
|
||||
data, errMarshal := json.Marshal(openAIResponsesStreamFailedChunk{
|
||||
Type: "response.failed",
|
||||
SequenceNumber: sequenceNumber,
|
||||
Response: openAIResponsesStreamFailedResponse{
|
||||
Status: "failed",
|
||||
Error: openAIResponsesStreamFailedErrorDetail(status, errText, legacyPayload.Code, legacyPayload.Message),
|
||||
},
|
||||
})
|
||||
if errMarshal == nil {
|
||||
return data
|
||||
}
|
||||
|
||||
return []byte(`{"type":"response.failed","sequence_number":0,"response":{"status":"failed","error":{"type":"server_error","code":"internal_server_error","message":"internal error"}}}`)
|
||||
}
|
||||
|
|
@ -0,0 +1,90 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildOpenAIResponsesStreamErrorChunk(t *testing.T) {
|
||||
chunk := BuildOpenAIResponsesStreamErrorChunk(http.StatusInternalServerError, "unexpected EOF", 0)
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(chunk, &payload); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if payload["type"] != "error" {
|
||||
t.Fatalf("type = %v, want %q", payload["type"], "error")
|
||||
}
|
||||
if payload["code"] != "internal_server_error" {
|
||||
t.Fatalf("code = %v, want %q", payload["code"], "internal_server_error")
|
||||
}
|
||||
if payload["message"] != "unexpected EOF" {
|
||||
t.Fatalf("message = %v, want %q", payload["message"], "unexpected EOF")
|
||||
}
|
||||
if payload["sequence_number"] != float64(0) {
|
||||
t.Fatalf("sequence_number = %v, want %v", payload["sequence_number"], 0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIResponsesStreamErrorChunkExtractsHTTPErrorBody(t *testing.T) {
|
||||
chunk := BuildOpenAIResponsesStreamErrorChunk(
|
||||
http.StatusInternalServerError,
|
||||
`{"error":{"message":"oops","type":"server_error","code":"internal_server_error"}}`,
|
||||
0,
|
||||
)
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(chunk, &payload); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if payload["type"] != "error" {
|
||||
t.Fatalf("type = %v, want %q", payload["type"], "error")
|
||||
}
|
||||
if payload["code"] != "internal_server_error" {
|
||||
t.Fatalf("code = %v, want %q", payload["code"], "internal_server_error")
|
||||
}
|
||||
if payload["message"] != "oops" {
|
||||
t.Fatalf("message = %v, want %q", payload["message"], "oops")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIResponsesStreamFailedChunkPreservesNestedError(t *testing.T) {
|
||||
chunk := BuildOpenAIResponsesStreamFailedChunk(
|
||||
http.StatusBadRequest,
|
||||
`{"error":{"type":"invalid_request","code":"cyber_policy","message":"blocked","param":null}}`,
|
||||
0,
|
||||
)
|
||||
|
||||
var payload struct {
|
||||
Type string `json:"type"`
|
||||
SequenceNumber int `json:"sequence_number"`
|
||||
Response struct {
|
||||
Status string `json:"status"`
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
} `json:"response"`
|
||||
}
|
||||
if err := json.Unmarshal(chunk, &payload); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if payload.Type != "response.failed" {
|
||||
t.Fatalf("type = %q, want %q", payload.Type, "response.failed")
|
||||
}
|
||||
if payload.SequenceNumber != 0 {
|
||||
t.Fatalf("sequence_number = %d, want 0", payload.SequenceNumber)
|
||||
}
|
||||
if payload.Response.Status != "failed" {
|
||||
t.Fatalf("response.status = %q, want %q", payload.Response.Status, "failed")
|
||||
}
|
||||
if payload.Response.Error.Type != "invalid_request" {
|
||||
t.Fatalf("response.error.type = %q, want %q", payload.Response.Error.Type, "invalid_request")
|
||||
}
|
||||
if payload.Response.Error.Code != "cyber_policy" {
|
||||
t.Fatalf("response.error.code = %q, want %q", payload.Response.Error.Code, "cyber_policy")
|
||||
}
|
||||
if payload.Response.Error.Message != "blocked" {
|
||||
t.Fatalf("response.error.message = %q, want %q", payload.Response.Error.Message, "blocked")
|
||||
}
|
||||
}
|
||||
73
backend/sdk/api/handlers/request_body.go
Normal file
73
backend/sdk/api/handlers/request_body.go
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
// ReadRequestBody reads the incoming request body and decodes supported
|
||||
// Content-Encoding values before handlers inspect JSON fields.
|
||||
func ReadRequestBody(c *gin.Context) ([]byte, error) {
|
||||
raw, err := c.GetRawData()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
encoding := ""
|
||||
if c != nil && c.Request != nil {
|
||||
encoding = strings.TrimSpace(c.Request.Header.Get("Content-Encoding"))
|
||||
}
|
||||
if encoding == "" || strings.EqualFold(encoding, "identity") {
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
decoded, err := decodeRequestBody(raw, encoding)
|
||||
if err != nil {
|
||||
if json.Valid(raw) {
|
||||
return raw, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return decoded, nil
|
||||
}
|
||||
|
||||
func decodeRequestBody(raw []byte, encoding string) ([]byte, error) {
|
||||
parts := strings.Split(encoding, ",")
|
||||
body := raw
|
||||
for i := len(parts) - 1; i >= 0; i-- {
|
||||
enc := strings.ToLower(strings.TrimSpace(parts[i]))
|
||||
switch enc {
|
||||
case "", "identity":
|
||||
continue
|
||||
case "zstd":
|
||||
decoded, err := decodeZstdRequestBody(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body = decoded
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported request content encoding: %s", enc)
|
||||
}
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func decodeZstdRequestBody(raw []byte) ([]byte, error) {
|
||||
decoder, err := zstd.NewReader(bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create zstd request decoder: %w", err)
|
||||
}
|
||||
defer decoder.Close()
|
||||
|
||||
decoded, err := io.ReadAll(decoder)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode zstd request body: %w", err)
|
||||
}
|
||||
return decoded, nil
|
||||
}
|
||||
168
backend/sdk/api/handlers/stream_forwarder.go
Normal file
168
backend/sdk/api/handlers/stream_forwarder.go
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
)
|
||||
|
||||
// PendingStreamError returns an immediately available non-nil stream error.
|
||||
func PendingStreamError(errs <-chan *interfaces.ErrorMessage) (*interfaces.ErrorMessage, bool) {
|
||||
if errs == nil {
|
||||
return nil, false
|
||||
}
|
||||
select {
|
||||
case errMsg, ok := <-errs:
|
||||
if ok && errMsg != nil {
|
||||
return errMsg, true
|
||||
}
|
||||
default:
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
type StreamForwardOptions struct {
|
||||
// KeepAliveInterval overrides the configured streaming keep-alive interval.
|
||||
// If nil, the configured default is used. If set to <= 0, keep-alives are disabled.
|
||||
KeepAliveInterval *time.Duration
|
||||
|
||||
// WriteChunk writes a single data chunk to the response body. It should not flush.
|
||||
WriteChunk func(chunk []byte)
|
||||
|
||||
// ChunkError optionally reports that WriteChunk emitted a terminal failure.
|
||||
// The failure is passed to cancel without writing another terminal payload.
|
||||
ChunkError func() *interfaces.ErrorMessage
|
||||
|
||||
// NormalizeTerminalError optionally replaces an upstream error before it is
|
||||
// written or passed to cancel.
|
||||
NormalizeTerminalError func(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage
|
||||
|
||||
// WriteTerminalError writes an error payload to the response body when streaming fails
|
||||
// after headers have already been committed. It should not flush.
|
||||
WriteTerminalError func(errMsg *interfaces.ErrorMessage)
|
||||
|
||||
// CloseError optionally validates a clean upstream channel close before WriteDone.
|
||||
// Returning an error surfaces it through WriteTerminalError instead of completing the stream.
|
||||
CloseError func() *interfaces.ErrorMessage
|
||||
|
||||
// WriteDone optionally writes a terminal marker when the upstream data channel closes
|
||||
// without an error (e.g. OpenAI's `[DONE]`). It should not flush.
|
||||
WriteDone func()
|
||||
|
||||
// WriteKeepAlive optionally writes a keep-alive heartbeat. It should not flush.
|
||||
// When nil, a standard SSE comment heartbeat is used.
|
||||
WriteKeepAlive func()
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) ForwardStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage, opts StreamForwardOptions) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if cancel == nil {
|
||||
return
|
||||
}
|
||||
|
||||
writeChunk := opts.WriteChunk
|
||||
if writeChunk == nil {
|
||||
writeChunk = func([]byte) {}
|
||||
}
|
||||
|
||||
writeKeepAlive := opts.WriteKeepAlive
|
||||
if writeKeepAlive == nil {
|
||||
writeKeepAlive = func() {
|
||||
_, _ = c.Writer.Write([]byte(": keep-alive\n\n"))
|
||||
}
|
||||
}
|
||||
|
||||
keepAliveInterval := StreamingKeepAliveInterval(h.Cfg)
|
||||
if opts.KeepAliveInterval != nil {
|
||||
keepAliveInterval = *opts.KeepAliveInterval
|
||||
}
|
||||
var keepAlive *time.Ticker
|
||||
var keepAliveC <-chan time.Time
|
||||
if keepAliveInterval > 0 {
|
||||
keepAlive = time.NewTicker(keepAliveInterval)
|
||||
defer keepAlive.Stop()
|
||||
keepAliveC = keepAlive.C
|
||||
}
|
||||
|
||||
var terminalErr *interfaces.ErrorMessage
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
// Prefer surfacing a terminal error if one is pending.
|
||||
if terminalErr == nil {
|
||||
if errMsg, ok := PendingStreamError(errs); ok {
|
||||
terminalErr = errMsg
|
||||
if opts.NormalizeTerminalError != nil {
|
||||
terminalErr = opts.NormalizeTerminalError(terminalErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
if terminalErr == nil && opts.CloseError != nil {
|
||||
terminalErr = opts.CloseError()
|
||||
}
|
||||
if terminalErr != nil {
|
||||
if opts.WriteTerminalError != nil {
|
||||
opts.WriteTerminalError(terminalErr)
|
||||
}
|
||||
flusher.Flush()
|
||||
cancel(terminalErr.Error)
|
||||
return
|
||||
}
|
||||
if opts.WriteDone != nil {
|
||||
opts.WriteDone()
|
||||
}
|
||||
flusher.Flush()
|
||||
cancel(nil)
|
||||
return
|
||||
}
|
||||
writeChunk(chunk)
|
||||
flusher.Flush()
|
||||
if opts.ChunkError != nil {
|
||||
chunkErr := opts.ChunkError()
|
||||
if chunkErr != nil {
|
||||
if opts.NormalizeTerminalError != nil {
|
||||
chunkErr = opts.NormalizeTerminalError(chunkErr)
|
||||
}
|
||||
if chunkErr != nil {
|
||||
cancel(chunkErr.Error)
|
||||
} else {
|
||||
cancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
errs = nil
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
terminalErr = errMsg
|
||||
if opts.NormalizeTerminalError != nil {
|
||||
terminalErr = opts.NormalizeTerminalError(terminalErr)
|
||||
}
|
||||
if opts.WriteTerminalError != nil {
|
||||
opts.WriteTerminalError(terminalErr)
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
var execErr error
|
||||
if terminalErr != nil {
|
||||
execErr = terminalErr.Error
|
||||
}
|
||||
cancel(execErr)
|
||||
return
|
||||
case <-keepAliveC:
|
||||
writeKeepAlive()
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
84
backend/sdk/api/handlers/stream_forwarder_test.go
Normal file
84
backend/sdk/api/handlers/stream_forwarder_test.go
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
)
|
||||
|
||||
func TestPendingStreamErrorReturnsBufferedError(t *testing.T) {
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
want := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("upstream failed")}
|
||||
errs <- want
|
||||
close(errs)
|
||||
|
||||
got, ok := PendingStreamError(errs)
|
||||
if !ok || got != want {
|
||||
t.Fatalf("PendingStreamError() = (%#v, %t), want (%#v, true)", got, ok, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateSSEDataJSONAllowsMultilinePayload(t *testing.T) {
|
||||
chunk := []byte("event: response.completed\n" +
|
||||
"data: {\"type\":\"response.completed\",\n" +
|
||||
"data: \"response\":{\"status\":\"completed\"}}\n\n")
|
||||
if err := validateSSEDataJSON(chunk); err != nil {
|
||||
t.Fatalf("validateSSEDataJSON() error = %v, want nil", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardStreamNormalizesErrorBeforeWriteAndCancel(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
data := make(chan []byte)
|
||||
close(data)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("raw secret")}
|
||||
close(errs)
|
||||
|
||||
var written, canceled string
|
||||
disabledKeepAlive := time.Duration(0)
|
||||
h := &BaseAPIHandler{}
|
||||
h.ForwardStream(c, recorder, func(err error) {
|
||||
if err != nil {
|
||||
canceled = err.Error()
|
||||
}
|
||||
}, data, errs, StreamForwardOptions{
|
||||
KeepAliveInterval: &disabledKeepAlive,
|
||||
NormalizeTerminalError: func(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage {
|
||||
return &interfaces.ErrorMessage{StatusCode: errMsg.StatusCode, Error: errors.New("safe error")}
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
written = errMsg.Error.Error()
|
||||
},
|
||||
})
|
||||
|
||||
if written != "safe error" || canceled != "safe error" {
|
||||
t.Fatalf("written=%q canceled=%q, want sanitized error", written, canceled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPendingStreamErrorIgnoresUnavailableErrors(t *testing.T) {
|
||||
closed := make(chan *interfaces.ErrorMessage)
|
||||
close(closed)
|
||||
|
||||
for name, errs := range map[string]<-chan *interfaces.ErrorMessage{
|
||||
"nil": nil,
|
||||
"closed empty": closed,
|
||||
"open empty": make(chan *interfaces.ErrorMessage),
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if got, ok := PendingStreamError(errs); ok || got != nil {
|
||||
t.Fatalf("PendingStreamError() = (%#v, %t), want (nil, false)", got, ok)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue