Add files via upload

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)

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