[codex] Implement KPI dataset v2 instrumentation

agent/core/session.py

@@ -560,6 +560,7 @@ def get_trajectory(self) -> dict:
             "session_id": self.session_id,
             "user_id": self.user_id,
             "hf_username": self.hf_username,
+            "user_plan": self.user_plan,
             "session_start_time": self.session_start_time,
             "session_end_time": datetime.now().isoformat(),
             "model_name": self.config.model_name,

agent/core/session_uploader.py

@@ -290,6 +290,7 @@ def _write_row_payload(data: dict, tmp_path: str) -> None:
     session_row = {
         "session_id": data["session_id"],
         "user_id": data.get("user_id"),
+        "user_plan": data.get("user_plan") or "unknown",
         "session_start_time": data["session_start_time"],
         "session_end_time": data["session_end_time"],
         "model_name": data["model_name"],

agent/core/telemetry.py

@@ -17,7 +17,10 @@
 from __future__ import annotations
 
 import asyncio
+import hashlib
+import hmac
 import logging
+import os
 import time
 from typing import Any
 
@@ -150,6 +153,92 @@ def _infer_push_to_hub(script_or_cmd: Any) -> bool:
     )
 
 
+def _kpi_hash_identifier(value: Any) -> str | None:
+    salt = os.environ.get("KPI_USER_HASH_SALT")
+    if not salt or value is None:
+        return None
+    raw = str(value)
+    if not raw:
+        return None
+    return hmac.new(
+        salt.encode("utf-8"), raw.encode("utf-8"), hashlib.sha256
+    ).hexdigest()
+
+
+def _sanitize_expected_hub_artifacts(artifacts: Any) -> list[dict[str, Any]]:
+    sanitized: list[dict[str, Any]] = []
+    if not isinstance(artifacts, list):
+        return sanitized
+    for artifact in artifacts:
+        if not isinstance(artifact, dict):
+            continue
+        repo_type = str(artifact.get("repo_type") or "model").strip().lower()
+        repo_id = artifact.get("repo_id")
+        if repo_type not in {"model", "dataset", "space"} or not repo_id:
+            continue
+        artifact_hash = _kpi_hash_identifier(f"{repo_type}:{repo_id}")
+        if artifact_hash is None:
+            continue
+        sanitized.append(
+            {
+                "repo_type": repo_type,
+                "artifact_hash": artifact_hash,
+                "private": artifact.get("private"),
+                "is_sandbox": bool(artifact.get("is_sandbox")),
+            }
+        )
+    return sanitized
+
+
+async def record_hub_artifact(
+    session: Any,
+    *,
+    repo_type: str,
+    repo_id: str,
+    source: str,
+    is_sandbox: bool | None = None,
+    private: bool | None = None,
+    success: bool = True,
+) -> dict[str, Any]:
+    """Emit a sanitized Hub artifact event.
+
+    The raw repo id never leaves this function. Consumers dedupe artifacts by
+    ``artifact_hash``.
+    """
+    from agent.core.session import Event
+
+    try:
+        normalized_type = str(repo_type or "").strip().lower()
+        if normalized_type not in {"model", "dataset", "space"}:
+            return {}
+        if is_sandbox is None:
+            try:
+                from agent.core.hub_artifacts import is_sandbox_hub_repo
+
+                is_sandbox = is_sandbox_hub_repo(repo_id, normalized_type)
+            except Exception:
+                is_sandbox = False
+        artifact_hash = _kpi_hash_identifier(f"{normalized_type}:{repo_id}")
+        if artifact_hash is None:
+            logger.debug(
+                "record_hub_artifact skipped because KPI_USER_HASH_SALT is unset"
+            )
+            return {}
+        payload = {
+            "repo_type": normalized_type,
+            "artifact_hash": artifact_hash,
+            "source": str(source or "unknown"),
+            "is_sandbox": bool(is_sandbox),
+            "private": private,
+            "success": bool(success),
+        }
+        await session.send_event(Event(event_type="hub_artifact", data=payload))
+        return payload
+    except Exception as e:
+        logger.debug("record_hub_artifact failed (non-fatal): %s", e)
+    return {}
+
+
 async def record_hf_job_submit(
     session: Any,
     job: Any,
@@ -165,6 +254,9 @@ async def record_hf_job_submit(
     t_start = time.monotonic()
     try:
         script_text = args.get("script") or args.get("command") or ""
+        expected_artifacts = _sanitize_expected_hub_artifacts(
+            args.get("expected_hub_artifacts")
+        )
         await session.send_event(
             Event(
                 event_type="hf_job_submit",
@@ -177,6 +269,8 @@ async def record_hf_job_submit(
                     "image": image,
                     "namespace": args.get("namespace"),
                     "push_to_hub": _infer_push_to_hub(script_text),
+                    "expected_hub_artifacts": expected_artifacts,
+                    "expected_hub_artifacts_count": len(expected_artifacts),
                 },
             )
         )
@@ -231,6 +325,34 @@ async def record_hf_job_complete(
     return {}
 
 
+async def record_hf_job_cancel(
+    session: Any,
+    *,
+    job_id: str,
+    namespace: str | None = None,
+    flavor: str | None = None,
+) -> dict:
+    from agent.core.session import Event
+
+    try:
+        payload = {
+            "job_id": job_id,
+            "namespace": namespace,
+            "flavor": flavor,
+            "final_status": "cancelled",
+            "wall_time_s": 0,
+            "billable_seconds_estimate": 0,
+            "price_usd_per_hour": None,
+            "estimated_cost_usd": None,
+            "cost_estimate_source": "manual_cancel",
+        }
+        await session.send_event(Event(event_type="hf_job_complete", data=payload))
+        return payload
+    except Exception as e:
+        logger.debug("record_hf_job_cancel failed (non-fatal): %s", e)
+    return {}
+
+
 # ── sandbox ─────────────────────────────────────────────────────────────────
 
 

agent/tools/hf_repo_git_tool.py

@@ -548,6 +548,17 @@ async def _create_repo(self, args: Dict[str, Any]) -> ToolResult:
             session=self.session,
             extra_metadata=extra_metadata,
         )
+        if self.session:
+            from agent.core import telemetry
+
+            await telemetry.record_hub_artifact(
+                self.session,
+                repo_type=repo_type,
+                repo_id=repo_id,
+                source="hf_repo_git",
+                private=private,
+                success=True,
+            )
 
         return {
             "formatted": f"**Repository created:** {repo_id}\n**Private:** {private}\n{result}",

agent/tools/jobs_tool.py

@@ -148,6 +148,43 @@ def _add_environment_variables(
     return result
 
 
+def _normalize_expected_hub_artifacts(raw: Any) -> list[dict[str, Any]]:
+    artifacts: list[dict[str, Any]] = []
+    if not isinstance(raw, list):
+        return artifacts
+    for item in raw:
+        if not isinstance(item, dict):
+            continue
+        repo_type = str(item.get("repo_type") or "model").strip().lower()
+        repo_id = (
+            item.get("repo_id")
+            or item.get("hub_model_id")
+            or item.get("hub_dataset_id")
+            or item.get("space_id")
+        )
+        if repo_type not in {"model", "dataset", "space"} or not repo_id:
+            continue
+        artifacts.append(
+            {
+                "repo_type": repo_type,
+                "repo_id": str(repo_id),
+                "private": item.get("private"),
+                "is_sandbox": bool(item.get("is_sandbox")),
+            }
+        )
+    return artifacts
+
+
+def _job_status_completed(status: Any) -> bool:
+    return str(status or "").strip().upper() in {
+        "COMPLETED",
+        "COMPLETE",
+        "SUCCEEDED",
+        "SUCCESS",
+        "DONE",
+    }
+
+
 def _build_uv_command(
     script: str,
     with_deps: list[str] | None = None,
@@ -389,7 +426,7 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult:
                 "isError": True,
             }
 
-    async def _seed_trackio_dashboard(self, space_id: str) -> None:
+    async def _seed_trackio_dashboard(self, space_id: str) -> bool:
         """Idempotently install trackio dashboard files into *space_id* before
         the job runs. Surfaces seed progress as tool_log events but never
         raises — a seed failure should not block job submission, since trackio
@@ -410,9 +447,41 @@ def _log(msg: str) -> None:
             await asyncio.to_thread(
                 ensure_trackio_dashboard, space_id, self.hf_token, _log
             )
+            if self.session:
+                from agent.core import telemetry
+
+                await telemetry.record_hub_artifact(
+                    self.session,
+                    repo_type="space",
+                    repo_id=space_id,
+                    source="trackio",
+                    is_sandbox=False,
+                    private=True,
+                    success=True,
+                )
+            return True
         except Exception as e:
             logger.warning(f"trackio dashboard seed failed for {space_id}: {e}")
             _log(f"trackio dashboard seed failed: {e}")
+        return False
+
+    async def _record_successful_job_artifacts(
+        self, artifacts: list[dict[str, Any]]
+    ) -> None:
+        if not self.session:
+            return
+        from agent.core import telemetry
+
+        for artifact in artifacts:
+            await telemetry.record_hub_artifact(
+                self.session,
+                repo_type=artifact["repo_type"],
+                repo_id=artifact["repo_id"],
+                source="hf_job",
+                is_sandbox=artifact["is_sandbox"],
+                private=artifact.get("private"),
+                success=True,
+            )
 
     async def _wait_for_job_completion(
         self, job_id: str, namespace: Optional[str] = None
@@ -569,6 +638,9 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult:
             # Run the job
             flavor = args.get("hardware_flavor", "cpu-basic")
             timeout_str = args.get("timeout", "30m")
+            expected_hub_artifacts = _normalize_expected_hub_artifacts(
+                args.get("expected_hub_artifacts")
+            )
 
             # Trackio: agent-declared space + project become env vars on the job
             # so trackio.init() picks them up automatically. We also surface them
@@ -658,6 +730,7 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult:
                         "hardware_flavor": flavor,
                         "timeout": timeout_str,
                         "namespace": self.namespace,
+                        "expected_hub_artifacts": expected_hub_artifacts,
                     },
                     image=image,
                     job_type=job_type,
@@ -711,6 +784,8 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult:
                         else None,
                         allow_zero_actual=True,
                     )
+                if _job_status_completed(final_status):
+                    await self._record_successful_job_artifacts(expected_hub_artifacts)
 
             # Untrack job ID (completed or failed, no longer needs cancellation)
             if self.session:
@@ -880,6 +955,14 @@ async def _cancel_job(self, args: Dict[str, Any]) -> ToolResult:
             job_id=job_id,
             namespace=self.namespace,
         )
+        if self.session:
+            from agent.core import telemetry
+
+            await telemetry.record_hf_job_cancel(
+                self.session,
+                job_id=job_id,
+                namespace=self.namespace,
+            )
 
         response = f"""✓ Job {job_id} has been cancelled.
 
@@ -1249,6 +1332,26 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult:
                     "the embedded dashboard to this project."
                 ),
             },
+            "expected_hub_artifacts": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "repo_type": {
+                            "type": "string",
+                            "enum": ["model", "dataset", "space"],
+                        },
+                        "repo_id": {"type": "string"},
+                        "private": {"type": "boolean"},
+                        "is_sandbox": {"type": "boolean"},
+                    },
+                    "required": ["repo_type", "repo_id"],
+                },
+                "description": (
+                    "Optional. Hub artifacts the job is expected to create. "
+                    "They are only counted after the job completes successfully."
+                ),
+            },
             "namespace": {
                 "type": "string",
                 "description": (

agent/tools/sandbox_tool.py

@@ -289,6 +289,17 @@ def _log(msg: str) -> None:
         await asyncio.to_thread(
             ensure_trackio_dashboard, space_id, session.hf_token, _log
         )
+        from agent.core import telemetry
+
+        await telemetry.record_hub_artifact(
+            session,
+            repo_type="space",
+            repo_id=space_id,
+            source="trackio",
+            is_sandbox=False,
+            private=True,
+            success=True,
+        )
     except Exception as e:
         _log(f"trackio dashboard seed failed: {e}")