diff --git a/src/wrapper/WebsocketsClient.ts b/src/wrapper/WebsocketsClient.ts index 5e130b4..e36ade0 100644 --- a/src/wrapper/WebsocketsClient.ts +++ b/src/wrapper/WebsocketsClient.ts @@ -3,38 +3,50 @@ import type { WebsocketsSocket } from "../api/resources/websockets/client/Socket import * as core from "../core/index.js"; import * as environments from "../environments.js"; -export type GetPaymentCredentials = (wsUrl: string) => Promise>; +export type GetPaymentCredentials = ( + wsUrl: string +) => Promise>; export class WebsocketsClient extends FernWebsocketsClient { - private readonly _getPaymentCredentials: GetPaymentCredentials | undefined; + private readonly _getPaymentCredentials: GetPaymentCredentials | undefined; - constructor(options: FernWebsocketsClient.Options, getPaymentCredentials?: GetPaymentCredentials) { - super(options); - this._getPaymentCredentials = getPaymentCredentials; - } - - public override async connect(args: FernWebsocketsClient.ConnectArgs = {}): Promise { - let connectArgs = args; + constructor( + options: FernWebsocketsClient.Options, + getPaymentCredentials?: GetPaymentCredentials + ) { + super(options); + this._getPaymentCredentials = getPaymentCredentials; + } - if (this._getPaymentCredentials) { - const wsUrl = core.url.join( - (await core.Supplier.get(this._options.baseUrl)) ?? - ((await core.Supplier.get(this._options.environment)) ?? environments.AgentMailEnvironment.Prod) - .websockets, - "/v0", - ); - const credentials = await this._getPaymentCredentials(wsUrl); - connectArgs = { - ...args, - queryParams: { ...credentials, ...args.queryParams }, - }; - } else if (!args.apiKey) { - const apiKey = (await core.Supplier.get(this._options.apiKey)) ?? process.env.AGENTMAIL_API_KEY; - connectArgs = { ...args, apiKey }; - } + public override async connect( + args: FernWebsocketsClient.ConnectArgs & { waitForOpen?: boolean } = {} + ): Promise { + const { waitForOpen = true, ...rest } = args; + let connectArgs = rest; - const socket = await super.connect(connectArgs); - await socket.waitForOpen(); - return socket; + if (this._getPaymentCredentials) { + const wsUrl = core.url.join( + (await core.Supplier.get(this._options.baseUrl)) ?? + ( + (await core.Supplier.get(this._options.environment)) ?? + environments.AgentMailEnvironment.Prod + ).websockets, + "/v0" + ); + const credentials = await this._getPaymentCredentials(wsUrl); + connectArgs = { + ...rest, + queryParams: { ...credentials, ...rest.queryParams }, + }; + } else if (!rest.apiKey) { + const apiKey = + (await core.Supplier.get(this._options.apiKey)) ?? + process.env.AGENTMAIL_API_KEY; + connectArgs = { ...rest, apiKey }; } + + const socket = await super.connect(connectArgs); + if (waitForOpen) await socket.waitForOpen(); + return socket; + } }