fix: preserve datasource response files (#159)

* fix: preserve datasource response files

* fix: restrict datasource response file writes

* test: avoid async http handler

.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.

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<unknown> {
+  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<void> {
+  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<typeof createServer>): Promise<void> {
+  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<typeof createServer>): Promise<void> {
+  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<JsonRpcResponse> {
+      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;
+}

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) {