From 7fcc6d9d4c02cf8ebc7a701a3dc660d873f8fcd1 Mon Sep 17 00:00:00 2001 From: Govinda Vashishtha <57435703+govindavashishtha@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:38:18 +0530 Subject: [PATCH 1/2] Keep composer interactive and preserve superseded in-flight turns. Send during running/pauses supersedes the prior client stream without cancelSession, stacks pause chrome above the composer, and commits half-baked turns as cancelled so they stay in the transcript. Co-authored-by: Cursor --- .changeset/composer-pause-stack.md | 6 + .../src/collectPending.ts | 144 ++++++------ .../src/convertTurnMessages.ts | 92 +++++++- packages/assistant-ui-runtime/src/index.ts | 1 + .../src/requiredActionInputs.ts | 26 ++- .../src/useTrueForgeAgentMessages.ts | 212 ++++++++++++++--- .../test/collectPending.test.ts | 94 +++++++- .../test/requiredActionInputs.test.ts | 48 +++- .../test/useTrueForgeAgentMessages.test.tsx | 214 ++++++++++++++++++ .../src/atoms/ComposerSections.tsx | 8 +- .../trueforge-ui/src/atoms/ComposerShell.tsx | 4 + .../atoms/adapters/AskUserPromptAdapter.tsx | 4 +- .../src/containers/ComposerContainer.tsx | 66 +++--- .../src/containers/TrueForgeUIShell.tsx | 2 +- .../test/atoms/ComposerSections.test.tsx | 54 ++++- .../test/atoms/ComposerShell.test.tsx | 31 ++- .../containers/ComposerContainer.test.tsx | 63 +++++- 17 files changed, 910 insertions(+), 159 deletions(-) create mode 100644 .changeset/composer-pause-stack.md diff --git a/.changeset/composer-pause-stack.md b/.changeset/composer-pause-stack.md new file mode 100644 index 000000000..e228c68e2 --- /dev/null +++ b/.changeset/composer-pause-stack.md @@ -0,0 +1,6 @@ +--- +'@truefoundry/trueforge-ui': patch +'@truefoundry/trueforge-assistant-ui-runtime': patch +--- + +Keep the composer interactive during running and paused turns: Cancel only when empty, Send supersedes the prior client stream without cancelSession (preserving superseded in-flight turns in the transcript), and stacked pause chrome is abandoned after a later user message. Hosts that override ComposerSendButton and still branch only on isRunning should update to use hasContent. diff --git a/packages/assistant-ui-runtime/src/collectPending.ts b/packages/assistant-ui-runtime/src/collectPending.ts index eef9475b1..a67a83631 100644 --- a/packages/assistant-ui-runtime/src/collectPending.ts +++ b/packages/assistant-ui-runtime/src/collectPending.ts @@ -3,6 +3,7 @@ import type { McpAuthRequiredEvent } from './server/index.js'; import { ROOT_THREAD_ID } from './constants.js'; import { isMcpServerAuthInfoList, isUnknownRecord } from './messageCustomMetadata.js'; +import { findCurrentPausedAssistantMessage } from './requiredActionInputs.js'; import { getToolApprovalThreadId, hasPendingToolApproval } from './toolApproval.js'; import { getToolResponseThreadId, @@ -55,97 +56,90 @@ function walkToolCallParts( export function collectPendingApprovals(messages: readonly ThreadMessage[]): PendingApproval[] { const pending: PendingApproval[] = []; - - for (const message of messages) { - if (message.role !== 'assistant') { - continue; - } - const rootThreadId = getToolApprovalThreadId(message) ?? ROOT_THREAD_ID; - walkToolCallParts( - message.content, - (part, threadId) => { - const approval = part.approval; - if (approval == null || !hasPendingToolApproval(approval)) { - return; - } - pending.push({ - approvalId: approval.id, - threadId, - toolName: part.toolName, - args: { ...part.args }, - argsText: part.argsText, - }); - }, - rootThreadId, - ); + const paused = findCurrentPausedAssistantMessage(messages); + if (paused == null) { + return pending; } + const rootThreadId = getToolApprovalThreadId(paused) ?? ROOT_THREAD_ID; + walkToolCallParts( + paused.content, + (part, threadId) => { + const approval = part.approval; + if (approval == null || !hasPendingToolApproval(approval)) { + return; + } + pending.push({ + approvalId: approval.id, + threadId, + toolName: part.toolName, + args: { ...part.args }, + argsText: part.argsText, + }); + }, + rootThreadId, + ); + return pending; } export function collectPendingToolResponses(messages: readonly ThreadMessage[]): PendingToolResponse[] { const pending: PendingToolResponse[] = []; - - for (const message of messages) { - if (message.role !== 'assistant') { - continue; - } - const rootThreadId = getToolResponseThreadId(message) ?? ROOT_THREAD_ID; - walkToolCallParts( - message.content, - (part, threadId) => { - if (!hasPendingToolResponse(part)) { - return; - } - const payload: AskUserQuestionInterruptPayload | undefined = isUnknownRecord(part.interrupt?.payload) - ? { - ...(typeof part.interrupt.payload['question'] === 'string' - ? { question: part.interrupt.payload['question'] } - : {}), - ...(Array.isArray(part.interrupt.payload['options']) && - part.interrupt.payload['options'].every(option => typeof option === 'string') - ? { options: part.interrupt.payload['options'] } - : {}), - } - : undefined; - pending.push({ - toolCallId: part.toolCallId, - threadId, - toolName: part.toolName, - args: { ...part.args }, - argsText: part.argsText, - ...(payload?.question != null ? { question: payload.question } : {}), - ...(payload?.options != null ? { options: payload.options } : {}), - }); - }, - rootThreadId, - ); + const paused = findCurrentPausedAssistantMessage(messages); + if (paused == null) { + return pending; } + const rootThreadId = getToolResponseThreadId(paused) ?? ROOT_THREAD_ID; + walkToolCallParts( + paused.content, + (part, threadId) => { + if (!hasPendingToolResponse(part)) { + return; + } + const payload: AskUserQuestionInterruptPayload | undefined = isUnknownRecord(part.interrupt?.payload) + ? { + ...(typeof part.interrupt.payload['question'] === 'string' + ? { question: part.interrupt.payload['question'] } + : {}), + ...(Array.isArray(part.interrupt.payload['options']) && + part.interrupt.payload['options'].every(option => typeof option === 'string') + ? { options: part.interrupt.payload['options'] } + : {}), + } + : undefined; + pending.push({ + toolCallId: part.toolCallId, + threadId, + toolName: part.toolName, + args: { ...part.args }, + argsText: part.argsText, + ...(payload?.question != null ? { question: payload.question } : {}), + ...(payload?.options != null ? { options: payload.options } : {}), + }); + }, + rootThreadId, + ); + return pending; } export function derivePendingMcpAuth( messages: readonly ThreadMessage[], ): { mcpServers: McpAuthRequiredEvent['mcpServers'] } | null { - for (const message of messages.toReversed()) { - if (message.role !== 'assistant') { - continue; - } - if (message.status.type !== 'requires-action') { - continue; - } - const custom = message.metadata.custom; - if (custom['pendingMcpAuth'] !== true) { - continue; - } - const servers = custom['mcpServers']; - if (!isMcpServerAuthInfoList(servers)) { - return { mcpServers: [] }; - } - return { mcpServers: servers }; + const paused = findCurrentPausedAssistantMessage(messages); + if (paused == null) { + return null; + } + const custom = paused.metadata.custom; + if (custom['pendingMcpAuth'] !== true) { + return null; + } + const servers = custom['mcpServers']; + if (!isMcpServerAuthInfoList(servers)) { + return { mcpServers: [] }; } - return null; + return { mcpServers: servers }; } /** diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index 58d5fdec4..42d6fc70e 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -48,6 +48,7 @@ import { applyApprovalDecisionsToContent, collectApprovalDecisionsFromTurnInput, collectSubsequentApprovalDecisions, + hasPendingToolApproval, messageHasPendingApprovals, TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY, toolApprovalMessageCustom, @@ -58,6 +59,7 @@ import { applyUserToolResponsesToFold, collectSubsequentToolResponses, collectToolResponsesFromTurnInput, + hasPendingToolResponse, messageHasPendingResponses, TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY, toolResponseMessageCustom, @@ -1047,6 +1049,94 @@ function applyRequiredActionsOverlayToMessages( }); } +type ThreadAssistantMessagePart = Extract['content'][number]; +type AssistantToolCallPart = Extract; + +function stripInteractivePendingFromContent( + content: readonly ThreadAssistantMessagePart[], +): ThreadAssistantMessagePart[] { + return content.map(part => { + if (part.type !== 'tool-call') { + return part; + } + + let next: AssistantToolCallPart = part; + if (hasPendingToolResponse(part)) { + const { interrupt: _interrupt, ...rest } = part; + next = rest; + } + const approval = next.approval; + if (approval != null && hasPendingToolApproval(approval)) { + next = { + ...next, + approval: { + ...approval, + approved: false, + reason: 'Superseded by a later message', + }, + }; + } + if (next.messages != null) { + const nested = next.messages.map(message => { + if (message.role !== 'assistant' || message.status.type !== 'requires-action') { + return message; + } + return stripInteractivePendingFromAssistant(message); + }); + next = { ...next, messages: nested }; + } + return next; + }); +} + +function stripInteractivePendingFromAssistant( + message: Extract, +): Extract { + const content = stripInteractivePendingFromContent(message.content); + const custom = Object.fromEntries( + Object.entries(message.metadata.custom).filter( + ([key]) => + key !== 'pendingMcpAuth' && + key !== 'mcpServers' && + key !== TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY && + key !== TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY, + ), + ); + return { + ...message, + content, + status: { type: 'incomplete', reason: 'cancelled' }, + metadata: { + ...message.metadata, + custom, + }, + }; +} + +/** Clear interactive pause chrome on assistants that a later user message abandoned. */ +function abandonSupersededPausedMessages(messages: readonly ThreadMessage[]): ThreadMessage[] { + let lastUserIndex = -1; + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i]?.role === 'user') { + lastUserIndex = i; + break; + } + } + if (lastUserIndex < 0) { + return [...messages]; + } + + return messages.map((message, index) => { + if (index >= lastUserIndex) { + return message; + } + if (message.role !== 'assistant' || message.status.type !== 'requires-action') { + return message; + } + return stripInteractivePendingFromAssistant(message); + }); +} + function projectHistoryTurns(snapshot: SessionSnapshot, options?: ProjectSessionMessagesOptions): ThreadMessage[] { const messages: ThreadMessage[] = []; let lastAssistantIndex: number | undefined; @@ -1200,7 +1290,7 @@ export function projectSessionMessages( } } - return applyRequiredActionsOverlayToMessages(messages, snapshot.requiredActions); + return abandonSupersededPausedMessages(applyRequiredActionsOverlayToMessages(messages, snapshot.requiredActions)); } const DEFAULT_LIST_EVENTS_CONCURRENCY = 5; diff --git a/packages/assistant-ui-runtime/src/index.ts b/packages/assistant-ui-runtime/src/index.ts index cfdd8e993..d6b7f14a3 100644 --- a/packages/assistant-ui-runtime/src/index.ts +++ b/packages/assistant-ui-runtime/src/index.ts @@ -49,6 +49,7 @@ export type { } from './messageCustomMetadata.js'; export { collectRequiredActionInputs, + findCurrentPausedAssistantMessage, findPausedAssistantMessage, messageHasPendingRequiredActions, } from './requiredActionInputs.js'; diff --git a/packages/assistant-ui-runtime/src/requiredActionInputs.ts b/packages/assistant-ui-runtime/src/requiredActionInputs.ts index 13dd69ce2..3957aeaba 100644 --- a/packages/assistant-ui-runtime/src/requiredActionInputs.ts +++ b/packages/assistant-ui-runtime/src/requiredActionInputs.ts @@ -25,14 +25,36 @@ export function isRequiredActionInput(item: TurnInputItem): item is RequiredActi return item.type === 'user.tool_approval' || item.type === 'user.tool_response'; } -export function findPausedAssistantMessage( +/** + * Most recent assistant that is still awaiting action, ignoring a trailing + * non-paused assistant (optimistic/running bubble). Returns undefined when a + * later user message has abandoned the pause. + */ +export function findCurrentPausedAssistantMessage( messages: readonly ThreadMessage[], ): Extract | undefined { for (let i = messages.length - 1; i >= 0; i--) { const candidate = messages[i]; - if (candidate?.role === 'assistant' && candidate.status.type === 'requires-action') { + if (candidate == null) { + continue; + } + if (candidate.role === 'user') { + return undefined; + } + if (candidate.role !== 'assistant') { + continue; + } + if (candidate.status.type === 'requires-action') { return candidate; } + // Skip a trailing non-paused assistant (streaming/complete bubble). } return undefined; } + +/** @deprecated Prefer {@link findCurrentPausedAssistantMessage}. */ +export function findPausedAssistantMessage( + messages: readonly ThreadMessage[], +): Extract | undefined { + return findCurrentPausedAssistantMessage(messages); +} diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index 587bae9af..be5f64c19 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -7,6 +7,7 @@ import type { ToolResponseRequiredEvent, Turn, TurnInputItem, + TurnStateCancelled, TurnStateDone, } from './server/index.js'; import type { AgentChatServer } from './server/types.js'; @@ -25,11 +26,14 @@ import { type UserMessageContent, } from './convertTurnMessages.js'; import { extractTurnUserText } from './extractTurnUserText.js'; +import { ingestTurnEvent } from './foldPeerThreads.js'; import { loadSessionSnapshot } from './loadSessionSnapshot.js'; import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from './mcpAuth.js'; import { isMcpServerAuthInfoList } from './messageCustomMetadata.js'; +import type { AssistantContentPart } from './modelMessageContent.js'; import { collectRequiredActionInputs, + findCurrentPausedAssistantMessage, findPausedAssistantMessage, messageHasPendingRequiredActions, type RequiredActionInput, @@ -99,6 +103,142 @@ function buildCompletedTurnState( }; } +function buildCancelledTurnState(completedAt: string): TurnStateCancelled { + return { + status: 'cancelled', + reason: 'Superseded by a later message', + completedAt, + }; +} + +/** + * Custom stream adapters may yield projected content without fold events. + * Materialize that content so a cancelled commit still projects after + * `activeStream` is cleared for the next user turn. + */ +function materializeAbandonedStreamRootIds(options: { + fold: SessionSnapshot['fold']; + turnId: string; + content: readonly AssistantContentPart[]; + existingRootIds: readonly string[]; +}): string[] { + if (options.existingRootIds.length > 0) { + return [...options.existingRootIds]; + } + if (options.content.length === 0) { + return []; + } + + const text = options.content + .filter((part): part is Extract => part.type === 'text') + .map(part => part.text) + .join(''); + const toolCalls = options.content + .filter((part): part is Extract => part.type === 'tool-call') + .map(part => ({ + id: part.toolCallId, + type: 'function' as const, + function: { + name: part.toolName, + arguments: part.argsText, + }, + })); + const modelId = `client-abandoned-${options.turnId}`; + ingestTurnEvent(options.fold, { + type: 'model.message', + id: modelId, + threadId: ROOT_THREAD_ID, + createdAt: new Date().toISOString(), + content: text.length > 0 ? text : null, + ...(toolCalls.length > 0 ? { toolCalls } : {}), + }); + return [modelId]; +} + +/** + * Force-commits an in-flight client turn (optimistic user and/or incomplete + * stream) as cancelled so a superseding user send keeps that half-baked chat + * in history. Does not call cancelSession — the backend supersedes on its own. + */ +function abandonInFlightClientTurn(snapshot: SessionSnapshot): SessionSnapshot { + const active = snapshot.activeStream; + const hasIncompleteStream = active != null && active.streamComplete !== true; + const hasPendingUser = snapshot.pendingUser != null; + + if (!hasIncompleteStream && !hasPendingUser) { + return snapshot; + } + + const completedAt = new Date().toISOString(); + const cancelledState = buildCancelledTurnState(completedAt); + + if (hasIncompleteStream && active != null) { + const activeSandboxIdValue = active.update.metadata?.custom?.['sandboxId']; + const activeSandboxId = typeof activeSandboxIdValue === 'string' ? activeSandboxIdValue : undefined; + const baseline = snapshot.groupRootBaseline ?? computeGroupRootBaseline(snapshot.turns); + const rootModelMessageIds = materializeAbandonedStreamRootIds({ + fold: snapshot.fold, + turnId: active.turnId, + content: active.update.content, + existingRootIds: rootModelMessageIdsSinceBaseline(snapshot.fold, baseline), + }); + + const lastTurn = snapshot.turns.at(-1); + if (lastTurn?.id === active.turnId) { + return replaceSessionSnapshot(snapshot, { + turns: snapshot.turns.map(turn => + turn.id === active.turnId + ? { + ...turn, + state: cancelledState, + rootModelMessageIds, + ...(activeSandboxId != null ? { sandboxId: activeSandboxId } : {}), + } + : turn, + ), + pendingUser: undefined, + activeStream: undefined, + }); + } + + const record: SessionTurnRecord = { + id: active.turnId, + createdAt: snapshot.pendingUser?.createdAt.toISOString() ?? completedAt, + state: cancelledState, + input: snapshot.pendingUser ? [buildUserTurnInput(snapshot.pendingUser.content)] : [], + ...(snapshot.pendingUser ? { userText: userMessageContentToText(snapshot.pendingUser.content) } : {}), + rootModelMessageIds, + ...(activeSandboxId != null ? { sandboxId: activeSandboxId } : {}), + }; + + return replaceSessionSnapshot(snapshot, { + turns: [...snapshot.turns, record], + pendingUser: undefined, + activeStream: undefined, + }); + } + + const pending = snapshot.pendingUser; + if (pending == null) { + return snapshot; + } + + const record: SessionTurnRecord = { + id: pending.turnId, + createdAt: pending.createdAt.toISOString(), + state: cancelledState, + input: [buildUserTurnInput(pending.content)], + userText: userMessageContentToText(pending.content), + rootModelMessageIds: [], + }; + + return replaceSessionSnapshot(snapshot, { + turns: [...snapshot.turns, record], + pendingUser: undefined, + activeStream: undefined, + }); +} + /** * Reconstructs the pending required actions a paused in-flight update carried, * so the pause survives `commitActiveStream`'s synthetic "done" state. @@ -612,6 +752,14 @@ export function useTrueForgeAgentMessages({ let runStreamStarted = false; let pendingUserTurnId: string | undefined; + // Invalidate the prior client stream before any await so a buffered RAF + // cannot restore stale activeStream after the optimistic user message. + // Continuations keep today's runStream abort; they must not race a user send. + if ('userMessage' in options) { + streamGenerationRef.current += 1; + abortControllerRef.current?.abort(); + } + try { let activeSessionId = sessionId; if (activeSessionId == null) { @@ -694,9 +842,12 @@ export function useTrueForgeAgentMessages({ if (branchBase != null && 'userMessage' in options) { // Atomic apply: never merge pendingUser onto a stale React `prev` // that still holds pre-branch turns (edit would show old + new). - const rootBucket = branchBase.fold.threads.get(ROOT_THREAD_ID); + // Commit any completed stream, then abandon incomplete in-flight chat + // into history before the new optimistic user message. + const abandoned = abandonInFlightClientTurn(commitActiveStream(branchBase)); + const rootBucket = abandoned.fold.threads.get(ROOT_THREAD_ID); groupRootBaseline = [...(rootBucket?.modelMessageIds ?? [])]; - const nextSnapshot = replaceSessionSnapshot(branchBase, { + const nextSnapshot = replaceSessionSnapshot(abandoned, { pendingUser: { turnId, content: options.userMessage, @@ -704,36 +855,43 @@ export function useTrueForgeAgentMessages({ }, activeStream: undefined, groupRootBaseline, + requiredActions: { + approvals: new Map(), + toolResponses: new Map(), + }, }); snapshotRef.current = nextSnapshot; setSnapshot(nextSnapshot); pendingUserWasSet = true; pendingUserTurnId = turnId; + } else if ('userMessage' in options) { + // Sync from snapshotRef so groupRootBaseline is set before runStream + // (setState updaters after await are not guaranteed to flush first). + const abandoned = abandonInFlightClientTurn(commitActiveStream(snapshotRef.current)); + const rootBucket = abandoned.fold.threads.get(ROOT_THREAD_ID); + groupRootBaseline = [...(rootBucket?.modelMessageIds ?? [])]; + const next = replaceSessionSnapshot(abandoned, { + pendingUser: { + turnId, + content: options.userMessage, + createdAt: new Date(), + }, + activeStream: undefined, + groupRootBaseline, + // Drop staged pause answers so they cannot resume behind this turn. + requiredActions: { + approvals: new Map(), + toolResponses: new Map(), + }, + }); + snapshotRef.current = next; + setSnapshot(next); + pendingUserWasSet = true; + pendingUserTurnId = turnId; } else { setSnapshot(prev => commitActiveStream(prev, 'inputs' in options ? options.inputs : undefined)); - - if ('userMessage' in options) { - const rootBucket = snapshotRef.current.fold.threads.get(ROOT_THREAD_ID); - groupRootBaseline = [...(rootBucket?.modelMessageIds ?? [])]; - setSnapshot(prev => { - const next = replaceSessionSnapshot(prev, { - pendingUser: { - turnId, - content: options.userMessage, - createdAt: new Date(), - }, - activeStream: undefined, - groupRootBaseline, - }); - snapshotRef.current = next; - return next; - }); - pendingUserWasSet = true; - pendingUserTurnId = turnId; - } else { - groupRootBaseline = - snapshotRef.current.groupRootBaseline ?? computeGroupRootBaseline(snapshotRef.current.turns); - } + groupRootBaseline = + snapshotRef.current.groupRootBaseline ?? computeGroupRootBaseline(snapshotRef.current.turns); } runStreamStarted = true; @@ -853,7 +1011,7 @@ export function useTrueForgeAgentMessages({ return; } const projected = projectSessionMessages(nextSnapshot, projectOptions); - const paused = findPausedAssistantMessage(projected); + const paused = findCurrentPausedAssistantMessage(projected); if (paused == null || messageHasPendingRequiredActions(paused)) { return; } @@ -1135,4 +1293,4 @@ export function useTrueForgeAgentMessages({ }; } -export { findPausedAssistantMessage, MCP_AUTH_RESUME_RUN_CUSTOM_KEY }; +export { findCurrentPausedAssistantMessage, findPausedAssistantMessage, MCP_AUTH_RESUME_RUN_CUSTOM_KEY }; diff --git a/packages/assistant-ui-runtime/test/collectPending.test.ts b/packages/assistant-ui-runtime/test/collectPending.test.ts index 7eea9e64e..a84a92179 100644 --- a/packages/assistant-ui-runtime/test/collectPending.test.ts +++ b/packages/assistant-ui-runtime/test/collectPending.test.ts @@ -1,21 +1,35 @@ import type { ThreadMessage } from '@assistant-ui/core'; import { describe, expect, it } from 'vitest'; -import { deriveSandboxId } from '../src/collectPending.js'; +import { + collectPendingApprovals, + collectPendingToolResponses, + derivePendingMcpAuth, + deriveSandboxId, +} from '../src/collectPending.js'; +import { ROOT_THREAD_ID } from '../src/constants.js'; +import { TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY } from '../src/toolResponse.js'; -function assistantMessage(id: string, custom: Record = {}): ThreadMessage { +function assistantMessage( + id: string, + options: { + custom?: Record; + status?: Extract['status']; + content?: Extract['content']; + } = {}, +): Extract { return { id, role: 'assistant' as const, - content: [], - status: { type: 'complete' as const, reason: 'stop' as const }, + content: options.content ?? [], + status: options.status ?? { type: 'complete' as const, reason: 'stop' as const }, createdAt: new Date(), metadata: { unstable_state: null, unstable_annotations: [], unstable_data: [], steps: [], - custom, + custom: options.custom ?? {}, }, }; } @@ -40,7 +54,7 @@ describe('deriveSandboxId', () => { it('returns the sandboxId from the most recent message that has one', () => { const messages = [ userMessage('u1'), - assistantMessage('a1', { sandboxId: 'sbx-1' }), + assistantMessage('a1', { custom: { sandboxId: 'sbx-1' } }), userMessage('u2'), assistantMessage('a2'), ]; @@ -48,12 +62,76 @@ describe('deriveSandboxId', () => { }); it('prefers a later sandboxId over an earlier one', () => { - const messages = [assistantMessage('a1', { sandboxId: 'sbx-1' }), assistantMessage('a2', { sandboxId: 'sbx-2' })]; + const messages = [ + assistantMessage('a1', { custom: { sandboxId: 'sbx-1' } }), + assistantMessage('a2', { custom: { sandboxId: 'sbx-2' } }), + ]; expect(deriveSandboxId(messages)).toBe('sbx-2'); }); it('keeps returning the sandboxId from earlier turns even when later assistant messages lack it', () => { - const messages = [assistantMessage('a1', { sandboxId: 'sbx-1' }), assistantMessage('a2'), assistantMessage('a3')]; + const messages = [ + assistantMessage('a1', { custom: { sandboxId: 'sbx-1' } }), + assistantMessage('a2'), + assistantMessage('a3'), + ]; expect(deriveSandboxId(messages)).toBe('sbx-1'); }); }); + +describe('current pause collection', () => { + const pausedAskUser = assistantMessage('paused', { + status: { type: 'requires-action', reason: 'tool-calls' }, + custom: { [TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY]: ROOT_THREAD_ID }, + content: [ + { + type: 'tool-call', + toolCallId: 'question-1', + toolName: 'ask_user_question', + args: {}, + argsText: '{}', + interrupt: { type: 'human', payload: { question: 'Pick?' } }, + }, + ], + }); + + it('collects pending tool responses from the current pause only', () => { + expect(collectPendingToolResponses([pausedAskUser])).toHaveLength(1); + }); + + it('ignores an abandoned pause after a later user message', () => { + const messages = [pausedAskUser, userMessage('u-followup')]; + expect(collectPendingToolResponses(messages)).toEqual([]); + expect(collectPendingApprovals(messages)).toEqual([]); + expect(derivePendingMcpAuth(messages)).toBeNull(); + }); + + it('derives MCP auth only from the current pause', () => { + const pausedMcp = assistantMessage('mcp', { + status: { type: 'requires-action', reason: 'interrupt' }, + custom: { + pendingMcpAuth: true, + mcpServers: [{ id: 's1', name: 'GitHub', authUrl: 'https://example.com' }], + }, + }); + expect(derivePendingMcpAuth([pausedMcp])?.mcpServers).toHaveLength(1); + expect(derivePendingMcpAuth([pausedMcp, userMessage('u2')])).toBeNull(); + }); + + it('collects nested approvals from the current paused message', () => { + const pausedApproval = assistantMessage('paused-approval', { + status: { type: 'requires-action', reason: 'tool-calls' }, + content: [ + { + type: 'tool-call', + toolCallId: 'approval-1', + toolName: 'bash', + args: {}, + argsText: '{}', + approval: { id: 'approval-1' }, + }, + ], + }); + expect(collectPendingApprovals([pausedApproval])).toEqual([expect.objectContaining({ approvalId: 'approval-1' })]); + }); +}); diff --git a/packages/assistant-ui-runtime/test/requiredActionInputs.test.ts b/packages/assistant-ui-runtime/test/requiredActionInputs.test.ts index bff0ccba0..ce71895e9 100644 --- a/packages/assistant-ui-runtime/test/requiredActionInputs.test.ts +++ b/packages/assistant-ui-runtime/test/requiredActionInputs.test.ts @@ -1,10 +1,56 @@ import { describe, expect, it } from 'vitest'; import { ROOT_THREAD_ID } from '../src/constants.js'; -import { collectRequiredActionInputs, messageHasPendingRequiredActions } from '../src/requiredActionInputs.js'; +import { + collectRequiredActionInputs, + findCurrentPausedAssistantMessage, + messageHasPendingRequiredActions, +} from '../src/requiredActionInputs.js'; import { applyApprovalDecisionsToMessage, collectApprovalInputs } from '../src/toolApproval.js'; import { applyToolResponseToMessage } from '../src/toolResponse.js'; +describe('findCurrentPausedAssistantMessage', () => { + const paused = { + id: 'paused', + role: 'assistant' as const, + content: [], + status: { type: 'requires-action' as const, reason: 'tool-calls' as const }, + createdAt: new Date(), + metadata: { + unstable_state: null, + unstable_annotations: [], + unstable_data: [], + steps: [], + custom: {}, + }, + }; + + it('returns the current requires-action assistant', () => { + expect(findCurrentPausedAssistantMessage([paused])).toBe(paused); + }); + + it('skips a trailing non-paused assistant', () => { + const running = { + ...paused, + id: 'running', + status: { type: 'running' as const }, + }; + expect(findCurrentPausedAssistantMessage([paused, running])).toBe(paused); + }); + + it('returns undefined when a later user message abandoned the pause', () => { + const user = { + id: 'u1', + role: 'user' as const, + content: [{ type: 'text' as const, text: 'next' }], + attachments: [], + createdAt: new Date(), + metadata: { custom: {} }, + }; + expect(findCurrentPausedAssistantMessage([paused, user])).toBeUndefined(); + }); +}); + describe('requiredActionInputs', () => { describe('batched resume invariant', () => { it('bundles approvals and responses once nothing is pending', () => { diff --git a/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx b/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx index 0822a7d4b..eba6ba950 100644 --- a/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx +++ b/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx @@ -902,6 +902,220 @@ describe('useTrueForgeAgentMessages', () => { }); }); + it('aborts the prior client stream on a new user send without cancelSession', async () => { + let firstSignal: AbortSignal | undefined; + let releaseFirstStream: (() => void) | undefined; + let releaseSecondStream: (() => void) | undefined; + const runningValues: boolean[] = []; + + vi.mocked(streamTurnContent) + .mockImplementationOnce(async function* (_server, _sessionId, _fold, _options, signal) { + firstSignal = signal; + yield { content: [{ type: 'text' as const, text: 'first reply' }] }; + await new Promise(resolve => { + releaseFirstStream = resolve; + }); + }) + .mockImplementationOnce(async function* () { + yield { content: [{ type: 'text' as const, text: 'second reply' }] }; + await new Promise(resolve => { + releaseSecondStream = resolve; + }); + }); + + const { result } = renderHook(() => useTrueForgeAgentMessages({ server: mockServer, sessionId: 'session-1' })); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + const firstSend = result.current.sendTurn({ userMessage: 'first' }); + await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(1)); + expect(result.current.isRunning).toBe(true); + runningValues.push(result.current.isRunning); + + const secondSend = result.current.sendTurn({ userMessage: 'second' }); + await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(2)); + runningValues.push(result.current.isRunning); + + expect(firstSignal?.aborted).toBe(true); + expect(mockServer.cancelSession).not.toHaveBeenCalled(); + expect(runningValues.every(value => value)).toBe(true); + + await act(async () => { + releaseFirstStream?.(); + await firstSend.catch(() => undefined); + }); + expect(result.current.isRunning).toBe(true); + + await act(async () => { + releaseSecondStream?.(); + await secondSend; + }); + }); + + it('keeps mid-stream user and assistant content when a later user message supersedes', async () => { + let releaseFirstStream: (() => void) | undefined; + let releaseSecondStream: (() => void) | undefined; + const runningValues: boolean[] = []; + + vi.stubGlobal( + 'requestAnimationFrame', + vi.fn((callback: FrameRequestCallback) => { + callback(performance.now()); + return 1; + }), + ); + vi.stubGlobal('cancelAnimationFrame', vi.fn()); + + vi.mocked(streamTurnContent) + .mockImplementationOnce(async function* () { + yield { content: [{ type: 'text' as const, text: 'first partial' }] }; + await new Promise(resolve => { + releaseFirstStream = resolve; + }); + }) + .mockImplementationOnce(async function* () { + yield { content: [{ type: 'text' as const, text: 'second reply' }] }; + await new Promise(resolve => { + releaseSecondStream = resolve; + }); + }); + + const { result } = renderHook(() => useTrueForgeAgentMessages({ server: mockServer, sessionId: 'session-1' })); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + const firstSend = result.current.sendTurn({ userMessage: 'first' }); + await waitFor(() => + expect( + result.current.messages.some( + message => + message.role === 'assistant' && + message.content.some(part => part.type === 'text' && part.text === 'first partial'), + ), + ).toBe(true), + ); + runningValues.push(result.current.isRunning); + + const secondSend = result.current.sendTurn({ userMessage: 'second' }); + await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(2)); + runningValues.push(result.current.isRunning); + + expect(mockServer.cancelSession).not.toHaveBeenCalled(); + expect(runningValues.every(value => value)).toBe(true); + + const userTexts = result.current.messages + .filter(message => message.role === 'user') + .map(message => + message.content + .filter((part): part is { type: 'text'; text: string } => part.type === 'text') + .map(part => part.text) + .join(''), + ); + expect(userTexts).toEqual(['first', 'second']); + + const cancelledAssistant = result.current.messages.find( + message => + message.role === 'assistant' && + message.content.some(part => part.type === 'text' && part.text === 'first partial'), + ); + expect(cancelledAssistant).toMatchObject({ + role: 'assistant', + status: { type: 'incomplete', reason: 'cancelled' }, + }); + + await act(async () => { + releaseFirstStream?.(); + await firstSend.catch(() => undefined); + }); + expect(result.current.isRunning).toBe(true); + + await act(async () => { + releaseSecondStream?.(); + await secondSend; + }); + }); + + it('keeps the first user message when a later send arrives before any assistant token', async () => { + let releaseFirstStream: (() => void) | undefined; + let releaseSecondStream: (() => void) | undefined; + + vi.mocked(streamTurnContent) + .mockImplementationOnce(async function* () { + await new Promise(resolve => { + releaseFirstStream = resolve; + }); + yield { content: [{ type: 'text' as const, text: 'late first reply' }] }; + }) + .mockImplementationOnce(async function* () { + yield { content: [{ type: 'text' as const, text: 'second reply' }] }; + await new Promise(resolve => { + releaseSecondStream = resolve; + }); + }); + + const { result } = renderHook(() => useTrueForgeAgentMessages({ server: mockServer, sessionId: 'session-1' })); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + const firstSend = result.current.sendTurn({ userMessage: 'first before tokens' }); + await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(1)); + await waitFor(() => + expect( + result.current.messages.some( + message => + message.role === 'user' && + message.content.some(part => part.type === 'text' && part.text === 'first before tokens'), + ), + ).toBe(true), + ); + + const secondSend = result.current.sendTurn({ userMessage: 'second' }); + await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(2)); + + expect(mockServer.cancelSession).not.toHaveBeenCalled(); + const userTexts = result.current.messages + .filter(message => message.role === 'user') + .map(message => + message.content + .filter((part): part is { type: 'text'; text: string } => part.type === 'text') + .map(part => part.text) + .join(''), + ); + expect(userTexts).toEqual(['first before tokens', 'second']); + expect(result.current.isRunning).toBe(true); + + await act(async () => { + releaseFirstStream?.(); + await firstSend.catch(() => undefined); + }); + expect(result.current.isRunning).toBe(true); + + await act(async () => { + releaseSecondStream?.(); + await secondSend; + }); + }); + + it('drops ask-user pause chrome after a superseding user send', async () => { + vi.mocked(loadSessionSnapshot).mockResolvedValue(snapshotWithAskUserPendingInFold()); + vi.mocked(streamTurnContent).mockReturnValue( + (async function* () { + yield { content: [{ type: 'text' as const, text: 'new reply' }] }; + })(), + ); + + const { result } = renderHook(() => useTrueForgeAgentMessages({ server: mockServer, sessionId: 'session-1' })); + await waitFor(() => expect(collectPendingToolResponses(result.current.messages)).toHaveLength(1)); + + await act(async () => { + await result.current.sendTurn({ userMessage: 'ignore the question' }); + }); + + expect(collectPendingToolResponses(result.current.messages)).toHaveLength(0); + expect(mockServer.cancelSession).not.toHaveBeenCalled(); + const paused = result.current.messages.find(message => message.role === 'assistant' && message.id.includes('ask')); + if (paused?.role === 'assistant') { + expect(paused.status.type).not.toBe('requires-action'); + } + }); + it('carries a streamed sandboxId through commit so it survives after the stream completes', async () => { vi.mocked(streamTurnContent).mockReturnValue( (async function* () { diff --git a/packages/trueforge-ui/src/atoms/ComposerSections.tsx b/packages/trueforge-ui/src/atoms/ComposerSections.tsx index 3daae1a35..d3e937cc3 100644 --- a/packages/trueforge-ui/src/atoms/ComposerSections.tsx +++ b/packages/trueforge-ui/src/atoms/ComposerSections.tsx @@ -22,6 +22,8 @@ export type ComposerSendButtonProps = { disabled: boolean; canSubmit: boolean; isRunning: boolean; + /** True when the composer has trimmed text or attachments. */ + hasContent: boolean; onSubmit: () => void; onCancel?: () => void; }; @@ -50,8 +52,10 @@ export function ComposerRightSection(_: ComposerRightSectionProps): ReactNode { return null; } -export function ComposerSendButton({ canSubmit, isRunning, onSubmit, onCancel }: ComposerSendButtonProps) { - if (isRunning) { +export function ComposerSendButton({ canSubmit, isRunning, hasContent, onSubmit, onCancel }: ComposerSendButtonProps) { + // Cancel only while busy with an empty composer. Content (even when Send is + // blocked by a missing model) keeps the Send affordance. + if (isRunning && !hasContent) { return ( diff --git a/packages/trueforge-ui/src/atoms/ComposerShell.tsx b/packages/trueforge-ui/src/atoms/ComposerShell.tsx index d711c7837..cc7e6bf51 100644 --- a/packages/trueforge-ui/src/atoms/ComposerShell.tsx +++ b/packages/trueforge-ui/src/atoms/ComposerShell.tsx @@ -11,6 +11,8 @@ export type ComposerShellProps = { disabled: boolean; canSubmit: boolean; isRunning?: boolean; + /** True when the composer has trimmed text or attachments. */ + hasContent?: boolean; attachments?: ReactNode; modelLabel?: string; modelIcon?: ReactNode; @@ -26,6 +28,7 @@ export function ComposerShell({ disabled, canSubmit, isRunning = false, + hasContent = false, attachments, modelLabel, modelIcon, @@ -73,6 +76,7 @@ export function ComposerShell({ disabled={disabled} canSubmit={canSubmit} isRunning={isRunning} + hasContent={hasContent} onSubmit={onSubmit} onCancel={onCancel} /> diff --git a/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx b/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx index 925ab0891..96d942bd0 100644 --- a/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx +++ b/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx @@ -177,7 +177,7 @@ export function AskUserPrompt({ )} -
+
{isMultiQuestion && (
@@ -244,7 +244,7 @@ export function AskUserPrompt({
{isMultiQuestion && isLastQuestion && !allQuestionsAnswered ? ( diff --git a/packages/trueforge-ui/src/containers/ComposerContainer.tsx b/packages/trueforge-ui/src/containers/ComposerContainer.tsx index a8db0b14b..3c6021566 100644 --- a/packages/trueforge-ui/src/containers/ComposerContainer.tsx +++ b/packages/trueforge-ui/src/containers/ComposerContainer.tsx @@ -2,7 +2,7 @@ import { ComposerPrimitive, useAui, useAuiState } from '@assistant-ui/react'; import { useTrueForgeAgentSpec, useTrueForgeCancel } from '@truefoundry/trueforge-assistant-ui-runtime'; -import { useRef } from 'react'; +import { useRef, type KeyboardEvent } from 'react'; import { DraftCatalogProvider } from '../atoms/draft/DraftCatalogProvider.js'; import { useComposerBusyState } from '../hooks/useComposerBusyState.js'; @@ -38,12 +38,10 @@ export function canSubmitComposer({ function ComposerBody({ placeholder, - forceDisabled = false, connectedToBanner = false, }: { placeholder: string; - forceDisabled?: boolean; - /** Flatten top radius/border so the approval banner sits flush above. */ + /** Flatten top radius/border so pause chrome sits flush above. */ connectedToBanner?: boolean; }) { const ComposerShell = useSlot('ComposerShell'); @@ -51,6 +49,7 @@ function ComposerBody({ const shell = useOptionalShellMode(); const hasText = useAuiState(s => s.composer.text.trim().length > 0); const hasAttachments = useAuiState(s => s.composer.attachments.length > 0); + const hasContent = hasText || hasAttachments; const { agentSpec } = useTrueForgeAgentSpec(); // Named (immutable) agents use a server-side model; only draft/mutable composers pick one here. const requiresModel = shell == null || (shell.mode.status === 'active' && shell.mode.isMutable); @@ -59,13 +58,24 @@ function ComposerBody({ const canManageSession = useActiveSessionCanManage(); const cancel = useTrueForgeCancel(); const fileInputRef = useRef(null); - const disabled = isBusy || forceDisabled || !canManageSession; + // Running/pause no longer lock the input — only session permissions do. + const disabled = !canManageSession; const canSubmit = canSubmitComposer({ disabled, hasText, hasAttachments, requiresModel, hasModel }); const submit = () => { if (!canSubmit) return; + // Do not cancelSession here; sendTurn detaches the prior client stream. send(() => aui.composer().send()); }; + // assistant-ui blocks Enter while the thread is running without queue support. + // Intercept before the primitive handler so Send can supersede a running turn. + const onInputKeyDown = (event: KeyboardEvent) => { + if (event.key !== 'Enter' || event.shiftKey || event.nativeEvent.isComposing) return; + if (!canSubmit) return; + event.preventDefault(); + submit(); + }; + return ( <> } disabled={disabled} canSubmit={canSubmit} - isRunning={isBusy && !forceDisabled} + hasContent={hasContent} + isRunning={isBusy} onSubmit={submit} onCancel={ canManageSession @@ -137,11 +149,9 @@ function ComposerBody({ function ComposerWithOptionalDraft({ placeholder, - forceDisabled = false, connectedToBanner = false, }: { placeholder: string; - forceDisabled?: boolean; connectedToBanner?: boolean; }) { const shell = useOptionalShellMode(); @@ -162,13 +172,13 @@ function ComposerWithOptionalDraft({ ComposerRightSection: usesDefaultRightSection ? DraftComposerRightSection : parentRightSection, }} > - + ); } - return ; + return ; } export function ComposerContainer({ @@ -177,23 +187,25 @@ export function ComposerContainer({ const pauseView = useComposerPauseView(); const canManageSession = useActiveSessionCanManage(); - if (pauseView.kind === 'mcp') { - return ; - } - if (pauseView.kind === 'custom') { - return ; - } - if (pauseView.kind === 'ask-user') { - return ; - } - if (pauseView.kind === 'approval') { - return ( -
- - -
- ); + const pauseChrome = + pauseView.kind === 'mcp' ? ( + + ) : pauseView.kind === 'custom' ? ( + + ) : pauseView.kind === 'ask-user' ? ( + + ) : pauseView.kind === 'approval' ? ( + + ) : null; + + if (pauseChrome == null) { + return ; } - return ; + return ( +
+ {pauseChrome} + +
+ ); } diff --git a/packages/trueforge-ui/src/containers/TrueForgeUIShell.tsx b/packages/trueforge-ui/src/containers/TrueForgeUIShell.tsx index 8c7c74699..5ba00ec4d 100644 --- a/packages/trueforge-ui/src/containers/TrueForgeUIShell.tsx +++ b/packages/trueforge-ui/src/containers/TrueForgeUIShell.tsx @@ -61,7 +61,7 @@ export type TrueForgeUIProps = { initialSettingsOpen?: boolean; /** * Host UI for pending client-side tools, keyed by tool name. - * When a pending tool matches, the composer mounts that component instead of Ask User. + * When a pending tool matches, that component is shown above the composer instead of Ask User. */ customActionRenderers?: CustomActionRenderers; /** diff --git a/packages/trueforge-ui/test/atoms/ComposerSections.test.tsx b/packages/trueforge-ui/test/atoms/ComposerSections.test.tsx index 275582750..351fde4ed 100644 --- a/packages/trueforge-ui/test/atoms/ComposerSections.test.tsx +++ b/packages/trueforge-ui/test/atoms/ComposerSections.test.tsx @@ -35,7 +35,13 @@ describe('ComposerSendButton', () => { it('disables submission until the composer can submit', () => { const onSubmit = vi.fn(); const { rerender } = render( - , + , ); const unavailableButton = screen.getByRole('button', { name: 'Send message' }); @@ -43,7 +49,7 @@ describe('ComposerSendButton', () => { fireEvent.click(unavailableButton); expect(onSubmit).not.toHaveBeenCalled(); - rerender(); + rerender(); const availableButton = screen.getByRole('button', { name: 'Send message' }); expect(availableButton).toBeEnabled(); expect(availableButton).toHaveAttribute('title', 'Send message'); @@ -51,10 +57,19 @@ describe('ComposerSendButton', () => { expect(onSubmit).toHaveBeenCalledOnce(); }); - it('switches to a cancel action while running', () => { + it('shows cancel only while running with an empty composer', () => { const onSubmit = vi.fn(); const onCancel = vi.fn(); - render(); + render( + , + ); const cancelButton = screen.getByRole('button', { name: 'Cancel' }); expect(cancelButton).toHaveTextContent('Cancel'); @@ -64,8 +79,37 @@ describe('ComposerSendButton', () => { expect(onSubmit).not.toHaveBeenCalled(); }); + it('shows send while running when the composer has content', () => { + const onSubmit = vi.fn(); + const onCancel = vi.fn(); + render( + , + ); + + expect(screen.queryByRole('button', { name: 'Cancel' })).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole('button', { name: 'Send message' })); + expect(onSubmit).toHaveBeenCalledOnce(); + expect(onCancel).not.toHaveBeenCalled(); + }); + + it('shows disabled send while running with content that cannot submit', () => { + render( + {}} + onCancel={() => {}} + />, + ); + + expect(screen.queryByRole('button', { name: 'Cancel' })).not.toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Send message' })).toBeDisabled(); + }); + it('disables cancellation when no cancel handler is supplied', () => { - render( {}} />); + render( {}} />); expect(screen.getByRole('button', { name: 'Cancel' })).toBeDisabled(); }); diff --git a/packages/trueforge-ui/test/atoms/ComposerShell.test.tsx b/packages/trueforge-ui/test/atoms/ComposerShell.test.tsx index 25ba7ca9a..37a3ec6e2 100644 --- a/packages/trueforge-ui/test/atoms/ComposerShell.test.tsx +++ b/packages/trueforge-ui/test/atoms/ComposerShell.test.tsx @@ -13,6 +13,7 @@ describe('ComposerShell', () => { attachments={
design.png
} disabled={false} canSubmit + hasContent modelLabel="GPT 5" modelIcon={M} connectorStatusLabel="Connected" @@ -42,8 +43,9 @@ describe('ComposerShell', () => { render( } - disabled + disabled={false} canSubmit={false} + hasContent={false} isRunning onSubmit={onSubmit} onCancel={onCancel} @@ -58,9 +60,34 @@ describe('ComposerShell', () => { expect(onSubmit).not.toHaveBeenCalled(); }); + it('shows send while running when hasContent is true', () => { + const onSubmit = vi.fn(); + render( + } + disabled={false} + canSubmit + hasContent + isRunning + onSubmit={onSubmit} + onCancel={() => {}} + />, + ); + + expect(screen.queryByRole('button', { name: 'Cancel' })).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole('button', { name: 'Send message' })); + expect(onSubmit).toHaveBeenCalledOnce(); + }); + it('omits optional composer chrome when labels and attachment handling are absent', () => { render( - } disabled={false} canSubmit={false} onSubmit={() => {}} />, + } + disabled={false} + canSubmit={false} + hasContent={false} + onSubmit={() => {}} + />, ); expect(screen.queryByRole('button', { name: 'Attach' })).not.toBeInTheDocument(); diff --git a/packages/trueforge-ui/test/containers/ComposerContainer.test.tsx b/packages/trueforge-ui/test/containers/ComposerContainer.test.tsx index 01be1f138..052005f63 100644 --- a/packages/trueforge-ui/test/containers/ComposerContainer.test.tsx +++ b/packages/trueforge-ui/test/containers/ComposerContainer.test.tsx @@ -32,8 +32,10 @@ const approvalsState = vi.hoisted(() => ({ }>, })); +const cancelMock = vi.hoisted(() => vi.fn()); + vi.mock('@truefoundry/trueforge-assistant-ui-runtime', () => ({ - useTrueForgeCancel: () => vi.fn(), + useTrueForgeCancel: () => cancelMock, useTrueForgeToolResponses: () => toolResponsesState, useTrueForgeApprovals: () => approvalsState, useTrueForgeAgentSpec: () => ({ agentSpec: agentSpecState.agentSpec }), @@ -102,6 +104,7 @@ describe('ComposerContainer', () => { toolResponsesState.pending = []; toolResponsesState.respond = vi.fn(); approvalsState.pending = []; + cancelMock.mockReset(); }); it('wraps the composer in an attachment dropzone by default', () => { renderComposer(); @@ -212,7 +215,7 @@ describe('ComposerContainer', () => { expect(screen.getByText('Custom right')).toBeInTheDocument(); }); - it('mounts a registered custom action renderer instead of the composer', () => { + it('keeps the composer enabled under a custom action renderer', () => { toolResponsesState.pending = [{ toolCallId: 'tc-1', toolName: 'secret_select', args: { secrets: ['a'] } }]; render( @@ -226,7 +229,8 @@ describe('ComposerContainer', () => { ); expect(screen.getByRole('button', { name: 'Secret selector' })).toBeInTheDocument(); - expect(screen.queryByRole('textbox', { name: 'Message input' })).not.toBeInTheDocument(); + expect(screen.getByRole('textbox', { name: 'Message input' })).toBeEnabled(); + expect(document.querySelector('[data-slot="aui_composer-pause"]')).toHaveAttribute('data-pause-kind', 'custom'); fireEvent.click(screen.getByRole('button', { name: 'Secret selector' })); expect(toolResponsesState.respond).toHaveBeenCalledWith({ @@ -235,7 +239,7 @@ describe('ComposerContainer', () => { }); }); - it('shows the approval banner above a disabled composer while approvals are pending', () => { + it('keeps the composer enabled under the approval banner', () => { approvalsState.pending = [ { approvalId: 'appr-1', @@ -257,7 +261,54 @@ describe('ComposerContainer', () => { expect(screen.getByText('2 tools need your input')).toBeInTheDocument(); expect(screen.getByText('(1/2)')).toBeInTheDocument(); - expect(screen.getByRole('textbox', { name: 'Message input' })).toBeDisabled(); - expect(document.querySelector('[data-slot="aui_composer-approval-pause"]')).toBeInTheDocument(); + expect(screen.getByRole('textbox', { name: 'Message input' })).toBeEnabled(); + expect(document.querySelector('[data-slot="aui_composer-pause"]')).toHaveAttribute('data-pause-kind', 'approval'); + }); + + it('shows cancel while running with an empty composer', () => { + render( + + + + + , + ); + + expect(screen.getByRole('textbox', { name: 'Message input' })).toBeEnabled(); + expect(screen.getByRole('button', { name: 'Cancel' })).toBeEnabled(); + expect(screen.queryByRole('button', { name: 'Send message' })).not.toBeInTheDocument(); + }); + + it('sends while running without calling cancel', async () => { + const onNew = vi.fn(async () => {}); + render( + + + + + , + ); + + const input = screen.getByRole('textbox', { name: 'Message input' }); + fireEvent.change(input, { target: { value: 'follow up' } }); + expect(screen.getByRole('button', { name: 'Send message' })).toBeEnabled(); + expect(screen.queryByRole('button', { name: 'Cancel' })).not.toBeInTheDocument(); + + fireEvent.keyDown(input, { key: 'Enter' }); + await waitFor(() => expect(onNew).toHaveBeenCalledTimes(1)); + expect(cancelMock).not.toHaveBeenCalled(); + }); + + it('does not cancel on empty Enter while running', () => { + render( + + + + + , + ); + + fireEvent.keyDown(screen.getByRole('textbox', { name: 'Message input' }), { key: 'Enter' }); + expect(cancelMock).not.toHaveBeenCalled(); }); }); From fb30c0551b7afb0d866032a29838f1424c844a9b Mon Sep 17 00:00:00 2001 From: Govinda Vashishtha <57435703+govindavashishtha@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:51:05 +0530 Subject: [PATCH 2/2] feat: polish MCP auth and Ask User interface - Introduced lettered options (A/B/C) for Ask User prompts, enhancing user interaction. - Updated button labels to "Continue" for clarity. - Added new color variable `--color-gray-850` for improved styling. - Adjusted layout for better alignment of components in the MCP auth and Ask User prompts. This update aims to enhance the user experience by providing clearer options and a more polished interface. --- .changeset/ask-user-mcp-ui-polish.md | 5 ++ .../atoms/adapters/AskUserPromptAdapter.tsx | 81 +++++++++++++++---- .../atoms/adapters/McpAuthPromptAdapter.tsx | 20 +++-- .../src/containers/ComposerContainer.tsx | 26 +++++- packages/trueforge-ui/src/styles.css | 3 + .../adapters/AskUserPromptAdapter.test.tsx | 24 +++++- 6 files changed, 133 insertions(+), 26 deletions(-) create mode 100644 .changeset/ask-user-mcp-ui-polish.md diff --git a/.changeset/ask-user-mcp-ui-polish.md b/.changeset/ask-user-mcp-ui-polish.md new file mode 100644 index 000000000..d283c215a --- /dev/null +++ b/.changeset/ask-user-mcp-ui-polish.md @@ -0,0 +1,5 @@ +--- +'@truefoundry/trueforge-ui': patch +--- + +Polish MCP auth and Ask User pause chrome: lettered A/B/C options with Continue, outlined Connect, and screenshot-matching stack above the composer. diff --git a/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx b/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx index 96d942bd0..826583d59 100644 --- a/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx +++ b/packages/trueforge-ui/src/atoms/adapters/AskUserPromptAdapter.tsx @@ -22,6 +22,17 @@ export type AskUserAnswerDraft = { export const ASK_USER_CUSTOM_OPTION = '__tfy_ask_user_question_custom__'; +/** A, B, … Z, then AA, AB, … for option indices. */ +export function askUserOptionLetter(index: number): string { + let n = index; + let label = ''; + do { + label = String.fromCharCode(65 + (n % 26)) + label; + n = Math.floor(n / 26) - 1; + } while (n >= 0); + return label; +} + export type AskUserPromptProps = { questions: Question[]; answeredQuestions?: AnsweredQuestion[]; @@ -42,6 +53,23 @@ export type AskUserPromptProps = { className?: string; }; +function OptionLetterBadge({ letter, selected, className }: { letter: string; selected: boolean; className?: string }) { + return ( + + {letter} + + ); +} + export function AskUserPrompt({ questions, answeredQuestions = [], @@ -92,6 +120,7 @@ export function AskUserPrompt({ const isCustomSelected = currentAnswer.radioValue === ASK_USER_CUSTOM_OPTION; const hasOptions = currentQuestion.options.length > 0; + const customOptionLetter = askUserOptionLetter(currentQuestion.options.length); const handleOptionSelect = (value: string) => { if (readOnly) return; @@ -138,11 +167,15 @@ export function AskUserPrompt({ return (
-
-
+
+
{totalQuestions > 1 ? 'Questions' : currentQuestion.question}
{readOnly ? ( @@ -151,7 +184,7 @@ export function AskUserPrompt({
) : null} {!readOnly && isMultiQuestion && ( -
+
-
-
+
+
{isMultiQuestion && (
{currentQuestion.question || 'What would you like to do?'}
)}
{currentQuestion.options.map((opt, index) => { const isSelected = currentAnswer.radioValue === opt; + const letter = askUserOptionLetter(index); return (
@@ -253,12 +290,24 @@ export function AskUserPrompt({ )} {isNextStep ? ( - + Next ) : ( - - Submit + + Continue )}
diff --git a/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx b/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx index 6b4e799a1..ef6acc179 100644 --- a/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx +++ b/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx @@ -36,7 +36,11 @@ export function McpAuthPrompt({ return (
@@ -46,6 +50,7 @@ export function McpAuthPrompt({ {servers.map(server => (
+ MCP Server Name : {server.name} @@ -56,7 +61,7 @@ export function McpAuthPrompt({ Connected ) : ( - onConnect(server.id)} @@ -64,13 +69,18 @@ export function McpAuthPrompt({ > Connect - + )}
))} {onContinue && ( -
- +
+ {continueLoading ? : null} Continue diff --git a/packages/trueforge-ui/src/containers/ComposerContainer.tsx b/packages/trueforge-ui/src/containers/ComposerContainer.tsx index 3c6021566..b6f98e7aa 100644 --- a/packages/trueforge-ui/src/containers/ComposerContainer.tsx +++ b/packages/trueforge-ui/src/containers/ComposerContainer.tsx @@ -202,10 +202,34 @@ export function ComposerContainer({ return ; } + // Approval sits flush on the composer. Ask User keeps a small gap. MCP tucks the + // composer over the card’s bottom edge (screenshot stack). + if (pauseView.kind === 'approval') { + return ( +
+ {pauseChrome} + +
+ ); + } + + if (pauseView.kind === 'mcp') { + return ( +
+
{pauseChrome}
+
+ +
+
+ ); + } + return (
{pauseChrome} - +
+ +
); } diff --git a/packages/trueforge-ui/src/styles.css b/packages/trueforge-ui/src/styles.css index ea369e41b..350b0c359 100644 --- a/packages/trueforge-ui/src/styles.css +++ b/packages/trueforge-ui/src/styles.css @@ -44,6 +44,7 @@ --color-gray-700: #3f3f46; --color-gray-750: #303036; --color-gray-800: #27272a; + --color-gray-850: #18181b; --color-gray-1000: #09090b; --color-indigo-500: #6366f1; --color-red-500: #ef4444; @@ -115,6 +116,7 @@ --color-gray-700: #d4d4d8; --color-gray-750: #dfdfe2; --color-gray-800: #ececee; + --color-gray-850: #f5f5f7; --color-gray-1000: #fdfdfd; --sidebar-bg: var(--color-gray-200); @@ -337,6 +339,7 @@ --color-gray-700: var(--color-gray-700); --color-gray-750: var(--color-gray-750); --color-gray-800: var(--color-gray-800); + --color-gray-850: var(--color-gray-850); --color-gray-1000: var(--color-gray-1000); --color-indigo-500: var(--color-indigo-500); --color-red-500: var(--color-red-500); diff --git a/packages/trueforge-ui/test/atoms/adapters/AskUserPromptAdapter.test.tsx b/packages/trueforge-ui/test/atoms/adapters/AskUserPromptAdapter.test.tsx index f4bf4c994..b7085b3f2 100644 --- a/packages/trueforge-ui/test/atoms/adapters/AskUserPromptAdapter.test.tsx +++ b/packages/trueforge-ui/test/atoms/adapters/AskUserPromptAdapter.test.tsx @@ -2,7 +2,12 @@ import { fireEvent, render, screen } from '@testing-library/react'; import { describe, expect, it, vi } from 'vitest'; -import { ASK_USER_CUSTOM_OPTION, AskUserPrompt, type Question } from '@/atoms/adapters/AskUserPromptAdapter.js'; +import { + ASK_USER_CUSTOM_OPTION, + AskUserPrompt, + askUserOptionLetter, + type Question, +} from '@/atoms/adapters/AskUserPromptAdapter.js'; const firstQuestion: Question = { id: 'deployment', @@ -16,6 +21,14 @@ const secondQuestion: Question = { options: ['US', 'EU'], }; +describe('askUserOptionLetter', () => { + it('maps indices to A…Z then AA…', () => { + expect(askUserOptionLetter(0)).toBe('A'); + expect(askUserOptionLetter(25)).toBe('Z'); + expect(askUserOptionLetter(26)).toBe('AA'); + }); +}); + describe('AskUserPrompt', () => { it('renders nothing without questions or answers, then summarizes answered questions', () => { const onSubmit = vi.fn(); @@ -56,6 +69,9 @@ describe('AskUserPrompt', () => { fireEvent.click(screen.getByRole('radio', { name: 'Immediately' })); expect(onCurrentAnswerChange).toHaveBeenCalledWith('deployment', { radioValue: 'Immediately' }); + expect(screen.getByText('A')).toBeInTheDocument(); + expect(screen.getByText('B')).toBeInTheDocument(); + expect(screen.getByText('C')).toBeInTheDocument(); const customAnswer = screen.getByRole('textbox', { name: 'Other (custom answer)' }); fireEvent.focus(customAnswer); @@ -69,7 +85,7 @@ describe('AskUserPrompt', () => { radioValue: ASK_USER_CUSTOM_OPTION, }); - fireEvent.click(screen.getByRole('button', { name: 'Submit' })); + fireEvent.click(screen.getByRole('button', { name: 'Continue' })); expect(onSubmit).toHaveBeenCalledOnce(); }); @@ -90,7 +106,7 @@ describe('AskUserPrompt', () => { expect(screen.getByText('Which region?')).toBeInTheDocument(); expect(screen.getByText('Answer all questions to submit')).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'Next question' })).toBeDisabled(); - expect(screen.getByRole('button', { name: 'Submit' })).toBeDisabled(); + expect(screen.getByRole('button', { name: 'Continue' })).toBeDisabled(); fireEvent.click(screen.getByRole('button', { name: 'Previous question' })); expect(onCurrentQuestionIndexChange).toHaveBeenCalledWith(0); @@ -174,7 +190,7 @@ describe('AskUserPrompt', () => { expect(radio).toBeDisabled(); } expect(screen.getByRole('textbox', { name: 'Other (custom answer)' })).toBeDisabled(); - expect(screen.queryByRole('button', { name: 'Submit' })).not.toBeInTheDocument(); + expect(screen.queryByRole('button', { name: 'Continue' })).not.toBeInTheDocument(); fireEvent.click(screen.getByRole('radio', { name: 'Immediately' })); expect(onCurrentAnswerChange).not.toHaveBeenCalled();