jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix CODE_ASSIST_ENDPOINT env var. (#2712)

This commit is contained in:
Tommaso Sciortino 2025-06-30 15:41:14 -07:00 committed by GitHub
parent 36e099ac22
commit 505a5d617b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 28 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
packages/core/src/code_assist/server.test.ts
View File

@ -35,14 +35,14 @@ describe('CodeAssistServer', () => {
], ],
}, },
}; };
vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse); vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
const response = await server.generateContent({ const response = await server.generateContent({
model: 'test-model', model: 'test-model',
contents: [{ role: 'user', parts: [{ text: 'request' }] }], contents: [{ role: 'user', parts: [{ text: 'request' }] }],
}); });
expect(server.callEndpoint).toHaveBeenCalledWith( expect(server.requestPost).toHaveBeenCalledWith(
'generateContent', 'generateContent',
expect.any(Object), expect.any(Object),
undefined, undefined,
@ -72,7 +72,7 @@ describe('CodeAssistServer', () => {
}, },
}; };
})(); })();
vi.spyOn(server, 'streamEndpoint').mockResolvedValue(mockResponse); vi.spyOn(server, 'requestStreamingPost').mockResolvedValue(mockResponse);
const stream = await server.generateContentStream({ const stream = await server.generateContentStream({
model: 'test-model', model: 'test-model',
@ -80,7 +80,7 @@ describe('CodeAssistServer', () => {
}); });
for await (const res of stream) { for await (const res of stream) {
expect(server.streamEndpoint).toHaveBeenCalledWith( expect(server.requestStreamingPost).toHaveBeenCalledWith(
'streamGenerateContent', 'streamGenerateContent',
expect.any(Object), expect.any(Object),
undefined, undefined,
@ -96,7 +96,7 @@ describe('CodeAssistServer', () => {
name: 'operations/123', name: 'operations/123',
done: true, done: true,
}; };
vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse); vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
const response = await server.onboardUser({ const response = await server.onboardUser({
tierId: 'test-tier', tierId: 'test-tier',
@ -104,7 +104,7 @@ describe('CodeAssistServer', () => {
metadata: {}, metadata: {},
}); });
expect(server.callEndpoint).toHaveBeenCalledWith( expect(server.requestPost).toHaveBeenCalledWith(
'onboardUser', 'onboardUser',
expect.any(Object), expect.any(Object),
); );
@ -117,13 +117,13 @@ describe('CodeAssistServer', () => {
const mockResponse = { const mockResponse = {
// TODO: Add mock response // TODO: Add mock response
}; };
vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse); vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
const response = await server.loadCodeAssist({ const response = await server.loadCodeAssist({
metadata: {}, metadata: {},
}); });
expect(server.callEndpoint).toHaveBeenCalledWith( expect(server.requestPost).toHaveBeenCalledWith(
'loadCodeAssist', 'loadCodeAssist',
expect.any(Object), expect.any(Object),
); );
@ -136,7 +136,7 @@ describe('CodeAssistServer', () => {
const mockResponse = { const mockResponse = {
totalTokens: 100, totalTokens: 100,
}; };
vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse); vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
const response = await server.countTokens({ const response = await server.countTokens({
model: 'test-model', model: 'test-model',

34
packages/core/src/code_assist/server.ts
View File

@ -40,8 +40,7 @@ export interface HttpOptions {
} }
// TODO: Use production endpoint once it supports our methods. // TODO: Use production endpoint once it supports our methods.
export const CODE_ASSIST_ENDPOINT = export const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com';
process.env.CODE_ASSIST_ENDPOINT ?? 'https://cloudcode-pa.googleapis.com';
export const CODE_ASSIST_API_VERSION = 'v1internal'; export const CODE_ASSIST_API_VERSION = 'v1internal';
export class CodeAssistServer implements ContentGenerator { export class CodeAssistServer implements ContentGenerator {
@ -54,7 +53,7 @@ export class CodeAssistServer implements ContentGenerator {
async generateContentStream( async generateContentStream(
req: GenerateContentParameters, req: GenerateContentParameters,
): Promise<AsyncGenerator<GenerateContentResponse>> { ): Promise<AsyncGenerator<GenerateContentResponse>> {
const resps = await this.streamEndpoint<CaGenerateContentResponse>( const resps = await this.requestStreamingPost<CaGenerateContentResponse>(
'streamGenerateContent', 'streamGenerateContent',
toGenerateContentRequest(req, this.projectId), toGenerateContentRequest(req, this.projectId),
req.config?.abortSignal, req.config?.abortSignal,
@ -69,7 +68,7 @@ export class CodeAssistServer implements ContentGenerator {
async generateContent( async generateContent(
req: GenerateContentParameters, req: GenerateContentParameters,
): Promise<GenerateContentResponse> { ): Promise<GenerateContentResponse> {
const resp = await this.callEndpoint<CaGenerateContentResponse>( const resp = await this.requestPost<CaGenerateContentResponse>(
'generateContent', 'generateContent',
toGenerateContentRequest(req, this.projectId), toGenerateContentRequest(req, this.projectId),
req.config?.abortSignal, req.config?.abortSignal,
@ -80,7 +79,7 @@ export class CodeAssistServer implements ContentGenerator {
async onboardUser( async onboardUser(
req: OnboardUserRequest, req: OnboardUserRequest,
): Promise<LongrunningOperationResponse> { ): Promise<LongrunningOperationResponse> {
return await this.callEndpoint<LongrunningOperationResponse>( return await this.requestPost<LongrunningOperationResponse>(
'onboardUser', 'onboardUser',
req, req,
); );
@ -89,14 +88,14 @@ export class CodeAssistServer implements ContentGenerator {
async loadCodeAssist( async loadCodeAssist(
req: LoadCodeAssistRequest, req: LoadCodeAssistRequest,
): Promise<LoadCodeAssistResponse> { ): Promise<LoadCodeAssistResponse> {
return await this.callEndpoint<LoadCodeAssistResponse>( return await this.requestPost<LoadCodeAssistResponse>(
'loadCodeAssist', 'loadCodeAssist',
req, req,
); );
} }
async getCodeAssistGlobalUserSetting(): Promise<CodeAssistGlobalUserSettingResponse> { async getCodeAssistGlobalUserSetting(): Promise<CodeAssistGlobalUserSettingResponse> {
return await this.getEndpoint<CodeAssistGlobalUserSettingResponse>( return await this.requestGet<CodeAssistGlobalUserSettingResponse>(
'getCodeAssistGlobalUserSetting', 'getCodeAssistGlobalUserSetting',
); );
} }
@ -104,14 +103,14 @@ export class CodeAssistServer implements ContentGenerator {
async setCodeAssistGlobalUserSetting( async setCodeAssistGlobalUserSetting(
req: SetCodeAssistGlobalUserSettingRequest, req: SetCodeAssistGlobalUserSettingRequest,
): Promise<CodeAssistGlobalUserSettingResponse> { ): Promise<CodeAssistGlobalUserSettingResponse> {
return await this.callEndpoint<CodeAssistGlobalUserSettingResponse>( return await this.requestPost<CodeAssistGlobalUserSettingResponse>(
'setCodeAssistGlobalUserSetting', 'setCodeAssistGlobalUserSetting',
req, req,
); );
} }
async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> { async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {
const resp = await this.callEndpoint<CaCountTokenResponse>( const resp = await this.requestPost<CaCountTokenResponse>(
'countTokens', 'countTokens',
toCountTokenRequest(req), toCountTokenRequest(req),
); );
@ -124,13 +123,13 @@ export class CodeAssistServer implements ContentGenerator {
throw Error(); throw Error();
} }
async callEndpoint<T>( async requestPost<T>(
method: string, method: string,
req: object, req: object,
signal?: AbortSignal, signal?: AbortSignal,
): Promise<T> { ): Promise<T> {
const res = await this.client.request({ const res = await this.client.request({
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`, url: this.getMethodUrl(method),
method: 'POST', method: 'POST',
headers: { headers: {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
@ -143,9 +142,9 @@ export class CodeAssistServer implements ContentGenerator {
return res.data as T; return res.data as T;
} }
async getEndpoint<T>(method: string, signal?: AbortSignal): Promise<T> { async requestGet<T>(method: string, signal?: AbortSignal): Promise<T> {
const res = await this.client.request({ const res = await this.client.request({
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`, url: this.getMethodUrl(method),
method: 'GET', method: 'GET',
headers: { headers: {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
@ -157,13 +156,13 @@ export class CodeAssistServer implements ContentGenerator {
return res.data as T; return res.data as T;
} }
async streamEndpoint<T>( async requestStreamingPost<T>(
method: string, method: string,
req: object, req: object,
signal?: AbortSignal, signal?: AbortSignal,
): Promise<AsyncGenerator<T>> { ): Promise<AsyncGenerator<T>> {
const res = await this.client.request({ const res = await this.client.request({
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`, url: this.getMethodUrl(method),
method: 'POST', method: 'POST',
params: { params: {
alt: 'sse', alt: 'sse',
@ -200,4 +199,9 @@ export class CodeAssistServer implements ContentGenerator {
} }
})(); })();
} }
getMethodUrl(method: string): string {
const endpoint = process.env.CODE_ASSIST_ENDPOINT ?? CODE_ASSIST_ENDPOINT;
return `${endpoint}/${CODE_ASSIST_API_VERSION}:${method}`;
}
} }