From 8daf705cf99ea7972567a9fe8907f8c5d548a2aa Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:45:04 -0700 Subject: [PATCH 01/20] add main ck queries --- queries/banned_coldkey.py | 75 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 queries/banned_coldkey.py diff --git a/queries/banned_coldkey.py b/queries/banned_coldkey.py new file mode 100644 index 00000000..c56a7861 --- /dev/null +++ b/queries/banned_coldkey.py @@ -0,0 +1,75 @@ +from typing import Optional +from uuid import UUID + +from models.banned_coldkey import BannedColdkey +from utils.database import DatabaseConnection, db_operation + +COLDKEY_BAN_LOCK_NAMESPACE = -1729 + + +async def lock_coldkey_ban_state(conn: DatabaseConnection, miner_coldkey: str) -> None: + await conn.execute( + "SELECT pg_advisory_xact_lock($1, hashtext($2))", + COLDKEY_BAN_LOCK_NAMESPACE, + miner_coldkey, + ) + + +@db_operation +async def get_banned_coldkey( + conn: DatabaseConnection, + miner_coldkey: str, +) -> Optional[BannedColdkey]: + row = await conn.fetchrow( + "SELECT * FROM banned_coldkeys WHERE miner_coldkey = $1", + miner_coldkey, + ) + return BannedColdkey(**row) if row is not None else None + + +@db_operation +async def is_agent_coldkey_banned(conn: DatabaseConnection, agent_id: UUID) -> bool: + return await conn.fetchval( + """ + SELECT EXISTS ( + SELECT 1 + FROM agents + INNER JOIN banned_coldkeys USING (miner_coldkey) + WHERE agents.agent_id = $1 + ) + """, + agent_id, + ) + + +@db_operation +async def ban_coldkey( + conn: DatabaseConnection, + miner_coldkey: str, + banned_reason: str, +) -> BannedColdkey: + async with conn.conn.transaction(): + await lock_coldkey_ban_state(conn, miner_coldkey) + row = await conn.fetchrow( + """ + INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) + VALUES ($1, $2) + ON CONFLICT (miner_coldkey) DO UPDATE + SET banned_reason = EXCLUDED.banned_reason + RETURNING * + """, + miner_coldkey, + banned_reason, + ) + return BannedColdkey(**row) + + +@db_operation +async def unban_coldkey(conn: DatabaseConnection, miner_coldkey: str) -> bool: + async with conn.conn.transaction(): + await lock_coldkey_ban_state(conn, miner_coldkey) + result = await conn.execute( + "DELETE FROM banned_coldkeys WHERE miner_coldkey = $1", + miner_coldkey, + ) + return result == "DELETE 1" From 86c1b55394403cefce1bfc37628791dcdf537499 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:45:44 -0700 Subject: [PATCH 02/20] check: approval judge --- api/endpoints/validator.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/api/endpoints/validator.py b/api/endpoints/validator.py index 9a3476ed..468f71bb 100644 --- a/api/endpoints/validator.py +++ b/api/endpoints/validator.py @@ -54,6 +54,7 @@ update_agent_status, ) from queries.approval import finish_agent_and_enqueue_approval +from queries.banned_coldkey import is_agent_coldkey_banned from queries.evaluation import ( create_new_evaluation_and_evaluation_runs, get_approved_leader_ranking_for_set, @@ -1193,6 +1194,10 @@ async def handle_evaluation_if_finished(evaluation_id: UUID) -> None: async def _should_run_auto_approval_judge(*, agent_id: UUID, set_id: int) -> bool: + if await is_agent_coldkey_banned(agent_id): + logger.info(f"Skipping auto approval for agent_id={agent_id}: miner coldkey is banned") + return False + candidate = await get_validator_agent_score_for_set( agent_id, set_id, From 06025fc0b700cd580c8b96b40e2746b860780d7e Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:47:45 -0700 Subject: [PATCH 03/20] block: upload path --- api/src/endpoints/upload.py | 55 +++++++++++++++++++++---------------- 1 file changed, 32 insertions(+), 23 deletions(-) diff --git a/api/src/endpoints/upload.py b/api/src/endpoints/upload.py index 870f5379..3348d9eb 100644 --- a/api/src/endpoints/upload.py +++ b/api/src/endpoints/upload.py @@ -13,7 +13,7 @@ from api.src.utils.request_cache import hourly_cache from api.src.utils.upload_agent_helpers import ( as_utc, - check_agent_banned, + check_coldkey_banned, check_file_size, check_hotkey_registered, check_if_extrinsic_failed, @@ -34,8 +34,8 @@ get_latest_agent_for_miner_hotkey, record_upload_attempt, ) -from queries.banned_hotkey import get_banned_hotkey -from queries.errors import DuplicateAgentIDError +from queries.banned_coldkey import get_banned_coldkey +from queries.errors import ColdkeyBannedError, DuplicateAgentIDError from queries.payments import ( complete_payment, create_payment_quote, @@ -86,6 +86,7 @@ async def check_agent_post( if config.DISALLOW_UPLOADS: raise HTTPException(status_code=503, detail=config.DISALLOW_UPLOADS_REASON) miner_hotkey = get_miner_hotkey(file_info) + is_owner_upload = miner_hotkey == config.OWNER_HOTKEY latest_agent_created_at_in_latest_set_id = await get_latest_agent_created_at_for_miner_hotkey_in_latest_set_id( miner_hotkey=miner_hotkey ) @@ -93,12 +94,13 @@ async def check_agent_post( check_rate_limit(latest_agent_created_at_in_latest_set_id) check_signature(public_key, file_info, signature, miner_hotkey) await check_hotkey_registered(miner_hotkey) - await check_agent_banned(miner_hotkey=miner_hotkey) - check_if_python_file(agent_file.filename) - await check_file_size(agent_file) coldkey = await subtensor_client.get_hotkey_owner(miner_hotkey) if coldkey is None: raise HTTPException(status_code=400, detail="Hotkey owner not found") + if not is_owner_upload: + await check_coldkey_banned(coldkey) + check_if_python_file(agent_file.filename) + await check_file_size(agent_file) try: alpha_stake = await subtensor_client.get_alpha_stake_availability( @@ -186,6 +188,7 @@ async def post_agent( prod = config.ENV == "prod" miner_hotkey = get_miner_hotkey(file_info) + coldkey: Optional[str] = None # Extract upload attempt data for tracking agent_file.file.seek(0, 2) @@ -214,7 +217,6 @@ async def post_agent( if prod: await check_hotkey_registered(miner_hotkey) - await check_agent_banned(miner_hotkey=miner_hotkey) check_if_python_file(agent_file.filename) agent_bytes, agent_text = await check_file_size(agent_file) @@ -271,12 +273,15 @@ async def post_agent( except (ValueError, TypeError, IndexError, AttributeError): raise HTTPException(status_code=402, detail="Burn extrinsic could not be decoded") from None + coldkey = await subtensor_client.get_hotkey_owner(miner_hotkey, block=int(payment_block_info.number)) + if coldkey is None: + raise HTTPException(status_code=402, detail="Hotkey owner not found at payment block") + await check_coldkey_banned(coldkey) + events = await subtensor_client.get_events(block_hash=payment_block_hash) if await check_if_extrinsic_failed(payment_extrinsic_index_int, events): raise HTTPException(status_code=402, detail="Burn extrinsic failed on-chain") - coldkey = await subtensor_client.get_hotkey_owner(miner_hotkey, block=int(payment_block_info.number)) - # Cross-check the extrinsic: recognized burn call signed by the miner coldkey. verify_burn_extrinsic(payment_extrinsic, expected_coldkey=coldkey) @@ -352,18 +357,22 @@ async def post_agent( payment_block_hash=payment_block_hash, payment_extrinsic_index=payment_extrinsic_index, ) - agent_id = await create_agent( - agent, - agent_text, - source_sha256=source_sha256, - runtime_openrouter_api_key_ciphertext=encrypted_openrouter_api_key, - management_openrouter_api_key_ciphertext=encrypted_openrouter_management_key, - openrouter_workspace_id=validated_openrouter_keys.workspace_id, - openrouter_api_key_label=validated_openrouter_keys.api_key_label, - openrouter_api_key_creator_user_id=validated_openrouter_keys.api_key_creator_user_id, - openrouter_validated_at=validated_openrouter_keys.validated_at, - create_pre_screening_job=config.PRE_SCREENING_JUDGE_ENABLED, - ) + try: + agent_id = await create_agent( + agent, + agent_text, + source_sha256=source_sha256, + runtime_openrouter_api_key_ciphertext=encrypted_openrouter_api_key, + management_openrouter_api_key_ciphertext=encrypted_openrouter_management_key, + openrouter_workspace_id=validated_openrouter_keys.workspace_id, + openrouter_api_key_label=validated_openrouter_keys.api_key_label, + openrouter_api_key_creator_user_id=validated_openrouter_keys.api_key_creator_user_id, + openrouter_validated_at=validated_openrouter_keys.validated_at, + miner_coldkey=coldkey if prod else None, + create_pre_screening_job=config.PRE_SCREENING_JUDGE_ENABLED, + ) + except ColdkeyBannedError as e: + raise HTTPException(status_code=403, detail="Your miner coldkey has been banned") from e if prod and not is_owner_upload: await complete_payment( @@ -398,7 +407,7 @@ async def post_agent( if e.status_code == 429 else "validation_error" ) - banned_hotkey = await get_banned_hotkey(miner_hotkey) if error_type == "banned" and miner_hotkey else None + banned_coldkey = await get_banned_coldkey(coldkey) if error_type == "banned" and coldkey else None # Record failed upload attempt await record_upload_attempt( @@ -406,7 +415,7 @@ async def post_agent( success=False, error_type=error_type, error_message=e.detail, - ban_reason=banned_hotkey.banned_reason if banned_hotkey else None, + ban_reason=banned_coldkey.banned_reason if banned_coldkey else None, http_status_code=e.status_code, **upload_data, ) From e09989f56e5b1d4cfac118a73106de2cc8e19987 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:48:11 -0700 Subject: [PATCH 04/20] update helper --- api/src/utils/upload_agent_helpers.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/api/src/utils/upload_agent_helpers.py b/api/src/utils/upload_agent_helpers.py index 18e189da..f148f319 100644 --- a/api/src/utils/upload_agent_helpers.py +++ b/api/src/utils/upload_agent_helpers.py @@ -8,7 +8,7 @@ from fastapi import HTTPException, UploadFile from api.config import MINER_AGENT_UPLOAD_RATE_LIMIT_SECONDS -from queries.banned_hotkey import get_banned_hotkey +from queries.banned_coldkey import get_banned_coldkey from utils.bittensor import subtensor_client logger = logging.getLogger(__name__) @@ -47,17 +47,17 @@ def check_if_python_file(filename: str) -> None: logger.debug("The file is a python file.") -async def check_agent_banned(miner_hotkey: str) -> None: - logger.debug(f"Checking if miner hotkey {miner_hotkey} is banned...") +async def check_coldkey_banned(miner_coldkey: str) -> None: + logger.debug(f"Checking if miner coldkey {miner_coldkey} is banned...") - if await get_banned_hotkey(miner_hotkey) is not None: - logger.error(f"A miner attempted to upload an agent with a banned hotkey: {miner_hotkey}.") + if await get_banned_coldkey(miner_coldkey) is not None: + logger.warning(f"A miner attempted to upload an agent with a banned coldkey: {miner_coldkey}.") raise HTTPException( status_code=403, - detail="Your miner hotkey has been banned for attempting to obfuscate code or otherwise cheat. If this is in error, please contact us on Discord", + detail="Your miner coldkey has been banned. If this is in error, please contact us on Discord", ) - logger.debug(f"Miner hotkey {miner_hotkey} is not banned.") + logger.debug(f"Miner coldkey {miner_coldkey} is not banned.") def check_rate_limit( From dfc67ca8cbed1204cfc142d7a27bb62207e8c9b2 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:48:36 -0700 Subject: [PATCH 05/20] add banned ck model --- db/models/__init__.py | 2 ++ db/models/agent.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) diff --git a/db/models/__init__.py b/db/models/__init__.py index a7266847..e564eca9 100644 --- a/db/models/__init__.py +++ b/db/models/__init__.py @@ -2,6 +2,7 @@ Agent, AgentOpenRouterSecret, AgentScore, + BannedColdkey, BannedHotkey, BenchmarkAgentId, UnapprovedAgentId, @@ -35,6 +36,7 @@ "PreScreeningResult", "AgentScore", "ApprovedAgent", + "BannedColdkey", "BannedHotkey", "BenchmarkAgentId", "Embedding", diff --git a/db/models/agent.py b/db/models/agent.py index 0179e2af..7f1a2151 100644 --- a/db/models/agent.py +++ b/db/models/agent.py @@ -15,6 +15,7 @@ class Agent(Base): agent_id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True) miner_hotkey: Mapped[str] = mapped_column(sa.Text, nullable=False) + miner_coldkey: Mapped[Optional[str]] = mapped_column(sa.Text) name: Mapped[str] = mapped_column(sa.Text, nullable=False) version_num: Mapped[int] = mapped_column(sa.Integer, nullable=False) status: Mapped[Optional[AgentStatus]] = mapped_column(sa.Enum(AgentStatus, name="agentstatus")) @@ -55,6 +56,18 @@ class BannedHotkey(Base): __table_args__ = (sa.Index("idx_banned_hotkeys_miner_hotkey", "miner_hotkey"),) +class BannedColdkey(Base): + __tablename__ = "banned_coldkeys" + + miner_coldkey: Mapped[str] = mapped_column(sa.Text, primary_key=True) + banned_reason: Mapped[str] = mapped_column(sa.Text, nullable=False) + banned_at: Mapped[datetime] = mapped_column( + sa.TIMESTAMP(timezone=True), + nullable=False, + server_default=sa.text("NOW()"), + ) + + class BenchmarkAgentId(Base): __tablename__ = "benchmark_agent_ids" From 1432606c681a437c92ef3606d35cdb8c5884c941 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:52:32 -0700 Subject: [PATCH 06/20] update: exclude banned from top_agents + check before creating agent --- queries/agent.py | 31 ++++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/queries/agent.py b/queries/agent.py index 7033b001..77794297 100644 --- a/queries/agent.py +++ b/queries/agent.py @@ -16,7 +16,8 @@ from models.evaluation import EvaluationStatus from models.evaluation_set import EvaluationSetGroup from models.queue import QueueStage -from queries.errors import DuplicateAgentIDError +from queries.banned_coldkey import get_banned_coldkey, lock_coldkey_ban_state +from queries.errors import ColdkeyBannedError, DuplicateAgentIDError from utils.agent_secrets import decrypt_agent_secret from utils.database import DatabaseConnection, db_operation from utils.s3 import upload_text_file_to_s3 @@ -192,6 +193,7 @@ async def create_agent( openrouter_api_key_label: str, openrouter_api_key_creator_user_id: str, openrouter_validated_at: datetime, + miner_coldkey: Optional[str] = None, create_pre_screening_job: bool = False, pre_screening_policy_version: str = DEFAULT_PRE_SCREENING_POLICY_VERSION, ) -> UUID: @@ -224,15 +226,32 @@ async def create_agent( source_sha256, ) + if miner_coldkey is not None: + await lock_coldkey_ban_state(conn, miner_coldkey) + if await get_banned_coldkey(miner_coldkey) is not None: + raise ColdkeyBannedError(miner_coldkey) + result = await conn.fetchval( """ - INSERT INTO agents (agent_id, miner_hotkey, name, version_num, created_at, status, ip_address, source_sha256, set_id) - VALUES ($1, $2, $3, $4, NOW(), $5, $6, $7, $8) + INSERT INTO agents ( + agent_id, + miner_hotkey, + miner_coldkey, + name, + version_num, + created_at, + status, + ip_address, + source_sha256, + set_id + ) + VALUES ($1, $2, $3, $4, $5, NOW(), $6, $7, $8, $9) ON CONFLICT (agent_id) DO NOTHING RETURNING agent_id """, agent_id, agent.miner_hotkey, + miner_coldkey, agent.name, agent.version_num, agent.status.value, @@ -489,6 +508,7 @@ async def get_top_agents(conn: DatabaseConnection, number_of_agents: int = 10, p ass.*, review.approval_review_status AS approval_review_status from agent_scores ass + join agents a on a.agent_id = ass.agent_id left join agent_final_review_statuses review on review.agent_id = ass.agent_id and review.set_id = ass.set_id @@ -502,6 +522,11 @@ async def get_top_agents(conn: DatabaseConnection, number_of_agents: int = 10, p ) rt on true where ass.set_id = (select max(set_id) from evaluation_sets) and ass.agent_id not in (select agent_id from benchmark_agent_ids) + and not exists ( + select 1 + from banned_coldkeys bc + where bc.miner_coldkey = a.miner_coldkey + ) and ass.status::text <> 'cancelled' and ( ass.approved is true From 1cbb474dfc99282b5d98c5a4525e541ea68cb170 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:53:34 -0700 Subject: [PATCH 07/20] finish_agent_and_enqueue_approval: check ban state --- queries/approval.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/queries/approval.py b/queries/approval.py index f5d936c4..ea595431 100644 --- a/queries/approval.py +++ b/queries/approval.py @@ -14,6 +14,7 @@ ) from models.evaluation import EvaluationStatus from models.evaluation_set import EvaluationSetGroup +from queries.banned_coldkey import get_banned_coldkey, lock_coldkey_ban_state from utils.database import DatabaseConnection, db_operation logger = logging.getLogger(__name__) @@ -58,21 +59,34 @@ async def finish_agent_and_enqueue_approval( set_id: int, policy_version: str, ) -> bool: - """Atomically move an agent to finished and enqueue its approval job.""" + """Finish an agent and enqueue approval unless its coldkey is banned.""" async with conn.conn.transaction(): - update_result = await conn.execute( + agent = await conn.fetchrow( """ - UPDATE agents - SET status = $2 + SELECT miner_coldkey + FROM agents WHERE agent_id = $1 - AND status = $3 + AND status = $2 + FOR UPDATE """, agent_id, - AgentStatus.finished.value, AgentStatus.evaluating.value, ) - if update_result == "UPDATE 0": + if agent is None: + return False + + miner_coldkey = agent["miner_coldkey"] + if miner_coldkey is not None: + await lock_coldkey_ban_state(conn, miner_coldkey) + + await conn.execute( + "UPDATE agents SET status = $2 WHERE agent_id = $1", + agent_id, + AgentStatus.finished.value, + ) + + if miner_coldkey is not None and await get_banned_coldkey(miner_coldkey) is not None: return False snapshot = await _build_approval_input_snapshot(conn, agent_id=agent_id, set_id=set_id) From 28fc20cecab1f81627292914aec988ac8294e4be Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:54:14 -0700 Subject: [PATCH 08/20] wip: models --- models/banned_coldkey.py | 9 +++++++++ queries/errors.py | 8 ++++++++ 2 files changed, 17 insertions(+) create mode 100644 models/banned_coldkey.py diff --git a/models/banned_coldkey.py b/models/banned_coldkey.py new file mode 100644 index 00000000..465de5a4 --- /dev/null +++ b/models/banned_coldkey.py @@ -0,0 +1,9 @@ +from datetime import datetime + +from pydantic import BaseModel + + +class BannedColdkey(BaseModel): + miner_coldkey: str + banned_reason: str + banned_at: datetime diff --git a/queries/errors.py b/queries/errors.py index 18e35592..004c2882 100644 --- a/queries/errors.py +++ b/queries/errors.py @@ -10,3 +10,11 @@ class DuplicateAgentIDError(Exception): def __init__(self, agent_id: "UUID"): self.agent_id = agent_id super().__init__(f"Agent {agent_id} already exists") + + +class ColdkeyBannedError(Exception): + """Raised when an agent insert loses a race with a coldkey ban.""" + + def __init__(self, miner_coldkey: str): + self.miner_coldkey = miner_coldkey + super().__init__(f"Coldkey {miner_coldkey} is banned") From fb01bafe73b14808ac744499c7279210788ca1ad Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:55:04 -0700 Subject: [PATCH 09/20] update: metrics query to exclude ck --- queries/evaluation_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/queries/evaluation_run.py b/queries/evaluation_run.py index f599129a..9c6abb9f 100644 --- a/queries/evaluation_run.py +++ b/queries/evaluation_run.py @@ -137,12 +137,12 @@ async def _get_evaluation_run_metrics_by_ids( JOIN evaluation_runs er2 ON er2.evaluation_id = e2.evaluation_id AND er2.problem_name = dp.problem_name JOIN agents a ON a.agent_id = e2.agent_id - LEFT JOIN banned_hotkeys bh ON bh.miner_hotkey = a.miner_hotkey + LEFT JOIN banned_coldkeys bc ON bc.miner_coldkey = a.miner_coldkey LEFT JOIN unapproved_agent_ids uai ON uai.agent_id = a.agent_id LEFT JOIN benchmark_agent_ids bai ON bai.agent_id = a.agent_id WHERE e2.set_id = dp.set_id - AND bh.miner_hotkey IS NULL + AND bc.miner_coldkey IS NULL AND uai.agent_id IS NULL AND bai.agent_id IS NULL ) agg From 76ed3519f676257b4a3f6b32bf6ecf85b092d4a4 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 16:57:25 -0700 Subject: [PATCH 10/20] updates: - approved agents for set: Exclude banned - get_evaluation_set_score_stats: exclude banned cks from comparision - get_evaluation_set_submission_stats: Exclude banned ck - _sql_agents_in_window_cte: exclude banned ck --- queries/evaluation_set.py | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/queries/evaluation_set.py b/queries/evaluation_set.py index ce206c79..b40b9b21 100644 --- a/queries/evaluation_set.py +++ b/queries/evaluation_set.py @@ -44,7 +44,13 @@ def _sql_agents_in_window_cte(select_columns: str) -> str: SELECT {select_columns}, CASE - WHEN review.approval_review_status = 'rejected' THEN true + WHEN review.approval_review_status = 'rejected' + OR EXISTS ( + SELECT 1 + FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) + THEN true ELSE false END AS disqualified FROM agents a @@ -385,7 +391,7 @@ async def get_evaluation_set_submission_stats(conn: DatabaseConnection, set_id: and aiw.status = 'finished' then aiw.agent_id end ) :: int as finished_at_validator_count, - COUNT(distinct aa.agent_id) :: int as approved_emission_count + COUNT(distinct aa.agent_id) FILTER (WHERE NOT aiw.disqualified) :: int as approved_emission_count from agents_in_window aiw left join last_evaluation_per_agent e on e.agent_id = aiw.agent_id @@ -424,11 +430,12 @@ async def get_evaluation_set_score_stats(conn: DatabaseConnection, set_id: int) {_sql_agents_in_window_cte("a.agent_id, a.status")}, prev_best AS ( SELECT - MAX(final_score) AS score + MAX(previous_score.final_score) AS score FROM - agent_scores + agent_scores previous_score + JOIN agents previous_agent ON previous_agent.agent_id = previous_score.agent_id WHERE - set_id = ( + previous_score.set_id = ( SELECT MAX(set_id) FROM @@ -436,12 +443,17 @@ async def get_evaluation_set_score_stats(conn: DatabaseConnection, set_id: int) WHERE set_id < $1 ) - AND agent_id NOT IN ( + AND previous_score.agent_id NOT IN ( SELECT agent_id FROM benchmark_agent_ids ) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys bc + WHERE bc.miner_coldkey = previous_agent.miner_coldkey + ) ) SELECT MAX(s.final_score) AS best, @@ -702,6 +714,11 @@ async def get_approved_agents_for_set(conn: DatabaseConnection, set_id: int) -> LEFT JOIN validator_metrics vm ON vm.agent_id = aa.agent_id WHERE aa.set_id = $1 AND aa.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND ass.status = 'finished' AND review.approval_review_status is distinct from 'rejected' ORDER BY aa.approved_at DESC From 0ed6faa65cc912541abc7d0b049c84b53ccae8b7 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:01:31 -0700 Subject: [PATCH 11/20] Updated: - get_approved_validator_leader_score_for_set: excludes banned coldkey scores - get_approved_leader_ranking_for_set: excludes banned coldkey leaders --- queries/evaluation.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/queries/evaluation.py b/queries/evaluation.py index 5e7edc34..4f6160de 100644 --- a/queries/evaluation.py +++ b/queries/evaluation.py @@ -187,6 +187,7 @@ async def get_approved_validator_leader_score_for_set( """ SELECT MAX(agent_score.final_score) FROM agent_scores agent_score + INNER JOIN agents agent ON agent.agent_id = agent_score.agent_id WHERE agent_score.set_id = $1 AND agent_score.approved IS TRUE AND agent_score.approved_at <= NOW() @@ -198,6 +199,11 @@ async def get_approved_validator_leader_score_for_set( FROM benchmark_agent_ids benchmark_agent WHERE benchmark_agent.agent_id = agent_score.agent_id ) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys banned_coldkey + WHERE banned_coldkey.miner_coldkey = agent.miner_coldkey + ) """, set_id, required_validator_count, @@ -266,6 +272,7 @@ async def get_approved_leader_ranking_for_set( rt.avg_cost_usd, ass.created_at FROM agent_scores ass + INNER JOIN agents agent ON agent.agent_id = ass.agent_id LEFT JOIN LATERAL ( SELECT AVG(eh.avg_cost_usd) AS avg_cost_usd FROM evaluations_hydrated eh @@ -285,6 +292,11 @@ async def get_approved_leader_ranking_for_set( FROM benchmark_agent_ids benchmark_agent WHERE benchmark_agent.agent_id = ass.agent_id ) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys banned_coldkey + WHERE banned_coldkey.miner_coldkey = agent.miner_coldkey + ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 """, From 0c596da2290d595ae2a26c31fcbfbb20e1eb2f9b Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:02:39 -0700 Subject: [PATCH 12/20] problem statistics: banned_hks -> banned_cks --- queries/problem_statistics.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/queries/problem_statistics.py b/queries/problem_statistics.py index 22c10487..03193782 100644 --- a/queries/problem_statistics.py +++ b/queries/problem_statistics.py @@ -174,7 +174,10 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ JOIN evaluations e ON erh.evaluation_id = e.evaluation_id JOIN agents a on e.agent_id = a.agent_id WHERE e.set_id = $1 - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY erh.problem_name @@ -209,7 +212,10 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ WHERE er.status = 'finished' AND e.set_id = $1 - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY @@ -268,7 +274,10 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ JOIN agents a ON e.agent_id = a.agent_id WHERE er.status = 'error' AND e.set_id = $1 - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) ) @@ -293,7 +302,10 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ JOIN evaluations e ON er.evaluation_id = e.evaluation_id JOIN agents a ON e.agent_id = a.agent_id WHERE e.set_id = $1 - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY er.problem_name, i.model @@ -332,7 +344,10 @@ async def get_problem_statistics(conn: DatabaseConnection, set_id: int) -> List[ WHERE erh.status = 'finished' AND erh.solved AND e.set_id = $1 - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY erh.problem_name, a.agent_id From 7f249c0960f77f1f6f2c27abeefbd4eb9867b9ef Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:03:34 -0700 Subject: [PATCH 13/20] updates: - get_weight_receiving_agent_hotkey: exclude banned cks from getting weights - get_weight_receiving_agent_info: excluded banned cks --- queries/scores.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/queries/scores.py b/queries/scores.py index 2851cd1e..db837151 100644 --- a/queries/scores.py +++ b/queries/scores.py @@ -14,6 +14,7 @@ async def get_weight_receiving_agent_hotkey(conn: DatabaseConnection) -> Optiona ass.approved AS approved, ass.approved_at AS approved_at FROM agent_scores ass + INNER JOIN agents a ON a.agent_id = ass.agent_id LEFT JOIN LATERAL ( SELECT AVG(eh.avg_cost_usd) AS avg_cost_usd FROM evaluations_hydrated eh @@ -28,6 +29,11 @@ async def get_weight_receiving_agent_hotkey(conn: DatabaseConnection) -> Optiona AND ass.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND ass.status::text <> 'cancelled' AND ass.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 ) @@ -55,6 +61,7 @@ async def get_weight_receiving_agent_info(conn: DatabaseConnection) -> Optional[ ass.approved AS approved, ass.approved_at AS approved_at FROM agent_scores ass + INNER JOIN agents a ON a.agent_id = ass.agent_id LEFT JOIN LATERAL ( SELECT AVG(eh.avg_cost_usd) AS avg_cost_usd FROM evaluations_hydrated eh @@ -69,6 +76,11 @@ async def get_weight_receiving_agent_info(conn: DatabaseConnection) -> Optional[ AND ass.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND ass.status::text <> 'cancelled' AND ass.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) ORDER BY ass.final_score DESC, rt.avg_cost_usd ASC NULLS LAST, ass.created_at ASC LIMIT 1 ) From 39b438b1de72fe2f1f4e8271a47378a0faa5603d Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:04:25 -0700 Subject: [PATCH 14/20] update: statistics (banned hks -> banned cks) --- queries/statistics.py | 67 +++++++++++++++++++++++++++++++++---------- 1 file changed, 52 insertions(+), 15 deletions(-) diff --git a/queries/statistics.py b/queries/statistics.py index e48d7b36..01259117 100644 --- a/queries/statistics.py +++ b/queries/statistics.py @@ -11,9 +11,15 @@ @db_operation async def top_score(conn: DatabaseConnection) -> Optional[float]: return await conn.fetchval(""" - SELECT MAX(final_score) FROM agent_scores - WHERE set_id = (SELECT MAX(set_id) FROM evaluation_sets) - AND agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + SELECT MAX(agent_scores.final_score) + FROM agent_scores + JOIN agents a ON a.agent_id = agent_scores.agent_id + WHERE agent_scores.set_id = (SELECT MAX(set_id) FROM evaluation_sets) + AND agent_scores.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) """) @@ -22,7 +28,10 @@ async def agents_created_24_hrs(conn: DatabaseConnection) -> int: return await conn.fetchval(""" SELECT COUNT(*) FROM agents WHERE created_at >= NOW() - INTERVAL '24 hours' - AND miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = agents.miner_coldkey + ) AND agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) """) @@ -34,11 +43,18 @@ async def score_improvement_24_hrs(conn: DatabaseConnection) -> float: """ WITH score_data AS ( SELECT - MAX(final_score) as max_score, - MAX(final_score) FILTER (WHERE created_at <= NOW() - INTERVAL '24 hours') as max_score_24_hrs_ago + MAX(agent_scores.final_score) as max_score, + MAX(agent_scores.final_score) FILTER ( + WHERE agent_scores.created_at <= NOW() - INTERVAL '24 hours' + ) as max_score_24_hrs_ago FROM agent_scores - WHERE set_id = (SELECT MAX(set_id) FROM evaluation_sets) - AND agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + JOIN agents a ON a.agent_id = agent_scores.agent_id + WHERE agent_scores.set_id = (SELECT MAX(set_id) FROM evaluation_sets) + AND agent_scores.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) ) SELECT COALESCE(max_score - max_score_24_hrs_ago, 0) @@ -72,7 +88,10 @@ async def get_top_scores_over_time(conn: DatabaseConnection) -> list[TopScoreOve WHERE agent_scores.final_score IS NOT NULL AND agent_scores.set_id = (SELECT set_id FROM max_set) - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) ), @@ -94,7 +113,10 @@ async def get_top_scores_over_time(conn: DatabaseConnection) -> list[TopScoreOve agent_scores.final_score IS NOT NULL AND agent_scores.created_at <= ts.hour AND agent_scores.set_id = (SELECT set_id FROM max_set) - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) ), @@ -139,7 +161,10 @@ async def get_perfectly_solved_over_time(conn: DatabaseConnection) -> list[Perfe AND erh.benchmark_family IS NOT NULL AND erh.benchmark_family <> '' AND erh.benchmark_family <> 'custom' - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND e.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND e.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY erh.benchmark_family, erh.problem_name @@ -196,7 +221,10 @@ async def get_average_score_per_evaluation_set_group( JOIN agents a on a.agent_id = eh.agent_id WHERE eh.status = 'success' AND eh.set_id = (SELECT MAX(set_id) FROM evaluation_sets) - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND eh.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND eh.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) GROUP BY validator_type @@ -232,7 +260,10 @@ async def get_average_wait_time_per_evaluation_set_group( WHERE e.status = 'success' AND e.evaluation_set_group = '{EvaluationSetGroup.screener_1.value}'::EvaluationSetGroup AND e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND e.finished_at >= NOW() - INTERVAL '6 hours' @@ -251,7 +282,10 @@ async def get_average_wait_time_per_evaluation_set_group( AND sc2_e.evaluation_set_group = '{EvaluationSetGroup.screener_2.value}'::EvaluationSetGroup AND sc1_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND sc2_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND sc2_e.finished_at >= NOW() - INTERVAL '6 hours' @@ -279,7 +313,10 @@ async def get_average_wait_time_per_evaluation_set_group( AND sc2_e.evaluation_set_group = '{EvaluationSetGroup.screener_2.value}'::EvaluationSetGroup AND sc2_e.set_id = (SELECT MAX(set_id) FROM evaluation_sets) AND v_e.validator_count = {config.NUM_EVALS_PER_AGENT} - AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys) + AND NOT EXISTS ( + SELECT 1 FROM banned_coldkeys bc + WHERE bc.miner_coldkey = a.miner_coldkey + ) AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) AND a.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) AND v_e.finished_at >= NOW() - INTERVAL '6 hours' From 662ed827b078b96f6199cf738b57033726443e23 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:06:01 -0700 Subject: [PATCH 15/20] new: add admin API endpoints for ban/unban --- api/config.py | 4 +++ api/endpoints/admin.py | 82 ++++++++++++++++++++++++++++++++++++++++++ api/src/main.py | 2 ++ 3 files changed, 88 insertions(+) create mode 100644 api/endpoints/admin.py diff --git a/api/config.py b/api/config.py index a6b5c977..05dacbf5 100644 --- a/api/config.py +++ b/api/config.py @@ -40,6 +40,10 @@ if not OWNER_HOTKEY: logger.fatal("OWNER_HOTKEY is not set in .env") +COLDKEY_BAN_ADMIN_API_KEY = os.getenv("COLDKEY_BAN_ADMIN_API_KEY") +if not COLDKEY_BAN_ADMIN_API_KEY: + logger.warning("COLDKEY_BAN_ADMIN_API_KEY is not set; coldkey ban administration will be unavailable") + UPLOAD_SEND_ADDRESS = os.getenv("UPLOAD_SEND_ADDRESS") if not UPLOAD_SEND_ADDRESS: logger.fatal("UPLOAD_SEND_ADDRESS is not set in .env") diff --git a/api/endpoints/admin.py b/api/endpoints/admin.py new file mode 100644 index 00000000..0bf457e8 --- /dev/null +++ b/api/endpoints/admin.py @@ -0,0 +1,82 @@ +import secrets +from typing import Annotated + +from bittensor_wallet.keypair import Keypair +from fastapi import APIRouter, Depends, HTTPException, Response, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import BaseModel, StringConstraints + +import api.config as config +from api.endpoints import evaluation_sets, retrieval, scoring, statistics +from models.banned_coldkey import BannedColdkey +from queries.banned_coldkey import ban_coldkey, unban_coldkey + +router = APIRouter(tags=["admin"]) +admin_bearer = HTTPBearer(auto_error=False) + + +class ColdkeyBanRequest(BaseModel): + reason: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=1000)] + + +def require_coldkey_ban_admin( + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(admin_bearer)], +) -> None: + expected = config.COLDKEY_BAN_ADMIN_API_KEY + if not expected: + raise HTTPException(status_code=503, detail="Coldkey ban administration is not configured") + if credentials is None or not secrets.compare_digest(credentials.credentials, expected): + raise HTTPException( + status_code=401, + detail="Invalid admin credentials", + headers={"WWW-Authenticate": "Bearer"}, + ) + + +def validate_coldkey(miner_coldkey: str) -> None: + try: + Keypair(ss58_address=miner_coldkey) + except Exception: + raise HTTPException(status_code=400, detail="Invalid coldkey SS58 address") from None + + +def clear_coldkey_sensitive_caches() -> None: + + cached_functions = ( + evaluation_sets._cached_build_detail, + evaluation_sets._cached_build_leaderboard, + evaluation_sets._cached_build_approved_agents, + retrieval.queue, + retrieval.top_agents, + retrieval.top_scores_over_time, + retrieval.perfectly_solved_over_time, + retrieval.network_statistics, + scoring.screener_info, + statistics.problem_statistics, + ) + for cached_function in cached_functions: + cached_function.cache_clear() + + +@router.put( + "/banned-coldkeys/{miner_coldkey}", + response_model=BannedColdkey, + dependencies=[Depends(require_coldkey_ban_admin)], +) +async def put_banned_coldkey(miner_coldkey: str, request: ColdkeyBanRequest) -> BannedColdkey: + validate_coldkey(miner_coldkey) + banned_coldkey = await ban_coldkey(miner_coldkey, request.reason) + clear_coldkey_sensitive_caches() + return banned_coldkey + + +@router.delete( + "/banned-coldkeys/{miner_coldkey}", + status_code=status.HTTP_204_NO_CONTENT, + dependencies=[Depends(require_coldkey_ban_admin)], +) +async def delete_banned_coldkey(miner_coldkey: str) -> Response: + validate_coldkey(miner_coldkey) + await unban_coldkey(miner_coldkey) + clear_coldkey_sensitive_caches() + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/api/src/main.py b/api/src/main.py index ad514707..583bd068 100644 --- a/api/src/main.py +++ b/api/src/main.py @@ -12,6 +12,7 @@ from fastapi.middleware.cors import CORSMiddleware import api.config as config +from api.endpoints.admin import router as admin_router from api.endpoints.agent import router as agent_router from api.endpoints.debug import router as debug_router from api.endpoints.evaluation_run import router as evaluation_run_router @@ -123,6 +124,7 @@ async def lifespan(app: FastAPI): app.add_middleware(CorrelationIdMiddleware) app.include_router(upload_router, prefix="/upload") +app.include_router(admin_router, prefix="/admin") app.include_router(retrieval_router, prefix="/retrieval") app.include_router(scoring_router, prefix="/scoring") app.include_router(validator_router, prefix="/validator") From 46639f9e0f56d12aa50c75b8dddc6c195b83447f Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:08:25 -0700 Subject: [PATCH 16/20] alembic migration: - banned_cks table - agents.miner_ck added - queue views: exclude banned cks - remove hk legacy --- .../versions/2026_07_10_add_coldkey_bans.py | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 alembic/versions/2026_07_10_add_coldkey_bans.py diff --git a/alembic/versions/2026_07_10_add_coldkey_bans.py b/alembic/versions/2026_07_10_add_coldkey_bans.py new file mode 100644 index 00000000..8f686238 --- /dev/null +++ b/alembic/versions/2026_07_10_add_coldkey_bans.py @@ -0,0 +1,267 @@ +"""Add coldkey bans and agent coldkey provenance. + +Revision ID: b7e4d2c9a106 +Revises: c8f41e9a2b37 +Create Date: 2026-07-10 00:00:00.000000 + +""" + +import re +from typing import Match, Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +revision: str = "b7e4d2c9a106" +down_revision: Union[str, Sequence[str], None] = "c8f41e9a2b37" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +_HOTKEY_BAN_CONDITION = "agents.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)" +_COLDKEY_BAN_CONDITION = """NOT EXISTS ( + SELECT 1 + FROM banned_coldkeys + WHERE banned_coldkeys.miner_coldkey = agents.miner_coldkey + )""" + +_REFRESH_AGENT_HOTKEY_FILTER = " AND a.miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)\n" +_REFRESH_AGENT_UNAPPROVED_FILTER = " AND a.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)\n" +_POPULATE_HOTKEY_FILTER = ( + " WHERE miner_hotkey NOT IN (SELECT miner_hotkey FROM banned_hotkeys)\n" + " AND agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)" +) +_POPULATE_UNAPPROVED_FILTER = " WHERE agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids)" + + +def _replace_queue_views(ban_condition: str) -> None: + op.execute(f""" + CREATE OR REPLACE VIEW pre_screening_queue AS + SELECT agents.agent_id, agents.status + FROM agents + WHERE agents.status IN ('pre_screening', 'pre_screening_needs_review') + AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND {ban_condition} + AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) + ORDER BY agents.created_at ASC; + """) + + for queue_name, status, set_group in ( + ("screener_1_queue", "screening_1", "screener_1"), + ("screener_2_queue", "screening_2", "screener_2"), + ): + op.execute(f""" + CREATE OR REPLACE VIEW {queue_name} AS + SELECT agents.agent_id, agents.status + FROM agents + WHERE agents.status = '{status}' + AND NOT EXISTS ( + SELECT 1 FROM evaluations e + WHERE e.agent_id = agents.agent_id + AND e.evaluation_set_group = '{set_group}'::evaluationsetgroup + AND ( + SELECT (CASE + WHEN COUNT(*) = 0 THEN NULL + WHEN EVERY(erh.status = 'finished' OR (erh.status = 'error' AND erh.error_code BETWEEN 1000 AND 1999)) THEN 'success' + WHEN EVERY(erh.status IN ('finished', 'error')) THEN 'failure' + ELSE 'running' + END)::evaluationstatus + FROM evaluation_runs_hydrated erh + WHERE erh.evaluation_id = e.evaluation_id + ) IN ('success', 'running') + ) + AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND {ban_condition} + AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) + ORDER BY agents.created_at ASC; + """) + + op.execute(f""" + CREATE OR REPLACE VIEW validator_queue AS + WITH + validator_eval_counts AS ( + SELECT + agent_id, + COUNT(*) FILTER (WHERE status = 'running') AS num_running_evals, + COUNT(*) FILTER (WHERE status = 'success') AS num_finished_evals + FROM evaluations_hydrated + WHERE evaluations_hydrated.status IN ('success', 'running') + AND evaluations_hydrated.evaluation_set_group = 'validator'::evaluationsetgroup + GROUP BY agent_id + ), + screener_2_scores AS ( + SELECT agent_id, MAX(score) AS score FROM evaluations_hydrated + WHERE evaluations_hydrated.evaluation_set_group = 'screener_2'::evaluationsetgroup + AND evaluations_hydrated.status = 'success' + GROUP BY agent_id + ) + SELECT + agent_id, + status, + COALESCE(num_running_evals, 0) as num_running_evals, + COALESCE(num_finished_evals, 0) as num_finished_evals + FROM agents + INNER JOIN screener_2_scores USING (agent_id) + LEFT JOIN validator_eval_counts USING (agent_id) + WHERE + agents.status = 'evaluating' + AND COALESCE(num_running_evals, 0) + COALESCE(num_finished_evals, 0) < 3 + AND agents.agent_id NOT IN (SELECT agent_id FROM benchmark_agent_ids) + AND {ban_condition} + AND agents.agent_id NOT IN (SELECT agent_id FROM unapproved_agent_ids) + ORDER BY + screener_2_scores.score DESC, + agents.created_at ASC, + num_finished_evals DESC; + """) + + +def _replace_function_fragment(signature: str, old: str, new: str) -> None: + definition = ( + op.get_bind() + .execute( + sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"), + {"signature": signature}, + ) + .scalar_one() + ) + if definition is None or definition.count(old) != 1: + raise RuntimeError(f"Unexpected definition for {signature}") + op.execute(definition.replace(old, new)) + + +def _replace_function_pattern(signature: str, pattern: str, replacement: str) -> None: + definition = ( + op.get_bind() + .execute( + sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"), + {"signature": signature}, + ) + .scalar_one() + ) + if definition is None: + raise RuntimeError(f"Missing function {signature}") + updated, replacement_count = re.subn(pattern, replacement, definition, flags=re.MULTILINE) + if replacement_count != 1: + raise RuntimeError(f"Unexpected definition for {signature}") + op.execute(updated) + + +def _hotkey_trigger_branch_pattern(row: str) -> str: + return ( + rf"^(?P[ \t]*)ELSIF TG_TABLE_NAME = 'banned_hotkeys' THEN\n" + rf"(?P=indent) DELETE FROM agent_scores\n" + rf"(?P=indent) WHERE agent_id IN " + rf"\(SELECT agent_id FROM agents WHERE miner_hotkey = {row}\.miner_hotkey\);\n" + rf"(?P=indent) RETURN {row};\n" + ) + + +def _restore_hotkey_trigger_branch(row: str) -> None: + signature = "refresh_agent_scores()" + definition = ( + op.get_bind() + .execute( + sa.text("SELECT pg_get_functiondef(to_regprocedure(:signature))"), + {"signature": signature}, + ) + .scalar_one() + ) + if definition is None: + raise RuntimeError(f"Missing function {signature}") + + pattern = ( + rf"^(?P[ \t]*)ELSIF TG_TABLE_NAME = 'approved_agents' THEN\n" + rf"(?P=indent) affected_agent_id := {row}\.agent_id;\n" + rf"(?P=indent)ELSIF TG_TABLE_NAME = 'unapproved_agent_ids' THEN\n" + ) + + def add_branch(match: Match[str]) -> str: + indent = match.group("indent") + return ( + f"{indent}ELSIF TG_TABLE_NAME = 'approved_agents' THEN\n" + f"{indent} affected_agent_id := {row}.agent_id;\n" + f"{indent}ELSIF TG_TABLE_NAME = 'banned_hotkeys' THEN\n" + f"{indent} DELETE FROM agent_scores\n" + f"{indent} WHERE agent_id IN " + f"(SELECT agent_id FROM agents WHERE miner_hotkey = {row}.miner_hotkey);\n" + f"{indent} RETURN {row};\n" + f"{indent}ELSIF TG_TABLE_NAME = 'unapproved_agent_ids' THEN\n" + ) + + updated, replacement_count = re.subn(pattern, add_branch, definition, flags=re.MULTILINE) + if replacement_count != 1: + raise RuntimeError(f"Unexpected definition for {signature}") + op.execute(updated) + + +def _remove_hotkey_score_behavior() -> None: + _replace_function_fragment( + "refresh_agent_scores_for_agent(uuid)", + _REFRESH_AGENT_HOTKEY_FILTER, + "", + ) + _replace_function_fragment( + "populate_agent_scores()", + _POPULATE_HOTKEY_FILTER, + _POPULATE_UNAPPROVED_FILTER, + ) + _replace_function_pattern( + "refresh_agent_scores()", + _hotkey_trigger_branch_pattern("OLD"), + "", + ) + _replace_function_pattern( + "refresh_agent_scores()", + _hotkey_trigger_branch_pattern("NEW"), + "", + ) + + +def _restore_hotkey_score_behavior() -> None: + _replace_function_fragment( + "refresh_agent_scores_for_agent(uuid)", + _REFRESH_AGENT_UNAPPROVED_FILTER, + _REFRESH_AGENT_HOTKEY_FILTER + _REFRESH_AGENT_UNAPPROVED_FILTER, + ) + _replace_function_fragment( + "populate_agent_scores()", + _POPULATE_UNAPPROVED_FILTER, + _POPULATE_HOTKEY_FILTER, + ) + _restore_hotkey_trigger_branch("OLD") + _restore_hotkey_trigger_branch("NEW") + + +def upgrade() -> None: + op.create_table( + "banned_coldkeys", + sa.Column("miner_coldkey", sa.Text(), primary_key=True), + sa.Column("banned_reason", sa.Text(), nullable=False), + sa.Column( + "banned_at", + sa.TIMESTAMP(timezone=True), + server_default=sa.text("NOW()"), + nullable=False, + ), + ) + + op.add_column("agents", sa.Column("miner_coldkey", sa.Text(), nullable=True)) + + _replace_queue_views(_COLDKEY_BAN_CONDITION) + _remove_hotkey_score_behavior() + op.execute("DROP TRIGGER IF EXISTS tr_refresh_agent_scores_banned_hotkeys ON banned_hotkeys;") + + +def downgrade() -> None: + _restore_hotkey_score_behavior() + op.execute(""" + CREATE TRIGGER tr_refresh_agent_scores_banned_hotkeys + AFTER INSERT OR UPDATE OR DELETE ON banned_hotkeys + FOR EACH ROW EXECUTE PROCEDURE refresh_agent_scores(); + """) + _replace_queue_views(_HOTKEY_BAN_CONDITION) + + op.drop_column("agents", "miner_coldkey") + op.drop_table("banned_coldkeys") From b1775565da3139784d4c4d4a82582c63cf43b1ff Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:10:29 -0700 Subject: [PATCH 17/20] update util --- api/.env.example | 1 + utils/ttl.py | 36 +++++++++++++++++++++++++----------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/api/.env.example b/api/.env.example index 1335e5dc..5a15d9f1 100644 --- a/api/.env.example +++ b/api/.env.example @@ -8,6 +8,7 @@ SUBTENSOR_ADDRESS=wss://entrypoint-finney.opentensor.ai SUBTENSOR_NETWORK=finney OWNER_HOTKEY= +COLDKEY_BAN_ADMIN_API_KEY= UPLOAD_SEND_ADDRESS= BURN=False DISALLOW_UPLOADS=False diff --git a/utils/ttl.py b/utils/ttl.py index e6b1d3e9..8a9571a0 100644 --- a/utils/ttl.py +++ b/utils/ttl.py @@ -51,6 +51,7 @@ def ttl_cache(ttl_seconds: int, max_entries: int = 200): def decorator(func: Callable): cache: Dict[TTLCacheKey, TTLCacheEntry] = {} recalculating_locks: Dict[TTLCacheKey, asyncio.Lock] = {} + cache_generation = 0 def _evict_expired(): """Remove all expired entries and their locks.""" @@ -98,25 +99,38 @@ async def wrapper(*args, **kwargs): logger.debug( f"[TTLCache] {func.__name__}(): First request, already calculating, stopped waiting" ) + if key in cache: return cache[key].value - else: - logger.debug(f"[TTLCache] {func.__name__}(): First request, triggering calculation") - await _recalculate(args, kwargs) - return cache[key].value + + logger.debug(f"[TTLCache] {func.__name__}(): First request, triggering calculation") + return await _recalculate(args, kwargs) + + def cache_clear() -> None: + nonlocal cache_generation + cache_generation += 1 + cache.clear() + for key, lock in list(recalculating_locks.items()): + if not lock.locked(): + recalculating_locks.pop(key, None) async def _recalculate(args: Tuple, kwargs: Dict): key = _args_and_kwargs_to_ttl_cache_key(args, kwargs) + lock = recalculating_locks.setdefault(key, asyncio.Lock()) - async with recalculating_locks[key]: + async with lock: if key in cache and datetime.now(timezone.utc) < cache[key].expires_at: - return + return cache[key].value + calculation_generation = cache_generation value = await func(*args, **kwargs) - cache[key] = TTLCacheEntry( - expires_at=datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds), value=value - ) - logger.debug(f"[TTLCache] {func.__name__}(): Calculation completed") - + if calculation_generation == cache_generation: + cache[key] = TTLCacheEntry( + expires_at=datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds), value=value + ) + logger.debug(f"[TTLCache] {func.__name__}(): Calculation completed") + return value + + wrapper.cache_clear = cache_clear return wrapper return decorator From e00eb70ebb79456c3debd3db12e4b60644d99508 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:12:04 -0700 Subject: [PATCH 18/20] tests: update existing --- tests/api/test_auto_approval.py | 25 +++++ tests/api/test_evaluation_sets.py | 82 +++++++++++++++- tests/api/test_openrouter_secrets.py | 7 +- tests/api/test_upload.py | 141 ++++++++++++++++++++++++++- tests/conftest.py | 1 + tests/queries/test_agent.py | 77 +++++++++++++-- tests/queries/test_scores.py | 137 +++++++++++++++++++++++++- 7 files changed, 451 insertions(+), 19 deletions(-) diff --git a/tests/api/test_auto_approval.py b/tests/api/test_auto_approval.py index 6aaa82c2..13b28c61 100644 --- a/tests/api/test_auto_approval.py +++ b/tests/api/test_auto_approval.py @@ -12,6 +12,14 @@ from queries.evaluation import AgentRankingProfile +@pytest.fixture(autouse=True) +def default_candidate_is_not_banned(monkeypatch): + async def not_banned(_agent_id) -> bool: + return False + + monkeypatch.setattr(validator_endpoint, "is_agent_coldkey_banned", not_banned) + + def _hydrated_evaluation(*, agent_id, set_id: int = 7) -> HydratedEvaluation: return HydratedEvaluation( evaluation_id=uuid4(), @@ -169,6 +177,23 @@ async def fake_get_approved_leader_ranking_for_set(_set_id, _excluded_agent_id, assert await validator_endpoint._should_run_auto_approval_judge(agent_id=agent_id, set_id=11) is True +@pytest.mark.anyio +async def test_should_run_auto_approval_judge_rejects_banned_candidate_before_scoring(monkeypatch) -> None: + agent_id = uuid4() + + async def is_banned(_agent_id) -> bool: + assert _agent_id == agent_id + return True + + async def unexpected_score_lookup(*_args): + raise AssertionError("A banned candidate should be rejected before score lookup") + + monkeypatch.setattr(validator_endpoint, "is_agent_coldkey_banned", is_banned) + monkeypatch.setattr(validator_endpoint, "get_validator_agent_score_for_set", unexpected_score_lookup) + + assert await validator_endpoint._should_run_auto_approval_judge(agent_id=agent_id, set_id=11) is False + + @pytest.mark.anyio async def test_should_run_auto_approval_judge_rejects_full_tie_with_leader(monkeypatch) -> None: agent_id = uuid4() diff --git a/tests/api/test_evaluation_sets.py b/tests/api/test_evaluation_sets.py index 7dd4ec37..956cd641 100644 --- a/tests/api/test_evaluation_sets.py +++ b/tests/api/test_evaluation_sets.py @@ -13,7 +13,8 @@ async def clean_tables(postgres_db): yield async with _db.pool.acquire() as conn: await conn.execute( - "TRUNCATE evaluation_sets, agents, agent_scores, evaluations, approved_agents, competitions, benchmark_agent_ids RESTART IDENTITY CASCADE" + "TRUNCATE evaluation_sets, agents, agent_scores, evaluations, approved_agents, competitions, " + "benchmark_agent_ids, banned_coldkeys RESTART IDENTITY CASCADE" ) @@ -46,13 +47,24 @@ async def _insert_eval_set(conn, set_id: int, created_at: datetime) -> None: async def _insert_agent( - conn, *, agent_id, miner_hotkey: str, status: str, created_at: datetime, set_id: int | None = None + conn, + *, + agent_id, + miner_hotkey: str, + status: str, + created_at: datetime, + set_id: int | None = None, + miner_coldkey: str | None = None, ) -> None: await conn.execute( - """INSERT INTO agents (agent_id, miner_hotkey, name, version_num, status, created_at, ip_address, set_id) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8)""", + """INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address, set_id + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)""", agent_id, miner_hotkey, + miner_coldkey, miner_hotkey, 1, status, @@ -467,6 +479,68 @@ async def test_evaluation_set_leaderboard_ranks_by_score_cost_then_submission_ti assert result.efficiency.average_agent_runtime_seconds == 36 +@pytest.mark.anyio +async def test_coldkey_ban_disqualifies_competitor_and_removes_approved_output(): + banned_agent = uuid4() + eligible_agent = uuid4() + + async with _db.pool.acquire() as conn: + await _insert_eval_set(conn, set_id=2, created_at=SET_2_CREATED) + await _insert_agent( + conn, + agent_id=banned_agent, + miner_hotkey="banned-hotkey", + miner_coldkey="banned-coldkey", + status="finished", + created_at=AGENT_TS_SET_2, + set_id=2, + ) + await _insert_agent( + conn, + agent_id=eligible_agent, + miner_hotkey="eligible-hotkey", + miner_coldkey="eligible-coldkey", + status="finished", + created_at=AGENT_TS_SET_2 + timedelta(minutes=1), + set_id=2, + ) + await _insert_approved_agent(conn, agent_id=banned_agent, set_id=2) + await _insert_approved_agent(conn, agent_id=eligible_agent, set_id=2) + await _insert_agent_score( + conn, + agent_id=banned_agent, + miner_hotkey="banned-hotkey", + set_id=2, + final_score=0.9, + ) + await _insert_agent_score( + conn, + agent_id=eligible_agent, + miner_hotkey="eligible-hotkey", + set_id=2, + final_score=0.8, + ) + await conn.execute( + "INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) VALUES ('banned-coldkey', 'test ban')" + ) + + leaderboard = await evaluation_sets_endpoint.evaluation_set_leaderboard(set_id=2) + by_id = {agent.agent_id: agent for agent in leaderboard} + assert by_id[banned_agent].disqualified is True + assert by_id[banned_agent].rank is None + assert by_id[eligible_agent].disqualified is False + assert by_id[eligible_agent].rank == 1 + + detail = await evaluation_sets_endpoint.evaluation_set_detail(set_id=2) + assert detail.scores.best == 0.8 + assert detail.submissions.approved_emission_count == 1 + assert detail.top_agent is not None + assert detail.top_agent.agent_id == eligible_agent + + approved_agents = await evaluation_sets_endpoint.evaluation_set_approved_agents(set_id=2) + assert [agent.id for agent in approved_agents] == [eligible_agent] + + @pytest.mark.anyio async def test_evaluation_set_detail_efficiency_uses_all_ranked_agents_not_top_25_only(): async with _db.pool.acquire() as conn: diff --git a/tests/api/test_openrouter_secrets.py b/tests/api/test_openrouter_secrets.py index 60c85a23..714ac6bc 100644 --- a/tests/api/test_openrouter_secrets.py +++ b/tests/api/test_openrouter_secrets.py @@ -160,7 +160,7 @@ async def fake_record_upload_attempt(*_args, **_kwargs) -> None: async def fake_check_hotkey_registered(_hotkey: str) -> None: return None - async def fake_check_agent_banned(*, miner_hotkey: str) -> None: + async def fake_check_coldkey_banned(_miner_coldkey: str) -> None: return None async def fake_get_upload_price(*args, **kwargs): @@ -184,7 +184,7 @@ async def fake_create_payment_quote(*, miner_hotkey: str, amount_alpha_rao: int, ) monkeypatch.setattr(upload_endpoint, "record_upload_attempt", fake_record_upload_attempt) monkeypatch.setattr(upload_endpoint, "check_hotkey_registered", fake_check_hotkey_registered) - monkeypatch.setattr(upload_endpoint, "check_agent_banned", fake_check_agent_banned) + monkeypatch.setattr(upload_endpoint, "check_coldkey_banned", fake_check_coldkey_banned) monkeypatch.setattr(upload_endpoint, "get_upload_price", fake_get_upload_price) monkeypatch.setattr(upload_endpoint, "create_payment_quote", fake_create_payment_quote) monkeypatch.setattr(upload_endpoint, "create_agent", create_agent_impl) @@ -422,6 +422,7 @@ async def fake_create_agent( openrouter_api_key_label=None, openrouter_api_key_creator_user_id=None, openrouter_validated_at=None, + miner_coldkey=None, create_pre_screening_job=False, ) -> None: captured["agent"] = agent @@ -432,6 +433,7 @@ async def fake_create_agent( captured["api_key_label"] = openrouter_api_key_label captured["api_key_creator_user_id"] = openrouter_api_key_creator_user_id captured["validated_at"] = openrouter_validated_at + captured["miner_coldkey"] = miner_coldkey captured["create_pre_screening_job"] = create_pre_screening_job upload_endpoint = _patch_upload_dependencies( @@ -461,6 +463,7 @@ async def fake_create_agent( assert captured["api_key_label"] == validated_keys.api_key_label assert captured["api_key_creator_user_id"] == validated_keys.api_key_creator_user_id assert captured["validated_at"] == validated_keys.validated_at + assert captured["miner_coldkey"] is None assert captured["create_pre_screening_job"] is False diff --git a/tests/api/test_upload.py b/tests/api/test_upload.py index a0c23622..f1aec214 100644 --- a/tests/api/test_upload.py +++ b/tests/api/test_upload.py @@ -5,6 +5,7 @@ One container starts per module; tables are truncated between tests. """ +import asyncio import uuid from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -15,7 +16,10 @@ import utils.database as _db from api.src.endpoints import upload as upload_module from api.src.endpoints.upload import AgentUploadResponse -from queries.agent import _derive_agent_id +from models.agent import AgentCreate, AgentStatus +from queries.agent import _derive_agent_id, create_agent +from queries.banned_coldkey import COLDKEY_BAN_LOCK_NAMESPACE, ban_coldkey +from queries.errors import ColdkeyBannedError from queries.payments import retrieve_payment_by_hash # ── constants ───────────────────────────────────────────────────────────────── @@ -45,7 +49,8 @@ async def clean_tables(postgres_db): yield async with _db.pool.acquire() as conn: await conn.execute( - "TRUNCATE evaluation_payments, upload_payment_quotes, agents, failed_upload_refunds, upload_attempts RESTART IDENTITY CASCADE" + "TRUNCATE evaluation_payments, upload_payment_quotes, agents, banned_coldkeys, " + "failed_upload_refunds, upload_attempts RESTART IDENTITY CASCADE" ) @@ -128,7 +133,6 @@ def _install_mocks(monkeypatch) -> None: """Patch blockchain + S3. prod flag is set by upload_prod_mode.""" monkeypatch.setattr(upload_module, "check_signature", MagicMock()) monkeypatch.setattr(upload_module, "check_hotkey_registered", AsyncMock()) - monkeypatch.setattr(upload_module, "check_agent_banned", AsyncMock()) monkeypatch.setattr( upload_module, "check_if_extrinsic_failed", @@ -318,6 +322,47 @@ async def test_check_agent_persists_payment_quote(): assert row["expires_at"] > row["created_at"] +@pytest.mark.anyio +async def test_check_agent_rejects_banned_coldkey_before_stake_lookup(): + from fastapi import HTTPException + + await ban_coldkey(FAKE_COLDKEY, "test ban") + + with pytest.raises(HTTPException) as exc_info: + await upload_module.check_agent_post( + request=_make_request(), + agent_file=_make_upload_file(), + public_key="deadbeef", + file_info=f"{FAKE_HOTKEY}:0", + signature="fakesig", + name="test-agent", + openrouter_api_key="sk-or-v1-runtime", + openrouter_management_key="sk-or-v1-management", + ) + + assert exc_info.value.status_code == 403 + upload_module.subtensor_client.get_alpha_stake_availability.assert_not_awaited() + + +@pytest.mark.anyio +async def test_check_agent_owner_bypasses_coldkey_ban(monkeypatch): + monkeypatch.setattr(upload_module.config, "OWNER_HOTKEY", FAKE_HOTKEY) + await ban_coldkey(FAKE_COLDKEY, "test ban") + + response = await upload_module.check_agent_post( + request=_make_request(), + agent_file=_make_upload_file(), + public_key="deadbeef", + file_info=f"{FAKE_HOTKEY}:0", + signature="fakesig", + name="owner-agent", + openrouter_api_key="sk-or-v1-runtime", + openrouter_management_key="sk-or-v1-management", + ) + + assert response.status == "success" + + @pytest.mark.anyio @pytest.mark.parametrize( ("position_rao", "total_rao", "locked_rao", "burnable_rao"), @@ -385,6 +430,93 @@ async def test_fresh_upload_creates_completed_payment(): assert payment is not None assert payment.agent_id == _deterministic_id() assert payment.quote_id == quote_id + async with _db.pool.acquire() as conn: + stored_coldkey = await conn.fetchval( + "SELECT miner_coldkey FROM agents WHERE agent_id = $1", + _deterministic_id(), + ) + assert stored_coldkey == FAKE_COLDKEY + + +@pytest.mark.anyio +async def test_final_upload_rejects_coldkey_banned_after_quote(): + from fastapi import HTTPException + + quote_id = await _insert_quote() + await ban_coldkey(FAKE_COLDKEY, "banned after quote") + + with pytest.raises(HTTPException) as exc_info: + await _call_post_agent(quote_id=quote_id) + + assert exc_info.value.status_code == 403 + upload_module.subtensor_client.get_events.assert_not_awaited() + async with _db.pool.acquire() as conn: + assert await conn.fetchval("SELECT count(*) FROM agents") == 0 + + +@pytest.mark.anyio +async def test_agent_insert_rechecks_coldkey_ban_transactionally(monkeypatch): + from fastapi import HTTPException + + quote_id = await _insert_quote() + await ban_coldkey(FAKE_COLDKEY, "authoritative ban") + monkeypatch.setattr(upload_module, "check_coldkey_banned", AsyncMock()) + + with pytest.raises(HTTPException) as exc_info: + await _call_post_agent(quote_id=quote_id) + + assert exc_info.value.status_code == 403 + async with _db.pool.acquire() as conn: + assert await conn.fetchval("SELECT count(*) FROM agents") == 0 + + +@pytest.mark.anyio +async def test_agent_insert_waits_for_concurrent_coldkey_ban(): + agent = AgentCreate( + miner_hotkey=FAKE_HOTKEY, + name="test-agent", + version_num=0, + status=AgentStatus.screening_1, + created_at=datetime.now(timezone.utc), + ip_address="127.0.0.1", + payment_block_hash="concurrent-ban-block", + payment_extrinsic_index="1", + ) + + async with _db.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute( + "SELECT pg_advisory_xact_lock($1, hashtext($2))", + COLDKEY_BAN_LOCK_NAMESPACE, + FAKE_COLDKEY, + ) + create_task = asyncio.create_task( + create_agent( + agent, + "print('test')", + source_sha256="concurrent-ban-source", + runtime_openrouter_api_key_ciphertext=b"runtime", + management_openrouter_api_key_ciphertext=b"management", + openrouter_workspace_id="workspace", + openrouter_api_key_label="label", + openrouter_api_key_creator_user_id="creator", + openrouter_validated_at=datetime.now(timezone.utc), + miner_coldkey=FAKE_COLDKEY, + ) + ) + await asyncio.sleep(0.05) + assert not create_task.done() + await conn.execute( + "INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) VALUES ($1, $2)", + FAKE_COLDKEY, + "concurrent ban", + ) + + with pytest.raises(ColdkeyBannedError): + await asyncio.wait_for(create_task, timeout=2) + + async with _db.pool.acquire() as conn: + assert await conn.fetchval("SELECT count(*) FROM agents") == 0 @pytest.mark.anyio @@ -645,10 +777,11 @@ async def test_owner_bypasses_payment_creates_agent_without_payment_row(): async with _db.pool.acquire() as conn: row = await conn.fetchrow( - "SELECT agent_id FROM agents WHERE miner_hotkey = $1", + "SELECT agent_id, miner_coldkey FROM agents WHERE miner_hotkey = $1", FAKE_OWNER_HOTKEY, ) assert row is not None + assert row["miner_coldkey"] is None @pytest.mark.anyio diff --git a/tests/conftest.py b/tests/conftest.py index 9313b4ab..9af59c9a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,6 +26,7 @@ "AWS_REGION": "us-east-1", "AWS_SECRET_ACCESS_KEY": "test", "BURN": "false", + "COLDKEY_BAN_ADMIN_API_KEY": "test-coldkey-ban-admin-key", "DATABASE_HOST": "localhost", "DATABASE_NAME": "ridges_test", "DATABASE_PASSWORD": "ridges", diff --git a/tests/queries/test_agent.py b/tests/queries/test_agent.py index b34363f7..fd68a403 100644 --- a/tests/queries/test_agent.py +++ b/tests/queries/test_agent.py @@ -23,25 +23,37 @@ async def clean_tables(postgres_db): async with _db.pool.acquire() as conn: await conn.execute( - "TRUNCATE evaluation_runs, evaluations, benchmark_agent_ids, agents RESTART IDENTITY CASCADE" + "TRUNCATE evaluation_runs, evaluations, benchmark_agent_ids, banned_coldkeys, banned_hotkeys, " + "agents RESTART IDENTITY CASCADE" ) yield async with _db.pool.acquire() as conn: await conn.execute( - "TRUNCATE evaluation_runs, evaluations, benchmark_agent_ids, agents RESTART IDENTITY CASCADE" + "TRUNCATE evaluation_runs, evaluations, benchmark_agent_ids, banned_coldkeys, banned_hotkeys, " + "agents RESTART IDENTITY CASCADE" ) -async def _insert_agent(*, status: str = "evaluating", created_at: datetime | None = None) -> uuid.UUID: +async def _insert_agent( + *, + status: str = "evaluating", + created_at: datetime | None = None, + miner_coldkey: str | None = None, + miner_hotkey: str = "test-hotkey", +) -> uuid.UUID: agent_id = uuid.uuid4() async with _db.pool.acquire() as conn: await conn.execute( """ - INSERT INTO agents (agent_id, miner_hotkey, name, version_num, status, created_at, ip_address) - VALUES ($1, $2, $3, $4, $5::agentstatus, $6, $7) + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, $4, $5, $6::agentstatus, $7, $8) """, agent_id, - "test-hotkey", + miner_hotkey, + miner_coldkey, "test-agent", 1, status, @@ -119,6 +131,59 @@ async def test_returns_agent_with_no_evaluations(): assert result == agent_id +@pytest.mark.anyio +@pytest.mark.parametrize( + ("status", "queue_view"), + [ + ("pre_screening", "pre_screening_queue"), + ("screening_1", "screener_1_queue"), + ("screening_2", "screener_2_queue"), + ], +) +async def test_stage_queues_exclude_banned_coldkeys_but_keep_null_coldkeys(status: str, queue_view: str): + banned_agent = await _insert_agent(status=status, miner_coldkey="banned-coldkey", miner_hotkey="banned-hotkey") + null_coldkey_agent = await _insert_agent(status=status, miner_coldkey=None, miner_hotkey="owner-hotkey") + async with _db.pool.acquire() as conn: + await conn.execute( + "INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) VALUES ('banned-coldkey', 'test ban')" + ) + queued = await conn.fetch(f"SELECT agent_id FROM {queue_view}") + + assert {row["agent_id"] for row in queued} == {null_coldkey_agent} + assert banned_agent not in {row["agent_id"] for row in queued} + + +@pytest.mark.anyio +async def test_validator_queue_excludes_banned_coldkey(): + banned_agent = await _insert_agent(miner_coldkey="banned-coldkey", miner_hotkey="banned-hotkey") + eligible_agent = await _insert_agent(miner_coldkey="eligible-coldkey", miner_hotkey="eligible-hotkey") + for agent_id in (banned_agent, eligible_agent): + evaluation_id = await _insert_evaluation(agent_id, group="screener_2") + await _insert_run(evaluation_id, status="finished", verifier_reward=1.0) + + async with _db.pool.acquire() as conn: + await conn.execute( + "INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) VALUES ('banned-coldkey', 'test ban')" + ) + queued = await conn.fetch("SELECT agent_id FROM validator_queue") + + assert {row["agent_id"] for row in queued} == {eligible_agent} + + +@pytest.mark.anyio +async def test_legacy_hotkey_ban_no_longer_removes_agent_from_validator_queue(): + agent_id = await _insert_agent(miner_coldkey="eligible-coldkey", miner_hotkey="legacy-banned-hotkey") + evaluation_id = await _insert_evaluation(agent_id, group="screener_2") + await _insert_run(evaluation_id, status="finished", verifier_reward=1.0) + + async with _db.pool.acquire() as conn: + await conn.execute( + "INSERT INTO banned_hotkeys (miner_hotkey, banned_reason) VALUES ('legacy-banned-hotkey', 'legacy ban')" + ) + + assert await get_next_agent_id_awaiting_evaluation_for_validator_hotkey(HOTKEY) == agent_id + + @pytest.mark.anyio async def test_skips_benchmark_agent(): """Agent in benchmark_agent_ids is excluded even when evaluating.""" diff --git a/tests/queries/test_scores.py b/tests/queries/test_scores.py index dc52c0c0..f22c06d7 100644 --- a/tests/queries/test_scores.py +++ b/tests/queries/test_scores.py @@ -6,13 +6,17 @@ import pytest import utils.database as _db +from queries.agent import get_top_agents +from queries.banned_coldkey import ban_coldkey, unban_coldkey +from queries.evaluation import get_approved_leader_ranking_for_set, get_approved_validator_leader_score_for_set from queries.scores import get_weight_receiving_agent_hotkey, get_weight_receiving_agent_info +from queries.statistics import score_improvement_24_hrs, top_score SET_ID = 23 SET_CREATED_AT = datetime(2026, 6, 1, tzinfo=timezone.utc) TRUNCATE_SCORE_TEST_TABLES = ( "TRUNCATE evaluation_runs, evaluations, agent_scores, evaluation_sets, " - "benchmark_agent_ids, agents RESTART IDENTITY CASCADE" + "benchmark_agent_ids, banned_coldkeys, banned_hotkeys, agents RESTART IDENTITY CASCADE" ) # TODO: Add more edge cases to scoring tests @@ -47,15 +51,17 @@ async def _insert_scored_agent( approved: bool = True, approved_at: datetime | None = None, created_at: datetime, + miner_coldkey: str | None = None, ) -> UUID: agent_id = uuid4() await conn.execute( """ - INSERT INTO agents (agent_id, miner_hotkey, name, version_num, status, created_at, ip_address) - VALUES ($1, $2, $3, 0, 'finished', $4, '127.0.0.1') + INSERT INTO agents (agent_id, miner_hotkey, miner_coldkey, name, version_num, status, created_at, ip_address) + VALUES ($1, $2, $3, $4, 0, 'finished', $5, '127.0.0.1') """, agent_id, miner_hotkey, + miner_coldkey, miner_hotkey, created_at, ) @@ -158,6 +164,131 @@ async def test_weight_receiver_is_top_scored_agent_when_eligible(): assert info["agent_id"] == leader_id +@pytest.mark.anyio +async def test_banned_coldkey_is_skipped_for_incentive(): + now = datetime.now(timezone.utc) + async with _db.pool.acquire() as conn: + await _insert_eval_set(conn) + await _insert_scored_agent( + conn, + miner_hotkey="banned-leader-hotkey", + miner_coldkey="banned-leader-coldkey", + final_score=0.50, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=1), + ) + second_id = await _insert_scored_agent( + conn, + miner_hotkey="eligible-second-hotkey", + miner_coldkey="eligible-second-coldkey", + final_score=0.49, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=2), + ) + + await ban_coldkey("banned-leader-coldkey", "test ban") + + assert await get_weight_receiving_agent_hotkey() == "eligible-second-hotkey" + assert await top_score() == 0.49 + assert await score_improvement_24_hrs() == 0 + info = await get_weight_receiving_agent_info() + assert info is not None + assert info["agent_id"] == second_id + + await unban_coldkey("banned-leader-coldkey") + assert await get_weight_receiving_agent_hotkey() == "banned-leader-hotkey" + + +@pytest.mark.anyio +async def test_top_agents_uses_coldkey_bans_at_read_time(): + now = datetime.now(timezone.utc) + async with _db.pool.acquire() as conn: + await _insert_eval_set(conn) + await _insert_scored_agent( + conn, + miner_hotkey="banned-leader-hotkey", + miner_coldkey="banned-leader-coldkey", + final_score=0.50, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=1), + ) + eligible_id = await _insert_scored_agent( + conn, + miner_hotkey="eligible-second-hotkey", + miner_coldkey="eligible-second-coldkey", + final_score=0.49, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=2), + ) + + await ban_coldkey("banned-leader-coldkey", "test ban") + assert [agent.agent_id for agent in await get_top_agents()] == [eligible_id] + + await unban_coldkey("banned-leader-coldkey") + assert [agent.miner_hotkey for agent in await get_top_agents()] == [ + "banned-leader-hotkey", + "eligible-second-hotkey", + ] + + +@pytest.mark.anyio +async def test_legacy_hotkey_ban_does_not_remove_top_agent_or_delete_score(): + now = datetime.now(timezone.utc) + async with _db.pool.acquire() as conn: + await _insert_eval_set(conn) + leader_id = await _insert_scored_agent( + conn, + miner_hotkey="legacy-banned-hotkey", + miner_coldkey="eligible-coldkey", + final_score=0.50, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=1), + ) + await conn.execute( + "INSERT INTO banned_hotkeys (miner_hotkey, banned_reason) VALUES ('legacy-banned-hotkey', 'legacy ban')" + ) + + assert [agent.agent_id for agent in await get_top_agents()] == [leader_id] + + +@pytest.mark.anyio +async def test_banned_coldkey_is_not_used_as_validator_leader_bar(): + now = datetime.now(timezone.utc) + async with _db.pool.acquire() as conn: + await _insert_eval_set(conn) + await _insert_scored_agent( + conn, + miner_hotkey="banned-leader-hotkey", + miner_coldkey="banned-leader-coldkey", + final_score=0.50, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=1), + ) + await _insert_scored_agent( + conn, + miner_hotkey="eligible-second-hotkey", + miner_coldkey="eligible-second-coldkey", + final_score=0.49, + cost_usd=0.10, + approved_at=now - timedelta(hours=1), + created_at=SET_CREATED_AT + timedelta(hours=2), + ) + + await ban_coldkey("banned-leader-coldkey", "test ban") + + excluded_agent_id = uuid4() + assert await get_approved_validator_leader_score_for_set(SET_ID, excluded_agent_id, 1) == 0.49 + leader = await get_approved_leader_ranking_for_set(SET_ID, excluded_agent_id, 1) + assert leader is not None + assert leader.final_score == 0.49 + + @pytest.mark.anyio async def test_expired_top_scored_agent_burns_instead_of_falling_through(): now = datetime.now(timezone.utc) From 6c8a009d049a2439f6b44366ab49fea263d685e5 Mon Sep 17 00:00:00 2001 From: ibraheem-latent Date: Mon, 13 Jul 2026 17:12:51 -0700 Subject: [PATCH 19/20] tests: add new --- tests/queries/test_coldkey_ban_approval.py | 108 +++++++++++++++++++++ tests/test_ttl.py | 46 +++++++++ 2 files changed, 154 insertions(+) create mode 100644 tests/queries/test_coldkey_ban_approval.py create mode 100644 tests/test_ttl.py diff --git a/tests/queries/test_coldkey_ban_approval.py b/tests/queries/test_coldkey_ban_approval.py new file mode 100644 index 00000000..4cd18f7a --- /dev/null +++ b/tests/queries/test_coldkey_ban_approval.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import asyncio +from uuid import UUID, uuid4 + +import pytest + +import utils.database as _db +from queries.approval import finish_agent_and_enqueue_approval +from queries.banned_coldkey import COLDKEY_BAN_LOCK_NAMESPACE + + +@pytest.fixture(autouse=True) +async def clean_tables(postgres_db): + async with _db.pool.acquire() as conn: + await conn.execute( + "TRUNCATE approval_jobs, agent_approval_states, banned_coldkeys, agents RESTART IDENTITY CASCADE" + ) + yield + async with _db.pool.acquire() as conn: + await conn.execute( + "TRUNCATE approval_jobs, agent_approval_states, banned_coldkeys, agents RESTART IDENTITY CASCADE" + ) + + +async def _insert_evaluating_agent(*, miner_coldkey: str | None) -> UUID: + agent_id = uuid4() + async with _db.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO agents ( + agent_id, miner_hotkey, miner_coldkey, name, version_num, + status, created_at, ip_address + ) + VALUES ($1, $2, $3, 'test-agent', 0, 'evaluating', NOW(), '127.0.0.1') + """, + agent_id, + f"hotkey-{agent_id}", + miner_coldkey, + ) + return agent_id + + +async def _assert_finished_without_job(agent_id: UUID) -> None: + async with _db.pool.acquire() as conn: + status = await conn.fetchval("SELECT status::text FROM agents WHERE agent_id = $1", agent_id) + job_count = await conn.fetchval("SELECT COUNT(*) FROM approval_jobs WHERE agent_id = $1", agent_id) + assert status == "finished" + assert job_count == 0 + + +@pytest.mark.anyio +async def test_ban_winning_publication_race_finishes_agent_without_approval_job() -> None: + miner_coldkey = "race-coldkey" + agent_id = await _insert_evaluating_agent(miner_coldkey=miner_coldkey) + + async with _db.pool.acquire() as ban_conn: + transaction = ban_conn.transaction() + await transaction.start() + await ban_conn.execute( + "SELECT pg_advisory_xact_lock($1, hashtext($2))", + COLDKEY_BAN_LOCK_NAMESPACE, + miner_coldkey, + ) + await ban_conn.execute( + "INSERT INTO banned_coldkeys (miner_coldkey, banned_reason) VALUES ($1, 'test ban')", + miner_coldkey, + ) + + approval_task = asyncio.create_task( + finish_agent_and_enqueue_approval( + agent_id=agent_id, + set_id=7, + policy_version="test-policy", + ) + ) + + for _ in range(100): + if any("pg_advisory_xact_lock" in entry["query"] for entry in _db.DEBUG_QUERIES["running"]): + break + await asyncio.sleep(0.01) + else: + approval_task.cancel() + raise AssertionError("Approval did not reach the coldkey publication lock") + + assert not approval_task.done() + await transaction.commit() + + assert await approval_task is False + await _assert_finished_without_job(agent_id) + + +@pytest.mark.anyio +async def test_null_coldkey_agent_can_still_enqueue_approval() -> None: + agent_id = await _insert_evaluating_agent(miner_coldkey=None) + + enqueued = await finish_agent_and_enqueue_approval( + agent_id=agent_id, + set_id=7, + policy_version="test-policy", + ) + + assert enqueued is True + async with _db.pool.acquire() as conn: + status = await conn.fetchval("SELECT status::text FROM agents WHERE agent_id = $1", agent_id) + job_count = await conn.fetchval("SELECT COUNT(*) FROM approval_jobs WHERE agent_id = $1", agent_id) + assert status == "finished" + assert job_count == 1 diff --git a/tests/test_ttl.py b/tests/test_ttl.py new file mode 100644 index 00000000..53c4411d --- /dev/null +++ b/tests/test_ttl.py @@ -0,0 +1,46 @@ +import asyncio + +import pytest + +from utils.ttl import ttl_cache + + +@pytest.mark.anyio +async def test_cache_clear_forces_recalculation() -> None: + current_value = 1 + + @ttl_cache(ttl_seconds=60) + async def cached_value() -> int: + return current_value + + assert await cached_value() == 1 + current_value = 2 + assert await cached_value() == 1 + + cached_value.cache_clear() + + assert await cached_value() == 2 + + +@pytest.mark.anyio +async def test_cache_clear_does_not_restore_in_flight_stale_value() -> None: + started = asyncio.Event() + release = asyncio.Event() + call_count = 0 + + @ttl_cache(ttl_seconds=60) + async def cached_value() -> int: + nonlocal call_count + call_count += 1 + value = call_count + started.set() + await release.wait() + return value + + first_call = asyncio.create_task(cached_value()) + await started.wait() + cached_value.cache_clear() + release.set() + + assert await first_call == 1 + assert await cached_value() == 2 From 6b6b1c56868149298b23034f2d1c9ae901640093 Mon Sep 17 00:00:00 2001 From: jmnmv12 Date: Tue, 14 Jul 2026 15:08:09 +0100 Subject: [PATCH 20/20] feat: :sparkles: Update the cache invalidation logic to auto register eligible cache invalidation methods --- api/endpoints/admin.py | 24 +++--------------------- utils/ttl.py | 17 ++++++++++++++++- 2 files changed, 19 insertions(+), 22 deletions(-) diff --git a/api/endpoints/admin.py b/api/endpoints/admin.py index 0bf457e8..469f6d92 100644 --- a/api/endpoints/admin.py +++ b/api/endpoints/admin.py @@ -7,9 +7,9 @@ from pydantic import BaseModel, StringConstraints import api.config as config -from api.endpoints import evaluation_sets, retrieval, scoring, statistics from models.banned_coldkey import BannedColdkey from queries.banned_coldkey import ban_coldkey, unban_coldkey +from utils.ttl import clear_all_ttl_caches router = APIRouter(tags=["admin"]) admin_bearer = HTTPBearer(auto_error=False) @@ -40,24 +40,6 @@ def validate_coldkey(miner_coldkey: str) -> None: raise HTTPException(status_code=400, detail="Invalid coldkey SS58 address") from None -def clear_coldkey_sensitive_caches() -> None: - - cached_functions = ( - evaluation_sets._cached_build_detail, - evaluation_sets._cached_build_leaderboard, - evaluation_sets._cached_build_approved_agents, - retrieval.queue, - retrieval.top_agents, - retrieval.top_scores_over_time, - retrieval.perfectly_solved_over_time, - retrieval.network_statistics, - scoring.screener_info, - statistics.problem_statistics, - ) - for cached_function in cached_functions: - cached_function.cache_clear() - - @router.put( "/banned-coldkeys/{miner_coldkey}", response_model=BannedColdkey, @@ -66,7 +48,7 @@ def clear_coldkey_sensitive_caches() -> None: async def put_banned_coldkey(miner_coldkey: str, request: ColdkeyBanRequest) -> BannedColdkey: validate_coldkey(miner_coldkey) banned_coldkey = await ban_coldkey(miner_coldkey, request.reason) - clear_coldkey_sensitive_caches() + clear_all_ttl_caches() return banned_coldkey @@ -78,5 +60,5 @@ async def put_banned_coldkey(miner_coldkey: str, request: ColdkeyBanRequest) -> async def delete_banned_coldkey(miner_coldkey: str) -> Response: validate_coldkey(miner_coldkey) await unban_coldkey(miner_coldkey) - clear_coldkey_sensitive_caches() + clear_all_ttl_caches() return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/utils/ttl.py b/utils/ttl.py index 8a9571a0..06113689 100644 --- a/utils/ttl.py +++ b/utils/ttl.py @@ -2,13 +2,27 @@ import logging from datetime import datetime, timedelta, timezone from functools import wraps -from typing import Any, Callable, Dict, Tuple, TypeAlias +from typing import Any, Callable, Dict, List, Tuple, TypeAlias from pydantic import BaseModel, ConfigDict logger = logging.getLogger(__name__) TTLCacheKey: TypeAlias = Tuple[Any, ...] +# Registry of every ttl_cache's cache_clear, so callers can invalidate all caches +# at once after a mutation that can change any cached, network-wide data. +_ALL_CACHE_CLEARERS: List[Callable[[], None]] = [] + + +def clear_all_ttl_caches() -> None: + """Invalidate every ttl_cache in the process. + + Call after admin-triggered mutations (e.g. coldkey bans) that can change any + cached data. This removes the need to maintain a hand-curated list of ban-sensitive caches. + """ + for clear in _ALL_CACHE_CLEARERS: + clear() + def _args_and_kwargs_to_ttl_cache_key(args: Tuple, kwargs: Dict) -> TTLCacheKey: return (args, tuple(sorted(kwargs.items()))) @@ -131,6 +145,7 @@ async def _recalculate(args: Tuple, kwargs: Dict): return value wrapper.cache_clear = cache_clear + _ALL_CACHE_CLEARERS.append(cache_clear) return wrapper return decorator