From 4d7340bf4911e8fa1101bf1f3be64912a50bd2d4 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Sat, 27 Dec 2025 15:51:28 +0200 Subject: [PATCH] feat: update message handling to use 'finish' property for step completion logic --- CHANGELOG.md | 3 ++ packages/desktop/src-tauri/Cargo.lock | 2 +- .../desktop/src-tauri/src/session_activity.rs | 13 +++---- .../ui/src/components/chat/ChatMessage.tsx | 26 +++++--------- .../components/chat/hooks/useTurnGrouping.ts | 30 ++++++---------- .../components/chat/message/MessageBody.tsx | 26 +++----------- packages/ui/src/hooks/useEventStream.ts | 34 +++++------------- packages/ui/src/lib/debug.ts | 2 +- packages/ui/src/lib/messageCompletion.ts | 5 ++- packages/ui/src/stores/messageStore.ts | 35 +++++++++---------- packages/web/server/index.js | 18 ---------- 11 files changed, 60 insertions(+), 134 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index def3baa5..0f01df87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +- Optimized session activity status handling. + + ## [1.3.6] - 2025-12-27 - Added the ability to manage (connect/disconnect) providers in settings. diff --git a/packages/desktop/src-tauri/Cargo.lock b/packages/desktop/src-tauri/Cargo.lock index 8f3d3413..b3447221 100644 --- a/packages/desktop/src-tauri/Cargo.lock +++ b/packages/desktop/src-tauri/Cargo.lock @@ -2847,7 +2847,7 @@ dependencies = [ [[package]] name = "openchamber-desktop" -version = "1.3.5" +version = "1.3.6" dependencies = [ "anyhow", "axum", diff --git a/packages/desktop/src-tauri/src/session_activity.rs b/packages/desktop/src-tauri/src/session_activity.rs index 68bbe703..e7cdd762 100644 --- a/packages/desktop/src-tauri/src/session_activity.rs +++ b/packages/desktop/src-tauri/src/session_activity.rs @@ -329,8 +329,8 @@ async fn handle_event( set_phase(app, &id, ActivityPhase::Busy, phases.clone(), cooldowns.clone()).await; } - // Derive cooldown from "step-finish reason=stop" marker when present. - if is_stop_step_finish_part(&event.properties) { + // Derive cooldown from info.finish === 'stop' when present. + if has_finish_stop(info) { enter_cooldown_if_busy(app, &id, phases.clone(), cooldowns.clone()).await; } } @@ -349,13 +349,8 @@ fn is_streaming_assistant_part(properties: &Value) -> bool { ) } -fn is_stop_step_finish_part(properties: &Value) -> bool { - let Some(part) = properties.get("part") else { - return false; - }; - let part_type = part.get("type").and_then(Value::as_str); - let reason = part.get("reason").and_then(Value::as_str); - part_type == Some("step-finish") && reason == Some("stop") +fn has_finish_stop(info: &Value) -> bool { + info.get("finish").and_then(Value::as_str) == Some("stop") } async fn enter_cooldown_if_busy( diff --git a/packages/ui/src/components/chat/ChatMessage.tsx b/packages/ui/src/components/chat/ChatMessage.tsx index d0af45e1..ddd7b255 100644 --- a/packages/ui/src/components/chat/ChatMessage.tsx +++ b/packages/ui/src/components/chat/ChatMessage.tsx @@ -287,6 +287,11 @@ const ChatMessage: React.FC = ({ return Boolean(messageCompletedAt && messageCompletedAt > 0); }, [isUser, messageCompletedAt]); + const messageFinish = React.useMemo(() => { + const finish = (message.info as { finish?: string }).finish; + return typeof finish === 'string' ? finish : undefined; + }, [message.info]); + const visibleParts = React.useMemo( () => filterVisibleParts(normalizedParts, { @@ -382,23 +387,8 @@ const ChatMessage: React.FC = ({ return { name, token: rawValue } satisfies AgentMentionInfo; }, [isUser, message.parts]); - const stepState = React.useMemo(() => { - - let stepStarts = 0; - let stepFinishes = 0; - visibleParts.forEach((part) => { - if (part.type === 'step-start') { - stepStarts += 1; - } else if (part.type === 'step-finish') { - stepFinishes += 1; - } - }); - return { - hasOpenStep: stepStarts > stepFinishes, - }; - }, [visibleParts]); - - const hasOpenStep = stepState.hasOpenStep; + // Message is considered to have an "open step" if info.finish is not yet present + const hasOpenStep = typeof messageFinish !== 'string'; const shouldCoordinateRendering = React.useMemo(() => { if (isUser) { @@ -756,6 +746,7 @@ const ChatMessage: React.FC = ({ parts={visibleParts} isUser={isUser} isMessageCompleted={isMessageCompleted} + messageFinish={messageFinish} syntaxTheme={syntaxTheme} isMobile={isMobile} hasTouchInput={hasTouchInput} @@ -797,6 +788,7 @@ const ChatMessage: React.FC = ({ parts={visibleParts} isUser={isUser} isMessageCompleted={isMessageCompleted} + messageFinish={messageFinish} syntaxTheme={syntaxTheme} isMobile={isMobile} hasTouchInput={hasTouchInput} diff --git a/packages/ui/src/components/chat/hooks/useTurnGrouping.ts b/packages/ui/src/components/chat/hooks/useTurnGrouping.ts index 57f5c730..521c3e0f 100644 --- a/packages/ui/src/components/chat/hooks/useTurnGrouping.ts +++ b/packages/ui/src/components/chat/hooks/useTurnGrouping.ts @@ -108,22 +108,15 @@ export const detectTurns = (messages: ChatMessageEntry[]): Turn[] => { const extractFinalAssistantText = (turn: Turn): string | undefined => { for (const assistantMsg of turn.assistantMessages) { + const infoFinish = (assistantMsg.info as { finish?: string | null | undefined }).finish; - for (const part of assistantMsg.parts) { - if (part.type === 'step-finish') { - const finishReason = (part as { reason?: string | null | undefined }).reason; - const infoFinish = (assistantMsg.info as { finish?: string | null | undefined }).finish; - - if (finishReason === 'stop' || infoFinish === 'stop') { - - const textPart = assistantMsg.parts.find(p => p.type === 'text'); - if (textPart) { - const textContent = (textPart as { text?: string | null | undefined }).text ?? - (textPart as { content?: string | null | undefined }).content; - if (typeof textContent === 'string' && textContent.trim().length > 0) { - return textContent; - } - } + if (infoFinish === 'stop') { + const textPart = assistantMsg.parts.find(p => p.type === 'text'); + if (textPart) { + const textContent = (textPart as { text?: string | null | undefined }).text ?? + (textPart as { content?: string | null | undefined }).content; + if (typeof textContent === 'string' && textContent.trim().length > 0) { + return textContent; } } } @@ -194,12 +187,9 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => { turn.assistantMessages.forEach((msg) => { const messageId = msg.info.id; + const infoFinish = (msg.info as { finish?: string | null | undefined }).finish; const hasStopFinishInMessage = ENABLE_TEXT_JUSTIFICATION_ACTIVITY - ? msg.parts.some((part) => { - if (part.type !== 'step-finish') return false; - const reason = (part as { reason?: string | null | undefined }).reason; - return reason === 'stop'; - }) + ? infoFinish === 'stop' : false; msg.parts.forEach((part) => { diff --git a/packages/ui/src/components/chat/message/MessageBody.tsx b/packages/ui/src/components/chat/message/MessageBody.tsx index cd3ff85d..d45db717 100644 --- a/packages/ui/src/components/chat/message/MessageBody.tsx +++ b/packages/ui/src/components/chat/message/MessageBody.tsx @@ -105,6 +105,7 @@ interface MessageBodyProps { parts: Part[]; isUser: boolean; isMessageCompleted: boolean; + messageFinish?: string; syntaxTheme: { [key: string]: React.CSSProperties }; @@ -300,6 +301,7 @@ const AssistantMessageBody: React.FC> = ({ messageId, parts, isMessageCompleted, + messageFinish, syntaxTheme, isMobile, @@ -347,9 +349,7 @@ const AssistantMessageBody: React.FC> = ({ const createSessionFromAssistantMessage = useSessionStore((state) => state.createSessionFromAssistantMessage); const isLastAssistantInTurn = turnGroupingContext?.isLastAssistantInTurn ?? false; - const hasStopFinish = React.useMemo(() => { - return parts.some((part) => part.type === 'step-finish' && (part as { reason?: string | null | undefined }).reason === 'stop'); - }, [parts]); + const hasStopFinish = messageFinish === 'stop'; const hasTools = toolParts.length > 0; @@ -414,24 +414,8 @@ const AssistantMessageBody: React.FC> = ({ }); }, [reasoningParts]); - const stepState = React.useMemo(() => { - let stepStarts = 0; - let stepFinishes = 0; - visibleParts.forEach((part) => { - if (part.type === 'step-start') { - stepStarts += 1; - } else if (part.type === 'step-finish') { - stepFinishes += 1; - } - }); - return { - stepStarts, - stepFinishes, - hasOpenStep: stepStarts > stepFinishes, - }; - }, [visibleParts]); - - const hasOpenStep = stepState.hasOpenStep; + // Message is considered to have an "open step" if info.finish is not yet present + const hasOpenStep = typeof messageFinish !== 'string'; const shouldHoldForReasoning = reasoningParts.length > 0 && diff --git a/packages/ui/src/hooks/useEventStream.ts b/packages/ui/src/hooks/useEventStream.ts index 3638fcc4..6ca4d7ac 100644 --- a/packages/ui/src/hooks/useEventStream.ts +++ b/packages/ui/src/hooks/useEventStream.ts @@ -698,9 +698,7 @@ export const useEventStream = () => { const existingMessage = getMessageFromStore(sessionId, messageId); const existingLen = computeTextLength(existingMessage?.parts || []); - const existingStopMarker = existingMessage?.parts?.some( - (part) => part?.type === 'step-finish' && (part as { reason?: string }).reason === 'stop' - ) ?? false; + const existingStopMarker = (existingMessage?.info as { finish?: string } | undefined)?.finish === 'stop'; const serverParts = (props as { parts?: unknown }).parts || (messageExt as { parts?: unknown }).parts; const partsArray = Array.isArray(serverParts) ? (serverParts as Part[]) : []; @@ -710,12 +708,13 @@ export const useEventStream = () => { if (!hasParts && !completedFromServer) break; + const finishCandidate = (message as { finish?: unknown }).finish; + const finish = typeof finishCandidate === 'string' ? finishCandidate : null; + const eventHasStopFinish = finish === 'stop'; + if ((messageExt as { role?: unknown }).role === 'assistant' && hasParts) { const incomingLen = computeTextLength(partsArray); const wouldShrink = existingLen > 0 && incomingLen + TEXT_SHRINK_TOLERANCE < existingLen; - const eventHasStopFinish = partsArray.some( - (p) => p?.type === 'step-finish' && (p as { reason?: string }).reason === 'stop' - ); if (wouldShrink && !eventHasStopFinish) { trackMessage(messageId, 'skipped_shrinking_update', { incomingLen, existingLen }); @@ -772,16 +771,12 @@ export const useEventStream = () => { const messageTime = (message as { time?: { completed?: unknown } }).time; const completedCandidate = (messageTime as { completed?: unknown } | undefined)?.completed; const hasCompletedTimestamp = typeof completedCandidate === 'number' && Number.isFinite(completedCandidate); - const finishCandidate = (message as { finish?: unknown }).finish; - const finish = typeof finishCandidate === 'string' ? finishCandidate : null; - const stopMarkerPresent = partsArray.some( - (p) => p?.type === 'step-finish' && (p as { reason?: string }).reason === 'stop' - ) || existingStopMarker; + const stopMarkerPresent = finish === 'stop' || existingStopMarker; const shouldFinalizeAssistantMessage = (message as { role?: string }).role === 'assistant' && - (hasCompletedTimestamp || finish === 'stop' || stopMarkerPresent); + (hasCompletedTimestamp || stopMarkerPresent); if (shouldFinalizeAssistantMessage && (message as { role?: string }).role === 'assistant') { @@ -894,20 +889,9 @@ export const useEventStream = () => { completeStreamingMessage(sessionId, messageId); // For web/vscode: trigger cooldown only when assistant message has finish === "stop" - // (or we can infer a stop marker) to match desktop backend semantics. + // to match desktop backend semantics. if (!isDesktopRuntimeRef.current) { - const finishCandidate = (message as { finish?: unknown }).finish; - const finish = typeof finishCandidate === 'string' ? finishCandidate : null; - - const inferredStopMarkerPresent = - Array.isArray(partsArray) && - partsArray.some((part) => { - if (!part || typeof part !== 'object') return false; - const partAny = part as { type?: string; reason?: string }; - return partAny.type === 'step-finish' && partAny.reason === 'stop'; - }); - - if (finish === 'stop' || inferredStopMarkerPresent) { + if (finish === 'stop') { const currentPhase = useSessionStore.getState().sessionActivityPhase?.get(sessionId); if (currentPhase === 'busy') { updateSessionActivityPhase(sessionId, 'cooldown'); diff --git a/packages/ui/src/lib/debug.ts b/packages/ui/src/lib/debug.ts index a715fbb4..06c9b181 100644 --- a/packages/ui/src/lib/debug.ts +++ b/packages/ui/src/lib/debug.ts @@ -590,7 +590,7 @@ export const debugUtils = { const lastMessage = assistantMessages[assistantMessages.length - 1]; const stepFinishParts = lastMessage.parts.filter((p: any) => p.type === 'step-finish'); - const hasStopReason = lastMessage.parts.some((p: any) => p.type === 'step-finish' && p.reason === 'stop'); + const hasStopReason = (lastMessage.info as { finish?: string }).finish === 'stop'; const timeInfo = lastMessage.info.time as any; const completedAt = timeInfo?.completed; diff --git a/packages/ui/src/lib/messageCompletion.ts b/packages/ui/src/lib/messageCompletion.ts index 6418a5f1..00ff9b4c 100644 --- a/packages/ui/src/lib/messageCompletion.ts +++ b/packages/ui/src/lib/messageCompletion.ts @@ -11,6 +11,7 @@ export interface MessageInfo { }; status?: string; streaming?: boolean; + finish?: string; } export interface MessageRecord { @@ -27,9 +28,7 @@ export function isMessageComplete(messageInfo: MessageInfo, parts: Part[] = []): const completedAt = typeof timeInfo?.completed === 'number' ? timeInfo.completed : undefined; const messageStatus = messageInfo?.status; - const hasStopFinish = parts.some(part => - part.type === 'step-finish' && (part as any).reason === 'stop' - ); + const hasStopFinish = messageInfo.finish === 'stop'; const hasCompletedFlag = (typeof completedAt === 'number' && completedAt > 0) || messageStatus === 'completed'; if (!hasCompletedFlag || !hasStopFinish) { diff --git a/packages/ui/src/stores/messageStore.ts b/packages/ui/src/stores/messageStore.ts index f379427d..04394619 100644 --- a/packages/ui/src/stores/messageStore.ts +++ b/packages/ui/src/stores/messageStore.ts @@ -101,13 +101,8 @@ const computePartsTextLength = (parts: Part[] | undefined): number => { }, 0); }; -const hasStopReasonStop = (parts: Part[] | undefined): boolean => { - if (!Array.isArray(parts)) { - return false; - } - return parts.some( - (part) => part?.type === "step-finish" && (part as Record)?.reason === "stop" - ); +const hasFinishStop = (info: { finish?: string } | undefined): boolean => { + return info?.finish === "stop"; }; const getPartKey = (part: Part | undefined): string | undefined => { @@ -159,8 +154,8 @@ const mergeDuplicateMessage = ( const incomingParts = Array.isArray(incoming.parts) ? incoming.parts : []; const existingLen = computePartsTextLength(existingParts); const incomingLen = computePartsTextLength(incomingParts); - const existingStop = hasStopReasonStop(existingParts); - const incomingStop = hasStopReasonStop(incomingParts); + const existingStop = hasFinishStop(existing.info); + const incomingStop = hasFinishStop(incoming.info); let parts = incomingParts; if (existingStop && existingLen >= incomingLen) { @@ -449,7 +444,7 @@ export const useMessageStore = create()( const existingParts = Array.isArray(existingEntry.parts) ? existingEntry.parts : []; const existingLen = computePartsTextLength(existingParts); const serverLen = computePartsTextLength(serverParts); - const storeHasStop = hasStopReasonStop(existingParts); + const storeHasStop = hasFinishStop(existingEntry.info); if (storeHasStop && existingLen > serverLen) { const mergedParts = mergePreferExistingParts(existingParts, serverParts); @@ -1388,14 +1383,6 @@ export const useMessageStore = create()( return finalizeAbortState({ messages: newMessages, ...updates }); } }); - - const partType = (part as any)?.type; - if (partType === 'step-finish' && actualRole !== 'user') { - setTimeout(() => { - const store = get(); - store.completeStreamingMessage(sessionId, messageId); - }, 0); - } }, addStreamingPart: (sessionId: string, messageId: string, part: Part, role?: string, currentSessionId?: string) => { @@ -1887,6 +1874,16 @@ export const useMessageStore = create()( return updates; }); + + // Trigger completion when info.finish is present for assistant messages + const infoFinish = (messageInfo as { finish?: string })?.finish; + const messageRole = (messageInfo as { role?: string })?.role; + if (typeof infoFinish === 'string' && messageRole !== 'user') { + setTimeout(() => { + const store = get(); + store.completeStreamingMessage(sessionId, messageId); + }, 0); + } }, completeStreamingMessage: (sessionId: string, messageId: string) => { @@ -2041,7 +2038,7 @@ export const useMessageStore = create()( const existingParts = Array.isArray(existingEntry.parts) ? existingEntry.parts : []; const existingLen = computePartsTextLength(existingParts); const serverLen = computePartsTextLength(serverParts); - const storeHasStop = hasStopReasonStop(existingParts); + const storeHasStop = hasFinishStop(existingEntry.info); if (storeHasStop && existingLen > serverLen) { const mergedParts = mergePreferExistingParts(existingParts, serverParts); diff --git a/packages/web/server/index.js b/packages/web/server/index.js index 95bc70b1..3d1d3fb7 100644 --- a/packages/web/server/index.js +++ b/packages/web/server/index.js @@ -864,24 +864,6 @@ function deriveSessionActivity(payload) { } } - if (payload.type === 'message.part.updated') { - const info = payload.properties?.info; - const part = payload.properties?.part; - const sessionId = info?.sessionID ?? part?.sessionID ?? payload.properties?.sessionID; - const role = info?.role; - const partType = part?.type; - const reason = part?.reason; - if ( - typeof sessionId === 'string' && - sessionId.length > 0 && - role === 'assistant' && - partType === 'step-finish' && - reason === 'stop' - ) { - return { sessionId, phase: 'cooldown' }; - } - } - if (payload.type === 'session.idle') { const sessionId = payload.properties?.sessionID; if (typeof sessionId === 'string' && sessionId.length > 0) {