diff --git a/internal/multiagent/model_output_guard_middleware.go b/internal/multiagent/model_output_guard_middleware.go index 5962166f..5b56b4e1 100644 --- a/internal/multiagent/model_output_guard_middleware.go +++ b/internal/multiagent/model_output_guard_middleware.go @@ -96,7 +96,17 @@ func (m *modelOutputGuardMiddleware) AfterModelRewriteState( badIndex := -1 argumentBytes := 0 if strings.EqualFold(strings.TrimSpace(finishReason), "length") { - reason = "output_limit" + if len(last.ToolCalls) == 0 { + reason = "output_limit" + } else { + for i, tc := range last.ToolCalls { + r, n := validateGeneratedToolCall(tc, m.cfg) + if r != "" { + reason, badIndex, argumentBytes = "output_limit", i, n + break + } + } + } } else { for i, tc := range last.ToolCalls { r, n := validateGeneratedToolCall(tc, m.cfg) diff --git a/internal/multiagent/model_output_guard_middleware_test.go b/internal/multiagent/model_output_guard_middleware_test.go index 609199f4..43a08997 100644 --- a/internal/multiagent/model_output_guard_middleware_test.go +++ b/internal/multiagent/model_output_guard_middleware_test.go @@ -50,6 +50,18 @@ func TestModelOutputGuardRejectsTruncatedToolCallBeforeExecution(t *testing.T) { } } +func TestModelOutputGuardAllowsValidToolCallDespiteLengthFinish(t *testing.T) { + original := `{"command":"echo ok"}` + state, err := runModelOutputGuard(t, []adk.Message{schema.UserMessage("run"), guardedAssistant(original, "length")}, config.MultiAgentEinoMiddlewareConfig{}) + if err != nil { + t.Fatal(err) + } + got := state.Messages[len(state.Messages)-1].ToolCalls[0].Function.Arguments + if got != original { + t.Fatalf("valid arguments should pass unchanged: %q", got) + } +} + func TestModelOutputGuardRejectsInvalidJSONShapes(t *testing.T) { for _, arguments := range []string{"", `[]`, `{"command":`} { t.Run(arguments, func(t *testing.T) {