@pikku/assistant-ui 0.12.0 → 0.12.1

package/dist/esm/use-pikku-agent-runtime.js

@@ -1,6 +1,332 @@
-import { useMemo, useEffect, useRef, useCallback } from 'react';
-import { useDataStreamRuntime } from '@assistant-ui/react-data-stream';
-const convertDbMessages = (dbMessages) => {
+import { createContext, useContext, useMemo, useRef, useState, useCallback, } from 'react';
+import { useLocalRuntime, } from '@assistant-ui/react';
+export const PikkuApprovalContext = createContext({
+    pendingApprovals: [],
+    handleApproval: () => { },
+});
+export const usePikkuApproval = () => useContext(PikkuApprovalContext);
+async function* parseSSEStream(reader) {
+    const decoder = new TextDecoder();
+    let buffer = '';
+    while (true) {
+        const { done, value } = await reader.read();
+        if (done)
+            break;
+        buffer += decoder.decode(value, { stream: true });
+        const lines = buffer.split('\n');
+        buffer = lines.pop();
+        for (const line of lines) {
+            if (line.startsWith('data: ')) {
+                const data = line.slice(6);
+                if (data) {
+                    try {
+                        yield JSON.parse(data);
+                    }
+                    catch {
+                        // skip unparseable lines
+                    }
+                }
+            }
+        }
+    }
+}
+/**
+ * Shared helper: consume an SSE stream and populate text/toolCalls.
+ * Returns an array of PendingApprovals when the stream requests them, or empty when done.
+ */
+async function processStream(reader, text, toolCalls, yieldContent, onFinish) {
+    const pendingApprovals = [];
+    for await (const event of parseSSEStream(reader)) {
+        switch (event.type) {
+            case 'text-delta':
+                text.value += event.text;
+                break;
+            case 'tool-call': {
+                let parsedArgs = event.args;
+                if (typeof event.args === 'string') {
+                    try {
+                        parsedArgs = JSON.parse(event.args);
+                    }
+                    catch {
+                        parsedArgs = {};
+                    }
+                }
+                toolCalls.push({
+                    type: 'tool-call',
+                    toolCallId: event.toolCallId,
+                    toolName: event.toolName,
+                    args: parsedArgs,
+                });
+                break;
+            }
+            case 'tool-result': {
+                // Skip tool-results that contain __approvalRequired
+                const resultObj = typeof event.result === 'object' ? event.result : null;
+                if (resultObj && '__approvalRequired' in resultObj) {
+                    break;
+                }
+                const tc = toolCalls.find((t) => t.toolCallId === event.toolCallId);
+                if (tc) {
+                    tc.result =
+                        typeof event.result === 'string'
+                            ? event.result
+                            : JSON.stringify(event.result);
+                    if (event.isError) {
+                        tc.isError = true;
+                    }
+                }
+                break;
+            }
+            case 'approval-request':
+                pendingApprovals.push({
+                    toolCallId: event.toolCallId,
+                    toolName: event.toolName,
+                    args: event.args,
+                    reason: event.reason,
+                    runId: event.runId,
+                });
+                break;
+            case 'error':
+                text.value += `\n\nError: ${event.message || event.errorText || 'Unknown error'}`;
+                break;
+            case 'done':
+                onFinish?.();
+                continue;
+        }
+        yieldContent();
+    }
+    return pendingApprovals;
+}
+export function isDeniedResult(result) {
+    if (result == null)
+        return false;
+    try {
+        const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+        return parsed && typeof parsed === 'object' && parsed.approved === false;
+    }
+    catch {
+        return false;
+    }
+}
+export function resolvePikkuToolStatus(status, result) {
+    if (status.type === 'running')
+        return { type: 'running' };
+    if (status.type === 'requires-action')
+        return { type: 'requires-action' };
+    if (isDeniedResult(result))
+        return { type: 'denied' };
+    if (typeof result === 'string' && result.startsWith('Error:'))
+        return { type: 'error' };
+    return { type: 'completed' };
+}
+function buildContent(text, toolCalls) {
+    const content = [];
+    if (text.value)
+        content.push({ type: 'text', text: text.value });
+    content.push(...toolCalls);
+    return content;
+}
+function createPikkuStreamingAdapter(optionsRef, pendingApprovalsRef, approvalDecisionsRef, setPendingApprovalsRef, onFinishRef) {
+    return {
+        async *run({ messages, abortSignal }) {
+            const opts = optionsRef.current;
+            // Check if this run() is a continuation after approval decisions.
+            // assistant-ui calls run() again after addResult provides tool results for ALL tool calls.
+            const pendingApprovals = pendingApprovalsRef.current;
+            if (pendingApprovals.length > 0) {
+                // Read the decisions accumulated by handleApproval() from click handlers.
+                const decisions = approvalDecisionsRef.current;
+                approvalDecisionsRef.current = [];
+                if (decisions.length === 0) {
+                    // No decisions set — shouldn't happen if composer is disabled.
+                    return;
+                }
+                // Clear pending approvals state
+                pendingApprovalsRef.current = [];
+                setPendingApprovalsRef.current([]);
+                // Send /resume for each decision sequentially.
+                // All but the last resume will return quickly (just tool-result + done).
+                // The last resume triggers continuation (next LLM step).
+                let lastText = { value: '' };
+                let lastToolCalls = [];
+                let nextApprovals = [];
+                for (let i = 0; i < decisions.length; i++) {
+                    const decision = decisions[i];
+                    // Find the runId from the matching pending approval
+                    const matchingApproval = pendingApprovals.find((p) => p.toolCallId === decision.toolCallId);
+                    const runId = matchingApproval?.runId ?? pendingApprovals[0]?.runId;
+                    const resumeResponse = await fetch(`${opts.api}/${opts.agentName}/resume`, {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json',
+                            ...opts.headers,
+                        },
+                        body: JSON.stringify({
+                            runId,
+                            toolCallId: decision.toolCallId,
+                            approved: decision.approved,
+                        }),
+                        signal: abortSignal,
+                        credentials: opts.credentials,
+                    });
+                    if (!resumeResponse.ok || !resumeResponse.body) {
+                        const errorText = resumeResponse.body
+                            ? await resumeResponse.text().catch(() => '')
+                            : '';
+                        throw new Error(`Resume failed: ${resumeResponse.status}${errorText ? ` - ${errorText}` : ''}`);
+                    }
+                    const text = { value: '' };
+                    const toolCalls = [];
+                    const reader = resumeResponse.body.getReader();
+                    const streamApprovals = await processStream(reader, text, toolCalls, () => { }, i === decisions.length - 1
+                        ? (onFinishRef.current ?? undefined)
+                        : undefined);
+                    // Keep the last resume's output (it has continuation content)
+                    lastText = text;
+                    lastToolCalls = toolCalls;
+                    if (streamApprovals.length > 0) {
+                        nextApprovals = streamApprovals;
+                    }
+                }
+                // Build content from the last resume's output
+                const content = buildContent(lastText, lastToolCalls);
+                if (nextApprovals.length > 0) {
+                    // More approvals from continuation — show them
+                    pendingApprovalsRef.current = nextApprovals;
+                    setPendingApprovalsRef.current(nextApprovals);
+                    // Add approval tool calls to content
+                    for (const approval of nextApprovals) {
+                        const approvalToolCall = {
+                            type: 'tool-call',
+                            toolCallId: approval.toolCallId,
+                            toolName: approval.toolName,
+                            args: {
+                                ...(typeof approval.args === 'object' && approval.args !== null
+                                    ? approval.args
+                                    : {}),
+                                ...(approval.reason
+                                    ? { __approvalReason: approval.reason }
+                                    : {}),
+                            },
+                        };
+                        const idx = lastToolCalls.findIndex((tc) => tc.toolCallId === approval.toolCallId && !tc.result);
+                        if (idx !== -1) {
+                            lastToolCalls[idx] = approvalToolCall;
+                        }
+                        else {
+                            lastToolCalls.push(approvalToolCall);
+                        }
+                    }
+                    const updatedContent = buildContent(lastText, lastToolCalls);
+                    yield {
+                        content: updatedContent,
+                        status: {
+                            type: 'requires-action',
+                            reason: 'tool-calls',
+                        },
+                    };
+                }
+                else if (content.length > 0) {
+                    yield { content };
+                }
+                return;
+            }
+            // Normal flow: new user message → stream
+            const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
+            if (!lastUserMsg)
+                return;
+            let messageText = '';
+            if (lastUserMsg.content) {
+                for (const part of lastUserMsg.content) {
+                    if ('text' in part && part.type === 'text') {
+                        messageText += part.text;
+                    }
+                }
+            }
+            const response = await fetch(`${opts.api}/${opts.agentName}/stream`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json', ...opts.headers },
+                body: JSON.stringify({
+                    agentName: opts.agentName,
+                    message: messageText,
+                    threadId: opts.threadId,
+                    resourceId: opts.resourceId,
+                    model: opts.model,
+                    temperature: opts.temperature,
+                }),
+                signal: abortSignal,
+                credentials: opts.credentials,
+            });
+            if (!response.ok || !response.body) {
+                throw new Error(`Agent stream failed: ${response.status}`);
+            }
+            const text = { value: '' };
+            const toolCalls = [];
+            let pendingContent = null;
+            const yieldContent = () => {
+                const content = buildContent(text, toolCalls);
+                if (content.length > 0)
+                    pendingContent = content;
+            };
+            const reader = response.body.getReader();
+            const approvals = await processStream(reader, text, toolCalls, yieldContent, onFinishRef.current ?? undefined);
+            if (approvals.length === 0) {
+                // No approval needed — yield final content and done
+                if (pendingContent) {
+                    yield { content: pendingContent };
+                }
+                return;
+            }
+            // Approvals requested: store them for the next run() call
+            pendingApprovalsRef.current = approvals;
+            setPendingApprovalsRef.current(approvals);
+            // Each approval tool call needs to be shown without a result
+            // so assistant-ui renders them as requires-action
+            for (const approval of approvals) {
+                const approvalToolCall = {
+                    type: 'tool-call',
+                    toolCallId: approval.toolCallId,
+                    toolName: approval.toolName,
+                    args: {
+                        ...(typeof approval.args === 'object' && approval.args !== null
+                            ? approval.args
+                            : {}),
+                        ...(approval.reason ? { __approvalReason: approval.reason } : {}),
+                    },
+                };
+                // Replace the existing tool call (if any) with the approval version
+                const parentIdx = toolCalls.findIndex((tc) => tc.toolCallId === approval.toolCallId && !tc.result);
+                if (parentIdx !== -1) {
+                    toolCalls[parentIdx] = approvalToolCall;
+                }
+                else {
+                    toolCalls.push(approvalToolCall);
+                }
+            }
+            // Remove any forwarded sub-agent tool calls that duplicate approval tool names
+            const approvalToolCallIds = new Set(approvals.map((a) => a.toolCallId));
+            const approvalToolNames = new Set(approvals.map((a) => a.toolName));
+            for (let i = toolCalls.length - 1; i >= 0; i--) {
+                if (approvalToolNames.has(toolCalls[i].toolName) &&
+                    !approvalToolCallIds.has(toolCalls[i].toolCallId)) {
+                    toolCalls.splice(i, 1);
+                }
+            }
+            const content = buildContent(text, toolCalls);
+            yield {
+                content,
+                status: {
+                    type: 'requires-action',
+                    reason: 'tool-calls',
+                },
+            };
+            // Generator returns here. assistant-ui will show the approval UI for each tool call.
+            // When the user clicks Approve/Deny on ALL tools → handleApproval accumulates decisions →
+            // addResult for each → run() is called again when all have results.
+        },
+    };
+}
+export const convertDbMessages = (dbMessages) => {
     const result = [];
     let currentAssistant = null;
     for (const msg of dbMessages) {
@@ -97,67 +423,221 @@ const convertDbMessages = (dbMessages) => {
     return result;
 };
 export function usePikkuAgentRuntime(options) {
-    const { api, threadId = null, initialMessages: rawInitialMessages, onThreadCreated, onFinish, onApprovalRequest, credentials, headers, } = options;
-    const threadIdRef = useRef(threadId);
-    threadIdRef.current = threadId;
-    const justCreatedThreadRef = useRef(false);
-    const bodyFn = useCallback(() => {
-        let currentThreadId = threadIdRef.current;
-        if (!currentThreadId) {
-            currentThreadId = crypto.randomUUID();
-            justCreatedThreadRef.current = true;
-            onThreadCreated?.(currentThreadId);
-        }
-        return { threadId: currentThreadId };
-    }, [onThreadCreated]);
-    const onData = useCallback((event) => {
-        if (event.name === 'approval-request') {
-            const approval = event.data;
-            onApprovalRequest?.({
-                toolCallId: approval.toolCallId,
+    const [pendingApprovals, setPendingApprovals] = useState([]);
+    const optionsRef = useRef(options);
+    optionsRef.current = options;
+    const onFinishRef = useRef(options.onFinish);
+    onFinishRef.current = options.onFinish;
+    const pendingApprovalsRef = useRef([]);
+    const approvalDecisionsRef = useRef([]);
+    const setPendingApprovalsRef = useRef(setPendingApprovals);
+    setPendingApprovalsRef.current = setPendingApprovals;
+    const adapter = useMemo(() => createPikkuStreamingAdapter(optionsRef, pendingApprovalsRef, approvalDecisionsRef, setPendingApprovalsRef, onFinishRef), []);
+    const initialMessages = useMemo(() => options.initialMessages
+        ? convertDbMessages(options.initialMessages)
+        : undefined, [options.initialMessages]);
+    const runtime = useLocalRuntime(adapter, { initialMessages });
+    // handleApproval is called from the Approve/Deny button click handler.
+    // It accumulates the decision in the ref. assistant-ui will call run()
+    // when ALL tool calls have results (via addResult).
+    const handleApproval = useCallback((toolCallId, approved) => {
+        approvalDecisionsRef.current.push({ toolCallId, approved });
+    }, []);
+    const isAwaitingApproval = pendingApprovals.length > 0;
+    return {
+        runtime,
+        pendingApprovals,
+        isAwaitingApproval,
+        handleApproval,
+    };
+}
+function createPikkuNonStreamingAdapter(optionsRef, pendingApprovalsRef, approvalDecisionsRef, setPendingApprovalsRef, onFinishRef) {
+    return {
+        async *run({ messages, abortSignal }) {
+            const opts = optionsRef.current;
+            // Continuation after approval decisions
+            const pendingApprovals = pendingApprovalsRef.current;
+            if (pendingApprovals.length > 0) {
+                const decisions = approvalDecisionsRef.current;
+                approvalDecisionsRef.current = [];
+                if (decisions.length === 0)
+                    return;
+                pendingApprovalsRef.current = [];
+                setPendingApprovalsRef.current([]);
+                // Resume uses SSE (same as streaming mode)
+                let lastText = { value: '' };
+                let lastToolCalls = [];
+                let nextApprovals = [];
+                for (let i = 0; i < decisions.length; i++) {
+                    const decision = decisions[i];
+                    const matchingApproval = pendingApprovals.find((p) => p.toolCallId === decision.toolCallId);
+                    const runId = matchingApproval?.runId ?? pendingApprovals[0]?.runId;
+                    const resumeResponse = await fetch(`${opts.api}/${opts.agentName}/resume`, {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json',
+                            ...opts.headers,
+                        },
+                        body: JSON.stringify({
+                            runId,
+                            toolCallId: decision.toolCallId,
+                            approved: decision.approved,
+                        }),
+                        signal: abortSignal,
+                        credentials: opts.credentials,
+                    });
+                    if (!resumeResponse.ok || !resumeResponse.body) {
+                        const errorText = resumeResponse.body
+                            ? await resumeResponse.text().catch(() => '')
+                            : '';
+                        throw new Error(`Resume failed: ${resumeResponse.status}${errorText ? ` - ${errorText}` : ''}`);
+                    }
+                    const text = { value: '' };
+                    const toolCalls = [];
+                    const reader = resumeResponse.body.getReader();
+                    const streamApprovals = await processStream(reader, text, toolCalls, () => { }, i === decisions.length - 1
+                        ? (onFinishRef.current ?? undefined)
+                        : undefined);
+                    lastText = text;
+                    lastToolCalls = toolCalls;
+                    if (streamApprovals.length > 0) {
+                        nextApprovals = streamApprovals;
+                    }
+                }
+                const content = buildContent(lastText, lastToolCalls);
+                if (nextApprovals.length > 0) {
+                    pendingApprovalsRef.current = nextApprovals;
+                    setPendingApprovalsRef.current(nextApprovals);
+                    for (const approval of nextApprovals) {
+                        const approvalToolCall = {
+                            type: 'tool-call',
+                            toolCallId: approval.toolCallId,
+                            toolName: approval.toolName,
+                            args: {
+                                ...(typeof approval.args === 'object' && approval.args !== null
+                                    ? approval.args
+                                    : {}),
+                                ...(approval.reason
+                                    ? { __approvalReason: approval.reason }
+                                    : {}),
+                            },
+                        };
+                        const idx = lastToolCalls.findIndex((tc) => tc.toolCallId === approval.toolCallId && !tc.result);
+                        if (idx !== -1) {
+                            lastToolCalls[idx] = approvalToolCall;
+                        }
+                        else {
+                            lastToolCalls.push(approvalToolCall);
+                        }
+                    }
+                    const updatedContent = buildContent(lastText, lastToolCalls);
+                    yield {
+                        content: updatedContent,
+                        status: {
+                            type: 'requires-action',
+                            reason: 'tool-calls',
+                        },
+                    };
+                }
+                else if (content.length > 0) {
+                    yield { content };
+                }
+                return;
+            }
+            // Normal flow: new user message → non-streaming POST
+            const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
+            if (!lastUserMsg)
+                return;
+            let messageText = '';
+            if (lastUserMsg.content) {
+                for (const part of lastUserMsg.content) {
+                    if ('text' in part && part.type === 'text') {
+                        messageText += part.text;
+                    }
+                }
+            }
+            const response = await fetch(`${opts.api}/${opts.agentName}`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json', ...opts.headers },
+                body: JSON.stringify({
+                    agentName: opts.agentName,
+                    message: messageText,
+                    threadId: opts.threadId,
+                    resourceId: opts.resourceId,
+                    model: opts.model,
+                    temperature: opts.temperature,
+                }),
+                signal: abortSignal,
+                credentials: opts.credentials,
             });
-        }
-    }, [onApprovalRequest]);
-    const onFinishCb = useCallback(() => {
-        onFinish?.();
-    }, [onFinish]);
-    const initialMessages = useMemo(() => (rawInitialMessages ? convertDbMessages(rawInitialMessages) : []), [rawInitialMessages]);
-    const runtime = useDataStreamRuntime({
-        api,
-        protocol: 'ui-message-stream',
-        body: bodyFn,
-        onData,
-        onFinish: onFinishCb,
-        initialMessages,
-        credentials,
-        headers,
-    });
-    const prevThreadIdRef = useRef(threadId);
-    const hasResetRef = useRef(false);
-    useEffect(() => {
-        if (prevThreadIdRef.current !== threadId) {
-            prevThreadIdRef.current = threadId;
-            if (justCreatedThreadRef.current) {
-                justCreatedThreadRef.current = false;
-                hasResetRef.current = true;
+            if (!response.ok) {
+                throw new Error(`Agent run failed: ${response.status}`);
+            }
+            const json = await response.json();
+            if (json.status === 'suspended' && json.pendingApprovals?.length > 0) {
+                const approvals = json.pendingApprovals;
+                pendingApprovalsRef.current = approvals;
+                setPendingApprovalsRef.current(approvals);
+                const toolCalls = approvals.map((approval) => ({
+                    type: 'tool-call',
+                    toolCallId: approval.toolCallId,
+                    toolName: approval.toolName,
+                    args: {
+                        ...(typeof approval.args === 'object' && approval.args !== null
+                            ? approval.args
+                            : {}),
+                        ...(approval.reason ? { __approvalReason: approval.reason } : {}),
+                    },
+                }));
+                const content = [];
+                if (json.result) {
+                    content.push({ type: 'text', text: json.result });
+                }
+                content.push(...toolCalls);
+                yield {
+                    content,
+                    status: {
+                        type: 'requires-action',
+                        reason: 'tool-calls',
+                    },
+                };
                 return;
             }
-            hasResetRef.current = false;
-            if (rawInitialMessages) {
-                ;
-                runtime.thread.reset(convertDbMessages(rawInitialMessages));
-                hasResetRef.current = true;
+            // No approvals — yield complete content
+            onFinishRef.current?.();
+            const content = [];
+            if (json.result) {
+                content.push({ type: 'text', text: String(json.result) });
             }
-            else {
-                ;
-                runtime.thread.reset([]);
+            if (content.length > 0) {
+                yield { content };
             }
-        }
-        else if (!hasResetRef.current && rawInitialMessages) {
-            ;
-            runtime.thread.reset(convertDbMessages(rawInitialMessages));
-            hasResetRef.current = true;
-        }
-    }, [threadId, rawInitialMessages, runtime]);
-    return runtime;
+        },
+    };
+}
+export function usePikkuAgentNonStreamingRuntime(options) {
+    const [pendingApprovals, setPendingApprovals] = useState([]);
+    const optionsRef = useRef(options);
+    optionsRef.current = options;
+    const onFinishRef = useRef(options.onFinish);
+    onFinishRef.current = options.onFinish;
+    const pendingApprovalsRef = useRef([]);
+    const approvalDecisionsRef = useRef([]);
+    const setPendingApprovalsRef = useRef(setPendingApprovals);
+    setPendingApprovalsRef.current = setPendingApprovals;
+    const adapter = useMemo(() => createPikkuNonStreamingAdapter(optionsRef, pendingApprovalsRef, approvalDecisionsRef, setPendingApprovalsRef, onFinishRef), []);
+    const initialMessages = useMemo(() => options.initialMessages
+        ? convertDbMessages(options.initialMessages)
+        : undefined, [options.initialMessages]);
+    const runtime = useLocalRuntime(adapter, { initialMessages });
+    const handleApproval = useCallback((toolCallId, approved) => {
+        approvalDecisionsRef.current.push({ toolCallId, approved });
+    }, []);
+    const isAwaitingApproval = pendingApprovals.length > 0;
+    return {
+        runtime,
+        pendingApprovals,
+        isAwaitingApproval,
+        handleApproval,
+    };
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pikku/assistant-ui",
-  "version": "0.12.0",
+  "version": "0.12.1",
   "author": "yasser.fadl@gmail.com",
   "license": "MIT",
   "type": "module",
@@ -20,7 +20,7 @@
   },
   "dependencies": {
     "@assistant-ui/react": "^0.12.12",
-    "@assistant-ui/react-data-stream": "^0.12.5"
+    "react-markdown": "^10.1.0"
   },
   "peerDependencies": {
     "react": "^18 || ^19",

package/src/index.ts

@@ -1,4 +1,18 @@
-export { usePikkuAgentRuntime } from './use-pikku-agent-runtime.js'
-export type { PikkuAgentRuntimeOptions } from './use-pikku-agent-runtime.js'
+export {
+  usePikkuAgentRuntime,
+  usePikkuAgentNonStreamingRuntime,
+  PikkuApprovalContext,
+  usePikkuApproval,
+  convertDbMessages,
+  isDeniedResult,
+  resolvePikkuToolStatus,
+} from './use-pikku-agent-runtime.js'
+export type {
+  PikkuAgentRuntimeOptions,
+  PendingApproval,
+  PikkuApprovalContextValue,
+  PikkuToolStatusType,
+  PikkuToolStatus,
+} from './use-pikku-agent-runtime.js'
 export { PikkuAgentChat } from './pikku-agent-chat.js'
 export type { PikkuAgentChatProps } from './pikku-agent-chat.js'