feat: improve typing throughout the codebase

AGENTS.md

@@ -4,11 +4,13 @@ This file guides agentic coding tools working in this repo.
 Keep it current when commands or conventions change.
 
 ## Scope
+
 - Project: Kinitro robotics evaluation subnet (Python 3.12).
 - Primary entry point: `kinitro` CLI (see `kinitro/cli/`).
 - Services: API, scheduler, executor, validator, miner tooling.
 
 ## Repo Map
+
 - `demos/` demonstration scripts and examples.
 - `docs/` operator, validator, and miner guides.
 - `environments/` evaluation environments (MetaWorld, Genesis).
@@ -27,44 +29,54 @@ Keep it current when commands or conventions change.
 - `kinitro/validator/` validator workflows.
 
 ### Key Modules
+
 - `kinitro/cli/` CLI entry point and command groups.
 - `kinitro/config.py` runtime settings and configuration.
 - `kinitro/crypto.py` cryptography helpers.
 - `kinitro/rl_interface.py` RL interface utilities (Observation, Action, ProprioKeys, ActionKeys).
 
 ## Setup / Install
+
 - Recommended: `uv sync`
 - Editable install: `pip install -e .`
 - Dev extras (if using pip): `pip install -e ".[dev]"`
 - Basilica CLI (for miner deployments): `curl -sSL https://basilica.ai/install.sh | bash`
 
 ## Common Commands
+
 ### Lint
+
 - `ruff check .` (checks entire project: kinitro/, tests/, environments/, scripts/, demos/)
 
 ### Type Check
+
 - `ty check .`
 
 ### Tests
+
 - `pytest tests/`
 - With MuJoCo: `MUJOCO_GL=egl pytest tests/`
 
 ### Single Test (examples)
+
 - `pytest tests/unit/test_pareto.py::TestEpsilonDominates::test_clear_dominance`
 - `pytest tests/unit/test_pareto.py -k test_clear_dominance`
 - `pytest -k pareto tests/unit`
 
 ### CLI Examples
+
 - List environments: `uv run kinitro env list`
 - Test an env: `uv run kinitro env test metaworld/pick-place-v3`
 - Build env image: `uv run kinitro env build --env-id metaworld/pick-place-v3 --tag my-env:v1`
 
 ## Git Hooks
+
 - Hook script: `.githooks/pre-commit` invokes the `pre-commit` tool for ruff formatting/linting, then runs `ty` type checking.
 - Setup: `git config core.hooksPath .githooks && uv tool install pre-commit`
 - The hook uses `.pre-commit-config.yaml` for ruff rules via the pre-commit framework.
 
 ## Services (local dev)
+
 - API: `uv run kinitro api --database-url postgresql://user:pass@host/db`
 - Scheduler: `uv run kinitro scheduler --netuid <id> --network finney --database-url postgresql://user:pass@host/db`
   - Filter to specific environment families: `--env-families metaworld` or `--env-families metaworld,genesis`
@@ -74,63 +86,80 @@ Keep it current when commands or conventions change.
 - Validator: `uv run kinitro validate --backend-url https://api.kinitro.ai --netuid <id> --network finney`
 
 ## Backend Setup (operator quick start)
+
 - Init DB: `uv run kinitro db init --database-url postgresql://user:pass@host/db`
 - DB status: `uv run kinitro db status --database-url postgresql://user:pass@host/db`
 
 ## Environment Config
+
 - See `.env.example` for common env vars.
 - Runtime settings are read via Pydantic settings classes in `kinitro/config.py`.
 - Keep secrets out of the repo; do not commit `.env` files or any files listed in `.gitignore`.
 
 ## Code Style Guidelines
+
 ### Imports
+
 - Group imports: standard library, third-party, then local.
 - Keep imports at the top of the file (after module docstring if present).
 - Sort with Ruff `I` rules (no separate isort config).
 - Prefer explicit imports over wildcard imports.
 
 ### Formatting
+
 - 4-space indentation.
 - Line length target: 100 (`ruff` ignores `E501` but still keep lines reasonable).
 - Use trailing commas in multi-line literals and call arguments.
 
 ### Typing
+
 - Python 3.12 syntax (`list[int]`, `dict[str, float]`, `str | None`).
-- Add return types to public functions and methods.
+- Add type hints to all function signatures (parameters and return types), not just public ones.
 - Prefer `BaseSettings` and `BaseModel` type annotations for config/DTOs.
 - `ty` runs in CI; avoid `Any` unless required and explain why in code.
+- Use `NewType` to distinguish domain identifiers that share an underlying primitive (e.g., `MinerUID`, `BlockNumber`, `EnvironmentId`, `Hotkey`). This prevents accidental misuse such as passing a `MinerUID` where a `BlockNumber` is expected.
+- Centralize shared newtypes, type aliases, enums, and `TypedDict` definitions in `kinitro/types.py`. Import from there rather than re-defining types locally.
+- When introducing a new domain concept that is fundamentally a `str`, `int`, or other primitive, create a `NewType` for it in `kinitro/types.py` and use it consistently across signatures, models, and data structures.
+- Prefer `TypedDict` or `dataclasses.dataclass` over plain `dict` for structured data with known keys.
 
 ### Naming
+
 - `snake_case` for functions, variables, modules.
 - `PascalCase` for classes, Pydantic models, and custom types.
 - `UPPER_CASE` for constants.
 - Test files: `test_*.py`; test classes: `Test*`.
 
 ### Error Handling
+
 - CLI commands use `typer.Exit(1)` for user-facing failures.
 - Provide actionable error messages; log details with `structlog`.
 - Avoid broad `except Exception` unless you re-raise or return a clear error.
 
 ### Logging
+
 - Use `structlog.get_logger()` and structured log fields.
 - Log at info/debug/warn/error levels consistently with existing patterns.
 
 ### Async and IO
+
 - Use `async def` for I/O and DB operations.
 - Use `asyncio.run(...)` at CLI boundaries.
 - Use `asynccontextmanager` and session managers for DB work.
 
 ### Pydantic / API Models
+
 - Pydantic models live in `kinitro/backend/models.py`.
 - Prefer `Field(...)` for validation constraints and docs.
 - Use `from_attributes = True` for ORM-to-model conversion.
 
 ### SQLAlchemy
+
 - ORM models are in `kinitro/backend/models.py`.
 - Keep DB writes inside the storage layer (`kinitro/backend/storage.py`).
 - Use `select(...).where(...).limit(...)` patterns for queries.
 
 ### Tests
+
 - Keep tests deterministic; use fixed seeds (`np.random.default_rng(42)`).
 - Use clear asserts; prefer `np.testing` for array comparisons.
 - Keep unit tests fast; integration tests go in `tests/integration`.
@@ -140,6 +169,7 @@ Keep it current when commands or conventions change.
 For detailed testing docs and troubleshooting, see `docs/e2e-testing.md`.
 
 ### Quick Reference
+
 ```bash
 # Start all services with mock miner (uses default ports)
 ./scripts/services.sh start --mock-miner
@@ -156,13 +186,16 @@ For detailed testing docs and troubleshooting, see `docs/e2e-testing.md`.
 ```
 
 For multi-worktree development (avoids port collisions between parallel checkouts):
+
 ```bash
 ./scripts/worktree-env.sh           # Generate isolated ports/database (once per worktree)
 ./scripts/services.sh start --all   # Uses worktree-specific ports
 ```
 
 ### API Endpoints
+
 Note: The API prefix is `/v1/` (not `/api/v1/`).
+
 - Health: `GET /health`
 - Miners: `GET /v1/miners`
 - Environments: `GET /v1/environments`
@@ -171,6 +204,7 @@ Note: The API prefix is `/v1/` (not `/api/v1/`).
 - Task stats: `GET /v1/tasks/stats`
 
 ## Docs and References
+
 - Developer overview: `README.md`.
 - Backend operator guide: `docs/backend-guide.md`.
 - Miner guide: `docs/miner-guide.md`.
@@ -179,6 +213,7 @@ Note: The API prefix is `/v1/` (not `/api/v1/`).
 - **Adding new environments**: `environments/README.md` - Complete guide for integrating new robotics environments.
 
 ## Change Hygiene for Agents
+
 - Do not modify files outside the task scope.
 - Avoid editing generated files and caches (e.g., `.ruff_cache/`).
 - Keep commits small and focused; update docs if behavior changes.

environments/_template/env.py

@@ -23,8 +23,10 @@
 import structlog
 
 from kinitro.environments import get_environment
+from kinitro.environments.base import RoboticsEnvironment
 from kinitro.environments.registry import get_environments_by_family
 from kinitro.rl_interface import Action
+from kinitro.types import EnvironmentId
 
 logger = structlog.get_logger()
 
@@ -36,7 +38,7 @@ def __init__(self):
         """Initialize the evaluation actor."""
         self._env_cache = {}
 
-    def _get_env(self, env_id: str):
+    def _get_env(self, env_id: EnvironmentId) -> RoboticsEnvironment:
         """Get or create a robotics environment (lazy loading)."""
         if env_id not in self._env_cache:
             self._env_cache[env_id] = get_environment(env_id)
@@ -67,18 +69,18 @@ async def _call_miner(
             resp.raise_for_status()
             return resp.json()
 
-    async def list_environments(self) -> list[str]:
+    async def list_environments(self) -> list[EnvironmentId]:
         """List available environments in this family."""
         # TODO: Change "myenv" to your environment family prefix
         return get_environments_by_family("myenv")
 
     async def evaluate(
         self,
         task_id: int,
+        base_url: str,
         seed: int | None = None,
         model: str | None = None,
-        base_url: str | None = None,
-        env_id: str = "myenv/v0",  # TODO: Change default env_id
+        env_id: EnvironmentId = EnvironmentId("myenv/v0"),  # TODO: Change default env_id
         max_timesteps: int = 500,
         action_timeout: float = 0.5,
         use_images: bool = True,
@@ -108,9 +110,6 @@ async def evaluate(
                 error=f"Invalid env_id: {env_id}. Must start with 'myenv/'",
             )
 
-        if base_url is None:
-            raise ValueError("base_url (miner endpoint) is required")
-
         seed = seed if seed is not None else task_id
         start_time = time.time()
 
@@ -207,7 +206,7 @@ async def evaluate(
 
     def _build_error_result(
         self,
-        env_id: str,
+        env_id: EnvironmentId,
         task_id: int,
         seed: int,
         start_time: float,

environments/genesis/env.py

@@ -32,6 +32,7 @@
 from kinitro.environments.base import RoboticsEnvironment
 from kinitro.environments.registry import get_environments_by_family
 from kinitro.rl_interface import Action
+from kinitro.types import EnvironmentId
 
 logger = structlog.get_logger()
 
@@ -58,12 +59,12 @@ class Actor:
     def __init__(self) -> None:
         """Initialize the evaluation actor."""
         self._env_cache = {}
-        self._env_locks: dict[str, asyncio.Lock] = {}
+        self._env_locks: dict[EnvironmentId, asyncio.Lock] = {}
         # Don't cache the HTTP client - create fresh for each evaluation
         # to avoid event loop binding issues when affinetes calls methods
         # from different event loops
 
-    def _get_env(self, env_id: str) -> RoboticsEnvironment:
+    def _get_env(self, env_id: EnvironmentId) -> RoboticsEnvironment:
         """Get or create a robotics environment."""
         if env_id not in self._env_cache:
             self._env_cache[env_id] = get_environment(env_id)
@@ -99,7 +100,7 @@ async def _call_miner(
             resp.raise_for_status()
             return resp.json()
 
-    async def list_environments(self) -> list[str]:
+    async def list_environments(self) -> list[EnvironmentId]:
         """List available Genesis environments."""
         return get_environments_by_family("genesis")
 
@@ -109,7 +110,7 @@ async def evaluate(
         base_url: str,
         seed: int | None = None,
         model: str | None = None,
-        env_id: str = "genesis/g1-v0",
+        env_id: EnvironmentId = EnvironmentId("genesis/g1-v0"),
         max_timesteps: int = 500,
         action_timeout: float = 0.5,
         use_images: bool = True,
@@ -188,7 +189,7 @@ async def _run_evaluation(
         seed: int,
         model: str | None,
         base_url: str,
-        env_id: str,
+        env_id: EnvironmentId,
         max_timesteps: int,
         action_timeout: float,
         use_images: bool,
@@ -322,7 +323,7 @@ async def _run_evaluation(
 
     def _build_error_result(
         self,
-        env_id: str,
+        env_id: EnvironmentId,
         task_id: int,
         seed: int,
         start_time: float,

environments/metaworld/env.py

@@ -37,8 +37,10 @@
 
 # Import from kinitro package (installed in container via PYTHONPATH)
 from kinitro.environments import get_environment
+from kinitro.environments.base import RoboticsEnvironment
 from kinitro.environments.registry import get_environments_by_family
 from kinitro.rl_interface import Action
+from kinitro.types import EnvironmentId
 
 logger = structlog.get_logger()
 
@@ -58,7 +60,7 @@ def __init__(self):
         # to avoid event loop binding issues when affinetes calls methods
         # from different event loops
 
-    def _get_env(self, env_id: str):
+    def _get_env(self, env_id: EnvironmentId) -> RoboticsEnvironment:
         """Get or create a robotics environment."""
         if env_id not in self._env_cache:
             self._env_cache[env_id] = get_environment(env_id)
@@ -91,17 +93,17 @@ async def _call_miner(
             resp.raise_for_status()
             return resp.json()
 
-    async def list_environments(self) -> list[str]:
+    async def list_environments(self) -> list[EnvironmentId]:
         """List available MetaWorld environments."""
         return get_environments_by_family("metaworld")
 
     async def evaluate(
         self,
         task_id: int,
+        base_url: str,
         seed: int | None = None,
         model: str | None = None,
-        base_url: str | None = None,
-        env_id: str = "metaworld/pick-place-v3",
+        env_id: EnvironmentId = EnvironmentId("metaworld/pick-place-v3"),
         max_timesteps: int = 500,
         action_timeout: float = 0.5,
         use_images: bool = True,
@@ -148,9 +150,6 @@ async def evaluate(
                 error=f"Invalid env_id for MetaWorld container: {env_id}. Must start with 'metaworld/'",
             )
 
-        if base_url is None:
-            raise ValueError("base_url (miner endpoint) is required")
-
         if seed is None:
             seed = task_id
 
@@ -184,7 +183,7 @@ async def _run_evaluation(
         seed: int,
         model: str | None,
         base_url: str,
-        env_id: str,
+        env_id: EnvironmentId,
         max_timesteps: int,
         action_timeout: float,
         use_images: bool,
@@ -318,7 +317,7 @@ async def _run_evaluation(
 
     def _build_error_result(
         self,
-        env_id: str,
+        env_id: EnvironmentId,
         task_id: int,
         seed: int,
         start_time: float,
@@ -342,7 +341,7 @@ def _build_error_result(
             "extra": extra_fields,
         }
 
-    async def cleanup(self):
+    async def cleanup(self) -> None:
         """Cleanup resources."""
         # HTTP clients are now created per-request, no need to close here
 

kinitro/__init__.py

@@ -6,4 +6,9 @@
 environments earn rewards via ε-Pareto dominance scoring.
 """
 
+import os
+
+# https://docs.learnbittensor.org/sdk/migration-guide#disabling-cli-argument-parsing
+os.environ.setdefault("BT_NO_PARSE_CLI_ARGS", "1")
+
 __version__ = "0.1.0"