gproxy_protocol/transform/identity/
state.rs

1use super::{
2    error::IdentityError,
3    types::{DialectId, IdentityRole},
4};
5use crate::capability::{CasResult, StateEntry, StateStore, StateWrite, Version};
6use bytes::Bytes;
7use serde::{Deserialize, Serialize};
8use std::time::SystemTime;
9
10const STATE_SCHEMA: u16 = 2;
11
12/// Target semantics are persisted with an identity record. A record from a
13/// different model or target dialect cannot be restored into a new flow.
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub struct IdentityTarget {
16    pub model: String,
17    pub dialect: DialectId,
18    /// Host-owned upstream/principal origin for signature binding. It is
19    /// intentionally separate from the wire dialect.
20    pub origin: Option<String>,
21}
22
23impl IdentityTarget {
24    pub fn new(model: impl Into<String>, dialect: DialectId) -> Result<Self, IdentityError> {
25        let model = model.into();
26        if model.trim().is_empty() {
27            return Err(IdentityError::InvalidIdentity(
28                "target model is invalid".into(),
29            ));
30        }
31        Ok(Self {
32            model,
33            dialect,
34            origin: None,
35        })
36    }
37
38    pub fn with_origin(mut self, origin: impl Into<String>) -> Result<Self, IdentityError> {
39        let origin = origin.into();
40        if origin.trim().is_empty() {
41            return Err(IdentityError::InvalidIdentity(
42                "identity origin is empty".into(),
43            ));
44        }
45        self.origin = Some(origin);
46        Ok(self)
47    }
48}
49
50/// The exact declared upstream field that owns an opaque value. A matching
51/// origin/model does not permit moving bytes between these fields.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
53pub enum OpaqueField {
54    ClaudeThinkingSignature,
55    ClaudeRedactedThinkingData,
56    GeminiPartThoughtSignature,
57    ResponsesReasoningEncryptedContent,
58}
59
60impl OpaqueField {
61    pub fn dialect(self) -> DialectId {
62        match self {
63            Self::ClaudeThinkingSignature | Self::ClaudeRedactedThinkingData => DialectId::Claude,
64            Self::GeminiPartThoughtSignature => DialectId::Gemini,
65            Self::ResponsesReasoningEncryptedContent => DialectId::OpenAi,
66        }
67    }
68}
69
70/// A declared opaque signature can be restored only for its original origin,
71/// model and field. It is intentionally just a small token, never a wire DTO.
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73pub struct OpaqueSignature {
74    pub field: OpaqueField,
75    pub value: String,
76    pub origin: String,
77    pub model: String,
78}
79
80impl OpaqueSignature {
81    pub fn new(
82        field: OpaqueField,
83        value: impl Into<String>,
84        origin: impl Into<String>,
85        model: impl Into<String>,
86    ) -> Result<Self, IdentityError> {
87        let value = value.into();
88        let origin = origin.into();
89        let model = model.into();
90        if value.is_empty() || origin.trim().is_empty() || model.trim().is_empty() {
91            return Err(IdentityError::InvalidIdentity(
92                "opaque signature declaration is invalid".into(),
93            ));
94        }
95        Ok(Self {
96            field,
97            value,
98            origin,
99            model,
100        })
101    }
102}
103
104/// The only identity facts allowed to cross a request boundary.
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
106pub struct IdentityStateRecord {
107    pub schema: u16,
108    pub role: IdentityRole,
109    pub target: IdentityTarget,
110    pub original_item_id: Option<String>,
111    pub original_call_id: Option<String>,
112    pub client_item_id: Option<String>,
113    pub client_call_id: Option<String>,
114    pub tool_name: Option<String>,
115    pub response_id: Option<String>,
116    pub conversation_id: Option<String>,
117    pub opaque_signature: Option<OpaqueSignature>,
118}
119
120impl IdentityStateRecord {
121    pub fn new(role: IdentityRole, target: IdentityTarget) -> Self {
122        Self {
123            schema: STATE_SCHEMA,
124            role,
125            target,
126            original_item_id: None,
127            original_call_id: None,
128            client_item_id: None,
129            client_call_id: None,
130            tool_name: None,
131            response_id: None,
132            conversation_id: None,
133            opaque_signature: None,
134        }
135    }
136}
137
138/// Fields that may be learned after the first stream event. Existing values
139/// are immutable; a contradictory late fact is rejected.
140#[derive(Debug, Clone, Default, PartialEq, Eq)]
141pub struct LateIdentityFacts {
142    pub original_item_id: Option<String>,
143    pub original_call_id: Option<String>,
144    pub client_item_id: Option<String>,
145    pub client_call_id: Option<String>,
146    pub tool_name: Option<String>,
147    pub response_id: Option<String>,
148    pub conversation_id: Option<String>,
149    pub opaque_signature: Option<OpaqueSignature>,
150}
151
152impl LateIdentityFacts {
153    fn apply(self, record: &mut IdentityStateRecord) -> Result<(), IdentityError> {
154        merge(
155            &mut record.original_item_id,
156            self.original_item_id,
157            "original_item_id",
158        )?;
159        merge(
160            &mut record.original_call_id,
161            self.original_call_id,
162            "original_call_id",
163        )?;
164        merge(
165            &mut record.client_item_id,
166            self.client_item_id,
167            "client_item_id",
168        )?;
169        merge(
170            &mut record.client_call_id,
171            self.client_call_id,
172            "client_call_id",
173        )?;
174        merge(&mut record.tool_name, self.tool_name, "tool_name")?;
175        merge(&mut record.response_id, self.response_id, "response_id")?;
176        merge(
177            &mut record.conversation_id,
178            self.conversation_id,
179            "conversation_id",
180        )?;
181        if let Some(signature) = self.opaque_signature {
182            if let Some(existing) = &record.opaque_signature {
183                if existing != &signature {
184                    return Err(IdentityError::InvalidIdentity(
185                        "opaque signature changed after emission".into(),
186                    ));
187                }
188            } else {
189                record.opaque_signature = Some(signature);
190            }
191        }
192        Ok(())
193    }
194}
195
196fn merge(
197    slot: &mut Option<String>,
198    incoming: Option<String>,
199    field: &str,
200) -> Result<(), IdentityError> {
201    let Some(incoming) = incoming else {
202        return Ok(());
203    };
204    if let Some(existing) = slot {
205        if existing != &incoming {
206            return Err(IdentityError::InvalidIdentity(format!(
207                "{field} changed after emission"
208            )));
209        }
210    } else {
211        *slot = Some(incoming);
212    }
213    Ok(())
214}
215
216#[derive(Debug, Clone, PartialEq, Eq)]
217pub struct IdentityStateSnapshot {
218    pub record: IdentityStateRecord,
219    pub version: Version,
220    pub expires_at: Option<SystemTime>,
221}
222
223/// Adapter-neutral persistence helper over the protocol's scoped CAS store.
224pub struct IdentityStateStore<S> {
225    store: S,
226}
227
228impl<S> IdentityStateStore<S> {
229    pub fn new(store: S) -> Self {
230        Self { store }
231    }
232    pub fn inner(&self) -> &S {
233        &self.store
234    }
235}
236
237impl<S: StateStore> IdentityStateStore<S> {
238    /// `scope` must bind tenant, upstream and principal; keys must also bind
239    /// the conversation/invocation. A wire dialect is not an upstream identity.
240    pub async fn read(
241        &self,
242        scope: &S::Scope,
243        key: &str,
244    ) -> Result<IdentityStateSnapshot, IdentityError> {
245        let Some(entry) = self
246            .store
247            .get(scope, key)
248            .await
249            .map_err(IdentityError::StateStore)?
250        else {
251            return Err(IdentityError::MissingState);
252        };
253        decode_entry(entry)
254    }
255
256    /// Save a record with compare-and-exchange. `expected = None` creates a
257    /// fresh scope key; callers should perform this before exposing an alias.
258    pub async fn save(
259        &self,
260        scope: &S::Scope,
261        key: &str,
262        expected: Option<Version>,
263        record: &IdentityStateRecord,
264        expires_at: Option<SystemTime>,
265    ) -> Result<Version, IdentityError> {
266        let payload = serde_json::to_vec(record).map_err(IdentityError::StateEncoding)?;
267        let result = self
268            .store
269            .compare_exchange(
270                scope,
271                key,
272                expected,
273                Some(StateWrite {
274                    payload: Bytes::from(payload),
275                    expires_at,
276                }),
277            )
278            .await
279            .map_err(IdentityError::StateStore)?;
280        match result {
281            CasResult::Applied(Some(version)) => Ok(version),
282            CasResult::Applied(None) | CasResult::Conflict => Err(IdentityError::Conflict),
283        }
284    }
285
286    /// CAS update for source ids/signatures that arrive after an alias was
287    /// emitted. Existing aliases remain untouched in the record.
288    pub async fn update_late(
289        &self,
290        scope: &S::Scope,
291        key: &str,
292        expected: Version,
293        facts: LateIdentityFacts,
294    ) -> Result<IdentityStateSnapshot, IdentityError> {
295        let current = self.read(scope, key).await?;
296        if current.version != expected {
297            return Err(IdentityError::Conflict);
298        }
299        let mut record = current.record;
300        facts.apply(&mut record)?;
301
302        let version = self
303            .save(scope, key, Some(expected), &record, current.expires_at)
304            .await?;
305        Ok(IdentityStateSnapshot {
306            record,
307            version,
308            expires_at: current.expires_at,
309        })
310    }
311}
312
313fn decode_entry(entry: StateEntry) -> Result<IdentityStateSnapshot, IdentityError> {
314    // Expiration is enforced atomically by StateStore::get. Do not introduce
315    // a second clock (or native-only SystemTime::now) at the protocol layer.
316    let record: IdentityStateRecord =
317        serde_json::from_slice(&entry.payload).map_err(IdentityError::StateEncoding)?;
318
319    Ok(IdentityStateSnapshot {
320        record,
321        version: entry.version,
322        expires_at: entry.expires_at,
323    })
324}
gproxy_protocol/transform/identity/state.rs

gproxy_protocol/transform/identity/
state.rs
