Coverage for gco_mcp/tools/storage.py: 98.00%
80 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +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# =============================================================================
64# Read-only inspection tools (async)
65# =============================================================================
68@mcp.tool(tags={"safe", "files"})
69@audit_logged
70async def files_get(region: str, fs_type: str = "efs") -> str:
71 """`gco files get` — get file system details for a region.
73 Returns the EFS (or FSx) file system's ID, lifecycle state, throughput mode,
74 encryption flags, and mount targets. To browse or fetch file contents, use
75 list_storage_contents (``gco files ls``) or ``gco files download``.
77 Args:
78 region: AWS region.
79 fs_type: File system type — "efs" (default) or "fsx".
80 """
81 return await asyncio.to_thread(cli_runner._run_cli, "files", "get", region, "-t", fs_type)
84@mcp.tool(tags={"safe", "files"})
85@audit_logged
86async def files_access_points(region: str | None = None) -> str:
87 """`gco files access-points` — list EFS access points.
89 Args:
90 region: AWS region.
91 """
92 args = ["files", "access-points"]
93 if region:
94 args += ["-r", region]
95 return await asyncio.to_thread(cli_runner._run_cli, *args)
98# =============================================================================
99# Regional bucket upload (low-risk write)
100# =============================================================================
103if is_enabled(FLAG_MODEL_UPLOAD):
105 @mcp.tool(tags={"data-upload", "storage", "local-filesystem"})
106 @audit_logged
107 async def upload_to_regional_bucket(
108 local_path: str, region: str, prefix: str = "uploads"
109 ) -> str:
110 """[gated by GCO_ENABLE_MODEL_UPLOAD] Upload local data to a regional bucket.
112 The source must resolve beneath ``GCO_STORAGE_LOCAL_ROOT``. Relative
113 paths such as ``model.bin`` and ``./datasets`` are interpreted relative
114 to that root, never relative to the MCP process working directory. The
115 CLI receives a private descriptor-backed snapshot; descendant links,
116 special files, hard links, and filesystem crossings fail closed.
118 Args:
119 local_path: Root-relative local file or directory to upload.
120 region: Target region whose regional bucket receives the objects.
121 prefix: S3 prefix for uploaded objects (default: ``uploads``).
122 """
123 try:
124 local_contract = _resolve_upload_local_path(local_path)
125 except (OSError, ValueError) as exc:
126 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"})
128 def _upload_from_staged_path() -> str:
129 with stage_upload_path(local_contract) as staged:
130 return cli_runner._run_cli(
131 "models",
132 "upload-regional",
133 staged.argument,
134 "-r",
135 region,
136 "--prefix",
137 prefix,
138 pass_fds=(staged.directory_fd,),
139 )
141 try:
142 # Run the staging context in the worker so cancellation cannot
143 # unlink its descriptor-backed snapshot while the CLI still reads.
144 return await asyncio.to_thread(_upload_from_staged_path)
145 except (OSError, ValueError) as exc:
146 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"})
149# Backward-compatible private name retained for focused tests and callers.
150_SyncLocalContract = LocalPathContract
153def _resolve_sync_local_path(
154 local_path: str,
155 *,
156 require_exists: bool,
157) -> LocalPathContract:
158 """Issue an identity-bound confinement contract for an MCP sync path."""
159 return resolve_local_path(
160 local_path,
161 require_exists=require_exists,
162 purpose="Local sync",
163 )
166def _resolve_upload_local_path(local_path: str) -> LocalPathContract:
167 """Resolve an existing short-upload source beneath the shared local root."""
168 return resolve_local_path(local_path, require_exists=True, purpose="Local upload")
171# Storage sync reads or writes the MCP host's filesystem and may transfer a
172# large amount of data, so the tool is absent unless the operator opts in.
173if is_enabled(FLAG_LOCAL_STORAGE_SYNC):
175 @mcp.tool(tags={"low-risk", "storage", "local-filesystem", "data-upload"})
176 @audit_logged
177 async def sync_storage_bucket(
178 bucket_alias: str,
179 local_dir: str,
180 direction: Literal["download", "upload"] = "download",
181 region: str | None = None,
182 prefix: str = "",
183 dry_run: bool = False,
184 force: bool = False,
185 ) -> str:
186 """Sync between a GCO S3 bucket and the MCP host in one direction.
188 [gated by GCO_ENABLE_LOCAL_STORAGE_SYNC] The local path is confined
189 beneath ``GCO_STORAGE_LOCAL_ROOT`` before the CLI is invoked. This
190 confinement requires POSIX descriptor-relative filesystem APIs and
191 fails closed on unsupported hosts. Download is the default; upload
192 reads a local file or directory and writes S3. Neither direction
193 deletes destination-only data.
195 Args:
196 bucket_alias: Human-friendly alias returned by
197 ``list_storage_buckets``.
198 local_dir: Local path relative to ``GCO_STORAGE_LOCAL_ROOT`` (or an
199 absolute path contained by that root).
200 direction: ``download`` for S3-to-local or ``upload`` for local-to-S3.
201 region: Region for an unqualified ``regional-shared`` alias.
202 prefix: Remote S3 key prefix to download from or upload into.
203 dry_run: Return the transfer summary without writing files or S3 objects.
204 force: Transfer all matching files even if the destination is current.
205 """
206 normalized_direction = direction.strip().lower()
207 try:
208 local_contract = _resolve_sync_local_path(
209 local_dir,
210 require_exists=normalized_direction == "upload",
211 )
212 except (OSError, ValueError) as exc:
213 return json.dumps(
214 {
215 "error": str(exc),
216 "code": "local_storage_path_rejected",
217 }
218 )
220 args = [
221 "storage",
222 "sync",
223 "--direction",
224 normalized_direction,
225 "--_gco-storage-root",
226 str(local_contract.root),
227 "--_gco-storage-root-device",
228 str(local_contract.device),
229 "--_gco-storage-root-inode",
230 str(local_contract.inode),
231 ]
232 if region:
233 args += ["--region", region]
234 if prefix:
235 args += ["--prefix", prefix]
236 if dry_run:
237 args.append("--dry-run")
238 if force:
239 args.append("--force")
240 # End option parsing before untrusted positional values. In particular,
241 # a confined root child named ``--prefix`` must remain a local path.
242 args += ["--", bucket_alias, local_contract.local_argument]
243 return await cli_runner._run_cli_async(
244 *args,
245 timeout_seconds=3600,
246 terminate_grace_seconds=30,
247 )