Skip to content

Commit 4ae6759

Browse files
committed
fix(cloud): #96 review — a request stays with its session and gateway; waiting on a renewal is free
1. A token from another session, or for another host, could be put on a request. authRetry took "a different gateway token is stored" to mean "the token was renewed", and copied it onto the original request. Two ways that is not true. Someone else signs in while the request is out — or half-way through a Sketch run, which sends one request for minutes — and the old account's prompt was sent again on the new account's token. And when the cloud host changes, the token stored for the new host went to the old URL. A request now stays with what it was sent on. The host hands authRetry its sessionGeneration: the count the refresh already keeps, which moves on a sign-in, a sign-out and an expiry and stays put through a renewal. It is noted per REQUEST, the first time one is sent, so the later nodes of a run are held to the session the run started in. A stored token is a renewal only for the same generation, provider and gateway URL. Otherwise recovery stops — before refreshing, so the new session's refresh token is not spent on the old request, and again before re-sending — and the request's own error stands. An ENDED session still outranks that and gets its sentence. What a renewal achieved is now read from what is stored afterwards, not from its answer. 2. The ration on ghost-text renewals also turned away requests that only had to wait. A keystroke cancels the completion that started a renewal; the next pause is refused on the old token while that renewal is still out, and was answered "one renewal a minute": no suggestion, though the token arrived a moment later. The limit is on STARTING a renewal. Waiting on one already out starts nothing, and is never held back. Also: authRetryCallers.test.js slices storedToken and catchUpWithStoredSession when extension.js has them. #92's round four (f4f38c1) has the refresh call both, and this file is stacked on it. A slice that is out of date still says so by name. Verified: 46 suites, 884 cases, 0 failing. authRetry.test.js goes from 31 to 39 cases and authRetryCallers.test.js from 68 to 72: a real storeSession() sign-in while an inline edit is out, and half-way through a Sketch run; the cloud host changed under a request; the next completion arriving while the cancelled one's renewal is still pending. On the module as pushed, the cases added or tightened here fail and the rest pass. 48 breakages — each rule removed, each caller reverted, each option dropped — fail a case each. node --check is clean, and tsc --checkJs reports nothing new. Trial-merged with no conflicts and 0 failing: develop (4e49295), 48 suites; #92's head (f4f38c1), 46; #95's new head (4ef3ee0), 46. Not run in the editor.
1 parent 7fefaa3 commit 4ae6759

4 files changed

Lines changed: 303 additions & 51 deletions

File tree

‎extensions/levelcode-ai/extension.js‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -628,9 +628,12 @@ async function refreshGatewayToken() {
628628
* never refreshed or retried. The chat and the agent keep their own retry (handleSend; the agent's
629629
* refreshAuth hook): they have a transcript and a sign-in card to keep in step with it.
630630
*
631-
* ONE instance for the window, so that requests which fail together wait on a single renewal.
631+
* ONE instance for the window, so that requests which fail together wait on a single renewal. And it
632+
* is told which session this window is on (sessionGeneration), for the reason the refresh keeps that
633+
* count: a request belongs to the session it was sent in. A token stored by a LATER sign-in is not a
634+
* renewal of the old one, and is never put on a request the earlier session sent.
632635
*/
633-
const authRetry = createAuthRetry({ prepProviderRequest, refreshGatewayToken, isAuthError, dbg });
636+
const authRetry = createAuthRetry({ prepProviderRequest, refreshGatewayToken, isAuthError, sessionGeneration: () => sessionGeneration, dbg });
634637

635638
/**
636639
* Resolve everything needed to call the active provider: id, key, model, baseURL, maxTokens.

‎extensions/levelcode-ai/providers/authRetry.js‎

Lines changed: 72 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,16 @@
3636
* - A request refused on a token that has been renewed SINCE it was sent — by a sibling node, by
3737
* the chat — is just sent again on the stored one. And the renewed token is written onto the
3838
* request itself, so whatever else holds that request (the rest of a Sketch run) is on it too.
39-
* - A request nobody asked for (ghost text) may start a renewal once a minute at most. It is
39+
* - A request stays with the session and the gateway it was sent on. "A different token is
40+
* stored" is only a renewal if it is the same session's token for the same URL: after a
41+
* sign-in as someone else, or a change of cloud host, the stored token is not this request's.
42+
* Putting it on the request would send one host's credential to another, or replay one
43+
* account's prompt as another's — so recovery stops there, before refreshing and again before
44+
* re-sending, and the request's own error stands.
45+
* - A request nobody asked for (ghost text) may START a renewal once a minute at most. It is
4046
* sent on every pause in typing; a gateway that keeps answering 401 must not turn typing into
41-
* a stream of refreshes, each one rotating the session's credentials.
47+
* a stream of refreshes, each one rotating the session's credentials. Waiting on a renewal
48+
* that is already out starts nothing, and is never held back.
4249
*--------------------------------------------------------------------------------------------*/
4350
// @ts-check
4451
'use strict';
@@ -52,22 +59,21 @@ const BACKGROUND_RENEWAL_INTERVAL_MS = 60 * 1000;
5259
function aborted(signal) { return !!(signal && signal.aborted); }
5360

5461
/**
55-
* `promise`'s value — or `undefined` the moment `signal` aborts, whichever comes first. The renewal
56-
* itself is not cancelled (other requests may be waiting on it, and its answer is worth having
57-
* either way); the caller just stops waiting, so Stop and Cancel stay as quick as they were.
58-
* `promise` must not reject.
59-
* @template T
60-
* @param {Promise<T>} promise
62+
* Wait for `promise` — or until `signal` aborts, whichever comes first. The renewal itself is not
63+
* cancelled (other requests may be waiting on it, and what it stores is worth having either way);
64+
* the caller just stops waiting, so Stop and Cancel stay as quick as they were. `promise` must not
65+
* reject.
66+
* @param {Promise<void>} promise
6167
* @param {AbortSignal|undefined} signal
62-
* @returns {Promise<T|undefined>}
68+
* @returns {Promise<void>}
6369
*/
6470
function unlessAborted(promise, signal) {
6571
if (!signal) { return promise; }
6672
return new Promise((resolve) => {
67-
if (signal.aborted) { resolve(undefined); return; }
68-
const stop = () => resolve(undefined);
73+
if (signal.aborted) { resolve(); return; }
74+
const stop = () => resolve();
6975
signal.addEventListener('abort', stop, { once: true });
70-
promise.then((v) => { signal.removeEventListener('abort', stop); resolve(v); });
76+
promise.then(() => { signal.removeEventListener('abort', stop); resolve(); });
7177
});
7278
}
7379

@@ -83,13 +89,16 @@ function unlessAborted(promise, signal) {
8389
* that fail together share.
8490
*
8591
* prepProviderRequest — the host's; asked again (never prompting) for what is stored NOW
86-
* refreshGatewayToken — the host's; true when a usable access token is in place afterwards
92+
* refreshGatewayToken — the host's; what it achieved is read back from what is stored, not from its answer
8793
* isAuthError — the host's test for "the provider refused the credentials"
94+
* sessionGeneration — the host's count of the sessions this window has been through: it moves
95+
* on a sign-in, a sign-out and an expiry, and stays put through a renewal
8896
*
8997
* @param {{
9098
* prepProviderRequest: (opts?: any) => Promise<any>,
91-
* refreshGatewayToken: () => Promise<boolean>,
99+
* refreshGatewayToken: () => Promise<any>,
92100
* isAuthError: (e: any) => boolean,
101+
* sessionGeneration?: () => any,
93102
* dbg?: (label: string, data?: any) => void,
94103
* now?: () => number
95104
* }} host
@@ -98,64 +107,93 @@ function unlessAborted(promise, signal) {
98107
function createAuthRetry(host) {
99108
const now = host.now || Date.now;
100109
const dbg = host.dbg || (() => { });
101-
/** @type {Promise<boolean>|null} the renewal that is out, if one is */
110+
const generation = host.sessionGeneration || (() => 0);
111+
/** @type {Promise<void>|null} the renewal that is out, if one is */
102112
let renewing = null;
103113
let lastBackgroundRenewal = -Infinity;
114+
/**
115+
* The session each request's credentials came from: the host's generation, noted the first time
116+
* the request is sent through here — for every caller, the moment after it was prepared. Kept
117+
* per REQUEST, not per send: a Sketch run sends its one request many times over minutes, and a
118+
* sign-in half-way through must not put the new account's token on the nodes still to come.
119+
* @type {WeakMap<object, any>}
120+
*/
121+
const sessionOf = new WeakMap();
104122

105-
/** Renew the token — or wait on the renewal already out. Never rejects: a refresh that throws has failed. */
123+
/**
124+
* Renew the token — or wait on the renewal already out. Never rejects, and says nothing: a
125+
* refresh that throws has failed, and whether one worked is read from what is stored afterwards.
126+
*/
106127
function renew() {
107128
if (!renewing) {
108129
renewing = Promise.resolve().then(() => host.refreshGatewayToken())
109-
.then(Boolean, () => false)
110-
.then((ok) => { renewing = null; return ok; });
130+
.then(() => { }, () => { })
131+
.then(() => { renewing = null; });
111132
}
112133
return renewing;
113134
}
114135

115136
/**
116-
* What the host would send a request on NOW, read against the token that was just refused:
117-
* a different cloud token (`apiKey`), or the news that the session has ended.
137+
* What the host would send this request on NOW, read against what it was sent on:
138+
*
139+
* 'ended' — the session is over: the host answers `signedOut`
140+
* 'superseded' — the request is no longer this session's to recover. The session it was sent
141+
* in has been replaced (a sign-in, a sign-out), or the host would now send it
142+
* somewhere else: another gateway URL, the user's own provider, nowhere.
143+
* Whatever token is stored for THAT is not put on this request.
144+
* 'renewed' — same session, same gateway, a different token in place: `apiKey`
145+
* 'same' — the token that was refused is still the one stored
146+
*
147+
* @param {any} req
118148
* @param {string} sent
119-
* @returns {Promise<{apiKey?: string, ended?: boolean}>}
149+
* @returns {Promise<{state: 'ended'|'superseded'|'renewed'|'same', apiKey?: string}>}
120150
*/
121-
async function stored(sent) {
151+
async function look(req, sent) {
122152
const cur = await host.prepProviderRequest({ prompt: false }); // a retry never opens a key dialog
123-
if (cur && cur.ok && cur.gateway && cur.apiKey && cur.apiKey !== sent) { return { apiKey: cur.apiKey }; }
124-
return { ended: !!cur && !cur.ok && cur.reason === 'signedOut' };
153+
if (cur && !cur.ok && cur.reason === 'signedOut') { return { state: 'ended' }; }
154+
const sameGateway = !!cur && cur.ok && cur.gateway && cur.providerId === req.providerId && cur.baseURL === req.baseURL;
155+
if (!sameGateway || generation() !== sessionOf.get(req)) { return { state: 'superseded' }; }
156+
return cur.apiKey && cur.apiKey !== sent ? { state: 'renewed', apiKey: cur.apiKey } : { state: 'same' };
125157
}
126158

127159
/**
128160
* A gateway request carrying `sent` was refused. Find the token to send it again on.
161+
* @param {any} req
129162
* @param {string} sent
130163
* @param {AuthRetryOptions} o
131164
* @returns {Promise<{outcome: string, apiKey?: string}>}
132165
*/
133-
async function recover(sent, o) {
134-
let cur = await stored(sent);
135-
if (cur.apiKey) { return { outcome: 'already-renewed', apiKey: cur.apiKey }; }
136-
if (cur.ended) { return { outcome: 'ended' }; }
137-
if (o.background) {
166+
async function recover(req, sent, o) {
167+
let cur = await look(req, sent);
168+
if (cur.state === 'renewed') { return { outcome: 'already-renewed', apiKey: cur.apiKey }; }
169+
if (cur.state !== 'same') { return { outcome: cur.state }; } // ended, or no longer this request's: nothing to refresh
170+
// The limit is on STARTING a renewal. One that is already out costs nothing to wait on — and the
171+
// request that started it may be gone by now (the next keystroke cancelled it) with this one
172+
// the only one left to use what it brings back.
173+
if (o.background && !renewing) {
138174
if (now() - lastBackgroundRenewal < BACKGROUND_RENEWAL_INTERVAL_MS) { return { outcome: 'throttled' }; }
139175
lastBackgroundRenewal = now();
140176
}
141-
const renewed = await unlessAborted(renew(), o.signal);
177+
await unlessAborted(renew(), o.signal);
142178
if (aborted(o.signal)) { return { outcome: 'aborted' }; }
143-
cur = await stored(sent);
144-
if (cur.ended) { return { outcome: 'ended' }; }
145-
return renewed && cur.apiKey ? { outcome: 'renewed', apiKey: cur.apiKey } : { outcome: 'failed' };
179+
// Asked again, not carried over: a sign-in or a sign-out can land while the renewal is out.
180+
cur = await look(req, sent);
181+
if (cur.state === 'renewed') { return { outcome: 'renewed', apiKey: cur.apiKey }; }
182+
return { outcome: cur.state === 'same' ? 'failed' : cur.state };
146183
}
147184

148185
return async function authRetry(req, send, opts) {
149186
const o = opts || {};
150187
const sent = req.apiKey;
188+
if (req.gateway && !sessionOf.has(req)) { sessionOf.set(req, generation()); }
151189
try {
152190
return await send(req);
153191
} catch (e) {
154192
if (!req.gateway || !host.isAuthError(e) || aborted(o.signal) || (o.streamed && o.streamed())) { throw e; }
155193
/** @type {{outcome: string, apiKey?: string}} */
156194
let found;
157195
// Whatever goes wrong while looking for a way to recover, the request's own failure is the news.
158-
try { found = await recover(sent, o); } catch { found = { outcome: 'failed' }; }
196+
try { found = await recover(req, sent, o); } catch { found = { outcome: 'failed' }; }
159197
dbg('auth.retry', { outcome: found.outcome });
160198
if (found.outcome === 'ended') {
161199
throw Object.assign(new Error(SESSION_EXPIRED_MESSAGE), { code: 'session_expired', cause: e });

0 commit comments

Comments
 (0)