Add opaque identity and make available to thread through to remote contracts

codex-rs/thread-store/src/lib.rs

@@ -22,6 +22,7 @@ pub use in_memory::InMemoryThreadStoreCalls;
 pub use live_thread::LiveThread;
 pub use live_thread::LiveThreadInitGuard;
 pub use local::LocalThreadStore;
+pub use remote::IDENTITY_KEY_HEADER;
 pub use remote::RemoteThreadStore;
 pub use store::ThreadStore;
 pub use types::AppendThreadItemsParams;

codex-rs/thread-store/src/remote/list_threads.rs

@@ -66,6 +66,8 @@ pub(super) async fn list_threads(
 #[cfg(test)]
 mod tests {
     use std::path::PathBuf;
+    use std::sync::Arc;
+    use std::sync::Mutex;
 
     use codex_protocol::openai_models::ReasoningEffort;
     use codex_protocol::protocol::SessionSource;
@@ -79,11 +81,14 @@ mod tests {
     use super::super::proto::thread_store_server;
     use super::super::proto::thread_store_server::ThreadStoreServer;
     use super::*;
+    use crate::IDENTITY_KEY_HEADER;
     use crate::ThreadSortKey;
     use crate::ThreadStore;
 
     #[derive(Default)]
-    struct TestServer;
+    struct TestServer {
+        captured_identity_key: Option<Arc<Mutex<Option<Vec<u8>>>>>,
+    }
 
     #[tonic::async_trait]
     impl thread_store_server::ThreadStore for TestServer {
@@ -91,6 +96,16 @@ mod tests {
             &self,
             request: Request<proto::ListThreadsRequest>,
         ) -> Result<Response<proto::ListThreadsResponse>, Status> {
+            if let Some(captured_identity_key) = &self.captured_identity_key {
+                let identity_key = request
+                    .metadata()
+                    .get_bin(IDENTITY_KEY_HEADER)
+                    .and_then(|value| value.to_bytes().ok())
+                    .map(|value| value.to_vec());
+                *captured_identity_key
+                    .lock()
+                    .expect("captured identity key mutex poisoned") = identity_key;
+            }
             let request = request.into_inner();
             assert_eq!(request.page_size, 2);
             assert_eq!(request.cursor.as_deref(), Some("cursor-1"));
@@ -170,7 +185,7 @@ mod tests {
         let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
         let server = tokio::spawn(async move {
             Server::builder()
-                .add_service(ThreadStoreServer::new(TestServer))
+                .add_service(ThreadStoreServer::new(TestServer::default()))
                 .serve_with_incoming_shutdown(
                     tokio_stream::wrappers::TcpListenerStream::new(listener),
                     async {
@@ -225,6 +240,58 @@ mod tests {
         server.await.expect("join server").expect("server");
     }
 
+    #[tokio::test]
+    async fn list_threads_forwards_identity_key_as_metadata() {
+        let captured_identity_key = Arc::new(Mutex::new(None));
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind test server");
+        let addr = listener.local_addr().expect("test server addr");
+        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+        let server_identity_key = captured_identity_key.clone();
+        let server = tokio::spawn(async move {
+            Server::builder()
+                .add_service(ThreadStoreServer::new(TestServer {
+                    captured_identity_key: Some(server_identity_key),
+                }))
+                .serve_with_incoming_shutdown(
+                    tokio_stream::wrappers::TcpListenerStream::new(listener),
+                    async {
+                        let _ = shutdown_rx.await;
+                    },
+                )
+                .await
+        });
+
+        let store = RemoteThreadStore::new_with_identity_key(
+            format!("http://{addr}"),
+            Some(b"tenant-key-\x00\xff".to_vec()),
+        );
+        store
+            .list_threads(ListThreadsParams {
+                page_size: 2,
+                cursor: Some("cursor-1".to_string()),
+                sort_key: ThreadSortKey::UpdatedAt,
+                allowed_sources: vec![SessionSource::Cli],
+                model_providers: Some(vec!["openai".to_string()]),
+                archived: true,
+                search_term: Some("needle".to_string()),
+            })
+            .await
+            .expect("list threads");
+
+        assert_eq!(
+            captured_identity_key
+                .lock()
+                .expect("captured identity key mutex poisoned")
+                .as_deref(),
+            Some(&b"tenant-key-\x00\xff"[..])
+        );
+
+        let _ = shutdown_tx.send(());
+        server.await.expect("join server").expect("server");
+    }
+
     #[test]
     fn stored_thread_proto_roundtrips_through_domain_type() {
         let thread = proto::StoredThread {

codex-rs/thread-store/src/remote/mod.rs

@@ -20,10 +20,41 @@ use crate::ThreadStoreError;
 use crate::ThreadStoreResult;
 use crate::UpdateThreadMetadataParams;
 use proto::thread_store_client::ThreadStoreClient;
+use tonic::codegen::InterceptedService;
+use tonic::metadata::BinaryMetadataValue;
+use tonic::service::Interceptor;
+use tonic::transport::Channel;
+use tonic::transport::Endpoint;
 
 #[path = "proto/codex.thread_store.v1.rs"]
 mod proto;
 
+/// Metadata key used to forward the app-server's opaque identity key to remote contracts.
+pub const IDENTITY_KEY_HEADER: &str = "x-codex-app-server-identity-key-bin";
+
+#[derive(Clone, Debug)]
+struct IdentityKeyInterceptor {
+    identity_key: Option<Vec<u8>>,
+}
+
+impl Interceptor for IdentityKeyInterceptor {
+    fn call(
+        &mut self,
+        mut request: tonic::Request<()>,
+    ) -> Result<tonic::Request<()>, tonic::Status> {
+        if let Some(identity_key) = &self.identity_key {
+            request.metadata_mut().insert_bin(
+                IDENTITY_KEY_HEADER,
+                BinaryMetadataValue::from_bytes(identity_key),
+            );
+        }
+        Ok(request)
+    }
+}
+
+type RemoteThreadStoreClient =
+    ThreadStoreClient<InterceptedService<Channel, IdentityKeyInterceptor>>;
+
 /// gRPC-backed [`ThreadStore`] implementation for deployments whose durable thread data lives
 /// outside the app-server process.
 ///
@@ -33,21 +64,43 @@ mod proto;
 #[derive(Clone, Debug)]
 pub struct RemoteThreadStore {
     endpoint: String,
+    identity_key: Option<Vec<u8>>,
 }
 
 impl RemoteThreadStore {
     pub fn new(endpoint: impl Into<String>) -> Self {
         Self {
             endpoint: endpoint.into(),
+            identity_key: None,
         }
     }
 
-    async fn client(&self) -> ThreadStoreResult<ThreadStoreClient<tonic::transport::Channel>> {
-        ThreadStoreClient::connect(self.endpoint.clone())
+    pub fn new_with_identity_key(
+        endpoint: impl Into<String>,
+        identity_key: Option<Vec<u8>>,
+    ) -> Self {
+        Self {
+            endpoint: endpoint.into(),
+            identity_key,
+        }
+    }
+
+    async fn client(&self) -> ThreadStoreResult<RemoteThreadStoreClient> {
+        let channel = Endpoint::new(self.endpoint.clone())
+            .map_err(|err| ThreadStoreError::InvalidRequest {
+                message: format!("invalid remote thread store endpoint: {err}"),
+            })?
+            .connect()
             .await
             .map_err(|err| ThreadStoreError::Internal {
                 message: format!("failed to connect to remote thread store: {err}"),
-            })
+            })?;
+        Ok(ThreadStoreClient::with_interceptor(
+            channel,
+            IdentityKeyInterceptor {
+                identity_key: self.identity_key.clone(),
+            },
+        ))
     }
 }