Coverage for gco_mcp / tools / storage.py: 100.00%
87 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""File and object storage MCP tools."""
3import asyncio
4import json
5from typing import Literal
7import cli_runner
8from audit import audit_logged
9from feature_flags import FLAG_LOCAL_STORAGE_SYNC, FLAG_MODEL_UPLOAD, is_enabled
10from local_data import LocalPathContract, resolve_local_path, stage_upload_path
11from server import mcp
14@mcp.tool(tags={"safe", "storage"})
15@audit_logged
16def list_storage_contents(region: str, path: str = "/") -> str:
17 """List contents of shared EFS storage.
19 Args:
20 region: AWS region.
21 path: Directory path to list (default: root).
22 """
23 args = ["files", "ls", "-r", region]
24 if path != "/":
25 args.append(path)
26 return cli_runner._run_cli(*args)
29@mcp.tool(tags={"safe", "storage"})
30@audit_logged
31def list_file_systems(region: str | None = None) -> str:
32 """List EFS and FSx file systems.
34 Args:
35 region: Specific region, or omit for all.
36 """
37 args = ["files", "list"]
38 if region:
39 args += ["-r", region]
40 return cli_runner._run_cli(*args)
43@mcp.tool(tags={"safe", "storage"})
44@audit_logged
45async def list_storage_buckets(region: str | None = None) -> str:
46 """List deployed GCO S3 buckets and their human-friendly aliases.
48 Returns user-facing buckets such as ``cluster-shared``, ``model-weights``,
49 ``regional-shared:<region>``, and the optional ``analytics-studio`` bucket.
50 Physical names are resolved from the deployment's SSM and CloudFormation
51 metadata rather than reconstructed.
53 Args:
54 region: Optionally limit regional-bucket discovery to one AWS region.
55 Global and analytics buckets are still included.
56 """
57 args = ["storage", "list"]
58 if region:
59 args += ["--region", region]
60 return await asyncio.to_thread(cli_runner._run_cli, *args)
63@mcp.tool(tags={"safe", "storage"})
64@audit_logged
65async def s3_inventory(region: str | None = None) -> str:
66 """Describe every S3 bucket the deployment creates, with its contract.
68 Broader than list_storage_buckets, which returns only the four buckets
69 addressable by ``storage sync``. This covers the always-on central
70 (``Cluster_Shared_Bucket``) and per-region (``Regional_Shared_Bucket``)
71 buckets, model weights, cost reports, the optional analytics Studio bucket,
72 and every server-access-log sink.
74 Each entry carries the owning stack and region, the bucket's purpose,
75 reserved object-key prefixes, whether job pods have read-write / read-only /
76 no access and how they discover the name (ConfigMap key or SSM path), the
77 teardown removal policy, and whether it is currently deployed. Buckets whose
78 stack is not deployed are included with ``status="not-deployed"`` so the
79 inventory is complete rather than silently partial.
81 Use this to answer "where can a job write?" — ``summary.pod_writable`` lists
82 exactly the buckets the job-pod role can write to.
84 Inventories buckets and their deployment contract; unrelated to the AWS
85 "S3 Inventory" feature, which reports the objects inside a bucket.
87 Args:
88 region: Limit regional entries to one AWS region. Global, monitoring,
89 and analytics entries are always included.
90 """
91 args = ["storage", "s3-inventory"]
92 if region:
93 args += ["--region", region]
94 return await asyncio.to_thread(cli_runner._run_cli, *args)
97# =============================================================================
98# Read-only inspection tools (async)
99# =============================================================================
102@mcp.tool(tags={"safe", "files"})
103@audit_logged
104async def files_get(region: str, fs_type: str = "efs") -> str:
105 """`gco files get` — get file system details for a region.
107 Returns the EFS (or FSx) file system's ID, lifecycle state, throughput mode,
108 encryption flags, and mount targets. To browse or fetch file contents, use
109 list_storage_contents (``gco files ls``) or ``gco files download``.
111 Args:
112 region: AWS region.
113 fs_type: File system type — "efs" (default) or "fsx".
114 """
115 return await asyncio.to_thread(cli_runner._run_cli, "files", "get", region, "-t", fs_type)
118@mcp.tool(tags={"safe", "files"})
119@audit_logged
120async def files_access_points(region: str | None = None) -> str:
121 """`gco files access-points` — list EFS access points.
123 Args:
124 region: AWS region.
125 """
126 args = ["files", "access-points"]
127 if region:
128 args += ["-r", region]
129 return await asyncio.to_thread(cli_runner._run_cli, *args)
132# =============================================================================
133# Regional bucket upload (low-risk write)
134# =============================================================================
137if is_enabled(FLAG_MODEL_UPLOAD):
139 @mcp.tool(tags={"data-upload", "storage", "local-filesystem"})
140 @audit_logged
141 async def upload_to_regional_bucket(
142 local_path: str, region: str, prefix: str = "uploads"
143 ) -> str:
144 """[gated by GCO_ENABLE_MODEL_UPLOAD] Upload local data to a regional bucket.
146 The source must resolve beneath ``GCO_STORAGE_LOCAL_ROOT``. Relative
147 paths such as ``model.bin`` and ``./datasets`` are interpreted relative
148 to that root, never relative to the MCP process working directory. The
149 CLI receives a private descriptor-backed snapshot; descendant links,
150 special files, hard links, and filesystem crossings fail closed.
152 Args:
153 local_path: Root-relative local file or directory to upload.
154 region: Target region whose regional bucket receives the objects.
155 prefix: S3 prefix for uploaded objects (default: ``uploads``).
156 """
157 try:
158 local_contract = _resolve_upload_local_path(local_path)
159 except (OSError, ValueError) as exc:
160 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"})
162 def _upload_from_staged_path() -> str:
163 with stage_upload_path(local_contract) as staged:
164 return cli_runner._run_cli(
165 "models",
166 "upload-regional",
167 staged.argument,
168 "-r",
169 region,
170 "--prefix",
171 prefix,
172 pass_fds=(staged.directory_fd,),
173 )
175 try:
176 # Run the staging context in the worker so cancellation cannot
177 # unlink its descriptor-backed snapshot while the CLI still reads.
178 return await asyncio.to_thread(_upload_from_staged_path)
179 except (OSError, ValueError) as exc:
180 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"})
183# Backward-compatible private name retained for focused tests and callers.
184_SyncLocalContract = LocalPathContract
187def _resolve_sync_local_path(
188 local_path: str,
189 *,
190 require_exists: bool,
191) -> LocalPathContract:
192 """Issue an identity-bound confinement contract for an MCP sync path."""
193 return resolve_local_path(
194 local_path,
195 require_exists=require_exists,
196 purpose="Local sync",
197 )
200def _resolve_upload_local_path(local_path: str) -> LocalPathContract:
201 """Resolve an existing short-upload source beneath the shared local root."""
202 return resolve_local_path(local_path, require_exists=True, purpose="Local upload")
205# Storage sync reads or writes the MCP host's filesystem and may transfer a
206# large amount of data, so the tool is absent unless the operator opts in.
207if is_enabled(FLAG_LOCAL_STORAGE_SYNC):
209 @mcp.tool(tags={"low-risk", "storage", "local-filesystem", "data-upload"})
210 @audit_logged
211 async def sync_storage_bucket(
212 bucket_alias: str,
213 local_dir: str,
214 direction: Literal["download", "upload"] = "download",
215 region: str | None = None,
216 prefix: str = "",
217 dry_run: bool = False,
218 force: bool = False,
219 ) -> str:
220 """Sync between a GCO S3 bucket and the MCP host in one direction.
222 [gated by GCO_ENABLE_LOCAL_STORAGE_SYNC] The local path is confined
223 beneath ``GCO_STORAGE_LOCAL_ROOT`` before the CLI is invoked. This
224 confinement requires POSIX descriptor-relative filesystem APIs and
225 fails closed on unsupported hosts. Download is the default; upload
226 reads a local file or directory and writes S3. Neither direction
227 deletes destination-only data.
229 Args:
230 bucket_alias: Human-friendly alias returned by
231 ``list_storage_buckets``.
232 local_dir: Local path relative to ``GCO_STORAGE_LOCAL_ROOT`` (or an
233 absolute path contained by that root).
234 direction: ``download`` for S3-to-local or ``upload`` for local-to-S3.
235 region: Region for an unqualified ``regional-shared`` alias.
236 prefix: Remote S3 key prefix to download from or upload into.
237 dry_run: Return the transfer summary without writing files or S3 objects.
238 force: Transfer all matching files even if the destination is current.
239 """
240 normalized_direction = direction.strip().lower()
241 try:
242 local_contract = _resolve_sync_local_path(
243 local_dir,
244 require_exists=normalized_direction == "upload",
245 )
246 except (OSError, ValueError) as exc:
247 return json.dumps(
248 {
249 "error": str(exc),
250 "code": "local_storage_path_rejected",
251 }
252 )
254 args = [
255 "storage",
256 "sync",
257 "--direction",
258 normalized_direction,
259 "--_gco-storage-root",
260 str(local_contract.root),
261 "--_gco-storage-root-device",
262 str(local_contract.device),
263 "--_gco-storage-root-inode",
264 str(local_contract.inode),
265 ]
266 if region:
267 args += ["--region", region]
268 if prefix:
269 args += ["--prefix", prefix]
270 if dry_run:
271 args.append("--dry-run")
272 if force:
273 args.append("--force")
274 # End option parsing before untrusted positional values. In particular,
275 # a confined root child named ``--prefix`` must remain a local path.
276 args += ["--", bucket_alias, local_contract.local_argument]
277 return await cli_runner._run_cli_async(
278 *args,
279 timeout_seconds=3600,
280 terminate_grace_seconds=30,
281 )