Refactor websocket logging and error handling
- Introduced new logging functions for websocket requests, handshakes, errors, and responses in `logging_helpers.go`. - Updated `CodexWebsocketsExecutor` to utilize the new logging functions for improved clarity and consistency in websocket operations. - Modified the handling of websocket upgrade rejections to log relevant metadata. - Changed the request body key to a timeline body key in `openai_responses_websocket.go` to better reflect its purpose. - Enhanced tests to verify the correct logging of websocket events and responses, including disconnect events and error handling scenarios.
This commit is contained in:
hkfires
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"html"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -19,9 +20,10 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
apiAttemptsKey = "API_UPSTREAM_ATTEMPTS"
|
||||
apiRequestKey = "API_REQUEST"
|
||||
apiResponseKey = "API_RESPONSE"
|
||||
apiAttemptsKey = "API_UPSTREAM_ATTEMPTS"
|
||||
apiRequestKey = "API_REQUEST"
|
||||
apiResponseKey = "API_RESPONSE"
|
||||
apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE"
|
||||
)
|
||||
|
||||
// UpstreamRequestLog captures the outbound upstream request details for logging.
|
||||
@@ -46,6 +48,7 @@ type upstreamAttempt struct {
|
||||
headersWritten bool
|
||||
bodyStarted bool
|
||||
bodyHasContent bool
|
||||
prevWasSSEEvent bool
|
||||
errorWritten bool
|
||||
}
|
||||
|
||||
@@ -173,15 +176,157 @@ func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
|
||||
attempt.response.WriteString("Body:\n")
|
||||
attempt.bodyStarted = true
|
||||
}
|
||||
currentChunkIsSSEEvent := bytes.HasPrefix(data, []byte("event:"))
|
||||
currentChunkIsSSEData := bytes.HasPrefix(data, []byte("data:"))
|
||||
if attempt.bodyHasContent {
|
||||
attempt.response.WriteString("\n\n")
|
||||
separator := "\n\n"
|
||||
if attempt.prevWasSSEEvent && currentChunkIsSSEData {
|
||||
separator = "\n"
|
||||
}
|
||||
attempt.response.WriteString(separator)
|
||||
}
|
||||
attempt.response.WriteString(string(data))
|
||||
attempt.bodyHasContent = true
|
||||
attempt.prevWasSSEEvent = currentChunkIsSSEEvent
|
||||
|
||||
updateAggregatedResponse(ginCtx, attempts)
|
||||
}
|
||||
|
||||
// RecordAPIWebsocketRequest stores an upstream websocket request event in Gin context.
|
||||
func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {
|
||||
if cfg == nil || !cfg.RequestLog {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
|
||||
builder.WriteString("Event: api.websocket.request\n")
|
||||
if info.URL != "" {
|
||||
builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))
|
||||
}
|
||||
if auth := formatAuthInfo(info); auth != "" {
|
||||
builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))
|
||||
}
|
||||
builder.WriteString("Headers:\n")
|
||||
writeHeaders(builder, info.Headers)
|
||||
builder.WriteString("\nBody:\n")
|
||||
if len(info.Body) > 0 {
|
||||
builder.Write(info.Body)
|
||||
} else {
|
||||
builder.WriteString("<empty>")
|
||||
}
|
||||
builder.WriteString("\n")
|
||||
|
||||
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
|
||||
}
|
||||
|
||||
// RecordAPIWebsocketHandshake stores the upstream websocket handshake response metadata.
|
||||
func RecordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, status int, headers http.Header) {
|
||||
if cfg == nil || !cfg.RequestLog {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
|
||||
builder.WriteString("Event: api.websocket.handshake\n")
|
||||
if status > 0 {
|
||||
builder.WriteString(fmt.Sprintf("Status: %d\n", status))
|
||||
}
|
||||
builder.WriteString("Headers:\n")
|
||||
writeHeaders(builder, headers)
|
||||
builder.WriteString("\n")
|
||||
|
||||
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
|
||||
}
|
||||
|
||||
// RecordAPIWebsocketUpgradeRejection stores a rejected websocket upgrade as an HTTP attempt.
|
||||
func RecordAPIWebsocketUpgradeRejection(ctx context.Context, cfg *config.Config, info UpstreamRequestLog, status int, headers http.Header, body []byte) {
|
||||
if cfg == nil || !cfg.RequestLog {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
|
||||
RecordAPIRequest(ctx, cfg, info)
|
||||
RecordAPIResponseMetadata(ctx, cfg, status, headers)
|
||||
AppendAPIResponseChunk(ctx, cfg, body)
|
||||
}
|
||||
|
||||
// WebsocketUpgradeRequestURL converts a websocket URL back to its HTTP handshake URL for logging.
|
||||
func WebsocketUpgradeRequestURL(rawURL string) string {
|
||||
trimmedURL := strings.TrimSpace(rawURL)
|
||||
if trimmedURL == "" {
|
||||
return ""
|
||||
}
|
||||
parsed, err := url.Parse(trimmedURL)
|
||||
if err != nil {
|
||||
return trimmedURL
|
||||
}
|
||||
switch strings.ToLower(parsed.Scheme) {
|
||||
case "ws":
|
||||
parsed.Scheme = "http"
|
||||
case "wss":
|
||||
parsed.Scheme = "https"
|
||||
}
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
// AppendAPIWebsocketResponse stores an upstream websocket response frame in Gin context.
|
||||
func AppendAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) {
|
||||
if cfg == nil || !cfg.RequestLog {
|
||||
return
|
||||
}
|
||||
data := bytes.TrimSpace(payload)
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
markAPIResponseTimestamp(ginCtx)
|
||||
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
|
||||
builder.WriteString("Event: api.websocket.response\n")
|
||||
builder.Write(data)
|
||||
builder.WriteString("\n")
|
||||
|
||||
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
|
||||
}
|
||||
|
||||
// RecordAPIWebsocketError stores an upstream websocket error event in Gin context.
|
||||
func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) {
|
||||
if cfg == nil || !cfg.RequestLog || err == nil {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
markAPIResponseTimestamp(ginCtx)
|
||||
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
|
||||
builder.WriteString("Event: api.websocket.error\n")
|
||||
if trimmed := strings.TrimSpace(stage); trimmed != "" {
|
||||
builder.WriteString(fmt.Sprintf("Stage: %s\n", trimmed))
|
||||
}
|
||||
builder.WriteString(fmt.Sprintf("Error: %s\n", err.Error()))
|
||||
|
||||
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
|
||||
}
|
||||
|
||||
func ginContextFrom(ctx context.Context) *gin.Context {
|
||||
ginCtx, _ := ctx.Value("gin").(*gin.Context)
|
||||
return ginCtx
|
||||
@@ -259,6 +404,40 @@ func updateAggregatedResponse(ginCtx *gin.Context, attempts []*upstreamAttempt)
|
||||
ginCtx.Set(apiResponseKey, []byte(builder.String()))
|
||||
}
|
||||
|
||||
func appendAPIWebsocketTimeline(ginCtx *gin.Context, chunk []byte) {
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
data := bytes.TrimSpace(chunk)
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
if existing, exists := ginCtx.Get(apiWebsocketTimelineKey); exists {
|
||||
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
|
||||
combined := make([]byte, 0, len(existingBytes)+len(data)+2)
|
||||
combined = append(combined, existingBytes...)
|
||||
if !bytes.HasSuffix(existingBytes, []byte("\n")) {
|
||||
combined = append(combined, '\n')
|
||||
}
|
||||
combined = append(combined, '\n')
|
||||
combined = append(combined, data...)
|
||||
ginCtx.Set(apiWebsocketTimelineKey, combined)
|
||||
return
|
||||
}
|
||||
}
|
||||
ginCtx.Set(apiWebsocketTimelineKey, bytes.Clone(data))
|
||||
}
|
||||
|
||||
func markAPIResponseTimestamp(ginCtx *gin.Context) {
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
if _, exists := ginCtx.Get("API_RESPONSE_TIMESTAMP"); exists {
|
||||
return
|
||||
}
|
||||
ginCtx.Set("API_RESPONSE_TIMESTAMP", time.Now())
|
||||
}
|
||||
|
||||
func writeHeaders(builder *strings.Builder, headers http.Header) {
|
||||
if builder == nil {
|
||||
return
|
||||
|
||||
Reference in New Issue
Block a user