TICKET-611: Autotester - ai-telemetry-gateway connection

client/autotest_client/__init__.py

@@ -242,6 +242,7 @@ def run_tests(settings_id, user):
 
     test_data = request.json["test_data"]
     categories = request.json["categories"]
+    batch_id = request.json.get("batch_id")
     high_priority = request.json.get("request_high_priority")
     queue_name = "batch" if len(test_data) > 1 else ("high" if high_priority else "low")
     queue = rq.Queue(queue_name, connection=REDIS_CONNECTION)
@@ -264,6 +265,7 @@ def run_tests(settings_id, user):
             "test_id": id_,
             "files_url": url,
             "categories": categories,
+            "batch_id": batch_id,
             "user": user,
             "test_env_vars": test_env_vars,
         }

server/autotest_server/__init__.py

@@ -223,6 +223,8 @@ def _run_test_specs(
     test_username: str,
     test_id: Union[int, str],
     test_env_vars: Dict[str, str],
+    files_url: Optional[str] = None,
+    batch_id: Optional[int] = None,
 ) -> List[ResultData]:
     """
     Run each test script in test_scripts in the tests_path directory using the
@@ -265,7 +267,11 @@ def _run_test_specs(
                         executable="/bin/bash",
                     )
                     try:
-                        settings_json = json.dumps({**settings, "test_data": test_data})
+                        # Raw attribution for telemetry-aware testers (e.g. AI).
+                        # The tester parses files_url into the spec fields; other
+                        # testers ignore this key.
+                        attribution = {"files_url": files_url, "categories": categories, "batch_id": batch_id}
+                        settings_json = json.dumps({**settings, "test_data": test_data, "_attribution": attribution})
                         out, err = proc.communicate(input=settings_json, timeout=timeout)
                     except subprocess.TimeoutExpired:
                         if test_username != getpass.getuser():
@@ -388,7 +394,7 @@ def tester_user() -> Tuple[str, str]:
     return user_name, user_workspace
 
 
-def run_test(settings_id, test_id, files_url, categories, user, test_env_vars):
+def run_test(settings_id, test_id, files_url, categories, user, test_env_vars, batch_id=None):
     results = []
     error = None
     try:
@@ -405,7 +411,9 @@ def run_test(settings_id, test_id, files_url, categories, user, test_env_vars):
             _clear_working_directory(tests_path, test_username)
             _setup_files(settings_id, user, files_url, tests_path, test_username)
             cmd = run_test_command(test_username=test_username)
-            results = _run_test_specs(cmd, settings, categories, tests_path, test_username, test_id, test_env_vars)
+            results = _run_test_specs(
+                cmd, settings, categories, tests_path, test_username, test_id, test_env_vars, files_url, batch_id
+            )
         finally:
             _stop_tester_processes(test_username)
             _clear_working_directory(tests_path, test_username)

server/autotest_server/settings.yml

@@ -5,6 +5,8 @@ worker_log_dir: !ENV ${WORKER_LOG_DIR}
 default_remote_url: https://polymouth.teach.cs.toronto.edu:443/chat
 remote_url_whitelist:
   - https://polymouth.teach.cs.toronto.edu:443/chat
+  - http://localhost:4000/v1
+  - http://host.docker.internal:4000/v1
 max_test_timeout: 3600
 workers:
   - user: !ENV ${USER}

server/autotest_server/testers/ai/ai_tester.py

@@ -1,15 +1,69 @@
 import json
 import os
+import re
 import sys
 
 from ..tester import Test, Tester
 from ..specs import TestSpecs
 import subprocess
-from typing import Type
+from typing import Optional, Type
+from urllib.parse import urlsplit
 from dotenv import load_dotenv
 from pathlib import Path
 import PyPDF2
 
+# Models the AI tester is allowed to invoke. "remote" routes to the
+# markus-ai-server proxy (local models); "openai-remote" routes to the
+# self-hosted LiteLLM gateway (cloud OpenAI, telemetry + budget enforced).
+# Any other model would reach a cloud provider untracked, so it is rejected.
+ALLOWED_MODELS = ("remote", "openai-remote")
+
+# Categories ordered most-privileged first. A job's categories array collapses
+# to a single role for telemetry; when more than one is present the most
+# privileged wins (see ai-telemetry-gateway decision-record §4).
+_ROLE_PRIORITY = ("instructor", "student")
+
+# The four spec attribution fields live inside the MarkUs file_url, shaped
+# <instance>/api/courses/<course_id>/assignments/<assignment_id>/groups/<group_id>/submission_files
+_FILE_URL_RE = re.compile(
+    r"/api/courses/(?P<course_id>\d+)"
+    r"/assignments/(?P<assignment_id>\d+)"
+    r"/groups/(?P<group_id>\d+)/submission_files"
+)
+
+
+def _resolve_category(categories: list) -> Optional[str]:
+    """Collapse the categories array to one role; most-privileged wins."""
+    for role in _ROLE_PRIORITY:
+        if role in categories:
+            return role
+    return categories[0] if categories else None
+
+
+def build_spend_metadata(files_url: Optional[str], categories: Optional[list], batch_id) -> dict:
+    """Extract the six attribution fields the gateway records on every call.
+
+    Raises ValueError when ``files_url`` does not match the expected MarkUs
+    shape. A malformed URL is a programming error in MarkUs, not a recoverable
+    runtime condition, so we fail loud rather than emit a ledger row with
+    nonsense identifiers.
+    """
+    parts = urlsplit(files_url or "")
+    match = _FILE_URL_RE.search(parts.path)
+    if not parts.netloc or match is None:
+        raise ValueError(
+            f"Cannot extract attribution from file_url {files_url!r}: expected "
+            ".../api/courses/<id>/assignments/<id>/groups/<id>/submission_files"
+        )
+    return {
+        "instance": parts.netloc,
+        "course_id": int(match["course_id"]),
+        "assignment_id": int(match["assignment_id"]),
+        "group_id": int(match["group_id"]),
+        "batch_id": batch_id,
+        "category": _resolve_category(categories or []),
+    }
+
 
 class AiTest(Test):
     def __init__(
@@ -72,12 +126,14 @@ def call_ai_feedback(self) -> dict:
         output_mode = test_group.get("output")
         cmd = [sys.executable, "-m", "ai_feedback"]
 
-        # Restrict to remote model only — prevent access to cloud AIs
-        if config.get("model", "") != "remote":
+        # Restrict to the gateway-fronted models — prevent untracked cloud access.
+        model = config.get("model", "")
+        if model not in ALLOWED_MODELS:
+            allowed = ", ".join(f"'{m}'" for m in ALLOWED_MODELS)
             results[test_label] = {
                 "title": test_label,
                 "status": "error",
-                "message": f"Unsupported model type: \"{config.get('model', '')}\". Only 'remote' model is allowed.",
+                "message": f'Unsupported model type: "{model}". Allowed models: {allowed}.',
             }
             return results
 
@@ -101,6 +157,23 @@ def call_ai_feedback(self) -> dict:
             }
             return results
 
+        # Gateway-bound calls carry MarkUs attribution so the ledger can record
+        # per-course spend and the gatekeeper can enforce budgets. The worker
+        # injects the raw pieces under "_attribution"; we parse them here, at
+        # the point the AI feedback library is about to be invoked.
+        if model == "openai-remote":
+            attribution = self.specs.get("_attribution", default={})
+            try:
+                metadata = build_spend_metadata(
+                    attribution.get("files_url"),
+                    attribution.get("categories"),
+                    attribution.get("batch_id"),
+                )
+            except ValueError as ve:
+                results[test_label] = {"title": test_label, "status": "error", "message": str(ve)}
+                return results
+            env["LITELLM_SPEND_METADATA"] = json.dumps(metadata)
+
         for key, value in config.items():
             cmd.extend(["--" + key, str(value)])
 

server/autotest_server/tests/testers/ai/test_ai_tester.py

@@ -4,7 +4,7 @@
 
 import subprocess
 import pytest
-from ....testers.ai.ai_tester import AiTester, AiTest
+from ....testers.ai.ai_tester import AiTester, AiTest, build_spend_metadata
 from ....testers.specs import TestSpecs
 
 DEFAULT_REMOTE_URL = "https://polymouth.teach.cs.toronto.edu:443/chat"
@@ -20,7 +20,10 @@ def set_required_env():
     os.makedirs(os.environ["WORKER_LOG_DIR"], exist_ok=True)
 
 
-def _make_spec(output="overall_comment", config_overrides=None):
+GATEWAY_FILE_URL = "https://markus.example.edu/api/courses/12/assignments/34/groups/56/submission_files?collected=true"
+
+
+def _make_spec(output="overall_comment", config_overrides=None, attribution=None):
     """Build a test spec with sensible defaults. Override only what varies."""
     config = {
         "model": "remote",
@@ -32,7 +35,7 @@ def _make_spec(output="overall_comment", config_overrides=None):
     }
     if config_overrides:
         config.update(config_overrides)
-    return {
+    spec = {
         "tester_type": "ai",
         "env_data": {"ai_feedback_version": "main"},
         "test_data": {
@@ -50,17 +53,35 @@ def _make_spec(output="overall_comment", config_overrides=None):
         },
         "_env": {"PYTHON": "/home/docker/.autotesting/scripts/128/ai_1/bin/python3"},
     }
+    if attribution is not None:
+        # Mirrors what the worker injects into the piped tester JSON.
+        spec["_attribution"] = attribution
+    return spec
 
 
 def _make_tester(**kwargs):
     return AiTester(specs=TestSpecs.from_json(json.dumps(_make_spec(**kwargs))))
 
 
-def _mock_subprocess(monkeypatch, *, stdout="OK", stderr=""):
+def _capture_subprocess(monkeypatch, *, stdout="OK", stderr=""):
+    """Mock subprocess.run with a successful result, recording the call's kwargs."""
+    captured = {}
     mocked = subprocess.CompletedProcess(
         args=["python", "-m", "ai_feedback"], returncode=0, stdout=stdout, stderr=stderr
     )
-    monkeypatch.setattr(subprocess, "run", lambda *a, **kw: mocked)
+
+    def fake_run(*a, **kw):
+        captured["cmd"] = a[0] if a else kw.get("args")
+        captured["env"] = kw.get("env")
+        return mocked
+
+    monkeypatch.setattr(subprocess, "run", fake_run)
+    return captured
+
+
+def _mock_subprocess(monkeypatch, *, stdout="OK", stderr=""):
+    """Mock subprocess.run with a successful result (call kwargs discarded)."""
+    _capture_subprocess(monkeypatch, stdout=stdout, stderr=stderr)
 
 
 @pytest.mark.parametrize(
@@ -125,6 +146,89 @@ def test_rejects_non_remote_model():
     assert "openai" in results["Test A"]["message"]
 
 
+def test_build_spend_metadata_parses_all_fields():
+    metadata = build_spend_metadata(GATEWAY_FILE_URL, ["student"], 99)
+    assert metadata == {
+        "instance": "markus.example.edu",
+        "course_id": 12,
+        "assignment_id": 34,
+        "group_id": 56,
+        "batch_id": 99,
+        "category": "student",
+    }
+
+
+def test_build_spend_metadata_tolerates_relative_url_root():
+    url = "https://example.edu/markus/api/courses/1/assignments/2/groups/3/submission_files"
+    metadata = build_spend_metadata(url, ["instructor"], None)
+    assert metadata["instance"] == "example.edu"
+    assert (metadata["course_id"], metadata["assignment_id"], metadata["group_id"]) == (1, 2, 3)
+
+
+def test_build_spend_metadata_most_privileged_category_wins():
+    metadata = build_spend_metadata(GATEWAY_FILE_URL, ["student", "instructor"], None)
+    assert metadata["category"] == "instructor"
+
+
+@pytest.mark.parametrize(
+    "bad_url",
+    [
+        "",
+        "not a url",
+        "/api/courses/1/assignments/2/groups/3/submission_files",  # no host
+        "https://m.edu/api/courses/1/assignments/2/submission_files",  # no groups segment
+    ],
+)
+def test_build_spend_metadata_rejects_malformed_url(bad_url):
+    with pytest.raises(ValueError, match="Cannot extract attribution"):
+        build_spend_metadata(bad_url, ["student"], None)
+
+
+def test_allows_openai_remote_model(monkeypatch):
+    tester = _make_tester(
+        config_overrides={"model": "openai-remote", "remote_url": "http://gateway:4000/v1"},
+        attribution={"files_url": GATEWAY_FILE_URL, "categories": ["student"], "batch_id": None},
+    )
+    _mock_subprocess(monkeypatch, stdout="Great job!")
+    results = tester.call_ai_feedback()
+    assert results["Test A"]["status"] == "success"
+
+
+def test_openai_remote_threads_metadata_env(monkeypatch):
+    tester = _make_tester(
+        config_overrides={"model": "openai-remote", "remote_url": "http://gateway:4000/v1"},
+        attribution={"files_url": GATEWAY_FILE_URL, "categories": ["student"], "batch_id": 7},
+    )
+    captured = _capture_subprocess(monkeypatch, stdout="Great job!")
+    tester.call_ai_feedback()
+    metadata = json.loads(captured["env"]["LITELLM_SPEND_METADATA"])
+    assert metadata == {
+        "instance": "markus.example.edu",
+        "course_id": 12,
+        "assignment_id": 34,
+        "group_id": 56,
+        "batch_id": 7,
+        "category": "student",
+    }
+
+
+def test_openai_remote_malformed_url_errors_without_calling(monkeypatch):
+    bad = {"files_url": "https://m.edu/api/courses/1/submission_files", "categories": ["student"], "batch_id": None}
+    tester = _make_tester(config_overrides={"model": "openai-remote"}, attribution=bad)
+    captured = _capture_subprocess(monkeypatch)
+    results = tester.call_ai_feedback()
+    assert results["Test A"]["status"] == "error"
+    assert "Cannot extract attribution" in results["Test A"]["message"]
+    assert captured == {}  # the AI feedback subprocess was never invoked
+
+
+def test_remote_model_does_not_set_metadata_env(monkeypatch):
+    tester = _make_tester()  # default model "remote"
+    captured = _capture_subprocess(monkeypatch, stdout="ok")
+    tester.call_ai_feedback()
+    assert "LITELLM_SPEND_METADATA" not in (captured["env"] or {})
+
+
 def test_missing_submission_file():
     tester = _make_tester(config_overrides={"submission": FIXTURES_DIR + "/nonexistent.py"})
     results = tester.call_ai_feedback()