Skip to content

Commit 4e3add3

Browse files
rtrompierrtrompier
authored
Merge branch 'main' into rtrompier-patch-1
2 parents 2e2229a + 8e5cf9f commit 4e3add3

File tree

6 files changed

+90
-50
lines changed

6 files changed

+90
-50
lines changed

‎src/lib/components/chat/ChatMessage.svelte‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,8 @@
128128
| { type: "text"; content: string }
129129
| { type: "tool"; uuid: string; updates: MessageToolUpdate[] };
130130
131+
type ToolBlock = Extract<Block, { type: "tool" }>;
132+
131133
let blocks = $derived.by(() => {
132134
const updates = message.updates ?? [];
133135
const res: Block[] = [];
@@ -155,9 +157,11 @@
155157
if (last?.type === "text") last.content += chunk;
156158
else res.push({ type: "text" as const, content: chunk });
157159
} else if (isMessageToolUpdate(update)) {
158-
const last = res.at(-1);
159-
if (last?.type === "tool" && last.uuid === update.uuid) {
160-
last.updates.push(update);
160+
const existingBlock = res.find(
161+
(b): b is ToolBlock => b.type === "tool" && b.uuid === update.uuid
162+
);
163+
if (existingBlock) {
164+
existingBlock.updates.push(update);
161165
} else {
162166
res.push({ type: "tool" as const, uuid: update.uuid, updates: [update] });
163167
}

‎src/lib/server/mcp/clientPool.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,3 +59,12 @@ export async function drainPool() {
5959
pool.delete(key);
6060
}
6161
}
62+
63+
export function evictFromPool(server: McpServerConfig): Client | undefined {
64+
const key = keyOf(server);
65+
const client = pool.get(key);
66+
if (client) {
67+
pool.delete(key);
68+
}
69+
return client;
70+
}

‎src/lib/server/mcp/httpClient.ts‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
import { Client } from "@modelcontextprotocol/sdk/client";
2-
import { getClient } from "./clientPool";
2+
import { getClient, evictFromPool } from "./clientPool";
3+
4+
function isConnectionClosedError(err: unknown): boolean {
5+
const message = err instanceof Error ? err.message : String(err);
6+
return message.includes("-32000") || message.toLowerCase().includes("connection closed");
7+
}
38

49
export interface McpServerConfig {
510
name: string;
@@ -34,20 +39,40 @@ export async function callMcpTool(
3439

3540
// Get a (possibly pooled) client. The client itself was connected with a signal
3641
// that already composes outer cancellation. We still enforce a per-call timeout here.
37-
const activeClient = client ?? (await getClient(server, signal));
42+
let activeClient = client ?? (await getClient(server, signal));
43+
44+
const callToolOptions = {
45+
signal,
46+
timeout: timeoutMs,
47+
// Enable progress tokens so long-running tools keep extending the timeout.
48+
onprogress: () => {},
49+
resetTimeoutOnProgress: true,
50+
};
3851

39-
// Prefer the SDK's built-in request controls (timeout, signal)
40-
const response = await activeClient.callTool(
41-
{ name: tool, arguments: normalizedArgs },
42-
undefined,
43-
{
44-
signal,
45-
timeout: timeoutMs,
46-
// Enable progress tokens so long-running tools keep extending the timeout.
47-
onprogress: () => {},
48-
resetTimeoutOnProgress: true,
52+
let response;
53+
try {
54+
response = await activeClient.callTool(
55+
{ name: tool, arguments: normalizedArgs },
56+
undefined,
57+
callToolOptions
58+
);
59+
} catch (err) {
60+
if (!isConnectionClosedError(err)) {
61+
throw err;
4962
}
50-
);
63+
64+
// Evict stale client and close it
65+
const stale = evictFromPool(server);
66+
stale?.close?.().catch(() => {});
67+
68+
// Retry with fresh client
69+
activeClient = await getClient(server, signal);
70+
response = await activeClient.callTool(
71+
{ name: tool, arguments: normalizedArgs },
72+
undefined,
73+
callToolOptions
74+
);
75+
}
5176

5277
const parts = Array.isArray(response?.content) ? (response.content as Array<unknown>) : [];
5378
const textParts = parts

‎src/lib/server/mcp/tools.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { Client } from "@modelcontextprotocol/sdk/client";
22
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
33
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
44
import type { McpServerConfig } from "./httpClient";
5+
import { logger } from "$lib/server/logger";
56
// use console.* for lightweight diagnostics in production logs
67

78
export type OpenAiTool = {
@@ -82,7 +83,7 @@ async function listServerTools(
8283
const response = await client.listTools({});
8384
const tools = Array.isArray(response?.tools) ? (response.tools as ListedTool[]) : [];
8485
try {
85-
console.debug(
86+
logger.debug(
8687
{
8788
server: server.name,
8889
url: server.url,

‎src/lib/server/textGeneration/mcp/runMcpFlow.ts‎

Lines changed: 25 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import { hasAuthHeader, isStrictHfMcpLogin, hasNonEmptyToken } from "$lib/server
2121
import { buildImageRefResolver } from "./fileRefs";
2222
import { prepareMessagesWithFiles } from "$lib/server/textGeneration/utils/prepareFiles";
2323
import { makeImageProcessor } from "$lib/server/endpoints/images";
24+
import { logger } from "$lib/server/logger";
2425

2526
export type RunMcpFlowContext = Pick<
2627
TextGenerationContext,
@@ -45,7 +46,7 @@ export async function* runMcpFlow({
4546
// Start from env-configured servers
4647
let servers = getMcpServers();
4748
try {
48-
console.debug(
49+
logger.debug(
4950
{ baseServers: servers.map((s) => ({ name: s.name, url: s.url })), count: servers.length },
5051
"[mcp] base servers loaded"
5152
);
@@ -74,7 +75,7 @@ export async function* runMcpFlow({
7475
for (const s of custom) byName.set(s.name, s);
7576
servers = [...byName.values()];
7677
try {
77-
console.debug(
78+
logger.debug(
7879
{
7980
customProvidedCount: custom.length,
8081
mergedServers: servers.map((s) => ({
@@ -96,7 +97,7 @@ export async function* runMcpFlow({
9697
const before = servers.map((s) => s.name);
9798
servers = servers.filter((s) => names.includes(s.name));
9899
try {
99-
console.debug(
100+
logger.debug(
100101
{ selectedNames: names, before, after: servers.map((s) => s.name) },
101102
"[mcp] applied name selection"
102103
);
@@ -108,7 +109,7 @@ export async function* runMcpFlow({
108109

109110
// If selection/merge yielded no servers, bail early with clearer log
110111
if (servers.length === 0) {
111-
console.warn("[mcp] no MCP servers selected after merge/name filter");
112+
logger.warn({}, "[mcp] no MCP servers selected after merge/name filter");
112113
return false;
113114
}
114115

@@ -125,15 +126,15 @@ export async function* runMcpFlow({
125126
try {
126127
const rejected = before.filter((b) => !servers.includes(b));
127128
if (rejected.length > 0) {
128-
console.warn(
129+
logger.warn(
129130
{ rejected: rejected.map((r) => ({ name: r.name, url: r.url })) },
130131
"[mcp] rejected servers by URL safety"
131132
);
132133
}
133134
} catch {}
134135
}
135136
if (servers.length === 0) {
136-
console.warn("[mcp] all selected MCP servers rejected by URL safety guard");
137+
logger.warn({}, "[mcp] all selected MCP servers rejected by URL safety guard");
137138
return false;
138139
}
139140

@@ -163,14 +164,14 @@ export async function* runMcpFlow({
163164
});
164165
if (overlayApplied.length > 0) {
165166
try {
166-
console.debug({ overlayApplied }, "[mcp] forwarded HF token to servers");
167+
logger.debug({ overlayApplied }, "[mcp] forwarded HF token to servers");
167168
} catch {}
168169
}
169170
}
170171
} catch {
171172
// best-effort overlay; continue if anything goes wrong
172173
}
173-
console.debug(
174+
logger.debug(
174175
{ count: servers.length, servers: servers.map((s) => s.name) },
175176
"[mcp] servers configured"
176177
);
@@ -182,7 +183,7 @@ export async function* runMcpFlow({
182183
try {
183184
const supportsTools = Boolean((model as unknown as { supportsTools?: boolean }).supportsTools);
184185
const toolsEnabled = Boolean(forceTools) || supportsTools;
185-
console.debug(
186+
logger.debug(
186187
{
187188
model: model.id ?? model.name,
188189
supportsTools,
@@ -192,7 +193,7 @@ export async function* runMcpFlow({
192193
"[mcp] tools gate evaluation"
193194
);
194195
if (!toolsEnabled) {
195-
console.info(
196+
logger.info(
196197
{ model: model.id ?? model.name },
197198
"[mcp] tools disabled for model; skipping MCP flow"
198199
);
@@ -226,7 +227,7 @@ export async function* runMcpFlow({
226227
});
227228

228229
if (!runMcp) {
229-
console.info(
230+
logger.info(
230231
{ model: targetModel.id ?? targetModel.name, resolvedRoute },
231232
"[mcp] runMcp=false (routing chose non-tools candidate)"
232233
);
@@ -235,13 +236,13 @@ export async function* runMcpFlow({
235236

236237
const { tools: oaTools, mapping } = await getOpenAiToolsForMcp(servers, { signal: abortSignal });
237238
try {
238-
console.info(
239+
logger.info(
239240
{ toolCount: oaTools.length, toolNames: oaTools.map((t) => t.function.name) },
240241
"[mcp] openai tool defs built"
241242
);
242243
} catch {}
243244
if (oaTools.length === 0) {
244-
console.warn("[mcp] zero tools available after listing; skipping MCP flow");
245+
logger.warn({}, "[mcp] zero tools available after listing; skipping MCP flow");
245246
return false;
246247
}
247248

@@ -273,7 +274,7 @@ export async function* runMcpFlow({
273274
});
274275

275276
const mmEnabled = (forceMultimodal ?? false) || targetModel.multimodal;
276-
console.info(
277+
logger.info(
277278
{
278279
targetModel: targetModel.id ?? targetModel.name,
279280
mmEnabled,
@@ -390,7 +391,7 @@ export async function* runMcpFlow({
390391
route: resolvedRoute,
391392
model: candidateModelId,
392393
};
393-
console.debug(
394+
logger.debug(
394395
{ route: resolvedRoute, model: candidateModelId },
395396
"[mcp] router metadata emitted"
396397
);
@@ -425,7 +426,7 @@ export async function* runMcpFlow({
425426
model: "",
426427
provider: providerHeader as unknown as import("@huggingface/inference").InferenceProvider,
427428
};
428-
console.debug({ provider: providerHeader }, "[mcp] provider metadata emitted");
429+
logger.debug({ provider: providerHeader }, "[mcp] provider metadata emitted");
429430
}
430431

431432
const toolCallState: Record<number, { id?: string; name?: string; arguments: string }> = {};
@@ -461,7 +462,7 @@ export async function* runMcpFlow({
461462
.map((k) => Number(k))
462463
.sort((a, b) => a - b)[0] ?? 0
463464
];
464-
console.info(
465+
logger.info(
465466
{ firstCallName: first?.name, hasId: Boolean(first?.id) },
466467
"[mcp] observed streamed tool_call delta"
467468
);
@@ -526,7 +527,7 @@ export async function* runMcpFlow({
526527
}
527528
}
528529
}
529-
console.info(
530+
logger.info(
530531
{ sawToolCalls: Object.keys(toolCallState).length > 0, tokens: tokenCount, loop },
531532
"[mcp] completion stream closed"
532533
);
@@ -536,7 +537,7 @@ export async function* runMcpFlow({
536537
const missingId = Object.values(toolCallState).some((c) => c?.name && !c?.id);
537538
let calls: NormalizedToolCall[];
538539
if (missingId) {
539-
console.debug(
540+
logger.debug(
540541
{ loop },
541542
"[mcp] missing tool_call id in stream; retrying non-stream to recover ids"
542543
);
@@ -611,7 +612,7 @@ export async function* runMcpFlow({
611612
];
612613
toolMsgCount = event.summary.toolMessages?.length ?? 0;
613614
toolRunCount = event.summary.toolRuns?.length ?? 0;
614-
console.info(
615+
logger.info(
615616
{ toolMsgCount, toolRunCount },
616617
"[mcp] tools executed; continuing loop for follow-up completion"
617618
);
@@ -635,13 +636,13 @@ export async function* runMcpFlow({
635636
text: lastAssistantContent,
636637
interrupted: false,
637638
};
638-
console.info(
639+
logger.info(
639640
{ length: lastAssistantContent.length, loop },
640641
"[mcp] final answer emitted (no tool_calls)"
641642
);
642643
return true;
643644
}
644-
console.warn("[mcp] exceeded tool-followup loops; falling back");
645+
logger.warn({}, "[mcp] exceeded tool-followup loops; falling back");
645646
} catch (err) {
646647
const msg = String(err ?? "");
647648
const isAbort =
@@ -651,10 +652,10 @@ export async function* runMcpFlow({
651652
msg.includes("Request was aborted");
652653
if (isAbort) {
653654
// Expected on user stop; keep logs quiet and do not treat as error
654-
console.debug("[mcp] aborted by user");
655+
logger.debug({}, "[mcp] aborted by user");
655656
return false;
656657
}
657-
console.warn({ err: msg }, "[mcp] flow failed, falling back to default endpoint");
658+
logger.warn({ err: msg }, "[mcp] flow failed, falling back to default endpoint");
658659
} finally {
659660
// ensure MCP clients are closed after the turn
660661
await drainPool();

‎src/routes/api/mcp/health/+server.ts‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -89,16 +89,16 @@ export const POST: RequestHandler = async ({ request, locals }) => {
8989

9090
// Try Streamable HTTP transport first
9191
try {
92-
logger.info(`[MCP Health] Trying HTTP transport for ${url}`);
92+
logger.info({}, `[MCP Health] Trying HTTP transport for ${url}`);
9393
client = new Client({
9494
name: "chat-ui-health-check",
9595
version: "1.0.0",
9696
});
9797

9898
const transport = new StreamableHTTPClientTransport(baseUrl, { requestInit });
99-
logger.info(`[MCP Health] Connecting to ${url}...`);
99+
logger.info({}, `[MCP Health] Connecting to ${url}...`);
100100
await client.connect(transport);
101-
logger.info(`[MCP Health] Connected successfully via HTTP`);
101+
logger.info({}, `[MCP Health] Connected successfully via HTTP`);
102102

103103
// Connection successful, get tools
104104
const toolsResponse = await client.listTools();
@@ -141,7 +141,7 @@ export const POST: RequestHandler = async ({ request, locals }) => {
141141
} catch (error) {
142142
httpError = error instanceof Error ? error : new Error(String(error));
143143
lastError = httpError;
144-
logger.info("Streamable HTTP failed, trying SSE transport...", lastError.message);
144+
logger.warn(lastError.message, "Streamable HTTP failed, trying SSE transport...");
145145

146146
// Close failed client
147147
try {
@@ -152,16 +152,16 @@ export const POST: RequestHandler = async ({ request, locals }) => {
152152

153153
// Try SSE transport
154154
try {
155-
logger.info(`[MCP Health] Trying SSE transport for ${url}`);
155+
logger.info({}, `[MCP Health] Trying SSE transport for ${url}`);
156156
client = new Client({
157157
name: "chat-ui-health-check",
158158
version: "1.0.0",
159159
});
160160

161161
const sseTransport = new SSEClientTransport(baseUrl, { requestInit });
162-
logger.info(`[MCP Health] Connecting via SSE...`);
162+
logger.info({}, `[MCP Health] Connecting via SSE...`);
163163
await client.connect(sseTransport);
164-
logger.info(`[MCP Health] Connected successfully via SSE`);
164+
logger.info({}, `[MCP Health] Connected successfully via SSE`);
165165

166166
// Connection successful, get tools
167167
const toolsResponse = await client.listTools();
@@ -211,7 +211,7 @@ export const POST: RequestHandler = async ({ request, locals }) => {
211211
{ cause: sseError instanceof Error ? sseError : undefined }
212212
);
213213
}
214-
logger.error("Both transports failed. Last error:", lastError);
214+
logger.error(lastError, "Both transports failed.");
215215
}
216216
}
217217

@@ -253,7 +253,7 @@ export const POST: RequestHandler = async ({ request, locals }) => {
253253
clearTimeout(timeoutId);
254254
return res;
255255
} catch (error) {
256-
logger.error("MCP health check failed:", error);
256+
logger.error(error, "MCP health check failed");
257257

258258
// Clean up client if it exists
259259
try {

0 commit comments

Comments
 (0)