Invalidate query cache on disconnect to fix stale auth tokens

[whoiskatrin]

Summary

• Invalidates the query cache entry when the WebSocket connection closes
• Ensures async query functions (e.g., for auth tokens) are re-executed on reconnect
• Eliminates auth failures caused by reusing stale cached tokens after disconnect

Problem

When using useAgent with an async query function for authentication:

query: async () => ({ token: await getIdToken(true) })

The cached query result was reused on reconnect even if the token had expired. Reconnects don't change the cache key or wait for TTL expiration, so stale tokens were sent to the server causing auth failures.

Solution

By calling deleteCacheEntry(cacheKey) in the onClose handler, we invalidate the cache on disconnect. When the socket reconnects, the async query will be re-executed to fetch a fresh token.

This is simpler than adding new API options (queryCacheMode, alwaysRefreshQuery, etc.) and provides the correct behavior by default - reconnects should always use fresh auth data.

Fixes #836

[changeset-bot]

🦋 Changeset detected

Latest commit: 8d3bdb8

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 1 package

Name	Type
agents	Patch

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

[claude]

Claude Code Review

Issue: Race condition when resolvedQuery is undefined during reconnection.

In react.tsx:560-567, the onClose handler invalidates the cache and sets awaitingQueryRefresh = true for async queries to prevent reconnection until fresh query params are ready. However, there's a potential issue:

The Problem:

1. onClose sets awaitingQueryRefresh = true and invalidates cache
2. Socket is disabled via socketEnabled = !awaitingQueryRefresh && ... (line 453)
3. The effect at lines 400-404 waits for resolvedQuery !== undefined to re-enable socket
4. But: If use(queryPromise!) at line 373 suspends waiting for the new query, resolvedQuery remains from the previous render
5. This means the effect at 400-404 might fire immediately with the stale resolvedQuery, re-enabling the socket before the new query resolves

The Race:

onClose → awaitingQueryRefresh=true, invalidate cache
  ↓
Component re-renders with awaitingQueryRefresh=true
  ↓
use(queryPromise!) may suspend OR return old cached value in this render cycle
  ↓  
Effect sees awaitingQueryRefresh=true && resolvedQuery=<stale value>
  ↓
Socket re-enables with stale auth token

Suggested Fix:
The effect should verify that resolvedQuery is actually fresh, not from before the disconnect. One approach:

// Track invalidation timestamp
const queryInvalidatedAtRef = useRef<number>(0);

// In onClose:
queryInvalidatedAtRef.current = Date.now();
deleteCacheEntry(cacheKeyRef.current);

// In the effect - only re-enable if query was fetched after invalidation:
useEffect(() => {
  if (awaitingQueryRefresh && resolvedQuery !== undefined) {
    const entry = getCacheEntry(cacheKey);
    if (entry && entry.createdAt > queryInvalidatedAtRef.current) {
      setAwaitingQueryRefresh(false);
    }
  }
}, [awaitingQueryRefresh, resolvedQuery, cacheKey]);

This requires tracking cache entry creation time (not currently stored).

Alternative: Key the effect off cacheInvalidatedAt instead:

useEffect(() => {
  if (awaitingQueryRefresh) {
    setAwaitingQueryRefresh(false);
  }
}, [cacheInvalidatedAt]);

Test Coverage: The test at cache-invalidation.test.tsx:169-198 pre-populates cache before connection, so doesn't exercise the actual async query re-execution path. A test that triggers real async query execution on reconnect would catch this race.

---

Otherwise solid implementation - good use of refs for cache key capture, comprehensive test coverage, and the core fix (cache invalidation on disconnect) is sound.

[pkg-pr-new]

Open in StackBlitz

npm i https://pkg.pr.new/cloudflare/agents@839

commit: 8d3bdb8

[aurewill-gavel]

Main remaining concern is a race - the socket’s reconnect loop can fire with the old URL/query before the re‑render (triggered by cache invalidation + state bump) completes. The only obvious way I see to eliminate the race is to pause reconnects while the query refresh happens.

Looks like we can use enabled (UsePartySocketOptions) to handle this? So a pattern like:

• onClose: setEnabled(false) + invalidate cache + bump state
• after query resolves / re-render: setEnabled(true)

Idea being to guarantee no reconnect attempt happens before the new query is ready.

[whoiskatrin]

yes, that's valid, we might just do that in the end

[aurewill-gavel]

Could we add an integration test where the query returns a different token on each invocation, then assert which token appears on the first reconnect? Probably makes sense to make it intentionally race‑y by:

1. Delaying the async query and
2. Forcing fast reconnects (set minReconnectionDelay/maxReconnectionDelay low).

Then assert which token the first reconnect used:

let calls = 0;
query: async () => { await delay(100); return { token: `t-${++calls}` } };

force reconnect();
assert first reconnect used "t-2" (not "t-1")

Without that timing pressure, the test could pass even if a race still exists, simply because the reconnect happens after the refresh render completes. This would validate the behavior the PR claims (fresh token on reconnect), beyond just cache deletion.

[whoiskatrin]

@aurewill-gavel give the latest version a go, please, if there will be more gaps, we will revisit

[aurewill-gavel]

@whoiskatrin Logic looks sound. That test I mentioned would be ideal but obv your decision whether or not to add it.

[threepointone]

let's land this, I suspect suspense might help out here a bit. and because it's not an api change, we can revisit and make it stronger if we need to.