mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-09 17:29:12 +00:00
fix: align MCP discovery and output with v1
This commit is contained in:
7Sageer
parent
bf886decb7
commit
b9362ccfe4
16 changed files with 680 additions and 99 deletions
|
|
@ -76,8 +76,26 @@ export interface CompressImageOptions {
|
|||
readonly byteBudget?: number;
|
||||
/** Override the raw-byte ceiling above which compression is skipped. */
|
||||
readonly maxDecodeBytes?: number;
|
||||
readonly telemetry?: ImageCompressionTelemetry;
|
||||
}
|
||||
|
||||
export interface ImageCompressionTelemetryClient {
|
||||
track(event: string, properties?: Readonly<Record<string, unknown>>): void;
|
||||
}
|
||||
|
||||
export interface ImageCompressionTelemetry {
|
||||
readonly client: ImageCompressionTelemetryClient;
|
||||
readonly source: string;
|
||||
}
|
||||
|
||||
type CompressOutcome =
|
||||
| 'compressed'
|
||||
| 'passthrough_fast'
|
||||
| 'passthrough_guard'
|
||||
| 'passthrough_unsupported'
|
||||
| 'passthrough_unhelpful'
|
||||
| 'passthrough_error';
|
||||
|
||||
export interface CompressImageResult {
|
||||
/** Bytes to send: the re-encoded image, or the original when unchanged. */
|
||||
readonly data: Uint8Array;
|
||||
|
|
@ -109,6 +127,7 @@ export async function compressImageForModel(
|
|||
mimeType: string,
|
||||
options: CompressImageOptions = {},
|
||||
): Promise<CompressImageResult> {
|
||||
const startedAt = Date.now();
|
||||
const maxEdge = options.maxEdge ?? MAX_IMAGE_EDGE_PX;
|
||||
const byteBudget = options.byteBudget ?? IMAGE_BYTE_BUDGET;
|
||||
const maxDecodeBytes = options.maxDecodeBytes ?? MAX_DECODE_BYTES;
|
||||
|
|
@ -126,23 +145,36 @@ export async function compressImageForModel(
|
|||
originalByteLength: bytes.length,
|
||||
finalByteLength: bytes.length,
|
||||
});
|
||||
const finish = (outcome: CompressOutcome, result: CompressImageResult): CompressImageResult => {
|
||||
reportCompressEvent(options.telemetry, {
|
||||
outcome,
|
||||
startedAt,
|
||||
inputMime: normalizedMime,
|
||||
result,
|
||||
});
|
||||
return result;
|
||||
};
|
||||
|
||||
if (bytes.length === 0) return passthrough();
|
||||
if (bytes.length === 0) return finish('passthrough_unsupported', passthrough());
|
||||
// Only re-encode formats the codec handles; everything else passes through.
|
||||
if (!RECODABLE_MIME.has(normalizedMime)) return passthrough();
|
||||
if (!RECODABLE_MIME.has(normalizedMime)) return finish('passthrough_unsupported', passthrough());
|
||||
|
||||
// Fast path: already within both budgets — no codec load, no allocation.
|
||||
const longestEdge = dims ? Math.max(dims.width, dims.height) : 0;
|
||||
const withinBytes = bytes.length <= byteBudget;
|
||||
const withinEdge = longestEdge > 0 && longestEdge <= maxEdge;
|
||||
if (withinBytes && (withinEdge || longestEdge === 0)) return passthrough();
|
||||
if (withinBytes && (withinEdge || longestEdge === 0)) {
|
||||
return finish('passthrough_fast', passthrough());
|
||||
}
|
||||
|
||||
// Decompression-bomb guard: refuse to decode absurd pixel counts. The sniff
|
||||
// above gave us the dimensions without decoding, so this costs nothing.
|
||||
if (dims && dims.width * dims.height > MAX_DECODE_PIXELS) return passthrough();
|
||||
if (dims && dims.width * dims.height > MAX_DECODE_PIXELS) {
|
||||
return finish('passthrough_guard', passthrough());
|
||||
}
|
||||
// Refuse to decode very large byte payloads (e.g. a huge or invalid image
|
||||
// from an MCP tool) that would be loaded just to be dropped downstream.
|
||||
if (bytes.length > maxDecodeBytes) return passthrough();
|
||||
if (bytes.length > maxDecodeBytes) return finish('passthrough_guard', passthrough());
|
||||
|
||||
try {
|
||||
const { Jimp } = await import('jimp');
|
||||
|
|
@ -166,9 +198,9 @@ export async function compressImageForModel(
|
|||
const finalPixels = encoded.width * encoded.height;
|
||||
const shrankBytes = encoded.data.length < bytes.length;
|
||||
const shrankPixels = originalPixels > 0 && finalPixels < originalPixels;
|
||||
if (!shrankBytes && !shrankPixels) return passthrough();
|
||||
if (!shrankBytes && !shrankPixels) return finish('passthrough_unhelpful', passthrough());
|
||||
|
||||
return {
|
||||
return finish('compressed', {
|
||||
data: encoded.data,
|
||||
mimeType: encoded.mimeType,
|
||||
width: encoded.width,
|
||||
|
|
@ -178,10 +210,10 @@ export async function compressImageForModel(
|
|||
changed: true,
|
||||
originalByteLength: bytes.length,
|
||||
finalByteLength: encoded.data.length,
|
||||
};
|
||||
});
|
||||
} catch {
|
||||
// Decode/encode failure — keep the original bytes.
|
||||
return passthrough();
|
||||
return finish('passthrough_error', passthrough());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -214,10 +246,11 @@ export async function compressBase64ForModel(
|
|||
// Skip very large payloads before allocating: base64 decodes to ~3/4 its
|
||||
// length, so a payload whose decoded size would exceed the cap is passed
|
||||
// through without the Buffer.from allocation (and without touching Jimp).
|
||||
const startedAt = Date.now();
|
||||
const maxDecodeBytes = options.maxDecodeBytes ?? MAX_DECODE_BYTES;
|
||||
const approxBytes = Math.floor((base64.length * 3) / 4);
|
||||
if (approxBytes > maxDecodeBytes) {
|
||||
return {
|
||||
const result: CompressBase64Result = {
|
||||
base64,
|
||||
mimeType,
|
||||
width: 0,
|
||||
|
|
@ -228,12 +261,19 @@ export async function compressBase64ForModel(
|
|||
originalByteLength: approxBytes,
|
||||
finalByteLength: approxBytes,
|
||||
};
|
||||
reportCompressEvent(options.telemetry, {
|
||||
outcome: 'passthrough_guard',
|
||||
startedAt,
|
||||
inputMime: normalizeMime(mimeType),
|
||||
result,
|
||||
});
|
||||
return result;
|
||||
}
|
||||
let bytes: Buffer;
|
||||
try {
|
||||
bytes = Buffer.from(base64, 'base64');
|
||||
} catch {
|
||||
return {
|
||||
const result: CompressBase64Result = {
|
||||
base64,
|
||||
mimeType,
|
||||
width: 0,
|
||||
|
|
@ -244,6 +284,13 @@ export async function compressBase64ForModel(
|
|||
originalByteLength: 0,
|
||||
finalByteLength: 0,
|
||||
};
|
||||
reportCompressEvent(options.telemetry, {
|
||||
outcome: 'passthrough_error',
|
||||
startedAt,
|
||||
inputMime: normalizeMime(mimeType),
|
||||
result,
|
||||
});
|
||||
return result;
|
||||
}
|
||||
const result = await compressImageForModel(bytes, mimeType, options);
|
||||
if (!result.changed) {
|
||||
|
|
@ -279,18 +326,21 @@ export async function compressBase64ForModel(
|
|||
* remote http(s) image) are passed through, as are non-image parts. Best
|
||||
* effort: a part that fails to compress is left unchanged.
|
||||
*
|
||||
* With `annotate` set, every image that was actually re-encoded gains a
|
||||
* {@link buildImageCompressionCaption} text part immediately before it, so the
|
||||
* model knows it is looking at a downsampled copy. `annotate.persistOriginal`
|
||||
* additionally saves the pre-compression bytes and puts the returned path in
|
||||
* the caption so the model can read the original back; persistence failures
|
||||
* degrade to a caption without a path.
|
||||
* With `annotate` set, every image that was actually re-encoded produces a
|
||||
* {@link buildImageCompressionCaption} string in the returned `captions` array
|
||||
* (in encounter order), so the caller can place it next to the image — e.g. as
|
||||
* a leading text part for the model, or on a `note` side channel kept out of
|
||||
* UI-visible output. `annotate.persistOriginal` additionally saves the
|
||||
* pre-compression bytes and puts the returned path in the caption so the model
|
||||
* can read the original back; persistence failures degrade to a caption
|
||||
* without a path.
|
||||
*/
|
||||
export async function compressImageContentParts(
|
||||
parts: readonly ContentPart[],
|
||||
options: CompressImageOptions & { readonly annotate?: CompressAnnotateOptions } = {},
|
||||
): Promise<ContentPart[]> {
|
||||
): Promise<{ parts: ContentPart[]; captions: readonly string[] }> {
|
||||
const { annotate, ...compressOptions } = options;
|
||||
const captions: string[] = [];
|
||||
const out: ContentPart[] = [];
|
||||
for (const part of parts) {
|
||||
if (part.type === 'image_url') {
|
||||
|
|
@ -310,9 +360,8 @@ export async function compressImageContentParts(
|
|||
originalPath = null;
|
||||
}
|
||||
}
|
||||
out.push({
|
||||
type: 'text',
|
||||
text: buildImageCompressionCaption({
|
||||
captions.push(
|
||||
buildImageCompressionCaption({
|
||||
original: {
|
||||
width: result.originalWidth,
|
||||
height: result.originalHeight,
|
||||
|
|
@ -327,7 +376,7 @@ export async function compressImageContentParts(
|
|||
},
|
||||
originalPath,
|
||||
}),
|
||||
});
|
||||
);
|
||||
}
|
||||
out.push({
|
||||
type: 'image_url',
|
||||
|
|
@ -339,7 +388,7 @@ export async function compressImageContentParts(
|
|||
}
|
||||
out.push(part);
|
||||
}
|
||||
return out;
|
||||
return { parts: out, captions };
|
||||
}
|
||||
|
||||
export interface CompressAnnotateOptions {
|
||||
|
|
@ -350,6 +399,46 @@ export interface CompressAnnotateOptions {
|
|||
readonly persistOriginal?: (bytes: Uint8Array, mimeType: string) => Promise<string | null>;
|
||||
}
|
||||
|
||||
interface CompressEventResult {
|
||||
readonly mimeType: string;
|
||||
readonly width: number;
|
||||
readonly height: number;
|
||||
readonly originalWidth: number;
|
||||
readonly originalHeight: number;
|
||||
readonly originalByteLength: number;
|
||||
readonly finalByteLength: number;
|
||||
}
|
||||
|
||||
function reportCompressEvent(
|
||||
telemetry: ImageCompressionTelemetry | undefined,
|
||||
input: {
|
||||
readonly outcome: CompressOutcome;
|
||||
readonly startedAt: number;
|
||||
readonly inputMime: string;
|
||||
readonly result: CompressEventResult;
|
||||
},
|
||||
): void {
|
||||
if (telemetry === undefined) return;
|
||||
try {
|
||||
telemetry.client.track('image_compress', {
|
||||
source: telemetry.source,
|
||||
outcome: input.outcome,
|
||||
input_mime: input.inputMime,
|
||||
output_mime: normalizeMime(input.result.mimeType),
|
||||
original_bytes: input.result.originalByteLength,
|
||||
final_bytes: input.result.finalByteLength,
|
||||
original_width: input.result.originalWidth,
|
||||
original_height: input.result.originalHeight,
|
||||
final_width: input.result.width,
|
||||
final_height: input.result.height,
|
||||
exif_transposed: false,
|
||||
duration_ms: Date.now() - input.startedAt,
|
||||
});
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// ── crop ─────────────────────────────────────────────────────────────
|
||||
|
||||
/** Crop rectangle in ORIGINAL-image pixel coordinates. */
|
||||
|
|
|
|||
|
|
@ -93,7 +93,11 @@ export type LoopRecordedEvent =
|
|||
| {
|
||||
readonly type: 'tool.result';
|
||||
readonly toolCallId: string;
|
||||
readonly result: { readonly output: string | readonly ContentPart[]; readonly isError?: boolean };
|
||||
readonly result: {
|
||||
readonly output: string | readonly ContentPart[];
|
||||
readonly isError?: boolean;
|
||||
readonly note?: string;
|
||||
};
|
||||
readonly parentUuid?: string;
|
||||
};
|
||||
|
||||
|
|
@ -250,24 +254,33 @@ function interruptedToolMessage(toolCallId: string): ContextMessage {
|
|||
function toolResultOutputForModel(result: {
|
||||
readonly output: string | readonly ContentPart[];
|
||||
readonly isError?: boolean;
|
||||
readonly note?: string;
|
||||
}): string | ContentPart[] {
|
||||
const { output, isError } = result;
|
||||
const { output, isError, note } = result;
|
||||
let base: string | ContentPart[];
|
||||
if (typeof output === 'string') {
|
||||
if (isError === true) {
|
||||
if (output.length === 0) return TOOL_EMPTY_ERROR_STATUS;
|
||||
if (output.trimStart().startsWith('<system>ERROR:')) return output;
|
||||
return `${TOOL_ERROR_STATUS}\n${output}`;
|
||||
if (output.length === 0) base = TOOL_EMPTY_ERROR_STATUS;
|
||||
else if (output.trimStart().startsWith('<system>ERROR:')) base = output;
|
||||
else base = `${TOOL_ERROR_STATUS}\n${output}`;
|
||||
} else if (output.length === 0 || output.trim() === TOOL_OUTPUT_EMPTY_TEXT) {
|
||||
base = TOOL_EMPTY_STATUS;
|
||||
} else {
|
||||
base = output;
|
||||
}
|
||||
if (output.length === 0 || output.trim() === TOOL_OUTPUT_EMPTY_TEXT) {
|
||||
return TOOL_EMPTY_STATUS;
|
||||
}
|
||||
return output;
|
||||
} else if (output.length === 0) {
|
||||
base = [{ type: 'text', text: isError === true ? TOOL_EMPTY_ERROR_STATUS : TOOL_EMPTY_STATUS }];
|
||||
} else if (isError === true) {
|
||||
base = [{ type: 'text', text: TOOL_ERROR_STATUS }, ...output];
|
||||
} else {
|
||||
base = [...output];
|
||||
}
|
||||
if (output.length === 0) {
|
||||
return [{ type: 'text', text: isError === true ? TOOL_EMPTY_ERROR_STATUS : TOOL_EMPTY_STATUS }];
|
||||
if (note === undefined || note.length === 0) return base;
|
||||
const notePart: ContentPart = { type: 'text', text: note };
|
||||
if (typeof base === 'string') return `${base}\n${note}`;
|
||||
const only = base[0];
|
||||
if (base.length === 1 && only?.type === 'text') {
|
||||
return [{ type: 'text', text: `${only.text}\n${note}` }];
|
||||
}
|
||||
if (isError === true) {
|
||||
return [{ type: 'text', text: TOOL_ERROR_STATUS }, ...output];
|
||||
}
|
||||
return [...output];
|
||||
return [...base, notePart];
|
||||
}
|
||||
|
|
|
|||
|
|
@ -234,7 +234,7 @@ export class AgentLoopService implements IAgentLoopService {
|
|||
type: 'tool.result',
|
||||
parentUuid: toolCallUuids.get(toolResult.toolCallId) ?? randomUUID(),
|
||||
toolCallId: toolResult.toolCallId,
|
||||
result: { output: result.output, isError: result.isError },
|
||||
result: { output: result.output, isError: result.isError, note: result.note },
|
||||
});
|
||||
if (result.stopTurn === true) stopTurn = true;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ import { SseMcpClient } from './client-sse';
|
|||
import type { UnexpectedCloseReason } from './client-shared';
|
||||
import { StdioMcpClient } from './client-stdio';
|
||||
import type { McpOAuthService } from '#/agent/mcp/oauth/service';
|
||||
import { assertMcpInputSchema, type MCPClient } from './types';
|
||||
import { assertMcpInputSchema, type MCPClient, type MCPToolDefinition } from './types';
|
||||
|
||||
export type McpServerStatus = 'pending' | 'connected' | 'failed' | 'disabled' | 'needs-auth';
|
||||
|
||||
|
|
@ -39,6 +39,7 @@ interface InternalEntry {
|
|||
attemptId: number;
|
||||
status: McpServerStatus;
|
||||
tools?: readonly Tool[];
|
||||
rawTools?: readonly MCPToolDefinition[];
|
||||
enabledNames?: ReadonlySet<string>;
|
||||
error?: string;
|
||||
client?: RuntimeMcpClient;
|
||||
|
|
@ -158,12 +159,18 @@ export class McpConnectionManager {
|
|||
resolved(
|
||||
name: string,
|
||||
):
|
||||
| { client: MCPClient; tools: readonly Tool[]; enabledNames: ReadonlySet<string> }
|
||||
| {
|
||||
client: MCPClient;
|
||||
tools: readonly Tool[];
|
||||
rawTools: readonly MCPToolDefinition[];
|
||||
enabledNames: ReadonlySet<string>;
|
||||
}
|
||||
| undefined {
|
||||
const entry = this.entries.get(name);
|
||||
if (
|
||||
entry?.status !== 'connected' ||
|
||||
entry.tools === undefined ||
|
||||
entry.rawTools === undefined ||
|
||||
entry.client === undefined
|
||||
) {
|
||||
return undefined;
|
||||
|
|
@ -171,6 +178,7 @@ export class McpConnectionManager {
|
|||
return {
|
||||
client: entry.client,
|
||||
tools: entry.tools,
|
||||
rawTools: entry.rawTools,
|
||||
enabledNames: entry.enabledNames ?? new Set(entry.tools.map((t) => t.name)),
|
||||
};
|
||||
}
|
||||
|
|
@ -214,6 +222,7 @@ export class McpConnectionManager {
|
|||
entry.status = 'disabled';
|
||||
entry.tools = undefined;
|
||||
entry.enabledNames = undefined;
|
||||
entry.rawTools = undefined;
|
||||
entry.error = undefined;
|
||||
this.emit(entry);
|
||||
this.entries.delete(name);
|
||||
|
|
@ -265,6 +274,7 @@ export class McpConnectionManager {
|
|||
entry.status = 'pending';
|
||||
entry.tools = undefined;
|
||||
entry.enabledNames = undefined;
|
||||
entry.rawTools = undefined;
|
||||
entry.error = undefined;
|
||||
this.emit(entry);
|
||||
await this.connectOne(entry, attemptId);
|
||||
|
|
@ -285,7 +295,7 @@ export class McpConnectionManager {
|
|||
const startupClient = await this.createClient(entry.config, entry.name);
|
||||
client = startupClient;
|
||||
entry.client = startupClient;
|
||||
const tools = await withTimeout(
|
||||
const discovered = await withTimeout(
|
||||
this.connectAndDiscoverTools(startupClient),
|
||||
timeoutMs,
|
||||
() => {
|
||||
|
|
@ -297,8 +307,9 @@ export class McpConnectionManager {
|
|||
await this.closeRuntimeClient(startupClient);
|
||||
return;
|
||||
}
|
||||
entry.tools = tools;
|
||||
entry.enabledNames = computeEnabledNames(entry.config, tools);
|
||||
entry.tools = discovered.tools;
|
||||
entry.rawTools = discovered.rawTools;
|
||||
entry.enabledNames = computeEnabledNames(entry.config, discovered.tools);
|
||||
entry.status = 'connected';
|
||||
this.watchForUnexpectedClose(entry, startupClient, attemptId);
|
||||
} catch (error) {
|
||||
|
|
@ -317,6 +328,7 @@ export class McpConnectionManager {
|
|||
}
|
||||
entry.tools = undefined;
|
||||
entry.enabledNames = undefined;
|
||||
entry.rawTools = undefined;
|
||||
// Drop the client reference so a later reconnect builds a fresh one.
|
||||
await this.closeClient(entry);
|
||||
}
|
||||
|
|
@ -338,6 +350,7 @@ export class McpConnectionManager {
|
|||
entry.error = formatUnexpectedCloseError(entry.name, reason);
|
||||
entry.tools = undefined;
|
||||
entry.enabledNames = undefined;
|
||||
entry.rawTools = undefined;
|
||||
entry.client = undefined;
|
||||
// Best-effort close; the transport is already gone, but this lets the
|
||||
// SDK release timers and pending request handlers.
|
||||
|
|
@ -398,14 +411,19 @@ export class McpConnectionManager {
|
|||
return isUnauthorizedLikeError(error);
|
||||
}
|
||||
|
||||
private async connectAndDiscoverTools(client: RuntimeMcpClient): Promise<Tool[]> {
|
||||
private async connectAndDiscoverTools(
|
||||
client: RuntimeMcpClient,
|
||||
): Promise<{ tools: Tool[]; rawTools: MCPToolDefinition[] }> {
|
||||
await client.connect();
|
||||
const mcpTools = await client.listTools();
|
||||
return mcpTools.map((mcpTool) => ({
|
||||
name: mcpTool.name,
|
||||
description: mcpTool.description,
|
||||
parameters: assertMcpInputSchema(mcpTool.name, mcpTool.inputSchema),
|
||||
}));
|
||||
return {
|
||||
rawTools: mcpTools,
|
||||
tools: mcpTools.map((mcpTool) => ({
|
||||
name: mcpTool.name,
|
||||
description: mcpTool.description,
|
||||
parameters: assertMcpInputSchema(mcpTool.name, mcpTool.inputSchema),
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
private async closeClient(entry: InternalEntry): Promise<void> {
|
||||
|
|
|
|||
|
|
@ -7,11 +7,12 @@ import type {
|
|||
McpServerEntry,
|
||||
} from './connection-manager';
|
||||
import type { McpOAuthService } from '#/agent/mcp/oauth/service';
|
||||
import type { MCPClient } from './types';
|
||||
import type { MCPClient, MCPToolDefinition } from './types';
|
||||
|
||||
export interface McpResolvedServer {
|
||||
readonly client: MCPClient;
|
||||
readonly tools: readonly KosongTool[];
|
||||
readonly rawTools: readonly MCPToolDefinition[];
|
||||
readonly enabledNames: ReadonlySet<string>;
|
||||
}
|
||||
|
||||
|
|
|
|||
48
packages/agent-core-v2/src/agent/mcp/mcpDiscoveryOps.ts
Normal file
48
packages/agent-core-v2/src/agent/mcp/mcpDiscoveryOps.ts
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
/**
|
||||
* `mcp` domain (L5) — MCP tool-discovery wire state.
|
||||
*
|
||||
* Restores the per-agent de-dup cursor for durable MCP discovery records,
|
||||
* keyed by `${serverName}\n${hash}` entries already present in this log.
|
||||
*/
|
||||
|
||||
import { defineModel } from '#/wire/model';
|
||||
import { defineOp } from '#/wire/op';
|
||||
import type { MCPToolDefinition } from './types';
|
||||
|
||||
export interface McpToolCollision {
|
||||
readonly qualified: string;
|
||||
readonly toolName: string;
|
||||
readonly collidesWith:
|
||||
| { readonly kind: 'same_server'; readonly toolName: string }
|
||||
| { readonly kind: 'other_server'; readonly serverName: string };
|
||||
}
|
||||
|
||||
export interface McpDiscoveryState {
|
||||
readonly seen: readonly string[];
|
||||
}
|
||||
|
||||
export const McpDiscoveryModel = defineModel<McpDiscoveryState>('mcp.discovery', () => ({
|
||||
seen: [],
|
||||
}));
|
||||
|
||||
export interface McpToolsDiscoveredPayload {
|
||||
readonly serverName: string;
|
||||
readonly hash: string;
|
||||
readonly tools: readonly MCPToolDefinition[];
|
||||
readonly enabledNames: readonly string[];
|
||||
readonly collisions?: readonly McpToolCollision[];
|
||||
}
|
||||
|
||||
export const mcpToolsDiscovered = defineOp(McpDiscoveryModel, 'mcp.tools_discovered', {
|
||||
apply: (s, p: McpToolsDiscoveredPayload): McpDiscoveryState => {
|
||||
const key = `${p.serverName}\n${p.hash}`;
|
||||
if (s.seen.includes(key)) return s;
|
||||
return { seen: [...s.seen, key] };
|
||||
},
|
||||
});
|
||||
|
||||
declare module '#/agent/wireRecord/wireRecord' {
|
||||
interface WireRecordMap {
|
||||
'mcp.tools_discovered': McpToolsDiscoveredPayload;
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,5 @@
|
|||
import { createHash } from 'node:crypto';
|
||||
|
||||
import { InstantiationType } from '#/_base/di/extensions';
|
||||
import { LifecycleScope, registerScopedService } from '#/_base/di/scope';
|
||||
import type { Tool as KosongTool } from '#/app/llmProtocol/tool';
|
||||
|
|
@ -17,7 +19,14 @@ import { createMcpTool } from '#/agent/mcp/tools/mcp';
|
|||
import type { McpServerEntry } from './connection-manager';
|
||||
import { IAgentMcpService, type McpServiceOptions } from './mcp';
|
||||
import { qualifyMcpToolName } from './tool-naming';
|
||||
import type { MCPClient } from './types';
|
||||
import type { MCPClient, MCPToolDefinition } from './types';
|
||||
import { IAgentWireService } from '#/wire/tokens';
|
||||
import type { IWireService } from '#/wire/wireService';
|
||||
import {
|
||||
McpDiscoveryModel,
|
||||
mcpToolsDiscovered,
|
||||
type McpToolCollision,
|
||||
} from './mcpDiscoveryOps';
|
||||
|
||||
declare module '#/app/event/eventBus' {
|
||||
interface DomainEventMap {
|
||||
|
|
@ -34,24 +43,19 @@ interface McpToolRegistration {
|
|||
readonly serverName: string;
|
||||
}
|
||||
|
||||
interface McpToolCollision {
|
||||
readonly qualified: string;
|
||||
readonly toolName: string;
|
||||
readonly collidesWith:
|
||||
| { readonly kind: 'same_server'; readonly toolName: string }
|
||||
| { readonly kind: 'other_server'; readonly serverName: string };
|
||||
}
|
||||
|
||||
export class AgentMcpService extends Disposable implements IAgentMcpService {
|
||||
declare readonly _serviceBrand: undefined;
|
||||
private readonly mcpTools = new Map<string, McpToolRegistration>();
|
||||
private readonly mcpToolsByServer = new Map<string, string[]>();
|
||||
private readonly pendingDiscoveries: Array<() => void> = [];
|
||||
private discoveryWritesReady = false;
|
||||
|
||||
constructor(
|
||||
private readonly options: McpServiceOptions = {},
|
||||
@IAgentToolRegistryService private readonly registry: IAgentToolRegistryService,
|
||||
@IEventBus private readonly eventBus: IEventBus,
|
||||
@IAgentToolExecutorService toolExecutor: IAgentToolExecutorService,
|
||||
@IAgentWireService private readonly wire: IWireService,
|
||||
) {
|
||||
super();
|
||||
this.attachMcpTools();
|
||||
|
|
@ -64,6 +68,8 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
|
|||
},
|
||||
),
|
||||
);
|
||||
this._register(this.wire.onRestored(() => this.flushPendingDiscoveries()));
|
||||
this._register(this.wire.onEmission(() => this.flushPendingDiscoveries()));
|
||||
}
|
||||
|
||||
get oauthService() {
|
||||
|
|
@ -164,6 +170,7 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
|
|||
resolved.enabledNames,
|
||||
);
|
||||
this.emitMcpToolCollisions(entry.name, result.collisions);
|
||||
this.recordDiscovery(entry.name, resolved.rawTools, resolved.enabledNames, result.collisions);
|
||||
this.eventBus.publish({
|
||||
type: 'tool.list.updated',
|
||||
reason: 'mcp.connected',
|
||||
|
|
@ -252,6 +259,44 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
|
|||
return true;
|
||||
}
|
||||
|
||||
private recordDiscovery(
|
||||
serverName: string,
|
||||
rawTools: readonly MCPToolDefinition[],
|
||||
enabledNames: ReadonlySet<string>,
|
||||
collisions: readonly McpToolCollision[],
|
||||
): void {
|
||||
const enabledNamesSnapshot = [...enabledNames].toSorted((a, b) => a.localeCompare(b));
|
||||
const work = (): void => {
|
||||
const hash = createHash('sha256')
|
||||
.update(JSON.stringify({ tools: rawTools, enabledNames: enabledNamesSnapshot, collisions }))
|
||||
.digest('hex');
|
||||
const key = `${serverName}\n${hash}`;
|
||||
if (this.wire.getModel(McpDiscoveryModel).seen.includes(key)) return;
|
||||
this.wire.dispatch(
|
||||
mcpToolsDiscovered({
|
||||
serverName,
|
||||
hash,
|
||||
tools: rawTools,
|
||||
enabledNames: enabledNamesSnapshot,
|
||||
collisions: collisions.length > 0 ? collisions : undefined,
|
||||
}),
|
||||
);
|
||||
};
|
||||
if (!this.discoveryWritesReady) {
|
||||
this.pendingDiscoveries.push(work);
|
||||
return;
|
||||
}
|
||||
work();
|
||||
}
|
||||
|
||||
private flushPendingDiscoveries(): void {
|
||||
this.discoveryWritesReady = true;
|
||||
const pending = this.pendingDiscoveries.splice(0);
|
||||
for (const work of pending) {
|
||||
work();
|
||||
}
|
||||
}
|
||||
|
||||
private emitMcpToolCollisions(
|
||||
serverName: string,
|
||||
collisions: readonly McpToolCollision[],
|
||||
|
|
|
|||
|
|
@ -14,7 +14,9 @@
|
|||
* would silently reintroduce the very degradation the caption reports.
|
||||
* 4. Compress oversized inline images, announcing each compression with a
|
||||
* caption (original vs. sent size, readback path to the persisted
|
||||
* original) so downsampling is never silent.
|
||||
* original) so downsampling is never silent. The captions ride the
|
||||
* result's `note` side channel — projected to the model at fold time, but
|
||||
* kept out of `output` so UIs never render them.
|
||||
* 5. Apply the per-part 10 MB binary cap: oversized binary parts
|
||||
* (image/audio/video URLs) collapse to a notice, so a single
|
||||
* screenshot cannot evict every text part.
|
||||
|
|
@ -26,6 +28,7 @@
|
|||
*/
|
||||
|
||||
import type { ContentPart } from '#/app/llmProtocol/message';
|
||||
import type { ITelemetryService } from '#/app/telemetry/telemetry';
|
||||
|
||||
import { compressImageContentParts } from '#/_base/tools/support/image-compress';
|
||||
import { persistOriginalImage } from '#/_base/tools/support/image-originals';
|
||||
|
|
@ -38,6 +41,7 @@ export interface McpOutputOptions {
|
|||
* Falls back to the shared temp-dir cache when absent.
|
||||
*/
|
||||
readonly originalsDir?: string;
|
||||
readonly telemetry?: ITelemetryService;
|
||||
}
|
||||
|
||||
// MCP servers can produce arbitrarily large outputs; cap what we feed back to
|
||||
|
|
@ -151,7 +155,12 @@ export async function mcpResultToExecutableOutput(
|
|||
result: MCPToolResult,
|
||||
qualifiedToolName: string,
|
||||
options: McpOutputOptions = {},
|
||||
): Promise<{ output: string | ContentPart[]; isError: boolean; truncated?: true }> {
|
||||
): Promise<{
|
||||
output: string | ContentPart[];
|
||||
isError: boolean;
|
||||
note?: string;
|
||||
truncated?: true;
|
||||
}> {
|
||||
const converted: ContentPart[] = [];
|
||||
for (const block of result.content) {
|
||||
const part = convertMCPContentBlock(block);
|
||||
|
|
@ -174,6 +183,10 @@ export async function mcpResultToExecutableOutput(
|
|||
// known) so the model can read detail back via ReadMediaFile + region.
|
||||
// Parts that cannot be compressed pass through.
|
||||
const compressed = await compressImageContentParts(budgeted.parts, {
|
||||
telemetry:
|
||||
options.telemetry === undefined
|
||||
? undefined
|
||||
: { client: options.telemetry, source: 'mcp_tool_result' },
|
||||
annotate: {
|
||||
persistOriginal: (bytes, mimeType) =>
|
||||
persistOriginalImage(
|
||||
|
|
@ -183,12 +196,16 @@ export async function mcpResultToExecutableOutput(
|
|||
),
|
||||
},
|
||||
});
|
||||
const capped = applyBinaryPartCap(compressed);
|
||||
const capped = applyBinaryPartCap(compressed.parts);
|
||||
const truncated = budgeted.truncated || capped.truncated;
|
||||
const output = collapseSingleText(capped.parts);
|
||||
return truncated
|
||||
? { output, isError: result.isError, truncated: true }
|
||||
: { output, isError: result.isError };
|
||||
const note = compressed.captions.length > 0 ? compressed.captions.join('\n') : undefined;
|
||||
return {
|
||||
output,
|
||||
isError: result.isError,
|
||||
note,
|
||||
truncated: truncated ? true : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -43,14 +43,15 @@ export function createMcpTool(
|
|||
function normalizeMcpToolResult(result: {
|
||||
readonly output: ExecutableToolResult['output'];
|
||||
readonly isError: boolean;
|
||||
readonly note?: string;
|
||||
readonly truncated?: true;
|
||||
}): ExecutableToolResult {
|
||||
if (result.isError) {
|
||||
return result.truncated === true
|
||||
? { output: result.output, isError: true, truncated: true }
|
||||
: { output: result.output, isError: true };
|
||||
? { output: result.output, isError: true, note: result.note, truncated: true }
|
||||
: { output: result.output, isError: true, note: result.note };
|
||||
}
|
||||
return result.truncated === true
|
||||
? { output: result.output, truncated: true }
|
||||
: { output: result.output };
|
||||
? { output: result.output, note: result.note, truncated: true }
|
||||
: { output: result.output, note: result.note };
|
||||
}
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ export interface ExecutableToolSuccessResult {
|
|||
readonly stopTurn?: boolean | undefined;
|
||||
readonly message?: string | undefined;
|
||||
readonly truncated?: boolean | undefined;
|
||||
readonly note?: string;
|
||||
readonly delivery?: ToolDelivery | undefined;
|
||||
}
|
||||
|
||||
|
|
@ -60,6 +61,7 @@ export interface ExecutableToolErrorResult {
|
|||
readonly message?: string | undefined;
|
||||
readonly stopTurn?: boolean | undefined;
|
||||
readonly truncated?: boolean | undefined;
|
||||
readonly note?: string;
|
||||
readonly delivery?: ToolDelivery | undefined;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -731,9 +731,20 @@ function normalizeToolResult(result: ExecutableToolResult): ToolResult {
|
|||
}
|
||||
}
|
||||
if (result.isError === true) {
|
||||
return { output, isError: true, stopTurn: result.stopTurn };
|
||||
return {
|
||||
output,
|
||||
isError: true,
|
||||
stopTurn: result.stopTurn,
|
||||
truncated: result.truncated,
|
||||
note: result.note,
|
||||
};
|
||||
}
|
||||
return { output, stopTurn: result.stopTurn };
|
||||
return {
|
||||
output,
|
||||
stopTurn: result.stopTurn,
|
||||
truncated: result.truncated,
|
||||
note: result.note,
|
||||
};
|
||||
}
|
||||
|
||||
function toolTelemetryOutcome(result: ToolResult): 'success' | 'error' | 'cancelled' {
|
||||
|
|
|
|||
|
|
@ -26,6 +26,8 @@ import {
|
|||
registerScopedService,
|
||||
} from '#/_base/di/scope';
|
||||
import { IBootstrapService } from '#/app/bootstrap/bootstrap';
|
||||
import { IEventBus } from '#/app/event/eventBus';
|
||||
import { ErrorCodes, makeErrorPayload } from '#/errors';
|
||||
import { ITelemetryService } from '#/app/telemetry/telemetry';
|
||||
import { ILogService } from '#/_base/log/log';
|
||||
import { IAgentProfileCatalogService } from '#/app/agentProfileCatalog/agentProfileCatalog';
|
||||
|
|
@ -245,6 +247,11 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
|
|||
const manager = this.getMcpManager();
|
||||
const initialLoad = this.connectMcpServers(manager).catch((error: unknown) => {
|
||||
this.log.error('mcp initial load failed', { error });
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
this.handles.get('main')?.accessor.get(IEventBus)?.publish({
|
||||
type: 'error',
|
||||
...makeErrorPayload(ErrorCodes.MCP_STARTUP_FAILED, message),
|
||||
});
|
||||
});
|
||||
this.mcpInitialLoad = initialLoad;
|
||||
return initialLoad;
|
||||
|
|
|
|||
|
|
@ -106,4 +106,45 @@ describe('loop-event fold parity', () => {
|
|||
|
||||
expect(folded).toEqual(baseline);
|
||||
});
|
||||
|
||||
it('folds a tool-result note as trailing model text without splitting text-only output', () => {
|
||||
context.append(
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [],
|
||||
toolCalls: [{ type: 'function', id: 'c3', name: 'Screenshot', arguments: '{}' }],
|
||||
},
|
||||
{
|
||||
role: 'tool',
|
||||
content: [{ type: 'text', text: 'result text\n<system>Image compressed.</system>' }],
|
||||
toolCalls: [],
|
||||
toolCallId: 'c3',
|
||||
isError: false,
|
||||
},
|
||||
);
|
||||
const baseline = comparable(context.get());
|
||||
context.clear();
|
||||
|
||||
context.appendLoopEvent({ type: 'step.begin', uuid: 's3' });
|
||||
context.appendLoopEvent({
|
||||
type: 'tool.call',
|
||||
stepUuid: 's3',
|
||||
toolCallId: 'c3',
|
||||
name: 'Screenshot',
|
||||
args: {},
|
||||
});
|
||||
context.appendLoopEvent({
|
||||
type: 'tool.result',
|
||||
toolCallId: 'c3',
|
||||
result: {
|
||||
output: 'result text',
|
||||
isError: false,
|
||||
note: '<system>Image compressed.</system>',
|
||||
},
|
||||
});
|
||||
context.appendLoopEvent({ type: 'step.end', uuid: 's3' });
|
||||
const folded = comparable(context.get());
|
||||
|
||||
expect(folded).toEqual(baseline);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -10,7 +10,12 @@ import type { McpConnectionManager, McpServerEntry } from '#/agent/mcp/connectio
|
|||
import { IAgentMcpService } from '#/agent/mcp/mcp';
|
||||
import { AgentMcpService } from '#/agent/mcp/mcpService';
|
||||
import type { McpOAuthService } from '#/agent/mcp/oauth/service';
|
||||
import type { MCPClient } from '#/agent/mcp/types';
|
||||
import type { MCPClient, MCPToolDefinition } from '#/agent/mcp/types';
|
||||
import { IAgentWireService } from '#/wire/tokens';
|
||||
import { WireService } from '#/wire/wireServiceImpl';
|
||||
import { McpDiscoveryModel } from '#/agent/mcp/mcpDiscoveryOps';
|
||||
import { AGENT_WIRE_PROTOCOL_VERSION } from '#/agent/wireRecord/wireRecord';
|
||||
import { wireMetadata } from '#/agent/wireRecord/metadataOps';
|
||||
import { AgentToolExecutorService } from '#/agent/toolExecutor/toolExecutorService';
|
||||
import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor';
|
||||
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
|
||||
|
|
@ -29,6 +34,7 @@ const MCP_OUTPUT_TRUNCATED_TEXT =
|
|||
interface ResolvedServer {
|
||||
readonly client: MCPClient;
|
||||
readonly tools: readonly KosongTool[];
|
||||
readonly rawTools: readonly MCPToolDefinition[];
|
||||
readonly enabledNames: ReadonlySet<string>;
|
||||
}
|
||||
|
||||
|
|
@ -74,8 +80,21 @@ class FakeMcpManager {
|
|||
client: MCPClient,
|
||||
tools: readonly KosongTool[],
|
||||
enabledNames = new Set(tools.map((tool) => tool.name)),
|
||||
rawTools?: readonly MCPToolDefinition[],
|
||||
): void {
|
||||
this.resolvedEntries.set(name, { client, tools, enabledNames });
|
||||
const resolvedRawTools =
|
||||
rawTools ??
|
||||
tools.map((tool) => ({
|
||||
name: tool.name,
|
||||
description: tool.description ?? '',
|
||||
inputSchema: (tool.parameters ?? {}) as MCPToolDefinition['inputSchema'],
|
||||
}));
|
||||
this.resolvedEntries.set(name, {
|
||||
client,
|
||||
tools,
|
||||
rawTools: resolvedRawTools,
|
||||
enabledNames,
|
||||
});
|
||||
}
|
||||
|
||||
connect(name: string, options: { readonly transport?: 'stdio' | 'http' | 'sse' } = {}): void {
|
||||
|
|
@ -128,6 +147,7 @@ describe('AgentMcpService', () => {
|
|||
let disposables: DisposableStore;
|
||||
let ix: TestInstantiationService;
|
||||
let events: DomainEvent[];
|
||||
let wire: WireService;
|
||||
|
||||
beforeEach(() => {
|
||||
disposables = new DisposableStore();
|
||||
|
|
@ -142,6 +162,8 @@ describe('AgentMcpService', () => {
|
|||
ix.set(IAgentToolRegistryService, new SyncDescriptor(AgentToolRegistryService));
|
||||
ix.set(IAgentToolExecutorService, new SyncDescriptor(AgentToolExecutorService));
|
||||
ix.stub(IAgentTurnService, stubTurnWithHooks());
|
||||
wire = disposables.add(new WireService({ logScope: 'mcp-test', logKey: 'wire.jsonl' }));
|
||||
ix.stub(IAgentWireService, wire);
|
||||
});
|
||||
afterEach(() => {
|
||||
disposables.dispose();
|
||||
|
|
@ -469,6 +491,200 @@ describe('AgentMcpService', () => {
|
|||
}),
|
||||
);
|
||||
});
|
||||
|
||||
const RAW_QUERY: MCPToolDefinition = {
|
||||
name: 'query_range',
|
||||
description: 'Query a metrics range',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: { query: { type: 'string' } },
|
||||
required: ['query'],
|
||||
},
|
||||
};
|
||||
|
||||
function collectDiscoveries(): {
|
||||
records: { type: string; [key: string]: unknown }[];
|
||||
off: { dispose(): void };
|
||||
} {
|
||||
const records: { type: string; [key: string]: unknown }[] = [];
|
||||
const off = wire.onEmission((e) => {
|
||||
if (e.record.type === 'mcp.tools_discovered') {
|
||||
records.push(e.record as { type: string; [key: string]: unknown });
|
||||
}
|
||||
});
|
||||
return { records, off };
|
||||
}
|
||||
|
||||
it('records tools/list once after restore and dedups unchanged reconnects', async () => {
|
||||
const manager = new FakeMcpManager();
|
||||
const client = fakeMcpClient([RAW_QUERY]);
|
||||
const rawTools = await client.listTools();
|
||||
manager.setResolved(
|
||||
'grafana',
|
||||
client,
|
||||
await discoverTools(client),
|
||||
new Set(['query_range']),
|
||||
rawTools,
|
||||
);
|
||||
createService(manager);
|
||||
|
||||
const { records, off } = collectDiscoveries();
|
||||
try {
|
||||
manager.connect('grafana');
|
||||
expect(records).toHaveLength(0); // parked until restore
|
||||
await wire.replay();
|
||||
expect(records).toHaveLength(1);
|
||||
expect(records[0]).toMatchObject({
|
||||
type: 'mcp.tools_discovered',
|
||||
serverName: 'grafana',
|
||||
tools: rawTools,
|
||||
enabledNames: ['query_range'],
|
||||
});
|
||||
expect(records[0]!['collisions']).toBeUndefined();
|
||||
|
||||
// identical content -> no second record
|
||||
manager.connect('grafana');
|
||||
expect(records).toHaveLength(1);
|
||||
|
||||
// allow-list change is a different gating decision -> record again
|
||||
manager.setResolved('grafana', client, await discoverTools(client), new Set(), rawTools);
|
||||
manager.connect('grafana');
|
||||
expect(records).toHaveLength(2);
|
||||
} finally {
|
||||
off.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('parks a discovery observed before restore and flushes it after replay', async () => {
|
||||
const manager = new FakeMcpManager();
|
||||
const client = fakeMcpClient([RAW_QUERY]);
|
||||
const rawTools = await client.listTools();
|
||||
manager.setResolved(
|
||||
'grafana',
|
||||
client,
|
||||
await discoverTools(client),
|
||||
new Set(['query_range']),
|
||||
rawTools,
|
||||
);
|
||||
createService(manager);
|
||||
|
||||
const { records, off } = collectDiscoveries();
|
||||
try {
|
||||
manager.connect('grafana');
|
||||
expect(records).toHaveLength(0); // parked, not yet durable
|
||||
await wire.replay();
|
||||
expect(records).toHaveLength(1);
|
||||
} finally {
|
||||
off.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('snapshots enabledNames when parking a discovery before restore', async () => {
|
||||
const manager = new FakeMcpManager();
|
||||
const client = fakeMcpClient([RAW_QUERY]);
|
||||
const rawTools = await client.listTools();
|
||||
const enabledNames = new Set(['query_range']);
|
||||
manager.setResolved(
|
||||
'grafana',
|
||||
client,
|
||||
await discoverTools(client),
|
||||
enabledNames,
|
||||
rawTools,
|
||||
);
|
||||
createService(manager);
|
||||
|
||||
const { records, off } = collectDiscoveries();
|
||||
try {
|
||||
manager.connect('grafana');
|
||||
enabledNames.clear();
|
||||
enabledNames.add('mutated_after_observation');
|
||||
await wire.replay();
|
||||
|
||||
expect(records).toHaveLength(1);
|
||||
expect(records[0]).toMatchObject({
|
||||
type: 'mcp.tools_discovered',
|
||||
serverName: 'grafana',
|
||||
tools: rawTools,
|
||||
enabledNames: ['query_range'],
|
||||
});
|
||||
} finally {
|
||||
off.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('flushes a parked discovery after the first live wire record on a fresh session', async () => {
|
||||
const manager = new FakeMcpManager();
|
||||
const client = fakeMcpClient([RAW_QUERY]);
|
||||
const rawTools = await client.listTools();
|
||||
manager.setResolved(
|
||||
'grafana',
|
||||
client,
|
||||
await discoverTools(client),
|
||||
new Set(['query_range']),
|
||||
rawTools,
|
||||
);
|
||||
createService(manager);
|
||||
|
||||
const { records, off } = collectDiscoveries();
|
||||
try {
|
||||
manager.connect('grafana');
|
||||
expect(records).toHaveLength(0);
|
||||
wire.dispatch(
|
||||
wireMetadata({
|
||||
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
|
||||
created_at: 1,
|
||||
}),
|
||||
);
|
||||
expect(records).toHaveLength(1);
|
||||
expect(records[0]).toMatchObject({
|
||||
type: 'mcp.tools_discovered',
|
||||
serverName: 'grafana',
|
||||
tools: rawTools,
|
||||
enabledNames: ['query_range'],
|
||||
});
|
||||
} finally {
|
||||
off.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('re-records when only the collision outcome changes', async () => {
|
||||
const manager = new FakeMcpManager();
|
||||
const occupant = fakeMcpClient([RAW_QUERY]);
|
||||
const occupantRaw = await occupant.listTools();
|
||||
manager.setResolved(
|
||||
'graf.ana',
|
||||
occupant,
|
||||
await discoverTools(occupant),
|
||||
new Set(['query_range']),
|
||||
occupantRaw,
|
||||
);
|
||||
createService(manager);
|
||||
manager.connect('graf.ana');
|
||||
await wire.replay(); // restore; occupant discovery recorded (before we subscribe)
|
||||
|
||||
const { records, off } = collectDiscoveries();
|
||||
try {
|
||||
const client = fakeMcpClient([RAW_QUERY]);
|
||||
const rawTools = await client.listTools();
|
||||
manager.setResolved(
|
||||
'graf_ana',
|
||||
client,
|
||||
await discoverTools(client),
|
||||
new Set(['query_range']),
|
||||
rawTools,
|
||||
);
|
||||
manager.connect('graf_ana'); // collides with the occupant's qualified name
|
||||
expect(records).toHaveLength(1);
|
||||
expect(records[0]!['collisions']).toHaveLength(1);
|
||||
|
||||
manager.disconnect('graf.ana'); // occupant gone
|
||||
manager.connect('graf_ana'); // same rawTools/allow-list, collision flipped
|
||||
expect(records).toHaveLength(2);
|
||||
expect(records[1]!['collisions']).toBeUndefined();
|
||||
} finally {
|
||||
off.dispose();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('AgentMcpService + AgentProfileService', () => {
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { tmpdir } from 'node:os';
|
|||
import { join } from 'node:path';
|
||||
import { describe, expect, test } from 'vitest';
|
||||
|
||||
import type { ITelemetryService } from '#/app/telemetry/telemetry';
|
||||
import { convertMCPContentBlock, mcpResultToExecutableOutput } from '#/agent/mcp/output';
|
||||
import { createMcpTool } from '#/agent/mcp/tools/mcp';
|
||||
import type { MCPClient, MCPContentBlock, MCPToolResult } from '#/agent/mcp/types';
|
||||
|
|
@ -28,6 +29,29 @@ function assertValidMcpBlock<T extends MCPContentBlock>(block: T): T {
|
|||
return block;
|
||||
}
|
||||
|
||||
interface TelemetryRecord {
|
||||
readonly event: string;
|
||||
readonly properties: Readonly<Record<string, unknown>> | undefined;
|
||||
}
|
||||
|
||||
function recordingTelemetry(records: TelemetryRecord[]): ITelemetryService {
|
||||
const telemetry: ITelemetryService = {
|
||||
_serviceBrand: undefined,
|
||||
track(event, properties) {
|
||||
records.push({ event, properties });
|
||||
},
|
||||
withContext: () => telemetry,
|
||||
setContext: () => {},
|
||||
addAppender: () => ({ dispose: () => {} }),
|
||||
removeAppender: () => {},
|
||||
setAppender: () => {},
|
||||
setEnabled: () => {},
|
||||
flush: async () => {},
|
||||
shutdown: async () => {},
|
||||
};
|
||||
return telemetry;
|
||||
}
|
||||
|
||||
describe('convertMCPContentBlock', () => {
|
||||
test('converts text block to TextPart', () => {
|
||||
const block: MCPContentBlock = { type: 'text', text: 'hello' };
|
||||
|
|
@ -321,7 +345,7 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
{ type: 'text', text: 'A'.repeat(100_000) },
|
||||
{ type: 'image_url', imageUrl: { url: 'data:image/png;base64,' + 'B'.repeat(500_000) } },
|
||||
]);
|
||||
expect(out).not.toHaveProperty('truncated');
|
||||
expect(out.truncated).toBeUndefined();
|
||||
});
|
||||
|
||||
test('downsamples an oversized real image instead of leaving it full-size', async () => {
|
||||
|
|
@ -358,15 +382,12 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
);
|
||||
|
||||
const parts = out.output as ContentPart[];
|
||||
const captionIndex = parts.findIndex(
|
||||
(p) => p.type === 'text' && p.text.includes('Image compressed'),
|
||||
);
|
||||
expect(captionIndex).toBeGreaterThanOrEqual(0);
|
||||
const caption = (parts[captionIndex] as { text: string }).text;
|
||||
const caption = out.note;
|
||||
expect(caption).toContain('Image compressed');
|
||||
expect(caption).toContain('2600x2600');
|
||||
expect(parts[captionIndex + 1]?.type).toBe('image_url');
|
||||
expect(parts.some((p) => p.type === 'image_url')).toBe(true);
|
||||
|
||||
const pathMatch = /saved at "([^"]+)"/.exec(caption);
|
||||
const pathMatch = /saved at "([^"]+)"/.exec(caption!);
|
||||
expect(pathMatch).not.toBeNull();
|
||||
const persisted = await readFile(pathMatch![1]!);
|
||||
expect(persisted.equals(bigBytes)).toBe(true);
|
||||
|
|
@ -383,9 +404,38 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
'mcp__s__shot',
|
||||
);
|
||||
|
||||
const parts = out.output as ContentPart[];
|
||||
const joined = parts.map((p) => (p.type === 'text' ? p.text : '')).join('');
|
||||
expect(joined).not.toContain('Image compressed');
|
||||
expect(out.note).toBeUndefined();
|
||||
});
|
||||
|
||||
test('reports MCP image compression telemetry with the MCP tool-result source', async () => {
|
||||
const records: TelemetryRecord[] = [];
|
||||
const big = Buffer.from(
|
||||
await new Jimp({ width: 2600, height: 2600, color: 0x3366ccff }).getBuffer('image/png'),
|
||||
).toString('base64');
|
||||
|
||||
await mcpResultToExecutableOutput(
|
||||
result([{ type: 'image', data: big, mimeType: 'image/png' }]),
|
||||
'mcp__s__shot',
|
||||
{ telemetry: recordingTelemetry(records) },
|
||||
);
|
||||
|
||||
const events = records.filter((record) => record.event === 'image_compress');
|
||||
expect(events).toHaveLength(1);
|
||||
const properties = events[0]!.properties;
|
||||
expect(properties).toEqual(
|
||||
expect.objectContaining({
|
||||
source: 'mcp_tool_result',
|
||||
outcome: 'compressed',
|
||||
input_mime: 'image/png',
|
||||
output_mime: 'image/png',
|
||||
original_width: 2600,
|
||||
original_height: 2600,
|
||||
exif_transposed: false,
|
||||
}),
|
||||
);
|
||||
expect(properties?.['final_width']).toBeLessThanOrEqual(2000);
|
||||
expect(properties?.['final_height']).toBeLessThanOrEqual(2000);
|
||||
expect(properties?.['duration_ms']).toEqual(expect.any(Number));
|
||||
});
|
||||
|
||||
test('persists originals into the provided session originals dir', async () => {
|
||||
|
|
@ -400,10 +450,9 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
{ originalsDir: dir },
|
||||
);
|
||||
|
||||
const parts = out.output as ContentPart[];
|
||||
const caption = parts.find((p) => p.type === 'text' && p.text.includes('Image compressed'));
|
||||
if (caption?.type !== 'text') throw new Error('expected a compression caption');
|
||||
const pathMatch = /saved at "([^"]+)"/.exec(caption.text);
|
||||
const caption = out.note;
|
||||
expect(caption).toContain('Image compressed');
|
||||
const pathMatch = /saved at "([^"]+)"/.exec(caption!);
|
||||
expect(pathMatch).not.toBeNull();
|
||||
expect(pathMatch![1]!.startsWith(dir)).toBe(true);
|
||||
const persisted = await readFile(pathMatch![1]!);
|
||||
|
|
@ -432,10 +481,8 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
const toolText = parts[0];
|
||||
if (toolText?.type !== 'text') throw new Error('expected the tool text part first');
|
||||
expect(toolText.text).toContain('Output truncated');
|
||||
const caption = parts.find((p) => p.type === 'text' && p.text.includes('Image compressed'));
|
||||
if (caption?.type !== 'text') throw new Error('expected a compression caption');
|
||||
expect(caption.text).toMatch(/<\/system>$/);
|
||||
expect(caption.text).toContain('saved at');
|
||||
expect(out.note).toMatch(/<\/system>$/);
|
||||
expect(out.note).toContain('saved at');
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
|
|
@ -454,13 +501,11 @@ describe('mcpResultToExecutableOutput', () => {
|
|||
{ originalsDir: dir },
|
||||
);
|
||||
|
||||
const parts = out.output as ContentPart[];
|
||||
expect(out.truncated).toBeUndefined();
|
||||
const caption = parts.find((p) => p.type === 'text' && p.text.includes('Image compressed'));
|
||||
if (caption?.type !== 'text') throw new Error('expected a compression caption');
|
||||
expect(caption.text).toMatch(/^<system>Image compressed/);
|
||||
expect(caption.text).toMatch(/<\/system>$/);
|
||||
expect(caption.text).toContain('saved at');
|
||||
expect(out.note).toMatch(/^<system>Image compressed/);
|
||||
expect(out.note).toMatch(/<\/system>$/);
|
||||
expect(out.note).toContain('saved at');
|
||||
const parts = out.output as ContentPart[];
|
||||
const joined = parts.map((p) => (p.type === 'text' ? p.text : '')).join('');
|
||||
expect(joined).not.toContain('Output truncated');
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
|
|
|
|||
|
|
@ -100,6 +100,33 @@ describe('AgentToolExecutorService', () => {
|
|||
});
|
||||
});
|
||||
|
||||
it('preserves internal result notes without exposing them on protocol tool.result events', async () => {
|
||||
const tool = new TestTool('captioned', {
|
||||
result: {
|
||||
output: 'image sent',
|
||||
note: '<system>Image compressed.</system>',
|
||||
},
|
||||
});
|
||||
registry.register(tool);
|
||||
|
||||
const results = await execute([toolCall('call_captioned', 'captioned', {})]);
|
||||
|
||||
expect(results[0]).toMatchObject({
|
||||
output: 'image sent',
|
||||
note: '<system>Image compressed.</system>',
|
||||
});
|
||||
const protocolResult = protocolEvents.find(
|
||||
(event): event is Extract<AgentEvent, { type: 'tool.result' }> =>
|
||||
event.type === 'tool.result',
|
||||
);
|
||||
expect(protocolResult).toMatchObject({
|
||||
type: 'tool.result',
|
||||
toolCallId: 'call_captioned',
|
||||
output: 'image sent',
|
||||
});
|
||||
expect(protocolResult as unknown as Record<string, unknown>).not.toHaveProperty('note');
|
||||
});
|
||||
|
||||
it('records an error tool.result when the tool name is unknown', async () => {
|
||||
const results = await execute([toolCall('call_missing', 'missing', { text: 'hi' })]);
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue