diff --git a/src/routes/chat.ts b/src/routes/chat.ts index aac90b1..6856500 100644 --- a/src/routes/chat.ts +++ b/src/routes/chat.ts @@ -35,6 +35,7 @@ const payloadSchema = z.object({ session_id: z.string().max(128).optional(), model: z.enum(supportedModels).optional(), approval_mode: z.enum(["request", "always"]).optional().default("request"), + regenerate_from_message_index: z.coerce.number().int().min(0).optional(), }); const abortPayloadSchema = z.object({ @@ -955,6 +956,11 @@ export const buildChatRouter = ( const initialSessionState = await sessionUiStateStore.read( toSessionUiStateContext(activeSessionRecord), ); + const persistedMessages = initialSessionState?.messages ?? []; + const baseMessages = + parsed.data.regenerate_from_message_index !== undefined + ? persistedMessages.slice(0, parsed.data.regenerate_from_message_index) + : persistedMessages; if (activeRuns.get(activeSessionRecord.sessionId)?.status === "running") { res.status(409).json({ message: "session is already streaming", @@ -988,7 +994,7 @@ export const buildChatRouter = ( const abortController = new AbortController(); sessionBridge.registerAbortController(clientSessionId, abortController); const initialMessages = createInitialStreamingMessages( - initialSessionState?.messages ?? [], + baseMessages, parsed.data.message, ); const branchGroups = initialSessionState?.branchGroups ?? []; @@ -1122,13 +1128,17 @@ export const buildChatRouter = ( }; try { + if (parsed.data.regenerate_from_message_index !== undefined) { + await runtime.revertLastUserMessage(binding.sessionId); + } + const preparedMessage = await buildPromptWithLearningContext( memoryStore, requestContext.actorKey, requestContext.projectKey, { recentTurns, - persistedMessages: initialSessionState?.messages, + persistedMessages: baseMessages, message: parsed.data.message, restoreConversation: !hadExistingRuntimeSession, }, diff --git a/src/runtime/opencode.ts b/src/runtime/opencode.ts index 3b689f1..7da7e3a 100644 --- a/src/runtime/opencode.ts +++ b/src/runtime/opencode.ts @@ -99,6 +99,28 @@ export class OpencodeRuntimeAdapter { return requireData(messages.data, "session.messages"); } + async revertMessage(sessionId: string, messageId: string) { + const client = await this.ensureClient(); + const response = await client.session.revert({ + sessionID: sessionId, + messageID: messageId, + }); + return response.data; + } + + async revertLastUserMessage(sessionId: string) { + const messages = await this.messages(sessionId, 40); + const lastUserMessage = [...messages] + .reverse() + .find((message) => message.info.role === "user"); + + if (!lastUserMessage) { + throw new Error("no user message found to revert"); + } + + return this.revertMessage(sessionId, lastUserMessage.info.id); + } + async abortSession(sessionId: string) { const client = await this.ensureClient(); const response = await client.session.abort({