Skip to content

Commit b47d733

Browse files
yyyu-googlecopybara-github
yyyu-google
authored and
copybara-github
committed
feat: exposing customHeader in requestOptions to allow users pass in customer headers.
PiperOrigin-RevId: 623649229
1 parent ec3c45f commit b47d733

File tree

4 files changed

+214
-22
lines changed

4 files changed

+214
-22
lines changed

‎src/functions/post_request.ts

+58-21
Original file line numberDiff line numberDiff line change
@@ -58,30 +58,19 @@ export async function postRequest({
5858
if (resourceMethod === constants.STREAMING_GENERATE_CONTENT_METHOD) {
5959
vertexEndpoint += '?alt=sse';
6060
}
61-
62-
if (
63-
requestOptions?.apiClient &&
64-
(requestOptions?.apiClient.includes('\n') ||
65-
requestOptions?.apiClient.includes('\r'))
66-
) {
67-
throw new ClientError(
68-
'Found line break in apiClient request option field, please remove ' +
69-
'the line break and try again.'
70-
);
71-
}
72-
73-
const extraHeaders: HeadersInit = requestOptions?.apiClient
74-
? {'X-Goog-Api-Client': requestOptions?.apiClient}
75-
: {};
61+
const necessaryHeaders = new Headers({
62+
Authorization: `Bearer ${token}`,
63+
'Content-Type': 'application/json',
64+
'User-Agent': constants.USER_AGENT,
65+
});
66+
const totalHeaders: Headers = getExtraHeaders(
67+
necessaryHeaders,
68+
requestOptions
69+
);
7670
return fetch(vertexEndpoint, {
7771
...getFetchOptions(requestOptions),
7872
method: 'POST',
79-
headers: {
80-
Authorization: `Bearer ${token}`,
81-
'Content-Type': 'application/json',
82-
'User-Agent': constants.USER_AGENT,
83-
...extraHeaders,
84-
},
73+
headers: totalHeaders,
8574
body: JSON.stringify(data),
8675
});
8776
}
@@ -101,3 +90,51 @@ function getFetchOptions(requestOptions?: RequestOptions): RequestInit {
10190
fetchOptions.signal = signal;
10291
return fetchOptions;
10392
}
93+
94+
function stringHasLineBreak(header?: string | null): boolean {
95+
if (header === null || header === undefined) {
96+
return false;
97+
}
98+
return header.includes('\n') || header.includes('\r');
99+
}
100+
function headersHasLineBreak(customHeaders?: Headers): boolean {
101+
if (!customHeaders) {
102+
return false;
103+
}
104+
for (const [key, value] of customHeaders.entries()) {
105+
if (stringHasLineBreak(key) || stringHasLineBreak(value)) {
106+
return true;
107+
}
108+
}
109+
return false;
110+
}
111+
112+
function getExtraHeaders(
113+
necessaryHeaders: Headers,
114+
requestOptions?: RequestOptions
115+
): Headers {
116+
if (stringHasLineBreak(requestOptions?.apiClient)) {
117+
throw new ClientError(
118+
'Found line break in apiClient request option field, please remove ' +
119+
'the line break and try again.'
120+
);
121+
}
122+
if (headersHasLineBreak(requestOptions?.customHeaders)) {
123+
throw new ClientError(
124+
'Found line break in customerHeaders request option field, please remove ' +
125+
'the line break and try again.'
126+
);
127+
}
128+
const totalHeaders: Headers = new Headers();
129+
for (const [key, val] of necessaryHeaders.entries()) {
130+
totalHeaders.append(key, val);
131+
}
132+
const customHeaders = requestOptions?.customHeaders ?? new Headers();
133+
for (const [key, val] of customHeaders.entries()) {
134+
totalHeaders.append(key, val);
135+
}
136+
if (requestOptions?.apiClient) {
137+
totalHeaders.append('X-Goog-Api-Client', requestOptions?.apiClient);
138+
}
139+
return totalHeaders;
140+
}

‎src/functions/test/post_request_test.ts

+150
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import {postRequest} from '../post_request';
2+
import {GenerateContentRequest, RequestOptions} from '../../types';
3+
describe('postRequest', () => {
4+
const REGION = 'us-central1';
5+
const PROJECT = 'project-id';
6+
const RESOURCE_PATH = 'resource-path';
7+
const RESOURCE_METHOD = 'resource-method';
8+
const TOKEN = 'token';
9+
const API_ENDPOINT = 'api-endpoint';
10+
const data = {} as GenerateContentRequest;
11+
let fetchSpy: jasmine.Spy;
12+
13+
beforeEach(() => {
14+
fetchSpy = spyOn(global, 'fetch').and.resolveTo({} as Response);
15+
});
16+
it('apiClient header contains line break type 1, should throw', async () => {
17+
const requestOptions: RequestOptions = {
18+
apiClient: 'apiClient\n',
19+
};
20+
const expectedErrorMessage =
21+
'[VertexAI.ClientError]: Found line break in apiClient request option field, please remove ' +
22+
'the line break and try again.';
23+
await postRequest({
24+
region: REGION,
25+
project: PROJECT,
26+
resourcePath: RESOURCE_PATH,
27+
resourceMethod: RESOURCE_METHOD,
28+
token: TOKEN,
29+
apiEndpoint: API_ENDPOINT,
30+
data: data,
31+
requestOptions: requestOptions,
32+
}).catch(e => {
33+
expect(e.message).toEqual(expectedErrorMessage);
34+
});
35+
});
36+
it('apiClient header contains line break type 2, should throw', async () => {
37+
const requestOptions: RequestOptions = {
38+
apiClient: 'apiClient\r',
39+
};
40+
const expectedErrorMessage =
41+
'[VertexAI.ClientError]: Found line break in apiClient request option field, please remove ' +
42+
'the line break and try again.';
43+
await postRequest({
44+
region: REGION,
45+
project: PROJECT,
46+
resourcePath: RESOURCE_PATH,
47+
resourceMethod: RESOURCE_METHOD,
48+
token: TOKEN,
49+
apiEndpoint: API_ENDPOINT,
50+
data: data,
51+
requestOptions: requestOptions,
52+
}).catch(e => {
53+
expect(e.message).toEqual(expectedErrorMessage);
54+
});
55+
});
56+
it('apiClient header correct, should call through', async () => {
57+
const requestOptions: RequestOptions = {
58+
apiClient: 'apiClient',
59+
};
60+
await postRequest({
61+
region: REGION,
62+
project: PROJECT,
63+
resourcePath: RESOURCE_PATH,
64+
resourceMethod: RESOURCE_METHOD,
65+
token: TOKEN,
66+
apiEndpoint: API_ENDPOINT,
67+
data: data,
68+
requestOptions: requestOptions,
69+
});
70+
const actualHeaders = fetchSpy.calls.mostRecent().args[1].headers;
71+
expect(actualHeaders.get('X-Goog-Api-Client')).toEqual('apiClient');
72+
});
73+
it('customer object header contains line break type 1, should throw', async () => {
74+
const requestOptions: RequestOptions = {
75+
customHeaders: new Headers({customerHeader: 'apiClient\n'}),
76+
} as RequestOptions;
77+
const expectedErrorMessage =
78+
'[VertexAI.ClientError]: Found line break in customerHeaders request option field, please remove ' +
79+
'the line break and try again.';
80+
await postRequest({
81+
region: REGION,
82+
project: PROJECT,
83+
resourcePath: RESOURCE_PATH,
84+
resourceMethod: RESOURCE_METHOD,
85+
token: TOKEN,
86+
apiEndpoint: API_ENDPOINT,
87+
data: data,
88+
requestOptions: requestOptions,
89+
}).catch(e => {
90+
expect(e.message).toEqual(expectedErrorMessage);
91+
});
92+
});
93+
it('customer object header contains line break type 2, should throw', async () => {
94+
const requestOptions: RequestOptions = {
95+
customHeaders: new Headers({customerHeader: 'apiClient\r'}),
96+
} as RequestOptions;
97+
const expectedErrorMessage =
98+
'[VertexAI.ClientError]: Found line break in customerHeaders request option field, please remove ' +
99+
'the line break and try again.';
100+
await postRequest({
101+
region: REGION,
102+
project: PROJECT,
103+
resourcePath: RESOURCE_PATH,
104+
resourceMethod: RESOURCE_METHOD,
105+
token: TOKEN,
106+
apiEndpoint: API_ENDPOINT,
107+
data: data,
108+
requestOptions: requestOptions,
109+
}).catch(e => {
110+
expect(e.message).toEqual(expectedErrorMessage);
111+
});
112+
});
113+
it('set customer header correctly should call through', async () => {
114+
const requestOptions: RequestOptions = {
115+
customHeaders: new Headers({customerHeader: 'customerHeaderValue'}),
116+
} as RequestOptions;
117+
await postRequest({
118+
region: REGION,
119+
project: PROJECT,
120+
resourcePath: RESOURCE_PATH,
121+
resourceMethod: RESOURCE_METHOD,
122+
token: TOKEN,
123+
apiEndpoint: API_ENDPOINT,
124+
data: data,
125+
requestOptions: requestOptions,
126+
});
127+
const actualHeaders = fetchSpy.calls.mostRecent().args[1].headers;
128+
expect(actualHeaders.get('customerHeader')).toEqual('customerHeaderValue');
129+
});
130+
it('set X-Goog-Api-Client in custom header and apiClient, should call through', async () => {
131+
const requestOptions: RequestOptions = {
132+
customHeaders: new Headers({'X-Goog-Api-Client': 'apiClient1'}),
133+
apiClient: 'apiClient2',
134+
} as RequestOptions;
135+
await postRequest({
136+
region: REGION,
137+
project: PROJECT,
138+
resourcePath: RESOURCE_PATH,
139+
resourceMethod: RESOURCE_METHOD,
140+
token: TOKEN,
141+
apiEndpoint: API_ENDPOINT,
142+
data: data,
143+
requestOptions: requestOptions,
144+
});
145+
const actualHeaders = fetchSpy.calls.mostRecent().args[1].headers;
146+
expect(actualHeaders.get('X-Goog-Api-Client')).toEqual(
147+
'apiClient1, apiClient2'
148+
);
149+
});
150+
});

‎src/types/content.ts

+4
Original file line numberDiff line numberDiff line change
@@ -916,4 +916,8 @@ export interface RequestOptions {
916916
* backend.
917917
*/
918918
apiClient?: string;
919+
/**
920+
* Value for custom HTTP headers to set on the HTTP request.
921+
*/
922+
customHeaders?: Headers;
919923
}

‎tsconfig.json

+2-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
"resolveJsonModule": true,
77
"lib": [
88
"es2018",
9-
"dom"
9+
"dom",
10+
"dom.iterable"
1011
]
1112
},
1213
"include": [

0 commit comments

Comments
 (0)