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, + }; +}