From 6c11bd3a15614c7abb0751f6a560cafcd0daaf46 Mon Sep 17 00:00:00 2001 From: Jose Armesto Date: Tue, 22 Sep 2026 17:15:51 +0200 Subject: [PATCH 1/2] fix(adk): fail the task when the model stops at max_tokens instead of completing with a cut answer A model that stops at its output limit returns an answer cut mid-sentence. The A2A task ended as completed, so a client - a person's front end, or a parent agent calling this one through the remote agent tool - read the partial text as the whole answer with no notice. The turn's last model response is now watched for genai.FinishReasonMaxTokens, which every model adapter maps its own stop reason onto, and the terminal task status becomes failed with a message naming the limit. The streamed artifacts are left as they are, so the partial answer stays in the caller's hands. Signed-off-by: Jose Armesto --- go/adk/pkg/a2a/executor.go | 9 ++ go/adk/pkg/a2a/output_limit.go | 98 +++++++++++++++++ go/adk/pkg/a2a/output_limit_test.go | 161 ++++++++++++++++++++++++++++ 3 files changed, 268 insertions(+) create mode 100644 go/adk/pkg/a2a/output_limit.go create mode 100644 go/adk/pkg/a2a/output_limit_test.go diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 44da5b62d0..d35b4bdf2a 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -63,6 +63,7 @@ func NewKAgentExecutor(cfg KAgentExecutorConfig) *KAgentExecutor { A2APartConverter: a2aPartConverter, GenAIPartConverter: genAIPartConverter, AfterEventCallback: func(ctx adka2a.ExecutorContext, event *adksession.Event, processed *a2atype.TaskArtifactUpdateEvent) error { + outputLimitWatchFrom(ctx).observe(event) if event.InvocationID != "" { trace.SpanFromContext(ctx).SetAttributes(attribute.String("gcp.vertex.agent.invocation_id", event.InvocationID)) } @@ -78,6 +79,13 @@ func NewKAgentExecutor(cfg KAgentExecutorConfig) *KAgentExecutor { } return nil }, + AfterExecuteCallback: func(ctx adka2a.ExecutorContext, finalEvent *a2atype.TaskStatusUpdateEvent, _ error) error { + outputLimitWatchFrom(ctx).failTruncatedAnswer(finalEvent) + // Never return the execution error: ADK drops the terminal status + // event and reports the error instead, which would lose the failure + // it is already carrying. + return nil + }, OutputMode: adka2a.OutputArtifactPerEvent, }) @@ -134,6 +142,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon ctx = withBearerToken(ctx) ctx = auth.WithUserID(ctx, userID) + ctx = withOutputLimitWatch(ctx) spanAttributes := map[string]string{ "kagent.user_id": userID, "gen_ai.task.id": string(reqCtx.TaskID), diff --git a/go/adk/pkg/a2a/output_limit.go b/go/adk/pkg/a2a/output_limit.go new file mode 100644 index 0000000000..91af7a9e31 --- /dev/null +++ b/go/adk/pkg/a2a/output_limit.go @@ -0,0 +1,98 @@ +package a2a + +import ( + "context" + "fmt" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "google.golang.org/adk/v2/server/adka2a/v2" + adksession "google.golang.org/adk/v2/session" + "google.golang.org/genai" +) + +// isErrorMessageMetadataKey marks a status message as carrying an error. It is +// the key ADK sets on the failures it reports itself, so a client renders this +// one the same way. +var isErrorMessageMetadataKey = adka2a.ToA2AMetaKey("is_error_message") + +// outputLimitWatch remembers how the turn's last model response ended. +// +// A model that stops at its output limit returns an answer cut mid-sentence, +// and the task must not report that as completed: an A2A client — a person's, +// or a parent agent that calls this one through the remote agent tool — has no +// other way to tell a whole answer from a truncated one. The finish reason is +// the provider-neutral signal; every model adapter maps its own stop reason +// onto it. +type outputLimitWatch struct { + reason genai.FinishReason + // generated is the output the response was billed for. When the model + // stopped at the limit, that is the limit. Every adapter in pkg/models + // reports the total billed output in CandidatesTokenCount, reasoning + // included; ThoughtsTokenCount, where an adapter sets it, is a breakdown of + // that total and must not be added to it. + generated int32 +} + +type outputLimitWatchKey struct{} + +// withOutputLimitWatch gives one execution its own watch. The callbacks below +// are configured once per executor and shared by every concurrent execution, +// so the state they read and write has to travel in the request's context. +func withOutputLimitWatch(ctx context.Context) context.Context { + return context.WithValue(ctx, outputLimitWatchKey{}, &outputLimitWatch{}) +} + +func outputLimitWatchFrom(ctx context.Context) *outputLimitWatch { + watch, _ := ctx.Value(outputLimitWatchKey{}).(*outputLimitWatch) + return watch +} + +// observe records how a model response ended. Streaming chunks and the events +// tools produce carry no finish reason, so only a completed model response +// moves the watch, and the last one of the turn is the one that decides. +func (w *outputLimitWatch) observe(event *adksession.Event) { + if w == nil || event == nil || event.Partial { + return + } + if event.FinishReason == "" || event.FinishReason == genai.FinishReasonUnspecified { + return + } + w.reason = event.FinishReason + w.generated = 0 + if usage := event.UsageMetadata; usage != nil { + w.generated = usage.CandidatesTokenCount + } +} + +// failTruncatedAnswer turns the turn's completed terminal status into a failure +// when the last model response stopped at the output limit. The artifacts +// already streamed are left untouched, so the caller keeps the partial answer +// and reads the failure beside it. +// +// A status that is not completed is left alone: ADK is already reporting a +// failure of its own — a broken function call the truncation caused, say — or a +// pause waiting for input, and overwriting it would replace the cause with a +// symptom. +func (w *outputLimitWatch) failTruncatedAnswer(final *a2atype.TaskStatusUpdateEvent) { + if w == nil || final == nil || w.reason != genai.FinishReasonMaxTokens { + return + } + if final.Status.State != a2atype.TaskStateCompleted { + return + } + part := a2atype.NewTextPart(outputLimitMessage(w.generated)) + part.Metadata = map[string]any{isErrorMessageMetadataKey: true} + final.Status.State = a2atype.TaskStateFailed + final.Status.Message = a2atype.NewMessageForTask(a2atype.MessageRoleAgent, final, part) +} + +// outputLimitMessage names the limit when it is knowable. The response does not +// carry the request's max_tokens, but a model that stopped at the limit +// generated exactly that many tokens, so usage names it; without usage the +// number is left out rather than guessed. +func outputLimitMessage(generated int32) string { + if generated <= 0 { + return "model output limit reached: the answer is incomplete" + } + return fmt.Sprintf("model output limit reached (max_tokens=%d): the answer is incomplete", generated) +} diff --git a/go/adk/pkg/a2a/output_limit_test.go b/go/adk/pkg/a2a/output_limit_test.go new file mode 100644 index 0000000000..f72775b878 --- /dev/null +++ b/go/adk/pkg/a2a/output_limit_test.go @@ -0,0 +1,161 @@ +package a2a + +import ( + "context" + "iter" + "strings" + "testing" + + "log/slog" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + adkagent "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/runner" + adksession "google.golang.org/adk/v2/session" + "google.golang.org/genai" +) + +// runTurn runs one turn of an agent that yields the given responses and returns +// the artifact updates and the terminal task status the executor produced. +func runTurn(t *testing.T, responses []model.LLMResponse) ([]*a2atype.TaskArtifactUpdateEvent, *a2atype.TaskStatusUpdateEvent) { + t.Helper() + const appName = "output-limit-app" + + agent, err := adkagent.New(adkagent.Config{ + Name: "output-limit-agent", + Run: func(ic adkagent.InvocationContext) iter.Seq2[*adksession.Event, error] { + return func(yield func(*adksession.Event, error) bool) { + for _, response := range responses { + event := &adksession.Event{ + Author: ic.Agent().Name(), + InvocationID: ic.InvocationID(), + Branch: ic.Branch(), + LLMResponse: response, + } + if !yield(event, nil) { + return + } + } + } + }, + }) + if err != nil { + t.Fatalf("agent.New() error = %v", err) + } + + executor := NewKAgentExecutor(KAgentExecutorConfig{ + AppName: appName, + SessionService: adksession.InMemoryService(), + Logger: slog.New(slog.DiscardHandler), + RunnerConfig: runner.Config{ + AppName: appName, + Agent: agent, + }, + }) + reqCtx := &a2asrv.ExecutorContext{ + TaskID: "task-1", + ContextID: "context-1", + Message: a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hi")), + } + + var updates []*a2atype.TaskArtifactUpdateEvent + var terminal *a2atype.TaskStatusUpdateEvent + for event, err := range executor.Execute(context.Background(), reqCtx) { + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + switch event := event.(type) { + case *a2atype.TaskArtifactUpdateEvent: + updates = append(updates, event) + case *a2atype.TaskStatusUpdateEvent: + if event.Status.State.Terminal() { + terminal = event + } + } + } + if terminal == nil { + t.Fatal("Execute() produced no terminal task status") + } + return updates, terminal +} + +// statusText returns the text of a terminal status message. +func statusText(t *testing.T, status *a2atype.TaskStatusUpdateEvent) string { + t.Helper() + if status.Status.Message == nil || len(status.Status.Message.Parts) == 0 { + t.Fatalf("terminal status = %#v, want a status message", status) + } + return status.Status.Message.Parts[0].Text() +} + +func TestKAgentExecutorFailsTheTaskWhenTheModelStopsAtTheOutputLimit(t *testing.T) { + updates, terminal := runTurn(t, []model.LLMResponse{ + { + Content: genai.NewContentFromText("the answer starts", genai.RoleModel), + Partial: true, + }, + { + Content: genai.NewContentFromText("the answer starts and is cut", genai.RoleModel), + FinishReason: genai.FinishReasonMaxTokens, + // The reasoning tokens are a breakdown of the billed output, so the + // limit named is 8192 and not 14192. + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{ + CandidatesTokenCount: 8192, + ThoughtsTokenCount: 6000, + }, + }, + }) + + if terminal.Status.State != a2atype.TaskStateFailed { + t.Fatalf("terminal state = %q, want %q", terminal.Status.State, a2atype.TaskStateFailed) + } + text := statusText(t, terminal) + if !strings.Contains(text, "output limit") || !strings.Contains(text, "max_tokens=8192") { + t.Fatalf("status message = %q, want the output limit and the limit value", text) + } + if len(updates) == 0 { + t.Fatal("artifact updates = 0, want the partial answer kept") + } + last := updates[len(updates)-1] + if len(last.Artifact.Parts) != 1 || last.Artifact.Parts[0].Text() != "the answer starts and is cut" { + t.Fatalf("last artifact = %#v, want the streamed partial answer", last.Artifact) + } +} + +func TestKAgentExecutorOmitsTheLimitWhenUsageDoesNotNameIt(t *testing.T) { + _, terminal := runTurn(t, []model.LLMResponse{{ + Content: genai.NewContentFromText("cut", genai.RoleModel), + FinishReason: genai.FinishReasonMaxTokens, + }}) + + if terminal.Status.State != a2atype.TaskStateFailed { + t.Fatalf("terminal state = %q, want %q", terminal.Status.State, a2atype.TaskStateFailed) + } + text := statusText(t, terminal) + if !strings.Contains(text, "output limit") || strings.Contains(text, "max_tokens=") { + t.Fatalf("status message = %q, want the output limit without a number", text) + } +} + +func TestKAgentExecutorCompletesTheTaskWhenTheLastResponseStopsNormally(t *testing.T) { + _, terminal := runTurn(t, []model.LLMResponse{ + { + Content: genai.NewContentFromText("a chunk", genai.RoleModel), + Partial: true, + FinishReason: genai.FinishReasonMaxTokens, + }, + { + Content: genai.NewContentFromText("the whole answer", genai.RoleModel), + FinishReason: genai.FinishReasonStop, + }, + }) + + if terminal.Status.State != a2atype.TaskStateCompleted { + t.Fatalf("terminal state = %q, want %q", terminal.Status.State, a2atype.TaskStateCompleted) + } + if terminal.Status.Message != nil { + t.Fatalf("completed status message = %#v, want none", terminal.Status.Message) + } +} From 0616846c5402fda1613716d6d96a9a165e9f1b63 Mon Sep 17 00:00:00 2001 From: Jose Armesto Date: Mon, 19 Oct 2026 18:20:50 +0200 Subject: [PATCH 2/2] fix(adk): report a failed sub-agent's reason to the parent, beside its output The remote agent tool read a failed task's artifacts before its status message, which is right for a completed task and wrong for a failed one that kept its output: a parent agent got the cut answer as the error and never the reason the task failed. A failed task now yields the status message as the error and the artifacts as the result, in the Go tool and its Python twin. A failure without a status message still reads its text from the artifacts, and one with neither still names the agent. The Python tool returned a bare string for a failure, which ADK hands to the model as a plain result; it now returns the error/result/subagent_session_id dict the Go tool and its own HITL failure branch already use. The import block of output_limit_test.go follows gci, which main now lints with. Signed-off-by: Jose Armesto --- go/adk/pkg/a2a/output_limit_test.go | 3 +- go/adk/pkg/tools/remote_a2a_tool.go | 79 +++++++++++-------- go/adk/pkg/tools/remote_a2a_tool_test.go | 54 +++++++++++++ .../src/kagent/adk/_remote_a2a_tool.py | 63 ++++++++++----- .../tests/unittests/test_remote_a2a_tool.py | 38 ++++++++- 5 files changed, 178 insertions(+), 59 deletions(-) diff --git a/go/adk/pkg/a2a/output_limit_test.go b/go/adk/pkg/a2a/output_limit_test.go index 8620f7b4ff..23bc86a010 100644 --- a/go/adk/pkg/a2a/output_limit_test.go +++ b/go/adk/pkg/a2a/output_limit_test.go @@ -3,11 +3,10 @@ package a2a import ( "context" "iter" + "log/slog" "strings" "testing" - "log/slog" - a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2asrv" adkagent "google.golang.org/adk/v2/agent" diff --git a/go/adk/pkg/tools/remote_a2a_tool.go b/go/adk/pkg/tools/remote_a2a_tool.go index 704d4327f6..0af8fd8414 100644 --- a/go/adk/pkg/tools/remote_a2a_tool.go +++ b/go/adk/pkg/tools/remote_a2a_tool.go @@ -434,17 +434,29 @@ func (s *remoteA2AState) processResult(ctx adkagent.Context, contextID string, r case a2atype.TaskStateInputRequired: return s.handleInputRequired(ctx, r, contextID), nil case a2atype.TaskStateFailed: - text, err := extractTextFromTask(r) + // The status message names the failure. The artifacts, when the + // runtime kept them (an answer cut at the model's output limit, + // say), hold what the agent produced before it failed, and the + // parent reads them beside the reason. A failure without a + // message reads its text from the artifacts, as before. + reason, err := extractTextFromMessage(r.Status.Message) if err != nil { return remoteA2AResponse{}, fmt.Errorf("extract remote agent failure: %w", err) } - if text == "" { - text = fmt.Sprintf("Remote agent '%s' failed.", s.name) + partial, err := extractTextFromArtifacts(r) + if err != nil { + return remoteA2AResponse{}, fmt.Errorf("extract remote agent failure: %w", err) } - return remoteA2AResponse{ - Error: text, - SubagentSessionID: contextID, - }, nil + ret := remoteA2AResponse{SubagentSessionID: contextID} + switch { + case reason != "": + ret.Error, ret.Result = reason, partial + case partial != "": + ret.Error = partial + default: + ret.Error = fmt.Sprintf("Remote agent '%s' failed.", s.name) + } + return ret, nil default: // completed — include sub-agent's final LLM usage from task.metadata // so the parent can display it on the AgentCall card in the UI. @@ -540,36 +552,37 @@ func extractTextFromTask(task *a2atype.Task) (string, error) { return "", nil } // Prefer artifacts (canonical result). - if len(task.Artifacts) > 0 { - var texts []string - for _, artifact := range task.Artifacts { - for _, part := range artifact.Parts { - if part == nil { - continue - } - if text := part.Text(); text != "" { - texts = append(texts, text) - continue - } - if !kagenta2a.IsStructuredOutputPart(part) { - continue - } - text, err := kagenta2a.StructuredOutputJSON(part) - if err != nil { - return "", err - } + text, err := extractTextFromArtifacts(task) + if err != nil || text != "" { + return text, err + } + // Fall back to status message. + return extractTextFromMessage(task.Status.Message) +} + +// extractTextFromArtifacts extracts text or a structured result from a Task's artifacts. +func extractTextFromArtifacts(task *a2atype.Task) (string, error) { + var texts []string + for _, artifact := range task.Artifacts { + for _, part := range artifact.Parts { + if part == nil { + continue + } + if text := part.Text(); text != "" { texts = append(texts, text) + continue } - } - if len(texts) > 0 { - return strings.Join(texts, "\n"), nil + if !kagenta2a.IsStructuredOutputPart(part) { + continue + } + text, err := kagenta2a.StructuredOutputJSON(part) + if err != nil { + return "", err + } + texts = append(texts, text) } } - // Fall back to status message. - if task.Status.Message != nil { - return extractTextFromMessage(task.Status.Message) - } - return "", nil + return strings.Join(texts, "\n"), nil } // extractTextFromMessage extracts text or a structured result from a direct A2A Message response. diff --git a/go/adk/pkg/tools/remote_a2a_tool_test.go b/go/adk/pkg/tools/remote_a2a_tool_test.go index 15301ed383..f10f48c941 100644 --- a/go/adk/pkg/tools/remote_a2a_tool_test.go +++ b/go/adk/pkg/tools/remote_a2a_tool_test.go @@ -215,6 +215,60 @@ func TestProcessResult_SetsSubagentSessionIDOnEveryBranch(t *testing.T) { }) } +func TestProcessResult_FailedTaskReportsTheReasonBesideThePartialAnswer(t *testing.T) { + s := &remoteA2AState{name: "worker"} + status := func(text string) a2atype.TaskStatus { + return a2atype.TaskStatus{ + State: a2atype.TaskStateFailed, + Message: a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.NewTextPart(text)), + } + } + artifacts := []*a2atype.Artifact{{Parts: a2atype.ContentParts{a2atype.NewTextPart("the answer starts and is cut")}}} + + tests := []struct { + name string + task *a2atype.Task + wantError string + wantResult string + }{ + { + name: "a status message names the failure and the artifacts stay the result", + task: &a2atype.Task{Status: status("model output limit reached (max_tokens=8192): the answer is incomplete"), Artifacts: artifacts}, + wantError: "model output limit reached (max_tokens=8192): the answer is incomplete", + wantResult: "the answer starts and is cut", + }, + { + name: "a status message alone is the error", + task: &a2atype.Task{Status: status("something broke")}, + wantError: "something broke", + }, + { + name: "without a status message the artifacts are the error, as before", + task: &a2atype.Task{Status: a2atype.TaskStatus{State: a2atype.TaskStateFailed}, Artifacts: artifacts}, + wantError: "the answer starts and is cut", + }, + { + name: "without either the agent is named", + task: &a2atype.Task{Status: a2atype.TaskStatus{State: a2atype.TaskStateFailed}}, + wantError: "Remote agent 'worker' failed.", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp, err := s.processResult(nil, "ctx-1", tt.task) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Error != tt.wantError || resp.Result != tt.wantResult { + t.Fatalf("response = %+v, want Error %q and Result %q", resp, tt.wantError, tt.wantResult) + } + if resp.SubagentSessionID != "ctx-1" { + t.Errorf("SubagentSessionID = %q, want %q", resp.SubagentSessionID, "ctx-1") + } + }) + } +} + func TestHandleInputRequiredStoresPublicRemoteHitlState(t *testing.T) { s := &remoteA2AState{name: "worker"} task := &a2atype.Task{ diff --git a/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py b/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py index f6e74e5ec8..d84f37ac8c 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py +++ b/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py @@ -79,26 +79,45 @@ def _extract_text_from_task(task: Task) -> str: """Extract text content from a completed task's artifacts or status message.""" # Prefer artifacts (the canonical result) - if task.artifacts: - texts: list[str] = [] - for artifact in task.artifacts: - if artifact.parts: - for part in artifact.parts: - if part.HasField("text") and part.text: - texts.append(part.text) - if texts: - return "\n".join(texts) - - # Fall back to status message - if task.status and task.status.message and task.status.message.parts: - texts = [] - for part in task.status.message.parts: + return _extract_text_from_artifacts(task) or _extract_text_from_status(task) + + +def _extract_text_from_artifacts(task: Task) -> str: + """Extract text content from a task's artifacts.""" + texts: list[str] = [] + for artifact in task.artifacts: + for part in artifact.parts: if part.HasField("text") and part.text: texts.append(part.text) - if texts: - return "\n".join(texts) + return "\n".join(texts) + + +def _extract_text_from_status(task: Task) -> str: + """Extract text content from a task's status message.""" + if not (task.status and task.status.message): + return "" + texts = [part.text for part in task.status.message.parts if part.HasField("text") and part.text] + return "\n".join(texts) - return "" + +def _failure_response(task: Task, context_id: Optional[str], fallback: str) -> dict[str, Any]: + """Report a failed task to the parent. + + The status message names the failure. The artifacts, when the runtime kept + them (an answer cut at the model's output limit, say), hold what the agent + produced before it failed, and the parent reads them beside the reason. A + failure without a message reads its text from the artifacts, as before. + """ + reason = _extract_text_from_status(task) + partial = _extract_text_from_artifacts(task) + response: dict[str, Any] = {"subagent_session_id": context_id} + if reason: + response["error"] = reason + if partial: + response["result"] = partial + else: + response["error"] = partial or fallback + return response def _extract_usage_from_task(task: Task) -> Optional[dict]: @@ -325,8 +344,7 @@ async def _send_first_call(self, args: dict[str, Any], tool_context: ToolContext return self._handle_input_required(task, tool_context, context_id) if state == TaskState.TASK_STATE_FAILED: - error_text = _extract_text_from_task(task) - return error_text or f"Remote agent '{self.name}' failed." + return _failure_response(task, context_id, f"Remote agent '{self.name}' failed.") # completed — include the sub-agent's final LLM usage from task.metadata # so the parent can display it on the AgentCall card in the UI. @@ -433,8 +451,11 @@ async def _handle_resume(self, tool_context: ToolContext) -> Any: return self._handle_input_required(task, tool_context) if task_state == TaskState.TASK_STATE_FAILED: - error_text = _extract_text_from_task(task) - return error_text or f"Remote agent '{remote_state.subagent_name}' failed after resume." + return _failure_response( + task, + remote_state.context_id or self._last_context_id, + f"Remote agent '{remote_state.subagent_name}' failed after resume.", + ) result_text = _extract_text_from_task(task) usage = _extract_usage_from_task(task) diff --git a/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py b/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py index 961f7a7525..2b9cd67ef3 100644 --- a/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py +++ b/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py @@ -4,9 +4,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx -from a2a.types import Message as A2AMessage -from a2a.types import Part as A2APart from a2a.types import ( + Artifact, Role, SendMessageRequest, StreamResponse, @@ -14,6 +13,8 @@ TaskState, TaskStatus, ) +from a2a.types import Message as A2AMessage +from a2a.types import Part as A2APart from google.adk.tools.tool_confirmation import ToolConfirmation from google.adk.tools.tool_context import ToolContext from google.protobuf.json_format import MessageToDict @@ -280,7 +281,38 @@ async def test_failed_task_returns_error_text(self): finally: p.stop() - assert result == "something broke" + assert result == {"error": "something broke", "subagent_session_id": "ctx-1"} + + async def test_failed_task_keeps_the_partial_answer_beside_the_reason(self): + """A failure that kept its artifacts reports the reason as the error and the artifacts as the result.""" + tool = _make_tool() + reason = "model output limit reached (max_tokens=8192): the answer is incomplete" + task = _make_task(TaskState.TASK_STATE_FAILED, text=reason) + task.artifacts.append(Artifact(artifact_id="a-1", parts=[A2APart(text="the answer starts and is cut")])) + p, _ = _patch_client(tool, _async_yield((task, None))) + try: + result = await tool.run_async(args={"request": "go"}, tool_context=MockToolContext().as_tool_context()) + finally: + p.stop() + + assert result == { + "error": reason, + "result": "the answer starts and is cut", + "subagent_session_id": "ctx-1", + } + + async def test_failed_task_without_a_status_message_reports_the_artifacts(self): + """A failure without a status message reads its text from the artifacts, as before.""" + tool = _make_tool() + task = _make_task(TaskState.TASK_STATE_FAILED) + task.artifacts.append(Artifact(artifact_id="a-1", parts=[A2APart(text="the answer starts and is cut")])) + p, _ = _patch_client(tool, _async_yield((task, None))) + try: + result = await tool.run_async(args={"request": "go"}, tool_context=MockToolContext().as_tool_context()) + finally: + p.stop() + + assert result == {"error": "the answer starts and is cut", "subagent_session_id": "ctx-1"} async def test_context_id_sent_in_outgoing_message(self): """The receiving Agent assigns the first conversation ID."""