Degrade on unknown annotation/region discriminators instead of aborting the message parse

message/annotation.go

@@ -37,16 +37,40 @@ type annotationKind string
 type Annotations []Annotation
 
 func (as *Annotations) UnmarshalJSON(data []byte) error {
-	var err error
-	*as, err = jsonx.UnmarshalDiscriminatedUnionSlice[Annotation](data, supportedAnnotations)
-	return err
+	out, err := jsonx.UnmarshalDiscriminatedUnionSliceWithFallback(data, supportedAnnotations, unmarshalRawAnnotation)
+	if err != nil {
+		return err
+	}
+	*as = out
+	return nil
+}
+
+func unmarshalRawAnnotation(data json.RawMessage) (Annotation, error) {
+	return &RawAnnotation{RawRepresentation: append(json.RawMessage(nil), data...)}, nil
 }
 
 // Annotation represents an annotation on content.
 type Annotation interface {
 	kind() annotationKind
 }
 
+// RawAnnotation represents a provider-specific annotation that does not fit one
+// of the structured annotation types. The original JSON is preserved in
+// RawRepresentation and round-tripped on marshal so that newer provider-emitted
+// annotation subtypes do not fail the enclosing deserialization.
+type RawAnnotation struct {
+	RawRepresentation json.RawMessage
+}
+
+func (t *RawAnnotation) MarshalJSON() ([]byte, error) {
+	if len(t.RawRepresentation) > 0 {
+		return t.RawRepresentation, nil
+	}
+	return []byte("{}"), nil
+}
+
+func (t *RawAnnotation) kind() annotationKind { return "" }
+
 // CitationAnnotation represents an annotation that links content to source references,
 // such as documents, URLs, files, or tool outputs.
 type CitationAnnotation struct {
@@ -80,9 +104,16 @@ type annotatedRegionKind string
 type AnnotatedRegions []AnnotatedRegion
 
 func (as *AnnotatedRegions) UnmarshalJSON(data []byte) error {
-	var err error
-	*as, err = jsonx.UnmarshalDiscriminatedUnionSlice[AnnotatedRegion](data, supportedAnnotatedRegions)
-	return err
+	out, err := jsonx.UnmarshalDiscriminatedUnionSliceWithFallback(data, supportedAnnotatedRegions, unmarshalRawAnnotatedRegion)
+	if err != nil {
+		return err
+	}
+	*as = out
+	return nil
+}
+
+func unmarshalRawAnnotatedRegion(data json.RawMessage) (AnnotatedRegion, error) {
+	return &RawAnnotatedRegion{RawRepresentation: append(json.RawMessage(nil), data...)}, nil
 }
 
 // AnnotatedRegion describes the portion of an associated [Content]
@@ -91,6 +122,23 @@ type AnnotatedRegion interface {
 	kind() annotatedRegionKind
 }
 
+// RawAnnotatedRegion represents a provider-specific annotated region that does
+// not fit one of the structured region types. The original JSON is preserved in
+// RawRepresentation and round-tripped on marshal so that newer provider-emitted
+// region subtypes do not fail the enclosing deserialization.
+type RawAnnotatedRegion struct {
+	RawRepresentation json.RawMessage
+}
+
+func (t *RawAnnotatedRegion) MarshalJSON() ([]byte, error) {
+	if len(t.RawRepresentation) > 0 {
+		return t.RawRepresentation, nil
+	}
+	return []byte("{}"), nil
+}
+
+func (t *RawAnnotatedRegion) kind() annotatedRegionKind { return "" }
+
 // TextSpanAnnotatedRegion describes a location in the associated [Content]
 // based on starting and ending character indices.
 type TextSpanAnnotatedRegion struct {

message/annotation_test.go

@@ -38,6 +38,68 @@ func TestAnnotationEncoding_Roundtrip(t *testing.T) {
 	}
 }
 
+func TestAnnotationEncoding_UnknownTypesPreservedAsRaw(t *testing.T) {
+	const rawRegion = `{"Type":"futureRegion","Start":1,"End":2}`
+	const rawAnnotation = `{"Type":"futureAnnotation","Value":42}`
+	data := []byte(`[` +
+		`{"Type":"citation","AnnotatedRegions":[` + rawRegion + `]},` +
+		rawAnnotation +
+		`]`)
+
+	var annotations message.Annotations
+	if err := json.Unmarshal(data, &annotations); err != nil {
+		t.Fatal(err)
+	}
+	if len(annotations) != 2 {
+		t.Fatalf("expected 2 annotations, got %d", len(annotations))
+	}
+
+	citation, ok := annotations[0].(*message.CitationAnnotation)
+	if !ok {
+		t.Fatalf("annotations[0] = %T, want *message.CitationAnnotation", annotations[0])
+	}
+	if len(citation.AnnotatedRegions) != 1 {
+		t.Fatalf("expected 1 annotated region, got %d", len(citation.AnnotatedRegions))
+	}
+	rawRegionVal, ok := citation.AnnotatedRegions[0].(*message.RawAnnotatedRegion)
+	if !ok {
+		t.Fatalf("region = %T, want *message.RawAnnotatedRegion", citation.AnnotatedRegions[0])
+	}
+	if got := string(rawRegionVal.RawRepresentation); got != rawRegion {
+		t.Fatalf("region raw = %s, want %s", got, rawRegion)
+	}
+
+	rawAnnotationVal, ok := annotations[1].(*message.RawAnnotation)
+	if !ok {
+		t.Fatalf("annotations[1] = %T, want *message.RawAnnotation", annotations[1])
+	}
+	if got := string(rawAnnotationVal.RawRepresentation); got != rawAnnotation {
+		t.Fatalf("annotation raw = %s, want %s", got, rawAnnotation)
+	}
+
+	// The unknown entries round-trip back to their original JSON.
+	if b, err := json.Marshal(citation.AnnotatedRegions[0]); err != nil || string(b) != rawRegion {
+		t.Fatalf("region remarshal = %s (err %v), want %s", b, err, rawRegion)
+	}
+	if b, err := json.Marshal(annotations[1]); err != nil || string(b) != rawAnnotation {
+		t.Fatalf("annotation remarshal = %s (err %v), want %s", b, err, rawAnnotation)
+	}
+}
+
+func TestAnnotationEncoding_KnownTypeInvalidPayloadReturnsError(t *testing.T) {
+	// A known discriminator ("citation") with a payload that does not match the
+	// structured type must surface a decode error rather than being silently
+	// downgraded to a RawAnnotation.
+	if err := json.Unmarshal([]byte(`[{"Type":"citation","URL":123}]`), new(message.Annotations)); err == nil {
+		t.Fatal("expected error decoding known annotation with invalid payload, got nil")
+	}
+
+	// Likewise for a known annotated region type ("text_span").
+	if err := json.Unmarshal([]byte(`[{"Type":"text_span","Start":"nope"}]`), new(message.AnnotatedRegions)); err == nil {
+		t.Fatal("expected error decoding known region with invalid payload, got nil")
+	}
+}
+
 func TestAnnotatedRegionEncoding_Roundtrip(t *testing.T) {
 	regions := message.AnnotatedRegions{
 		&message.TextSpanAnnotatedRegion{