From da5d91d39cf0ed358cb9c0fe2de97ac6ba65e84a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 6 Aug 2026 13:18:52 -0400 Subject: [PATCH] Keep non-text response segments when a stream finalizes finalizeStreamedTranscript replaced the trailing response entry's whole segments array with a single text segment, so every segment ahead of the streamed text was dropped the moment the stream ended. appendStreamingResponse appends a new text segment whenever the trailing segment is not text, so a response that already carried a structure or image segment reaches finalization with two segments and loses the first one. Mirror appendStreamingResponse instead: replace the trailing text segment in place, carrying its ID over, or append a new one when the trailing segment is not text, and leave everything earlier untouched. Neither streaming mutator had test coverage, so cover both. Co-Authored-By: Claude Opus 5 (1M context) --- Sources/AnyLanguageModel/Transcript.swift | 36 ++++---- .../TranscriptTests.swift | 88 +++++++++++++++++++ 2 files changed, 104 insertions(+), 20 deletions(-) diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index f256d046..3d740441 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -60,15 +60,18 @@ public struct Transcript: Sendable, Equatable, Codable { replace(index: entries.count - 1, with: .response(response)) } - /// Replaces the trailing response entry's text with the final text, or appends a new response entry if the last entry isn't a response. + /// Replaces the trailing response entry's streamed text with the final text, or appends a new response entry if the last entry isn't a response. /// Prevents streamed responses from having duplicate entries on completion. /// + /// Segments ahead of the trailing text are left alone, so anything else the model + /// produced during the same response survives the end of the stream. + /// /// - Parameters: /// - text: The text to replace the final response with. /// - assetIDs: The assetIDs for the response. mutating func finalizeStreamedTranscript(_ text: String, assetIDs: [String]) { // Make sure the last entry in the transcript is a response. If it is not, create a new response and append it to the end of the transcript. - guard case .response(let response) = entries.last else { + guard case .response(var response) = entries.last else { append( Entry.response( Response( @@ -82,26 +85,19 @@ public struct Transcript: Sendable, Equatable, Codable { return } - // If the last segment is text we want to carry its ID over to the new text segment. Otherwise generate a new ID for it. - let id = - switch response.segments.last { - case .text(let last): - last.id - default: - UUID().uuidString - } - - let newResponse: Entry = Entry.response( - Response( - id: response.id, - assetIDs: assetIDs, - segments: [ - Transcript.Segment.text(Transcript.TextSegment(id: id, content: text)) - ] + // Streamed text only ever lands in the trailing text segment, so replace that one + // in place — carrying its ID over — and leave every earlier segment untouched. + if case .text(let last)? = response.segments.last { + response.segments[response.segments.count - 1] = Transcript.Segment.text( + Transcript.TextSegment(id: last.id, content: text) ) - ) + } else { + response.segments.append(Transcript.Segment.text(Transcript.TextSegment(content: text))) + } + response.assetIDs = assetIDs - replace(index: entries.count - 1, with: newResponse) + // Replace the latest entry with the one we just updated. + replace(index: entries.count - 1, with: .response(response)) } /// An entry in a transcript. diff --git a/Tests/AnyLanguageModelTests/TranscriptTests.swift b/Tests/AnyLanguageModelTests/TranscriptTests.swift index 037c70a0..325b8f44 100644 --- a/Tests/AnyLanguageModelTests/TranscriptTests.swift +++ b/Tests/AnyLanguageModelTests/TranscriptTests.swift @@ -108,6 +108,94 @@ struct TranscriptTests { } } + @Test func finalizeStreamedTranscriptKeepsSegmentsAheadOfTheStreamedText() throws { + let structured = Transcript.StructuredSegment( + id: "structured-id", + source: "source", + content: try GeneratedContent(json: #"{"ok":true}"#) + ) + var transcript = Transcript(entries: [ + .response( + .init( + id: "response-id", + assetIDs: [], + segments: [ + .structure(structured), + .text(.init(id: "text-id", content: "partial")), + ] + ) + ) + ]) + + transcript.finalizeStreamedTranscript("complete", assetIDs: ["asset"]) + + guard case .response(let response)? = transcript.last else { + Issue.record("Expected a trailing response entry") + return + } + #expect(response.id == "response-id") + #expect(response.assetIDs == ["asset"]) + #expect(response.segments.count == 2) + #expect(response.segments.first == .structure(structured)) + // The trailing text segment is replaced in place, keeping its ID stable. + #expect(response.segments.last == .text(.init(id: "text-id", content: "complete"))) + } + + @Test func finalizeStreamedTranscriptAppendsTextWhenTrailingSegmentIsNotText() throws { + let structured = Transcript.StructuredSegment( + id: "structured-id", + source: "source", + content: try GeneratedContent(json: #"{"ok":true}"#) + ) + var transcript = Transcript(entries: [ + .response(.init(id: "response-id", assetIDs: [], segments: [.structure(structured)])) + ]) + + transcript.finalizeStreamedTranscript("complete", assetIDs: []) + + guard case .response(let response)? = transcript.last else { + Issue.record("Expected a trailing response entry") + return + } + #expect(response.segments.count == 2) + #expect(response.segments.first == .structure(structured)) + #expect(response.segments.last?.description == "complete") + } + + @Test func finalizeStreamedTranscriptAppendsResponseWhenLastEntryIsNotAResponse() { + var transcript = Transcript(entries: [ + .prompt(.init(id: "prompt-id", segments: [.text(.init(content: "Hello"))])) + ]) + + transcript.finalizeStreamedTranscript("complete", assetIDs: ["asset"]) + + #expect(transcript.count == 2) + guard case .response(let response)? = transcript.last else { + Issue.record("Expected a trailing response entry") + return + } + #expect(response.assetIDs == ["asset"]) + #expect(response.segments.count == 1) + #expect(response.segments.first?.description == "complete") + } + + @Test func appendStreamingResponseGrowsTheTrailingTextSegmentInPlace() { + var transcript = Transcript(entries: [ + .prompt(.init(id: "prompt-id", segments: [.text(.init(content: "Hello"))])) + ]) + + transcript.appendStreamingResponse("He") + transcript.appendStreamingResponse("Hello") + + #expect(transcript.count == 2) + guard case .response(let response)? = transcript.last else { + Issue.record("Expected a trailing response entry") + return + } + #expect(response.segments.count == 1) + #expect(response.segments.first?.description == "Hello") + } + @Test func responseFormatNameExtractsRefTypeNameOrFallsBack() { let refFormat = Transcript.ResponseFormat(type: Person.self) #expect(refFormat.name.contains("Person"))