From 12bc954711575ca661d733083eda9cd6abdf5eb3 Mon Sep 17 00:00:00 2001
From: lowkeynicc <nick@lunoho.company>
Date: Wed, 13 Nov 2024 16:19:45 -0500
Subject: [PATCH 1/2] sdk: add webSocketTokenAccountSubscriber

---
 .../webSocketTokenAccountSubscriber.ts        | 103 ++++++++++++++++++
 sdk/src/index.ts                              |   1 +
 2 files changed, 104 insertions(+)
 create mode 100644 sdk/src/accounts/webSocketTokenAccountSubscriber.ts

diff --git a/sdk/src/accounts/webSocketTokenAccountSubscriber.ts b/sdk/src/accounts/webSocketTokenAccountSubscriber.ts
new file mode 100644
index 0000000000..cfca3d3149
--- /dev/null
+++ b/sdk/src/accounts/webSocketTokenAccountSubscriber.ts
@@ -0,0 +1,103 @@
+import {
+	DataAndSlot,
+	AccountSubscriber,
+	NotSubscribedError,
+	ResubOpts,
+	TokenAccountEvents,
+	TokenAccountSubscriber,
+} from './types';
+import { Program } from '@coral-xyz/anchor';
+import StrictEventEmitter from 'strict-event-emitter-types';
+import { EventEmitter } from 'events';
+import { Commitment, PublicKey } from '@solana/web3.js';
+import { WebSocketAccountSubscriber } from './webSocketAccountSubscriber';
+import { Account } from '@solana/spl-token';
+
+export class WebSocketTokenAccountSubscriber implements TokenAccountSubscriber {
+	isSubscribed: boolean;
+	resubOpts?: ResubOpts;
+	commitment?: Commitment;
+	program: Program;
+	eventEmitter: StrictEventEmitter<EventEmitter, TokenAccountEvents>;
+	tokenAccountPublicKey: PublicKey;
+
+	tokenAccountSubscriber: AccountSubscriber<Account>;
+
+	public constructor(
+		program: Program,
+		tokenAccountPublicKey: PublicKey,
+		resubOpts?: ResubOpts,
+		commitment?: Commitment
+	) {
+		this.isSubscribed = false;
+		this.program = program;
+		this.resubOpts = resubOpts;
+		this.tokenAccountPublicKey = tokenAccountPublicKey;
+		this.eventEmitter = new EventEmitter();
+		this.commitment = commitment;
+	}
+
+	async subscribe(tokenAccount?: Account): Promise<boolean> {
+		if (this.isSubscribed) {
+			return true;
+		}
+
+		this.tokenAccountSubscriber = new WebSocketAccountSubscriber(
+			'user',
+			this.program,
+			this.tokenAccountPublicKey,
+			undefined,
+			this.resubOpts,
+			this.commitment
+		);
+
+		if (tokenAccount) {
+			this.tokenAccountSubscriber.setData(tokenAccount);
+		}
+
+		await this.tokenAccountSubscriber.subscribe((data: Account) => {
+			this.eventEmitter.emit('tokenAccountUpdate', data);
+			this.eventEmitter.emit('update');
+		});
+
+		this.eventEmitter.emit('update');
+		this.isSubscribed = true;
+		return true;
+	}
+
+	async fetch(): Promise<void> {
+		await Promise.all([this.tokenAccountSubscriber.fetch()]);
+	}
+
+	async unsubscribe(): Promise<void> {
+		if (!this.isSubscribed) {
+			return;
+		}
+
+		await Promise.all([this.tokenAccountSubscriber.unsubscribe()]);
+
+		this.isSubscribed = false;
+	}
+
+	assertIsSubscribed(): void {
+		if (!this.isSubscribed) {
+			throw new NotSubscribedError(
+				'You must call `subscribe` before using this function'
+			);
+		}
+	}
+
+	public getTokenAccountAndSlot(): DataAndSlot<Account> {
+		this.assertIsSubscribed();
+		return this.tokenAccountSubscriber.dataAndSlot;
+	}
+
+	public updateData(tokenAccount: Account, slot: number) {
+		const currentDataSlot = this.tokenAccountSubscriber.dataAndSlot?.slot || 0;
+		if (currentDataSlot <= slot) {
+			this.tokenAccountSubscriber.setData(tokenAccount, slot);
+			this.eventEmitter.emit('tokenAccountUpdate', tokenAccount);
+			this.eventEmitter.emit('update');
+		}
+	}
+}
diff --git a/sdk/src/index.ts b/sdk/src/index.ts
index dbf61296ac..11a3290ae8 100644
--- a/sdk/src/index.ts
+++ b/sdk/src/index.ts
@@ -12,6 +12,7 @@ export * from './accounts/fetch';
 export * from './accounts/webSocketDriftClientAccountSubscriber';
 export * from './accounts/webSocketInsuranceFundStakeAccountSubscriber';
 export * from './accounts/webSocketHighLeverageModeConfigAccountSubscriber';
+export * from './accounts/webSocketTokenAccountSubscriber';
 export * from './accounts/bulkAccountLoader';
 export * from './accounts/bulkUserSubscription';
 export * from './accounts/bulkUserStatsSubscription';

From d9624d3c1651b9cb691b0280ea9ade3886d25308 Mon Sep 17 00:00:00 2001
From: lowkeynicc <nick@lunoho.company>
Date: Wed, 13 Nov 2024 18:00:50 -0500
Subject: [PATCH 2/2] fix buffer fn

---
 sdk/src/accounts/webSocketTokenAccountSubscriber.ts | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/sdk/src/accounts/webSocketTokenAccountSubscriber.ts b/sdk/src/accounts/webSocketTokenAccountSubscriber.ts
index cfca3d3149..fb33d3d426 100644
--- a/sdk/src/accounts/webSocketTokenAccountSubscriber.ts
+++ b/sdk/src/accounts/webSocketTokenAccountSubscriber.ts
@@ -12,6 +12,7 @@ import { EventEmitter } from 'events';
 import { Commitment, PublicKey } from '@solana/web3.js';
 import { WebSocketAccountSubscriber } from './webSocketAccountSubscriber';
 import { Account } from '@solana/spl-token';
+import { parseTokenAccount } from '../token';
 
 export class WebSocketTokenAccountSubscriber implements TokenAccountSubscriber {
 	isSubscribed: boolean;
@@ -43,10 +44,10 @@ export class WebSocketTokenAccountSubscriber implements TokenAccountSubscriber {
 		}
 
 		this.tokenAccountSubscriber = new WebSocketAccountSubscriber(
-			'user',
+			'token',
 			this.program,
 			this.tokenAccountPublicKey,
-			undefined,
+			(buffer) => parseTokenAccount(buffer, this.tokenAccountPublicKey),
 			this.resubOpts,
 			this.commitment
 		);