Skip to content

Commit f90ed1f

Browse files
shreys7shreys7
authored
[POA-4481] Do not obfuscate postman internal team witnesses for LLM calls (#147)
1 parent 4c413b3 commit f90ed1f

File tree

3 files changed

+218
-2
lines changed

3 files changed

+218
-2
lines changed

‎trace/backend_collector_test.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -802,7 +802,8 @@ func TestAlwaysCapturePayloads(t *testing.T) {
802802
},
803803
Host: "example.com",
804804
Header: map[string][]string{
805-
"Content-Type": {"application/json"},
805+
"Content-Type": {"application/json"},
806+
"x-portkey-metadata": {`{"custom_postot_conversation_id":"12345678-1234-b1234-1234-123456789012","custom_postot_interaction_id":"12345678-1234-b1234-1234-123456789012","_prompt":"agent_system","_agent":"Root","_workflow":"agent_mode_chat","_domain":"agent_mode","_environment":"beta","_user":"916b323","_nrTraceId":"0d495a1379560f5d14f7f9c43d57bb07","_organization":"1"}`},
806807
},
807808
Body: memview.New([]byte(`{"name": "prince", "number": 6119717375543385000}`)),
808809
},
@@ -879,6 +880,7 @@ func TestAlwaysCapturePayloads(t *testing.T) {
879880
ApiType: pb.ApiType_HTTP_REST,
880881
},
881882
Args: map[string]*pb.Data{
883+
"27YhIi6o-7o=": newTestHeaderSpec(dataFromPrimitive(spec_util.NewPrimitiveString(`{"custom_postot_conversation_id":"12345678-1234-b1234-1234-123456789012","custom_postot_interaction_id":"12345678-1234-b1234-1234-123456789012","_prompt":"agent_system","_agent":"Root","_workflow":"agent_mode_chat","_domain":"agent_mode","_environment":"beta","_user":"916b323","_nrTraceId":"0d495a1379560f5d14f7f9c43d57bb07","_organization":"1"}`)), "x-portkey-metadata", 0),
882884
"nxnOc5Qy3D4=": newTestBodySpecFromStruct(0, pb.HTTPBody_JSON, "application/json", map[string]*pb.Data{
883885
"name": dataFromPrimitive(spec_util.NewPrimitiveString("prince")),
884886
"number": dataFromPrimitive(spec_util.NewPrimitiveInt64(6119717375543385000)),

‎trace/util.go‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,21 @@
11
package trace
22

33
import (
4+
"encoding/json"
45
"regexp"
56
"strings"
67

78
pb "github.com/akitasoftware/akita-ir/go/api_spec"
89
"github.com/akitasoftware/akita-libs/http_rest_methods"
10+
"github.com/akitasoftware/akita-libs/spec_util"
911
)
1012

13+
var internalPostmanTeamIDMap = map[string]string{
14+
"beta": "1",
15+
"stage": "2482",
16+
"production": "6029",
17+
}
18+
1119
// Returns true if the witness should be excluded from Repro Mode.
1220
//
1321
// XXX This is a stop-gap hack to exclude certain endpoints for Cloud API from
@@ -53,6 +61,30 @@ func hasMatchingPath(witness *pb.Witness, alwaysCapturePayloadsPathsRegex []*reg
5361
return false
5462
}
5563

64+
// PortkeyMetadata represents the structure of x-portkey-metadata header
65+
type PortkeyMetadata struct {
66+
Organization string `json:"_organization"`
67+
Environment string `json:"_environment"`
68+
}
69+
70+
// Returns true if the witness is from Postman internal team
71+
func isPostmanInternalTeam(witness *pb.Witness) bool {
72+
// Look for x-portkey-metadata header in the request args
73+
for _, data := range witness.GetMethod().GetArgs() {
74+
if header := spec_util.HTTPHeaderFromData(data); header != nil && strings.ToLower(header.GetKey()) == "x-portkey-metadata" {
75+
if primitive := data.GetPrimitive(); primitive != nil && primitive.GetStringValue() != nil {
76+
var metadata PortkeyMetadata
77+
if err := json.Unmarshal([]byte(primitive.GetStringValue().Value), &metadata); err == nil {
78+
if teamID, ok := internalPostmanTeamIDMap[metadata.Environment]; ok {
79+
return metadata.Organization == teamID
80+
}
81+
}
82+
}
83+
}
84+
}
85+
return false
86+
}
87+
5688
// Determines whether a given method has only error (4xx or 5xx) response codes.
5789
// Returns true if the method has at least one response and all response codes are 4xx or 5xx.
5890
// Here method will only have single response object because in agent each witness stores single request-response pair.
@@ -80,8 +112,10 @@ func shouldCapturePayload(witness *pb.Witness, alwaysCapturePayloadsPathsRegex [
80112
}
81113

82114
// Step 2: Check if the method path matches the always capture payloads regex.
115+
// Hack: We are using `alwaysCapturePayloads` arg to always capture payloads for LLM calls.
116+
// But, since we are not allowed to do that on Portkey production, we add the postman internal team check.
83117
if hasMatchingPath(witness, alwaysCapturePayloadsPathsRegex) {
84-
return true
118+
return isPostmanInternalTeam(witness)
85119
}
86120

87121
// Step 3: Check if the method has only error responses.

‎trace/util_test.go‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package trace
2+
3+
import (
4+
"regexp"
5+
"testing"
6+
7+
pb "github.com/akitasoftware/akita-ir/go/api_spec"
8+
"github.com/akitasoftware/akita-libs/spec_util"
9+
)
10+
11+
func TestIsPostmanInternalTeam(t *testing.T) {
12+
tests := []struct {
13+
name string
14+
witness *pb.Witness
15+
expected bool
16+
}{
17+
{
18+
name: "Postman internal team with organization 1",
19+
witness: &pb.Witness{
20+
Method: &pb.Method{
21+
Args: map[string]*pb.Data{
22+
"test-key": {
23+
Value: &pb.Data_Primitive{
24+
Primitive: spec_util.NewPrimitiveString(`{"custom_postot_conversation_id":"12345678-1234-b1234-1234-123456789012","custom_postot_interaction_id":"12345678-1234-b1234-1234-123456789012","_prompt":"agent_system","_agent":"Root","_workflow":"agent_mode_chat","_domain":"agent_mode","_environment":"beta","_user":"916b323","_nrTraceId":"0d495a1379560f5d14f7f9c43d57bb07","_organization":"1"}`),
25+
},
26+
Meta: &pb.DataMeta{
27+
Meta: &pb.DataMeta_Http{
28+
Http: &pb.HTTPMeta{
29+
Location: &pb.HTTPMeta_Header{
30+
Header: &pb.HTTPHeader{
31+
Key: "x-portkey-metadata",
32+
},
33+
},
34+
},
35+
},
36+
},
37+
},
38+
},
39+
},
40+
},
41+
expected: true,
42+
},
43+
{
44+
name: "Non-Postman team with organization 2",
45+
witness: &pb.Witness{
46+
Method: &pb.Method{
47+
Args: map[string]*pb.Data{
48+
"test-key": {
49+
Value: &pb.Data_Primitive{
50+
Primitive: spec_util.NewPrimitiveString(`{"custom_postot_conversation_id":"12345678-1234-b1234-1234-123456789012","custom_postot_interaction_id":"12345678-1234-b1234-1234-123456789012","_prompt":"agent_system","_agent":"Root","_workflow":"agent_mode_chat","_domain":"agent_mode","_environment":"beta","_user":"916b323","_nrTraceId":"0d495a1379560f5d14f7f9c43d57bb07","_organization":"2"}`),
51+
},
52+
Meta: &pb.DataMeta{
53+
Meta: &pb.DataMeta_Http{
54+
Http: &pb.HTTPMeta{
55+
Location: &pb.HTTPMeta_Header{
56+
Header: &pb.HTTPHeader{
57+
Key: "x-portkey-metadata",
58+
},
59+
},
60+
},
61+
},
62+
},
63+
},
64+
},
65+
},
66+
},
67+
expected: false,
68+
},
69+
{
70+
name: "No x-portkey-metadata header",
71+
witness: &pb.Witness{
72+
Method: &pb.Method{
73+
Args: map[string]*pb.Data{
74+
"test-key": {
75+
Value: &pb.Data_Primitive{
76+
Primitive: spec_util.NewPrimitiveString("some-other-header-value"),
77+
},
78+
Meta: &pb.DataMeta{
79+
Meta: &pb.DataMeta_Http{
80+
Http: &pb.HTTPMeta{
81+
Location: &pb.HTTPMeta_Header{
82+
Header: &pb.HTTPHeader{
83+
Key: "authorization",
84+
},
85+
},
86+
},
87+
},
88+
},
89+
},
90+
},
91+
},
92+
},
93+
expected: false,
94+
},
95+
{
96+
name: "Invalid JSON in x-portkey-metadata header",
97+
witness: &pb.Witness{
98+
Method: &pb.Method{
99+
Args: map[string]*pb.Data{
100+
"test-key": {
101+
Value: &pb.Data_Primitive{
102+
Primitive: spec_util.NewPrimitiveString("invalid-json"),
103+
},
104+
Meta: &pb.DataMeta{
105+
Meta: &pb.DataMeta_Http{
106+
Http: &pb.HTTPMeta{
107+
Location: &pb.HTTPMeta_Header{
108+
Header: &pb.HTTPHeader{
109+
Key: "x-portkey-metadata",
110+
},
111+
},
112+
},
113+
},
114+
},
115+
},
116+
},
117+
},
118+
},
119+
expected: false,
120+
},
121+
{
122+
name: "Empty witness",
123+
witness: &pb.Witness{
124+
Method: &pb.Method{
125+
Args: map[string]*pb.Data{},
126+
},
127+
},
128+
expected: false,
129+
},
130+
}
131+
132+
for _, tt := range tests {
133+
t.Run(tt.name, func(t *testing.T) {
134+
result := isPostmanInternalTeam(tt.witness)
135+
if result != tt.expected {
136+
t.Errorf("isPostmanInternalTeam() = %v, expected %v", result, tt.expected)
137+
}
138+
})
139+
}
140+
}
141+
142+
func TestShouldCapturePayloadWithPostmanTeam(t *testing.T) {
143+
// Test that shouldCapturePayload returns true for Postman internal team
144+
witness := &pb.Witness{
145+
Method: &pb.Method{
146+
Meta: &pb.MethodMeta{
147+
Meta: &pb.MethodMeta_Http{
148+
Http: &pb.HTTPMethodMeta{
149+
Method: "POST",
150+
PathTemplate: "/v1/chat/completions",
151+
Host: "api.portkey.ai",
152+
},
153+
},
154+
},
155+
Args: map[string]*pb.Data{
156+
"test-key": {
157+
Value: &pb.Data_Primitive{
158+
Primitive: spec_util.NewPrimitiveString(`{"custom_postot_conversation_id":"12345678-1234-b1234-1234-123456789012","custom_postot_interaction_id":"12345678-1234-b1234-1234-123456789012","_prompt":"agent_system","_agent":"Root","_workflow":"agent_mode_chat","_domain":"agent_mode","_environment":"beta","_user":"916b323","_nrTraceId":"0d495a1379560f5d14f7f9c43d57bb07","_organization":"1"}`),
159+
},
160+
Meta: &pb.DataMeta{
161+
Meta: &pb.DataMeta_Http{
162+
Http: &pb.HTTPMeta{
163+
Location: &pb.HTTPMeta_Header{
164+
Header: &pb.HTTPHeader{
165+
Key: "x-portkey-metadata",
166+
},
167+
},
168+
},
169+
},
170+
},
171+
},
172+
},
173+
},
174+
}
175+
176+
result := shouldCapturePayload(witness, []*regexp.Regexp{regexp.MustCompile("/v1/chat/completions")})
177+
if !result {
178+
t.Errorf("shouldCapturePayload() = false, expected true for Postman internal team")
179+
}
180+
}

0 commit comments

Comments
 (0)