mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-23 23:56:41 +00:00
fix(mcp): canonicalize trust and refresh scheduling
This commit is contained in:
parent
ce7695cec7
commit
4024a5ef7e
5 changed files with 56 additions and 7 deletions
|
|
@ -3,6 +3,7 @@ import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
|
|||
|
||||
import { ErrorCodes, Error2 } from '#/errors';
|
||||
import { IBootstrapService } from '#/app/bootstrap/bootstrap';
|
||||
import { findGitWorkTree } from '#/app/git/workTree';
|
||||
import { loadMcpServersDetailed } from '#/app/mcpConfig/configLoader';
|
||||
import { IMcpConfigStore } from '#/app/mcpConfig/configStore';
|
||||
import { IPluginService } from '#/app/plugin/plugin';
|
||||
|
|
@ -43,7 +44,8 @@ export class McpRegistryService implements IMcpRegistryService {
|
|||
});
|
||||
}
|
||||
} else {
|
||||
if (!(await readWorkspaceTrust(this.docs, query.cwd))) {
|
||||
const workspaceRoot = (await findGitWorkTree(this.fs, query.cwd))?.root ?? query.cwd;
|
||||
if (!(await readWorkspaceTrust(this.docs, workspaceRoot))) {
|
||||
const userEntries = await this.store.list();
|
||||
for (const server of userEntries) {
|
||||
const { name, ...config } = server;
|
||||
|
|
|
|||
|
|
@ -474,7 +474,12 @@ export class McpOAuthService {
|
|||
now: () => this.scheduler.now(),
|
||||
onTokensSaved: (tokens) => {
|
||||
this.emit({ type: 'tokens-saved', serverName, serverUrl: canonicalUrl });
|
||||
if (typeof tokens.obtained_at === 'number' && typeof tokens.expires_in === 'number') {
|
||||
if (
|
||||
typeof tokens.obtained_at === 'number' &&
|
||||
typeof tokens.expires_in === 'number' &&
|
||||
typeof tokens.refresh_token === 'string' &&
|
||||
tokens.refresh_token.length > 0
|
||||
) {
|
||||
this.scheduleRefresh(
|
||||
serverName,
|
||||
canonicalUrl,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import { encodeWorkDirKey } from '#/_base/utils/workdir-slug';
|
||||
import { normalize } from 'pathe';
|
||||
|
||||
import { encodeWorkDirKey, workspaceRootKey } from '#/_base/utils/workdir-slug';
|
||||
import type { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore';
|
||||
|
||||
const TRUST_SCOPE = 'workspace-trust';
|
||||
|
|
@ -13,7 +15,7 @@ export async function readWorkspaceTrust(
|
|||
root: string,
|
||||
): Promise<boolean> {
|
||||
try {
|
||||
return (await docs.get<TrustRecord>(TRUST_SCOPE, encodeWorkDirKey(root))) !== undefined;
|
||||
return (await docs.get<TrustRecord>(TRUST_SCOPE, trustKey(root))) !== undefined;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
|
|
@ -24,12 +26,16 @@ export function writeWorkspaceTrust(
|
|||
root: string,
|
||||
trustedAt: number,
|
||||
): Promise<void> {
|
||||
return docs.set(TRUST_SCOPE, encodeWorkDirKey(root), { root, trustedAt });
|
||||
return docs.set(TRUST_SCOPE, trustKey(root), { root, trustedAt });
|
||||
}
|
||||
|
||||
export function deleteWorkspaceTrust(
|
||||
docs: IAtomicDocumentStore,
|
||||
root: string,
|
||||
): Promise<void> {
|
||||
return docs.delete(TRUST_SCOPE, encodeWorkDirKey(root));
|
||||
return docs.delete(TRUST_SCOPE, trustKey(root));
|
||||
}
|
||||
|
||||
function trustKey(root: string): string {
|
||||
return encodeWorkDirKey(workspaceRootKey(normalize(root)));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { join } from 'pathe';
|
|||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { DisposableStore } from '#/_base/di/lifecycle';
|
||||
import { encodeWorkDirKey } from '#/_base/utils/workdir-slug';
|
||||
import { createServices } from '#/_base/di/test';
|
||||
import { IBootstrapService } from '#/app/bootstrap/bootstrap';
|
||||
import {
|
||||
|
|
@ -49,6 +50,7 @@ describe('McpRegistryService', () => {
|
|||
let pluginEntries: PluginMcpServerEntry[];
|
||||
let pluginError: Error | undefined;
|
||||
let trusted: boolean;
|
||||
let trustedKey: string | undefined;
|
||||
let registry: IMcpRegistryService;
|
||||
|
||||
beforeEach(() => {
|
||||
|
|
@ -59,6 +61,7 @@ describe('McpRegistryService', () => {
|
|||
pluginEntries = [];
|
||||
pluginError = undefined;
|
||||
trusted = true;
|
||||
trustedKey = undefined;
|
||||
const ix = createServices(disposables, {
|
||||
additionalServices: (reg) => {
|
||||
reg.defineInstance(IFileSystemStorageService, new InMemoryStorageService());
|
||||
|
|
@ -72,7 +75,12 @@ describe('McpRegistryService', () => {
|
|||
});
|
||||
reg.defineInstance(IHostFileSystem, new HostFileSystem());
|
||||
reg.definePartialInstance(IAtomicDocumentStore, {
|
||||
get: async <T>() => (trusted ? ({} as T) : undefined),
|
||||
get: async <T>(_scope: string, key: string) => {
|
||||
if (!trusted || (trustedKey !== undefined && key !== encodeWorkDirKey(trustedKey))) {
|
||||
return undefined;
|
||||
}
|
||||
return {} as T;
|
||||
},
|
||||
});
|
||||
reg.define(IMcpRegistryService, McpRegistryService);
|
||||
},
|
||||
|
|
@ -198,6 +206,18 @@ describe('McpRegistryService', () => {
|
|||
]);
|
||||
});
|
||||
|
||||
it('uses the canonical git root when checking trust from a subdirectory', async () => {
|
||||
const { project, sub } = await makeProject();
|
||||
await writeJson(join(project, '.mcp.json'), {
|
||||
mcpServers: { projectOnly: { command: 'project-only' } },
|
||||
});
|
||||
trustedKey = project;
|
||||
|
||||
const entries = await registry.list({ cwd: sub });
|
||||
|
||||
expect(entries.map((entry) => entry.name)).toContain('projectOnly');
|
||||
});
|
||||
|
||||
it('exposes plugin servers as read-only entries with their effective config', async () => {
|
||||
pluginEntries = [
|
||||
pluginEntry('demo', 'finance', {
|
||||
|
|
|
|||
|
|
@ -762,6 +762,22 @@ describe('McpOAuthService proactive refresh scheduling', () => {
|
|||
await fixture.scheduler.advanceBy(10_000);
|
||||
expect(refreshSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not schedule an expiring grant without a refresh token', async () => {
|
||||
const fixture = makeFixture();
|
||||
cleanups.push(() => fixture.service.dispose());
|
||||
const refreshSpy = vi.spyOn(fixture.service, 'refresh');
|
||||
|
||||
await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({
|
||||
access_token: 'a',
|
||||
token_type: 'Bearer',
|
||||
expires_in: 60,
|
||||
});
|
||||
|
||||
await fixture.scheduler.advanceBy(60_000);
|
||||
expect(refreshSpy).not.toHaveBeenCalled();
|
||||
expect(fixture.events.some((event) => event.type === 'refresh-failed')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('McpOAuthService shutdown', () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue