From 3d8787200c648124e666552cb9f40e8ff99b6977 Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Mon, 8 Jun 2026 13:45:37 +0800 Subject: [PATCH 1/2] Handle tool argument errors in graceful exit resume --- adk/cancel/graceful-exit/main.go | 176 +++++++++++++++++++++++++------ 1 file changed, 143 insertions(+), 33 deletions(-) diff --git a/adk/cancel/graceful-exit/main.go b/adk/cancel/graceful-exit/main.go index e2d4cfc0..165a6f27 100644 --- a/adk/cancel/graceful-exit/main.go +++ b/adk/cancel/graceful-exit/main.go @@ -91,6 +91,7 @@ import ( "log" "os" "os/signal" + "strings" "syscall" "time" @@ -123,6 +124,20 @@ 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" // --- Mock tools for the nested agent topology --- @@ -144,6 +159,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 +190,47 @@ 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, "failed to unmarshal arguments") || + strings.Contains(msg, "failed to unmarshal arguments in json") || + strings.Contains(msg, "invalid type") +} + +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 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 +265,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 +366,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 +406,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,19 +439,20 @@ 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 { @@ -380,26 +462,54 @@ func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) bool { // StreamCanceledError is expected when CancelImmediate fires. break } - 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 } - if !first { - fmt.Print(colorReset + "\n") + + eventNum++ + printEventHeader(eventNum, "stream", event.AgentName, messageTypeLabel(contentRole)) + if content.Len() > 0 { + fmt.Printf("%s%s%s\n", colorCyan, content.String(), colorReset) + } + for idx := 0; idx < len(toolCallArgs); idx++ { + 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) } } } From 56a47a3ed3b0a3bf5b8b0ef1d083d45a7fd451b6 Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Mon, 8 Jun 2026 14:05:04 +0800 Subject: [PATCH 2/2] Address review feedback on graceful exit logging --- adk/cancel/graceful-exit/main.go | 37 ++++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/adk/cancel/graceful-exit/main.go b/adk/cancel/graceful-exit/main.go index 165a6f27..de396345 100644 --- a/adk/cancel/graceful-exit/main.go +++ b/adk/cancel/graceful-exit/main.go @@ -89,8 +89,10 @@ import ( "fmt" "io" "log" + "maps" "os" "os/signal" + "slices" "strings" "syscall" "time" @@ -140,6 +142,16 @@ func printEventHeader(eventNum int, mode, agentName, role string) { 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. @@ -196,9 +208,11 @@ func shouldReturnToolArgumentError(err error) bool { } msg := err.Error() - return strings.Contains(msg, "failed to unmarshal arguments") || - strings.Contains(msg, "failed to unmarshal arguments in json") || - strings.Contains(msg, "invalid type") + 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 { @@ -206,6 +220,12 @@ func toolArgumentErrorAsResult(ctx context.Context, in *compose.ToolInput, err e 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 { @@ -459,8 +479,13 @@ func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) runStatus { 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 !hasRole { contentRole = chunk.Role @@ -491,7 +516,7 @@ func drainEvents(iter *adk.AsyncIterator[*adk.AgentEvent]) runStatus { if content.Len() > 0 { fmt.Printf("%s%s%s\n", colorCyan, content.String(), colorReset) } - for idx := 0; idx < len(toolCallArgs); idx++ { + for _, idx := range sortedToolCallIndexes(toolCallArgs) { args, ok := toolCallArgs[idx] if !ok { continue