import { ForbiddenException, Injectable, Logger, NotFoundException } from '@nestjs/common'; import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent'; import { AgentService, type AgentSession } from '../agent/agent.service.js'; import type { ActorTenantScope } from '../auth/session-scope.js'; import type { ChatRuntime, LegacyBrowserMessagePayload, LegacyEmbeddedChatPort, LegacyRuntimeEvent, LegacyRuntimeResult, LegacySessionPresentation, LegacySocketTurnLease, LegacyUsage, OwnedConversationContext, VerifiedDiscordIngressContext, VerifiedDiscordTurnLease, LegacyRuntimeStream, } from './chat-runtime.js'; /** Fixed timeout for a synchronous REST turn, matching the historical controller budget. */ const REST_TURN_TIMEOUT_MS = 120_000; /** * The `legacy` chat runtime and the sole implementation of {@link LegacyEmbeddedChatPort}. * * It owns the embedded in-process execution path — the `AgentService` stack that the * `ChatController` and `ChatGateway` drove directly before Task Five. Once the * {@link import('./chat-runtime-router.js').ChatRuntimeRouter} fronts it, the browser * HTTP/WebSocket legacy path and verified-Discord ingress route through THIS runtime, so * neither the controller nor the gateway retains `AgentService`, `piSession`, session, * listener, channel, or metric access. Ownership (`userId`/`tenantId`) is re-checked by * `AgentService` on every operation; a missing, foreign, or no-longer-owned conversation * collapses to `conversation_unavailable` and never throws out of the port. */ @Injectable() export class EmbeddedChatRuntime implements ChatRuntime, LegacyEmbeddedChatPort { readonly kind = 'embedded' as const; private readonly logger = new Logger(EmbeddedChatRuntime.name); constructor(readonly agentService: AgentService) {} // ------------------------------------------------------------------------- // Legacy REST completion (op A) // ------------------------------------------------------------------------- async completeLegacyRestTurn( context: OwnedConversationContext, input: Readonly<{ content: string }>, ): Promise< LegacyRuntimeResult> > { const scope = toScope(context.scope); const { conversationId } = context; const resolved = await this.resolveOrCreate(conversationId, scope, {}); if (!resolved.ok) return resolved; let responseText = ''; let timer: ReturnType | undefined; let detach: (() => void) | undefined; let disposed = false; // One idempotent teardown owned OUTSIDE the completion promise: it clears the timeout and // detaches the event listener exactly once, whichever of agent_end, timeout, or a prompt // rejection fires first. Without this, a prompt() rejection surfaced through the catch below // would return while leaving the listener attached (free to consume a later turn's events) and // the 120s timer live (its rejection later going unobserved). const dispose = (): void => { if (disposed) return; disposed = true; if (timer !== undefined) clearTimeout(timer); detach?.(); }; const done = new Promise((resolve, reject) => { timer = setTimeout(() => { dispose(); reject(new Error('Agent response timed out')); }, REST_TURN_TIMEOUT_MS); detach = this.agentService.onEvent( conversationId, (event: AgentSessionEvent) => { if ( event.type === 'message_update' && event.assistantMessageEvent.type === 'text_delta' ) { responseText += event.assistantMessageEvent.delta; } if (event.type === 'agent_end') { dispose(); resolve(); } }, scope, ); }); // Attach the prompt and the completion promise CONCURRENTLY. Awaiting prompt() first left the // timeout unobservable until prompt settled (a hung prompt could never time out) and, worse, // let the 120s timer reject `done` while nothing yet awaited it — a transient unhandledRejection // window. Promise.all installs handlers on BOTH synchronously, so the timeout bounds the whole // turn even while prompt is pending, and neither promise can reject unobserved. Success still // requires both prompt() to resolve AND agent_end to arrive (identical to the prior sequential // await). The idempotent dispose() clears the timer + detaches on whichever settles first. const prompting = this.agentService.prompt(conversationId, input.content, scope); try { await Promise.all([prompting, done]); } catch (err) { dispose(); const message = err instanceof Error ? err.message : String(err); if (message.includes('timed out')) { return { ok: false, code: 'timeout', retryable: true }; } this.logger.error(`Legacy REST turn failed for conversation=${conversationId}`, message); return { ok: false, code: 'operation_failed', retryable: false }; } const presentation = this.presentationFor(conversationId, scope) ?? resolved.presentation; return { ok: true, value: { text: responseText, presentation } }; } // ------------------------------------------------------------------------- // Legacy Socket streaming (op B) // ------------------------------------------------------------------------- async prepareLegacySocketTurn( context: OwnedConversationContext, input: LegacyBrowserMessagePayload, stream: LegacyRuntimeStream, ): Promise> { const scope = toScope(context.scope); const { conversationId } = context; const resolved = await this.resolveOrCreate(conversationId, scope, { ...(input.provider ? { provider: input.provider } : {}), ...(input.modelId ? { modelId: input.modelId } : {}), ...(input.agentId ? { agentConfigId: input.agentId } : {}), }); if (!resolved.ok) return resolved; let detach: () => void; try { detach = this.subscribe(conversationId, scope, stream); } catch (err) { // A partial listener/channel setup rolled itself back inside subscribe(); surface a total // safe failure instead of throwing out of the port. Retryable — the attach is transient. this.logger.error( `Embedded socket subscription failed for conversation=${conversationId}`, err instanceof Error ? err.message : String(err), ); return { ok: false, code: 'runtime_unavailable', retryable: true }; } return { ok: true, value: this.buildLease( conversationId, scope, input.content, input.attachments, detach, resolved.presentation, ), }; } // ------------------------------------------------------------------------- // Thinking level (op C) — synchronous, total // ------------------------------------------------------------------------- setLegacyThinking( context: OwnedConversationContext, level: string, ): LegacyRuntimeResult { const scope = toScope(context.scope); const session = this.agentService.getSession(context.conversationId, scope); if (!session) return CONVERSATION_UNAVAILABLE; const availableThinkingLevels = session.piSession.getAvailableThinkingLevels(); if (!(availableThinkingLevels as readonly string[]).includes(level)) { return { ok: false, code: 'thinking_level_invalid', retryable: false, availableThinkingLevels, }; } session.piSession.setThinkingLevel(level as never); return { ok: true, value: this.presentationForSession(session) }; } // ------------------------------------------------------------------------- // Abort (op D) // ------------------------------------------------------------------------- async abortLegacyTurn(context: OwnedConversationContext): Promise> { const scope = toScope(context.scope); const session = this.agentService.getSession(context.conversationId, scope); if (!session) return CONVERSATION_UNAVAILABLE; try { await session.piSession.abort(); } catch (err) { this.logger.error( `Legacy abort failed for conversation=${context.conversationId}`, err instanceof Error ? err.message : String(err), ); return { ok: false, code: 'operation_failed', retryable: false }; } return { ok: true, value: undefined }; } // ------------------------------------------------------------------------- // Model override (synchronous, total) // ------------------------------------------------------------------------- applyLegacyModelOverride( context: OwnedConversationContext, modelId: string, ): LegacyRuntimeResult { const scope = toScope(context.scope); const session = this.agentService.getSession(context.conversationId, scope); if (!session) return CONVERSATION_UNAVAILABLE; this.agentService.updateSessionModel(context.conversationId, modelId, scope); const refreshed = this.agentService.getSession(context.conversationId, scope) ?? session; return { ok: true, value: this.presentationForSession(refreshed) }; } // ------------------------------------------------------------------------- // Presentation read (synchronous, total) // ------------------------------------------------------------------------- readLegacySessionPresentation( context: OwnedConversationContext, ): LegacyRuntimeResult { const scope = toScope(context.scope); const session = this.agentService.getSession(context.conversationId, scope); if (!session) return CONVERSATION_UNAVAILABLE; return { ok: true, value: this.presentationForSession(session) }; } // ------------------------------------------------------------------------- // Verified Discord ingress (embedded-only in both modes) // ------------------------------------------------------------------------- async dispatchVerifiedDiscordIngress( context: VerifiedDiscordIngressContext, stream: LegacyRuntimeStream, ): Promise> { const scope = toScope(context.scope); const { conversationId } = context; const resolved = await this.resolveOrCreate( conversationId, scope, { agentConfigId: context.configuredAgent.agentConfigId }, { agentConfigId: context.configuredAgent.agentConfigId, instanceId: context.configuredAgent.instanceId, }, ); if (!resolved.ok) return resolved; let detach: () => void; try { detach = this.subscribe(conversationId, scope, stream); } catch (err) { // A partial listener/channel setup rolled itself back inside subscribe(); surface a total // safe failure instead of throwing out of the port. Retryable — the attach is transient. this.logger.error( `Embedded Discord subscription failed for conversation=${conversationId}`, err instanceof Error ? err.message : String(err), ); return { ok: false, code: 'runtime_unavailable', retryable: true }; } return { ok: true, value: this.buildLease( conversationId, scope, context.content, context.attachments, detach, resolved.presentation, ), }; } // ------------------------------------------------------------------------- // Shared helpers // ------------------------------------------------------------------------- /** * Resolves the owned session, creating it on first use. Ownership/scope rejections * (`Forbidden`/`NotFound`) collapse to `conversation_unavailable`; any other creation * failure surfaces as the retryable `runtime_unavailable`. On success returns the * session presentation so callers avoid a redundant `getSession`. */ private async resolveOrCreate( conversationId: string, scope: ActorTenantScope, extraOptions: Readonly<{ provider?: string; modelId?: string; agentConfigId?: string }>, expectedAgent?: Readonly<{ agentConfigId: string; instanceId: string }>, ): Promise< | { readonly ok: true; readonly presentation: LegacySessionPresentation } | Exclude, { ok: true }> > { // A verified-Discord turn may only run under a session whose configured identity matches the // reconciled agent record EXACTLY (config id + resolved name). This holds for BOTH a reused // pre-existing session AND a freshly created one: a session carrying a different configured // agent — however it arose — is rejected rather than executed under the verified label, so we // never silently run a different prompt/model/tool policy. A plain (non-verified) turn passes // no expectedAgent and skips the check. const identityMatches = (candidate: AgentSession): boolean => expectedAgent === undefined || (candidate.agentConfigId === expectedAgent.agentConfigId && candidate.agentName === expectedAgent.instanceId); let session = this.agentService.getSession(conversationId, scope); if (session && !identityMatches(session)) { // Reused same-scope session minted under a different configured identity — reject with zero // effects rather than dispatch a verified turn onto a foreign agent's session. return CONVERSATION_UNAVAILABLE; } if (!session) { try { session = await this.agentService.createSession(conversationId, { userId: scope.userId, tenantId: scope.tenantId, ...extraOptions, }); } catch (err) { if (err instanceof ForbiddenException || err instanceof NotFoundException) { return CONVERSATION_UNAVAILABLE; } this.logger.error( `Embedded session creation failed for conversation=${conversationId}`, err instanceof Error ? err.stack : String(err), ); return { ok: false, code: 'runtime_unavailable', retryable: true }; } // The just-created session must ALSO carry the reconciled identity before any effect. A // createSession that returns a session under a different configured agent (misconfiguration // or a substituted factory) is rejected here, before subscribe/persist/ack/prompt. if (!identityMatches(session)) { return CONVERSATION_UNAVAILABLE; } } return { ok: true, presentation: this.presentationForSession(session) }; } /** Installs a normalizing event listener that forwards to the server-owned stream. */ private subscribe( conversationId: string, scope: ActorTenantScope, stream: LegacyRuntimeStream, ): () => void { const unsubscribe = this.agentService.onEvent( conversationId, (event: AgentSessionEvent) => { const normalized = this.normalizeEvent(conversationId, scope, event); if (normalized) stream.onEvent(normalized); }, scope, ); try { this.agentService.addChannel(conversationId, stream.channelId, scope); } catch (err) { // Partial setup: the listener was acquired but the channel attach failed. Roll back // exactly what was acquired (the listener) before the failure escapes, so no leaked // subscription survives; the caller converts the rethrow into a total safe failure. try { unsubscribe(); } catch { /* idempotent teardown */ } throw err; } return () => { try { unsubscribe(); } catch { /* idempotent teardown */ } try { this.agentService.removeChannel(conversationId, stream.channelId, scope); } catch { /* idempotent teardown */ } }; } /** Builds an atomically one-shot, scope-rechecking dispatch lease. */ private buildLease( conversationId: string, scope: ActorTenantScope, content: string, attachments: VerifiedDiscordIngressContext['attachments'], detach: () => void, presentation: LegacySessionPresentation, ): LegacySocketTurnLease & VerifiedDiscordTurnLease { let dispatched = false; let disposed = false; return { presentation, dispatch: async (): Promise> => { if (dispatched) { return { ok: false, code: 'turn_already_dispatched', retryable: false }; } dispatched = true; try { await this.agentService.prompt(conversationId, content, scope, attachments); } catch (err) { this.logger.error( `Legacy dispatch failed for conversation=${conversationId}`, err instanceof Error ? err.message : String(err), ); return { ok: false, code: 'operation_failed', retryable: false }; } return { ok: true, value: undefined }; }, dispose: async (): Promise => { if (disposed) return; disposed = true; detach(); }, }; } /** Normalizes a raw agent event into the redaction-agnostic transport event, or drops it. */ private normalizeEvent( conversationId: string, scope: ActorTenantScope, event: AgentSessionEvent, ): LegacyRuntimeEvent | undefined { switch (event.type) { case 'agent_start': return { type: 'started' }; case 'agent_end': return { type: 'settled', ...this.usageFor(conversationId, scope) }; case 'message_update': { const assistant = event.assistantMessageEvent; if (assistant.type === 'text_delta') return { type: 'text_delta', text: assistant.delta }; if (assistant.type === 'thinking_delta') { return { type: 'thinking_delta', text: assistant.delta }; } return undefined; } case 'tool_execution_start': return { type: 'tool_started', toolCallId: event.toolCallId, toolName: event.toolName }; case 'tool_execution_end': return { type: 'tool_finished', toolCallId: event.toolCallId, toolName: event.toolName, isError: event.isError, }; default: return undefined; } } /** * Gathers terminal usage from the Pi session and records it into session metrics. * Embedded owns AgentService metrics; the gateway never touches `piSession` stats. */ private usageFor(conversationId: string, scope: ActorTenantScope): { usage?: LegacyUsage } { const session = this.agentService.getSession(conversationId, scope); const piSession = session?.piSession; const stats = piSession?.getSessionStats(); if (!session || !stats) return {}; const contextUsage = piSession?.getContextUsage(); const tokens = { input: stats.tokens?.input ?? 0, output: stats.tokens?.output ?? 0, cacheRead: stats.tokens?.cacheRead ?? 0, cacheWrite: stats.tokens?.cacheWrite ?? 0, total: stats.tokens?.total ?? 0, }; this.agentService.recordTokenUsage(conversationId, { ...tokens }); return { usage: { provider: session.provider, modelId: session.modelId, thinkingLevel: piSession?.thinkingLevel ?? 'off', tokens, cost: stats.cost ?? 0, context: { percent: contextUsage?.percent ?? null, window: contextUsage?.contextWindow ?? 0, }, }, }; } /** Presentation from a live session id, or undefined when no owned session exists. */ private presentationFor( conversationId: string, scope: ActorTenantScope, ): LegacySessionPresentation | undefined { const session = this.agentService.getSession(conversationId, scope); return session ? this.presentationForSession(session) : undefined; } /** User-facing projection carrying no session handle, credential, or raw stats. */ private presentationForSession(session: AgentSession): LegacySessionPresentation { return { provider: session.provider, modelId: session.modelId, thinkingLevel: session.piSession.thinkingLevel, availableThinkingLevels: session.piSession.getAvailableThinkingLevels(), ...(session.agentName ? { agentName: session.agentName } : {}), }; } } /** The shared terminal `conversation_unavailable` failure (missing/foreign/lost ownership). */ const CONVERSATION_UNAVAILABLE = { ok: false as const, code: 'conversation_unavailable' as const, retryable: false as const, }; /** Narrows a branded context scope to the `AgentService` actor/tenant scope (identical shape). */ function toScope(scope: Readonly<{ userId: string; tenantId: string }>): ActorTenantScope { return { userId: scope.userId, tenantId: scope.tenantId }; }