diff --git a/src/services/__tests__/omni-bridge.test.ts b/src/services/__tests__/omni-bridge.test.ts index abff1c8fd..0508c8036 100644 --- a/src/services/__tests__/omni-bridge.test.ts +++ b/src/services/__tests__/omni-bridge.test.ts @@ -705,3 +705,309 @@ describe('OmniBridge — session lifecycle (Group 2)', () => { } }); }); + +// ============================================================================ +// Session reset subscription — issue #1089 +// ============================================================================ + +describe('OmniBridge — session reset (#1089)', () => { + /** Pre-populate a live session entry on the bridge for reset tests. */ + function injectSession( + bridge: OmniBridge, + key: string, + instanceId: string, + session: OmniSession, + ): { entry: { idleTimer: ReturnType | null } } { + const idleTimer = setTimeout(() => {}, 60_000); + const entry = { + session, + instanceId, + spawning: false, + buffer: [], + idleTimer, + }; + (bridge as any).sessions.set(key, entry); + return { entry }; + } + + it('shuts down the executor and removes the session on reset for a hot chat', async () => { + const { executor, calls, makeSession } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + const session = makeSession('test-agent', 'chat-1'); + injectSession(bridge, 'test-agent:chat-1', 'inst-1', session); + + await (bridge as any).handleSessionReset('inst-1', 'chat-1', 'kill'); + + expect(calls.shutdown.length).toBe(1); + expect(calls.shutdown[0].chatId).toBe('chat-1'); + expect((bridge as any).sessions.size).toBe(0); + } finally { + await bridge.stop(); + } + }); + + it('no-ops on reset for a cold chat (no live session)', async () => { + const { executor, calls } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + // Sessions map is empty — reset must not throw and must not call shutdown. + await (bridge as any).handleSessionReset('inst-1', 'chat-cold'); + + expect(calls.shutdown.length).toBe(0); + expect((bridge as any).sessions.size).toBe(0); + } finally { + await bridge.stop(); + } + }); + + it('clears the idle timer when evicting a reset session', async () => { + const { executor, makeSession } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + const session = makeSession('test-agent', 'chat-1'); + const { entry } = injectSession(bridge, 'test-agent:chat-1', 'inst-1', session); + const timerBefore = entry.idleTimer; + expect(timerBefore).not.toBeNull(); + + await (bridge as any).handleSessionReset('inst-1', 'chat-1'); + + // Session removed → idle timer no longer reachable from sessions map + expect((bridge as any).sessions.has('test-agent:chat-1')).toBe(false); + } finally { + await bridge.stop(); + } + }); + + it('parses subject with dotted chatId (e.g. WhatsApp +5511...@s.whatsapp.net)', async () => { + const { executor, calls, makeSession } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + const dottedChat = '+5511999999999@s.whatsapp.net'; + const session = makeSession('test-agent', dottedChat); + injectSession(bridge, `test-agent:${dottedChat}`, 'inst-x', session); + + // Build a fake NATS message and feed it through the dispatch path the way + // processSessionResetEvents would: subject + JSON-encoded payload bytes. + const fakeMsg = { + subject: `omni.session.reset.inst-x.${dottedChat}`, + data: new TextEncoder().encode(JSON.stringify({ action: 'kill' })), + }; + // Drive a single iteration through processSessionResetEvents using a one-shot iterator. + const oneShot = { + [Symbol.asyncIterator]: async function* () { + yield fakeMsg; + }, + } as any; + await (bridge as any).processSessionResetEvents(oneShot); + + expect(calls.shutdown.length).toBe(1); + expect(calls.shutdown[0].chatId).toBe(dottedChat); + } finally { + await bridge.stop(); + } + }); + + it('tolerates malformed JSON payload by routing on subject alone', async () => { + const { executor, calls, makeSession } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + const session = makeSession('test-agent', 'chat-1'); + injectSession(bridge, 'test-agent:chat-1', 'inst-1', session); + + const fakeMsg = { + subject: 'omni.session.reset.inst-1.chat-1', + data: new TextEncoder().encode('not-json-at-all'), + }; + const oneShot = { + [Symbol.asyncIterator]: async function* () { + yield fakeMsg; + }, + } as any; + await (bridge as any).processSessionResetEvents(oneShot); + + // Subject-only routing still kills the session. + expect(calls.shutdown.length).toBe(1); + } finally { + await bridge.stop(); + } + }); + + it('warns and skips on malformed subject with too few segments', async () => { + const { executor, calls } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + const fakeMsg = { + subject: 'omni.session.reset', + data: new TextEncoder().encode('{}'), + }; + const oneShot = { + [Symbol.asyncIterator]: async function* () { + yield fakeMsg; + }, + } as any; + await (bridge as any).processSessionResetEvents(oneShot); + + expect(calls.shutdown.length).toBe(0); + } finally { + await bridge.stop(); + } + }); + + it('cancels a spawning session on reset and tears down the freshly-spawned executor', async () => { + // Hold the spawn promise open so we can fire reset mid-spawn. + let releaseSpawn!: (s: OmniSession) => void; + const spawnGate = new Promise((resolve) => { + releaseSpawn = resolve; + }); + + const { executor, calls, makeSession } = makeMockExecutor({ + spawnFn: async (agentName, chatId) => { + // Block until the test releases the spawn. + const session = await spawnGate; + return session ?? makeSession(agentName, chatId); + }, + }); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + await bridge.start(); + + try { + // Kick off the spawn — routeMessage will block on spawnGate. + const routePromise = (bridge as any).routeMessage(makeMsg({ content: 'first' })); + + // Yield so spawnSession installs the placeholder before we reset. + await new Promise((r) => setTimeout(r, 5)); + expect((bridge as any).sessions.has('test-agent:chat-1')).toBe(true); + expect((bridge as any).sessions.get('test-agent:chat-1').spawning).toBe(true); + + // Fire the reset while spawn is still in flight. + await (bridge as any).handleSessionReset('inst-1', 'chat-1', 'kill'); + + // Placeholder is gone — the spawn-in-flight has been logically cancelled. + expect((bridge as any).sessions.has('test-agent:chat-1')).toBe(false); + + // Now release the spawn — spawnSession should detect cancelled and tear it down. + releaseSpawn(makeSession('test-agent', 'chat-1')); + await routePromise; + + // executor.shutdown was called for the freshly-spawned (cancelled) session. + expect(calls.shutdown.length).toBe(1); + expect(calls.shutdown[0].chatId).toBe('chat-1'); + + // Buffered triggering message must NOT be delivered to the killed session. + expect(calls.deliver.length).toBe(0); + } finally { + await bridge.stop(); + } + }); + + it('drains the message queue after evicting a reset session', async () => { + const { executor, calls, makeSession } = makeMockExecutor(); + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => makeFakeNats()) as any, + }); + (bridge as any).executor = executor; + // Force the bridge to look fully saturated so a reset opens a slot. + (bridge as any).maxConcurrent = 1; + await bridge.start(); + + try { + const session = makeSession('test-agent', 'chat-1'); + injectSession(bridge, 'test-agent:chat-1', 'inst-1', session); + + // Park a queued message that's waiting for a free slot. + (bridge as any).messageQueue.push(makeMsg({ chatId: 'chat-2', agent: 'agent-b' })); + + await (bridge as any).handleSessionReset('inst-1', 'chat-1', 'kill'); + + // Original session is gone, queued message picked up its slot, drainQueue spawned it. + expect((bridge as any).sessions.has('test-agent:chat-1')).toBe(false); + expect((bridge as any).messageQueue.length).toBe(0); + expect(calls.spawn.length).toBe(1); + expect(calls.spawn[0].chatId).toBe('chat-2'); + } finally { + await bridge.stop(); + } + }); + + it('subscribes to omni.session.reset.> on start()', async () => { + const subscribeCalls: string[] = []; + const fakeSub: Partial & AsyncIterable = { + unsubscribe: () => {}, + [Symbol.asyncIterator]: async function* () {}, + }; + const nc: Partial = { + info: undefined, + closed: async () => undefined, + close: async () => undefined, + drain: async () => undefined, + publish: () => {}, + subscribe: (subject: string) => { + subscribeCalls.push(subject); + return fakeSub as Subscription; + }, + }; + + const bridge = new OmniBridge({ + natsUrl: 'test://fake', + pgProvider: degradedPgProvider, + natsConnectFn: (async () => nc as NatsConnection) as any, + }); + + try { + await bridge.start(); + expect(subscribeCalls).toContain('omni.session.reset.>'); + } finally { + await bridge.stop(); + } + }); +}); diff --git a/src/services/omni-bridge.ts b/src/services/omni-bridge.ts index 3d4bd634c..86eb0bc29 100644 --- a/src/services/omni-bridge.ts +++ b/src/services/omni-bridge.ts @@ -63,6 +63,12 @@ interface SessionEntry { spawning: boolean; buffer: OmniMessage[]; idleTimer: ReturnType | null; + /** + * Set when a reset arrives while the entry is still in the spawning state. + * `spawnSession` checks this flag after `executor.spawn` resolves and tears + * down the freshly-created session instead of putting it into rotation. + */ + cancelled?: boolean; } export interface BridgeStatus { @@ -263,6 +269,13 @@ export class OmniBridge { this.processTurnEvents(sub); } + // Session reset events from Omni — kill the agent's executor session when the user + // sends a reset token (e.g. 🗑️). Paired with omni-side publisher in + // automagik-dev/omni#361. Subject hierarchy: omni.session.reset.{instance}.{chat} + // (recursive `>` because WhatsApp chat ids contain dots). + const sessionResetSub = this.nc.subscribe('omni.session.reset.>'); + this.processSessionResetEvents(sessionResetSub); + // Start idle session checker this.idleCheckTimer = setInterval(() => this.checkIdleSessions(), IDLE_CHECK_INTERVAL_MS); @@ -550,9 +563,14 @@ export class OmniBridge { */ private findSessionKey(instanceId: string, chatId: string): string | undefined { for (const [key, entry] of this.sessions) { - if (entry.instanceId === instanceId && entry.session?.chatId === chatId) { - return key; - } + if (entry.instanceId !== instanceId) continue; + // Live entry — match against the spawned session's chatId. + if (entry.session?.chatId === chatId) return key; + // Spawning entry — session is null, so match against the map key suffix + // (`${agent}:${chatId}`). Without this fallback, a reset arriving in the + // narrow window between placeholder insertion and spawn completion would + // be misclassified as a cold chat and ignored. + if (entry.spawning && key.endsWith(`:${chatId}`)) return key; } return undefined; } @@ -570,6 +588,93 @@ export class OmniBridge { } } + /** + * Process incoming session reset events from NATS. + * + * Subject shape: `omni.session.reset.{instanceId}.{chatId}` — chatId may + * contain dots (WhatsApp ids look like `+5511...@s.whatsapp.net`), so we + * splice from index 4 onward. + * + * Payload is best-effort: malformed JSON is tolerated and treated as `{}`, + * because the routing decision lives in the subject, not the body. The + * `action` field is read for observability only — the only behavior today + * is "kill the session". + */ + private async processSessionResetEvents(sub: Subscription): Promise { + for await (const msg of sub) { + try { + const parts = msg.subject.split('.'); + if (parts.length < 5) { + console.warn(`[omni-bridge] Malformed session-reset subject: ${msg.subject}`); + continue; + } + const instanceId = parts[3]; + const chatId = parts.slice(4).join('.'); + + let action: string | undefined; + try { + const payload = JSON.parse(this.sc.decode(msg.data)) as { action?: string }; + action = payload.action; + } catch { + // Malformed/empty payload — proceed with subject-only routing. + } + + await this.handleSessionReset(instanceId, chatId, action); + } catch (err) { + console.warn('[omni-bridge] Error processing session reset event:', err); + } + } + } + + /** + * Handle a session reset request — evict the session and shut down the executor. + * + * Defensive: no-op when the session is unknown (cold chat), so a user can + * tap reset on a chat that has no live agent without producing an error. + * + * Mirrors `handleTurnTimeout`'s cleanup so the session map, idle timer, and + * executor stay coherent. + */ + private async handleSessionReset(instanceId: string, chatId: string, action?: string): Promise { + const sessionKey = this.findSessionKey(instanceId, chatId); + if (!sessionKey) { + // Cold chat — nothing to reset, but acknowledge in logs for traceability. + console.log(`[omni-bridge] Session reset for cold chat ${instanceId}/${chatId} — no-op`); + return; + } + + const entry = this.sessions.get(sessionKey); + if (!entry) return; + + const actionTag = action ? ` (action=${action})` : ''; + + // Spawning sessions: we cannot interrupt the in-flight executor.spawn call, + // but we can flag the entry so spawnSession tears down the freshly-created + // session as soon as the await resolves. The placeholder is removed from + // the map immediately so subsequent messages spawn a fresh session. + if (entry.spawning) { + console.log(`[omni-bridge] Session reset for spawning ${sessionKey}${actionTag}, marking cancelled`); + entry.cancelled = true; + entry.buffer = []; // Drop buffered messages — user explicitly reset. + this.turnTracker.close(sessionKey, 'reset'); + this.removeSession(sessionKey); + await this.drainQueue(); + return; + } + + if (!entry.session) return; + + console.log(`[omni-bridge] Session reset for ${sessionKey}${actionTag}, evicting`); + this.turnTracker.close(sessionKey, 'reset'); + try { + await this.executor.shutdown(entry.session); + } catch (err) { + console.warn(`[omni-bridge] Error shutting down reset session ${sessionKey}:`, err); + } + this.removeSession(sessionKey); + await this.drainQueue(); + } + /** * Handle a turn timeout event — evict the session and shut down the executor. */ @@ -667,6 +772,19 @@ export class OmniBridge { console.log(`[omni-bridge] Spawning session for ${key}...`); const session = await this.executor.spawn(message.agent, message.chatId, spawnEnv); + // Reset arrived while spawn was in flight — tear down the freshly-created + // session and bail. The placeholder was already removed from the map by + // handleSessionReset, so we don't need to clean it up here. + if (placeholder.cancelled) { + console.log(`[omni-bridge] Spawn for ${key} completed but was cancelled by reset, shutting down`); + try { + await this.executor.shutdown(session); + } catch (err) { + console.warn(`[omni-bridge] Error shutting down cancelled spawn for ${key}:`, err); + } + return; + } + placeholder.session = session; placeholder.spawning = false;