From b00c7b4b6411d9d52a0a15cc784c6f8eff1fc8e9 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Fri, 24 Jul 2026 08:05:27 +0530 Subject: [PATCH 1/2] Propagate ContinuationToken through Response.ToUpdates ToUpdates copied response-level FinishReason, CreatedAt, AdditionalProperties and Usage onto emitted updates but omitted ContinuationToken. Since Collect clears the token on any empty-token update, a Response carrying a token lost it after a ToUpdates/Collect round-trip. Set it on the trailing update, matching .NET AsChatResponseUpdate. --- agent/response.go | 3 ++- agent/response_test.go | 25 +++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/agent/response.go b/agent/response.go index 6bd8aa4b..bf94923c 100644 --- a/agent/response.go +++ b/agent/response.go @@ -149,11 +149,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 { diff --git a/agent/response_test.go b/agent/response_test.go index 6ed913fd..46a885a1 100644 --- a/agent/response_test.go +++ b/agent/response_test.go @@ -788,3 +788,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) + } +} From 514d8cb0c9a3f6ab53f6da1960736bcc1861841b Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Fri, 24 Jul 2026 10:36:35 +0530 Subject: [PATCH 2/2] Skip nil raw representations in Response.Update and document ToUpdates metadata update --- agent/response.go | 23 +++++++++++++++-------- agent/response_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/agent/response.go b/agent/response.go index bf94923c..071b0f85 100644 --- a/agent/response.go +++ b/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 @@ -197,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. diff --git a/agent/response_test.go b/agent/response_test.go index 46a885a1..bd1f00c2 100644 --- a/agent/response_test.go +++ b/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{