diff --git a/src/bridge.ts b/src/bridge.ts new file mode 100644 index 0000000..198b50a --- /dev/null +++ b/src/bridge.ts @@ -0,0 +1,270 @@ +import type { Context } from '@deepseek-ai/cordis' +import type { AgentHandle } from '@deepseek-ai/dsh-agent' +import { createUserMessage } from '@deepseek-ai/dsh-llm/message' +import type { ContentBlock, TextBlock } from '@deepseek-ai/dsh-llm/types' +import { SessionId, type SessionEvent } from '@deepseek-ai/dsh-session/types' +import { isAuthorized } from './auth.ts' +import { TelegramClient, type TelegramClientLike, type TelegramUpdate } from './client.ts' +import { MSG, parseCommand } from './commands.ts' +import { markdownToHtml, splitMessage } from './format.ts' + +export interface TelegramBridgeOptions { + token: string + allowedUserIds: number[] + allowAllUsers: boolean + client?: TelegramClientLike + sleep?: (ms: number) => Promise + provider?: string + model?: string + maxMessageLength?: number + cwd?: string +} + +interface ChatState { + chatId: number + handle: AgentHandle + sessionId: ReturnType +} + +interface SessionLike { + id: ReturnType +} + +const defaultSleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) + +function contentToText(content: readonly ContentBlock[]): string { + return content + .filter((block): block is TextBlock => block.type === 'text') + .map((block) => block.text) + .join('') +} + +export class TelegramBridge { + private readonly ctx: Context + private readonly token: string + private readonly allowedUserIds: number[] + private readonly allowAllUsers: boolean + private readonly client: TelegramClientLike + private readonly sleep: (ms: number) => Promise + private readonly provider: string + private readonly model: string + private readonly maxMessageLength: number + private readonly cwd: string + + private readonly chats = new Map() + private polling = false + private offset: number | undefined + private pollPromise: Promise | undefined + private pollAbort: AbortController | undefined + private disposeSessionListener: (() => void) | undefined + + constructor(ctx: Context, options: TelegramBridgeOptions) { + this.ctx = ctx + this.token = options.token + this.allowedUserIds = options.allowedUserIds + this.allowAllUsers = options.allowAllUsers + this.client = options.client ?? new TelegramClient(options.token) + this.sleep = options.sleep ?? defaultSleep + this.provider = options.provider ?? 'deepseek-official' + this.model = options.model ?? 'deepseek-v4-flash' + this.maxMessageLength = options.maxMessageLength ?? 4096 + this.cwd = options.cwd ?? process.cwd() + } + + start(): void { + this.disposeSessionListener?.() + this.disposeSessionListener = this.ctx.on('session/event', (session, event) => { + void this.onSessionEvent(session, event) + }) + if (!this.polling) { + this.polling = true + this.pollAbort = new AbortController() + this.pollPromise = this.pollLoop() + } + } + + async stop(): Promise { + this.polling = false + this.pollAbort?.abort() + this.pollAbort = undefined + this.disposeSessionListener?.() + this.disposeSessionListener = undefined + const disposals = [...this.chats.values()].map((chat) => chat.handle.dispose()) + this.chats.clear() + await Promise.all(disposals) + if (this.pollPromise) { + await this.pollPromise.catch(() => {}) + this.pollPromise = undefined + } + } + + async processUpdate(update: TelegramUpdate): Promise { + const message = update.message + if (!message?.text) return + + const chatId = message.chat.id + const userId = message.from?.id + + if (!isAuthorized({ allowAllUsers: this.allowAllUsers, allowedUserIds: this.allowedUserIds, userId })) { + await this.client.sendMessage(chatId, MSG.DENIED) + return + } + + const parsed = parseCommand(message.text) + switch (parsed.type) { + case 'start': + await this.client.sendMessage(chatId, MSG.WELCOME) + return + case 'help': + await this.client.sendMessage(chatId, MSG.HELP) + return + case 'new': + await this.resetChat(chatId) + await this.client.sendMessage(chatId, MSG.NEW_SESSION) + return + case 'unknown': + await this.client.sendMessage(chatId, MSG.unknown(parsed.command)) + return + case 'plain': + await this.ensureChat(chatId) + await this.followup(chatId, parsed.text) + return + } + } + + private async pollLoop(): Promise { + const signal = this.pollAbort?.signal + let errorCount = 0 + while (this.polling) { + try { + const updates = await this.client.getUpdates(this.offset) + if (!this.polling) break + errorCount = 0 + if (updates.length === 0) { + await this.interruptibleDelay(50, signal) + continue + } + for (const update of updates) { + if (!this.polling) break + await this.processUpdate(update) + this.offset = update.update_id + 1 + } + } catch (err) { + if (!this.polling) break + errorCount += 1 + this.ctx.logger.error(this.redact(err)) + await this.interruptibleSleep(Math.min(1000 * errorCount, 10_000), signal) + } + } + } + + private interruptibleDelay(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted || !this.polling) return Promise.resolve() + return new Promise((resolve) => { + const timer = setTimeout(resolve, ms) + signal?.addEventListener('abort', () => { + clearTimeout(timer) + resolve() + }, { once: true }) + }) + } + + private async interruptibleSleep(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted || !this.polling) return + await Promise.race([ + this.sleep(ms), + new Promise((resolve) => { + if (signal?.aborted) { + resolve() + return + } + signal?.addEventListener('abort', () => resolve(), { once: true }) + }), + ]) + } + + private async onSessionEvent(session: SessionLike, event: SessionEvent): Promise { + const chat = this.findChatBySessionId(session.id) + if (!chat) return + + if (event.type === 'turn/start') { + await this.client.sendChatAction(chat.chatId, 'typing') + return + } + + if (event.type === 'assistant/message') { + const text = contentToText(event.data.message.content) + if (text) { + await this.deliver(chat.chatId, text) + } + } + } + + private findChatBySessionId(sessionId: ReturnType): ChatState | undefined { + for (const chat of this.chats.values()) { + if (chat.sessionId === sessionId) return chat + } + return undefined + } + + private async ensureChat(chatId: number): Promise { + const key = String(chatId) + const existing = this.chats.get(key) + if (existing) return existing + + const sessionId = SessionId(`telegram:${chatId}`) + const handle = await this.createAgent(sessionId) + const state: ChatState = { chatId, handle, sessionId } + this.chats.set(key, state) + return state + } + + private async resetChat(chatId: number): Promise { + const key = String(chatId) + const existing = this.chats.get(key) + if (existing) { + await existing.handle.dispose() + this.chats.delete(key) + } + const sessionId = SessionId(`telegram:${chatId}:${Date.now()}`) + const handle = await this.createAgent(sessionId) + this.chats.set(key, { chatId, handle, sessionId }) + } + + private async createAgent(sessionId: ReturnType): Promise { + return this.ctx.agents.create({ + sessionId, + meta: { cwd: this.cwd }, + agentOptions: { + provider: this.provider, + model: this.model, + }, + }) + } + + private async followup(chatId: number, text: string): Promise { + const chat = await this.ensureChat(chatId) + const message = createUserMessage({ + content: [{ type: 'text', text }], + source: { kind: 'user' }, + }) + chat.handle.agent.followup(message) + } + + private async deliver(chatId: number, markdown: string): Promise { + const chunks = splitMessage(markdown, this.maxMessageLength) + for (const chunk of chunks) { + const html = markdownToHtml(chunk) + try { + await this.client.sendMessage(chatId, html, 'HTML') + } catch { + await this.client.sendMessage(chatId, chunk) + } + } + } + + private redact(value: unknown): string { + const message = value instanceof Error ? value.message : String(value) + return message.split(this.token).join('***') + } +} diff --git a/tests/bridge.test.ts b/tests/bridge.test.ts new file mode 100644 index 0000000..dd7934f --- /dev/null +++ b/tests/bridge.test.ts @@ -0,0 +1,317 @@ +import assert from 'node:assert/strict' +import test from 'node:test' +import type { CreateAgentOptions } from '@deepseek-ai/dsh-agent' +import type { UserMessage } from '@deepseek-ai/dsh-llm/types' +import { SessionId } from '@deepseek-ai/dsh-session/types' +import { TelegramBridge } from '../src/bridge.ts' +import type { TelegramClientLike, TelegramUpdate } from '../src/client.ts' +import { MSG } from '../src/commands.ts' + +type SentMessage = { chatId: number; text: string; parseMode?: string } + +function fakeClient( + sent: SentMessage[], + overrides: Partial = {}, +): TelegramClientLike { + return { + getMe: async () => ({ id: 1 }), + getUpdates: async () => [], + sendMessage: async (chatId, text, parseMode) => { + sent.push({ chatId, text, parseMode }) + return { message_id: 1, date: 0, chat: { id: chatId, type: 'private' }, text } + }, + sendChatAction: async () => true, + ...overrides, + } +} + +function messageUpdate(chatId: number, userId: number, text: string, updateId = 1): TelegramUpdate { + return { + update_id: updateId, + message: { + message_id: updateId, + date: 0, + chat: { id: chatId, type: 'private' }, + from: { id: userId }, + text, + }, + } +} + +function fakeAgents(track: { + createCount: number + disposeCount: number + sessionIds: string[] + followups: UserMessage[][] +}) { + return { + create: async (options: CreateAgentOptions) => { + track.createCount += 1 + track.sessionIds.push(String(options.sessionId)) + const followups: UserMessage[] = [] + track.followups.push(followups) + return { + agent: { + followup(message: UserMessage) { + followups.push(message) + }, + }, + dispose: async () => { + track.disposeCount += 1 + }, + } + }, + } +} + +function bridgeCtx( + agents: ReturnType, + onSessionEvent?: (session: { id: ReturnType }, event: unknown) => void, +) { + return { + logger: { warn() {}, error() {} }, + agents, + on(event: string, listener: (session: { id: ReturnType }, event: unknown) => void) { + if (event === 'session/event' && onSessionEvent) { + onSessionEvent = listener + } + return () => {} + }, + _fireSessionEvent(session: { id: ReturnType }, ev: unknown) { + onSessionEvent?.(session, ev) + }, + } +} + +test('unauthorized user gets denied and no agent', async () => { + const sent: SentMessage[] = [] + let createCount = 0 + const ctx = { + logger: { warn() {}, error() {} }, + agents: { + create: async () => { + createCount += 1 + throw new Error('should not create') + }, + }, + on() { + return () => {} + }, + } + const bridge = new TelegramBridge(ctx as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(99, 99, 'hi')) + assert.equal(createCount, 0) + assert.equal(sent[0]?.text, MSG.DENIED) +}) + +test('/help sends help without creating agent', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + const bridge = new TelegramBridge(bridgeCtx(fakeAgents(track)) as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(42, 1, '/help')) + assert.equal(track.createCount, 0) + assert.equal(sent[0]?.text, MSG.HELP) +}) + +test('first plain text creates agent then followups', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + const bridge = new TelegramBridge(bridgeCtx(fakeAgents(track)) as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + assert.equal(track.createCount, 1) + assert.equal(track.sessionIds[0], 'telegram:42') + assert.equal(track.followups.length, 1) + assert.equal((track.followups[0]![0]!.content[0] as { text: string }).text, 'hello') +}) + +test('second plain text in same chat only followups', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + const bridge = new TelegramBridge(bridgeCtx(fakeAgents(track)) as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + await bridge.processUpdate(messageUpdate(42, 1, 'again', 2)) + assert.equal(track.createCount, 1) + assert.equal(track.followups.length, 1) + assert.equal(track.followups[0]!.length, 2) + assert.equal((track.followups[0]![1]!.content[0] as { text: string }).text, 'again') +}) + +test('/new disposes previous handle and creates fresh session', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + const bridge = new TelegramBridge(bridgeCtx(fakeAgents(track)) as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + await bridge.processUpdate(messageUpdate(42, 1, '/new', 2)) + assert.equal(track.disposeCount, 1) + assert.equal(track.createCount, 2) + assert.equal(track.sessionIds[0], 'telegram:42') + assert.match(track.sessionIds[1]!, /^telegram:42:\d+$/) + assert.equal(sent.at(-1)?.text, MSG.NEW_SESSION) +}) + +test('turn/start sends typing for matching session', async () => { + const sent: SentMessage[] = [] + const actions: Array<{ chatId: number; action: string }> = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + let fire: ((session: { id: ReturnType }, event: unknown) => void) | undefined + const ctx = { + logger: { warn() {}, error() {} }, + agents: fakeAgents(track), + on(_event: string, listener: (session: { id: ReturnType }, event: unknown) => void) { + fire = listener + return () => {} + }, + } + const bridge = new TelegramBridge(ctx as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent, { + sendChatAction: async (chatId, action) => { + actions.push({ chatId, action }) + return true + }, + }), + sleep: async () => {}, + }) + bridge.start() + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + fire!({ id: SessionId(track.sessionIds[0]!) }, { type: 'turn/start', data: { turn: 1 } }) + await new Promise((resolve) => setImmediate(resolve)) + assert.deepEqual(actions, [{ chatId: 42, action: 'typing' }]) + await bridge.stop() +}) + +test('assistant/message delivers HTML chunks for matching session', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + let fire: ((session: { id: ReturnType }, event: unknown) => void) | undefined + const ctx = { + logger: { warn() {}, error() {} }, + agents: fakeAgents(track), + on(_event: string, listener: (session: { id: ReturnType }, event: unknown) => void) { + fire = listener + return () => {} + }, + } + const bridge = new TelegramBridge(ctx as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + bridge.start() + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + fire!( + { id: SessionId(track.sessionIds[0]!) }, + { + type: 'assistant/message', + data: { + turn: 1, + step: 1, + message: { + content: [{ type: 'text', text: '**bold**' }], + }, + }, + }, + ) + await new Promise((resolve) => setImmediate(resolve)) + const reply = sent.find((m) => m.parseMode === 'HTML') + assert.ok(reply) + assert.match(reply!.text, /bold<\/b>/) + await bridge.stop() +}) + +test('HTML send failure retries plain text for that chunk', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + let fire: ((session: { id: ReturnType }, event: unknown) => void) | undefined + const ctx = { + logger: { warn() {}, error() {} }, + agents: fakeAgents(track), + on(_event: string, listener: (session: { id: ReturnType }, event: unknown) => void) { + fire = listener + return () => {} + }, + } + const bridge = new TelegramBridge(ctx as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent, { + sendMessage: async (chatId, text, parseMode) => { + if (parseMode === 'HTML') throw new Error('bad html') + sent.push({ chatId, text, parseMode }) + return { message_id: 1, date: 0, chat: { id: chatId, type: 'private' }, text } + }, + }), + sleep: async () => {}, + }) + bridge.start() + await bridge.processUpdate(messageUpdate(42, 1, 'hello')) + fire!( + { id: SessionId(track.sessionIds[0]!) }, + { + type: 'assistant/message', + data: { + turn: 1, + step: 1, + message: { + content: [{ type: 'text', text: 'plain chunk' }], + }, + }, + }, + ) + await new Promise((resolve) => setImmediate(resolve)) + assert.equal(sent.length, 1) + assert.equal(sent[0]?.parseMode, undefined) + assert.equal(sent[0]?.text, 'plain chunk') + await bridge.stop() +}) + +test('stop disposes all chat handles', async () => { + const sent: SentMessage[] = [] + const track = { createCount: 0, disposeCount: 0, sessionIds: [] as string[], followups: [] as UserMessage[][] } + const bridge = new TelegramBridge(bridgeCtx(fakeAgents(track)) as never, { + token: 't', + allowedUserIds: [1], + allowAllUsers: false, + client: fakeClient(sent), + sleep: async () => {}, + }) + await bridge.processUpdate(messageUpdate(42, 1, 'one')) + await bridge.processUpdate(messageUpdate(43, 1, 'two', 2)) + await bridge.stop() + assert.equal(track.disposeCount, 2) +})