@datalayer/agent-runtimes 1.0.3 → 1.0.4

package/lib/examples/AgentToolApprovalsExample.js

@@ -7,8 +7,9 @@ import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
 import { useCallback, useEffect, useMemo, useRef, useState, } from 'react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { Box } from '@datalayer/primer-addons';
+import { ErrorView } from './components';
 import { Button, Heading, SegmentedControl, Spinner, Text, } from '@primer/react';
-import { AlertIcon, SignOutIcon, ToolsIcon } from '@primer/octicons-react';
+import { SignOutIcon, ToolsIcon } from '@primer/octicons-react';
 import { useCoreStore } from '@datalayer/core';
 import { DEFAULT_SERVICE_URLS } from '@datalayer/core/lib/api/constants';
 import { useSimpleAuthStore } from '@datalayer/core/lib/views/otel';
@@ -111,6 +112,8 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
     const [toolApprovalState, setToolApprovalState] = useState({});
     const [wsState, setWsState] = useState('closed');
     const createdApprovalSignatures = useRef(new Set());
+    const toolRespondersRef = useRef(new Map());
+    const respondedToolCallsRef = useRef(new Set());
     const resolveLocalRef = useRef(async () => null);
     const authFetchRef = useRef((url, opts) => fetch(url, opts));
     const chatAuthToken = token === null ? undefined : token;
@@ -126,6 +129,29 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
             ...(opts.headers ?? {}),
         },
     }), [token]);
+    const emitServerToolDecision = useCallback((toolName, toolArgs, approved, approvalId, message) => {
+        const signature = approvalSignature(toolName, toolArgs);
+        const responder = toolRespondersRef.current.get(signature);
+        if (!responder?.respond) {
+            return false;
+        }
+        if (respondedToolCallsRef.current.has(responder.toolCallId)) {
+            return false;
+        }
+        setToolApprovalState(prev => ({
+            ...prev,
+            [responder.toolCallId]: approved ? 'approved' : 'denied',
+        }));
+        respondedToolCallsRef.current.add(responder.toolCallId);
+        responder.respond({
+            type: 'tool-approval-decision',
+            approved,
+            approvalId,
+            toolName: responder.toolName || toolName,
+            ...(message ? { message } : {}),
+        });
+        return true;
+    }, []);
     useEffect(() => {
         let isCancelled = false;
         const createLocalAgent = async () => {
@@ -315,6 +341,8 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
                     const action = wsEvent === 'tool_approval_approved' ? 'approve' : 'reject';
                     void (async () => {
                         try {
+                            const approved = wsEvent === 'tool_approval_approved';
+                            emitServerToolDecision(approval.tool_name, approval.tool_args ?? {}, approved, approval.id, approval.note);
                             const localMatch = await resolveLocalRef.current(approval.tool_name, approval.tool_args ?? {});
                             if (localMatch) {
                                 const localBaseUrl = `${agentBaseUrl}/api/v1/tool-approvals/${localMatch.id}/${action}`;
@@ -356,6 +384,7 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
         chatAuthToken,
         agentId,
         pollApprovals,
+        emitServerToolDecision,
     ]);
     // Keep refs in sync so the WS handler always has the latest functions.
     useEffect(() => {
@@ -384,8 +413,9 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
                 throw new Error(errorText ||
                     `Failed to approve request (${response.status} ${response.statusText})`);
             }
-            // In server mode the WS event will remove from queue and bridge to
-            // the local runtime — don't do it here to keep flow server-driven.
+            // In server mode, UI/continuation updates are websocket-driven only.
+            // The click only sends approve to the server; local chat reacts when
+            // tool_approval_approved is received.
             if (mode !== 'server') {
                 setApprovals(prev => prev.filter(a => a.id !== requestId));
                 void pollApprovals();
@@ -402,7 +432,14 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
         finally {
             setApprovalLoading(null);
         }
-    }, [approvals, mode, aiAgentsBaseUrl, agentBaseUrl, authFetch, pollApprovals]);
+    }, [
+        approvals,
+        mode,
+        aiAgentsBaseUrl,
+        agentBaseUrl,
+        authFetch,
+        pollApprovals,
+    ]);
     const reject = useCallback(async (requestId, note) => {
         setApprovalLoading(requestId);
         setApprovalError(null);
@@ -423,8 +460,9 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
                 throw new Error(errorText ||
                     `Failed to reject request (${response.status} ${response.statusText})`);
             }
-            // In server mode the WS event will remove from queue and bridge to
-            // the local runtime — don't do it here to keep flow server-driven.
+            // In server mode, UI/continuation updates are websocket-driven only.
+            // The click only sends reject to the server; local chat reacts when
+            // tool_approval_rejected is received.
             if (mode !== 'server') {
                 setApprovals(prev => prev.filter(a => a.id !== requestId));
                 void pollApprovals();
@@ -441,7 +479,14 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
         finally {
             setApprovalLoading(null);
         }
-    }, [approvals, mode, aiAgentsBaseUrl, agentBaseUrl, authFetch, pollApprovals]);
+    }, [
+        approvals,
+        mode,
+        aiAgentsBaseUrl,
+        agentBaseUrl,
+        authFetch,
+        pollApprovals,
+    ]);
     const ensureServerApproval = useCallback(async (toolName, args) => {
         if (mode !== 'server' || !aiAgentsBaseUrl || !isReady) {
             return;
@@ -508,19 +553,46 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
             return stableStringify(approval.tool_args ?? {}) === argsSig;
         }) || null);
     }, [approvals]);
-    const handleToolLevelApprove = useCallback(async (toolCallId, requestId) => {
+    const handleToolLevelApprove = useCallback(async (toolCallId, requestId, toolName, respond) => {
         const ok = await approve(requestId, 'Approved from tool message card');
         if (ok) {
-            setToolApprovalState(prev => ({ ...prev, [toolCallId]: 'approved' }));
+            if (mode !== 'server') {
+                setToolApprovalState(prev => ({ ...prev, [toolCallId]: 'approved' }));
+                respond?.({
+                    type: 'tool-approval-decision',
+                    approved: true,
+                    approvalId: requestId,
+                    toolName,
+                });
+            }
         }
-    }, [approve]);
-    const handleToolLevelDeny = useCallback(async (toolCallId, requestId) => {
+    }, [approve, mode]);
+    const handleToolLevelDeny = useCallback(async (toolCallId, requestId, toolName, respond) => {
         const ok = await reject(requestId, 'Rejected from tool message card');
         if (ok) {
-            setToolApprovalState(prev => ({ ...prev, [toolCallId]: 'denied' }));
+            if (mode !== 'server') {
+                setToolApprovalState(prev => ({ ...prev, [toolCallId]: 'denied' }));
+                respond?.({
+                    type: 'tool-approval-decision',
+                    approved: false,
+                    approvalId: requestId,
+                    toolName,
+                });
+            }
+        }
+    }, [reject, mode]);
+    const renderToolResult = useCallback(({ toolCallId, toolName, args, result, status, error, respond }) => {
+        const signature = approvalSignature(toolName, args);
+        if (respond && status === 'inProgress') {
+            toolRespondersRef.current.set(signature, {
+                toolCallId,
+                toolName,
+                respond,
+            });
+        }
+        else if (status === 'complete' || status === 'error') {
+            toolRespondersRef.current.delete(signature);
         }
-    }, [reject]);
-    const renderToolResult = useCallback(({ toolCallId, toolName, args, result, status, error }) => {
         const matchedApproval = findMatchingApproval(toolName, args);
         const resultObject = result && typeof result === 'object'
             ? result
@@ -539,9 +611,9 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
                 ? 'pending'
                 : undefined);
         return (_jsx(ToolCallDisplay, { toolCallId: toolCallId, toolName: toolName, args: args, result: result, status: status, error: error, approvalRequired: !!approvalState, approvalState: approvalState, approvalLoading: loadingThisApproval, onApprove: matchedApproval
-                ? () => void handleToolLevelApprove(toolCallId, matchedApproval.id)
+                ? () => void handleToolLevelApprove(toolCallId, matchedApproval.id, toolName, respond)
                 : undefined, onDeny: matchedApproval
-                ? () => void handleToolLevelDeny(toolCallId, matchedApproval.id)
+                ? () => void handleToolLevelDeny(toolCallId, matchedApproval.id, toolName, respond)
                 : undefined }));
     }, [
         mode,
@@ -563,14 +635,7 @@ const AgentToolApprovalsInner = ({ onLogout, }) => {
             }, children: [_jsx(Spinner, { size: "large" }), _jsx(Text, { sx: { color: 'fg.muted' }, children: "Launching tool approval demo..." })] }));
     }
     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 serverPanel = mode === 'server' ? (_jsxs(Box, { sx: {
             width: 320,