From c1db64dcb43c3396deb0ee2d9f4e4f2b9d60e1e6 Mon Sep 17 00:00:00 2001
From: achingbrain <alex@achingbrain.net>
Date: Wed, 19 Jun 2024 07:18:56 +0100
Subject: [PATCH] fix: retain invocation context

---
 packages/it-rpc/src/index.ts       |  8 ++++----
 packages/it-rpc/test/index.spec.ts | 15 +++++++++++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/packages/it-rpc/src/index.ts b/packages/it-rpc/src/index.ts
index 6f8e6e0e..bf84f97f 100644
--- a/packages/it-rpc/src/index.ts
+++ b/packages/it-rpc/src/index.ts
@@ -509,7 +509,7 @@ class DuplexRPC implements Duplex<AsyncGenerator<Uint8Array, void, unknown>> {
     const target = this.lookupInvocationTarget(message.path)
 
     // invoke the method
-    const val = await target.apply(context, message.args.map(arg => this.values.fromValue(arg, this.output, invocation)))
+    const val = await target.fn.apply(target.context, message.args.map(arg => this.values.fromValue(arg, this.output, invocation)))
 
     this.output.push(RPCMessage.encode({
       type: MessageType.methodResolved,
@@ -550,7 +550,7 @@ class DuplexRPC implements Duplex<AsyncGenerator<Uint8Array, void, unknown>> {
     const target = this.lookupInvocationTarget(message.path)
 
     // invoke the method
-    const gen = target.apply(context, message.args.map(arg => this.values.fromValue(arg, this.output, invocation)))
+    const gen = target.fn.apply(target.context, message.args.map(arg => this.values.fromValue(arg, this.output, invocation)))
 
     if (typeof gen.next !== 'function') {
       throw new InvalidReturnTypeError(`${message.path} did not return an async generator`)
@@ -601,7 +601,7 @@ class DuplexRPC implements Duplex<AsyncGenerator<Uint8Array, void, unknown>> {
     }))
   }
 
-  private lookupInvocationTarget (path: string): (...args: any[]) => any {
+  private lookupInvocationTarget (path: string): { context: any, fn(...args: any[]): any } {
     // look up the requested method on the target
     const pathParts = path.split('.')
     let target = this.targets.get(pathParts[0])
@@ -627,7 +627,7 @@ class DuplexRPC implements Duplex<AsyncGenerator<Uint8Array, void, unknown>> {
       throw new InvalidInvocationTypeError('Invocation target was not a function')
     }
 
-    return target
+    return { context, fn: target }
   }
 
   private async handleMethodResolved (message: MethodResolvedMessage, invocation: Invocation): Promise<void> {
diff --git a/packages/it-rpc/test/index.spec.ts b/packages/it-rpc/test/index.spec.ts
index de5baeca..79a402e6 100644
--- a/packages/it-rpc/test/index.spec.ts
+++ b/packages/it-rpc/test/index.spec.ts
@@ -1,4 +1,5 @@
 import { expect } from 'aegir/chai'
+import all from 'it-all'
 import { rpc } from '../src/index.js'
 import type { RPC } from '../src/index.js'
 
@@ -14,6 +15,12 @@ const target = {
   },
   async supportSimpleArguments (...args: any[]): Promise<any[]> {
     return args
+  },
+  async contextAccess (): Promise<boolean> {
+    return this.prop
+  },
+  async * generatorContextAccess (): AsyncGenerator<boolean> {
+    yield this.prop
   }
 }
 
@@ -78,4 +85,12 @@ describe('basics', () => {
 
     await expect(sender.supportSimpleArguments(...input)).to.eventually.deep.equal(input)
   })
+
+  it('should retain object context when invoking a method', async () => {
+    await expect(sender.contextAccess()).to.eventually.be.true()
+  })
+
+  it('should retain object context when invoking a generator', async () => {
+    await expect(all(sender.generatorContextAccess())).to.eventually.deep.equal([true])
+  })
 })