diff --git a/provider/a2aprovider/a2a.go b/provider/a2aprovider/a2a.go index 110be087..4799f5dd 100644 --- a/provider/a2aprovider/a2a.go +++ b/provider/a2aprovider/a2a.go @@ -285,7 +285,7 @@ func updateSessionContextID(session *agent.Session, contextID, taskID string, ta // Surface cases where the A2A agent responds with a response that // has a different context ID than the session's context ID. currentContextID := getContextID(session) - if currentContextID != "" && currentContextID != contextID { + if currentContextID != "" && contextID != "" && currentContextID != contextID { return fmt.Errorf("mismatched context ID: session has %q but A2A response has %q", currentContextID, contextID) } setContextID(session, contextID) diff --git a/provider/a2aprovider/a2a_test.go b/provider/a2aprovider/a2a_test.go index 0b47df58..046eedc2 100644 --- a/provider/a2aprovider/a2a_test.go +++ b/provider/a2aprovider/a2a_test.go @@ -30,6 +30,10 @@ type mockA2ATransport struct { sendStreamingMessageCalled bool subscribeToTaskCalled bool getTaskCalled bool + // rawStreamingResponse yields streamingResponseToReturn verbatim, without + // backfilling an empty ContextID from the request. It lets tests exercise + // bare streamed messages that carry no context ID. + rawStreamingResponse bool } func (m *mockA2ATransport) SendMessage(ctx context.Context, _ a2aclient.ServiceParams, params *a2a.SendMessageRequest) (a2a.SendMessageResult, error) { @@ -50,6 +54,22 @@ func (m *mockA2ATransport) SendStreamingMessage(ctx context.Context, _ a2aclient m.sendStreamingMessageCalled = true m.capturedMessageSendParams = params responseToYield := m.streamingResponseToReturn + if m.rawStreamingResponse { + if responseToYield == nil { + // Default to a non-nil event so a raw test that forgets to set + // streamingResponseToReturn doesn't yield a nil event (which would + // panic when production code calls TaskInfo()). Deliberately leave + // ContextID empty — the "raw" mode exists to exercise bare messages + // that carry no context ID. + responseToYield = &a2a.Message{ + ID: "default-stream-id", + Role: a2a.MessageRoleAgent, + } + } + return func(yield func(a2a.Event, error) bool) { + yield(responseToYield, nil) + } + } if responseToYield == nil { // Return default empty message with context ID from request responseToYield = &a2a.Message{ @@ -547,6 +567,69 @@ func TestRunStreamingWithSessionHavingDifferentContextID(t *testing.T) { } } +// TestRunStreamingWithEmptyContextIDKeepsSessionContext verifies that a bare +// streamed message carrying an empty context ID neither errors the run nor +// clobbers the context ID already stored in the session. This mirrors the .NET +// behavior where ContextId is only assigned when currently unset +// (ContextId ??= contextId). +func TestRunStreamingWithEmptyContextIDKeepsSessionContext(t *testing.T) { + transport := &mockA2ATransport{ + rawStreamingResponse: true, + streamingResponseToReturn: &a2a.Message{ + ID: "stream-1", + Role: a2a.MessageRoleAgent, + Parts: a2a.ContentParts{a2a.NewTextPart("Response")}, + // No ContextID: a bare streamed message chunk. + }, + } + a := newTestAgent(transport, agent.Config{}) + + session, err := a.CreateSession(t.Context(), agent.WithServiceID("ctx-1")) + if err != nil { + t.Fatal(err) + } + + for _, err := range a.RunText(t.Context(), "Test streaming", agent.WithSession(session), agent.Stream(true)) { + if err != nil { + t.Fatalf("error = %v, want nil", err) + } + } + + if got := session.ServiceID(); got != "ctx-1" { + t.Errorf("session.ServiceID = %q, want %q", got, "ctx-1") + } +} + +// TestRunStreamingWithEmptyInitialContextStoresResponseContext verifies that when +// the session has no context ID yet, the first streamed event's context ID is stored. +func TestRunStreamingWithEmptyInitialContextStoresResponseContext(t *testing.T) { + transport := &mockA2ATransport{ + rawStreamingResponse: true, + streamingResponseToReturn: &a2a.Message{ + ID: "stream-1", + Role: a2a.MessageRoleAgent, + Parts: a2a.ContentParts{a2a.NewTextPart("Response")}, + ContextID: "ctx-1", + }, + } + a := newTestAgent(transport, agent.Config{}) + + session, err := a.CreateSession(t.Context()) + if err != nil { + t.Fatal(err) + } + + for _, err := range a.RunText(t.Context(), "Test streaming", agent.WithSession(session), agent.Stream(true)) { + if err != nil { + t.Fatalf("error = %v, want nil", err) + } + } + + if got := session.ServiceID(); got != "ctx-1" { + t.Errorf("session.ServiceID = %q, want %q", got, "ctx-1") + } +} + // TestRunStreamingAllowsNonUserRoleMessages tests that streaming allows non-user messages func TestRunStreamingAllowsNonUserRoleMessages(t *testing.T) { transport := &mockA2ATransport{ diff --git a/provider/a2aprovider/session.go b/provider/a2aprovider/session.go index 2096a273..a7ca05db 100644 --- a/provider/a2aprovider/session.go +++ b/provider/a2aprovider/session.go @@ -15,6 +15,9 @@ const ( ) func setContextID(session *agent.Session, contextID string) { + if contextID == "" { + return + } session.SetServiceID(contextID) }