diff --git a/internal/api/handlers/management/api_tools.go b/internal/api/handlers/management/api_tools.go index 3690100f7..63bdf183b 100644 --- a/internal/api/handlers/management/api_tools.go +++ b/internal/api/handlers/management/api_tools.go @@ -175,7 +175,13 @@ func (h *Handler) APICall(c *gin.Context) { return } if token != "" { - body.Data = strings.ReplaceAll(body.Data, "$TOKEN$", token) + replacement := token + if json.Valid([]byte(body.Data)) && strings.ContainsAny(token, "\"\\\r\n\t") { + if b, errMarshal := json.Marshal(token); errMarshal == nil && len(b) >= 2 { + replacement = string(b[1 : len(b)-1]) + } + } + body.Data = strings.ReplaceAll(body.Data, "$TOKEN$", replacement) } } @@ -254,6 +260,9 @@ func tokenValueForAuth(auth *coreauth.Auth) string { if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { return v } + if v := strings.TrimSpace(auth.Attributes["session_token"]); v != "" { + return v + } } return "" } diff --git a/internal/registry/devin_models.go b/internal/registry/devin_models.go index 56ea8bfdd..a0a1824be 100644 --- a/internal/registry/devin_models.go +++ b/internal/registry/devin_models.go @@ -163,12 +163,12 @@ func sanitizeAndValidateDevinModels(models []*ModelInfo) ([]*ModelInfo, error) { if !strings.HasPrefix(strings.ToLower(id), "devin/") { id = "devin/" + id } + id = strings.ToLower(id) m.ID = id - seenID := strings.ToLower(id) - if _, exists := seen[seenID]; exists { + if _, exists := seen[id]; exists { return nil, fmt.Errorf("duplicate model id: %q", id) } - seen[seenID] = struct{}{} + seen[id] = struct{}{} // Ensure proper default fields if m.Type == "" { diff --git a/internal/runtime/executor/devin_executor.go b/internal/runtime/executor/devin_executor.go index 88b88313c..79f50ab67 100644 --- a/internal/runtime/executor/devin_executor.go +++ b/internal/runtime/executor/devin_executor.go @@ -568,13 +568,47 @@ func (e *DevinExecutor) streamDevinFrames( } if frameRes.Usage != nil { - finalUsage = frameRes.Usage - } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil { - if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok { - finalUsage = &helps.DevinUsage{ - PromptTokens: inTok, - CompletionTokens: outTok, - CachedTokens: cachedTok, + if finalUsage == nil { + finalUsage = frameRes.Usage + } else { + if frameRes.Usage.PromptTokens > 0 { + finalUsage.PromptTokens = frameRes.Usage.PromptTokens + } + if frameRes.Usage.CompletionTokens > 0 { + finalUsage.CompletionTokens = frameRes.Usage.CompletionTokens + } + if frameRes.Usage.CachedTokens > 0 { + finalUsage.CachedTokens = frameRes.Usage.CachedTokens + } + if frameRes.Usage.RequestID != "" { + finalUsage.RequestID = frameRes.Usage.RequestID + } + if frameRes.Usage.ModelName != "" { + finalUsage.ModelName = frameRes.Usage.ModelName + } + if len(frameRes.Usage.Headers) > 0 { + if finalUsage.Headers == nil { + finalUsage.Headers = make(map[string]string, len(frameRes.Usage.Headers)) + } + for hk, hv := range frameRes.Usage.Headers { + finalUsage.Headers[hk] = hv + } + } + } + } + if len(frameRes.ResponseDimensionGroups) > 0 && (finalUsage == nil || (finalUsage.PromptTokens == 0 && finalUsage.CompletionTokens == 0)) { + if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups...); ok { + if finalUsage == nil { + finalUsage = &helps.DevinUsage{} + } + if finalUsage.PromptTokens == 0 { + finalUsage.PromptTokens = inTok + } + if finalUsage.CompletionTokens == 0 { + finalUsage.CompletionTokens = outTok + } + if finalUsage.CachedTokens == 0 { + finalUsage.CachedTokens = cachedTok } } } @@ -913,13 +947,47 @@ func consumeDevinFramesToInteractions(body io.Reader, model, chatModelUID string } if frameRes.Usage != nil { - finalUsage = frameRes.Usage - } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil { - if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok { - finalUsage = &helps.DevinUsage{ - PromptTokens: inTok, - CompletionTokens: outTok, - CachedTokens: cachedTok, + if finalUsage == nil { + finalUsage = frameRes.Usage + } else { + if frameRes.Usage.PromptTokens > 0 { + finalUsage.PromptTokens = frameRes.Usage.PromptTokens + } + if frameRes.Usage.CompletionTokens > 0 { + finalUsage.CompletionTokens = frameRes.Usage.CompletionTokens + } + if frameRes.Usage.CachedTokens > 0 { + finalUsage.CachedTokens = frameRes.Usage.CachedTokens + } + if frameRes.Usage.RequestID != "" { + finalUsage.RequestID = frameRes.Usage.RequestID + } + if frameRes.Usage.ModelName != "" { + finalUsage.ModelName = frameRes.Usage.ModelName + } + if len(frameRes.Usage.Headers) > 0 { + if finalUsage.Headers == nil { + finalUsage.Headers = make(map[string]string, len(frameRes.Usage.Headers)) + } + for hk, hv := range frameRes.Usage.Headers { + finalUsage.Headers[hk] = hv + } + } + } + } + if len(frameRes.ResponseDimensionGroups) > 0 && (finalUsage == nil || (finalUsage.PromptTokens == 0 && finalUsage.CompletionTokens == 0)) { + if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups...); ok { + if finalUsage == nil { + finalUsage = &helps.DevinUsage{} + } + if finalUsage.PromptTokens == 0 { + finalUsage.PromptTokens = inTok + } + if finalUsage.CompletionTokens == 0 { + finalUsage.CompletionTokens = outTok + } + if finalUsage.CachedTokens == 0 { + finalUsage.CachedTokens = cachedTok } } } diff --git a/internal/runtime/executor/helps/devin_wire.go b/internal/runtime/executor/helps/devin_wire.go index 79b96ab03..820678470 100644 --- a/internal/runtime/executor/helps/devin_wire.go +++ b/internal/runtime/executor/helps/devin_wire.go @@ -117,7 +117,7 @@ type DevinFrameResult struct { Latency float64 MessageID string Usage *DevinUsage - ResponseDimensionGroups []byte + ResponseDimensionGroups [][]byte UnknownFieldNumbers []int } @@ -588,7 +588,7 @@ func ParseDevinFrame(payload []byte) (DevinFrameResult, error) { case 21: res.DeltaSignatureType = string(val) case 28: - res.ResponseDimensionGroups = val + res.ResponseDimensionGroups = append(res.ResponseDimensionGroups, val) default: res.UnknownFieldNumbers = append(res.UnknownFieldNumbers, int(num)) } @@ -795,10 +795,10 @@ func parseDevinUsageField(data []byte) *DevinUsage { u.Headers = make(map[string]string) } u.Headers[k] = v - if strings.EqualFold(k, "x-request-id") || strings.EqualFold(k, "request-id") { + if (strings.EqualFold(k, "x-request-id") || strings.EqualFold(k, "request-id")) && v != "" { u.RequestID = v } - } else if isPrintableASCII(val) { + } else if len(val) > 0 && isPrintableASCII(val) && u.RequestID == "" { u.RequestID = string(val) } case 9: @@ -827,31 +827,19 @@ func parseDevinUsageField(data []byte) *DevinUsage { return u } -// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) to extract Token Usage metrics: +// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) entries to extract Token Usage metrics: // input_tokens, output_tokens, cached_input_tokens. -func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTokens, cachedTokens int64, found bool) { - pos := 0 - for pos < len(data) { - num, typ, n := protowire.ConsumeTag(data[pos:]) - if n <= 0 { - break +// Accepts one or more group payloads (each corresponding to a Field 28 value), or an outer envelope containing Tag 28. +func ParseDevinResponseDimensionGroups(groups ...[]byte) (promptTokens, completionTokens, cachedTokens int64, found bool) { + for _, gBytes := range groups { + if len(gBytes) == 0 { + continue } - pos += n - if typ != protowire.BytesType { - nSkip := protowire.ConsumeFieldValue(num, typ, data[pos:]) - if nSkip <= 0 { - break + // If outer envelope carries Tag 28, unwrap it to get inner group bytes. + if num, typ, n := protowire.ConsumeTag(gBytes); n > 0 && num == 28 && typ == protowire.BytesType { + if inner, bn := protowire.ConsumeBytes(gBytes[n:]); bn > 0 { + gBytes = inner } - pos += nSkip - continue - } - gBytes, bn := protowire.ConsumeBytes(data[pos:]) - if bn <= 0 { - break - } - pos += bn - if num != 28 { - continue } gPos := 0 @@ -954,6 +942,9 @@ func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTok found = true } } + if found { + return promptTokens, completionTokens, cachedTokens, true + } } } return promptTokens, completionTokens, cachedTokens, found diff --git a/internal/runtime/executor/helps/devin_wire_test.go b/internal/runtime/executor/helps/devin_wire_test.go index 71d924336..00b962b0b 100644 --- a/internal/runtime/executor/helps/devin_wire_test.go +++ b/internal/runtime/executor/helps/devin_wire_test.go @@ -2,6 +2,7 @@ package helps import ( "bytes" + "math" "strings" "testing" @@ -705,3 +706,84 @@ func TestParseDevinUsageField_AnthropicRequestId(t *testing.T) { t.Errorf("CachedTokens = %d, want 577", usage.CachedTokens) } } + +func TestParseDevinResponseDimensionGroups(t *testing.T) { + buildMetric := func(key string, val float32) []byte { + // Dimension submessage (Tag 4 of Metric) + var dim []byte + dim = protowire.AppendTag(dim, 2, protowire.Fixed32Type) + dim = protowire.AppendFixed32(dim, math.Float32bits(val)) + + // Metric submessage (Tag 2 of Group) + var metric []byte + metric = protowire.AppendTag(metric, 4, protowire.BytesType) + metric = protowire.AppendBytes(metric, dim) + metric = protowire.AppendTag(metric, 5, protowire.BytesType) + metric = protowire.AppendString(metric, key) + return metric + } + + // Build Group (Tag 28) + var group []byte + group = protowire.AppendTag(group, 1, protowire.BytesType) + group = protowire.AppendString(group, "Token Usage") + + group = protowire.AppendTag(group, 2, protowire.BytesType) + group = protowire.AppendBytes(group, buildMetric("input_tokens", 575.0)) + + group = protowire.AppendTag(group, 2, protowire.BytesType) + group = protowire.AppendBytes(group, buildMetric("output_tokens", 5.0)) + + group = protowire.AppendTag(group, 2, protowire.BytesType) + group = protowire.AppendBytes(group, buildMetric("cached_input_tokens", 128.0)) + + // Envelope Tag 28 + var root []byte + root = protowire.AppendTag(root, 28, protowire.BytesType) + root = protowire.AppendBytes(root, group) + + promptTokens, completionTokens, cachedTokens, found := ParseDevinResponseDimensionGroups(root) + if !found { + t.Fatal("expected found = true") + } + if promptTokens != 575 { + t.Errorf("promptTokens = %d, want 575", promptTokens) + } + if completionTokens != 5 { + t.Errorf("completionTokens = %d, want 5", completionTokens) + } + if cachedTokens != 128 { + t.Errorf("cachedTokens = %d, want 128", cachedTokens) + } + + // Verify inner group directly (as extracted by ParseDevinFrame case 28) + p2, c2, ca2, found2 := ParseDevinResponseDimensionGroups(group) + if !found2 || p2 != 575 || c2 != 5 || ca2 != 128 { + t.Errorf("inner group ParseDevinResponseDimensionGroups = (%d,%d,%d,%t), want (575,5,128,true)", p2, c2, ca2, found2) + } + + // Verify multi-group where unrelated group precedes Token Usage + var latencyGroup []byte + latencyGroup = protowire.AppendTag(latencyGroup, 1, protowire.BytesType) + latencyGroup = protowire.AppendString(latencyGroup, "Latency Metrics") + + p3, c3, ca3, found3 := ParseDevinResponseDimensionGroups(latencyGroup, group) + if !found3 || p3 != 575 || c3 != 5 || ca3 != 128 { + t.Errorf("multi-group ParseDevinResponseDimensionGroups = (%d,%d,%d,%t), want (575,5,128,true)", p3, c3, ca3, found3) + } +} + +func TestParseDevinResponseDimensionGroups_UnrelatedGroup(t *testing.T) { + var group []byte + group = protowire.AppendTag(group, 1, protowire.BytesType) + group = protowire.AppendString(group, "Latency Metrics") + + var root []byte + root = protowire.AppendTag(root, 28, protowire.BytesType) + root = protowire.AppendBytes(root, group) + + promptTokens, completionTokens, cachedTokens, found := ParseDevinResponseDimensionGroups(root) + if found { + t.Errorf("expected found = false for unrelated group, got true with prompt=%d, comp=%d, cached=%d", promptTokens, completionTokens, cachedTokens) + } +}