@datalayer/agent-runtimes 1.0.4 → 1.0.5

package/lib/hooks/useAIAgentsWebSocket.js

@@ -12,7 +12,7 @@
  *
  * @module hooks/useAIAgentsWebSocket
  */
-import { useEffect, useRef } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { useQueryClient } from '@tanstack/react-query';
 import { useCoreStore, useIAMStore } from '@datalayer/core/lib/state';
 import { DEFAULT_SERVICE_URLS, API_BASE_PATHS, } from '@datalayer/core/lib/api/constants';
@@ -27,6 +27,8 @@ function useBaseUrl() {
 }
 // ─── Hook ────────────────────────────────────────────────────────────
 const RECONNECT_DELAY_MS = 3_000;
+const WS_DEFAULT_PATH = `${API_BASE_PATHS.AI_AGENTS}/ws`;
+const isObject = (value) => !!value && typeof value === 'object';
 /**
  * Connect to the AI Agents generic WebSocket.
  *
@@ -44,28 +46,68 @@ const RECONNECT_DELAY_MS = 3_000;
 export function useAIAgentsWebSocket(options) {
     const token = useAuthToken();
     const baseUrl = useBaseUrl();
+    const configuredBaseUrl = options?.baseUrl ?? baseUrl;
+    const enabled = options?.enabled ?? true;
+    const wsPath = options?.path ?? WS_DEFAULT_PATH;
+    const queryParamsKey = JSON.stringify(options?.queryParams ?? {});
     const queryClient = useQueryClient();
+    const [connectionState, setConnectionState] = useState('closed');
+    const [lastClose, setLastClose] = useState(null);
+    const [reconnectAttempt, setReconnectAttempt] = useState(0);
     const wsRef = useRef(null);
     const reconnectTimer = useRef(null);
     // Keep a ref of channels so we can re-subscribe on reconnect without
     // tearing down the socket when the array reference changes.
     const channelsRef = useRef(options?.channels ?? []);
     channelsRef.current = options?.channels ?? [];
+    const autoReconnectRef = useRef(options?.autoReconnect ?? true);
+    autoReconnectRef.current = options?.autoReconnect ?? true;
+    const maxReconnectAttemptsRef = useRef(options?.maxReconnectAttempts);
+    maxReconnectAttemptsRef.current = options?.maxReconnectAttempts;
+    const reconnectDelayRef = useRef(options?.reconnectDelayMs);
+    reconnectDelayRef.current = options?.reconnectDelayMs;
+    const onOpenRef = useRef(options?.onOpen);
+    onOpenRef.current = options?.onOpen;
     const onMessageRef = useRef(options?.onMessage);
     onMessageRef.current = options?.onMessage;
+    const onCloseRef = useRef(options?.onClose);
+    onCloseRef.current = options?.onClose;
     useEffect(() => {
-        if (!token)
+        if (!enabled || !token) {
+            setConnectionState('closed');
             return;
+        }
         let disposed = false;
+        let reconnectAttempts = 0;
+        const toWsUrl = () => {
+            const httpUrl = wsPath.startsWith('http://') || wsPath.startsWith('https://')
+                ? wsPath
+                : `${configuredBaseUrl}${wsPath.startsWith('/') ? '' : '/'}${wsPath}`;
+            const url = new URL(httpUrl.replace(/^http/, 'ws'));
+            url.searchParams.set('token', token);
+            const queryParams = JSON.parse(queryParamsKey);
+            Object.entries(queryParams).forEach(([key, value]) => {
+                if (value === null || value === undefined) {
+                    return;
+                }
+                url.searchParams.set(key, String(value));
+            });
+            return url.toString();
+        };
         function connect() {
             if (disposed)
                 return;
-            // Build the WebSocket URL.  Replace http(s) with ws(s).
-            const httpUrl = `${baseUrl}${API_BASE_PATHS.AI_AGENTS}/ws`;
-            const wsUrl = httpUrl.replace(/^http/, 'ws') + `?token=${encodeURIComponent(token)}`;
+            const wsUrl = toWsUrl();
+            setConnectionState('connecting');
             const ws = new WebSocket(wsUrl);
             wsRef.current = ws;
             ws.onopen = () => {
+                reconnectAttempts = 0;
+                setReconnectAttempt(0);
+                setConnectionState('connected');
+                setLastClose(null);
+                console.debug('[ws:connect] url=%s', wsUrl);
+                onOpenRef.current?.();
                 // Subscribe to extra channels.
                 const channels = channelsRef.current;
                 if (channels.length > 0) {
@@ -73,17 +115,33 @@ export function useAIAgentsWebSocket(options) {
                 }
             };
             ws.onmessage = ev => {
-                let msg;
+                let raw;
                 try {
-                    msg = JSON.parse(ev.data);
+                    raw = JSON.parse(String(ev.data));
                 }
                 catch {
                     return;
                 }
+                const msg = isObject(raw)
+                    ? {
+                        channel: typeof raw.channel === 'string' ? raw.channel : undefined,
+                        event: typeof raw.event === 'string' ? raw.event : undefined,
+                        data: isObject(raw.data)
+                            ? raw.data
+                            : undefined,
+                        type: typeof raw.type === 'string' ? raw.type : undefined,
+                        payload: raw.payload,
+                        raw,
+                    }
+                    : { raw };
                 // Fire optional callback.
+                console.debug('[ws:recv] type=%s', msg.type ?? msg.event ?? 'unknown');
                 onMessageRef.current?.(msg);
                 // Invalidate React Query caches based on the event type.
                 const { event } = msg;
+                if (typeof event !== 'string') {
+                    return;
+                }
                 if (event.startsWith('event_')) {
                     queryClient.invalidateQueries({ queryKey: ['agent-events'] });
                 }
@@ -94,15 +152,43 @@ export function useAIAgentsWebSocket(options) {
                     queryClient.invalidateQueries({ queryKey: ['agent-notifications'] });
                 }
             };
-            ws.onclose = () => {
+            ws.onclose = event => {
                 wsRef.current = null;
-                if (!disposed) {
-                    reconnectTimer.current = setTimeout(connect, RECONNECT_DELAY_MS);
+                setConnectionState('closed');
+                const closeInfo = {
+                    code: event.code,
+                    reason: event.reason,
+                    wasClean: event.wasClean,
+                    detail: `code ${event.code}${event.reason ? `: ${event.reason}` : ''}${event.wasClean ? ' (clean)' : ' (unclean)'}`,
+                };
+                setLastClose(closeInfo);
+                console.debug('[ws:disconnect] code=%d reason=%s', event.code, event.reason || '(none)');
+                onCloseRef.current?.(closeInfo);
+                if (disposed || !autoReconnectRef.current) {
+                    return;
                 }
+                reconnectAttempts += 1;
+                setReconnectAttempt(reconnectAttempts);
+                const maxAttempts = maxReconnectAttemptsRef.current;
+                if (typeof maxAttempts === 'number' &&
+                    Number.isFinite(maxAttempts) &&
+                    reconnectAttempts > maxAttempts) {
+                    return;
+                }
+                const configuredDelay = reconnectDelayRef.current;
+                const delay = typeof configuredDelay === 'function'
+                    ? configuredDelay(reconnectAttempts)
+                    : typeof configuredDelay === 'number'
+                        ? configuredDelay
+                        : RECONNECT_DELAY_MS;
+                reconnectTimer.current = setTimeout(connect, Math.max(0, delay));
             };
             ws.onerror = () => {
                 // onclose will fire after onerror; reconnect happens there.
-                ws.close();
+                if (ws.readyState === WebSocket.CONNECTING ||
+                    ws.readyState === WebSocket.OPEN) {
+                    ws.close();
+                }
             };
         }
         connect();
@@ -110,11 +196,13 @@ export function useAIAgentsWebSocket(options) {
             disposed = true;
             if (reconnectTimer.current) {
                 clearTimeout(reconnectTimer.current);
+                reconnectTimer.current = null;
             }
             wsRef.current?.close();
             wsRef.current = null;
+            setConnectionState('closed');
         };
-    }, [token, baseUrl, queryClient]);
+    }, [token, configuredBaseUrl, wsPath, queryParamsKey, enabled, queryClient]);
     // When the channel list changes, send subscribe/unsubscribe diffs.
     const prevChannelsRef = useRef([]);
     useEffect(() => {
@@ -133,4 +221,22 @@ export function useAIAgentsWebSocket(options) {
         }
         prevChannelsRef.current = [...channelsRef.current];
     }, [options?.channels]);
+    return {
+        connectionState,
+        lastClose,
+        reconnectAttempt,
+        send: (payload) => {
+            const ws = wsRef.current;
+            if (!ws || ws.readyState !== WebSocket.OPEN) {
+                return false;
+            }
+            try {
+                ws.send(typeof payload === 'string' ? payload : JSON.stringify(payload));
+                return true;
+            }
+            catch {
+                return false;
+            }
+        },
+    };
 }

package/lib/hooks/useAcp.d.ts

@@ -1,5 +1,4 @@
-import { type StopReason, type SessionUpdate, type AgentCapabilities, type ClientCapabilities, type PermissionOption, type ToolCallUpdate } from '@agentclientprotocol/sdk';
-export type { StopReason, AgentCapabilities, ClientCapabilities, ToolCallUpdate, PermissionOption, };
+import { type StopReason, type SessionUpdate, type AgentCapabilities, type PermissionOption, type ToolCallUpdate } from '@agentclientprotocol/sdk';
 export type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'error';
 export type SessionUpdateType = SessionUpdate extends {
     sessionUpdate: infer T;

package/lib/hooks/useAgUi.d.ts

@@ -5,7 +5,7 @@ export interface Message {
 }
 export interface UseAGUIOptions {
     /**
-     * Base URL for the AG-UI endpoint (e.g., 'http://localhost:8000')
+     * Base URL for the AG-UI endpoint (e.g., `http://localhost:8000`).
      */
     baseUrl?: string;
     /**

package/lib/hooks/{useAgents.d.ts → useAgentRuntimes.d.ts}

@@ -98,14 +98,14 @@ export declare const agentQueryKeys: {
  * });
  * ```
  */
-export declare function useAgents(options?: UseAgentOptions): UseAgentReturn;
+export declare function useAgentRuntimes(options?: UseAgentOptions): UseAgentReturn;
 /**
  * Hook to fetch user's agent runtimes (running agent instances).
  *
  * The backend returns active runtimes from the operator **plus** paused
  * runtimes synthesised from Solr checkpoint records (with ``status="paused"``).
  */
-export declare function useAgentRuntimes(): import("@tanstack/react-query").UseQueryResult<AgentRuntimeData[], Error>;
+export declare function useAgentRuntimesQuery(): import("@tanstack/react-query").UseQueryResult<AgentRuntimeData[], Error>;
 /**
  * Hook to fetch a single agent runtime by pod name.
  */
@@ -150,3 +150,40 @@ export interface UseAgentsRuntimesReturn {
  * Consolidated runtime list and mutations.
  */
 export declare function useAgentsRuntimes(): UseAgentsRuntimesReturn;
+export interface UseAgentRuntimeWebSocketOptions {
+    /** Enable/disable the connection. Defaults to `true`. */
+    enabled?: boolean;
+    /**
+     * Base URL of the agent-runtime server
+     * (e.g. `http://localhost:8765`). The WS path is appended automatically.
+     */
+    baseUrl: string;
+    /** Auth token passed as `?token=` query parameter. */
+    authToken?: string;
+    /** Optional `agent_id` query parameter to scope the stream. */
+    agentId?: string;
+    /** Auto-reconnect on unexpected disconnects. Defaults to `true`. */
+    autoReconnect?: boolean;
+    /** Delay between reconnection attempts (ms). Defaults to 3 000. */
+    reconnectDelayMs?: number | ((attempt: number) => number);
+    /** Maximum reconnect attempts. Unbounded by default. */
+    maxReconnectAttempts?: number;
+    /** Additional callback fired for every incoming WS message. */
+    onMessage?: (msg: {
+        type?: string;
+        payload?: unknown;
+        raw: unknown;
+    }) => void;
+}
+/**
+ * Connect to the agent-runtime monitoring WebSocket.
+ *
+ * The hook writes all incoming data into the `useAgentRuntimeStore` Zustand
+ * store. Components that need approvals, MCP status, context snapshots, or
+ * full-context data simply read from the store.
+ *
+ * Mount this hook **once** near the top of your component tree (e.g. in
+ * the example root or in `ChatBase`). All other components read from the
+ * store — no extra WebSocket connections needed.
+ */
+export declare function useAgentRuntimeWebSocket(options: UseAgentRuntimeWebSocketOptions): void;

package/lib/hooks/{useAgents.js → useAgentRuntimes.js}

@@ -3,13 +3,13 @@
  * Distributed under the terms of the Modified BSD License.
  */
 /**
- * Unified hook for managing agents.
+ * Unified hook for managing agent runtimes.
  *
  * Combines agent lifecycle management (ephemeral/durable),
  * runtime catalog (React Query CRUD), lifecycle/catalog stores,
- * and the AI Agents REST API.
+ * the AI Agents REST API, and the agent-runtime WebSocket stream.
  *
- * @module hooks/useAgents
+ * @module hooks/useAgentRuntimes
  */
 import { useState, useCallback, useMemo, useEffect, useRef } from 'react';
 import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
@@ -17,7 +17,8 @@ import { create } from 'zustand';
 import { persist, createJSONStorage } from 'zustand/middleware';
 import { useCoreStore, useDatalayer } from '@datalayer/core';
 import { useIAMStore } from '@datalayer/core/lib/state';
-import { useAgentStore, useAgentRuntime, useAgentStatus, useAgentError, useIsLaunching, } from '../stores/agentsStore';
+import { agentRuntimeStore, useAgentRuntimeStore, useAgentRuntimeConnection, useAgentRuntimeStatus, useAgentRuntimeError, useAgentRuntimeIsLaunching, } from '../stores/agentRuntimeStore';
+import { parseAgentStreamMessage, } from '../types/stream';
 import { DEFAULT_AGENT_CONFIG } from '../types/config';
 // ═══════════════════════════════════════════════════════════════════════════
 // Constants
@@ -103,18 +104,18 @@ function toAgentRuntimeData(raw) {
  * });
  * ```
  */
-export function useAgents(options = {}) {
+export function useAgentRuntimes(options = {}) {
     const { agentSpecId, agentConfig, autoCreateAgent = true, autoStart = false, agentSpec, } = options;
     // Base store state
-    const runtime = useAgentRuntime();
-    const baseStatus = useAgentStatus();
-    const storeError = useAgentError();
-    const isLaunching = useIsLaunching();
+    const runtime = useAgentRuntimeConnection();
+    const baseStatus = useAgentRuntimeStatus();
+    const storeError = useAgentRuntimeError();
+    const isLaunching = useAgentRuntimeIsLaunching();
     // Store actions
-    const storeLaunchAgent = useAgentStore(state => state.launchAgent);
-    const storeConnectAgent = useAgentStore(state => state.connectAgent);
-    const storeCreateAgent = useAgentStore(state => state.createAgent);
-    const storeDisconnect = useAgentStore(state => state.disconnect);
+    const storeLaunchAgent = useAgentRuntimeStore(state => state.launchAgent);
+    const storeConnectAgent = useAgentRuntimeStore(state => state.connectAgent);
+    const storeCreateAgent = useAgentRuntimeStore(state => state.createAgent);
+    const storeDisconnect = useAgentRuntimeStore(state => state.disconnect);
     // Lifecycle local state
     const [lifecycleStatus, setLifecycleStatus] = useState('idle');
     const [lifecycleError, setLifecycleError] = useState(null);
@@ -408,7 +409,7 @@ export function useAgents(options = {}) {
  * The backend returns active runtimes from the operator **plus** paused
  * runtimes synthesised from Solr checkpoint records (with ``status="paused"``).
  */
-export function useAgentRuntimes() {
+export function useAgentRuntimesQuery() {
     const { configuration } = useCoreStore();
     const { requestDatalayer } = useDatalayer({ notifyOnError: false });
     const { user } = useIAMStore();
@@ -607,7 +608,7 @@ export const useAgentLifecycleStore = create()(persist((set, get) => ({
  * Consolidated runtime list and mutations.
  */
 export function useAgentsRuntimes() {
-    const runtimesQuery = useAgentRuntimes();
+    const runtimesQuery = useAgentRuntimesQuery();
     const createRuntimeMutation = useCreateAgentRuntime();
     const deleteRuntimeMutation = useDeleteAgentRuntime();
     const refreshRuntimes = useRefreshAgentRuntimes();
@@ -631,3 +632,112 @@ export function useAgentsRuntimes() {
         deleteRuntimeMutation,
     ]);
 }
+const DEFAULT_WS_PATH = '/api/v1/tool-approvals/ws';
+const DEFAULT_RECONNECT_DELAY_MS = 3_000;
+/**
+ * Connect to the agent-runtime monitoring WebSocket.
+ *
+ * The hook writes all incoming data into the `useAgentRuntimeStore` Zustand
+ * store. Components that need approvals, MCP status, context snapshots, or
+ * full-context data simply read from the store.
+ *
+ * Mount this hook **once** near the top of your component tree (e.g. in
+ * the example root or in `ChatBase`). All other components read from the
+ * store — no extra WebSocket connections needed.
+ */
+export function useAgentRuntimeWebSocket(options) {
+    const { enabled = true, baseUrl, authToken, agentId, autoReconnect = true, reconnectDelayMs = DEFAULT_RECONNECT_DELAY_MS, maxReconnectAttempts, } = options;
+    const onMessageRef = useRef(options.onMessage);
+    onMessageRef.current = options.onMessage;
+    useEffect(() => {
+        if (!enabled || !baseUrl) {
+            agentRuntimeStore.getState().setWsState('closed');
+            return;
+        }
+        let disposed = false;
+        let reconnectAttempts = 0;
+        let reconnectTimer = null;
+        function buildWsUrl() {
+            const httpUrl = `${baseUrl}${DEFAULT_WS_PATH}`;
+            const url = new URL(httpUrl.replace(/^http/, 'ws'));
+            if (authToken) {
+                url.searchParams.set('token', authToken);
+            }
+            if (agentId) {
+                url.searchParams.set('agent_id', agentId);
+            }
+            return url.toString();
+        }
+        function connect() {
+            if (disposed)
+                return;
+            const wsUrl = buildWsUrl();
+            agentRuntimeStore.getState().setWsState('connecting');
+            const ws = new WebSocket(wsUrl);
+            agentRuntimeStore.getState().setWs(ws, agentId);
+            ws.onopen = () => {
+                reconnectAttempts = 0;
+                agentRuntimeStore.getState().setWsState('connected');
+            };
+            ws.onmessage = (ev) => {
+                let raw;
+                try {
+                    raw = JSON.parse(String(ev.data));
+                }
+                catch {
+                    return;
+                }
+                const parsed = parseAgentStreamMessage(raw);
+                onMessageRef.current?.({
+                    type: parsed?.type,
+                    payload: parsed?.payload,
+                    raw,
+                });
+                if (!parsed)
+                    return;
+                const state = agentRuntimeStore.getState();
+                if (parsed.type === 'agent.snapshot') {
+                    state.applySnapshot(parsed.payload);
+                    return;
+                }
+            };
+            ws.onclose = () => {
+                agentRuntimeStore.getState().setWs(null, agentId);
+                agentRuntimeStore.getState().setWsState('closed');
+                if (disposed || !autoReconnect)
+                    return;
+                reconnectAttempts += 1;
+                if (typeof maxReconnectAttempts === 'number' &&
+                    reconnectAttempts > maxReconnectAttempts) {
+                    return;
+                }
+                const delay = typeof reconnectDelayMs === 'function'
+                    ? reconnectDelayMs(reconnectAttempts)
+                    : reconnectDelayMs;
+                reconnectTimer = setTimeout(connect, Math.max(0, delay));
+            };
+            ws.onerror = () => {
+                if (ws.readyState === WebSocket.CONNECTING ||
+                    ws.readyState === WebSocket.OPEN) {
+                    ws.close();
+                }
+            };
+        }
+        connect();
+        return () => {
+            disposed = true;
+            if (reconnectTimer)
+                clearTimeout(reconnectTimer);
+            agentRuntimeStore.getState().setWs(null, agentId);
+            agentRuntimeStore.getState().resetWs();
+        };
+    }, [
+        enabled,
+        baseUrl,
+        authToken,
+        agentId,
+        autoReconnect,
+        reconnectDelayMs,
+        maxReconnectAttempts,
+    ]);
+}

package/lib/hooks/useAgentsCatalog.js

@@ -3,7 +3,7 @@
  * Distributed under the terms of the Modified BSD License.
  */
 /**
- * Agent catalog hooks.
+ * Agents catalog hooks.
  *
  * Includes the Zustand catalog store and the agent registry hook.
  *

package/lib/hooks/useAgentsService.d.ts

@@ -4,7 +4,7 @@ export type RequestOptions = {
 };
 export type RoomType = 'notebook_persist' | 'notebook_memory' | 'doc_memory';
 /**
- * @deprecated Use useAgents instead
+ * @deprecated Use useAgentRuntimes instead
  */
 export declare const useAgentsService: (baseUrlOverride?: string) => {
     createAgent: (documentId: string, documentType: RoomType, ingress?: string, token?: string, kernelId?: string, { signal, baseUrl }?: RequestOptions) => Promise<any>;
@@ -17,6 +17,6 @@ export declare const useAgentsService: (baseUrlOverride?: string) => {
  * Get the notebook AI agent if any.
  *
  * This performs a periodic liveness check and keeps the local store in sync.
- * @deprecated Use useAgents instead
+ * @deprecated Use useAgentRuntimes instead
  */
 export declare function useNotebookAgents(notebookId: string): import("..").AgentRegistryEntry | undefined;

package/lib/hooks/useAgentsService.js

@@ -11,12 +11,12 @@
 import { useEffect } from 'react';
 import { useCoreStore, useDatalayer } from '@datalayer/core';
 import { URLExt } from '@jupyterlab/coreutils';
-import { useAgentStore } from '../stores/agentsStore';
+import { useAgentRuntimeStore } from '../stores/agentRuntimeStore';
 // ═══════════════════════════════════════════════════════════════════════════
 // Agents Service REST API hook.
 // ═══════════════════════════════════════════════════════════════════════════
 /**
- * @deprecated Use useAgents instead
+ * @deprecated Use useAgentRuntimes instead
  */
 export const useAgentsService = (baseUrlOverride = 'api/ai-agents/v1') => {
     const { configuration } = useCoreStore();
@@ -86,14 +86,14 @@ export const useAgentsService = (baseUrlOverride = 'api/ai-agents/v1') => {
  * Get the notebook AI agent if any.
  *
  * This performs a periodic liveness check and keeps the local store in sync.
- * @deprecated Use useAgents instead
+ * @deprecated Use useAgentRuntimes instead
  */
 export function useNotebookAgents(notebookId) {
     const { getAgent } = useAgentsService();
-    const agents = useAgentStore(state => state.agents);
-    const upsertAgent = useAgentStore(state => state.upsertAgent);
-    const deleteAgent = useAgentStore(state => state.deleteAgent);
-    const getAgentById = useAgentStore(state => state.getAgentById);
+    const agents = useAgentRuntimeStore(state => state.agents);
+    const upsertAgent = useAgentRuntimeStore(state => state.upsertAgent);
+    const deleteAgent = useAgentRuntimeStore(state => state.deleteAgent);
+    const getAgentById = useAgentRuntimeStore(state => state.getAgentById);
     useEffect(() => {
         let abortController;
         const refreshAIAgent = async () => {

package/lib/hooks/useCheckpoints.js

@@ -14,7 +14,7 @@ import { useMemo, useState, useCallback } from 'react';
 import { useQuery, useQueryClient, useMutation } from '@tanstack/react-query';
 import { useIAMStore } from '@datalayer/core/lib/state';
 import { useCoreStore, useDatalayer } from '@datalayer/core';
-import { agentQueryKeys, AGENT_QUERY_OPTIONS } from './useAgents';
+import { agentQueryKeys, AGENT_QUERY_OPTIONS } from './useAgentRuntimes';
 // ═══════════════════════════════════════════════════════════════════════════
 // Query hooks
 // ═══════════════════════════════════════════════════════════════════════════

package/lib/hooks/useConfig.d.ts

@@ -3,9 +3,12 @@ import type { RemoteConfig } from '../types/config';
  * Hook to safely use query when QueryClient is available.
  * Returns a mock result if no QueryClientProvider is present.
  */
-export declare function useConfig(enabled: boolean, configEndpoint?: string, authToken?: string): import("@tanstack/query-core").QueryObserverRefetchErrorResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverSuccessResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverLoadingErrorResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverPendingResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverPlaceholderResult<RemoteConfig, Error> | {
+export declare function useConfig(enabled: boolean, configEndpoint?: string, authToken?: string, agentId?: string): import("@tanstack/query-core").QueryObserverRefetchErrorResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverSuccessResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverLoadingErrorResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverPendingResult<RemoteConfig, Error> | import("@tanstack/query-core").QueryObserverPlaceholderResult<RemoteConfig, Error> | {
     data: undefined;
     isLoading: boolean;
     isError: boolean;
     error: null;
+    refetch: () => Promise<{
+        data: undefined;
+    }>;
 };

package/lib/hooks/useConfig.js

@@ -9,7 +9,7 @@ import { requestAPI } from '../api/handler';
  * Hook to safely use query when QueryClient is available.
  * Returns a mock result if no QueryClientProvider is present.
  */
-export function useConfig(enabled, configEndpoint, authToken) {
+export function useConfig(enabled, configEndpoint, authToken, agentId) {
     const queryClient = useContext(QueryClientContext);
     if (!queryClient) {
         return {
@@ -17,6 +17,7 @@ export function useConfig(enabled, configEndpoint, authToken) {
             isLoading: false,
             isError: false,
             error: null,
+            refetch: () => Promise.resolve({ data: undefined }),
         };
     }
     // eslint-disable-next-line react-hooks/rules-of-hooks
@@ -30,7 +31,13 @@ export function useConfig(enabled, configEndpoint, authToken) {
                 if (authToken) {
                     headers['Authorization'] = `Bearer ${authToken}`;
                 }
-                const response = await fetch(configEndpoint, { headers });
+                let endpoint = configEndpoint;
+                if (agentId) {
+                    const url = new URL(configEndpoint, window.location.origin);
+                    url.searchParams.set('agent_id', agentId);
+                    endpoint = url.toString();
+                }
+                const response = await fetch(endpoint, { headers });
                 if (!response.ok) {
                     throw new Error(`Config fetch failed: ${response.statusText}`);
                 }
@@ -39,7 +46,7 @@ export function useConfig(enabled, configEndpoint, authToken) {
             // Otherwise use Jupyter requestAPI.
             return requestAPI('configure');
         },
-        queryKey: ['models', configEndpoint || 'jupyter'],
+        queryKey: ['models', configEndpoint || 'jupyter', agentId || 'global'],
         enabled,
         retry: 1,
     });

package/lib/hooks/useContextSnapshot.d.ts

@@ -1,10 +1,15 @@
 import type { ContextSnapshotData } from '../types/context';
 /**
- * Hook to poll agent context-snapshot from the backend.
- * Returns cumulative token usage (input/output breakdown) tracked by the agent server.
+ * Hook that previously polled agent context-snapshot from the backend.
+ *
+ * The REST endpoint has been removed — context snapshot data is now
+ * delivered via the WebSocket stream (`agent.snapshot` messages).
+ * This hook is kept as a no-op so existing call-sites compile without
+ * changes; the token-usage bar simply stays hidden until a WS-based
+ * replacement is wired in.
  */
-export declare function useContextSnapshot(enabled: boolean, configEndpoint?: string, agentId?: string, authToken?: string): import("@tanstack/query-core").QueryObserverRefetchErrorResult<ContextSnapshotData, Error> | import("@tanstack/query-core").QueryObserverSuccessResult<ContextSnapshotData, Error> | import("@tanstack/query-core").QueryObserverLoadingErrorResult<ContextSnapshotData, Error> | import("@tanstack/query-core").QueryObserverPendingResult<ContextSnapshotData, Error> | import("@tanstack/query-core").QueryObserverPlaceholderResult<ContextSnapshotData, Error> | {
-    data: undefined;
+export declare function useContextSnapshot(_enabled: boolean, _configEndpoint?: string, _agentId?: string, _authToken?: string): {
+    data: ContextSnapshotData | undefined;
     isLoading: boolean;
     isError: boolean;
     error: null;