diff --git a/.changeset/fork-at-turn.md b/.changeset/fork-at-turn.md new file mode 100644 index 00000000000..957cb2c0302 --- /dev/null +++ b/.changeset/fork-at-turn.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": minor +--- + +The `:fork` server endpoint accepts a `prompt_id` to fork the session from an earlier prompt instead of the latest state. diff --git a/docs/en/reference/server-api.md b/docs/en/reference/server-api.md index 0bdd3f54067..1db2c917d00 100644 --- a/docs/en/reference/server-api.md +++ b/docs/en/reference/server-api.md @@ -708,16 +708,20 @@ Session actions are dispatched through one route: the path tail is parsed as `{s #### `POST /api/v1/sessions/{session_id}:fork` -Copies the session — its transcript, agent state, and files — into a new session in the same workspace, and broadcasts `event.session.created`. Forking is rejected while any of the session's agents has an active turn. +Copies the session — its transcript, agent state, and files — into a new session in the same workspace, and broadcasts `event.session.created`. Pass `prompt_id` to fork from an earlier point instead of the latest state: the copy keeps only the conversation through that prompt's turn segment and discards everything after it. + +Forking is rejected while the session has an active or queued turn, or while another operation (compaction, undo, another fork) is in progress; retry once the session is idle. | Parameter | In | Type | Description | | --- | --- | --- | --- | | `title` | body | string | Title for the fork (at least 1 character). Default `Fork: ` | | `metadata` | body | object | Custom metadata for the fork | +| `prompt_id` | body | string | Id of the user prompt to fork from. Default: the latest state | On success, `data` is [the session object](#the-session-object) of the new session. -- `40901`: the session has an active turn and cannot be forked +- `40001`: `prompt_id` is unknown in the session +- `40901`: a turn is active or another operation is in progress, so the session cannot be forked right now #### `POST /api/v1/sessions/{session_id}:compact` @@ -787,7 +791,7 @@ On success, `data` is `{ items, has_more }` where each item is [the session obje #### `POST /api/v1/sessions/{session_id}/children` -Creates a child session: a fork of this session recorded as its child, so it shows up under `GET /api/v1/sessions/{session_id}/children`. The same active-turn restriction as `:fork` applies. +Creates a child session: a fork of this session recorded as its child, so it shows up under `GET /api/v1/sessions/{session_id}/children`. The same rules as `:fork` apply: it is rejected while a turn is active or another operation is in progress. | Parameter | In | Type | Description | | --- | --- | --- | --- | @@ -797,7 +801,7 @@ Creates a child session: a fork of this session recorded as its child, so it sho On success, `data` is [the session object](#the-session-object) of the new session, and the server broadcasts `event.session.created`. -- `40901`: the session has an active turn and cannot be forked +- `40901`: a turn is active or another operation is in progress, so the session cannot be forked right now #### `GET /api/v1/sessions/{session_id}/status` diff --git a/docs/zh/reference/server-api.md b/docs/zh/reference/server-api.md index 32cad85a455..6e0be2516e0 100644 --- a/docs/zh/reference/server-api.md +++ b/docs/zh/reference/server-api.md @@ -708,16 +708,20 @@ schema 还接受 `agent_config` 内的 `system_prompt`、`tools`、`mcp_servers` #### `POST /api/v1/sessions/{session_id}:fork` -将会话——其转录、Agent 状态与文件——复制到同一工作区中的新会话,并广播 `event.session.created`。当会话中任一 Agent 有进行中的轮次时,fork 会被拒绝。 +将会话——其转录、Agent 状态与文件——复制到同一工作区中的新会话,并广播 `event.session.created`。传入 `prompt_id` 可从更早的位置分叉,而不是复制最新状态:副本只保留到该提示词所在轮次段结束为止的对话,其后的内容全部丢弃。 + +会话有进行中或排队中的轮次,或者有其他操作(压缩、撤销、另一个 fork)在进行时,fork 会被拒绝;等会话空闲后重试。 | 参数 | 位置 | 类型 | 说明 | | --- | --- | --- | --- | | `title` | body | string | fork 的标题(至少 1 个字符)。默认 `Fork: ` | | `metadata` | body | object | fork 的自定义元数据 | +| `prompt_id` | body | string | 要从中分叉的用户提示词 id。默认为最新状态 | 成功时,`data` 为新会话的 [session 对象](#session-对象)。 -- `40901`:会话有进行中的轮次,无法 fork +- `40001`:`prompt_id` 在会话中不存在 +- `40901`:有轮次正在进行或有其他操作在进行,暂时无法 fork #### `POST /api/v1/sessions/{session_id}:compact` @@ -787,7 +791,7 @@ schema 还接受 `agent_config` 内的 `system_prompt`、`tools`、`mcp_servers` #### `POST /api/v1/sessions/{session_id}/children` -创建子会话:fork 当前会话并记录为其子会话,因此会出现在 `GET /api/v1/sessions/{session_id}/children` 下。适用与 `:fork` 相同的进行中轮次限制。 +创建子会话:fork 当前会话并记录为其子会话,因此会出现在 `GET /api/v1/sessions/{session_id}/children` 下。适用与 `:fork` 相同的规则:有轮次正在进行或有其他操作在进行时会被拒绝。 | 参数 | 位置 | 类型 | 说明 | | --- | --- | --- | --- | @@ -797,7 +801,7 @@ schema 还接受 `agent_config` 内的 `system_prompt`、`tools`、`mcp_servers` 成功时,`data` 为新会话的 [session 对象](#session-对象),并且服务端广播 `event.session.created`。 -- `40901`:会话有进行中的轮次,无法 fork +- `40901`:有轮次正在进行或有其他操作在进行,暂时无法 fork #### `GET /api/v1/sessions/{session_id}/status` diff --git a/packages/agent-core-v2/src/agent/loop/loopService.ts b/packages/agent-core-v2/src/agent/loop/loopService.ts index 2c2d69c6def..5a4ef78591b 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -151,6 +151,8 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { private machineTurnSuppressed = false; private readonly settleWaiters: Array<() => void> = []; private quiescenceDepth = 0; + private pendingTurnCloses = 0; + private pendingFinalizations = 0; private activeRequestTrace: LLMRequestTrace | undefined; private engine: MachineEngine | undefined; @@ -768,7 +770,8 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { this.quiescenceDepth > 0 || this.active !== undefined || this.hasPendingRequests() || - this.pendingMachineTurn !== undefined + this.pendingMachineTurn !== undefined || + this.pendingTurnCloses > 0 ) { return undefined; } @@ -911,7 +914,8 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { if ( this.active === undefined && !this.hasPendingRequests() && - this.pendingMachineTurn === undefined + this.pendingMachineTurn === undefined && + this.pendingFinalizations === 0 ) { return Promise.resolve(); } @@ -924,6 +928,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { if ( this.active !== undefined || this.pendingMachineTurn !== undefined || + this.pendingFinalizations > 0 || this.hasPendingRequests() ) return; if (this.settleWaiters.length === 0) return; @@ -2030,75 +2035,86 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { private async endTurn(turn: ActiveTurn, result: TurnResult): Promise { if (this.active !== turn) return; this.active = undefined; - await this.wire.drainPersisted().catch(() => undefined); - for (const nudge of this.nudges.slice(this.nudgeCursor)) { - if (nudge.turnScoped && !nudge.dropped) { - nudge.dropped = true; - nudge.onDrop?.(); + this.pendingTurnCloses += 1; + this.pendingFinalizations += 1; + try { + await this.wire.drainPersisted().catch(() => undefined); + for (const nudge of this.nudges.slice(this.nudgeCursor)) { + if (nudge.turnScoped && !nudge.dropped) { + nudge.dropped = true; + nudge.onDrop?.(); + } } - } - turn.turn.state = result.type; - if (!turn.readyResolved) { - if (result.type === 'failed') { - turn.ready.reject(result.error); - } else if (result.type === 'cancelled') { - turn.ready.reject( - result.reason instanceof Error ? result.reason : abortError('Turn cancelled'), - ); - } else { - turn.ready.reject(new Error2(ErrorCodes.INTERNAL, 'Turn ended before first step')); + turn.turn.state = result.type; + if (!turn.readyResolved) { + if (result.type === 'failed') { + turn.ready.reject(result.error); + } else if (result.type === 'cancelled') { + turn.ready.reject( + result.reason instanceof Error ? result.reason : abortError('Turn cancelled'), + ); + } else { + turn.ready.reject(new Error2(ErrorCodes.INTERNAL, 'Turn ended before first step')); + } } - } - const durationMs = Date.now() - turn.startedAt; - const traceId = - result.type === 'completed' ? this.lastRequestTraceId : this.activeRequestTrace?.traceId; - const error = result.type === 'failed' ? toKimiErrorPayload(result.error) : undefined; - const interruptReason = result.type === 'completed' ? undefined : interruptReasonFor(result); - void this.dispatcher.dispatch( - new TurnEnded({ - agentId: this.scopeContext.agentId, - turnId: turn.id, - reason: result.type, - error, - durationMs, - interruptReason, - stopReason: result.type === 'completed' ? result.stopReason : undefined, - traceId, - }), - ); - if (error !== undefined) { + const durationMs = Date.now() - turn.startedAt; + const traceId = + result.type === 'completed' ? this.lastRequestTraceId : this.activeRequestTrace?.traceId; + const error = result.type === 'failed' ? toKimiErrorPayload(result.error) : undefined; + const interruptReason = result.type === 'completed' ? undefined : interruptReasonFor(result); void this.dispatcher.dispatch( - new AgentErrorEvent({ ...error, agentId: this.scopeContext.agentId }), + new TurnEnded({ + agentId: this.scopeContext.agentId, + turnId: turn.id, + reason: result.type, + error, + durationMs, + interruptReason, + stopReason: result.type === 'completed' ? result.stopReason : undefined, + traceId, + }), ); - } - if (interruptReason !== undefined) { - const interrupted: TurnInterruptedEvent = { + if (error !== undefined) { + void this.dispatcher.dispatch( + new AgentErrorEvent({ ...error, agentId: this.scopeContext.agentId }), + ); + } + if (interruptReason !== undefined) { + const interrupted: TurnInterruptedEvent = { + turn_id: turn.id, + at_step: result.steps, + mode: turn.mode ?? 'agent', + interrupt_reason: interruptReason, + provider_type: turn.providerType, + protocol: turn.protocol, + trace_id: traceId, + }; + this.telemetry.track2('turn_interrupted', interrupted); + } + const ended: TurnEndedTelemetryEvent = { turn_id: turn.id, - at_step: result.steps, + reason: result.type, + duration_ms: durationMs, mode: turn.mode ?? 'agent', - interrupt_reason: interruptReason, + error_type: error?.code, provider_type: turn.providerType, protocol: turn.protocol, trace_id: traceId, + enabled_plugins: this.plugins.enabledPluginIds()?.join(','), }; - this.telemetry.track2('turn_interrupted', interrupted); - } - const ended: TurnEndedTelemetryEvent = { - turn_id: turn.id, - reason: result.type, - duration_ms: durationMs, - mode: turn.mode ?? 'agent', - error_type: error?.code, - provider_type: turn.providerType, - protocol: turn.protocol, - trace_id: traceId, - enabled_plugins: this.plugins.enabledPluginIds()?.join(','), - }; - this.telemetry.track2('turn_ended', ended); - this.telemetry.setContext({ turn_id: undefined, trace_id: undefined, thinking_effort: undefined }); - this.activeRequestTrace = undefined; - this.lastRequestTraceId = undefined; - turn.result.resolve(result); + this.telemetry.track2('turn_ended', ended); + this.telemetry.setContext({ turn_id: undefined, trace_id: undefined, thinking_effort: undefined }); + this.activeRequestTrace = undefined; + this.lastRequestTraceId = undefined; + } finally { + this.pendingTurnCloses -= 1; + } + try { + await this.dispatcher.flush().catch(() => undefined); + turn.result.resolve(result); + } finally { + this.pendingFinalizations -= 1; + } this.maybeSettle(); } diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/internal/forkTurnSlice.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/internal/forkTurnSlice.ts index 559346caba2..eb9d93ca91d 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/internal/forkTurnSlice.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/internal/forkTurnSlice.ts @@ -11,6 +11,7 @@ export interface MainTurnSlice { readonly records: readonly WireRecord[]; readonly cutoffTime?: number; readonly lastPrompt?: string; + readonly lastTurnReason?: 'completed' | 'cancelled' | 'failed'; } export function assertForkTurnIndex(turnIndex: number | undefined): void { @@ -28,10 +29,7 @@ export function sliceMainRecordsAtTurn( sourceSessionId: string, turnIndex: number, ): MainTurnSlice { - const turnStarts: number[] = []; - for (let index = 0; index < records.length; index += 1) { - if (isUserVisibleTurnRecord(records[index]!)) turnStarts.push(index); - } + const turnStarts = userVisibleTurnStartIndices(records); const start = turnStarts[turnIndex]; if (start === undefined) { throw new Error2( @@ -58,6 +56,7 @@ export function sliceMainRecordsAtTurn( records: retained, cutoffTime: cutoffTimes.length === 0 ? undefined : Math.max(...cutoffTimes), lastPrompt, + lastTurnReason: turnOutcomeOfRecords(retained), }; } @@ -77,6 +76,45 @@ export function sliceSubagentRecordsAtTime( return records.slice(0, end); } +export type ForkPromptResolution = + | { readonly status: 'found'; readonly index: number } + | { readonly status: 'unknown' } + | { readonly status: 'ambiguous' }; + +export function resolveForkPromptIndex( + records: readonly WireRecord[], + promptId: string, +): ForkPromptResolution { + const turnStarts = userVisibleTurnStartIndices(records); + const targetIds = new Set([promptId]); + for (const record of records) { + if (record.type !== 'turn.steer') continue; + const promptIds = record['promptIds']; + if (!Array.isArray(promptIds) || !promptIds.includes(promptId)) continue; + const messageId = record['messageId']; + if (typeof messageId === 'string' && messageId.length > 0) targetIds.add(messageId); + } + let found: number | undefined; + let ambiguous = false; + for (let start = 0; start < turnStarts.length; start += 1) { + const message = asRecord(records[turnStarts[start]!]!['message']); + const id = message?.['id']; + if (typeof id !== 'string' || !targetIds.has(id)) continue; + if (found !== undefined && found !== start) ambiguous = true; + found = found ?? start; + } + if (ambiguous) return { status: 'ambiguous' }; + return found === undefined ? { status: 'unknown' } : { status: 'found', index: found }; +} + +function userVisibleTurnStartIndices(records: readonly WireRecord[]): number[] { + const turnStarts: number[] = []; + for (let index = 0; index < records.length; index += 1) { + if (isUserVisibleTurnRecord(records[index]!)) turnStarts.push(index); + } + return turnStarts; +} + function isUserVisibleTurnRecord(record: WireRecord): boolean { if (record.type !== 'context.append_message') return false; const message = asRecord(record['message']); @@ -203,6 +241,30 @@ function recordTime(record: WireRecord): number | undefined { return undefined; } +function turnOutcomeOfRecords( + records: readonly WireRecord[], +): 'completed' | 'cancelled' | 'failed' | undefined { + for (let index = records.length - 1; index >= 0; index -= 1) { + const record = records[index]!; + if (record.type === 'turn.ended') { + const reason = record['reason']; + if (reason === 'completed' || reason === 'cancelled') return reason; + if (reason === 'failed' || reason === 'blocked') return 'failed'; + return undefined; + } + if (isEngineTurnStartRecord(record) || isUserVisibleTurnRecord(record)) return undefined; + } + return undefined; +} + +function isEngineTurnStartRecord(record: WireRecord): boolean { + return ( + record.type === 'turn.prompt' || + record.type === 'turn.steer' || + record.type === 'agent.turn.started' + ); +} + function promptMetadataFromTurnRecord(record: WireRecord): string | undefined { if (record.type !== 'context.append_message') return undefined; const message = asRecord(record['message']); diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycle.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycle.ts index 1d0a7fcb9f8..d019fd1d327 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycle.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycle.ts @@ -23,6 +23,7 @@ export interface ForkSessionOptions { readonly title?: string; readonly metadata?: Record; readonly turnIndex?: number; + readonly promptId?: string; } export interface ResumeSessionOptions { diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts index 0db40969805..cf2c1b96bf1 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts @@ -101,6 +101,8 @@ import { agentScopeOf, sessionDirOf, sessionScopeOf } from './internal/addressin import { SessionArchived, SessionDeleted } from './sessionLifecycleEvents'; import { assertForkTurnIndex, + type MainTurnSlice, + resolveForkPromptIndex, sliceMainRecordsAtTurn, sliceSubagentRecordsAtTime, } from './internal/forkTurnSlice'; @@ -511,18 +513,18 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec for (const agent of sourceAgents.list()) { const agentHandle = sourceAgents.handleOf(agent.agentId); if (agentHandle === undefined) continue; - if (agentHandle.accessor.get(IAgentLoopService).snapshot().state === 'running') { - throw new Error2( - ErrorCodes.SESSION_FORK_ACTIVE_TURN, - `Session "${sourceId}" cannot be forked while a turn is running`, - { details: { sessionId: sourceId } }, - ); - } await agentHandle.accessor.get(IEventDispatcher).flush(); } await this.appendLogStore.flush(); } assertForkTurnIndex(opts.turnIndex); + if (opts.turnIndex !== undefined && opts.promptId !== undefined) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + 'forkSession turnIndex and promptId are mutually exclusive', + { details: { turnIndex: opts.turnIndex, promptId: opts.promptId } }, + ); + } let targetId: string | undefined; let targetSessionDir: string | undefined; @@ -535,10 +537,11 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec if (agentHandle === undefined) continue; const hold = agentHandle.accessor.get(IAgentLoopService).tryAcquireQuiescence(); if (hold === undefined) { + for (const acquired of quiescenceHolds) acquired.dispose(); throw new Error2( ErrorCodes.SESSION_FORK_ACTIVE_TURN, - `Session "${sourceId}" cannot be forked while a turn is running or queued, or while another fork is copying it`, - { details: { sessionId: sourceId, agentId: agent.agentId } }, + `Session "${sourceId}" cannot be forked while another operation is in progress`, + { details: { sessionId: sourceId } }, ); } quiescenceHolds.push(hold); @@ -566,14 +569,12 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec ); } - const turnSlice = - opts.turnIndex === undefined - ? undefined - : sliceMainRecordsAtTurn( - flattenChain(await this.readSourceWireRecords(sourceHandle, sourceId, MAIN_AGENT_ID)), - sourceId, - opts.turnIndex, - ); + const turnSlice = await this.resolveForkTurnSlice( + sourceHandle, + sourceId, + opts.turnIndex, + opts.promptId, + ); targetSessionDir = sessionDirOf(this.bootstrap.homeDir, this.handlerScope, targetId); await this.copySessionFiles( @@ -651,7 +652,8 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec agents, custom: forkCustomMetadata(sourceMeta?.custom, opts.metadata), lastPrompt: turnSlice === undefined ? sourceMeta?.lastPrompt : turnSlice.lastPrompt, - lastTurnReason: sourceMeta?.lastTurnReason, + lastTurnReason: + turnSlice === undefined ? sourceMeta?.lastTurnReason : turnSlice.lastTurnReason, }; await this.docs.set( sessionScopeOf(this.handlerScope, targetId), @@ -702,6 +704,33 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec }); } + private async resolveForkTurnSlice( + sourceHandle: ISessionScopeHandle | undefined, + sourceId: string, + turnIndex: number | undefined, + promptId: string | undefined, + ): Promise { + if (turnIndex === undefined && promptId === undefined) return undefined; + const records = flattenChain( + await this.readSourceWireRecords(sourceHandle, sourceId, MAIN_AGENT_ID), + ); + let resolved = turnIndex; + if (promptId !== undefined) { + const resolution = resolveForkPromptIndex(records, promptId); + if (resolution.status !== 'found') { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + resolution.status === 'ambiguous' + ? `Prompt "${promptId}" matches multiple turns in session "${sourceId}"` + : `Prompt "${promptId}" was not found in session "${sourceId}"`, + { details: { promptId, sessionId: sourceId } }, + ); + } + resolved = resolution.index; + } + return sliceMainRecordsAtTurn(records, sourceId, resolved!); + } + private async resolveSourceTitle(sourceId: string): Promise { const live = this.sessions.get(sourceId); if (live !== undefined) { diff --git a/packages/agent-core-v2/test/agent/loop/loop.test.ts b/packages/agent-core-v2/test/agent/loop/loop.test.ts index 6781551d8d1..72346638cf2 100644 --- a/packages/agent-core-v2/test/agent/loop/loop.test.ts +++ b/packages/agent-core-v2/test/agent/loop/loop.test.ts @@ -950,6 +950,41 @@ describe('Agent loop', () => { await expect(active.result).resolves.toMatchObject({ type: 'completed' }); }); + it('refuses a quiescence lease while a turn is finalizing', async () => { + let releaseDrain!: () => void; + const drainGate = new Promise((resolve) => { + releaseDrain = resolve; + }); + const drain = vi + .spyOn(ctx.wire, 'drainPersisted') + .mockImplementationOnce(() => drainGate); + + ctx.mockNextResponse({ type: 'text', text: 'done' }); + const active = submitTurn(loop, 'go').turn; + await vi.waitFor(() => { + expect(drain).toHaveBeenCalled(); + }); + + expect(loop.snapshot().state).toBe('idle'); + expect(loop.tryAcquireQuiescence()).toBeUndefined(); + + releaseDrain(); + await expect(active.result).resolves.toMatchObject({ type: 'completed' }); + drain.mockRestore(); + + expect(loop.tryAcquireQuiescence()).toBeDefined(); + }); + + it('exposes a quiescence lease once the turn.ended event is observable', async () => { + ctx.mockNextResponse({ type: 'text', text: 'done' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'go' }] }); + await ctx.untilTurnEnd(); + + const lease = loop.tryAcquireQuiescence(); + expect(lease).toBeDefined(); + lease?.dispose(); + }); + it('holds new admissions until an idle quiescence lease is released', async () => { const lease = loop.tryAcquireQuiescence(); expect(lease).toBeDefined(); diff --git a/packages/agent-core-v2/test/workspace/sessionLifecycle/forkTurnSlice.test.ts b/packages/agent-core-v2/test/workspace/sessionLifecycle/forkTurnSlice.test.ts index a67eada7e47..14afd5094c3 100644 --- a/packages/agent-core-v2/test/workspace/sessionLifecycle/forkTurnSlice.test.ts +++ b/packages/agent-core-v2/test/workspace/sessionLifecycle/forkTurnSlice.test.ts @@ -1,12 +1,16 @@ import { describe, expect, it } from 'vitest'; -import { sliceMainRecordsAtTurn } from '#/workspace/sessionLifecycle/internal/forkTurnSlice'; +import { + resolveForkPromptIndex, + sliceMainRecordsAtTurn, +} from '#/workspace/sessionLifecycle/internal/forkTurnSlice'; import type { WireRecord } from '#/wire/record'; -function userTurnRecord(text: string, time: number): WireRecord { +function userTurnRecord(text: string, time: number, id?: string): WireRecord { return { type: 'context.append_message', message: { + id, role: 'user', content: [{ type: 'text', text }], origin: { kind: 'user' }, @@ -15,6 +19,14 @@ function userTurnRecord(text: string, time: number): WireRecord { }; } +function turnEndedRecord( + turnId: number, + time: number, + reason: 'completed' | 'cancelled' | 'failed' | 'blocked' = 'completed', +): WireRecord { + return { type: 'turn.ended', agentId: 'main', turnId, reason, time }; +} + describe('sliceMainRecordsAtTurn', () => { it('derives a fork last prompt from readable client metadata while retaining the original message', () => { const record: WireRecord = { type: 'context.append_message', message: { role: 'user', content: [{ type: 'text', text: 'serialized' }], origin: { kind: 'user', clientMetadata: [{ display_text: 'Save button · Rename it' }] } }, time: 2 }; @@ -49,3 +61,133 @@ describe('sliceMainRecordsAtTurn', () => { expect(types).toContain('context.append_message'); }); }); + +describe('sliceMainRecordsAtTurn lastTurnReason', () => { + it('derives the outcome from the last retained turn.ended record', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('first', 2), + turnEndedRecord(0, 3, 'completed'), + userTurnRecord('second', 4), + turnEndedRecord(1, 5, 'cancelled'), + userTurnRecord('third', 6), + turnEndedRecord(2, 7, 'failed'), + ]; + expect(sliceMainRecordsAtTurn(records, 'ses_source', 0).lastTurnReason).toBe('completed'); + expect(sliceMainRecordsAtTurn(records, 'ses_source', 1).lastTurnReason).toBe('cancelled'); + expect(sliceMainRecordsAtTurn(records, 'ses_source', 2).lastTurnReason).toBe('failed'); + }); + + it('maps a blocked outcome to failed', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('blocked turn', 2), + turnEndedRecord(0, 3, 'blocked'), + ]; + expect(sliceMainRecordsAtTurn(records, 'ses_source', 0).lastTurnReason).toBe('failed'); + }); + + it('returns undefined when the slice retains no turn.ended record', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('original prompt', 2), + { type: 'turn.steer', turnId: 0, messageId: 'msg_steer', origin: { kind: 'user' }, time: 3 }, + userTurnRecord('steered input', 4), + turnEndedRecord(0, 5), + ]; + expect(sliceMainRecordsAtTurn(records, 'ses_source', 0).lastTurnReason).toBeUndefined(); + expect(sliceMainRecordsAtTurn(records, 'ses_source', 1).lastTurnReason).toBe('completed'); + }); + + it('returns undefined when a later steered turn is cut before its terminal record', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('first', 2), + turnEndedRecord(0, 3, 'completed'), + userTurnRecord('original prompt', 4), + { type: 'turn.steer', turnId: 1, messageId: 'msg_steer', origin: { kind: 'user' }, time: 5 }, + userTurnRecord('steered input', 6), + turnEndedRecord(1, 7), + ]; + expect(sliceMainRecordsAtTurn(records, 'ses_source', 0).lastTurnReason).toBe('completed'); + expect(sliceMainRecordsAtTurn(records, 'ses_source', 1).lastTurnReason).toBeUndefined(); + expect(sliceMainRecordsAtTurn(records, 'ses_source', 2).lastTurnReason).toBe('completed'); + }); +}); + +describe('resolveForkPromptIndex', () => { + it('resolves the visible turn index of a prompt message id', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('first', 2, 'p1'), + turnEndedRecord(0, 3), + userTurnRecord('second', 4, 'p2'), + turnEndedRecord(1, 5), + userTurnRecord('third still running', 6, 'p3'), + ]; + expect(resolveForkPromptIndex(records, 'p1')).toEqual({ status: 'found', index: 0 }); + expect(resolveForkPromptIndex(records, 'p2')).toEqual({ status: 'found', index: 1 }); + expect(resolveForkPromptIndex(records, 'p3')).toEqual({ status: 'found', index: 2 }); + }); + + it('resolves the index of a steered prompt message id', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('original prompt', 2, 'p1'), + { type: 'turn.steer', turnId: 0, messageId: 'p2', origin: { kind: 'user' }, time: 3 }, + userTurnRecord('steered input', 4, 'p2'), + turnEndedRecord(0, 5), + ]; + expect(resolveForkPromptIndex(records, 'p2')).toEqual({ status: 'found', index: 1 }); + }); + + it('resolves an original prompt id of a merged steer through the steer record', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('original prompt', 2, 'p1'), + { + type: 'turn.steer', + turnId: 0, + messageId: 'merged-1', + promptIds: ['p2', 'p3'], + origin: { kind: 'user' }, + time: 3, + }, + userTurnRecord('merged steered input', 4, 'merged-1'), + turnEndedRecord(0, 5), + ]; + expect(resolveForkPromptIndex(records, 'p2')).toEqual({ status: 'found', index: 1 }); + expect(resolveForkPromptIndex(records, 'p3')).toEqual({ status: 'found', index: 1 }); + }); + + it('reports ambiguous when a prompt id matches multiple visible turns', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('first', 2, 'p1'), + turnEndedRecord(0, 3), + userTurnRecord('second', 4, 'p1'), + turnEndedRecord(1, 5), + ]; + expect(resolveForkPromptIndex(records, 'p1')).toEqual({ status: 'ambiguous' }); + }); + + it('reports unknown for a missing prompt id or a non-visible append', () => { + const records: WireRecord[] = [ + { type: 'metadata', protocol_version: '1.5', created_at: 1 }, + userTurnRecord('real prompt', 2, 'p1'), + turnEndedRecord(0, 3), + { + type: 'context.append_message', + message: { + id: 'cron-1', + role: 'user', + content: [{ type: 'text', text: 'cron follow-up' }], + origin: { kind: 'cron_job', jobId: 'job-1' }, + }, + time: 4, + }, + ]; + expect(resolveForkPromptIndex(records, 'missing')).toEqual({ status: 'unknown' }); + expect(resolveForkPromptIndex(records, 'cron-1')).toEqual({ status: 'unknown' }); + }); +}); diff --git a/packages/kap-server/src/protocol/session.ts b/packages/kap-server/src/protocol/session.ts index f408383451b..cef6b3732c6 100644 --- a/packages/kap-server/src/protocol/session.ts +++ b/packages/kap-server/src/protocol/session.ts @@ -75,6 +75,7 @@ export type SessionCreate = z.infer; export const sessionForkSchema = z.object({ title: z.string().min(1).optional(), metadata: z.record(z.string(), z.unknown()).optional(), + prompt_id: z.string().min(1).optional(), }); export type SessionFork = z.infer; diff --git a/packages/kap-server/src/routes/sessions.ts b/packages/kap-server/src/routes/sessions.ts index 386b01e9201..d5844000347 100644 --- a/packages/kap-server/src/routes/sessions.ts +++ b/packages/kap-server/src/routes/sessions.ts @@ -162,6 +162,7 @@ const sessionActionRequestSchema = z.preprocess( instruction: z.string().optional(), count: z.number().int().positive().optional(), page_size: z.number().int().min(1).max(100).optional(), + prompt_id: z.string().min(1).optional(), }), ); @@ -900,6 +901,7 @@ async function forkSessionAction( sourceSessionId: id, title: body.title, metadata: body.metadata, + promptId: body.prompt_id, }); const session = toWireSession( { ...meta, workspaceId: forkHandler.workspaceId }, diff --git a/packages/kap-server/test/prompts.test.ts b/packages/kap-server/test/prompts.test.ts index acc157cc36b..85d9ff1746f 100644 --- a/packages/kap-server/test/prompts.test.ts +++ b/packages/kap-server/test/prompts.test.ts @@ -365,6 +365,7 @@ describe('server-v2 /api/v1 prompts', () => { await vi.waitFor(() => { expect(ended).toHaveLength(1); }); + await main.accessor.get(IAgentLoopService).settled(); subscription.dispose(); expect(ended[0]).toMatchObject({ reason: 'failed', @@ -427,6 +428,7 @@ describe('server-v2 /api/v1 prompts', () => { }, { timeout: 15000 }, ); + await main.accessor.get(IAgentLoopService).settled(); subscription.dispose(); expect(ended[0]?.reason).toBe('failed'); expect(ended[0]?.error?.message).toContain('stub'); @@ -1713,6 +1715,9 @@ describe('server-v2 /api/v1 prompts', () => { content: [{ type: 'text', text: 'hello' }], }); const promptId = submitted.body.data.prompt_id; + const session = getLiveSessionById(server!.core.accessor, id); + const main = session!.accessor.get(IAgentLifecycleService).handleOf('main')!; + await main.accessor.get(IAgentLoopService).settled(); const aborted = await call<{ aborted: boolean }>( 'POST', diff --git a/packages/kap-server/test/sessions.test.ts b/packages/kap-server/test/sessions.test.ts index 9b558a1f0ca..425de90b0fd 100644 --- a/packages/kap-server/test/sessions.test.ts +++ b/packages/kap-server/test/sessions.test.ts @@ -22,6 +22,7 @@ import { IAgentConversationUndoService, IAgentCronService, IAgentLifecycleService, + IAgentLoopService, IEventBus, IEventDispatcher, IEventService, @@ -38,7 +39,7 @@ import { type ScopeSeed, } from '@moonshot-ai/agent-core-v2'; import { SessionMetaUpdated } from '@moonshot-ai/agent-core-v2/session/sessionMetadata/sessionMetaEvents'; -import { TurnSteer } from '@moonshot-ai/agent-core-v2/agent/loop/turnOps'; +import { TurnSteer, TurnEnded } from '@moonshot-ai/agent-core-v2/agent/loop/turnOps'; import type { AgentTranscriptSnapshot } from '@moonshot-ai/transcript'; import { TurnStarted } from '@moonshot-ai/agent-core-v2/agent/loop/turnEvents'; import { sessionWarningsResponseSchema } from '@moonshot-ai/agent-core-v2/app/sessionLegacy/sessionProtocol'; @@ -1282,6 +1283,146 @@ describe('server-v2 /api/v1/sessions', () => { expect(forkedState.titleKind).toBe('custom'); }); + it('passes prompt_id through :fork to the session manager fork', async () => { + const cwd = home as string; + const parent = await postJson('/api/v1/sessions', { metadata: { cwd } }); + const parentId = parent.body.data.id; + const manager = (server as RunningServer).core.accessor.get(ISessionManager); + const fork = vi + .spyOn(manager, 'fork') + .mockRejectedValue(new Error2(ErrorCodes.SESSION_BUSY, 'session is busy')); + + try { + const response = await postJson(`/api/v1/sessions/${parentId}:fork`, { + prompt_id: 'prompt-1', + }); + + expect(response.body.code).toBe(40901); + expect(fork).toHaveBeenCalledWith({ + sourceSessionId: parentId, + title: undefined, + metadata: undefined, + promptId: 'prompt-1', + }); + } finally { + fork.mockRestore(); + } + }); + + it('forks through :fork without prompt_id as before', async () => { + const cwd = home as string; + const parent = await postJson('/api/v1/sessions', { metadata: { cwd } }); + const parentId = parent.body.data.id; + const manager = (server as RunningServer).core.accessor.get(ISessionManager); + const fork = vi.spyOn(manager, 'fork'); + + try { + const forked = await postJson(`/api/v1/sessions/${parentId}:fork`, {}); + + expect(forked.body.code).toBe(0); + expect(fork).toHaveBeenCalledWith({ + sourceSessionId: parentId, + title: undefined, + metadata: undefined, + promptId: undefined, + }); + } finally { + fork.mockRestore(); + } + }); + + it.each([{ prompt_id: 1 }, { prompt_id: '' }])( + 'rejects an out-of-contract prompt_id on :fork (%s)', + async (body) => { + const cwd = home as string; + const parent = await postJson('/api/v1/sessions', { metadata: { cwd } }); + const response = await postJson( + `/api/v1/sessions/${parent.body.data.id}:fork`, + body, + ); + expect(response.body.code).toBe(40001); + expect(response.body.details?.[0]?.path).toBe('prompt_id'); + }, + ); + + it('rejects :fork while a non-turn operation holds the session quiescence', async () => { + const cwd = home as string; + const parent = await postJson('/api/v1/sessions', { metadata: { cwd } }); + const parentId = parent.body.data.id; + const session = getLiveSessionById((server as RunningServer).core.accessor, parentId); + expect(session).toBeDefined(); + await session!.accessor.get(IAgentLifecycleService).create({ agentId: MAIN_AGENT_ID }); + const agent = session!.accessor.get(IAgentLifecycleService).handleOf(MAIN_AGENT_ID)!; + agent.accessor.get(IAgentContextMemoryService).append( + { role: 'user', content: [{ type: 'text', text: 'first prompt' }], toolCalls: [], origin: { kind: 'user' } }, + { role: 'assistant', content: [{ type: 'text', text: 'first answer' }], toolCalls: [] }, + ); + await agent.accessor.get(IEventDispatcher).dispatch( + new TurnEnded({ agentId: MAIN_AGENT_ID, turnId: 0, reason: 'completed' }), + ); + await agent.accessor.get(IWireService).flush(); + const hold = agent.accessor.get(IAgentLoopService).tryAcquireQuiescence(); + expect(hold).toBeDefined(); + + try { + const response = await postJson(`/api/v1/sessions/${parentId}:fork`, {}); + expect(response.body.code).toBe(40901); + } finally { + hold!.dispose(); + } + + const forked = await postJson(`/api/v1/sessions/${parentId}:fork`, {}); + expect(forked.body.code).toBe(0); + }); + + it('derives the fork outcome from the sliced wire instead of the source metadata', async () => { + const cwd = home as string; + const parent = await postJson('/api/v1/sessions', { metadata: { cwd } }); + const parentId = parent.body.data.id; + const session = getLiveSessionById((server as RunningServer).core.accessor, parentId); + expect(session).toBeDefined(); + await session!.accessor.get(IAgentLifecycleService).create({ agentId: MAIN_AGENT_ID }); + const agent = session!.accessor.get(IAgentLifecycleService).handleOf(MAIN_AGENT_ID)!; + const context = agent.accessor.get(IAgentContextMemoryService); + const user = (text: string, promptId?: string): ContextMessage => ({ + role: 'user', + content: [{ type: 'text', text }], + toolCalls: [], + origin: { kind: 'user' }, + id: promptId, + }); + const assistant = (text: string): ContextMessage => ({ + role: 'assistant', + content: [{ type: 'text', text }], + toolCalls: [], + }); + context.append(user('first prompt', 'p1'), assistant('first answer')); + await agent.accessor.get(IEventDispatcher).dispatch( + new TurnEnded({ agentId: MAIN_AGENT_ID, turnId: 0, reason: 'completed' }), + ); + context.append(user('second prompt', 'p2'), assistant('second answer')); + await agent.accessor.get(IEventDispatcher).dispatch( + new TurnEnded({ agentId: MAIN_AGENT_ID, turnId: 1, reason: 'failed' }), + ); + await session!.accessor.get(ISessionMetadata).update({ lastTurnReason: 'failed' }); + await agent.accessor.get(IWireService).flush(); + + const forked = await postJson(`/api/v1/sessions/${parentId}:fork`, { + prompt_id: 'p1', + }); + expect(forked.body.code).toBe(0); + const forkedId = forked.body.data.id; + + const metaFiles = (await readdir(home as string, { recursive: true })).filter( + (path) => path.endsWith(`${forkedId}/state.json`), + ); + expect(metaFiles).toHaveLength(1); + const diskMeta = JSON.parse( + await readFile(join(home as string, metaFiles[0]!), 'utf8'), + ) as { lastTurnReason?: unknown }; + expect(diskMeta.lastTurnReason).toBe('completed'); + }); + it.each([ { count: 1, code: 0, prompts: [] as string[], texts: [] as string[] }, { count: 2, code: 40911, prompts: ['original prompt'], texts: ['answer before steer', 'steered prompt', 'answer after steer', 'second steer', 'answer after second steer'] }, diff --git a/packages/kap-server/test/transcriptContract.e2e.test.ts b/packages/kap-server/test/transcriptContract.e2e.test.ts index 042e5cb09a9..3b9829d54d1 100644 --- a/packages/kap-server/test/transcriptContract.e2e.test.ts +++ b/packages/kap-server/test/transcriptContract.e2e.test.ts @@ -8,6 +8,7 @@ import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest'; import { WebSocket, type RawData } from 'ws'; import { IAgentLifecycleService, + IAgentLoopService, IConfigService, closeSessionById, MAIN_AGENT_ID, @@ -268,8 +269,13 @@ describe('transcript contract e2e', () => { await server!.core.accessor.get(IConfigService).reload(); } - const idle = (server: RunningServer, base: string, sid: string) => - until('session idle', async () => !(await getSessionFacts(server, base, sid)).busy); + const idle = async (server: RunningServer, base: string, sid: string): Promise => { + await until('session idle', async () => !(await getSessionFacts(server, base, sid)).busy); + const main = getLiveSessionById(server.core.accessor, sid) + ?.accessor.get(IAgentLifecycleService) + .handleOf('main'); + await main?.accessor.get(IAgentLoopService).settled(); + }; function dumpState(tx: TxSnapshot, hits: string[]): string { const turns = tx.items diff --git a/packages/klient/src/contract/session/lifecycle.ts b/packages/klient/src/contract/session/lifecycle.ts index f50ed2cb7b2..c056aa24e7a 100644 --- a/packages/klient/src/contract/session/lifecycle.ts +++ b/packages/klient/src/contract/session/lifecycle.ts @@ -41,10 +41,14 @@ export const forkSessionOptionsSchema = z.object({ title: z.string().optional(), metadata: z.record(z.string(), z.unknown()).optional(), turnIndex: z.number().optional(), + promptId: z.string().optional(), }); /** Same fields as `ForkSessionOptions` in the engine, minus the fork-only truncation. */ -export const createChildSessionOptionsSchema = forkSessionOptionsSchema.omit({ turnIndex: true }); +export const createChildSessionOptionsSchema = forkSessionOptionsSchema.omit({ + turnIndex: true, + promptId: true, +}); /** `IScopeHandle` as it survives JSON — `{ id, kind }` plus extras. */ export const handleWireSchema = z.looseObject({ diff --git a/packages/node-sdk/src/kimi-harness.ts b/packages/node-sdk/src/kimi-harness.ts index 288b4efc182..07bb294f37a 100644 --- a/packages/node-sdk/src/kimi-harness.ts +++ b/packages/node-sdk/src/kimi-harness.ts @@ -266,6 +266,7 @@ export class KimiHarness { title: input.title, metadata: input.metadata, turnIndex: input.turnIndex, + promptId: input.promptId, }); const session = new Session({ id: summary.id, diff --git a/packages/node-sdk/src/sdk-rpc-client-v2.ts b/packages/node-sdk/src/sdk-rpc-client-v2.ts index 5c463473159..6f9441fbef6 100644 --- a/packages/node-sdk/src/sdk-rpc-client-v2.ts +++ b/packages/node-sdk/src/sdk-rpc-client-v2.ts @@ -1438,13 +1438,14 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { /** * Through `engineAccessor` (the handler chain's `ISessionLifecycleService.fork`) because the - * klient facade fork takes no explicit target id. `turnIndex` truncation and - * the live-source busy rejection (v1's `SESSION_FORK_ACTIVE_TURN`) are the - * engine's own now, so their failures cross the in-process call with v1's - * codes and details (`request.invalid` with `{turnIndex, availableTurns}` / - * `session.fork_active_turn`). The default title still differs by design - * (v1: "New Session", v2: "Fork: ") — pass an explicit title for - * identical results. + * klient facade fork takes no explicit target id. `promptId` truncation and + * the busy-source handling (a live session mid-turn forks through its last + * completed turn; v1's `SESSION_FORK_ACTIVE_TURN` only when no turn has + * completed yet) are the engine's own now, so their failures cross the + * in-process call with v1's codes and details (`request.invalid` / + * `session.fork_active_turn`). The default + * title still differs by design (v1: "New Session", v2: "Fork: ") — + * pass an explicit title for identical results. */ override async forkSession(input: ForkSessionInput): Promise { // The source session's reads (metadata, wire flush) stay atomic against @@ -1463,6 +1464,7 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { title: input.title, metadata: input.metadata, turnIndex: input.turnIndex, + promptId: input.promptId, }); const handle = await resumeSessionById(this.engineAccessor, meta.id); if (handle === undefined) throw SDKRpcClientV2.sessionNotFound(meta.id); diff --git a/packages/node-sdk/src/types.ts b/packages/node-sdk/src/types.ts index 0f05c076011..b12346ce566 100644 --- a/packages/node-sdk/src/types.ts +++ b/packages/node-sdk/src/types.ts @@ -302,9 +302,16 @@ export interface ForkSessionInput { readonly metadata?: JsonObject; /** * Zero-based index of the user-visible turn to retain through. Omit it to - * preserve the existing full-session fork behavior. + * preserve the existing full-session fork behavior. Mutually exclusive with + * `promptId`. */ readonly turnIndex?: number; + /** + * Id of the user prompt to retain through (the fork keeps everything up to + * the end of that prompt's turn segment), resolved by the engine against the + * full wire history. Mutually exclusive with `turnIndex`. + */ + readonly promptId?: string; } export interface ExportSessionInput { diff --git a/packages/node-sdk/test/session-prompt-events.test.ts b/packages/node-sdk/test/session-prompt-events.test.ts index e0eea9e7d96..e20409cf081 100644 --- a/packages/node-sdk/test/session-prompt-events.test.ts +++ b/packages/node-sdk/test/session-prompt-events.test.ts @@ -7,8 +7,25 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { createKimiHarness, type Event, type KimiHarness } from '#/index'; +import { + getLiveSessionById, + IAgentLifecycleService, + IAgentLoopService, + MAIN_AGENT_ID, + type ServicesAccessor, +} from '@moonshot-ai/agent-core-v2'; + import { TEST_IDENTITY } from './test-identity'; +async function settleSessionLoop(harness: KimiHarness, sessionId: string): Promise { + const accessor = (harness as unknown as { rpc: { engineAccessor: ServicesAccessor } }).rpc + .engineAccessor; + const agent = getLiveSessionById(accessor, sessionId) + ?.accessor.get(IAgentLifecycleService) + .handleOf(MAIN_AGENT_ID); + await agent?.accessor.get(IAgentLoopService).settled(); +} + const MODEL_URL = 'https://model.example.test/v1/chat/completions'; const fakeProviderState = vi.hoisted(() => ({ @@ -541,7 +558,7 @@ describe('Session.prompt events', () => { } }); - it('persists only conversation through the selected turn across resume', async () => { + it('persists only conversation through the selected prompt across resume', async () => { const homeDir = await makeTempDir(); const workDir = await makeTempDir(); const harness = createKimiHarness({ identity: TEST_IDENTITY, homeDir }); @@ -549,14 +566,14 @@ describe('Session.prompt events', () => { try { await configureFakeProvider(harness); const source = await harness.createSession({ id: 'ses_turn_fork_source', workDir }); - await runPrompt(source, 'first question', 'first answer'); - await runPrompt(source, 'second question', 'second answer'); - await runPrompt(source, 'third question', 'third answer'); + await runPrompt(source, 'first question', 'first answer', 'p1'); + await runPrompt(source, 'second question', 'second answer', 'p2'); + await runPrompt(source, 'third question', 'third answer', 'p3'); const fork = await harness.forkSession({ id: source.id, forkId: 'ses_turn_fork_child', - turnIndex: 1, + promptId: 'p2', }); await fork.close(); const resumed = await harness.resumeSession({ id: fork.id }); @@ -585,15 +602,15 @@ describe('Session.prompt events', () => { workDir, metadata: { source: 'vscode' }, }); - await runPrompt(source, 'branch here', 'kept answer'); - await runPrompt(source, 'future prompt', 'discarded answer'); + await runPrompt(source, 'branch here', 'kept answer', 'p1'); + await runPrompt(source, 'future prompt', 'discarded answer', 'p2'); const fork = await harness.forkSession({ id: source.id, forkId: 'ses_turn_fork_metadata_child', title: 'Historical branch', metadata: { branch: 'historical' }, - turnIndex: 0, + promptId: 'p1', }); const state = fork.getResumeState(); @@ -623,16 +640,17 @@ describe('Session.prompt events', () => { try { await configureFakeProvider(harness); const source = await harness.createSession({ id: 'ses_turn_fork_undo_source', workDir }); - await runPrompt(source, 'kept prompt one', 'kept answer one'); - await runPrompt(source, 'kept prompt two', 'kept answer two'); - await runPrompt(source, 'retracted prompt', 'retracted answer'); + await runPrompt(source, 'kept prompt one', 'kept answer one', 'p1'); + await runPrompt(source, 'kept prompt two', 'kept answer two', 'p2'); + await runPrompt(source, 'retracted prompt', 'retracted answer', 'p3'); + await settleSessionLoop(harness, source.id); await source.undoHistory(1); - await runPrompt(source, 'follow-up prompt', 'follow-up answer'); + await runPrompt(source, 'follow-up prompt', 'follow-up answer', 'p4'); const fork = await harness.forkSession({ id: source.id, forkId: 'ses_turn_fork_undo_child', - turnIndex: 2, + promptId: 'p4', }); await fork.close(); const resumed = await harness.resumeSession({ id: fork.id }); @@ -688,9 +706,9 @@ describe('Session.prompt events', () => { try { await configureFakeProvider(harness); const source = await harness.createSession({ id: 'ses_turn_fork_id_source', workDir }); - await runPrompt(source, 'kept prompt', 'kept answer'); - await runPrompt(source, 'future prompt', 'future answer'); - const fork = await harness.forkSession({ id: source.id, turnIndex: 0 }); + await runPrompt(source, 'kept prompt', 'kept answer', 'p1'); + await runPrompt(source, 'future prompt', 'future answer', 'p2'); + const fork = await harness.forkSession({ id: source.id, promptId: 'p1' }); const started = waitForEvent(fork, (event) => event.type === 'turn.started'); const ended = waitForEvent(fork, (event) => event.type === 'turn.ended'); @@ -711,11 +729,11 @@ describe('Session.prompt events', () => { try { await configureFakeProvider(harness); const source = await harness.createSession({ id: 'ses_turn_fork_agents_source', workDir }); - await runPrompt(source, 'kept prompt', 'kept answer'); - await runPrompt(source, 'future prompt', 'future answer'); + await runPrompt(source, 'kept prompt', 'kept answer', 'p1'); + await runPrompt(source, 'future prompt', 'future answer', 'p2'); await source.init(); - const fork = await harness.forkSession({ id: source.id, turnIndex: 0 }); + const fork = await harness.forkSession({ id: source.id, promptId: 'p1' }); expect(Object.keys(fork.getResumeState()?.sessionMetadata.agents ?? {})).toEqual(['main']); } finally { @@ -771,6 +789,89 @@ describe('Session.prompt events', () => { } }); + it('rejects an unknown fork prompt id without creating the fork', async () => { + const homeDir = await makeTempDir(); + const workDir = await makeTempDir(); + const harness = createKimiHarness({ identity: TEST_IDENTITY, homeDir }); + + try { + await configureFakeProvider(harness); + const source = await harness.createSession({ id: 'ses_prompt_fork_unknown_source', workDir }); + await runPrompt(source, 'only question', 'only answer', 'p1'); + + await expect( + harness.forkSession({ + id: source.id, + forkId: 'ses_prompt_fork_unknown_child', + promptId: 'missing', + }), + ).rejects.toMatchObject({ + name: 'KimiError', + code: 'request.invalid', + details: { promptId: 'missing', sessionId: source.id }, + }); + await expect( + harness.listSessions({ sessionId: 'ses_prompt_fork_unknown_child' }), + ).resolves.toEqual([]); + } finally { + await harness.close(); + } + }); + + it('rejects a fork with both turnIndex and promptId', async () => { + const homeDir = await makeTempDir(); + const workDir = await makeTempDir(); + const harness = createKimiHarness({ identity: TEST_IDENTITY, homeDir }); + + try { + await configureFakeProvider(harness); + const source = await harness.createSession({ id: 'ses_prompt_fork_both_source', workDir }); + await runPrompt(source, 'only question', 'only answer', 'p1'); + + await expect( + harness.forkSession({ id: source.id, turnIndex: 0, promptId: 'p1' }), + ).rejects.toMatchObject({ + name: 'KimiError', + code: 'request.invalid', + }); + } finally { + await harness.close(); + } + }); + + it('rejects forking a session while a turn is running', async () => { + const homeDir = await makeTempDir(); + const workDir = await makeTempDir(); + const harness = createKimiHarness({ identity: TEST_IDENTITY, homeDir }); + + try { + await configureFakeProvider(harness); + const source = await harness.createSession({ id: 'ses_turn_fork_busy_source', workDir }); + await runPrompt(source, 'first question', 'first answer', 'p1'); + + const gate = deferredResponse(); + fetchStub!.mockImplementationOnce(() => gate.promise); + const started = waitForEvent(source, (event) => event.type === 'turn.started'); + const ended = waitForEvent(source, (event) => event.type === 'turn.ended'); + await source.prompt('second question'); + await started; + + await expect(harness.forkSession({ id: source.id })).rejects.toMatchObject({ + name: 'KimiError', + code: 'session.fork_active_turn', + }); + await expect(harness.forkSession({ id: source.id, promptId: 'p1' })).rejects.toMatchObject({ + name: 'KimiError', + code: 'session.fork_active_turn', + }); + + gate.resolve(sseResponse('second answer')); + await ended; + } finally { + await harness.close(); + } + }); + it('rejects empty prompt input', async () => { const homeDir = await makeTempDir(); const workDir = await makeTempDir(); @@ -792,16 +893,37 @@ describe('Session.prompt events', () => { }); async function runPrompt( - session: Parameters[0] & { prompt(input: string): Promise }, + session: Parameters[0] & { + prompt(input: string, options?: { promptId?: string }): Promise; + }, input: string, response: string, + promptId?: string, ): Promise { fakeProviderState.responseText = response; const done = waitForEvent(session, (event) => event.type === 'turn.ended'); - await session.prompt(input); + await session.prompt(input, promptId === undefined ? undefined : { promptId }); await done; } +function sseResponse(text: string): Response { + return new Response(sseBody(text), { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }); +} + +function deferredResponse(): { + promise: Promise; + resolve: (response: Response) => void; +} { + let resolve!: (response: Response) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + function visibleReplayText( records: readonly { readonly type: string;