Skip to content

fix(artifacts): load artifacts removed by a rewind as None in file and GCS services - #7483

Open
spandankeche wants to merge 1 commit into
google:mainfrom
spandankeche:fix/rewind-artifact-tombstone
Open

spandankeche wants to merge 1 commit into
google:mainfrom
spandankeche:fix/rewind-artifact-tombstone

Conversation

@spandankeche

Copy link
Copy Markdown

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:
Runner.rewind_async removes an artifact that did not exist at the rewind point by saving an empty application/octet-stream part as a new version. Only InMemoryArtifactService loaded that part as None. FileArtifactService, the default store for adk web and adk api_server, and GcsArtifactService returned the empty part, so after a rewind load_artifact returned a 0-byte artifact, the artifact endpoint answered 200 instead of 404, and load_artifacts told the model the artifact existed and was empty.

Solution:

  • Define the rewind marker once in artifact_util (_new_rewind_tombstone, _is_rewind_tombstone, both private). _rewind_utils saves it through the helper, and InMemoryArtifactService recognizes it through the helper, with no change in behavior.
  • FileArtifactService and GcsArtifactService load a version holding the marker as None, the same as InMemoryArtifactService.
  • The marker itself is unchanged, so sessions rewound by earlier releases, which already store it, are read correctly.

Not changed, to keep this to one concern: all three services still list the removed artifact in list_artifact_keys after a rewind (noted in #7482).

Testing Plan

Unit Tests:

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

New tests:

  • test_artifact_service.py::test_artifact_removed_by_rewind_loads_as_none: on all three services, a version holding the marker loads as None, by latest and by explicit version.
  • test_artifact_service.py::test_version_before_rewind_tombstone_still_loads: on all three services, the version saved before the marker still loads its content.
  • test_runner_rewind.py::test_rewind_removes_artifact_created_after_rewind_point: a Runner.rewind_async round trip on the in-memory and file services.

The tests catch the bug: with only the file and GCS change reverted, the 4 file and GCS cases of the first test and the file case of the runner test fail; with only the in-memory check removed, its 3 in-memory cases fail.

  • New tests: 11 passed (3 tests, parametrized over the artifact services and the version argument).
  • Full suite, pytest tests/unittests, installed the way the CI unit-test job does (uv sync --extra test, UV_EXCLUDE_NEWER one hour back), on main at 3d11f9fe7 plus this change: 18,367 / 18,355 / 18,357 / 18,356 passed on Python 3.11 / 3.12 / 3.13 / 3.14. The 4 to 7 failures per version are unrelated to this change; each was rerun three times on this branch and three times on untouched main (see Additional context).
  • pre-commit run --all-files: 14/14 hooks passed.
  • Mypy Check replica (mypy ., base vs. this change, comm -13) on Python 3.11, 3.12 and 3.13: 831 errors on both sides, 0 new.

Manual End-to-End (E2E) Tests:

Setup: wheel built from this branch with uv build and installed into a fresh virtualenv (ADK imported from site-packages), Python 3.11. The script is the one from #7482: a Runner on InMemorySessionService, a session where inv2 first saves report.txt, then rewind_async(rewind_before_invocation_id="inv2") and load_artifact("report.txt").

Before (google-adk 2.11.0):

$ python repro.py
memory  load_artifact after rewind: None
file    load_artifact after rewind: Part(
  inline_data=Blob(
    data=b'',
    mime_type='application/octet-stream'
  )
)

After (wheel from this branch):

$ python repro.py
memory  load_artifact after rewind: None
file    load_artifact after rewind: None

The same steps with GcsArtifactService (on the test suite's MockClient) also print None. With an LlmAgent that saves the artifact from a tool and later calls load_artifacts (scripted model), the model previously received Artifact report.txt is: followed by an empty text part on the file service and on adk web's per-agent file storage; with this change it receives no artifact content, the same as with the in-memory service.

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. (There are none.)

Additional context

  • Docs: no change needed. The adk-docs rewind page already says rewind restores session artifacts to their condition before the rewind point, and this change makes the file and GCS services do that. No guide in docs/guides/ covers rewind.
  • Local full-suite failures unrelated to this change. Each failing test was rerun three times on this branch and three times on untouched main (3d11f9fe7):
    • The 4 TestParseToolCallArguments log tests in models/test_litellm.py fail only inside the full run and pass in every rerun on both branches.
    • test_import_loading.py::test_entry_point_loads_only_allowlisted_packages[agent] and [runner] fail on Python 3.12 on both branches with identical messages: Ubuntu's system Python 3.12 provides a sitecustomize module that the check reports.
    • Golden test_telemetry_schema cases in telemetry/test_functional.py fail intermittently on both branches (1 of 9 reruns here, 2 of 9 on main).
    • test_bigquery_agent_analytics_plugin.py::TestSafetyLifecycleHardening::test_completed_constructor_close_dispatched_off_loop failed once on 3.14 under full-suite load and passed in all six reruns.

Runner.rewind_async removes an artifact that did not exist at the
rewind point by saving an empty application/octet-stream part as a new
version. Only InMemoryArtifactService loaded that part as None.
FileArtifactService, the default store for adk web and adk api_server,
and GcsArtifactService returned it as a 0-byte artifact, so the
artifact endpoint answered 200 instead of 404 and load_artifacts told
the model the artifact existed and was empty.

Define the marker once in artifact_util, save it through that helper
in the rewind code, and have the file and GCS services load it as None
like the in-memory service. The marker itself is unchanged, so sessions
rewound by earlier releases are read correctly.

Fixes google#7482
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.

After Runner.rewind_async, FileArtifactService and GcsArtifactService load a removed artifact as an empty file instead of None

2 participants