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
175 return175 return
176 }176 }
177 if token != "" {177 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)
179 }185 }
180 }186 }
181 187
@@ -254,6 +260,9 @@ func tokenValueForAuth(auth *coreauth.Auth) string { · 72 unmodified lines
254 if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" {260 if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" {
255 return v261 return v
256 }262 }
263 if v := strings.TrimSpace(auth.Attributes["session_token"]); v != "" {
264 return v
265 }
257 }266 }
258 return ""267 return ""
259}268}
internal/registry/devin_models.go+3−3
@@ -163,12 +163,12 @@ func sanitizeAndValidateDevinModels(models []*ModelInfo) ([]*ModelInfo, error) { · 162 unmodified lines
163 if !strings.HasPrefix(strings.ToLower(id), "devin/") {163 if !strings.HasPrefix(strings.ToLower(id), "devin/") {
164 id = "devin/" + id164 id = "devin/" + id
165 }165 }
166 id = strings.ToLower(id)
166 m.ID = id167 m.ID = id
167 seenID := strings.ToLower(id)168 if _, exists := seen[id]; exists {
168 if _, exists := seen[seenID]; exists {
169 return nil, fmt.Errorf("duplicate model id: %q", id)169 return nil, fmt.Errorf("duplicate model id: %q", id)
170 }170 }
171 seen[seenID] = struct{}{}171 seen[id] = struct{}{}
172 172
173 // Ensure proper default fields173 // Ensure proper default fields
174 if m.Type == "" {174 if m.Type == "" {
internal/runtime/executor/devin_executor.go+82−14
@@ -568,13 +568,47 @@ func (e *DevinExecutor) streamDevinFrames( · 567 unmodified lines
568 }568 }
569 569
570 if frameRes.Usage != nil {570 if frameRes.Usage != nil {
571 finalUsage = frameRes.Usage571 if finalUsage == nil {
572 } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil {572 finalUsage = frameRes.Usage
573 if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok {573 } else {
574 finalUsage = &helps.DevinUsage{574 if frameRes.Usage.PromptTokens > 0 {
575 PromptTokens: inTok,575 finalUsage.PromptTokens = frameRes.Usage.PromptTokens
576 CompletionTokens: outTok,576 }
577 CachedTokens: cachedTok,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
578 }612 }
579 }613 }
580 }614 }
@@ -913,13 +947,47 @@ func consumeDevinFramesToInteractions(body io.Reader, model, chatModelUID string · 332 unmodified lines
913 }947 }
914 948
915 if frameRes.Usage != nil {949 if frameRes.Usage != nil {
916 finalUsage = frameRes.Usage950 if finalUsage == nil {
917 } else if len(frameRes.ResponseDimensionGroups) > 0 && finalUsage == nil {951 finalUsage = frameRes.Usage
918 if inTok, outTok, cachedTok, ok := helps.ParseDevinResponseDimensionGroups(frameRes.ResponseDimensionGroups); ok {952 } else {
919 finalUsage = &helps.DevinUsage{953 if frameRes.Usage.PromptTokens > 0 {
920 PromptTokens: inTok,954 finalUsage.PromptTokens = frameRes.Usage.PromptTokens
921 CompletionTokens: outTok,955 }
922 CachedTokens: cachedTok,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
923 }991 }
924 }992 }
925 }993 }
internal/runtime/executor/helps/devin_wire.go+17−26
@@ -117,7 +117,7 @@ type DevinFrameResult struct { · 116 unmodified lines
117 Latency float64117 Latency float64
118 MessageID string118 MessageID string
119 Usage *DevinUsage119 Usage *DevinUsage
120 ResponseDimensionGroups []byte120 ResponseDimensionGroups [][]byte
121 UnknownFieldNumbers []int121 UnknownFieldNumbers []int
122}122}
123 123
@@ -588,7 +588,7 @@ func ParseDevinFrame(payload []byte) (DevinFrameResult, error) { · 464 unmodified lines
588 case 21:588 case 21:
589 res.DeltaSignatureType = string(val)589 res.DeltaSignatureType = string(val)
590 case 28:590 case 28:
591 res.ResponseDimensionGroups = val591 res.ResponseDimensionGroups = append(res.ResponseDimensionGroups, val)
592 default:592 default:
593 res.UnknownFieldNumbers = append(res.UnknownFieldNumbers, int(num))593 res.UnknownFieldNumbers = append(res.UnknownFieldNumbers, int(num))
594 }594 }
@@ -795,10 +795,10 @@ func parseDevinUsageField(data []byte) *DevinUsage { · 200 unmodified lines
795 u.Headers = make(map[string]string)795 u.Headers = make(map[string]string)
796 }796 }
797 u.Headers[k] = v797 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 != "" {
799 u.RequestID = v799 u.RequestID = v
800 }800 }
801 } else if isPrintableASCII(val) {801 } else if len(val) > 0 && isPrintableASCII(val) && u.RequestID == "" {
802 u.RequestID = string(val)802 u.RequestID = string(val)
803 }803 }
804 case 9:804 case 9:
@@ -827,31 +827,19 @@ func parseDevinUsageField(data []byte) *DevinUsage { · 22 unmodified lines
827 return u827 return u
828}828}
829 829
830// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) to extract Token Usage metrics:830// ParseDevinResponseDimensionGroups parses Field 28 (ResponseDimensionGroups) entries to extract Token Usage metrics:
831// input_tokens, output_tokens, cached_input_tokens.831// input_tokens, output_tokens, cached_input_tokens.
832func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTokens, cachedTokens int64, found bool) {832// Accepts one or more group payloads (each corresponding to a Field 28 value), or an outer envelope containing Tag 28.
833 pos := 0833func ParseDevinResponseDimensionGroups(groups ...[]byte) (promptTokens, completionTokens, cachedTokens int64, found bool) {
834 for pos < len(data) {834 for _, gBytes := range groups {
835 num, typ, n := protowire.ConsumeTag(data[pos:])835 if len(gBytes) == 0 {
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
846 continue836 continue
847 }837 }
848 gBytes, bn := protowire.ConsumeBytes(data[pos:])838 // If outer envelope carries Tag 28, unwrap it to get inner group bytes.
849 if bn <= 0 {839 if num, typ, n := protowire.ConsumeTag(gBytes); n > 0 && num == 28 && typ == protowire.BytesType {
850 break840 if inner, bn := protowire.ConsumeBytes(gBytes[n:]); bn > 0 {
851 }841 gBytes = inner
852 pos += bn842 }
853 if num != 28 {
854 continue
855 }843 }
856 844
857 gPos := 0845 gPos := 0
@@ -954,6 +942,9 @@ func ParseDevinResponseDimensionGroups(data []byte) (promptTokens, completionTok · 96 unmodified lines
954 found = true942 found = true
955 }943 }
956 }944 }
945 if found {
946 return promptTokens, completionTokens, cachedTokens, true
947 }
957 }948 }
958 }949 }
959 return promptTokens, completionTokens, cachedTokens, found950 return promptTokens, completionTokens, cachedTokens, found
internal/runtime/executor/helps/devin_wire_test.go+82−0
@@ -2,6 +2,7 @@ package helps · 1 unmodified line
2 2
3import (3import (
4 "bytes"4 "bytes"
5 "math"
5 "strings"6 "strings"
6 "testing"7 "testing"
7 8
@@ -705,3 +706,84 @@ func TestParseDevinUsageField_AnthropicRequestId(t *testing.T) { · 697 unmodified lines
705 t.Errorf("CachedTokens = %d, want 577", usage.CachedTokens)706 t.Errorf("CachedTokens = %d, want 577", usage.CachedTokens)
706 }707 }
707}708}
709
710func 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
776func 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