From f88aea5bbea07da6ffa6b293d800db88ec5e0de0 Mon Sep 17 00:00:00 2001
From: ciaranbor <ciaranborourke-dev@proton.me>
Date: Fri, 6 Mar 2026 18:24:24 +0000
Subject: [PATCH] Add SetStorageConfig command

---
 src/exo/shared/types/commands.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index a6c988a5..c5041e07 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -7,6 +7,8 @@ from exo.api.types import (
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
+from exo.shared.types.memory import Memory
+from exo.shared.types.storage import StoragePolicy
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -92,6 +94,12 @@ class DeleteCustomModelCard(BaseCommand):
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
+class SetStorageConfig(BaseCommand):
+    target_node_id: NodeId
+    max_storage: Memory | None
+    storage_policy: StoragePolicy
+
+
 Command = (
     TestCommand
     | RequestEventLog
@@ -106,6 +114,7 @@ Command = (
     | SendInputChunk
     | AddCustomModelCard
     | DeleteCustomModelCard
+    | SetStorageConfig
 )