diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index dd6469b7cf..fa33d93679 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -99,6 +99,7 @@ func NewKAgentExecutor(cfg KAgentExecutorConfig) (*KAgentExecutor, error) { A2APartConverter: a2aPartConverter, GenAIPartConverter: structuredOutputPartConverter(output, rootName), AfterEventCallback: func(ctx adka2a.ExecutorContext, event *adksession.Event, processed *a2atype.TaskArtifactUpdateEvent) error { + outputLimitWatchFrom(ctx).observe(event) // Preserve the artifact's protocol type and assign kagent's ordering key. if processed.Artifact != nil { position := event.Timestamp @@ -113,7 +114,11 @@ func NewKAgentExecutor(cfg KAgentExecutorConfig) (*KAgentExecutor, error) { // input-required or failed), whose metadata a2a-go merges into the // stored task. AfterExecuteCallback: func(ctx adka2a.ExecutorContext, finalEvent *a2atype.TaskStatusUpdateEvent, _ error) error { + outputLimitWatchFrom(ctx).failTruncatedAnswer(finalEvent) turnUsageFrom(ctx).stampEvent(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, @@ -267,6 +272,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon // The synthetic ADK user ID is only a native session lookup key. Memory // and outgoing credentials must receive only the passed-through caller. ctx = auth.WithUserID(ctx, trustedUserID) + ctx = withOutputLimitWatch(ctx) // Resumed tasks (HITL cycles, follow-up messages) carry the previously // persisted total, so the usage total stays a task-lifetime sum. usage := &turnUsage{} 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..23bc86a010 --- /dev/null +++ b/go/adk/pkg/a2a/output_limit_test.go @@ -0,0 +1,163 @@ +package a2a + +import ( + "context" + "iter" + "log/slog" + "strings" + "testing" + + 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, err := NewKAgentExecutor(KAgentExecutorConfig{ + AppName: appName, + SessionService: adksession.InMemoryService(), + Logger: slog.New(slog.DiscardHandler), + RunnerConfig: runner.Config{ + AppName: appName, + Agent: agent, + }, + }) + if err != nil { + t.Fatalf("NewKAgentExecutor() error = %v", err) + } + 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) + } +} 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."""