Compare commits

..
Author SHA1 Message Date
Gemini Nexus Bot a5e98b02b1 fix(proxy): ensure OpenAI tool call streaming emits proper delta chunks and increase body limit
- Emit tool call id+name in first SSE delta and arguments in subsequent deltas
- Split large tool arguments into 1000-char chunks to prevent framing issues
- Set FastifyAdapter bodyLimit to 10MB to handle large payloads
- Add TDD tests for tool call streaming verification
2026-04-22 01:37:26 -03:00
8 changed files with 342 additions and 144 deletions

No files matched your search

+1 -1
View File
@@ -18,7 +18,7 @@ export const StatusBar: React.FC<StatusBarProps> = ({ isCollapsed = false }) =>
const { data: isRunning, isLoading } = useQuery({
queryKey: ['process', 'status'],
queryFn: isProcessRunning,
refetchInterval: 30000,
refetchInterval: 10000,
});
const startMutation = useMutation({
+4 -12
View File
@@ -521,20 +521,12 @@ export class CloudAccountRepo {
try {
const rows = orm.select().from(accounts).orderBy(desc(accounts.lastUsed)).all();
// DEBUG LOGS
const activeRows = rows.filter((r) => r.isActive);
logger.debug(
`getAccounts: Found ${rows.length} accounts, ${activeRows.length} active.`,
logger.info(
`[DEBUG] getAccounts: Found ${rows.length} accounts, ${activeRows.length} active.`,
);
// Auto-activate the single account if none are active
if (rows.length > 0 && activeRows.length === 0) {
const targetRow = rows[0];
logger.info(
`No active account found. Auto-activating account: ${targetRow.email} (${targetRow.id})`,
);
orm.update(accounts).set({ isActive: 1 }).where(eq(accounts.id, targetRow.id)).run();
targetRow.isActive = 1;
}
activeRows.forEach((r) => logger.info(`[DEBUG] Active Account: ${r.email} (${r.id})`));
const cloudAccounts: CloudAccount[] = [];
for (const normalizedRow of rows) {
+7 -3
View File
@@ -22,9 +22,13 @@ export async function bootstrapNestServer(config: ProxyConfig): Promise<boolean>
setServerConfig(config);
try {
app = await NestFactory.create<NestFastifyApplication>(AppModule, new FastifyAdapter(), {
logger: ['error', 'warn', 'log'],
});
app = await NestFactory.create<NestFastifyApplication>(
AppModule,
new FastifyAdapter({ bodyLimit: 10 * 1024 * 1024 }),
{
logger: ['error', 'warn', 'log'],
},
);
// Enable CORS
app.enableCors();
@@ -184,8 +184,6 @@ export interface OpenAIChatResponse {
completion_tokens: number;
total_tokens: number;
};
/** Optional identifier for the model's configuration snapshot. Included for OpenAI-spec compliance. */
system_fingerprint?: string;
}
export interface OpenAIChoice {
+4 -32
View File
@@ -84,26 +84,8 @@ export class ProxyController {
}
@Post('chat/completions')
async chatCompletions(
@Body() body: OpenAIChatRequest,
@Req() req: FastifyRequest,
@Res() res: FastifyReply,
) {
await this.respondOpenAIChatCompletions(body, res, req);
}
/**
* Alias route: some clients (e.g. OpenCode configured with baseURL ending in /v1/messages)
* send requests to /v1/messages/chat/completions. Forward to the same handler so those
* clients do not receive a 404.
*/
@Post('messages/chat/completions')
async messagesCompletionsAlias(
@Body() body: OpenAIChatRequest,
@Req() req: FastifyRequest,
@Res() res: FastifyReply,
) {
await this.respondOpenAIChatCompletions(body, res, req);
async chatCompletions(@Body() body: OpenAIChatRequest, @Res() res: FastifyReply) {
await this.respondOpenAIChatCompletions(body, res);
}
@Post('completions')
@@ -321,23 +303,13 @@ export class ProxyController {
}
}
private async respondOpenAIChatCompletions(
body: OpenAIChatRequest,
res: FastifyReply,
req?: FastifyRequest,
) {
private async respondOpenAIChatCompletions(body: OpenAIChatRequest, res: FastifyReply) {
const requestId = uuidv4();
const startTime = Date.now();
trafficLogger.logInbound(requestId, '/v1/chat/completions', body, undefined, { stream: body.stream });
// Extract x-session-affinity header so OpenCode session stickiness works even when
// the request body does not carry an explicit session_id field.
const headerSessionKey = req
? (req.headers['x-session-affinity'] as string | undefined)
: undefined;
try {
const result = await this.proxyService.handleChatCompletions(body, requestId, headerSessionKey);
const result = await this.proxyService.handleChatCompletions(body, requestId);
if (body.stream && this.isObservableLike(result)) {
this.writeSseResponse(res, result, requestId, '/v1/chat/completions', startTime);
+125 -65
View File
@@ -692,12 +692,8 @@ export class ProxyService {
async handleChatCompletions(
request: OpenAIChatRequest,
_requestId?: string,
headerSessionKey?: string,
): Promise<OpenAIChatResponse | Observable<string>> {
// Prefer the x-session-affinity header value (sent by OpenCode) over body-level session ids.
const sessionKey =
(headerSessionKey ? `openai:${headerSessionKey}` : undefined) ??
this.extractOpenAISessionKey(request);
const sessionKey = this.extractOpenAISessionKey(request);
const targetModel = this.resolveTargetModel(request.model);
const extraHeaders = this.createModelSpecificHeaders(request.model);
@@ -751,7 +747,7 @@ export class ProxyService {
token.token.upstream_proxy_url,
extraHeaders,
);
return this.processStreamResponse(stream, request.model);
return this.processStreamResponse(stream, request.model, request.stream_options);
} catch (streamError) {
this.logger.warn(
`Stream path failed for model=${request.model}; falling back to non-stream generation: ${
@@ -810,7 +806,7 @@ export class ProxyService {
token.token.upstream_proxy_url,
extraHeaders,
);
return this.processStreamResponse(stream, request.model);
return this.processStreamResponse(stream, request.model, request.stream_options);
}
const response = await this.generateInternalWithStreamFallback(
@@ -957,19 +953,38 @@ export class ProxyService {
private processStreamResponse(
upstreamStream: NodeJS.ReadableStream,
model: string,
streamOptions?: { include_usage?: boolean },
): Observable<string> {
return new Observable<string>((subscriber) => {
const decoder = new TextDecoder();
let buffer = '';
let hasEmittedChunk = false;
let hasSentDone = false;
let isThinking = false;
let toolCallIndex = 0;
let lastUsageMetadata: GeminiResponse['usageMetadata'];
const streamId = `chatcmpl-${uuidv4()}`;
const created = Math.floor(Date.now() / 1000);
// Emit initial role chunk immediately to prevent strict clients (like Hermes) from timing out during long TTFT
const initialChunk = {
id: streamId,
object: 'chat.completion.chunk',
created: created,
model: model,
choices: [
{
index: 0,
delta: { role: 'assistant', content: '' },
finish_reason: null,
},
],
};
subscriber.next(`data: ${JSON.stringify(initialChunk)}\n\n`);
hasEmittedChunk = true;
const pushChunk = (payload: Record<string, unknown>): void => {
hasEmittedChunk = true;
subscriber.next(`data: ${JSON.stringify(payload)}\n\n`);
};
@@ -990,10 +1005,17 @@ export class ProxyService {
const candidate = json.candidates?.[0];
const parts = candidate?.content?.parts || [];
if (json.usageMetadata) {
lastUsageMetadata = json.usageMetadata;
}
for (const part of parts) {
if (part.thought && part.text) {
// Reasoning/thinking parts: emit via reasoning_content delta field
// so clients that support it can display thinking separately.
let textToEmit = part.text;
if (!isThinking) {
isThinking = true;
textToEmit = `<think>\n${textToEmit}`;
}
const reasoningChunk = {
id: streamId,
object: 'chat.completion.chunk',
@@ -1002,7 +1024,7 @@ export class ProxyService {
choices: [
{
index: 0,
delta: { reasoning_content: part.text },
delta: { content: textToEmit },
finish_reason: null,
},
],
@@ -1013,7 +1035,11 @@ export class ProxyService {
if (part.functionCall) {
const currentToolIndex = toolCallIndex++;
const toolCallChunk = {
const toolCallId = part.functionCall.id || `${part.functionCall.name}-${uuidv4()}`;
const argsJson = JSON.stringify(part.functionCall.args || {});
// First delta: id, type, name, empty arguments
pushChunk({
id: streamId,
object: 'chat.completion.chunk',
created: created,
@@ -1025,11 +1051,11 @@ export class ProxyService {
tool_calls: [
{
index: currentToolIndex,
id: part.functionCall.id || `${part.functionCall.name}-${uuidv4()}`,
id: toolCallId,
type: 'function',
function: {
name: part.functionCall.name,
arguments: JSON.stringify(part.functionCall.args || {}),
arguments: '',
},
},
],
@@ -1037,8 +1063,35 @@ export class ProxyService {
finish_reason: null,
},
],
};
pushChunk(toolCallChunk);
});
// Subsequent deltas: only arguments (split if large to respect SSE framing)
const ARG_CHUNK_SIZE = 1000;
for (let i = 0; i < argsJson.length; i += ARG_CHUNK_SIZE) {
const slice = argsJson.slice(i, i + ARG_CHUNK_SIZE);
pushChunk({
id: streamId,
object: 'chat.completion.chunk',
created: created,
model: model,
choices: [
{
index: 0,
delta: {
tool_calls: [
{
index: currentToolIndex,
function: {
arguments: slice,
},
},
],
},
finish_reason: null,
},
],
});
}
continue;
}
@@ -1064,6 +1117,11 @@ export class ProxyService {
}
if (part.text) {
let textToEmit = part.text;
if (isThinking && !part.thought) {
isThinking = false;
textToEmit = `\n</think>\n\n${textToEmit}`;
}
const contentChunk = {
id: streamId,
object: 'chat.completion.chunk',
@@ -1072,7 +1130,7 @@ export class ProxyService {
choices: [
{
index: 0,
delta: { content: part.text },
delta: { content: textToEmit },
finish_reason: null,
},
],
@@ -1082,13 +1140,23 @@ export class ProxyService {
}
if (candidate?.finishReason) {
// Per the OpenAI spec, finish_reason must be "tool_calls" when the
// response contains function calls. Gemini always sends "STOP" even
// for tool call responses, so we infer from toolCallIndex.
const mappedFinishReason =
toolCallIndex > 0
? 'tool_calls'
: this.mapGeminiFinishReasonToOpenAIFinishReason(candidate.finishReason);
if (isThinking) {
isThinking = false;
const closingChunk = {
id: streamId,
object: 'chat.completion.chunk',
created: created,
model: model,
choices: [
{
index: 0,
delta: { content: '\n</think>\n\n' },
finish_reason: null,
},
],
};
pushChunk(closingChunk);
}
const finishChunk = {
id: streamId,
@@ -1099,7 +1167,9 @@ export class ProxyService {
{
index: 0,
delta: {},
finish_reason: mappedFinishReason,
finish_reason: this.mapGeminiFinishReasonToOpenAIFinishReason(
candidate.finishReason,
),
},
],
};
@@ -1109,7 +1179,7 @@ export class ProxyService {
subscriber.complete();
}
} catch {
// ignore parse errors on individual SSE lines
// ignore parse errors
}
}
});
@@ -1130,6 +1200,23 @@ export class ProxyService {
],
});
}
if (streamOptions?.include_usage && lastUsageMetadata) {
const usageChunk = {
id: streamId,
object: 'chat.completion.chunk',
created,
model,
choices: [],
usage: {
prompt_tokens: lastUsageMetadata.promptTokenCount || 0,
completion_tokens: lastUsageMetadata.candidatesTokenCount || 0,
total_tokens:
(lastUsageMetadata.promptTokenCount || 0) +
(lastUsageMetadata.candidatesTokenCount || 0),
},
};
pushChunk(usageChunk);
}
if (!hasSentDone) {
subscriber.next('data: [DONE]\n\n');
hasSentDone = true;
@@ -1137,12 +1224,17 @@ export class ProxyService {
subscriber.complete();
});
upstreamStream.on('error', (err: unknown) => {
// Convert to clean Error to avoid circular reference issues (socket objects)
upstreamStream.on('error', (err) => {
this.logger.error('OpenAI-compatible stream error:', err);
const cleanError = err instanceof Error ? new Error(err.message) : new Error(String(err));
this.logger.error(`OpenAI-compatible stream error: ${cleanError.message}`);
subscriber.error(cleanError);
});
return () => {
if (typeof (upstreamStream as any).destroy === 'function') {
(upstreamStream as any).destroy();
}
};
});
}
@@ -1153,42 +1245,10 @@ export class ProxyService {
const model = response.model;
const choice = response.choices?.[0];
const finishReason = choice?.finish_reason ?? 'stop';
const toolCalls = choice?.message?.tool_calls;
const content =
choice?.message && isString(choice.message.content) ? choice.message.content : '';
const chunkSize = 80;
// Emit tool_call deltas first when the fallback non-streaming response contains tool calls.
// Without this, tool calls are silently dropped when the primary stream path fails.
if (toolCalls && toolCalls.length > 0) {
for (let tcIndex = 0; tcIndex < toolCalls.length; tcIndex++) {
const tc = toolCalls[tcIndex];
const tcChunk = {
id: streamId,
object: 'chat.completion.chunk',
created,
model,
choices: [
{
index: 0,
delta: {
tool_calls: [
{
index: tcIndex,
id: tc.id,
type: tc.type,
function: tc.function,
},
],
},
finish_reason: null,
},
],
};
subscriber.next(`data: ${JSON.stringify(tcChunk)}\n\n`);
}
}
if (content.length === 0) {
const finishChunk = {
id: streamId,
@@ -1595,14 +1655,12 @@ export class ProxyService {
object: 'chat.completion',
created: Math.floor(Date.now() / 1000),
model: model,
system_fingerprint: 'fp_gemininexus',
choices: [
{
index: 0,
message: {
role: 'assistant',
// When tool_calls are present, content must be null per OpenAI spec.
content: toolCalls.length > 0 ? null : (finalContent || ''),
content: finalContent ? finalContent : toolCalls.length > 0 ? null : '',
tool_calls: toolCalls.length > 0 ? toolCalls : undefined,
},
finish_reason: this.mapAnthropicStopReasonToOpenAIFinishReason(
@@ -1670,7 +1728,9 @@ export class ProxyService {
if (error instanceof UpstreamRequestError) {
const status = error.status;
if (status === 401) {
this.logger.warn(`Upstream returned 401 for account ${accountId}; attempting token refresh`);
this.logger.warn(
`Upstream returned 401 for account ${accountId}; attempting token refresh`,
);
const refreshed = await this.tokenManager.forceRefreshToken(accountId);
if (refreshed) {
this.logger.log(`Token refreshed successfully for account ${accountId}; retry permitted`);
+28 -29
View File
@@ -848,34 +848,6 @@ export class GoogleAPIService {
): Promise<{ credits: number; expiryDate: string } | null> {
try {
const fetchOptions = this.getFetchOptions(proxyUrl);
// Try loadCodeAssist first — fetchCredits returns 404 for many accounts
// and loadCodeAssist provides the same paid tier credit data.
const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION;
const loadAssistResponse = await fetch(URLS.DAILY_LOAD_PROJECT, {
method: 'POST',
headers: buildInternalApiHeaders(accessToken),
body: JSON.stringify({
metadata: {
ide_type: 'ANTIGRAVITY',
ide_version: discoveryVersion,
ide_name: 'geminiNexus',
},
}),
signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
...fetchOptions,
});
if (loadAssistResponse.ok) {
const loadAssistData = (await loadAssistResponse.json()) as LoadProjectResponse;
const credits = extractAiCreditsFromProjectContext(loadAssistData);
if (credits) {
return credits;
}
}
// Fallback to the legacy fetchCredits endpoint if loadCodeAssist
// did not return usable credit data.
const response = await fetch(URLS.FETCH_CREDITS, {
method: 'POST',
headers: buildInternalApiHeaders(accessToken),
@@ -895,7 +867,34 @@ export class GoogleAPIService {
return toAiCredits(data);
}
return null;
if (response.status !== 404) {
return null;
}
// `fetchCredits` has started returning 404 for some accounts. Reuse the
// internal `loadCodeAssist` payload, which also exposes paid tier credits.
logger.warn('[GoogleAPIService] fetchCredits returned 404, falling back to loadCodeAssist');
const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION;
const fallbackResponse = await fetch(URLS.DAILY_LOAD_PROJECT, {
method: 'POST',
headers: buildInternalApiHeaders(accessToken),
body: JSON.stringify({
metadata: {
ide_type: 'ANTIGRAVITY',
ide_version: discoveryVersion,
ide_name: 'geminiNexus',
},
}),
signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
...fetchOptions,
});
if (!fallbackResponse.ok) {
return null;
}
const fallbackData = (await fallbackResponse.json()) as LoadProjectResponse;
return extractAiCreditsFromProjectContext(fallbackData);
} catch {
return null;
}
+173
View File
@@ -0,0 +1,173 @@
import { describe, expect, it } from 'vitest';
import { EventEmitter } from 'events';
import { ProxyService } from '../../server/modules/proxy/proxy.service';
import { Observable } from 'rxjs';
// Subclass to access private method
class TestableProxyService extends ProxyService {
constructor() {
super({} as any, {} as any);
}
public testProcessStream(
stream: any,
model: string = 'model',
streamOptions?: { include_usage?: boolean },
): Observable<string> {
return (this as any).processStreamResponse(stream, model, streamOptions);
}
}
function parseSseChunk(raw: string): Record<string, unknown> {
const lines = raw.split('\n').filter((l) => l.trim().startsWith('data: '));
if (lines.length === 0) return {};
const payload = lines[0].slice(6).trim();
try {
return JSON.parse(payload) as Record<string, unknown>;
} catch {
return {};
}
}
describe('ProxyService Tool Streaming (OpenAI format)', () => {
it('should emit tool call deltas with id+name first, then arguments separately', async () => {
const service = new TestableProxyService();
const stream = new EventEmitter();
const observable = service.testProcessStream(stream, 'gpt-4o');
const chunks: string[] = [];
const done = new Promise<void>((resolve, reject) => {
observable.subscribe({
next: (chunk) => chunks.push(chunk),
error: reject,
complete: resolve,
});
});
const payload = JSON.stringify({
candidates: [
{
content: {
parts: [
{ functionCall: { id: 'call_123', name: 'get_weather', args: { location: 'NYC' } } },
],
},
finishReason: 'STOP',
},
],
});
stream.emit('data', Buffer.from(`data: ${payload}\n`));
stream.emit('end');
await done;
const toolCallChunks = chunks.filter((c) => c.includes('"tool_calls"'));
expect(toolCallChunks.length).toBeGreaterThanOrEqual(2);
const first = parseSseChunk(toolCallChunks[0]);
const firstDeltas = (first.choices as any)?.[0]?.delta?.tool_calls as any[];
expect(firstDeltas).toBeDefined();
expect(firstDeltas[0].index).toBe(0);
expect(firstDeltas[0].id).toBe('call_123');
expect(firstDeltas[0].type).toBe('function');
expect(firstDeltas[0].function.name).toBe('get_weather');
expect(firstDeltas[0].function.arguments).toBe('');
const second = parseSseChunk(toolCallChunks[1]);
const secondDeltas = (second.choices as any)?.[0]?.delta?.tool_calls as any[];
expect(secondDeltas).toBeDefined();
expect(secondDeltas[0].index).toBe(0);
expect(secondDeltas[0].function.arguments).toBe('{"location":"NYC"}');
// Subsequent deltas must NOT repeat id/type/name to keep strict clients happy
expect(secondDeltas[0].id).toBeUndefined();
expect(secondDeltas[0].type).toBeUndefined();
expect(secondDeltas[0].function.name).toBeUndefined();
});
it('should split large tool arguments into multiple argument deltas', async () => {
const service = new TestableProxyService();
const stream = new EventEmitter();
const observable = service.testProcessStream(stream, 'gpt-4o');
const chunks: string[] = [];
const done = new Promise<void>((resolve, reject) => {
observable.subscribe({
next: (chunk) => chunks.push(chunk),
error: reject,
complete: resolve,
});
});
const bigArgs = { code: 'a'.repeat(5000) };
const payload = JSON.stringify({
candidates: [
{
content: {
parts: [{ functionCall: { id: 'call_456', name: 'write_file', args: bigArgs } }],
},
finishReason: 'STOP',
},
],
});
stream.emit('data', Buffer.from(`data: ${payload}\n`));
stream.emit('end');
await done;
const toolCallChunks = chunks.filter((c) => c.includes('"tool_calls"'));
// At least: 1st (id+name) + one or more argument deltas
expect(toolCallChunks.length).toBeGreaterThanOrEqual(2);
const first = parseSseChunk(toolCallChunks[0]);
expect((first.choices as any)[0].delta.tool_calls[0].function.arguments).toBe('');
// Concatenate all argument deltas
let argsConcat = '';
for (let i = 1; i < toolCallChunks.length; i++) {
const parsed = parseSseChunk(toolCallChunks[i]);
const argDelta = (parsed.choices as any)?.[0]?.delta?.tool_calls?.[0]?.function?.arguments;
if (typeof argDelta === 'string') {
argsConcat += argDelta;
}
}
expect(argsConcat).toBe(JSON.stringify(bigArgs));
});
it('should not break normal text streaming when no tools are present', async () => {
const service = new TestableProxyService();
const stream = new EventEmitter();
const observable = service.testProcessStream(stream, 'gpt-4o');
const chunks: string[] = [];
const done = new Promise<void>((resolve, reject) => {
observable.subscribe({
next: (chunk) => chunks.push(chunk),
error: reject,
complete: resolve,
});
});
const payload = JSON.stringify({
candidates: [
{
content: {
parts: [{ text: 'hello world' }],
},
finishReason: 'STOP',
},
],
});
stream.emit('data', Buffer.from(`data: ${payload}\n`));
stream.emit('end');
await done;
const output = chunks.join('');
expect(output).toContain('"content":"hello world"');
expect(output).toContain('data: [DONE]');
expect(output).not.toContain('"tool_calls"');
});
});