Skip to content

Commit 137fa37

Browse files
xinyi-gongxinyi-gong
authored
feat - Persist thinking blocks to conversation history and restore on session switch (#210)
1 parent 75f108f commit 137fa37

11 files changed

Lines changed: 783 additions & 48 deletions

File tree

‎com.microsoft.copilot.eclipse.core.test/src/com/microsoft/copilot/eclipse/core/persistence/ConversationPersistenceManagerTests.java‎

Lines changed: 147 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import static org.mockito.ArgumentMatchers.any;
1111
import static org.mockito.ArgumentMatchers.anyString;
1212
import static org.mockito.ArgumentMatchers.eq;
13+
import static org.mockito.ArgumentMatchers.isNull;
1314
import static org.mockito.Mockito.mock;
1415
import static org.mockito.Mockito.never;
1516
import static org.mockito.Mockito.verify;
@@ -35,10 +36,15 @@
3536

3637
import com.microsoft.copilot.eclipse.core.AuthStatusManager;
3738
import com.microsoft.copilot.eclipse.core.logger.CopilotForEclipseLogger;
39+
import com.microsoft.copilot.eclipse.core.lsp.protocol.AgentRound;
3840
import com.microsoft.copilot.eclipse.core.lsp.protocol.ChatProgressValue;
3941
import com.microsoft.copilot.eclipse.core.lsp.protocol.CopilotModel;
42+
import com.microsoft.copilot.eclipse.core.lsp.protocol.Thinking;
4043
import com.microsoft.copilot.eclipse.core.lsp.protocol.Turn;
44+
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.EditAgentRoundData;
4145
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ReplyData;
46+
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ThinkingBlockData;
47+
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ThinkingBlockState;
4248
import com.microsoft.copilot.eclipse.core.persistence.UserTurnData.MessageData;
4349

4450
@ExtendWith(MockitoExtension.class)
@@ -187,7 +193,7 @@ void testCacheConversationProgress_Success() throws Exception {
187193

188194
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
189195

190-
CompletableFuture<Void> result = persistenceManager.cacheConversationProgress(conversationId, progress);
196+
CompletableFuture<Void> result = persistenceManager.cacheConversationProgress(conversationId, progress, null);
191197

192198
result.get(); // Wait for completion
193199

@@ -198,6 +204,98 @@ void testCacheConversationProgress_Success() throws Exception {
198204
assertTrue(cache.containsKey(conversationId));
199205
}
200206

207+
@Test
208+
void testUpdateThinkingBlockTitle_updatesCachedThinkingBlockById() throws Exception {
209+
ConversationPersistenceManager manager = createManagerWithRealDataFactory();
210+
String conversationId = "00000000-0000-0000-0000-000000000000";
211+
String turnId = "00000000-0000-0000-0000-000000000002";
212+
String thinkingBlockId = "thinking-block-1";
213+
ConversationData conversationData = createTestConversationData(conversationId);
214+
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
215+
216+
manager.cacheConversationProgress(conversationId,
217+
createThinkingProgressValue(conversationId, turnId, "thinking content"), thinkingBlockId).get();
218+
manager.updateThinkingBlockTitle(conversationId, turnId, thinkingBlockId, "Generated title").get();
219+
220+
ThinkingBlockData block = getCachedCopilotTurn(manager, conversationId, turnId).getReply()
221+
.getEditAgentRounds().get(0).getThinkingBlock();
222+
assertNotNull(block);
223+
assertEquals(thinkingBlockId, block.getId());
224+
assertEquals("Generated title", block.getTitle());
225+
}
226+
227+
@Test
228+
void testCancelThinkingBlock_updatesCachedThinkingBlockById() throws Exception {
229+
ConversationPersistenceManager manager = createManagerWithRealDataFactory();
230+
String conversationId = "00000000-0000-0000-0000-000000000000";
231+
String turnId = "00000000-0000-0000-0000-000000000002";
232+
String thinkingBlockId = "thinking-block-1";
233+
ConversationData conversationData = createTestConversationData(conversationId);
234+
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
235+
236+
manager.cacheConversationProgress(conversationId,
237+
createThinkingProgressValue(conversationId, turnId, "thinking content"), thinkingBlockId).get();
238+
manager.cancelThinkingBlock(conversationId, turnId, thinkingBlockId).get();
239+
240+
ThinkingBlockData block = getCachedCopilotTurn(manager, conversationId, turnId).getReply()
241+
.getEditAgentRounds().get(0).getThinkingBlock();
242+
assertNotNull(block);
243+
assertEquals(thinkingBlockId, block.getId());
244+
assertEquals(ThinkingBlockState.CANCELLED, block.getState());
245+
}
246+
247+
@Test
248+
void testCacheConversationProgress_withThinkingBlockId_updatesMatchingPlaceholderRound() throws Exception {
249+
ConversationPersistenceManager manager = createManagerWithRealDataFactory();
250+
String conversationId = "00000000-0000-0000-0000-000000000000";
251+
String turnId = "00000000-0000-0000-0000-000000000002";
252+
String firstThinkingBlockId = "thinking-block-1";
253+
String secondThinkingBlockId = "thinking-block-2";
254+
ConversationData conversationData = createTestConversationData(conversationId);
255+
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
256+
257+
manager.cacheConversationProgress(conversationId,
258+
createThinkingProgressValue(conversationId, turnId, "first"), firstThinkingBlockId).get();
259+
manager.cacheConversationProgress(conversationId,
260+
createThinkingProgressValue(conversationId, turnId, "second"), secondThinkingBlockId).get();
261+
manager.cacheConversationProgress(conversationId,
262+
createAgentRoundProgressValue(conversationId, turnId, 1, "round reply"), secondThinkingBlockId).get();
263+
264+
ReplyData reply = getCachedCopilotTurn(manager, conversationId, turnId).getReply();
265+
List<EditAgentRoundData> rounds = reply.getEditAgentRounds();
266+
assertEquals(2, rounds.size());
267+
EditAgentRoundData updatedRound = rounds.stream()
268+
.filter(round -> round.getRoundId() == 1)
269+
.findFirst()
270+
.orElseThrow();
271+
assertEquals(secondThinkingBlockId, updatedRound.getThinkingBlock().getId());
272+
assertEquals("round reply", updatedRound.getReply());
273+
assertTrue(rounds.stream()
274+
.anyMatch(round -> firstThinkingBlockId.equals(round.getThinkingBlock().getId())));
275+
}
276+
277+
@Test
278+
void testCacheConversationProgress_preservesWhitespaceOnlyThinkingFragments() throws Exception {
279+
ConversationPersistenceManager manager = createManagerWithRealDataFactory();
280+
String conversationId = "00000000-0000-0000-0000-000000000000";
281+
String turnId = "00000000-0000-0000-0000-000000000002";
282+
String thinkingBlockId = "thinking-block-1";
283+
ConversationData conversationData = createTestConversationData(conversationId);
284+
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
285+
286+
manager.cacheConversationProgress(conversationId,
287+
createThinkingProgressValue(conversationId, turnId, "before title."), thinkingBlockId).get();
288+
manager.cacheConversationProgress(conversationId,
289+
createThinkingProgressValue(conversationId, turnId, "\n"), thinkingBlockId).get();
290+
manager.cacheConversationProgress(conversationId,
291+
createThinkingProgressValue(conversationId, turnId, "**Next title**\n\nbody"), thinkingBlockId).get();
292+
293+
ThinkingBlockData block = getCachedCopilotTurn(manager, conversationId, turnId).getReply()
294+
.getEditAgentRounds().get(0).getThinkingBlock();
295+
assertNotNull(block);
296+
assertEquals("before title.\n**Next title**\n\nbody", block.getContent());
297+
}
298+
201299
@Test
202300
void testPersistConversationProgress_Success() throws Exception {
203301
String conversationId = "00000000-0000-0000-0000-000000000000";
@@ -206,7 +304,7 @@ void testPersistConversationProgress_Success() throws Exception {
206304

207305
when(mockPersistenceService.loadConversationFromPersistedJsonFile(conversationId)).thenReturn(conversationData);
208306

209-
CompletableFuture<Void> result = persistenceManager.persistConversationProgress(conversationId, progress);
307+
CompletableFuture<Void> result = persistenceManager.persistConversationProgress(conversationId, progress, null);
210308

211309
result.get(); // Wait for completion
212310

@@ -229,7 +327,7 @@ void testUpdateConversationProgress_NewConversation() throws Exception {
229327
assertNotNull(result);
230328
assertEquals(conversationId, result.getConversationId());
231329
verify(mockDataFactory).updateConversationMetadata(newConversationData, progress);
232-
verify(mockDataFactory).updateReplyFromProgress(any(), eq(progress));
330+
verify(mockDataFactory).updateReplyFromProgress(any(), eq(progress), isNull());
233331
}
234332

235333
// Helper methods to create test data
@@ -294,4 +392,50 @@ private ChatProgressValue createTestChatProgressValue() {
294392
progress.setSuggestedTitle("Test Suggested Title");
295393
return progress;
296394
}
395+
396+
private ConversationPersistenceManager createManagerWithRealDataFactory() throws Exception {
397+
ConversationPersistenceManager manager = new ConversationPersistenceManager(mockAuthStatusManager);
398+
setPrivateField(manager, "persistenceService", mockPersistenceService);
399+
return manager;
400+
}
401+
402+
private CopilotTurnData getCachedCopilotTurn(ConversationPersistenceManager manager, String conversationId,
403+
String turnId) throws Exception {
404+
ConversationData conversationData = manager.loadConversation(conversationId).get();
405+
assertNotNull(conversationData);
406+
return conversationData.getTurns().stream()
407+
.filter(CopilotTurnData.class::isInstance)
408+
.map(CopilotTurnData.class::cast)
409+
.filter(turn -> turnId.equals(turn.getTurnId()))
410+
.findFirst()
411+
.orElseThrow();
412+
}
413+
414+
private ChatProgressValue createThinkingProgressValue(String conversationId, String turnId, String thinkingText) {
415+
ChatProgressValue progress = new ChatProgressValue();
416+
progress.setKind(WorkDoneProgressKind.report);
417+
progress.setConversationId(conversationId);
418+
progress.setTurnId(turnId);
419+
progress.setThinking(new Thinking("server-thinking-id", thinkingText, null));
420+
return progress;
421+
}
422+
423+
private ChatProgressValue createAgentRoundProgressValue(String conversationId, String turnId, int roundId,
424+
String reply) throws Exception {
425+
ChatProgressValue progress = new ChatProgressValue();
426+
progress.setKind(WorkDoneProgressKind.report);
427+
progress.setConversationId(conversationId);
428+
progress.setTurnId(turnId);
429+
AgentRound round = new AgentRound();
430+
setPrivateField(round, "roundId", roundId);
431+
setPrivateField(round, "reply", reply);
432+
setPrivateField(progress, "editAgentRounds", List.of(round));
433+
return progress;
434+
}
435+
436+
private void setPrivateField(Object target, String fieldName, Object value) throws Exception {
437+
var field = target.getClass().getDeclaredField(fieldName);
438+
field.setAccessible(true);
439+
field.set(target, value);
440+
}
297441
}

‎com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/persistence/ConversationDataFactory.java‎

Lines changed: 84 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,14 @@
1919
import com.microsoft.copilot.eclipse.core.lsp.protocol.ChatCompletionContentPart;
2020
import com.microsoft.copilot.eclipse.core.lsp.protocol.ChatProgressValue;
2121
import com.microsoft.copilot.eclipse.core.lsp.protocol.ConversationError;
22+
import com.microsoft.copilot.eclipse.core.lsp.protocol.Thinking;
2223
import com.microsoft.copilot.eclipse.core.lsp.protocol.Turn;
2324
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.EditAgentRoundData;
2425
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ErrorData;
2526
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ErrorMessageData;
2627
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ReplyData;
28+
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ThinkingBlockData;
29+
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ThinkingBlockState;
2730
import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData.ToolCallData;
2831
import com.microsoft.copilot.eclipse.core.persistence.UserTurnData.MessageData;
2932

@@ -32,6 +35,8 @@
3235
* no business logic.
3336
*/
3437
public class ConversationDataFactory {
38+
private static final int SYNTHETIC_ROUND_ID = -1;
39+
3540
private final AuthStatusManager authStatusManager;
3641

3742
/**
@@ -95,12 +100,16 @@ public CopilotTurnData createCopilotTurnData(String turnId) {
95100
*
96101
* @param reply the reply data to update
97102
* @param progress the progress value to extract data from
103+
* @param thinkingBlockId the UI-generated thinking block ID, when the progress belongs to a thinking round
98104
*/
99-
public void updateReplyFromProgress(ReplyData reply, ChatProgressValue progress) {
105+
public void updateReplyFromProgress(ReplyData reply, ChatProgressValue progress, String thinkingBlockId) {
100106
ensureReplyInitialized(reply);
101107

108+
// Accumulate thinking content
109+
appendThinkingContent(reply, progress.getThinking(), thinkingBlockId);
110+
102111
// Merge agent rounds and append tool calls
103-
mergeAgentRounds(reply, progress.getAgentRounds());
112+
mergeAgentRounds(reply, progress.getAgentRounds(), thinkingBlockId);
104113

105114
// Handle plain reply streaming (no agentRounds)
106115
appendProgressReplyText(reply, progress.getReply());
@@ -117,27 +126,25 @@ public void updateReplyFromProgress(ReplyData reply, ChatProgressValue progress)
117126
reply.setHideText(progress.isHideText());
118127
}
119128

120-
private void mergeAgentRounds(ReplyData reply, List<AgentRound> agentRounds) {
129+
private void mergeAgentRounds(ReplyData reply, List<AgentRound> agentRounds, String thinkingBlockId) {
121130
if (agentRounds == null || agentRounds.isEmpty()) {
122131
return;
123132
}
133+
boolean thinkingRoundUpdated = false;
124134
for (AgentRound round : agentRounds) {
125-
EditAgentRoundData existingRound = findRoundById(reply.getEditAgentRounds(), round.getRoundId());
135+
EditAgentRoundData existingRound = thinkingRoundUpdated
136+
? null : findRoundByThinkingBlockId(reply.getEditAgentRounds(), thinkingBlockId);
137+
if (existingRound != null) {
138+
thinkingRoundUpdated = true;
139+
}
140+
if (existingRound == null) {
141+
existingRound = findRoundById(reply.getEditAgentRounds(), round.getRoundId());
142+
}
126143
if (existingRound == null) {
127144
EditAgentRoundData er = convertAgentRoundToEditAgentRoundData(round);
128145
reply.getEditAgentRounds().add(er);
129146
} else {
130-
appendReplyToAgentRound(existingRound, round.getReply());
131-
if (round.getToolCalls() != null && !round.getToolCalls().isEmpty()) {
132-
for (AgentToolCall tc : round.getToolCalls()) {
133-
ToolCallData existingToolCall = findToolCallById(existingRound.getToolCalls(), tc.getId());
134-
if (existingToolCall == null) {
135-
existingRound.getToolCalls().add(convertAgentToolCallToToolCallData(tc));
136-
} else {
137-
updateToolCallData(existingToolCall, tc);
138-
}
139-
}
140-
}
147+
updateAgentRoundData(existingRound, round);
141148
}
142149
}
143150
}
@@ -172,6 +179,19 @@ private void appendReplyToAgentRound(EditAgentRoundData round, String addition)
172179
round.setReply(existingReply + addition);
173180
}
174181

182+
private void appendThinkingContent(ReplyData reply, Thinking thinking, String thinkingBlockId) {
183+
// Preserve whitespace-only thinking fragments; they can carry markdown boundaries between sections.
184+
if (thinking == null || StringUtils.isEmpty(thinking.text())) {
185+
return;
186+
}
187+
if (StringUtils.isBlank(thinkingBlockId)) {
188+
return;
189+
}
190+
EditAgentRoundData round = getOrCreateThinkingRound(reply, thinkingBlockId);
191+
ThinkingBlockData block = round.getThinkingBlock();
192+
block.setContent(StringUtils.defaultString(block.getContent()) + thinking.text());
193+
}
194+
175195
private void applyConversationError(ReplyData reply, ConversationError error) {
176196
if (error == null) {
177197
return;
@@ -278,6 +298,55 @@ private String extractResponseFromCopilotTurnData(CopilotTurnData copilotTurnDat
278298
}
279299

280300
// Private helper methods for data transformation
301+
private EditAgentRoundData getOrCreateThinkingRound(ReplyData reply, String thinkingBlockId) {
302+
EditAgentRoundData existingRound = findRoundByThinkingBlockId(reply.getEditAgentRounds(), thinkingBlockId);
303+
if (existingRound != null) {
304+
return existingRound;
305+
}
306+
EditAgentRoundData round = new EditAgentRoundData();
307+
round.setRoundId(SYNTHETIC_ROUND_ID);
308+
round.setToolCalls(new ArrayList<>());
309+
round.setThinkingBlock(new ThinkingBlockData(thinkingBlockId, ""));
310+
reply.getEditAgentRounds().add(round);
311+
return round;
312+
}
313+
314+
private void updateAgentRoundData(EditAgentRoundData target, AgentRound source) {
315+
target.setRoundId(source.getRoundId());
316+
appendReplyToAgentRound(target, source.getReply());
317+
ThinkingBlockData thinkingBlock = target.getThinkingBlock();
318+
if (thinkingBlock != null && !thinkingBlock.isFinalized()) {
319+
thinkingBlock.setState(ThinkingBlockState.COMPLETED);
320+
}
321+
if (source.getToolCalls() == null || source.getToolCalls().isEmpty()) {
322+
return;
323+
}
324+
if (target.getToolCalls() == null) {
325+
target.setToolCalls(new ArrayList<>());
326+
}
327+
for (AgentToolCall toolCall : source.getToolCalls()) {
328+
ToolCallData existingToolCall = findToolCallById(target.getToolCalls(), toolCall.getId());
329+
if (existingToolCall == null) {
330+
target.getToolCalls().add(convertAgentToolCallToToolCallData(toolCall));
331+
} else {
332+
updateToolCallData(existingToolCall, toolCall);
333+
}
334+
}
335+
}
336+
337+
private EditAgentRoundData findRoundByThinkingBlockId(List<EditAgentRoundData> rounds, String thinkingBlockId) {
338+
if (rounds == null || StringUtils.isBlank(thinkingBlockId)) {
339+
return null;
340+
}
341+
for (EditAgentRoundData round : rounds) {
342+
ThinkingBlockData thinkingBlock = round.getThinkingBlock();
343+
if (thinkingBlock != null && thinkingBlockId.equals(thinkingBlock.getId())) {
344+
return round;
345+
}
346+
}
347+
return null;
348+
}
349+
281350
private EditAgentRoundData findRoundById(List<EditAgentRoundData> rounds, int roundId) {
282351
if (rounds == null) {
283352
return null;

0 commit comments

Comments
 (0)