From c0c9f155da51db90f06862a4845e63b148cc0f08 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Sat, 4 Jul 2026 17:14:50 +0800 Subject: [PATCH] fix: restore prompt id in server-v2 snapshots --- .../composables/client/useWorkspaceState.ts | 7 +- apps/kimi-web/test/workspace-state.test.ts | 26 +++++ packages/server-v2/src/routes/snapshot.ts | 37 +++++- packages/server-v2/test/snapshot.test.ts | 106 ++++++++++++++++++ 4 files changed, 171 insertions(+), 5 deletions(-) diff --git a/apps/kimi-web/src/composables/client/useWorkspaceState.ts b/apps/kimi-web/src/composables/client/useWorkspaceState.ts index c1a71fc11..e0b400d78 100644 --- a/apps/kimi-web/src/composables/client/useWorkspaceState.ts +++ b/apps/kimi-web/src/composables/client/useWorkspaceState.ts @@ -1266,11 +1266,12 @@ export function useWorkspaceState(rawState: ExtendedState, deps: UseWorkspaceSta // 1. Authoritative id captured at submit time. let promptId = rawState.promptIdBySession[sid]; - // 2. Fallback to projector-derived id only when it is a real daemon prompt_id - // (synthetic `pr_...` ids are rejected by the daemon). + // 2. Fallback to projector-derived id only when it is a real daemon prompt_id. + // The v1 daemon uses `prompt_...`, server-v2 legacy uses `msg_...`; + // only local synthetic `pr_...` ids are rejected by the daemon. if (promptId === undefined) { const candidate = session?.currentPromptId; - if (candidate?.startsWith('prompt_')) { + if (candidate !== undefined && candidate.length > 0 && !candidate.startsWith('pr_')) { promptId = candidate; } } diff --git a/apps/kimi-web/test/workspace-state.test.ts b/apps/kimi-web/test/workspace-state.test.ts index ee1abfe63..6afbd7159 100644 --- a/apps/kimi-web/test/workspace-state.test.ts +++ b/apps/kimi-web/test/workspace-state.test.ts @@ -236,6 +236,32 @@ describe('useWorkspaceState — abortCurrentPrompt', () => { expect(apiMock.abortPrompt).toHaveBeenCalledWith('sess_1', 'prompt_stale'); expect(apiMock.abortSession).not.toHaveBeenCalled(); }); + + it('uses a server-v2 msg prompt id recovered from session state', async () => { + apiMock.abortPrompt.mockResolvedValue({ aborted: true }); + const state = createState(); + state.promptIdBySession = {}; + state.sessions = [{ ...state.sessions[0]!, currentPromptId: 'msg_live' }]; + const workspace = useWorkspaceState(state, createDeps()); + + await workspace.abortCurrentPrompt(); + + expect(apiMock.abortPrompt).toHaveBeenCalledWith('sess_1', 'msg_live'); + expect(apiMock.abortSession).not.toHaveBeenCalled(); + }); + + it('does not send synthetic projector prompt ids to per-prompt abort', async () => { + apiMock.abortSession.mockResolvedValue({ aborted: true }); + const state = createState(); + state.promptIdBySession = {}; + state.sessions = [{ ...state.sessions[0]!, currentPromptId: 'pr_synthetic' }]; + const workspace = useWorkspaceState(state, createDeps()); + + await workspace.abortCurrentPrompt(); + + expect(apiMock.abortPrompt).not.toHaveBeenCalled(); + expect(apiMock.abortSession).toHaveBeenCalledWith('sess_1'); + }); }); describe('mergeWorkspaces', () => { diff --git a/packages/server-v2/src/routes/snapshot.ts b/packages/server-v2/src/routes/snapshot.ts index 4d270ce16..6ea09db33 100644 --- a/packages/server-v2/src/routes/snapshot.ts +++ b/packages/server-v2/src/routes/snapshot.ts @@ -16,15 +16,22 @@ import { IAgentContextMemoryService, IAgentLifecycleService, + IAgentPromptLegacyService, ISessionInteractionService, ISessionContext, ISessionLifecycleService, ISessionMetadata, IWorkspaceRegistry, toProtocolMessage, + type IAgentScopeHandle, type Scope, } from '@moonshot-ai/agent-core-v2'; -import { ErrorCode, sessionSnapshotResponseSchema, type Message } from '@moonshot-ai/protocol'; +import { + ErrorCode, + sessionSnapshotResponseSchema, + type InFlightTurn, + type Message, +} from '@moonshot-ai/protocol'; import { z } from 'zod'; import { errEnvelope, okEnvelope } from '../envelope'; @@ -111,6 +118,14 @@ export function registerSnapshotRoutes(app: SnapshotRouteHost, deps: SnapshotRou const offset = history.length - page.length; items = page.map((msg, i) => toProtocolMessage(session_id, offset + i, msg, meta.createdAt)); } + const currentPromptId = + snapState.inFlightTurn === null + ? undefined + : readCurrentPromptId(main); + const inFlightTurn = attachCurrentPromptIdToInFlight( + snapState.inFlightTurn, + currentPromptId, + ); // Pending approvals / questions. const interaction = handle.accessor.get(ISessionInteractionService); @@ -128,7 +143,7 @@ export function registerSnapshotRoutes(app: SnapshotRouteHost, deps: SnapshotRou epoch: snapState.epoch, session, messages: { items, has_more: hasMore }, - in_flight_turn: snapState.inFlightTurn, + in_flight_turn: inFlightTurn, pending_approvals: pendingApprovals, pending_questions: pendingQuestions, }, @@ -139,3 +154,21 @@ export function registerSnapshotRoutes(app: SnapshotRouteHost, deps: SnapshotRou ); app.get(route.path, route.options, route.handler as Parameters[2]); } + +function readCurrentPromptId(main: IAgentScopeHandle | undefined): string | undefined { + if (main === undefined) return undefined; + try { + return main.accessor.get(IAgentPromptLegacyService).list().active?.prompt_id; + } catch { + // Auxiliary reconnect metadata must not make the whole snapshot fail. + return undefined; + } +} + +function attachCurrentPromptIdToInFlight( + inFlightTurn: InFlightTurn | null, + currentPromptId: string | undefined, +): InFlightTurn | null { + if (inFlightTurn === null || currentPromptId === undefined) return inFlightTurn; + return { ...inFlightTurn, current_prompt_id: currentPromptId }; +} diff --git a/packages/server-v2/test/snapshot.test.ts b/packages/server-v2/test/snapshot.test.ts index 21d6b3559..342b29cd4 100644 --- a/packages/server-v2/test/snapshot.test.ts +++ b/packages/server-v2/test/snapshot.test.ts @@ -11,15 +11,121 @@ import { IAgentContextMemoryService, IAgentEventSinkService, IAgentLifecycleService, + IAgentPromptLegacyService, + ISessionInteractionService, ISessionContext, ISessionLifecycleService, + ISessionMetadata, + IWorkspaceRegistry, } from '@moonshot-ai/agent-core-v2'; import { sessionSnapshotResponseSchema, type AgentEvent } from '@moonshot-ai/protocol'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { registerSnapshotRoutes } from '../src/routes/snapshot'; import { type RunningServer, startServer } from '../src/start'; import { authHeaders } from './helpers/auth'; +function fakeAccessor(entries: ReadonlyArray) { + const services = new Map(entries); + return { + get(id: unknown): T { + if (!services.has(id)) { + throw new Error(`unexpected service request: ${String(id)}`); + } + return services.get(id) as T; + }, + }; +} + +describe('server-v2 snapshot route enrichment', () => { + it('attaches current_prompt_id to an in-flight turn from promptLegacy active state', async () => { + const sessionId = 'sess_snapshot'; + const promptId = 'msg_snapshot_prompt'; + const workspaceId = 'wd_snapshot_012345abcdef'; + const now = Date.parse('2026-01-01T00:00:00.000Z'); + const main = { + accessor: fakeAccessor([ + [IAgentContextMemoryService, { get: () => [] }], + [ + IAgentPromptLegacyService, + { list: () => ({ active: { prompt_id: promptId }, queued: [] }) }, + ], + ]), + }; + const session = { + accessor: fakeAccessor([ + [ISessionContext, { workspaceId }], + [ + ISessionMetadata, + { + read: async () => ({ + id: sessionId, + title: 'Snapshot', + createdAt: now, + updatedAt: now, + archived: false, + }), + }, + ], + [IAgentLifecycleService, { getHandle: () => main }], + [ISessionInteractionService, { listPending: () => [] }], + ]), + }; + const core = { + accessor: fakeAccessor([ + [ISessionLifecycleService, { resume: async () => session }], + [IWorkspaceRegistry, { get: async () => ({ root: '/workspace' }) }], + ]), + }; + const broadcaster = { + getSnapshotState: async () => ({ + seq: 1, + epoch: 'ep_snapshot', + inFlightTurn: { + turn_id: 7, + assistant_text: 'Hello', + thinking_text: '', + running_tools: [], + }, + }), + }; + + let routeHandler: + | (( + req: { id: string; params: { session_id: string } }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void) + | undefined; + registerSnapshotRoutes( + { + get: (_path, _options, handler) => { + routeHandler = handler; + }, + }, + { core: core as never, broadcaster: broadcaster as never }, + ); + + let payload: unknown; + await routeHandler?.( + { id: 'req_snapshot', params: { session_id: sessionId } }, + { + send: (value) => { + payload = value; + }, + }, + ); + + const body = payload as { code: number; data: unknown }; + expect(body.code).toBe(0); + const snap = sessionSnapshotResponseSchema.parse(body.data); + expect(snap.in_flight_turn).toMatchObject({ + turn_id: 7, + assistant_text: 'Hello', + current_prompt_id: promptId, + }); + }); +}); + describe('server-v2 GET /api/v1/sessions/:id/snapshot', () => { let server: RunningServer | undefined; let home: string | undefined;