Skip to content

Commit 6ba9b3f

Browse files
Pigbibicodex
Pigbibi
and
codex
committed
fix: include pending runs in automation decisions
Co-Authored-By: Codex <noreply@openai.com>
1 parent 3f5fc85 commit 6ba9b3f

2 files changed

Lines changed: 90 additions & 30 deletions

File tree

‎service/ai_gateway_service.py‎

Lines changed: 53 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
from service.auth import authenticate
3737
from service.contracts import (
38+
MODE_REVIEW_AND_FIX,
3839
MODE_REVIEW_ONLY,
3940
TASK_ANALYZE,
4041
TASK_EXECUTE,
@@ -508,7 +509,13 @@ def _public_job_payload(job: dict[str, Any]) -> dict[str, object]:
508509
return payload
509510

510511

511-
def _automation_control_snapshot(repo: str, *, task_name: str = "", requested_mode: str = MODE_REVIEW_ONLY) -> dict[str, Any]:
512+
def _automation_control_snapshot(
513+
repo: str,
514+
*,
515+
task_name: str = "",
516+
requested_mode: str = MODE_REVIEW_ONLY,
517+
pending_run: dict[str, Any] | None = None,
518+
) -> dict[str, Any]:
512519
try:
513520
org_health = read_org_health()
514521
except Exception:
@@ -524,6 +531,8 @@ def _automation_control_snapshot(repo: str, *, task_name: str = "", requested_mo
524531
except Exception:
525532
recent_runs = []
526533
ledger_unavailable = True
534+
if pending_run is not None:
535+
recent_runs = [pending_run, *recent_runs]
527536
execution = decide_automation_execution(
528537
repo=repo or "unknown",
529538
task_name=task_name,
@@ -549,7 +558,11 @@ def _automation_control_snapshot(repo: str, *, task_name: str = "", requested_mo
549558
strict_action = CONTROL_ESCALATE
550559
elif execution.get("action") == EXECUTION_DEFER:
551560
strict_action = CONTROL_PAUSE_AUTO_FIX
552-
elif execution.get("effective_mode") == MODE_REVIEW_ONLY and strict_action == CONTROL_CONTINUE:
561+
elif (
562+
execution.get("requested_mode") == MODE_REVIEW_AND_FIX
563+
and execution.get("effective_mode") == MODE_REVIEW_ONLY
564+
and strict_action == CONTROL_CONTINUE
565+
):
553566
strict_action = CONTROL_REVIEW_ONLY
554567
if strict_action != original_action:
555568
control["action"] = strict_action
@@ -691,24 +704,32 @@ def _automation_triage_snapshot(
691704
def _record_job_automation_run(job: dict[str, Any]) -> None:
692705
try:
693706
repo = str(job.get("source_repository") or job.get("repository") or "unknown")
694-
control = _automation_control_snapshot(repo, task_name=str(job.get("task") or ""), requested_mode=str(job.get("mode") or MODE_REVIEW_ONLY))
707+
task_name = str(job.get("task") or "")
708+
task_state = job_task_state(job)
709+
metadata = {
710+
"origin": "service_job",
711+
"repository": repo,
712+
"source_repository": str(job.get("source_repository") or ""),
713+
"caller_repository": str(job.get("repository") or ""),
714+
"source_ref": str(job.get("source_ref") or ""),
715+
"mode": str(job.get("mode") or ""),
716+
"failure_category": str(job.get("failure_category") or ""),
717+
}
718+
control = _automation_control_snapshot(
719+
repo,
720+
task_name=task_name,
721+
requested_mode=str(job.get("mode") or MODE_REVIEW_ONLY),
722+
pending_run={"task_name": task_name, "task_state": task_state, "metadata": metadata},
723+
)
695724
get_automation_run_ledger().record(
696725
str(job.get("job_id") or ""),
697-
job_task_state(job),
698-
task_name=str(job.get("task") or ""),
726+
task_state,
727+
task_name=task_name,
699728
suggested_action=str(control.get("action") or ""),
700729
service_health=str(control.get("service_health") or ""),
701730
quota_status=str(control.get("quota_status") or ""),
702731
org_health_status=str(control.get("org_health_status") or ""),
703-
metadata={
704-
"origin": "service_job",
705-
"repository": repo,
706-
"source_repository": str(job.get("source_repository") or ""),
707-
"caller_repository": str(job.get("repository") or ""),
708-
"source_ref": str(job.get("source_ref") or ""),
709-
"mode": str(job.get("mode") or ""),
710-
"failure_category": str(job.get("failure_category") or ""),
711-
},
732+
metadata=metadata,
712733
owner_repository=repo,
713734
)
714735
except Exception as exc:
@@ -1548,11 +1569,6 @@ def _handle_record_automation_run(self, claims: dict[str, Any], payload: dict[st
15481569
_validate_source_repo_org(claims, source_repo)
15491570
_assert_source_repository_owner_or_operator(claims, source_repo)
15501571
repo = source_repo or str(claims.get("repository") or "unknown")
1551-
control = _automation_control_snapshot(
1552-
repo,
1553-
task_name=str(payload.get("task") or payload.get("task_name") or ""),
1554-
requested_mode=str(payload.get("mode") or MODE_REVIEW_ONLY),
1555-
)
15561572
metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {}
15571573
ledger = get_automation_run_ledger()
15581574
run_id = str(payload.get("run_id") or payload.get("job_id") or "")
@@ -1567,21 +1583,30 @@ def _handle_record_automation_run(self, claims: dict[str, Any], payload: dict[st
15671583
if existing_metadata.get("origin") == "service_job":
15681584
raise PermissionError("automation run is service-owned")
15691585
_assert_automation_run_access(existing, claims)
1586+
task_name = str(payload.get("task") or payload.get("task_name") or "")
1587+
task_state = str(payload.get("task_state") or payload.get("state") or "running")
1588+
run_metadata = {
1589+
**metadata,
1590+
"origin": "external_workflow",
1591+
"repository": repo,
1592+
"source_repository": source_repo,
1593+
"caller_repository": str(claims.get("repository") or ""),
1594+
}
1595+
control = _automation_control_snapshot(
1596+
repo,
1597+
task_name=task_name,
1598+
requested_mode=str(payload.get("mode") or MODE_REVIEW_ONLY),
1599+
pending_run={"task_name": task_name, "task_state": task_state, "metadata": run_metadata},
1600+
)
15701601
record = get_automation_run_ledger().record(
15711602
run_id,
1572-
str(payload.get("task_state") or payload.get("state") or "running"),
1573-
task_name=str(payload.get("task") or payload.get("task_name") or ""),
1603+
task_state,
1604+
task_name=task_name,
15741605
suggested_action=str(control.get("action") or ""),
15751606
service_health=str(control.get("service_health") or ""),
15761607
quota_status=control.get("quota_status") or "",
15771608
org_health_status=str(control.get("org_health_status") or ""),
1578-
metadata={
1579-
**metadata,
1580-
"origin": "external_workflow",
1581-
"repository": repo,
1582-
"source_repository": source_repo,
1583-
"caller_repository": str(claims.get("repository") or ""),
1584-
},
1609+
metadata=run_metadata,
15851610
owner_repository=repo,
15861611
)
15871612
_json_response(self, HTTPStatus.OK, {"status": "ok", "run": record, "control": control})

‎tests/test_ai_gateway_automation_control.py‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def test_control_snapshot_defaults_to_review_only_for_healthy_repo(self) -> None
2727
):
2828
control = _automation_control_snapshot("QuantStrategyLab/TargetRepo")
2929

30-
self.assertEqual(control["action"], "review_only")
30+
self.assertEqual(control["action"], "continue")
3131
self.assertEqual(control["execution"]["effective_mode"], "review_only")
3232
self.assertFalse(control["execution"]["auto_fix_allowed"])
3333

@@ -76,7 +76,7 @@ def test_control_snapshot_applies_service_owned_execution_policy(self) -> None:
7676
patch("service.ai_gateway_service.get_quota_manager", return_value=quota),
7777
patch("service.ai_gateway_service.get_automation_run_ledger", return_value=ledger),
7878
):
79-
control = _automation_control_snapshot("QuantStrategyLab/TargetRepo")
79+
control = _automation_control_snapshot("QuantStrategyLab/TargetRepo", requested_mode="review_and_fix")
8080

8181
self.assertEqual(control["action"], "review_only")
8282
self.assertEqual(control["execution"]["effective_mode"], "review_only")
@@ -131,6 +131,41 @@ def snapshot(self, limit=100):
131131
self.assertEqual(control["execution"]["action"], "human_review")
132132
self.assertEqual(control["execution"]["consecutive_failures"], 2)
133133

134+
def test_control_snapshot_counts_pending_run_for_failure_threshold(self) -> None:
135+
runs = [
136+
{
137+
"task_name": "monthly",
138+
"task_state": "failed",
139+
"metadata": {"source_repository": "QuantStrategyLab/TargetRepo"},
140+
}
141+
]
142+
health = type("Health", (), {"status": "healthy"})()
143+
quota = type("Quota", (), {"runtime_status": lambda self, repo: {"status": "ok"}})()
144+
ledger = type("Ledger", (), {"snapshot": lambda self, limit=None: {"runs": runs}})()
145+
pending_run = {
146+
"task_name": "monthly",
147+
"task_state": "failed",
148+
"metadata": {"source_repository": "QuantStrategyLab/TargetRepo"},
149+
}
150+
151+
with (
152+
patch("service.ai_gateway_service.read_org_health", return_value={"status": "ok"}),
153+
patch("service.ai_gateway_service.get_health_monitor", return_value=health),
154+
patch("service.ai_gateway_service.get_quota_manager", return_value=quota),
155+
patch("service.ai_gateway_service.get_automation_run_ledger", return_value=ledger),
156+
patch("service.ai_gateway_service.load_execution_policy", return_value={"default": {"max_consecutive_failures": 2}}),
157+
):
158+
control = _automation_control_snapshot(
159+
"QuantStrategyLab/TargetRepo",
160+
task_name="monthly",
161+
requested_mode="review_and_fix",
162+
pending_run=pending_run,
163+
)
164+
165+
self.assertEqual(control["action"], "escalate")
166+
self.assertEqual(control["execution"]["action"], "human_review")
167+
self.assertEqual(control["execution"]["consecutive_failures"], 2)
168+
134169
def test_control_snapshot_fails_closed_when_ledger_is_unavailable(self) -> None:
135170
health = type("Health", (), {"status": "healthy"})()
136171
quota = type("Quota", (), {"runtime_status": lambda self, repo: {"status": "ok"}})()

0 commit comments

Comments
 (0)