Skip to content

fix(vertexai): preserve ExecutableCode and CodeExecutionResult in session storage - #1729

Merged
foxfrikses merged 5 commits into
google:mainfrom
karthik-0306:fix-vertexai-code-execution
Oct 7, 2026
Merged

foxfrikses merged 5 commits into
google:mainfrom
karthik-0306:fix-vertexai-code-execution

Conversation

@karthik-0306

Copy link
Copy Markdown
Contributor

Linked issue

Closes: Vertex AI Session Storage drops ExecutableCode and CodeExecutionResult parts #1728

2. Describe the change:

Problem:
When storing and retrieving session history using the Vertex AI backend (session/vertexai), the aiplatformToGenaiContent and createAiplatformpbContent mappers silently drop ExecutableCode, CodeExecutionResult, and FileData parts. There's a TODO in both functions acknowledging this omission. When an agent uses Gemini's Code Execution tool, the session history successfully saves the model's text thought but completely loses the code block and its result, causing the model to hallucinate on subsequent turns because it lacks the execution context it expects.

Solution:
Added mapping logic for Part_ExecutableCode, Part_CodeExecutionResult, and Part_FileData in both aiplatformToGenaiContent and createAiplatformpbContent to ensure these fields are preserved round-trip in session storage.

Behavior change

Agents utilizing Vertex AI for session storage will now correctly preserve code execution history and file data references across conversation turns. Before, these parts were silently dropped during Protobuf conversion.

Testing Plan

Unit Tests:
All unit tests pass locally.
With your source change reverted and your tests kept, which test fails?
go test ./session/vertexai/ -run TestAiplatformToGenaiContent_OtherParts

--- FAIL: TestAiplatformToGenaiContent_OtherParts (0.00s)
    vertexai_test.go:1016: incorrect ExecutableCode: <nil>
FAIL

With the change, go test -race ./... passes, and gofmt / golangci-lint are clean.

Manual End-to-End (E2E) Tests:
Ran an agent with the Code Execution tool enabled using Vertex AI for session storage.
before: The session event stored in Vertex AI lacked the code snippet, causing the model to immediately repeat itself on turn 2.
after: The session event in Vertex AI successfully stored and retrieved the executed code and outcome, allowing a multi-turn code-generation conversation.

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.

@google-cla

google-cla Bot commented Oct 5, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@foxfrikses
foxfrikses self-requested a review October 7, 2026 08:48

@foxfrikses foxfrikses 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.

Thank you for this fix. I think it's good. I'm planning on approving and merging this PR in after you address the one comment I left.

Comment thread session/vertexai/vertexai_test.go Resolved
@karthik-0306

Copy link
Copy Markdown
Contributor Author

Thank you for this fix. I think it's good. I'm planning on approving and merging this PR in after you address the one comment I left.

Thank you! The feedback has been addressed and the new commit is pushed. Let me know if everything looks good to merge now.

@foxfrikses foxfrikses 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.

Good. Thank you for the contribution!

@foxfrikses
foxfrikses merged commit 44d41c9 into google:main Oct 7, 2026
19 of 20 checks passed
wolo-lab added a commit that referenced this pull request Oct 14, 2026
…#1756) (#1832)

* fix(session/inmemory): holding sess.mu accross the eventCopy creation in AppendEvent (#1575)

(cherry picked from commit b8fecdf)

* fix(vertexai): preserve ExecutableCode and CodeExecutionResult in session storage (#1729)

(cherry picked from commit 44d41c9)

* fix(memory): normalize punctuation in in-memory search (#1565)

(cherry picked from commit 0d23d2e)

* fix(agent): synchronize concurrent artifact saves (#1756)

(cherry picked from commit ca91ac0)
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.

Bug: Vertex AI drops executed code from history

2 participants