Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -532,7 +532,7 @@ describe('ensureModelSecret', () => {
expect(secretManager.create).not.toHaveBeenCalled();
});

test('skips when credentials map has multiple entries (Vertex AI)', async () => {
test('applies Vertex AI workspace configuration instead of creating a secret', async () => {
vi.mocked(providerRegistry.getInferenceConnectionCredentials).mockReturnValue({
credentials: { projectId: 'my-project', region: 'us-east5', credentialsFile: '/path/to/creds.json' },
llmMetadataName: 'vertexai',
Expand All @@ -543,6 +543,83 @@ describe('ensureModelSecret', () => {
await manager.ensureModelSecret(options);

expect(secretManager.create).not.toHaveBeenCalled();
expect(options.workspaceConfiguration?.environment).toEqual([
{ name: 'CLAUDE_CODE_USE_VERTEX', value: '1' },
{ name: 'CLOUD_ML_REGION', value: 'us-east5' },
{ name: 'ANTHROPIC_VERTEX_PROJECT_ID', value: 'my-project' },
]);
expect(options.workspaceConfiguration?.mounts).toEqual([
{ host: '/path/to/creds.json', target: '$HOME/.config/gcloud/application_default_credentials.json', ro: true },
]);
});

test('deduplicates Vertex AI env vars when workspaceConfiguration already has entries', async () => {
vi.mocked(providerRegistry.getInferenceConnectionCredentials).mockReturnValue({
credentials: { projectId: 'my-project', region: 'us-east5', credentialsFile: '/path/to/creds.json' },
llmMetadataName: 'vertexai',
endpoint: undefined,
});

const options = {
...baseOptions,
model: 'vertexai::claude-sonnet-4-20250514::',
workspaceConfiguration: {
environment: [
{ name: 'CLOUD_ML_REGION', value: 'old-region' },
{ name: 'OTHER_VAR', value: 'keep' },
],
mounts: [
{ host: '/old/creds.json', target: '$HOME/.config/gcloud/application_default_credentials.json', ro: true },
],
},
} as AgentWorkspaceCreateOptions;
await manager.ensureModelSecret(options);

expect(options.workspaceConfiguration?.environment).toEqual([
{ name: 'OTHER_VAR', value: 'keep' },
{ name: 'CLAUDE_CODE_USE_VERTEX', value: '1' },
{ name: 'CLOUD_ML_REGION', value: 'us-east5' },
{ name: 'ANTHROPIC_VERTEX_PROJECT_ID', value: 'my-project' },
]);
expect(options.workspaceConfiguration?.mounts).toEqual([
{ host: '/path/to/creds.json', target: '$HOME/.config/gcloud/application_default_credentials.json', ro: true },
]);
});

test('replaces tilde with $HOME in Vertex AI credentials file path', async () => {
vi.mocked(providerRegistry.getInferenceConnectionCredentials).mockReturnValue({
credentials: {
projectId: 'my-project',
region: 'us-east5',
credentialsFile: '~/.config/gcloud/application_default_credentials.json',
},
llmMetadataName: 'vertexai',
endpoint: undefined,
});

const options = { ...baseOptions, model: 'vertexai::claude-sonnet-4-20250514::' };
await manager.ensureModelSecret(options);

expect(options.workspaceConfiguration?.mounts).toEqual([
{
host: '$HOME/.config/gcloud/application_default_credentials.json',
target: '$HOME/.config/gcloud/application_default_credentials.json',
ro: true,
},
]);
});

test('skips Vertex AI config when credentials are incomplete', async () => {
vi.mocked(providerRegistry.getInferenceConnectionCredentials).mockReturnValue({
credentials: { projectId: 'my-project', region: '', credentialsFile: '' },
llmMetadataName: 'vertexai',
endpoint: undefined,
});

const options = { ...baseOptions, model: 'vertexai::claude-sonnet-4-20250514::' };
await manager.ensureModelSecret(options);

expect(options.workspaceConfiguration).toBeUndefined();
});

test('skips when llmMetadataName is unknown', async () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,11 @@ export class AgentWorkspaceManager implements Disposable {
const connectionInfo = this.providerRegistry.getInferenceConnectionCredentials(options.model);
if (!connectionInfo) return;

if (connectionInfo.llmMetadataName === 'vertexai') {
this.applyVertexAiConfiguration(options, connectionInfo.credentials);
return;
}

const entries = Object.entries(connectionInfo.credentials);
if (entries.length !== 1) return;

Expand All @@ -163,6 +168,32 @@ export class AgentWorkspaceManager implements Disposable {
}
}

private applyVertexAiConfiguration(options: AgentWorkspaceCreateOptions, credentials: Record<string, string>): void {
const { projectId, region, credentialsFile } = credentials;
if (!projectId || !region || !credentialsFile) return;

options.workspaceConfiguration ??= {};
options.workspaceConfiguration.environment ??= [];
options.workspaceConfiguration.mounts ??= [];

const envVars: Array<{ name: string; value: string }> = [
{ name: 'CLAUDE_CODE_USE_VERTEX', value: '1' },
{ name: 'CLOUD_ML_REGION', value: region },
{ name: 'ANTHROPIC_VERTEX_PROJECT_ID', value: projectId },
];
for (const env of envVars) {
options.workspaceConfiguration.environment = options.workspaceConfiguration.environment.filter(
e => e.name !== env.name,
);
options.workspaceConfiguration.environment.push(env);
}

const adcTarget = '$HOME/.config/gcloud/application_default_credentials.json';
const hostPath = credentialsFile.startsWith('~/') ? `$HOME/${credentialsFile.slice(2)}` : credentialsFile;
options.workspaceConfiguration.mounts = options.workspaceConfiguration.mounts.filter(m => m.target !== adcTarget);
options.workspaceConfiguration.mounts.push({ host: hostPath, target: adcTarget, ro: true });
}

/**
* Maps provider metadata to the kdn secret create options.
*
Expand Down
Loading