From e2fb9c86c472cd94268482728635d6364d96019f Mon Sep 17 00:00:00 2001 From: Pigbibi <20649888+Pigbibi@users.noreply.github.com> Date: Thu, 10 Sep 2026 01:20:05 +0800 Subject: [PATCH] Reuse terminal research AI jobs for the same request authority Co-Authored-By: Codex --- service/ai_gateway_service.py | 96 ++++++++++++- tests/test_subscription_execution.py | 206 +++++++++++++++++++++++++++ 2 files changed, 295 insertions(+), 7 deletions(-) diff --git a/service/ai_gateway_service.py b/service/ai_gateway_service.py index b29ce968..0f2f6ede 100644 --- a/service/ai_gateway_service.py +++ b/service/ai_gateway_service.py @@ -28,6 +28,7 @@ import tempfile import threading import time +from collections.abc import Mapping from collections import Counter from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -114,6 +115,16 @@ def try_record_platform_execution(*_args, **_kwargs): DEFAULT_JOB_TTL_SECONDS = 86_400 DEFAULT_JOB_MAX_ACTIVE = 10 ACTIVE_JOB_STATUSES = frozenset({"queued", "running"}) +REUSABLE_RESEARCH_JOB_STATUSES = frozenset({"queued", "running", "succeeded", "failed"}) +REQUEST_AUTHORITY_FIELDS = ( + "repository", + "run_id", + "run_attempt", + "actor", + "ref", + "workflow_ref", + "job_workflow_ref", +) WRITE_AUTH_METHODS = frozenset({"github_oidc", "none"}) TRUSTED_AUTOMATION_PROOF_PATH_ENV = "CODEX_AUDIT_SERVICE_TRUSTED_AUTOMATION_PROOF_PATH" DASHBOARD_REPOSITORIES_ENV = "CODEX_AUDIT_SERVICE_DASHBOARD_REPOSITORIES" @@ -1114,6 +1125,7 @@ def _job_dedupe_key( def _request_job_dedupe_key(claims: dict[str, Any], payload: dict[str, Any]) -> str: """Bind the original request before selecting a provider/model from quota.""" + authority = _request_authority(claims) identity = { "request": _job_dedupe_key( {**payload, "provider": ""}, @@ -1121,6 +1133,7 @@ def _request_job_dedupe_key(claims: dict[str, Any], payload: dict[str, Any]) -> run_id=str(claims.get("run_id") or ""), run_attempt=str(claims.get("run_attempt") or ""), ), + "authority": authority, "allowed_providers": payload.get("allowed_providers", ["codex"]), "prompt_sha256": hashlib.sha256(str(payload.get("prompt") or "").encode()).hexdigest(), **{key: payload.get(key) for key in ("research_stage", "model", "reasoning_effort", "complexity", "changed_files", "changed_lines", "sandbox", "timeout_seconds")}, @@ -1128,6 +1141,41 @@ def _request_job_dedupe_key(claims: dict[str, Any], payload: dict[str, Any]) -> return hashlib.sha256(json.dumps(identity, sort_keys=True, separators=(",", ":")).encode()).hexdigest() +def _request_authority(claims: Mapping[str, Any]) -> dict[str, str]: + return {key: str(claims.get(key) or "") for key in REQUEST_AUTHORITY_FIELDS} + + +def _validate_reusable_research_job( + job: Mapping[str, Any], + *, + claims: Mapping[str, Any], + payload: Mapping[str, Any], +) -> None: + """Fail closed when a persisted result does not match its exact caller/task.""" + authority = _request_authority(claims) + if job.get("request_authority") != authority: + raise PermissionError("persisted research job authority does not match request") + for field in ("repository", "run_id", "run_attempt", "actor"): + if str(job.get(field) or "") != authority[field]: + raise PermissionError("persisted research job authority does not match request") + expected = { + "source_repository": str(payload.get("source_repository") or ""), + "source_ref": str(payload.get("source_ref") or ""), + "task": str(payload.get("task") or TASK_EXECUTE), + "mode": str(payload.get("mode") or MODE_REVIEW_ONLY), + "research_stage": str(payload.get("research_stage") or ""), + } + if any(str(job.get(key) or "") != value for key, value in expected.items()): + raise PermissionError("persisted research job task identity does not match request") + providers = payload.get("allowed_providers", ["codex"]) + if not isinstance(providers, list) or str(job.get("provider") or "") not in providers: + raise PermissionError("persisted research job route does not match request") + for field in ("model", "reasoning_effort"): + requested = str(payload.get(field) or "") + if requested not in {"", "auto"} and str(job.get(field) or "") != requested: + raise PermissionError("persisted research job route does not match request") + + def _classify_codex_exec_failure(text: str) -> str: if any(word in text for word in ("quota", "rate limit", "too many active", "budget")): return "quota_or_capacity_failure" @@ -1165,17 +1213,35 @@ def _classify_failure(error: str) -> str: return "unknown_failure" -def _find_active_job_by_dedupe_key(dedupe_key: str, *, field: str = "dedupe_key") -> dict[str, Any] | None: +def _find_job_by_dedupe_key( + dedupe_key: str, + *, + field: str = "dedupe_key", + statuses: frozenset[str] = ACTIVE_JOB_STATUSES, +) -> dict[str, Any] | None: if os.environ.get("CODEX_AUDIT_SERVICE_DEDUPE_JOBS", "true").strip().lower() in {"0", "false", "no", "off"}: return None + candidates: list[dict[str, Any]] = [] for path in _job_dir().glob("*.json"): try: job = json.loads(path.read_text(encoding="utf-8")) except Exception: continue - if job.get(field) == dedupe_key and job.get("status") in ACTIVE_JOB_STATUSES: - return _mark_stale_job_failed(job) - return None + if job.get(field) != dedupe_key or job.get("status") not in statuses: + continue + job = _mark_stale_job_failed(job) + if job.get("status") in statuses: + candidates.append(job) + if not candidates: + return None + return min( + candidates, + key=lambda job: ( + 0 if job.get("status") in ACTIVE_JOB_STATUSES else 1, + float(job.get("created_at") or 0), + str(job.get("job_id") or ""), + ), + ) def _run_job(job_id: str, payload: dict[str, Any]) -> None: @@ -1270,7 +1336,7 @@ def _submit_job(claims: dict[str, Any], payload: dict[str, Any], *, request_dedu # Admission check/dedupe/cap and create must be atomic: concurrent callers can # otherwise both pass the unlocked checks and start duplicate jobs. with _JOB_WRITE_LOCK: - existing_job = _find_active_job_by_dedupe_key( + existing_job = _find_job_by_dedupe_key( request_dedupe_key or dedupe_key, field="request_dedupe_key" if request_dedupe_key else "dedupe_key", ) @@ -1298,6 +1364,7 @@ def _submit_job(claims: dict[str, Any], payload: dict[str, Any], *, request_dedu "run_id": str(claims.get("run_id") or ""), "run_attempt": str(claims.get("run_attempt") or ""), "actor": str(claims.get("actor") or ""), + "request_authority": _request_authority(claims), "source_repository": str(payload.get("source_repository") or ""), "source_ref": str(payload.get("source_ref") or ""), "task": str(payload.get("task") or TASK_EXECUTE), @@ -1585,8 +1652,23 @@ def _handle_execute_async(self, claims: dict[str, Any], payload: dict[str, Any]) request_key = _request_job_dedupe_key(claims, payload) with _JOB_WRITE_LOCK: _cleanup_expired_jobs() - existing = _find_active_job_by_dedupe_key(request_key, field="request_dedupe_key") - if existing is not None and existing.get("status") in ACTIVE_JOB_STATUSES: + reusable_statuses = ( + REUSABLE_RESEARCH_JOB_STATUSES + if payload.get("research_stage") + else ACTIVE_JOB_STATUSES + ) + existing = _find_job_by_dedupe_key( + request_key, + field="request_dedupe_key", + statuses=reusable_statuses, + ) + if existing is not None: + if payload.get("research_stage"): + _validate_reusable_research_job( + existing, + claims=claims, + payload=payload, + ) job = _public_job_payload(existing) job["deduped"] = True _json_response(self, HTTPStatus.ACCEPTED, job) diff --git a/tests/test_subscription_execution.py b/tests/test_subscription_execution.py index 5f25689e..93ab46ea 100644 --- a/tests/test_subscription_execution.py +++ b/tests/test_subscription_execution.py @@ -264,6 +264,212 @@ def cursor_route(payload, usage, **kwargs): assert cursor.call_count == (1 if first_provider == 'cursor' else 0) +@pytest.mark.parametrize('terminal_status', ['succeeded', 'failed']) +def test_repeated_research_request_reuses_persisted_terminal_before_quota(tmp_path, terminal_status): + from service.quota import QuotaManager + + claims = { + 'repository': 'Synthetic/caller', + 'run_id': 'run-1', + 'run_attempt': '1', + 'actor': 'researcher', + 'ref': 'refs/heads/main', + 'workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + 'job_workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + } + request = { + 'prompt': 'same bounded task', + 'mode': 'review_only', + 'research_stage': 'optimization', + 'allowed_providers': ['codex'], + } + request_key = gateway._request_job_dedupe_key(claims, request) + existing = { + 'job_id': 'p' * 24, + 'status': terminal_status, + 'created_at': 1000, + 'updated_at': 1001, + 'expires_at': 9999, + 'repository': claims['repository'], + 'run_id': claims['run_id'], + 'run_attempt': claims['run_attempt'], + 'actor': claims['actor'], + 'request_authority': {key: claims[key] for key in ( + 'repository', 'run_id', 'run_attempt', 'actor', 'ref', 'workflow_ref', 'job_workflow_ref' + )}, + 'request_dedupe_key': request_key, + 'source_repository': '', + 'source_ref': '', + 'task': 'execute', + 'mode': 'review_only', + 'provider': 'codex', + 'research_stage': 'optimization', + 'model': 'gpt-5.6-sol', + 'reasoning_effort': 'high', + 'output': 'bounded result', + 'failure_category': 'unknown_failure', + 'error': 'sanitized failure', + } + quota = QuotaManager() + with patch.dict('os.environ', { + 'CODEX_AUDIT_SERVICE_JOB_DIR': str(tmp_path / 'jobs'), + 'CODEX_AUDIT_SERVICE_QUOTA_STORE': str(tmp_path / 'quota.json'), + }, clear=True), patch.object(gateway.time, 'time', return_value=2000): + gateway._write_job(existing) + with patch.object(gateway, 'get_quota_manager', return_value=quota), patch.object( + gateway, '_admit_codex_execute' + ) as admission, patch.object(quota, 'record_execute') as record, patch.object( + gateway, '_submit_job' + ) as submit, patch.object(gateway, '_json_response') as response: + gateway.AiGatewayRequestHandler._handle_execute_async(object(), claims, dict(request)) + + admission.assert_not_called() + record.assert_not_called() + submit.assert_not_called() + assert response.call_args.args[1] == 202 + body = response.call_args.args[2] + assert body['job_id'] == 'p' * 24 + assert body['status'] == terminal_status + assert body['deduped'] is True + + +@pytest.mark.parametrize('mismatch', ['authority', 'task', 'route']) +def test_matching_terminal_key_with_mismatched_persisted_identity_fails_closed(tmp_path, mismatch): + from service.quota import QuotaManager + + claims = { + 'repository': 'Synthetic/caller', 'run_id': 'run-1', 'run_attempt': '1', + 'actor': 'researcher', 'ref': 'refs/heads/main', + 'workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + 'job_workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + } + request = {'prompt': 'bounded task', 'mode': 'review_only', 'research_stage': 'optimization'} + job = { + 'job_id': 'm' * 24, 'status': 'succeeded', 'created_at': 1000, + 'updated_at': 1001, 'expires_at': 9999, + 'repository': claims['repository'], 'run_id': claims['run_id'], + 'run_attempt': claims['run_attempt'], 'actor': claims['actor'], + 'request_authority': gateway._request_authority(claims), + 'request_dedupe_key': gateway._request_job_dedupe_key(claims, request), + 'source_repository': '', 'source_ref': '', 'task': 'execute', + 'mode': 'review_only', 'provider': 'codex', 'research_stage': 'optimization', + 'model': 'gpt-5.6-sol', 'reasoning_effort': 'high', 'output': 'bounded result', + } + if mismatch == 'authority': + job['request_authority'] = dict(job['request_authority'], actor='different') + elif mismatch == 'task': + job['research_stage'] = 'drift_analysis' + else: + job['provider'] = 'cursor' + quota = QuotaManager() + with patch.dict('os.environ', { + 'CODEX_AUDIT_SERVICE_JOB_DIR': str(tmp_path / 'jobs'), + 'CODEX_AUDIT_SERVICE_QUOTA_STORE': str(tmp_path / 'quota.json'), + }, clear=True), patch.object(gateway.time, 'time', return_value=2000): + gateway._write_job(job) + with patch.object(gateway, 'get_quota_manager', return_value=quota), patch.object( + gateway, '_admit_codex_execute' + ) as admission, patch.object(quota, 'record_execute') as record, patch.object( + gateway, '_submit_job' + ) as submit: + with pytest.raises(PermissionError, match='persisted research job'): + gateway.AiGatewayRequestHandler._handle_execute_async(object(), claims, dict(request)) + + admission.assert_not_called() + record.assert_not_called() + submit.assert_not_called() + + +@pytest.mark.parametrize('difference', ['run_id', 'actor', 'workflow_ref', 'prompt']) +def test_terminal_reuse_requires_same_authority_and_request(tmp_path, difference): + from service.quota import QuotaManager + + claims = { + 'repository': 'Synthetic/caller', 'run_id': 'run-1', 'run_attempt': '1', + 'actor': 'researcher', 'ref': 'refs/heads/main', + 'workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + 'job_workflow_ref': 'Synthetic/caller/.github/workflows/research.yml@refs/heads/main', + } + request = {'prompt': 'bounded task', 'mode': 'review_only', 'research_stage': 'optimization'} + stored_key = gateway._request_job_dedupe_key(claims, request) + changed_claims = dict(claims) + changed_request = dict(request) + if difference == 'prompt': + changed_request['prompt'] = 'different bounded task' + else: + changed_claims[difference] = 'different' + quota = QuotaManager() + with patch.dict('os.environ', { + 'CODEX_AUDIT_SERVICE_JOB_DIR': str(tmp_path / 'jobs'), + 'CODEX_AUDIT_SERVICE_QUOTA_STORE': str(tmp_path / 'quota.json'), + }, clear=True), patch.object(gateway.time, 'time', return_value=2000): + gateway._write_job({ + 'job_id': 'o' * 24, 'status': 'succeeded', 'created_at': 1000, + 'updated_at': 1001, 'expires_at': 9999, 'request_dedupe_key': stored_key, + }) + with patch.object(gateway, 'get_quota_manager', return_value=quota), patch.object( + gateway, '_admit_codex_execute', return_value=None + ) as admission, patch.object(quota, 'record_execute') as record, patch.object( + gateway, '_submit_job', return_value={'job_id': 'new-job'} + ) as submit, patch.object(gateway, 'get_health_monitor'), patch.object( + gateway, '_json_response' + ) as response: + gateway.AiGatewayRequestHandler._handle_execute_async( + object(), changed_claims, changed_request + ) + + admission.assert_called_once() + record.assert_called_once() + submit.assert_called_once() + assert response.call_args.args[2]['job_id'] == 'new-job' + + +def test_expired_or_recovered_terminal_has_explicit_reuse_semantics(tmp_path): + from service.quota import QuotaManager + + claims = {'repository': 'Synthetic/caller', 'run_id': 'run-1', 'run_attempt': '1'} + request = {'prompt': 'bounded task', 'mode': 'review_only', 'research_stage': 'optimization'} + request_key = gateway._request_job_dedupe_key(claims, request) + quota = QuotaManager() + with patch.dict('os.environ', { + 'CODEX_AUDIT_SERVICE_JOB_DIR': str(tmp_path / 'jobs'), + 'CODEX_AUDIT_SERVICE_QUOTA_STORE': str(tmp_path / 'quota.json'), + }, clear=True), patch.object(gateway.time, 'time', return_value=2000): + gateway._write_job({ + 'job_id': 'e' * 24, 'status': 'succeeded', 'created_at': 900, + 'updated_at': 901, 'expires_at': 1999, 'request_dedupe_key': request_key, + }) + gateway._write_job({ + 'job_id': 'r' * 24, 'status': 'running', 'created_at': 1000, + 'updated_at': 1001, 'expires_at': 9999, 'timeout_seconds': 2700, + 'request_dedupe_key': request_key, + 'repository': claims['repository'], 'run_id': claims['run_id'], + 'run_attempt': claims['run_attempt'], 'actor': '', + 'request_authority': gateway._request_authority(claims), + 'source_repository': '', 'source_ref': '', 'task': 'execute', + 'mode': 'review_only', 'provider': 'codex', + 'research_stage': 'optimization', 'model': 'gpt-5.6-sol', + 'reasoning_effort': 'high', + }) + with patch.object(gateway, '_record_job_automation_run'), patch.object(gateway, '_audit_log'): + assert gateway._recover_orphaned_jobs() == 1 + with patch.object(gateway, 'get_quota_manager', return_value=quota), patch.object( + gateway, '_admit_codex_execute' + ) as admission, patch.object(quota, 'record_execute') as record, patch.object( + gateway, '_submit_job' + ) as submit, patch.object(gateway, '_json_response') as response: + gateway.AiGatewayRequestHandler._handle_execute_async(object(), claims, dict(request)) + + admission.assert_not_called() + record.assert_not_called() + submit.assert_not_called() + body = response.call_args.args[2] + assert body['job_id'] == 'r' * 24 + assert body['status'] == 'failed' + assert body['deduped'] is True + assert not (tmp_path / 'jobs' / f"{'e' * 24}.json").exists() + + def test_original_request_identity_binds_route_inputs_and_ignores_claimed_provider(): claims = {'repository': 'Synthetic/caller', 'run_id': '1'} payload = {'prompt': 'synthetic', 'research_stage': 'drift_analysis', 'allowed_providers': ['codex', 'cursor']}