From 94a7264cbf623dfde297257a156fa3afaa62ce1b Mon Sep 17 00:00:00 2001
From: ciaranbor <ciaranborourke-dev@proton.me>
Date: Fri, 6 Mar 2026 18:30:05 +0000
Subject: [PATCH] Add cli arguments to set max storage and eviction policy

---
 src/exo/download/shard_downloader.py |  5 +----
 src/exo/main.py                      | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/exo/download/shard_downloader.py b/src/exo/download/shard_downloader.py
index 2addda80..259a0c41 100644
--- a/src/exo/download/shard_downloader.py
+++ b/src/exo/download/shard_downloader.py
@@ -1,6 +1,5 @@
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable
-from copy import copy
 from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
@@ -77,9 +76,7 @@ class NoopShardDownloader(ShardDownloader):
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata
     ) -> RepoDownloadProgress:
-        dp = copy(NOOP_DOWNLOAD_PROGRESS)
-        dp.shard = shard
-        return dp
+        return NOOP_DOWNLOAD_PROGRESS.model_copy(update={"shard": shard})
 
 
 NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
diff --git a/src/exo/main.py b/src/exo/main.py
index 68cd70d4..31100b7a 100644
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -317,6 +317,8 @@ class Args(CamelCaseModel):
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
     bootstrap_peers: list[str] = []
     libp2p_port: int
+    max_storage_gb: float | None = None
+    storage_policy: StoragePolicy = "manual"
 
     @classmethod
     def parse(cls) -> Self:
@@ -404,6 +406,20 @@ class Args(CamelCaseModel):
             dest="fast_synch",
             help="Force MLX FAST_SYNCH off",
         )
+        parser.add_argument(
+            "--max-storage-gb",
+            type=float,
+            dest="max_storage_gb",
+            default=None,
+            help="Maximum storage for downloaded models in GB (default: unlimited)",
+        )
+        parser.add_argument(
+            "--storage-policy",
+            choices=["manual", "auto-evict"],
+            dest="storage_policy",
+            default="manual",
+            help="Storage policy: 'manual' rejects on exceed, 'auto-evict' removes LRU models",
+        )
 
         args = parser.parse_args()
         return cls(**vars(args))  # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically