Skip to content

fix(plugins): drop per-invocation retry counts when a run ends - #7638

Open
Somay-kousis wants to merge 1 commit into
google:mainfrom
Somay-kousis:fix/r5-c2
Open

Somay-kousis wants to merge 1 commit into
google:mainfrom
Somay-kousis:fix/r5-c2

Conversation

@Somay-kousis

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

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

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

Problem:

ReflectAndRetryToolPlugin and ReflectAndRetryModelPlugin track failures per invocation by default. An invocation's entry was created on its first failure and never removed. A success only cleared the tool or model name inside it, so a long-running server kept one entry for every invocation that ever saw a failure.

Solution:

  • ScopedFailureTracker.clear_scope(scope_key) drops a whole scope under the tracker's lock.
  • Both plugins implement after_run_callback and on_run_error_callback. Under TrackingScope.INVOCATION they clear that invocation's scope. The runner calls one or the other at the end of every run, so the entry goes away whether the run finished, ran out of retries, or failed because the plugin re-raised the tool error.
  • TrackingScope.GLOBAL is untouched, so its counts still carry across runs.

Counting inside a run doesn't change.

One behavior note: a resumable invocation that pauses and is resumed later now starts its counts from zero in the resumed run, because the paused run ended. Happy to skip clearing for paused invocations instead if you prefer.

Testing Plan

Unit Tests:

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

New tests:

  • test_reflect_retry_tool_plugin.py, through InMemoryRunner: the tracker is empty after runs where the tool fails then succeeds, after runs where retries run out, and after a run that fails because the plugin re-raises. A GLOBAL scope test checks that its counts survive the run.
  • test_reflect_retry_model_plugin.py, through InMemoryRunner: the tracker is empty after runs with a malformed model response.
  • test_reflect_retry_utils.py: clear_scope drops only the named scope, and is a no-op for an unknown one.
main: 6 failed, 54 passed   (the GLOBAL guard passes on main, as it should)
fix:  60 passed

Removing each new cleanup call makes the matching test fail. tests/unittests/plugins and tests/unittests/runners pass (1514 passed, 2 skipped, 3 xfailed). isort, pyink and mypy are clean on the changed files.

Manual End-to-End (E2E) Tests:

The repro in the issue (App + Runner, scripted model, flaky tool):

before: after 200 invocations: plugin holds 200 invocation scopes
after:  after 200 invocations: plugin holds 0 invocation scopes

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.

ReflectAndRetryToolPlugin and ReflectAndRetryModelPlugin track failures
per invocation by default. The entry for an invocation was created on its
first failure and never removed: a success only cleared the tool or model
name inside it. A long-running server therefore kept one entry for every
invocation that ever saw a failure.

Add ScopedFailureTracker.clear_scope and have both plugins clear the
invocation's scope from after_run_callback and on_run_error_callback when
tracking per invocation. GLOBAL scope counts are kept.
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.

ReflectAndRetryToolPlugin and ReflectAndRetryModelPlugin keep failure counts for every finished invocation

2 participants