Skip to content

Commit 4e9d221

Browse files
allenhutchisonallenhutchison
committed
feat: Add robust refresh token preservation during credential updates and introduce a manual refresh tool.
1 parent bf5e217 commit 4e9d221

3 files changed

Lines changed: 160 additions & 6 deletions

File tree

‎workspace-mcp-server/src/__tests__/auth/AuthManager.test.ts‎

Lines changed: 107 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,12 @@ describe('AuthManager', () => {
2323

2424
// Setup mock OAuth2 client
2525
mockOAuth2Client = {
26-
setCredentials: jest.fn(),
26+
setCredentials: jest.fn().mockImplementation((creds) => {
27+
mockOAuth2Client.credentials = creds;
28+
}),
2729
generateAuthUrl: jest.fn(),
2830
on: jest.fn(),
31+
refreshAccessToken: jest.fn(),
2932
credentials: {}
3033
};
3134

@@ -69,11 +72,112 @@ describe('AuthManager', () => {
6972
await tokensCallback(newTokens);
7073

7174
// Verify saveCredentials was called with merged tokens
75+
// New tokens take precedence, but refresh_token is preserved from old credentials
7276
expect(OAuthCredentialStorage.saveCredentials).toHaveBeenCalledWith({
7377
access_token: 'new_token',
74-
refresh_token: 'old_refresh', // Should be preserved
75-
expiry_date: 123456789,
78+
refresh_token: 'old_refresh', // Preserved from old credentials
79+
expiry_date: 123456789
80+
// Note: scope is NOT preserved because newTokens didn't include it
81+
});
82+
});
83+
84+
it('should preserve refresh token during manual refresh if not returned', async () => {
85+
// Setup initial state with a refresh token
86+
(OAuthCredentialStorage.loadCredentials as jest.Mock).mockResolvedValue({
87+
access_token: 'old_token',
88+
refresh_token: 'old_refresh_token',
89+
scope: 'scope1'
90+
});
91+
92+
// Initialize client to populate this.client
93+
await authManager.getAuthenticatedClient();
94+
95+
// Mock refresh to return ONLY access token (no refresh token)
96+
// We need to update the mock to actually update credentials, similar to real OAuth2Client
97+
mockOAuth2Client.refreshAccessToken.mockImplementation(async () => {
98+
const newCreds = {
99+
access_token: 'new_access_token',
100+
expiry_date: 999999999
101+
};
102+
mockOAuth2Client.credentials = newCreds;
103+
return { credentials: newCreds };
104+
});
105+
106+
await authManager.refreshToken();
107+
108+
// Verify saveCredentials was called with BOTH new access token AND old refresh token
109+
expect(OAuthCredentialStorage.saveCredentials).toHaveBeenCalledWith(expect.objectContaining({
110+
access_token: 'new_access_token',
111+
refresh_token: 'old_refresh_token'
112+
}));
113+
});
114+
115+
it('should preserve refresh token when refreshAccessToken mutates credentials in-place', async () => {
116+
// Setup initial state with a refresh token
117+
(OAuthCredentialStorage.loadCredentials as jest.Mock).mockResolvedValue({
118+
access_token: 'old_token',
119+
refresh_token: 'old_refresh_token',
76120
scope: 'scope1'
77121
});
122+
123+
// Initialize client to populate this.client
124+
await authManager.getAuthenticatedClient();
125+
126+
// This test simulates the REAL OAuth2Client behavior where refreshAccessToken
127+
// mutates the credentials object IN-PLACE before returning
128+
mockOAuth2Client.refreshAccessToken.mockImplementation(async () => {
129+
// CRITICAL: Mutate the existing credentials object in-place
130+
// This is what the real OAuth2Client does!
131+
mockOAuth2Client.credentials.access_token = 'new_access_token';
132+
mockOAuth2Client.credentials.expiry_date = 999999999;
133+
// Note: refresh_token is NOT included in the refresh response
134+
delete mockOAuth2Client.credentials.refresh_token;
135+
delete mockOAuth2Client.credentials.scope;
136+
137+
// Return the new credentials (which are the SAME object reference)
138+
return { credentials: mockOAuth2Client.credentials };
139+
});
140+
141+
await authManager.refreshToken();
142+
143+
// This test will FAIL if the bug exists, because:
144+
// 1. Line 146 captures a reference to mockOAuth2Client.credentials
145+
// 2. Line 148 calls refreshAccessToken which mutates that same object
146+
// 3. The merge logic sees currentCredentials.refresh_token is undefined (it was deleted)
147+
// 4. The refresh_token is lost
148+
expect(OAuthCredentialStorage.saveCredentials).toHaveBeenCalledWith(expect.objectContaining({
149+
access_token: 'new_access_token',
150+
refresh_token: 'old_refresh_token'
151+
}));
152+
});
153+
154+
it('should preserve refresh token in tokens event handler', async () => {
155+
// Setup initial state with a refresh token in storage
156+
(OAuthCredentialStorage.loadCredentials as jest.Mock).mockResolvedValue({
157+
access_token: 'old_token',
158+
refresh_token: 'stored_refresh_token',
159+
scope: 'scope1'
160+
});
161+
162+
await authManager.getAuthenticatedClient();
163+
164+
// Get the registered callback
165+
const tokensCallback = mockOAuth2Client.on.mock.calls.find((call: any[]) => call[0] === 'tokens')[1];
166+
167+
// Simulate automatic refresh that doesn't include refresh_token
168+
const newTokens = {
169+
access_token: 'auto_refreshed_token',
170+
expiry_date: 999999999
171+
// Note: no refresh_token
172+
};
173+
174+
await tokensCallback(newTokens);
175+
176+
// Verify saveCredentials was called with BOTH new access token AND stored refresh token
177+
expect(OAuthCredentialStorage.saveCredentials).toHaveBeenCalledWith({
178+
access_token: 'auto_refreshed_token',
179+
expiry_date: 999999999,
180+
refresh_token: 'stored_refresh_token'
181+
});
78182
});
79183
});

‎workspace-mcp-server/src/auth/AuthManager.ts‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ export class AuthManager {
6767
// Check if we have a cached client with valid credentials
6868
if (this.client && this.client.credentials && this.client.credentials.refresh_token) {
6969
logToFile('Returning existing cached client with valid credentials');
70+
logToFile(`Access token exists: ${!!this.client.credentials.access_token}`);
71+
logToFile(`Expiry date: ${this.client.credentials.expiry_date}`);
72+
logToFile(`Current time: ${Date.now()}`);
73+
logToFile(`Token expired: ${this.client.credentials.expiry_date ? this.client.credentials.expiry_date < Date.now() : 'unknown'}`);
7074
return this.client;
7175
}
7276

@@ -83,10 +87,11 @@ export class AuthManager {
8387
}
8488

8589
try {
86-
const current = await OAuthCredentialStorage.loadCredentials();
90+
// Create a copy to preserve refresh_token from storage
91+
const current = await OAuthCredentialStorage.loadCredentials() || {};
8792
const merged = {
88-
...current,
89-
...tokens
93+
...tokens,
94+
refresh_token: tokens.refresh_token || current.refresh_token
9095
};
9196
await OAuthCredentialStorage.saveCredentials(merged);
9297
logToFile('Credentials saved after refresh');
@@ -132,6 +137,34 @@ export class AuthManager {
132137
logToFile('Authentication cleared.');
133138
}
134139

140+
public async refreshToken(): Promise<void> {
141+
logToFile('Manual token refresh triggered');
142+
if (!this.client) {
143+
logToFile('No client available to refresh, getting new client');
144+
this.client = await this.getAuthenticatedClient();
145+
}
146+
try {
147+
// Create a DEEP COPY of credentials before refresh to preserve refresh_token
148+
// (refreshAccessToken mutates this.client.credentials in-place)
149+
const currentCredentials = { ...this.client.credentials };
150+
151+
const { credentials } = await this.client.refreshAccessToken();
152+
153+
// Merge with existing credentials to preserve refresh_token if not returned
154+
const mergedCredentials = {
155+
...credentials,
156+
refresh_token: credentials.refresh_token || currentCredentials.refresh_token
157+
};
158+
159+
this.client.setCredentials(mergedCredentials);
160+
await OAuthCredentialStorage.saveCredentials(mergedCredentials);
161+
logToFile('Token refreshed and saved successfully');
162+
} catch (error) {
163+
logToFile(`Error during manual token refresh: ${error}`);
164+
throw error;
165+
}
166+
}
167+
135168
private async getAvailablePort(): Promise<number> {
136169
return new Promise((resolve, reject) => {
137170
let port = 0;

‎workspace-mcp-server/src/index.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,23 @@ async function main() {
8989
}
9090
);
9191

92+
server.registerTool(
93+
"auth.refreshToken",
94+
{
95+
description: 'Manually triggers the token refresh process.',
96+
inputSchema: {}
97+
},
98+
async () => {
99+
await authManager.refreshToken();
100+
return {
101+
content: [{
102+
type: "text",
103+
text: "Token refresh process triggered successfully."
104+
}]
105+
};
106+
}
107+
);
108+
92109
server.registerTool(
93110
"docs.create",
94111
{

0 commit comments

Comments
 (0)