Skip to main content

tact_memory/store/remote/
client.rs

1use crate::{
2    MemoryCandidate, MemoryError, MemoryKey, MemoryLimits, MemoryRecord, MemoryScan, MemoryStore,
3    server::protocol,
4};
5use protocol::{
6    DeleteRequest, ErrorResponse, ExportRequest, ExportResponse, ListResponse, PutRequest,
7    PutResponse, ReadRequest, ReadResponse, RemoteErrorCode, RemoteRole, ScanRequest, ScanResponse,
8    SessionResponse, SyncReport, SyncRequest,
9};
10use reqwest::{Client, Response, StatusCode, Url};
11use serde::{Serialize, de::DeserializeOwned};
12use std::{collections::HashSet, fmt, sync::Arc, time::Duration};
13use thiserror::Error;
14use tokio::time::sleep;
15use zeroize::{Zeroize, Zeroizing};
16
17const ATTEMPTS: usize = 3;
18const CONNECT_TIMEOUT: Duration = Duration::from_millis(750);
19const REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
20const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
21const RETRY_BACKOFFS: [Option<Duration>; ATTEMPTS] = [
22    Some(Duration::from_millis(100)),
23    Some(Duration::from_millis(250)),
24    None,
25];
26
27/// Secret bearer token for a remote memory service.
28pub struct RemoteToken(Zeroizing<String>);
29
30impl RemoteToken {
31    /// Wraps a non-empty token in zeroizing storage.
32    pub fn new(token: String) -> Result<Self, RemoteClientError> {
33        if token.trim().is_empty() {
34            return Err(RemoteClientError::EmptyToken);
35        }
36        Ok(Self(Zeroizing::new(token)))
37    }
38
39    fn expose(&self) -> &str {
40        self.0.as_str()
41    }
42}
43
44impl fmt::Debug for RemoteToken {
45    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46        formatter.write_str("RemoteToken([REDACTED])")
47    }
48}
49
50impl Drop for RemoteToken {
51    fn drop(&mut self) {
52        self.zeroize();
53    }
54}
55
56impl Zeroize for RemoteToken {
57    fn zeroize(&mut self) {
58        self.0.zeroize();
59    }
60}
61
62/// Failure while configuring or calling a remote memory service.
63#[derive(Debug, Error)]
64pub enum RemoteClientError {
65    /// Endpoint is not an HTTP(S) URL without embedded credentials.
66    #[error("remote memory endpoint is invalid")]
67    InvalidEndpoint,
68    /// Namespace violates the wire-format contract.
69    #[error("remote memory namespace is invalid")]
70    InvalidNamespace,
71    /// Bearer token is empty.
72    #[error("remote memory token is empty")]
73    EmptyToken,
74    /// Request failed before a valid response arrived.
75    #[error("remote memory request could not reach the server")]
76    Transport,
77    /// Server rejected the bearer token.
78    #[error("remote memory server rejected authentication")]
79    Unauthorized,
80    /// Credential cannot mutate its namespace.
81    #[error("remote memory credential is read-only")]
82    ReadOnly,
83    /// Server returned a different authenticated namespace.
84    #[error("remote memory namespace does not match the credential")]
85    NamespaceMismatch,
86    /// Server speaks a different protocol version.
87    #[error("remote memory protocol is incompatible")]
88    IncompatibleProtocol,
89    /// Server rejected the request with a protocol error code.
90    #[error("remote memory server rejected the operation: {code:?}")]
91    Rejected {
92        /// Stable server failure category.
93        code: RemoteErrorCode,
94    },
95    /// Response violated bounds, ordering, or ownership constraints.
96    #[error("remote memory server returned an invalid response")]
97    InvalidResponse,
98    /// Service remained unavailable after bounded retries.
99    #[error("remote memory service is unavailable")]
100    Unavailable,
101}
102
103/// Authenticated HTTP implementation of [`MemoryStore`].
104#[derive(Clone)]
105pub struct RemoteMemoryClient {
106    inner: Arc<RemoteClientInner>,
107}
108
109impl fmt::Debug for RemoteMemoryClient {
110    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
111        formatter
112            .debug_struct("RemoteMemoryClient")
113            .field("endpoint", &self.inner.endpoint)
114            .field("namespace", &self.inner.namespace)
115            .field("token", &"[REDACTED]")
116            .finish()
117    }
118}
119
120struct RemoteClientInner {
121    endpoint: Url,
122    namespace: String,
123    token: RemoteToken,
124    client: Client,
125    role: tokio::sync::OnceCell<RemoteRole>,
126    bookmark: tokio::sync::Mutex<Option<String>>,
127}
128
129impl RemoteMemoryClient {
130    /// Creates a client bound to one validated namespace.
131    pub fn new(
132        endpoint: &str,
133        namespace: String,
134        token: RemoteToken,
135    ) -> Result<Self, RemoteClientError> {
136        let _ = rustls::crypto::ring::default_provider().install_default();
137        let mut endpoint = Url::parse(endpoint).map_err(|_| RemoteClientError::InvalidEndpoint)?;
138        if !matches!(endpoint.scheme(), "http" | "https")
139            || !endpoint.username().is_empty()
140            || endpoint.password().is_some()
141        {
142            return Err(RemoteClientError::InvalidEndpoint);
143        }
144        if !protocol::is_valid_namespace(&namespace) {
145            return Err(RemoteClientError::InvalidNamespace);
146        }
147        if !endpoint.path().ends_with('/') {
148            let mut path = endpoint.path().to_owned();
149            path.push('/');
150            endpoint.set_path(&path);
151        }
152        let client = Client::builder()
153            .connect_timeout(CONNECT_TIMEOUT)
154            .timeout(REQUEST_TIMEOUT)
155            .build()
156            .map_err(|_| RemoteClientError::InvalidEndpoint)?;
157        Ok(Self {
158            inner: Arc::new(RemoteClientInner {
159                endpoint,
160                namespace,
161                token,
162                client,
163                role: tokio::sync::OnceCell::new(),
164                bookmark: tokio::sync::Mutex::new(None),
165            }),
166        })
167    }
168
169    /// Returns the configured namespace without network I/O.
170    pub fn namespace(&self) -> &str {
171        &self.inner.namespace
172    }
173
174    /// Negotiates and caches the authenticated remote role.
175    pub async fn session(&self) -> Result<RemoteRole, RemoteClientError> {
176        self.inner
177            .role
178            .get_or_try_init(|| async {
179                let response: SessionResponse =
180                    self.get(protocol::SESSION_PATH, Replay::Safe).await?;
181                if response.protocol_version != crate::VERSION {
182                    return Err(RemoteClientError::IncompatibleProtocol);
183                }
184                if response.namespace != self.inner.namespace {
185                    return Err(RemoteClientError::NamespaceMismatch);
186                }
187                Ok(response.role)
188            })
189            .await
190            .copied()
191    }
192
193    async fn scan(
194        &self,
195        query: &str,
196        limit: usize,
197    ) -> Result<Vec<MemoryCandidate>, RemoteClientError> {
198        let limit = limit.min(MemoryLimits::PRODUCTION.scan_results);
199        if limit == 0 {
200            return Ok(Vec::new());
201        }
202        let response: ScanResponse = self
203            .post(
204                protocol::SCAN_PATH,
205                &ScanRequest {
206                    query: query.to_owned(),
207                    limit,
208                },
209                Replay::ConnectOnly,
210            )
211            .await?;
212        if response.candidates.len() > limit {
213            return Err(RemoteClientError::InvalidResponse);
214        }
215        let mut seen = HashSet::new();
216        let mut candidates = Vec::new();
217        let mut previous_score = None;
218        for candidate in response.candidates {
219            if !self.valid_candidate(&candidate) {
220                continue;
221            }
222            if previous_score.is_some_and(|score| candidate.score > score) {
223                return Err(RemoteClientError::InvalidResponse);
224            }
225            previous_score = Some(candidate.score);
226            let Some(namespace) = candidate.key.namespace.clone() else {
227                continue;
228            };
229            if !seen.insert((namespace, candidate.key.id)) {
230                return Err(RemoteClientError::InvalidResponse);
231            }
232            candidates.push(candidate);
233        }
234        Ok(candidates)
235    }
236
237    async fn read(
238        &self,
239        ids: &[i64],
240        keys: &[MemoryKey],
241    ) -> Result<Vec<MemoryRecord>, RemoteClientError> {
242        let keys = keys
243            .iter()
244            .filter(|key| Self::valid_key(key) && key.namespace.is_some())
245            .cloned()
246            .collect::<Vec<_>>();
247        let ids = ids.iter().copied().filter(|id| *id > 0).collect::<Vec<_>>();
248        if keys.is_empty() && ids.is_empty() {
249            return Ok(Vec::new());
250        }
251        let requested = keys.iter().cloned().collect::<HashSet<_>>();
252        let requested_ids = ids.iter().copied().collect::<HashSet<_>>();
253        let response: ReadResponse = self
254            .post(
255                protocol::READ_PATH,
256                &ReadRequest { ids, keys },
257                Replay::ConnectOnly,
258            )
259            .await?;
260        let requested_records = requested
261            .iter()
262            .filter_map(|key| key.namespace.clone().map(|namespace| (namespace, key.id)))
263            .chain(
264                requested_ids
265                    .iter()
266                    .map(|id| (self.namespace().to_owned(), *id)),
267            )
268            .collect::<HashSet<_>>();
269        if response.memories.len() > requested_records.len() {
270            return Err(RemoteClientError::InvalidResponse);
271        }
272
273        let mut seen = HashSet::new();
274        let mut memories = Vec::new();
275        for memory in response.memories {
276            if !(requested.contains(&memory.key)
277                || (memory.key.namespace.as_deref() == Some(self.namespace())
278                    && requested_ids.contains(&memory.key.id)))
279                || !Self::valid_record(&memory)
280            {
281                continue;
282            }
283            let logical_key = (memory.key.namespace.clone().unwrap(), memory.key.id);
284            if !seen.insert(logical_key) {
285                return Err(RemoteClientError::InvalidResponse);
286            }
287            memories.push(memory);
288        }
289        Ok(memories)
290    }
291
292    async fn list(&self) -> Result<Vec<MemoryRecord>, RemoteClientError> {
293        let response: ListResponse = self.post(protocol::LIST_PATH, &(), Replay::Safe).await?;
294        if response.memories.len() > MemoryLimits::PRODUCTION.records {
295            return Err(RemoteClientError::InvalidResponse);
296        }
297        let mut seen = HashSet::new();
298        let mut memories = Vec::new();
299        for memory in response.memories {
300            if !Self::valid_record(&memory) {
301                continue;
302            }
303            let Some(namespace) = memory.key.namespace.clone() else {
304                continue;
305            };
306            let logical_key = (namespace, memory.key.id);
307            if !seen.insert(logical_key) {
308                return Err(RemoteClientError::InvalidResponse);
309            }
310            memories.push(memory);
311        }
312        Ok(memories)
313    }
314
315    async fn put(
316        &self,
317        content: &str,
318        replacement: Option<&MemoryKey>,
319    ) -> Result<MemoryRecord, RemoteClientError> {
320        let response: PutResponse = self
321            .post(
322                protocol::PUT_PATH,
323                &PutRequest {
324                    content: content.to_owned(),
325                    replacement: replacement.cloned(),
326                },
327                Replay::ConnectOnly,
328            )
329            .await?;
330        if !Self::valid_record(&response.memory)
331            || response.memory.key.namespace.as_deref() != Some(self.namespace())
332            || response.memory.content != content
333            || match replacement {
334                Some(replacement) => {
335                    response.memory.key.id != replacement.id
336                        || replacement
337                            .version
338                            .checked_add(1)
339                            .is_none_or(|version| response.memory.key.version != version)
340                }
341                None => response.memory.key.version != 1,
342            }
343        {
344            return Err(RemoteClientError::InvalidResponse);
345        }
346        Ok(response.memory)
347    }
348
349    async fn delete(&self, key: &MemoryKey) -> Result<(), RemoteClientError> {
350        if key.namespace.as_deref() != Some(self.namespace()) {
351            return Err(RemoteClientError::NamespaceMismatch);
352        }
353        self.post::<_, serde_json::Value>(
354            protocol::DELETE_PATH,
355            &DeleteRequest { key: key.clone() },
356            Replay::Safe,
357        )
358        .await?;
359        Ok(())
360    }
361
362    async fn sync(&self, memories: &[MemoryRecord]) -> Result<SyncReport, RemoteClientError> {
363        let report: SyncReport = self
364            .post(
365                protocol::SYNC_PATH,
366                &SyncRequest {
367                    memories: memories.to_vec(),
368                },
369                Replay::Safe,
370            )
371            .await?;
372        let applied = report
373            .inserted
374            .checked_add(report.replaced)
375            .and_then(|count| count.checked_add(report.unchanged));
376        if applied != Some(memories.len()) || report.deleted > MemoryLimits::PRODUCTION.records {
377            return Err(RemoteClientError::InvalidResponse);
378        }
379        Ok(report)
380    }
381
382    fn validate_export_page(
383        namespaces: Option<&[String]>,
384        cursor: Option<&protocol::ExportCursor>,
385        accumulated_records: usize,
386        accumulated_content_bytes: usize,
387        response: &ExportResponse,
388    ) -> Result<usize, RemoteClientError> {
389        if response.memories.len() > protocol::MAX_EXPORT_PAGE_RECORDS
390            || accumulated_records
391                .checked_add(response.memories.len())
392                .is_none_or(|count| count > MemoryLimits::PRODUCTION.records)
393        {
394            return Err(RemoteClientError::InvalidResponse);
395        }
396
397        let mut previous = cursor.cloned();
398        let mut page_content_bytes = 0usize;
399        for memory in &response.memories {
400            let Some(namespace) = memory.key.namespace.as_deref() else {
401                return Err(RemoteClientError::InvalidResponse);
402            };
403            let selected = namespaces
404                .is_none_or(|selected| selected.iter().any(|candidate| candidate == namespace));
405            let ordered = previous.as_ref().is_none_or(|previous| {
406                (namespace, memory.key.id) > (previous.namespace.as_str(), previous.id)
407            });
408            if !Self::valid_record(memory) || !selected || !ordered {
409                return Err(RemoteClientError::InvalidResponse);
410            }
411
412            page_content_bytes = page_content_bytes
413                .checked_add(memory.content.len())
414                .ok_or(RemoteClientError::InvalidResponse)?;
415            if accumulated_content_bytes
416                .checked_add(page_content_bytes)
417                .is_none_or(|bytes| bytes > MemoryLimits::PRODUCTION.total_content_bytes)
418            {
419                return Err(RemoteClientError::InvalidResponse);
420            }
421            previous = Some(protocol::ExportCursor {
422                namespace: namespace.to_owned(),
423                id: memory.key.id,
424            });
425        }
426
427        if let Some(next_cursor) = &response.next_cursor {
428            let exact_last_key = response.memories.last().is_some_and(|memory| {
429                memory.key.namespace.as_deref() == Some(next_cursor.namespace.as_str())
430                    && memory.key.id == next_cursor.id
431            });
432            if !exact_last_key {
433                return Err(RemoteClientError::InvalidResponse);
434            }
435        }
436        Ok(page_content_bytes)
437    }
438
439    fn valid_key(key: &MemoryKey) -> bool {
440        key.namespace
441            .as_deref()
442            .is_none_or(protocol::is_valid_namespace)
443            && key.id > 0
444            && key.version > 0
445    }
446
447    fn valid_candidate(&self, candidate: &MemoryCandidate) -> bool {
448        Self::valid_key(&candidate.key)
449            && candidate.key.namespace.is_some()
450            && candidate.preview.len() <= 64
451            && candidate.score.is_finite()
452            && candidate.score >= 0.0
453            && !crate::secrets::contains_likely_secret(&candidate.preview)
454    }
455
456    fn valid_record(memory: &MemoryRecord) -> bool {
457        Self::valid_key(&memory.key)
458            && memory.key.namespace.is_some()
459            && !memory.content.trim().is_empty()
460            && memory.content.len() <= MemoryLimits::PRODUCTION.content_bytes
461            && memory.created_at_ms >= 0
462            && memory.updated_at_ms >= memory.created_at_ms
463            && !crate::secrets::contains_likely_secret(&memory.content)
464    }
465
466    async fn get<Response>(&self, path: &str, replay: Replay) -> Result<Response, RemoteClientError>
467    where
468        Response: DeserializeOwned,
469    {
470        self.send(path, None::<&()>, replay).await
471    }
472
473    async fn post<Request, Response>(
474        &self,
475        path: &str,
476        body: &Request,
477        replay: Replay,
478    ) -> Result<Response, RemoteClientError>
479    where
480        Request: Serialize + ?Sized,
481        Response: DeserializeOwned,
482    {
483        self.send(path, Some(body), replay).await
484    }
485
486    async fn send<Request, Response>(
487        &self,
488        path: &str,
489        body: Option<&Request>,
490        replay: Replay,
491    ) -> Result<Response, RemoteClientError>
492    where
493        Request: Serialize + ?Sized,
494        Response: DeserializeOwned,
495    {
496        // Bookmarks are opaque, so concurrent responses cannot be merged safely. Holding this guard
497        // across the exchange makes this client and its clones one monotonic session.
498        let mut bookmark = self.inner.bookmark.lock().await;
499        let url = self
500            .inner
501            .endpoint
502            .join(path)
503            .map_err(|_| RemoteClientError::InvalidEndpoint)?;
504        for backoff in RETRY_BACKOFFS {
505            let request = match body {
506                Some(body) => self.inner.client.post(url.clone()).json(body),
507                None => self.inner.client.get(url.clone()),
508            }
509            // Reqwest owns a transient, non-zeroizing header copy for the request lifetime.
510            .bearer_auth(self.inner.token.expose())
511            .header(protocol::NAMESPACE_HEADER, &self.inner.namespace);
512            let request = match bookmark.as_deref() {
513                Some(bookmark) => request.header(protocol::BOOKMARK_HEADER, bookmark),
514                None => request,
515            };
516
517            match request.send().await {
518                Ok(response) if response.status().is_success() => {
519                    let response_bookmark = bookmark_from_headers(response.headers())?;
520                    let decoded = decode_response(response).await?;
521                    if let Some(response_bookmark) = response_bookmark {
522                        *bookmark = Some(response_bookmark);
523                    }
524                    return Ok(decoded);
525                }
526                Ok(response) if replay == Replay::Safe && retryable_status(response.status()) => {
527                    let Some(backoff) = backoff else {
528                        return Err(response_error(response).await);
529                    };
530                    sleep(response_retry_delay(response.headers(), backoff)).await;
531                }
532                Ok(response) => return Err(response_error(response).await),
533                Err(error) if error.is_connect() => {
534                    let Some(backoff) = backoff else {
535                        return Err(RemoteClientError::Transport);
536                    };
537                    sleep(backoff).await;
538                }
539                Err(_) => return Err(RemoteClientError::Transport),
540            }
541        }
542        Err(RemoteClientError::Unavailable)
543    }
544}
545
546impl MemoryStore for RemoteMemoryClient {
547    fn scan(
548        &self,
549        query: &str,
550        limit: usize,
551    ) -> impl std::future::Future<Output = Result<MemoryScan, MemoryError>> + Send {
552        async move {
553            let candidates = RemoteMemoryClient::scan(self, query, limit).await?;
554            Ok(MemoryScan {
555                abstained: candidates.is_empty(),
556                candidates,
557            })
558        }
559    }
560    fn read(
561        &self,
562        ids: &[i64],
563        keys: &[MemoryKey],
564    ) -> impl std::future::Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
565        async move { Ok(RemoteMemoryClient::read(self, ids, keys).await?) }
566    }
567    fn list(
568        &self,
569    ) -> impl std::future::Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
570        async move { Ok(RemoteMemoryClient::list(self).await?) }
571    }
572    fn put(
573        &self,
574        content: &str,
575        replacement: Option<MemoryKey>,
576    ) -> impl std::future::Future<Output = Result<MemoryRecord, MemoryError>> + Send {
577        async move {
578            if content.trim().is_empty() {
579                return Err(MemoryError::EmptyContent);
580            }
581            Ok(RemoteMemoryClient::put(self, content, replacement.as_ref()).await?)
582        }
583    }
584    fn delete(
585        &self,
586        key: MemoryKey,
587    ) -> impl std::future::Future<Output = Result<(), MemoryError>> + Send {
588        async move { Ok(RemoteMemoryClient::delete(self, &key).await?) }
589    }
590    fn sync(
591        &self,
592        memories: &[MemoryRecord],
593    ) -> impl std::future::Future<Output = Result<SyncReport, MemoryError>> + Send {
594        async move { Ok(RemoteMemoryClient::sync(self, memories).await?) }
595    }
596    fn export_page(
597        &self,
598        namespaces: Option<&[String]>,
599        cursor: Option<&protocol::ExportCursor>,
600        limit: usize,
601    ) -> impl std::future::Future<
602        Output = Result<(Vec<MemoryRecord>, Option<protocol::ExportCursor>), MemoryError>,
603    > + Send {
604        let namespaces = namespaces.map(<[String]>::to_vec);
605        let cursor = cursor.cloned();
606        async move {
607            let limit = limit.clamp(1, protocol::MAX_EXPORT_PAGE_RECORDS);
608            let response: ExportResponse = self
609                .post(
610                    protocol::EXPORT_PATH,
611                    &ExportRequest {
612                        namespaces: namespaces.clone(),
613                        cursor: cursor.clone(),
614                        limit,
615                    },
616                    Replay::Safe,
617                )
618                .await?;
619            if response.memories.len() > limit {
620                return Err(RemoteClientError::InvalidResponse.into());
621            }
622            Self::validate_export_page(namespaces.as_deref(), cursor.as_ref(), 0, 0, &response)?;
623            Ok((response.memories, response.next_cursor))
624        }
625    }
626}
627
628#[derive(Clone, Copy, Eq, PartialEq)]
629enum Replay {
630    Safe,
631    ConnectOnly,
632}
633
634fn retryable_status(status: StatusCode) -> bool {
635    matches!(
636        status,
637        StatusCode::TOO_MANY_REQUESTS
638            | StatusCode::BAD_GATEWAY
639            | StatusCode::SERVICE_UNAVAILABLE
640            | StatusCode::GATEWAY_TIMEOUT
641    )
642}
643
644fn response_retry_delay(headers: &reqwest::header::HeaderMap, fallback: Duration) -> Duration {
645    headers
646        .get(reqwest::header::RETRY_AFTER)
647        .and_then(|value| value.to_str().ok())
648        .and_then(|value| value.parse::<u64>().ok())
649        .map(|seconds| Duration::from_secs(seconds.min(2)))
650        .unwrap_or(fallback)
651}
652
653fn bookmark_from_headers(
654    headers: &reqwest::header::HeaderMap,
655) -> Result<Option<String>, RemoteClientError> {
656    let mut values = headers.get_all(protocol::BOOKMARK_HEADER).iter();
657    let Some(value) = values.next() else {
658        return Ok(None);
659    };
660    if values.next().is_some() {
661        return Err(RemoteClientError::InvalidResponse);
662    }
663
664    let bookmark = value
665        .to_str()
666        .map_err(|_| RemoteClientError::InvalidResponse)?;
667    if bookmark.is_empty() {
668        return Ok(None);
669    }
670    if !protocol::is_valid_bookmark(bookmark) {
671        return Err(RemoteClientError::InvalidResponse);
672    }
673    Ok(Some(bookmark.to_owned()))
674}
675
676async fn response_error(response: Response) -> RemoteClientError {
677    let status = response.status();
678    let code = decode_response::<ErrorResponse>(response)
679        .await
680        .ok()
681        .map(|response| response.code);
682    match (status, code) {
683        (StatusCode::UNAUTHORIZED, _) | (_, Some(RemoteErrorCode::Unauthorized)) => {
684            RemoteClientError::Unauthorized
685        }
686        (StatusCode::FORBIDDEN, Some(RemoteErrorCode::NamespaceMismatch)) => {
687            RemoteClientError::NamespaceMismatch
688        }
689        (StatusCode::FORBIDDEN, _) | (_, Some(RemoteErrorCode::Forbidden)) => {
690            RemoteClientError::ReadOnly
691        }
692        (_, Some(RemoteErrorCode::UnsupportedProtocol)) => RemoteClientError::IncompatibleProtocol,
693        (_, Some(code)) => RemoteClientError::Rejected { code },
694        (StatusCode::NOT_FOUND, None) => RemoteClientError::IncompatibleProtocol,
695        (_, None) if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS => {
696            RemoteClientError::Unavailable
697        }
698        _ => RemoteClientError::InvalidResponse,
699    }
700}
701
702async fn decode_response<Decoded>(mut response: Response) -> Result<Decoded, RemoteClientError>
703where
704    Decoded: DeserializeOwned,
705{
706    let mut bytes = Vec::new();
707    while let Some(chunk) = response
708        .chunk()
709        .await
710        .map_err(|_| RemoteClientError::InvalidResponse)?
711    {
712        if bytes.len().saturating_add(chunk.len()) > MAX_RESPONSE_BYTES {
713            return Err(RemoteClientError::InvalidResponse);
714        }
715        bytes.extend_from_slice(&chunk);
716    }
717    serde_json::from_slice(&bytes).map_err(|_| RemoteClientError::InvalidResponse)
718}
719
720#[cfg(test)]
721mod tests {
722    use super::*;
723    use protocol::ExportCursor;
724
725    fn memory(namespace: &str, id: i64, content: &str) -> MemoryRecord {
726        MemoryRecord {
727            key: MemoryKey::remote(namespace.to_owned(), id, 1),
728            content: content.to_owned(),
729            created_at_ms: 1,
730            updated_at_ms: 1,
731            last_scanned_at_ms: None,
732            scan_count: 0,
733            last_used_at_ms: None,
734            use_count: 0,
735            probation_until_ms: None,
736        }
737    }
738
739    fn response(memories: Vec<MemoryRecord>, next: Option<(&str, i64)>) -> ExportResponse {
740        ExportResponse {
741            memories,
742            next_cursor: next.map(|(namespace, id)| ExportCursor {
743                namespace: namespace.to_owned(),
744                id,
745            }),
746        }
747    }
748
749    fn invalid(result: Result<usize, RemoteClientError>) -> bool {
750        matches!(result, Err(RemoteClientError::InvalidResponse))
751    }
752
753    #[tokio::test]
754    async fn delete_rejects_a_key_without_the_authenticated_namespace() {
755        let token = RemoteToken::new("test-token".to_owned()).unwrap();
756        let client =
757            RemoteMemoryClient::new("http://127.0.0.1:1/", "alice".to_owned(), token).unwrap();
758
759        assert!(matches!(
760            client.delete(&MemoryKey::local(1, 1)).await,
761            Err(RemoteClientError::NamespaceMismatch)
762        ));
763    }
764
765    #[test]
766    fn client_failures_preserve_memory_error_semantics() {
767        assert!(MemoryError::from(RemoteClientError::Transport).is_retryable());
768        assert!(MemoryError::from(RemoteClientError::Unavailable).is_retryable());
769        assert!(
770            MemoryError::from(RemoteClientError::Rejected {
771                code: RemoteErrorCode::Unavailable,
772            })
773            .is_retryable()
774        );
775        assert!(!MemoryError::from(RemoteClientError::InvalidResponse).is_retryable());
776        assert!(matches!(
777            MemoryError::from(RemoteClientError::Rejected {
778                code: RemoteErrorCode::Conflict,
779            }),
780            MemoryError::Conflict
781        ));
782    }
783
784    #[test]
785    fn export_page_enforces_namespace_and_exact_cursor() {
786        let selected = ["alpha".to_owned()];
787
788        assert_eq!(
789            RemoteMemoryClient::validate_export_page(
790                Some(&selected),
791                None,
792                0,
793                0,
794                &response(vec![memory("alpha", 1, "one")], Some(("alpha", 1))),
795            )
796            .expect("valid page"),
797            3
798        );
799        assert!(invalid(RemoteMemoryClient::validate_export_page(
800            Some(&selected),
801            None,
802            0,
803            0,
804            &response(vec![memory("beta", 1, "one")], None),
805        )));
806        assert!(invalid(RemoteMemoryClient::validate_export_page(
807            Some(&selected),
808            None,
809            0,
810            0,
811            &response(vec![memory("alpha", 1, "one")], Some(("alpha", 2))),
812        )));
813    }
814
815    #[test]
816    fn export_page_rejects_duplicate_and_out_of_order_keys() {
817        assert!(invalid(RemoteMemoryClient::validate_export_page(
818            None,
819            None,
820            0,
821            0,
822            &response(
823                vec![memory("alpha", 1, "one"), memory("alpha", 1, "two")],
824                None,
825            ),
826        )));
827        assert!(invalid(RemoteMemoryClient::validate_export_page(
828            None,
829            Some(&ExportCursor {
830                namespace: "beta".to_owned(),
831                id: 2,
832            }),
833            2,
834            6,
835            &response(vec![memory("alpha", 1, "one")], Some(("alpha", 1))),
836        )));
837    }
838
839    #[test]
840    fn export_page_rejects_multi_cursor_cycles() {
841        let first = response(vec![memory("alpha", 1, "one")], Some(("alpha", 1)));
842        let second = response(vec![memory("beta", 1, "two")], Some(("beta", 1)));
843        let cycle = response(vec![memory("alpha", 1, "one")], Some(("alpha", 1)));
844
845        assert_eq!(
846            RemoteMemoryClient::validate_export_page(None, None, 0, 0, &first)
847                .expect("valid first page"),
848            3
849        );
850        assert_eq!(
851            RemoteMemoryClient::validate_export_page(
852                None,
853                first.next_cursor.as_ref(),
854                1,
855                3,
856                &second,
857            )
858            .expect("valid second page"),
859            3
860        );
861        assert!(invalid(RemoteMemoryClient::validate_export_page(
862            None,
863            second.next_cursor.as_ref(),
864            2,
865            6,
866            &cycle,
867        )));
868    }
869
870    #[test]
871    fn export_page_rejects_aggregate_limits_before_accumulation() {
872        let page = response(vec![memory("alpha", 1, "x")], None);
873
874        assert!(invalid(RemoteMemoryClient::validate_export_page(
875            None,
876            None,
877            MemoryLimits::PRODUCTION.records,
878            0,
879            &page,
880        )));
881        assert!(invalid(RemoteMemoryClient::validate_export_page(
882            None,
883            None,
884            0,
885            MemoryLimits::PRODUCTION.total_content_bytes,
886            &page,
887        )));
888    }
889}