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.
- {dh_cli-0.13.3 → dh_cli-0.14.0}/PKG-INFO +4 -1
- {dh_cli-0.13.3 → dh_cli-0.14.0}/pyproject.toml +2 -1
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/__init__.py +38 -0
- dh_cli-0.14.0/src/dh_cli/hz/auth.py +162 -0
- dh_cli-0.14.0/src/dh_cli/hz/mcp_server.py +200 -0
- dh_cli-0.14.0/tests/hz/live_smoke.py +49 -0
- dh_cli-0.14.0/tests/hz/test_auth.py +134 -0
- dh_cli-0.14.0/tests/hz/test_mcp_server.py +268 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/.gitignore +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/LICENSE +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/README.md +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/_identity.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/aws_batch.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/boltz.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/cancel.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/clean.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/cost.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/disco.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/equilibrator.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/esmfold2.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/finalize.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/local.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/logs.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/orca.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/retry.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/rfd3.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/status.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/submit.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/cost_report.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/fasta_utils.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/h5_utils.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/job_id.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/manifest.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/batch/s3_transport.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/cost_report.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/bedrock/pricing.yaml +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/cloud_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/codeartifact.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/api_client.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/auth.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/progress.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/retry.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/engines_studios/studio_resolve.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/github_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/deploy.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/local.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/test.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/tf.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/hz/users.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/main.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/sweep.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/src/dh_cli/warehouse.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/cur_fixtures.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_aws_batch_resources.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_filters.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_reader.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_reconcile.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_batch_cost_report.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_image_override.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_no_train_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_cpu_only.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_image_validation.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_job_tags.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_merge.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/batch/test_submit_mlflow_env.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/conftest.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/A_cache_write.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/B_cache_read.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/C_plain.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/E_service_role.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_build_report.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_classify_arn.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cli_exit_codes.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cost_calc.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cost_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_cur_reconciliation.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_key_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_render_formats.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_resolve_base_model.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_rotate.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/bedrock/test_s3_walker.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_engine_retry_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_retry.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_ssh_config.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_studio_resolution_commands.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/engines_studios/test_studio_resolve.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/__init__.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/conftest.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_identity.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login_error_paths.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_login_security.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_logout.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_rotate.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/github/test_status.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_init.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_suites.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/hz/test_users.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_aws_status.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_cloud_gcp.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_disco_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_disco_jobdef.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_equilibrator_command.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_esmfold2_cli.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_boltz_tar.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_disco.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_disco_dispatch.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_equilibrator.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_esmfold2.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_protmpnn.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_finalize_rfd3.py +0 -0
- {dh_cli-0.13.3 → dh_cli-0.14.0}/tests/test_rfd3_command.py +0 -0
- {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.
|
|
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.
|
|
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)
|