`useQuery` (cache-first) stays stale when an optimistic mutation is written during its own result delivery

import { ApolloClient, ApolloLink, InMemoryCache, Observable, gql } from '@apollo/client';
import { ApolloProvider, useMutation, useQuery } from '@apollo/client/react';
import { render, screen, waitFor } from '@testing-library/react';
import { useEffect } from 'react';
import { describe, expect, it } from 'vitest';

const ITEM = gql`
  query Item {
    item(id: "1") {
      id
      read
    }
  }
`;
const MARK_READ = gql`
  mutation MarkRead {
    markRead(id: "1") {
      id
      read
    }
  }
`;

// A link that answers after 10ms. `completeInSameTick` controls whether it
// completes in the same tick it emits, or one task later (as a fetch-backed
// HttpLink does: the body stream ends after the result is emitted).
function makeClient(completeInSameTick: boolean) {
  let serverRead = false;
  let mutationsSettled = 0;
  const link = new ApolloLink(
    (operation) =>
      new Observable((observer) => {
        const timer = setTimeout(() => {
          if (operation.operationName === 'MarkRead') {
            serverRead = true;
            observer.next({ data: { markRead: { __typename: 'Item', id: '1', read: true } } });
          } else {
            observer.next({ data: { item: { __typename: 'Item', id: '1', read: serverRead } } });
          }
          const done = () => {
            if (operation.operationName === 'MarkRead') mutationsSettled += 1;
            observer.complete();
          };
          if (completeInSameTick) done();
          else setTimeout(done, 0);
        }, 10);
        return () => clearTimeout(timer);
      })
  );
  const client = new ApolloClient({ cache: new InMemoryCache(), link });
  return { client, settled: () => mutationsSettled };
}

const VARIANT = { optimistic: true, fetchPolicy: 'cache-first' as 'cache-first' | 'cache-and-network' };

// Mark the item read as soon as the query reports it unread.
function App() {
  const { data } = useQuery<{ item: { id: string; read: boolean } }>(ITEM, {
    fetchPolicy: VARIANT.fetchPolicy,
  });
  const [markRead] = useMutation(
    MARK_READ,
    VARIANT.optimistic
      ? { optimisticResponse: { markRead: { __typename: 'Item', id: '1', read: true } } }
      : {}
  );
  const read = data?.item.read;
  useEffect(() => {
    if (read === false) void markRead();
  }, [read, markRead]);
  return <output aria-label="read">{String(read)}</output>;
}

async function run(completeInSameTick: boolean) {
  const { client, settled } = makeClient(completeInSameTick);
  render(
    <ApolloProvider client={client}>
      <App />
    </ApolloProvider>
  );
  await waitFor(() => expect(settled()).toBe(1));
  await new Promise((resolve) => setTimeout(resolve, 50));
  const cached = client.readQuery<{ item: { read: boolean } }>({ query: ITEM });
  return { rendered: screen.getByLabelText('read').textContent, cached: cached?.item.read };
}

describe('optimistic write during result delivery', () => {
  it('control: link completes in the same tick it emits', async () => {
    const r = await run(true);
    expect(r).toEqual({ rendered: 'true', cached: true });
  });

  it('bug: link completes one task after it emits', async () => {
    const r = await run(false);
    expect(r).toEqual({ rendered: 'true', cached: true });
  });

  it('variant: same timing, no optimisticResponse', async () => {
    VARIANT.optimistic = false;
    try {
      const r = await run(false);
      expect(r).toEqual({ rendered: 'true', cached: true });
    } finally {
      VARIANT.optimistic = true;
    }
  });

  it('variant: same timing, cache-and-network', async () => {
    VARIANT.fetchPolicy = 'cache-and-network';
    try {
      const r = await run(false);
      expect(r).toEqual({ rendered: 'true', cached: true });
    } finally {
      VARIANT.fetchPolicy = 'cache-first';
    }
  });
});