wehub-resource-sync
144 lines
4.4 KiB
TypeScript
144 lines
4.4 KiB
TypeScript
import { afterEach, beforeAll, beforeEach, describe, expect, test, vi } from "bun:test";
|
|
import * as fs from "node:fs/promises";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import type { SourceMeta } from "@oh-my-pi/pi-coding-agent/capability/types";
|
|
import type { MCPServerConfig } from "@oh-my-pi/pi-coding-agent/mcp/types";
|
|
import { MCPCommandController } from "@oh-my-pi/pi-coding-agent/modes/controllers/mcp-command-controller";
|
|
import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
|
|
import {
|
|
getConfigRootDir,
|
|
getMCPConfigPath,
|
|
getProjectDir,
|
|
removeWithRetries,
|
|
setAgentDir,
|
|
setProjectDir,
|
|
} from "@oh-my-pi/pi-utils";
|
|
|
|
const originalProjectDir = getProjectDir();
|
|
const originalAgentDir = process.env.PI_CODING_AGENT_DIR;
|
|
const fallbackAgentDir = path.join(getConfigRootDir(), "agent");
|
|
|
|
function restoreAgentDir(): void {
|
|
if (originalAgentDir) {
|
|
setAgentDir(originalAgentDir);
|
|
process.env.PI_CODING_AGENT_DIR = originalAgentDir;
|
|
Bun.env.PI_CODING_AGENT_DIR = originalAgentDir;
|
|
return;
|
|
}
|
|
setAgentDir(fallbackAgentDir);
|
|
delete process.env.PI_CODING_AGENT_DIR;
|
|
delete Bun.env.PI_CODING_AGENT_DIR;
|
|
}
|
|
|
|
function createController() {
|
|
const refreshMCPTools = vi.fn(async () => {});
|
|
const mcpManager = {
|
|
disconnectAll: vi.fn(async () => {}),
|
|
discoverAndConnect: vi.fn(async () => ({ errors: new Map<string, string>() })),
|
|
disconnectServer: vi.fn(async () => {}),
|
|
connectServers: vi.fn(
|
|
async (_configs: Record<string, MCPServerConfig>, _sources: Record<string, SourceMeta>) => ({
|
|
errors: new Map<string, string>(),
|
|
connectedServers: [],
|
|
tools: [],
|
|
exaApiKeys: [],
|
|
}),
|
|
),
|
|
getTools: vi.fn(() => []),
|
|
waitForConnection: vi.fn(async () => ({})),
|
|
getConnectionStatus: vi.fn(() => "connected"),
|
|
getSource: vi.fn(() => undefined),
|
|
};
|
|
const controller = new MCPCommandController({
|
|
chatContainer: { addChild: vi.fn() },
|
|
present: vi.fn(),
|
|
ui: { requestRender: vi.fn() },
|
|
editor: {},
|
|
showError: vi.fn(),
|
|
showStatus: vi.fn(),
|
|
oauthManualInput: {
|
|
hasPending: vi.fn(() => false),
|
|
pendingProviderId: undefined,
|
|
tryClaimInput: vi.fn(),
|
|
},
|
|
session: {
|
|
refreshMCPTools,
|
|
modelRegistry: { authStorage: undefined },
|
|
},
|
|
mcpManager,
|
|
} as never);
|
|
|
|
return { controller, mcpManager, refreshMCPTools };
|
|
}
|
|
|
|
async function writeProjectConfig(projectDir: string, servers: Record<string, MCPServerConfig>): Promise<void> {
|
|
await Bun.write(
|
|
getMCPConfigPath("project", projectDir),
|
|
`${JSON.stringify(
|
|
{
|
|
mcpServers: servers,
|
|
},
|
|
null,
|
|
2,
|
|
)}\n`,
|
|
);
|
|
}
|
|
|
|
describe("/mcp enable and disable", () => {
|
|
let projectDir = "";
|
|
let agentDir = "";
|
|
|
|
beforeAll(() => {
|
|
initTheme();
|
|
});
|
|
|
|
beforeEach(async () => {
|
|
projectDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-mcp-toggle-project-"));
|
|
agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-mcp-toggle-agent-"));
|
|
setProjectDir(projectDir);
|
|
setAgentDir(agentDir);
|
|
});
|
|
|
|
afterEach(async () => {
|
|
vi.restoreAllMocks();
|
|
setProjectDir(originalProjectDir);
|
|
restoreAgentDir();
|
|
await removeWithRetries(projectDir);
|
|
await removeWithRetries(agentDir);
|
|
});
|
|
|
|
test("disabling one configured server does not reload other MCP servers", async () => {
|
|
await writeProjectConfig(projectDir, {
|
|
mcp1: { type: "stdio", command: "mcp-one" },
|
|
mcp2: { type: "stdio", command: "mcp-two" },
|
|
});
|
|
const { controller, mcpManager, refreshMCPTools } = createController();
|
|
|
|
await controller.handle("/mcp disable mcp1");
|
|
|
|
expect(mcpManager.disconnectServer).toHaveBeenCalledWith("mcp1");
|
|
expect(refreshMCPTools).toHaveBeenCalledWith([]);
|
|
expect(mcpManager.disconnectAll).not.toHaveBeenCalled();
|
|
expect(mcpManager.discoverAndConnect).not.toHaveBeenCalled();
|
|
expect(mcpManager.connectServers).not.toHaveBeenCalled();
|
|
});
|
|
|
|
test("enabling one configured server connects only that MCP server", async () => {
|
|
await writeProjectConfig(projectDir, {
|
|
mcp1: { type: "stdio", command: "mcp-one", enabled: false },
|
|
mcp2: { type: "stdio", command: "mcp-two" },
|
|
});
|
|
const { controller, mcpManager } = createController();
|
|
|
|
await controller.handle("/mcp enable mcp1");
|
|
|
|
expect(mcpManager.disconnectAll).not.toHaveBeenCalled();
|
|
expect(mcpManager.discoverAndConnect).not.toHaveBeenCalled();
|
|
expect(mcpManager.connectServers).toHaveBeenCalledTimes(1);
|
|
const [configs] = mcpManager.connectServers.mock.calls[0]!;
|
|
expect(Object.keys(configs)).toEqual(["mcp1"]);
|
|
expect(configs.mcp1).toEqual({ type: "stdio", command: "mcp-one", enabled: true });
|
|
});
|
|
});
|