84 lines
2.4 KiB
Go
84 lines
2.4 KiB
Go
package grokbuild
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"net/http"
|
|
"slices"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
var keepaliveSSEComment = []byte(": keepalive\n\n")
|
|
|
|
// KeepaliveSSEComment returns the standard SSE comment used for keepalive.
|
|
func KeepaliveSSEComment() []byte {
|
|
return bytes.Clone(keepaliveSSEComment)
|
|
}
|
|
|
|
// IsGrokClientUserAgent checks if the user agent contains "grok-pager" or "grok-shell".
|
|
func IsGrokClientUserAgent(userAgent string) bool {
|
|
ua := strings.ToLower(userAgent)
|
|
return strings.Contains(ua, "grok-pager") || strings.Contains(ua, "grok-shell")
|
|
}
|
|
|
|
// IsGrokClientHeaders checks if the provided HTTP headers indicate a Grok client.
|
|
func IsGrokClientHeaders(headers http.Header) bool {
|
|
if headers == nil {
|
|
return false
|
|
}
|
|
for key, values := range headers {
|
|
if strings.EqualFold(key, "User-Agent") {
|
|
if slices.ContainsFunc(values, IsGrokClientUserAgent) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// IsGrokClientContext checks if either the context (e.g. Gin context) or headers indicate a Grok client.
|
|
func IsGrokClientContext(ctx context.Context, headers http.Header) bool {
|
|
if ctx != nil {
|
|
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
|
|
if IsGrokClientHeaders(ginCtx.Request.Header) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return IsGrokClientHeaders(headers)
|
|
}
|
|
|
|
// IsKeepalivePayload reports whether a JSON payload has type "keepalive".
|
|
func IsKeepalivePayload(payload []byte) bool {
|
|
return gjson.GetBytes(payload, "type").String() == "keepalive"
|
|
}
|
|
|
|
// IsKeepaliveSSELine reports whether an SSE line represents a keepalive event or data frame.
|
|
func IsKeepaliveSSELine(line []byte) bool {
|
|
trimmed := bytes.TrimSpace(line)
|
|
if bytes.HasPrefix(trimmed, []byte("event:")) {
|
|
eventName := bytes.TrimSpace(trimmed[6:])
|
|
return bytes.Equal(eventName, []byte("keepalive"))
|
|
}
|
|
if bytes.HasPrefix(trimmed, []byte("data:")) {
|
|
data := bytes.TrimSpace(trimmed[5:])
|
|
return IsKeepalivePayload(data)
|
|
}
|
|
return false
|
|
}
|
|
|
|
// TransformKeepaliveSSELine transforms a keepalive SSE line into an SSE comment line
|
|
// when isGrokClient is true. If the line is not a keepalive line or isGrokClient is false,
|
|
// it returns the original line and false.
|
|
func TransformKeepaliveSSELine(line []byte, isGrokClient bool) ([]byte, bool) {
|
|
if !isGrokClient {
|
|
return line, false
|
|
}
|
|
if IsKeepaliveSSELine(line) {
|
|
return bytes.Clone(keepaliveSSEComment), true
|
|
}
|
|
return line, false
|
|
}
|