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>
 
#include "livekit/visibility.h"
 
namespace livekit {
 
class LocalParticipant;
 
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;
  std::optional<std::string> reply_to_stream_id;
};
 
struct ByteStreamInfo : BaseStreamInfo {
  std::string name;
};
 
// Readers
//   - TextStreamReader: yields UTF-8 text chunks (std::string)
//   - ByteStreamReader: yields raw bytes (std::vector<uint8_t>)
 
class LIVEKIT_API TextStreamReader {
public:
  explicit TextStreamReader(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 LIVEKIT_API ByteStreamReader {
public:
  explicit ByteStreamReader(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 LIVEKIT_API 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, std::string topic = "",
                   std::map<std::string, std::string> attributes = {}, std::string stream_id = "",
                   std::optional<std::size_t> total_size = std::nullopt, std::string mime_type = "",
                   std::vector<std::string> destination_identities = {}, 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 LIVEKIT_API 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 LIVEKIT_API 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_;
};
 
using TextStreamHandler =
    std::function<void(std::shared_ptr<TextStreamReader>, const std::string& participant_identity)>;
 
using ByteStreamHandler =
    std::function<void(std::shared_ptr<ByteStreamReader>, const std::string& participant_identity)>;
 
} // namespace livekit