chalk-remote-call-python 1.7.0__tar.gz → 1.7.1__tar.gz

{chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.7.0
+Version: 1.7.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.auth.v1.rs

@@ -39,6 +39,9 @@ pub enum Permission {
     InfrastructureWrite = 31,
     EnvironmentCreate = 32,
     InternalDataplaneStatusUpdate = 33,
+    InternalWorkingTokenExchange = 34,
+    BillingWrite = 35,
+    QueryOfflineRead = 36,
 }
 impl Permission {
     /// String value of the enum field names used in the ProtoBuf definition.
@@ -81,6 +84,9 @@ impl Permission {
             Permission::InfrastructureWrite => "PERMISSION_INFRASTRUCTURE_WRITE",
             Permission::EnvironmentCreate => "PERMISSION_ENVIRONMENT_CREATE",
             Permission::InternalDataplaneStatusUpdate => "PERMISSION_INTERNAL_DATAPLANE_STATUS_UPDATE",
+            Permission::InternalWorkingTokenExchange => "PERMISSION_INTERNAL_WORKING_TOKEN_EXCHANGE",
+            Permission::BillingWrite => "PERMISSION_BILLING_WRITE",
+            Permission::QueryOfflineRead => "PERMISSION_QUERY_OFFLINE_READ",
         }
     }
     /// Creates an enum from field names used in the ProtoBuf definition.
@@ -120,6 +126,9 @@ impl Permission {
             "PERMISSION_INFRASTRUCTURE_WRITE" => Some(Self::InfrastructureWrite),
             "PERMISSION_ENVIRONMENT_CREATE" => Some(Self::EnvironmentCreate),
             "PERMISSION_INTERNAL_DATAPLANE_STATUS_UPDATE" => Some(Self::InternalDataplaneStatusUpdate),
+            "PERMISSION_INTERNAL_WORKING_TOKEN_EXCHANGE" => Some(Self::InternalWorkingTokenExchange),
+            "PERMISSION_BILLING_WRITE" => Some(Self::BillingWrite),
+            "PERMISSION_QUERY_OFFLINE_READ" => Some(Self::QueryOfflineRead),
             _ => None,
         }
     }

{chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.rs

@@ -86,6 +86,38 @@ pub struct PollRemoteCallResponse {
     #[prost(message, repeated, tag="4")]
     pub errors: ::prost::alloc::vec::Vec<super::super::common::v1::ChalkError>,
 }
+/// Drop pending items from one or all per-function queues for the tenant.
+/// Already-popped (in-flight) calls continue to completion; only items still
+/// queued in Redis are removed.
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct PurgeQueueRequest {
+    #[prost(oneof="purge_queue_request::Request", tags="1, 2")]
+    pub request: ::core::option::Option<purge_queue_request::Request>,
+}
+/// Nested message and enum types in `PurgeQueueRequest`.
+pub mod purge_queue_request {
+    #[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Request {
+        /// Drop pending items from the named function's queue.
+        #[prost(string, tag="1")]
+        FunctionName(::prost::alloc::string::String),
+        /// If true, drop pending items from every per-function queue for the tenant.
+        #[prost(bool, tag="2")]
+        All(bool),
+    }
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct PurgeQueueResponse {
+    /// Number of pending items dropped per function. For a single-function
+    /// purge, contains exactly one entry (or none, if the queue was already
+    /// empty). For a tenant-wide purge, contains every queue that had at
+    /// least one pending item. The total dropped is `sum(values)`.
+    #[prost(map="string, uint64", tag="1")]
+    pub items_removed_by_function: ::std::collections::HashMap<::prost::alloc::string::String, u64>,
+}
 // ── Function Queue Meta Service ───────────────────────────────────
 
 /// Information about a single function call.
@@ -130,6 +162,32 @@ pub struct GetRecentCallsResponse {
     #[prost(message, repeated, tag="1")]
     pub calls: ::prost::alloc::vec::Vec<FunctionCallInfo>,
 }
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct GetCallResultsRequest {
+    /// Call IDs to fetch from the function queue.
+    #[prost(string, repeated, tag="1")]
+    pub call_ids: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct CallResult {
+    /// The requested call ID.
+    #[prost(string, tag="1")]
+    pub call_id: ::prost::alloc::string::String,
+    /// Present when the function queue returned a result for this call.
+    #[prost(message, optional, tag="2")]
+    pub response: ::core::option::Option<PollRemoteCallResponse>,
+    /// Populated when the queue lookup failed for this specific call.
+    #[prost(string, tag="3")]
+    pub error_message: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct GetCallResultsResponse {
+    #[prost(message, repeated, tag="1")]
+    pub results: ::prost::alloc::vec::Vec<CallResult>,
+}
 /// Request the number of calls submitted in the past hour.
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

{chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.1}/chalk-remote-call-rs/chalk-remote-call-proto/src/gen/chalk.runtime.v1.tonic.rs

@@ -443,6 +443,36 @@ pub mod async_remote_call_service_client {
                 );
             self.inner.unary(req, path, codec).await
         }
+        pub async fn purge_queue(
+            &mut self,
+            request: impl tonic::IntoRequest<super::PurgeQueueRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::PurgeQueueResponse>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/chalk.runtime.v1.AsyncRemoteCallService/PurgeQueue",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(
+                    GrpcMethod::new(
+                        "chalk.runtime.v1.AsyncRemoteCallService",
+                        "PurgeQueue",
+                    ),
+                );
+            self.inner.unary(req, path, codec).await
+        }
     }
 }
 /// Generated server implementations.
@@ -466,6 +496,13 @@ pub mod async_remote_call_service_server {
             tonic::Response<super::PollRemoteCallResponse>,
             tonic::Status,
         >;
+        async fn purge_queue(
+            &self,
+            request: tonic::Request<super::PurgeQueueRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::PurgeQueueResponse>,
+            tonic::Status,
+        >;
     }
     #[derive(Debug)]
     pub struct AsyncRemoteCallServiceServer<T: AsyncRemoteCallService> {
@@ -642,6 +679,52 @@ pub mod async_remote_call_service_server {
                     };
                     Box::pin(fut)
                 }
+                "/chalk.runtime.v1.AsyncRemoteCallService/PurgeQueue" => {
+                    #[allow(non_camel_case_types)]
+                    struct PurgeQueueSvc<T: AsyncRemoteCallService>(pub Arc<T>);
+                    impl<
+                        T: AsyncRemoteCallService,
+                    > tonic::server::UnaryService<super::PurgeQueueRequest>
+                    for PurgeQueueSvc<T> {
+                        type Response = super::PurgeQueueResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::PurgeQueueRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as AsyncRemoteCallService>::purge_queue(&inner, request)
+                                    .await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let max_decoding_message_size = self.max_decoding_message_size;
+                    let max_encoding_message_size = self.max_encoding_message_size;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let method = PurgeQueueSvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            )
+                            .apply_max_message_size_config(
+                                max_decoding_message_size,
+                                max_encoding_message_size,
+                            );
+                        let res = grpc.unary(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 _ => {
                     Box::pin(async move {
                         Ok(
@@ -797,6 +880,38 @@ pub mod function_queue_meta_service_client {
                 );
             self.inner.unary(req, path, codec).await
         }
+        /** Return current status and accumulated result chunks for specific calls.
+*/
+        pub async fn get_call_results(
+            &mut self,
+            request: impl tonic::IntoRequest<super::GetCallResultsRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::GetCallResultsResponse>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/chalk.runtime.v1.FunctionQueueMetaService/GetCallResults",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(
+                    GrpcMethod::new(
+                        "chalk.runtime.v1.FunctionQueueMetaService",
+                        "GetCallResults",
+                    ),
+                );
+            self.inner.unary(req, path, codec).await
+        }
         /** Return the number of calls submitted to a function in the past hour.
 */
         pub async fn get_call_count(
@@ -847,6 +962,15 @@ pub mod function_queue_meta_service_server {
             tonic::Response<super::GetRecentCallsResponse>,
             tonic::Status,
         >;
+        /** Return current status and accumulated result chunks for specific calls.
+*/
+        async fn get_call_results(
+            &self,
+            request: tonic::Request<super::GetCallResultsRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::GetCallResultsResponse>,
+            tonic::Status,
+        >;
         /** Return the number of calls submitted to a function in the past hour.
 */
         async fn get_call_count(
@@ -984,6 +1108,55 @@ pub mod function_queue_meta_service_server {
                     };
                     Box::pin(fut)
                 }
+                "/chalk.runtime.v1.FunctionQueueMetaService/GetCallResults" => {
+                    #[allow(non_camel_case_types)]
+                    struct GetCallResultsSvc<T: FunctionQueueMetaService>(pub Arc<T>);
+                    impl<
+                        T: FunctionQueueMetaService,
+                    > tonic::server::UnaryService<super::GetCallResultsRequest>
+                    for GetCallResultsSvc<T> {
+                        type Response = super::GetCallResultsResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::GetCallResultsRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as FunctionQueueMetaService>::get_call_results(
+                                        &inner,
+                                        request,
+                                    )
+                                    .await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let max_decoding_message_size = self.max_decoding_message_size;
+                    let max_encoding_message_size = self.max_encoding_message_size;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let method = GetCallResultsSvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            )
+                            .apply_max_message_size_config(
+                                max_decoding_message_size,
+                                max_encoding_message_size,
+                            );
+                        let res = grpc.unary(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 "/chalk.runtime.v1.FunctionQueueMetaService/GetCallCount" => {
                     #[allow(non_camel_case_types)]
                     struct GetCallCountSvc<T: FunctionQueueMetaService>(pub Arc<T>);

chalk_remote_call_python-1.7.1/chalk-remote-call-rs/chalk-remote-call-server/src/async_service.rs

@@ -0,0 +1,238 @@
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+
+use chalk_remote_call_proto::chalk::common::v1::ChalkError;
+use chalk_remote_call_proto::chalk::runtime::v1::async_remote_call_service_server::AsyncRemoteCallService;
+use chalk_remote_call_proto::chalk::runtime::v1::{
+    remote_call_args, EnqueueRemoteCallRequest, EnqueueRemoteCallResponse, PollRemoteCallRequest,
+    PollRemoteCallResponse, PurgeQueueRequest, PurgeQueueResponse, RemoteCallStatus,
+};
+use tonic::{Request, Response, Status};
+use tracing::{debug, error, info, instrument};
+
+use crate::python_bridge::{ChunkBuffer, PythonHandler};
+use crate::service::extract_metadata;
+
+/// In-memory execution record for one enqueued call.
+///
+/// Lives only on the sg pod that accepted the `EnqueueRemoteCall`, and is lost
+/// on pod restart. That is the accepted limitation of the base async path:
+/// arbitrary user handlers cannot be checkpointed, so an sg-pod death loses the
+/// computation regardless. A reclaiming consumer will re-enqueue (and the
+/// function re-runs). Durable reattach across *consumer* restarts is a later
+/// phase and lives outside this struct.
+struct CallState {
+    status: RemoteCallStatus,
+    /// Result chunks, written directly by the handler's `BufferEmitter` as it
+    /// runs and read by `PollRemoteCall` by cursor. Held behind its own lock
+    /// (not the registry map lock) so the hot emit path only contends per-call,
+    /// and `status`/`error` can be updated without blocking emits.
+    chunks: ChunkBuffer,
+    /// Populated only when `status == Failed`; surfaced as a `ChalkError`.
+    error: Option<String>,
+}
+
+impl CallState {
+    fn new() -> Self {
+        Self {
+            status: RemoteCallStatus::Pending,
+            chunks: Arc::new(Mutex::new(Vec::new())),
+            error: None,
+        }
+    }
+}
+
+type Registry = Arc<Mutex<HashMap<String, CallState>>>;
+
+/// Async (enqueue + poll) executor for the scaling-group hop.
+///
+/// `EnqueueRemoteCall` runs the handler in a detached background task that
+/// writes result chunks straight into the call's buffer, and returns
+/// immediately; `PollRemoteCall` reads those chunks by cursor. This decouples
+/// function execution from any single connection's lifetime, so the path is
+/// immune to gateway request/stream timeouts — every RPC here is sub-second
+/// regardless of how long the function runs.
+pub struct AsyncRemoteCallServiceImpl {
+    python_handler: Arc<PythonHandler>,
+    registry: Registry,
+    next_id: AtomicU64,
+    /// True when the server was started with coalescing enabled. In that mode
+    /// `python_handler`'s bridge expects the batched calling convention, which
+    /// the per-call dispatch does not satisfy, so async enqueue is refused
+    /// rather than silently calling the handler with the wrong shape.
+    batching_enabled: bool,
+}
+
+impl AsyncRemoteCallServiceImpl {
+    pub fn new(python_handler: Arc<PythonHandler>, batching_enabled: bool) -> Self {
+        Self {
+            python_handler,
+            registry: Arc::new(Mutex::new(HashMap::new())),
+            next_id: AtomicU64::new(1),
+            batching_enabled,
+        }
+    }
+}
+
+#[tonic::async_trait]
+impl AsyncRemoteCallService for AsyncRemoteCallServiceImpl {
+    #[instrument(skip_all)]
+    async fn enqueue_remote_call(
+        &self,
+        request: Request<EnqueueRemoteCallRequest>,
+    ) -> Result<Response<EnqueueRemoteCallResponse>, Status> {
+        if self.batching_enabled {
+            return Err(Status::unimplemented(
+                "async enqueue/poll is not supported when request coalescing is enabled",
+            ));
+        }
+
+        let metadata = extract_metadata(request.metadata());
+        let peer = request
+            .remote_addr()
+            .map(|a| a.to_string())
+            .unwrap_or_default();
+        let req = request.into_inner();
+        let function_name = req.name;
+
+        // Only inline feather bytes are supported today; the object-store
+        // variant is reserved for large payloads and not wired up here.
+        let ipc_bytes = match req.args.and_then(|a| a.args) {
+            Some(remote_call_args::Args::FeatherBytes(b)) => b.to_vec(),
+            Some(remote_call_args::Args::StorageObjectId(_)) => {
+                return Err(Status::unimplemented(
+                    "storage_object_id args are not supported by the scaling-group executor",
+                ));
+            }
+            None => Vec::new(),
+        };
+
+        // Pod-scoped id: the registry only knows about calls accepted by this
+        // pod, so a monotonic per-process counter is sufficient for uniqueness.
+        let call_id = format!("sg-{}", self.next_id.fetch_add(1, Ordering::Relaxed));
+
+        let state = CallState::new();
+        // The handler writes chunks into this buffer directly; the registry
+        // keeps the same Arc so polls can read them back.
+        let buffer = state.chunks.clone();
+        self.registry
+            .lock()
+            .expect("registry mutex poisoned")
+            .insert(call_id.clone(), state);
+
+        info!(call_id = %call_id, function_name = %function_name, "enqueued async call");
+
+        let handler = self.python_handler.clone();
+        let registry = self.registry.clone();
+        let id = call_id.clone();
+        tokio::spawn(async move {
+            set_status(&registry, &id, RemoteCallStatus::Running);
+
+            // The handler appends each chunk straight into `buffer`; completion
+            // is signalled purely by this returning, so there is no channel to
+            // drain and nothing that waits on Python object teardown.
+            let result = handler
+                .call_into_buffer(ipc_bytes, function_name, metadata, peer, buffer)
+                .await;
+
+            let mut reg = registry.lock().expect("registry mutex poisoned");
+            if let Some(s) = reg.get_mut(&id) {
+                match result {
+                    Ok(()) => {
+                        s.status = RemoteCallStatus::Completed;
+                        let n = s.chunks.lock().map(|b| b.len()).unwrap_or(0);
+                        info!(call_id = %id, chunks = n, "async call completed");
+                    }
+                    Err(e) => {
+                        s.status = RemoteCallStatus::Failed;
+                        s.error = Some(e.to_string());
+                        error!(call_id = %id, error = %e, "async call failed");
+                    }
+                }
+            }
+        });
+
+        Ok(Response::new(EnqueueRemoteCallResponse { call_id }))
+    }
+
+    #[instrument(skip_all)]
+    async fn poll_remote_call(
+        &self,
+        request: Request<PollRemoteCallRequest>,
+    ) -> Result<Response<PollRemoteCallResponse>, Status> {
+        let req = request.into_inner();
+        let cursor: usize = if req.cursor.is_empty() {
+            0
+        } else {
+            req.cursor
+                .parse()
+                .map_err(|_| Status::invalid_argument(format!("invalid cursor: {}", req.cursor)))?
+        };
+
+        // Read status/error and grab the chunk-buffer handle under the map lock,
+        // then release it before locking the buffer. `status` is set to a
+        // terminal value only after the handler returns (all chunks already
+        // appended), so seeing `Completed` guarantees the buffer is complete.
+        let (status, error, chunks) = {
+            let reg = self.registry.lock().expect("registry mutex poisoned");
+            let state = reg
+                .get(&req.call_id)
+                .ok_or_else(|| Status::not_found(format!("unknown call_id: {}", req.call_id)))?;
+            (state.status, state.error.clone(), state.chunks.clone())
+        };
+
+        let buf = chunks.lock().expect("chunk buffer poisoned");
+        // Clamp so a stale/over-large cursor returns an empty suffix rather than
+        // panicking on the slice.
+        let total = buf.len();
+        let from = cursor.min(total);
+        let results = buf[from..].to_vec();
+        drop(buf);
+
+        let errors = if status == RemoteCallStatus::Failed {
+            vec![ChalkError {
+                message: error.unwrap_or_else(|| "remote call failed".to_string()),
+                ..Default::default()
+            }]
+        } else {
+            Vec::new()
+        };
+
+        debug!(
+            call_id = %req.call_id,
+            status = ?status,
+            from,
+            returned = results.len(),
+            "poll"
+        );
+
+        Ok(Response::new(PollRemoteCallResponse {
+            status: status as i32,
+            results,
+            cursor: total.to_string(),
+            errors,
+        }))
+    }
+
+    async fn purge_queue(
+        &self,
+        _request: Request<PurgeQueueRequest>,
+    ) -> Result<Response<PurgeQueueResponse>, Status> {
+        // PurgeQueue operates on the tenant's Redis-backed function queues,
+        // which live in the function-queue server — not on an individual
+        // scaling-group executor pod.
+        Err(Status::unimplemented(
+            "PurgeQueue is a function-queue-server operation; not supported on the scaling-group executor",
+        ))
+    }
+}
+
+/// Set the status of an in-flight call, if it still exists in the registry.
+fn set_status(registry: &Registry, call_id: &str, status: RemoteCallStatus) {
+    if let Ok(mut reg) = registry.lock() {
+        if let Some(s) = reg.get_mut(call_id) {
+            s.status = status;
+        }
+    }
+}

{chalk_remote_call_python-1.7.0 → chalk_remote_call_python-1.7.1}/chalk-remote-call-rs/chalk-remote-call-server/src/lib.rs

@@ -1,3 +1,4 @@
+mod async_service;
 mod coalesce;
 mod python_bridge;
 mod server;