Coverage for gco_mcp/tools/models.py: 95.92%

1"""Model weight management MCP tools.""" 

2 

3import asyncio 

4import contextlib 

5import json 

6 

7import cli_runner 

8from audit import audit_logged 

9from feature_flags import FLAG_DESTRUCTIVE_OPERATIONS, FLAG_MODEL_UPLOAD, is_enabled 

10from local_data import resolve_local_path, stage_upload_path 

11from server import mcp 

12 

13 

14@mcp.tool(tags={"safe", "models"}) 

15@audit_logged 

16def list_models() -> str: 

17 """List all uploaded model weights in the S3 bucket.""" 

18 return cli_runner._run_cli("models", "list") 

19 

20 

21@mcp.tool(tags={"safe", "models"}) 

22@audit_logged 

23def get_model_uri(model_name: str) -> str: 

24 """Get the S3 URI for a model (for use with --model-source). 

25 

26 Args: 

27 model_name: Name of the model. 

28 """ 

29 return cli_runner._run_cli("models", "uri", model_name) 

30 

31 

32async def _ctx_warning(message: str) -> None: 

33 """Emit ``ctx.warning(...)`` from inside a tool body, no-op when no Context.""" 

34 try: 

35 from fastmcp.server.dependencies import get_context 

36 

37 ctx = get_context() 

38 except Exception: 

39 return 

40 with contextlib.suppress(Exception): 

41 await ctx.warning(message) 

42 

43 

44# ============================================================================= 

45# Model upload — gated by GCO_ENABLE_MODEL_UPLOAD 

46# ============================================================================= 

47 

48 

49if is_enabled(FLAG_MODEL_UPLOAD): 

50 

51 @mcp.tool(tags={"data-upload", "models"}) 

52 @audit_logged 

53 async def models_upload( 

54 model_name: str, 

55 source_path: str, 

56 ) -> str: 

57 """[gated by GCO_ENABLE_MODEL_UPLOAD] data-upload. 

58 

59 Upload model weights from a source confined beneath 

60 ``GCO_STORAGE_LOCAL_ROOT``. Relative paths are resolved beneath that 

61 root rather than the MCP process's working directory. The source is 

62 exposed to the CLI through a private descriptor-backed snapshot; 

63 descendant links, special files, hard links, and filesystem crossings 

64 fail closed. 

65 

66 Args: 

67 model_name: Model name in the registry. 

68 source_path: Root-relative local file or directory to upload. 

69 """ 

70 try: 

71 local_contract = resolve_local_path( 

72 source_path, 

73 require_exists=True, 

74 purpose="Model upload", 

75 ) 

76 except (OSError, ValueError) as exc: 

77 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"}) 

78 

79 def _upload_from_staged_path() -> str: 

80 with stage_upload_path(local_contract) as staged: 

81 return cli_runner._run_cli( 

82 "models", 

83 "upload", 

84 staged.argument, 

85 "--name", 

86 model_name, 

87 pass_fds=(staged.directory_fd,), 

88 ) 

89 

90 try: 

91 # Keep staging, descriptor inheritance, subprocess execution, and 

92 # cleanup in one worker. If the awaiting MCP request is cancelled, 

93 # the worker retains the private snapshot until the CLI exits. 

94 return await asyncio.to_thread(_upload_from_staged_path) 

95 except (OSError, ValueError) as exc: 

96 return json.dumps({"error": str(exc), "code": "local_data_path_rejected"}) 

97 

98 

99# ============================================================================= 

100# Destructive tools — gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS 

101# ============================================================================= 

102 

103 

104if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS): 

105 

106 @mcp.tool(tags={"destructive", "models"}) 

107 @audit_logged 

108 async def delete_model(model_name: str) -> str: 

109 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive. 

110 

111 `gco models delete` — delete a model from the central S3 bucket. 

112 Cannot be undone — every file under the model's S3 prefix is 

113 permanently removed. 

114 

115 Args: 

116 model_name: Name of the model to delete. 

117 """ 

118 await _ctx_warning(f"Deleting model {model_name!r} — this cannot be undone.") 

119 return await asyncio.to_thread(cli_runner._run_cli, "models", "delete", model_name, "-y")