wip: almost working background remote sdk

apps/browser/src/background/main.background.ts

@@ -321,6 +321,7 @@ import { BackgroundPlatformUtilsService } from "../platform/services/platform-ut
 import { BrowserPlatformUtilsService } from "../platform/services/platform-utils/browser-platform-utils.service";
 import { PopupViewCacheBackgroundService } from "../platform/services/popup-view-cache-background.service";
 import { BrowserSdkLoadService } from "../platform/services/sdk/browser-sdk-load.service";
+import { RemoteSdkServerService } from "../platform/services/sdk/remote-sdk-server.service";
 import { BackgroundTaskSchedulerService } from "../platform/services/task-scheduler/background-task-scheduler.service";
 import { BackgroundMemoryStorageService } from "../platform/storage/background-memory-storage.service";
 import { BrowserStorageServiceProvider } from "../platform/storage/browser-storage-service.provider";
@@ -474,6 +475,8 @@ export default class MainBackground {
   onReplacedRan: boolean;
   loginToAutoFill: CipherView = null;
 
+  remoteSdkServerService: RemoteSdkServerService;
+
   private commandsBackground: CommandsBackground;
   private contextMenusBackground: ContextMenusBackground;
   private idleBackground: IdleBackground;
@@ -1477,6 +1480,12 @@ export default class MainBackground {
       this.authService,
     );
 
+    this.remoteSdkServerService = new RemoteSdkServerService(
+      this.accountService,
+      this.authService,
+      this.sdkService,
+    );
+
     // Synchronous startup
     if (this.webPushConnectionService instanceof WorkerWebPushConnectionService) {
       this.webPushConnectionService.start();
@@ -1526,6 +1535,7 @@ export default class MainBackground {
     this.webRequestBackground?.startListening();
     this.syncServiceListener?.listener$().subscribe();
     await this.autoSubmitLoginBackground.init();
+    this.remoteSdkServerService.init();
 
     // If the user is logged out, switch to the next account
     const active = await firstValueFrom(this.accountService.activeAccount$);

apps/browser/src/platform/browser/browser-api.ts

@@ -470,7 +470,11 @@ export class BrowserApi {
 
   static messageListener$() {
     return new Observable<unknown>((subscriber) => {
-      const handler = (message: unknown) => {
+      const handler = (
+        message: unknown,
+        sender: chrome.runtime.MessageSender,
+        sendResponse: any,
+      ) => {
         subscriber.next(message);
       };
 

apps/browser/src/platform/popup/services/remote-sdk.service.ts

@@ -0,0 +1,57 @@
+import { filter, map, shareReplay } from "rxjs";
+
+import { Rc } from "@bitwarden/common/platform/misc/reference-counting/rc";
+import { RemoteSdkService } from "@bitwarden/common/platform/services/sdk/remote-sdk.service";
+import { RpcClient } from "@bitwarden/common/platform/services/sdk/rpc/client";
+import { Response } from "@bitwarden/common/platform/services/sdk/rpc/protocol";
+import { BitwardenClient } from "@bitwarden/sdk-internal";
+
+import { LogService } from "../../../../../../libs/logging/src";
+import { BrowserApi } from "../../browser/browser-api";
+import {
+  isRemoteSdkRoot,
+  RemoteSdkRequest,
+  RemoteSdkResendRootRequest,
+} from "../../services/sdk/messages";
+
+const root$ = BrowserApi.messageListener$().pipe(
+  // tap((message) => console.log("RemoteSdkService: Received message", message)),
+  filter(isRemoteSdkRoot),
+  map((message) => message.result),
+  shareReplay({ bufferSize: 1, refCount: false }),
+);
+
+export class BrowserRemoteSdkService implements RemoteSdkService {
+  private client = new RpcClient<Rc<BitwardenClient>>({
+    subscribeToRoot: () => {
+      void BrowserApi.sendMessage("sdk.request", {
+        type: "RemoteSdkResendRootRequest",
+      } satisfies RemoteSdkResendRootRequest);
+      return root$.pipe(
+        map((result) => ({
+          status: "success",
+          result,
+        })),
+      );
+    },
+    sendCommand: async (command) => {
+      this.logService.debug("[RemoteSdkService]: Sending command", command);
+      const response = (await BrowserApi.sendMessageWithResponse("sdk.request", {
+        type: "RemoteSdkRequest",
+        command,
+      } satisfies RemoteSdkRequest)) as Response;
+      this.logService.debug("[RemoteSdkService]: Received response", response);
+      return response;
+    },
+  });
+
+  constructor(private logService: LogService) {
+    // Eagerly subscribe to the remote client to make sure we catch the value.
+    // Also a hack, the server should send the value on request.
+    this.remoteClient$.subscribe();
+  }
+
+  get remoteClient$() {
+    return this.client.getRoot();
+  }
+}

apps/browser/src/platform/services/sdk/messages.ts

@@ -0,0 +1,57 @@
+import { Command, Response, Result } from "@bitwarden/common/platform/services/sdk/rpc/protocol";
+
+export type RemoteSdkRequest = {
+  type: "RemoteSdkRequest";
+  command: Command;
+};
+
+export type RemoteSdkResponse = {
+  type: "RemoteSdkResponse";
+  response: Response;
+};
+
+export type RemoteSdkResendRootRequest = {
+  type: "RemoteSdkResendRootRequest";
+};
+
+export type RemoteSdkRoot = {
+  type: "RemoteSdkRoot";
+  result: Result;
+};
+
+export function isRemoteSdkRequest(message: unknown): message is RemoteSdkRequest {
+  return (
+    typeof message === "object" &&
+    message !== null &&
+    (message as any).type === "RemoteSdkRequest" &&
+    typeof (message as any).command === "object"
+  );
+}
+
+export function isRemoteSdkResponse(message: unknown): message is RemoteSdkResponse {
+  return (
+    typeof message === "object" &&
+    message !== null &&
+    (message as any).type === "RemoteSdkResponse" &&
+    typeof (message as any).response === "object"
+  );
+}
+
+export function isRemoteSdkResendRootRequest(
+  message: unknown,
+): message is RemoteSdkResendRootRequest {
+  return (
+    typeof message === "object" &&
+    message !== null &&
+    (message as any).type === "RemoteSdkResendRootRequest"
+  );
+}
+
+export function isRemoteSdkRoot(message: unknown): message is RemoteSdkRoot {
+  return (
+    typeof message === "object" &&
+    message !== null &&
+    (message as any).type === "RemoteSdkRoot" &&
+    typeof (message as any).result === "object"
+  );
+}

apps/browser/src/platform/services/sdk/remote-sdk-server.service.ts

@@ -0,0 +1,77 @@
+import { combineLatest, map, of, startWith, Subject, switchMap } from "rxjs";
+
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { AuthService } from "@bitwarden/common/auth/abstractions/auth.service";
+import { AuthenticationStatus } from "@bitwarden/common/auth/enums/authentication-status";
+import { SdkService } from "@bitwarden/common/platform/abstractions/sdk/sdk.service";
+import { RpcServer } from "@bitwarden/common/platform/services/sdk/rpc/server";
+import { Rc } from "@bitwarden/common/src/platform/misc/reference-counting/rc";
+import { BitwardenClient } from "@bitwarden/sdk-internal";
+
+import { BrowserApi } from "../../browser/browser-api";
+
+import { isRemoteSdkRequest, isRemoteSdkResendRootRequest, RemoteSdkRoot } from "./messages";
+
+export class RemoteSdkServerService {
+  server = new RpcServer<Rc<BitwardenClient> | null>();
+
+  constructor(
+    private accountService: AccountService,
+    private authService: AuthService,
+    private sdkService: SdkService,
+  ) {}
+
+  init() {
+    // TODO: This is hacky because we don't support multiple roots.
+    // Needs to be fixed.
+    this.accountService.activeAccount$
+      .pipe(
+        switchMap((account) => {
+          if (!account) {
+            return of(null);
+          }
+
+          return this.authService
+            .authStatusFor$(account.id)
+            .pipe(map((status) => ({ account, status })));
+        }),
+        switchMap((accountStatus) => {
+          if (accountStatus == null || accountStatus.status !== AuthenticationStatus.Unlocked) {
+            return of(null);
+          }
+          return this.sdkService.userClient$(accountStatus.account.id);
+        }),
+      )
+      .subscribe((client) => {
+        this.server.setValue(client);
+      });
+
+    const resendRequest$ = new Subject<void>();
+
+    combineLatest({
+      root: this.server.value$,
+      resendRequest: resendRequest$.pipe(startWith(null)),
+    }).subscribe(({ root }) => {
+      void BrowserApi.sendMessage("sdk.root", {
+        type: "RemoteSdkRoot",
+        result: root,
+      } satisfies RemoteSdkRoot);
+    });
+
+    BrowserApi.messageListener(
+      "sdk.request",
+      (message: any, sender: chrome.runtime.MessageSender, sendResponse: any) => {
+        if (isRemoteSdkRequest(message)) {
+          void this.server.handle(message.command).then((response) => {
+            sendResponse(response);
+          });
+          return true; // Indicate that we will send a response asynchronously
+        }
+
+        if (isRemoteSdkResendRootRequest(message)) {
+          resendRequest$.next();
+        }
+      },
+    );
+  }
+}

apps/browser/src/popup/services/services.module.ts

@@ -112,6 +112,7 @@ import { ConsoleLogService } from "@bitwarden/common/platform/services/console-l
 import { ContainerService } from "@bitwarden/common/platform/services/container.service";
 import { DefaultSdkClientFactory } from "@bitwarden/common/platform/services/sdk/default-sdk-client-factory";
 import { NoopSdkClientFactory } from "@bitwarden/common/platform/services/sdk/noop-sdk-client-factory";
+import { RemoteSdkService } from "@bitwarden/common/platform/services/sdk/remote-sdk.service";
 import { StorageServiceProvider } from "@bitwarden/common/platform/services/storage-service.provider";
 import { PrimarySecondaryStorageService } from "@bitwarden/common/platform/storage/primary-secondary-storage.service";
 import { WindowStorageService } from "@bitwarden/common/platform/storage/window-storage.service";
@@ -178,6 +179,7 @@ import { OffscreenDocumentService } from "../../platform/offscreen-document/abst
 import { DefaultOffscreenDocumentService } from "../../platform/offscreen-document/offscreen-document.service";
 import { PopupCompactModeService } from "../../platform/popup/layout/popup-compact-mode.service";
 import { BrowserFileDownloadService } from "../../platform/popup/services/browser-file-download.service";
+import { BrowserRemoteSdkService } from "../../platform/popup/services/remote-sdk.service";
 import { PopupViewCacheService } from "../../platform/popup/view-cache/popup-view-cache.service";
 import { ScriptInjectorService } from "../../platform/services/abstractions/script-injector.service";
 import { BrowserEnvironmentService } from "../../platform/services/browser-environment.service";
@@ -312,7 +314,7 @@ const safeProviders: SafeProvider[] = [
   safeProvider({
     provide: TotpServiceAbstraction,
     useClass: TotpService,
-    deps: [SdkService],
+    deps: [SdkService, RemoteSdkService],
   }),
   safeProvider({
     provide: OffscreenDocumentService,
@@ -713,6 +715,11 @@ const safeProviders: SafeProvider[] = [
     useClass: ExtensionNewDeviceVerificationComponentService,
     deps: [],
   }),
+  safeProvider({
+    provide: RemoteSdkService,
+    useClass: BrowserRemoteSdkService,
+    deps: [LogService],
+  }),
 ];
 
 @NgModule({

libs/common/src/billing/services/account/billing-account-profile-state.service.ts

@@ -38,15 +38,14 @@ export class DefaultBillingAccountProfileStateService implements BillingAccountP
   }
 
   hasPremiumFromAnySource$(userId: UserId): Observable<boolean> {
-    return this.stateProvider
-      .getUser(userId, BILLING_ACCOUNT_PROFILE_KEY_DEFINITION)
-      .state$.pipe(
-        map(
-          (profile) =>
-            profile?.hasPremiumFromAnyOrganization === true ||
-            profile?.hasPremiumPersonally === true,
-        ),
-      );
+    return this.stateProvider.getUser(userId, BILLING_ACCOUNT_PROFILE_KEY_DEFINITION).state$.pipe(
+      map(() => true),
+      // map(
+      //   (profile) =>
+      //     profile?.hasPremiumFromAnyOrganization === true ||
+      //     profile?.hasPremiumPersonally === true,
+      // ),
+    );
   }
 
   async setHasPremium(

libs/common/src/platform/services/sdk/remote-sdk.service.ts

@@ -1,11 +1,11 @@
-import { BitwardenClient, FolderView } from "@bitwarden/sdk-internal";
+import { Observable } from "rxjs";
+
+import { BitwardenClient } from "@bitwarden/sdk-internal";
+
+import { Rc } from "../../misc/reference-counting/rc";
 
 import { Remote } from "./remote";
 
-export type RemoteSdk = Remote<BitwardenClient>;
-
-const remoteClient: RemoteSdk = {} as RemoteSdk;
-
-export async function test(): Promise<FolderView[]> {
-  return await remoteClient.vault().await.folders().await.list();
+export abstract class RemoteSdkService {
+  abstract remoteClient$: Observable<Remote<Rc<BitwardenClient>>>;
 }

libs/common/src/platform/services/sdk/remote.ts

@@ -4,24 +4,41 @@ export type Remote<T> = {
   [K in keyof T]: RemoteProperty<T[K]>;
 };
 
+type Resolved<T> = T extends Promise<infer U> ? U : T;
+
+/**
+ * Maps remote object fields to RPC-exposed types.
+ *
+ * Property access (non-function):
+ * - If the value is serializable (see IsSerializable), returns Promise<Resolved<T>>.
+ * - If not serializable (e.g., class instance, Wasm object), returns Remote<Resolved<T>> (a live reference).
+ *   Note: properties do NOT expose `.await`; they are direct remote references.
+ *
+ * Function call:
+ * - If the return value is serializable, returns Promise<Resolved<R>>.
+ * - If not serializable, returns ChainablePromise<Remote<Resolved<R>>> so callers can use `.await`
+ *   for ergonomic chaining, e.g. remote.vault().await.totp().await.generate(...).
+ */
 export type RemoteProperty<T> = T extends (...args: any[]) => any
   ? RemoteFunction<T>
-  : RemoteValue<T>;
+  : IsSerializable<Resolved<T>> extends true
+    ? Promise<Resolved<T>>
+    : Remote<Resolved<T>>;
 
 export type RemoteReference<T> = Remote<T>;
 
-export type RemoteValue<T> = T extends { free(): void }
-  ? ChainablePromise<RemoteReference<T>>
-  : T extends Promise<infer R>
-    ? Promise<R>
-    : Promise<T>;
-
+/**
+ * RemoteFunction arguments must be Serializable at compile time. For non-serializable
+ * return types, we expose ChainablePromise<Remote<...>> to enable Rust-like `.await` chaining.
+ */
 export type RemoteFunction<T extends (...args: any[]) => any> = <A extends Parameters<T>>(
   // Enforce serializability of RPC arguments here.
   // If we wanted to we could allow for remote references as arguments, we could do that here.
   // In that case the client would also need to maintain a ReferenceStore for outgoing references.
   ...args: SerializableArgs<A>
-) => RemoteValue<ReturnType<T>>;
+) => IsSerializable<Resolved<ReturnType<T>>> extends true
+  ? Promise<Resolved<ReturnType<T>>>
+  : ChainablePromise<Remote<Resolved<ReturnType<T>>>>;
 
 // Serializable type rules to mirror `isSerializable` from rpc/server.ts
 // - Primitives: string | number | boolean | null

libs/common/src/platform/services/sdk/rpc/proxies.ts

@@ -1,3 +1,5 @@
+import { chain } from "../chainable-promise";
+
 import { RpcRequestChannel } from "./client";
 import { Command, PropertySymbol, ReferenceId, Response, serializeSymbol } from "./protocol";
 
@@ -86,16 +88,22 @@ function RpcPropertyReference(channel: RpcRequestChannel, reference: RpcProperty
       //   `Accessing ${reference.objectReference.objectType}.${reference.propertyName}.${propertyName}`,
       // );
 
-      // Allow Function.prototype.call to be used by TS helpers (e.g., disposables)
+  // Allow Function.prototype.call/apply/bind to be used by TS helpers and wrappers (e.g., disposables, chainable await)
       if (propertyName === "call") {
         return Function.prototype.call;
       }
+      if (propertyName === "apply") {
+        return Function.prototype.apply;
+      }
+      if (propertyName === "bind") {
+        return Function.prototype.bind;
+      }
 
       if (propertyName !== "then") {
         // Support chained call like: (await obj.prop).method() AND obj.prop.method()
         // by lazily resolving obj.prop first, then invoking method/property on the resolved reference.
         return (...argArray: unknown[]) => {
-          return (async () => {
+          const p = (async () => {
             // First resolve the original referenced property value via GET
             const getResult = await sendAndUnwrap(channel, buildGetCommand(reference));
             if (getResult.type !== "reference") {
@@ -109,12 +117,20 @@ function RpcPropertyReference(channel: RpcRequestChannel, reference: RpcProperty
               channel,
               buildCallCommand(getResult.referenceId, propertyName, argArray),
             );
-            return unwrapResult(channel, callResult);
+            if (callResult.type === "value") {
+              return callResult.value;
+            }
+            return RpcObjectReference.create(
+              channel,
+              callResult.referenceId,
+              callResult.objectType,
+            );
           })();
+          return chain(p as Promise<any>);
         };
       }
 
-      return (onFulfilled: (value: any) => void, onRejected: (error: any) => void) => {
+  return (onFulfilled: (value: any) => void, onRejected: (error: any) => void) => {
         (async () => {
           const result = await sendAndUnwrap(channel, buildGetCommand(reference));
           return unwrapResult(channel, result);
@@ -129,7 +145,14 @@ function RpcPropertyReference(channel: RpcRequestChannel, reference: RpcProperty
         reference.property,
         argArray,
       );
-      return sendAndUnwrap(channel, command).then((result) => unwrapResult(channel, result));
+      const p = (async () => {
+        const result = await sendAndUnwrap(channel, command);
+        if (result.type === "value") {
+          return result.value;
+        }
+        return RpcObjectReference.create(channel, result.referenceId, result.objectType);
+      })();
+  return chain(p as Promise<any>);
     },
   });
 }

libs/common/src/platform/services/sdk/rpc/rpc.spec.ts

@@ -76,7 +76,7 @@ describe("RpcServer", () => {
     {
       using ref = await remoteRc.take();
       const remoteWasmObj = ref.value;
-      const greeting = await remoteWasmObj.greet();
+      const greeting = await remoteWasmObj.testClass().await.greet("RC");
       expect(greeting).toBe("Hello, RC!");
       expect(wasmObj.freed).toBe(false);
     }
@@ -121,6 +121,10 @@ class WasmLikeObject {
   async greet(): Promise<string> {
     return `Hello, ${this.name}!`;
   }
+
+  async testClass(): Promise<TestClass> {
+    return new TestClass();
+  }
 }
 
 class InMemoryChannel<T> implements RpcRequestChannel {

libs/common/src/platform/services/sdk/rpc/server.ts

@@ -49,7 +49,7 @@ export class RpcServer<T> {
         if (typeof method !== "function") {
           return {
             status: "error",
-            error: `[RPC] Property ${String(propertyKey)} is not a function`,
+            error: `[RPC] Property ${String(propertyKey)} is not a function of ${target.constructor.name}`,
           };
         }
 

libs/common/src/vault/services/totp.service.ts

@@ -1,8 +1,9 @@
-import { Observable, map, shareReplay, switchMap, timer } from "rxjs";
+import { map, Observable, shareReplay, switchMap, timer } from "rxjs";
 
 import { TotpResponse } from "@bitwarden/sdk-internal";
 
 import { SdkService } from "../../platform/abstractions/sdk/sdk.service";
+import { RemoteSdkService } from "../../platform/services/sdk/remote-sdk.service";
 import { TotpService as TotpServiceAbstraction } from "../abstractions/totp.service";
 
 /**
@@ -26,18 +27,33 @@ export type TotpInfo = {
 };
 
 export class TotpService implements TotpServiceAbstraction {
-  constructor(private sdkService: SdkService) {}
+  constructor(
+    private sdkService: SdkService,
+    private remoteSdkService?: RemoteSdkService,
+  ) {}
 
   getCode$(key: string): Observable<TotpResponse> {
     return timer(0, 1000).pipe(
-      switchMap(() =>
-        this.sdkService.client$.pipe(
-          map((sdk) => {
-            return sdk.vault().totp().generate_totp(key);
-          }),
-        ),
-      ),
-      shareReplay({ refCount: true, bufferSize: 1 }),
+      switchMap(() => {
+        if (this.remoteSdkService) {
+          console.log("[TOTP] Using remote SDK service to generate TOTP");
+          return this.remoteSdkService.remoteClient$.pipe(
+            switchMap(async (sdk) => {
+              using ref = await sdk.take();
+              console.log("TOTP", await ref.value.vault().await.totp());
+              return ref.value.vault().await.totp().await.generate_totp(key);
+            }),
+            shareReplay({ bufferSize: 1, refCount: true }),
+          );
+        } else {
+          return this.sdkService.client$.pipe(
+            map((sdk) => {
+              return sdk.vault().totp().generate_totp(key);
+            }),
+            shareReplay({ bufferSize: 1, refCount: true }),
+          );
+        }
+      }),
     );
   }
 }