Skip to content

Commit 1c19ed5

Browse files
authored
feat: pass auth tokens by reference, not value (#1362)
## Problem We've split up the auth store in the previous PR into a service, but now we're not actually using it properly. A lot of services use their own (stale) local auth token value. We instead want to use the auth services token. <!-- Who is this for and what problem does it solve? --> <!-- Closes #ISSUE_ID --> ## Changes - update main services to et auth token through AuthService instead of storing it away somewhere locally - agent package now accepts a get token and refresh token function instead of a token value - clean up a bunch of auth code <!-- What did you change and why? --> <!-- If there are frontend changes, include screenshots. -->
1 parent ae1b192 commit 1c19ed5

File tree

35 files changed

+778
-638
lines changed

35 files changed

+778
-638
lines changed

apps/code/src/main/di/container.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { SuspensionRepositoryImpl } from "../db/repositories/suspension-reposito
88
import { WorkspaceRepository } from "../db/repositories/workspace-repository";
99
import { WorktreeRepository } from "../db/repositories/worktree-repository";
1010
import { DatabaseService } from "../db/service";
11+
import { AgentAuthAdapter } from "../services/agent/auth-adapter";
1112
import { AgentService } from "../services/agent/service";
1213
import { AppLifecycleService } from "../services/app-lifecycle/service";
1314
import { ArchiveService } from "../services/archive/service";
@@ -57,6 +58,7 @@ container.bind(MAIN_TOKENS.WorkspaceRepository).to(WorkspaceRepository);
5758
container.bind(MAIN_TOKENS.WorktreeRepository).to(WorktreeRepository);
5859
container.bind(MAIN_TOKENS.ArchiveRepository).to(ArchiveRepository);
5960
container.bind(MAIN_TOKENS.SuspensionRepository).to(SuspensionRepositoryImpl);
61+
container.bind(MAIN_TOKENS.AgentAuthAdapter).to(AgentAuthAdapter);
6062
container.bind(MAIN_TOKENS.AgentService).to(AgentService);
6163
container.bind(MAIN_TOKENS.AuthService).to(AuthService);
6264
container.bind(MAIN_TOKENS.AuthProxyService).to(AuthProxyService);

apps/code/src/main/di/tokens.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ export const MAIN_TOKENS = Object.freeze({
1818
SuspensionRepository: Symbol.for("Main.SuspensionRepository"),
1919

2020
// Services
21+
AgentAuthAdapter: Symbol.for("Main.AgentAuthAdapter"),
2122
AgentService: Symbol.for("Main.AgentService"),
2223
AuthService: Symbol.for("Main.AuthService"),
2324
AuthProxyService: Symbol.for("Main.AuthProxyService"),

apps/code/src/main/menu.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ function buildFileMenu(): MenuItemConstructorOptions {
127127
{
128128
label: "Invalidate OAuth token",
129129
click: () => {
130-
container.get<UIService>(MAIN_TOKENS.UIService).invalidateToken();
130+
void container
131+
.get<UIService>(MAIN_TOKENS.UIService)
132+
.invalidateToken();
131133
},
132134
},
133135
{
Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+
3+
const mockFetch = vi.hoisted(() => vi.fn());
4+
5+
vi.mock("../../utils/logger.js", () => ({
6+
logger: {
7+
scope: () => ({
8+
info: vi.fn(),
9+
error: vi.fn(),
10+
warn: vi.fn(),
11+
debug: vi.fn(),
12+
}),
13+
},
14+
}));
15+
16+
vi.mock("@posthog/agent/posthog-api", () => ({
17+
getLlmGatewayUrl: vi.fn(() => "https://gateway.example.com"),
18+
}));
19+
20+
vi.stubGlobal("fetch", mockFetch);
21+
22+
import { AgentAuthAdapter } from "./auth-adapter";
23+
24+
const baseCredentials = {
25+
apiHost: "https://app.posthog.com",
26+
projectId: 1,
27+
};
28+
29+
function createDependencies() {
30+
return {
31+
authService: {
32+
getValidAccessToken: vi.fn().mockResolvedValue({
33+
accessToken: "test-access-token",
34+
apiHost: "https://app.posthog.com",
35+
}),
36+
refreshAccessToken: vi.fn().mockResolvedValue({
37+
accessToken: "fresh-access-token",
38+
apiHost: "https://app.posthog.com",
39+
}),
40+
authenticatedFetch: vi
41+
.fn()
42+
.mockImplementation(
43+
async (
44+
fetchImpl: typeof fetch,
45+
input: string | Request,
46+
init?: RequestInit,
47+
) => fetchImpl(input, init),
48+
),
49+
},
50+
authProxy: {
51+
start: vi.fn().mockResolvedValue("http://127.0.0.1:9999"),
52+
},
53+
};
54+
}
55+
56+
describe("AgentAuthAdapter", () => {
57+
let adapter: AgentAuthAdapter;
58+
let deps: ReturnType<typeof createDependencies>;
59+
60+
beforeEach(() => {
61+
vi.clearAllMocks();
62+
mockFetch.mockResolvedValue({
63+
ok: true,
64+
json: () => Promise.resolve({ results: [] }),
65+
});
66+
67+
deps = createDependencies();
68+
adapter = new AgentAuthAdapter(
69+
deps.authService as never,
70+
deps.authProxy as never,
71+
);
72+
});
73+
74+
afterEach(() => {
75+
vi.restoreAllMocks();
76+
});
77+
78+
it("builds the default PostHog MCP server", async () => {
79+
const servers = await adapter.buildMcpServers(baseCredentials);
80+
81+
expect(servers).toEqual(
82+
expect.arrayContaining([
83+
expect.objectContaining({
84+
name: "posthog",
85+
type: "http",
86+
url: "https://mcp.posthog.com/mcp",
87+
headers: expect.arrayContaining([
88+
{
89+
name: "Authorization",
90+
value: "Bearer test-access-token",
91+
},
92+
]),
93+
}),
94+
]),
95+
);
96+
});
97+
98+
it("includes enabled user-installed MCP servers from backend", async () => {
99+
mockFetch.mockResolvedValue({
100+
ok: true,
101+
json: () =>
102+
Promise.resolve({
103+
results: [
104+
{
105+
id: "inst-1",
106+
url: "https://custom-mcp.example.com",
107+
proxy_url: "https://proxy.posthog.com/inst-1/",
108+
name: "custom-server",
109+
display_name: "Custom Server",
110+
auth_type: "none",
111+
is_enabled: true,
112+
pending_oauth: false,
113+
needs_reauth: false,
114+
},
115+
],
116+
}),
117+
});
118+
119+
const servers = await adapter.buildMcpServers(baseCredentials);
120+
121+
expect(servers).toEqual(
122+
expect.arrayContaining([
123+
expect.objectContaining({
124+
name: "custom-server",
125+
url: "https://custom-mcp.example.com",
126+
headers: [],
127+
}),
128+
]),
129+
);
130+
});
131+
132+
it("routes authenticated installed MCP servers through the proxy URL", async () => {
133+
mockFetch.mockResolvedValue({
134+
ok: true,
135+
json: () =>
136+
Promise.resolve({
137+
results: [
138+
{
139+
id: "inst-2",
140+
url: "https://remote-mcp.example.com",
141+
proxy_url: "https://proxy.posthog.com/inst-2/",
142+
name: "secure-server",
143+
display_name: "Secure Server",
144+
auth_type: "oauth",
145+
is_enabled: true,
146+
pending_oauth: false,
147+
needs_reauth: false,
148+
},
149+
],
150+
}),
151+
});
152+
153+
const servers = await adapter.buildMcpServers(baseCredentials);
154+
155+
expect(servers).toEqual(
156+
expect.arrayContaining([
157+
expect.objectContaining({
158+
name: "secure-server",
159+
url: "https://proxy.posthog.com/inst-2/",
160+
headers: [
161+
{ name: "Authorization", value: "Bearer test-access-token" },
162+
],
163+
}),
164+
]),
165+
);
166+
});
167+
168+
it("configures environment using the gateway proxy and current token", async () => {
169+
await adapter.configureProcessEnv({
170+
credentials: baseCredentials,
171+
mockNodeDir: "/mock/node",
172+
proxyUrl: "http://127.0.0.1:9999",
173+
claudeCliPath: "/mock/claude-cli.js",
174+
});
175+
176+
expect(process.env.POSTHOG_API_KEY).toBe("test-access-token");
177+
expect(process.env.POSTHOG_AUTH_HEADER).toBe("Bearer test-access-token");
178+
expect(process.env.LLM_GATEWAY_URL).toBe("http://127.0.0.1:9999");
179+
expect(process.env.CLAUDE_CODE_EXECUTABLE).toBe("/mock/claude-cli.js");
180+
expect(process.env.POSTHOG_PROJECT_ID).toBe("1");
181+
});
182+
});

0 commit comments

Comments
 (0)