refactor(gateway): retain Session contexts through dispatch

packages/api/gateway/src/client/remote-events.ts

@@ -8,8 +8,10 @@ import type {
 } from '@deepseek-ai/dsh-client-connection/client'
 import type {
   TypertClientEventListener,
+  TypertOwnedValue,
   TypertRemoteEvent,
 } from '@deepseek-ai/dsh-typert-protocol'
+import { isTypertOwnedValue } from '@deepseek-ai/dsh-typert-protocol'
 import { randomUUID } from '@deepseek-ai/dsh-util-crypto'
 import {
   REMOTE_EVENT_RESULT_ENDPOINT,
@@ -188,36 +190,42 @@ export class ClientRemoteEvents {
     signal: AbortSignal,
   ): Promise<void> {
     const adapter = this.ownerCtx.typert.contexts.getClient('agent')
-    let target: Context | undefined
+    let resolved: Context | TypertOwnedValue<Context> | undefined
     try {
-      target = adapter?.resolve(frame.agentId)
+      resolved = adapter?.resolve(frame.agentId)
     } catch (error) {
       this.reportError(frame.event, error)
     }
-    let outcome: RemoteEventReplyOutcome = { kind: 'next' }
-    if (target !== undefined) {
-      try {
-        outcome = await this.dispatchWaterfall(target, frame, signal)
-      } catch (error) {
-        if (signal.aborted) return
-        outcome = { kind: 'rejected', error: projectRemoteEventRejection(error) }
+    const owned = isTypertOwnedValue(resolved) ? resolved : undefined
+    try {
+      const target = isTypertOwnedValue(resolved) ? resolved.value : resolved
+      let outcome: RemoteEventReplyOutcome = { kind: 'next' }
+      if (target !== undefined) {
+        try {
+          outcome = await this.dispatchWaterfall(target, frame, signal)
+        } catch (error) {
+          if (signal.aborted) return
+          outcome = { kind: 'rejected', error: projectRemoteEventRejection(error) }
+        }
       }
+      if (signal.aborted) return
+      const result: RemoteEventResult = {
+        clientId,
+        eventId: frame.eventId,
+        outcome: outcome.kind === 'result' && outcome.value === undefined
+          ? { kind: 'result' }
+          : outcome,
+      }
+      const response = await this.connection.rpc.call(
+        '/api',
+        REMOTE_EVENT_RESULT_ENDPOINT,
+        { args: result },
+        signal,
+      )
+      if (!response.ok) throw new Error(response.error.message)
+    } finally {
+      owned?.[Symbol.dispose]()
     }
-    if (signal.aborted) return
-    const result: RemoteEventResult = {
-      clientId,
-      eventId: frame.eventId,
-      outcome: outcome.kind === 'result' && outcome.value === undefined
-        ? { kind: 'result' }
-        : outcome,
-    }
-    const response = await this.connection.rpc.call(
-      '/api',
-      REMOTE_EVENT_RESULT_ENDPOINT,
-      { args: result },
-      signal,
-    )
-    if (!response.ok) throw new Error(response.error.message)
   }
 
   private async dispatchWaterfall(
@@ -230,14 +238,12 @@ export class ClientRemoteEvents {
       agent: target,
       signal,
     }
-    const value = await abortable(
-      Promise.resolve(privateEvents(target).waterfall(
-        target,
-        this.eventKey(frame.event),
-        request,
-        () => Promise.resolve(REMOTE_EVENT_NEXT),
-      )),
-      signal,
+    // Cancellation reaches the handler, whose Context remains owned until it settles.
+    const value = await privateEvents(target).waterfall(
+      target,
+      this.eventKey(frame.event),
+      request,
+      () => Promise.resolve(REMOTE_EVENT_NEXT),
     )
     if (value !== REMOTE_EVENT_NEXT && value !== undefined && !isRemoteJsonValue(value)) {
       throw new TypeError('Remote event listener result is not lossless JSON data')
@@ -327,20 +333,6 @@ function invalidRemoteEventFrame(): never {
   throw new TypeError('client api: invalid forwarded Remote event frame')
 }
 
-/** Race listener completion against its delivery lifetime. */
-async function abortable<T>(value: T | PromiseLike<T>, signal: AbortSignal): Promise<T> {
-  signal.throwIfAborted()
-  let rejectAbort: ((reason: unknown) => void) | undefined
-  const aborted = new Promise<never>((_resolve, reject) => { rejectAbort = reject })
-  const onAbort = (): void => { rejectAbort?.(signal.reason) }
-  signal.addEventListener('abort', onAbort, { once: true })
-  try {
-    return await Promise.race([Promise.resolve(value), aborted])
-  } finally {
-    signal.removeEventListener('abort', onAbort)
-  }
-}
-
 function privateEvents(ctx: Context): PrivateEventContext {
   return ctx
 }