fix: align datasource plugin environment (#595)

* fix: align datasource plugin environment

* refactor: inject managed Kimi env into all stdio plugins

Pin the datasource credential-name test to the canonical
resolveKimiCodeOAuthKey so a digest drift in the standalone plugin
fails CI, and drop the hardcoded plugin-name special case so every
stdio plugin receives the active managed Kimi base URL / OAuth host
consistently (process.env and KIMI_CODE_HOME are already shared with
all plugins).

.changeset/datasource-env-oauth.md

@@ -0,0 +1,6 @@
+---
+"@moonshot-ai/kimi-code": patch
+"@moonshot-ai/agent-core": patch
+---
+
+Fix Kimi Datasource to use the matching OAuth credentials and service endpoint for the active Kimi Code environment.

apps/kimi-code/test/utils/kimi-datasource-plugin.test.ts

@@ -5,6 +5,7 @@ import { tmpdir } from 'node:os';
 import { join } from 'node:path';
 import { createInterface } from 'node:readline';
 
+import { resolveKimiCodeOAuthKey } from '@moonshot-ai/kimi-code-oauth';
 import { describe, expect, it } from 'vitest';
 
 const REPO_ROOT = join(import.meta.dirname, '../../../..');
@@ -114,12 +115,98 @@ describe('kimi-datasource MCP server', () => {
       await expect(readFile(blockedFile, 'utf8')).rejects.toMatchObject({ code: 'ENOENT' });
       expect(requests).toEqual([
         {
+          authorization: 'Bearer test-token',
           method: 'call_data_source_tool',
           params: {
             data_source_name: 'world_bank_open_data',
             api_name: 'world_bank_open_data',
             params: { filepath: textFile },
           },
+          url: '/',
+        },
+      ]);
+    } finally {
+      child?.stdin.end();
+      child?.kill();
+      await closeServer(server);
+      await rm(tempDir, { recursive: true, force: true });
+    }
+  });
+
+  it('uses env-scoped credentials and derives the datasource URL from KIMI_CODE_BASE_URL', async () => {
+    const tempDir = await mkdtemp(join(tmpdir(), 'kimi-datasource-plugin-'));
+    const kimiHome = join(tempDir, 'kimi-home');
+    const requests: unknown[] = [];
+    let child: ChildProcessWithoutNullStreams | undefined;
+
+    const server = createServer((request, response) => {
+      void handleMockDatasourceRequest(request, response, {
+        requests,
+        textFile: join(tempDir, 'unused.csv'),
+        binaryFile: join(tempDir, 'unused_payload.csv'),
+        blockedFile: join(tempDir, 'blocked.csv'),
+      });
+    });
+
+    try {
+      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.');
+      }
+
+      const baseUrl = `http://127.0.0.1:${address.port}/coding/v1`;
+      const oauthHost = 'https://auth.dev.kimi.team';
+      const scopedCredential = kimiCodeEnvCredentialName({ oauthHost, baseUrl });
+
+      await mkdir(join(kimiHome, 'credentials'), { recursive: true });
+      await writeFile(
+        join(kimiHome, 'credentials', 'kimi-code.json'),
+        JSON.stringify({ access_token: 'expired-prod-token', expires_at: 1 }),
+        'utf8',
+      );
+      await writeFile(
+        join(kimiHome, 'credentials', `${scopedCredential}.json`),
+        JSON.stringify({ access_token: 'scoped-token', expires_at: 4_102_444_800 }),
+        'utf8',
+      );
+
+      child = spawn(process.execPath, [SERVER_ENTRY], {
+        cwd: REPO_ROOT,
+        env: {
+          ...process.env,
+          KIMI_CODE_HOME: kimiHome,
+          KIMI_CODE_BASE_URL: baseUrl,
+          KIMI_CODE_OAUTH_HOST: oauthHost,
+          KIMI_DATASOURCE_API_URL: undefined,
+        },
+        stdio: ['pipe', 'pipe', 'pipe'],
+      });
+      const client = createRpcClient(child);
+
+      await client.request('initialize', {});
+      const result = await client.request('tools/call', {
+        name: 'get_data_source_desc',
+        arguments: {
+          name: 'arxiv',
+        },
+      });
+
+      expect(result.error).toBeUndefined();
+      expect(result.result).toEqual({
+        content: [
+          {
+            type: 'text',
+            text: expect.stringContaining('assistant complete result'),
+          },
+        ],
+      });
+      expect(requests).toEqual([
+        {
+          authorization: 'Bearer scoped-token',
+          method: 'get_data_source_desc',
+          params: { name: 'arxiv' },
+          url: '/coding/v1/tools',
         },
       ]);
     } finally {
@@ -131,6 +218,17 @@ describe('kimi-datasource MCP server', () => {
   });
 });
 
+// Pin the expected credential file name to the canonical OAuth-key resolver so
+// this test fails if the plugin's standalone digest drifts from the source of
+// truth in @moonshot-ai/kimi-code-oauth. The credential file name is the OAuth
+// key with its `oauth/` prefix stripped.
+function kimiCodeEnvCredentialName(options: {
+  readonly oauthHost: string;
+  readonly baseUrl: string;
+}): string {
+  return resolveKimiCodeOAuthKey(options).replace(/^oauth\//, '');
+}
+
 async function readJson(request: IncomingMessage): Promise<unknown> {
   let body = '';
   for await (const chunk of request) {
@@ -150,7 +248,11 @@ async function handleMockDatasourceRequest(
   },
 ): Promise<void> {
   try {
-    options.requests.push(await readJson(request));
+    options.requests.push({
+      ...(await readJson(request) as Record<string, unknown>),
+      authorization: request.headers.authorization,
+      url: request.url,
+    });
     response.setHeader('Content-Type', 'application/json');
     response.end(
       JSON.stringify({

docs/en/customization/plugins.md

@@ -59,7 +59,7 @@ Kimi Datasource is the official Kimi Code data plugin. It lets you query financi
 
 ### Installation
 
-You must first complete OAuth login with a Kimi Code account via `/login`. The plugin relies on local credentials to access data services.
+You must first complete OAuth login with a Kimi Code account via `/login`. The plugin relies on local credentials to access data services. When you run Kimi Code with `KIMI_CODE_OAUTH_HOST` or `KIMI_CODE_BASE_URL`, log in under the same environment so the datasource plugin uses the matching credentials and service endpoint.
 
 1. Run `/plugins` and select **Marketplace**
 2. Find **Kimi Datasource** and press `Space` to install

docs/zh/customization/plugins.md

@@ -59,7 +59,7 @@ Kimi Datasource 是 Kimi Code 官方数据插件，让你通过自然语言直
 
 ### 安装
 
-需先通过 `/login` 完成 Kimi Code 账号 OAuth 登录，插件依赖本地凭据访问数据服务。
+需先通过 `/login` 完成 Kimi Code 账号 OAuth 登录，插件依赖本地凭据访问数据服务。当你通过 `KIMI_CODE_OAUTH_HOST` 或 `KIMI_CODE_BASE_URL` 运行 Kimi Code 时，需要在同一环境下登录，这样 datasource 插件会使用匹配的凭据和服务端点。
 
 1. 运行 `/plugins`，选择 **Marketplace**
 2. 找到 **Kimi Datasource**，按 `Space` 安装

packages/agent-core/src/rpc/core-impl.ts

@@ -20,6 +20,7 @@ import {
   resolveKimiHome,
   writeConfigFile,
   type KimiConfig,
+  type McpServerConfig,
   type MoonshotServiceConfig,
 } from '../config';
 import {
@@ -98,6 +99,9 @@ import { KaosShellNotFoundError, LocalKaos, type Kaos } from '@moonshot-ai/kaos'
 import type { ToolServices } from '../tools/support/services';
 
 const KIMI_CODE_PROVIDER_NAME = 'managed:kimi-code';
+const KIMI_CODE_BASE_URL_ENV = 'KIMI_CODE_BASE_URL';
+const KIMI_CODE_OAUTH_HOST_ENV = 'KIMI_CODE_OAUTH_HOST';
+const KIMI_OAUTH_HOST_ENV = 'KIMI_OAUTH_HOST';
 type AgentScopedPayload<T> = T & { readonly agentId: string };
 type SessionScopedPayload<T> = T & { readonly sessionId: string };
 type SessionAgentPayload<T> = SessionScopedPayload<AgentScopedPayload<T>>;
@@ -797,7 +801,7 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
   }
 
   private mergePluginMcpConfig(base: SessionMcpConfig | undefined): SessionMcpConfig | undefined {
-    const pluginServers = this.plugins.enabledMcpServers();
+    const pluginServers = this.withManagedKimiPluginEnv(this.plugins.enabledMcpServers());
     if (Object.keys(pluginServers).length === 0) return base;
     return {
       servers: {
@@ -807,6 +811,36 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
     };
   }
 
+  private withManagedKimiPluginEnv(
+    pluginServers: Record<string, McpServerConfig>,
+  ): Record<string, McpServerConfig> {
+    const managedEnv = this.managedKimiCodeEnvForPlugins();
+    if (Object.keys(managedEnv).length === 0) return pluginServers;
+
+    const out: Record<string, McpServerConfig> = {};
+    for (const [name, server] of Object.entries(pluginServers)) {
+      out[name] =
+        server.transport === 'stdio'
+          ? { ...server, env: { ...server.env, ...managedEnv } }
+          : server;
+    }
+    return out;
+  }
+
+  private managedKimiCodeEnvForPlugins(): Record<string, string> {
+    const provider = this.config.providers[KIMI_CODE_PROVIDER_NAME];
+    const envBaseUrl = process.env[KIMI_CODE_BASE_URL_ENV];
+    const envOAuthHost = process.env[KIMI_CODE_OAUTH_HOST_ENV] ?? process.env[KIMI_OAUTH_HOST_ENV];
+    const hasEnvOverride = envBaseUrl !== undefined || envOAuthHost !== undefined;
+    const baseUrl =
+      envBaseUrl !== undefined ? envBaseUrl.replace(/\/+$/, '') : provider?.baseUrl;
+    const oauthHost = hasEnvOverride ? envOAuthHost : provider?.oauth?.oauthHost;
+    const env: Record<string, string> = {};
+    if (baseUrl !== undefined) env[KIMI_CODE_BASE_URL_ENV] = baseUrl;
+    if (oauthHost !== undefined) env[KIMI_CODE_OAUTH_HOST_ENV] = oauthHost;
+    return env;
+  }
+
   private sessionApi(sessionId: string): SessionAPIImpl {
     const session = this.sessions.get(sessionId);
     if (session === undefined) {

packages/agent-core/test/rpc/plugins-rpc.test.ts

@@ -83,6 +83,64 @@ describe('KimiCore plugin RPCs', () => {
     );
   });
 
+  it('injects persisted managed Kimi Code environment into the datasource plugin MCP server', async () => {
+    const previousBaseUrl = process.env['KIMI_CODE_BASE_URL'];
+    const previousCodeOAuthHost = process.env['KIMI_CODE_OAUTH_HOST'];
+    const previousOAuthHost = process.env['KIMI_OAUTH_HOST'];
+    delete process.env['KIMI_CODE_BASE_URL'];
+    delete process.env['KIMI_CODE_OAUTH_HOST'];
+    delete process.env['KIMI_OAUTH_HOST'];
+
+    const home = await mkdtemp(path.join(tmpdir(), 'kimi-home-'));
+    const pluginRoot = await mkdtemp(path.join(tmpdir(), 'plugin-'));
+    try {
+      await writeFile(
+        path.join(home, 'config.toml'),
+        `
+[providers."managed:kimi-code"]
+type = "kimi"
+base_url = "https://coding.deva.msh.team/coding/v1"
+api_key = ""
+oauth = { storage = "file", key = "oauth/kimi-code-env-1234", oauth_host = "https://auth.dev.kimi.team" }
+`,
+        'utf8',
+      );
+      await writeFile(
+        path.join(pluginRoot, 'kimi.plugin.json'),
+        JSON.stringify({
+          name: 'kimi-datasource',
+          mcpServers: {
+            data: { command: 'node', args: ['./bin/kimi-datasource.mjs'] },
+          },
+        }),
+        'utf8',
+      );
+
+      const core = new KimiCore(async () => ({}) as never, { homeDir: home });
+      await new Promise((r) => setImmediate(r));
+      await core.installPlugin({ source: pluginRoot });
+
+      const mcpConfig = (
+        core as unknown as {
+          mergePluginMcpConfig(base: undefined): {
+            servers: Record<string, { env?: Record<string, string> }>;
+          };
+        }
+      ).mergePluginMcpConfig(undefined);
+
+      expect(mcpConfig.servers['plugin-kimi-datasource:data']?.env).toEqual(
+        expect.objectContaining({
+          KIMI_CODE_BASE_URL: 'https://coding.deva.msh.team/coding/v1',
+          KIMI_CODE_OAUTH_HOST: 'https://auth.dev.kimi.team',
+        }),
+      );
+    } finally {
+      restoreEnv('KIMI_CODE_BASE_URL', previousBaseUrl);
+      restoreEnv('KIMI_CODE_OAUTH_HOST', previousCodeOAuthHost);
+      restoreEnv('KIMI_OAUTH_HOST', previousOAuthHost);
+    }
+  });
+
   it('throws PLUGIN_LOAD_FAILED on every RPC when installed.json is corrupt', async () => {
     const home = await mkdtemp(path.join(tmpdir(), 'kimi-home-'));
     await mkdir(path.join(home, 'plugins'), { recursive: true });
@@ -149,3 +207,11 @@ describe('KimiCore plugin RPCs', () => {
     );
   });
 });
+
+function restoreEnv(name: string, value: string | undefined): void {
+  if (value === undefined) {
+    delete process.env[name];
+    return;
+  }
+  process.env[name] = value;
+}

plugins/marketplace.json

@@ -5,7 +5,7 @@
       "id": "kimi-datasource",
       "tier": "official",
       "displayName": "Kimi Datasource",
-      "version": "3.1.1",
+      "version": "3.1.2",
       "description": "Official datasource workflows.",
       "keywords": ["data", "mcp"],
       "source": "./official/kimi-datasource"

plugins/official/kimi-datasource/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 3.1.2 - 2026-06-09
+
+- Use OAuth credentials and datasource endpoints that match the active Kimi Code environment.
+
 ## 3.1.1 - 2026-06-02
 
 - Refine skill activation wording and answer-language guidance.

plugins/official/kimi-datasource/SKILL.md

@@ -16,7 +16,7 @@ description: |
 
 这两个工具由 Kimi Code 托管执行，参数直接按 tool schema 传 JSON。
 
-工具会读取 `$KIMI_CODE_HOME/credentials/kimi-code.json`。如果没有登录凭据，让用户先在 Kimi Code 里执行 `/login`。
+工具会读取当前 Kimi Code 环境对应的本地 OAuth 登录凭据；当设置了 `KIMI_CODE_OAUTH_HOST` / `KIMI_CODE_BASE_URL` 时，会使用对应环境的隔离凭据。如果没有登录凭据，让用户先在 Kimi Code 里执行 `/login`。
 
 ## 1. 这个 skill 提供什么能力
 

plugins/official/kimi-datasource/bin/kimi-datasource.mjs

@@ -9,17 +9,19 @@
 //   - tools/call
 //   - ping
 //
-// Business logic (API call, credentials, headers) is unchanged from the
-// previous one-shot CLI; only the transport changed.
+// Business logic is kept self-contained so the plugin can run from a zipped
+// marketplace install without workspace package dependencies.
 
-import { randomUUID } from 'node:crypto';
+import { createHash, randomUUID } from 'node:crypto';
 import { mkdir, readFile, writeFile } from 'node:fs/promises';
 import { arch, homedir, hostname, release, type } from 'node:os';
 import path from 'node:path';
 import readline from 'node:readline';
 
-const VERSION = '3.1.1';
-const API_URL = process.env.KIMI_DATASOURCE_API_URL ?? 'https://api.kimi.com/coding/v1/tools';
+const VERSION = '3.1.2';
+const DEFAULT_KIMI_CODE_OAUTH_HOST = 'https://auth.kimi.com';
+const DEFAULT_KIMI_CODE_BASE_URL = 'https://api.kimi.com/coding/v1';
+const API_URL = datasourceApiUrl();
 const REQUEST_TIMEOUT_MS = 30_000;
 const PROTOCOL_VERSION = '2025-06-18';
 
@@ -209,9 +211,53 @@ function resolveKimiHome() {
   return explicit && explicit.length > 0 ? explicit : path.join(homedir(), '.kimi-code');
 }
 
+function datasourceApiUrl() {
+  const explicit = process.env.KIMI_DATASOURCE_API_URL?.trim();
+  if (explicit !== undefined && explicit.length > 0) return explicit;
+  return `${kimiCodeBaseUrl()}/tools`;
+}
+
+function kimiCodeBaseUrl() {
+  return (process.env.KIMI_CODE_BASE_URL ?? DEFAULT_KIMI_CODE_BASE_URL).replace(/\/+$/, '');
+}
+
+function kimiCodeOAuthHost() {
+  return normalizeEndpoint(
+    process.env.KIMI_CODE_OAUTH_HOST ??
+      process.env.KIMI_OAUTH_HOST ??
+      DEFAULT_KIMI_CODE_OAUTH_HOST,
+  );
+}
+
+function normalizeEndpoint(value) {
+  return value.trim().replace(/\/+$/, '');
+}
+
+function resolveKimiCodeCredentialName() {
+  const oauthHost = kimiCodeOAuthHost();
+  const baseUrl = kimiCodeBaseUrl();
+  if (
+    oauthHost === normalizeEndpoint(DEFAULT_KIMI_CODE_OAUTH_HOST) &&
+    baseUrl === DEFAULT_KIMI_CODE_BASE_URL
+  ) {
+    return 'kimi-code';
+  }
+
+  // Keep this in sync with packages/oauth/src/managed-kimi-code.ts.
+  const digest = createHash('sha256')
+    .update(JSON.stringify({ oauthHost, baseUrl }))
+    .digest('hex')
+    .slice(0, 16);
+  return `kimi-code-env-${digest}`;
+}
+
 async function loadAccessToken() {
   const kimiHome = resolveKimiHome();
-  const credentialsFile = path.join(kimiHome, 'credentials', 'kimi-code.json');
+  const credentialsFile = path.join(
+    kimiHome,
+    'credentials',
+    `${resolveKimiCodeCredentialName()}.json`,
+  );
   let parsed;
   try {
     parsed = JSON.parse(await readFile(credentialsFile, 'utf8'));

plugins/official/kimi-datasource/kimi.plugin.json

@@ -1,6 +1,6 @@
 {
   "name": "kimi-datasource",
-  "version": "3.1.1",
+  "version": "3.1.2",
   "description": "Finance, macro, enterprise, and academic data tools for Kimi Code.",
   "keywords": ["finance", "data-source", "mcp"],
   "mcpServers": {