Coverage for gco_mcp/server.py: 84.48%

1""" 

2FastMCP server instance and instructions for the GCO MCP server. 

3 

4This module creates the shared ``mcp`` FastMCP instance that all tool and 

5resource modules register against. Import ``mcp`` from here — never create 

6a second instance. 

7""" 

8 

9import os 

10import sys 

11 

12from fastmcp import FastMCP 

13 

14# Code Mode lives under fastmcp.experimental — the import path itself signals 

15# the API can move between minor versions. The fastmcp pin in pyproject.toml is 

16# intentionally an `==` to keep that surface stable for a release. 

17from fastmcp.experimental.transforms.code_mode import ( 

18 CodeMode, 

19 GetSchemas, 

20 GetTags, 

21 MontySandboxProvider, 

22 Search, 

23) 

24from fastmcp.server.transforms import ResourcesAsTools 

25from fastmcp.server.transforms.search import BM25SearchTransform, RegexSearchTransform 

26 

27# ``run_mcp`` supports both legacy top-level imports (``import server`` after 

28# adding gco_mcp/ to sys.path) and package imports. Bind both names to the first 

29# loaded module before constructing the FastMCP object so the two routes can 

30# never create independent registries. 

31_THIS_MODULE = sys.modules[__name__] 

32if __name__ == "server": 32 ↛ 34line 32 didn't jump to line 34 because the condition on line 32 was always true

33 sys.modules.setdefault("gco_mcp.server", _THIS_MODULE) 

34elif __name__ == "gco_mcp.server": 

35 sys.modules.setdefault("server", _THIS_MODULE) 

36 

37mcp = FastMCP( 

38 "GCO", 

39 instructions=( 

40 "Multi-region EKS Auto Mode platform for AI/ML workload orchestration. " 

41 "Submit jobs, manage inference endpoints, check capacity, track costs, " 

42 "and manage infrastructure across AWS regions.\n\n" 

43 "Resources available:\n" 

44 "- docs:// — Documentation, architecture guides, and example job/inference manifests\n" 

45 "- k8s:// — Kubernetes manifests deployed to the cluster (RBAC, deployments, NodePools, etc.)\n" 

46 "- iam:// — IAM policy templates for access control\n" 

47 "- infra:// — Dockerfiles, Helm charts, CI/CD config\n" 

48 "- ci:// — GitHub Actions workflows, composite actions, scripts, issue/PR templates\n" 

49 "- source:// — Full source code of the platform\n" 

50 "- demos:// — Demo walkthroughs, live demo scripts, and presentation materials\n" 

51 "- clients:// — API client examples (Python, curl, AWS CLI)\n" 

52 "- scripts:// — Utility scripts for cluster access, versioning, testing\n" 

53 "- tests:// — Test suite documentation, patterns, and configuration\n" 

54 "- config:// — CDK configuration and environment variables\n" 

55 "- images:// — ECR repositories, tags, image details, and replication state\n" 

56 "- gco:// — Live regional jobs, Kubernetes objects, cluster topology, and inference state\n" 

57 "- costs:// and tasks:// — Windowed cost and background-task status views\n" 

58 "- mission:// — Mission sessions, reports, and audit replay (feature-gated)\n" 

59 "- mcp:// — Live tool/resource indexes and feature-flag mappings\n\n" 

60 "Start with docs://gco/index or mcp://gco/resources/index to explore." 

61 ), 

62 # NOTE on background-task support: ``tasks=True`` is intentionally NOT set 

63 # here. FastMCP's ``tasks=True`` at the server level applies a default 

64 # ``TaskConfig(mode="optional")`` to every tool, which requires every tool 

65 # function to be async (FastMCP raises ValueError at registration time 

66 # otherwise). The async migration of existing sync tools lands in a 

67 # later phase; until then, the long-running tools that genuinely need 

68 # background-task support set ``task=TaskConfig(mode=...)`` on their 

69 # individual ``@mcp.tool(...)`` decorators rather than relying on the 

70 # server-wide default. The ``fastmcp[tasks]`` extra is still pulled in 

71 # via ``pyproject.toml`` so pydocket is available when those per-tool 

72 # decorators run. 

73) 

74 

75# Always-on: tool-only clients (Cursor) get list_resources/read_resource synthetic tools. 

76# Registered AFTER the catalog-replacement transform below so the synthetic 

77# resource tools survive even when BM25/Regex/Code Mode replace the catalog. 

78 

79 

80def _int_env(name: str, default: int) -> int: 

81 """Parse an integer env var; fall back to default on missing/empty/non-numeric.""" 

82 raw = os.environ.get(name, "").strip() 

83 if not raw: 

84 return default 

85 try: 

86 return int(raw) 

87 except ValueError: 

88 return default 

89 

90 

91def _float_env(name: str, default: float) -> float: 

92 """Parse a float env var; fall back to default on missing/empty/non-numeric.""" 

93 raw = os.environ.get(name, "").strip() 

94 if not raw: 

95 return default 

96 try: 

97 return float(raw) 

98 except ValueError: 

99 return default 

100 

101 

102# Catalog-replacement transform. Mutually exclusive between the four values. 

103# Default is "bm25" so a brand-new install gets relevance-ranked tool search 

104# without any extra configuration. An unknown value (typo, etc.) also falls 

105# back to "bm25" so a misconfigured client doesn't accidentally drop into the 

106# full-catalog listing. 

107_TOOL_SEARCH = os.environ.get("GCO_MCP_TOOL_SEARCH", "bm25").strip().lower() 

108_ALWAYS_VISIBLE = [ 

109 "find_examples", 

110 "find_docs", 

111 "list_jobs", 

112 "submit_job_sqs", 

113 "list_inference_endpoints", 

114 "check_capacity", 

115 "task_status", 

116] 

117if _TOOL_SEARCH == "bm25": 

118 mcp.add_transform(BM25SearchTransform(always_visible=_ALWAYS_VISIBLE)) 

119elif _TOOL_SEARCH == "regex": 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 mcp.add_transform(RegexSearchTransform(always_visible=_ALWAYS_VISIBLE)) 

121elif _TOOL_SEARCH == "code_mode": 

122 # Four-stage discovery: GetTags → Search → GetSchemas → execute. Tags are 

123 # mandatory on every tool, so GetTags as the first stage gives the LLM 

124 # cheap browse-by-category before searching. 

125 mcp.add_transform( 

126 CodeMode( 

127 discovery_tools=[GetTags(), Search(), GetSchemas()], 

128 sandbox_provider=MontySandboxProvider( 

129 limits={ 

130 "max_duration_secs": _float_env("GCO_MCP_CODE_MODE_MAX_DURATION_SECS", 30.0), 

131 "max_memory": _int_env("GCO_MCP_CODE_MODE_MAX_MEMORY", 200_000_000), 

132 }, 

133 ), 

134 ) 

135 ) 

136elif _TOOL_SEARCH == "off": 

137 pass # legacy: list_tools returns the full catalog 

138else: 

139 # Unknown value → behave as the default (bm25). 

140 mcp.add_transform(BM25SearchTransform(always_visible=_ALWAYS_VISIBLE)) 

141 

142 

143# Resources As Tools is added AFTER the catalog-replacement transform so its 

144# synthetic ``list_resources`` / ``read_resource`` tools are appended to the 

145# catalog the search transform produced. Tool-only clients (Cursor, etc.) 

146# always see the resource surface even under search-mode. 

147mcp.add_transform(ResourcesAsTools(mcp)) 

148 

149 

150# Audit-capture middleware. Installs once after the transforms so every 

151# tool invocation gets fresh per-call buffers for ctx.warning/info/error 

152# and ctx.elicit. The patched Context methods are a no-op outside an 

153# active middleware scope, so this has no effect on non-MCP callers. 

154from audit_middleware import AuditCaptureMiddleware # noqa: E402 

155 

156mcp.add_middleware(AuditCaptureMiddleware())