dh-cli 0.13.3__tar.gz → 0.14.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (134) hide show
  1. {dh_cli-0.13.3 → dh_cli-0.14.0}/PKG-INFO +4 -1
  2. {dh_cli-0.13.3 → dh_cli-0.14.0}/pyproject.toml +2 -1
  3. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/__init__.py +38 -0
  4. dh_cli-0.14.0/src/dh_cli/hz/auth.py +162 -0
  5. dh_cli-0.14.0/src/dh_cli/hz/mcp_server.py +200 -0
  6. dh_cli-0.14.0/tests/hz/live_smoke.py +49 -0
  7. dh_cli-0.14.0/tests/hz/test_auth.py +134 -0
  8. dh_cli-0.14.0/tests/hz/test_mcp_server.py +268 -0
  9. {dh_cli-0.13.3 → dh_cli-0.14.0}/.gitignore +0 -0
  10. {dh_cli-0.13.3 → dh_cli-0.14.0}/LICENSE +0 -0
  11. {dh_cli-0.13.3 → dh_cli-0.14.0}/README.md +0 -0
  12. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/__init__.py +0 -0
  13. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/_identity.py +0 -0
  14. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/__init__.py +0 -0
  15. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/aws_batch.py +0 -0
  16. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/__init__.py +0 -0
  17. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/boltz.py +0 -0
  18. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/cancel.py +0 -0
  19. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/clean.py +0 -0
  20. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/cost.py +0 -0
  21. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/disco.py +0 -0
  22. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
  23. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/equilibrator.py +0 -0
  24. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/esmfold2.py +0 -0
  25. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/finalize.py +0 -0
  26. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
  27. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/local.py +0 -0
  28. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/logs.py +0 -0
  29. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/orca.py +0 -0
  30. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
  31. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
  32. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/retry.py +0 -0
  33. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/rfd3.py +0 -0
  34. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/status.py +0 -0
  35. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/submit.py +0 -0
  36. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
  37. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/cost_report.py +0 -0
  38. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/fasta_utils.py +0 -0
  39. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/h5_utils.py +0 -0
  40. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/job_id.py +0 -0
  41. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/manifest.py +0 -0
  42. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/s3_transport.py +0 -0
  43. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/__init__.py +0 -0
  44. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/commands.py +0 -0
  45. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/cost_report.py +0 -0
  46. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/pricing.yaml +0 -0
  47. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/cloud_commands.py +0 -0
  48. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/codeartifact.py +0 -0
  49. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/__init__.py +0 -0
  50. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/api_client.py +0 -0
  51. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/auth.py +0 -0
  52. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
  53. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/progress.py +0 -0
  54. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/retry.py +0 -0
  55. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
  56. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
  57. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/studio_resolve.py +0 -0
  58. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/github_commands.py +0 -0
  59. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/deploy.py +0 -0
  60. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/local.py +0 -0
  61. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/test.py +0 -0
  62. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/tf.py +0 -0
  63. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/users.py +0 -0
  64. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/main.py +0 -0
  65. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/sweep.py +0 -0
  66. {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/warehouse.py +0 -0
  67. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/__init__.py +0 -0
  68. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/cur_fixtures.py +0 -0
  69. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_aws_batch_resources.py +0 -0
  70. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_command.py +0 -0
  71. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_filters.py +0 -0
  72. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_reader.py +0 -0
  73. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_reconcile.py +0 -0
  74. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_report.py +0 -0
  75. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_image_override.py +0 -0
  76. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_no_train_command.py +0 -0
  77. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_cpu_only.py +0 -0
  78. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_image_validation.py +0 -0
  79. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_job_tags.py +0 -0
  80. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_merge.py +0 -0
  81. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_mlflow_env.py +0 -0
  82. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/conftest.py +0 -0
  83. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/A_cache_write.json +0 -0
  84. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/B_cache_read.json +0 -0
  85. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/C_plain.json +0 -0
  86. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
  87. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/E_service_role.json +0 -0
  88. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
  89. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
  90. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_build_report.py +0 -0
  91. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_classify_arn.py +0 -0
  92. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cli_exit_codes.py +0 -0
  93. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cost_calc.py +0 -0
  94. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cost_command.py +0 -0
  95. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cur_reconciliation.py +0 -0
  96. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_key_command.py +0 -0
  97. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_render_formats.py +0 -0
  98. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_resolve_base_model.py +0 -0
  99. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_rotate.py +0 -0
  100. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_s3_walker.py +0 -0
  101. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/__init__.py +0 -0
  102. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_engine_retry_commands.py +0 -0
  103. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_retry.py +0 -0
  104. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_ssh_config.py +0 -0
  105. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_studio_resolution_commands.py +0 -0
  106. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_studio_resolve.py +0 -0
  107. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/__init__.py +0 -0
  108. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/conftest.py +0 -0
  109. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
  110. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_identity.py +0 -0
  111. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login.py +0 -0
  112. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login_error_paths.py +0 -0
  113. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login_security.py +0 -0
  114. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_logout.py +0 -0
  115. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_rotate.py +0 -0
  116. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_status.py +0 -0
  117. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_init.py +0 -0
  118. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_suites.py +0 -0
  119. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_users.py +0 -0
  120. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_aws_status.py +0 -0
  121. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_cloud_gcp.py +0 -0
  122. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_disco_command.py +0 -0
  123. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_disco_jobdef.py +0 -0
  124. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_equilibrator_command.py +0 -0
  125. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_esmfold2_cli.py +0 -0
  126. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_boltz_tar.py +0 -0
  127. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_disco.py +0 -0
  128. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_disco_dispatch.py +0 -0
  129. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_equilibrator.py +0 -0
  130. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_esmfold2.py +0 -0
  131. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_protmpnn.py +0 -0
  132. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_rfd3.py +0 -0
  133. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_rfd3_command.py +0 -0
  134. {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_sweep.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: dh-cli
3
- Version: 0.13.3
3
+ Version: 0.14.0
4
4
  Summary: Dayhoff Labs developer CLI
5
5
  Author-email: Dayhoff Labs <dev@dayhofflabs.com>
6
6
  License: # PolyForm Noncommercial License 1.0.0
@@ -135,6 +135,9 @@ Requires-Dist: typer>=0.9.0
135
135
  Provides-Extra: h5
136
136
  Requires-Dist: h5py>=3.0; extra == 'h5'
137
137
  Requires-Dist: numpy>=1.24; extra == 'h5'
138
+ Provides-Extra: mcp
139
+ Requires-Dist: httpx>=0.27; extra == 'mcp'
140
+ Requires-Dist: mcp>=2.2; extra == 'mcp'
138
141
  Description-Content-Type: text/markdown
139
142
 
140
143
  # dh-cli
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "dh-cli"
7
- version = "0.13.3"
7
+ version = "0.14.0"
8
8
  description = "Dayhoff Labs developer CLI"
9
9
  requires-python = ">=3.11"
10
10
  readme = "README.md"
@@ -25,6 +25,7 @@ dependencies = [
25
25
 
26
26
  [project.optional-dependencies]
27
27
  h5 = ["h5py>=3.0", "numpy>=1.24"]
28
+ mcp = ["mcp>=2.2", "httpx>=0.27"]
28
29
 
29
30
  [project.scripts]
30
31
  dh = "dh_cli.main:app"
@@ -95,6 +95,43 @@ def run_script(script: Path, args: list[str] | None = None, cwd: Path | None = N
95
95
  raise typer.Exit(result.returncode)
96
96
 
97
97
 
98
+ @hz_app.command()
99
+ def login(
100
+ email: str | None = typer.Option(None, "--email", help="Login email (prompted if omitted)."),
101
+ ):
102
+ """Log in to Horizyn (prod) with an emailed code; saves credentials for `dh hz mcp serve`."""
103
+ from dh_cli.hz.auth import credentials_path
104
+ from dh_cli.hz.auth import login as _login
105
+
106
+ try:
107
+ creds = _login(email)
108
+ except RuntimeError as exc:
109
+ typer.echo(f"Error: {exc}", err=True)
110
+ raise typer.Exit(1)
111
+ typer.echo(f"Logged in as {creds.email}. Credentials saved to {credentials_path()}", err=True)
112
+
113
+
114
+ mcp_app = typer.Typer(help="Horizyn MCP server for agent clients (Cursor, Codex, Claude Code).")
115
+
116
+
117
+ @mcp_app.command()
118
+ def serve():
119
+ """Run the Horizyn MCP server over stdio. Run `dh hz login` first."""
120
+ try:
121
+ import httpx # noqa: F401
122
+ import mcp # noqa: F401
123
+ except ImportError:
124
+ typer.echo(
125
+ "The Horizyn MCP server needs the 'mcp' extra: pip install 'dh-cli[mcp]'",
126
+ err=True,
127
+ )
128
+ raise typer.Exit(1)
129
+
130
+ from dh_cli.hz.mcp_server import server
131
+
132
+ server.run(transport="stdio")
133
+
134
+
98
135
  from dh_cli.hz.deploy import deploy_app # noqa: E402
99
136
  from dh_cli.hz.local import local_app # noqa: E402
100
137
  from dh_cli.hz.test import test_app # noqa: E402
@@ -106,3 +143,4 @@ hz_app.add_typer(deploy_app, name="deploy", help="Deploy API, docking, or fronte
106
143
  hz_app.add_typer(local_app, name="local", help="Run local development servers.")
107
144
  hz_app.add_typer(test_app, name="test", help="Run integration test suites.")
108
145
  hz_app.add_typer(tf_app, name="tf", help="Terraform apply for Horizyn infrastructure.")
146
+ hz_app.add_typer(mcp_app, name="mcp", help="Horizyn MCP server for agent clients.")
@@ -0,0 +1,162 @@
1
+ """Horizyn credentials: Cognito passwordless login, storage, and silent refresh.
2
+
3
+ Prod only. Credentials live in ``~/.horizyn/credentials_prod.json`` (0600) and
4
+ hold the ID token, the refresh token, and the login email. Port of
5
+ ``horizyn-api/test/server/deployed/get_jwt_token.py`` without the env switch.
6
+ """
7
+
8
+ import base64
9
+ import json
10
+ import os
11
+ import sys
12
+ import time
13
+ from dataclasses import asdict, dataclass
14
+ from pathlib import Path
15
+
16
+ import boto3
17
+ from botocore.exceptions import BotoCoreError, ClientError
18
+
19
+ REGION = "us-east-1"
20
+ CLIENT_ID = "4bqudnqan63iur9fi8224m2pee"
21
+ REFRESH_MARGIN_S = 300
22
+
23
+
24
+ def credentials_path() -> Path:
25
+ return Path.home() / ".horizyn" / "credentials_prod.json"
26
+
27
+
28
+ class NotLoggedIn(Exception):
29
+ def __init__(self, detail: str = ""):
30
+ msg = "Not logged in to Horizyn: run `dh hz login`"
31
+ super().__init__(f"{msg} ({detail})" if detail else msg)
32
+
33
+
34
+ @dataclass
35
+ class Credentials:
36
+ id_token: str
37
+ refresh_token: str
38
+ email: str
39
+
40
+ @classmethod
41
+ def load(cls, path: Path | None = None) -> "Credentials":
42
+ path = path or credentials_path()
43
+ try:
44
+ data = json.loads(path.read_text())
45
+ except FileNotFoundError:
46
+ raise NotLoggedIn("no credentials file") from None
47
+ return cls(
48
+ id_token=data.get("id_token", ""),
49
+ refresh_token=data.get("refresh_token", ""),
50
+ email=data.get("email", ""),
51
+ )
52
+
53
+ def save(self, path: Path | None = None) -> None:
54
+ path = path or credentials_path()
55
+ path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
56
+ path.parent.chmod(0o700)
57
+ fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
58
+ with os.fdopen(fd, "w") as f:
59
+ json.dump(asdict(self), f, indent=2)
60
+ path.chmod(0o600)
61
+
62
+
63
+ def _jwt_claims(token: str) -> dict:
64
+ """Decode a JWT payload without verifying it. Returns {} if undecodable."""
65
+ try:
66
+ payload = token.split(".")[1]
67
+ payload += "=" * (-len(payload) % 4)
68
+ claims = json.loads(base64.urlsafe_b64decode(payload))
69
+ except (IndexError, ValueError):
70
+ return {}
71
+ return claims if isinstance(claims, dict) else {}
72
+
73
+
74
+ def _cognito():
75
+ return boto3.client("cognito-idp", region_name=REGION)
76
+
77
+
78
+ def get_id_token(path: Path | None = None) -> str:
79
+ """Return a valid ID token, refreshing via Cognito when it is about to expire."""
80
+ creds = Credentials.load(path)
81
+ exp = _jwt_claims(creds.id_token).get("exp", 0)
82
+ if exp - REFRESH_MARGIN_S > time.time():
83
+ return creds.id_token
84
+ if not creds.refresh_token:
85
+ raise NotLoggedIn("no refresh token")
86
+
87
+ try:
88
+ response = _cognito().initiate_auth(
89
+ ClientId=CLIENT_ID,
90
+ AuthFlow="REFRESH_TOKEN_AUTH",
91
+ AuthParameters={"REFRESH_TOKEN": creds.refresh_token},
92
+ )
93
+ except (ClientError, BotoCoreError) as exc:
94
+ raise NotLoggedIn(f"refresh failed: {type(exc).__name__}") from None
95
+
96
+ result = response.get("AuthenticationResult") or {}
97
+ id_token = result.get("IdToken")
98
+ if not id_token:
99
+ raise NotLoggedIn("refresh returned no token")
100
+ creds.id_token = id_token
101
+ if result.get("RefreshToken"):
102
+ creds.refresh_token = result["RefreshToken"]
103
+ creds.save(path)
104
+ return id_token
105
+
106
+
107
+ def _prompt(text: str) -> str:
108
+ print(text, end="", file=sys.stderr, flush=True)
109
+ return input().strip()
110
+
111
+
112
+ def login(email: str | None = None, path: Path | None = None) -> Credentials:
113
+ """Interactive CUSTOM_AUTH login: email → emailed code → tokens saved to disk.
114
+
115
+ Prompts go to stderr. Raises RuntimeError with a user-facing message on failure.
116
+ """
117
+ if not email:
118
+ email = _prompt("Email address: ")
119
+ if not email:
120
+ raise RuntimeError("Email address is required")
121
+
122
+ cognito = _cognito()
123
+ print(f"Sending verification code to {email}...", file=sys.stderr)
124
+ try:
125
+ response = cognito.initiate_auth(
126
+ ClientId=CLIENT_ID,
127
+ AuthFlow="CUSTOM_AUTH",
128
+ AuthParameters={"USERNAME": email},
129
+ )
130
+ except ClientError as exc:
131
+ raise RuntimeError(f"Could not start login: {_client_error(exc)}") from None
132
+
133
+ code = _prompt("Enter the 6-digit verification code: ")
134
+ if len(code) != 6 or not code.isdigit():
135
+ raise RuntimeError("Please enter a valid 6-digit code")
136
+
137
+ try:
138
+ response = cognito.respond_to_auth_challenge(
139
+ ClientId=CLIENT_ID,
140
+ ChallengeName="CUSTOM_CHALLENGE",
141
+ Session=response["Session"],
142
+ ChallengeResponses={"ANSWER": code, "USERNAME": email},
143
+ )
144
+ except ClientError as exc:
145
+ raise RuntimeError(f"Verification failed: {_client_error(exc)}") from None
146
+
147
+ result = response.get("AuthenticationResult") or {}
148
+ if not result.get("IdToken") or not result.get("RefreshToken"):
149
+ raise RuntimeError("Authentication did not complete; try again")
150
+
151
+ creds = Credentials(
152
+ id_token=result["IdToken"],
153
+ refresh_token=result["RefreshToken"],
154
+ email=_jwt_claims(result["IdToken"]).get("email") or email,
155
+ )
156
+ creds.save(path)
157
+ return creds
158
+
159
+
160
+ def _client_error(exc: ClientError) -> str:
161
+ err = exc.response.get("Error", {})
162
+ return f"{err.get('Code', 'Unknown')} - {err.get('Message', str(exc))}"
@@ -0,0 +1,200 @@
1
+ """Horizyn MCP server over stdio (`dh hz mcp serve`). Prod only.
2
+
3
+ Every request carries the logged-in user's own Cognito ID token, so queries
4
+ land in that user's Horizyn history. Requires the ``mcp`` extra:
5
+ ``pip install 'dh-cli[mcp]'``.
6
+ """
7
+
8
+ from importlib.metadata import version
9
+
10
+ import httpx
11
+ from mcp.server.mcpserver import MCPServer
12
+ from mcp.server.mcpserver.exceptions import ToolError
13
+ from mcp.types import ToolAnnotations
14
+
15
+ from dh_cli.hz import API_URLS
16
+ from dh_cli.hz.auth import NotLoggedIn, get_id_token
17
+
18
+ CLIENT_HEADER = f"mcp/{version('dh-cli')}"
19
+ _READ_ONLY = ToolAnnotations(read_only_hint=True)
20
+
21
+ server = MCPServer(
22
+ "horizyn",
23
+ instructions=(
24
+ "Horizyn: Dayhoff Labs' enzyme and reaction similarity search. Tools query "
25
+ "production as the logged-in user (`dh hz login`). Similarity scores are "
26
+ "cosine similarities in [0, 1]; higher means more similar."
27
+ ),
28
+ )
29
+
30
+ _client = httpx.Client(base_url=API_URLS["prod"], timeout=120)
31
+
32
+
33
+ def _request(method: str, path: str, **kwargs) -> dict:
34
+ try:
35
+ token = get_id_token()
36
+ except NotLoggedIn as exc:
37
+ raise ToolError(str(exc)) from None
38
+
39
+ headers = {"Authorization": f"Bearer {token}", "X-Horizyn-Client": CLIENT_HEADER}
40
+ try:
41
+ resp = _client.request(method, path, headers=headers, **kwargs)
42
+ except httpx.HTTPError as exc:
43
+ raise ToolError(f"Horizyn request failed: {type(exc).__name__}: {exc}") from None
44
+
45
+ if resp.status_code == 503:
46
+ raise ToolError("Horizyn is starting or unavailable, retry shortly")
47
+ if resp.status_code >= 400:
48
+ raise ToolError(_error_text(resp))
49
+ return resp.json()
50
+
51
+
52
+ def _error_text(resp: httpx.Response) -> str:
53
+ try:
54
+ body = resp.json()
55
+ except ValueError:
56
+ body = {}
57
+ if not isinstance(body, dict):
58
+ body = {}
59
+ parts = [p for p in (body.get("error"), body.get("message")) if p]
60
+ detail = body.get("detail")
61
+ if isinstance(detail, str):
62
+ parts.append(detail)
63
+ elif isinstance(detail, list):
64
+ # FastAPI validation errors; take loc + msg only, never the echoed `input`.
65
+ for err in detail:
66
+ if isinstance(err, dict) and err.get("msg"):
67
+ loc = ".".join(str(x) for x in err.get("loc", []) if x != "body")
68
+ parts.append(f"{loc}: {err['msg']}" if loc else err["msg"])
69
+ text = "; ".join(parts) if parts else resp.reason_phrase
70
+ return f"Horizyn returned {resp.status_code}: {text}"
71
+
72
+
73
+ def _query(path: str, body: dict, filters: dict) -> dict:
74
+ body.update({k: v for k, v in filters.items() if v is not None})
75
+ data = _request("POST", path, json=body)
76
+ data.pop("timings_ms", None)
77
+ data.pop("tsne", None)
78
+ return data
79
+
80
+
81
+ @server.tool(annotations=_READ_ONLY)
82
+ def horizyn_whoami() -> dict:
83
+ """Identify the logged-in Horizyn user.
84
+
85
+ Returns email, tier (admin, alpha, or beta) and the feature flags that tier
86
+ grants. Use this first to confirm the server is authenticated.
87
+ """
88
+ return _request("GET", "/me")
89
+
90
+
91
+ @server.tool(annotations=_READ_ONLY)
92
+ def horizyn_list_screening_sets(screening_set_type: str | None = None) -> dict:
93
+ """List the screening sets available for similarity queries.
94
+
95
+ Each entry gives the set's id, type (protein or reaction), size, model
96
+ version and whether it is the current default (`is_default`).
97
+ `screening_set_type` optionally filters to "protein" or "reaction".
98
+ """
99
+ params = {"screening_set_type": screening_set_type} if screening_set_type else None
100
+ return _request("GET", "/screening_sets", params=params)
101
+
102
+
103
+ @server.tool(annotations=_READ_ONLY)
104
+ def horizyn_query_reaction(
105
+ smiles: str,
106
+ screening_set: str | None = None,
107
+ top_k: int = 50,
108
+ page: int = 1,
109
+ page_size: int = 20,
110
+ ec_include: list[str] | None = None,
111
+ ec_exclude: list[str] | None = None,
112
+ tc_include: list[str] | None = None,
113
+ tc_exclude: list[str] | None = None,
114
+ cofactor_include: list[str] | None = None,
115
+ cofactor_exclude: list[str] | None = None,
116
+ min_length: int | None = None,
117
+ max_length: int | None = None,
118
+ ) -> dict:
119
+ """Find proteins (or reactions) similar to a reaction SMILES.
120
+
121
+ `smiles` is a reaction SMILES such as "CC=O.O>>CC(O)O". When
122
+ `screening_set` is omitted Horizyn searches its current default protein set
123
+ (the response's `screening_set` names it; see horizyn_list_screening_sets).
124
+ Pass a reaction set such as horizyn_v1_1_rxns to find similar reactions
125
+ instead. `score` on each result is the cosine similarity between the query
126
+ and the hit in Horizyn's embedding space (0-1, higher is more similar).
127
+ `top_k` (max 1000) bounds the similarity search; `page`/`page_size`
128
+ paginate the result list.
129
+
130
+ All filters are optional. `ec_include`/`ec_exclude` and
131
+ `tc_include`/`tc_exclude` take hierarchical prefixes ("1", "1.1", "1.1.1",
132
+ "1.1.1.1"; "__NONE__" matches unannotated). `cofactor_include`/
133
+ `cofactor_exclude` take cofactor names ("NAD", "Zn"; "__ANY__"/"__NONE__").
134
+ `min_length`/`max_length` bound sequence length. `filter_stats` in the
135
+ response reports how many hits the filters kept.
136
+ """
137
+ body = {"smiles": smiles, "top_k": top_k, "page": page, "page_size": page_size}
138
+ filters = {
139
+ "screening_set": screening_set,
140
+ "ec_include": ec_include,
141
+ "ec_exclude": ec_exclude,
142
+ "tc_include": tc_include,
143
+ "tc_exclude": tc_exclude,
144
+ "cofactor_include": cofactor_include,
145
+ "cofactor_exclude": cofactor_exclude,
146
+ "min_length": min_length,
147
+ "max_length": max_length,
148
+ }
149
+ return _query("/query/reaction", body, filters)
150
+
151
+
152
+ @server.tool(annotations=_READ_ONLY)
153
+ def horizyn_query_protein(
154
+ sequence: str,
155
+ screening_set: str | None = None,
156
+ top_k: int = 50,
157
+ page: int = 1,
158
+ page_size: int = 20,
159
+ ec_include: list[str] | None = None,
160
+ ec_exclude: list[str] | None = None,
161
+ tc_include: list[str] | None = None,
162
+ tc_exclude: list[str] | None = None,
163
+ cofactor_include: list[str] | None = None,
164
+ cofactor_exclude: list[str] | None = None,
165
+ min_length: int | None = None,
166
+ max_length: int | None = None,
167
+ ) -> dict:
168
+ """Find proteins (or reactions) similar to a protein sequence.
169
+
170
+ `sequence` is a one-letter amino-acid string. Same screening sets, scoring,
171
+ pagination and optional filters as horizyn_query_reaction: `score` is the
172
+ cosine similarity (0-1) in Horizyn's shared reaction/protein embedding
173
+ space, omitting `screening_set` searches Horizyn's current default protein
174
+ set, and every filter argument may be omitted.
175
+ """
176
+ body = {"sequence": sequence, "top_k": top_k, "page": page, "page_size": page_size}
177
+ filters = {
178
+ "screening_set": screening_set,
179
+ "ec_include": ec_include,
180
+ "ec_exclude": ec_exclude,
181
+ "tc_include": tc_include,
182
+ "tc_exclude": tc_exclude,
183
+ "cofactor_include": cofactor_include,
184
+ "cofactor_exclude": cofactor_exclude,
185
+ "min_length": min_length,
186
+ "max_length": max_length,
187
+ }
188
+ return _query("/query/protein", body, filters)
189
+
190
+
191
+ @server.tool(annotations=_READ_ONLY)
192
+ def horizyn_protein_details(protein_id: str, screening_set: str) -> dict:
193
+ """Get full details for one protein from a screening set.
194
+
195
+ `protein_id` is the `id` returned by a query and `screening_set` is the
196
+ `screening_set` echoed in that query's response (the id is only meaningful
197
+ within its set). Returns name, organism, EC numbers, cofactors, expression
198
+ score, full sequence, length, AlphaFold id and any extra metadata.
199
+ """
200
+ return _request("GET", f"/protein/{protein_id}/details", params={"screening_set": screening_set})
@@ -0,0 +1,49 @@
1
+ """Live protocol smoke for `dh hz mcp serve`. Not collected by pytest.
2
+
3
+ Spawns the server over stdio via pixi, lists tools, and calls horizyn_whoami
4
+ against prod using the credentials from `dh hz login`.
5
+
6
+ cd dayhoff-tools && pixi run -e cpu python packages/dh-cli/tests/hz/live_smoke.py
7
+ """
8
+
9
+ import json
10
+ import sys
11
+ from pathlib import Path
12
+
13
+ import anyio
14
+ from mcp.client.client import Client
15
+ from mcp.client.stdio import StdioServerParameters
16
+
17
+ MANIFEST = Path(__file__).resolve().parents[4] / "pixi.toml"
18
+ EXPECTED = {
19
+ "horizyn_whoami",
20
+ "horizyn_list_screening_sets",
21
+ "horizyn_query_reaction",
22
+ "horizyn_query_protein",
23
+ "horizyn_protein_details",
24
+ }
25
+
26
+
27
+ async def main() -> int:
28
+ params = StdioServerParameters(
29
+ command="pixi",
30
+ args=["run", "--manifest-path", str(MANIFEST), "-e", "cpu", "dh", "hz", "mcp", "serve"],
31
+ )
32
+ async with Client(params) as client:
33
+ tools = await client.list_tools()
34
+ names = {t.name for t in tools.tools}
35
+ print("tools:", sorted(names))
36
+ assert names == EXPECTED, (names - EXPECTED, EXPECTED - names)
37
+
38
+ result = await client.call_tool("horizyn_whoami", {})
39
+ text = result.content[0].text
40
+ print("whoami:", text, "is_error:", result.is_error)
41
+ if result.is_error:
42
+ return 1
43
+ me = json.loads(text)
44
+ assert me.get("email"), me
45
+ return 0
46
+
47
+
48
+ if __name__ == "__main__":
49
+ sys.exit(anyio.run(main))
@@ -0,0 +1,134 @@
1
+ """Tests for dh hz auth: credential storage and token refresh."""
2
+
3
+ import base64
4
+ import json
5
+ import stat
6
+ import time
7
+ from unittest.mock import MagicMock, patch
8
+
9
+ import pytest
10
+ from botocore.exceptions import ClientError
11
+
12
+ from dh_cli.hz.auth import Credentials, NotLoggedIn, get_id_token
13
+
14
+
15
+ def _jwt(exp: float, email: str = "user@dayhofflabs.com") -> str:
16
+ header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode()
17
+ payload = json.dumps({"exp": int(exp), "email": email}).encode()
18
+ body = base64.urlsafe_b64encode(payload).rstrip(b"=").decode()
19
+ return f"{header}.{body}.sig"
20
+
21
+
22
+ @pytest.fixture
23
+ def creds_path(tmp_path):
24
+ return tmp_path / ".horizyn" / "credentials_prod.json"
25
+
26
+
27
+ class TestCredentials:
28
+ def test_save_load_round_trip(self, creds_path):
29
+ creds = Credentials(id_token="id", refresh_token="refresh", email="a@b.com")
30
+ creds.save(creds_path)
31
+
32
+ assert Credentials.load(creds_path) == creds
33
+
34
+ def test_save_sets_restrictive_modes(self, creds_path):
35
+ Credentials(id_token="id", refresh_token="refresh", email="a@b.com").save(creds_path)
36
+
37
+ assert stat.S_IMODE(creds_path.stat().st_mode) == 0o600
38
+ assert stat.S_IMODE(creds_path.parent.stat().st_mode) == 0o700
39
+
40
+ def test_save_overwrite_keeps_mode(self, creds_path):
41
+ creds = Credentials(id_token="id", refresh_token="refresh", email="a@b.com")
42
+ creds.save(creds_path)
43
+ creds.id_token = "id2"
44
+ creds.save(creds_path)
45
+
46
+ assert stat.S_IMODE(creds_path.stat().st_mode) == 0o600
47
+ assert Credentials.load(creds_path).id_token == "id2"
48
+
49
+ def test_load_missing_raises_not_logged_in(self, creds_path):
50
+ with pytest.raises(NotLoggedIn, match="dh hz login"):
51
+ Credentials.load(creds_path)
52
+
53
+
54
+ class TestGetIdToken:
55
+ def test_returns_cached_when_fresh(self, creds_path):
56
+ fresh = _jwt(time.time() + 3600)
57
+ Credentials(id_token=fresh, refresh_token="r", email="a@b.com").save(creds_path)
58
+
59
+ with patch("dh_cli.hz.auth.boto3.client") as client:
60
+ assert get_id_token(creds_path) == fresh
61
+ client.assert_not_called()
62
+
63
+ def test_refreshes_when_within_margin(self, creds_path):
64
+ stale = _jwt(time.time() + 60)
65
+ Credentials(id_token=stale, refresh_token="r", email="a@b.com").save(creds_path)
66
+ new = _jwt(time.time() + 3600)
67
+ cognito = MagicMock()
68
+ cognito.initiate_auth.return_value = {"AuthenticationResult": {"IdToken": new}}
69
+
70
+ with patch("dh_cli.hz.auth.boto3.client", return_value=cognito):
71
+ assert get_id_token(creds_path) == new
72
+
73
+ cognito.initiate_auth.assert_called_once_with(
74
+ ClientId="4bqudnqan63iur9fi8224m2pee",
75
+ AuthFlow="REFRESH_TOKEN_AUTH",
76
+ AuthParameters={"REFRESH_TOKEN": "r"},
77
+ )
78
+ saved = Credentials.load(creds_path)
79
+ assert saved.id_token == new
80
+ assert saved.refresh_token == "r"
81
+
82
+ def test_refresh_stores_rotated_refresh_token(self, creds_path):
83
+ Credentials(id_token=_jwt(0), refresh_token="r", email="a@b.com").save(creds_path)
84
+ cognito = MagicMock()
85
+ cognito.initiate_auth.return_value = {
86
+ "AuthenticationResult": {"IdToken": _jwt(time.time() + 3600), "RefreshToken": "r2"}
87
+ }
88
+
89
+ with patch("dh_cli.hz.auth.boto3.client", return_value=cognito):
90
+ get_id_token(creds_path)
91
+
92
+ assert Credentials.load(creds_path).refresh_token == "r2"
93
+
94
+ def test_undecodable_token_is_treated_as_expired(self, creds_path):
95
+ Credentials(id_token="garbage", refresh_token="r", email="a@b.com").save(creds_path)
96
+ new = _jwt(time.time() + 3600)
97
+ cognito = MagicMock()
98
+ cognito.initiate_auth.return_value = {"AuthenticationResult": {"IdToken": new}}
99
+
100
+ with patch("dh_cli.hz.auth.boto3.client", return_value=cognito):
101
+ assert get_id_token(creds_path) == new
102
+
103
+ def test_no_file_raises_not_logged_in(self, creds_path):
104
+ with pytest.raises(NotLoggedIn, match="dh hz login"):
105
+ get_id_token(creds_path)
106
+
107
+ def test_missing_refresh_token_raises_not_logged_in(self, creds_path):
108
+ Credentials(id_token=_jwt(0), refresh_token="", email="a@b.com").save(creds_path)
109
+
110
+ with patch("dh_cli.hz.auth.boto3.client") as client:
111
+ with pytest.raises(NotLoggedIn, match="dh hz login"):
112
+ get_id_token(creds_path)
113
+ client.assert_not_called()
114
+
115
+ def test_cognito_error_raises_not_logged_in(self, creds_path):
116
+ Credentials(id_token=_jwt(0), refresh_token="r", email="a@b.com").save(creds_path)
117
+ cognito = MagicMock()
118
+ cognito.initiate_auth.side_effect = ClientError(
119
+ {"Error": {"Code": "NotAuthorizedException", "Message": "Refresh Token has expired"}},
120
+ "InitiateAuth",
121
+ )
122
+
123
+ with patch("dh_cli.hz.auth.boto3.client", return_value=cognito):
124
+ with pytest.raises(NotLoggedIn, match="dh hz login"):
125
+ get_id_token(creds_path)
126
+
127
+ def test_refresh_response_without_token_raises_not_logged_in(self, creds_path):
128
+ Credentials(id_token=_jwt(0), refresh_token="r", email="a@b.com").save(creds_path)
129
+ cognito = MagicMock()
130
+ cognito.initiate_auth.return_value = {"ChallengeName": "SOMETHING"}
131
+
132
+ with patch("dh_cli.hz.auth.boto3.client", return_value=cognito):
133
+ with pytest.raises(NotLoggedIn, match="dh hz login"):
134
+ get_id_token(creds_path)