alexion
6081e3e325
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
56 lines
2.7 KiB
TypeScript
56 lines
2.7 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import { InMemoryAuthStorageBackend } from "../src/core/auth-storage.ts";
|
|
import { McpOAuthCredentialStore } from "../src/extensions/mcp/oauth.ts";
|
|
|
|
const SERVER_URL = "https://mcp.example.com/mcp";
|
|
|
|
function state(accessToken: string) {
|
|
return { serverUrl: SERVER_URL, tokens: { access_token: accessToken, token_type: "Bearer" } };
|
|
}
|
|
|
|
function storedKeys(backend: InMemoryAuthStorageBackend): string[] {
|
|
return Object.keys(JSON.parse(backend.withLock((current) => ({ result: current ?? "{}" }))));
|
|
}
|
|
|
|
describe("MCP OAuth credential store", () => {
|
|
// https://github.com/earendil-works/pi/issues/10252
|
|
it("keeps separate credentials for servers sharing a SERVER_URL", async () => {
|
|
const store = new McpOAuthCredentialStore(new InMemoryAuthStorageBackend());
|
|
await store.forServer("work", SERVER_URL).save(state("work-token"));
|
|
await store.forServer("personal", SERVER_URL).save(state("personal-token"));
|
|
|
|
expect((await store.forServer("work", SERVER_URL).load())?.tokens?.access_token).toBe("work-token");
|
|
expect((await store.forServer("personal", SERVER_URL).load())?.tokens?.access_token).toBe("personal-token");
|
|
|
|
expect(store.remove("work", SERVER_URL)).toBe(true);
|
|
expect(await store.forServer("work", SERVER_URL).load()).toBeUndefined();
|
|
expect(store.tokens("personal", SERVER_URL)?.access_token).toBe("personal-token");
|
|
});
|
|
|
|
it("moves credentials stored by SERVER_URL to the first server that loads them", async () => {
|
|
const backend = new InMemoryAuthStorageBackend();
|
|
backend.withLock(() => ({ result: undefined, next: JSON.stringify({ [SERVER_URL]: state("legacy-token") }) }));
|
|
const store = new McpOAuthCredentialStore(backend);
|
|
|
|
// Reading tokens does not take the legacy state over.
|
|
expect(store.tokens("work", SERVER_URL)?.access_token).toBe("legacy-token");
|
|
expect(storedKeys(backend)).toEqual([SERVER_URL]);
|
|
|
|
expect((await store.forServer("my_work", SERVER_URL).load())?.tokens?.access_token).toBe("legacy-token");
|
|
// Names differing only in `-` and `_` are the same server.
|
|
expect(store.tokens("my-work", SERVER_URL)?.access_token).toBe("legacy-token");
|
|
expect(await store.forServer("personal", SERVER_URL).load()).toBeUndefined();
|
|
expect(storedKeys(backend)).toEqual([`mcp__my_work|${SERVER_URL}`]);
|
|
});
|
|
|
|
it("signs out of credentials stored by SERVER_URL", () => {
|
|
const backend = new InMemoryAuthStorageBackend();
|
|
backend.withLock(() => ({ result: undefined, next: JSON.stringify({ [SERVER_URL]: state("legacy-token") }) }));
|
|
const store = new McpOAuthCredentialStore(backend);
|
|
|
|
expect(store.remove("work", SERVER_URL)).toBe(true);
|
|
expect(storedKeys(backend)).toEqual([]);
|
|
expect(store.remove("work", SERVER_URL)).toBe(false);
|
|
});
|
|
});
|