Add applied_permission_level from the model claim; remove unbacked JW…

tests/recipes/test_recipes.py

@@ -149,6 +149,7 @@ def test_recipe_by_production_semantic_version(mocker):
             "production_version_id": "production-version-id"
         }
     )
+    mocker.patch("wrangles.data.model_claim", return_value={})
     model_content = mocker.patch(
         "wrangles.data.model_content",
         return_value={"recipe": "{}"}
@@ -173,6 +174,7 @@ def test_recipe_by_production_semantic_version_falls_back_to_latest(
         "wrangles.data.model",
         return_value={"purpose": "recipe"}
     )
+    mocker.patch("wrangles.data.model_claim", return_value={})
     model_content = mocker.patch(
         "wrangles.data.model_content",
         return_value={"recipe": "{}"}

tests/recipes/test_variables.py

@@ -367,25 +367,21 @@ def test_variables_variable_overwrite():
     assert isinstance(df['vars'][0], dict)
 
 
-def test_applied_permission_group_variable(monkeypatch):
+def test_applied_permission_group_variable_explicit(monkeypatch):
     """
-    Test that the authenticated user's effective permission group is available as a recipe variable.
+    Test that an explicitly-passed applied_permission_group is available as
+    a recipe variable for a non-model_id recipe (there is no server-side
+    default to fall back to in that case).
     """
-    token = wrangles.auth._jwt.encode(
-        {"applied_permission_group": "enterprise"},
-        "test-secret",
-        algorithm="HS256"
-    )
-    monkeypatch.setattr(wrangles.auth, "get_access_token", lambda: token)
-
     df = wrangles.recipe.run(
         """
         read:
         - test:
             rows: 1
             values:
                 group: ${applied_permission_group}
-        """
+        """,
+        variables={"applied_permission_group": "enterprise"}
     )
 
     assert df['group'][0] == 'enterprise'
@@ -395,8 +391,6 @@ def test_applied_permission_group_variable_if(monkeypatch):
     """
     Test that applied_permission_group can be used in Python-style if conditions.
     """
-    monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "enterprise")
-
     df = wrangles.recipe.run(
         """
         read:
@@ -409,46 +403,68 @@ def test_applied_permission_group_variable_if(monkeypatch):
             output: allowed
             value: true
             if: applied_permission_group == 'enterprise'
-        """
+        """,
+        variables={"applied_permission_group": "enterprise"}
     )
 
     assert df['allowed'][0] == True
 
 
-def test_applied_permission_group_variable_user_override(monkeypatch):
+def test_applied_permission_group_variable_from_recipe_metadata(monkeypatch):
     """
-    Test that explicit variables still override the authenticated permission group.
+    Test that recipe metadata permission group is preferred for remote recipes.
     """
-    monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "enterprise")
-
-    df = wrangles.recipe.run(
-        """
-        read:
-        - test:
-            rows: 1
-            values:
-                group: ${applied_permission_group}
-        """,
-        variables={"applied_permission_group": "manual"}
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model",
+        lambda model_id: {
+            "purpose": "recipe",
+            "production_version_id": "v1",
+            "applied_permission_group": "metadata-group",
+        }
     )
+    # No model claim available - the metadata-derived value above should be
+    # left untouched rather than overridden.
+    monkeypatch.setattr(wrangles.recipe._data, "model_claim", lambda model_id: {})
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_content",
+        lambda model_id, version_id=None: {
+            "recipe": """
+            read:
+            - test:
+                rows: 1
+                values:
+                    group: ${applied_permission_group}
+            """
+        }
+    )
+
+    df = wrangles.recipe.run("12345678-1234-1234")
 
-    assert df['group'][0] == 'manual'
+    assert df["group"][0] == "metadata-group"
 
 
-def test_applied_permission_group_variable_from_recipe_metadata(monkeypatch):
+def test_applied_permission_group_variable_metadata_overrides_explicit(monkeypatch, caplog):
     """
-    Test that recipe metadata permission group is preferred for remote recipes.
+    A model_id-addressed recipe's real permission group (resolved
+    server-side from the model's metadata) must override an explicit
+    variables={"applied_permission_group": ...} too - otherwise a caller
+    could simply claim a higher role than the model's database actually
+    grants them.
     """
-    monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "token-group")
     monkeypatch.setattr(
         wrangles.recipe._data,
         "model",
         lambda model_id: {
             "purpose": "recipe",
             "production_version_id": "v1",
-            "applied_permission_group": "metadata-group",
+            "applied_permission_group": "editor",
         }
     )
+    # No model claim available - only the metadata-derived override (from
+    # data.model above) is exercised by this test.
+    monkeypatch.setattr(wrangles.recipe._data, "model_claim", lambda model_id: {})
     monkeypatch.setattr(
         wrangles.recipe._data,
         "model_content",
@@ -463,6 +479,133 @@ def test_applied_permission_group_variable_from_recipe_metadata(monkeypatch):
         }
     )
 
+    with caplog.at_level("WARNING"):
+        df = wrangles.recipe.run(
+            "12345678-1234-1234",
+            variables={"applied_permission_group": "admin"}
+        )
+
+    assert df["group"][0] == "editor"
+    assert "does not match this model's actual permission group" in caplog.text
+
+
+def test_applied_permission_level_variable_from_model_claim(monkeypatch):
+    """
+    Test that applied_permission_level is filled from the model claim's role
+    when running a model_id directly, e.g. from Python.
+    """
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model",
+        lambda model_id: {"purpose": "recipe", "production_version_id": "v1"}
+    )
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_claim",
+        lambda model_id: {
+            "model_id": model_id,
+            "role": "viewer",
+            "applied_group": "Dev (WrangleWorks)",
+        }
+    )
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_content",
+        lambda model_id, version_id=None: {
+            "recipe": """
+            read:
+            - test:
+                rows: 1
+                values:
+                    level: ${applied_permission_level}
+                    group: ${applied_permission_group}
+            """
+        }
+    )
+
     df = wrangles.recipe.run("12345678-1234-1234")
 
-    assert df["group"][0] == "metadata-group"
+    assert df["level"][0] == "viewer"
+    assert df["group"][0] == "Dev (WrangleWorks)"
+
+
+def test_applied_permission_level_variable_claim_overrides_explicit(monkeypatch, caplog):
+    """
+    Like applied_permission_group, an explicit
+    variables={"applied_permission_level": ...} must not let a caller claim
+    a higher role than the model claim actually grants - run(model,
+    variables={"applied_permission_level": "admin"}) when the real claim
+    says "viewer" must use "viewer".
+    """
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model",
+        lambda model_id: {"purpose": "recipe", "production_version_id": "v1"}
+    )
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_claim",
+        lambda model_id: {
+            "model_id": model_id,
+            "role": "viewer",
+            "applied_group": "Dev (WrangleWorks)",
+        }
+    )
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_content",
+        lambda model_id, version_id=None: {
+            "recipe": """
+            read:
+            - test:
+                rows: 1
+                values:
+                    level: ${applied_permission_level}
+            """
+        }
+    )
+
+    with caplog.at_level("WARNING"):
+        df = wrangles.recipe.run(
+            "12345678-1234-1234",
+            variables={"applied_permission_level": "admin"}
+        )
+
+    assert df["level"][0] == "viewer"
+    assert "does not match this model's actual permission level" in caplog.text
+
+
+def test_model_claim_failure_does_not_block_recipe_load(monkeypatch, caplog):
+    """
+    A failure resolving the model claim (e.g. network issue) must not block
+    the recipe from loading - applied_permission_level is simply left unset.
+    """
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model",
+        lambda model_id: {"purpose": "recipe", "production_version_id": "v1"}
+    )
+
+    def _raise(model_id):
+        raise RuntimeError("boom")
+
+    monkeypatch.setattr(wrangles.recipe._data, "model_claim", _raise)
+    monkeypatch.setattr(
+        wrangles.recipe._data,
+        "model_content",
+        lambda model_id, version_id=None: {
+            "recipe": """
+            read:
+            - test:
+                rows: 1
+                values:
+                    result: kept
+            """
+        }
+    )
+
+    with caplog.at_level("WARNING"):
+        df = wrangles.recipe.run("12345678-1234-1234")
+
+    assert df["result"][0] == "kept"
+    assert "Could not resolve model claim" in caplog.text

tests/recipes/wrangles/test_extract.py

@@ -4413,8 +4413,11 @@ def test_ai_invalid_model_per_row_error(self):
                 "data": ["wrench 25mm", "6m cable"],
             })
         )
+        # OpenAI may report an unknown model as any 4xx client error
+        # (e.g. 400 invalid_request_error or 404 model_not_found)
+        # depending on the API version, so don't pin to one exact code.
         assert all(
-            "OpenAI API error" in value and "status=400" in value
+            "OpenAI API error" in value and "status=4" in value
             for value in df['length']
         )
 

tests/recipes/wrangles/test_main.py

@@ -4049,7 +4049,8 @@ def fake_model_content(model_id, version_id=None):
         """
 
         with patch('wrangles.recipe._data.model', side_effect=fake_model), \
-             patch('wrangles.recipe._data.model_content', side_effect=fake_model_content):
+             patch('wrangles.recipe._data.model_content', side_effect=fake_model_content), \
+             patch('wrangles.recipe._data.model_claim', return_value={}):
             with pytest.raises(Exception) as info:
                 wrangles.recipe.run(outer_recipe)
 

tests/test_data.py

@@ -27,6 +27,7 @@ def _mock_model_response(monkeypatch, response):
         lambda: data.model(MODEL_ID),
         lambda: data.model_update(MODEL_ID, {"name": "Updated model"}),
         lambda: data.model_content(MODEL_ID),
+        lambda: data.model_claim(MODEL_ID),
     ],
 )
 def test_model_endpoints_raise_authentication_error_for_401(monkeypatch, call_model_endpoint):
@@ -48,6 +49,7 @@ def test_model_endpoints_raise_authentication_error_for_401(monkeypatch, call_mo
         lambda: data.model(MODEL_ID),
         lambda: data.model_update(MODEL_ID, {"name": "Updated model"}),
         lambda: data.model_content(MODEL_ID),
+        lambda: data.model_claim(MODEL_ID),
     ],
 )
 def test_model_endpoints_raise_authorization_error_for_403(monkeypatch, call_model_endpoint):
@@ -80,3 +82,17 @@ def test_model_content_success_returns_content(monkeypatch):
     _mock_model_response(monkeypatch, FakeResponse(200, content))
 
     assert data.model_content(MODEL_ID) == content
+
+
+def test_model_claim_success_returns_claim(monkeypatch):
+    claim = {
+        "model_id": MODEL_ID,
+        "role": "admin",
+        "organization_id": "team-id",
+        "applied_group": "Dev (WrangleWorks)",
+        "applied_group_type": "group",
+        "applied_group_id": "team-id",
+    }
+    _mock_model_response(monkeypatch, FakeResponse(200, claim))
+
+    assert data.model_claim(MODEL_ID) == claim

wrangles/auth.py

@@ -88,29 +88,31 @@ def get_access_token():
 
 def extract_applied_permission_group(source: dict):
     """
-    Extract the effective permission group from a metadata or token payload.
+    Extract the effective permission group from a model metadata payload.
     """
     if not isinstance(source, dict):
         return None
 
     return source.get("applied_permission_group")
 
 
-def get_applied_permission_group():
+def extract_applied_permission_group_from_claim(claim: dict):
     """
-    Return the authenticated user's effective permission group from the current access token.
-
-    If no user is authenticated or the token does not contain the claim,
-    return None so recipes can still run without backend credentials.
+    Extract the applied permission group (the group/org/user display name a
+    model claim is granted through) from a /model/claim response.
     """
-    try:
-        token = get_access_token()
-    except Exception:
+    if not isinstance(claim, dict):
         return None
 
-    try:
-        claims = _jwt.decode(token, options={"verify_signature": False})
-    except Exception:
+    return claim.get("applied_group")
+
+
+def extract_applied_permission_level(claim: dict):
+    """
+    Extract the applied permission level (the user's role on a model - e.g.
+    admin, editor, viewer) from a /model/claim response.
+    """
+    if not isinstance(claim, dict):
         return None
 
-    return extract_applied_permission_group(claims)
+    return claim.get("role")