From fa4a2640112b6847257fd3dc0dc1a8c097eaa4dc Mon Sep 17 00:00:00 2001
From: Andreas Coroiu <andreas.coroiu@gmail.com>
Date: Fri, 31 Oct 2025 15:04:43 +0100
Subject: [PATCH] feat: support all responses

---
 .../services/sdk/rpc/batching-proxies.spec.ts | 65 ++++++++++++++++---
 .../services/sdk/rpc/batching-proxies.ts      |  9 ++-
 .../src/platform/services/sdk/rpc/error.ts    |  8 +++
 3 files changed, 71 insertions(+), 11 deletions(-)
 create mode 100644 libs/common/src/platform/services/sdk/rpc/error.ts

diff --git a/libs/common/src/platform/services/sdk/rpc/batching-proxies.spec.ts b/libs/common/src/platform/services/sdk/rpc/batching-proxies.spec.ts
index a400dd304a9..0811c119d45 100644
--- a/libs/common/src/platform/services/sdk/rpc/batching-proxies.spec.ts
+++ b/libs/common/src/platform/services/sdk/rpc/batching-proxies.spec.ts
@@ -20,16 +20,17 @@ describe("Batching proxies", () => {
     });
   });
 
-  // Not sure what await itself should do yet
-  // it("should allow awaiting the proxy itself", async () => {
-  //   const reference = {
-  //     referenceId: 2,
-  //     objectType: "AwaitableObject",
-  //   };
-  //   const proxy = RpcObjectReference(channel, reference);
-  //   const awaited = await proxy;
-  //   expect(awaited).toBe(proxy);
-  // });
+  it("awaiting the proxy returns the proxy itself", async () => {
+    const reference = {
+      referenceId: 1,
+      objectType: "TestObject",
+    };
+    const proxy = RpcObjectReference(channel, reference);
+
+    const awaited = await proxy;
+
+    expect(awaited).toBe(proxy);
+  });
 
   it("returns a pending object reference proxy when accesing a property", async () => {
     const reference = { referenceId: 1, objectType: "TestObject" };
@@ -133,6 +134,50 @@ describe("Batching proxies", () => {
       ],
     } satisfies Command);
   });
+
+  it("returns value when receiving a value response", async () => {
+    const reference = { referenceId: 1, objectType: "TestObject" };
+    const proxy = RpcObjectReference(channel, reference) as any;
+
+    channel.responses.push({ status: "success", result: { type: "value", value: 42 } });
+
+    const result = await proxy.prop;
+
+    expect(result).toBe(42);
+  });
+
+  it("returns reference when receiving a reference response", async () => {
+    const reference = { referenceId: 1, objectType: "TestObject" };
+    const proxy = RpcObjectReference(channel, reference) as any;
+
+    channel.responses.push({ status: "success", result: { type: "reference", referenceId: 2 } });
+
+    const result = await proxy.prop;
+
+    expect(result[ProxyInfo].referenceId).toBe(2);
+  });
+
+  it("throws error when receiving an error response", async () => {
+    class TestError extends Error {
+      constructor(
+        message: string,
+        public someProperty: string,
+      ) {
+        super(message);
+      }
+    }
+
+    const reference = { referenceId: 1, objectType: "TestObject" };
+    const proxy = RpcObjectReference(channel, reference) as any;
+    const error = new TestError("Something went wrong", "someValue");
+
+    channel.responses.push({ status: "error", error });
+
+    // Note: We are getting the actual error instance here because our
+    // channel is a local mock. In a real RPC scenario, the error would be
+    // serialized and deserialized, losing its prototype and properties.
+    await expect(proxy.prop).rejects.toThrow(error);
+  });
 });
 
 class RpcChannel implements RpcRequestChannel {
diff --git a/libs/common/src/platform/services/sdk/rpc/batching-proxies.ts b/libs/common/src/platform/services/sdk/rpc/batching-proxies.ts
index 06c33d3b659..7b55ee40de7 100644
--- a/libs/common/src/platform/services/sdk/rpc/batching-proxies.ts
+++ b/libs/common/src/platform/services/sdk/rpc/batching-proxies.ts
@@ -65,7 +65,14 @@ function proxyHandler(
         return (target as any)[ProxyInfo];
       }
 
-      if (property === "then") {
+      if (property === "then" && commands.length === 0) {
+        // This means we awaited a RpcObjectReference which resolves to itself
+        // We don't support transfering references to Promises themselves, we'll
+        // automatically await them before returning
+        return undefined;
+      }
+
+      if (property === "then" && commands.length > 0) {
         return BatchCommandExecutor(channel, reference.referenceId, commands);
       }
 
diff --git a/libs/common/src/platform/services/sdk/rpc/error.ts b/libs/common/src/platform/services/sdk/rpc/error.ts
new file mode 100644
index 00000000000..b421a6451f7
--- /dev/null
+++ b/libs/common/src/platform/services/sdk/rpc/error.ts
@@ -0,0 +1,8 @@
+/**
+ * RpcError represents an error that occurs during RPC communication.
+ */
+export class RpcError extends Error {
+  constructor(message: string) {
+    super(message);
+  }
+}