refactor(adk): address pr review comments

- Propagate ProcessState error instead of ignoring it
- Remove unused parameter from applyDecisionForRetry
- Use named return to eliminate repeated var zero M in generateLegacy
- Use isNilMessage helper consistently in runctx and utils
- Reorder StreamErr/concatenatedMessage checks for consistency
- Rename TypedRootInput to AgenticRootInput for clarity

Change-Id: Id179317d22d98ecd292f5caac80b303ecc4ce58f

Author: shentongmartin

adk/chatmodel.go

@@ -713,7 +713,7 @@ func extractTextContent[M messageType](msg M) string {
 				texts = append(texts, block.AssistantGenText.Text)
 			}
 		}
-		return strings.Join(texts, "")
+		return strings.Join(texts, "\n")
 	default:
 		return ""
 	}
@@ -950,10 +950,12 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu
 			if err != nil {
 				return nil, err
 			}
-			_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
+			if err := compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 				st.Messages = append(st.Messages, messages...)
 				return nil
-			})
+			}); err != nil {
+				return nil, err
+			}
 			return messages, nil
 		}))
 

adk/handler.go

@@ -281,7 +281,7 @@ func (b *TypedBaseChatModelAgentMiddleware[M]) AfterToolCallsRewriteState(ctx co
 
 func processTypedState(ctx context.Context, fn func(extra map[string]any) map[string]any) error {
 	runCtx := getRunCtx(ctx)
-	if runCtx != nil && runCtx.TypedRootInput != nil {
+	if runCtx != nil && runCtx.AgenticRootInput != nil {
 		return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.AgenticMessage]) error {
 			st.Extra = fn(st.Extra)
 			return nil

adk/retry_chatmodel.go

@@ -312,7 +312,7 @@ func (r *typedRetryModelWrapper[M]) Generate(ctx context.Context, input []M, opt
 	return r.generateLegacy(ctx, input, opts...)
 }
 
-func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []M, opts ...model.Option) (M, error) {
+func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []M, opts ...model.Option) (zero M, _ error) {
 	isRetryAble := r.config.IsRetryAble
 	if isRetryAble == nil {
 		isRetryAble = defaultIsRetryAble
@@ -330,30 +330,25 @@ func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []
 		}
 
 		if _, ok := compose.ExtractInterruptInfo(err); ok {
-			var zero M
 			return zero, err
 		}
 
 		if errors.Is(err, ErrStreamCanceled) {
-			var zero M
 			return zero, err
 		}
 
 		if !isRetryAble(ctx, err) {
-			var zero M
 			return zero, err
 		}
 
 		lastErr = err
 		if attempt < r.config.MaxRetries {
 			if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil {
-				var zero M
 				return zero, err
 			}
 		}
 	}
 
-	var zero M
 	return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
 }
 
@@ -439,7 +434,7 @@ func generateWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx con
 			break
 		}
 
-		applyDecisionForRetry(r, &currentInput, &currentOpts, ctx, decision)
+		applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 
 		delay := decision.Backoff
 		if delay == 0 {
@@ -576,7 +571,7 @@ func streamWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx conte
 
 			lastErr = err
 			if attempt < r.config.MaxRetries {
-				applyDecisionForRetry(r, &currentInput, &currentOpts, ctx, decision)
+				applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 				delay := decision.Backoff
 				if delay == 0 {
 					delay = backoffFunc(ctx, attempt+1)
@@ -646,7 +641,7 @@ func streamWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx conte
 		lastErr = verdictErr
 
 		if attempt < r.config.MaxRetries {
-			applyDecisionForRetry(r, &currentInput, &currentOpts, ctx, decision)
+			applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 			delay := decision.Backoff
 			if delay == 0 {
 				delay = backoffFunc(ctx, attempt+1)
@@ -660,7 +655,7 @@ func streamWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx conte
 	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
 }
 
-func applyDecisionForRetry(r *typedRetryModelWrapper[*schema.Message], currentInput *[]*schema.Message, currentOpts *[]model.Option, ctx context.Context, decision *RetryDecision) {
+func applyDecisionForRetry(currentInput *[]*schema.Message, currentOpts *[]model.Option, ctx context.Context, decision *RetryDecision) {
 	if decision.ModifiedInputMessages != nil {
 		*currentInput = decision.ModifiedInputMessages
 		if decision.PersistModifiedInputMessages {

adk/runctx.go

@@ -87,8 +87,7 @@ type typedAgentEventWrapperForGob[M messageType] struct {
 func (e *typedAgentEventWrapper[M]) GobEncode() ([]byte, error) {
 	if e.event != nil && e.event.Output != nil && e.event.Output.MessageOutput != nil && e.event.Output.MessageOutput.IsStreaming {
 		// Materialize the stream before encoding.
-		var zero M
-		if any(e.concatenatedMessage) == any(zero) && e.StreamErr == nil {
+		if isNilMessage(e.concatenatedMessage) && e.StreamErr == nil {
 			e.consumeStream()
 		}
 	}
@@ -127,8 +126,7 @@ func (e *typedAgentEventWrapper[M]) consumeStream() {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
-	var zero M
-	if any(e.concatenatedMessage) != any(zero) {
+	if !isNilMessage(e.concatenatedMessage) {
 		return
 	}
 
@@ -396,7 +394,7 @@ type runContext struct {
 	RootInput *AgentInput
 	RunPath   []RunStep
 
-	TypedRootInput any
+	AgenticRootInput any
 
 	Session *runSession
 }
@@ -407,10 +405,10 @@ func (rc *runContext) isRoot() bool {
 
 func (rc *runContext) deepCopy() *runContext {
 	copied := &runContext{
-		RootInput:      rc.RootInput,
-		TypedRootInput: rc.TypedRootInput,
-		RunPath:        make([]RunStep, len(rc.RunPath)),
-		Session:        rc.Session,
+		RootInput:        rc.RootInput,
+		AgenticRootInput: rc.AgenticRootInput,
+		RunPath:          make([]RunStep, len(rc.RunPath)),
+		Session:          rc.Session,
 	}
 
 	copy(copied.RunPath, rc.RunPath)
@@ -462,7 +460,7 @@ func initTypedRunCtx[M messageType](ctx context.Context, agentName string, input
 		if _, ok := any(zero).(*schema.Message); ok {
 			runCtx.RootInput = any(input).(*AgentInput)
 		} else {
-			runCtx.TypedRootInput = input
+			runCtx.AgenticRootInput = input
 		}
 	}
 
@@ -601,7 +599,7 @@ func ctxWithNewTypedRunCtx[M messageType](ctx context.Context, input *TypedAgent
 	if _, ok := any(zero).(*schema.Message); ok {
 		rc.RootInput = any(input).(*AgentInput)
 	} else {
-		rc.TypedRootInput = input
+		rc.AgenticRootInput = input
 	}
 	return setRunCtx(ctx, rc)
 }

adk/utils.go

@@ -129,14 +129,14 @@ func getMessageFromTypedWrappedEvent[M messageType](e *typedAgentEventWrapper[M]
 		return e.event.Output.MessageOutput.Message, nil
 	}
 
-	if any(e.concatenatedMessage) != any(zero) {
-		return e.concatenatedMessage, nil
-	}
-
 	if e.StreamErr != nil {
 		return zero, e.StreamErr
 	}
 
+	if !isNilMessage(e.concatenatedMessage) {
+		return e.concatenatedMessage, nil
+	}
+
 	e.consumeStream()
 
 	if e.StreamErr != nil {