@datalayer/agent-runtimes 1.0.1 → 1.0.3

package/lib/hooks/useAgentsCatalog.d.ts

@@ -0,0 +1,25 @@
+import type { AgentRuntimeData, AgentSpec } from '../types';
+/**
+ * Centralised Zustand store for agents.
+ *
+ * Two collections are maintained:
+ *
+ * 1. **agentSpecs** – the static catalogue of available agent blueprints
+ *    (from `@datalayer/agent-runtimes/lib/specs`).
+ *    Populated once at import time; call `refreshSpecs()` to re-read.
+ *
+ * 2. **runningAgents** – live agent runtimes fetched from the runtimes
+ *    service.  Updated via `setRunningAgents()` whenever the TanStack
+ *    query refreshes.
+ */
+export type AgentCatalogStoreState = {
+    /** Static catalogue of agent blueprints. */
+    agentSpecs: AgentSpec[];
+    /** Live agent runtimes (running / starting). */
+    runningAgents: AgentRuntimeData[];
+    /** Re-read agent specs from the config. */
+    refreshSpecs: () => void;
+    /** Replace the running agents list (call from TanStack query effect). */
+    setRunningAgents: (agents: AgentRuntimeData[]) => void;
+};
+export declare const useAgentCatalogStore: import("zustand").UseBoundStore<import("zustand").StoreApi<AgentCatalogStoreState>>;

package/lib/hooks/useAgentsCatalog.js

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+/**
+ * Agent catalog hooks.
+ *
+ * Includes the Zustand catalog store and the agent registry hook.
+ *
+ * @module hooks/useAgentsCatalog
+ */
+import { create } from 'zustand';
+import { listAgentSpecs } from '../specs';
+export const useAgentCatalogStore = create()(set => ({
+    agentSpecs: listAgentSpecs(),
+    runningAgents: [],
+    refreshSpecs: () => set({ agentSpecs: listAgentSpecs() }),
+    setRunningAgents: agents => set(state => ({
+        runningAgents: agents.map(agent => {
+            if (agent.agentSpec)
+                return agent;
+            if (!agent.agent_spec_id)
+                return agent;
+            const spec = state.agentSpecs.find(s => s.id === agent.agent_spec_id);
+            return spec ? { ...agent, agentSpec: spec } : agent;
+        }),
+    })),
+}));

package/lib/hooks/useAgentsRegistry.d.ts

@@ -0,0 +1,10 @@
+export declare function useAgentRegistry(): {
+    agents: readonly import("..").AgentRegistryEntry[];
+    upsertAgent: (agent: Partial<import("..").AgentRegistryEntry> & {
+        id: string;
+        baseUrl: string;
+        protocol: import("../types").Protocol;
+    }) => void;
+    deleteAgent: (id: string) => void;
+    getAgentById: (id: string) => import("..").AgentRegistryEntry | undefined;
+};

package/lib/hooks/useAgentsRegistry.js

@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+/**
+ * Hook exposing the agent registry from the Zustand store.
+ *
+ * Provides CRUD operations on the in-memory agent map without coupling
+ * consumers directly to the low-level store.
+ *
+ * @module hooks/useAgentsRegistry
+ */
+import { useAgentStore } from '../stores/agentsStore';
+export function useAgentRegistry() {
+    const agents = useAgentStore(state => state.agents);
+    const upsertAgent = useAgentStore(state => state.upsertAgent);
+    const deleteAgent = useAgentStore(state => state.deleteAgent);
+    const getAgentById = useAgentStore(state => state.getAgentById);
+    return { agents, upsertAgent, deleteAgent, getAgentById };
+}

package/lib/hooks/useAgentsService.d.ts

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

package/lib/hooks/useAgentsService.js

@@ -0,0 +1,146 @@
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+/**
+ * Agents Service REST API hook and notebook agent management.
+ *
+ * @deprecated useAgents instead
+ * @module hooks/useAgentsService
+ */
+import { useEffect } from 'react';
+import { useCoreStore, useDatalayer } from '@datalayer/core';
+import { URLExt } from '@jupyterlab/coreutils';
+import { useAgentStore } from '../stores/agentsStore';
+// ═══════════════════════════════════════════════════════════════════════════
+// Agents Service REST API hook.
+// ═══════════════════════════════════════════════════════════════════════════
+/**
+ * @deprecated Use useAgents instead
+ */
+export const useAgentsService = (baseUrlOverride = 'api/ai-agents/v1') => {
+    const { configuration } = useCoreStore();
+    const { requestDatalayer } = useDatalayer({ notifyOnError: false });
+    const createAgent = (documentId, documentType, ingress, token, kernelId, { signal, baseUrl = baseUrlOverride } = {}) => {
+        return requestDatalayer({
+            url: URLExt.join(configuration.aiagentsRunUrl, baseUrl, 'agents'),
+            method: 'POST',
+            body: {
+                document_id: documentId,
+                document_type: documentType,
+                runtime: {
+                    ingress,
+                    token,
+                    kernel_id: kernelId,
+                },
+            },
+            signal,
+        });
+    };
+    const getAgents = ({ signal, baseUrl = baseUrlOverride, } = {}) => {
+        return requestDatalayer({
+            url: URLExt.join(configuration.aiagentsRunUrl, baseUrl, 'agents'),
+            method: 'GET',
+            signal,
+        });
+    };
+    const deleteAgent = (documentId, { signal, baseUrl = baseUrlOverride } = {}) => {
+        return requestDatalayer({
+            url: URLExt.join(configuration.aiagentsRunUrl, baseUrl, 'agents', documentId),
+            method: 'DELETE',
+            signal,
+        });
+    };
+    const getAgent = (documentId, { signal, baseUrl = baseUrlOverride } = {}) => {
+        return requestDatalayer({
+            url: URLExt.join(configuration.aiagentsRunUrl, baseUrl, 'agents', documentId),
+            method: 'GET',
+            signal,
+        });
+    };
+    const patchAgent = (documentId, ingress, token, kernelId, { signal, baseUrl = baseUrlOverride } = {}) => {
+        return requestDatalayer({
+            url: URLExt.join(configuration.aiagentsRunUrl, baseUrl, 'agents', documentId),
+            method: 'PATCH',
+            body: {
+                runtime: ingress && token && kernelId
+                    ? {
+                        ingress,
+                        token,
+                        kernel_id: kernelId,
+                    }
+                    : null,
+            },
+            signal,
+        });
+    };
+    return {
+        createAgent,
+        getAgents,
+        deleteAgent,
+        getAgent,
+        patchAgent,
+    };
+};
+/**
+ * Get the notebook AI agent if any.
+ *
+ * This performs a periodic liveness check and keeps the local store in sync.
+ * @deprecated Use useAgents 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);
+    useEffect(() => {
+        let abortController;
+        const refreshAIAgent = async () => {
+            abortController = new AbortController();
+            try {
+                const response = await getAgent(notebookId, {
+                    signal: abortController.signal,
+                });
+                if (!response.success) {
+                    deleteAgent(notebookId);
+                    return;
+                }
+                const currentAgent = getAgentById(notebookId);
+                const runtimeId = response.agent.runtime?.id;
+                if (currentAgent) {
+                    if (currentAgent.runtimeId !== runtimeId) {
+                        upsertAgent({
+                            id: notebookId,
+                            baseUrl: currentAgent.baseUrl,
+                            protocol: currentAgent.protocol,
+                            runtimeId,
+                            status: 'running',
+                        });
+                    }
+                }
+                else {
+                    upsertAgent({
+                        id: notebookId,
+                        name: `Notebook ${notebookId}`,
+                        description: 'AI agent for notebook',
+                        baseUrl: '',
+                        protocol: 'vercel-ai',
+                        documentId: notebookId,
+                        runtimeId,
+                        status: 'running',
+                    });
+                }
+            }
+            catch {
+                deleteAgent(notebookId);
+            }
+        };
+        const refreshInterval = setInterval(refreshAIAgent, 60_000);
+        return () => {
+            abortController?.abort('Component unmounted');
+            clearInterval(refreshInterval);
+        };
+    }, [agents, notebookId, getAgent, deleteAgent, getAgentById, upsertAgent]);
+    return getAgentById(notebookId);
+}

package/lib/hooks/useChat.d.ts

@@ -1,5 +1,5 @@
-import type { ChatMessage } from '../components/chat/types/message';
-import type { InferenceRequestOptions } from '../components/chat/types/inference';
+import type { ChatMessage } from '../types/messages';
+import type { InferenceRequestOptions } from '../types/inference';
 /**
  * Return type for useChat hook
  */

package/lib/hooks/useChat.js

@@ -6,11 +6,12 @@
  * Main chat hook for chat component.
  * Provides chat state, message sending, and streaming functionality.
  *
- * @module components/chat/hooks/useChat
+ * @module chat/hooks/useChat
  */
 import { useCallback, useMemo } from 'react';
-import { useChatStore } from '../components/chat/store';
-import { createUserMessage, createAssistantMessage, generateMessageId, } from '../components/chat/types/message';
+import { useChatStore } from '../stores';
+import { createUserMessage, createAssistantMessage, generateMessageId, } from '../types/messages';
+import { sanitizeAssistantContent } from '../utils';
 /**
  * Main hook for chat functionality
  *
@@ -104,10 +105,11 @@ export function useChat() {
                         case 'message':
                             // AG-UI sends incremental messages with accumulating content
                             if (event.message) {
+                                const rawContent = typeof event.message.content === 'string'
+                                    ? event.message.content
+                                    : '';
                                 currentAssistantContent =
-                                    typeof event.message.content === 'string'
-                                        ? event.message.content
-                                        : '';
+                                    sanitizeAssistantContent(rawContent);
                                 // Update the assistant message with the current content
                                 updateMessage(assistantMessageId, {
                                     content: currentAssistantContent,
@@ -197,9 +199,10 @@ export function useChat() {
                         // Subscribe to continuation events
                         const unsubscribeContinuation = protocolAdapter.subscribe(event => {
                             if (event.type === 'message' && event.message) {
-                                const content = typeof event.message.content === 'string'
+                                const rawContent = typeof event.message.content === 'string'
                                     ? event.message.content
                                     : '';
+                                const content = sanitizeAssistantContent(rawContent);
                                 updateMessage(continuationAssistantId, {
                                     content,
                                     toolCalls: event.message.toolCalls,
@@ -255,8 +258,11 @@ export function useChat() {
                     }
                 });
                 // Update final message with complete content and tool calls
+                const finalAssistantContent = typeof response.message.content === 'string'
+                    ? sanitizeAssistantContent(response.message.content)
+                    : response.message.content;
                 updateMessage(assistantMessageId, {
-                    content: response.message.content,
+                    content: finalAssistantContent,
                     toolCalls: response.message.toolCalls,
                 });
                 // Handle tool calls if present

package/lib/hooks/useCheckpoints.d.ts

@@ -0,0 +1,176 @@
+import type { CheckpointMode, AgentStatus, AgentConnection, CheckpointRecord } from '../types';
+/**
+ * Checkpoint data returned by the runtime-checkpoints API.
+ */
+export type CheckpointData = {
+    id: string;
+    name: string;
+    description: string;
+    runtime_uid: string;
+    agent_spec_id: string;
+    agentspec: Record<string, unknown>;
+    metadata: Record<string, unknown>;
+    checkpoint_mode?: 'criu' | 'light';
+    messages?: string[];
+    status: string;
+    status_message: string;
+    updated_at: string;
+    start_date?: string;
+    end_date?: string;
+};
+export type PauseAgentParams = {
+    podName: string;
+    mode?: CheckpointMode;
+    agentSpecId?: string;
+    agentSpec?: Record<string, any>;
+    messages?: string[];
+};
+export type ResumeAgentParams = {
+    podName: string;
+    agentSpecId?: string;
+    mode?: CheckpointMode;
+    checkpointId?: string;
+};
+export type CheckpointAgentParams = {
+    podName: string;
+    name?: string;
+    mode?: CheckpointMode;
+    agentSpecId?: string;
+    agentSpec?: Record<string, any>;
+    messages?: string[];
+};
+export type TerminateAgentParams = {
+    podName: string;
+};
+/**
+ * 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 declare function useCheckpointsQuery(): import("@tanstack/react-query").UseQueryResult<CheckpointData[], Error>;
+/**
+ * Hook to refresh the checkpoints list.
+ */
+export declare function useRefreshCheckpoints(): () => void;
+/**
+ * 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 declare function useDeletePausedAgentRuntime(): import("@tanstack/react-query").UseMutationResult<any, Error, string, unknown>;
+/**
+ * Hook to resume a paused agent runtime via checkpoint restore.
+ */
+export declare function useResumePausedAgentRuntime(): import("@tanstack/react-query").UseMutationResult<any, Error, string, unknown>;
+/**
+ * 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 declare function usePauseAgent(): import("@tanstack/react-query").UseMutationResult<import("@datalayer/core/lib/api/runtimes/runtimes").PauseResumeResponse, Error, PauseAgentParams, unknown>;
+/**
+ * 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 declare function useResumeAgent(): import("@tanstack/react-query").UseMutationResult<import("@datalayer/core/lib/api/runtimes/runtimes").PauseResumeResponse, Error, ResumeAgentParams, unknown>;
+/**
+ * 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 declare function useCheckpointAgent(): import("@tanstack/react-query").UseMutationResult<{
+    checkpoint: import("@datalayer/core/lib/api/runtimes/checkpoints").RuntimeCheckpointData;
+    success: boolean;
+    message: string;
+    checkpoint_id: string | null;
+}, Error, CheckpointAgentParams, unknown>;
+/**
+ * Hook to terminate an agent runtime (delete it).
+ */
+export declare function useTerminateAgent(): import("@tanstack/react-query").UseMutationResult<void, Error, TerminateAgentParams, unknown>;
+/**
+ * Options for {@link useAgentLifecycle}.
+ */
+export interface AgentLifecycleOptions {
+    /** Agent spec ID (identifies the agent for lifecycle management) */
+    agentSpecId?: string;
+    /** Full agent spec object (persisted with checkpoints) */
+    agentSpec?: Record<string, any>;
+    /** Current runtime connection from the agent store */
+    runtime: AgentConnection | null;
+    /** Callback to connect to a new runtime pod */
+    connectToRuntime: (opts: {
+        podName: string;
+        environmentName: string;
+        jupyterBaseUrl?: string;
+    }) => void;
+    /** Callback to disconnect from the runtime */
+    disconnect: () => void;
+    /** Callback to launch a new runtime */
+    launchRuntime: () => Promise<AgentConnection>;
+    /**
+     * Callback fired when the hook needs to signal
+     * that the agent should be re-created on a (possibly new) pod.
+     */
+    onResetAgentCreation?: () => void;
+}
+/**
+ * Return type for {@link useAgentLifecycle}.
+ */
+export interface AgentLifecycleReturn {
+    /** Pause the agent (checkpoint-mode aware) */
+    pause: (mode?: CheckpointMode, messages?: string[]) => Promise<void>;
+    /** Resume a paused agent (checkpoint restore) */
+    resume: (mode?: CheckpointMode, checkpointId?: string, podName?: string) => Promise<void>;
+    /** Terminate the agent (delete runtime) */
+    terminate: () => Promise<void>;
+    /** Create a named checkpoint (pause → wait → stay paused) */
+    checkpoint: (name?: string, mode?: CheckpointMode, messages?: string[]) => Promise<void>;
+    /** Refresh the checkpoints list */
+    refreshCheckpoints: () => void;
+    /** List of persisted checkpoints */
+    checkpoints: CheckpointRecord[];
+    /** Lifecycle status */
+    status: AgentStatus;
+    /** Lifecycle error */
+    error: string | null;
+    /** Update the lifecycle status directly */
+    setStatus: (status: AgentStatus) => void;
+    /** Update the lifecycle error directly */
+    setError: (error: string | null) => void;
+}
+/**
+ * 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 declare function useAgentLifecycle(options: AgentLifecycleOptions): AgentLifecycleReturn;
+/**
+ * Composite hook for checkpoint operations.
+ *
+ * Provides query access (list + refresh) and mutation hooks
+ * for pause, resume, and checkpoint lifecycle.
+ */
+export declare function useCheckpoints(): {
+    checkpointsQuery: import("@tanstack/react-query").UseQueryResult<CheckpointData[], Error>;
+    refresh: () => void;
+    pauseAgent: import("@tanstack/react-query").UseMutationResult<import("@datalayer/core/lib/api/runtimes/runtimes").PauseResumeResponse, Error, PauseAgentParams, unknown>;
+    resumeAgent: import("@tanstack/react-query").UseMutationResult<import("@datalayer/core/lib/api/runtimes/runtimes").PauseResumeResponse, Error, ResumeAgentParams, unknown>;
+    checkpointAgent: import("@tanstack/react-query").UseMutationResult<{
+        checkpoint: import("@datalayer/core/lib/api/runtimes/checkpoints").RuntimeCheckpointData;
+        success: boolean;
+        message: string;
+        checkpoint_id: string | null;
+    }, Error, CheckpointAgentParams, unknown>;
+};