From b5f9577f98bd473064ab7f25e0f598fb121abe7c Mon Sep 17 00:00:00 2001 From: pood1e Date: Sat, 31 Oct 2026 14:14:44 +0800 Subject: [PATCH] fix(adk): bridge workflow input through native HITL Signed-off-by: pood1e --- go/adk/pkg/a2a/executor.go | 7 +- go/adk/pkg/a2a/hitl.go | 21 ++-- go/adk/pkg/a2a/hitl_test.go | 33 ++++-- go/adk/pkg/a2a/workflow_input.go | 75 ++++++++++++ go/adk/pkg/a2a/workflow_input_test.go | 161 ++++++++++++++++++++++++++ 5 files changed, 276 insertions(+), 21 deletions(-) create mode 100644 go/adk/pkg/a2a/workflow_input.go create mode 100644 go/adk/pkg/a2a/workflow_input_test.go diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index dd6469b7cf..bf3dea3125 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -344,7 +344,12 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon // If the event is a task status update event and the status is input required, build the HITL status message if update, ok := event.(*a2atype.TaskStatusUpdateEvent); ok && update.Status.State == a2atype.TaskStateInputRequired && update.Status.Message != nil { - update.Status.Message = BuildHITLStatusMessage(update.Status.Message, hitlActivated) + public, buildErr := BuildHITLStatusMessage(update.Status.Message, hitlActivated) + if buildErr != nil { + yield(nil, fmt.Errorf("failed to build HITL status message: %w", buildErr)) + return + } + update.Status.Message = public update.Status.Message.TaskID = update.TaskID update.Status.Message.ContextID = update.ContextID } diff --git a/go/adk/pkg/a2a/hitl.go b/go/adk/pkg/a2a/hitl.go index 9ec9e83f06..2a718d8b07 100644 --- a/go/adk/pkg/a2a/hitl.go +++ b/go/adk/pkg/a2a/hitl.go @@ -460,10 +460,13 @@ func hitlStatusText(tools []apia2a.HITLTool, hints []string) string { return genericHITLText } -// BuildHITLStatusMessage: ADK confirmation DataParts → public HITL Message extension. -func BuildHITLStatusMessage(message *a2atype.Message, activated bool) *a2atype.Message { +// BuildHITLStatusMessage maps ADK pause DataParts to the public HITL Message extension. +func BuildHITLStatusMessage(message *a2atype.Message, activated bool) (*a2atype.Message, error) { if message == nil { - return nil + return nil, nil + } + if public, err := buildWorkflowInputStatus(message, activated); public != nil || err != nil { + return public, err } var tools []apia2a.HITLTool var remote *RemoteHitlState @@ -488,7 +491,7 @@ func BuildHITLStatusMessage(message *a2atype.Message, activated bool) *a2atype.M } } if len(tools) == 0 { - return message + return message, nil } // The text part carries the same information as the typed payload, so a client @@ -498,7 +501,7 @@ func BuildHITLStatusMessage(message *a2atype.Message, activated bool) *a2atype.M public := a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.NewTextPart(text)) public.TaskID, public.ContextID = message.TaskID, message.ContextID if !activated { - return public + return public, nil } var nested *apia2a.NestedHITLRequest @@ -515,7 +518,7 @@ func BuildHITLStatusMessage(message *a2atype.Message, activated bool) *a2atype.M return AttachHitlExtension(public, &apia2a.AskUserRequest{ Type: HITLTypeAskUserRequest, ID: tools[0].ID, Questions: remote.AskUserRequest.Questions, Nested: nested, - }) + }), nil } // An ask_user call with no answerable question becomes an approval: an empty // question list gives a request that no response can satisfy, because resume @@ -524,11 +527,11 @@ func BuildHITLStatusMessage(message *a2atype.Message, activated bool) *a2atype.M return AttachHitlExtension(public, &apia2a.AskUserRequest{ Type: HITLTypeAskUserRequest, ID: tools[0].ID, Questions: questions, - }) + }), nil } return AttachHitlExtension(public, &apia2a.ToolApprovalRequest{ Type: HITLTypeToolApprovalRequest, Hint: text, Tools: tools, Nested: nested, - }) + }), nil } func publicAskUserQuestions(value any) []apia2a.HITLQuestion { @@ -587,6 +590,8 @@ func BuildResumeHITLMessage(storedTask *a2atype.Task, incoming *a2atype.Message) var parts []*a2atype.Part var err error switch { + case askReq != nil && isWorkflowInput(storedTask.Status.Message, askReq.ID): + parts, err = workflowInputResponse(askReq, incoming) case askReq != nil && askReq.Nested != nil: parts, err = processNestedAskUser(askReq, incoming) case approvalReq != nil && approvalReq.Nested != nil: diff --git a/go/adk/pkg/a2a/hitl_test.go b/go/adk/pkg/a2a/hitl_test.go index d8f790e337..83227fdb48 100644 --- a/go/adk/pkg/a2a/hitl_test.go +++ b/go/adk/pkg/a2a/hitl_test.go @@ -138,7 +138,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { t.Run("tool approval", func(t *testing.T) { internal := a2atype.NewMessage(a2atype.MessageRoleAgent, confirmationPart("confirm-1", "delete_file", "call-1", map[string]any{"path": "/tmp/x"}, nil)) - payload := GetToolApprovalRequest(BuildHITLStatusMessage(internal, true)) + payload := GetToolApprovalRequest(mustBuildHITLStatusMessage(t, internal, true)) if payload == nil || payload.Tools[0].ID != "confirm-1" { t.Fatalf("payload = %#v", payload) } @@ -150,7 +150,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { }} internal := a2atype.NewMessage(a2atype.MessageRoleAgent, confirmationPart("confirm-2", "ask_user", "call-2", map[string]any{"questions": questions}, nil)) - payload := GetAskUserRequest(BuildHITLStatusMessage(internal, true)) + payload := GetAskUserRequest(mustBuildHITLStatusMessage(t, internal, true)) if payload == nil || len(payload.Questions) != 1 || payload.Questions[0].Question != "Which databases?" || !slices.Equal(payload.Questions[0].Choices, []string{"PostgreSQL", "MySQL"}) || !payload.Questions[0].Multiple { t.Fatalf("payload = %#v", payload) @@ -160,7 +160,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { t.Run("ask user without questions asks for approval", func(t *testing.T) { internal := a2atype.NewMessage(a2atype.MessageRoleAgent, hintlessConfirmationPart("confirm-2", "ask_user", "call-2")) - public := BuildHITLStatusMessage(internal, true) + public := mustBuildHITLStatusMessage(t, internal, true) if GetAskUserRequest(public) != nil { t.Fatal("an ask_user call without questions must not become an ask_user request") } @@ -177,7 +177,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { questions := []any{map[string]any{"question": ""}, map[string]any{"question": "Which cluster?"}} internal := a2atype.NewMessage(a2atype.MessageRoleAgent, confirmationPart("confirm-2", "ask_user", "call-2", map[string]any{"questions": questions}, nil)) - public := BuildHITLStatusMessage(internal, true) + public := mustBuildHITLStatusMessage(t, internal, true) payload := GetAskUserRequest(public) if payload == nil || len(payload.Questions) != 1 || payload.Questions[0].Question != "Which cluster?" { t.Fatalf("payload = %#v", payload) @@ -198,7 +198,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { }}, }, } - payload := GetToolApprovalRequest(BuildHITLStatusMessage( + payload := GetToolApprovalRequest(mustBuildHITLStatusMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, confirmationPart("parent-confirm", "k8s_agent", "parent-call", nil, remote.ToMap())), true, @@ -209,7 +209,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { }) t.Run("not activated", func(t *testing.T) { - public := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, + public := mustBuildHITLStatusMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, confirmationPart("confirm-1", "delete_file", "call-1", nil, nil)), false) if GetToolApprovalRequest(public) != nil { t.Fatalf("unexpected payload on inactive client") @@ -220,7 +220,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { }) t.Run("not activated without a tool hint names the tool", func(t *testing.T) { - public := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, + public := mustBuildHITLStatusMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, hintlessConfirmationPart("confirm-1", "delete_file", "call-1")), false) if want := "Approval is required for tool(s): delete_file"; messageText(public) != want { t.Errorf("text = %q, want %q", messageText(public), want) @@ -232,7 +232,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { confirmationPart("confirm-1", "delete_file", "call-1", nil, nil), hintlessConfirmationPart("confirm-2", "restart_pod", "call-2")) want := "Please confirm (delete_file, restart_pod)" - if text := messageText(BuildHITLStatusMessage(internal, false)); text != want { + if text := messageText(mustBuildHITLStatusMessage(t, internal, false)); text != want { t.Errorf("text = %q, want %q", text, want) } }) @@ -242,13 +242,13 @@ func TestBuildHITLStatusMessage(t *testing.T) { hintedConfirmationPart("confirm-1", "delete_file", "call-1", "First hint"), hintedConfirmationPart("confirm-2", "restart_pod", "call-2", "Second hint")) want := "First hint; Second hint (delete_file, restart_pod)" - if text := messageText(BuildHITLStatusMessage(internal, false)); text != want { + if text := messageText(mustBuildHITLStatusMessage(t, internal, false)); text != want { t.Errorf("text = %q, want %q", text, want) } }) t.Run("activated payload hint repeats the text", func(t *testing.T) { - public := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, + public := mustBuildHITLStatusMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, hintlessConfirmationPart("confirm-1", "delete_file", "call-1")), true) payload := GetToolApprovalRequest(public) if payload == nil { @@ -263,7 +263,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { questions := []any{map[string]any{"question": "Which database?"}} internal := a2atype.NewMessage(a2atype.MessageRoleAgent, askUserConfirmationPart("confirm-2", "call-2", questions, "Which database?")) - public := BuildHITLStatusMessage(internal, false) + public := mustBuildHITLStatusMessage(t, internal, false) if GetAskUserRequest(public) != nil { t.Fatal("unexpected payload on inactive client") } @@ -273,7 +273,7 @@ func TestBuildHITLStatusMessage(t *testing.T) { }) t.Run("non-confirmation long-running call", func(t *testing.T) { - public := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, + public := mustBuildHITLStatusMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, dataPart(map[string]any{"name": "auth_required"}, map[string]any{ "adk_type": "function_call", "adk_is_long_running": true, })), true) @@ -469,3 +469,12 @@ func TestBuildRemoteHitlStateAndHintAskUserNested(t *testing.T) { t.Fatalf("hint = %q, want %q", got, want) } } + +func mustBuildHITLStatusMessage(t *testing.T, message *a2atype.Message, activated bool) *a2atype.Message { + t.Helper() + public, err := BuildHITLStatusMessage(message, activated) + if err != nil { + t.Fatal(err) + } + return public +} diff --git a/go/adk/pkg/a2a/workflow_input.go b/go/adk/pkg/a2a/workflow_input.go new file mode 100644 index 0000000000..24b6900c9e --- /dev/null +++ b/go/adk/pkg/a2a/workflow_input.go @@ -0,0 +1,75 @@ +package a2a + +import ( + "fmt" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "google.golang.org/adk/v2/server/adka2a/v2" + "google.golang.org/adk/v2/workflow" + + apia2a "github.com/kagent-dev/kagent/go/api/a2a" +) + +// Keep the canonical call so resume reaches the workflow input dispatcher. +func buildWorkflowInputStatus(message *a2atype.Message, activated bool) (*a2atype.Message, error) { + var public *a2atype.Message + var pending int + for _, part := range message.Parts { + data := asDataPart(part) + if data == nil || + part.Metadata[adka2a.ToA2AMetaKey(A2ADataPartMetadataTypeKey)] != A2ADataPartMetadataTypeFunctionCall || + part.Metadata[adka2a.ToA2AMetaKey(A2ADataPartMetadataIsLongRunningKey)] != true { + continue + } + pending++ + if data["name"] != workflow.WorkflowInputFunctionCallName { + continue + } + if public != nil { + return nil, fmt.Errorf("workflow pause requires exactly one input request") + } + converted, err := adka2a.ToGenAIPart(part) + if err != nil { + return nil, fmt.Errorf("failed to decode workflow input request: %w", err) + } + if converted == nil || converted.FunctionCall == nil || converted.FunctionCall.ID == "" { + return nil, fmt.Errorf("workflow input request requires a correlation id") + } + call := converted.FunctionCall + text, _ := call.Args["message"].(string) + if text == "" { + text = "Provide input to continue the workflow." + } + public = a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.NewTextPart(text), apia2a.NewToolCallPart(call.ID, call.Name, call.Args)) + public.TaskID, public.ContextID = message.TaskID, message.ContextID + if activated { + AttachHitlExtension(public, &apia2a.AskUserRequest{ + Type: HITLTypeAskUserRequest, ID: call.ID, + Questions: []apia2a.HITLQuestion{{Question: text}}, + }) + } + } + if public != nil && pending != 1 { + return nil, fmt.Errorf("workflow pause requires exactly one input request") + } + return public, nil +} + +func isWorkflowInput(message *a2atype.Message, id string) bool { + for _, part := range message.Parts { + call, ok := apia2a.ParseToolActivity(part) + if ok && call.Kind == apia2a.ToolCallKind && call.Name == workflow.WorkflowInputFunctionCallName && call.ID == id { + return true + } + } + return false +} + +func workflowInputResponse(request *apia2a.AskUserRequest, message *a2atype.Message) ([]*a2atype.Part, error) { + response := GetAskUserResponse(message) + if response == nil || response.ID != request.ID || len(response.Answers) != 1 || len(response.Answers[0].Answer) != 1 { + return nil, fmt.Errorf("workflow input response requires one correlated answer") + } + return []*a2atype.Part{apia2a.NewToolResultPart(request.ID, workflow.WorkflowInputFunctionCallName, + map[string]any{"response": response.Answers[0].Answer[0]})}, nil +} diff --git a/go/adk/pkg/a2a/workflow_input_test.go b/go/adk/pkg/a2a/workflow_input_test.go new file mode 100644 index 0000000000..bec1eb26f3 --- /dev/null +++ b/go/adk/pkg/a2a/workflow_input_test.go @@ -0,0 +1,161 @@ +package a2a + +import ( + "fmt" + "log/slog" + "testing" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/agent/workflowagent" + "google.golang.org/adk/v2/runner" + "google.golang.org/adk/v2/server/adka2a/v2" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/workflow" + "google.golang.org/genai" + + apia2a "github.com/kagent-dev/kagent/go/api/a2a" +) + +func TestWorkflowInputHITLRoundTrip(t *testing.T) { + for _, answer := range []string{"approve", "reject", `{"approved":true}`} { + t.Run(answer, func(t *testing.T) { + var received any + rerun := true + ask := workflow.NewEmittingFunctionNode("confirm", func(ctx agent.Context, _ any, emit func(*session.Event) error) (any, error) { + return workflow.ResumeOrRequestInput(ctx, emit, session.RequestInput{ + InterruptID: "confirmation", Message: "Confirm the proposal", + }) + }, workflow.NodeConfig{RerunOnResume: &rerun}) + finish := workflow.NewFunctionNode("finish", func(_ agent.Context, input any) (any, error) { + received = input + return input, nil + }, workflow.NodeConfig{}) + root, err := workflowagent.New(workflowagent.Config{ + Name: "workflow", Edges: workflow.Chain(workflow.Start, ask, finish), + }) + if err != nil { + t.Fatal(err) + } + executor, err := NewKAgentExecutor(KAgentExecutorConfig{ + AppName: "test-app", SessionService: session.InMemoryService(), + Logger: slog.New(slog.DiscardHandler), + RunnerConfig: runner.Config{AppName: "test-app", Agent: root}, + }) + if err != nil { + t.Fatal(err) + } + ctx, call := a2asrv.NewCallContext(t.Context(), a2asrv.NewServiceParams(map[string][]string{ + a2atype.SvcParamExtensions: {HITLExtensionURI}, + })) + if _, _, err := HITLActivationInterceptor().Before(ctx, call, &a2asrv.Request{}); err != nil { + t.Fatal(err) + } + request := &a2asrv.ExecutorContext{ + TaskID: "task", ContextID: "context", + Message: a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("begin")), + } + var pause *a2atype.TaskStatusUpdateEvent + for event, err := range executor.Execute(ctx, request) { + if err != nil { + t.Fatal(err) + } + if update, ok := event.(*a2atype.TaskStatusUpdateEvent); ok && update.Status.State == a2atype.TaskStateInputRequired { + pause = update + } + } + if pause == nil || received != nil { + t.Fatalf("pause = %#v, downstream input = %v", pause, received) + } + question := GetAskUserRequest(pause.Status.Message) + if question == nil || question.ID != "confirmation" || len(question.Questions) != 1 { + t.Fatalf("workflow input is not a public answerable question: %#v", question) + } + request.StoredTask = &a2atype.Task{ID: request.TaskID, ContextID: request.ContextID, Status: pause.Status} + request.Message = hitlDecisionMessage(&apia2a.AskUserResponse{ + Type: HITLTypeAskUserResponse, ID: question.ID, + Answers: []apia2a.AskUserAnswer{{Answer: []string{answer}}}, + }) + var completed bool + for event, err := range executor.Execute(ctx, request) { + if err != nil { + t.Fatal(err) + } + if update, ok := event.(*a2atype.TaskStatusUpdateEvent); ok && update.Status.State == a2atype.TaskStateCompleted { + completed = true + } + } + if !completed || received == nil { + t.Fatalf("completed = %v, downstream input = %v", completed, received) + } + if answer != `{"approved":true}` && received != answer { + t.Fatalf("downstream input = %v, want %q", received, answer) + } + if answer == `{"approved":true}` && fmt.Sprint(received) != "map[approved:true]" { + t.Fatalf("structured input = %#v", received) + } + }) + } +} + +func TestWorkflowInputResponseRequiresCorrelation(t *testing.T) { + message := AttachHitlExtension(a2atype.NewMessage(a2atype.MessageRoleAgent, + apia2a.NewToolCallPart("input", workflow.WorkflowInputFunctionCallName, nil)), &apia2a.AskUserRequest{ + Type: HITLTypeAskUserRequest, ID: "input", + Questions: []apia2a.HITLQuestion{{Question: "Confirm"}}, + }) + stored := &a2atype.Task{Status: a2atype.TaskStatus{State: a2atype.TaskStateInputRequired, Message: message}} + for _, test := range []struct { + name string + id string + answers []apia2a.AskUserAnswer + }{ + {"stale request", "stale", []apia2a.AskUserAnswer{{Answer: []string{"approve"}}}}, + {"missing answer", "input", nil}, + {"multiple values", "input", []apia2a.AskUserAnswer{{Answer: []string{"approve", "reject"}}}}, + {"multiple answers", "input", []apia2a.AskUserAnswer{{Answer: []string{"approve"}}, {Answer: []string{"reject"}}}}, + } { + t.Run(test.name, func(t *testing.T) { + incoming := hitlDecisionMessage(&apia2a.AskUserResponse{ + Type: HITLTypeAskUserResponse, ID: test.id, Answers: test.answers, + }) + if _, err := BuildResumeHITLMessage(stored, incoming); err == nil { + t.Fatal("invalid workflow input was accepted") + } + }) + } +} + +func TestWorkflowInputRejectsMixedPendingCalls(t *testing.T) { + for _, other := range []string{workflow.WorkflowInputFunctionCallName, "adk_request_confirmation"} { + t.Run(other, func(t *testing.T) { + parts, err := adka2a.ToA2AParts([]*genai.Part{ + {FunctionCall: &genai.FunctionCall{ID: "input", Name: workflow.WorkflowInputFunctionCallName}}, + {FunctionCall: &genai.FunctionCall{ID: "other", Name: other}}, + }, []string{"input", "other"}) + if err != nil { + t.Fatal(err) + } + if _, err := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, parts...), true); err == nil { + t.Fatal("mixed pending calls were silently discarded") + } + }) + } +} + +func TestWorkflowInputWithoutExtensionRetainsReadablePrompt(t *testing.T) { + parts, err := adka2a.ToA2AParts([]*genai.Part{{FunctionCall: &genai.FunctionCall{ + ID: "input", Name: workflow.WorkflowInputFunctionCallName, Args: map[string]any{"message": "Confirm the proposal"}, + }}}, []string{"input"}) + if err != nil { + t.Fatal(err) + } + public, err := BuildHITLStatusMessage(a2atype.NewMessage(a2atype.MessageRoleAgent, parts...), false) + if err != nil { + t.Fatal(err) + } + if GetAskUserRequest(public) != nil || messageText(public) != "Confirm the proposal" || !isWorkflowInput(public, "input") { + t.Fatalf("invalid non-negotiated workflow input: %#v", public) + } +}