Handle recoverable tool argument errors in graceful exit resume

adk/cancel/graceful-exit/main.go

@@ -89,8 +89,11 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"maps"
 	"os"
 	"os/signal"
+	"slices"
+	"strings"
 	"syscall"
 	"time"
 
@@ -123,8 +126,32 @@ func sysMsgf(color, format string, args ...any) {
 	fmt.Printf("\n%s"+format+"%s\n", append([]any{color}, append(args, colorReset)...)...)
 }
 
+func messageTypeLabel(role schema.RoleType) string {
+	if role == schema.Tool {
+		return "tool result"
+	}
+	if role == "" {
+		return "message"
+	}
+	return string(role)
+}
+
+func printEventHeader(eventNum int, mode, agentName, role string) {
+	fmt.Printf("\n%s********event-%d(%s) [%s] [%s] *******%s\n", colorDim, eventNum, mode, agentName, role, colorReset)
+}
+
 const checkpointID = "graceful-exit-demo"
 
+// Eino local tool wrappers currently report argument/schema failures using these
+// stable error prefixes in components/tool/utils/{invokable,streamable}_func.go.
+const (
+	localFuncUnmarshalPrefix       = "[LocalFunc] failed to unmarshal arguments"
+	localFuncUnmarshalJSONPrefix   = "[LocalFunc] failed to unmarshal arguments in json"
+	localFuncInvalidTypePrefix     = "[LocalFunc] invalid type"
+	localStreamUnmarshalPrefix     = "[LocalStreamFunc] failed to unmarshal arguments"
+	localStreamUnmarshalJSONPrefix = "[LocalStreamFunc] failed to unmarshal arguments in json"
+)
+
 // --- Mock tools for the nested agent topology ---
 // The root agent uses search_web, then delegates to the analyst AgentTool.
 // The analyst sub-agent uses analyze_data and summarize_findings.
@@ -144,6 +171,14 @@ type summarizeInput struct {
 	Findings []string `json:"findings" jsonschema:"description=List of findings to summarize"`
 }
 
+type runStatus int
+
+const (
+	runCompleted runStatus = iota
+	runCanceled
+	runFailed
+)
+
 func searchWeb(_ context.Context, input *searchInput) (string, error) {
 	// Simulate a slow network call — gives the user time to press Ctrl-C.
 	time.Sleep(500 * time.Millisecond)
@@ -167,6 +202,55 @@ func summarizeFindings(_ context.Context, input *summarizeInput) (string, error)
 		"Recommended action: escalate to priority observation queue.", len(input.Findings)), nil
 }
 
+func shouldReturnToolArgumentError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	msg := err.Error()
+	return strings.Contains(msg, localFuncUnmarshalJSONPrefix) ||
+		strings.Contains(msg, localFuncUnmarshalPrefix) ||
+		strings.Contains(msg, localFuncInvalidTypePrefix) ||
+		strings.Contains(msg, localStreamUnmarshalJSONPrefix) ||
+		strings.Contains(msg, localStreamUnmarshalPrefix)
+}
+
+func toolArgumentErrorAsResult(ctx context.Context, in *compose.ToolInput, err error) string {
+	return fmt.Sprintf("tool execution failed for %q due to invalid arguments; please fix the arguments to satisfy the tool schema and try again: %v",
+		in.Name, err)
+}
+
+func sortedToolCallIndexes(toolCallArgs map[int]*strings.Builder) []int {
+	indexes := slices.Collect(maps.Keys(toolCallArgs))
+	slices.Sort(indexes)
+	return indexes
+}
+
+func returnArgumentErrorsAsToolResults() compose.ToolMiddleware {
+	return compose.ToolMiddleware{
+		Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
+			return func(ctx context.Context, in *compose.ToolInput) (*compose.ToolOutput, error) {
+				out, err := next(ctx, in)
+				if err == nil || !shouldReturnToolArgumentError(err) {
+					return out, err
+				}
+				return &compose.ToolOutput{Result: toolArgumentErrorAsResult(ctx, in, err)}, nil
+			}
+		},
+		Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
+			return func(ctx context.Context, in *compose.ToolInput) (*compose.StreamToolOutput, error) {
+				out, err := next(ctx, in)
+				if err == nil || !shouldReturnToolArgumentError(err) {
+					return out, err
+				}
+				return &compose.StreamToolOutput{
+					Result: schema.StreamReaderFromArray([]string{toolArgumentErrorAsResult(ctx, in, err)}),
+				}, nil
+			}
+		},
+	}
+}
+
 func main() {
 	ctx := context.Background()
 
@@ -201,13 +285,16 @@ func main() {
 			"1. Use analyze_data to analyze the content.\n" +
 			"2. Use summarize_findings to produce a final summary.\n" +
 			"Always use BOTH tools before giving your answer.",
-		Model: chatModel,
-		ToolsConfig: adk.ToolsConfig{
-			ToolsNodeConfig: compose.ToolsNodeConfig{
-				Tools: []einotool.BaseTool{analyzeTool, summarizeTool},
+			Model: chatModel,
+			ToolsConfig: adk.ToolsConfig{
+				ToolsNodeConfig: compose.ToolsNodeConfig{
+					Tools: []einotool.BaseTool{analyzeTool, summarizeTool},
+					ToolCallMiddlewares: []compose.ToolMiddleware{
+						returnArgumentErrorsAsToolResults(),
+					},
+				},
 			},
-		},
-	})
+		})
 	if err != nil {
 		log.Fatalf("create analyst agent: %v", err)
 	}
@@ -299,12 +386,17 @@ func main() {
 	}()
 
 	// Consume the event stream.
-	canceled := drainEvents(iter)
+	initialStatus := drainEvents(iter)
 
 	// Stop listening for signals — Phase 2 should not be interrupted.
 	signal.Stop(sigCh)
 
-	if !canceled {
+	if initialStatus == runFailed {
+		sysMsg(colorRed, "Phase 1 failed before cancellation or completion.")
+		return
+	}
+
+	if initialStatus != runCanceled {
 		sysMsg(colorGreen, "Agent completed without cancellation.")
 		return
 	}
@@ -334,20 +426,29 @@ func main() {
 		log.Fatalf("resume failed: %v", err)
 	}
 
-	drainEvents(resumeIter)
+	resumeStatus := drainEvents(resumeIter)
 
 	fmt.Println()
-	sysMsg(colorGreen, "✓ Done — agent resumed and completed successfully.")
+	if resumeStatus == runCompleted {
+		sysMsg(colorGreen, "✓ Done — agent resumed and completed successfully.")
+		return
+	}
+	if resumeStatus == runCanceled {
+		sysMsg(colorYellow, "Resume stopped by cancellation before completion.")
+		return
+	}
+	sysMsg(colorRed, "Resume failed before completion.")
 }
 
 // drainEvents consumes all events from the iterator, printing output and
-// detecting CancelError. Returns true if a CancelError was encountered.
-// Events are prefixed with the agent name to distinguish top-level vs nested output.
-func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) bool {
+// detecting terminal status.
+// Each visible event gets a numbered header showing stream mode, agent, and role.
+func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) runStatus {
+	eventNum := 0
 	for {
 		event, ok := iter.Next()
 		if !ok {
-			return false
+			return runCompleted
 		}
 
 		if event.Err != nil {
@@ -358,48 +459,82 @@ func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) bool {
 				sysMsgf(colorYellow, "    Mode:      %v", cancelErr.Info.Mode)
 				sysMsgf(colorYellow, "    Escalated: %v", cancelErr.Info.Escalated)
 				sysMsg(colorDim, "────────────────────────────────────────────────────────────")
-				return true
+				return runCanceled
 			}
 			log.Printf("unexpected error: %v", event.Err)
-			return false
+			return runFailed
 		}
 
-		// Determine display prefix based on which agent emitted the event.
-		prefix := fmt.Sprintf("[%s] ", event.AgentName)
-
 		// Print streamed/non-streamed message content.
 		if event.Output != nil && event.Output.MessageOutput != nil {
 			if s := event.Output.MessageOutput.MessageStream; s != nil {
-				first := true
+				var content strings.Builder
+				toolCallArgs := make(map[int]*strings.Builder)
+				toolCallNames := make(map[int]string)
+				contentRole := schema.Assistant
+				hasRole := false
 				for {
 					chunk, recvErr := s.Recv()
 					if recvErr != nil {
 						if recvErr == io.EOF {
 							break
 						}
-						// StreamCanceledError is expected when CancelImmediate fires.
-						break
+						var streamCanceledErr *adk.StreamCanceledError
+						if errors.As(recvErr, &streamCanceledErr) {
+							// StreamCanceledError is expected when CancelImmediate fires.
+							break
+						}
+						log.Printf("unexpected stream recv error: %v", recvErr)
+						return runFailed
 					}
-					if first {
-						fmt.Print(colorDim + prefix + colorReset + colorCyan)
-						first = false
+					if !hasRole {
+						contentRole = chunk.Role
+						hasRole = true
 					}
-					fmt.Print(chunk.Content)
+					for _, tc := range chunk.ToolCalls {
+						idx := 0
+						if tc.Index != nil {
+							idx = *tc.Index
+						}
+						if _, ok := toolCallArgs[idx]; !ok {
+							toolCallArgs[idx] = &strings.Builder{}
+						}
+						if tc.Function.Name != "" {
+							toolCallNames[idx] = tc.Function.Name
+						}
+						toolCallArgs[idx].WriteString(tc.Function.Arguments)
+					}
+					content.WriteString(chunk.Content)
+				}
+
+				if content.Len() == 0 && len(toolCallArgs) == 0 {
+					continue
+				}
+
+				eventNum++
+				printEventHeader(eventNum, "stream", event.AgentName, messageTypeLabel(contentRole))
+				if content.Len() > 0 {
+					fmt.Printf("%s%s%s\n", colorCyan, content.String(), colorReset)
 				}
-				if !first {
-					fmt.Print(colorReset + "\n")
+				for _, idx := range sortedToolCallIndexes(toolCallArgs) {
+					args, ok := toolCallArgs[idx]
+					if !ok {
+						continue
+					}
+					fmt.Printf("%s[tool call] %s(%s)%s\n", colorDim, toolCallNames[idx], args.String(), colorReset)
 				}
 			} else if m := event.Output.MessageOutput.Message; m != nil {
+				if m.Content == "" && len(m.ToolCalls) == 0 {
+					continue
+				}
+
+				eventNum++
+				printEventHeader(eventNum, "non-stream", event.AgentName, messageTypeLabel(m.Role))
 				if m.Content != "" {
-					fmt.Printf("%s%s%s%s%s%s\n", colorDim, prefix, colorReset, colorCyan, m.Content, colorReset)
+					fmt.Printf("%s%s%s\n", colorCyan, m.Content, colorReset)
 				}
-				// Show tool call invocations for visibility.
 				for _, tc := range m.ToolCalls {
-					sysMsgf(colorDim, "%s  → tool call: %s(%s)", prefix, tc.Function.Name, tc.Function.Arguments)
-				}
-				// Show tool results.
-				if m.Role == schema.Tool {
-					sysMsgf(colorDim, "%s  ← tool result: %.100s...", prefix, m.Content)
+					fmt.Printf("%s[tool call] %s(%s)%s\n", colorDim, tc.Function.Name, tc.Function.Arguments, colorReset)
 				}
 			}
 		}