data_stream.h

/*
 * Copyright 2025 LiveKit
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an “AS IS” BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <vector>
 
namespace livekit {
 
class LocalParticipant;
 
// Chunk size for data streams (matches Python STREAM_CHUNK_SIZE).
// Chosen to balance throughput and latency, and to work well with WebRTC data
// channels.
constexpr std::size_t kStreamChunkSize = 15'000; // 15 KB
 
struct BaseStreamInfo {
  std::string stream_id;
 
  std::string mime_type;
 
  std::string topic;
 
  std::int64_t timestamp = 0;
 
  std::optional<std::size_t> size;
 
  std::map<std::string, std::string> attributes;
};
 
struct TextStreamInfo : BaseStreamInfo {
  std::vector<std::string> attachments;
};
 
struct ByteStreamInfo : BaseStreamInfo {
  std::string name;
};
 
// Readers
//   - TextStreamReader: yields UTF-8 text chunks (std::string)
//   - ByteStreamReader: yields raw bytes (std::vector<uint8_t>)
 
class TextStreamReader {
public:
  explicit TextStreamReader(const TextStreamInfo &info);
 
  TextStreamReader(const TextStreamReader &) = delete;
  TextStreamReader &operator=(const TextStreamReader &) = delete;
 
  bool readNext(std::string &out);
 
  std::string readAll();
 
  const TextStreamInfo &info() const noexcept { return info_; }
 
private:
  friend class Room;
 
  void onChunkUpdate(const std::string &text);
 
  void onStreamClose(const std::map<std::string, std::string> &trailer_attrs);
 
  TextStreamInfo info_;
 
  // Queue of text chunks; empty string with closed_==true means EOS.
  std::deque<std::string> queue_;
  bool closed_ = false;
 
  std::mutex mutex_;
  std::condition_variable cv_;
};
 
class ByteStreamReader {
public:
  explicit ByteStreamReader(const ByteStreamInfo &info);
 
  ByteStreamReader(const ByteStreamReader &) = delete;
  ByteStreamReader &operator=(const ByteStreamReader &) = delete;
 
  bool readNext(std::vector<std::uint8_t> &out);
 
  const ByteStreamInfo &info() const noexcept { return info_; }
 
private:
  friend class Room;
 
  void onChunkUpdate(const std::vector<std::uint8_t> &bytes);
 
  void onStreamClose(const std::map<std::string, std::string> &trailer_attrs);
 
  ByteStreamInfo info_;
 
  std::deque<std::vector<std::uint8_t>> queue_;
  bool closed_ = false;
 
  std::mutex mutex_;
  std::condition_variable cv_;
};
 
class BaseStreamWriter {
public:
  virtual ~BaseStreamWriter() = default;
 
  const std::string &streamId() const noexcept { return stream_id_; }
 
  const std::string &topic() const noexcept { return topic_; }
 
  const std::string &mimeType() const noexcept { return mime_type_; }
 
  std::int64_t timestampMs() const noexcept { return timestamp_ms_; }
 
  bool isClosed() const noexcept { return closed_; }
 
  void close(const std::string &reason = "",
             const std::map<std::string, std::string> &attributes = {});
 
protected:
  BaseStreamWriter(LocalParticipant &local_participant,
                   const std::string &topic = "",
                   const std::map<std::string, std::string> &attributes = {},
                   const std::string &stream_id = "",
                   std::optional<std::size_t> total_size = std::nullopt,
                   const std::string &mime_type = "",
                   const std::vector<std::string> &destination_identities = {},
                   const std::string &sender_identity = "");
 
  enum class StreamKind { kUnknown, kText, kByte };
 
  LocalParticipant &local_participant_;
 
  // Public-ish metadata (mirrors BaseStreamInfo, but kept simple here)
  std::string stream_id_;
  std::string mime_type_;
  std::string topic_;
  std::int64_t timestamp_ms_ = 0;
  std::optional<std::size_t> total_size_;
  std::map<std::string, std::string> attributes_;
  std::vector<std::string> destination_identities_;
  std::string sender_identity_;
 
  bool closed_ = false;
  bool header_sent_ = false;
  std::uint64_t next_chunk_index_ = 0;
  StreamKind kind_ = StreamKind::kUnknown;
  std::string reply_to_id_;
  std::string byte_name_; // Used by ByteStreamWriter
 
  void ensureHeaderSent();
 
  void sendChunk(const std::vector<std::uint8_t> &content);
 
  void sendTrailer(const std::string &reason,
                   const std::map<std::string, std::string> &attributes);
};
 
class TextStreamWriter : public BaseStreamWriter {
public:
  TextStreamWriter(LocalParticipant &local_participant,
                   const std::string &topic = "",
                   const std::map<std::string, std::string> &attributes = {},
                   const std::string &stream_id = "",
                   std::optional<std::size_t> total_size = std::nullopt,
                   const std::string &reply_to_id = "",
                   const std::vector<std::string> &destination_identities = {},
                   const std::string &sender_identity = "");
 
  void write(const std::string &text);
 
  const TextStreamInfo &info() const noexcept { return info_; }
 
private:
  TextStreamInfo info_;
  std::mutex write_mutex_;
};
 
class ByteStreamWriter : public BaseStreamWriter {
public:
  ByteStreamWriter(LocalParticipant &local_participant, const std::string &name,
                   const std::string &topic = "",
                   const std::map<std::string, std::string> &attributes = {},
                   const std::string &stream_id = "",
                   std::optional<std::size_t> total_size = std::nullopt,
                   const std::string &mime_type = "application/octet-stream",
                   const std::vector<std::string> &destination_identities = {},
                   const std::string &sender_identity = "");
 
  void write(const std::vector<std::uint8_t> &data);
 
  const ByteStreamInfo &info() const noexcept { return info_; }
 
private:
  ByteStreamInfo info_;
  std::mutex write_mutex_;
};
 
/* Callback invoked when a new incoming text stream is opened.
 *
 * The TextStreamReader is provided as a shared_ptr to ensure it remains
 * alive for the duration of asynchronous reads (for example, when the
 * user spawns a background thread to consume the stream).
 */
using TextStreamHandler =
    std::function<void(std::shared_ptr<TextStreamReader>,
                       const std::string &participant_identity)>;
 
/* Callback invoked when a new incoming byte stream is opened.
 *
 * The ByteStreamReader is provided as a shared_ptr to ensure it remains
 * alive for the duration of asynchronous reads (for example, file writes
 * or background processing).
 */
using ByteStreamHandler =
    std::function<void(std::shared_ptr<ByteStreamReader>,
                       const std::string &participant_identity)>;
 
} // namespace livekit