fix(adk): fail the task when the model stops at max_tokens instead of completing with a cut answer

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{}

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)
+}

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)
+	}
+}

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.