From 0ed602d0642378cbbeb3a6db069561438aa94e5b Mon Sep 17 00:00:00 2001 From: erauner Date: Thu, 30 Jul 2026 23:15:52 -0500 Subject: [PATCH] fix(adk): preserve empty MCP arrays in task history --- go/adk/pkg/a2a/converter.go | 62 +++++++++ go/adk/pkg/a2a/converter_test.go | 125 ++++++++++++++++++ go/adk/pkg/a2a/executor.go | 1 + go/adk/pkg/mcp/structured_content_test.go | 146 ++++++++++++++++++++++ 4 files changed, 334 insertions(+) create mode 100644 go/adk/pkg/mcp/structured_content_test.go diff --git a/go/adk/pkg/a2a/converter.go b/go/adk/pkg/a2a/converter.go index e1fef222f..3a5fae3c4 100644 --- a/go/adk/pkg/a2a/converter.go +++ b/go/adk/pkg/a2a/converter.go @@ -2,6 +2,7 @@ package a2a import ( "context" + "encoding/gob" "encoding/json" "maps" @@ -11,6 +12,67 @@ import ( "google.golang.org/genai" ) +// emptyJSONArray survives a2a-go's gob-based task deep copies while retaining +// the original JSON representation. Gob decodes ordinary zero-length slices as +// nil, which would otherwise change structured MCP output from [] to null. +type emptyJSONArray struct{} + +func (emptyJSONArray) MarshalJSON() ([]byte, error) { + return []byte("[]"), nil +} + +func init() { + gob.Register(emptyJSONArray{}) +} + +// preserveEmptyJSONArrays replaces only zero-length JSON arrays inside DataPart +// data. It does not infer arrays from nil values or alter non-empty values. +func preserveEmptyJSONArrays(part a2atype.Part) a2atype.Part { + switch p := part.(type) { + case *a2atype.DataPart: + cp := *p + cp.Data = preserveEmptyJSONArraysInMap(p.Data) + return &cp + case a2atype.DataPart: + p.Data = preserveEmptyJSONArraysInMap(p.Data) + return p + default: + return part + } +} + +func preserveEmptyJSONArraysInMap(input map[string]any) map[string]any { + if input == nil { + return nil + } + output := make(map[string]any, len(input)) + for key, value := range input { + output[key] = preserveEmptyJSONArraysInValue(value) + } + return output +} + +func preserveEmptyJSONArraysInValue(value any) any { + switch typed := value.(type) { + case map[string]any: + return preserveEmptyJSONArraysInMap(typed) + case []any: + if typed == nil { + return typed + } + if len(typed) == 0 { + return emptyJSONArray{} + } + output := make([]any, len(typed)) + for i, item := range typed { + output[i] = preserveEmptyJSONArraysInValue(item) + } + return output + default: + return value + } +} + // isEmptyDataPart returns true if the part is a DataPart with nil or empty Data. // The ADK processor emits such parts as cleanup signals for streaming partial // artifacts and as a fallback for unrecognized GenAI part types. diff --git a/go/adk/pkg/a2a/converter_test.go b/go/adk/pkg/a2a/converter_test.go index 806e39482..68cd91fe4 100644 --- a/go/adk/pkg/a2a/converter_test.go +++ b/go/adk/pkg/a2a/converter_test.go @@ -1,12 +1,17 @@ package a2a import ( + "bytes" "context" + "encoding/gob" + "encoding/json" "testing" a2atype "github.com/a2aproject/a2a-go/a2a" + "github.com/kagent-dev/kagent/go/core/pkg/a2acompat/trpcv0" "google.golang.org/adk/server/adka2a" //nolint:staticcheck // kagent still uses a2a-go v1; this ADK package is the compatibility adapter. "google.golang.org/genai" + trpc "trpc.group/trpc-go/trpc-a2a-go/protocol" ) // --------------------------------------------------------------------------- @@ -289,3 +294,123 @@ func TestToA2AMetadataMap_nil(t *testing.T) { t.Fatalf("expected nil map, got %#v", m) } } + +func TestPreserveEmptyJSONArraysAcrossTaskStatusGobCopy(t *testing.T) { + part := a2atype.DataPart{Data: map[string]any{ + "response": map[string]any{ + "output": map[string]any{ + "messages": []any{}, + "explanation_codes": []any{}, + "matches": []any{map[string]any{"matched_on": []any{}}}, + "unchanged_null": nil, + "unchanged_typed_null": []any(nil), + "unchanged_values": []any{"one"}, + }, + }, + }} + + unprotected := gobCopyMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, part)) + unprotectedOutput := functionResponseOutput(t, unprotected) + unprotectedMessages, ok := unprotectedOutput["messages"].([]any) + if !ok || unprotectedMessages != nil { + t.Fatalf("gob root-cause probe: messages = %#v (%T), want typed nil slice", unprotectedOutput["messages"], unprotectedOutput["messages"]) + } + + protectedPart := preserveEmptyJSONArrays(part) + protected := gobCopyMessage(t, a2atype.NewMessage(a2atype.MessageRoleAgent, protectedPart)) + protected = gobCopyMessage(t, protected) + encoded, err := json.Marshal(protected) + if err != nil { + t.Fatalf("marshal preserved message: %v", err) + } + var wire map[string]any + if err := json.Unmarshal(encoded, &wire); err != nil { + t.Fatalf("unmarshal preserved message: %v", err) + } + parts := wire["parts"].([]any) + data := parts[0].(map[string]any)["data"].(map[string]any) + output := data["response"].(map[string]any)["output"].(map[string]any) + assertNonNilEmptyJSONArray(t, output, "messages") + assertNonNilEmptyJSONArray(t, output, "explanation_codes") + matches := output["matches"].([]any) + assertNonNilEmptyJSONArray(t, matches[0].(map[string]any), "matched_on") + if output["unchanged_null"] != nil { + t.Fatalf("unchanged_null = %#v, want nil", output["unchanged_null"]) + } + if output["unchanged_typed_null"] != nil { + t.Fatalf("unchanged_typed_null = %#v, want JSON null", output["unchanged_typed_null"]) + } + values := output["unchanged_values"].([]any) + if len(values) != 1 || values[0] != "one" { + t.Fatalf("unchanged_values = %#v, want [one]", values) + } + + legacyJSON, err := json.Marshal(&a2atype.Task{ + ID: "task-1", + ContextID: "context-1", + History: []*a2atype.Message{protected}, + Status: a2atype.TaskStatus{State: a2atype.TaskStateCompleted}, + }) + if err != nil { + t.Fatalf("marshal legacy task: %v", err) + } + var legacyTask trpc.Task + if err := json.Unmarshal(legacyJSON, &legacyTask); err != nil { + t.Fatalf("unmarshal legacy task: %v", err) + } + v1Task, err := trpcv0.ToV1Task(&legacyTask) + if err != nil { + t.Fatalf("convert legacy task to v1: %v", err) + } + roundTripTask, err := trpcv0.ToLegacyTask(v1Task) + if err != nil { + t.Fatalf("convert v1 task to legacy: %v", err) + } + roundTripJSON, err := json.Marshal(roundTripTask) + if err != nil { + t.Fatalf("marshal compatibility round trip: %v", err) + } + var roundTripWire map[string]any + if err := json.Unmarshal(roundTripJSON, &roundTripWire); err != nil { + t.Fatalf("unmarshal compatibility round trip: %v", err) + } + history := roundTripWire["history"].([]any) + roundTripParts := history[0].(map[string]any)["parts"].([]any) + roundTripData := roundTripParts[0].(map[string]any)["data"].(map[string]any) + roundTripOutput := roundTripData["response"].(map[string]any)["output"].(map[string]any) + assertNonNilEmptyJSONArray(t, roundTripOutput, "messages") + assertNonNilEmptyJSONArray(t, roundTripOutput, "explanation_codes") +} + +func gobCopyMessage(t *testing.T, message *a2atype.Message) *a2atype.Message { + t.Helper() + var buffer bytes.Buffer + if err := gob.NewEncoder(&buffer).Encode(message); err != nil { + t.Fatalf("gob encode message: %v", err) + } + var copied a2atype.Message + if err := gob.NewDecoder(&buffer).Decode(&copied); err != nil { + t.Fatalf("gob decode message: %v", err) + } + return &copied +} + +func functionResponseOutput(t *testing.T, message *a2atype.Message) map[string]any { + t.Helper() + if len(message.Parts) != 1 { + t.Fatalf("parts = %d, want 1", len(message.Parts)) + } + part := asDataPart(message.Parts[0]) + if part == nil { + t.Fatalf("part type = %T, want DataPart", message.Parts[0]) + } + return part.Data["response"].(map[string]any)["output"].(map[string]any) +} + +func assertNonNilEmptyJSONArray(t *testing.T, values map[string]any, key string) { + t.Helper() + value, ok := values[key].([]any) + if !ok || value == nil || len(value) != 0 { + t.Fatalf("%s = %#v (%T), want non-nil empty JSON array", key, values[key], values[key]) + } +} diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index e92fc7df1..765b2db17 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -326,6 +326,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont if isEmptyDataPart(a2aPart) { continue } + a2aPart = preserveEmptyJSONArrays(a2aPart) // Stamp kagent_subagent_session_id onto function_call DataParts. if len(subagentSessionIDs) > 0 { a2aPart = stampSubagentSessionID(a2aPart, subagentSessionIDs) diff --git a/go/adk/pkg/mcp/structured_content_test.go b/go/adk/pkg/mcp/structured_content_test.go new file mode 100644 index 000000000..a7708c01b --- /dev/null +++ b/go/adk/pkg/mcp/structured_content_test.go @@ -0,0 +1,146 @@ +package mcp + +import ( + "context" + "encoding/json" + "iter" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/model" + "google.golang.org/adk/runner" + adksession "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/mcptoolset" + "google.golang.org/genai" + + kagentsession "github.com/kagent-dev/kagent/go/adk/pkg/session" +) + +const emptyArraysToolName = "empty_arrays" + +type emptyArraysOutput struct { + Messages []string `json:"messages"` + ExplanationCodes []string `json:"explanation_codes"` +} + +type toolCallModel struct { + calls atomic.Int32 +} + +func (*toolCallModel) Name() string { return "tool-call-model" } + +func (m *toolCallModel) GenerateContent(_ context.Context, _ *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + if m.calls.Add(1) == 1 { + part := genai.NewPartFromFunctionCall(emptyArraysToolName, map[string]any{}) + part.FunctionCall.ID = "call-1" + yield(&model.LLMResponse{ + Content: &genai.Content{Role: genai.RoleModel, Parts: []*genai.Part{part}}, + TurnComplete: true, + }, nil) + return + } + yield(&model.LLMResponse{ + Content: genai.NewContentFromText("done", genai.RoleModel), + TurnComplete: true, + }, nil) + } +} + +func TestMCPStructuredContentPreservesNestedEmptyArraysInFunctionResponseEvent(t *testing.T) { + clientTransport, serverTransport := mcpsdk.NewInMemoryTransports() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "empty-arrays", Version: "test"}, nil) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: emptyArraysToolName}, func(context.Context, *mcpsdk.CallToolRequest, struct{}) (*mcpsdk.CallToolResult, emptyArraysOutput, error) { + return nil, emptyArraysOutput{ + Messages: []string{}, + ExplanationCodes: []string{}, + }, nil + }) + serverSession, err := server.Connect(t.Context(), serverTransport, nil) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, serverSession.Close()) }) + + toolset, err := mcptoolset.New(mcptoolset.Config{Transport: clientTransport}) + require.NoError(t, err) + adkAgent, err := llmagent.New(llmagent.Config{ + Name: "structured_content_agent", + Model: &toolCallModel{}, + Toolsets: []tool.Toolset{toolset}, + }) + require.NoError(t, err) + + controller := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodPost && r.URL.Path == "/api/sessions": + require.NoError(t, json.NewEncoder(w).Encode(map[string]any{ + "data": map[string]any{"id": "session-1", "user_id": "user-1"}, + })) + case r.Method == http.MethodGet && r.URL.Path == "/api/sessions/session-1": + require.NoError(t, json.NewEncoder(w).Encode(map[string]any{ + "data": map[string]any{ + "session": map[string]any{"id": "session-1", "user_id": "user-1"}, + "events": []any{}, + }, + })) + case r.Method == http.MethodPost && r.URL.Path == "/api/sessions/session-1/events": + w.WriteHeader(http.StatusCreated) + default: + http.Error(w, "unexpected request", http.StatusNotFound) + } + })) + t.Cleanup(controller.Close) + + sessionService := kagentsession.NewKAgentSessionService(controller.URL, controller.Client()) + _, err = sessionService.Create(t.Context(), &adksession.CreateRequest{ + AppName: "test-app", + UserID: "user-1", + SessionID: "session-1", + }) + require.NoError(t, err) + + adkRunner, err := runner.New(runner.Config{ + AppName: "test-app", + Agent: adkAgent, + SessionService: sessionService, + }) + require.NoError(t, err) + + var response map[string]any + for event, runErr := range adkRunner.Run( + t.Context(), + "user-1", + "session-1", + genai.NewContentFromText("call the tool", genai.RoleUser), + agent.RunConfig{}, + ) { + require.NoError(t, runErr) + if event.Content == nil { + continue + } + for _, part := range event.Content.Parts { + if part.FunctionResponse != nil && part.FunctionResponse.Name == emptyArraysToolName { + response = part.FunctionResponse.Response + } + } + } + + require.NotNil(t, response) + output, ok := response["output"].(map[string]any) + require.True(t, ok, "output type = %T", response["output"]) + messages, ok := output["messages"].([]any) + require.True(t, ok, "messages type = %T", output["messages"]) + require.NotNil(t, messages) + require.Empty(t, messages) + codes, ok := output["explanation_codes"].([]any) + require.True(t, ok, "explanation_codes type = %T", output["explanation_codes"]) + require.NotNil(t, codes) + require.Empty(t, codes) +}