@datalayer/agent-runtimes 1.0.2 → 1.0.4

package/lib/hooks/useCheckpoints.js

@@ -0,0 +1,466 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+/**
+ * Agent checkpoint hooks.
+ *
+ * Includes query hooks for listing checkpoints and mutation hooks
+ * for pause, resume, and checkpoint lifecycle operations.
+ *
+ * @module hooks/useCheckpoints
+ */
+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';
+// ═══════════════════════════════════════════════════════════════════════════
+// Query hooks
+// ═══════════════════════════════════════════════════════════════════════════
+/**
+ * Fetch all runtime checkpoints for the current user.
+ *
+ * Calls ``GET /api/runtimes/v1/runtime-checkpoints`` and returns
+ * the list of checkpoint records in visible states.
+ */
+export function useCheckpointsQuery() {
+    const { configuration } = useCoreStore();
+    const { requestDatalayer } = useDatalayer({ notifyOnError: false });
+    const { user } = useIAMStore();
+    return useQuery({
+        queryKey: agentQueryKeys.checkpoints.lists(),
+        queryFn: async () => {
+            const resp = await requestDatalayer({
+                url: `${configuration.runtimesRunUrl}/api/runtimes/v1/runtime-checkpoints`,
+                method: 'GET',
+            });
+            if (resp.success && resp.checkpoints) {
+                return resp.checkpoints;
+            }
+            return [];
+        },
+        ...AGENT_QUERY_OPTIONS,
+        refetchInterval: 15000,
+        enabled: !!user,
+    });
+}
+/**
+ * Hook to refresh the checkpoints list.
+ */
+export function useRefreshCheckpoints() {
+    const queryClient = useQueryClient();
+    return () => {
+        queryClient.invalidateQueries({
+            queryKey: agentQueryKeys.checkpoints.all(),
+        });
+    };
+}
+/**
+ * Hook to delete a paused agent runtime.
+ *
+ * Paused agents have no K8s pod — their state lives entirely in Solr
+ * checkpoint records. This calls the dedicated
+ * ``DELETE /runtimes/{podName}/paused`` endpoint which removes those
+ * Solr records.
+ */
+export function useDeletePausedAgentRuntime() {
+    const { configuration } = useCoreStore();
+    const { requestDatalayer } = useDatalayer({ notifyOnError: false });
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (podName) => {
+            return requestDatalayer({
+                url: `${configuration.runtimesRunUrl}/api/runtimes/v1/runtimes/${podName}/paused`,
+                method: 'DELETE',
+            });
+        },
+        onSuccess: (_data, podName) => {
+            queryClient.cancelQueries({
+                queryKey: agentQueryKeys.agentRuntimes.detail(podName),
+            });
+            queryClient.removeQueries({
+                queryKey: agentQueryKeys.agentRuntimes.detail(podName),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.lists(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+/**
+ * Hook to resume a paused agent runtime via checkpoint restore.
+ */
+export function useResumePausedAgentRuntime() {
+    const { configuration } = useCoreStore();
+    const { requestDatalayer } = useDatalayer({ notifyOnError: false });
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (podName) => {
+            return requestDatalayer({
+                url: `${configuration.runtimesRunUrl}/api/runtimes/v1/runtimes/${podName}/resume`,
+                method: 'POST',
+            });
+        },
+        onSuccess: () => {
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.all(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+// ═══════════════════════════════════════════════════════════════════════════
+// Mutation hooks (pause / resume / checkpoint)
+// ═══════════════════════════════════════════════════════════════════════════
+/**
+ * Hook to pause an agent runtime (creates a checkpoint).
+ *
+ * Calls ``POST .../pause`` then polls ``waitForCheckpointStatus``
+ * until the checkpoint reaches ``paused`` or ``failed``.
+ * On success, both runtimes and checkpoints queries are invalidated.
+ */
+export function usePauseAgent() {
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (params) => {
+            const { iamStore, coreStore } = await import('@datalayer/core/lib/state');
+            const token = iamStore.getState().token || '';
+            const runtimesRunUrl = coreStore.getState().configuration?.runtimesRunUrl || '';
+            const { pauseRuntime } = await import('@datalayer/core/lib/api/runtimes/runtimes');
+            const mode = params.mode || 'light';
+            const resp = await pauseRuntime(token, params.podName, runtimesRunUrl, {
+                agent_spec_id: params.agentSpecId,
+                checkpoint_mode: mode,
+                ...(params.messages && mode === 'light'
+                    ? { messages: params.messages }
+                    : {}),
+                ...(params.agentSpec ? { agentspec: params.agentSpec } : {}),
+            });
+            if (resp.checkpoint_id) {
+                const { waitForCheckpointStatus } = await import('@datalayer/core/lib/api/runtimes/checkpoints');
+                const ckpt = await waitForCheckpointStatus(token, params.podName, resp.checkpoint_id, ['paused', 'failed'], runtimesRunUrl);
+                if (ckpt.status === 'failed') {
+                    throw new Error(`Checkpoint ${resp.checkpoint_id} failed during ${mode.toUpperCase()} pause`);
+                }
+            }
+            return resp;
+        },
+        onSuccess: () => {
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.all(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+/**
+ * Hook to resume a paused agent runtime from a checkpoint.
+ *
+ * Calls ``POST .../resume`` then invalidates both runtimes and
+ * checkpoints queries so the UI reflects the new state.
+ */
+export function useResumeAgent() {
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (params) => {
+            const { iamStore, coreStore } = await import('@datalayer/core/lib/state');
+            const token = iamStore.getState().token || '';
+            const runtimesRunUrl = coreStore.getState().configuration?.runtimesRunUrl || '';
+            const { resumeRuntime } = await import('@datalayer/core/lib/api/runtimes/runtimes');
+            return resumeRuntime(token, params.podName, runtimesRunUrl, {
+                agent_spec_id: params.agentSpecId,
+                ...(params.mode ? { checkpoint_mode: params.mode } : {}),
+                ...(params.checkpointId ? { checkpoint_id: params.checkpointId } : {}),
+            });
+        },
+        onSuccess: () => {
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.all(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+/**
+ * Hook to create a named checkpoint (pause → wait → stay paused).
+ *
+ * Unlike ``usePauseAgent`` which is intended for a full pause lifecycle,
+ * this hook creates a named checkpoint record for later restore.
+ */
+export function useCheckpointAgent() {
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (params) => {
+            const { iamStore, coreStore } = await import('@datalayer/core/lib/state');
+            const token = iamStore.getState().token || '';
+            const runtimesRunUrl = coreStore.getState().configuration?.runtimesRunUrl || '';
+            const { pauseRuntime } = await import('@datalayer/core/lib/api/runtimes/runtimes');
+            const { waitForCheckpointStatus } = await import('@datalayer/core/lib/api/runtimes/checkpoints');
+            const mode = params.mode || 'criu';
+            const pauseResp = await pauseRuntime(token, params.podName, runtimesRunUrl, {
+                name: params.name || `checkpoint-${Date.now()}`,
+                description: `${mode.toUpperCase()} checkpoint for ${params.agentSpecId}`,
+                checkpoint_mode: mode,
+                ...(params.messages && mode === 'light'
+                    ? { messages: params.messages }
+                    : {}),
+                agent_spec_id: params.agentSpecId,
+                agentspec: params.agentSpec || {},
+            });
+            const checkpointId = pauseResp.checkpoint_id;
+            if (!checkpointId) {
+                throw new Error('Pause did not return a checkpoint_id');
+            }
+            const ckpt = await waitForCheckpointStatus(token, params.podName, checkpointId, ['paused', 'failed'], runtimesRunUrl);
+            if (ckpt.status === 'failed') {
+                throw new Error(`Checkpoint ${checkpointId} failed during ${mode.toUpperCase()} pause`);
+            }
+            return { ...pauseResp, checkpoint: ckpt };
+        },
+        onSuccess: () => {
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.all(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+// ═══════════════════════════════════════════════════════════════════════════
+// Terminate mutation hook
+// ═══════════════════════════════════════════════════════════════════════════
+/**
+ * Hook to terminate an agent runtime (delete it).
+ */
+export function useTerminateAgent() {
+    const queryClient = useQueryClient();
+    return useMutation({
+        mutationFn: async (params) => {
+            const { iamStore, coreStore } = await import('@datalayer/core/lib/state');
+            const token = iamStore.getState().token || '';
+            const runtimesRunUrl = coreStore.getState().configuration?.runtimesRunUrl || '';
+            const { deleteRuntime } = await import('@datalayer/core/lib/api/runtimes/runtimes');
+            return deleteRuntime(token, params.podName, runtimesRunUrl);
+        },
+        onSuccess: () => {
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.agentRuntimes.all(),
+            });
+            queryClient.invalidateQueries({
+                queryKey: agentQueryKeys.checkpoints.all(),
+            });
+        },
+    });
+}
+/**
+ * High-level hook that manages the full agent lifecycle:
+ * pause, resume, checkpoint, terminate, and status tracking.
+ *
+ * Designed to be composed inside `useAgents` or used standalone
+ * by advanced consumers who need direct access.
+ */
+export function useAgentLifecycle(options) {
+    const { agentSpecId, agentSpec, runtime, connectToRuntime, disconnect, launchRuntime, onResetAgentCreation, } = options;
+    // Local lifecycle state
+    const [lifecycleStatus, setLifecycleStatus] = useState('idle');
+    const [lifecycleError, setLifecycleError] = useState(null);
+    // Mutation hooks
+    const pauseAgentMutation = usePauseAgent();
+    const resumeAgentMutation = useResumeAgent();
+    const checkpointAgentMutation = useCheckpointAgent();
+    const terminateAgentMutation = useTerminateAgent();
+    const checkpointsQuery = useCheckpointsQuery();
+    const refreshCheckpointsList = useRefreshCheckpoints();
+    // ─── Pause ──────────────────────────────────────────────────────────
+    const pause = useCallback(async (mode = 'light', messages) => {
+        if (!runtime) {
+            setLifecycleError('No runtime to pause');
+            return;
+        }
+        if (lifecycleStatus === 'resumed') {
+            setLifecycleError('Resumed agents cannot be paused');
+            return;
+        }
+        try {
+            await pauseAgentMutation.mutateAsync({
+                podName: runtime.podName,
+                mode,
+                agentSpecId,
+                agentSpec,
+                messages,
+            });
+            setLifecycleStatus('paused');
+        }
+        catch (err) {
+            const msg = err instanceof Error ? err.message : String(err);
+            setLifecycleError(msg);
+        }
+    }, [runtime, lifecycleStatus, agentSpecId, agentSpec, pauseAgentMutation]);
+    // ─── Resume ─────────────────────────────────────────────────────────
+    const resume = useCallback(async (mode = 'criu', checkpointId, podName) => {
+        setLifecycleStatus('resuming');
+        setLifecycleError(null);
+        try {
+            const checkpoints = checkpointsQuery.data || [];
+            const targetPodName = podName ||
+                runtime?.podName ||
+                (checkpointId
+                    ? checkpoints.find((c) => c.id === checkpointId)?.runtime_uid
+                    : undefined);
+            if (targetPodName) {
+                await resumeAgentMutation.mutateAsync({
+                    podName: targetPodName,
+                    agentSpecId,
+                    mode,
+                    checkpointId,
+                });
+                // Refresh and rebind runtime connection so downstream calls
+                // target the restored runtime URL.
+                try {
+                    const { iamStore, coreStore } = await import('@datalayer/core/lib/state');
+                    const token = iamStore.getState().token || '';
+                    const runtimesRunUrl = coreStore.getState().configuration?.runtimesRunUrl || '';
+                    const { listRuntimes } = await import('@datalayer/core/lib/api/runtimes/runtimes');
+                    const runtimesResponse = await listRuntimes(token, runtimesRunUrl);
+                    const runtimes = runtimesResponse.runtimes || [];
+                    const aiAgentRuntimes = runtimes.filter(rt => {
+                        if (rt.environment_name !== 'ai-agents-env') {
+                            return false;
+                        }
+                        if (!agentSpecId) {
+                            return true;
+                        }
+                        const runtimeAgentSpecId = rt
+                            .agent_spec_id;
+                        return runtimeAgentSpecId === agentSpecId;
+                    });
+                    const latestRuntime = aiAgentRuntimes.slice().sort((a, b) => {
+                        const aTs = Number(a.started_at || 0);
+                        const bTs = Number(b.started_at || 0);
+                        return bTs - aTs;
+                    })[0];
+                    if (latestRuntime?.pod_name && latestRuntime?.ingress) {
+                        connectToRuntime({
+                            podName: latestRuntime.pod_name,
+                            environmentName: latestRuntime.environment_name,
+                            jupyterBaseUrl: latestRuntime.ingress,
+                        });
+                    }
+                }
+                catch (refreshError) {
+                    console.warn('Failed to refresh runtime binding after resume:', refreshError);
+                }
+                onResetAgentCreation?.();
+                setLifecycleStatus('resumed');
+            }
+            else {
+                await launchRuntime();
+            }
+        }
+        catch (err) {
+            const msg = err instanceof Error ? err.message : String(err);
+            setLifecycleError(msg);
+            setLifecycleStatus('error');
+        }
+    }, [
+        runtime,
+        checkpointsQuery.data,
+        agentSpecId,
+        resumeAgentMutation,
+        launchRuntime,
+        connectToRuntime,
+        onResetAgentCreation,
+    ]);
+    // ─── Checkpoint ─────────────────────────────────────────────────────
+    const checkpoint = useCallback(async (name, mode = 'criu', messages) => {
+        if (!runtime) {
+            setLifecycleError('No runtime to checkpoint');
+            return;
+        }
+        try {
+            await checkpointAgentMutation.mutateAsync({
+                podName: runtime.podName,
+                name,
+                mode,
+                agentSpecId,
+                agentSpec,
+                messages,
+            });
+            setLifecycleStatus('paused');
+        }
+        catch (err) {
+            const msg = err instanceof Error ? err.message : String(err);
+            setLifecycleError(msg);
+        }
+    }, [runtime, agentSpecId, agentSpec, checkpointAgentMutation]);
+    // ─── Terminate ──────────────────────────────────────────────────────
+    const terminate = useCallback(async () => {
+        if (!runtime) {
+            disconnect();
+            setLifecycleStatus('idle');
+            return;
+        }
+        try {
+            await terminateAgentMutation.mutateAsync({
+                podName: runtime.podName,
+            });
+            disconnect();
+            setLifecycleStatus('idle');
+            setLifecycleError(null);
+        }
+        catch (err) {
+            const msg = err instanceof Error ? err.message : String(err);
+            setLifecycleError(msg);
+        }
+    }, [runtime, terminateAgentMutation, disconnect]);
+    // ─── Refresh Checkpoints ────────────────────────────────────────────
+    const refreshCheckpoints = useCallback(() => {
+        refreshCheckpointsList();
+    }, [refreshCheckpointsList]);
+    return {
+        pause,
+        resume,
+        terminate,
+        checkpoint,
+        refreshCheckpoints,
+        checkpoints: (checkpointsQuery.data || []),
+        status: lifecycleStatus,
+        error: lifecycleError,
+        setStatus: setLifecycleStatus,
+        setError: setLifecycleError,
+    };
+}
+// ═══════════════════════════════════════════════════════════════════════════
+// Composite hook
+// ═══════════════════════════════════════════════════════════════════════════
+/**
+ * Composite hook for checkpoint operations.
+ *
+ * Provides query access (list + refresh) and mutation hooks
+ * for pause, resume, and checkpoint lifecycle.
+ */
+export function useCheckpoints() {
+    const checkpointsQuery = useCheckpointsQuery();
+    const refresh = useRefreshCheckpoints();
+    const pauseAgent = usePauseAgent();
+    const resumeAgent = useResumeAgent();
+    const checkpointAgent = useCheckpointAgent();
+    return useMemo(() => ({
+        checkpointsQuery,
+        refresh,
+        pauseAgent,
+        resumeAgent,
+        checkpointAgent,
+    }), [checkpointsQuery, refresh, pauseAgent, resumeAgent, checkpointAgent]);
+}

package/lib/hooks/useConfig.d.ts

@@ -0,0 +1,11 @@
+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> | {
+    data: undefined;
+    isLoading: boolean;
+    isError: boolean;
+    error: null;
+};

package/lib/hooks/useConfig.js

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+import { useContext } from 'react';
+import { useQuery, QueryClientContext } from '@tanstack/react-query';
+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) {
+    const queryClient = useContext(QueryClientContext);
+    if (!queryClient) {
+        return {
+            data: undefined,
+            isLoading: false,
+            isError: false,
+            error: null,
+        };
+    }
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    return useQuery({
+        queryFn: async () => {
+            // If configEndpoint is provided, use direct fetch (for FastAPI).
+            if (configEndpoint) {
+                const headers = {
+                    'Content-Type': 'application/json',
+                };
+                if (authToken) {
+                    headers['Authorization'] = `Bearer ${authToken}`;
+                }
+                const response = await fetch(configEndpoint, { headers });
+                if (!response.ok) {
+                    throw new Error(`Config fetch failed: ${response.statusText}`);
+                }
+                return response.json();
+            }
+            // Otherwise use Jupyter requestAPI.
+            return requestAPI('configure');
+        },
+        queryKey: ['models', configEndpoint || 'jupyter'],
+        enabled,
+        retry: 1,
+    });
+}

package/lib/hooks/useContextSnapshot.d.ts

@@ -0,0 +1,11 @@
+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.
+ */
+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;
+    isLoading: boolean;
+    isError: boolean;
+    error: null;
+};

package/lib/hooks/useContextSnapshot.js

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+import { useContext } from 'react';
+import { useQuery, QueryClientContext } from '@tanstack/react-query';
+import { getApiBaseFromConfig } from '../utils';
+/**
+ * Hook to poll agent context-snapshot from the backend.
+ * Returns cumulative token usage (input/output breakdown) tracked by the agent server.
+ */
+export function useContextSnapshot(enabled, configEndpoint, agentId, authToken) {
+    const queryClient = useContext(QueryClientContext);
+    if (!queryClient) {
+        return { data: undefined, isLoading: false, isError: false, error: null };
+    }
+    const snapshotUrl = configEndpoint && agentId
+        ? `${getApiBaseFromConfig(configEndpoint)}/configure/agents/${encodeURIComponent(agentId)}/context-snapshot`
+        : undefined;
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    return useQuery({
+        queryKey: ['context-snapshot-header', agentId, snapshotUrl],
+        queryFn: async () => {
+            if (!snapshotUrl) {
+                throw new Error('No context-snapshot URL available');
+            }
+            const headers = { 'Content-Type': 'application/json' };
+            if (authToken) {
+                headers['Authorization'] = `Bearer ${authToken}`;
+            }
+            const response = await fetch(snapshotUrl, { headers });
+            if (!response.ok) {
+                throw new Error(`Context snapshot fetch failed: ${response.statusText}`);
+            }
+            return response.json();
+        },
+        enabled: enabled && !!snapshotUrl,
+        // Poll every 10s, but stop after an error (e.g. runtime terminated).
+        refetchInterval: query => (query.state.status === 'error' ? false : 10_000),
+        refetchOnMount: 'always',
+        staleTime: 0,
+        retry: 1,
+    });
+}

package/lib/hooks/useMonitoring.d.ts

@@ -0,0 +1,24 @@
+export interface OtelQueryOptions {
+    metric: string;
+    serviceName?: string;
+    runUrl?: string;
+    apiKey?: string;
+    limit?: number;
+}
+interface MetricValueRow {
+    value?: unknown;
+    value_double?: unknown;
+    value_int?: unknown;
+}
+export declare function toMetricValue(row: MetricValueRow): number;
+export declare function fetchOtelMetricRows({ metric, serviceName, runUrl, apiKey, limit, }: OtelQueryOptions): Promise<MetricValueRow[]>;
+export declare function fetchOtelMetricTotal(options: OtelQueryOptions): Promise<number>;
+export interface OtelTotalTokensOptions {
+    serviceName?: string;
+    runUrl?: string;
+    apiKey?: string;
+    limit?: number;
+}
+export declare function fetchOtelTotalTokens({ serviceName, runUrl, apiKey, limit, }: OtelTotalTokensOptions): Promise<number>;
+export declare function useOtelTotalTokens({ serviceName, runUrl, apiKey, limit, }: OtelTotalTokensOptions): string;
+export {};