From 59d3c81f9f5aeea23f22203dcdf6c65db804941a Mon Sep 17 00:00:00 2001 From: David Kaya Date: Mon, 30 Mar 2026 16:24:30 +0100 Subject: [PATCH] fix: classify unstreamed sub-agent messages as thinking Sub-agent messages bypass the streaming path (turn-delta) entirely due to SDK batching behavior. They arrive only at turn-complete time via FinalizeCompletedMessages. Without classification, they appear as separate chat bubbles cluttering the transcript. Two-pronged fix: Sidecar: FinalizeCompletedMessages now tags messages from _reclassifiedMessageIds with MessageKind='thinking'. This covers messages that WERE streamed and reclassified during the turn. Added MessageKind property to ChatMessageDto. Main process: finalizeTurn detects unstreamed assistant messages (not in existing session.messages) and classifies them as thinking when a visible response was already streamed. Emits message-reclassified events so the renderer can update incrementally, though the primary path is the workspace:updated broadcast which already includes the correct messageKind. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Contracts/ProtocolModels.cs | 1 + .../Services/CopilotTurnExecutionState.cs | 8 +++ .../CopilotTurnExecutionStateTests.cs | 51 +++++++++++++++++++ src/main/AryxAppService.ts | 33 +++++++++++- 4 files changed, 92 insertions(+), 1 deletion(-) diff --git a/sidecar/src/Aryx.AgentHost/Contracts/ProtocolModels.cs b/sidecar/src/Aryx.AgentHost/Contracts/ProtocolModels.cs index ca74389..ec5a6a4 100644 --- a/sidecar/src/Aryx.AgentHost/Contracts/ProtocolModels.cs +++ b/sidecar/src/Aryx.AgentHost/Contracts/ProtocolModels.cs @@ -85,6 +85,7 @@ public sealed class ChatMessageDto public string AuthorName { get; init; } = string.Empty; public string Content { get; init; } = string.Empty; public string CreatedAt { get; init; } = string.Empty; + public string? MessageKind { get; set; } public IReadOnlyList Attachments { get; init; } = []; } diff --git a/sidecar/src/Aryx.AgentHost/Services/CopilotTurnExecutionState.cs b/sidecar/src/Aryx.AgentHost/Services/CopilotTurnExecutionState.cs index 5822912..16438df 100644 --- a/sidecar/src/Aryx.AgentHost/Services/CopilotTurnExecutionState.cs +++ b/sidecar/src/Aryx.AgentHost/Services/CopilotTurnExecutionState.cs @@ -313,6 +313,14 @@ internal sealed class CopilotTurnExecutionState ActiveAgent); } + foreach (ChatMessageDto message in CompletedMessages) + { + if (_reclassifiedMessageIds.Contains(message.Id)) + { + message.MessageKind = "thinking"; + } + } + return CompletedMessages; } diff --git a/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs b/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs index 53ea0b1..d8b378d 100644 --- a/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs +++ b/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs @@ -1,6 +1,7 @@ using Aryx.AgentHost.Contracts; using Aryx.AgentHost.Services; using GitHub.Copilot.SDK; +using Microsoft.Extensions.AI; namespace Aryx.AgentHost.Tests; @@ -601,6 +602,56 @@ public sealed class CopilotTurnExecutionStateTests """); } + [Fact] + public void FinalizeCompletedMessages_TagsReclassifiedMessagesAsThinking() + { + RunTurnCommandDto command = CreateCommand(); + CopilotTurnExecutionState state = new(command); + + // Simulate assistant message with tool requests → triggers reclassification + state.ObserveSessionEvent( + command.Pattern.Agents[0], + SessionEvent.FromJson( + """ + { + "type": "assistant.message", + "data": { + "messageId": "msg-intermediate", + "content": "Let me search...", + "toolRequests": [ + { + "toolCallId": "tool-call-1", + "name": "grep", + "arguments": {} + } + ] + }, + "id": "11111111-1111-1111-1111-111111111111", + "timestamp": "2026-03-27T00:00:00Z" + } + """)); + + state.DrainPendingEvents(); + + // Build completed messages with a reclassified and a non-reclassified message + ChatMessage intermediateMsg = new(ChatRole.Assistant, "Let me search..."); + intermediateMsg.MessageId = "msg-intermediate"; + intermediateMsg.AuthorName = "Primary"; + + ChatMessage finalMsg = new(ChatRole.Assistant, "Here are the results."); + finalMsg.MessageId = "msg-final"; + finalMsg.AuthorName = "Primary"; + + state.UpdateCompletedMessages([intermediateMsg, finalMsg], []); + IReadOnlyList messages = state.FinalizeCompletedMessages(); + + ChatMessageDto intermediate = Assert.Single(messages, m => m.Id == "msg-intermediate"); + Assert.Equal("thinking", intermediate.MessageKind); + + ChatMessageDto final_ = Assert.Single(messages, m => m.Id == "msg-final"); + Assert.Null(final_.MessageKind); + } + private static RunTurnCommandDto CreateCommand() { return new RunTurnCommandDto diff --git a/src/main/AryxAppService.ts b/src/main/AryxAppService.ts index 61d85ea..c170d84 100644 --- a/src/main/AryxAppService.ts +++ b/src/main/AryxAppService.ts @@ -1824,6 +1824,15 @@ export class AryxAppService extends EventEmitter { const pattern = this.requirePattern(workspace, session.patternId); const incomingIds = new Set(messages.map((message) => message.id)); + // Messages that were streamed during the turn already exist in session.messages + // (possibly with messageKind: 'thinking' from reclassification). Unstreamed messages + // (e.g. from sub-agents) only appear now. Classify them as thinking when a visible + // response was already streamed, since they are intermediate tool-driving steps. + const existingIds = new Set(session.messages.map((m) => m.id)); + const hasVisibleResponse = session.messages.some( + (m) => m.role === 'assistant' && m.messageKind !== 'thinking', + ); + for (const message of messages) { const occurredAt = nowIso(); const existing = session.messages.find((current) => current.id === message.id); @@ -1832,9 +1841,22 @@ export class AryxAppService extends EventEmitter { existing.content = message.content; existing.pending = false; } else { - session.messages.push({ ...message, pending: false }); + const isUnstreamedIntermediate = + message.role === 'assistant' + && hasVisibleResponse + && !message.messageKind; + session.messages.push({ + ...message, + pending: false, + messageKind: message.messageKind ?? (isUnstreamedIntermediate ? 'thinking' : undefined), + }); } + const reclassifiedAsThinking = + !existingIds.has(message.id) + && (message.messageKind === 'thinking' + || (message.role === 'assistant' && hasVisibleResponse && !message.messageKind)); + const nextRun = this.updateSessionRun(session, requestId, (run) => upsertRunMessageEvent(run, { messageId: message.id, @@ -1851,6 +1873,15 @@ export class AryxAppService extends EventEmitter { authorName: message.authorName, content: message.content, }); + if (reclassifiedAsThinking) { + this.emitSessionEvent({ + sessionId, + kind: 'message-reclassified', + occurredAt, + messageId: message.id, + messageKind: 'thinking', + }); + } if (nextRun) { this.emitRunUpdated(sessionId, occurredAt, nextRun); }