@@ -135,3 +135,151 @@ def test_json_string_not_matching_schema_raises_error(self):
135135
136136 with pytest .raises (ValueError , match = "doesn't match expected schema" ):
137137 parse_json_or_text (data , SampleResponse )
138+
139+
140+ class ProductResponse (BaseModel ):
141+ """Response type without protocol fields for testing protocol field stripping."""
142+
143+ products : list [str ]
144+ total : int = 0
145+
146+
147+ class TestProtocolFieldExtraction :
148+ """Tests for protocol field extraction from A2A responses.
149+
150+ A2A servers may include protocol-level fields (message, context_id, data)
151+ that are not part of task-specific response schemas. These are separated
152+ for task data validation, but preserved at the TaskResult level.
153+
154+ See: https://github.com/adcontextprotocol/adcp-client-python/issues/109
155+ """
156+
157+ def test_response_with_message_field_separated (self ):
158+ """Test that protocol 'message' field is separated before validation."""
159+ # A2A server returns task data with protocol message mixed in
160+ data = {
161+ "message" : "No products matched your requirements." ,
162+ "products" : ["product-1" , "product-2" ],
163+ "total" : 2 ,
164+ }
165+
166+ result = parse_json_or_text (data , ProductResponse )
167+
168+ assert isinstance (result , ProductResponse )
169+ assert result .products == ["product-1" , "product-2" ]
170+ assert result .total == 2
171+
172+ def test_response_with_context_id_separated (self ):
173+ """Test that protocol 'context_id' field is separated before validation."""
174+ data = {
175+ "context_id" : "session-123" ,
176+ "products" : ["product-1" ],
177+ "total" : 1 ,
178+ }
179+
180+ result = parse_json_or_text (data , ProductResponse )
181+
182+ assert isinstance (result , ProductResponse )
183+ assert result .products == ["product-1" ]
184+
185+ def test_response_with_multiple_protocol_fields_separated (self ):
186+ """Test that multiple protocol fields are separated."""
187+ data = {
188+ "message" : "Found products" ,
189+ "context_id" : "session-456" ,
190+ "products" : ["a" , "b" , "c" ],
191+ "total" : 3 ,
192+ }
193+
194+ result = parse_json_or_text (data , ProductResponse )
195+
196+ assert isinstance (result , ProductResponse )
197+ assert result .products == ["a" , "b" , "c" ]
198+ assert result .total == 3
199+
200+ def test_response_with_data_wrapper_extracted (self ):
201+ """Test that ProtocolResponse 'data' wrapper is extracted."""
202+ # Full ProtocolResponse format: {"message": "...", "data": {...task_data...}}
203+ data = {
204+ "message" : "Operation completed" ,
205+ "context_id" : "ctx-789" ,
206+ "data" : {
207+ "products" : ["wrapped-product" ],
208+ "total" : 1 ,
209+ },
210+ }
211+
212+ result = parse_json_or_text (data , ProductResponse )
213+
214+ assert isinstance (result , ProductResponse )
215+ assert result .products == ["wrapped-product" ]
216+ assert result .total == 1
217+
218+ def test_response_with_payload_wrapper_extracted (self ):
219+ """Test that ProtocolEnvelope 'payload' wrapper is extracted."""
220+ # Full ProtocolEnvelope format
221+ data = {
222+ "message" : "Operation completed" ,
223+ "status" : "completed" ,
224+ "task_id" : "task-123" ,
225+ "timestamp" : "2025-01-01T00:00:00Z" ,
226+ "payload" : {
227+ "products" : ["envelope-product" ],
228+ "total" : 1 ,
229+ },
230+ }
231+
232+ result = parse_json_or_text (data , ProductResponse )
233+
234+ assert isinstance (result , ProductResponse )
235+ assert result .products == ["envelope-product" ]
236+ assert result .total == 1
237+
238+ def test_exact_match_still_works (self ):
239+ """Test that responses exactly matching schema still work."""
240+ data = {
241+ "products" : ["exact-match" ],
242+ "total" : 1 ,
243+ }
244+
245+ result = parse_json_or_text (data , ProductResponse )
246+
247+ assert result .products == ["exact-match" ]
248+ assert result .total == 1
249+
250+ def test_json_string_with_protocol_fields (self ):
251+ """Test JSON string with protocol fields is parsed correctly."""
252+ data = json .dumps (
253+ {
254+ "message" : "Success" ,
255+ "products" : ["from-json-string" ],
256+ "total" : 1 ,
257+ }
258+ )
259+
260+ result = parse_json_or_text (data , ProductResponse )
261+
262+ assert result .products == ["from-json-string" ]
263+
264+ def test_invalid_data_after_separation_raises_error (self ):
265+ """Test that invalid data still raises error after separation."""
266+ data = {
267+ "message" : "Some message" ,
268+ "wrong_field" : "value" ,
269+ }
270+
271+ with pytest .raises (ValueError , match = "doesn't match expected schema" ):
272+ parse_json_or_text (data , ProductResponse )
273+
274+ def test_model_with_message_field_validates_directly (self ):
275+ """Test that models containing 'message' field validate without separation."""
276+ # SampleResponse has a 'message' field, so it should validate directly
277+ data = {
278+ "message" : "Hello" ,
279+ "count" : 42 ,
280+ }
281+
282+ result = parse_json_or_text (data , SampleResponse )
283+
284+ assert result .message == "Hello"
285+ assert result .count == 42
0 commit comments