use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use toolkit_utils::SecretString;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct TenantId(String);
impl TenantId {
#[must_use]
pub fn new(s: impl Into<String>) -> Self {
let s = s.into();
assert!(!s.is_empty(), "TenantId must not be empty");
Self(s)
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn into_inner(self) -> String {
self.0
}
}
impl From<String> for TenantId {
fn from(s: String) -> Self {
Self::new(s)
}
}
impl From<&str> for TenantId {
fn from(s: &str) -> Self {
Self::new(s)
}
}
impl AsRef<str> for TenantId {
fn as_ref(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for TenantId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct UserId(String);
impl UserId {
#[must_use]
pub fn new(s: impl Into<String>) -> Self {
let s = s.into();
assert!(!s.is_empty(), "UserId must not be empty");
Self(s)
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn into_inner(self) -> String {
self.0
}
}
impl From<String> for UserId {
fn from(s: String) -> Self {
Self::new(s)
}
}
impl From<&str> for UserId {
fn from(s: &str) -> Self {
Self::new(s)
}
}
impl AsRef<str> for UserId {
fn as_ref(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for UserId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Clone, Serialize, Deserialize)]
pub struct Session {
pub session_id: Uuid,
pub tenant_id: TenantId,
pub user_id: UserId,
pub client_id: Option<String>,
pub session_type_id: Option<Uuid>,
pub enabled_capabilities: Option<serde_json::Value>,
pub metadata: Option<serde_json::Value>,
pub lifecycle_state: LifecycleState,
#[serde(
default,
serialize_with = "toolkit_utils::secret_string::serialize_option_exposed"
)]
pub share_token: Option<SecretString>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
}
impl std::fmt::Debug for Session {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let share_token_redacted: Option<&'static str> =
self.share_token.as_ref().map(|_| "<redacted>");
f.debug_struct("Session")
.field("session_id", &self.session_id)
.field("tenant_id", &self.tenant_id)
.field("user_id", &self.user_id)
.field("client_id", &self.client_id)
.field("session_type_id", &self.session_type_id)
.field("enabled_capabilities", &self.enabled_capabilities)
.field("metadata", &self.metadata)
.field("lifecycle_state", &self.lifecycle_state)
.field("share_token", &share_token_redacted)
.field("created_at", &self.created_at)
.field("updated_at", &self.updated_at)
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LifecycleState {
Active,
Archived,
SoftDeleted,
HardDeleted,
}
impl LifecycleState {
#[must_use]
pub fn as_str(&self) -> &'static str {
match self {
Self::Active => "active",
Self::Archived => "archived",
Self::SoftDeleted => "soft_deleted",
Self::HardDeleted => "hard_deleted",
}
}
#[must_use]
pub fn from_str_value(s: &str) -> Option<Self> {
match s {
"active" => Some(Self::Active),
"archived" => Some(Self::Archived),
"soft_deleted" => Some(Self::SoftDeleted),
"hard_deleted" => Some(Self::HardDeleted),
_ => None,
}
}
#[must_use]
pub fn can_transition_to(&self, target: &Self) -> bool {
matches!(
(self, target),
(Self::Active, Self::Archived)
| (Self::Active, Self::SoftDeleted)
| (Self::Active, Self::HardDeleted)
| (Self::Archived, Self::Active)
| (Self::Archived, Self::SoftDeleted)
| (Self::Archived, Self::HardDeleted)
| (Self::SoftDeleted, Self::Active)
| (Self::SoftDeleted, Self::HardDeleted)
)
}
}
impl std::fmt::Display for LifecycleState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[cfg(test)]
mod tests {
use super::LifecycleState;
#[test]
fn every_documented_valid_transition_is_allowed() {
let valid_edges = [
(LifecycleState::Active, LifecycleState::Archived),
(LifecycleState::Active, LifecycleState::SoftDeleted),
(LifecycleState::Active, LifecycleState::HardDeleted),
(LifecycleState::Archived, LifecycleState::Active),
(LifecycleState::Archived, LifecycleState::SoftDeleted),
(LifecycleState::Archived, LifecycleState::HardDeleted),
(LifecycleState::SoftDeleted, LifecycleState::Active),
(LifecycleState::SoftDeleted, LifecycleState::HardDeleted),
];
for (from, to) in valid_edges {
assert!(
from.can_transition_to(&to),
"{from:?} -> {to:?} should be a valid transition"
);
}
}
#[test]
fn representative_invalid_transitions_are_rejected() {
assert!(!LifecycleState::HardDeleted.can_transition_to(&LifecycleState::Active));
assert!(!LifecycleState::HardDeleted.can_transition_to(&LifecycleState::Archived));
assert!(!LifecycleState::Active.can_transition_to(&LifecycleState::Active));
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionType {
pub session_type_id: Uuid,
pub name: String,
pub plugin_instance_id: Option<String>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub message_id: Uuid,
pub session_id: Uuid,
#[serde(default)]
pub tenant_id: Option<TenantId>,
#[serde(default)]
pub user_id: Option<UserId>,
pub parent_message_id: Option<Uuid>,
#[serde(default)]
pub variant_index: u32,
#[serde(default)]
pub is_active: bool,
pub role: MessageRole,
#[serde(default)]
pub parts: Vec<MessagePart>,
#[serde(default)]
pub file_ids: Vec<Uuid>,
pub metadata: Option<serde_json::Value>,
#[serde(default = "default_true")]
pub is_complete: bool,
#[serde(default)]
pub is_hidden_from_user: bool,
#[serde(default)]
pub is_hidden_from_backend: bool,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
}
fn default_true() -> bool {
true
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum MessageRole {
User,
Assistant,
System,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum MessagePartType {
Text,
Code,
Images,
Videos,
Links,
Statuses,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessagePartInput {
#[serde(rename = "type")]
pub part_type: MessagePartType,
pub content: serde_json::Value,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub file_citations: Vec<FileCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub link_citations: Vec<LinkCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub references: Vec<LinkReference>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessagePart {
pub id: Uuid,
pub message_id: Uuid,
#[serde(rename = "type")]
pub part_type: MessagePartType,
pub content: serde_json::Value,
pub number: u32,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub file_citations: Vec<FileCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub link_citations: Vec<LinkCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub references: Vec<LinkReference>,
}
impl MessagePart {
#[must_use]
pub fn text(id: Uuid, message_id: Uuid, number: u32, text: impl Into<String>) -> Self {
Self {
id,
message_id,
part_type: MessagePartType::Text,
content: serde_json::json!({ "text": text.into() }),
number,
file_citations: Vec::new(),
link_citations: Vec::new(),
references: Vec::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TextPositionAnchor {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_start: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_end: Option<i64>,
#[serde(default)]
pub quote: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_preview: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileCitation {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub citation_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub index: Option<u32>,
pub document_id: String,
pub document_name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub document_title: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub source: Option<String>,
#[serde(default)]
pub quote: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_start: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_end: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_preview: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_content: Option<String>,
#[serde(default = "default_chunk_type")]
pub chunk_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub page: Option<i32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timestamp: Option<f64>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub highlights: Vec<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reference_type: Option<String>,
#[serde(default)]
pub text_positions: Vec<u32>,
#[serde(default)]
pub text_position_anchors: Vec<TextPositionAnchor>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub meta: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LinkCitation {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub citation_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub index: Option<u32>,
pub url: String,
pub title: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub preview_text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub favicon_url: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub quote: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_start: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub char_end: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reference_type: Option<String>,
#[serde(default)]
pub text_positions: Vec<u32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LinkReference {
#[serde(default)]
pub title: String,
pub url: String,
#[serde(default)]
pub preview_text: String,
#[serde(default)]
pub position: Vec<u32>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub preview_highlights: Vec<serde_json::Value>,
#[serde(default = "default_ref_type")]
pub ref_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ref_meta: Option<serde_json::Value>,
#[serde(default)]
pub idx: u32,
}
fn default_chunk_type() -> String {
String::from("text")
}
fn default_ref_type() -> String {
String::from("url")
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Capability {
pub name: String,
pub value: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CapabilityValue {
pub name: String,
pub value: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VariantInfo {
pub message_id: Uuid,
pub variant_index: u32,
pub total_variants: u32,
pub is_active: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MemoryStrategy {
Full,
SlidingWindow {
window_size: u32,
},
Summarized {
recent_messages_to_keep: u32,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RetentionPolicy {
None,
AgeBased {
max_age_days: u32,
},
CountBased {
max_message_count: u32,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamingEvent {
Start(StreamingStartEvent),
Chunk(StreamingChunkEvent),
Status(StreamingStatusEvent),
Part(StreamingPartEvent),
Citation(StreamingCitationEvent),
State(StreamingStateEvent),
SessionMeta(StreamingSessionMetaEvent),
Tool(StreamingToolEvent),
Complete(StreamingCompleteEvent),
Error(StreamingErrorEvent),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingStartEvent {
pub message_id: Uuid,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingChunkEvent {
pub message_id: Uuid,
pub chunk: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingCompleteEvent {
pub message_id: Uuid,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub file_citations: Vec<FileCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub link_citations: Vec<LinkCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub references: Vec<LinkReference>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingErrorEvent {
pub message_id: Uuid,
pub error: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingStatusEvent {
pub message_id: Uuid,
pub code: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub detail: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingPartEvent {
pub message_id: Uuid,
pub part: MessagePartInput,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingCitationEvent {
pub message_id: Uuid,
#[serde(default)]
pub part_number: i32,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub file_citations: Vec<FileCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub link_citations: Vec<LinkCitation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub references: Vec<LinkReference>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingStateEvent {
pub message_id: Uuid,
pub state: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingSessionMetaEvent {
pub message_id: Uuid,
pub patch: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct StreamingToolEvent {
pub message_id: Uuid,
pub tool: String,
pub payload: serde_json::Value,
}
#[cfg(test)]
mod streaming_event_wire_format_tests {
use super::{
StreamingChunkEvent, StreamingCompleteEvent, StreamingErrorEvent, StreamingEvent,
StreamingStartEvent,
};
use uuid::Uuid;
fn fixed_id() -> Uuid {
Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap()
}
#[test]
fn start_event_serializes_with_snake_case() {
let json = serde_json::to_value(StreamingEvent::Start(StreamingStartEvent {
message_id: fixed_id(),
}))
.unwrap();
assert_eq!(
json,
serde_json::json!({
"type": "start",
"message_id": "00000000-0000-0000-0000-000000000001",
})
);
}
#[test]
fn chunk_event_serializes_with_snake_case() {
let json = serde_json::to_value(StreamingEvent::Chunk(StreamingChunkEvent {
message_id: fixed_id(),
chunk: "hello".into(),
}))
.unwrap();
assert_eq!(
json,
serde_json::json!({
"type": "chunk",
"message_id": "00000000-0000-0000-0000-000000000001",
"chunk": "hello",
})
);
}
#[test]
fn complete_event_serializes_with_snake_case() {
let json = serde_json::to_value(StreamingEvent::Complete(StreamingCompleteEvent {
message_id: fixed_id(),
metadata: Some(serde_json::json!({ "usage": { "input_units": 1 } })),
file_citations: vec![],
link_citations: vec![],
references: vec![],
}))
.unwrap();
assert_eq!(
json,
serde_json::json!({
"type": "complete",
"message_id": "00000000-0000-0000-0000-000000000001",
"metadata": { "usage": { "input_units": 1 } },
})
);
}
#[test]
fn complete_event_omits_metadata_when_none() {
let json = serde_json::to_value(StreamingEvent::Complete(StreamingCompleteEvent {
message_id: fixed_id(),
metadata: None,
file_citations: vec![],
link_citations: vec![],
references: vec![],
}))
.unwrap();
assert_eq!(
json,
serde_json::json!({
"type": "complete",
"message_id": "00000000-0000-0000-0000-000000000001",
})
);
}
#[test]
fn error_event_serializes_with_snake_case() {
let json = serde_json::to_value(StreamingEvent::Error(StreamingErrorEvent {
message_id: fixed_id(),
error: "upstream timeout".into(),
}))
.unwrap();
assert_eq!(
json,
serde_json::json!({
"type": "error",
"message_id": "00000000-0000-0000-0000-000000000001",
"error": "upstream timeout",
})
);
}
}
#[cfg(test)]
mod id_validation_tests {
use super::{TenantId, UserId};
#[test]
fn tenant_id_accepts_non_empty() {
assert_eq!(TenantId::new("t").as_str(), "t");
assert_eq!(TenantId::from(String::from("t")).as_str(), "t");
assert_eq!(TenantId::from("t").as_str(), "t");
}
#[test]
#[should_panic(expected = "TenantId must not be empty")]
fn tenant_id_new_rejects_empty() {
drop(TenantId::new(""));
}
#[test]
#[should_panic(expected = "TenantId must not be empty")]
fn tenant_id_from_string_rejects_empty() {
drop(TenantId::from(String::new()));
}
#[test]
#[should_panic(expected = "TenantId must not be empty")]
fn tenant_id_from_str_rejects_empty() {
drop(TenantId::from(""));
}
#[test]
fn user_id_accepts_non_empty() {
assert_eq!(UserId::new("u").as_str(), "u");
assert_eq!(UserId::from(String::from("u")).as_str(), "u");
assert_eq!(UserId::from("u").as_str(), "u");
}
#[test]
#[should_panic(expected = "UserId must not be empty")]
fn user_id_new_rejects_empty() {
drop(UserId::new(""));
}
#[test]
#[should_panic(expected = "UserId must not be empty")]
fn user_id_from_string_rejects_empty() {
drop(UserId::from(String::new()));
}
#[test]
#[should_panic(expected = "UserId must not be empty")]
fn user_id_from_str_rejects_empty() {
drop(UserId::from(""));
}
}