309 lines
9.3 KiB
TypeScript
309 lines
9.3 KiB
TypeScript
import type { IncomingMessage, ServerResponse } from 'http';
|
|
import type { HostApiContext } from '../context';
|
|
import { flushStreamingHeaders, sendJson, writeStreamingChunk } from '../route-utils';
|
|
import {
|
|
getFreshWorksSquareAIGatewayCredential,
|
|
markWorksSquareAIGatewayCredentialExpired,
|
|
type WorksSquareAIGatewayCredential,
|
|
} from '../../services/works-square-ai-gateway';
|
|
import { proxyAwareFetch } from '../../utils/proxy-fetch';
|
|
import { logger } from '../../utils/logger';
|
|
import { getOpencodeErrorKind } from '../../../shared/opencode-error-kind';
|
|
|
|
const AI_PROXY_PREFIX = '/api/ai-proxy/v1';
|
|
const MAX_ERROR_LOG_MESSAGE_LENGTH = 600;
|
|
const SKIPPED_REQUEST_HEADERS = new Set([
|
|
'accept-encoding',
|
|
'authorization',
|
|
'connection',
|
|
'content-length',
|
|
'host',
|
|
'keep-alive',
|
|
'proxy-authenticate',
|
|
'proxy-authorization',
|
|
'te',
|
|
'trailer',
|
|
'transfer-encoding',
|
|
'upgrade',
|
|
]);
|
|
const SKIPPED_RESPONSE_HEADERS = new Set([
|
|
'connection',
|
|
'content-encoding',
|
|
'content-length',
|
|
'keep-alive',
|
|
'proxy-authenticate',
|
|
'proxy-authorization',
|
|
'te',
|
|
'trailer',
|
|
'transfer-encoding',
|
|
'upgrade',
|
|
]);
|
|
|
|
async function readRequestBody(req: IncomingMessage): Promise<Buffer | undefined> {
|
|
const chunks: Buffer[] = [];
|
|
for await (const chunk of req) {
|
|
chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
|
|
}
|
|
if (chunks.length === 0) return undefined;
|
|
return Buffer.concat(chunks);
|
|
}
|
|
|
|
function buildForwardHeaders(
|
|
req: IncomingMessage,
|
|
credential: WorksSquareAIGatewayCredential,
|
|
): Record<string, string> {
|
|
const headers: Record<string, string> = {};
|
|
for (const [name, value] of Object.entries(req.headers)) {
|
|
const lowerName = name.toLowerCase();
|
|
if (SKIPPED_REQUEST_HEADERS.has(lowerName) || value == null) continue;
|
|
headers[lowerName] = Array.isArray(value) ? value.join(', ') : String(value);
|
|
}
|
|
headers['x-works-square-ai-token'] = credential.accessToken;
|
|
return headers;
|
|
}
|
|
|
|
function buildTargetUrl(credential: WorksSquareAIGatewayCredential, url: URL): string {
|
|
const suffix = url.pathname.slice(AI_PROXY_PREFIX.length) || '/';
|
|
return `${credential.oneApiBaseUrl}${suffix}${url.search}`;
|
|
}
|
|
|
|
function isExpiredGatewayTokenResponse(status: number, bodyText: string): boolean {
|
|
if (status !== 401) return false;
|
|
const normalized = bodyText.toLowerCase();
|
|
return normalized.includes('expired') && (
|
|
normalized.includes('token')
|
|
|| normalized.includes('gateway')
|
|
|| normalized.includes('ai access')
|
|
);
|
|
}
|
|
|
|
function getForwardedOneApiStatus(status: number, bodyText: string): number {
|
|
if (status === 429 && getOpencodeErrorKind(bodyText) === 'quota_exhausted') {
|
|
return 402;
|
|
}
|
|
return status;
|
|
}
|
|
|
|
function asRecord(value: unknown): Record<string, unknown> | null {
|
|
return value && typeof value === 'object' && !Array.isArray(value)
|
|
? value as Record<string, unknown>
|
|
: null;
|
|
}
|
|
|
|
function truncateForLog(value: string): string {
|
|
return value.length > MAX_ERROR_LOG_MESSAGE_LENGTH
|
|
? `${value.slice(0, MAX_ERROR_LOG_MESSAGE_LENGTH)}...`
|
|
: value;
|
|
}
|
|
|
|
function extractOneApiErrorSummary(bodyText: string): Record<string, unknown> {
|
|
const trimmed = bodyText.trim();
|
|
if (!trimmed) return {};
|
|
|
|
try {
|
|
const parsed = JSON.parse(trimmed) as unknown;
|
|
const record = asRecord(parsed);
|
|
const errorRecord = asRecord(record?.error);
|
|
if (typeof record?.error === 'string' && record.error.trim()) {
|
|
return { message: truncateForLog(record.error.trim()) };
|
|
}
|
|
if (errorRecord) {
|
|
return {
|
|
...(typeof errorRecord.message === 'string' && errorRecord.message.trim()
|
|
? { message: truncateForLog(errorRecord.message.trim()) }
|
|
: {}),
|
|
...(typeof errorRecord.code === 'string' && errorRecord.code.trim()
|
|
? { code: errorRecord.code.trim() }
|
|
: {}),
|
|
...(typeof errorRecord.type === 'string' && errorRecord.type.trim()
|
|
? { type: errorRecord.type.trim() }
|
|
: {}),
|
|
...(typeof errorRecord.param === 'string' && errorRecord.param.trim()
|
|
? { param: errorRecord.param.trim() }
|
|
: {}),
|
|
};
|
|
}
|
|
return {};
|
|
} catch {
|
|
return { bodyPreview: truncateForLog(trimmed) };
|
|
}
|
|
}
|
|
|
|
function logNonSuccessOneApiResponseFromBody(upstream: Response, targetUrl: string, bodyText: string): void {
|
|
if (upstream.ok) return;
|
|
logger.warn('[ai-proxy] One-api returned non-success response', {
|
|
status: upstream.status,
|
|
statusText: upstream.statusText,
|
|
url: targetUrl,
|
|
...extractOneApiErrorSummary(bodyText),
|
|
});
|
|
}
|
|
|
|
function copyResponseHeaders(upstream: Response, res: ServerResponse): void {
|
|
upstream.headers.forEach((value, name) => {
|
|
if (!SKIPPED_RESPONSE_HEADERS.has(name.toLowerCase())) {
|
|
res.setHeader(name, value);
|
|
}
|
|
});
|
|
}
|
|
|
|
async function writeBufferedResponse(
|
|
res: ServerResponse,
|
|
upstream: Response,
|
|
bodyText: string,
|
|
statusCode = upstream.status,
|
|
): Promise<void> {
|
|
res.statusCode = statusCode;
|
|
copyResponseHeaders(upstream, res);
|
|
res.end(bodyText);
|
|
}
|
|
|
|
async function writeUpstreamResponse(res: ServerResponse, upstream: Response): Promise<void> {
|
|
res.statusCode = upstream.status;
|
|
copyResponseHeaders(upstream, res);
|
|
flushStreamingHeaders(res);
|
|
if (!upstream.body) {
|
|
res.end();
|
|
return;
|
|
}
|
|
|
|
const startedAt = Date.now();
|
|
let lastChunkAt = startedAt;
|
|
let chunkCount = 0;
|
|
let byteCount = 0;
|
|
const reader = upstream.body.getReader();
|
|
try {
|
|
while (true) {
|
|
const { done, value } = await reader.read();
|
|
if (done) break;
|
|
if (value) {
|
|
const now = Date.now();
|
|
const gapMs = now - lastChunkAt;
|
|
lastChunkAt = now;
|
|
chunkCount += 1;
|
|
byteCount += value.byteLength;
|
|
if (chunkCount === 1) {
|
|
logger.debug('[ai-proxy] First upstream response chunk received', {
|
|
status: upstream.status,
|
|
bytes: value.byteLength,
|
|
elapsedMs: now - startedAt,
|
|
});
|
|
} else if (gapMs >= 1_000) {
|
|
logger.debug('[ai-proxy] Slow upstream response chunk gap', {
|
|
status: upstream.status,
|
|
chunk: chunkCount,
|
|
bytes: value.byteLength,
|
|
gapMs,
|
|
elapsedMs: now - startedAt,
|
|
});
|
|
}
|
|
if (!await writeStreamingChunk(res, value)) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
} finally {
|
|
reader.releaseLock();
|
|
}
|
|
logger.debug('[ai-proxy] Upstream response stream completed', {
|
|
status: upstream.status,
|
|
chunks: chunkCount,
|
|
bytes: byteCount,
|
|
elapsedMs: Date.now() - startedAt,
|
|
});
|
|
res.end();
|
|
}
|
|
|
|
async function forwardToOneApi(
|
|
req: IncomingMessage,
|
|
url: URL,
|
|
requestBody: Buffer | undefined,
|
|
): Promise<{ response: Response; targetUrl: string } | null> {
|
|
const credential = await getFreshWorksSquareAIGatewayCredential();
|
|
if (!credential) return null;
|
|
|
|
const targetUrl = buildTargetUrl(credential, url);
|
|
logger.info('[ai-proxy] Forwarding request to one-api', {
|
|
method: req.method ?? 'GET',
|
|
url: targetUrl,
|
|
});
|
|
|
|
return {
|
|
response: await proxyAwareFetch(targetUrl, {
|
|
method: req.method ?? 'GET',
|
|
headers: buildForwardHeaders(req, credential),
|
|
body: requestBody as BodyInit | undefined,
|
|
}),
|
|
targetUrl,
|
|
};
|
|
}
|
|
|
|
export async function handleAiProxyRoutes(
|
|
req: IncomingMessage,
|
|
res: ServerResponse,
|
|
url: URL,
|
|
_ctx: HostApiContext,
|
|
): Promise<boolean> {
|
|
if (!url.pathname.startsWith(AI_PROXY_PREFIX)) {
|
|
return false;
|
|
}
|
|
|
|
try {
|
|
const requestBody = await readRequestBody(req);
|
|
let upstreamResult = await forwardToOneApi(req, url, requestBody);
|
|
if (!upstreamResult) {
|
|
sendJson(res, 401, {
|
|
success: false,
|
|
error: 'Works Square login session is missing or expired',
|
|
});
|
|
return true;
|
|
}
|
|
let upstream = upstreamResult.response;
|
|
let targetUrl = upstreamResult.targetUrl;
|
|
|
|
if (upstream.status === 401) {
|
|
const bodyText = await upstream.text();
|
|
if (isExpiredGatewayTokenResponse(upstream.status, bodyText)) {
|
|
markWorksSquareAIGatewayCredentialExpired();
|
|
logger.info('[ai-proxy] Retrying after expired Works Square AI gateway token');
|
|
upstreamResult = await forwardToOneApi(req, url, requestBody);
|
|
if (!upstreamResult) {
|
|
sendJson(res, 401, {
|
|
success: false,
|
|
error: 'Works Square login session is missing or expired',
|
|
});
|
|
return true;
|
|
}
|
|
upstream = upstreamResult.response;
|
|
targetUrl = upstreamResult.targetUrl;
|
|
} else {
|
|
logNonSuccessOneApiResponseFromBody(upstream, targetUrl, bodyText);
|
|
await writeBufferedResponse(res, upstream, bodyText, getForwardedOneApiStatus(upstream.status, bodyText));
|
|
return true;
|
|
}
|
|
}
|
|
|
|
if (!upstream.ok) {
|
|
const bodyText = await upstream.text();
|
|
logNonSuccessOneApiResponseFromBody(upstream, targetUrl, bodyText);
|
|
await writeBufferedResponse(res, upstream, bodyText, getForwardedOneApiStatus(upstream.status, bodyText));
|
|
return true;
|
|
}
|
|
|
|
await writeUpstreamResponse(res, upstream);
|
|
} catch (error) {
|
|
logger.error('[ai-proxy] Request failed', error);
|
|
if (res.headersSent) {
|
|
if (!res.destroyed && !res.writableEnded) {
|
|
res.end();
|
|
}
|
|
} else {
|
|
sendJson(res, 502, {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : String(error),
|
|
});
|
|
}
|
|
}
|
|
return true;
|
|
}
|