feat: initial implementation for users and markets WS improvements

LukasDeco · LukasDeco · commit f57f5c995f10 · 2025-08-06T13:24:40.000-06:00
diff --git a/sdk/src/accounts/webSocketDriftClientAccountSubscriberV2.ts b/sdk/src/accounts/webSocketDriftClientAccountSubscriberV2.ts
diff --git a/sdk/src/accounts/websocketProgramUserAccountSubscriber.ts b/sdk/src/accounts/websocketProgramUserAccountSubscriber.ts
@@ -0,0 +1,94 @@
+import {
+	DataAndSlot,
+	AccountSubscriber,
+	NotSubscribedError,
+	UserAccountEvents,
+	UserAccountSubscriber,
+	ResubOpts,
+} from './types';
+import { Program } from '@coral-xyz/anchor';
+import StrictEventEmitter from 'strict-event-emitter-types';
+import { EventEmitter } from 'events';
+import { Commitment, Context, PublicKey } from '@solana/web3.js';
+import { WebSocketProgramAccountSubscriber } from './webSocketProgramAccountSubscriber';
+import { UserAccount } from '../types';
+
+export class WebSocketProgramUserAccountSubscriber implements UserAccountSubscriber {
+	isSubscribed: boolean;
+	eventEmitter: StrictEventEmitter<EventEmitter, UserAccountEvents>;
+	
+	private userAccountPublicKey: PublicKey;
+	private program: Program;
+	private programSubscriber: WebSocketProgramAccountSubscriber<UserAccount>;
+	private userAccountAndSlot?: DataAndSlot<UserAccount>;
+
+	public constructor(
+		program: Program,
+		userAccountPublicKey: PublicKey,
+		programSubscriber: WebSocketProgramAccountSubscriber<UserAccount>
+	) {
+		this.isSubscribed = false;
+		this.program = program;
+		this.userAccountPublicKey = userAccountPublicKey;
+		this.eventEmitter = new EventEmitter();
+		this.programSubscriber = programSubscriber;
+	}
+
+	async subscribe(userAccount?: UserAccount): Promise<boolean> {
+		if (this.isSubscribed) {
+			return true;
+		}
+
+		if (userAccount) {
+			this.updateData(userAccount, 0);
+		}
+
+		this.programSubscriber.onChange = (
+			accountId: PublicKey,
+			data: UserAccount,
+			context: Context
+		) => {
+			if (accountId.equals(this.userAccountPublicKey)) {
+				this.updateData(data, context.slot);
+				this.eventEmitter.emit('userAccountUpdate', data);
+				this.eventEmitter.emit('update');
+			}
+		};
+
+		this.isSubscribed = true;
+		return true;
+	}
+
+	async fetch(): Promise<void> {
+		if (!this.isSubscribed) {
+			throw new NotSubscribedError(
+				'Must subscribe before fetching account updates'
+			);
+		}
+
+		const account = await this.program.account.user.fetch(
+			this.userAccountPublicKey
+		);
+		this.updateData(account as UserAccount, 0);
+	}
+
+	updateData(userAccount: UserAccount, slot: number): void {
+		this.userAccountAndSlot = {
+			data: userAccount,
+			slot,
+		};
+	}
+
+	async unsubscribe(): Promise<void> {
+		this.isSubscribed = false;
+	}
+
+	getUserAccountAndSlot(): DataAndSlot<UserAccount> {
+		if (!this.userAccountAndSlot) {
+			throw new NotSubscribedError(
+				'Must subscribe before getting user account data'
+			);
+		}
+		return this.userAccountAndSlot;
+	}
+}
diff --git a/sdk/src/driftClient.ts b/sdk/src/driftClient.ts
@@ -193,6 +193,7 @@ import { getOracleId } from './oracles/oracleId';
 import { SignedMsgOrderParams } from './types';
 import { sha256 } from '@noble/hashes/sha256';
 import { getOracleConfidenceFromMMOracleData } from './oracles/utils';
+import { WebSocketDriftClientAccountSubscriberV2 } from './accounts/webSocketDriftClientAccountSubscriberV2';
 
 type RemainingAccountParams = {
 	userAccounts: UserAccount[];
@@ -370,6 +371,7 @@ export class DriftClient {
 				resubTimeoutMs: config.accountSubscription?.resubTimeoutMs,
 				logResubMessages: config.accountSubscription?.logResubMessages,
 				commitment: config.accountSubscription?.commitment,
+				programUserAccountSubscriber: config.accountSubscription?.programUserAccountSubscriber,
 			};
 			this.userStatsAccountSubscriptionConfig = {
 				type: 'websocket',
@@ -435,7 +437,7 @@ export class DriftClient {
 				}
 			);
 		} else {
-			this.accountSubscriber = new WebSocketDriftClientAccountSubscriber(
+			this.accountSubscriber = new WebSocketDriftClientAccountSubscriberV2(
 				this.program,
 				config.perpMarketIndexes ?? [],
 				config.spotMarketIndexes ?? [],
@@ -608,7 +610,7 @@ export class DriftClient {
 	public getSpotMarketAccount(
 		marketIndex: number
 	): SpotMarketAccount | undefined {
-		return this.accountSubscriber.getSpotMarketAccountAndSlot(marketIndex).data;
+		return this.accountSubscriber.getSpotMarketAccountAndSlot(marketIndex)?.data;
 	}
 
 	/**
@@ -619,7 +621,7 @@ export class DriftClient {
 		marketIndex: number
 	): Promise<SpotMarketAccount | undefined> {
 		await this.accountSubscriber.fetch();
-		return this.accountSubscriber.getSpotMarketAccountAndSlot(marketIndex).data;
+		return this.accountSubscriber.getSpotMarketAccountAndSlot(marketIndex)?.data;
 	}
 
 	public getSpotMarketAccounts(): SpotMarketAccount[] {
@@ -928,6 +930,8 @@ export class DriftClient {
 		authority?: PublicKey,
 		userAccount?: UserAccount
 	): Promise<boolean> {
+
+		
 		authority = authority ?? this.authority;
 		const userKey = this.getUserMapKey(subAccountId, authority);
 
@@ -955,6 +959,7 @@ export class DriftClient {
 	 * Adds and subscribes to users based on params set by the constructor or by updateWallet.
 	 */
 	public async addAndSubscribeToUsers(authority?: PublicKey): Promise<boolean> {
+		console.log('adding and subscribing to users', this.users.size);
 		// save the rpc calls if driftclient is initialized without a real wallet
 		if (this.skipLoadUsers) return true;
 
diff --git a/sdk/src/driftClientConfig.ts b/sdk/src/driftClientConfig.ts
@@ -5,7 +5,7 @@ import {
 	PublicKey,
 	TransactionVersion,
 } from '@solana/web3.js';
-import { IWallet, TxParams } from './types';
+import { IWallet, TxParams, UserAccount } from './types';
 import { OracleInfo } from './oracles/types';
 import { BulkAccountLoader } from './accounts/bulkAccountLoader';
 import { DriftEnv } from './config';
@@ -19,6 +19,7 @@ import {
 import { Coder, Program } from '@coral-xyz/anchor';
 import { WebSocketAccountSubscriber } from './accounts/webSocketAccountSubscriber';
 import { WebSocketAccountSubscriberV2 } from './accounts/webSocketAccountSubscriberV2';
+import { WebSocketProgramAccountSubscriber } from './accounts/webSocketProgramAccountSubscriber';
 
 export type DriftClientConfig = {
 	connection: Connection;
@@ -71,6 +72,7 @@ export type DriftClientSubscriptionConfig =
 				resubOpts?: ResubOpts,
 				commitment?: Commitment
 			) => WebSocketAccountSubscriberV2<any> | WebSocketAccountSubscriber<any>;
+			programUserAccountSubscriber?: WebSocketProgramAccountSubscriber<UserAccount>;
 	  }
 	| {
 			type: 'polling';
diff --git a/sdk/src/index.ts b/sdk/src/index.ts
@@ -12,6 +12,8 @@ export * from './accounts/webSocketDriftClientAccountSubscriber';
 export * from './accounts/webSocketInsuranceFundStakeAccountSubscriber';
 export * from './accounts/webSocketHighLeverageModeConfigAccountSubscriber';
 export { WebSocketAccountSubscriberV2 } from './accounts/webSocketAccountSubscriberV2';
+export { WebSocketProgramAccountSubscriber } from './accounts/webSocketProgramAccountSubscriber';
+export { WebSocketProgramUserAccountSubscriber } from './accounts/websocketProgramUserAccountSubscriber';
 export * from './accounts/bulkAccountLoader';
 export * from './accounts/bulkUserSubscription';
 export * from './accounts/bulkUserStatsSubscription';
@@ -131,5 +133,6 @@ export * from './clock/clockSubscriber';
 export * from './math/userStatus';
 export * from './indicative-quotes/indicativeQuotesSender';
 export * from './constants';
+export { default as driftIDL } from './idl/drift.json';
 
 export { BN, PublicKey, pyth };
diff --git a/sdk/src/memcmp.ts b/sdk/src/memcmp.ts
@@ -112,3 +112,20 @@ export function getSignedMsgUserOrdersFilter(): MemcmpFilter {
 		},
 	};
 }
+
+export function getPerpMarketAccountsFilter(): MemcmpFilter {
+	return {
+		memcmp: {
+			offset: 0,
+			bytes: bs58.encode(BorshAccountsCoder.accountDiscriminator('PerpMarket')),
+		},
+	};
+}
+export function getSpotMarketAccountsFilter(): MemcmpFilter {
+	return {
+		memcmp: {
+			offset: 0,
+			bytes: bs58.encode(BorshAccountsCoder.accountDiscriminator('SpotMarket')),
+		},
+	};
+}
diff --git a/sdk/src/oracles/oracleId.ts b/sdk/src/oracles/oracleId.ts
@@ -24,9 +24,43 @@ export function getOracleSourceNum(source: OracleSource): number {
 	throw new Error('Invalid oracle source');
 }
 
+export function getOracleSourceFromNum(sourceNum: number): OracleSource {
+	if (sourceNum === OracleSourceNum.PYTH) return 'pyth';
+	if (sourceNum === OracleSourceNum.PYTH_1K) return 'pyth1K';
+	if (sourceNum === OracleSourceNum.PYTH_1M) return 'pyth1M';
+	if (sourceNum === OracleSourceNum.PYTH_PULL) return 'pythPull';
+	if (sourceNum === OracleSourceNum.PYTH_1K_PULL) return 'pyth1KPull';
+	if (sourceNum === OracleSourceNum.PYTH_1M_PULL) return 'pyth1MPull';
+	if (sourceNum === OracleSourceNum.SWITCHBOARD) return 'switchboard';
+	if (sourceNum === OracleSourceNum.QUOTE_ASSET) return 'quoteAsset';
+	if (sourceNum === OracleSourceNum.PYTH_STABLE_COIN) return 'pythStableCoin';
+	if (sourceNum === OracleSourceNum.PYTH_STABLE_COIN_PULL)
+		return 'pythStableCoinPull';
+	if (sourceNum === OracleSourceNum.PRELAUNCH) return 'prelaunch';
+	if (sourceNum === OracleSourceNum.SWITCHBOARD_ON_DEMAND)
+		return 'switchboardOnDemand';
+	if (sourceNum === OracleSourceNum.PYTH_LAZER) return 'pythLazer';
+	if (sourceNum === OracleSourceNum.PYTH_LAZER_1K) return 'pythLazer1K';
+	if (sourceNum === OracleSourceNum.PYTH_LAZER_1M) return 'pythLazer1M';
+	if (sourceNum === OracleSourceNum.PYTH_LAZER_STABLE_COIN)
+		return 'pythLazerStableCoin';
+	throw new Error('Invalid oracle source');
+}
+
 export function getOracleId(
 	publicKey: PublicKey,
 	source: OracleSource
 ): string {
 	return `${publicKey.toBase58()}-${getOracleSourceNum(source)}`;
 }
+
+export function getPublicKeyAndSourceFromOracleId(oracleId: string): {
+	publicKey: PublicKey;
+	source: OracleSource;
+} {
+	const [publicKey, source] = oracleId.split('-');
+	return {
+		publicKey: new PublicKey(publicKey),
+		source: getOracleSourceFromNum(parseInt(source)),
+	};
+}
diff --git a/sdk/src/user.ts b/sdk/src/user.ts
@@ -75,6 +75,7 @@ import {
 } from './types';
 import { standardizeBaseAssetAmount } from './math/orders';
 import { UserStats } from './userStats';
+import { WebSocketProgramUserAccountSubscriber } from './accounts/websocketProgramUserAccountSubscriber';
 import {
 	calculateAssetWeight,
 	calculateLiabilityWeight,
@@ -149,15 +150,23 @@ export class User {
 				}
 			);
 		} else {
-			this.accountSubscriber = new WebSocketUserAccountSubscriber(
-				config.driftClient.program,
-				config.userAccountPublicKey,
-				{
-					resubTimeoutMs: config.accountSubscription?.resubTimeoutMs,
-					logResubMessages: config.accountSubscription?.logResubMessages,
-				},
-				config.accountSubscription?.commitment
-			);
+			if (config.accountSubscription?.type === 'websocket' && config.accountSubscription?.programUserAccountSubscriber) {
+				this.accountSubscriber = new WebSocketProgramUserAccountSubscriber(
+					config.driftClient.program,
+					config.userAccountPublicKey,
+					config.accountSubscription.programUserAccountSubscriber
+				);
+			} else {
+				this.accountSubscriber = new WebSocketUserAccountSubscriber(
+					config.driftClient.program,
+					config.userAccountPublicKey,
+					{
+						resubTimeoutMs: config.accountSubscription?.resubTimeoutMs,
+						logResubMessages: config.accountSubscription?.logResubMessages,
+					},
+					config.accountSubscription?.commitment
+				);
+			}
 		}
 		this.eventEmitter = this.accountSubscriber.eventEmitter;
 	}
diff --git a/sdk/src/userConfig.ts b/sdk/src/userConfig.ts
@@ -2,6 +2,9 @@ import { DriftClient } from './driftClient';
 import { Commitment, PublicKey } from '@solana/web3.js';
 import { BulkAccountLoader } from './accounts/bulkAccountLoader';
 import { GrpcConfigs, UserAccountSubscriber } from './accounts/types';
+import { WebSocketProgramAccountSubscriber } from './accounts/webSocketProgramAccountSubscriber';
+import { UserAccount } from '@drift-labs/sdk';
+import { WebSocketProgramUserAccountSubscriber } from './accounts/websocketProgramUserAccountSubscriber';
 
 export type UserConfig = {
 	accountSubscription?: UserSubscriptionConfig;
@@ -21,6 +24,7 @@ export type UserSubscriptionConfig =
 			resubTimeoutMs?: number;
 			logResubMessages?: boolean;
 			commitment?: Commitment;
+			programUserAccountSubscriber?: WebSocketProgramAccountSubscriber<UserAccount>;
 	  }
 	| {
 			type: 'polling';
