Skip to content

fix(workflow): give each retry of a dynamic function node fresh arguments - #7556

Open
arjunprabhulal wants to merge 2 commits into
google:mainfrom
arjunprabhulal:fix/function-node-retry-fresh-args
Open

arjunprabhulal wants to merge 2 commits into
google:mainfrom
arjunprabhulal:fix/function-node-retry-fresh-args

Conversation

@arjunprabhulal

Copy link
Copy Markdown

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

No existing issue. The bug is described below.

2. Or, if no issue exists, describe the change:

Problem:

When a FunctionNode that is run through ctx.run_node or a NodeTool is retried, the retry receives the same argument objects as the attempt that failed. If the function changed a list or dict argument before raising, the retry starts from the changed value.

Steps to reproduce:

  1. A FunctionNode with parameter_binding='node_input' and retry_config wraps add_item(items: list[int]), which appends 99 to items and raises on its first call.
  2. A parent runs it with ctx.run_node(add_item_node, node_input={'items': [1, 2]}).

Expected: both attempts receive [1, 2].

Observed: the second attempt receives [1, 2, 99].

This is a regression from 156f007. The scheduler and NodeTool resolve the arguments before the node runs and pass the resolved _ResolvedParameters in as node_input. NodeRunner hands that same object to every attempt, and _bind_parameters now uses it as is with a shallow dict() copy, so every attempt shares the coerced lists and dicts. Before that change each attempt coerced again and got new containers. Arguments that come in over a static graph edge are not affected, because they arrive as a plain dict and are resolved on every attempt.

Solution:

_ResolvedParameters now remembers the input it was resolved from. The first bind uses the resolved values directly, so the normal path still coerces only once, which test_node_input_binding_coerces_parameters_once pins. Any later bind of the same object, which only happens on a retry, resolves again from the original input and gets its own argument objects.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

New test test_dynamic_function_node_retry_gets_fresh_arguments in tests/unittests/workflow/test_node_runner_failure.py. The existing coerce once tests in test_function_node.py still pass.

Without the fix:

$ pytest tests/unittests/workflow/test_node_runner_failure.py -k fresh_arguments
E     assert [[1, 2], [1, 2, 99]] == [[1, 2], [1, 2]]
1 failed, 24 deselected

With the fix:

$ pytest tests/unittests/workflow/test_node_runner_failure.py tests/unittests/workflow/test_function_node.py -k "fresh_arguments or coerces"
7 passed, 98 deselected

$ pytest tests/unittests/workflow tests/unittests/agents tests/unittests/tools tests/unittests/flows tests/unittests/runners
6121 passed, 2 skipped, 6 xfailed

The same subset passes on Python 3.11 and 3.13. The full unit suite (pytest tests/unittests) shows the same single failure as an unmodified checkout of main, test_reading_directory_raises in tests/unittests/integrations/vmaas, which fails there without this change.

I also ran the same function as a NodeTool called by an LlmAgent. Before the fix the retry received [1, 2, 99], and after it the retry receives [1, 2].

pre-commit run --files passes on both changed files.

Manual End-to-End (E2E) Tests:

A script runs the steps above through Runner with InMemorySessionService.

def add_item(items: list[int]) -> list[int]:
  global attempt
  attempt += 1
  print(f"  attempt {attempt} received items={items}")
  items.append(99)
  if attempt == 1:
    raise RuntimeError("failed after appending")
  return items


add_item_node = FunctionNode(
    func=add_item,
    parameter_binding="node_input",
    retry_config=RetryConfig(max_attempts=2, initial_delay=0.0, jitter=0.0),
)


@node(rerun_on_resume=True)
async def parent(ctx: Context):
  return await ctx.run_node(add_item_node, node_input={"items": [1, 2]})

Before:

  attempt 1 received items=[1, 2]
  attempt 2 received items=[1, 2, 99]
  workflow result: [1, 2, 99, 99]

After:

  attempt 1 received items=[1, 2]
  attempt 2 received items=[1, 2]
  workflow result: [1, 2, 99]

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

…ents

When a FunctionNode run through ctx.run_node or a NodeTool was retried,
every attempt received the same argument objects. If an attempt changed
a list or dict argument and then failed, the retry started from the
changed value. The arguments are resolved once before the node runs and
that resolved input was reused as is by each attempt.

The resolved input now remembers what it was resolved from. The first
attempt still uses it directly, so arguments are coerced only once on
the normal path, and any later attempt resolves again from the original
input. A new test changes a list argument, fails once, and checks that
the retry sees the original list.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request ensures that retried dynamic function nodes receive fresh arguments instead of mutated ones from previous attempts by storing the original source input and a bound flag in _ResolvedParameters. A unit test has been added to verify this behavior. The feedback suggests adding class-level type annotations for source and bound in _ResolvedParameters to improve type safety and static analysis.

Comment thread src/google/adk/workflow/_function_node.py
Declare source and bound at class level so type checkers and IDEs see
them on every _ResolvedParameters instance.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant