Skip to content

Commit f4f38c1

Browse files
committed
fix(cloud): #92 review, round four — a replay that asks twice, windows that catch up
1. The replay decided before its read and posted after it. replaySessionExpired asked "is an expiry waiting?", awaited SecretStorage, then posted on that answer. "Use my own key instead", or a switch to BYOK, landing while the read was out left the answer stale, and the card came back after the user had dismissed it. It asks again after the read, immediately before posting. It also no longer rejects: it runs at the end of a resume or a checkpoint restore, and a store that will not read is no reason to cut those short. 2. A window did not notice another window ending the session they share. The focus check returned at "no token", so this window's flag, popover and footer stayed on "signed in", with no card, until the next message failed — a chat that is already open sends no second `ready`. The check now starts by comparing what is stored with what this window is showing, and catches up when they differ (catchUpWithStoredSession): ended elsewhere -> the card and a signed-out popover; signed out elsewhere -> the popover; signed IN elsewhere -> the account message that takes a card still on screen back down. Two things had to move for that to work. The ten-minute ration lived in the focus hook, so a window that had checked a minute ago did not look at all; it is inside the check now, on the half that can cost a request, and every focus does the local half. And the token is read between session changes (storedToken), so this window's own half-finished sign-out — flag already cleared, token not yet deleted — is not mistaken for another window's doing. On `ready` nothing changes: the handler has just read the same token, and the card is still replayed last. checkCloudSession no longer rejects either. 3. Not changed here: what inline edit and Agent Sketch say for an ended session. #95, stacked on this branch, already does that — the shared sentence in both, and a Sign in action on the edit. Verified: 43 suites, 768 cases, 0 failing. sessionExpiredHost.test.js goes from 57 to 67 cases; its SecretStorage stand-in can now hold a read open or refuse one. Each defect put back fails its own case: no second look -> "no card for an expiry that is no longer waiting", for both ways of answering it; no catch-up on the check -> the three other-window cases; catching up only on "gone" -> the sign-in direction; the ration back in the hook, or ahead of the catch-up -> the wiring guard and "found inside the ten minutes"; the check or the replay allowed to reject -> "the keychain is locked" escapes; the token read mid-change -> a second account resync during this window's own sign-out. tsc --checkJs reports nothing new.
1 parent 29ee664 commit f4f38c1

3 files changed

Lines changed: 199 additions & 30 deletions

File tree

‎extensions/levelcode-ai/extension.js‎

Lines changed: 69 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -425,20 +425,37 @@ async function renewSession(endpoint) {
425425
// Not applied: the stored session is no longer the one this request was about — a sign-in, a
426426
// sign-out, another refresh that finished first, or another window. What the caller needs is
427427
// whether a token is in place to retry with.
428-
const token = await ctx.secrets.get(ACCOUNT_TOKEN_KEY);
428+
const token = await storedToken();
429429
dbg('cloud.refresh', { superseded: true, outcome, token: !!token });
430-
if (!token && cloudSignedIn) {
431-
// Gone, and not by this window's hand — a sign-out or an expiry HERE clears the flag before it
432-
// deletes anything. Another window got there first, so this one catches up: the card if the
433-
// session ended, and a popover and footer that stop claiming it is live either way.
434-
cloudSignedIn = false;
435-
if (sessionExpiredPending()) { postSessionExpired(); }
436-
await postAccount(false);
437-
sendConfigToWebview();
438-
}
430+
if (!!token !== cloudSignedIn) { await catchUpWithStoredSession(token); }
439431
return !!token;
440432
}
441433

434+
/**
435+
* The stored access token, read between session changes rather than in the middle of one. Every
436+
* change this window makes moves the token and `cloudSignedIn` together inside the lock, so a token
437+
* read this way that disagrees with the flag was changed by ANOTHER window.
438+
*/
439+
function storedToken() {
440+
return withSessionLock(() => ctx.secrets.get(ACCOUNT_TOKEN_KEY));
441+
}
442+
443+
/**
444+
* The stored session belongs to every window, and a window only hears about the changes it makes
445+
* itself. When what is stored stops matching what this window is showing, another window has signed
446+
* out, signed in, or found the expiry — and a chat that is already open sends no second `ready` to
447+
* notice. Bring this window level: the flag, the popover and footer, and the card. An expiry that
448+
* is now waiting gets its card; if none is (the other window signed in), the account message takes
449+
* a card that is still up back down.
450+
* @param {string|undefined} token what SecretStorage holds now
451+
*/
452+
async function catchUpWithStoredSession(token) {
453+
cloudSignedIn = !!token;
454+
if (!token && sessionExpiredPending()) { postSessionExpired(); }
455+
await postAccount(false);
456+
sendConfigToWebview();
457+
}
458+
442459
/**
443460
* Every change to the stored session — a sign-in, a sign-out, a refresh's new tokens, an expiry —
444461
* runs through here, one at a time.
@@ -579,11 +596,21 @@ async function clearSessionExpired() {
579596
* resumed session (both post `reset`, which empties the log) and a checkpoint restore (which
580597
* drops every node after the restored turn). Each takes the card with it, and the next message
581598
* would be stopped by an expiry nothing on screen mentions any more.
599+
*
600+
* Never rejects. It runs at the end of things that have already done their work — a resume, a
601+
* restore — and a store that will not read is no reason to cut those short.
582602
*/
583603
async function replaySessionExpired() {
584-
if (!sessionExpiredPending() || await ctx.secrets.get(ACCOUNT_TOKEN_KEY)) { return; }
585-
dbg('cloud.sessionExpired.replay', {});
586-
postSessionExpired();
604+
if (!sessionExpiredPending()) { return; }
605+
try {
606+
const token = await ctx.secrets.get(ACCOUNT_TOKEN_KEY);
607+
// Asked AGAIN, because the answer can change while that read is out: "Use my own key instead"
608+
// clears the marker, a Settings change takes the session out of play. Posting on the earlier
609+
// answer would put back a card the user has just dismissed.
610+
if (token || !sessionExpiredPending()) { return; }
611+
dbg('cloud.sessionExpired.replay', {});
612+
postSessionExpired();
613+
} catch (e) { dbg('cloud.sessionExpired.replay', { error: String((e && e.message) || e) }); }
587614
}
588615

589616
/**
@@ -592,19 +619,32 @@ async function replaySessionExpired() {
592619
*
593620
* Cheap by design: the access token's own `exp` is read locally and the network is only touched
594621
* when it is expired or about to be. A session with hours left costs nothing here. Called on the
595-
* webview's `ready` and again when the window regains focus after a while away.
622+
* webview's `ready` and again every time the window regains focus.
623+
*
624+
* Two halves. Catching up with the other windows is one local read, so it happens on every call.
625+
* The half that can cost a request is rationed on focus. Never rejects: `ready` awaits this before
626+
* it restores the chat.
596627
*/
628+
const SESSION_CHECK_EVERY_MS = 10 * 60 * 1000;
597629
let lastSessionCheck = 0;
598630
async function checkCloudSession(reason) {
599631
if (!ctx || providerMode() !== 'gateway') { return; }
600-
const token = await ctx.secrets.get(ACCOUNT_TOKEN_KEY);
601-
// No token, so no session to check. One that already ENDED — found while no chat was open to hear
602-
// about it — is not announced from here: `ready` replays it last, via replaySessionExpired().
603-
if (!token) { return; }
604-
lastSessionCheck = Date.now();
605-
if (!session.accessNeedsRefresh(token)) { return; }
606-
dbg('cloud.sessionCheck', { reason, expiresAt: session.jwtExpiresAt(token) });
607-
await refreshCloudToken(); // an expired refresh token lands in sessionExpired() from inside
632+
try {
633+
const token = await storedToken();
634+
// First, whether this window is still showing the session that is actually stored. On `ready`
635+
// there is nothing to catch up on — the handler has just read the same token — and an expiry
636+
// found while no chat was open is replayed LAST there, by replaySessionExpired(). On focus this
637+
// is what tells an open chat that another window has ended, left or renewed the session.
638+
if (!!token !== cloudSignedIn) { await catchUpWithStoredSession(token); }
639+
if (!token) { return; }
640+
// From here on it can cost a request, so focus gets a ration: a window clicked in and out of
641+
// while offline must not retry the refresh on every click.
642+
if (reason === 'focus' && Date.now() - lastSessionCheck <= SESSION_CHECK_EVERY_MS) { return; }
643+
lastSessionCheck = Date.now();
644+
if (!session.accessNeedsRefresh(token)) { return; }
645+
dbg('cloud.sessionCheck', { reason, expiresAt: session.jwtExpiresAt(token) });
646+
await refreshCloudToken(); // an expired refresh token lands in sessionExpired() from inside
647+
} catch (e) { dbg('cloud.sessionCheck', { reason, error: String((e && e.message) || e) }); }
608648
}
609649

610650
/** Gateway-mode token refresh (the streaming 401 retry path). Delegates to refreshCloudToken. */
@@ -1432,7 +1472,7 @@ function newChat() {
14321472
post({ type: 'reset' });
14331473
postContextFiles();
14341474
postMemoryDigest(); // the fresh empty state shows the welcome-back strip
1435-
replaySessionExpired().catch(() => { }); // `reset` emptied the log, and an unanswered expiry's card with it
1475+
replaySessionExpired(); // `reset` emptied the log, and an unanswered expiry's card with it
14361476
}
14371477

14381478
/** The currently open file as a context block (capped), or null. */
@@ -3293,15 +3333,17 @@ function onConfigChanged(e) {
32933333
if (e.affectsConfiguration('levelcode.ai')) { sendConfigToWebview(); }
32943334
if (e.affectsConfiguration('levelcode.ai.providerMode') || e.affectsConfiguration('levelcode.cloud')) {
32953335
postAccount();
3296-
replaySessionExpired().catch(() => { });
3336+
replaySessionExpired();
32973337
}
32983338
}
32993339

33003340
function activate(context) {
3301-
// A window that comes back after a while away may have outlived its access token (8 h). Re-check
3302-
// on focus, throttled, so the expiry is found before the next message rather than by it.
3341+
// A window that comes back after a while away may have outlived its access token (8 h) — or
3342+
// another window may have ended, left or renewed the session they share. Re-check on every focus,
3343+
// so either is found before the next message rather than by it. The look is one local read;
3344+
// checkCloudSession rations the part that can cost a request.
33033345
context.subscriptions.push(vscode.window.onDidChangeWindowState((st) => {
3304-
if (st.focused && Date.now() - lastSessionCheck > 10 * 60 * 1000) { checkCloudSession('focus').catch(() => {}); }
3346+
if (st.focused) { checkCloudSession('focus'); }
33053347
}));
33063348
ctx = context;
33073349
// Constructed directly rather than by registerWebviewViewProvider: the chat is no longer a

‎extensions/levelcode-ai/test/sessionExpiredHost.test.js‎

Lines changed: 119 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
* - the account message tells the webview whether an expiry is still waiting, so the card leaves
2323
* when BYOK is chosen in Settings and comes back if gateway mode is
2424
* - a 2xx refresh reply that cannot be stored changes nothing, and nothing escapes into `ready`
25+
* - the replay asks whether an expiry is waiting AFTER its read, not only before it
26+
* - a window catches up with a session another window ended, left or renewed — on every focus
2527
*
2628
* The functions are sliced out of the shipped extension.js and run against stand-ins for what they
2729
* touch (SecretStorage, globalState, fetch, the provider adapter), the way mcpManage.test.js does
@@ -92,7 +94,7 @@ function decl(name) {
9294

9395
const FUNCTIONS = [
9496
'post', 'providerMode', 'providerErrorMessage', 'isAuthError', 'cloudEndpoint', 'cloudApiUrl',
95-
'refreshCloudToken', 'renewSession', 'withSessionLock', 'sameSession', 'sessionExpired', 'postSessionExpired',
97+
'refreshCloudToken', 'renewSession', 'storedToken', 'catchUpWithStoredSession', 'withSessionLock', 'sameSession', 'sessionExpired', 'postSessionExpired',
9698
'sessionExpiredPending', 'isEndedSessionError', 'clearSessionExpired', 'replaySessionExpired', 'checkCloudSession',
9799
'refreshGatewayToken', 'prepProviderRequest', 'handleSend', 'currentAccount', 'postAccount', 'accountSignOut',
98100
'storeSession', 'newChat', 'resumeSession', 'restoreCheckpoint', 'onConfigChanged'
@@ -109,7 +111,7 @@ const makeHost = new Function('env', [
109111
'let activeWebview = env.webview, cloudSignedIn = env.signedIn, agentMode = false, conversation = [], pendingContext = null, conversationEpoch = 0, abort = null;',
110112
'let agentMessages = [], currentCheckpoint = null, lastAgentGoal = null;',
111113
decl('ACCOUNT_TOKEN_KEY'), decl('ACCOUNT_REFRESH_KEY'), decl('ACCOUNT_PROFILE_KEY'), decl('ACCOUNT_EXPIRED_KEY'),
112-
decl('sessionQueue'), decl('sessionGeneration'), decl('lastSessionCheck'),
114+
decl('sessionQueue'), decl('sessionGeneration'), decl('SESSION_CHECK_EVERY_MS'), decl('lastSessionCheck'),
113115
...FUNCTIONS.map(extract),
114116
'return { ' + FUNCTIONS.join(', ') + ',',
115117
' KEY: { token: ACCOUNT_TOKEN_KEY, refresh: ACCOUNT_REFRESH_KEY, profile: ACCOUNT_PROFILE_KEY, expired: ACCOUNT_EXPIRED_KEY },',
@@ -197,6 +199,10 @@ function boot(over) {
197199
/** @type {Error|null} */ keyError: null,
198200
/** Makes SecretStorage refuse every write — a locked keychain. */
199201
/** @type {Error|null} */ storeError: null,
202+
/** …and every read. */
203+
/** @type {Error|null} */ readError: null,
204+
/** While set, every SecretStorage read waits on it — a read that is slow to come back. */
205+
/** @type {Promise<any>|null} */ readGate: null,
200206
/** The account messages the popover (and the sign-in card) were sent. */
201207
accounts: () => posted.filter((m) => m.type === 'account'),
202208
types: () => posted.map((m) => m.type),
@@ -207,7 +213,12 @@ function boot(over) {
207213
session: o.refreshTimeoutMs ? Object.assign({}, session, { REFRESH_TIMEOUT_MS: o.refreshTimeoutMs }) : session,
208214
ctx: {
209215
secrets: {
210-
get: async (k) => { await tick(); return secrets.get(k); },
216+
get: async (k) => {
217+
await tick();
218+
if (t.readGate) { await t.readGate; }
219+
if (t.readError) { throw t.readError; }
220+
return secrets.get(k);
221+
},
211222
store: async (k, v) => {
212223
const op = 'store ' + k; ops.push(op); t.onOp(op); await tick();
213224
if (t.storeError) { throw t.storeError; }
@@ -915,5 +926,110 @@ async function test(name, fn) { await fn(); n++; console.log(' ok - ' + name);
915926
assert.strictEqual(t.state.get(t.K.expired), undefined);
916927
});
917928

929+
// ── the replay asks again after its read ─────────────────────────────────────────────────────
930+
for (const [how, answer] of [
931+
['"Use my own key instead"', (t) => t.host.clearSessionExpired()],
932+
['BYOK chosen in Settings', async (t) => { t.settings['levelcode.ai.providerMode'] = 'byok'; }]
933+
]) {
934+
await test('replay: an expiry answered by ' + how + ' WHILE the token read is out is not shown again', async () => {
935+
const t = boot({ expired: true, profile: { name: 'Ada' } });
936+
const read = deferred();
937+
t.readGate = read.promise; // SecretStorage is slow to answer…
938+
const replay = t.host.replaySessionExpired();
939+
await settle(); // …so the replay is parked on that read…
940+
await answer(t); // …and the user answers the card meanwhile.
941+
read.resolve();
942+
await replay;
943+
assert.deepStrictEqual(t.posted, [], 'no card for an expiry that is no longer waiting');
944+
});
945+
}
946+
947+
await test('replay: the same slow read with nothing answered still ends in the card', async () => {
948+
const t = boot({ expired: true, profile: { name: 'Ada' } });
949+
const read = deferred();
950+
t.readGate = read.promise;
951+
const replay = t.host.replaySessionExpired();
952+
await settle();
953+
read.resolve();
954+
await replay;
955+
assert.deepStrictEqual(t.types(), ['sessionExpired']);
956+
});
957+
958+
await test('a store that will not READ: the check and the replay finish quietly — neither can break what awaits them', async () => {
959+
const t = boot({ expired: true, access: deadAccess(), refresh: 'r1' });
960+
t.readError = new Error('the keychain is locked');
961+
await t.host.checkCloudSession('ready');
962+
await t.host.replaySessionExpired();
963+
assert.deepStrictEqual(t.posted, []);
964+
});
965+
966+
// ── a window catches up with what another window did to the session they share ───────────────
967+
await test('another window ENDED the session before this window\'s focus check: flag, popover and card catch up', async () => {
968+
const t = boot({ access: liveAccess(), refresh: 'r1', profile: { name: 'Ada' } });
969+
assert.strictEqual(t.host.signedIn(), true, 'this window is showing a live session');
970+
t.secrets.delete(t.K.token); t.secrets.delete(t.K.refresh); t.state.set(t.K.expired, true); // the other window's sessionExpired()
971+
await t.host.checkCloudSession('focus');
972+
assert.strictEqual(t.host.signedIn(), false);
973+
assert.deepStrictEqual(t.posted.filter((m) => m.type === 'sessionExpired'), [{ type: 'sessionExpired', name: 'Ada', message: session.SESSION_EXPIRED_MESSAGE }]);
974+
assert.deepStrictEqual(t.accounts().map(brief), [{ signedIn: false, mode: 'gateway', expired: true, status: 'signedout' }]);
975+
assert.strictEqual(t.calls.refresh.length, 0, 'nothing to renew, nothing requested');
976+
assert.deepStrictEqual(t.ops, [], 'and nothing written');
977+
const before = t.posted.length;
978+
await t.host.checkCloudSession('focus'); // clicking back in again
979+
assert.strictEqual(t.posted.length, before, 'said once, not on every focus');
980+
});
981+
982+
await test('another window SIGNED OUT: the popover catches up, and there is no card — nothing expired', async () => {
983+
const t = boot({ access: liveAccess(), refresh: 'r1', profile: { name: 'Ada' } });
984+
t.secrets.delete(t.K.token); t.secrets.delete(t.K.refresh); t.state.delete(t.K.profile); // the other window's accountSignOut()
985+
await t.host.checkCloudSession('focus');
986+
assert.strictEqual(t.host.signedIn(), false);
987+
assert.deepStrictEqual(t.accounts().map(brief), [{ signedIn: false, mode: 'gateway', expired: false, status: 'signedout' }]);
988+
assert.ok(!t.types().includes('sessionExpired'));
989+
});
990+
991+
await test('another window SIGNED IN while this one showed the card: the account message is its cue to go', async () => {
992+
const t = boot({ expired: true, profile: { name: 'Ada' } });
993+
assert.strictEqual(t.host.signedIn(), false);
994+
t.secrets.set(t.K.token, liveAccess()); t.secrets.set(t.K.refresh, 'r2'); t.state.delete(t.K.expired); // the other window's storeSession()
995+
await t.host.checkCloudSession('focus');
996+
assert.strictEqual(t.host.signedIn(), true);
997+
assert.deepStrictEqual(t.accounts().map(brief), [{ signedIn: true, mode: 'gateway', expired: false, status: undefined }]);
998+
assert.ok(!t.types().includes('sessionExpired'));
999+
});
1000+
1001+
await test('a focus check that lands in the MIDDLE of this window\'s own sign-out does not take it for another window\'s', async () => {
1002+
// Mid-sign-out the flag already says "signed out" while the token is still stored. Read then,
1003+
// the two disagree, and the catch-up would announce a session this window is busy ending.
1004+
const t = boot({ access: liveAccess(), refresh: 'r1', profile: { name: 'Ada' } });
1005+
let check = null;
1006+
t.onOp = (op) => { if (!check && op === 'forget ' + t.K.token) { check = t.host.checkCloudSession('focus'); } };
1007+
await t.host.accountSignOut();
1008+
assert.ok(check, 'the focus check was started mid-sign-out');
1009+
await check;
1010+
assert.deepStrictEqual(t.accounts().map(brief), [{ signedIn: false, mode: 'gateway', expired: false, status: 'signedout' }], 'one resync: the sign-out\'s own');
1011+
assert.strictEqual(t.host.signedIn(), false);
1012+
});
1013+
1014+
await test('the ration is on the refresh, not on the look: a second focus skips the request but still catches up', async () => {
1015+
const t = boot({ access: deadAccess(), refresh: 'r1', profile: { name: 'Ada' } });
1016+
t.refreshReply = 'offline';
1017+
await t.host.checkCloudSession('focus');
1018+
await t.host.checkCloudSession('focus');
1019+
assert.strictEqual(t.calls.refresh.length, 1, 'offline: one attempt, not one per click');
1020+
t.secrets.delete(t.K.token); t.secrets.delete(t.K.refresh); t.state.set(t.K.expired, true); // meanwhile, in another window
1021+
await t.host.checkCloudSession('focus');
1022+
assert.strictEqual(t.posted.filter((m) => m.type === 'sessionExpired').length, 1, 'found inside the ten minutes');
1023+
assert.strictEqual(t.calls.refresh.length, 1);
1024+
});
1025+
1026+
await test('`ready` does not say it twice: the handler has just read the same token, and the card is replayed last', async () => {
1027+
const t = boot({ expired: true, profile: { name: 'Ada' } });
1028+
await t.host.checkCloudSession('ready');
1029+
assert.deepStrictEqual(t.posted, [], 'the check is quiet');
1030+
await t.host.replaySessionExpired();
1031+
assert.deepStrictEqual(t.types(), ['sessionExpired'], 'the last step speaks');
1032+
});
1033+
9181034
console.log('\nsessionExpiredHost: ' + n + ' tests passed.');
9191035
})().catch((e) => { console.error(e); process.exit(1); });

‎extensions/levelcode-ai/test/sessionExpiredUi.test.js‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,17 @@ test('host: the webview ready handler checks the session before the first messag
8585
assert.ok(/onDidChangeWindowState/.test(ext), 'focus re-check');
8686
});
8787

88+
test('host: every focus runs the check — the ten-minute ration is on the refresh, inside it', () => {
89+
// With the ration in the hook, a window that had checked a minute ago did not look again, and so
90+
// did not notice another window ending the session they share until the next message failed.
91+
assert.ok(/onDidChangeWindowState\(\(st\) => \{\s*if \(st\.focused\) \{ checkCloudSession\('focus'\); \}/.test(ext), 'no time condition in the hook');
92+
const fn = ext.slice(ext.indexOf('async function checkCloudSession('), ext.indexOf('async function refreshGatewayToken('));
93+
const catchUp = fn.indexOf('await catchUpWithStoredSession(token)');
94+
const ration = fn.indexOf("reason === 'focus' && Date.now() - lastSessionCheck <= SESSION_CHECK_EVERY_MS");
95+
assert.ok(catchUp > 0 && ration > 0, 'both halves present');
96+
assert.ok(catchUp < ration, 'catching up comes before the ration');
97+
});
98+
8899
test('host: ready replays an unanswered expiry LAST — after the check, and under a replayed transcript', () => {
89100
// The focus check can end the session with no chat open; post() then has nowhere to send the card,
90101
// and the next `ready` finds no token to check. The replay is what shows it — and it has to come

0 commit comments

Comments
 (0)