Feat: add VertexAiCodeExecutor (Vertex AI Code Interpreter Extension) - #553
Feat: add VertexAiCodeExecutor (Vertex AI Code Interpreter Extension)#553AmaadMartin wants to merge 4 commits into
Conversation
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 = ` |
There was a problem hiding this comment.
We should support TS as well.
There was a problem hiding this comment.
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
AgentEngineSandboxCodeExecutorfor JavaScript, orUnsafeLocalCodeExecutor
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.
- 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.
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):
2. Or, if no issue exists, describe the change:
Problem:
adk-pythonships aVertexAiCodeExecutorthat runs model-generated code on themanaged Vertex AI Code Interpreter Extension, but
adk-jshas 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), aBaseCodeExecutorthat:
resourceName->CODE_INTERPRETER_EXTENSION_NAMEenv var -> create a new onefrom the public hub manifest
(
gs://vertex-extension-public/code_interpreter.yaml) and write the createdresource name back to the env var for reuse.
the extension's
executeoperation (operationId: 'execute', snake_caseoperationParams), and reuses thesession_idacross calls so state persists.stdout/stderrandFile[], applying the referenceMIME rules:
png/jpg/jpeg->image/<ext>(faithful, non-normalized so.jpg->image/jpg),csv->text/csv, everything else -> a best-effortguessMimeTypefallback.Because
@google-cloud/vertexaiexposes no Extensions API surface, the defaultclient talks to the extension over its public REST endpoints
(
<location>-aiplatform.googleapis.com/v1beta1) usingGoogleAuth(ADC,cloud-platformscope) +fetch, mirroring the approach already used by otherVertex-backed
adk-jsservices. No new package dependencies are added. Theextension client is injectable for hermetic testing, the class sets
stateful = true/optimizeDataFile = true(so the request processorpopulates
executionId, enabling session reuse), and it is marked@experimentaland exported from the core barrel.Testing Plan
Unit Tests:
core/test/code_executors/vertex_ai_code_executor_test.tscovers the executorwith an injected mock client, and the default REST client with mocked
fetch+GoogleAuth, reaching 100% line and branch coverage of the new file: extensionresolution order (explicit / env var / auto-create + env write-back / concurrent
single import),
stateful/optimizeDataFileparity, theProject ID is required.guard, session-id forwarding and reuse, input-fileforwarding, the imported-library preamble, all MIME mapping rules, empty/missing
output and
stdout/stderrdefaults, the correct:executeURL andcontentenvelope parsing, LRO import polling, and error paths (non-2xx, invalid resource
name, import timeout).
Run locally:
Result:
npm run lint,npm run format:check, andnpm run docs:checkare also cleanfor 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 bydefault so CI stays hermetic, and only runs when explicitly opted in. It runs
Python twice in one session and asserts stdout, a generated
.png(withimage/pngMIME type), and persisted session state (a variable defined in thefirst call is used in the second).
Checklist
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.