From b40ab7d5858e1522eaf48d538ed88b78886a0552 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Fri, 18 Sep 2026 17:51:59 +0530 Subject: [PATCH] Replay citation annotations on the Responses input path Building a Responses assistant turn from client-side history always sent an empty output_text annotations array, dropping any citations the text arrived with (a standing TODO). populateAnnotations maps inbound Responses citations onto CitationAnnotation, but the reverse was never implemented, so a prior assistant turn round-tripped through history forwarding with its url/file/ container citations stripped. Add annotationsToOutputText, the inverse of populateAnnotations, inferring each output-text annotation variant from the populated citation fields (as the Python client does). Regions without both span bounds are skipped, and a citation with multiple regions fans out to one annotation per region. --- provider/openaiprovider/responses.go | 75 ++++++++++++++++++++++- provider/openaiprovider/responses_test.go | 66 ++++++++++++++++++++ 2 files changed, 139 insertions(+), 2 deletions(-) diff --git a/provider/openaiprovider/responses.go b/provider/openaiprovider/responses.go index cb8859ed..e26a060a 100644 --- a/provider/openaiprovider/responses.go +++ b/provider/openaiprovider/responses.go @@ -618,8 +618,7 @@ func responsesBuildMessageParam(msg *message.Message, resp responses.ResponseInp case *message.TextContent: outputContents = append(outputContents, responses.ResponseOutputMessageContentUnionParam{ OfOutputText: &responses.ResponseOutputTextParam{ - // TODO: Convert message annotations back to Responses output-text annotations. - Annotations: []responses.ResponseOutputTextAnnotationUnionParam{}, + Annotations: annotationsToOutputText(c.Annotations), Text: c.Text, }, }) @@ -1901,6 +1900,78 @@ func populateAnnotations(anns []responses.ResponseOutputTextAnnotationUnion, con } } +// annotationsToOutputText is the inverse of populateAnnotations: it converts +// framework citation annotations back into Responses output-text annotation +// params so that an assistant turn replayed from client-side history carries +// the citations it originally arrived with. Which variant to emit is inferred +// from the populated fields, mirroring how populateAnnotations maps each one. +// Annotations that are not citations are skipped. +func annotationsToOutputText(anns []message.Annotation) []responses.ResponseOutputTextAnnotationUnionParam { + out := make([]responses.ResponseOutputTextAnnotationUnionParam, 0, len(anns)) + for _, ann := range anns { + cit, ok := ann.(*message.CitationAnnotation) + if !ok { + continue + } + containerID, _ := cit.AdditionalProperties["ContainerId"].(string) + switch { + case containerID != "" && cit.FileID != "": + for _, span := range citationSpans(cit.AnnotatedRegions) { + out = append(out, responses.ResponseOutputTextAnnotationUnionParam{ + OfContainerFileCitation: &responses.ResponseOutputTextAnnotationContainerFileCitationParam{ + ContainerID: containerID, + FileID: cit.FileID, + Filename: cit.Title, + StartIndex: span[0], + EndIndex: span[1], + }, + }) + } + case cit.URL != "": + for _, span := range citationSpans(cit.AnnotatedRegions) { + out = append(out, responses.ResponseOutputTextAnnotationUnionParam{ + OfURLCitation: &responses.ResponseOutputTextAnnotationURLCitationParam{ + URL: cit.URL, + Title: cit.Title, + StartIndex: span[0], + EndIndex: span[1], + }, + }) + } + case cit.FileID != "" && cit.Title != "": + out = append(out, responses.ResponseOutputTextAnnotationUnionParam{ + OfFileCitation: &responses.ResponseOutputTextAnnotationFileCitationParam{ + FileID: cit.FileID, + Filename: cit.Title, + }, + }) + case cit.FileID != "": + out = append(out, responses.ResponseOutputTextAnnotationUnionParam{ + OfFilePath: &responses.ResponseOutputTextAnnotationFilePathParam{ + FileID: cit.FileID, + }, + }) + } + } + return out +} + +// citationSpans extracts the [start, end) index pairs from a citation's +// annotated regions, skipping regions without both bounds set. Each Responses +// annotation carries a single span, so a citation with multiple regions is +// fanned out into one annotation per region. +func citationSpans(regions message.AnnotatedRegions) [][2]int64 { + var spans [][2]int64 + for _, region := range regions { + span, ok := region.(*message.TextSpanAnnotatedRegion) + if !ok || span.StartIndex == nil || span.EndIndex == nil { + continue + } + spans = append(spans, [2]int64{int64(*span.StartIndex), int64(*span.EndIndex)}) + } + return spans +} + func textSpanAnnotatedRegion(start, end int64) *message.TextSpanAnnotatedRegion { startIndex, endIndex := int(start), int(end) return &message.TextSpanAnnotatedRegion{StartIndex: &startIndex, EndIndex: &endIndex} diff --git a/provider/openaiprovider/responses_test.go b/provider/openaiprovider/responses_test.go index 5b1d7a2d..ada9dc77 100644 --- a/provider/openaiprovider/responses_test.go +++ b/provider/openaiprovider/responses_test.go @@ -1036,6 +1036,72 @@ func TestResponsesAssistantReplayPreservesHostedToolItems(t *testing.T) { } } +// A replayed assistant turn from client-side history must carry back the +// citations its text originally arrived with, so the Responses API sees the +// same annotations on round-trip. This is the inverse of populateAnnotations. +func TestResponsesAssistantReplayRoundTripsCitationAnnotations(t *testing.T) { + start, end := 0, 5 + contents := message.Contents{ + &message.TextContent{ + Text: "hello world", + ContentHeader: message.ContentHeader{ + Annotations: []message.Annotation{ + &message.CitationAnnotation{ + Title: "Example", + URL: "https://example.com", + AnnotatedRegions: message.AnnotatedRegions{&message.TextSpanAnnotatedRegion{StartIndex: &start, EndIndex: &end}}, + }, + &message.CitationAnnotation{ + FileID: "file_123", + Title: "doc.pdf", + }, + }, + }, + }, + } + + var captured map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&captured); err != nil { + t.Fatal(err) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"resp","object":"response","created_at":1,"status":"completed","model":"gpt-4o-mini","output":[]}`) + })) + defer server.Close() + + a := newTestResponsesClient(server, "gpt-4o-mini") + if _, err := a.Run(t.Context(), []*message.Message{{Role: message.RoleAssistant, Contents: contents}}).Collect(); err != nil { + t.Fatal(err) + } + + input, ok := captured["input"].([]any) + if !ok || len(input) != 1 { + t.Fatalf("input = %#v", captured["input"]) + } + item, _ := input[0].(map[string]any) + textContents, _ := item["content"].([]any) + if len(textContents) != 1 { + t.Fatalf("assistant content = %#v, want one output_text", item["content"]) + } + outputText, _ := textContents[0].(map[string]any) + anns, _ := outputText["annotations"].([]any) + if len(anns) != 2 { + t.Fatalf("annotations = %#v, want url_citation and file_citation", outputText["annotations"]) + } + url, _ := anns[0].(map[string]any) + if url["type"] != "url_citation" || url["url"] != "https://example.com" || url["title"] != "Example" { + t.Errorf("url citation = %#v", url) + } + if url["start_index"] != float64(0) || url["end_index"] != float64(5) { + t.Errorf("url citation span = [%v, %v), want [0, 5)", url["start_index"], url["end_index"]) + } + file, _ := anns[1].(map[string]any) + if file["type"] != "file_citation" || file["file_id"] != "file_123" || file["filename"] != "doc.pdf" { + t.Errorf("file citation = %#v", file) + } +} + func TestResponsesAssistantReplayReconstructsPersistedMCPContents(t *testing.T) { original := message.Contents{ &message.ToolApprovalRequestContent{