From 537824fe3b27c35d79f7d53c5193dcd9ef50ccca Mon Sep 17 00:00:00 2001 From: Remon Oldenbeuving Date: Fri, 6 Mar 2026 15:16:00 +0100 Subject: [PATCH] fix(bot): replay original message after account linking Store the original mention per link token so Kilo can resume it after the user links their account. This avoids the confusing silent stop when the ephemeral link prompt disappears. --- src/app/api/chat/link-account/route.ts | 47 +++++++++++- src/lib/bot-identity.ts | 4 +- src/lib/bot.ts | 74 +++++++++--------- src/lib/bot/handle-linked-message.ts | 77 +++++++++++++++++++ src/lib/bot/link-account.ts | 28 +++++-- src/lib/bot/pending-link-replay.ts | 101 +++++++++++++++++++++++++ 6 files changed, 282 insertions(+), 49 deletions(-) create mode 100644 src/lib/bot/handle-linked-message.ts create mode 100644 src/lib/bot/pending-link-replay.ts diff --git a/src/app/api/chat/link-account/route.ts b/src/app/api/chat/link-account/route.ts index 3118819af..9123432d2 100644 --- a/src/app/api/chat/link-account/route.ts +++ b/src/app/api/chat/link-account/route.ts @@ -1,6 +1,13 @@ +import { after } from 'next/server'; +import { captureException } from '@sentry/nextjs'; import { bot } from '@/lib/bot'; +import { replayLinkedBotMessage } from '@/lib/bot/handle-linked-message'; import { APP_URL } from '@/lib/constants'; import { linkKiloUser, verifyLinkToken } from '@/lib/bot-identity'; +import { + consumePendingLinkReplay, + deserializePendingLinkReplay, +} from '@/lib/bot/pending-link-replay'; import { db } from '@/lib/drizzle'; import { isOrganizationMember } from '@/lib/organizations/organizations'; import { getUserFromAuth } from '@/lib/user.server'; @@ -77,7 +84,8 @@ async function verifyIntegrationAccess( * 2. Authenticate the user via NextAuth session (redirect to sign-in if needed). * 3. Verify the user belongs to the org that owns the integration. * 4. Write the platform identity → Kilo user mapping into Redis. - * 5. Show a success page. + * 5. Replay the original chat message when pending context exists. + * 6. Show a success page. */ export async function GET(request: Request) { const url = new URL(request.url); @@ -119,14 +127,47 @@ export async function GET(request: Request) { await linkKiloUser(bot.getState(), identity, user.id); + const pendingReplayPayload = await consumePendingLinkReplay(bot.getState(), token); + let replayQueued = false; + + if (pendingReplayPayload) { + try { + const replayContext = deserializePendingLinkReplay(pendingReplayPayload, bot.reviver()); + replayQueued = true; + + after(async () => { + await replayLinkedBotMessage({ + thread: replayContext.thread, + message: replayContext.message, + user, + }); + }); + } catch (error) { + captureException(error, { + tags: { component: 'kilo-bot', op: 'deserialize-pending-link-replay' }, + extra: { + platform: identity.platform, + teamId: identity.teamId, + userId: identity.userId, + }, + }); + } + } + + const successMessage = replayQueued + ? `Your ${identity.platform} account has been linked to your Kilo account.
+ Kilo is now handling your original chat message. You can close this tab.
+ If you do not see a reply in chat in a few moments, @mention Kilo again.` + : `Your ${identity.platform} account has been linked to your Kilo account.
+ You can close this tab and @mention Kilo again in your chat.`; + return new Response( ` Account Linked

Account linked

-

Your ${identity.platform} account has been linked to your Kilo account.
- You can close this tab and @mention Kilo again in your chat.

+

${successMessage}

`, { headers: { 'content-type': 'text/html; charset=utf-8' } } diff --git a/src/lib/bot-identity.ts b/src/lib/bot-identity.ts index 01a44cf13..e4bf487d6 100644 --- a/src/lib/bot-identity.ts +++ b/src/lib/bot-identity.ts @@ -66,7 +66,7 @@ export async function unlinkKiloUser( const HMAC_ALGORITHM = 'sha256'; -const TOKEN_TTL_SECONDS = 30 * 60; +export const LINK_TOKEN_TTL_SECONDS = 30 * 60; const NONCE_BYTES = 16; @@ -117,7 +117,7 @@ export function verifyLinkToken(token: string): PlatformIdentity | null { if (typeof data.iat !== 'number') return null; const age = Math.floor(Date.now() / 1000) - data.iat; - if (age < 0 || age > TOKEN_TTL_SECONDS) return null; + if (age < 0 || age > LINK_TOKEN_TTL_SECONDS) return null; if (typeof data.nonce !== 'string' || data.nonce.length === 0) return null; diff --git a/src/lib/bot.ts b/src/lib/bot.ts index fc7ac39b1..e7a56d3f4 100644 --- a/src/lib/bot.ts +++ b/src/lib/bot.ts @@ -1,14 +1,18 @@ -import { Chat, emoji, type ActionEvent, type Message, type Thread } from 'chat'; +import { Chat, type ActionEvent, type Message, type Thread } from 'chat'; import { createSlackAdapter } from '@chat-adapter/slack'; import { createRedisState } from '@chat-adapter/state-redis'; import { createMemoryState } from '@chat-adapter/state-memory'; import { captureException } from '@sentry/nextjs'; -import { resolveKiloUserId, unlinkKiloUser } from '@/lib/bot-identity'; +import { resolveKiloUserId, type PlatformIdentity, unlinkKiloUser } from '@/lib/bot-identity'; +import { handleLinkedBotMessage } from '@/lib/bot/handle-linked-message'; import { getPlatformIdentity, getPlatformIntegration } from '@/lib/bot/platform-helpers'; -import { LINK_ACCOUNT_ACTION_PREFIX, promptLinkAccount } from '@/lib/bot/link-account'; -import { createBotRequest, updateBotRequest } from '@/lib/bot/request-logging'; +import { + createLinkAccountTarget, + LINK_ACCOUNT_ACTION_PREFIX, + promptLinkAccount, +} from '@/lib/bot/link-account'; +import { storePendingLinkReplay } from '@/lib/bot/pending-link-replay'; import { findUserById } from '@/lib/user'; -import { processMessage } from '@/lib/bot/run'; const slackAdapter = createSlackAdapter({ clientId: process.env.SLACK_NEXT_CLIENT_ID, @@ -25,6 +29,31 @@ export const bot = new Chat({ state: process.env.REDIS_URL ? createRedisState() : createMemoryState(), }); +async function promptLinkAccountForMessage( + thread: Thread, + message: Message, + identity: PlatformIdentity +): Promise { + const linkAccountTarget = createLinkAccountTarget(identity); + + try { + await storePendingLinkReplay(bot.getState(), linkAccountTarget.token, thread, message); + } catch (error) { + captureException(error, { + tags: { component: 'kilo-bot', op: 'store-pending-link-replay' }, + extra: { + messageId: message.id, + platform: identity.platform, + teamId: identity.teamId, + threadId: thread.id, + userId: identity.userId, + }, + }); + } + + await promptLinkAccount(thread, message, linkAccountTarget); +} + bot.onNewMention(async function handleIncomingMessage( thread: Thread, message: Message @@ -43,7 +72,7 @@ bot.onNewMention(async function handleIncomingMessage( } if (!kiloUserId) { - await promptLinkAccount(thread, message, identity); + await promptLinkAccountForMessage(thread, message, identity); return; } @@ -51,40 +80,11 @@ bot.onNewMention(async function handleIncomingMessage( if (!user) { await unlinkKiloUser(bot.getState(), identity); - await promptLinkAccount(thread, message, identity); + await promptLinkAccountForMessage(thread, message, identity); return; } - const platform = thread.id.split(':')[0]; - const botRequestId = await createBotRequest({ - createdBy: user.id, - organizationId: platformIntegration.owned_by_organization_id ?? null, - platformIntegrationId: platformIntegration.id, - platform, - platformThreadId: thread.id, - platformMessageId: message.id, - userMessage: message.text, - modelUsed: undefined, - }); - - const received = thread.createSentMessageFromMessage(message); - await received.addReaction(emoji.eyes); - - try { - await processMessage({ thread, message, platformIntegration, user, botRequestId }); - } catch (error) { - console.error('[Bot] Unhandled error in message handler:', error); - if (botRequestId) { - const errMsg = error instanceof Error ? error.message : String(error); - updateBotRequest(botRequestId, { - status: 'error', - errorMessage: errMsg.slice(0, 2000), - }); - } - await thread.post({ markdown: 'Sorry, something went wrong while processing your message.' }); - } finally { - await Promise.all([received.removeReaction(emoji.eyes), received.addReaction(emoji.check)]); - } + await handleLinkedBotMessage({ thread, message, platformIntegration, user }); }); // When the user clicks the "Link Account" LinkButton, Slack fires a diff --git a/src/lib/bot/handle-linked-message.ts b/src/lib/bot/handle-linked-message.ts new file mode 100644 index 000000000..46ee696a1 --- /dev/null +++ b/src/lib/bot/handle-linked-message.ts @@ -0,0 +1,77 @@ +import { emoji, type Message, type Thread } from 'chat'; +import type { PlatformIntegration, User } from '@kilocode/db'; +import { captureException } from '@sentry/nextjs'; +import { getPlatformIntegration } from '@/lib/bot/platform-helpers'; +import { createBotRequest, updateBotRequest } from '@/lib/bot/request-logging'; +import { processMessage } from '@/lib/bot/run'; + +type LinkedBotMessageContext = { + message: Message; + platformIntegration: PlatformIntegration; + thread: Thread; + user: User; +}; + +export async function handleLinkedBotMessage({ + thread, + message, + platformIntegration, + user, +}: LinkedBotMessageContext): Promise { + const platform = thread.id.split(':')[0]; + const botRequestId = await createBotRequest({ + createdBy: user.id, + organizationId: platformIntegration.owned_by_organization_id ?? null, + platformIntegrationId: platformIntegration.id, + platform, + platformThreadId: thread.id, + platformMessageId: message.id, + userMessage: message.text, + modelUsed: undefined, + }); + + const received = thread.createSentMessageFromMessage(message); + await received.addReaction(emoji.eyes); + + try { + await processMessage({ thread, message, platformIntegration, user, botRequestId }); + } catch (error) { + console.error('[Bot] Unhandled error in message handler:', error); + if (botRequestId) { + const errMsg = error instanceof Error ? error.message : String(error); + updateBotRequest(botRequestId, { + status: 'error', + errorMessage: errMsg.slice(0, 2000), + }); + } + await thread.post({ markdown: 'Sorry, something went wrong while processing your message.' }); + } finally { + await Promise.all([received.removeReaction(emoji.eyes), received.addReaction(emoji.check)]); + } +} + +export async function replayLinkedBotMessage({ + thread, + message, + user, +}: { + thread: Thread; + message: Message; + user: User; +}): Promise { + const platformIntegration = await getPlatformIntegration(thread, message); + + if (!platformIntegration) { + captureException(new Error('No active platform integration found during link replay'), { + extra: { + messageId: message.id, + platform: thread.id.split(':')[0], + threadId: thread.id, + userId: user.id, + }, + }); + return; + } + + await handleLinkedBotMessage({ thread, message, platformIntegration, user }); +} diff --git a/src/lib/bot/link-account.ts b/src/lib/bot/link-account.ts index 8f677ff4a..0f0df83d3 100644 --- a/src/lib/bot/link-account.ts +++ b/src/lib/bot/link-account.ts @@ -7,23 +7,37 @@ const LINK_ACCOUNT_PATH = '/api/chat/link-account'; export const LINK_ACCOUNT_ACTION_PREFIX = `link-${APP_URL}${LINK_ACCOUNT_PATH}`; -function buildLinkAccountUrl(identity: PlatformIdentity): string { +export type LinkAccountTarget = { + token: string; + url: string; +}; + +function buildLinkAccountUrl(token: string): string { const url = new URL(LINK_ACCOUNT_PATH, APP_URL); - url.searchParams.set('token', createLinkToken(identity)); + url.searchParams.set('token', token); return url.toString(); } -function linkAccountCard(linkUrl: string) { +export function createLinkAccountTarget(identity: PlatformIdentity): LinkAccountTarget { + const token = createLinkToken(identity); + return { + token, + url: buildLinkAccountUrl(token), + }; +} + +function linkAccountCard(linkTarget: LinkAccountTarget) { return Card({ title: 'Link your Kilo account', children: [ Section([ CardText( 'To use Kilo from this workspace you first need to link your chat account. ' + - 'Click the button below to sign in and link your account.' + 'Click the button below to sign in and link your account. ' + + 'After linking, Kilo will continue with your original message automatically.' ), ]), - Actions([LinkButton({ label: 'Link Account', url: linkUrl, style: 'primary' })]), + Actions([LinkButton({ label: 'Link Account', url: linkTarget.url, style: 'primary' })]), ], }); } @@ -31,12 +45,12 @@ function linkAccountCard(linkUrl: string) { export async function promptLinkAccount( thread: Thread, message: Message, - identity: PlatformIdentity + linkTarget: LinkAccountTarget ): Promise { // Post to the channel when the @mention is top-level, otherwise into the thread. const target = isChannelLevelMessage(thread, message) ? thread.channel : thread; - await target.postEphemeral(message.author, linkAccountCard(buildLinkAccountUrl(identity)), { + await target.postEphemeral(message.author, linkAccountCard(linkTarget), { fallbackToDM: true, }); } diff --git a/src/lib/bot/pending-link-replay.ts b/src/lib/bot/pending-link-replay.ts new file mode 100644 index 000000000..1baedcc79 --- /dev/null +++ b/src/lib/bot/pending-link-replay.ts @@ -0,0 +1,101 @@ +import type { Message, SerializedThread, StateAdapter, Thread } from 'chat'; +import { LINK_TOKEN_TTL_SECONDS } from '@/lib/bot-identity'; + +const PENDING_LINK_REPLAY_KEY_PREFIX = 'bot:pending-link-replay'; +const PENDING_LINK_REPLAY_TTL_MS = LINK_TOKEN_TTL_SECONDS * 1000; +const PENDING_LINK_REPLAY_LOCK_TTL_MS = 10_000; + +export type PendingLinkReplayContext = { + message: Message; + thread: Thread; +}; + +function getPendingLinkReplayKey(token: string): string { + return `${PENDING_LINK_REPLAY_KEY_PREFIX}:${token}`; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function isThread(value: unknown): value is Thread { + return ( + isRecord(value) && + typeof value.channelId === 'string' && + typeof value.createSentMessageFromMessage === 'function' && + typeof value.id === 'string' && + typeof value.post === 'function' + ); +} + +function isMessage(value: unknown): value is Message { + return ( + isRecord(value) && + typeof value.id === 'string' && + typeof value.threadId === 'string' && + typeof value.text === 'string' + ); +} + +function hasToJSON(value: unknown): value is { toJSON(): TSerialized } { + return isRecord(value) && typeof value.toJSON === 'function'; +} + +function serializeThread(thread: Thread): SerializedThread { + if (!hasToJSON(thread)) { + throw new Error('Expected thread to support serialization'); + } + + return thread.toJSON(); +} + +export async function storePendingLinkReplay( + state: StateAdapter, + token: string, + thread: Thread, + message: Message +): Promise { + await state.set( + getPendingLinkReplayKey(token), + JSON.stringify({ + thread: serializeThread(thread), + message: message.toJSON(), + }), + PENDING_LINK_REPLAY_TTL_MS + ); +} + +export async function consumePendingLinkReplay( + state: StateAdapter, + token: string +): Promise { + const key = getPendingLinkReplayKey(token); + const lock = await state.acquireLock(key, PENDING_LINK_REPLAY_LOCK_TTL_MS); + if (!lock) return null; + + try { + const payload = await state.get(key); + if (!payload) return null; + + await state.delete(key); + return payload; + } finally { + await state.releaseLock(lock); + } +} + +export function deserializePendingLinkReplay( + payload: string, + reviver: (key: string, value: unknown) => unknown +): PendingLinkReplayContext { + const parsed: unknown = JSON.parse(payload, reviver); + + if (!isRecord(parsed) || !isThread(parsed.thread) || !isMessage(parsed.message)) { + throw new Error('Invalid pending link replay payload'); + } + + return { + thread: parsed.thread, + message: parsed.message, + }; +}