History

fix(devin): unwrap repeated field 28 groups, merge partial field 7 usage, and harden APICall escaping

sususu committed 2 months ago

Commit 4c331bb1 parent 07195205 files · +194 −44Browse files

Unified · Split

5 files
  1. internal/api/handlers/management/api_tools.go +10 −1
  2. internal/registry/devin_models.go +3 −3
  3. internal/runtime/executor/devin_executor.go +82 −14
  4. internal/runtime/executor/helps/devin_wire.go +17 −26
  5. internal/runtime/executor/helps/devin_wire_test.go +82 −0
internal/api/handlers/management/api_tools.go+10−1
@@ -175,7 +175,13 @@ func (h *Handler) APICall(c *gin.Context) { · 174 unmodified lines
175175 return
176176 }
177177 if token != "" {
178− body.Data = strings.ReplaceAll(body.Data, "$TOKEN$", token)
178+ replacement := token
179+ if json.Valid([]byte(body.Data)) && strings.ContainsAny(token, "\"\\\r\n\t") {
180+ if b, errMarshal := json.Marshal(token); errMarshal == nil && len(b) >= 2 {
181+ replacement = string(b[1 : len(b)-1])
182+ }
183+ }
184+ body.Data = strings.ReplaceAll(body.Data, "$TOKEN$", replacement)
179185 }
180186 }
181187
@@ -254,6 +260,9 @@ func tokenValueForAuth(auth *coreauth.Auth) string { · 72 unmodified lines
254260 if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" {
255261 return v
256262 }
263+ if v := strings.TrimSpace(auth.Attributes["session_token"]); v != "" {
264+ return v
265+ }
257266 }
258267 return ""
259268 }
internal/registry/devin_models.go+3−3
@@ -163,12 +163,12 @@ func sanitizeAndValidateDevinModels(models []*ModelInfo) ([]*ModelInfo, error) { · 162 unmodified lines
163163 if !strings.HasPrefix(strings.ToLower(id), "devin/") {
164164 id = "devin/" + id
165165 }
166+ id = strings.ToLower(id)
166167 m.ID = id
167− seenID := strings.ToLower(id)
168− if _, exists := seen[seenID]; exists {
168+ if _, exists := seen[id]; exists {
169169 return nil, fmt.Errorf("duplicate model id: %q", id)
170170 }
171− seen[seenID] = struct{}{}
171+ seen[id] = struct{}{}
172172
173173 // Ensure proper default fields
174174 if m.Type == "" {
internal/runtime/executor/devin_executor.go+82−14
@@ -568,13 +568,47 @@ func (e *DevinExecutor) streamDevinFrames( · 567 unmodified lines
568568 }
569569
570570 if frameRes.Usage != nil {
571− finalUsage = frameRes.Usage
572− } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil {
573− if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok {
574− finalUsage = &helps.DevinUsage{
575− PromptTokens: inTok,
576− CompletionTokens: outTok,
577− CachedTokens: cachedTok,
571+ if finalUsage == nil {
572+ finalUsage = frameRes.Usage
573+ } else {
574+ if frameRes.Usage.PromptTokens > 0 {
575+ finalUsage.PromptTokens = frameRes.Usage.PromptTokens
576+ }
577+ if frameRes.Usage.CompletionTokens > 0 {
578+ finalUsage.CompletionTokens = frameRes.Usage.CompletionTokens
579+ }
580+ if frameRes.Usage.CachedTokens > 0 {
581+ finalUsage.CachedTokens = frameRes.Usage.CachedTokens
582+ }
583+ if frameRes.Usage.RequestID != "" {
584+ finalUsage.RequestID = frameRes.Usage.RequestID
585+ }
586+ if frameRes.Usage.ModelName != "" {
587+ finalUsage.ModelName = frameRes.Usage.ModelName
588+ }
589+ if len(frameRes.Usage.Headers) > 0 {
590+ if finalUsage.Headers == nil {
591+ finalUsage.Headers = make(map[string]string, len(frameRes.Usage.Headers))
592+ }
593+ for hk, hv := range frameRes.Usage.Headers {
594+ finalUsage.Headers[hk] = hv
595+ }
596+ }
597+ }
598+ }
599+ if len(frameRes.ResponseDimensionGroups) > 0 && (finalUsage == nil || (finalUsage.PromptTokens == 0 && finalUsage.CompletionTokens == 0)) {
600+ if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups...); ok {
601+ if finalUsage == nil {
602+ finalUsage = &helps.DevinUsage{}
603+ }
604+ if finalUsage.PromptTokens == 0 {
605+ finalUsage.PromptTokens = inTok
606+ }
607+ if finalUsage.CompletionTokens == 0 {
608+ finalUsage.CompletionTokens = outTok
609+ }
610+ if finalUsage.CachedTokens == 0 {
611+ finalUsage.CachedTokens = cachedTok
578612 }
579613 }
580614 }
@@ -913,13 +947,47 @@ func consumeDevinFramesToInteractions(body io.Reader, model, chatModelUID string · 332 unmodified lines
913947 }
914948
915949 if frameRes.Usage != nil {
916− finalUsage = frameRes.Usage
917− } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil {
918− if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok {
919− finalUsage = &helps.DevinUsage{
920− PromptTokens: inTok,
921− CompletionTokens: outTok,
922− CachedTokens: cachedTok,
950+ if finalUsage == nil {
951+ finalUsage = frameRes.Usage
952+ } else {
953+ if frameRes.Usage.PromptTokens > 0 {
954+ finalUsage.PromptTokens = frameRes.Usage.PromptTokens
955+ }
956+ if frameRes.Usage.CompletionTokens > 0 {
957+ finalUsage.CompletionTokens = frameRes.Usage.CompletionTokens
958+ }
959+ if frameRes.Usage.CachedTokens > 0 {
960+ finalUsage.CachedTokens = frameRes.Usage.CachedTokens
961+ }
962+ if frameRes.Usage.RequestID != "" {
963+ finalUsage.RequestID = frameRes.Usage.RequestID
964+ }
965+ if frameRes.Usage.ModelName != "" {
966+ finalUsage.ModelName = frameRes.Usage.ModelName
967+ }
968+ if len(frameRes.Usage.Headers) > 0 {
969+ if finalUsage.Headers == nil {
970+ finalUsage.Headers = make(map[string]string, len(frameRes.Usage.Headers))
971+ }
972+ for hk, hv := range frameRes.Usage.Headers {
973+ finalUsage.Headers[hk] = hv
974+ }
975+ }
976+ }
977+ }
978+ if len(frameRes.ResponseDimensionGroups) > 0 && (finalUsage == nil || (finalUsage.PromptTokens == 0 && finalUsage.CompletionTokens == 0)) {
979+ if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups...); ok {
980+ if finalUsage == nil {
981+ finalUsage = &helps.DevinUsage{}
982+ }
983+ if finalUsage.PromptTokens == 0 {
984+ finalUsage.PromptTokens = inTok
985+ }
986+ if finalUsage.CompletionTokens == 0 {
987+ finalUsage.CompletionTokens = outTok
988+ }
989+ if finalUsage.CachedTokens == 0 {
990+ finalUsage.CachedTokens = cachedTok
923991 }
924992 }
925993 }
internal/runtime/executor/helps/devin_wire.go+17−26
@@ -117,7 +117,7 @@ type DevinFrameResult struct { · 116 unmodified lines
117117 Latency float64
118118 MessageID string
119119 Usage *DevinUsage
120− ResponseDimensionGroups []byte
120+ ResponseDimensionGroups [][]byte
121121 UnknownFieldNumbers []int
122122 }
123123
@@ -588,7 +588,7 @@ func ParseDevinFrame(payload []byte) (DevinFrameResult, error) { · 464 unmodified lines
588588 case 21:
589589 res.DeltaSignatureType = string(val)
590590 case 28:
591− res.ResponseDimensionGroups = val
591+ res.ResponseDimensionGroups = append(res.ResponseDimensionGroups, val)
592592 default:
593593 res.UnknownFieldNumbers = append(res.UnknownFieldNumbers, int(num))
594594 }
@@ -795,10 +795,10 @@ func parseDevinUsageField(data []byte) *DevinUsage { · 200 unmodified lines
795795 u.Headers = make(map[string]string)
796796 }
797797 u.Headers[k] = v
798− if strings.EqualFold(k, "x-request-id") || strings.EqualFold(k, "request-id") {
798+ if (strings.EqualFold(k, "x-request-id") || strings.EqualFold(k, "request-id")) && v != "" {
799799 u.RequestID = v
800800 }
801− } else if isPrintableASCII(val) {
801+ } else if len(val) > 0 && isPrintableASCII(val) && u.RequestID == "" {
802802 u.RequestID = string(val)
803803 }
804804 case 9:
@@ -827,31 +827,19 @@ func parseDevinUsageField(data []byte) *DevinUsage { · 22 unmodified lines
827827 return u
828828 }
829829
830−// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) to extract Token Usage metrics:
830+// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) entries to extract Token Usage metrics:
831831 // input_tokens, output_tokens, cached_input_tokens.
832−func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTokens, cachedTokens int64, found bool) {
833− pos := 0
834− for pos < len(data) {
835− num, typ, n := protowire.ConsumeTag(data[pos:])
836− if n <= 0 {
837− break
838− }
839− pos += n
840− if typ != protowire.BytesType {
841− nSkip := protowire.ConsumeFieldValue(num, typ, data[pos:])
842− if nSkip <= 0 {
843− break
844− }
845− pos += nSkip
832+// Accepts one or more group payloads (each corresponding to a Field 28 value), or an outer envelope containing Tag 28.
833+func ParseDevinResponseDimensionGroups(groups ...[]byte) (promptTokens, completionTokens, cachedTokens int64, found bool) {
834+ for _, gBytes := range groups {
835+ if len(gBytes) == 0 {
846836 continue
847837 }
848− gBytes, bn := protowire.ConsumeBytes(data[pos:])
849− if bn <= 0 {
850− break
851− }
852− pos += bn
853− if num != 28 {
854− continue
838+ // If outer envelope carries Tag 28, unwrap it to get inner group bytes.
839+ if num, typ, n := protowire.ConsumeTag(gBytes); n > 0 && num == 28 && typ == protowire.BytesType {
840+ if inner, bn := protowire.ConsumeBytes(gBytes[n:]); bn > 0 {
841+ gBytes = inner
842+ }
855843 }
856844
857845 gPos := 0
@@ -954,6 +942,9 @@ func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTok · 96 unmodified lines
954942 found = true
955943 }
956944 }
945+ if found {
946+ return promptTokens, completionTokens, cachedTokens, true
947+ }
957948 }
958949 }
959950 return promptTokens, completionTokens, cachedTokens, found
internal/runtime/executor/helps/devin_wire_test.go+82−0
@@ -2,6 +2,7 @@ package helps · 1 unmodified line
22
33 import (
44 "bytes"
5+ "math"
56 "strings"
67 "testing"
78
@@ -705,3 +706,84 @@ func TestParseDevinUsageField_AnthropicRequestId(t *testing.T) { · 697 unmodified lines
705706 t.Errorf("CachedTokens = %d, want 577", usage.CachedTokens)
706707 }
707708 }
709+
710+func TestParseDevinResponseDimensionGroups(t *testing.T) {
711+ buildMetric := func(key string, val float32) []byte {
712+ // Dimension submessage (Tag 4 of Metric)
713+ var dim []byte
714+ dim = protowire.AppendTag(dim, 2, protowire.Fixed32Type)
715+ dim = protowire.AppendFixed32(dim, math.Float32bits(val))
716+
717+ // Metric submessage (Tag 2 of Group)
718+ var metric []byte
719+ metric = protowire.AppendTag(metric, 4, protowire.BytesType)
720+ metric = protowire.AppendBytes(metric, dim)
721+ metric = protowire.AppendTag(metric, 5, protowire.BytesType)
722+ metric = protowire.AppendString(metric, key)
723+ return metric
724+ }
725+
726+ // Build Group (Tag 28)
727+ var group []byte
728+ group = protowire.AppendTag(group, 1, protowire.BytesType)
729+ group = protowire.AppendString(group, "Token Usage")
730+
731+ group = protowire.AppendTag(group, 2, protowire.BytesType)
732+ group = protowire.AppendBytes(group, buildMetric("input_tokens", 575.0))
733+
734+ group = protowire.AppendTag(group, 2, protowire.BytesType)
735+ group = protowire.AppendBytes(group, buildMetric("output_tokens", 5.0))
736+
737+ group = protowire.AppendTag(group, 2, protowire.BytesType)
738+ group = protowire.AppendBytes(group, buildMetric("cached_input_tokens", 128.0))
739+
740+ // Envelope Tag 28
741+ var root []byte
742+ root = protowire.AppendTag(root, 28, protowire.BytesType)
743+ root = protowire.AppendBytes(root, group)
744+
745+ promptTokens, completionTokens, cachedTokens, found := ParseDevinResponseDimensionGroups(root)
746+ if !found {
747+ t.Fatal("expected found = true")
748+ }
749+ if promptTokens != 575 {
750+ t.Errorf("promptTokens = %d, want 575", promptTokens)
751+ }
752+ if completionTokens != 5 {
753+ t.Errorf("completionTokens = %d, want 5", completionTokens)
754+ }
755+ if cachedTokens != 128 {
756+ t.Errorf("cachedTokens = %d, want 128", cachedTokens)
757+ }
758+
759+ // Verify inner group directly (as extracted by ParseDevinFrame case 28)
760+ p2, c2, ca2, found2 := ParseDevinResponseDimensionGroups(group)
761+ if !found2 || p2 != 575 || c2 != 5 || ca2 != 128 {
762+ t.Errorf("inner group ParseDevinResponseDimensionGroups = (%d,%d,%d,%t), want (575,5,128,true)", p2, c2, ca2, found2)
763+ }
764+
765+ // Verify multi-group where unrelated group precedes Token Usage
766+ var latencyGroup []byte
767+ latencyGroup = protowire.AppendTag(latencyGroup, 1, protowire.BytesType)
768+ latencyGroup = protowire.AppendString(latencyGroup, "Latency Metrics")
769+
770+ p3, c3, ca3, found3 := ParseDevinResponseDimensionGroups(latencyGroup, group)
771+ if !found3 || p3 != 575 || c3 != 5 || ca3 != 128 {
772+ t.Errorf("multi-group ParseDevinResponseDimensionGroups = (%d,%d,%d,%t), want (575,5,128,true)", p3, c3, ca3, found3)
773+ }
774+}
775+
776+func TestParseDevinResponseDimensionGroups_UnrelatedGroup(t *testing.T) {
777+ var group []byte
778+ group = protowire.AppendTag(group, 1, protowire.BytesType)
779+ group = protowire.AppendString(group, "Latency Metrics")
780+
781+ var root []byte
782+ root = protowire.AppendTag(root, 28, protowire.BytesType)
783+ root = protowire.AppendBytes(root, group)
784+
785+ promptTokens, completionTokens, cachedTokens, found := ParseDevinResponseDimensionGroups(root)
786+ if found {
787+ t.Errorf("expected found = false for unrelated group, got true with prompt=%d, comp=%d, cached=%d", promptTokens, completionTokens, cachedTokens)
788+ }
789+}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t