Skip to content

Feat: add VertexAiCodeExecutor (Vertex AI Code Interpreter Extension) - #553

Open
AmaadMartin wants to merge 4 commits into
google:mainfrom
AmaadMartin:feat/vertex-ai-code-executor
Open

Feat: add VertexAiCodeExecutor (Vertex AI Code Interpreter Extension)#553
AmaadMartin wants to merge 4 commits into
google:mainfrom
AmaadMartin:feat/vertex-ai-code-executor

Conversation

@AmaadMartin

Copy link
Copy Markdown
Collaborator

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

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

  • No existing issue; see the description below.

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

Problem:
adk-python ships a VertexAiCodeExecutor that runs model-generated code on the
managed Vertex AI Code Interpreter Extension, but adk-js has no equivalent.
Users who want a managed, stateful, sandboxed Python execution backend
(returning stdout plus generated image/data files) have no option in the
TypeScript SDK. This is the last remaining code-executor parity gap (the
Container and GKE executors are tracked separately and are out of scope).

Solution:
Add VertexAiCodeExecutor
(core/src/code_executors/vertex_ai_code_executor.ts), a BaseCodeExecutor
that:

  • Resolves a Code Interpreter extension in the reference order — explicit
    resourceName -> CODE_INTERPRETER_EXTENSION_NAME env var -> create a new one
    from the public hub manifest
    (gs://vertex-extension-public/code_interpreter.yaml) and write the created
    resource name back to the env var for reuse.
  • Prepends the ported Python imported-library preamble to every code block, runs
    the extension's execute operation (operationId: 'execute', snake_case
    operationParams), and reuses the session_id across calls so state persists.
  • Maps the response to stdout/stderr and File[], applying the reference
    MIME rules: png/jpg/jpeg -> image/<ext> (faithful, non-normalized so
    .jpg -> image/jpg), csv -> text/csv, everything else -> a best-effort
    guessMimeType fallback.

Because @google-cloud/vertexai exposes no Extensions API surface, the default
client talks to the extension over its public REST endpoints
(<location>-aiplatform.googleapis.com/v1beta1) using GoogleAuth (ADC,
cloud-platform scope) + fetch, mirroring the approach already used by other
Vertex-backed adk-js services. No new package dependencies are added. The
extension client is injectable for hermetic testing, the class sets
stateful = true / optimizeDataFile = true (so the request processor
populates executionId, enabling session reuse), and it is marked
@experimental and exported from the core barrel.

Testing Plan

Unit Tests:

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

core/test/code_executors/vertex_ai_code_executor_test.ts covers the executor
with an injected mock client, and the default REST client with mocked fetch +
GoogleAuth, reaching 100% line and branch coverage of the new file: extension
resolution order (explicit / env var / auto-create + env write-back / concurrent
single import), stateful / optimizeDataFile parity, the
Project ID is required. guard, session-id forwarding and reuse, input-file
forwarding, the imported-library preamble, all MIME mapping rules, empty/missing
output and stdout/stderr defaults, the correct :execute URL and content
envelope parsing, LRO import polling, and error paths (non-2xx, invalid resource
name, import timeout).

Run locally:

npm run build
npx vitest run --project unit:core core/test/code_executors/vertex_ai_code_executor_test.ts

Result:

 ✓ |unit:core| core/test/code_executors/vertex_ai_code_executor_test.ts (33 tests)

 Test Files  1 passed (1)
      Tests  33 passed (33)

npm run lint, npm run format:check, and npm run docs:check are also clean
for the touched files.

Manual End-to-End (E2E) Tests:

An env-gated live test lives at
tests/e2e/code_executors/vertex_ai_code_executor_e2e_test.ts. It is skipped by
default so CI stays hermetic, and only runs when explicitly opted in. It runs
Python twice in one session and asserts stdout, a generated .png (with
image/png MIME type), and persisted session state (a variable defined in the
first call is used in the second).

gcloud auth application-default login
export GOOGLE_CLOUD_PROJECT=<your-project>
# optional: export CODE_INTERPRETER_EXTENSION_NAME=<existing-extension>
RUN_VERTEX_AI_CODE_EXECUTOR_E2E=true npx vitest run --project e2e tests/e2e/code_executors

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.

Additional context

The live E2E test above is provided but left unchecked: it requires a real
Google Cloud project with the Vertex AI Code Interpreter Extension available, so
it is opt-in rather than part of the default suite. No dependent changes are
required — this PR is additive and introduces no new package dependencies.

Amaad Martin added 2 commits July 28, 2026 13:55
Port the adk-python VertexAiCodeExecutor to adk-js. It runs model-generated
Python code on the managed Vertex AI Code Interpreter Extension, reusing a
session id across calls so state persists, and maps the extension response to
stdout/stderr plus output File[] (png/jpg/jpeg images, csv data files, and a
best-effort MIME guess for anything else).

Since @google-cloud/vertexai exposes no Extensions surface, the default client
talks to the extension over its public REST endpoints using GoogleAuth (ADC) +
fetch, matching the approach used by other Vertex-backed adk-js services. The
extension client is injectable for testing. Marked @experimental and exported
from the core barrel.
Add unit tests (injected client + mocked fetch/GoogleAuth for the default REST
client) reaching 100% line and branch coverage of the new executor: extension
resolution order, session reuse, MIME mapping rules, output/stdout/stderr
defaults, LRO import polling, and error paths. Add an env-gated live e2e test
(opt-in) that runs Python twice in one session and asserts stdout, a png
output file, and persisted session state.
* the adk-python `VertexAiCodeExecutor` (the Code Interpreter runs Python, so
* the preamble stays Python).
*/
const IMPORTED_LIBRARIES = `

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

We should support TS as well.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Thanks — this surfaced a real bug, though not one I can fix by adding a TS
runtime here.

The managed Code Interpreter extension cannot run TypeScript. It is provisioned
from the public gs://vertex-extension-public/code_interpreter.yaml manifest,
which exposes a Python runtime and nothing else, so the language set is a
property of the Vertex service rather than a choice this executor makes. That is
also why the preamble on these lines is Python (pandas, numpy, matplotlib,
scipy) and why adk-python's VertexAiCodeExecutor hardcodes Python too. There
is no request field to select a language.

The actual defect you pointed at: executeCode was ignoring
codeExecutionInput.language entirely, so TypeScript would have been wrapped in
the Python preamble and shipped to the Python interpreter, coming back as an
opaque SyntaxError with errorRetryAttempts = 2 feeding it to the model twice
first. Fixed in f7903b6a:

    // Report an unrunnable language the way UnsafeLocalCodeExecutor does — as
    // stderr rather than a throw — so the model can retry in Python instead of
    // seeing a Python SyntaxError for, say, TypeScript.
    if (codeExecutionInput.language !== CodeExecutionLanguage.PYTHON) {
      return {
        stdout: '',
        stderr: `Unsupported language: ${codeExecutionInput.language}. The Vertex AI Code Interpreter extension only runs ${CodeExecutionLanguage.PYTHON}.`,
        outputFiles: [],
      };
    }

Returning it on stderr rather than throwing matches
UnsafeLocalCodeExecutor.executeCode (unsafe_local_code_executor.ts:128-143),
which keeps the failure model-recoverable. No behavior change in the normal
flow: code_execution_request_processor.ts hardcodes
language: CodeExecutionLanguage.PYTHON at both call sites, so this only fires
for direct programmatic use. Covered by a parameterized test over
TYPESCRIPT/JAVASCRIPT/SHELL/UNSPECIFIED asserting the extension is never
called.

The class doc now points elsewhere for other languages, since the JS/TS path
already exists in the repo:

Python is the only language this executor can run [...] Use
AgentEngineSandboxCodeExecutor for JavaScript, or UnsafeLocalCodeExecutor
for the shell family.

AgentEngineSandboxCodeExecutor maps CodeExecutionLanguage.JAVASCRIPT to
Language.LANGUAGE_JAVASCRIPT (agent_engine_sandbox_code_executor.ts:76-85),
so TS/JS execution belongs there. Happy to open a follow-up adding TypeScript to
that executor's mapLanguage if the Agent Engine sandbox supports it — that
would be a change to a different file than this PR touches. Let me know if you
would rather I do it here.

Amaad Martin added 2 commits July 28, 2026 19:40
- Fail the extension import when the LRO completes with an error instead of
  dereferencing an absent `response`; model `error` on the operation type.
- Validate the configured resource name (option or
  CODE_INTERPRETER_EXTENSION_NAME) in the constructor and derive the endpoint
  region from it, so a malformed name fails fast instead of mid agent-run.
- Require a project ID whether or not a client is injected, and drop the
  `this.projectId!` non-null assertion in getResourceName().
- Correct the constructor comment: setting `stateful`/`optimizeDataFile` is a
  deliberate divergence from adk-python, not parity.
- Document that `image/jpg` is intentionally non-normalized.
- Drop the unused `generated_code` field from ExtensionExecuteResponse.
- Tests: type the mock as the real CodeInterpreterExtensionClient (removes the
  `as never` and `as string` casts), drop the dead `withClient` option, assert
  the endpoint URL instead of reaching into the private `location` field,
  restore fake timers in afterEach, and cover the failed-import branches.
The managed Code Interpreter extension is provisioned from the public
code_interpreter.yaml manifest and exposes a Python runtime only, but
executeCode() ignored codeExecutionInput.language and prepended the Python
preamble to whatever it was given, so a TypeScript or shell snippet reached
the Python interpreter and came back as an opaque SyntaxError.

Report it as `Unsupported language: <lang>` on stderr, matching
UnsafeLocalCodeExecutor, so the model can retry in a supported language, and
document the Python-only constraint on the class.
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.

2 participants