4 files changed, 146 insertions, 16 deletions

diff --git a/jb/api/magic_token.py b/jb/api/magic_token.py
index 136e1b4..e0a1cca 100644
--- a/jb/api/magic_token.py
+++ b/jb/api/magic_token.py
@@ -4,15 +4,17 @@ import secrets
 from fastapi import HTTPException, status
 
 from jb.decorators import REDIS
+from jb.models.auth import AmtAccountLink, User
 
 MAGIC_TOKEN_PREFIX = "auth:magic:"
+AMT_ACCOUNT_LINK_TOKEN_PREFIX = "auth:amt-account-link:"
 MAGIC_TOKEN_TTL: int = 5 * 60  # 5 minutes, in seconds
 
 
-def redis_token_key(token: str) -> str:
+def redis_token_key(token: str, prefix: str = MAGIC_TOKEN_PREFIX) -> str:
     # Redis never contains a usable credential, even if its keys are exposed.
     digest = hashlib.sha256(token.encode("utf-8")).hexdigest()
-    return f"{MAGIC_TOKEN_PREFIX}{digest}"
+    return f"{prefix}{digest}"
 
 
 def create_magic_token(user_email: str) -> str:
@@ -39,3 +41,31 @@ def consume_magic_token(token: str) -> str:
             detail="Invalid or expired magic token",
         )
     return user_email
+
+
+def create_amt_account_link_token(user: User, amt_worker_id: str) -> str:
+    """Bind an email and AMT worker ID to an opaque, short-lived token."""
+    data = AmtAccountLink(
+        email=user.email,
+        amt_worker_id=amt_worker_id,
+    )
+    token = secrets.token_urlsafe(32)
+    REDIS.set(
+        redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX),
+        data.model_dump_json(),
+        ex=MAGIC_TOKEN_TTL,
+    )
+    return token
+
+
+def consume_amt_account_link_token(token: str) -> AmtAccountLink:
+    """Atomically consume and validate an AMT account-link token."""
+    raw_data = REDIS.getdel(
+        redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX)
+    )
+    if raw_data is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired account-link token",
+        )
+    return AmtAccountLink.model_validate_json(raw_data)
diff --git a/jb/managers/gr_api.py b/jb/managers/gr_api.py
index 8ffeaf0..a31f7da 100644
--- a/jb/managers/gr_api.py
+++ b/jb/managers/gr_api.py
@@ -1,16 +1,23 @@
 """Client for General Research's product-user API."""
 
+import logging
 from typing import Any
 
 import requests
 
 from jb.models.auth import User
 
+logger = logging.getLogger(__name__)
+
 
 class GRApiError(RuntimeError):
     """The General Research API could not satisfy a request."""
 
 
+class GRApiNotFoundError(GRApiError):
+    """The requested General Research API resource does not exist."""
+
+
 class GRApiManager:
     def __init__(
         self,
@@ -40,6 +47,14 @@ class GRApiManager:
             response = self.session.request(method, url, timeout=self.timeout, **kwargs)
             response.raise_for_status()
             return response
+        except requests.HTTPError as exc:
+            if exc.response is not None and exc.response.status_code == 404:
+                raise GRApiNotFoundError(
+                    f"General Research API resource not found: {method} {url}"
+                ) from exc
+            raise GRApiError(
+                f"General Research API request failed: {method} {url}"
+            ) from exc
         except requests.RequestException as exc:
             raise GRApiError(
                 f"General Research API request failed: {method} {url}"
@@ -101,12 +116,46 @@ class GRApiManager:
         ).json()
         return self.get_user(user.product_user_id)
 
-    def transition_product_user_id(self, user: User, amt_worker_id: str) -> User:
+    def transition_user_from_amt(self, user: User, amt_worker_id: str) -> User:
         """This should only be called once upon transition from an
         AMT account to a General Research account."""
         url = f"{self.base_url}/{self.product_id}/user/{amt_worker_id}/"
-        res = self._request(
-            "PATCH", url, json={"product_user_id": user.product_user_id}
-        )
+        try:
+            self._request(
+                "PATCH", url, json={"product_user_id": user.product_user_id}
+            )
+        except GRApiNotFoundError as exc:
+            raise ValueError(f"User {amt_worker_id} does not exist") from exc
+        except GRApiError as exc:
+            http_error = exc.__cause__
+            response = (
+                http_error.response
+                if isinstance(http_error, requests.HTTPError)
+                else None
+            )
+            if response is not None and response.status_code == 400:
+                try:
+                    detail = response.json().get("detail")
+                except (ValueError, AttributeError):
+                    detail = None
+                if detail == "Unable to update User":
+                    raise ValueError(
+                        "unable to update user, probably another user already "
+                        "exists with this email"
+                    ) from exc
+            raise
         self.set_user_email(user)
-        return self.get_user(user.product_user_id)
+        transitioned_user = self.get_user(user.product_user_id)
+        logger.warning(
+            "Transitioned product user from AMT worker %s to %s with email %s",
+            amt_worker_id,
+            transitioned_user.product_user_id,
+            transitioned_user.email,
+            extra={
+                "event": "transition_user_from_amt",
+                "amt_worker_id": amt_worker_id,
+                "product_user_id": transitioned_user.product_user_id,
+                "email": str(transitioned_user.email),
+            },
+        )
+        return transitioned_user
diff --git a/jb/models/auth.py b/jb/models/auth.py
index 8b48088..d0e8386 100644
--- a/jb/models/auth.py
+++ b/jb/models/auth.py
@@ -68,7 +68,9 @@ class User(BaseModel):
         if not isinstance(provided_id, str) or not hmac.compare_digest(
             provided_id, expected_id
         ):
-            raise ValueError(f"product_user_id {provided_id} does not match email {email}")
+            raise ValueError(
+                f"product_user_id {provided_id} does not match email {email}"
+            )
 
         # The computed field is authoritative; do not retain the input value.
         validated_data = dict(data)
@@ -92,6 +94,13 @@ class MagicLinkExchangeRequest(BaseModel):
     token: str = Field(min_length=1)
 
 
+class AmtAccountLink(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    email: EmailStr
+    amt_worker_id: str = Field(min_length=3, max_length=50)
+
+
 class SessionResponse(BaseModel):
     session_token: str
     token_type: str = "bearer"
diff --git a/jb/views/auth.py b/jb/views/auth.py
index 33c9452..dc63ca3 100644
--- a/jb/views/auth.py
+++ b/jb/views/auth.py
@@ -1,10 +1,3 @@
-"""Redis-backed magic-link authentication.
-
-The user service is deliberately not coupled to this module. Once that service
-has resolved an email address to its stable user identifier, call
-``create_magic_token`` and put the returned token in the emailed login URL.
-"""
-
 from typing import Annotated
 from urllib.parse import urlencode
 
@@ -16,12 +9,18 @@ from jb.api.auth import (
     create_session,
     get_authenticated_user,
 )
-from jb.api.magic_token import consume_magic_token, create_magic_token
+from jb.api.magic_token import (
+    consume_amt_account_link_token,
+    consume_magic_token,
+    create_amt_account_link_token,
+    create_magic_token,
+)
 from jb.config import settings
 from jb.dependencies import get_gr_api_manager
 from jb.managers.gr_api import GRApiManager
 from jb.models.auth import (
     AccountLogin,
+    AmtAccountLink,
     MagicLinkExchangeRequest,
     User,
 )
@@ -44,6 +43,20 @@ def request_mock_magic_link(body: AccountLogin) -> dict[str, str]:
     return {"magic_link": f"/auth/magic-link/?{query}"}
 
 
+@auth_router.post("/link-amt/request")
+def link_amt_account(body: AmtAccountLink) -> dict[str, str]:
+    """Create a mock AMT account-link email in development."""
+    if not settings.debug:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+
+    # TODO: Derive amt_worker_id from a server-validated AMT assignment and
+    # send this link by email instead of returning it.
+    user = User(email=body.email)
+    token = create_amt_account_link_token(user, body.amt_worker_id)
+    query = urlencode({"token": token})
+    return {"magic_link": f"/auth/link-amt/?{query}"}
+
+
 @auth_router.get("/magic-link/", response_class=HTMLResponse, include_in_schema=False)
 def magic_link_landing_page() -> HTMLResponse:
     """Serve the SPA without redeeming the token; email prefetches are harmless."""
@@ -57,6 +70,12 @@ def magic_link_landing_page() -> HTMLResponse:
     )
 
 
+@auth_router.get("/link-amt/", response_class=HTMLResponse, include_in_schema=False)
+def link_amt_account_landing_page() -> HTMLResponse:
+    """Serve the account-link SPA without consuming the one-time token."""
+    return magic_link_landing_page()
+
+
 @auth_router.post("/magic-link/exchange", status_code=status.HTTP_204_NO_CONTENT)
 def exchange_magic_link(
     body: MagicLinkExchangeRequest,
@@ -82,6 +101,29 @@ def exchange_magic_link(
     )
 
 
+@auth_router.post("/link-amt/exchange", status_code=status.HTTP_204_NO_CONTENT)
+def exchange_amt_account_link(
+    body: MagicLinkExchangeRequest,
+    response: Response,
+    gr_api: Annotated[GRApiManager, Depends(get_gr_api_manager)],
+) -> None:
+    """Validate the email link, then transition the bound AMT account."""
+    token_data = consume_amt_account_link_token(body.token)
+    user = User(email=token_data.email)
+    user = gr_api.transition_user_from_amt(user, token_data.amt_worker_id)
+
+    session_token = create_session(user.product_user_id)
+    response.set_cookie(
+        key=SESSION_COOKIE_NAME,
+        value=session_token,
+        max_age=settings.session_token_ttl_seconds,
+        httponly=True,
+        secure=not settings.debug,
+        samesite="lax",
+        path="/",
+    )
+
+
 @auth_router.get("/session", response_model=User)
 def get_session(
     user: Annotated[User, Depends(get_authenticated_user)],