@datalayer/agent-runtimes 1.0.3 → 1.0.5

package/lib/examples/AgentEvalsExample.js

@@ -18,14 +18,14 @@ import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
 import { useEffect, useState, useCallback, useRef } from 'react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { Text, Button, Spinner, Heading, Label, Flash, ProgressBar, } from '@primer/react';
-import { AlertIcon, BeakerIcon, CheckCircleIcon, XCircleIcon, PlayIcon, SignOutIcon, } from '@primer/octicons-react';
+import { BeakerIcon, CheckCircleIcon, XCircleIcon, PlayIcon, } from '@primer/octicons-react';
 import { Box } from '@datalayer/primer-addons';
+import { AuthRequiredView, ErrorView } from './components';
 import { ThemedProvider } from './utils/themedProvider';
+import { uniqueAgentId } from './utils/agentId';
 import { useSimpleAuthStore } from '@datalayer/core/lib/views/otel';
-import { SignInSimple } from '@datalayer/core/lib/views/iam';
-import { UserBadge } from '@datalayer/core/lib/views/profile';
 import { Chat } from '../chat';
-import { useAgents } from '../hooks/useAgents';
+import { useAgentRuntimes } from '../hooks/useAgentRuntimes';
 const queryClient = new QueryClient();
 // ─── Constants ─────────────────────────────────────────────────────────────
 const AGENT_NAME = 'eval-demo-agent';
@@ -33,12 +33,14 @@ const AGENT_SPEC_ID = 'monitor-sales-kpis';
 // ─── Inner component (rendered after auth) ─────────────────────────────────
 const AgentEvalsInner = ({ onLogout }) => {
     const { token } = useSimpleAuthStore();
-    const { runtime, status: runtimeStatus, isReady, error: hookError, } = useAgents({
+    const agentName = useRef(uniqueAgentId(AGENT_NAME)).current;
+    const { runtime, status: runtimeStatus, isReady, error: hookError, } = useAgentRuntimes({
         agentSpecId: AGENT_SPEC_ID,
         autoStart: true,
         agentConfig: {
-            name: AGENT_NAME,
-            protocol: 'ag-ui',
+            name: agentName,
+            model: 'bedrock:us.anthropic.claude-3-5-haiku-20241022-v1:0',
+            protocol: 'vercel-ai',
             description: 'Agent with evaluation and quality scoring',
         },
     });
@@ -113,14 +115,7 @@ const AgentEvalsInner = ({ onLogout }) => {
                         : 'Creating eval demo agent…' })] }));
     }
     if (runtimeStatus === 'error' || hookError) {
-        return (_jsxs(Box, { sx: {
-                display: 'flex',
-                flexDirection: 'column',
-                alignItems: 'center',
-                justifyContent: 'center',
-                height: '100vh',
-                gap: 3,
-            }, children: [_jsx(AlertIcon, { size: 48 }), _jsx(Text, { sx: { color: 'danger.fg' }, children: hookError || 'Agent failed to start' })] }));
+        return _jsx(ErrorView, { error: hookError, onLogout: onLogout });
     }
     const latestScore = evalRuns.length > 0 ? evalRuns[0].score : null;
     return (_jsxs(Box, { sx: {
@@ -136,7 +131,7 @@ const AgentEvalsInner = ({ onLogout }) => {
                     borderBottom: '1px solid',
                     borderColor: 'border.default',
                     flexShrink: 0,
-                }, children: [_jsx(BeakerIcon, { size: 16 }), _jsxs(Heading, { as: "h3", sx: { fontSize: 2, flex: 1 }, children: ["Evaluation \u2014 ", podName] }), token && _jsx(UserBadge, { token: token, variant: "small" }), _jsx(Button, { size: "small", variant: "invisible", onClick: onLogout, leadingVisual: SignOutIcon, sx: { color: 'fg.muted' }, children: "Sign out" })] }), _jsxs(Box, { sx: { flex: 1, minHeight: 0, display: 'flex' }, children: [_jsx(Box, { sx: { flex: 1, minWidth: 0 }, children: _jsx(Chat, { protocol: "ag-ui", baseUrl: agentBaseUrl, agentId: agentId, title: "Eval Agent", placeholder: "Chat with the agent, then run evaluations\u2026", description: latestScore != null
+                }, children: [_jsx(BeakerIcon, { size: 16 }), _jsxs(Heading, { as: "h3", sx: { fontSize: 2, flex: 1 }, children: ["Evaluation \u2014 ", podName] })] }), _jsxs(Box, { sx: { flex: 1, minHeight: 0, display: 'flex' }, children: [_jsx(Box, { sx: { flex: 1, minWidth: 0 }, children: _jsx(Chat, { protocol: "vercel-ai", baseUrl: agentBaseUrl, agentId: agentId, title: "Eval Agent", placeholder: "Chat with the agent, then run evaluations\u2026", description: latestScore != null
                                 ? `Last score: ${(latestScore * 100).toFixed(0)}%`
                                 : 'No evaluations run yet', showHeader: true, autoFocus: true, height: "100%", runtimeId: podName, historyEndpoint: `${agentBaseUrl}/api/v1/history`, suggestions: [
                                 {
@@ -188,7 +183,7 @@ const syncTokenToIamStore = (token) => {
 };
 // ─── Main component with auth gate ─────────────────────────────────────────
 const AgentEvalsExample = () => {
-    const { token, setAuth, clearAuth } = useSimpleAuthStore();
+    const { token, clearAuth } = useSimpleAuthStore();
     const hasSynced = useRef(false);
     useEffect(() => {
         if (token && !hasSynced.current) {
@@ -196,11 +191,6 @@ const AgentEvalsExample = () => {
             syncTokenToIamStore(token);
         }
     }, [token]);
-    const handleSignIn = useCallback((newToken, handle) => {
-        setAuth(newToken, handle);
-        hasSynced.current = true;
-        syncTokenToIamStore(newToken);
-    }, [setAuth]);
     const handleLogout = useCallback(() => {
         clearAuth();
         hasSynced.current = false;
@@ -209,7 +199,7 @@ const AgentEvalsExample = () => {
         });
     }, [clearAuth]);
     if (!token) {
-        return (_jsx(ThemedProvider, { children: _jsx(SignInSimple, { onSignIn: handleSignIn, onApiKeySignIn: apiKey => handleSignIn(apiKey, 'api-key-user'), title: "Agent Evaluation", description: "Sign in to evaluate agent quality and review scores.", leadingIcon: _jsx(BeakerIcon, { size: 24 }) }) }));
+        return (_jsx(ThemedProvider, { children: _jsx(AuthRequiredView, {}) }));
     }
     return (_jsx(QueryClientProvider, { client: queryClient, children: _jsx(ThemedProvider, { children: _jsx(AgentEvalsInner, { onLogout: handleLogout }) }) }));
 };

package/lib/examples/AgentGuardrailsExample.js

@@ -18,40 +18,213 @@ import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
 import { useEffect, useState, useCallback, useRef } from 'react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { Text, Button, Spinner, Heading, Label, Flash, ProgressBar, } from '@primer/react';
-import { AlertIcon, ShieldCheckIcon, CheckIcon, XIcon, SignOutIcon, } from '@primer/octicons-react';
+import { ShieldCheckIcon, CheckIcon, XIcon, DotFillIcon, } from '@primer/octicons-react';
 import { Box } from '@datalayer/primer-addons';
+import { buildOtelWebSocketUrl } from '@datalayer/core/lib/otel';
+import { useCoreStore } from '@datalayer/core/lib/state';
+import { AuthRequiredView, ErrorView } from './components';
 import { ThemedProvider } from './utils/themedProvider';
+import { uniqueAgentId } from './utils/agentId';
+import { parseAgentStreamMessage } from '../types/stream';
+import { getAgentSpecs } from '../specs/agents';
+import { subscribeOtelWs } from '../context/otelWsPool';
+import { toMetricValue } from '../hooks/useMonitoring';
+import { useAIAgentsWebSocket } from '../hooks';
 const queryClient = new QueryClient();
 import { useSimpleAuthStore } from '@datalayer/core/lib/views/otel';
-import { SignInSimple } from '@datalayer/core/lib/views/iam';
-import { UserBadge } from '@datalayer/core/lib/views/profile';
 import { Chat } from '../chat';
-import { useAgents } from '../hooks/useAgents';
 // ─── Constants ─────────────────────────────────────────────────────────────
 const AGENT_NAME = 'guardrails-demo-agent';
-const AGENT_SPEC_ID = 'monitor-sales-kpis';
-const COST_LIMIT_USD = 5.0;
+const AGENT_SPEC_ID = 'demo-guardrails';
+const DEFAULT_LOCAL_BASE_URL = import.meta.env.VITE_BASE_URL || 'http://localhost:8765';
+const OTEL_BASE_URL_ENV = import.meta.env.VITE_OTEL_BASE_URL;
+const DATALAYER_RUN_URL_ENV = import.meta.env.DATALAYER_RUN_URL;
+const OTEL_SERVICE_NAME = 'agent-runtimes';
+const COST_RUN_METRIC = 'agent_runtimes.capability.cost.run.usd';
+const COST_CUMULATIVE_METRIC = 'agent_runtimes.capability.cost.cumulative.usd';
+const toNumberOrNull = (value) => {
+    if (typeof value === 'number' && Number.isFinite(value)) {
+        return value;
+    }
+    if (typeof value === 'string' && value.trim().length > 0) {
+        const parsed = Number(value);
+        if (Number.isFinite(parsed)) {
+            return parsed;
+        }
+    }
+    return null;
+};
+const resolvePerRunBudgetFromSpec = (agentSpecId) => {
+    const spec = getAgentSpecs(agentSpecId);
+    if (!spec?.guardrails?.length) {
+        return null;
+    }
+    for (const guardrail of spec.guardrails) {
+        if (!guardrail || typeof guardrail !== 'object') {
+            continue;
+        }
+        const costBudget = guardrail.cost_budget;
+        if (!costBudget || typeof costBudget !== 'object') {
+            continue;
+        }
+        const perRunUsd = toNumberOrNull(costBudget.per_run_usd);
+        if (perRunUsd != null) {
+            return perRunUsd;
+        }
+    }
+    return null;
+};
+const DEFAULT_RUN_BUDGET_USD = resolvePerRunBudgetFromSpec(AGENT_SPEC_ID);
+const toApprovalRequest = (payload) => ({
+    id: payload.id,
+    tool_name: payload.tool_name,
+    tool_args: payload.tool_args ?? {},
+    created_at: payload.created_at ?? new Date().toISOString(),
+});
+const normalizeBaseUrl = (rawBaseUrl) => {
+    if (rawBaseUrl.startsWith('http://') ||
+        rawBaseUrl.startsWith('https://') ||
+        rawBaseUrl.startsWith('ws://') ||
+        rawBaseUrl.startsWith('wss://')) {
+        return rawBaseUrl;
+    }
+    const protocol = typeof window !== 'undefined' && window.location.protocol === 'https:'
+        ? 'https:'
+        : 'http:';
+    const host = typeof window !== 'undefined' ? window.location.host : '';
+    return `${protocol}//${host}${rawBaseUrl}`;
+};
+const parseAttributes = (attrs) => {
+    if (attrs && typeof attrs === 'object' && !Array.isArray(attrs)) {
+        return attrs;
+    }
+    if (typeof attrs === 'string') {
+        try {
+            return JSON.parse(attrs);
+        }
+        catch {
+            return {};
+        }
+    }
+    return {};
+};
+const extractServiceName = (row) => {
+    const directCandidates = [row.service_name, row.service, row.serviceName];
+    for (const candidate of directCandidates) {
+        if (typeof candidate === 'string' && candidate.length > 0) {
+            return candidate;
+        }
+    }
+    const resourceAttributes = row.resource_attributes;
+    if (resourceAttributes && typeof resourceAttributes === 'object') {
+        const nested = resourceAttributes['service.name'];
+        if (typeof nested === 'string' && nested.length > 0) {
+            return nested;
+        }
+    }
+    return undefined;
+};
+const extractAgentId = (row) => {
+    const attrs = parseAttributes(row.attributes);
+    const aid = attrs['agent.id'];
+    return typeof aid === 'string' ? aid : undefined;
+};
+const rowTimestampMs = (row) => {
+    const nanoTs = row.timestamp_unix_nano ?? row.observed_timestamp_unix_nano;
+    if (typeof nanoTs === 'number' && nanoTs > 0) {
+        return nanoTs / 1_000_000;
+    }
+    if (typeof nanoTs === 'string' && nanoTs.length > 0) {
+        const parsed = Number(nanoTs);
+        if (Number.isFinite(parsed) && parsed > 0) {
+            return parsed / 1_000_000;
+        }
+    }
+    const isoTs = row.timestamp;
+    if (typeof isoTs === 'string' && isoTs.length > 0) {
+        const ms = new Date(isoTs).getTime();
+        if (Number.isFinite(ms) && ms > 0) {
+            return ms;
+        }
+    }
+    return 0;
+};
+const selectLatestOtelCost = (rows, agentId) => {
+    const filtered = rows.filter(row => {
+        const metricName = row.metric_name;
+        if (metricName !== COST_RUN_METRIC &&
+            metricName !== COST_CUMULATIVE_METRIC) {
+            return false;
+        }
+        if (extractServiceName(row) !== OTEL_SERVICE_NAME) {
+            return false;
+        }
+        return extractAgentId(row) === agentId;
+    });
+    if (filtered.length === 0) {
+        return null;
+    }
+    let latestTimestampMs = 0;
+    let runUsd = 0;
+    let cumulativeUsd = 0;
+    for (const row of filtered) {
+        const ts = rowTimestampMs(row);
+        if (!Number.isFinite(ts) || ts <= 0) {
+            continue;
+        }
+        if (ts > latestTimestampMs) {
+            latestTimestampMs = ts;
+            runUsd = 0;
+            cumulativeUsd = 0;
+        }
+        if (ts !== latestTimestampMs) {
+            continue;
+        }
+        const value = Math.max(0, toMetricValue(row));
+        if (row.metric_name === COST_RUN_METRIC) {
+            runUsd += value;
+        }
+        else if (row.metric_name === COST_CUMULATIVE_METRIC) {
+            cumulativeUsd = Math.max(cumulativeUsd, value);
+        }
+    }
+    if (latestTimestampMs <= 0) {
+        return null;
+    }
+    return {
+        timestampMs: latestTimestampMs,
+        runUsd,
+        cumulativeUsd,
+    };
+};
 // ─── Inner component (rendered after auth) ─────────────────────────────────
 const AgentGuardrailsInner = ({ onLogout, }) => {
     const { token } = useSimpleAuthStore();
-    const { runtime, status: runtimeStatus, isReady, error: hookError, } = useAgents({
-        agentSpecId: AGENT_SPEC_ID,
-        autoStart: true,
-        agentConfig: {
-            name: AGENT_NAME,
-            protocol: 'ag-ui',
-            description: 'Agent with cost budget and tool approval guardrails',
-        },
-    });
+    const { configuration } = useCoreStore();
+    const agentName = useRef(uniqueAgentId(AGENT_NAME)).current;
+    const [runtimeStatus, setRuntimeStatus] = useState('launching');
+    const [isReady, setIsReady] = useState(false);
+    const [hookError, setHookError] = useState(null);
+    const [agentId, setAgentId] = useState(agentName);
+    const [isReconnectedAgent, setIsReconnectedAgent] = useState(false);
     // Cost tracking
-    const [costUsd, setCostUsd] = useState(0);
+    const [snapshotRunCostUsd, setSnapshotRunCostUsd] = useState(0);
+    const [runBudgetUsd, setRunBudgetUsd] = useState(DEFAULT_RUN_BUDGET_USD);
     const [totalTokens, setTotalTokens] = useState(0);
+    const [otelRunCostUsd, setOtelRunCostUsd] = useState(null);
+    const [otelCumulativeCostUsd, setOtelCumulativeCostUsd] = useState(null);
+    const [otelSampleTimestamp, setOtelSampleTimestamp] = useState(null);
     // Tool approval queue
     const [approvals, setApprovals] = useState([]);
     const [approvalLoading, setApprovalLoading] = useState(null);
-    const agentBaseUrl = runtime?.agentBaseUrl || '';
-    const agentId = runtime?.agentId || AGENT_NAME;
-    const podName = runtime?.podName || '(launching…)';
+    const agentBaseUrl = DEFAULT_LOCAL_BASE_URL;
+    const otelBaseUrl = configuration?.otelRunUrl ||
+        configuration?.runUrl ||
+        OTEL_BASE_URL_ENV ||
+        DATALAYER_RUN_URL_ENV ||
+        'https://prod1.datalayer.run';
+    const podName = agentId;
+    const chatAuthToken = token === null ? undefined : token;
     // Authenticated fetch helper (for sidecar endpoints)
     const authFetch = useCallback((url, opts = {}) => fetch(url, {
         ...opts,
@@ -61,46 +234,158 @@ const AgentGuardrailsInner = ({ onLogout, }) => {
             ...(opts.headers ?? {}),
         },
     }), [token]);
-    // ── Poll cost + tool approvals ───────────────────────────────────────────
     useEffect(() => {
-        if (!isReady || !agentBaseUrl)
-            return;
-        const poll = async () => {
+        let isCancelled = false;
+        const createLocalAgent = async () => {
+            setRuntimeStatus('launching');
+            setIsReady(false);
+            setHookError(null);
+            setIsReconnectedAgent(false);
             try {
-                // Cost usage
-                const costRes = await authFetch(`${agentBaseUrl}/api/v1/agents/${agentId}/cost`);
-                if (costRes.ok) {
-                    const d = await costRes.json();
-                    setCostUsd(d.total_cost_usd ?? d.cost ?? 0);
-                    setTotalTokens(d.total_tokens ?? 0);
+                const response = await authFetch(`${agentBaseUrl}/api/v1/agents`, {
+                    method: 'POST',
+                    body: JSON.stringify({
+                        name: agentName,
+                        description: 'Agent with cost budget and tool approval guardrails',
+                        agent_library: 'pydantic-ai',
+                        transport: 'vercel-ai',
+                        agent_spec_id: AGENT_SPEC_ID,
+                        enable_skills: true,
+                        tools: [],
+                    }),
+                });
+                let resolvedAgentId = agentName;
+                let isAlreadyRunning = false;
+                if (response.ok) {
+                    const data = await response.json();
+                    resolvedAgentId = data?.id || agentName;
+                }
+                else {
+                    const contentType = response.headers.get('content-type') || '';
+                    let detail = '';
+                    if (contentType.includes('application/json')) {
+                        const data = await response.json().catch(() => null);
+                        detail =
+                            (typeof data?.detail === 'string' && data.detail) ||
+                                (typeof data?.message === 'string' && data.message) ||
+                                '';
+                    }
+                    else {
+                        detail = await response.text();
+                    }
+                    if (response.status === 409 || /already exists/i.test(detail || '')) {
+                        isAlreadyRunning = true;
+                    }
+                    else {
+                        throw new Error(detail || `Failed to create local agent: ${response.status}`);
+                    }
+                }
+                if (!isCancelled) {
+                    setAgentId(resolvedAgentId);
+                    setIsReconnectedAgent(isAlreadyRunning);
+                    setIsReady(true);
+                    setRuntimeStatus('ready');
                 }
             }
-            catch {
-                /* endpoint may not be wired */
+            catch (error) {
+                if (!isCancelled) {
+                    setHookError(error instanceof Error ? error.message : 'Agent failed to start');
+                    setRuntimeStatus('error');
+                }
             }
-            try {
-                // Tool approvals
-                const apprRes = await authFetch(`${agentBaseUrl}/api/v1/agents/${agentId}/tool-approvals?status=pending`);
-                if (apprRes.ok) {
-                    const d = await apprRes.json();
-                    setApprovals(Array.isArray(d) ? d : (d.requests ?? []));
+        };
+        void createLocalAgent();
+        return () => {
+            isCancelled = true;
+        };
+    }, [agentBaseUrl, agentName, authFetch]);
+    const handleGuardrailsStreamMessage = useCallback((message) => {
+        try {
+            const stream = parseAgentStreamMessage(message?.raw ?? message);
+            if (!stream) {
+                return;
+            }
+            if (stream.type === 'agent.snapshot') {
+                const payload = stream.payload;
+                setApprovals((payload.approvals ?? []).map(toApprovalRequest));
+                const snapshotCost = payload.contextSnapshot?.costUsage ?? payload.costUsage;
+                if (snapshotCost) {
+                    setSnapshotRunCostUsd(Number(snapshotCost.cumulativeCostUsd ??
+                        snapshotCost.lastTurnCostUsd ??
+                        0));
+                    setRunBudgetUsd(prev => snapshotCost.perRunBudgetUsd == null
+                        ? (prev ?? DEFAULT_RUN_BUDGET_USD)
+                        : Number(snapshotCost.perRunBudgetUsd));
+                    setTotalTokens(Number(snapshotCost.totalTokensUsed ?? 0));
                 }
+                return;
             }
-            catch {
-                /* ok */
+            if (stream.type === 'tool_approval_created') {
+                const approval = toApprovalRequest(stream.payload);
+                setApprovals(prev => {
+                    const next = prev.filter(item => item.id !== approval.id);
+                    next.unshift(approval);
+                    return next;
+                });
+                return;
             }
-        };
-        poll();
-        const interval = setInterval(poll, 5_000);
-        return () => clearInterval(interval);
-    }, [isReady, agentBaseUrl, agentId, authFetch]);
+            if (stream.type === 'tool_approval_approved' ||
+                stream.type === 'tool_approval_rejected') {
+                const approval = stream.payload;
+                setApprovals(prev => prev.filter(item => item.id !== approval.id));
+            }
+        }
+        catch {
+            // Ignore malformed stream payloads.
+        }
+    }, []);
+    const approvalSocket = useAIAgentsWebSocket({
+        enabled: isReady && Boolean(agentBaseUrl),
+        baseUrl: agentBaseUrl,
+        path: '/api/v1/tool-approvals/ws',
+        queryParams: { agent_id: agentId },
+        onMessage: handleGuardrailsStreamMessage,
+        reconnectDelayMs: attempt => Math.min(1000 * 2 ** Math.max(0, attempt - 1), 10000),
+    });
+    const wsState = approvalSocket.connectionState;
+    useEffect(() => {
+        setOtelRunCostUsd(null);
+        setOtelCumulativeCostUsd(null);
+        setOtelSampleTimestamp(null);
+        if (!isReady || !token) {
+            return;
+        }
+        let wsUrl;
+        try {
+            wsUrl = buildOtelWebSocketUrl({
+                baseUrl: normalizeBaseUrl(otelBaseUrl),
+                token,
+            });
+        }
+        catch {
+            return;
+        }
+        const unsubscribe = subscribeOtelWs(wsUrl, msg => {
+            if (msg.signal !== 'metrics' || !Array.isArray(msg.data)) {
+                return;
+            }
+            const sample = selectLatestOtelCost(msg.data, agentId);
+            if (!sample) {
+                return;
+            }
+            setOtelRunCostUsd(sample.runUsd);
+            setOtelCumulativeCostUsd(sample.cumulativeUsd);
+            setOtelSampleTimestamp(sample.timestampMs);
+        });
+        return unsubscribe;
+    }, [agentId, isReady, otelBaseUrl, token]);
     // ── Approve / Reject ─────────────────────────────────────────────────────
     const handleApprove = useCallback(async (requestId) => {
         if (!agentBaseUrl)
             return;
         setApprovalLoading(requestId);
         try {
-            await authFetch(`${agentBaseUrl}/api/v1/agents/${agentId}/tool-approvals/${requestId}/approve`, { method: 'POST' });
+            await authFetch(`${agentBaseUrl}/api/v1/tool-approvals/${requestId}/approve`, { method: 'POST' });
             setApprovals(prev => prev.filter(a => a.id !== requestId));
         }
         catch {
@@ -109,13 +394,13 @@ const AgentGuardrailsInner = ({ onLogout, }) => {
         finally {
             setApprovalLoading(null);
         }
-    }, [agentBaseUrl, agentId, authFetch]);
+    }, [agentBaseUrl, authFetch]);
     const handleReject = useCallback(async (requestId) => {
         if (!agentBaseUrl)
             return;
         setApprovalLoading(requestId);
         try {
-            await authFetch(`${agentBaseUrl}/api/v1/agents/${agentId}/tool-approvals/${requestId}/reject`, {
+            await authFetch(`${agentBaseUrl}/api/v1/tool-approvals/${requestId}/reject`, {
                 method: 'POST',
                 body: JSON.stringify({ reason: 'User rejected' }),
             });
@@ -127,7 +412,7 @@ const AgentGuardrailsInner = ({ onLogout, }) => {
         finally {
             setApprovalLoading(null);
         }
-    }, [agentBaseUrl, agentId, authFetch]);
+    }, [agentBaseUrl, authFetch]);
     // ── Loading / Error ──────────────────────────────────────────────────────
     if (!isReady && runtimeStatus !== 'error') {
         return (_jsxs(Box, { sx: {
@@ -137,21 +422,31 @@ const AgentGuardrailsInner = ({ onLogout, }) => {
                 justifyContent: 'center',
                 height: '100vh',
                 gap: 3,
-            }, children: [_jsx(Spinner, { size: "large" }), _jsx(Text, { sx: { color: 'fg.muted' }, children: runtimeStatus === 'launching'
-                        ? 'Launching runtime for guardrails agent…'
-                        : 'Creating guardrails demo agent…' })] }));
+            }, children: [_jsx(Spinner, { size: "large" }), _jsx(Text, { sx: { color: 'fg.muted' }, children: "Launching guardrails demo agent..." })] }));
     }
     if (runtimeStatus === 'error' || hookError) {
-        return (_jsxs(Box, { sx: {
-                display: 'flex',
-                flexDirection: 'column',
-                alignItems: 'center',
-                justifyContent: 'center',
-                height: '100vh',
-                gap: 3,
-            }, children: [_jsx(AlertIcon, { size: 48 }), _jsx(Text, { sx: { color: 'danger.fg' }, children: hookError || 'Agent failed to start' })] }));
+        return _jsx(ErrorView, { error: hookError, onLogout: onLogout });
     }
-    const costPercent = Math.min((costUsd / COST_LIMIT_USD) * 100, 100);
+    const runCostUsd = Math.max(snapshotRunCostUsd, otelRunCostUsd ?? 0);
+    const cumulativeCostUsd = otelCumulativeCostUsd;
+    const isOverRunBudget = runBudgetUsd != null && runBudgetUsd > 0 && runCostUsd > runBudgetUsd;
+    const runBudgetDisplayUsd = runBudgetUsd != null ? runBudgetUsd.toFixed(2) : '0.00';
+    const overBudgetBanner = isOverRunBudget
+        ? {
+            variant: 'danger',
+            message: `Run budget exceeded: $${runCostUsd.toFixed(4)} / $${runBudgetDisplayUsd}. Start a new run or adjust the budget before sending more messages.`,
+        }
+        : undefined;
+    const budgetForProgress = runBudgetUsd && runBudgetUsd > 0 ? runBudgetUsd : 1;
+    const usagePercentRaw = (runCostUsd / budgetForProgress) * 100;
+    const costPercent = Math.min(usagePercentRaw, 100);
+    const usageSafePercent = Math.min(costPercent, 50);
+    const usageWatchPercent = Math.min(Math.max(costPercent - 50, 0), 30);
+    const usageDangerPercent = Math.min(Math.max(costPercent - 80, 0), 20);
+    const overBudgetPercent = Math.max(usagePercentRaw - 100, 0);
+    const overBudgetAmountUsd = runBudgetUsd != null && runBudgetUsd > 0
+        ? Math.max(runCostUsd - runBudgetUsd, 0)
+        : 0;
     const costColor = costPercent > 80
         ? 'danger.fg'
         : costPercent > 50
@@ -170,14 +465,24 @@ const AgentGuardrailsInner = ({ onLogout, }) => {
                     borderBottom: '1px solid',
                     borderColor: 'border.default',
                     flexShrink: 0,
-                }, children: [_jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [_jsx(ShieldCheckIcon, { size: 16 }), _jsxs(Heading, { as: "h3", sx: { fontSize: 2 }, children: ["Guardrails Demo \u2014 ", podName] })] }), _jsxs(Box, { sx: { flex: 1, maxWidth: 300 }, children: [_jsxs(Box, { sx: {
+                }, children: [isReconnectedAgent && (_jsx(Label, { variant: "secondary", size: "small", children: "Reconnected" })), _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [_jsx(ShieldCheckIcon, { size: 16 }), _jsxs(Heading, { as: "h3", sx: { fontSize: 2 }, children: ["Guardrails Demo \u2014 ", podName] })] }), _jsx(Label, { variant: wsState === 'connected' ? 'success' : 'secondary', children: wsState }), _jsxs(Box, { sx: { flex: 1, maxWidth: 300 }, children: [_jsxs(Box, { sx: {
                                     display: 'flex',
                                     justifyContent: 'space-between',
                                     fontSize: 0,
                                     mb: 1,
-                                }, children: [_jsxs(Text, { sx: { color: costColor, fontWeight: 'semibold' }, children: ["$", costUsd.toFixed(4)] }), _jsxs(Text, { sx: { color: 'fg.muted' }, children: ["/ $", COST_LIMIT_USD.toFixed(2), " limit"] })] }), _jsx(ProgressBar, { progress: costPercent, sx: { height: 6 } })] }), _jsxs(Label, { variant: "secondary", children: [totalTokens.toLocaleString(), " tokens"] }), token && _jsx(UserBadge, { token: token, variant: "small" }), _jsx(Button, { size: "small", variant: "invisible", onClick: onLogout, leadingVisual: SignOutIcon, sx: { color: 'fg.muted' }, children: "Sign out" })] }), approvals.map(req => (_jsx(Flash, { variant: "warning", sx: { mx: 3, mt: 2 }, children: _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 2 }, children: [_jsxs(Text, { sx: { flex: 1, fontSize: 1 }, children: [_jsx("strong", { children: req.tool_name }), " requests approval", req.arguments
-                                    ? ` — ${JSON.stringify(req.arguments).slice(0, 120)}`
-                                    : ''] }), _jsx(Button, { size: "small", variant: "primary", leadingVisual: CheckIcon, onClick: () => handleApprove(req.id), disabled: approvalLoading === req.id, children: "Approve" }), _jsx(Button, { size: "small", variant: "danger", leadingVisual: XIcon, onClick: () => handleReject(req.id), disabled: approvalLoading === req.id, children: "Reject" })] }) }, req.id))), _jsx(Box, { sx: { flex: 1, minHeight: 0 }, children: _jsx(Chat, { protocol: "ag-ui", baseUrl: agentBaseUrl, agentId: agentId, title: "Guardrails Agent", placeholder: "Ask something that triggers tools\u2026", description: "Agent with $5 cost limit and tool approval gates", showHeader: false, showTokenUsage: true, autoFocus: true, height: "100%", runtimeId: podName, historyEndpoint: `${agentBaseUrl}/api/v1/history`, suggestions: [
+                                }, children: [_jsxs(Text, { sx: { color: costColor, fontWeight: 'semibold' }, children: ["$", runCostUsd.toFixed(4)] }), _jsx(Text, { sx: { color: 'fg.muted' }, children: runBudgetUsd != null
+                                            ? ` / $${runBudgetUsd.toFixed(2)} run budget`
+                                            : ' / no run budget' })] }), _jsxs(ProgressBar, { "aria-label": "Run budget usage", "aria-valuenow": Math.max(0, costPercent), sx: { height: 6 }, children: [_jsx(ProgressBar.Item, { progress: usageSafePercent, style: { backgroundColor: 'var(--bgColor-success-emphasis)' }, "aria-label": `Healthy usage ${usageSafePercent.toFixed(1)}%` }), _jsx(ProgressBar.Item, { progress: usageWatchPercent, style: { backgroundColor: 'var(--bgColor-accent-emphasis)' }, "aria-label": `Watch usage ${usageWatchPercent.toFixed(1)}%` }), _jsx(ProgressBar.Item, { progress: usageDangerPercent, style: { backgroundColor: 'var(--bgColor-danger-emphasis)' }, "aria-label": `Critical usage ${usageDangerPercent.toFixed(1)}%` })] }), _jsxs(Box, { sx: {
+                                    display: 'flex',
+                                    alignItems: 'center',
+                                    gap: 2,
+                                    flexWrap: 'wrap',
+                                    mt: 1,
+                                }, role: "presentation", children: [_jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [_jsx(DotFillIcon, { size: 12, fill: "var(--bgColor-success-emphasis)" }), _jsx(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: "0-50%" })] }), _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [_jsx(DotFillIcon, { size: 12, fill: "var(--bgColor-accent-emphasis)" }), _jsx(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: "50-80%" })] }), _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [_jsx(DotFillIcon, { size: 12, fill: "var(--bgColor-danger-emphasis)" }), _jsx(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: "80-100%" })] }), overBudgetPercent > 0 && (_jsxs(Label, { variant: "danger", size: "small", children: ["+", overBudgetPercent.toFixed(1), "% over"] }))] })] }), _jsxs(Label, { variant: otelSampleTimestamp == null ? 'secondary' : 'success', children: ["OTEL ", otelSampleTimestamp == null ? 'waiting' : 'live'] }), cumulativeCostUsd != null && (_jsxs(Label, { variant: "secondary", children: ["Total $", cumulativeCostUsd.toFixed(4)] })), _jsxs(Label, { variant: "secondary", children: [totalTokens.toLocaleString(), " tokens"] })] }), isOverRunBudget && (_jsx(Flash, { variant: "danger", sx: { mx: 3, mt: 2 }, children: _jsxs(Text, { sx: { fontSize: 1 }, children: [_jsx("strong", { children: "Run budget exceeded." }), " Current run cost is $", runCostUsd.toFixed(4), " against a budget of $", runBudgetDisplayUsd, overBudgetAmountUsd > 0
+                            ? ` (over by $${overBudgetAmountUsd.toFixed(4)}).`
+                            : '.', ' ', "Start a new run or increase the run budget before continuing."] }) })), approvals.map(req => (_jsx(Flash, { variant: "warning", sx: { mx: 3, mt: 2 }, children: _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 2 }, children: [_jsxs(Text, { sx: { flex: 1, fontSize: 1 }, children: [_jsx("strong", { children: req.tool_name }), " requests approval", req.tool_args
+                                    ? ` — ${JSON.stringify(req.tool_args).slice(0, 120)}`
+                                    : ''] }), _jsx(Button, { size: "small", variant: "primary", leadingVisual: CheckIcon, onClick: () => handleApprove(req.id), disabled: approvalLoading === req.id, children: "Approve" }), _jsx(Button, { size: "small", variant: "danger", leadingVisual: XIcon, onClick: () => handleReject(req.id), disabled: approvalLoading === req.id, children: "Reject" })] }) }, req.id))), _jsx(Box, { sx: { flex: 1, minHeight: 0 }, children: _jsx(Chat, { protocol: "vercel-ai", baseUrl: agentBaseUrl, agentId: agentId, authToken: chatAuthToken, title: "Guardrails Agent", placeholder: "Ask something that triggers tools\u2026", description: "Cost guardrail with OTEL-backed gauge and manual tool approval gates", showHeader: false, showTokenUsage: true, errorBanner: overBudgetBanner, disableInputPrompt: isOverRunBudget, autoFocus: true, height: "100%", runtimeId: agentId, historyEndpoint: `${agentBaseUrl}/api/v1/history`, suggestions: [
                         { title: 'Update CRM', message: 'Update the CRM records for Q3' },
                         { title: 'Report', message: 'Generate the weekly KPI report' },
                     ], submitOnSuggestionClick: true }) })] }));
@@ -190,7 +495,7 @@ const syncTokenToIamStore = (token) => {
 };
 // ─── Main component with auth gate ─────────────────────────────────────────
 const AgentGuardrailsExample = () => {
-    const { token, setAuth, clearAuth } = useSimpleAuthStore();
+    const { token, clearAuth } = useSimpleAuthStore();
     const hasSynced = useRef(false);
     useEffect(() => {
         if (token && !hasSynced.current) {
@@ -198,11 +503,6 @@ const AgentGuardrailsExample = () => {
             syncTokenToIamStore(token);
         }
     }, [token]);
-    const handleSignIn = useCallback((newToken, handle) => {
-        setAuth(newToken, handle);
-        hasSynced.current = true;
-        syncTokenToIamStore(newToken);
-    }, [setAuth]);
     const handleLogout = useCallback(() => {
         clearAuth();
         hasSynced.current = false;
@@ -211,7 +511,7 @@ const AgentGuardrailsExample = () => {
         });
     }, [clearAuth]);
     if (!token) {
-        return (_jsx(ThemedProvider, { children: _jsx(SignInSimple, { onSignIn: handleSignIn, onApiKeySignIn: apiKey => handleSignIn(apiKey, 'api-key-user'), title: "Guardrails Agent", description: "Sign in to use agents with cost and tool guardrails.", leadingIcon: _jsx(ShieldCheckIcon, { size: 24 }) }) }));
+        return (_jsx(ThemedProvider, { children: _jsx(AuthRequiredView, {}) }));
     }
     return (_jsx(QueryClientProvider, { client: queryClient, children: _jsx(ThemedProvider, { children: _jsx(AgentGuardrailsInner, { onLogout: handleLogout }) }) }));
 };

package/lib/examples/AgentHooksExample.d.ts

@@ -0,0 +1,3 @@
+import React from 'react';
+declare const AgentHooksExample: React.FC;
+export default AgentHooksExample;