diff --git a/src/main/ai/agentSession/AgentSessionRuntimeService.ts b/src/main/ai/agentSession/AgentSessionRuntimeService.ts index 256ad7cf44c..13b31388e88 100644 --- a/src/main/ai/agentSession/AgentSessionRuntimeService.ts +++ b/src/main/ai/agentSession/AgentSessionRuntimeService.ts @@ -60,7 +60,7 @@ export interface BeginAgentSessionTurnInput { agentId: string agentType: string modelId: UniqueModelId - assistantMessageId?: string + assistantMessageId: string userMessage?: AgentSessionMessageEntity /** Container-level OTel trace id (one trace per session); cached on the entry. */ traceId?: string @@ -69,6 +69,7 @@ export interface BeginAgentSessionTurnInput { export interface AgentSessionRuntimeHandle { listeners: StreamListener[] turnId: string + abortController: AbortController } export interface OpenAgentSessionTurnStreamInput { @@ -90,10 +91,11 @@ export interface AgentSessionRuntimeSnapshot { type AgentSessionTurn = { turnId: string - assistantMessageId?: string + assistantMessageId: string userMessage: AgentSessionMessageEntity modelId: UniqueModelId admitted: boolean + abortController: AbortController terminalStatus?: AgentSessionRuntimeTerminalStatus controller?: ReadableStreamDefaultController activeToolIds: Set @@ -211,6 +213,7 @@ export class AgentSessionRuntimeService extends BaseService { userMessage, modelId: input.modelId, admitted: false, + abortController: new AbortController(), activeToolIds: new Set() } @@ -231,7 +234,8 @@ export class AgentSessionRuntimeService extends BaseService { new AgentSessionRuntimeTerminalListener(this, input.sessionId), new TraceFlushListener(input.topicId) ], - turnId + turnId, + abortController: turn.abortController } } @@ -256,7 +260,8 @@ export class AgentSessionRuntimeService extends BaseService { new AgentSessionRuntimeTerminalListener(this, input.sessionId), new TraceFlushListener(input.topicId) ], - turnId + turnId, + abortController: turn.abortController } } @@ -572,6 +577,13 @@ export class AgentSessionRuntimeService extends BaseService { return toolApprovalRegistry.dispatch(approvalId, decision) } + abortPendingTurn(sessionId: string, reason: string): boolean { + const turn = this.entries.get(sessionId)?.currentTurn + if (!turn || turn.terminalStatus || turn.abortController.signal.aborted) return false + turn.abortController.abort(reason) + return true + } + protected onStop(): void { this.closeAll() toolApprovalRegistry.clear('agent-session-runtime-stop') @@ -915,6 +927,7 @@ export class AgentSessionRuntimeService extends BaseService { userMessage: nextMessage, modelId: entry.modelId, admitted: false, + abortController: new AbortController(), activeToolIds: new Set() } @@ -939,6 +952,7 @@ export class AgentSessionRuntimeService extends BaseService { messages, runtime: { kind: 'agent-session', sessionId: entry.sessionId, turnId } }, + abortController: entry.currentTurn.abortController, listeners: [ this.createPersistenceListener(entry, nextMessage), new AgentSessionRuntimeTerminalListener(this, entry.sessionId), @@ -1005,6 +1019,7 @@ export class AgentSessionRuntimeService extends BaseService { modelId: entry.modelId, // Pre-admitted: the steer was already delivered via the hook, so `admitTurn` must NOT re-send it. admitted: true, + abortController: new AbortController(), activeToolIds: new Set() } @@ -1029,6 +1044,7 @@ export class AgentSessionRuntimeService extends BaseService { messages, runtime: { kind: 'agent-session', sessionId: entry.sessionId, turnId } }, + abortController: entry.currentTurn.abortController, listeners: [ this.createPersistenceListener(entry, steerMessage), new AgentSessionRuntimeTerminalListener(this, entry.sessionId), @@ -1089,12 +1105,18 @@ export class AgentSessionRuntimeService extends BaseService { entry: AgentSessionRuntimeEntry, userMessage: AgentSessionMessageEntity ): StreamListener { + const currentTurn = entry.currentTurn + if (!currentTurn) { + throw new Error(`Cannot create persistence listener without an active turn: ${entry.sessionId}`) + } + const { assistantMessageId } = currentTurn const userText = extractMessageText(userMessage) return new PersistenceListener({ topicId: entry.topicId, modelId: entry.modelId, backend: new AgentSessionMessageBackend({ sessionId: entry.sessionId, + assistantMessageId, modelId: entry.modelId, runtimeResumeToken: () => entry.lastResumeToken, afterPersist: async (finalMessage) => { diff --git a/src/main/ai/agentSession/__tests__/AgentSessionRuntimeService.test.ts b/src/main/ai/agentSession/__tests__/AgentSessionRuntimeService.test.ts index e0ee565d76a..2376b1d983d 100644 --- a/src/main/ai/agentSession/__tests__/AgentSessionRuntimeService.test.ts +++ b/src/main/ai/agentSession/__tests__/AgentSessionRuntimeService.test.ts @@ -238,6 +238,32 @@ describe('AgentSessionRuntimeService', () => { }) }) + it('aborts the current turn controller before the stream starts', () => { + const service = new AgentSessionRuntimeService() + const handle = service.beginTurn(baseTurnInput) + + expect(service.abortPendingTurn('session-1', 'user-requested')).toBe(true) + expect(handle.abortController.signal.aborted).toBe(true) + expect(handle.abortController.signal.reason).toBe('user-requested') + }) + + it('does not reuse an aborted controller for a later turn', () => { + const service = new AgentSessionRuntimeService() + const first = service.beginTurn(baseTurnInput) + + expect(service.abortPendingTurn('session-1', 'user-requested')).toBe(true) + void terminalListener(first).onPaused({ status: 'paused', isTopicDone: true }) + + const second = service.beginTurn({ + ...baseTurnInput, + assistantMessageId: 'assistant-2', + userMessage: userMessage('user-2') + }) + + expect(first.abortController.signal.aborted).toBe(true) + expect(second.abortController.signal.aborted).toBe(false) + }) + it('marks the runtime idle when the terminal listener observes done', () => { const service = new AgentSessionRuntimeService() const handle = service.beginTurn(baseTurnInput) @@ -616,6 +642,30 @@ describe('AgentSessionRuntimeService', () => { }) }) + it('persists empty paused terminals to the active assistant placeholder', async () => { + const service = new AgentSessionRuntimeService() + const handle = service.beginTurn({ ...baseTurnInput, userMessage: userMessage('user-1') }) + getEntry(service).lastResumeToken = 'resume-1' + + await persistenceListener(handle).onPaused({ + status: 'paused', + isTopicDone: true, + finalMessage: undefined + }) + + expect(mocks.saveMessage).toHaveBeenCalledWith({ + sessionId: 'session-1', + runtimeResumeToken: 'resume-1', + message: { + id: 'assistant-1', + role: 'assistant', + status: 'paused', + data: { parts: [] }, + modelId: 'claude-code::claude-sonnet-4-5' + } + }) + }) + it('routes runtime events from the selected driver into the active turn', async () => { const events = createAsyncQueue() const connection = { @@ -1627,6 +1677,7 @@ describe('AgentSessionRuntimeService', () => { ], runtime: { kind: 'agent-session', sessionId: 'session-1', turnId: expect.any(String) } }, + abortController: expect.any(AbortController), listeners: [ expect.objectContaining({ id: expect.stringContaining('persistence:agents-db:') }), expect.objectContaining({ id: 'agent-runtime:session-1' }), diff --git a/src/main/ai/agentSession/persistence/AgentSessionMessageBackend.ts b/src/main/ai/agentSession/persistence/AgentSessionMessageBackend.ts index ec54827205c..26c61ef9b19 100644 --- a/src/main/ai/agentSession/persistence/AgentSessionMessageBackend.ts +++ b/src/main/ai/agentSession/persistence/AgentSessionMessageBackend.ts @@ -10,13 +10,14 @@ import { agentSessionMessageService } from '@data/services/AgentSessionMessageService' import type { CherryMessagePart, CherryUIMessage } from '@shared/data/types/message' import type { UniqueModelId } from '@shared/data/types/model' -import { v7 as uuidv7 } from 'uuid' import { finalizeInterruptedParts, type PersistAssistantInput, type PersistenceBackend } from '../../streamManager' export interface AgentSessionMessageBackendOptions { /** Cherry Studio agent-session id. */ sessionId: string + /** Existing assistant placeholder id to finalize. */ + assistantMessageId: string /** Model id used for this assistant message. */ modelId?: UniqueModelId /** Opaque runtime resume token persisted for future recovery; `undefined` when unknown. */ @@ -27,6 +28,7 @@ export interface AgentSessionMessageBackendOptions { export class AgentSessionMessageBackend implements PersistenceBackend { readonly kind = 'agents-db' + readonly canPersistEmptyTerminal = true readonly afterPersist?: (finalMessage: CherryUIMessage) => Promise constructor(private readonly opts: AgentSessionMessageBackendOptions) { @@ -41,7 +43,7 @@ export class AgentSessionMessageBackend implements PersistenceBackend { sessionId: this.opts.sessionId, ...(runtimeResumeToken ? { runtimeResumeToken } : {}), message: { - id: finalMessage?.id ?? uuidv7(), + id: finalMessage?.id ?? this.opts.assistantMessageId, role: 'assistant', status, data: { parts }, diff --git a/src/main/ai/streamManager/AiStreamManager.ts b/src/main/ai/streamManager/AiStreamManager.ts index 8c1158ee8cf..806c3ca7068 100644 --- a/src/main/ai/streamManager/AiStreamManager.ts +++ b/src/main/ai/streamManager/AiStreamManager.ts @@ -19,7 +19,7 @@ import type { UniqueModelId } from '@shared/data/types/model' import type { SerializedError } from '@shared/types/error' import { type UIMessageChunk } from 'ai' -import { isAgentSessionTopic } from '../agentSession/topic' +import { extractAgentSessionId, isAgentSessionTopic } from '../agentSession/topic' import { applyTurnOutputAttributes } from '../observability' import type { AiStreamRequest, CallOverrides } from '../types' import { buildCompactReplay } from './buildCompactReplay' @@ -79,6 +79,7 @@ export interface SendModelSpec { modelId: UniqueModelId request: AiStreamRequest rootSpan?: Span + abortController?: AbortController } export interface SendInput { @@ -105,6 +106,7 @@ export interface StartRuntimeTurnInput { request: AiStreamRequest listeners: StreamListener[] rootSpan?: Span + abortController?: AbortController } // ── Inspection snapshots ──────────────────────────────────────────── @@ -356,11 +358,18 @@ export class AiStreamManager extends BaseService { const isMultiModel = input.models.length > 1 const executions = new Map() - for (const { modelId, request, rootSpan } of input.models) { + for (const { modelId, request, rootSpan, abortController } of input.models) { if (executions.has(modelId)) { throw new Error(`send() got duplicate modelId ${modelId} for topic ${input.topicId}`) } - const exec = this.createAndLaunchExecution(input.topicId, modelId, request, input.siblingsGroupId, rootSpan) + const exec = this.createAndLaunchExecution( + input.topicId, + modelId, + request, + input.siblingsGroupId, + rootSpan, + abortController + ) executions.set(modelId, exec) } @@ -380,6 +389,10 @@ export class AiStreamManager extends BaseService { // Chat broadcasts to SharedCache so `useChatWithHistory.resumeActiveStream` can attach; prompt is silent. stream.lifecycle.onCreated(stream) + if ([...executions.values()].every((exec) => exec.abortController.signal.aborted)) { + stream.status = 'aborted' + } + return { mode: 'started', executionIds: input.models.map((m) => m.modelId) @@ -437,7 +450,14 @@ export class AiStreamManager extends BaseService { return this.send({ topicId: input.topicId, - models: [{ modelId: input.modelId, request: input.request, rootSpan: input.rootSpan }], + models: [ + { + modelId: input.modelId, + request: input.request, + rootSpan: input.rootSpan, + abortController: input.abortController + } + ], listeners: [...carriedListeners, ...input.listeners] }) } @@ -541,7 +561,12 @@ export class AiStreamManager extends BaseService { /** Abort all executions in a topic. */ abort(topicId: string, reason: string): void { const stream = this.activeStreams.get(topicId) - if (!stream || !isLiveStatus(stream.status)) return + if (!stream || !isLiveStatus(stream.status)) { + if (isAgentSessionTopic(topicId)) { + application.get('AgentSessionRuntimeService').abortPendingTurn(extractAgentSessionId(topicId), reason) + } + return + } logger.info('Aborting stream', { topicId, reason }) for (const exec of stream.executions.values()) { if (exec.status === 'streaming') { @@ -961,14 +986,15 @@ export class AiStreamManager extends BaseService { modelId: UniqueModelId, request: AiStreamRequest, siblingsGroupId?: number, - rootSpan?: Span + rootSpan?: Span, + abortController?: AbortController ): StreamExecution { // `loopPromise` is overwritten right after launch; initialise to a resolved sentinel // so the `exec` object reference is stable inside the arrow function below. const exec: StreamExecution = { modelId, anchorMessageId: request.messageId, - abortController: new AbortController(), + abortController: abortController ?? new AbortController(), status: 'streaming', buffer: [], droppedChunks: 0, diff --git a/src/main/ai/streamManager/__tests__/AiStreamManager.test.ts b/src/main/ai/streamManager/__tests__/AiStreamManager.test.ts index 42e0b0b4ed8..350b50418af 100644 --- a/src/main/ai/streamManager/__tests__/AiStreamManager.test.ts +++ b/src/main/ai/streamManager/__tests__/AiStreamManager.test.ts @@ -58,6 +58,8 @@ class FakeListener implements StreamListener { // ── Mocks ─────────────────────────────────────────────────────────── +const mockAbortPendingTurn = vi.fn<(sessionId: string, reason: string) => boolean>(() => false) + vi.mock('@main/data/services/MessageService', () => ({ messageService: { create: vi.fn().mockResolvedValue({ id: 'msg-001' }) } })) @@ -129,7 +131,7 @@ vi.mock('@application', async () => { AiService: { streamText: mockStreamText }, CacheService: fakeCacheService, TraceStorageService: { saveSpans: mockSaveSpans }, - AgentSessionRuntimeService: { willContinueTopic: mockWillContinueTopic } + AgentSessionRuntimeService: { willContinueTopic: mockWillContinueTopic, abortPendingTurn: mockAbortPendingTurn } } as Parameters[0]) }) @@ -175,11 +177,12 @@ function startSingle( request: AiStreamRequest listeners: StreamListener[] siblingsGroupId?: number + abortController?: AbortController } ) { manager.send({ topicId: opts.topicId, - models: [{ modelId: opts.modelId, request: opts.request }], + models: [{ modelId: opts.modelId, request: opts.request, abortController: opts.abortController }], listeners: opts.listeners, siblingsGroupId: opts.siblingsGroupId }) @@ -202,6 +205,7 @@ describe('AiStreamManager', () => { ) mockSaveSpans.mockResolvedValue(undefined) mockWillContinueTopic.mockReturnValue(false) + mockAbortPendingTurn.mockReturnValue(false) sharedCacheStore.clear() }) @@ -258,6 +262,53 @@ describe('AiStreamManager', () => { expect(mgr.inspect('a')).toBeUndefined() }) + it('aborts the agent-session turn controller for a pre-stream stop request', async () => { + const turnAbortController = new AbortController() + mockAbortPendingTurn.mockImplementationOnce((_sessionId, reason) => { + turnAbortController.abort(reason) + return true + }) + const listener = new FakeListener('l:agent') + + mgr.abort('agent-session:session-1', 'user-requested') + const snap = startSingle(mgr, { + topicId: 'agent-session:session-1', + modelId: 'provider-a::model-a', + request: { ...req('agent-session:session-1'), messageId: 'assistant-paused' }, + listeners: [listener], + abortController: turnAbortController + }) + + expect(mockAbortPendingTurn).toHaveBeenCalledWith('session-1', 'user-requested') + expect(snap.status).toBe('aborted') + + await vi.advanceTimersByTimeAsync(0) + expect(listener.pausedResults).toHaveLength(1) + }) + + it('does not apply an old pre-stream stop request to a new agent-session turn controller', () => { + const oldTurnAbortController = new AbortController() + const newTurnAbortController = new AbortController() + mockAbortPendingTurn.mockImplementationOnce((_sessionId, reason) => { + oldTurnAbortController.abort(reason) + return true + }) + + mgr.abort('agent-session:session-1', 'user-requested') + const snap = startSingle(mgr, { + topicId: 'agent-session:session-1', + modelId: 'provider-a::model-a', + request: { ...req('agent-session:session-1'), messageId: 'assistant-new' }, + listeners: [new FakeListener('l:agent')], + abortController: newTurnAbortController + }) + + expect(snap.status).toBe('pending') + expect(snap.executions[0].abortSignal.aborted).toBe(false) + expect(oldTurnAbortController.signal.aborted).toBe(true) + expect(newTurnAbortController.signal.aborted).toBe(false) + }) + it('evicts finished stream and creates new one', async () => { startSingle(mgr, { topicId: 'a', diff --git a/src/main/ai/streamManager/context/AgentChatContextProvider.ts b/src/main/ai/streamManager/context/AgentChatContextProvider.ts index 3d26fc4087d..676a46a8665 100644 --- a/src/main/ai/streamManager/context/AgentChatContextProvider.ts +++ b/src/main/ai/streamManager/context/AgentChatContextProvider.ts @@ -204,7 +204,8 @@ export class AgentChatContextProvider implements ChatContextProvider { messageId: assistantMessageId, runtime: { kind: 'agent-session', sessionId, turnId: runtime.turnId } }, - rootSpan: turnTrace.rootSpan + rootSpan: turnTrace.rootSpan, + abortController: runtime.abortController } ], userMessageId, diff --git a/src/main/ai/streamManager/context/ChatContextProvider.ts b/src/main/ai/streamManager/context/ChatContextProvider.ts index 8089b0dd070..9cc9c71cafd 100644 --- a/src/main/ai/streamManager/context/ChatContextProvider.ts +++ b/src/main/ai/streamManager/context/ChatContextProvider.ts @@ -16,7 +16,12 @@ import type { MainDispatchRequest } from './dispatch' export interface PreparedDispatch { topicId: string - models: ReadonlyArray<{ modelId: UniqueModelId; request: AiStreamRequest; rootSpan?: Span }> + models: ReadonlyArray<{ + modelId: UniqueModelId + request: AiStreamRequest + rootSpan?: Span + abortController?: AbortController + }> listeners: StreamListener[] /** DB id of the user message row this dispatch created, surfaced back to renderer for optimistic join. */ userMessageId?: string diff --git a/src/main/ai/streamManager/listeners/PersistenceListener.ts b/src/main/ai/streamManager/listeners/PersistenceListener.ts index 2c0e6f58759..1b4a9d143fa 100644 --- a/src/main/ai/streamManager/listeners/PersistenceListener.ts +++ b/src/main/ai/streamManager/listeners/PersistenceListener.ts @@ -102,7 +102,7 @@ export class PersistenceListener implements StreamListener { status: 'success' | 'paused' | 'error', transportTimings: TransportTimings | undefined ): void { - if (!finalMessage && status !== 'error') { + if (!finalMessage && (status === 'success' || !this.opts.backend.canPersistEmptyTerminal)) { logger.warn('Terminal event without finalMessage, skipping persistence', { backend: this.opts.backend.kind, topicId: this.opts.topicId, diff --git a/src/main/ai/streamManager/listeners/__tests__/PersistenceListener.test.ts b/src/main/ai/streamManager/listeners/__tests__/PersistenceListener.test.ts index 3844afcf2fe..734707c6185 100644 --- a/src/main/ai/streamManager/listeners/__tests__/PersistenceListener.test.ts +++ b/src/main/ai/streamManager/listeners/__tests__/PersistenceListener.test.ts @@ -366,6 +366,18 @@ describe('PersistenceListener + TemporaryChatBackend', () => { expect(appendMessageMock).not.toHaveBeenCalled() }) + it('skips persistence when onPaused arrives without a finalMessage and there is no placeholder row', async () => { + const listener = makeListener() + + await listener.onPaused({ + finalMessage: undefined, + status: 'paused', + timings: { startedAt: 1000, completedAt: 2500.9 } + }) + + expect(appendMessageMock).not.toHaveBeenCalled() + }) + it('swallows append errors so stream teardown is not disrupted', async () => { appendMessageMock.mockImplementationOnce(() => { throw new Error('write failed') diff --git a/src/main/ai/streamManager/persistence/PersistenceBackend.ts b/src/main/ai/streamManager/persistence/PersistenceBackend.ts index 2c699ab280b..b6476ca19c9 100644 --- a/src/main/ai/streamManager/persistence/PersistenceBackend.ts +++ b/src/main/ai/streamManager/persistence/PersistenceBackend.ts @@ -93,6 +93,12 @@ export interface PersistenceBackend { /** Tag for logging (e.g. "sqlite", "temp", "agents-db"). */ readonly kind: string + /** + * True for backends that finalize a pre-created placeholder row. They must + * still write terminal status when a stream is paused before producing chunks. + */ + readonly canPersistEmptyTerminal?: boolean + persistAssistant(input: PersistAssistantInput): void /**