@truefoundry/assistant-ui-runtime 0.1.0-rc.1

package/src/toolApproval.test.ts

@@ -0,0 +1,137 @@
+import { describe, expect, it } from "vitest";
+
+import {
+    applyApprovalDecisionsToMessage,
+    collectApprovalInputs,
+    messageHasPendingApprovals,
+} from "./toolApproval.js";
+
+describe("toolApproval", () => {
+    it("detects pending nested approvals", () => {
+        const message = {
+            id: "a1",
+            role: "assistant" as const,
+            content: [
+                {
+                    type: "tool-call" as const,
+                    toolCallId: "tc-root",
+                    toolName: "create_sub_agent",
+                    args: {},
+                    argsText: "{}",
+                    messages: [
+                        {
+                            id: "sub",
+                            role: "assistant" as const,
+                            content: [
+                                {
+                                    type: "tool-call" as const,
+                                    toolCallId: "tc-sub",
+                                    toolName: "bash",
+                                    args: {},
+                                    argsText: "{}",
+                                    approval: { id: "approval-sub" },
+                                },
+                            ],
+                            status: { type: "requires-action" as const, reason: "tool-calls" as const },
+                            createdAt: new Date(),
+                            metadata: {
+                                unstable_state: null,
+                                unstable_annotations: [],
+                                unstable_data: [],
+                                steps: [],
+                                custom: {},
+                            },
+                        },
+                    ],
+                },
+            ],
+            status: { type: "requires-action" as const, reason: "tool-calls" as const },
+            createdAt: new Date(),
+            metadata: {
+                unstable_state: null,
+                unstable_annotations: [],
+                unstable_data: [],
+                steps: [],
+                custom: {},
+            },
+        };
+
+        expect(messageHasPendingApprovals(message)).toBe(true);
+    });
+
+    it("collects decided approvals for sdk resume", () => {
+        const pending = {
+            id: "a1",
+            role: "assistant" as const,
+            content: [
+                {
+                    type: "tool-call" as const,
+                    toolCallId: "approval-1",
+                    toolName: "bash",
+                    args: {},
+                    argsText: "{}",
+                    approval: { id: "approval-1" },
+                },
+            ],
+            status: { type: "requires-action" as const, reason: "tool-calls" as const },
+            createdAt: new Date(),
+            metadata: {
+                unstable_state: null,
+                unstable_annotations: [],
+                unstable_data: [],
+                steps: [],
+                custom: {},
+            },
+        };
+
+        const decided = applyApprovalDecisionsToMessage(pending, {
+            approvalId: "approval-1",
+            approved: true,
+        });
+
+        expect(messageHasPendingApprovals(decided)).toBe(false);
+        const inputs = collectApprovalInputs(decided, "main");
+        expect(inputs).toEqual([
+            {
+                type: "user.tool_approval",
+                threadId: "main",
+                toolCallId: "approval-1",
+                approval: { status: "allow" },
+            },
+        ]);
+    });
+
+    it("collects root approvals when custom thread id metadata is present", () => {
+        const pending = {
+            id: "turn-1-assistant",
+            role: "assistant" as const,
+            content: [
+                {
+                    type: "tool-call" as const,
+                    toolCallId: "approval-1",
+                    toolName: "bash",
+                    args: {},
+                    argsText: "{}",
+                    approval: { id: "approval-1" },
+                },
+            ],
+            status: { type: "requires-action" as const, reason: "tool-calls" as const },
+            createdAt: new Date(),
+            metadata: {
+                unstable_state: null,
+                unstable_annotations: [],
+                unstable_data: [],
+                steps: [],
+                custom: { toolApprovalThreadId: "main" },
+            },
+        };
+
+        const decided = applyApprovalDecisionsToMessage(pending, {
+            approvalId: "approval-1",
+            approved: true,
+        });
+
+        expect(messageHasPendingApprovals(decided)).toBe(false);
+        expect(collectApprovalInputs(decided, "main")).toHaveLength(1);
+    });
+});

package/src/toolApproval.ts

@@ -0,0 +1,459 @@
+import type {
+    MessageStatus,
+    ThreadAssistantMessage,
+    ThreadAssistantMessagePart,
+    ThreadMessage,
+} from "@assistant-ui/core";
+import type {
+    ToolApprovalRequiredEvent,
+    Turn,
+    UserToolApprovalEvent,
+} from "truefoundry-gateway-sdk/agents";
+
+import { ROOT_THREAD_ID } from "./constants.js";
+import type { ToolApprovalMessageCustomMetadata } from "./messageCustomMetadata.js";
+import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
+
+export { ROOT_THREAD_ID } from "./constants.js";
+
+export type StoredApprovalDecision = {
+    approved: boolean;
+    reason?: string;
+};
+
+export const TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY = "toolApprovalThreadId";
+
+export type RespondToToolApprovalOptions = {
+    approvalId: string;
+    approved: boolean;
+    optionId?: string;
+    reason?: string;
+};
+
+type ApprovalDecision = UserToolApprovalEvent["approval"];
+
+type ToolCallPart = Extract<
+    ThreadMessage["content"][number],
+    { type: "tool-call" }
+>;
+
+type AssistantToolCallPart = Extract<
+    ThreadAssistantMessagePart,
+    { type: "tool-call" }
+>;
+
+export function hasPendingToolApproval(
+    approval: ToolCallPart["approval"] | undefined,
+): boolean {
+    return (
+        approval != null &&
+        approval.approved === undefined &&
+        approval.resolution === undefined
+    );
+}
+
+function applyApprovalDecisionToToolCall(
+    part: AssistantToolCallPart,
+    options: RespondToToolApprovalOptions,
+): AssistantToolCallPart {
+    const { approved, optionId, reason } = options;
+    const targetApproval = part.approval!;
+    const approval = {
+        ...targetApproval,
+        approved,
+        ...(optionId != null ? { optionId } : {}),
+        ...(reason != null ? { reason } : {}),
+    };
+    if (approved) {
+        return { ...part, approval };
+    }
+    return {
+        ...part,
+        approval,
+        result: { error: reason || "Tool approval denied" },
+        isError: true,
+    };
+}
+
+function updateToolApprovalInContent(
+    content: readonly ThreadAssistantMessagePart[],
+    options: RespondToToolApprovalOptions,
+): { content: readonly ThreadAssistantMessagePart[]; found: boolean } {
+    let found = false;
+    const newContent = content.map((part) => {
+        if (part.type !== "tool-call") {
+            return part;
+        }
+
+        if (part.approval?.id === options.approvalId) {
+            found = true;
+            return applyApprovalDecisionToToolCall(part, options);
+        }
+
+        if (part.messages == null) {
+            return part;
+        }
+
+        const messages = part.messages.map((message) => {
+            if (message.role !== "assistant") {
+                return message;
+            }
+            const nested = updateToolApprovalInContent(message.content, options);
+            if (!nested.found) {
+                return message;
+            }
+            found = true;
+            return { ...message, content: nested.content };
+        });
+        return { ...part, messages };
+    });
+
+    return { content: newContent, found };
+}
+
+export function applyApprovalDecisionsToMessage(
+    message: ThreadAssistantMessage,
+    options: RespondToToolApprovalOptions,
+): ThreadAssistantMessage {
+    const { content } = updateToolApprovalInContent(message.content, options);
+    return { ...message, content: [...content] };
+}
+
+export function toolApprovalStatus(): MessageStatus {
+    return { type: "requires-action", reason: "tool-calls" };
+}
+
+export function toolApprovalMessageCustom(
+    threadId: string,
+): ToolApprovalMessageCustomMetadata {
+    return {
+        [TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY]:
+            threadId === ROOT_THREAD_ID ? ROOT_THREAD_ID : threadId,
+    };
+}
+
+export function getToolApprovalThreadId(
+    message: ThreadMessage | undefined,
+): string | undefined {
+    if (message?.role !== "assistant") {
+        return undefined;
+    }
+    const threadId = message.metadata.custom[TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY];
+    return typeof threadId === "string" ? threadId : undefined;
+}
+
+export function findApprovalRequiredInTurn(
+    turn: Pick<Turn, "state">,
+): ToolApprovalRequiredEvent | undefined {
+    if (turn.state.status !== "done") {
+        return undefined;
+    }
+    return turn.state.requiredActions?.find(
+        (action): action is ToolApprovalRequiredEvent =>
+            action.type === "tool.approval_required",
+    );
+}
+
+function toolCallPartHasPendingApproval(part: ToolCallPart): boolean {
+    return hasPendingToolApproval(part.approval);
+}
+
+function nestedMessagesHavePendingApprovals(
+    messages: readonly ThreadMessage[],
+): boolean {
+    for (const message of messages) {
+        if (messageHasPendingApprovals(message)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+export function messageHasPendingApprovals(message: ThreadMessage | undefined): boolean {
+    if (message?.role !== "assistant") {
+        return false;
+    }
+    for (const part of message.content) {
+        if (part.type !== "tool-call") {
+            continue;
+        }
+        if (toolCallPartHasPendingApproval(part)) {
+            return true;
+        }
+        if (part.messages != null && nestedMessagesHavePendingApprovals(part.messages)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+function walkAssistantToolCallParts(
+    content: readonly ThreadAssistantMessagePart[],
+    visit: (part: AssistantToolCallPart) => void,
+): void {
+    for (const part of content) {
+        if (part.type !== "tool-call") {
+            continue;
+        }
+        visit(part);
+        if (part.messages == null) {
+            continue;
+        }
+        for (const message of part.messages) {
+            if (message.role === "assistant") {
+                walkAssistantToolCallParts(message.content, visit);
+            }
+        }
+    }
+}
+
+export function collectDecidedApprovalsFromContent(
+    content: readonly ThreadAssistantMessagePart[],
+): Map<string, StoredApprovalDecision> {
+    const decisions = new Map<string, StoredApprovalDecision>();
+    walkAssistantToolCallParts(content, (part) => {
+        const { approval } = part;
+        if (approval?.id == null || approval.approved === undefined) {
+            return;
+        }
+        decisions.set(approval.id, {
+            approved: approval.approved,
+            ...(approval.reason != null ? { reason: approval.reason } : {}),
+        });
+    });
+    return decisions;
+}
+
+function applyApprovalDecisionsToMessages(
+    messages: readonly ThreadMessage[],
+    decisions: ReadonlyMap<string, StoredApprovalDecision>,
+): ThreadMessage[] {
+    return messages.map((message) => {
+        if (message.role !== "assistant") {
+            return message;
+        }
+        return {
+            ...message,
+            content: applyApprovalDecisionsToContent(message.content, decisions),
+        };
+    });
+}
+
+export function applyApprovalDecisionsToContent(
+    content: readonly ThreadAssistantMessagePart[],
+    decisions: ReadonlyMap<string, StoredApprovalDecision>,
+): ThreadAssistantMessagePart[] {
+    return content.map((part) => {
+        if (part.type !== "tool-call") {
+            return part;
+        }
+
+        const decision =
+            part.approval?.id != null ? decisions.get(part.approval.id) : undefined;
+        let nextPart: AssistantToolCallPart = part;
+
+        if (
+            decision != null &&
+            part.approval != null &&
+            part.approval.approved === undefined
+        ) {
+            nextPart = applyApprovalDecisionToToolCall(part, {
+                approvalId: part.approval.id,
+                approved: decision.approved,
+                reason: decision.reason,
+            });
+        }
+
+        if (nextPart.messages == null) {
+            return nextPart;
+        }
+
+        return {
+            ...nextPart,
+            messages: applyApprovalDecisionsToMessages(nextPart.messages, decisions),
+        };
+    });
+}
+
+export function mergeDecidedApprovalsIntoContent(
+    incoming: readonly ThreadAssistantMessagePart[],
+    existing: readonly ThreadAssistantMessagePart[],
+): ThreadAssistantMessagePart[] {
+    const decided = collectDecidedApprovalsFromContent(existing);
+    if (decided.size === 0) {
+        return [...incoming];
+    }
+    return applyApprovalDecisionsToContent(incoming, decided);
+}
+
+export function extractToolApprovalsFromTurnInput(
+    input: Turn["input"] | undefined,
+): UserToolApprovalEvent[] {
+    const events: UserToolApprovalEvent[] = [];
+    for (const item of input ?? []) {
+        if (item.type === "user.tool_approval") {
+            events.push(item);
+        }
+    }
+    return events;
+}
+
+export function collectSubsequentApprovalDecisions(
+    turns: readonly Pick<Turn, "input">[],
+    fromIndex: number,
+): Map<string, StoredApprovalDecision> {
+    const decisions = new Map<string, StoredApprovalDecision>();
+
+    for (let index = fromIndex + 1; index < turns.length; index++) {
+        const input = turns[index]?.input ?? [];
+        if (input.some((item) => item.type === "user.message")) {
+            break;
+        }
+
+        for (const event of extractToolApprovalsFromTurnInput(input)) {
+            decisions.set(event.toolCallId, {
+                approved: event.approval.status === "allow",
+                ...(event.approval.status === "deny" && event.approval.reason != null
+                    ? { reason: event.approval.reason }
+                    : {}),
+            });
+        }
+    }
+
+    return decisions;
+}
+
+export function collectApprovalDecisionsFromTurnInput(
+    input: Turn["input"] | undefined,
+): Map<string, StoredApprovalDecision> {
+    const decisions = new Map<string, StoredApprovalDecision>();
+    for (const event of extractToolApprovalsFromTurnInput(input)) {
+        decisions.set(event.toolCallId, {
+            approved: event.approval.status === "allow",
+            ...(event.approval.status === "deny" && event.approval.reason != null
+                ? { reason: event.approval.reason }
+                : {}),
+        });
+    }
+    return decisions;
+}
+
+function contentHasPendingApprovals(
+    content: readonly ThreadAssistantMessagePart[],
+): boolean {
+    return messageHasPendingApprovals({
+        id: "pending-check",
+        role: "assistant",
+        content,
+        status: { type: "complete", reason: "stop" },
+        createdAt: new Date(),
+        metadata: {
+            unstable_state: null,
+            unstable_annotations: [],
+            unstable_data: [],
+            steps: [],
+            custom: {},
+        },
+    });
+}
+
+export function resolveToolApprovalUpdate(
+    update: TurnStreamUpdate,
+    priorDecisions?: ReadonlyMap<string, StoredApprovalDecision>,
+): TurnStreamUpdate {
+    let { content } = update;
+    if (priorDecisions != null && priorDecisions.size > 0) {
+        content = applyApprovalDecisionsToContent(content, priorDecisions);
+    }
+
+    if (
+        contentHasPendingApprovals(content) ||
+        update.status?.type !== "requires-action" ||
+        update.status.reason !== "tool-calls"
+    ) {
+        return { ...update, content };
+    }
+
+    const { status: _status, metadata: _metadata, ...rest } = update;
+    return { ...rest, content };
+}
+
+export function mapApprovalDecision(
+    approved: boolean,
+    reason?: string,
+): ApprovalDecision {
+    if (approved) {
+        return { status: "allow" };
+    }
+    return { status: "deny", ...(reason != null ? { reason } : {}) };
+}
+
+function isDecidedApprovalAwaitingSdk(part: ToolCallPart): boolean {
+    const { approval, result, isError } = part;
+    if (approval?.id == null || approval.approved === undefined) {
+        return false;
+    }
+    if (approval.approved === true) {
+        return result === undefined;
+    }
+    return isError === true;
+}
+
+function collectApprovalInputsFromMessages(
+    messages: readonly ThreadMessage[],
+    defaultThreadId: string,
+): UserToolApprovalEvent[] {
+    const events: UserToolApprovalEvent[] = [];
+    for (const message of messages) {
+        events.push(...collectApprovalInputs(message, defaultThreadId));
+    }
+    return events;
+}
+
+export function collectApprovalInputs(
+    message: ThreadMessage,
+    threadId: string,
+): UserToolApprovalEvent[] {
+    if (message.role !== "assistant" || !threadId) {
+        return [];
+    }
+    if (messageHasPendingApprovals(message)) {
+        return [];
+    }
+
+    const scopedThreadId = getToolApprovalThreadId(message) ?? threadId;
+    const events: UserToolApprovalEvent[] = [];
+
+    for (const part of message.content) {
+        if (part.type !== "tool-call") {
+            continue;
+        }
+        if (isDecidedApprovalAwaitingSdk(part)) {
+            const { approval } = part;
+            if (approval == null) {
+                continue;
+            }
+            events.push({
+                type: "user.tool_approval",
+                threadId: scopedThreadId,
+                toolCallId: approval.id,
+                approval: mapApprovalDecision(approval.approved!, approval.reason),
+            });
+        }
+        if (part.messages != null) {
+            events.push(
+                ...collectApprovalInputsFromMessages(part.messages, scopedThreadId),
+            );
+        }
+    }
+    return events;
+}
+
+export function toTrueFoundryApprovalInputs(
+    message: Extract<ThreadMessage, { role: "assistant" }>,
+    response: RespondToToolApprovalOptions,
+    defaultThreadId: string = ROOT_THREAD_ID,
+): UserToolApprovalEvent[] {
+    const updated = applyApprovalDecisionsToMessage(message, response);
+    return collectApprovalInputs(updated, defaultThreadId);
+}