Allow passing compute fn to the consturctor

ui/src/base/async_memo.ts

diff --git a/ui/src/base/async_memo.ts b/ui/src/base/async_memo.ts
index 592775c..176cbc5 100644
--- a/ui/src/base/async_memo.ts
+++ b/ui/src/base/async_memo.ts
@@ -91,6 +91,10 @@
 // Simple alias for a function that returns a Promise of type T.
 type AsyncFunc<T> = () => Promise<T>;
 
+export type AsyncMemoCompute<T> = (
+  signal: CancellationSignal,
+) => Promise<T | typeof TASK_CANCELLED>;
+
 /**
  * Runs async functions one at a time with cancellation support.
  *
@@ -145,10 +149,9 @@
 
 export interface AsyncMemoOptions<T, K extends JSONCompatible<K>> {
   readonly key: K;
-  // Function to call when the key changes to fetch the memo data.
-  readonly compute: (
-    signal: CancellationSignal,
-  ) => Promise<T | typeof TASK_CANCELLED>;
+  // Function to call when the key changes to fetch the memo data. Overrides
+  // the compute function passed to the constructor, if present.
+  readonly compute?: AsyncMemoCompute<T>;
   // If provided, compute() only runs when this is truthy.
   readonly enabled?: boolean;
   // Keep returning the previous memo while the new one loads as long as only
@@ -182,6 +185,8 @@
  * results before running a new compute and when the memo is disposed.
  */
 export class AsyncMemo<T> {
+  private readonly queue: AtomicTaskQueue;
+  private readonly defaultCompute?: AsyncMemoCompute<T>;
   private cache?: Cache<T>;
   private pendingKey?: object;
   private disposed = false;
@@ -189,9 +194,22 @@
   // Stores error keyed by keyStr - thrown on next use() with same key
   private error?: {keyStr: string; error: Error};
 
+  constructor();
+  constructor(compute: AsyncMemoCompute<T>);
+  constructor(queue: AtomicTaskQueue, compute?: AsyncMemoCompute<T>);
   constructor(
-    private readonly queue: AtomicTaskQueue = new AtomicTaskQueue(),
-  ) {}
+    queueOrCompute:
+      AtomicTaskQueue | AsyncMemoCompute<T> = new AtomicTaskQueue(),
+    compute?: AsyncMemoCompute<T>,
+  ) {
+    if (typeof queueOrCompute === 'function') {
+      this.queue = new AtomicTaskQueue();
+      this.defaultCompute = queueOrCompute;
+    } else {
+      this.queue = queueOrCompute;
+      this.defaultCompute = compute;
+    }
+  }
 
   /**
    * Call every render cycle to get the current memoized result.
@@ -204,7 +222,8 @@
     if (this.disposed) {
       throw new Error('AsyncMemo.use() called after dispose()');
     }
-    const {key, compute, enabled, retainOn = []} = options;
+    const {key, enabled, retainOn = []} = options;
+    const compute = options.compute ?? this.defaultCompute;
     const keyStr = stringifyJsonWithBigints(key);
 
     // If we have a stored error for this key, throw it
@@ -226,6 +245,12 @@
     const canRun = enabled === undefined || enabled;
 
     if (isKeyDifferentFromPending && isKeyDifferentFromCache && canRun) {
+      if (compute === undefined) {
+        throw new Error(
+          'AsyncMemo requires a compute function in its constructor or use()',
+        );
+      }
+
       // Cancel any in-flight task
       if (this.currentSignal) {
         this.currentSignal.cancelled = true;

ui/src/base/async_memo_unittest.ts

diff --git a/ui/src/base/async_memo_unittest.ts b/ui/src/base/async_memo_unittest.ts
index 9308716..03b85f2 100644
--- a/ui/src/base/async_memo_unittest.ts
+++ b/ui/src/base/async_memo_unittest.ts
@@ -42,6 +42,41 @@
   expect(queryFn).toHaveBeenCalledTimes(1);
 });
 
+test('uses compute function passed to constructor', async () => {
+  const queryFn = vi.fn().mockResolvedValue(42);
+  const slot = new AsyncMemo<number>(queryFn);
+
+  slot.use({key: {id: 1}});
+  await flushPromises();
+
+  const result = slot.use({key: {id: 1}});
+  expect(result.data).toBe(42);
+  expect(result.isPending).toBe(false);
+  expect(queryFn).toHaveBeenCalledTimes(1);
+});
+
+test('use compute function overrides constructor function', async () => {
+  const defaultQueryFn = vi.fn().mockResolvedValue(1);
+  const overrideQueryFn = vi.fn().mockResolvedValue(2);
+  const slot = new AsyncMemo<number>(defaultQueryFn);
+
+  slot.use({key: {id: 1}, compute: overrideQueryFn});
+  await flushPromises();
+
+  const result = slot.use({key: {id: 1}});
+  expect(result.data).toBe(2);
+  expect(defaultQueryFn).not.toHaveBeenCalled();
+  expect(overrideQueryFn).toHaveBeenCalledTimes(1);
+});
+
+test('throws when no compute function is provided', () => {
+  const slot = new AsyncMemo<number>();
+
+  expect(() => slot.use({key: {id: 1}})).toThrow(
+    'AsyncMemo requires a compute function in its constructor or use()',
+  );
+});
+
 test('cached result is returned without re-query', async () => {
   const slot = new AsyncMemo<number>();