diff --git a/.changeset/fix-datasource-response-files.md b/.changeset/fix-datasource-response-files.md new file mode 100644 index 000000000..edd8eeecf --- /dev/null +++ b/.changeset/fix-datasource-response-files.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix official datasource tools to preserve complete responses and write returned result files. diff --git a/apps/kimi-code/test/utils/kimi-datasource-plugin.test.ts b/apps/kimi-code/test/utils/kimi-datasource-plugin.test.ts new file mode 100644 index 000000000..815e4f8d2 --- /dev/null +++ b/apps/kimi-code/test/utils/kimi-datasource-plugin.test.ts @@ -0,0 +1,224 @@ +import { spawn, type ChildProcessWithoutNullStreams } from 'node:child_process'; +import { mkdtemp, mkdir, readFile, rm, writeFile } from 'node:fs/promises'; +import { createServer, type IncomingMessage, type ServerResponse } from 'node:http'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { createInterface } from 'node:readline'; + +import { describe, expect, it } from 'vitest'; + +const REPO_ROOT = join(import.meta.dirname, '../../../..'); +const SERVER_ENTRY = join(REPO_ROOT, 'plugins/official/kimi-datasource/bin/kimi-datasource.mjs'); + +describe('kimi-datasource MCP server', () => { + it('prefers assistant text and writes response files', async () => { + const tempDir = await mkdtemp(join(tmpdir(), 'kimi-datasource-plugin-')); + const kimiHome = join(tempDir, 'kimi-home'); + const textFile = join(tempDir, 'world-bank.csv'); + const binaryFile = join(tempDir, 'world-bank_payload.csv'); + const blockedFile = join(tempDir, 'blocked.csv'); + const requests: unknown[] = []; + let child: ChildProcessWithoutNullStreams | undefined; + + const server = createServer((request, response) => { + void handleMockDatasourceRequest(request, response, { + requests, + textFile, + binaryFile, + blockedFile, + }); + }); + + try { + await mkdir(join(kimiHome, 'credentials'), { recursive: true }); + await writeFile( + join(kimiHome, 'credentials', 'kimi-code.json'), + JSON.stringify({ access_token: 'test-token', expires_at: 4_102_444_800 }), + 'utf8', + ); + await listen(server); + + const address = server.address(); + if (address === null || typeof address === 'string') { + throw new Error('Expected an ephemeral TCP port for the test server.'); + } + + child = spawn(process.execPath, [SERVER_ENTRY], { + cwd: REPO_ROOT, + env: { + ...process.env, + KIMI_CODE_HOME: kimiHome, + KIMI_DATASOURCE_API_URL: `http://127.0.0.1:${address.port}`, + }, + stdio: ['pipe', 'pipe', 'pipe'], + }); + const client = createRpcClient(child); + + await client.request('initialize', {}); + const result = await client.request('tools/call', { + name: 'call_data_source_tool', + arguments: { + data_source_name: 'world_bank_open_data', + api_name: 'world_bank_open_data', + params: { filepath: textFile }, + }, + }); + + expect(result.error).toBeUndefined(); + expect(result.result).toEqual({ + content: [ + { + type: 'text', + text: expect.stringContaining('assistant complete result'), + }, + ], + }); + expect(JSON.stringify(result.result)).toContain('skipped returned file'); + expect(await readFile(textFile, 'utf8')).toBe('country,value\nCN,1\n'); + expect(await readFile(binaryFile, 'utf8')).toBe('binary payload'); + await expect(readFile(blockedFile, 'utf8')).rejects.toMatchObject({ code: 'ENOENT' }); + expect(requests).toEqual([ + { + method: 'call_data_source_tool', + params: { + data_source_name: 'world_bank_open_data', + api_name: 'world_bank_open_data', + params: { filepath: textFile }, + }, + }, + ]); + } finally { + child?.stdin.end(); + child?.kill(); + await closeServer(server); + await rm(tempDir, { recursive: true, force: true }); + } + }); +}); + +async function readJson(request: IncomingMessage): Promise { + let body = ''; + for await (const chunk of request) { + body += chunk; + } + return JSON.parse(body); +} + +async function handleMockDatasourceRequest( + request: IncomingMessage, + response: ServerResponse, + options: { + readonly requests: unknown[]; + readonly textFile: string; + readonly binaryFile: string; + readonly blockedFile: string; + }, +): Promise { + try { + options.requests.push(await readJson(request)); + response.setHeader('Content-Type', 'application/json'); + response.end( + JSON.stringify({ + is_success: true, + result: { + assistant: [{ type: 'text', text: 'assistant complete result' }], + user: [{ type: 'text', text: '{"data_preview": null}' }], + }, + files: [ + { name: options.textFile, content: 'country,value\nCN,1\n' }, + { + name: options.binaryFile, + content: Buffer.from('binary payload').toString('base64'), + encoding: 'base64', + }, + { name: options.blockedFile, content: 'blocked\n' }, + ], + }), + ); + } catch (error) { + response.statusCode = 500; + response.end(error instanceof Error ? error.message : String(error)); + } +} + +function listen(server: ReturnType): Promise { + return new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', () => { + server.off('error', reject); + resolve(); + }); + }); +} + +function closeServer(server: ReturnType): Promise { + return new Promise((resolve, reject) => { + if (!server.listening) { + resolve(); + return; + } + server.close((err) => { + if (err) reject(err); + else resolve(); + }); + }); +} + +function createRpcClient(child: ChildProcessWithoutNullStreams) { + let nextId = 1; + const stderr: string[] = []; + const pending = new Map< + number, + { + resolve: (value: JsonRpcResponse) => void; + reject: (err: Error) => void; + timeout: NodeJS.Timeout; + } + >(); + + child.stderr.setEncoding('utf8'); + child.stderr.on('data', (chunk) => { + stderr.push(chunk); + }); + + const lines = createInterface({ input: child.stdout }); + lines.on('line', (line) => { + const message = JSON.parse(line) as JsonRpcResponse; + const id = typeof message.id === 'number' ? message.id : undefined; + if (id === undefined) return; + const waiter = pending.get(id); + if (waiter === undefined) return; + clearTimeout(waiter.timeout); + pending.delete(id); + waiter.resolve(message); + }); + + child.on('exit', (code, signal) => { + for (const [id, waiter] of pending) { + clearTimeout(waiter.timeout); + waiter.reject(new Error(`MCP server exited before response ${id}: code=${code}, signal=${signal}.`)); + } + pending.clear(); + }); + + return { + request(method: string, params: unknown): Promise { + const id = nextId++; + const payload = `${JSON.stringify({ jsonrpc: '2.0', id, method, params })}\n`; + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + pending.delete(id); + reject(new Error(`Timed out waiting for MCP response ${id}. stderr: ${stderr.join('')}`)); + }, 5_000); + pending.set(id, { resolve, reject, timeout }); + child.stdin.write(payload); + }); + }, + }; +} + +interface JsonRpcResponse { + id?: number; + result?: unknown; + error?: unknown; +} diff --git a/plugins/official/kimi-datasource/bin/kimi-datasource.mjs b/plugins/official/kimi-datasource/bin/kimi-datasource.mjs index a7199717a..df3c293c8 100755 --- a/plugins/official/kimi-datasource/bin/kimi-datasource.mjs +++ b/plugins/official/kimi-datasource/bin/kimi-datasource.mjs @@ -190,9 +190,10 @@ async function runTool(params) { try { const built = handler.buildParams(args); const response = await callKimiTool(handler.method, built); + const fileWarnings = await writeResponseFiles(response, expectedResponseFilePath(built)); const text = extractText(response); const formatted = (handler.format?.(text, built) ?? text).trim(); - return { content: [{ type: 'text', text: formatted }] }; + return { content: [{ type: 'text', text: appendWarnings(formatted, fileWarnings) }] }; } catch (err) { const message = err instanceof Error ? err.message : String(err); return { @@ -202,6 +203,73 @@ async function runTool(params) { } } +async function writeResponseFiles(response, expectedOutputPath) { + if (!isRecord(response) || !Array.isArray(response.files)) return []; + const warnings = []; + + for (const file of response.files) { + if (!isRecord(file)) continue; + const name = typeof file.name === 'string' ? file.name.trim() : ''; + if (name.length === 0 || file.content === undefined || file.content === null) continue; + + const writePath = allowedResponseFilePath(name, expectedOutputPath); + if (writePath === undefined) { + warnings.push(`Warning: skipped returned file ${name} because it is outside the requested output path.`); + continue; + } + + try { + await mkdir(path.dirname(writePath), { recursive: true }); + if (file.encoding === 'base64') { + await writeFile(writePath, Buffer.from(String(file.content), 'base64')); + } else { + await writeFile(writePath, String(file.content), 'utf8'); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + warnings.push(`Warning: failed to write file ${writePath}: ${message}`); + } + } + + return warnings; +} + +function expectedResponseFilePath(params) { + return outputPathField(params) ?? (isRecord(params) ? outputPathField(params.params) : undefined); +} + +function outputPathField(value) { + if (!isRecord(value)) return undefined; + for (const field of ['file_path', 'filepath']) { + const pathValue = value[field]; + if (typeof pathValue !== 'string') continue; + const trimmed = pathValue.trim(); + if (trimmed.length > 0) return trimmed; + } + return undefined; +} + +function allowedResponseFilePath(name, expectedOutputPath) { + if (expectedOutputPath === undefined) return undefined; + + const actual = path.resolve(name); + const expected = path.resolve(expectedOutputPath); + if (actual === expected) return actual; + + const actualParts = path.parse(actual); + const expectedParts = path.parse(expected); + if (actualParts.dir !== expectedParts.dir) return undefined; + if (actualParts.ext !== expectedParts.ext) return undefined; + if (!actualParts.name.startsWith(`${expectedParts.name}_`)) return undefined; + + return actual; +} + +function appendWarnings(text, warnings) { + if (warnings.length === 0) return text; + return `${text}\n\n${warnings.join('\n')}`; +} + function resolveKimiHome() { const explicit = process.env.KIMI_CODE_HOME?.trim(); return explicit && explicit.length > 0 ? explicit : path.join(homedir(), '.kimi-code'); @@ -319,23 +387,29 @@ function extractText(response) { if (!isRecord(response)) return String(response); if (response.is_success === false) { - const message = extractUserText(response.error) ?? JSON.stringify(response); + const message = extractChannelText(response.error) ?? JSON.stringify(response); throw new Error(`Tool API returned an error: ${message}`); } - const text = extractUserText(response.result); + const text = extractChannelText(response.result); if (text !== undefined) return text; return `Tool API succeeded but did not return user text. Raw response: ${JSON.stringify(response)}`; } -function extractUserText(value) { - if (!isRecord(value) || !Array.isArray(value.user)) return undefined; - const text = value.user - .filter((item) => isRecord(item) && item.type === 'text' && typeof item.text === 'string') - .map((item) => item.text) - .filter(Boolean) - .join('\n\n'); - return text.length > 0 ? text : undefined; +function extractChannelText(value) { + if (!isRecord(value)) return undefined; + for (const channel of ['assistant', 'user']) { + const items = value[channel]; + if (!Array.isArray(items)) continue; + const text = items + .filter((item) => isRecord(item) && item.type === 'text' && typeof item.text === 'string') + .map((item) => item.text) + .filter(Boolean) + .join('\n\n') + .trim(); + if (text.length > 0) return text; + } + return undefined; } function defaultStockFilePath(ticker, queryType) {