Unity SDK — realtime rooms in C#

// FlowFn Streams client for Unity (2021.3+) — protocol v1.
//
// One MonoBehaviour wraps the wire protocol: connect with a signed player
// token (production) or a bare user id (open-identity prototyping), join
// rooms (optionally a channel subset), publish JSON payloads, and react to
// C# events — all raised on the main thread, so Unity APIs are safe inside
// handlers.
//
// Requires Newtonsoft JSON (Unity package `com.unity.nuget.newtonsoft-json`).
// Uses System.Net.WebSockets.ClientWebSocket — supported on desktop, mobile,
// and consoles; NOT on WebGL (browser builds need a JS-bridge socket; ask us).
//
// Quickstart:
//     var fs = gameObject.AddComponent<FlowFnStream>();
//     fs.OnMessage += (room, channel, from, payload, ts) => {
//         if (channel == "ai") ShowNpcLine(payload);
//     };
//     fs.OnConnected += _ => fs.JoinRoom("match_8f3a");
//     fs.ConnectWithToken("strm-7f3a2c1d", tokenFromYourServer);
//     ...
//     fs.Publish("match_8f3a", "chat", new { text = "gg" });
//
// Tokens are minted by YOUR server with the stream's server key:
//     POST https://stream.flowfn.com/stream/v1/tokens
//     Authorization: Bearer <server key>
//     {"stream":"strm-…","user_id":"p1","name":"Ana","rooms":["match_8f3a"]}
// Never ship the server key inside a client build.

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using UnityEngine;

public class FlowFnStream : MonoBehaviour
{
    [Tooltip("Gateway base URL — override for self-hosted/staging.")]
    public string gatewayUrl = "wss://stream.flowfn.com/stream/v1";

    [Tooltip("Reconnect automatically (with backoff) after an unexpected close.")]
    public bool autoReconnect = true;

    /// Fired after the gateway's `welcome` frame — safe to join rooms.
    public event Action<List<JObject>> OnConnected;
    /// A join acked: members is the room roster including you.
    public event Action<string, List<JObject>, string> OnJoined;
    public event Action<string, string> OnLeft;
    public event Action<string, JObject> OnPeerJoined;
    public event Action<string, JObject> OnPeerLeft;
    /// Every subscribed message. from has user_id / name / role — role is
    /// "player", "server", or "flowfn" (engine output players can't spoof).
    public event Action<string, string, JObject, JToken, long> OnMessage;
    /// Protocol error frames: stable codes (rate_limited, room_full, …).
    public event Action<string, string, string> OnStreamError;
    /// Socket closed. 4001 unauthorized · 4003 stream disabled · 4004
    /// connection limit · 4008 rate-limit kick · 4009 server restarting.
    public event Action<int, string> OnDisconnected;

    ClientWebSocket _socket;
    CancellationTokenSource _cts;
    readonly ConcurrentQueue<string> _inbox = new ConcurrentQueue<string>();
    readonly Dictionary<string, List<string>> _rejoin = new Dictionary<string, List<string>>();
    string _url = "";
    bool _wasConnected;
    float _reconnectWait = 1f;
    float _reconnectTimer = -1f;
    int _refSeq;
    volatile bool _closedByUs;

    // ── Connect ────────────────────────────────────────────────────────

    /// Production: signed player token minted by YOUR server.
    public void ConnectWithToken(string streamCode, string token) =>
        Open($"{gatewayUrl}?stream={Uri.EscapeDataString(streamCode)}&token={Uri.EscapeDataString(token)}");

    /// Prototyping only: requires the stream's "open identity" toggle.
    public void ConnectOpen(string streamCode, string userId, string displayName = "")
    {
        var url = $"{gatewayUrl}?stream={Uri.EscapeDataString(streamCode)}&user_id={Uri.EscapeDataString(userId)}";
        if (!string.IsNullOrEmpty(displayName)) url += $"&name={Uri.EscapeDataString(displayName)}";
        Open(url);
    }

    /// Trusted game-server process only — NEVER ship the key in a client build.
    public void ConnectWithServerKey(string streamCode, string serverKey) =>
        Open($"{gatewayUrl}?stream={Uri.EscapeDataString(streamCode)}&key={Uri.EscapeDataString(serverKey)}");

    // ── Ops ────────────────────────────────────────────────────────────

    /// Join a room. Pass channels to subscribe to a subset (saves bandwidth);
    /// null/empty = every channel. Rooms are created by joining them.
    public void JoinRoom(string room, List<string> channels = null, string reference = "")
    {
        var frame = new JObject { ["op"] = "join", ["room"] = room };
        if (channels != null && channels.Count > 0) frame["channels"] = new JArray(channels);
        _rejoin[room] = channels != null && channels.Count > 0 ? new List<string>(channels) : null;
        Send(frame, reference);
    }

    public void LeaveRoom(string room, string reference = "")
    {
        _rejoin.Remove(room);
        Send(new JObject { ["op"] = "leave", ["room"] = room }, reference);
    }

    /// Publish a JSON-serializable payload on a channel (16 KB frame cap).
    public void Publish(string room, string channel, object payload, string reference = "")
    {
        var frame = new JObject
        {
            ["op"] = "publish",
            ["room"] = room,
            ["channel"] = channel,
            ["payload"] = payload == null ? JValue.CreateNull() : JToken.FromObject(payload),
        };
        Send(frame, reference);
    }

    public void Close()
    {
        autoReconnect = false;
        _rejoin.Clear();
        _closedByUs = true;
        _cts?.Cancel();
        try { _socket?.Abort(); } catch { /* already gone */ }
    }

    public bool IsConnected => _socket != null && _socket.State == WebSocketState.Open && _wasConnected;

    // ── Internals ──────────────────────────────────────────────────────

    void Open(string url)
    {
        // Tear down any previous socket WITHOUT touching autoReconnect or the
        // re-join map — reconnects come through here and must keep both.
        _closedByUs = true;
        _cts?.Cancel();
        try { _socket?.Abort(); } catch { /* already gone */ }
        _closedByUs = false;
        _url = url;
        _wasConnected = false;
        _cts = new CancellationTokenSource();
        _socket = new ClientWebSocket();
        _ = RunSocket(_cts.Token);
    }

    async Task RunSocket(CancellationToken ct)
    {
        int closeCode = 1006;
        string closeReason = "abnormal closure";
        try
        {
            await _socket.ConnectAsync(new Uri(_url), ct);
            var buffer = new byte[32 * 1024];
            var text = new StringBuilder();
            while (_socket.State == WebSocketState.Open && !ct.IsCancellationRequested)
            {
                var result = await _socket.ReceiveAsync(new ArraySegment<byte>(buffer), ct);
                if (result.MessageType == WebSocketMessageType.Close)
                {
                    closeCode = (int)(_socket.CloseStatus ?? WebSocketCloseStatus.Empty);
                    closeReason = _socket.CloseStatusDescription ?? "";
                    break;
                }
                text.Append(Encoding.UTF8.GetString(buffer, 0, result.Count));
                if (result.EndOfMessage)
                {
                    _inbox.Enqueue(text.ToString());
                    text.Clear();
                }
            }
        }
        catch (Exception e) when (!(e is OperationCanceledException))
        {
            closeReason = e.Message;
        }
        if (!_closedByUs) _inbox.Enqueue($"#close {closeCode} {closeReason}");
    }

    void Send(JObject frame, string reference)
    {
        if (_socket == null || _socket.State != WebSocketState.Open)
        {
            Debug.LogWarning($"FlowFnStream: not connected — dropped op {frame["op"]}");
            return;
        }
        frame["ref"] = string.IsNullOrEmpty(reference) ? $"r{++_refSeq}" : reference;
        var bytes = Encoding.UTF8.GetBytes(frame.ToString(Formatting.None));
        _ = _socket.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Text, true, _cts.Token);
    }

    void Update()
    {
        while (_inbox.TryDequeue(out var raw))
        {
            if (raw.StartsWith("#close ")) { HandleClose(raw); continue; }
            HandleFrame(raw);
        }
        if (_reconnectTimer > 0f)
        {
            _reconnectTimer -= Time.unscaledDeltaTime;
            if (_reconnectTimer <= 0f && _url != "") Open(_url);
        }
    }

    void HandleClose(string packed)
    {
        var parts = packed.Split(new[] { ' ' }, 3); // "#close", code, reason…
        int code = parts.Length > 1 && int.TryParse(parts[1], out var c) ? c : 1006;
        string reason = parts.Length > 2 ? parts[2] : "";
        OnDisconnected?.Invoke(code, reason);
        // 4001/4003 are deliberate refusals — reconnecting would just loop.
        if (autoReconnect && code != 4001 && code != 4003)
        {
            _reconnectTimer = _reconnectWait;
            _reconnectWait = Mathf.Min(_reconnectWait * 2f, 30f);
        }
    }

    void HandleFrame(string raw)
    {
        JObject frame;
        try { frame = JObject.Parse(raw); } catch { return; }
        switch ((string)frame["op"] ?? "")
        {
            case "welcome":
                _wasConnected = true;
                _reconnectWait = 1f;
                OnConnected?.Invoke(ToObjectList(frame["channels"]));
                foreach (var kv in new Dictionary<string, List<string>>(_rejoin))
                {
                    var f = new JObject { ["op"] = "join", ["room"] = kv.Key };
                    if (kv.Value != null) f["channels"] = new JArray(kv.Value);
                    Send(f, "");
                }
                break;
            case "ack":
                var room = (string)frame["room"] ?? "";
                if (frame["members"] != null) OnJoined?.Invoke(room, ToObjectList(frame["members"]), (string)frame["ref"] ?? "");
                else if (room != "") OnLeft?.Invoke(room, (string)frame["ref"] ?? "");
                break;
            case "message":
                OnMessage?.Invoke(
                    (string)frame["room"] ?? "",
                    (string)frame["channel"] ?? "",
                    frame["from"] as JObject ?? new JObject(),
                    frame["payload"],
                    (long?)frame["ts"] ?? 0L);
                break;
            case "peer_joined":
                OnPeerJoined?.Invoke((string)frame["room"] ?? "", frame["user"] as JObject ?? new JObject());
                break;
            case "peer_left":
                OnPeerLeft?.Invoke((string)frame["room"] ?? "", frame["user"] as JObject ?? new JObject());
                break;
            case "error":
                OnStreamError?.Invoke((string)frame["code"] ?? "", (string)frame["message"] ?? "", (string)frame["ref"] ?? "");
                break;
        }
    }

    static List<JObject> ToObjectList(JToken token)
    {
        var list = new List<JObject>();
        if (token is JArray arr)
            foreach (var item in arr)
                if (item is JObject o) list.Add(o);
        return list;
    }

    void OnDestroy() => Close();
}