Propagate ContinuationToken through Response.ToUpdates so it survives a ToUpdates/Collect round-trip

agent/response.go

@@ -119,8 +119,9 @@ func (resp *Response) Coalesce() {
 // ToUpdates converts this response into response updates suitable for streaming
 // scenarios.
 //
-// Each message in the response becomes a separate update. Response-level usage
-// and additional properties are included as an additional update when present.
+// Each message in the response becomes a separate update. Response-level usage,
+// additional properties, and a non-empty continuation token are included as an
+// additional metadata-only update when present.
 func (resp *Response) ToUpdates() []*ResponseUpdate {
 	if resp == nil {
 		return nil
@@ -149,11 +150,12 @@ func (resp *Response) ToUpdates() []*ResponseUpdate {
 		})
 	}
 
-	if hasUsage || hasAdditionalProperties {
+	if hasUsage || hasAdditionalProperties || resp.ContinuationToken != "" {
 		extra := &ResponseUpdate{
 			AdditionalProperties: resp.AdditionalProperties,
 			AgentID:              resp.AgentID,
 			ResponseID:           resp.ID,
+			ContinuationToken:    resp.ContinuationToken,
 			CreatedAt:            resp.CreatedAt,
 		}
 		if hasUsage {
@@ -196,12 +198,18 @@ func (resp *Response) Update(update *ResponseUpdate) {
 		}
 		maps.Copy(msg.AdditionalProperties, update.AdditionalProperties)
 	}
-	if msg.RawRepresentation == nil {
-		msg.RawRepresentation = update.RawRepresentation
-	} else if s, ok := msg.RawRepresentation.([]any); ok {
-		msg.RawRepresentation = append(s, update.RawRepresentation)
-	} else {
-		msg.RawRepresentation = []any{msg.RawRepresentation, update.RawRepresentation}
+	// A nil RawRepresentation carries no provider data, so treat it as a no-op.
+	// This keeps metadata-only updates (e.g. response-level usage or a
+	// continuation token emitted by ToUpdates) from mutating the message's raw
+	// data during a ToUpdates/Collect round-trip.
+	if update.RawRepresentation != nil {
+		if msg.RawRepresentation == nil {
+			msg.RawRepresentation = update.RawRepresentation
+		} else if s, ok := msg.RawRepresentation.([]any); ok {
+			msg.RawRepresentation = append(s, update.RawRepresentation)
+		} else {
+			msg.RawRepresentation = []any{msg.RawRepresentation, update.RawRepresentation}
+		}
 	}
 
 	// Other members on a ResponseUpdate map to members of the response.

agent/response_test.go

@@ -461,6 +461,35 @@ func TestResponse_Update_RawRepresentation(t *testing.T) {
 	}
 }
 
+func TestResponse_ToUpdates_RoundTripPreservesRawRepresentationWithContinuationToken(t *testing.T) {
+	// A response with a message raw representation and a continuation token emits
+	// a trailing metadata-only update (RawRepresentation nil). Collecting the
+	// updates back must not fold that nil into the message's raw data.
+	original := &agent.Response{
+		ContinuationToken: "token-123",
+		Messages: []*message.Message{
+			{
+				ID:                "msg1",
+				Role:              message.RoleAssistant,
+				RawRepresentation: "raw1",
+				Contents:          message.Contents{&message.TextContent{Text: "Hello"}},
+			},
+		},
+	}
+
+	var collected agent.Response
+	for _, update := range original.ToUpdates() {
+		collected.Update(update)
+	}
+
+	if got := collected.Messages[0].RawRepresentation; got != "raw1" {
+		t.Errorf("expected RawRepresentation to round-trip as 'raw1', got %v", got)
+	}
+	if collected.ContinuationToken != "token-123" {
+		t.Errorf("expected ContinuationToken 'token-123', got %q", collected.ContinuationToken)
+	}
+}
+
 func TestResponse_Coalesce_PreservesEmptyMessagesAndWhitespaceText(t *testing.T) {
 	resp := &agent.Response{}
 	resp.Update(&agent.ResponseUpdate{
@@ -788,3 +817,28 @@ func TestResponse_ToUpdates_WithAdditionalPropertiesOnlyProducesSingleUpdate(t *
 		t.Errorf("expected key value, got %v", updates[0].AdditionalProperties["key"])
 	}
 }
+
+func TestResponse_ToUpdates_PropagatesContinuationToken(t *testing.T) {
+	resp := &agent.Response{
+		ContinuationToken: "tok-123",
+		Messages: []*message.Message{
+			{
+				Role:     message.RoleAssistant,
+				Contents: message.Contents{&message.TextContent{Text: "Text"}},
+			},
+		},
+	}
+
+	updates := resp.ToUpdates()
+
+	// The token must survive a ToUpdates/Collect round-trip.
+	var roundTripped agent.Response
+	for _, update := range updates {
+		roundTripped.Update(update)
+	}
+	roundTripped.Coalesce()
+
+	if roundTripped.ContinuationToken != "tok-123" {
+		t.Errorf("expected ContinuationToken tok-123, got %q", roundTripped.ContinuationToken)
+	}
+}