fix(tempo): keep selected session rail sticky

Author: parvahuja

src/tempo/session/client/Session.test.ts

@@ -459,6 +459,145 @@ describe('precompile client session', () => {
     expect(payload.authorizationSignature).toBeDefined()
   })
 
+  test.each([
+    [
+      'the route becomes inactive',
+      { active: false },
+      'Machine-token channel is not bound to this merchant session challenge.',
+    ],
+    ['the route RPC fails', { routeError: true }, 'route RPC failed'],
+  ] as const)(
+    'keeps a selected machine-token channel sticky when %s',
+    async (_name, clientOptions, expectedError) => {
+      const machineDescriptor = getMachineDescriptor()
+      const machineKey = channelKey({
+        chainId,
+        escrow: tip20ChannelEscrow,
+        payee: machineDescriptor.payee,
+        token: machineDescriptor.token,
+      })
+      const channelId = Channel.computeId({
+        ...machineDescriptor,
+        chainId,
+        escrow: tip20ChannelEscrow,
+      })
+      const channelStore = createChannelStore()
+      await channelStore.set({
+        chainId,
+        channelId,
+        cumulativeAmount: 100n,
+        deposit: 1_000n,
+        descriptor: machineDescriptor,
+        escrow: tip20ChannelEscrow,
+        opened: true,
+      })
+      let currentClient = createMachineClient({ balance: 0n })
+      const resolveAccount = vi.fn(() => account)
+      const onChannelUpdate = vi.fn()
+      const method = session({
+        account,
+        channelStore,
+        decimals: 0,
+        getClient: () => currentClient,
+        onChannelUpdate,
+        resolveAccount,
+      })
+      const challenge = () =>
+        makeSessionChallenge({
+          methodDetails: { escrowContract: tip20ChannelEscrow, machineTokenEnabled: true },
+        })
+
+      expect(
+        deserialize(await method.createCredential({ challenge: challenge(), context: {} })),
+      ).toMatchObject({
+        action: 'voucher',
+        channelId,
+        cumulativeAmount: '200',
+      })
+      currentClient = createMachineClient(clientOptions)
+      resolveAccount.mockClear()
+
+      await expect(
+        method.createCredential({ challenge: challenge(), context: {} }),
+      ).rejects.toThrow(expectedError)
+      expect(resolveAccount).not.toHaveBeenCalled()
+      expect(await channelStore.get(defaultChannelKey)).toBeUndefined()
+      expect(await channelStore.get(machineKey)).toMatchObject({
+        channelId,
+        cumulativeAmount: 200n,
+        opened: true,
+      })
+      expect(onChannelUpdate).toHaveBeenCalledTimes(1)
+    },
+  )
+
+  test('keeps a selected direct channel direct when a machine channel later appears', async () => {
+    const directChannelId = Channel.computeId({
+      ...descriptor,
+      chainId,
+      escrow: tip20ChannelEscrow,
+    })
+    const channelStore = createChannelStore()
+    await channelStore.set({
+      chainId,
+      channelId: directChannelId,
+      cumulativeAmount: 100n,
+      deposit: 1_000n,
+      descriptor,
+      escrow: tip20ChannelEscrow,
+      opened: true,
+    })
+    const method = session({
+      account,
+      channelStore,
+      decimals: 0,
+      getClient: () => machineClient,
+    })
+    const challenge = () =>
+      makeSessionChallenge({
+        methodDetails: { escrowContract: tip20ChannelEscrow, machineTokenEnabled: true },
+      })
+
+    expect(
+      deserialize(await method.createCredential({ challenge: challenge(), context: {} })),
+    ).toMatchObject({
+      action: 'voucher',
+      channelId: directChannelId,
+      cumulativeAmount: '200',
+    })
+
+    const machineDescriptor = getMachineDescriptor()
+    const machineChannelId = Channel.computeId({
+      ...machineDescriptor,
+      chainId,
+      escrow: tip20ChannelEscrow,
+    })
+    const machineKey = channelKey({
+      chainId,
+      escrow: tip20ChannelEscrow,
+      payee: machineDescriptor.payee,
+      token: machineDescriptor.token,
+    })
+    await channelStore.set({
+      chainId,
+      channelId: machineChannelId,
+      cumulativeAmount: 100n,
+      deposit: 1_000n,
+      descriptor: machineDescriptor,
+      escrow: tip20ChannelEscrow,
+      opened: true,
+    })
+
+    expect(
+      deserialize(await method.createCredential({ challenge: challenge(), context: {} })),
+    ).toMatchObject({
+      action: 'voucher',
+      channelId: directChannelId,
+      cumulativeAmount: '300',
+    })
+    expect(await channelStore.get(machineKey)).toMatchObject({ cumulativeAmount: 100n })
+  })
+
   test('rejects an underfunded machine-token top-up without mutating channel state', async () => {
     const machineDescriptor = getMachineDescriptor()
     const channelId = Channel.computeId({

src/tempo/session/client/Session.ts

@@ -56,6 +56,11 @@ export { sessionContextSchema, type SessionContext } from './CredentialState.js'
 /** Tail promise for each payment scope, used to serialize automatic opens per store. */
 const channelTails = new WeakMap<ChannelStore, Map<string, Promise<void>>>()
 
+type SelectedRail = { type: 'direct' } | { type: 'machine'; route: MachineTokenSession.Route }
+
+/** Selected rail for each logical payment scope, shared by methods using one store. */
+const selectedRails = new WeakMap<ChannelStore, Map<string, SelectedRail>>()
+
 /** Serializes automatic opens for one payment scope across methods sharing a store. */
 async function lockChannel(store: ChannelStore, key: string) {
   const tails = channelTails.get(store) ?? new Map<string, Promise<void>>()
@@ -199,6 +204,8 @@ export function session(parameters: session.Parameters = {}) {
     topUpAmountParameter !== undefined ? parseUnits(topUpAmountParameter, decimals) : undefined
   const store = channelStore ?? createChannelStore()
   const sink = { store, notifyUpdate: (entry: ChannelEntry) => onChannelUpdate?.(entry) }
+  const rails = selectedRails.get(store) ?? new Map<string, SelectedRail>()
+  selectedRails.set(store, rails)
 
   type ResolvedSelection = {
     account: ViemAccount
@@ -254,6 +261,48 @@ export function session(parameters: session.Parameters = {}) {
     resolved,
   })
 
+  const rememberRail = (
+    direct: ChallengeContext,
+    resolved: ChallengeContext,
+    entry: ChannelEntry,
+  ) => {
+    if (!entry.opened) return
+    rails.set(
+      direct.key,
+      resolved.machineRoute
+        ? { type: 'machine', route: resolved.machineRoute }
+        : { type: 'direct' },
+    )
+  }
+
+  const resolveSelectedChannel = async (
+    direct: ChallengeContext,
+  ): Promise<{ channel: ChannelEntry; resolved: ChallengeContext } | undefined> => {
+    const rail = rails.get(direct.key)
+    if (!rail) return undefined
+    const cached = rail.type === 'direct' ? direct : applyMachineTokenRoute(direct, rail.route)
+    const channel = await store.get(cached.key)
+    if (!channel?.opened) {
+      rails.delete(direct.key)
+      return undefined
+    }
+    if (rail.type === 'direct') return { channel, resolved: direct }
+    if (
+      !MachineTokenSession.isEnabledChallenge(direct.challenge) ||
+      !isAddressEqual(direct.escrow, tip20ChannelEscrow)
+    )
+      throw new Error('Machine-token channel is not bound to this merchant session challenge.')
+    const route = await MachineTokenSession.matchRoute(direct.client, {
+      chainId: direct.chainId,
+      descriptor: rail.route,
+      merchant: direct.payee,
+      targetToken: direct.token,
+    })
+    if (!route)
+      throw new Error('Machine-token channel is not bound to this merchant session challenge.')
+    return { channel, resolved: applyMachineTokenRoute(direct, route) }
+  }
+
   /** Selects one rail for this channel lifecycle before any credential is signed. */
   const resolveSelection = async (
     direct: ChallengeContext,
@@ -267,7 +316,13 @@ export function session(parameters: session.Parameters = {}) {
       return select(resolved, context, await store.get(resolved.key))
     }
 
-    const selectDirect = async () => select(direct, context, await store.get(direct.key))
+    const selectDirect = async () => {
+      const entry = await store.get(direct.key)
+      if (entry?.opened) rememberRail(direct, direct, entry)
+      return select(direct, context, entry)
+    }
+    const selected = await resolveSelectedChannel(direct)
+    if (selected) return select(selected.resolved, context, selected.channel)
     if (
       !MachineTokenSession.isEnabledChallenge(direct.challenge) ||
       !isAddressEqual(direct.escrow, tip20ChannelEscrow)
@@ -283,10 +338,16 @@ export function session(parameters: session.Parameters = {}) {
 
     const machine = applyMachineTokenRoute(direct, route)
     const machineEntry = await store.get(machine.key)
-    if (machineEntry?.opened) return select(machine, context, machineEntry)
+    if (machineEntry?.opened) {
+      rememberRail(direct, machine, machineEntry)
+      return select(machine, context, machineEntry)
+    }
 
     const directEntry = await store.get(direct.key)
-    if (directEntry?.opened) return select(direct, context, directEntry)
+    if (directEntry?.opened) {
+      rememberRail(direct, direct, directEntry)
+      return select(direct, context, directEntry)
+    }
 
     const account = await resolveCredentialAccount(machine, context, undefined)
     const openingDeposit = resolveOpeningDeposit({
@@ -320,6 +381,9 @@ export function session(parameters: session.Parameters = {}) {
       return { channel: await store.get(resolved.key), resolved }
     }
 
+    const selected = await resolveSelectedChannel(direct)
+    if (selected) return selected
+
     if (
       MachineTokenSession.isEnabledChallenge(direct.challenge) &&
       isAddressEqual(direct.escrow, tip20ChannelEscrow)
@@ -449,8 +513,13 @@ export function session(parameters: session.Parameters = {}) {
           payload = { ...payload, authorizationSignature }
         }
         if (pendingEntry && !(attempt && plan.type === 'open')) {
-          if (pendingEntry.opened) await store.set(pendingEntry)
-          else await store.delete(resolved.key)
+          if (pendingEntry.opened) {
+            await store.set(pendingEntry)
+            rememberRail(direct, resolved, pendingEntry)
+          } else {
+            await store.delete(resolved.key)
+            rails.delete(direct.key)
+          }
           sink.notifyUpdate(pendingEntry)
         }
         const credential = await serializeCredential(
@@ -475,6 +544,7 @@ export function session(parameters: session.Parameters = {}) {
                   await store.set(opened)
                   sink.notifyUpdate(opened)
                 }
+                rememberRail(direct, resolved, opened)
               }
               return true
             } finally {