From a85c62c125495c58f86b2a2b1edc389a701c72a3 Mon Sep 17 00:00:00 2001 From: Antonio Cheong Date: Mon, 6 Jul 2026 15:34:15 +0800 Subject: [PATCH] fix gateway request cleanup and model listing --- AGENTS.md | 3 +- Dockerfile | 3 +- README.md | 4 + gateway/server.mjs | 163 +++++++++++++++++++++++++++-------- gateway/test/server.test.mjs | 34 ++++++++ 5 files changed, 171 insertions(+), 36 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 92fdada..9dc918b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -35,6 +35,7 @@ node --check gateway/server.mjs - Do not point SDK startup at `/usr/local/bin/codebuddy`; the SDK resolves `dist/codebuddy-headless.js` relative to the provided path and the symlink path resolves incorrectly. - Treat request model `codebuddy` as a gateway alias. Do not pass it through to the CodeBuddy CLI as a model id. - REST API built-in CodeBuddy tools are disabled by default. External OpenAI-style `tools` are exposed as request-scoped SDK MCP tools and returned to clients as OpenAI-compatible `tool_calls`. +- `/v1/models` uses bounded SDK model discovery and falls back to the gateway default model list after `CODEBUDDY_GATEWAY_MODELS_TIMEOUT`. ## Docker Verification @@ -62,7 +63,7 @@ Verify: curl -s http://127.0.0.1:11532/health ``` -Then test `/v1/chat/completions` with `model: "codebuddy"` and a simple exact-response prompt. For custom tools, verify the first response returns `finish_reason: "tool_calls"` and an OpenAI-compatible `message.tool_calls` array, then send the external tool result back as a `role: "tool"` message. +Then test `/v1/models`, `/v1/chat/completions` with `model: "codebuddy"` and a simple exact-response prompt, and a streaming chat response. For custom tools, verify the first response returns `finish_reason: "tool_calls"` and an OpenAI-compatible `message.tool_calls` array, then send the external tool result back as a `role: "tool"` message. Clean up only the disposable container, not the auth volume: diff --git a/Dockerfile b/Dockerfile index e30ff6e..4f793f6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -44,7 +44,8 @@ COPY sshd_config /etc/ssh/sshd_config COPY docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh COPY codebuddy-stdin.sh /usr/local/bin/codebuddy-stdin -RUN chmod +x /usr/local/bin/docker-entrypoint.sh /usr/local/bin/codebuddy-stdin +RUN sed -i 's/\r$//' /usr/local/bin/docker-entrypoint.sh /usr/local/bin/codebuddy-stdin \ + && chmod +x /usr/local/bin/docker-entrypoint.sh /usr/local/bin/codebuddy-stdin WORKDIR /home/codebuddy/workspace diff --git a/README.md b/README.md index bfe9dec..62b6ea6 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,8 @@ docker run -d --name codebuddy-gateway \ | `CODEBUDDY_GATEWAY_MODEL` | — | Model override | | `CODEBUDDY_GATEWAY_MAX_TURNS` | `30` | Max conversation turns | | `CODEBUDDY_GATEWAY_TIMEOUT` | `300000` | Request timeout in ms | +| `CODEBUDDY_GATEWAY_MODELS_TIMEOUT` | `5000` | `/v1/models` SDK discovery timeout in ms before returning fallback models | +| `CODEBUDDY_GATEWAY_MAX_BODY_BYTES` | `10485760` | Maximum JSON request body size in bytes | These env vars are auto-loaded on SSH login via `.bashrc`, `.profile`, and `.bash_profile`. If the home volume already has files, startup keeps existing files and only creates missing ones. @@ -271,6 +273,8 @@ Then call the API again with the original assistant `tool_calls` message and the curl -s http://127.0.0.1:10532/v1/models ``` +The endpoint first asks the SDK for model metadata. If discovery does not respond within `CODEBUDDY_GATEWAY_MODELS_TIMEOUT`, it returns the gateway fallback model list instead of leaving the HTTP request hanging. + ### Health Check ```bash diff --git a/gateway/server.mjs b/gateway/server.mjs index f050ede..72a788f 100644 --- a/gateway/server.mjs +++ b/gateway/server.mjs @@ -16,6 +16,8 @@ const DISALLOWED_TOOLS = parseCsv(process.env.CODEBUDDY_GATEWAY_DISALLOWED_TOOLS const MODEL = process.env.CODEBUDDY_GATEWAY_MODEL || undefined; const MAX_TURNS = parseInt(process.env.CODEBUDDY_GATEWAY_MAX_TURNS || '30', 10); const REQUEST_TIMEOUT_MS = parseInt(process.env.CODEBUDDY_GATEWAY_TIMEOUT || '300000', 10); +const MODELS_REFRESH_TIMEOUT_MS = parseInt(process.env.CODEBUDDY_GATEWAY_MODELS_TIMEOUT || '5000', 10); +const MAX_REQUEST_BODY_BYTES = parseInt(process.env.CODEBUDDY_GATEWAY_MAX_BODY_BYTES || String(10 * 1024 * 1024), 10); const EXTERNAL_TOOL_SERVER_NAME = 'openai_client_tools'; // Build the env object to pass through to the SDK. @@ -54,14 +56,59 @@ export function parseCsv(value) { } function json(res, status, body) { + if (res.destroyed || res.writableEnded) return; res.writeHead(status, { 'Content-Type': 'application/json' }); res.end(JSON.stringify(body)); } -function readBody(req) { +export class RequestBodyTooLargeError extends Error { + constructor(limit) { + super(`Request body exceeds ${limit} bytes`); + this.name = 'RequestBodyTooLargeError'; + this.statusCode = 413; + } +} + +function abortWith(controller, reason) { + if (!controller.signal.aborted) { + controller.abort(reason); + } +} + +export function attachRequestAbort(req, res, abortController, timeoutMs = REQUEST_TIMEOUT_MS) { + const onAborted = () => abortWith(abortController, new Error('HTTP request aborted by client')); + const onResponseClosed = () => { + if (!res.writableEnded) { + abortWith(abortController, new Error('HTTP response closed before completion')); + } + }; + const timeout = Number.isFinite(timeoutMs) && timeoutMs > 0 + ? setTimeout(() => abortWith(abortController, new Error('Gateway request timed out')), timeoutMs) + : null; + + req.once('aborted', onAborted); + res.once('close', onResponseClosed); + + return () => { + req.off('aborted', onAborted); + res.off('close', onResponseClosed); + if (timeout) clearTimeout(timeout); + }; +} + +export function readBody(req, maxBytes = MAX_REQUEST_BODY_BYTES) { return new Promise((resolve, reject) => { const chunks = []; - req.on('data', (c) => chunks.push(c)); + let received = 0; + req.on('data', (c) => { + received += c.length; + if (received > maxBytes) { + reject(new RequestBodyTooLargeError(maxBytes)); + req.destroy(); + return; + } + chunks.push(c); + }); req.on('end', () => resolve(Buffer.concat(chunks).toString('utf-8'))); req.on('error', reject); }); @@ -256,6 +303,7 @@ export function resolveEffectiveModel(requestModel, configuredModel) { } function streamToolCallChunk(res, id, model, collector) { + if (res.destroyed || res.writableEnded) return; const toolChunk = { id: id || `chatcmpl-${Date.now()}`, object: 'chat.completion.chunk', @@ -276,32 +324,54 @@ function streamToolCallChunk(res, id, model, collector) { let cachedModels = null; async function refreshModels() { - try { - const q = query({ - prompt: '', - options: { - maxTurns: 0, - permissionMode: 'bypassPermissions', - outputFormat: 'text', - settingSources: ['user'], - env: SDK_ENV, - }, - }); - // We just need the init message which includes model info. - for await (const msg of q) { - if (msg.type === 'init') { - const models = msg.models || []; - if (models.length > 0) { - return models.map((m) => ({ id: m, object: 'model', owned_by: 'codebuddy' })); + const abortController = new AbortController(); + let timeout = null; + + const discoverModels = async () => { + try { + const q = query({ + prompt: '', + options: { + abortController, + maxTurns: 0, + permissionMode: 'bypassPermissions', + outputFormat: 'text', + settingSources: ['user'], + env: SDK_ENV, + }, + }); + // We just need the init message which includes model info. + for await (const msg of q) { + if (msg.type === 'init') { + const models = msg.models || []; + if (models.length > 0) { + return models.map((m) => ({ id: m, object: 'model', owned_by: 'codebuddy' })); + } } } + } catch { + // fall through to defaults } - } catch { - // fall through to defaults + return null; + }; + + const discovery = discoverModels(); + const timeoutMs = Number.isFinite(MODELS_REFRESH_TIMEOUT_MS) ? MODELS_REFRESH_TIMEOUT_MS : 0; + if (timeoutMs <= 0) return discovery; + + const fallbackOnTimeout = new Promise((resolve) => { + timeout = setTimeout(() => { + abortWith(abortController, new Error('Model discovery timed out')); + resolve(null); + }, timeoutMs); + }); + + try { + return await Promise.race([discovery, fallbackOnTimeout]); + } finally { + if (timeout) clearTimeout(timeout); } - return null; } - async function handleModels(_req, res) { if (!cachedModels) { cachedModels = await refreshModels(); @@ -318,7 +388,15 @@ async function handleModels(_req, res) { // POST /v1/chat/completions async function handleChatCompletions(req, res) { - const body = await readBody(req); + let body; + try { + body = await readBody(req); + } catch (err) { + if (err instanceof RequestBodyTooLargeError) { + return json(res, 413, { error: { message: err.message } }); + } + throw err; + } let parsed; try { parsed = JSON.parse(body); @@ -336,6 +414,7 @@ async function handleChatCompletions(req, res) { const toolCallCollector = createToolCallCollector(); const externalToolServer = buildExternalToolServer(externalTools, toolCallCollector); const abortController = new AbortController(); + const detachRequestAbort = attachRequestAbort(req, res, abortController); const prompt = messagesToPrompt(messages); // Shared query options @@ -406,7 +485,9 @@ async function handleChatCompletions(req, res) { model: model || 'codebuddy', choices: [{ index: 0, delta: { content: block.text }, finish_reason: null }], }; - res.write(sseEvent(null, chunk)); + if (!res.destroyed && !res.writableEnded) { + res.write(sseEvent(null, chunk)); + } } } } @@ -434,7 +515,9 @@ async function handleChatCompletions(req, res) { model: model || 'codebuddy', choices: [{ index: 0, delta: {}, finish_reason: finishReason }], }; - res.write(sseEvent(null, finalChunk)); + if (!res.destroyed && !res.writableEnded) { + res.write(sseEvent(null, finalChunk)); + } } catch (err) { if (toolCallCollector.hasCalls() && isAbortError(err)) { finishReason = 'tool_calls'; @@ -449,18 +532,28 @@ async function handleChatCompletions(req, res) { model: model || 'codebuddy', choices: [{ index: 0, delta: {}, finish_reason: finishReason }], }; - res.write(sseEvent(null, finalChunk)); + if (!res.destroyed && !res.writableEnded) { + res.write(sseEvent(null, finalChunk)); + } } else if (isAuthError(err)) { - res.write(sseEvent('error', { - message: 'Authentication required. Set CODEBUDDY_API_KEY or run "codebuddy login" via SSH.', - detail: err.message, - })); + if (!res.destroyed && !res.writableEnded) { + res.write(sseEvent('error', { + message: 'Authentication required. Set CODEBUDDY_API_KEY or run "codebuddy login" via SSH.', + detail: err.message, + })); + } } else { - res.write(sseEvent('error', { message: err.message })); + if (!isAbortError(err) && !res.destroyed && !res.writableEnded) { + res.write(sseEvent('error', { message: err.message })); + } } + } finally { + detachRequestAbort(); + } + if (!res.destroyed && !res.writableEnded) { + res.write(sseEvent(null, '[DONE]')); + res.end(); } - res.write(sseEvent(null, '[DONE]')); - res.end(); } else { // ---- Non-streaming ---- try { @@ -574,6 +667,8 @@ async function handleChatCompletions(req, res) { } else { json(res, 500, { error: { message: err.message } }); } + } finally { + detachRequestAbort(); } } } diff --git a/gateway/test/server.test.mjs b/gateway/test/server.test.mjs index 3930c0e..7fd2281 100644 --- a/gateway/test/server.test.mjs +++ b/gateway/test/server.test.mjs @@ -1,13 +1,18 @@ import test from 'node:test'; import assert from 'node:assert/strict'; +import { EventEmitter } from 'node:events'; +import { PassThrough } from 'node:stream'; import { z } from 'zod'; import { + attachRequestAbort, captureToolUseBlocks, createToolCallCollector, jsonSchemaObjectToZodShape, messagesToPrompt, normalizeOpenAiTools, parseCsv, + readBody, + RequestBodyTooLargeError, resolveEffectiveModel, streamingToolCalls, } from '../server.mjs'; @@ -169,3 +174,32 @@ test('resolveEffectiveModel treats codebuddy as gateway default alias', () => { assert.equal(resolveEffectiveModel('minimax-m3', 'glm-5.2'), 'minimax-m3'); assert.equal(resolveEffectiveModel(undefined, 'glm-5.2'), 'glm-5.2'); }); + + +test('readBody rejects oversized request bodies', async () => { + const req = new PassThrough(); + const bodyPromise = readBody(req, 4); + + req.write(Buffer.from('123')); + req.write(Buffer.from('45')); + + await assert.rejects(bodyPromise, RequestBodyTooLargeError); + assert.equal(req.destroyed, true); +}); + +test('attachRequestAbort aborts on client disconnect and detaches listeners', () => { + const req = new EventEmitter(); + const res = new EventEmitter(); + res.writableEnded = false; + const abortController = new AbortController(); + + const detach = attachRequestAbort(req, res, abortController, 0); + + res.emit('close'); + assert.equal(abortController.signal.aborted, true); + assert.match(abortController.signal.reason.message, /response closed/); + + detach(); + assert.equal(req.listenerCount('aborted'), 0); + assert.equal(res.listenerCount('close'), 0); +});