Add PodSnapshot extension to Python client

clients/python/agentic-sandbox-client/k8s_agent_sandbox/constants.py

@@ -0,0 +1,37 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Constants for API Groups and Resources
+GATEWAY_API_GROUP = "gateway.networking.k8s.io"
+GATEWAY_API_VERSION = "v1"
+GATEWAY_PLURAL = "gateways"
+
+CLAIM_API_GROUP = "extensions.agents.x-k8s.io"
+CLAIM_API_VERSION = "v1alpha1"
+CLAIM_PLURAL_NAME = "sandboxclaims"
+
+SANDBOX_API_GROUP = "agents.x-k8s.io"
+SANDBOX_API_VERSION = "v1alpha1"
+SANDBOX_PLURAL_NAME = "sandboxes"
+
+POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name"
+
+PODSNAPSHOT_API_GROUP = "podsnapshot.gke.io"
+PODSNAPSHOT_API_VERSION = "v1alpha1"
+PODSNAPSHOT_PLURAL = "podsnapshots"
+PODSNAPSHOTMANUALTRIGGER_PLURAL = "podsnapshotmanualtriggers"
+PODSNAPSHOT_API_KIND = "PodSnapshotManualTrigger"
+
+SNAPSHOT_NAMESPACE_MANAGED = "gke-managed-pod-snapshots"
+SNAPSHOT_AGENT = "pod-snapshot-agent"

clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .podsnapshot_client import PodSnapshotSandboxClient

clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot.md

@@ -0,0 +1,81 @@
+# Agentic Sandbox Pod Snapshot Extension
+
+This directory contains the Python client extension for interacting with the Agentic Sandbox to manage Pod Snapshots. This extension allows you to trigger snapshots of a running sandbox and restore a new sandbox from the recently created snapshot.
+
+## `podsnapshot_client.py`
+
+This file defines the `PodSnapshotSandboxClient` class, which extend the base `SandboxClient` to provide snapshot capabilities.
+
+### `PodSnapshotSandboxClient`
+
+A specialized Sandbox client for interacting with the gke pod snapshot controller.
+
+### Key Features:
+
+*   **`PodSnapshotSandboxClient(template_name: str, podsnapshot_timeout: int = 180, server_port: int = 8080, ...)`**:
+    *   Initializes the client with optional podsnapshot timeout and server port.
+    *   If snapshot exists, the pod snapshot controller restores from the most recent snapshot matching the label of the `SandboxTemplate`, otherwise creates a new `Sandbox`.
+*   **`snapshot_controller_ready(self) -> bool`**:
+    *   Checks if the snapshot agent (GKE managed) is running and ready.
+*   **`snapshot(self, trigger_name: str) -> SnapshotResponse`**:
+    *   Triggers a manual snapshot of the current sandbox pod by creating a `PodSnapshotManualTrigger` resource.
+    *   The `trigger_name` is suffixed with unique hash.
+    *   Waits for the snapshot to be processed.
+    *   The pod snapshot controller creates a `PodSnapshot` resource automatically.
+    *   Returns the SnapshotResponse object(success, error_code, error_reason, trigger_name, snapshot_uid).
+*   **`is_restored_from_snapshot(self, snapshot_uid: str) -> RestoreResult`**:
+    *   Checks if the sandbox pod was restored from the specified snapshot.
+    *   Verifies restoration by checking the 'PodRestored' condition in the pod status and confirming the message contains the expected snapshot UID.
+    *   Returns RestoreResult object(success, error_code, error_reason).
+*   **`__exit__(self)`**:
+    *   Cleans up the `PodSnapshotManualTrigger` resources.
+    *   Cleans up the `SandboxClaim` resources.
+
+## `test_podsnapshot_extension.py`
+
+This file, located in the parent directory (`clients/python/agentic-sandbox-client/`), contains an integration test script for the `PodSnapshotSandboxClient` extension. It verifies the snapshot and restore functionality.
+
+### Test Phases:
+
+1.  **Phase 1: Starting Counter & Snapshotting**:
+    *   Starts a sandbox with a counter application.
+    *   Takes a snapshot (`test-snapshot-10`) after ~10 seconds.
+    *   Takes a second snapshot (`test-snapshot-20`) after another ~10 seconds.
+2.  **Phase 2: Restoring from Recent Snapshot**:
+    *   Restores a sandbox from the second snapshot.
+    *   Verifies that sandbox has been restored from the recent snapshot. 
+
+### Prerequisites
+
+1.  **Python Virtual Environment**:
+    ```bash
+    python3 -m venv .venv
+    source .venv/bin/activate
+    ```
+
+2.  **Install Dependencies**:
+    ```bash
+    pip install kubernetes
+    pip install -e clients/python/agentic-sandbox-client/
+    ```
+
+3.  **Pod Snapshot Controller**: The Pod Snapshot controller must be installed in a **GKE standard cluster** running with **gVisor**. 
+   * For detailed setup instructions, refer to the [GKE Pod Snapshots public documentation](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/pod-snapshots).
+   * Ensure a GCS bucket is configured to store the pod snapshot states and that the necessary IAM permissions are applied.
+
+4.  **CRDs**: `PodSnapshotStorageConfig`, `PodSnapshotPolicy` CRDs must be applied. `PodSnapshotPolicy` should specify the selector match labels.
+
+5.  **Sandbox Template**: A `SandboxTemplate` (e.g., `python-counter-template`) with runtime gVisor, appropriate KSA and label that matches that selector label in `PodSnapshotPolicy` must be available in the cluster.
+
+### Running Tests:
+
+To run the integration test, execute the script with the appropriate arguments:
+
+```bash
+python3 clients/python/agentic-sandbox-client/test_podsnapshot_extension.py \
+  --labels app=agent-sandbox-workload \
+  --template-name python-counter-template \
+  --namespace sandbox-test
+```
+
+Adjust the `--namespace`, `--template-name`, and `--labels` as needed for your environment.