From c0c5bcfd7e4baf9a0d3d1a5e2d2d01703af5a39a Mon Sep 17 00:00:00 2001 From: buravc Date: Fri, 14 Aug 2026 17:14:33 +0200 Subject: [PATCH 1/4] fix(golang-adk): resolve user identity from A2A call context Signed-off-by: buravc --- go/adk/pkg/a2a/executor.go | 14 ++++++++++++-- go/adk/pkg/taskstore/store.go | 13 ++++++++++--- go/adk/pkg/taskstore/store_test.go | 18 ++++++++++++++++++ 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 7b8fb95e6..3ecd08962 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -119,6 +119,16 @@ func (u *userIDInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont return ctx, nil, nil } +// CallerUserID returns the authenticated user name from the a2asrv +// CallContext attached to ctx, or "" if none is set. +func CallerUserID(ctx context.Context) string { + callCtx, ok := a2asrv.CallContextFrom(ctx) + if !ok || callCtx.User == nil { + return "" + } + return callCtx.User.Name +} + // Execute applies kagent-specific request setup and delegates event generation // to the upstream ADK executor, which streams output as artifact updates. func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2atype.Event, error] { @@ -129,8 +139,8 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon } userID := "A2A_USER_" + reqCtx.ContextID - if callCtx, ok := a2asrv.CallContextFrom(ctx); ok && callCtx.User != nil && callCtx.User.Name != "" { - userID = callCtx.User.Name + if id := CallerUserID(ctx); id != "" { + userID = id } sessionID := reqCtx.ContextID diff --git a/go/adk/pkg/taskstore/store.go b/go/adk/pkg/taskstore/store.go index 4f090a50a..013e45017 100644 --- a/go/adk/pkg/taskstore/store.go +++ b/go/adk/pkg/taskstore/store.go @@ -8,6 +8,7 @@ import ( a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" a2ataskstore "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" + "github.com/kagent-dev/kagent/go/adk/pkg/a2a" "github.com/kagent-dev/kagent/go/adk/pkg/controllerclient" apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1" "google.golang.org/grpc/codes" @@ -30,6 +31,12 @@ func NewKAgentTaskStore(client *controllerclient.Client) *KAgentTaskStore { return &KAgentTaskStore{client: client} } +// userID resolves the caller for a TaskStore call; auth.WithUserID's value +// doesn't reach this goroutine. +func userID(ctx context.Context) string { + return a2a.CallerUserID(ctx) +} + func (s *KAgentTaskStore) saveTask(ctx context.Context, task *a2atype.Task) (a2ataskstore.TaskVersion, error) { if task == nil { return a2ataskstore.TaskVersionMissing, fmt.Errorf("task cannot be nil") @@ -52,7 +59,7 @@ func (s *KAgentTaskStore) saveTask(ctx context.Context, task *a2atype.Task) (a2a if err != nil { return a2ataskstore.TaskVersionMissing, fmt.Errorf("encode task: %w", err) } - callContext, cancel := s.client.CallContext(ctx, "") + callContext, cancel := s.client.CallContext(ctx, userID(ctx)) defer cancel() _, err = s.client.TaskService().UpsertTask(callContext, &apiv1alpha1.UpsertTaskRequest{Task: encoded}) if err != nil { @@ -77,7 +84,7 @@ func (s *KAgentTaskStore) Update(ctx context.Context, update *a2ataskstore.Updat // Get implements taskstore.Store. func (s *KAgentTaskStore) Get(ctx context.Context, taskID a2atype.TaskID) (*a2ataskstore.StoredTask, error) { - callContext, cancel := s.client.CallContext(ctx, "") + callContext, cancel := s.client.CallContext(ctx, userID(ctx)) defer cancel() response, err := s.client.TaskService().GetTask(callContext, &apiv1alpha1.GetTaskRequest{TaskId: string(taskID)}) if err != nil { @@ -118,7 +125,7 @@ func (s *KAgentTaskStore) List(ctx context.Context, req *a2atype.ListTasksReques return &a2atype.ListTasksResponse{Tasks: []*a2atype.Task{}, PageSize: pageSize}, nil } - callContext, cancel := s.client.CallContext(ctx, "") + callContext, cancel := s.client.CallContext(ctx, userID(ctx)) defer cancel() response, err := s.client.TaskService().ListTasks(callContext, &apiv1alpha1.ListTasksRequest{SessionId: req.ContextID}) if err != nil { diff --git a/go/adk/pkg/taskstore/store_test.go b/go/adk/pkg/taskstore/store_test.go index 77795221c..873306878 100644 --- a/go/adk/pkg/taskstore/store_test.go +++ b/go/adk/pkg/taskstore/store_test.go @@ -8,6 +8,7 @@ import ( a2a "github.com/a2aproject/a2a-go/v2/a2a" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" + "github.com/a2aproject/a2a-go/v2/a2asrv" a2ataskstore "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" "github.com/kagent-dev/kagent/go/adk/pkg/auth" "github.com/kagent-dev/kagent/go/adk/pkg/controllerclient" @@ -124,6 +125,23 @@ func TestGetDecodesCanonicalTask(t *testing.T) { assert.Equal(t, "done", stored.Task.History[0].Parts[0].Text()) } +func TestGetPrefersCallContextUserOverContextValue(t *testing.T) { + encoded, err := pbconv.ToProtoTask(&a2a.Task{ID: a2a.TaskID("task-4"), ContextID: "session-4"}) + require.NoError(t, err) + service := newTaskStore(t, &taskTestServer{get: func(ctx context.Context, _ *apiv1alpha1.GetTaskRequest) (*apiv1alpha1.GetTaskResponse, error) { + values, _ := metadata.FromIncomingContext(ctx) + assert.Equal(t, []string{"call-context-user"}, values.Get("x-user-id")) + return &apiv1alpha1.GetTaskResponse{Task: encoded}, nil + }}) + + baseCtx, callCtx := a2asrv.NewCallContext(t.Context(), nil) + callCtx.User = a2asrv.NewAuthenticatedUser("call-context-user", nil) + ctx := auth.WithUserID(baseCtx, "context-value-user") + + _, err = service.Get(ctx, a2a.TaskID("task-4")) + require.NoError(t, err) +} + func TestGetMapsNotFound(t *testing.T) { service := newTaskStore(t, &taskTestServer{get: func(context.Context, *apiv1alpha1.GetTaskRequest) (*apiv1alpha1.GetTaskResponse, error) { return nil, status.Error(codes.NotFound, "missing") From b0e6d83bce360d8d196cd36c797a315a71bc2979 Mon Sep 17 00:00:00 2001 From: buravc Date: Sat, 15 Aug 2026 17:02:37 +0200 Subject: [PATCH 2/4] clarify unit test Signed-off-by: buravc --- go/adk/pkg/taskstore/store_test.go | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/go/adk/pkg/taskstore/store_test.go b/go/adk/pkg/taskstore/store_test.go index 873306878..bfbc35597 100644 --- a/go/adk/pkg/taskstore/store_test.go +++ b/go/adk/pkg/taskstore/store_test.go @@ -125,21 +125,23 @@ func TestGetDecodesCanonicalTask(t *testing.T) { assert.Equal(t, "done", stored.Task.History[0].Parts[0].Text()) } -func TestGetPrefersCallContextUserOverContextValue(t *testing.T) { +func TestGetResolvesUserFromA2ACallContext(t *testing.T) { encoded, err := pbconv.ToProtoTask(&a2a.Task{ID: a2a.TaskID("task-4"), ContextID: "session-4"}) require.NoError(t, err) + + var gotUserID []string service := newTaskStore(t, &taskTestServer{get: func(ctx context.Context, _ *apiv1alpha1.GetTaskRequest) (*apiv1alpha1.GetTaskResponse, error) { values, _ := metadata.FromIncomingContext(ctx) - assert.Equal(t, []string{"call-context-user"}, values.Get("x-user-id")) + gotUserID = values.Get("x-user-id") return &apiv1alpha1.GetTaskResponse{Task: encoded}, nil }}) - baseCtx, callCtx := a2asrv.NewCallContext(t.Context(), nil) + ctx, callCtx := a2asrv.NewCallContext(t.Context(), nil) callCtx.User = a2asrv.NewAuthenticatedUser("call-context-user", nil) - ctx := auth.WithUserID(baseCtx, "context-value-user") _, err = service.Get(ctx, a2a.TaskID("task-4")) require.NoError(t, err) + assert.Equal(t, []string{"call-context-user"}, gotUserID) } func TestGetMapsNotFound(t *testing.T) { From cc98ac64fb87c2e6c9f162ba9a422d7bbb0e92d9 Mon Sep 17 00:00:00 2001 From: buravc Date: Thu, 20 Aug 2026 10:58:43 +0200 Subject: [PATCH 3/4] fix(golang-adk): resolve user id in the pre-fork interceptor Move auth.WithUserID from KAgentExecutor.Execute into userIDCallInterceptor.Before, which runs before a2a-go forks task execution and persistence into separate goroutines. Both goroutines now inherit the user id from the same ctx. taskstore no longer needs to read a2asrv.CallContext directly; controllerclient.CallContext already falls back to auth.UserIDFromContext when passed an empty user id. Signed-off-by: buravc --- go/adk/pkg/a2a/executor.go | 14 ++------------ go/adk/pkg/taskstore/store.go | 13 +++---------- 2 files changed, 5 insertions(+), 22 deletions(-) diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 12149949f..405999756 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -101,17 +101,7 @@ func (u *userIDInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont } // Set the authenticated user so downstream code picks up the real identity. callCtx.User = a2asrv.NewAuthenticatedUser(vals[0], nil) - return ctx, nil, nil -} - -// CallerUserID returns the authenticated user name from the a2asrv -// CallContext attached to ctx, or "" if none is set. -func CallerUserID(ctx context.Context) string { - callCtx, ok := a2asrv.CallContextFrom(ctx) - if !ok || callCtx.User == nil { - return "" - } - return callCtx.User.Name + return auth.WithUserID(ctx, vals[0]), nil, nil } // Execute applies kagent-specific request setup and delegates event generation @@ -124,7 +114,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon } userID := "A2A_USER_" + reqCtx.ContextID - if id := CallerUserID(ctx); id != "" { + if id := auth.UserIDFromContext(ctx); id != "" { userID = id } sessionID := reqCtx.ContextID diff --git a/go/adk/pkg/taskstore/store.go b/go/adk/pkg/taskstore/store.go index 013e45017..4f090a50a 100644 --- a/go/adk/pkg/taskstore/store.go +++ b/go/adk/pkg/taskstore/store.go @@ -8,7 +8,6 @@ import ( a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" a2ataskstore "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" - "github.com/kagent-dev/kagent/go/adk/pkg/a2a" "github.com/kagent-dev/kagent/go/adk/pkg/controllerclient" apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1" "google.golang.org/grpc/codes" @@ -31,12 +30,6 @@ func NewKAgentTaskStore(client *controllerclient.Client) *KAgentTaskStore { return &KAgentTaskStore{client: client} } -// userID resolves the caller for a TaskStore call; auth.WithUserID's value -// doesn't reach this goroutine. -func userID(ctx context.Context) string { - return a2a.CallerUserID(ctx) -} - func (s *KAgentTaskStore) saveTask(ctx context.Context, task *a2atype.Task) (a2ataskstore.TaskVersion, error) { if task == nil { return a2ataskstore.TaskVersionMissing, fmt.Errorf("task cannot be nil") @@ -59,7 +52,7 @@ func (s *KAgentTaskStore) saveTask(ctx context.Context, task *a2atype.Task) (a2a if err != nil { return a2ataskstore.TaskVersionMissing, fmt.Errorf("encode task: %w", err) } - callContext, cancel := s.client.CallContext(ctx, userID(ctx)) + callContext, cancel := s.client.CallContext(ctx, "") defer cancel() _, err = s.client.TaskService().UpsertTask(callContext, &apiv1alpha1.UpsertTaskRequest{Task: encoded}) if err != nil { @@ -84,7 +77,7 @@ func (s *KAgentTaskStore) Update(ctx context.Context, update *a2ataskstore.Updat // Get implements taskstore.Store. func (s *KAgentTaskStore) Get(ctx context.Context, taskID a2atype.TaskID) (*a2ataskstore.StoredTask, error) { - callContext, cancel := s.client.CallContext(ctx, userID(ctx)) + callContext, cancel := s.client.CallContext(ctx, "") defer cancel() response, err := s.client.TaskService().GetTask(callContext, &apiv1alpha1.GetTaskRequest{TaskId: string(taskID)}) if err != nil { @@ -125,7 +118,7 @@ func (s *KAgentTaskStore) List(ctx context.Context, req *a2atype.ListTasksReques return &a2atype.ListTasksResponse{Tasks: []*a2atype.Task{}, PageSize: pageSize}, nil } - callContext, cancel := s.client.CallContext(ctx, userID(ctx)) + callContext, cancel := s.client.CallContext(ctx, "") defer cancel() response, err := s.client.TaskService().ListTasks(callContext, &apiv1alpha1.ListTasksRequest{SessionId: req.ContextID}) if err != nil { From 261ac6be358c5ece0d10b8c3031bfe214366a820 Mon Sep 17 00:00:00 2001 From: buravc Date: Thu, 20 Aug 2026 10:58:47 +0200 Subject: [PATCH 4/4] test(golang-adk): cover the pre-fork interceptor user id fix Add a test asserting userIDCallInterceptor.Before attaches the user id to the returned ctx. Remove the taskstore-level CallContext test, which no longer applies now that taskstore doesn't read a2asrv identity directly. Signed-off-by: buravc --- go/adk/pkg/a2a/executor_test.go | 19 +++++++++++++++++++ go/adk/pkg/taskstore/store_test.go | 20 -------------------- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go index 51574a32b..c21574f2f 100644 --- a/go/adk/pkg/a2a/executor_test.go +++ b/go/adk/pkg/a2a/executor_test.go @@ -8,6 +8,7 @@ import ( a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/go-logr/logr" + "github.com/kagent-dev/kagent/go/adk/pkg/auth" adkagent "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/model" "google.golang.org/adk/v2/runner" @@ -126,6 +127,24 @@ func TestKAgentExecutor_ForwardsCleanup(t *testing.T) { } } +func TestUserIDCallInterceptor_SetsCtxUserID(t *testing.T) { + ctx, callCtx := a2asrv.NewCallContext(context.Background(), a2asrv.NewServiceParams(map[string][]string{ + "x-user-id": {"real-user"}, + })) + + returned, _, err := UserIDCallInterceptor().Before(ctx, callCtx, &a2asrv.Request{}) + if err != nil { + t.Fatalf("Before() error = %v", err) + } + + if got := callCtx.User.Name; got != "real-user" { + t.Fatalf("callCtx.User.Name = %q, want %q", got, "real-user") + } + if got := auth.UserIDFromContext(returned); got != "real-user" { + t.Fatalf("auth.UserIDFromContext(returned) = %q, want %q", got, "real-user") + } +} + func TestKAgentExecutor_TranslatesADKPauseAtA2ABoundary(t *testing.T) { reqCtx := &a2asrv.ExecutorContext{ TaskID: "task-1", ContextID: "ctx-1", diff --git a/go/adk/pkg/taskstore/store_test.go b/go/adk/pkg/taskstore/store_test.go index bfbc35597..77795221c 100644 --- a/go/adk/pkg/taskstore/store_test.go +++ b/go/adk/pkg/taskstore/store_test.go @@ -8,7 +8,6 @@ import ( a2a "github.com/a2aproject/a2a-go/v2/a2a" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" - "github.com/a2aproject/a2a-go/v2/a2asrv" a2ataskstore "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" "github.com/kagent-dev/kagent/go/adk/pkg/auth" "github.com/kagent-dev/kagent/go/adk/pkg/controllerclient" @@ -125,25 +124,6 @@ func TestGetDecodesCanonicalTask(t *testing.T) { assert.Equal(t, "done", stored.Task.History[0].Parts[0].Text()) } -func TestGetResolvesUserFromA2ACallContext(t *testing.T) { - encoded, err := pbconv.ToProtoTask(&a2a.Task{ID: a2a.TaskID("task-4"), ContextID: "session-4"}) - require.NoError(t, err) - - var gotUserID []string - service := newTaskStore(t, &taskTestServer{get: func(ctx context.Context, _ *apiv1alpha1.GetTaskRequest) (*apiv1alpha1.GetTaskResponse, error) { - values, _ := metadata.FromIncomingContext(ctx) - gotUserID = values.Get("x-user-id") - return &apiv1alpha1.GetTaskResponse{Task: encoded}, nil - }}) - - ctx, callCtx := a2asrv.NewCallContext(t.Context(), nil) - callCtx.User = a2asrv.NewAuthenticatedUser("call-context-user", nil) - - _, err = service.Get(ctx, a2a.TaskID("task-4")) - require.NoError(t, err) - assert.Equal(t, []string{"call-context-user"}, gotUserID) -} - func TestGetMapsNotFound(t *testing.T) { service := newTaskStore(t, &taskTestServer{get: func(context.Context, *apiv1alpha1.GetTaskRequest) (*apiv1alpha1.GetTaskResponse, error) { return nil, status.Error(codes.NotFound, "missing")