@stigmer/react 1.0.3 → 1.0.4

package/src/agent/__tests__/useDefaultAgent.test.tsx

@@ -1,308 +1,183 @@
-import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
-import { renderHook, act } from "@testing-library/react";
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor, act } from "@testing-library/react";
 import type { ReactNode } from "react";
-import type { Agent } from "@stigmer/protos/ai/stigmer/agentic/agent/v1/api_pb";
-import type { Stigmer } from "@stigmer/sdk";
 import { StigmerContext } from "../../context";
+import { FetchCacheContext } from "../../internal/FetchCacheProvider";
 import { useDefaultAgent } from "../useDefaultAgent";
 
-const STALE_THRESHOLD_MS = 30_000;
-
-function fakeAgent(instanceId: string): Agent {
-  return {
-    status: { defaultInstanceId: instanceId },
-  } as unknown as Agent;
-}
-
-function buildMockClient(overrides: {
-  getDefault?: ReturnType<typeof vi.fn>;
+function createMockStigmer(overrides: {
+  getDefault?: () => Promise<unknown>;
 } = {}) {
   return {
     agent: {
-      getDefault: overrides.getDefault ?? vi.fn(),
+      getDefault: overrides.getDefault ?? vi.fn().mockResolvedValue(null),
     },
-  } as unknown as Stigmer;
+  } as never;
 }
 
-function makeWrapper(client: Stigmer) {
-  return ({ children }: { children: ReactNode }) => (
-    <StigmerContext.Provider value={client}>
-      {children}
-    </StigmerContext.Provider>
-  );
-}
-
-function fireVisibilityChange(state: DocumentVisibilityState) {
-  Object.defineProperty(document, "visibilityState", {
-    value: state,
-    writable: true,
-    configurable: true,
-  });
-  document.dispatchEvent(new Event("visibilitychange"));
+function wrapper(client: unknown) {
+  return function Wrapper({ children }: { children: ReactNode }) {
+    return (
+      <FetchCacheContext.Provider value={null}>
+        <StigmerContext.Provider value={client as never}>
+          {children}
+        </StigmerContext.Provider>
+      </FetchCacheContext.Provider>
+    );
+  };
 }
 
 describe("useDefaultAgent", () => {
-  let getDefaultMock: ReturnType<typeof vi.fn>;
-  let client: Stigmer;
-
   beforeEach(() => {
-    getDefaultMock = vi.fn();
-    client = buildMockClient({ getDefault: getDefaultMock });
-  });
-
-  afterEach(() => {
     vi.restoreAllMocks();
   });
 
-  // -----------------------------------------------------------------------
-  // Basic fetch behavior (no fake timers needed)
-  // -----------------------------------------------------------------------
+  it("returns loading state initially and resolves with agent data", async () => {
+    const agent = { metadata: { id: "agt-1", name: "Default Agent" }, status: { defaultInstanceId: "ain-1" } };
+    const getDefault = vi.fn().mockResolvedValue(agent);
+    const client = createMockStigmer({ getDefault });
 
-  it("fetches the default agent on mount", async () => {
-    const agent = fakeAgent("inst_1");
-    getDefaultMock.mockResolvedValueOnce(agent);
-
-    const { result } = renderHook(() => useDefaultAgent("acme"), {
-      wrapper: makeWrapper(client),
+    const { result } = renderHook(() => useDefaultAgent("test-org"), {
+      wrapper: wrapper(client),
     });
 
     expect(result.current.isLoading).toBe(true);
+    expect(result.current.agent).toBeNull();
 
-    await act(async () => {});
+    await waitFor(() => expect(result.current.isLoading).toBe(false));
 
-    expect(result.current.isLoading).toBe(false);
     expect(result.current.agent).toBe(agent);
     expect(result.current.error).toBeNull();
-    expect(getDefaultMock).toHaveBeenCalledOnce();
+    expect(getDefault).toHaveBeenCalledTimes(1);
   });
 
   it("skips fetching when org is null", () => {
+    const getDefault = vi.fn();
+    const client = createMockStigmer({ getDefault });
+
     const { result } = renderHook(() => useDefaultAgent(null), {
-      wrapper: makeWrapper(client),
+      wrapper: wrapper(client),
     });
 
-    expect(result.current.agent).toBeNull();
     expect(result.current.isLoading).toBe(false);
-    expect(getDefaultMock).not.toHaveBeenCalled();
+    expect(result.current.agent).toBeNull();
+    expect(result.current.error).toBeNull();
+    expect(getDefault).not.toHaveBeenCalled();
   });
 
-  // -----------------------------------------------------------------------
-  // Retry on transient failure
-  // -----------------------------------------------------------------------
-
-  it("retries once on transient failure then succeeds", async () => {
-    vi.useFakeTimers();
-    try {
-      const agent = fakeAgent("inst_retry");
-      getDefaultMock
-        .mockRejectedValueOnce(new Error("network timeout"))
-        .mockResolvedValueOnce(agent);
-
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
-
-      expect(result.current.isLoading).toBe(true);
-
-      // Flush the rejected first attempt.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
-
-      // Advance past the 1s retry delay.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(1_100);
-      });
+  it("exposes error when fetch fails after retry", async () => {
+    const apiError = new Error("Service unavailable");
+    const getDefault = vi.fn().mockRejectedValue(apiError);
+    const client = createMockStigmer({ getDefault });
 
-      expect(result.current.isLoading).toBe(false);
-      expect(result.current.agent).toBe(agent);
-      expect(result.current.error).toBeNull();
-      expect(getDefaultMock).toHaveBeenCalledTimes(2);
-    } finally {
-      vi.useRealTimers();
-    }
-  });
-
-  it("surfaces the error after exhausting retries", async () => {
-    vi.useFakeTimers();
-    try {
-      getDefaultMock
-        .mockRejectedValueOnce(new Error("fail 1"))
-        .mockRejectedValueOnce(new Error("fail 2"));
+    const { result } = renderHook(() => useDefaultAgent("test-org"), {
+      wrapper: wrapper(client),
+    });
 
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
+    // useDefaultAgent retries once (MAX_RETRIES=1) with a 1s delay.
+    // useFetch wraps the retry fn, so the final error surfaces after
+    // both attempts fail. Wait for the error to propagate.
+    await waitFor(() => expect(result.current.error).toBeTruthy(), { timeout: 10_000 });
 
-      // Flush first attempt + advance past retry delay + flush retry.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(1_100);
-      });
+    expect(result.current.error!.message).toBe("Service unavailable");
+    expect(result.current.agent).toBeNull();
+    expect(result.current.isLoading).toBe(false);
+    // 1 initial + 1 retry = 2 calls
+    expect(getDefault).toHaveBeenCalledTimes(2);
+  }, 15_000);
+
+  it("refetch triggers a new fetch", async () => {
+    let callCount = 0;
+    const agent = { metadata: { id: "agt-1" }, status: { defaultInstanceId: "ain-1" } };
+    const getDefault = vi.fn().mockImplementation(async () => {
+      callCount++;
+      return agent;
+    });
+    const client = createMockStigmer({ getDefault });
 
-      expect(result.current.isLoading).toBe(false);
-      expect(result.current.agent).toBeNull();
-      expect(result.current.error).toBeInstanceOf(Error);
-      expect(result.current.error!.message).toBe("fail 2");
-      expect(getDefaultMock).toHaveBeenCalledTimes(2);
-    } finally {
-      vi.useRealTimers();
-    }
-  });
+    const { result } = renderHook(() => useDefaultAgent("test-org"), {
+      wrapper: wrapper(client),
+    });
 
-  // -----------------------------------------------------------------------
-  // Visibility-aware refetch
-  // -----------------------------------------------------------------------
-
-  it("refetches when document becomes visible after stale threshold", async () => {
-    vi.useFakeTimers();
-    try {
-      const agent1 = fakeAgent("inst_old");
-      const agent2 = fakeAgent("inst_new");
-      getDefaultMock
-        .mockResolvedValueOnce(agent1)
-        .mockResolvedValueOnce(agent2);
-
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
+    await waitFor(() => expect(result.current.isLoading).toBe(false));
+    expect(callCount).toBe(1);
 
-      // Flush initial fetch.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
-      expect(result.current.agent).toBe(agent1);
-      expect(getDefaultMock).toHaveBeenCalledTimes(1);
+    act(() => result.current.refetch());
 
-      // Simulate idle period longer than the stale threshold.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(STALE_THRESHOLD_MS + 1_000);
-      });
+    await waitFor(() => expect(callCount).toBe(2), { timeout: 10_000 });
+  }, 15_000);
 
-      // Simulate app coming back to foreground.
-      act(() => {
-        fireVisibilityChange("visible");
-      });
+  describe("waitForResolution", () => {
+    it("resolves immediately when agent is already loaded", async () => {
+      const agent = { metadata: { id: "agt-1" }, status: { defaultInstanceId: "ain-1" } };
+      const getDefault = vi.fn().mockResolvedValue(agent);
+      const client = createMockStigmer({ getDefault });
 
-      // Flush the refetch.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
+      const { result } = renderHook(() => useDefaultAgent("test-org"), {
+        wrapper: wrapper(client),
       });
 
-      expect(result.current.agent).toBe(agent2);
-      expect(getDefaultMock).toHaveBeenCalledTimes(2);
-    } finally {
-      vi.useRealTimers();
-    }
-  });
+      await waitFor(() => expect(result.current.isLoading).toBe(false));
 
-  it("does NOT refetch when visible within the stale window", async () => {
-    vi.useFakeTimers();
-    try {
-      const agent = fakeAgent("inst_fresh");
-      getDefaultMock.mockResolvedValueOnce(agent);
+      const resolved = await result.current.waitForResolution();
+      expect(resolved).toBe(agent);
+    });
 
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
+    it("awaits and resolves when fetch completes", async () => {
+      let resolveExternal!: (value: unknown) => void;
+      const fetchPromise = new Promise((res) => { resolveExternal = res; });
+      const getDefault = vi.fn().mockReturnValue(fetchPromise);
+      const client = createMockStigmer({ getDefault });
 
-      // Flush initial fetch.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
+      const { result } = renderHook(() => useDefaultAgent("test-org"), {
+        wrapper: wrapper(client),
       });
-      expect(result.current.agent).toBe(agent);
-      expect(getDefaultMock).toHaveBeenCalledTimes(1);
 
-      // Only a short time passes — well within stale threshold.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(5_000);
-      });
+      expect(result.current.isLoading).toBe(true);
 
-      act(() => {
-        fireVisibilityChange("visible");
-      });
+      const waitPromise = result.current.waitForResolution();
 
-      expect(getDefaultMock).toHaveBeenCalledTimes(1);
-    } finally {
-      vi.useRealTimers();
-    }
-  });
+      const agent = { metadata: { id: "agt-1" }, status: { defaultInstanceId: "ain-1" } };
+      resolveExternal(agent);
 
-  it("does not refetch on hidden event", async () => {
-    vi.useFakeTimers();
-    try {
-      const agent = fakeAgent("inst_1");
-      getDefaultMock.mockResolvedValueOnce(agent);
+      await waitFor(() => expect(result.current.isLoading).toBe(false));
 
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
+      const resolved = await waitPromise;
+      expect(resolved).toBe(agent);
+    });
 
-      // Flush initial fetch.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
-      expect(result.current.agent).toBe(agent);
+    it("rejects when fetch fails", async () => {
+      const apiError = new Error("Service unavailable");
+      const getDefault = vi.fn().mockRejectedValue(apiError);
+      const client = createMockStigmer({ getDefault });
 
-      // Advance well past stale threshold.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(STALE_THRESHOLD_MS + 1_000);
+      const { result } = renderHook(() => useDefaultAgent("test-org"), {
+        wrapper: wrapper(client),
       });
 
-      // Hidden (not visible) should not trigger refetch.
-      act(() => {
-        fireVisibilityChange("hidden");
-      });
+      expect(result.current.isLoading).toBe(true);
 
-      expect(getDefaultMock).toHaveBeenCalledTimes(1);
-    } finally {
-      vi.useRealTimers();
-    }
-  });
+      const waitPromise = result.current.waitForResolution();
+      // Prevent unhandled rejection warning — assertion below verifies the value
+      waitPromise.catch(() => {});
 
-  // -----------------------------------------------------------------------
-  // Manual refetch
-  // -----------------------------------------------------------------------
-
-  it("recovers via manual refetch after initial failure", async () => {
-    vi.useFakeTimers();
-    try {
-      const agent = fakeAgent("inst_recovered");
-      getDefaultMock
-        .mockRejectedValueOnce(new Error("initial fail"))
-        .mockRejectedValueOnce(new Error("retry fail"))
-        .mockResolvedValueOnce(agent);
-
-      const { result } = renderHook(() => useDefaultAgent("acme"), {
-        wrapper: makeWrapper(client),
-      });
+      await waitFor(() => expect(result.current.error).toBeTruthy(), { timeout: 10_000 });
 
-      // Exhaust initial attempt + retry.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(1_100);
-      });
+      await expect(waitPromise).rejects.toThrow("Service unavailable");
+    }, 15_000);
 
-      expect(result.current.error).not.toBeNull();
+    it("rejects immediately when error is already set", async () => {
+      const apiError = new Error("Already failed");
+      const getDefault = vi.fn().mockRejectedValue(apiError);
+      const client = createMockStigmer({ getDefault });
 
-      // Manual refetch triggers recovery.
-      act(() => {
-        result.current.refetch();
+      const { result } = renderHook(() => useDefaultAgent("test-org"), {
+        wrapper: wrapper(client),
       });
 
-      // Flush the refetch.
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
+      await waitFor(() => expect(result.current.error).toBeTruthy(), { timeout: 10_000 });
 
-      expect(result.current.agent).toBe(agent);
-      expect(result.current.error).toBeNull();
-    } finally {
-      vi.useRealTimers();
-    }
+      await expect(result.current.waitForResolution()).rejects.toThrow("Already failed");
+    }, 15_000);
   });
 });

package/src/agent/useDefaultAgent.ts

@@ -1,6 +1,6 @@
 "use client";
 
-import { useEffect, useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 import { create } from "@bufbuild/protobuf";
 import type { Agent } from "@stigmer/protos/ai/stigmer/agentic/agent/v1/api_pb";
 import { GetDefaultAgentRequestSchema } from "@stigmer/protos/ai/stigmer/agentic/agent/v1/io_pb";
@@ -29,6 +29,15 @@ export interface UseDefaultAgentReturn {
   readonly error: Error | null;
   /** Discard cached data and re-fetch the default agent from the server. */
   readonly refetch: () => void;
+  /**
+   * Returns a promise that resolves with the Agent once the in-flight
+   * fetch settles. If the agent is already loaded, resolves immediately.
+   * If the fetch has already failed, rejects immediately.
+   *
+   * Use this to await the default agent in async flows (e.g. submit)
+   * instead of failing when `isLoading` is still true.
+   */
+  readonly waitForResolution: () => Promise<Agent>;
 }
 
 /**
@@ -74,6 +83,32 @@ export function useDefaultAgent(org: string | null): UseDefaultAgentReturn {
     null,
   );
 
+  // Deferred pattern: allows callers to await the in-flight fetch.
+  const deferredRef = useRef<Deferred<Agent> | null>(null);
+
+  useEffect(() => {
+    if (agent && deferredRef.current) {
+      deferredRef.current.resolve(agent);
+      deferredRef.current = null;
+    }
+  }, [agent]);
+
+  useEffect(() => {
+    if (error && deferredRef.current) {
+      deferredRef.current.reject(error);
+      deferredRef.current = null;
+    }
+  }, [error]);
+
+  const waitForResolution = useCallback((): Promise<Agent> => {
+    if (agent) return Promise.resolve(agent);
+    if (error) return Promise.reject(error);
+    if (!deferredRef.current) {
+      deferredRef.current = createDeferred<Agent>();
+    }
+    return deferredRef.current.promise;
+  }, [agent, error]);
+
   useEffect(() => {
     if (typeof document === "undefined") return;
 
@@ -88,13 +123,29 @@ export function useDefaultAgent(org: string | null): UseDefaultAgentReturn {
     return () => document.removeEventListener("visibilitychange", onVisible);
   }, [refetch]);
 
-  return { agent, isLoading, isRefetching, error, refetch };
+  return { agent, isLoading, isRefetching, error, refetch, waitForResolution };
 }
 
 // ---------------------------------------------------------------------------
 // Internal helpers
 // ---------------------------------------------------------------------------
 
+interface Deferred<T> {
+  promise: Promise<T>;
+  resolve: (value: T) => void;
+  reject: (reason: unknown) => void;
+}
+
+function createDeferred<T>(): Deferred<T> {
+  let resolve!: (value: T) => void;
+  let reject!: (reason: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject };
+}
+
 /**
  * Retries `fn` up to `retries` times with a fixed delay between attempts.
  */