@zofai/zo-sdk 0.2.16 → 0.2.18

package/src/storkClient.ts

@@ -326,3 +326,311 @@ export async function fetchStorkUpdateDataFromKronos(
   }
   return updates
 }
+
+export type StorkSignatureType = 'evm' | 'stark'
+
+export interface StorkStreamUpdate {
+  traceId?: string
+  /** Stork asset symbol (e.g. `BTCUSD`) → kronos price entry. */
+  prices: Record<string, IStorkKronosPriceEntry>
+}
+
+export type StorkStreamServerMessage
+  = | { type: 'connected', message?: string }
+    | { type: 'oracle_prices', trace_id?: string, data: Record<string, IStorkKronosPriceEntry> }
+    | { type: 'error', message?: string }
+    | { type: 'pong' }
+
+export type StorkStreamClientMessage
+  = | { type: 'subscribe', data: string[] }
+    | { type: 'unsubscribe', data: string[] }
+    | { type: 'ping' }
+
+/** Derive ws(s):// stream URL from kronos http(s) base URL. */
+export function apiBaseUrlToStorkWsUrl(
+  baseUrl: string,
+  signatureType: StorkSignatureType = 'evm',
+): string {
+  const url = new URL(baseUrl)
+  url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'
+  url.pathname = `/stork/stream/${signatureType}`
+  url.search = ''
+  url.hash = ''
+  return url.toString()
+}
+
+export interface IStorkStreamClientOptions {
+  /** Defaults to `evm` (`/stork/stream/evm`). Use `stark` for StarkEx-native signatures. */
+  signatureType?: StorkSignatureType
+  /** Defaults to apiBaseUrlToStorkWsUrl(baseUrl, signatureType). */
+  wsUrl?: string
+  /** Reconnect after disconnect. Default true. */
+  autoReconnect?: boolean
+  /** Initial reconnect delay in ms. Default 1000. */
+  reconnectDelayMs?: number
+  /** Ping interval in ms. Default 30000. Set 0 to disable. */
+  pingIntervalMs?: number
+  /** Custom WebSocket constructor (Node test mocks, etc.). */
+  WebSocketImpl?: typeof WebSocket
+}
+
+type StorkStreamUpdateHandler = (update: StorkStreamUpdate) => void
+type StorkStreamMessageHandler = (message: StorkStreamServerMessage) => void
+
+/**
+ * WebSocket client for kronos `/stork/stream/{evm|stark}`.
+ * Multiplexes asset subscriptions and re-subscribes after reconnect.
+ */
+export class StorkStreamClient {
+  private readonly wsUrl: string
+  private readonly autoReconnect: boolean
+  private readonly reconnectDelayMs: number
+  private readonly pingIntervalMs: number
+  private readonly WebSocketImpl: typeof WebSocket
+
+  private ws: WebSocket | null = null
+  private connectPromise: Promise<void> | null = null
+  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
+  private pingTimer: ReturnType<typeof setInterval> | null = null
+  private closedByUser = false
+
+  private readonly assetRefCount = new Map<string, number>()
+  private readonly updateHandlers = new Set<StorkStreamUpdateHandler>()
+  private readonly messageHandlers = new Set<StorkStreamMessageHandler>()
+
+  constructor(baseUrl: string, options?: IStorkStreamClientOptions) {
+    const signatureType = options?.signatureType ?? 'evm'
+    this.wsUrl = options?.wsUrl ?? apiBaseUrlToStorkWsUrl(baseUrl, signatureType)
+    this.autoReconnect = options?.autoReconnect ?? true
+    this.reconnectDelayMs = options?.reconnectDelayMs ?? 1000
+    this.pingIntervalMs = options?.pingIntervalMs ?? 30_000
+    this.WebSocketImpl = options?.WebSocketImpl ?? WebSocket
+  }
+
+  /** Open the WebSocket and wait for socket open. */
+  connect(): Promise<void> {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      return Promise.resolve()
+    }
+    if (this.connectPromise) {
+      return this.connectPromise
+    }
+
+    this.closedByUser = false
+    this.connectPromise = new Promise((resolve, reject) => {
+      const ws = new this.WebSocketImpl(this.wsUrl)
+      this.ws = ws
+
+      const onOpen = () => {
+        this.startPing()
+        this.resubscribeAll()
+        resolve()
+      }
+
+      const onError = () => {
+        if (ws.readyState !== WebSocket.OPEN) {
+          reject(new Error(`Stork WebSocket connection failed: ${this.wsUrl}`))
+        }
+        this.emitMessage({ type: 'error', message: `Stork WebSocket error: ${this.wsUrl}` })
+      }
+
+      ws.addEventListener('open', onOpen, { once: true })
+      ws.addEventListener('error', onError, { once: true })
+      ws.addEventListener('message', (event) => {
+        this.handleMessage(String(event.data))
+      })
+      ws.addEventListener('close', () => {
+        this.stopPing()
+        this.ws = null
+        this.connectPromise = null
+        if (!this.closedByUser && this.autoReconnect) {
+          this.scheduleReconnect()
+        }
+      })
+    })
+
+    return this.connectPromise
+  }
+
+  disconnect(): void {
+    this.closedByUser = true
+    this.clearReconnect()
+    this.stopPing()
+    if (this.ws) {
+      this.ws.close()
+      this.ws = null
+    }
+    this.connectPromise = null
+  }
+
+  /**
+   * Subscribe to live Stork price updates for asset symbols (e.g. `BTCUSD`).
+   * Returns an unsubscribe function.
+   */
+  subscribe(
+    assets: string[],
+    onUpdate: StorkStreamUpdateHandler,
+  ): () => void {
+    if (assets.length === 0) {
+      throw new Error('assets is required')
+    }
+
+    const newAssets = this.addAssetRefs(assets)
+    this.updateHandlers.add(onUpdate)
+
+    void this.connect().then(() => {
+      if (newAssets.length > 0) {
+        this.send({ type: 'subscribe', data: newAssets })
+      }
+    })
+
+    return () => {
+      this.unsubscribe(assets, onUpdate)
+    }
+  }
+
+  /** Remove one handler; unsubscribe upstream assets when no handlers remain. */
+  unsubscribe(assets: string[], onUpdate?: StorkStreamUpdateHandler): void {
+    if (onUpdate) {
+      this.updateHandlers.delete(onUpdate)
+    }
+    else {
+      this.updateHandlers.clear()
+    }
+
+    if (this.updateHandlers.size === 0) {
+      const removedAssets = this.clearAssetRefs()
+      if (removedAssets.length > 0 && this.ws?.readyState === WebSocket.OPEN) {
+        this.send({ type: 'unsubscribe', data: removedAssets })
+      }
+      return
+    }
+
+    const removedAssets = this.removeAssetRefs(assets)
+    if (removedAssets.length > 0 && this.ws?.readyState === WebSocket.OPEN) {
+      this.send({ type: 'unsubscribe', data: removedAssets })
+    }
+  }
+
+  onMessage(handler: StorkStreamMessageHandler): () => void {
+    this.messageHandlers.add(handler)
+    return () => {
+      this.messageHandlers.delete(handler)
+    }
+  }
+
+  private addAssetRefs(assets: string[]): string[] {
+    const newAssets: string[] = []
+    for (const asset of assets) {
+      const count = this.assetRefCount.get(asset) ?? 0
+      this.assetRefCount.set(asset, count + 1)
+      if (count === 0) {
+        newAssets.push(asset)
+      }
+    }
+    return newAssets
+  }
+
+  private removeAssetRefs(assets: string[]): string[] {
+    const removedAssets: string[] = []
+    for (const asset of assets) {
+      const count = this.assetRefCount.get(asset)
+      if (count === undefined) {
+        continue
+      }
+      if (count <= 1) {
+        this.assetRefCount.delete(asset)
+        removedAssets.push(asset)
+      }
+      else {
+        this.assetRefCount.set(asset, count - 1)
+      }
+    }
+    return removedAssets
+  }
+
+  private clearAssetRefs(): string[] {
+    const removedAssets = [...this.assetRefCount.keys()]
+    this.assetRefCount.clear()
+    return removedAssets
+  }
+
+  private send(message: StorkStreamClientMessage): void {
+    if (this.ws?.readyState !== WebSocket.OPEN) {
+      return
+    }
+    this.ws.send(JSON.stringify(message))
+  }
+
+  private resubscribeAll(): void {
+    const assets = [...this.assetRefCount.keys()]
+    if (assets.length > 0) {
+      this.send({ type: 'subscribe', data: assets })
+    }
+  }
+
+  private scheduleReconnect(): void {
+    if (this.reconnectTimer) {
+      return
+    }
+    this.reconnectTimer = setTimeout(() => {
+      this.reconnectTimer = null
+      void this.connect().catch(() => {
+        this.scheduleReconnect()
+      })
+    }, this.reconnectDelayMs)
+  }
+
+  private clearReconnect(): void {
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer)
+      this.reconnectTimer = null
+    }
+  }
+
+  private startPing(): void {
+    if (this.pingIntervalMs <= 0) {
+      return
+    }
+    this.stopPing()
+    this.pingTimer = setInterval(() => {
+      this.send({ type: 'ping' })
+    }, this.pingIntervalMs)
+  }
+
+  private stopPing(): void {
+    if (this.pingTimer) {
+      clearInterval(this.pingTimer)
+      this.pingTimer = null
+    }
+  }
+
+  private emitMessage(message: StorkStreamServerMessage): void {
+    for (const handler of this.messageHandlers) {
+      handler(message)
+    }
+  }
+
+  private handleMessage(raw: string): void {
+    let message: StorkStreamServerMessage
+    try {
+      message = JSON.parse(raw) as StorkStreamServerMessage
+    }
+    catch {
+      this.emitMessage({ type: 'error', message: `Invalid JSON: ${raw}` })
+      return
+    }
+
+    this.emitMessage(message)
+
+    if (message.type === 'oracle_prices') {
+      const update: StorkStreamUpdate = {
+        traceId: message.trace_id,
+        prices: message.data,
+      }
+      for (const handler of this.updateHandlers) {
+        handler(update)
+      }
+    }
+  }
+}

package/tests/storkClient.test.ts

@@ -1,11 +1,13 @@
 import { describe, expect, it, vi } from 'vitest'
 
 import {
+  apiBaseUrlToStorkWsUrl,
   buildStorkUpdateDataFromSignedPrice,
   fetchStorkUpdateDataFromKronos,
   getLatestStorkPrices,
   hexToBytes,
   parseStorkUpdateDataRaw,
+  StorkStreamClient,
 } from '../src/storkClient'
 import { appendStorkUpdatesToTransaction, estimateStorkUpdateFeeMist } from '../src/storkOracle'
 import { Transaction } from '@mysten/sui/transactions'
@@ -314,3 +316,89 @@ describe('storkOracle', () => {
     expect(tx.getData().commands.length).toBe(3)
   })
 })
+
+describe('StorkStreamClient', () => {
+  class MockWebSocket {
+    static OPEN = 1
+    readyState = MockWebSocket.OPEN
+    sent: string[] = []
+    private listeners: Record<string, Array<(event: { data?: string }) => void>> = {}
+
+    constructor(public url: string) {
+      queueMicrotask(() => {
+        this.listeners.open?.forEach(cb => cb({}))
+        this.emit({ data: JSON.stringify({ type: 'connected' }) })
+      })
+    }
+
+    addEventListener(event: string, cb: (event: { data?: string }) => void, opts?: { once?: boolean }) {
+      if (!this.listeners[event]) {
+        this.listeners[event] = []
+      }
+      if (opts?.once) {
+        const onceWrapper = (ev: { data?: string }) => {
+          this.listeners[event] = this.listeners[event].filter(h => h !== onceWrapper)
+          cb(ev)
+        }
+        this.listeners[event].push(onceWrapper)
+        return
+      }
+      this.listeners[event].push(cb)
+    }
+
+    send(data: string) {
+      this.sent.push(data)
+    }
+
+    close() {
+      this.readyState = 3
+      this.listeners.close?.forEach(cb => cb({}))
+    }
+
+    emit(event: { data?: string }) {
+      this.listeners.message?.forEach(cb => cb(event))
+    }
+  }
+
+  it('derives evm and stark ws urls from kronos base url', () => {
+    expect(apiBaseUrlToStorkWsUrl('https://api.zofinance.io')).toBe('wss://api.zofinance.io/stork/stream/evm')
+    expect(apiBaseUrlToStorkWsUrl('http://localhost:8080', 'stark')).toBe('ws://localhost:8080/stork/stream/stark')
+  })
+
+  it('subscribes and forwards oracle_prices', async () => {
+    vi.stubGlobal('WebSocket', MockWebSocket)
+
+    const client = new StorkStreamClient('http://localhost:8080', {
+      WebSocketImpl: MockWebSocket as unknown as typeof WebSocket,
+      pingIntervalMs: 0,
+    })
+
+    const updates: string[] = []
+    client.subscribe(['BTCUSD'], (update) => {
+      updates.push(update.prices.BTCUSD?.asset_id ?? '')
+    })
+
+    await client.connect()
+    await new Promise<void>(resolve => setTimeout(resolve, 0))
+
+    expect(JSON.parse((client as any).ws.sent[0]).type).toBe('subscribe')
+    expect(JSON.parse((client as any).ws.sent[0]).data).toEqual(['BTCUSD'])
+
+    ;(client as any).ws.emit({
+      data: JSON.stringify({
+        type: 'oracle_prices',
+        trace_id: 'trace-1',
+        data: {
+          BTCUSD: {
+            asset_id: 'BTCUSD',
+            price: '100000000000000000000',
+            stork_signed_price: { price: '100000000000000000000' },
+          },
+        },
+      }),
+    })
+
+    expect(updates).toEqual(['BTCUSD'])
+    client.disconnect()
+  })
+})