Reuse terminal research AI jobs without duplicate model execution

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,20 +1125,57 @@ 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": ""},
             repository=str(claims.get("repository") or ""),
             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")},
     }
     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)

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']}