Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
764bb2f08e | ||
|
|
37814da581 | ||
|
|
6de31e0f3e | ||
|
|
51d7a356ca | ||
|
|
19ce458a1f | ||
|
|
12a0936e65 | ||
|
|
68b95fea0d | ||
|
|
8a07817910 | ||
|
|
a642b04577 | ||
|
|
bb80f8cc2f | ||
|
|
0eb6e1a864 |
@@ -64,6 +64,7 @@ class AgentEventBus:
|
||||
conversation_id=(payload.get('conversation_id') if isinstance(payload, dict) else None),
|
||||
task_id=(payload.get('task_id') if isinstance(payload, dict) else None),
|
||||
run_id=(payload.get('run_id') if isinstance(payload, dict) else None),
|
||||
correlation_id=(payload.get('correlation_id') if isinstance(payload, dict) else None),
|
||||
causation_id=(payload.get('causation_id') if isinstance(payload, dict) else None),
|
||||
sequence=(payload.get('sequence') if isinstance(payload, dict) else None),
|
||||
idempotency_key=(payload.get('idempotency_key') if isinstance(payload, dict) else None),
|
||||
|
||||
@@ -7,7 +7,7 @@ class AgentEventObservabilityBridge:
|
||||
def __init__(self, hooks: AgentObservabilityHooks | None = None):
|
||||
self.hooks = hooks or AgentObservabilityHooks()
|
||||
|
||||
def consume(self, event: dict):
|
||||
def consume(self, event: dict, *, actor=None, trace=None):
|
||||
kind = event.get("kind")
|
||||
payload = event.get("payload", {})
|
||||
if kind == "agent_started":
|
||||
@@ -16,6 +16,8 @@ class AgentEventObservabilityBridge:
|
||||
visible_scope=payload.get("scope", "self"),
|
||||
entry_agent=payload.get("agent", "unknown"),
|
||||
execution_id=event.get("id"),
|
||||
actor=actor,
|
||||
trace=trace,
|
||||
)
|
||||
if kind == "agent_delegated":
|
||||
return self.hooks.delegate(event["execution_id"], payload.get("agent", "unknown"))
|
||||
|
||||
@@ -31,6 +31,18 @@ class AgentExecutionRecord:
|
||||
status: str = "completed"
|
||||
started_at: str = field(default_factory=utc_now)
|
||||
ended_at: str = field(default_factory=utc_now)
|
||||
trace_id: str | None = None
|
||||
span_id: str | None = None
|
||||
parent_span_id: str | None = None
|
||||
correlation_id: str | None = None
|
||||
causation_id: str | None = None
|
||||
tenant_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
task_id: str | None = None
|
||||
claim_id: str | None = None
|
||||
run_id: str | None = None
|
||||
message_id: str | None = None
|
||||
|
||||
def to_dict(self):
|
||||
return asdict(self)
|
||||
|
||||
@@ -9,25 +9,49 @@ try:
|
||||
except Exception:
|
||||
APIRouter = None
|
||||
|
||||
from .agent_observability import AgentExecutionRecord
|
||||
from .agent_observability import AgentExecutionRecord, AgentObservability
|
||||
from .agent_observability_dashboard import AgentObservabilityDashboard
|
||||
from .agent_observability_store import AgentObservabilityStore
|
||||
from .trace_tool_credential_governance import GovernanceActor, TraceContext, authorize_scope
|
||||
|
||||
store = AgentObservabilityStore()
|
||||
dashboard = AgentObservabilityDashboard(store.load())
|
||||
|
||||
|
||||
def record_agent_execution(payload: dict[str, Any]):
|
||||
record = AgentExecutionRecord(**payload)
|
||||
def record_agent_execution(payload: dict[str, Any], *, actor: GovernanceActor | dict[str, Any] | None = None, trace: TraceContext | None = None):
|
||||
scoped = any(payload.get(name) is not None for name in (
|
||||
"trace_id", "span_id", "tenant_id", "project_id", "conversation_id", "task_id", "run_id"
|
||||
))
|
||||
values = dict(payload)
|
||||
if scoped:
|
||||
if actor is None or trace is None:
|
||||
raise RuntimeError("TRACE_ACTOR_CONTEXT_REQUIRED")
|
||||
trusted = GovernanceActor.from_value(actor)
|
||||
authorize_scope(trusted, trace.scope)
|
||||
bound = {
|
||||
"trace_id": trace.trace_id, "span_id": trace.span_id,
|
||||
"parent_span_id": trace.parent_span_id, "correlation_id": trace.correlation_id,
|
||||
"causation_id": trace.causation_id, "tenant_id": trace.scope.tenant_id,
|
||||
"user_id": trace.scope.user_id, "project_id": trace.scope.project_id,
|
||||
"conversation_id": trace.scope.conversation_id, "task_id": trace.task_id,
|
||||
"claim_id": trace.claim_id, "run_id": trace.scope.run_id, "message_id": trace.message_id,
|
||||
}
|
||||
for name, authoritative in bound.items():
|
||||
if values.get(name) not in (None, authoritative):
|
||||
raise RuntimeError("TRACE_SCOPE_MISMATCH_DENY:" + name)
|
||||
values[name] = authoritative
|
||||
record = AgentExecutionRecord(**values)
|
||||
return {"status": "ok", "record": store.append(record)}
|
||||
|
||||
|
||||
def get_agent_observability(user_id: str | None = None, group_id: str | None = None, role: str = "user"):
|
||||
records = store.query(user_id, group_id, role)
|
||||
def get_agent_observability(user_id: str | None = None, group_id: str | None = None, role: str = "user", *, actor: GovernanceActor | dict[str, Any] | None = None):
|
||||
records = store.query(user_id, group_id, role, actor=actor)
|
||||
visible = AgentObservability()
|
||||
visible.records = [AgentExecutionRecord(**record) for record in records]
|
||||
return {
|
||||
"status": "ok",
|
||||
"records": records,
|
||||
"summary": store.load().summary(user_id if role != "owner" else None),
|
||||
"summary": visible.summary(None),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -14,6 +14,11 @@ from uuid import uuid4
|
||||
|
||||
from .agent_observability import AgentExecutionRecord
|
||||
from .agent_observability_store import AgentObservabilityStore
|
||||
from .trace_tool_credential_governance import (
|
||||
GovernanceActor,
|
||||
TraceContext,
|
||||
authorize_scope,
|
||||
)
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
@@ -36,6 +41,7 @@ class AgentRunContext:
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
status: str = "running"
|
||||
trace: TraceContext | None = None
|
||||
|
||||
def delegate(self, agent_name: str) -> None:
|
||||
if agent_name not in self.path:
|
||||
@@ -59,7 +65,19 @@ class AgentObservabilityHooks:
|
||||
visible_scope: str,
|
||||
entry_agent: str,
|
||||
execution_id: str | None = None,
|
||||
*,
|
||||
actor: GovernanceActor | dict[str, Any] | None = None,
|
||||
trace: TraceContext | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if execution_id is not None and execution_id in self.active:
|
||||
raise RuntimeError("OBSERVABILITY_EXECUTION_CONFLICT")
|
||||
if trace is not None:
|
||||
if actor is None:
|
||||
raise RuntimeError("TRACE_ACTOR_CONTEXT_REQUIRED")
|
||||
trusted_actor = GovernanceActor.from_value(actor)
|
||||
authorize_scope(trusted_actor, trace.scope)
|
||||
if trace.scope.user_id is not None and trace.scope.user_id != user_id:
|
||||
raise RuntimeError("TRACE_SCOPE_MISMATCH_DENY:user_id")
|
||||
run = AgentRunContext(
|
||||
execution_id=execution_id or f"exec_{uuid4().hex[:12]}",
|
||||
user_id=user_id,
|
||||
@@ -67,6 +85,7 @@ class AgentObservabilityHooks:
|
||||
entry_agent=entry_agent,
|
||||
path=[entry_agent],
|
||||
used_hermes=entry_agent.lower() == "hermes",
|
||||
trace=trace,
|
||||
)
|
||||
self.active[run.execution_id] = run
|
||||
return {"status": "started", "execution_id": run.execution_id, "run": run.__dict__}
|
||||
@@ -106,6 +125,18 @@ class AgentObservabilityHooks:
|
||||
status=status,
|
||||
started_at=run.started_at,
|
||||
ended_at=ended_at,
|
||||
trace_id=run.trace.trace_id if run.trace else None,
|
||||
span_id=run.trace.span_id if run.trace else None,
|
||||
parent_span_id=run.trace.parent_span_id if run.trace else None,
|
||||
correlation_id=run.trace.correlation_id if run.trace else None,
|
||||
causation_id=run.trace.causation_id if run.trace else None,
|
||||
tenant_id=run.trace.scope.tenant_id if run.trace else None,
|
||||
project_id=run.trace.scope.project_id if run.trace else None,
|
||||
conversation_id=run.trace.scope.conversation_id if run.trace else None,
|
||||
task_id=run.trace.task_id if run.trace else None,
|
||||
claim_id=run.trace.claim_id if run.trace else None,
|
||||
run_id=run.trace.scope.run_id if run.trace else None,
|
||||
message_id=run.trace.message_id if run.trace else None,
|
||||
)
|
||||
return {"status": "stored", "record": self.store.append(record)}
|
||||
|
||||
|
||||
@@ -12,6 +12,12 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .agent_observability import AgentExecutionRecord, AgentObservability
|
||||
from .trace_tool_credential_governance import (
|
||||
GovernanceActor,
|
||||
GovernanceError,
|
||||
GovernanceScope,
|
||||
authorize_scope,
|
||||
)
|
||||
|
||||
|
||||
class AgentObservabilityStore:
|
||||
@@ -23,6 +29,11 @@ class AgentObservabilityStore:
|
||||
def append(self, record: AgentExecutionRecord) -> dict[str, Any]:
|
||||
path = self.records_dir / f"{record.execution_id}.json"
|
||||
data = record.to_dict()
|
||||
if path.exists():
|
||||
existing = json.loads(path.read_text(encoding="utf-8"))
|
||||
if existing == data:
|
||||
return existing
|
||||
raise RuntimeError("OBSERVABILITY_EXECUTION_CONFLICT")
|
||||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return data
|
||||
|
||||
@@ -33,8 +44,36 @@ class AgentObservabilityStore:
|
||||
observability.records.append(AgentExecutionRecord(**raw))
|
||||
return observability
|
||||
|
||||
def query(self, user_id: str | None = None, group_id: str | None = None, role: str = "user") -> list[dict[str, Any]]:
|
||||
def query(self, user_id: str | None = None, group_id: str | None = None, role: str = "user", *, actor: GovernanceActor | dict[str, Any] | None = None) -> list[dict[str, Any]]:
|
||||
records = [record.to_dict() for record in self.load().records]
|
||||
scoped = [record for record in records if record.get("tenant_id") is not None]
|
||||
legacy = [record for record in records if record.get("tenant_id") is None]
|
||||
if actor is not None:
|
||||
trusted = GovernanceActor.from_value(actor)
|
||||
visible: list[dict[str, Any]] = []
|
||||
for record in scoped:
|
||||
try:
|
||||
authorize_scope(
|
||||
trusted,
|
||||
GovernanceScope(
|
||||
"tenant",
|
||||
tenant_id=record.get("tenant_id"),
|
||||
user_id=record.get("user_id"),
|
||||
project_id=record.get("project_id"),
|
||||
conversation_id=record.get("conversation_id"),
|
||||
run_id=record.get("run_id"),
|
||||
),
|
||||
)
|
||||
except GovernanceError:
|
||||
continue
|
||||
visible.append(record)
|
||||
if trusted.privileged:
|
||||
visible.extend(legacy)
|
||||
else:
|
||||
visible.extend(record for record in legacy if record.get("user_id") == trusted.scope.user_id)
|
||||
return visible
|
||||
# Caller-supplied role remains a legacy compatibility path only.
|
||||
records = legacy
|
||||
if role == "owner":
|
||||
return records
|
||||
if role == "group_admin" and group_id:
|
||||
|
||||
@@ -0,0 +1,722 @@
|
||||
"""Owner-controlled CAS-bound execution-envelope rebind for Online Agent Tasks.
|
||||
|
||||
This service exists for a narrow recovery case: a Task was projected/claimed with a read-only
|
||||
source_capture envelope before its already-registered deterministic plan became authoritative.
|
||||
|
||||
It never creates Tasks/Jobs/branches and never accepts arbitrary metadata patches. It only replaces
|
||||
metadata.execution + the derived execution_projection for the same governed binding, after an
|
||||
official NEEDS_REWORK transition has returned the Task to QUEUED with no active lease.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
import hashlib
|
||||
import hmac
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sqlite3
|
||||
from typing import Any, Callable, Mapping, Sequence
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from aios_online_agent_execution import (
|
||||
GovernedExecutionBinding,
|
||||
validate_execution_envelope,
|
||||
)
|
||||
|
||||
TASK_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")
|
||||
KEY_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,159}$")
|
||||
SHA_RE = re.compile(r"^[0-9a-f]{64}$")
|
||||
MAX_BODY = 64 * 1024
|
||||
AUTHORIZATION_ID_ENV = "AIOS_AGENT_TASK_EXECUTION_REBIND_AUTHORIZATION_ID"
|
||||
|
||||
|
||||
class ExecutionRebindError(RuntimeError):
|
||||
def __init__(self, code: str, status: int = 409) -> None:
|
||||
super().__init__(code)
|
||||
self.code = code
|
||||
self.status = status
|
||||
|
||||
|
||||
def _canonical(value: object) -> bytes:
|
||||
return json.dumps(
|
||||
value, sort_keys=True, separators=(",", ":"), ensure_ascii=False
|
||||
).encode("utf-8")
|
||||
|
||||
|
||||
def _sha(value: object) -> str:
|
||||
return hashlib.sha256(_canonical(value)).hexdigest()
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _require_string(
|
||||
value: object,
|
||||
code: str,
|
||||
pattern: re.Pattern[str] | None = None,
|
||||
) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ExecutionRebindError(code, 400)
|
||||
value = value.strip()
|
||||
if pattern is not None and pattern.fullmatch(value) is None:
|
||||
raise ExecutionRebindError(code, 400)
|
||||
return value
|
||||
|
||||
|
||||
def _required_authorization_id() -> str:
|
||||
value = os.environ.get(AUTHORIZATION_ID_ENV, "").strip()
|
||||
if not value:
|
||||
raise RuntimeError("TASK_EXECUTION_REBIND_AUTHORIZATION_ID_MISSING")
|
||||
if KEY_RE.fullmatch(value) is None:
|
||||
raise RuntimeError("TASK_EXECUTION_REBIND_AUTHORIZATION_ID_INVALID")
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecutionRebindRequest:
|
||||
expected_status: str
|
||||
expected_revision: int
|
||||
expected_attempt: int
|
||||
previous_verification_sha256: str
|
||||
expected_old_execution_digest: str
|
||||
expected_new_execution_digest: str
|
||||
authorization_id: str
|
||||
idempotency_key: str
|
||||
reason: str
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, value: Mapping[str, Any]) -> "ExecutionRebindRequest":
|
||||
allowed = {
|
||||
"expected_status",
|
||||
"expected_revision",
|
||||
"expected_attempt",
|
||||
"previous_verification_sha256",
|
||||
"expected_old_execution_digest",
|
||||
"expected_new_execution_digest",
|
||||
"authorization_id",
|
||||
"idempotency_key",
|
||||
"reason",
|
||||
}
|
||||
if set(value) != allowed:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_FIELDS_INVALID", 400)
|
||||
|
||||
status = _require_string(value["expected_status"], "EXPECTED_STATUS_INVALID")
|
||||
if status != "QUEUED":
|
||||
raise ExecutionRebindError("EXPECTED_STATUS_INVALID", 400)
|
||||
|
||||
revision = value["expected_revision"]
|
||||
attempt = value["expected_attempt"]
|
||||
if (
|
||||
not isinstance(revision, int)
|
||||
or isinstance(revision, bool)
|
||||
or revision < 0
|
||||
or not isinstance(attempt, int)
|
||||
or isinstance(attempt, bool)
|
||||
or attempt < 1
|
||||
):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_INTEGER_INVALID", 400)
|
||||
|
||||
verification = _require_string(
|
||||
value["previous_verification_sha256"], "VERIFICATION_SHA256_INVALID"
|
||||
)
|
||||
old_digest = _require_string(
|
||||
value["expected_old_execution_digest"], "OLD_EXECUTION_DIGEST_INVALID"
|
||||
)
|
||||
new_digest = _require_string(
|
||||
value["expected_new_execution_digest"], "NEW_EXECUTION_DIGEST_INVALID"
|
||||
)
|
||||
for digest, code in (
|
||||
(verification, "VERIFICATION_SHA256_INVALID"),
|
||||
(old_digest, "OLD_EXECUTION_DIGEST_INVALID"),
|
||||
(new_digest, "NEW_EXECUTION_DIGEST_INVALID"),
|
||||
):
|
||||
if SHA_RE.fullmatch(digest) is None:
|
||||
raise ExecutionRebindError(code, 400)
|
||||
if old_digest == new_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DIGEST_UNCHANGED", 400)
|
||||
|
||||
authorization = _require_string(
|
||||
value["authorization_id"], "AUTHORIZATION_ID_INVALID", KEY_RE
|
||||
)
|
||||
key = _require_string(value["idempotency_key"], "IDEMPOTENCY_KEY_INVALID", KEY_RE)
|
||||
reason = _require_string(value["reason"], "EXECUTION_REBIND_REASON_INVALID")
|
||||
if len(reason) > 512:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_REASON_INVALID", 400)
|
||||
|
||||
return cls(
|
||||
status,
|
||||
revision,
|
||||
attempt,
|
||||
verification,
|
||||
old_digest,
|
||||
new_digest,
|
||||
authorization,
|
||||
key,
|
||||
reason,
|
||||
)
|
||||
|
||||
def public(self) -> dict[str, Any]:
|
||||
return {
|
||||
"expected_status": self.expected_status,
|
||||
"expected_revision": self.expected_revision,
|
||||
"expected_attempt": self.expected_attempt,
|
||||
"previous_verification_sha256": self.previous_verification_sha256,
|
||||
"expected_old_execution_digest": self.expected_old_execution_digest,
|
||||
"expected_new_execution_digest": self.expected_new_execution_digest,
|
||||
"authorization_id": self.authorization_id,
|
||||
"idempotency_key": self.idempotency_key,
|
||||
"reason": self.reason,
|
||||
}
|
||||
|
||||
|
||||
class CanonicalBridgeExecutionResolver:
|
||||
"""Resolve the already-registered plan from the canonical governed Bridge."""
|
||||
|
||||
def __init__(self, base_url: str, control_token: str) -> None:
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.control_token = control_token
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
task_id: str,
|
||||
required_capabilities: Sequence[str],
|
||||
) -> dict[str, Any]:
|
||||
body = json.dumps(
|
||||
{
|
||||
"task_id": task_id,
|
||||
"required_capabilities": list(required_capabilities),
|
||||
},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
request = urllib.request.Request(
|
||||
self.base_url + "/v2/execution-envelope",
|
||||
method="POST",
|
||||
data=body,
|
||||
headers={
|
||||
"Authorization": "Bearer " + self.control_token,
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=15) as response:
|
||||
raw = response.read()
|
||||
status = response.status
|
||||
except urllib.error.HTTPError as exc:
|
||||
raw = exc.read()
|
||||
status = exc.code
|
||||
try:
|
||||
value = json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise ExecutionRebindError("CANONICAL_EXECUTION_RESPONSE_INVALID", 502) from exc
|
||||
if not isinstance(value, Mapping):
|
||||
raise ExecutionRebindError("CANONICAL_EXECUTION_RESPONSE_INVALID", 502)
|
||||
if status != 200:
|
||||
raise ExecutionRebindError(
|
||||
str(value.get("error") or "CANONICAL_EXECUTION_RESOLUTION_FAILED"),
|
||||
502 if status >= 500 else 409,
|
||||
)
|
||||
return dict(value)
|
||||
|
||||
|
||||
class ExecutionRebindStore:
|
||||
"""Official owner-controlled mutation with CAS, plan resolution and durable audit."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db_path: Path,
|
||||
resolver: Callable[[str, Sequence[str]], Mapping[str, Any]],
|
||||
) -> None:
|
||||
self.db_path = Path(db_path)
|
||||
self.resolver = resolver
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.db_path, timeout=30.0)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys=ON")
|
||||
connection.execute("PRAGMA busy_timeout=30000")
|
||||
return connection
|
||||
|
||||
@staticmethod
|
||||
def _columns(
|
||||
connection: sqlite3.Connection,
|
||||
table: str,
|
||||
) -> dict[str, sqlite3.Row]:
|
||||
return {
|
||||
str(row["name"]): row
|
||||
for row in connection.execute(f'PRAGMA table_info("{table}")')
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _json_value(raw: object, code: str) -> dict[str, Any]:
|
||||
if isinstance(raw, Mapping):
|
||||
return dict(raw)
|
||||
if not isinstance(raw, str):
|
||||
raise ExecutionRebindError(code, 409)
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ExecutionRebindError(code, 409) from exc
|
||||
if not isinstance(value, Mapping):
|
||||
raise ExecutionRebindError(code, 409)
|
||||
return dict(value)
|
||||
|
||||
@staticmethod
|
||||
def _verification_hash(raw: object) -> str:
|
||||
if raw in (None, ""):
|
||||
value: object = {}
|
||||
elif isinstance(raw, str):
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ExecutionRebindError("TASK_VERIFICATION_INVALID", 409) from exc
|
||||
else:
|
||||
value = raw
|
||||
return _sha(value)
|
||||
|
||||
@staticmethod
|
||||
def _metadata_column(columns: Mapping[str, sqlite3.Row]) -> str:
|
||||
for name in ("metadata_json", "metadata"):
|
||||
if name in columns:
|
||||
return name
|
||||
raise ExecutionRebindError("TASK_METADATA_SCHEMA_UNSUPPORTED", 503)
|
||||
|
||||
@staticmethod
|
||||
def _required_capabilities(
|
||||
task: sqlite3.Row,
|
||||
columns: Mapping[str, sqlite3.Row],
|
||||
old_execution: Mapping[str, Any],
|
||||
) -> list[str]:
|
||||
for field in ("required_capabilities_json", "required_capabilities"):
|
||||
if field not in columns:
|
||||
continue
|
||||
raw = task[field]
|
||||
if raw in (None, ""):
|
||||
continue
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
value = [item for item in raw.split(",") if item.strip()]
|
||||
else:
|
||||
value = raw
|
||||
if isinstance(value, list) and all(isinstance(item, str) for item in value):
|
||||
return [item.strip() for item in value if item.strip()]
|
||||
value = old_execution.get("capability_scope")
|
||||
if isinstance(value, list) and all(isinstance(item, str) for item in value):
|
||||
return [item for item in value if item]
|
||||
raise ExecutionRebindError("TASK_CAPABILITIES_UNAVAILABLE", 503)
|
||||
|
||||
@staticmethod
|
||||
def _event_table(
|
||||
connection: sqlite3.Connection,
|
||||
) -> tuple[str, dict[str, sqlite3.Row]]:
|
||||
tables = {
|
||||
row["name"]
|
||||
for row in connection.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table'"
|
||||
)
|
||||
}
|
||||
for name in ("task_events", "events"):
|
||||
if name not in tables:
|
||||
continue
|
||||
columns = ExecutionRebindStore._columns(connection, name)
|
||||
if "task_id" in columns and ("event_type" in columns or "type" in columns):
|
||||
return name, columns
|
||||
raise ExecutionRebindError("TASK_EVENT_TABLE_UNSUPPORTED", 503)
|
||||
|
||||
@staticmethod
|
||||
def _insert_event(
|
||||
connection: sqlite3.Connection,
|
||||
task_id: str,
|
||||
payload: Mapping[str, Any],
|
||||
) -> None:
|
||||
table, columns = ExecutionRebindStore._event_table(connection)
|
||||
values: dict[str, Any] = {"task_id": task_id}
|
||||
if "event_type" in columns:
|
||||
values["event_type"] = "TASK_EXECUTION_REBIND_AUTHORIZED"
|
||||
else:
|
||||
values["type"] = "TASK_EXECUTION_REBIND_AUTHORIZED"
|
||||
payload_text = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
||||
if "payload_json" in columns:
|
||||
values["payload_json"] = payload_text
|
||||
elif "payload" in columns:
|
||||
values["payload"] = payload_text
|
||||
else:
|
||||
raise ExecutionRebindError("TASK_EVENT_PAYLOAD_UNSUPPORTED", 503)
|
||||
if "created_at" in columns:
|
||||
values["created_at"] = _now()
|
||||
if "agent_id" in columns:
|
||||
values["agent_id"] = "owner-control"
|
||||
if "actor_id" in columns:
|
||||
values["actor_id"] = "owner-control"
|
||||
names = list(values)
|
||||
quoted = ",".join('"' + name + '"' for name in names)
|
||||
placeholders = ",".join("?" for _ in names)
|
||||
connection.execute(
|
||||
f'INSERT INTO "{table}" ({quoted}) VALUES ({placeholders})',
|
||||
tuple(values[name] for name in names),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _ensure_rebind_table(connection: sqlite3.Connection) -> None:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS task_execution_rebinds (
|
||||
task_id TEXT NOT NULL,
|
||||
idempotency_key TEXT NOT NULL,
|
||||
request_sha256 TEXT NOT NULL,
|
||||
authorization_id TEXT NOT NULL,
|
||||
previous_verification_sha256 TEXT NOT NULL,
|
||||
old_execution_digest TEXT NOT NULL,
|
||||
new_execution_digest TEXT NOT NULL,
|
||||
expected_revision INTEGER NOT NULL,
|
||||
applied_revision INTEGER NOT NULL,
|
||||
attempt INTEGER NOT NULL,
|
||||
reason TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY(task_id,idempotency_key)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
def _load_and_validate(
|
||||
self,
|
||||
connection: sqlite3.Connection,
|
||||
task_id: str,
|
||||
request: ExecutionRebindRequest,
|
||||
) -> tuple[sqlite3.Row, dict[str, sqlite3.Row], str, dict[str, Any], dict[str, Any]]:
|
||||
columns = self._columns(connection, "tasks")
|
||||
required = {
|
||||
"task_id",
|
||||
"status",
|
||||
"attempt",
|
||||
"revision",
|
||||
"claimed_by",
|
||||
"lease_token_hash",
|
||||
"lease_expires_at",
|
||||
"verification_json",
|
||||
}
|
||||
if not required.issubset(columns):
|
||||
raise ExecutionRebindError("TASK_SCHEMA_EXECUTION_REBIND_UNSUPPORTED", 503)
|
||||
metadata_column = self._metadata_column(columns)
|
||||
task = connection.execute(
|
||||
"SELECT * FROM tasks WHERE task_id=?",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
if task is None:
|
||||
raise ExecutionRebindError("TASK_NOT_FOUND", 404)
|
||||
|
||||
if (
|
||||
str(task["status"]) != request.expected_status
|
||||
or int(task["revision"]) != request.expected_revision
|
||||
or int(task["attempt"]) != request.expected_attempt
|
||||
):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_CAS_MISMATCH", 409)
|
||||
if task["lease_token_hash"] not in (None, "") or task["lease_expires_at"] not in (None, ""):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
|
||||
|
||||
verification_sha = self._verification_hash(task["verification_json"])
|
||||
if verification_sha != request.previous_verification_sha256:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_VERIFICATION_DRIFT", 409)
|
||||
|
||||
metadata = self._json_value(
|
||||
task[metadata_column],
|
||||
"TASK_METADATA_INVALID",
|
||||
)
|
||||
execution = metadata.get("execution")
|
||||
if not isinstance(execution, Mapping):
|
||||
raise ExecutionRebindError("TASK_EXECUTION_METADATA_MISSING", 409)
|
||||
old_execution = validate_execution_envelope(dict(execution))
|
||||
if old_execution["execution_digest"] != request.expected_old_execution_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_OLD_DIGEST_DRIFT", 409)
|
||||
if old_execution.get("action") != "source_capture":
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_SOURCE_CAPTURE_REQUIRED", 409)
|
||||
|
||||
projection = metadata.get("execution_projection")
|
||||
if isinstance(projection, Mapping):
|
||||
projected_digest = str(projection.get("execution_digest") or "")
|
||||
if projected_digest and projected_digest != request.expected_old_execution_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_PROJECTION_DRIFT", 409)
|
||||
return task, columns, metadata_column, metadata, old_execution
|
||||
|
||||
def rebind(
|
||||
self,
|
||||
task_id: str,
|
||||
request: ExecutionRebindRequest,
|
||||
) -> dict[str, Any]:
|
||||
task_id = _require_string(task_id, "TASK_ID_INVALID", TASK_RE)
|
||||
required_authorization = _required_authorization_id()
|
||||
if request.authorization_id != required_authorization:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_AUTHORIZATION_ID_MISMATCH", 403)
|
||||
|
||||
request_sha = _sha(request.public())
|
||||
|
||||
# Idempotent replay must win before CAS validation because a successful
|
||||
# rebind necessarily increments the Task revision and changes the persisted envelope.
|
||||
with self._connect() as connection:
|
||||
self._ensure_rebind_table(connection)
|
||||
replay = connection.execute(
|
||||
"SELECT * FROM task_execution_rebinds WHERE task_id=? AND idempotency_key=?",
|
||||
(task_id, request.idempotency_key),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
if replay["request_sha256"] != request_sha:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_IDEMPOTENCY_CONFLICT", 409)
|
||||
connection.commit()
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "QUEUED",
|
||||
"attempt": int(replay["attempt"]),
|
||||
"revision": int(replay["applied_revision"]),
|
||||
"old_execution_digest": replay["old_execution_digest"],
|
||||
"new_execution_digest": replay["new_execution_digest"],
|
||||
"authorization_id": replay["authorization_id"],
|
||||
"idempotent": True,
|
||||
}
|
||||
|
||||
# Read/CAS snapshot first. Canonical Bridge resolution happens outside the DB write lock.
|
||||
with self._connect() as connection:
|
||||
task, columns, _, _, old_execution = self._load_and_validate(
|
||||
connection, task_id, request
|
||||
)
|
||||
capabilities = self._required_capabilities(task, columns, old_execution)
|
||||
|
||||
binding = GovernedExecutionBinding.from_mapping(old_execution)
|
||||
resolved = dict(self.resolver(task_id, capabilities))
|
||||
new_execution = validate_execution_envelope(resolved, binding=binding)
|
||||
if new_execution.get("action") != "deterministic":
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DETERMINISTIC_PLAN_REQUIRED", 409)
|
||||
if new_execution["execution_digest"] != request.expected_new_execution_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_NEW_DIGEST_DRIFT", 409)
|
||||
if new_execution["execution_digest"] == old_execution["execution_digest"]:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DIGEST_UNCHANGED", 409)
|
||||
|
||||
# Re-read under BEGIN IMMEDIATE and apply only if the exact snapshot is unchanged.
|
||||
with self._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
self._ensure_rebind_table(connection)
|
||||
|
||||
replay = connection.execute(
|
||||
"SELECT * FROM task_execution_rebinds WHERE task_id=? AND idempotency_key=?",
|
||||
(task_id, request.idempotency_key),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
if replay["request_sha256"] != request_sha:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_IDEMPOTENCY_CONFLICT", 409)
|
||||
connection.commit()
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "QUEUED",
|
||||
"attempt": int(replay["attempt"]),
|
||||
"revision": int(replay["applied_revision"]),
|
||||
"old_execution_digest": replay["old_execution_digest"],
|
||||
"new_execution_digest": replay["new_execution_digest"],
|
||||
"authorization_id": replay["authorization_id"],
|
||||
"idempotent": True,
|
||||
}
|
||||
|
||||
task, columns, metadata_column, metadata, current_old = self._load_and_validate(
|
||||
connection, task_id, request
|
||||
)
|
||||
if current_old["execution_digest"] != old_execution["execution_digest"]:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_OLD_DIGEST_DRIFT", 409)
|
||||
|
||||
new_metadata = dict(metadata)
|
||||
new_metadata["execution"] = new_execution
|
||||
new_metadata["execution_projection"] = {
|
||||
"schema": "aios.online-agent.task-projection.v1",
|
||||
"task_id": binding.task_id,
|
||||
"job_id": binding.job_id,
|
||||
"execution_digest": new_execution["execution_digest"],
|
||||
}
|
||||
applied_revision = request.expected_revision + 1
|
||||
previous_claimed_by = task["claimed_by"]
|
||||
|
||||
cursor = connection.execute(
|
||||
f"""
|
||||
UPDATE tasks
|
||||
SET "{metadata_column}"=?, revision=?, claimed_by=NULL,
|
||||
lease_token_hash=NULL, lease_expires_at=NULL
|
||||
WHERE task_id=? AND status=? AND revision=? AND attempt=?
|
||||
AND claimed_by IS ? AND lease_token_hash IS NULL AND lease_expires_at IS NULL
|
||||
""",
|
||||
(
|
||||
json.dumps(new_metadata, sort_keys=True, separators=(",", ":")),
|
||||
applied_revision,
|
||||
task_id,
|
||||
request.expected_status,
|
||||
request.expected_revision,
|
||||
request.expected_attempt,
|
||||
previous_claimed_by,
|
||||
),
|
||||
)
|
||||
if cursor.rowcount != 1:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_CAS_MISMATCH", 409)
|
||||
|
||||
event_payload = {
|
||||
"schema": "aios.task-execution-rebind.event.v1",
|
||||
"task_id": task_id,
|
||||
"job_id": binding.job_id,
|
||||
"authorization_id": request.authorization_id,
|
||||
"idempotency_key": request.idempotency_key,
|
||||
"reason": request.reason,
|
||||
"previous_verification_sha256": request.previous_verification_sha256,
|
||||
"old_execution_digest": old_execution["execution_digest"],
|
||||
"new_execution_digest": new_execution["execution_digest"],
|
||||
"old_action": old_execution.get("action"),
|
||||
"new_action": new_execution.get("action"),
|
||||
"previous_revision": request.expected_revision,
|
||||
"applied_revision": applied_revision,
|
||||
"attempt": request.expected_attempt,
|
||||
"request_sha256": request_sha,
|
||||
"previous_claimed_by": previous_claimed_by,
|
||||
"stale_claim_normalized": previous_claimed_by not in (None, ""),
|
||||
}
|
||||
self._insert_event(connection, task_id, event_payload)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO task_execution_rebinds(
|
||||
task_id,idempotency_key,request_sha256,authorization_id,
|
||||
previous_verification_sha256,old_execution_digest,new_execution_digest,
|
||||
expected_revision,applied_revision,attempt,reason,created_at
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?)
|
||||
""",
|
||||
(
|
||||
task_id,
|
||||
request.idempotency_key,
|
||||
request_sha,
|
||||
request.authorization_id,
|
||||
request.previous_verification_sha256,
|
||||
old_execution["execution_digest"],
|
||||
new_execution["execution_digest"],
|
||||
request.expected_revision,
|
||||
applied_revision,
|
||||
request.expected_attempt,
|
||||
request.reason,
|
||||
_now(),
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "QUEUED",
|
||||
"attempt": request.expected_attempt,
|
||||
"revision": applied_revision,
|
||||
"old_execution_digest": old_execution["execution_digest"],
|
||||
"new_execution_digest": new_execution["execution_digest"],
|
||||
"authorization_id": request.authorization_id,
|
||||
"idempotent": False,
|
||||
}
|
||||
|
||||
|
||||
class ExecutionRebindHandler(BaseHTTPRequestHandler):
|
||||
store: ExecutionRebindStore | None = None
|
||||
control_token: str = ""
|
||||
|
||||
def log_message(self, *_args: Any) -> None:
|
||||
return
|
||||
|
||||
def _json(self, value: Mapping[str, Any], status: int) -> None:
|
||||
data = json.dumps(
|
||||
dict(value), sort_keys=True, separators=(",", ":")
|
||||
).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def _authorized(self) -> bool:
|
||||
header = self.headers.get("Authorization", "")
|
||||
return (
|
||||
header.startswith("Bearer ")
|
||||
and bool(self.control_token)
|
||||
and hmac.compare_digest(header[7:], self.control_token)
|
||||
)
|
||||
|
||||
def do_GET(self) -> None:
|
||||
if self.path == "/health":
|
||||
self._json(
|
||||
{"status": "ok", "service": "aios-agent-task-execution-rebind"},
|
||||
200,
|
||||
)
|
||||
else:
|
||||
self._json({"error": "NOT_FOUND"}, 404)
|
||||
|
||||
def do_POST(self) -> None:
|
||||
match = re.fullmatch(r"/v1/tasks/([^/]+)/execution-rebind", self.path)
|
||||
if match is None:
|
||||
self._json({"error": "NOT_FOUND"}, 404)
|
||||
return
|
||||
if not self._authorized():
|
||||
self._json({"error": "UNAUTHORIZED"}, 401)
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
except ValueError:
|
||||
self._json({"error": "CONTENT_LENGTH_INVALID"}, 400)
|
||||
return
|
||||
if length < 1 or length > MAX_BODY:
|
||||
self._json({"error": "REQUEST_SIZE_INVALID"}, 413)
|
||||
return
|
||||
try:
|
||||
value = json.loads(self.rfile.read(length).decode("utf-8"))
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError
|
||||
request = ExecutionRebindRequest.from_mapping(value)
|
||||
if self.store is None:
|
||||
raise ExecutionRebindError("RUNTIME_UNAVAILABLE", 503)
|
||||
result = self.store.rebind(match.group(1), request)
|
||||
self._json(result, 200 if result["idempotent"] else 201)
|
||||
except ExecutionRebindError as exc:
|
||||
self._json({"error": exc.code}, exc.status)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
|
||||
self._json({"error": "REQUEST_INVALID"}, 400)
|
||||
|
||||
|
||||
def load_control_token(path: Path) -> str:
|
||||
if path.is_symlink() or not path.is_file():
|
||||
raise RuntimeError("TASK_EXECUTION_REBIND_CONTROL_CREDENTIAL_UNAVAILABLE")
|
||||
value = path.read_text(encoding="utf-8").strip()
|
||||
if len(value) < 32:
|
||||
raise RuntimeError("TASK_EXECUTION_REBIND_CONTROL_CREDENTIAL_INVALID")
|
||||
return value
|
||||
|
||||
|
||||
def serve() -> None:
|
||||
db = Path(os.environ.get("AIOS_AGENT_TASK_DB", "/var/lib/aios-agent-task/state.sqlite3"))
|
||||
_required_authorization_id()
|
||||
token_path = Path(
|
||||
os.environ.get(
|
||||
"AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
|
||||
str(
|
||||
Path(os.environ.get("CREDENTIALS_DIRECTORY", "/run/credentials"))
|
||||
/ "agent_task_control_token"
|
||||
),
|
||||
)
|
||||
)
|
||||
token = load_control_token(token_path)
|
||||
bridge = os.environ.get(
|
||||
"AIOS_GOVERNED_WRITE_API_URL",
|
||||
"http://127.0.0.1:18892",
|
||||
)
|
||||
host = os.environ.get("AIOS_TASK_EXECUTION_REBIND_HOST", "127.0.0.1")
|
||||
port = int(os.environ.get("AIOS_TASK_EXECUTION_REBIND_PORT", "18942"))
|
||||
ExecutionRebindHandler.store = ExecutionRebindStore(
|
||||
db,
|
||||
CanonicalBridgeExecutionResolver(bridge, token),
|
||||
)
|
||||
ExecutionRebindHandler.control_token = token
|
||||
ThreadingHTTPServer((host, port), ExecutionRebindHandler).serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
serve()
|
||||
@@ -0,0 +1,427 @@
|
||||
"""Owner-controlled, CAS-bound retry extension for exhausted Online Agent Tasks."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
import hashlib
|
||||
import hmac
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sqlite3
|
||||
from typing import Any, Mapping
|
||||
|
||||
TASK_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")
|
||||
KEY_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,159}$")
|
||||
MAX_BODY = 64 * 1024
|
||||
AUTHORIZE_EXTENSION_ID_ENV = "AIOS_AGENT_TASK_AUTHORIZATION_ID"
|
||||
|
||||
|
||||
class RetryExtensionError(RuntimeError):
|
||||
def __init__(self, code: str, status: int = 409) -> None:
|
||||
super().__init__(code)
|
||||
self.code = code
|
||||
self.status = status
|
||||
|
||||
|
||||
def _canonical(value: object) -> bytes:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
|
||||
|
||||
def _sha(value: object) -> str:
|
||||
return hashlib.sha256(_canonical(value)).hexdigest()
|
||||
|
||||
|
||||
def _required_authorization_id() -> str:
|
||||
value = os.environ.get(AUTHORIZE_EXTENSION_ID_ENV, "").strip()
|
||||
if not value:
|
||||
raise RuntimeError("TASK_RETRY_AUTHORIZATION_ID_MISSING")
|
||||
if KEY_RE.fullmatch(value) is None:
|
||||
raise RuntimeError("TASK_RETRY_AUTHORIZATION_ID_INVALID")
|
||||
return value
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _require_string(value: object, code: str, pattern: re.Pattern[str] | None = None) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise RetryExtensionError(code, 400)
|
||||
value = value.strip()
|
||||
if pattern is not None and pattern.fullmatch(value) is None:
|
||||
raise RetryExtensionError(code, 400)
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RetryExtensionRequest:
|
||||
expected_status: str
|
||||
expected_revision: int
|
||||
expected_attempt: int
|
||||
expected_max_attempts: int
|
||||
next_max_attempts: int
|
||||
previous_verification_sha256: str
|
||||
authorization_id: str
|
||||
idempotency_key: str
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, value: Mapping[str, Any]) -> "RetryExtensionRequest":
|
||||
allowed = {
|
||||
"expected_status", "expected_revision", "expected_attempt", "expected_max_attempts",
|
||||
"next_max_attempts", "previous_verification_sha256", "authorization_id", "idempotency_key",
|
||||
}
|
||||
if set(value) != allowed:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_FIELDS_INVALID", 400)
|
||||
status = _require_string(value["expected_status"], "EXPECTED_STATUS_INVALID")
|
||||
if status != "FAILED":
|
||||
raise RetryExtensionError("EXPECTED_STATUS_INVALID", 400)
|
||||
ints: dict[str, int] = {}
|
||||
for field in ("expected_revision", "expected_attempt", "expected_max_attempts", "next_max_attempts"):
|
||||
item = value[field]
|
||||
if not isinstance(item, int) or isinstance(item, bool) or item < 0:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_INTEGER_INVALID", 400)
|
||||
ints[field] = item
|
||||
if ints["expected_attempt"] != ints["expected_max_attempts"]:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_NOT_EXHAUSTED", 409)
|
||||
if ints["next_max_attempts"] != ints["expected_max_attempts"] + 1:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_MUST_BE_BOUNDED_ONE_STEP", 400)
|
||||
if ints["next_max_attempts"] > 50:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_LIMIT_EXCEEDED", 400)
|
||||
verification = _require_string(value["previous_verification_sha256"], "VERIFICATION_SHA256_INVALID")
|
||||
if re.fullmatch(r"[0-9a-f]{64}", verification) is None:
|
||||
raise RetryExtensionError("VERIFICATION_SHA256_INVALID", 400)
|
||||
authorization = _require_string(value["authorization_id"], "AUTHORIZATION_ID_INVALID", KEY_RE)
|
||||
key = _require_string(value["idempotency_key"], "IDEMPOTENCY_KEY_INVALID", KEY_RE)
|
||||
return cls(
|
||||
status,
|
||||
ints["expected_revision"],
|
||||
ints["expected_attempt"],
|
||||
ints["expected_max_attempts"],
|
||||
ints["next_max_attempts"],
|
||||
verification,
|
||||
authorization,
|
||||
key,
|
||||
)
|
||||
|
||||
def public(self) -> dict[str, Any]:
|
||||
return {
|
||||
"expected_status": self.expected_status,
|
||||
"expected_revision": self.expected_revision,
|
||||
"expected_attempt": self.expected_attempt,
|
||||
"expected_max_attempts": self.expected_max_attempts,
|
||||
"next_max_attempts": self.next_max_attempts,
|
||||
"previous_verification_sha256": self.previous_verification_sha256,
|
||||
"authorization_id": self.authorization_id,
|
||||
"idempotency_key": self.idempotency_key,
|
||||
}
|
||||
|
||||
|
||||
class RetryExtensionStore:
|
||||
"""Owns the official mutation. Direct/manual DB edits remain forbidden."""
|
||||
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = Path(db_path)
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.db_path, timeout=30.0)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys=ON")
|
||||
connection.execute("PRAGMA busy_timeout=30000")
|
||||
return connection
|
||||
|
||||
@staticmethod
|
||||
def _columns(connection: sqlite3.Connection, table: str) -> dict[str, sqlite3.Row]:
|
||||
return {str(row["name"]): row for row in connection.execute(f'PRAGMA table_info("{table}")')}
|
||||
|
||||
@staticmethod
|
||||
def _event_table(connection: sqlite3.Connection) -> tuple[str, dict[str, sqlite3.Row]]:
|
||||
tables = {
|
||||
row["name"]
|
||||
for row in connection.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
for name in ("task_events", "events"):
|
||||
if name not in tables:
|
||||
continue
|
||||
columns = RetryExtensionStore._columns(connection, name)
|
||||
if "task_id" in columns and ("event_type" in columns or "type" in columns):
|
||||
return name, columns
|
||||
raise RetryExtensionError("TASK_EVENT_TABLE_UNSUPPORTED", 503)
|
||||
|
||||
@staticmethod
|
||||
def _verification_hash(raw: object) -> str:
|
||||
if raw in (None, ""):
|
||||
value: object = {}
|
||||
elif isinstance(raw, str):
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise RetryExtensionError("TASK_VERIFICATION_INVALID", 409) from exc
|
||||
else:
|
||||
value = raw
|
||||
return _sha(value)
|
||||
|
||||
@staticmethod
|
||||
def _ensure_extension_table(connection: sqlite3.Connection) -> None:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS task_retry_extensions (
|
||||
task_id TEXT NOT NULL,
|
||||
idempotency_key TEXT NOT NULL,
|
||||
request_sha256 TEXT NOT NULL,
|
||||
authorization_id TEXT NOT NULL,
|
||||
previous_verification_sha256 TEXT NOT NULL,
|
||||
previous_max_attempts INTEGER NOT NULL,
|
||||
next_max_attempts INTEGER NOT NULL,
|
||||
expected_revision INTEGER NOT NULL,
|
||||
applied_revision INTEGER NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY(task_id, idempotency_key)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _insert_event(
|
||||
connection: sqlite3.Connection,
|
||||
task_id: str,
|
||||
payload: Mapping[str, Any],
|
||||
) -> None:
|
||||
table, columns = RetryExtensionStore._event_table(connection)
|
||||
values: dict[str, Any] = {"task_id": task_id}
|
||||
if "event_type" in columns:
|
||||
values["event_type"] = "TASK_RETRY_EXTENSION_AUTHORIZED"
|
||||
else:
|
||||
values["type"] = "TASK_RETRY_EXTENSION_AUTHORIZED"
|
||||
payload_text = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
||||
if "payload_json" in columns:
|
||||
values["payload_json"] = payload_text
|
||||
elif "payload" in columns:
|
||||
values["payload"] = payload_text
|
||||
else:
|
||||
raise RetryExtensionError("TASK_EVENT_PAYLOAD_UNSUPPORTED", 503)
|
||||
if "created_at" in columns:
|
||||
values["created_at"] = _now()
|
||||
if "agent_id" in columns:
|
||||
values["agent_id"] = "owner-control"
|
||||
if "actor_id" in columns:
|
||||
values["actor_id"] = "owner-control"
|
||||
names = list(values)
|
||||
sql = f'INSERT INTO "{table}" ({",".join(f"""\"{name}\"""" for name in names)}) VALUES ({",".join("?" for _ in names)})'
|
||||
connection.execute(sql, tuple(values[name] for name in names))
|
||||
|
||||
def extend(self, task_id: str, request: RetryExtensionRequest) -> dict[str, Any]:
|
||||
task_id = _require_string(task_id, "TASK_ID_INVALID", TASK_RE)
|
||||
required_authorization_id = _required_authorization_id()
|
||||
if request.authorization_id != required_authorization_id:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_AUTHORIZATION_ID_MISMATCH", 403)
|
||||
request_payload = request.public()
|
||||
request_sha = _sha(request_payload)
|
||||
with self._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
self._ensure_extension_table(connection)
|
||||
replay = connection.execute(
|
||||
"SELECT * FROM task_retry_extensions WHERE task_id=? AND idempotency_key=?",
|
||||
(task_id, request.idempotency_key),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
if replay["request_sha256"] != request_sha:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_IDEMPOTENCY_CONFLICT", 409)
|
||||
connection.commit()
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "QUEUED",
|
||||
"attempt": request.expected_attempt,
|
||||
"max_attempts": int(replay["next_max_attempts"]),
|
||||
"revision": int(replay["applied_revision"]),
|
||||
"idempotent": True,
|
||||
"authorization_id": replay["authorization_id"],
|
||||
"previous_verification_sha256": replay["previous_verification_sha256"],
|
||||
}
|
||||
|
||||
columns = self._columns(connection, "tasks")
|
||||
required = {
|
||||
"task_id", "status", "attempt", "max_attempts", "revision", "claimed_by",
|
||||
"lease_token_hash", "lease_expires_at", "verification_json",
|
||||
}
|
||||
if not required.issubset(columns):
|
||||
raise RetryExtensionError("TASK_SCHEMA_RETRY_EXTENSION_UNSUPPORTED", 503)
|
||||
task = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
||||
if task is None:
|
||||
raise RetryExtensionError("TASK_NOT_FOUND", 404)
|
||||
actual_verification_sha = self._verification_hash(task["verification_json"])
|
||||
if actual_verification_sha != request.previous_verification_sha256:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_VERIFICATION_DRIFT", 409)
|
||||
if (
|
||||
task["status"] != request.expected_status
|
||||
or int(task["revision"]) != request.expected_revision
|
||||
or int(task["attempt"]) != request.expected_attempt
|
||||
or int(task["max_attempts"]) != request.expected_max_attempts
|
||||
):
|
||||
raise RetryExtensionError("RETRY_EXTENSION_CAS_MISMATCH", 409)
|
||||
for lease_field in ("lease_token_hash", "lease_expires_at"):
|
||||
if lease_field in columns and task[lease_field] not in (None, ""):
|
||||
raise RetryExtensionError("RETRY_EXTENSION_ACTIVE_LEASE_DENY", 409)
|
||||
|
||||
previous_claimed_by = task["claimed_by"]
|
||||
applied_revision = request.expected_revision + 1
|
||||
cursor = connection.execute(
|
||||
"""
|
||||
UPDATE tasks
|
||||
SET status='QUEUED', max_attempts=?, revision=?, claimed_by=NULL,
|
||||
lease_token_hash=NULL, lease_expires_at=NULL
|
||||
WHERE task_id=? AND status=? AND revision=? AND attempt=? AND max_attempts=?
|
||||
AND claimed_by IS ?
|
||||
AND lease_token_hash IS NULL AND lease_expires_at IS NULL
|
||||
""",
|
||||
(
|
||||
request.next_max_attempts,
|
||||
applied_revision,
|
||||
task_id,
|
||||
request.expected_status,
|
||||
request.expected_revision,
|
||||
request.expected_attempt,
|
||||
request.expected_max_attempts,
|
||||
previous_claimed_by,
|
||||
),
|
||||
)
|
||||
if cursor.rowcount != 1:
|
||||
raise RetryExtensionError("RETRY_EXTENSION_CAS_MISMATCH", 409)
|
||||
|
||||
event_payload = {
|
||||
"schema": "aios.task-retry-extension.event.v1",
|
||||
"task_id": task_id,
|
||||
"authorization_id": request.authorization_id,
|
||||
"idempotency_key": request.idempotency_key,
|
||||
"previous_verification_sha256": request.previous_verification_sha256,
|
||||
"previous_max_attempts": request.expected_max_attempts,
|
||||
"next_max_attempts": request.next_max_attempts,
|
||||
"previous_revision": request.expected_revision,
|
||||
"applied_revision": applied_revision,
|
||||
"request_sha256": request_sha,
|
||||
"previous_claimed_by": previous_claimed_by,
|
||||
"stale_terminal_claim_normalized": previous_claimed_by not in (None, ""),
|
||||
}
|
||||
self._insert_event(connection, task_id, event_payload)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO task_retry_extensions(
|
||||
task_id,idempotency_key,request_sha256,authorization_id,previous_verification_sha256,
|
||||
previous_max_attempts,next_max_attempts,expected_revision,applied_revision,created_at
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?)
|
||||
""",
|
||||
(
|
||||
task_id,
|
||||
request.idempotency_key,
|
||||
request_sha,
|
||||
request.authorization_id,
|
||||
request.previous_verification_sha256,
|
||||
request.expected_max_attempts,
|
||||
request.next_max_attempts,
|
||||
request.expected_revision,
|
||||
applied_revision,
|
||||
_now(),
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "QUEUED",
|
||||
"attempt": request.expected_attempt,
|
||||
"max_attempts": request.next_max_attempts,
|
||||
"revision": applied_revision,
|
||||
"idempotent": False,
|
||||
"authorization_id": request.authorization_id,
|
||||
"previous_verification_sha256": request.previous_verification_sha256,
|
||||
}
|
||||
|
||||
|
||||
class RetryExtensionHandler(BaseHTTPRequestHandler):
|
||||
store: RetryExtensionStore | None = None
|
||||
control_token: str = ""
|
||||
|
||||
def log_message(self, *_args: Any) -> None:
|
||||
return
|
||||
|
||||
def _json(self, value: Mapping[str, Any], status: int) -> None:
|
||||
data = json.dumps(dict(value), sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def _authorized(self) -> bool:
|
||||
header = self.headers.get("Authorization", "")
|
||||
return (
|
||||
header.startswith("Bearer ")
|
||||
and bool(self.control_token)
|
||||
and hmac.compare_digest(header[7:], self.control_token)
|
||||
)
|
||||
|
||||
def do_GET(self) -> None:
|
||||
if self.path == "/health":
|
||||
self._json({"status": "ok", "service": "aios-agent-task-retry-extension"}, 200)
|
||||
else:
|
||||
self._json({"error": "NOT_FOUND"}, 404)
|
||||
|
||||
def do_POST(self) -> None:
|
||||
match = re.fullmatch(r"/v1/tasks/([^/]+)/retry-extension", self.path)
|
||||
if match is None:
|
||||
self._json({"error": "NOT_FOUND"}, 404)
|
||||
return
|
||||
if not self._authorized():
|
||||
self._json({"error": "UNAUTHORIZED"}, 401)
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
except ValueError:
|
||||
self._json({"error": "CONTENT_LENGTH_INVALID"}, 400)
|
||||
return
|
||||
if length < 1 or length > MAX_BODY:
|
||||
self._json({"error": "REQUEST_SIZE_INVALID"}, 413)
|
||||
return
|
||||
try:
|
||||
value = json.loads(self.rfile.read(length).decode("utf-8"))
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError
|
||||
request = RetryExtensionRequest.from_mapping(value)
|
||||
if self.store is None:
|
||||
raise RetryExtensionError("RUNTIME_UNAVAILABLE", 503)
|
||||
result = self.store.extend(match.group(1), request)
|
||||
self._json(result, 200 if result["idempotent"] else 201)
|
||||
except RetryExtensionError as exc:
|
||||
self._json({"error": exc.code}, exc.status)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
|
||||
self._json({"error": "REQUEST_INVALID"}, 400)
|
||||
|
||||
|
||||
def load_control_token(path: Path) -> str:
|
||||
if path.is_symlink() or not path.is_file():
|
||||
raise RuntimeError("TASK_RETRY_CONTROL_CREDENTIAL_UNAVAILABLE")
|
||||
value = path.read_text(encoding="utf-8").strip()
|
||||
if len(value) < 32:
|
||||
raise RuntimeError("TASK_RETRY_CONTROL_CREDENTIAL_INVALID")
|
||||
return value
|
||||
|
||||
|
||||
def serve() -> None:
|
||||
db = Path(os.environ.get("AIOS_AGENT_TASK_DB", "/var/lib/aios-agent-task/state.sqlite3"))
|
||||
_required_authorization_id()
|
||||
token_path = Path(
|
||||
os.environ.get(
|
||||
"AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
|
||||
str(Path(os.environ.get("CREDENTIALS_DIRECTORY", "/run/credentials")) / "agent_task_control_token"),
|
||||
)
|
||||
)
|
||||
host = os.environ.get("AIOS_TASK_RETRY_HOST", "127.0.0.1")
|
||||
port = int(os.environ.get("AIOS_TASK_RETRY_PORT", "18941"))
|
||||
RetryExtensionHandler.store = RetryExtensionStore(db)
|
||||
RetryExtensionHandler.control_token = load_control_token(token_path)
|
||||
ThreadingHTTPServer((host, port), RetryExtensionHandler).serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
serve()
|
||||
@@ -12,8 +12,6 @@ BRIDGE_OPERATIONS = (
|
||||
"getRuntimeStatus",
|
||||
"getGovernedWriteCapabilities",
|
||||
"createGovernedTask",
|
||||
"projectOnlineAgentExecution",
|
||||
"registerOnlineAgentExecutionPlan",
|
||||
"acquireWriterLease",
|
||||
"bootstrapWriterLease",
|
||||
"renewWriterLease",
|
||||
@@ -34,11 +32,6 @@ BRIDGE_OPERATIONS = (
|
||||
|
||||
OPERATION_REQUIRED_FIELDS = {
|
||||
"createGovernedTask": ("task_id", "repository", "base_oid", "target_branch", "risk", "description", "idempotency_key"),
|
||||
"projectOnlineAgentExecution": ("task_id",),
|
||||
"registerOnlineAgentExecutionPlan": (
|
||||
"schema", "task_id", "job_id", "source_head_oid", "generation",
|
||||
"idempotency_key", "execution",
|
||||
),
|
||||
"acquireWriterLease": ("schema", "task_id", "repository", "branch_mode", "base_oid", "expected_branch_oid", "target_branch", "generation", "lease_seconds"),
|
||||
"bootstrapWriterLease": ("schema", "task_id", "repository", "branch_mode", "base_oid", "expected_branch_oid", "target_branch", "generation", "lease_seconds"),
|
||||
"renewWriterLease": ("schema", "task_id", "repository", "branch_mode", "base_oid", "expected_branch_oid", "target_branch", "generation", "lease_seconds"),
|
||||
@@ -58,25 +51,6 @@ OPERATION_REQUIRED_FIELDS = {
|
||||
}
|
||||
|
||||
BRIDGE_OPERATION_CONTRACTS = {
|
||||
"projectOnlineAgentExecution": {
|
||||
"method": "POST",
|
||||
"path": "/v2/execution-envelope",
|
||||
"required_fields": OPERATION_REQUIRED_FIELDS["projectOnlineAgentExecution"],
|
||||
"response_fields": (
|
||||
"schema", "task_id", "job_id", "repository", "base_oid",
|
||||
"expected_branch_oid", "target_branch", "branch_mode", "generation",
|
||||
"idempotency_key", "action", "execution_digest",
|
||||
),
|
||||
},
|
||||
"registerOnlineAgentExecutionPlan": {
|
||||
"method": "POST",
|
||||
"path": "/v2/execution-plan",
|
||||
"required_fields": OPERATION_REQUIRED_FIELDS["registerOnlineAgentExecutionPlan"],
|
||||
"response_fields": (
|
||||
"schema", "task_id", "job_id", "source_head_oid", "generation",
|
||||
"idempotency_key", "plan_digest", "execution", "execution_digest",
|
||||
),
|
||||
},
|
||||
"acquireWriterLease": {
|
||||
"method": "POST",
|
||||
"path": "/v2/writer-lease",
|
||||
@@ -231,8 +205,6 @@ class BridgeV2Adapter:
|
||||
return handler(dict(request))
|
||||
|
||||
def createGovernedTask(self, request: Mapping[str, Any]) -> Any: return self._call("createGovernedTask", request)
|
||||
def projectOnlineAgentExecution(self, request: Mapping[str, Any]) -> Any: return self._call("projectOnlineAgentExecution", request)
|
||||
def registerOnlineAgentExecutionPlan(self, request: Mapping[str, Any]) -> Any: return self._call("registerOnlineAgentExecutionPlan", request)
|
||||
def acquireWriterLease(self, request: Mapping[str, Any]) -> Any: return self._call("acquireWriterLease", request)
|
||||
def bootstrapWriterLease(self, request: Mapping[str, Any]) -> Any: return self._call("bootstrapWriterLease", request)
|
||||
def renewWriterLease(self, request: Mapping[str, Any]) -> Any: return self._call("renewWriterLease", request)
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
"""Loopback-only governed Bridge v2 HTTP gateway with durable handlers."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass, replace
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
import base64
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
@@ -13,8 +12,6 @@ import re
|
||||
import sqlite3
|
||||
import threading
|
||||
from typing import Any, Mapping
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from aios_authority_git_executor import AuthorityGitExecutor
|
||||
from aios_chatgpt_bridge_v2 import BridgeV2Adapter, BridgeV2Error, OPERATION_REQUIRED_FIELDS
|
||||
@@ -22,9 +19,7 @@ from aios_continuation_executor import AuthorityGitAdapter, ContinuationCommit,
|
||||
from aios_continuation_lease import WriterLeaseRequestV2
|
||||
from aios_control_plane_snapshot import build_control_plane_snapshot
|
||||
from aios_gitops_job_controller import GitOpsJobController, JobError, JobState, JobStore
|
||||
from aios_governed_task_bootstrap import BootstrapError, GovernedTaskBinding, GovernedTaskBootstrap, _legacy_safe, _task_key
|
||||
from aios_online_agent_execution import ExecutionProjectionError, OnlineTaskProjectionClient
|
||||
from aios_online_agent_plan_store import ExecutionPlanError, ExecutionPlanStore
|
||||
from aios_governed_task_bootstrap import BootstrapError, GovernedTaskBootstrap
|
||||
from governed_production_release import build_executor_config
|
||||
|
||||
MAX_BODY = 2 * 1024 * 1024
|
||||
@@ -32,8 +27,6 @@ ROUTES = {
|
||||
"/v2/runtime": "getRuntimeStatus",
|
||||
"/v2/capabilities": "getGovernedWriteCapabilities",
|
||||
"/v2/governed-task": "createGovernedTask",
|
||||
"/v2/execution-envelope": "projectOnlineAgentExecution",
|
||||
"/v2/execution-plan": "registerOnlineAgentExecutionPlan",
|
||||
"/v2/writer-lease": "acquireWriterLease",
|
||||
"/v2/writer-lease/bootstrap": "bootstrapWriterLease",
|
||||
"/v2/writer-lease/renew": "renewWriterLease",
|
||||
@@ -69,9 +62,6 @@ class RuntimeConfig:
|
||||
registry_root: Path = Path("/var/lib/aios-gitops/runtime/registry-projection")
|
||||
baseline_path: Path = Path("/var/lib/aios-gitops-runtime-anchor/baseline/current.json")
|
||||
health_url: str = "http://127.0.0.1:18927/health"
|
||||
agent_task_api_url: str = "http://127.0.0.1:18940"
|
||||
agent_task_control_token_path: Path = Path("/run/credentials/aios-chatgpt-bridge-v2-staging.service/agent_task_control_token")
|
||||
online_task_projection_enabled: bool = False
|
||||
|
||||
|
||||
def _safe(value: str) -> str:
|
||||
@@ -84,16 +74,6 @@ class GovernedRuntime:
|
||||
config.state_root.mkdir(parents=True, exist_ok=True)
|
||||
self.authority = AuthorityGitExecutor(config.authoritative_git_dir, config.state_root / "worktrees")
|
||||
self.bootstrap = GovernedTaskBootstrap(self.authority, config.state_root, remote_name=config.remote_name, token_path=config.token_path, ssh_key_path=config.ssh_key_path, api_base=config.api_base)
|
||||
self.online_task_projector = (
|
||||
OnlineTaskProjectionClient(config.agent_task_api_url, config.agent_task_control_token_path)
|
||||
if config.online_task_projection_enabled
|
||||
else None
|
||||
)
|
||||
self.projection_outbox_root = config.state_root / "online-task-projection-outbox"
|
||||
self.projection_outbox_root.mkdir(parents=True, exist_ok=True)
|
||||
self.execution_plan_store = ExecutionPlanStore(config.state_root / "online-agent-execution-plans")
|
||||
self._projection_locks_guard = threading.Lock()
|
||||
self._projection_locks: dict[str, threading.RLock] = {}
|
||||
self.runtimes: dict[tuple[str, int], tuple[Any, GitOpsJobController, ContinuationExecutor]] = {}
|
||||
self.controllers: dict[str, GitOpsJobController] = {}
|
||||
self._recovery_lock = threading.RLock()
|
||||
@@ -113,54 +93,6 @@ class GovernedRuntime:
|
||||
askpass.write_text('#!/bin/sh\ncase "$1" in *Username*) printf "%s\\n" "$GITEA_MACHINE_ACTOR" ;; *) cat "$CREDENTIAL_FILE" ;; esac\n', encoding="utf-8")
|
||||
askpass.chmod(0o700)
|
||||
os.environ.setdefault("GIT_ASKPASS", str(askpass)); os.environ.setdefault("GIT_ASKPASS_REQUIRE", "force"); os.environ.setdefault("GIT_TERMINAL_PROMPT", "0"); os.environ.setdefault("GITEA_MACHINE_ACTOR", "aios-autonomy-agent")
|
||||
if self.online_task_projector is not None:
|
||||
threading.Thread(target=self._projection_reconciler, name="online-task-projection-reconciler", daemon=True).start()
|
||||
|
||||
def _projection_lock(self, task_id: str) -> threading.RLock:
|
||||
guard = getattr(self, "_projection_locks_guard", None)
|
||||
if guard is None:
|
||||
guard = threading.Lock()
|
||||
self._projection_locks_guard = guard
|
||||
self._projection_locks = {}
|
||||
with guard:
|
||||
return self._projection_locks.setdefault(task_id, threading.RLock())
|
||||
|
||||
def _projection_outbox_path(self, task_id: str) -> Path:
|
||||
return self.projection_outbox_root / (_task_key(task_id) + ".json")
|
||||
|
||||
def _write_projection_outbox(self, task_id: str, projection: Mapping[str, Any], binding: Any) -> Path:
|
||||
path = self._projection_outbox_path(task_id)
|
||||
temporary = path.with_name(path.name + f".{os.getpid()}.{threading.get_ident()}.tmp")
|
||||
temporary.write_text(json.dumps({"task": dict(projection), "binding": asdict(binding)}, sort_keys=True, separators=(",", ":")), encoding="utf-8")
|
||||
temporary.replace(path)
|
||||
return path
|
||||
|
||||
def _deliver_projection_outbox(self, path: Path) -> dict[str, Any]:
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, Mapping) or not isinstance(value.get("task"), Mapping) or not isinstance(value.get("binding"), Mapping):
|
||||
raise RuntimeError("ONLINE_TASK_PROJECTION_OUTBOX_INVALID")
|
||||
binding = self.bootstrap.execution_binding(GovernedTaskBinding(**dict(value["binding"])))
|
||||
if self.online_task_projector is None:
|
||||
raise RuntimeError("ONLINE_TASK_PROJECTION_CONTROL_UNAVAILABLE")
|
||||
created = self.online_task_projector.create(dict(value["task"]), binding=binding)
|
||||
path.unlink(missing_ok=True)
|
||||
return created
|
||||
|
||||
def _projection_reconciler(self) -> None:
|
||||
while True:
|
||||
for path in sorted(self.projection_outbox_root.glob("*.json")):
|
||||
try:
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
binding = value.get("binding") if isinstance(value, Mapping) else None
|
||||
task_id = str(binding.get("task_id") or "") if isinstance(binding, Mapping) else ""
|
||||
if not task_id:
|
||||
continue
|
||||
with self._projection_lock(task_id):
|
||||
if path.exists():
|
||||
self._deliver_projection_outbox(path)
|
||||
except Exception:
|
||||
continue
|
||||
threading.Event().wait(5)
|
||||
|
||||
def runtime_status(self) -> dict[str, Any]:
|
||||
return {
|
||||
@@ -181,10 +113,7 @@ class GovernedRuntime:
|
||||
}
|
||||
|
||||
def _job_root(self, task_id: str) -> Path:
|
||||
parent = self.config.state_root / "jobs"
|
||||
keyed = parent / _task_key(task_id)
|
||||
legacy = parent / _legacy_safe(task_id)
|
||||
return keyed if keyed.exists() or not legacy.exists() else legacy
|
||||
return self.config.state_root / "jobs" / _safe(task_id)
|
||||
|
||||
def _build_runtime(
|
||||
self,
|
||||
@@ -514,77 +443,7 @@ class GovernedRuntime:
|
||||
return result
|
||||
|
||||
def create_task(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
online_request = request.get("online_task")
|
||||
if online_request is not None and not isinstance(online_request, Mapping):
|
||||
raise ValueError("ONLINE_TASK_REQUEST_INVALID")
|
||||
execution_request = request.get("execution")
|
||||
if execution_request is not None and not isinstance(execution_request, Mapping):
|
||||
raise ValueError("EXECUTION_REQUEST_INVALID")
|
||||
online_task_projector = getattr(self, "online_task_projector", None)
|
||||
projection_required = online_request is not None or execution_request is not None or request.get("require_online_task_projection") is True
|
||||
if online_task_projector is None and projection_required:
|
||||
raise RuntimeError("ONLINE_TASK_PROJECTION_CONTROL_UNAVAILABLE")
|
||||
result = self.bootstrap.create(request)
|
||||
self.bindings[result["task_id"]] = result
|
||||
if online_task_projector is None:
|
||||
return result
|
||||
online = dict(online_request or {})
|
||||
online.setdefault("goal", str(request.get("description") or "Governed engineering task"))
|
||||
online.setdefault("instructions", "Execute only the exact governed binding and return server-verifiable evidence.")
|
||||
online.setdefault("success_criteria", [{"id": "governed", "description": "Complete the exact governed task lifecycle."}])
|
||||
online.setdefault("required_capabilities", ["git.read", "git.test", "governed.write.branch", "governed.write.file"])
|
||||
online.setdefault("priority", int(request.get("priority") or 0))
|
||||
online.setdefault("max_attempts", int(request.get("max_attempts") or 10))
|
||||
projection = self.bootstrap.online_task_projection(
|
||||
result["task_id"],
|
||||
online,
|
||||
execution_request,
|
||||
expected_job_id=str(result.get("job_id") or "") or None,
|
||||
)
|
||||
durable_binding = self.bootstrap.get_binding(result["task_id"])
|
||||
with self._projection_lock(result["task_id"]):
|
||||
path = self._write_projection_outbox(result["task_id"], projection, durable_binding)
|
||||
created = self._deliver_projection_outbox(path)
|
||||
return {
|
||||
**result,
|
||||
"online_task_projection": {
|
||||
"status": "PROJECTED",
|
||||
"task_id": created["task_id"],
|
||||
"job_id": result["job_id"],
|
||||
"task_status": created["status"],
|
||||
"execution_digest": projection["metadata"]["execution"]["execution_digest"],
|
||||
},
|
||||
}
|
||||
|
||||
def project_online_agent_execution(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
capabilities = request.get("required_capabilities") or []
|
||||
if not isinstance(capabilities, list) or not all(isinstance(item, str) for item in capabilities):
|
||||
raise ValueError("CAPABILITY_SCOPE_INVALID")
|
||||
execution = request.get("execution")
|
||||
if execution is not None and not isinstance(execution, Mapping):
|
||||
raise ValueError("EXECUTION_REQUEST_INVALID")
|
||||
task_id = str(request["task_id"])
|
||||
durable = self.bootstrap.get_binding(task_id)
|
||||
binding = self.bootstrap.execution_binding(durable)
|
||||
expected_job_id = str(request.get("job_id") or "") or None
|
||||
if expected_job_id and expected_job_id != binding.job_id:
|
||||
raise BootstrapError("JOB_BINDING_CONFLICT")
|
||||
if execution is None:
|
||||
planned = self.execution_plan_store.resolve(binding)
|
||||
if planned is not None:
|
||||
execution = planned["execution"]
|
||||
return self.bootstrap.execution_envelope(
|
||||
task_id,
|
||||
execution,
|
||||
required_capabilities=tuple(capabilities),
|
||||
expected_job_id=expected_job_id,
|
||||
)
|
||||
|
||||
def register_online_agent_execution_plan(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
task_id = str(request["task_id"])
|
||||
durable = self.bootstrap.get_binding(task_id)
|
||||
binding = self.bootstrap.execution_binding(durable)
|
||||
return self.execution_plan_store.register(binding, request)
|
||||
result = self.bootstrap.create(request); self.bindings[result["task_id"]] = result; return result
|
||||
|
||||
def acquire_lease(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
binding, controller, _ = self._runtime(request)
|
||||
@@ -1481,19 +1340,7 @@ class GovernedRuntime:
|
||||
controller = self._controller_for_job(job_id)
|
||||
row = controller.store.read(job_id)
|
||||
if action == "get":
|
||||
response = controller.store._receipt(row, duplicate=True).as_dict()
|
||||
try:
|
||||
with controller.store._connect() as connection:
|
||||
connection.row_factory = sqlite3.Row
|
||||
leases = connection.execute(
|
||||
"SELECT lease_id,job_id,branch,lineage_anchor_oid,current_head_oid,generation,"
|
||||
"lease_expires_at,state,revision FROM branch_leases WHERE job_id=? ORDER BY revision DESC",
|
||||
(job_id,),
|
||||
).fetchall()
|
||||
except sqlite3.Error as exc:
|
||||
raise RuntimeError("JOB_DB_UNAVAILABLE") from exc
|
||||
response["branch_leases"] = [dict(item) for item in leases]
|
||||
return response
|
||||
return controller.store._receipt(row, duplicate=True).as_dict()
|
||||
if action == "resume":
|
||||
return controller.resume(job_id).as_dict()
|
||||
return self._executor_from_row(row, controller).reconcile(job_id)
|
||||
@@ -1558,8 +1405,6 @@ class GovernedRuntime:
|
||||
def handlers(self) -> dict[str, Any]:
|
||||
return {
|
||||
"createGovernedTask": self.create_task,
|
||||
"projectOnlineAgentExecution": self.project_online_agent_execution,
|
||||
"registerOnlineAgentExecutionPlan": self.register_online_agent_execution_plan,
|
||||
"acquireWriterLease": self.acquire_lease,
|
||||
"bootstrapWriterLease": self.bootstrap_writer_lease,
|
||||
"renewWriterLease": self.renew_writer_lease,
|
||||
@@ -1589,8 +1434,6 @@ class _HTTPError(RuntimeError): pass
|
||||
|
||||
class BridgeHandler(BaseHTTPRequestHandler):
|
||||
adapter: BridgeV2Adapter | None = None
|
||||
control_token: str | None = None
|
||||
agent_task_api_url: str = "http://127.0.0.1:18940"
|
||||
def log_message(self, *_args: Any) -> None: return
|
||||
def _json(self, value: Any, status: int = 200) -> None:
|
||||
raw = json.dumps(value, sort_keys=True, separators=(",", ":")).encode(); self.send_response(status); self.send_header("Content-Type", "application/json"); self.send_header("Content-Length", str(len(raw))); self.end_headers(); self.wfile.write(raw)
|
||||
@@ -1608,50 +1451,7 @@ class BridgeHandler(BaseHTTPRequestHandler):
|
||||
allowed = set(OPERATION_REQUIRED_FIELDS[operation]) | {"schema", "job_id", "lease_seconds", "create_target_branch", "worktree_mode", "cas_ref_update_required", "expected_branch_oid", "generation", "branch_mode"}
|
||||
if operation == "commitFilesGoverned": allowed.add("files")
|
||||
allowed.update({"owner_approval", "execute"})
|
||||
if operation == "projectOnlineAgentExecution":
|
||||
allowed.update({"execution", "required_capabilities"})
|
||||
if operation == "registerOnlineAgentExecutionPlan":
|
||||
allowed.update({"execution", "source_head_oid", "expected_previous_plan_digest"})
|
||||
if operation == "createGovernedTask":
|
||||
allowed.update({"online_task", "execution", "require_online_task_projection", "priority", "max_attempts", "parent_task_id", "parent_job_id"})
|
||||
if set(body) - allowed: raise _HTTPError("UNKNOWN_FIELDS")
|
||||
def _authorized(self, operation: str, body: Mapping[str, Any]) -> bool:
|
||||
header = self.headers.get("Authorization", "")
|
||||
if not header.startswith("Bearer "):
|
||||
return False
|
||||
supplied = header[7:]
|
||||
if self.control_token and hmac.compare_digest(supplied, self.control_token):
|
||||
return True
|
||||
poll = urllib.request.Request(
|
||||
self.agent_task_api_url.rstrip("/") + "/v1/tasks/poll",
|
||||
method="POST",
|
||||
data=b'{"limit":1}',
|
||||
headers={"Authorization": header, "Content-Type": "application/json", "Accept": "application/json"},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(poll, timeout=5) as response:
|
||||
principal = json.loads(response.read().decode("utf-8"))
|
||||
except (OSError, urllib.error.URLError, urllib.error.HTTPError, ValueError):
|
||||
return False
|
||||
agent_id = str(principal.get("agent_id") or "") if isinstance(principal, Mapping) else ""
|
||||
if not agent_id:
|
||||
return False
|
||||
claim_bound = {"acquireWriterLease", "bootstrapWriterLease", "renewWriterLease", "commitFileGoverned", "commitFilesGoverned"}
|
||||
if operation not in claim_bound:
|
||||
return operation in {"projectOnlineAgentExecution", "getGitOpsJob"}
|
||||
task_id = str(body.get("task_id") or "")
|
||||
if not task_id:
|
||||
return False
|
||||
task_request = urllib.request.Request(
|
||||
self.agent_task_api_url.rstrip("/") + "/v1/tasks/" + task_id,
|
||||
headers={"Authorization": header, "Accept": "application/json"},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(task_request, timeout=5) as response:
|
||||
task = json.loads(response.read().decode("utf-8"))
|
||||
except (OSError, urllib.error.URLError, urllib.error.HTTPError, ValueError):
|
||||
return False
|
||||
return isinstance(task, Mapping) and task.get("task_id") == task_id and task.get("claimed_by") == agent_id and task.get("status") in {"CLAIMED", "RUNNING"}
|
||||
def do_GET(self) -> None:
|
||||
if self.adapter is None: self._json({"error": "RUNTIME_UNAVAILABLE"}, 503); return
|
||||
if self.path == "/v2/runtime": self._json(self.adapter.getRuntimeStatus()); return
|
||||
@@ -1662,23 +1462,14 @@ class BridgeHandler(BaseHTTPRequestHandler):
|
||||
operation = ROUTES.get(self.path)
|
||||
if operation is None or operation in ("getRuntimeStatus", "getGovernedWriteCapabilities") or self.adapter is None: self._json({"error": "NOT_FOUND"}, 404); return
|
||||
try:
|
||||
body = self._body()
|
||||
except _HTTPError as exc: self._json({"error": str(exc)}, 413 if str(exc) == "REQUEST_TOO_LARGE" else 400); return
|
||||
if not self._authorized(operation, body): self._json({"error": "BRIDGE_CALLER_UNAUTHORIZED"}, 401); return
|
||||
try:
|
||||
self._allowed(operation, body); self._json(getattr(self.adapter, operation)(body))
|
||||
body = self._body(); self._allowed(operation, body); self._json(getattr(self.adapter, operation)(body))
|
||||
except _HTTPError as exc: self._json({"error": str(exc)}, 413 if str(exc) == "REQUEST_TOO_LARGE" else 400)
|
||||
except (BootstrapError, ExecutionProjectionError, ExecutionPlanError, ValueError, BridgeV2Error) as exc: self._json({"error": str(exc)}, 409 if any(token in str(exc) for token in ("CONFLICT", "DRIFT", "CAS")) else 400)
|
||||
except (BootstrapError, ValueError, BridgeV2Error) as exc: self._json({"error": str(exc)}, 409 if "CONFLICT" in str(exc) else 400)
|
||||
except Exception as exc: self._json({"error": str(exc) or type(exc).__name__}, 409 if any(token in str(exc) for token in ("CONFLICT", "EXPECTATION", "MONOTONIC")) else 503)
|
||||
|
||||
|
||||
def serve(adapter: BridgeV2Adapter, host: str = "127.0.0.1", port: int = 18892, *, control_token_path: Path, agent_task_api_url: str) -> None:
|
||||
if not control_token_path.is_file() or control_token_path.is_symlink():
|
||||
raise RuntimeError("BRIDGE_CONTROL_CREDENTIAL_UNAVAILABLE")
|
||||
control_token = control_token_path.read_text(encoding="utf-8").strip()
|
||||
if len(control_token) < 32:
|
||||
raise RuntimeError("BRIDGE_CONTROL_CREDENTIAL_INVALID")
|
||||
BridgeHandler.adapter = adapter; BridgeHandler.control_token = control_token; BridgeHandler.agent_task_api_url = agent_task_api_url; ThreadingHTTPServer((host, port), BridgeHandler).serve_forever()
|
||||
def serve(adapter: BridgeV2Adapter, host: str = "127.0.0.1", port: int = 18892) -> None:
|
||||
BridgeHandler.adapter = adapter; ThreadingHTTPServer((host, port), BridgeHandler).serve_forever()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@@ -1697,21 +1488,11 @@ def main() -> None:
|
||||
remote_name=os.environ.get("AIOS_GITEA_REMOTE", "gitea-http"),
|
||||
api_base=os.environ.get("AIOS_GITEA_API_BASE", "https://g.cnnetmkt.top/git/api/v1"),
|
||||
target_version=os.environ.get("AIOS_CONTROL_PLANE_TARGET_VERSION", "V2R2C1"),
|
||||
agent_task_api_url=os.environ.get("AIOS_AGENT_TASK_API_URL", "http://127.0.0.1:18940"),
|
||||
agent_task_control_token_path=Path(
|
||||
os.environ.get(
|
||||
"AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
|
||||
str(Path(os.environ.get("CREDENTIALS_DIRECTORY", "/run/credentials/aios-chatgpt-bridge-v2-staging.service")) / "agent_task_control_token"),
|
||||
)
|
||||
),
|
||||
online_task_projection_enabled=os.environ.get("AIOS_ONLINE_TASK_PROJECTION_ENABLED", "1") == "1",
|
||||
)
|
||||
runtime = GovernedRuntime(config)
|
||||
serve(
|
||||
BridgeV2Adapter(runtime_status=runtime.runtime_status(), handlers=runtime.handlers()),
|
||||
port=int(os.environ.get("AIOS_BRIDGE_V2_PORT", "18892")),
|
||||
control_token_path=config.agent_task_control_token_path,
|
||||
agent_task_api_url=config.agent_task_api_url,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -157,7 +157,7 @@ class DetachedGitWorkspace:
|
||||
def materialize_files(self, files: Mapping[str, bytes]) -> None:
|
||||
dirty = self._dirty_paths()
|
||||
expected_paths = set(files)
|
||||
if dirty and not dirty.issubset(expected_paths):
|
||||
if dirty and dirty != expected_paths:
|
||||
raise ContinuationExecutionError("WORKTREE_DIRTY_SCOPE_MISMATCH")
|
||||
for relative, content in files.items():
|
||||
target = (self.path / relative).resolve()
|
||||
@@ -499,19 +499,7 @@ class ContinuationExecutor:
|
||||
current = self.controller.store.read(binding.job_id)
|
||||
current_state = JobState(current["state"])
|
||||
if current_state == JobState.COMPLETED:
|
||||
cycle = self.controller.begin_continuation_cycle(
|
||||
binding.job_id,
|
||||
idempotency_key=commit.idempotency_key,
|
||||
expected_head_oid=binding.request.expected_branch_oid,
|
||||
generation=binding.request.generation,
|
||||
expected_revision=int(current["revision"]),
|
||||
)
|
||||
if cycle.get("completed"):
|
||||
return self.controller.store._receipt(
|
||||
self.controller.store.read(binding.job_id), duplicate=True
|
||||
).as_dict()
|
||||
current = self.controller.store.read(binding.job_id)
|
||||
current_state = JobState(current["state"])
|
||||
return self.controller.store._receipt(current, duplicate=True).as_dict()
|
||||
if current_state == JobState.REMOTE_VERIFIED:
|
||||
receipt = json.loads(str(current["receipt_json"] or "{}"))
|
||||
commit_oid = str(receipt.get("authority_oid") or "")
|
||||
@@ -519,18 +507,12 @@ class ContinuationExecutor:
|
||||
raise ContinuationExecutionError("GOVERNED_COMMIT_OID_INVALID")
|
||||
if before_complete is not None:
|
||||
before_complete(commit_oid)
|
||||
result = self.controller.advance(
|
||||
return self.controller.advance(
|
||||
binding.job_id,
|
||||
JobState.COMPLETED,
|
||||
side_effect_state="REMOTE_VERIFIED",
|
||||
receipt=receipt,
|
||||
).as_dict()
|
||||
self.controller.complete_continuation_cycle(
|
||||
binding.job_id,
|
||||
idempotency_key=commit.idempotency_key,
|
||||
commit_oid=commit_oid,
|
||||
)
|
||||
return result
|
||||
self._validate_preflight(binding, job_id=binding.job_id)
|
||||
request = binding.request
|
||||
projected = False
|
||||
@@ -594,19 +576,12 @@ class ContinuationExecutor:
|
||||
self.remote.push_exact(request.target_branch, request.expected_branch_oid, new_oid)
|
||||
except Exception:
|
||||
return self._remote_pending(binding.job_id, authority_oid, new_oid)
|
||||
result = self._finish_remote_verification(
|
||||
return self._finish_remote_verification(
|
||||
binding.job_id,
|
||||
authority_oid,
|
||||
new_oid,
|
||||
before_complete=before_complete,
|
||||
)
|
||||
if str(result.get("state") or "") == JobState.COMPLETED.value:
|
||||
self.controller.complete_continuation_cycle(
|
||||
binding.job_id,
|
||||
idempotency_key=commit.idempotency_key,
|
||||
commit_oid=new_oid,
|
||||
)
|
||||
return result
|
||||
finally:
|
||||
if projected:
|
||||
self.credential.revoke()
|
||||
|
||||
@@ -6,6 +6,7 @@ from datetime import datetime, timedelta, timezone
|
||||
import os
|
||||
from pathlib import Path
|
||||
import secrets
|
||||
import re
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
|
||||
@@ -21,9 +22,14 @@ class CredentialLease:
|
||||
scopes: tuple[str, ...]
|
||||
provider: str
|
||||
state: str = "ACTIVE"
|
||||
scope_type: str = "system"
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
run_id: str | None = None
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return {"lease_id": self.lease_id, "task_id": self.task_id, "repository": self.repository, "target_branch": self.target_branch, "generation": self.generation, "issued_at": self.issued_at, "expires_at": self.expires_at, "scopes": list(self.scopes), "provider": self.provider, "state": self.state}
|
||||
return {"lease_id": self.lease_id, "task_id": self.task_id, "repository": self.repository, "target_branch": self.target_branch, "generation": self.generation, "issued_at": self.issued_at, "expires_at": self.expires_at, "scopes": list(self.scopes), "provider": self.provider, "state": self.state, "scope_type": self.scope_type, "tenant_id": self.tenant_id, "user_id": self.user_id, "project_id": self.project_id, "run_id": self.run_id}
|
||||
|
||||
|
||||
class FileCredentialProvider:
|
||||
@@ -43,11 +49,19 @@ class CredentialBroker:
|
||||
self.now = now or (lambda: datetime.now(timezone.utc))
|
||||
self._leases: dict[str, CredentialLease] = {}
|
||||
|
||||
def issue(self, *, task_id: str, repository: str, target_branch: str, generation: int, scopes: tuple[str, ...] = ("git:push", "gitea:read"), lease_seconds: int = 900) -> dict[str, Any]:
|
||||
def issue(self, *, task_id: str, repository: str, target_branch: str, generation: int, scopes: tuple[str, ...] = ("git:push", "gitea:read"), lease_seconds: int = 900, scope_type: str = "system", tenant_id: str | None = None, user_id: str | None = None, project_id: str | None = None, run_id: str | None = None) -> dict[str, Any]:
|
||||
if repository != "aios/aios" or not target_branch.startswith("agent/") or generation < 1 or not 1 <= lease_seconds <= 86400:
|
||||
raise ValueError("CREDENTIAL_LEASE_BINDING_INVALID")
|
||||
if scope_type == "tenant":
|
||||
if not tenant_id:
|
||||
raise ValueError("CREDENTIAL_SCOPE_DENY")
|
||||
elif scope_type in {"system", "global"}:
|
||||
if any((tenant_id, user_id, project_id, run_id)):
|
||||
raise ValueError("CREDENTIAL_SCOPE_DENY")
|
||||
else:
|
||||
raise ValueError("CREDENTIAL_SCOPE_DENY")
|
||||
now = self.now().astimezone(timezone.utc)
|
||||
lease = CredentialLease("lease-" + secrets.token_hex(16), task_id, repository, target_branch, generation, now.isoformat(), (now + timedelta(seconds=lease_seconds)).isoformat(), tuple(scopes), type(self.provider).__name__)
|
||||
lease = CredentialLease("lease-" + secrets.token_hex(16), task_id, repository, target_branch, generation, now.isoformat(), (now + timedelta(seconds=lease_seconds)).isoformat(), tuple(scopes), type(self.provider).__name__, "ACTIVE", scope_type, tenant_id, user_id, project_id, run_id)
|
||||
self._leases[lease.lease_id] = lease
|
||||
return lease.as_dict()
|
||||
|
||||
@@ -64,7 +78,39 @@ class CredentialBroker:
|
||||
lease = self._get(lease_id)
|
||||
if lease.state != "ACTIVE":
|
||||
raise RuntimeError("CREDENTIAL_LEASE_NOT_ACTIVE")
|
||||
return self.provider.environment(lease)
|
||||
environment = self.provider.environment(lease)
|
||||
for key, value in environment.items():
|
||||
upper = str(key).upper()
|
||||
if not isinstance(value, str):
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
sensitive = any(marker in upper for marker in ("CREDENTIAL", "TOKEN", "SECRET", "PASSWORD", "PRIVATE_KEY", "API_KEY"))
|
||||
if re.search(r"(?i)bearer\s+\S+|(?:token|secret|password|api[_-]?key)\s*[:=]", value):
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
if sensitive and upper.endswith("_FILE"):
|
||||
candidate = Path(value)
|
||||
credential_root = Path("/run/credentials")
|
||||
try:
|
||||
resolved = candidate.resolve(strict=False)
|
||||
except OSError as exc:
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY") from exc
|
||||
if not candidate.is_absolute() or not (resolved == credential_root or credential_root in resolved.parents):
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
elif sensitive and upper.endswith("_REFERENCE"):
|
||||
if re.fullmatch(r"(?:systemd-credential|vault|secret-manager):[A-Za-z0-9][A-Za-z0-9._:/#-]{0,255}", value) is None:
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
elif sensitive:
|
||||
raise RuntimeError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
return environment
|
||||
|
||||
def authorize(self, lease_id: str, *, task_id: str, operation: str, scope_type: str = "system", tenant_id: str | None = None, user_id: str | None = None, project_id: str | None = None, run_id: str | None = None) -> Mapping[str, str]:
|
||||
lease = self._get(lease_id)
|
||||
if lease.state != "ACTIVE" or lease.task_id != task_id or operation not in lease.scopes:
|
||||
raise RuntimeError("CREDENTIAL_SCOPE_DENY")
|
||||
requested = (scope_type, tenant_id, user_id, project_id, run_id)
|
||||
actual = (lease.scope_type, lease.tenant_id, lease.user_id, lease.project_id, lease.run_id)
|
||||
if requested != actual:
|
||||
raise RuntimeError("CREDENTIAL_SCOPE_DENY")
|
||||
return self.environment(lease_id)
|
||||
|
||||
def status(self, lease_id: str) -> dict[str, Any]:
|
||||
return self._get(lease_id).as_dict()
|
||||
|
||||
@@ -230,15 +230,7 @@ class JobStore:
|
||||
kind TEXT NOT NULL, digest TEXT NOT NULL, payload_json TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL, FOREIGN KEY(job_id) REFERENCES jobs(job_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS continuation_cycles (
|
||||
job_id TEXT NOT NULL, idempotency_key TEXT NOT NULL,
|
||||
sequence INTEGER NOT NULL, expected_head_oid TEXT NOT NULL,
|
||||
generation INTEGER NOT NULL, state TEXT NOT NULL,
|
||||
commit_oid TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY(job_id,idempotency_key), UNIQUE(job_id,sequence),
|
||||
FOREIGN KEY(job_id) REFERENCES jobs(job_id)
|
||||
);
|
||||
INSERT OR REPLACE INTO metadata(key, value) VALUES('schema_version', '2');
|
||||
INSERT OR REPLACE INTO metadata(key, value) VALUES('schema_version', '1');
|
||||
"""
|
||||
)
|
||||
except sqlite3.Error as exc:
|
||||
@@ -539,183 +531,6 @@ class GitOpsJobController:
|
||||
except sqlite3.Error as exc:
|
||||
raise ServiceUnavailable("JOB_DB_UNAVAILABLE") from exc
|
||||
|
||||
def begin_continuation_cycle(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
idempotency_key: str,
|
||||
expected_head_oid: str,
|
||||
generation: int,
|
||||
expected_revision: int,
|
||||
) -> dict[str, Any]:
|
||||
if (
|
||||
not idempotency_key
|
||||
or len(idempotency_key) > 128
|
||||
or not all(character.isalnum() or character in "._:-" for character in idempotency_key)
|
||||
or len(expected_head_oid) != 40
|
||||
or any(character not in "0123456789abcdef" for character in expected_head_oid)
|
||||
or generation < 1
|
||||
):
|
||||
raise JobError("CONTINUATION_CYCLE_BINDING_INVALID")
|
||||
try:
|
||||
with self.store._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
row = connection.execute("SELECT * FROM jobs WHERE job_id=?", (job_id,)).fetchone()
|
||||
if row is None:
|
||||
raise JobError("JOB_NOT_FOUND")
|
||||
existing = connection.execute(
|
||||
"SELECT * FROM continuation_cycles WHERE job_id=? AND idempotency_key=?",
|
||||
(job_id, idempotency_key),
|
||||
).fetchone()
|
||||
if existing is not None:
|
||||
if str(existing["expected_head_oid"]) != expected_head_oid or int(existing["generation"]) != generation:
|
||||
raise JobError("CONTINUATION_CYCLE_BINDING_CONFLICT")
|
||||
if str(existing["state"]) == "COMPLETED":
|
||||
connection.commit()
|
||||
return {"completed": True, "commit_oid": str(existing["commit_oid"] or ""), "idempotent": True}
|
||||
if str(row["state"]) == JobState.COMPLETED.value:
|
||||
receipt = json.loads(str(row["receipt_json"] or "{}"))
|
||||
commit_oid = str(receipt.get("authority_oid") or "")
|
||||
lease = connection.execute(
|
||||
"SELECT current_head_oid FROM branch_leases WHERE job_id=? AND generation=? AND state='ACTIVE'",
|
||||
(job_id, generation),
|
||||
).fetchall()
|
||||
if (
|
||||
len(lease) == 1
|
||||
and len(commit_oid) == 40
|
||||
and str(lease[0]["current_head_oid"] or "") == commit_oid
|
||||
and receipt.get("remote_verified") is True
|
||||
):
|
||||
now = utc_now()
|
||||
connection.execute(
|
||||
"UPDATE continuation_cycles SET state='COMPLETED',commit_oid=?,updated_at=? "
|
||||
"WHERE job_id=? AND idempotency_key=? AND state='ACTIVE'",
|
||||
(commit_oid, now, job_id, idempotency_key),
|
||||
)
|
||||
connection.commit()
|
||||
return {"completed": True, "commit_oid": commit_oid, "idempotent": True}
|
||||
raise JobError("CONTINUATION_CYCLE_RECOVERY_REQUIRED")
|
||||
connection.commit()
|
||||
return {"completed": False, "commit_oid": None, "idempotent": True}
|
||||
if int(row["revision"]) != expected_revision:
|
||||
raise JobError("REVISION_CAS_CONFLICT")
|
||||
if str(row["state"]) != JobState.COMPLETED.value:
|
||||
raise JobError("CONTINUATION_CYCLE_JOB_STATE_INVALID")
|
||||
if str(row["branch"]) in {"main", "master"} or str(row["branch"]).startswith(("refs/", "tags/")):
|
||||
raise JobError("PROTECTED_BRANCH_REJECTED")
|
||||
receipt = json.loads(str(row["receipt_json"] or "{}"))
|
||||
if receipt.get("remote_verified") is not True or any(
|
||||
str(receipt.get(field) or "") != expected_head_oid
|
||||
for field in ("authority_oid", "api_branch_oid", "git_remote_oid")
|
||||
):
|
||||
raise JobError("CONTINUATION_CYCLE_RECEIPT_INVALID")
|
||||
leases = connection.execute(
|
||||
"SELECT * FROM branch_leases WHERE job_id=? AND branch=? AND generation=? AND state='ACTIVE'",
|
||||
(job_id, str(row["branch"]), generation),
|
||||
).fetchall()
|
||||
if len(leases) != 1 or str(leases[0]["current_head_oid"] or "") != expected_head_oid:
|
||||
raise JobError("CONTINUATION_CYCLE_LEASE_INVALID")
|
||||
if connection.execute(
|
||||
"SELECT 1 FROM continuation_cycles WHERE job_id=? AND state='ACTIVE' LIMIT 1",
|
||||
(job_id,),
|
||||
).fetchone() is not None:
|
||||
raise JobError("CONTINUATION_CYCLE_ACTIVE_CONFLICT")
|
||||
sequence = int(connection.execute(
|
||||
"SELECT COALESCE(MAX(sequence),0)+1 FROM continuation_cycles WHERE job_id=?",
|
||||
(job_id,),
|
||||
).fetchone()[0])
|
||||
now = utc_now()
|
||||
connection.execute(
|
||||
"INSERT INTO continuation_cycles(job_id,idempotency_key,sequence,expected_head_oid,generation,state,commit_oid,created_at,updated_at) "
|
||||
"VALUES(?,?,?,?,?,'ACTIVE',NULL,?,?)",
|
||||
(job_id, idempotency_key, sequence, expected_head_oid, generation, now, now),
|
||||
)
|
||||
next_revision = expected_revision + 1
|
||||
updated = connection.execute(
|
||||
"UPDATE jobs SET state=?,side_effect_state=?,revision=?,updated_at=? WHERE job_id=? AND revision=?",
|
||||
(JobState.LEASED.value, "CONTINUATION_LEASED", next_revision, now, job_id, expected_revision),
|
||||
)
|
||||
if updated.rowcount != 1:
|
||||
raise JobError("REVISION_CAS_CONFLICT")
|
||||
connection.execute(
|
||||
"INSERT INTO state_transitions(job_id,from_state,to_state,revision,side_effect_state,created_at) VALUES(?,?,?,?,?,?)",
|
||||
(job_id, JobState.COMPLETED.value, JobState.LEASED.value, next_revision, "CONTINUATION_LEASED", now),
|
||||
)
|
||||
cycle_receipt = {
|
||||
"idempotency_key": idempotency_key,
|
||||
"expected_head_oid": expected_head_oid,
|
||||
"generation": generation,
|
||||
"sequence": sequence,
|
||||
}
|
||||
connection.execute(
|
||||
"INSERT INTO receipts(receipt_id,job_id,kind,digest,payload_json,created_at) VALUES(?,?,?,?,?,?)",
|
||||
(
|
||||
f"receipt-{job_id}-{next_revision}",
|
||||
job_id,
|
||||
"CONTINUATION_LEASED",
|
||||
hashlib.sha256(canonical_json(cycle_receipt)).hexdigest(),
|
||||
json.dumps(cycle_receipt, sort_keys=True),
|
||||
now,
|
||||
),
|
||||
)
|
||||
result = dict(connection.execute("SELECT * FROM jobs WHERE job_id=?", (job_id,)).fetchone())
|
||||
connection.commit()
|
||||
except (JobError, BindingConflict):
|
||||
raise
|
||||
except sqlite3.Error as exc:
|
||||
raise ServiceUnavailable("JOB_DB_UNAVAILABLE") from exc
|
||||
result["request"] = json.loads(result["request_json"])
|
||||
self.store._mirror(result)
|
||||
return {"completed": False, "commit_oid": None, "idempotent": False}
|
||||
|
||||
def complete_continuation_cycle(self, job_id: str, *, idempotency_key: str, commit_oid: str) -> bool:
|
||||
if len(commit_oid) != 40 or any(character not in "0123456789abcdef" for character in commit_oid):
|
||||
raise JobError("CONTINUATION_CYCLE_COMMIT_OID_INVALID")
|
||||
try:
|
||||
with self.store._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
cycle = connection.execute(
|
||||
"SELECT * FROM continuation_cycles WHERE job_id=? AND idempotency_key=?",
|
||||
(job_id, idempotency_key),
|
||||
).fetchone()
|
||||
if cycle is None:
|
||||
connection.commit()
|
||||
return False
|
||||
if str(cycle["state"]) == "COMPLETED":
|
||||
if str(cycle["commit_oid"] or "") != commit_oid:
|
||||
raise JobError("CONTINUATION_CYCLE_COMMIT_CONFLICT")
|
||||
connection.commit()
|
||||
return True
|
||||
row = connection.execute("SELECT * FROM jobs WHERE job_id=?", (job_id,)).fetchone()
|
||||
leases = connection.execute(
|
||||
"SELECT current_head_oid FROM branch_leases WHERE job_id=? AND generation=? AND state='ACTIVE'",
|
||||
(job_id, int(cycle["generation"])),
|
||||
).fetchall()
|
||||
receipt = json.loads(str(row["receipt_json"] or "{}")) if row is not None else {}
|
||||
if (
|
||||
row is None
|
||||
or str(row["state"]) != JobState.COMPLETED.value
|
||||
or receipt.get("remote_verified") is not True
|
||||
or str(receipt.get("authority_oid") or "") != commit_oid
|
||||
or len(leases) != 1
|
||||
or str(leases[0]["current_head_oid"] or "") != commit_oid
|
||||
):
|
||||
raise JobError("CONTINUATION_CYCLE_COMPLETION_EVIDENCE_INVALID")
|
||||
now = utc_now()
|
||||
updated = connection.execute(
|
||||
"UPDATE continuation_cycles SET state='COMPLETED',commit_oid=?,updated_at=? "
|
||||
"WHERE job_id=? AND idempotency_key=? AND state='ACTIVE'",
|
||||
(commit_oid, now, job_id, idempotency_key),
|
||||
)
|
||||
if updated.rowcount != 1:
|
||||
raise JobError("CONTINUATION_CYCLE_CAS_CONFLICT")
|
||||
connection.commit()
|
||||
return True
|
||||
except (JobError, BindingConflict):
|
||||
raise
|
||||
except sqlite3.Error as exc:
|
||||
raise ServiceUnavailable("JOB_DB_UNAVAILABLE") from exc
|
||||
|
||||
def advance(self, job_id: str, state: JobState, *, side_effect_state: str, receipt: Mapping[str, Any] | None = None, expected_revision: int | None = None) -> JobReceipt:
|
||||
return self._advance(
|
||||
job_id,
|
||||
|
||||
@@ -2,8 +2,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, asdict
|
||||
from contextlib import contextmanager
|
||||
import fcntl
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
@@ -11,11 +9,6 @@ import re
|
||||
from typing import Any, Mapping
|
||||
|
||||
from aios_gitops_job_controller import GitOpsJobController, JobStore
|
||||
from aios_online_agent_execution import (
|
||||
GovernedExecutionBinding,
|
||||
project_execution_envelope,
|
||||
project_online_task,
|
||||
)
|
||||
|
||||
OID_RE = re.compile(r"^[0-9a-f]{40}$")
|
||||
BRANCH_RE = re.compile(r"^agent/[A-Za-z0-9][A-Za-z0-9._/-]{2,179}$")
|
||||
@@ -43,161 +36,35 @@ class GovernedTaskBinding:
|
||||
return {"task_id": self.task_id, "repository": self.repository, "base_oid": self.base_oid, "target_branch": self.target_branch, "expected_branch_oid": self.expected_branch_oid, "generation": self.generation, "branch_mode": self.branch_mode, "create_target_branch": self.branch_mode == "CREATE_NEW", "worktree_mode": "detached_at_base_oid" if self.branch_mode == "CREATE_NEW" else "detached_at_expected_head", "cas_ref_update_required": True, "idempotent": idempotent, "state": "BOUND", "job_id": self.job_id}
|
||||
|
||||
|
||||
def _legacy_safe(value: str) -> str:
|
||||
def _safe(value: str) -> str:
|
||||
return re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-")[:120] or "task"
|
||||
|
||||
|
||||
def _task_key(value: str) -> str:
|
||||
readable = _legacy_safe(value)[:72]
|
||||
return f"{readable}-{hashlib.sha256(value.encode('utf-8')).hexdigest()[:24]}"
|
||||
|
||||
|
||||
class GovernedTaskBootstrap:
|
||||
def __init__(self, authority: Any, root: Path, *, remote_name: str, token_path: Path, ssh_key_path: Path, api_base: str = "https://g.cnnetmkt.top/git/api/v1") -> None:
|
||||
self.authority = authority; self.root = Path(root); self.remote_name = remote_name; self.token_path = Path(token_path); self.ssh_key_path = Path(ssh_key_path); self.api_base = api_base; self.task_meta_root = self.root / "task-bindings"; self.job_root = self.root / "jobs"; self.lock_root = self.root / "task-locks"; self.task_meta_root.mkdir(parents=True, exist_ok=True); self.job_root.mkdir(parents=True, exist_ok=True); self.lock_root.mkdir(parents=True, exist_ok=True)
|
||||
self.authority = authority; self.root = Path(root); self.remote_name = remote_name; self.token_path = Path(token_path); self.ssh_key_path = Path(ssh_key_path); self.api_base = api_base; self.task_meta_root = self.root / "task-bindings"; self.job_root = self.root / "jobs"; self.task_meta_root.mkdir(parents=True, exist_ok=True); self.job_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _store(self, task_id: str) -> tuple[JobStore, GitOpsJobController]:
|
||||
legacy = self.job_root / _legacy_safe(task_id)
|
||||
keyed = self.job_root / _task_key(task_id)
|
||||
root = legacy if legacy.is_dir() and not keyed.exists() else keyed
|
||||
store = JobStore(root, payload_root=root / "payloads"); return store, GitOpsJobController(store)
|
||||
root = self.job_root / _safe(task_id); store = JobStore(root, payload_root=root / "payloads"); return store, GitOpsJobController(store)
|
||||
|
||||
def _metadata(self, task_id: str) -> Path:
|
||||
return self.task_meta_root / (_task_key(task_id) + ".json")
|
||||
|
||||
def _legacy_metadata(self, task_id: str) -> Path:
|
||||
return self.task_meta_root / (_legacy_safe(task_id) + ".json")
|
||||
|
||||
@contextmanager
|
||||
def _task_lock(self, task_id: str):
|
||||
path = self.lock_root / (_task_key(task_id) + ".lock")
|
||||
with path.open("a+b") as handle:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
@contextmanager
|
||||
def _branch_lock(self, repository: str, target_branch: str):
|
||||
key = hashlib.sha256((repository + "\0" + target_branch).encode("utf-8")).hexdigest()
|
||||
path = self.lock_root / ("branch-" + key + ".lock")
|
||||
with path.open("a+b") as handle:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
return self.task_meta_root / (_safe(task_id) + ".json")
|
||||
|
||||
def _read_existing(self, task_id: str) -> GovernedTaskBinding | None:
|
||||
path = self._metadata(task_id)
|
||||
if not path.exists():
|
||||
path = self._legacy_metadata(task_id)
|
||||
if not path.is_file() or path.is_symlink(): return None
|
||||
binding = GovernedTaskBinding(**json.loads(path.read_text(encoding="utf-8")))
|
||||
if binding.task_id != task_id:
|
||||
raise BootstrapError("TASK_BINDING_KEY_COLLISION")
|
||||
return binding
|
||||
|
||||
def get_binding(self, task_id: str) -> GovernedTaskBinding:
|
||||
"""Return the exact durable binding; never synthesize missing lineage."""
|
||||
existing = self._read_existing(str(task_id))
|
||||
if existing is None or not existing.job_id:
|
||||
raise BootstrapError("TASK_EXECUTION_PROJECTION_MISSING")
|
||||
return existing
|
||||
|
||||
def execution_binding(self, binding: GovernedTaskBinding) -> GovernedExecutionBinding:
|
||||
expected = binding.expected_branch_oid
|
||||
mode = binding.branch_mode
|
||||
if mode == "CREATE_NEW":
|
||||
current = self.authority.read_ref("refs/heads/" + binding.target_branch)
|
||||
if current is not None:
|
||||
store, _ = self._store(binding.task_id)
|
||||
with store._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT current_head_oid,generation,state,lineage_anchor_oid "
|
||||
"FROM branch_leases WHERE job_id=? AND branch=? AND state='ACTIVE' "
|
||||
"ORDER BY revision DESC",
|
||||
(binding.job_id, binding.target_branch),
|
||||
).fetchall()
|
||||
if len(rows) != 1:
|
||||
raise BootstrapError("SAME_JOB_CONTINUATION_LEASE_INVALID")
|
||||
row = rows[0]
|
||||
if (
|
||||
str(row["current_head_oid"] or "") != current
|
||||
or int(row["generation"] or 0) != binding.generation
|
||||
or str(row["lineage_anchor_oid"] or "") != binding.base_oid
|
||||
):
|
||||
raise BootstrapError("SAME_JOB_CONTINUATION_BINDING_INVALID")
|
||||
expected = current
|
||||
mode = "CONTINUE_EXISTING"
|
||||
return GovernedExecutionBinding.from_mapping(
|
||||
{
|
||||
"task_id": binding.task_id,
|
||||
"job_id": binding.job_id,
|
||||
"repository": binding.repository,
|
||||
"base_oid": binding.base_oid,
|
||||
"expected_branch_oid": expected,
|
||||
"target_branch": binding.target_branch,
|
||||
"branch_mode": mode,
|
||||
"generation": binding.generation,
|
||||
"idempotency_key": binding.idempotency_key,
|
||||
}
|
||||
)
|
||||
|
||||
def execution_envelope(
|
||||
self,
|
||||
task_id: str,
|
||||
execution_request: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
required_capabilities: tuple[str, ...] = (),
|
||||
expected_job_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
binding = self.get_binding(task_id)
|
||||
if expected_job_id and binding.job_id != expected_job_id:
|
||||
raise BootstrapError("JOB_BINDING_CONFLICT")
|
||||
return project_execution_envelope(
|
||||
self.execution_binding(binding),
|
||||
execution_request,
|
||||
required_capabilities=required_capabilities,
|
||||
)
|
||||
|
||||
def online_task_projection(
|
||||
self,
|
||||
task_id: str,
|
||||
task_request: Mapping[str, Any],
|
||||
execution_request: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
expected_job_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
binding = self.get_binding(task_id)
|
||||
if expected_job_id and binding.job_id != expected_job_id:
|
||||
raise BootstrapError("JOB_BINDING_CONFLICT")
|
||||
return project_online_task(self.execution_binding(binding), task_request, execution_request)
|
||||
return GovernedTaskBinding(**json.loads(path.read_text(encoding="utf-8")))
|
||||
|
||||
def _write(self, binding: GovernedTaskBinding) -> None:
|
||||
path = self._metadata(binding.task_id); temporary = path.with_suffix(".tmp"); temporary.write_text(json.dumps(asdict(binding), sort_keys=True), encoding="utf-8"); temporary.replace(path)
|
||||
|
||||
def create(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
task_id = str(request.get("task_id") or "")
|
||||
if not task_id:
|
||||
raise BootstrapError("TASK_BINDING_INVALID")
|
||||
repository = str(request.get("repository") or "")
|
||||
target_branch = str(request.get("target_branch") or "")
|
||||
with self._task_lock(task_id):
|
||||
with self._branch_lock(repository, target_branch):
|
||||
return self._create_locked(request)
|
||||
|
||||
def _create_locked(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
||||
values = {key: str(request[key]) for key in ("task_id", "repository", "base_oid", "target_branch", "risk", "description", "idempotency_key")}
|
||||
if values["repository"] != "aios/aios" or not OID_RE.fullmatch(values["base_oid"]): raise BootstrapError("TASK_BINDING_INVALID")
|
||||
if not BRANCH_RE.fullmatch(values["target_branch"]): raise BootstrapError("TARGET_BRANCH_INVALID")
|
||||
existing = self._read_existing(values["task_id"])
|
||||
if existing is not None:
|
||||
if any(getattr(existing, key) != values[key] for key in ("repository", "base_oid", "target_branch", "risk", "description", "idempotency_key")): raise BootstrapError("TASK_BINDING_CONFLICT")
|
||||
optional = {"expected_branch_oid": existing.expected_branch_oid, "branch_mode": existing.branch_mode, "generation": existing.generation}
|
||||
for key, durable in optional.items():
|
||||
if key in request and str(request[key]) != str(durable): raise BootstrapError("TASK_BINDING_CONFLICT")
|
||||
if any(getattr(existing, key) != values[key] for key in ("repository", "base_oid", "target_branch", "risk", "description")): raise BootstrapError("TASK_BINDING_CONFLICT")
|
||||
return existing.response(idempotent=True)
|
||||
if not self.authority.commit_exists(values["base_oid"]): raise BootstrapError("BASE_COMMIT_NOT_FOUND")
|
||||
current = self.authority.read_ref("refs/heads/" + values["target_branch"])
|
||||
@@ -207,24 +74,8 @@ class GovernedTaskBootstrap:
|
||||
if expected != values["base_oid"]: raise BootstrapError("EXPECTED_BASE_CONFLICT")
|
||||
else:
|
||||
if expected != current: raise BootstrapError("EXISTING_BRANCH_OID_CONFLICT")
|
||||
parent_task_id = str(request.get("parent_task_id") or "")
|
||||
parent_job_id = str(request.get("parent_job_id") or "")
|
||||
if not parent_task_id or not parent_job_id:
|
||||
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_REQUIRED")
|
||||
parent = self._read_existing(parent_task_id)
|
||||
if parent is None or parent.job_id != parent_job_id or parent.repository != values["repository"] or parent.target_branch != values["target_branch"]:
|
||||
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_CONFLICT")
|
||||
store, _ = self._store(parent_task_id)
|
||||
with store._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT current_head_oid,generation,state FROM branch_leases WHERE job_id=? AND branch=? AND state='ACTIVE' ORDER BY revision DESC",
|
||||
(parent_job_id, values["target_branch"]),
|
||||
).fetchall()
|
||||
if len(rows) != 1 or str(rows[0]["current_head_oid"] or "") != current:
|
||||
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_CONFLICT")
|
||||
mode, generation = "CONTINUE_EXISTING", int(rows[0]["generation"]) + 1
|
||||
if "generation" in request and int(request["generation"]) != generation:
|
||||
raise BootstrapError("GENERATION_BINDING_CONFLICT")
|
||||
mode, generation = "CONTINUE_EXISTING", int(request.get("generation", 2))
|
||||
if generation < 1: raise BootstrapError("GENERATION_INVALID")
|
||||
payload = json.dumps(values, sort_keys=True, separators=(",", ":")).encode()
|
||||
_, controller = self._store(values["task_id"])
|
||||
receipt = controller.submit({"idempotency_key": values["idempotency_key"], "campaign_id": values["task_id"], "task_id": values["task_id"], "repository": values["repository"], "operation": "GOVERNED_TASK_CREATE", "branch": values["target_branch"], "expected_branch_oid": expected, "payload_digest": hashlib.sha256(payload).hexdigest(), "lineage_anchor_oid": values["base_oid"]}, payload=payload)
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Durable checkpoint journal for the generic Online Agent worker.
|
||||
|
||||
The JSON task state remains a fast cache. This SQLite journal is the durable,
|
||||
append-only evidence used for restart/resume/replay decisions.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
from typing import Any, Mapping
|
||||
|
||||
PHASE_ORDER = {
|
||||
"claimed": 10,
|
||||
"workspace_prepared": 20,
|
||||
"verification_passed": 30,
|
||||
"source_capture_done": 40,
|
||||
"commit_done": 40,
|
||||
"reconciliation_required": 45,
|
||||
"result_pending": 50,
|
||||
"result_submitted": 60,
|
||||
}
|
||||
|
||||
|
||||
class CheckpointError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def canonical(value: object) -> bytes:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str).encode("utf-8")
|
||||
|
||||
|
||||
def sha256_json(value: object) -> str:
|
||||
return hashlib.sha256(canonical(value)).hexdigest()
|
||||
|
||||
|
||||
def now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CheckpointBinding:
|
||||
task_id: str
|
||||
attempt: int
|
||||
execution_digest: str
|
||||
job_id: str
|
||||
repository: str
|
||||
base_oid: str
|
||||
expected_branch_oid: str
|
||||
target_branch: str
|
||||
generation: int
|
||||
|
||||
def material(self) -> dict[str, Any]:
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"attempt": self.attempt,
|
||||
"execution_digest": self.execution_digest,
|
||||
"job_id": self.job_id,
|
||||
"repository": self.repository,
|
||||
"base_oid": self.base_oid,
|
||||
"expected_branch_oid": self.expected_branch_oid,
|
||||
"target_branch": self.target_branch,
|
||||
"generation": self.generation,
|
||||
}
|
||||
|
||||
@property
|
||||
def binding_sha256(self) -> str:
|
||||
return sha256_json(self.material())
|
||||
|
||||
|
||||
class CheckpointStore:
|
||||
def __init__(self, path: Path) -> None:
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._connect() as connection:
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS checkpoints(
|
||||
task_id TEXT NOT NULL,
|
||||
attempt INTEGER NOT NULL,
|
||||
sequence INTEGER NOT NULL,
|
||||
phase TEXT NOT NULL,
|
||||
phase_order INTEGER NOT NULL,
|
||||
execution_digest TEXT NOT NULL,
|
||||
job_id TEXT NOT NULL,
|
||||
repository TEXT NOT NULL,
|
||||
base_oid TEXT NOT NULL,
|
||||
expected_branch_oid TEXT NOT NULL,
|
||||
target_branch TEXT NOT NULL,
|
||||
generation INTEGER NOT NULL,
|
||||
binding_sha256 TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL,
|
||||
payload_sha256 TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY(task_id, attempt, sequence)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_checkpoints_task_latest
|
||||
ON checkpoints(task_id, attempt DESC, sequence DESC);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS idx_checkpoints_exact_replay
|
||||
ON checkpoints(task_id, attempt, phase, payload_sha256);
|
||||
CREATE TABLE IF NOT EXISTS checkpoint_heads(
|
||||
task_id TEXT PRIMARY KEY,
|
||||
attempt INTEGER NOT NULL,
|
||||
sequence INTEGER NOT NULL,
|
||||
phase TEXT NOT NULL,
|
||||
phase_order INTEGER NOT NULL,
|
||||
execution_digest TEXT NOT NULL,
|
||||
binding_sha256 TEXT NOT NULL,
|
||||
payload_sha256 TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.path, timeout=30.0)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys=ON")
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
connection.execute("PRAGMA synchronous=FULL")
|
||||
connection.execute("PRAGMA busy_timeout=30000")
|
||||
return connection
|
||||
|
||||
@staticmethod
|
||||
def _phase_order(phase: str) -> int:
|
||||
try:
|
||||
return PHASE_ORDER[phase]
|
||||
except KeyError as exc:
|
||||
raise CheckpointError("CHECKPOINT_PHASE_INVALID") from exc
|
||||
|
||||
def append(
|
||||
self,
|
||||
binding: CheckpointBinding,
|
||||
phase: str,
|
||||
payload: Mapping[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
order = self._phase_order(phase)
|
||||
payload_value = dict(payload)
|
||||
payload_sha = sha256_json(payload_value)
|
||||
with self._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
replay = connection.execute(
|
||||
"""
|
||||
SELECT * FROM checkpoints
|
||||
WHERE task_id=? AND attempt=? AND phase=? AND payload_sha256=?
|
||||
""",
|
||||
(binding.task_id, binding.attempt, phase, payload_sha),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
connection.commit()
|
||||
return dict(replay)
|
||||
|
||||
head = connection.execute(
|
||||
"SELECT * FROM checkpoint_heads WHERE task_id=?",
|
||||
(binding.task_id,),
|
||||
).fetchone()
|
||||
if head is not None:
|
||||
head_attempt = int(head["attempt"])
|
||||
head_order = int(head["phase_order"])
|
||||
if binding.attempt < head_attempt:
|
||||
raise CheckpointError("CHECKPOINT_ATTEMPT_REGRESSION")
|
||||
if binding.attempt == head_attempt:
|
||||
if head["execution_digest"] != binding.execution_digest:
|
||||
raise CheckpointError("CHECKPOINT_EXECUTION_DIGEST_CONFLICT")
|
||||
if head["binding_sha256"] != binding.binding_sha256:
|
||||
raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
|
||||
if order < head_order and phase != "reconciliation_required":
|
||||
raise CheckpointError("CHECKPOINT_PHASE_REGRESSION")
|
||||
sequence = int(head["sequence"]) + 1
|
||||
else:
|
||||
sequence = 1
|
||||
else:
|
||||
sequence = 1
|
||||
|
||||
created_at = now()
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO checkpoints(
|
||||
task_id,attempt,sequence,phase,phase_order,execution_digest,job_id,
|
||||
repository,base_oid,expected_branch_oid,target_branch,generation,
|
||||
binding_sha256,payload_json,payload_sha256,created_at
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
||||
""",
|
||||
(
|
||||
binding.task_id,
|
||||
binding.attempt,
|
||||
sequence,
|
||||
phase,
|
||||
order,
|
||||
binding.execution_digest,
|
||||
binding.job_id,
|
||||
binding.repository,
|
||||
binding.base_oid,
|
||||
binding.expected_branch_oid,
|
||||
binding.target_branch,
|
||||
binding.generation,
|
||||
binding.binding_sha256,
|
||||
canonical(payload_value).decode("utf-8"),
|
||||
payload_sha,
|
||||
created_at,
|
||||
),
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO checkpoint_heads(
|
||||
task_id,attempt,sequence,phase,phase_order,execution_digest,
|
||||
binding_sha256,payload_sha256,updated_at
|
||||
) VALUES (?,?,?,?,?,?,?,?,?)
|
||||
ON CONFLICT(task_id) DO UPDATE SET
|
||||
attempt=excluded.attempt,
|
||||
sequence=excluded.sequence,
|
||||
phase=excluded.phase,
|
||||
phase_order=excluded.phase_order,
|
||||
execution_digest=excluded.execution_digest,
|
||||
binding_sha256=excluded.binding_sha256,
|
||||
payload_sha256=excluded.payload_sha256,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
binding.task_id,
|
||||
binding.attempt,
|
||||
sequence,
|
||||
phase,
|
||||
order,
|
||||
binding.execution_digest,
|
||||
binding.binding_sha256,
|
||||
payload_sha,
|
||||
created_at,
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
row = connection.execute(
|
||||
"SELECT * FROM checkpoints WHERE task_id=? AND attempt=? AND sequence=?",
|
||||
(binding.task_id, binding.attempt, sequence),
|
||||
).fetchone()
|
||||
return dict(row)
|
||||
|
||||
def latest(self, task_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT * FROM checkpoints
|
||||
WHERE task_id=?
|
||||
ORDER BY attempt DESC, sequence DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
value = dict(row)
|
||||
value["payload"] = json.loads(value.pop("payload_json"))
|
||||
return value
|
||||
|
||||
def latest_for_attempt(self, task_id: str, attempt: int) -> dict[str, Any] | None:
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT * FROM checkpoints
|
||||
WHERE task_id=? AND attempt=?
|
||||
ORDER BY sequence DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(task_id, attempt),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
value = dict(row)
|
||||
value["payload"] = json.loads(value.pop("payload_json"))
|
||||
return value
|
||||
|
||||
def verify_latest_binding(self, binding: CheckpointBinding) -> dict[str, Any] | None:
|
||||
value = self.latest(binding.task_id)
|
||||
if value is None:
|
||||
return None
|
||||
if int(value["attempt"]) > binding.attempt:
|
||||
raise CheckpointError("CHECKPOINT_FUTURE_ATTEMPT")
|
||||
# A previous attempt may be reused only when the immutable execution/scope binding matches.
|
||||
if value["execution_digest"] != binding.execution_digest:
|
||||
return None
|
||||
expected = CheckpointBinding(
|
||||
task_id=value["task_id"],
|
||||
attempt=int(value["attempt"]),
|
||||
execution_digest=value["execution_digest"],
|
||||
job_id=value["job_id"],
|
||||
repository=value["repository"],
|
||||
base_oid=value["base_oid"],
|
||||
expected_branch_oid=value["expected_branch_oid"],
|
||||
target_branch=value["target_branch"],
|
||||
generation=int(value["generation"]),
|
||||
)
|
||||
current_scope = binding.material()
|
||||
historical_scope = expected.material()
|
||||
current_scope.pop("attempt", None)
|
||||
historical_scope.pop("attempt", None)
|
||||
if sha256_json(current_scope) != sha256_json(historical_scope):
|
||||
raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
|
||||
return value
|
||||
@@ -0,0 +1,463 @@
|
||||
"""Bounded read-only repository interface map for governed Online Agent analysis.
|
||||
|
||||
The compact map intentionally returns symbols and short source excerpts rather than a
|
||||
source archive. It is designed for control-plane consumers that cannot safely ingest a
|
||||
multi-megabyte repository snapshot.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import Any, Iterable, Mapping, Sequence
|
||||
|
||||
|
||||
SCHEMA = "aios.engineering.interface-map.v1"
|
||||
MAX_RESULT_BYTES = 120 * 1024
|
||||
TARGET_RESULT_BYTES = 108 * 1024
|
||||
MAX_MODULES = 40
|
||||
MAX_SYMBOLS = 80
|
||||
MAX_TESTS = 40
|
||||
MAX_EXCERPTS = 25
|
||||
MAX_EXCERPT_CHARS = 1200
|
||||
MAX_FILE_BYTES = 2 * 1024 * 1024
|
||||
|
||||
ROOTS = (
|
||||
"active-runtime/app",
|
||||
"runtime/agent-os",
|
||||
"tests",
|
||||
"db/migrations",
|
||||
"migrations",
|
||||
"deploy/systemd",
|
||||
"docs/task-plane",
|
||||
"docs/superpowers/specs",
|
||||
)
|
||||
SUFFIXES = {".py", ".sql", ".service", ".md", ".json", ".yaml", ".yml"}
|
||||
|
||||
DEFAULT_TERMS = {
|
||||
"agent", "worker", "fleet", "scheduler", "scheduling", "task", "claim", "lease",
|
||||
"heartbeat", "capacity", "concurrency", "priority", "fairness", "queue", "retry",
|
||||
"backoff", "circuit", "health", "liveness", "scope", "tenant", "project", "trace",
|
||||
"message", "checkpoint", "recovery", "idempotency", "dispatch", "assignment",
|
||||
}
|
||||
FLEET_CONTRACT_TERMS: dict[str, tuple[str, ...]] = {
|
||||
"FLEET_AGENT_REGISTRATION": ("agent", "register", "worker", "capability", "version"),
|
||||
"FLEET_CAPACITY_SCHEDULING": ("capacity", "slot", "concurrency", "scheduler", "queue"),
|
||||
"FLEET_FAIRNESS": ("fair", "priority", "queue", "tenant", "project", "starvation"),
|
||||
"FLEET_FAULT_ISOLATION": ("failure", "health", "liveness", "timeout", "circuit", "poison"),
|
||||
"FLEET_WORKER_REASSIGNMENT": ("reassign", "assignment", "claim", "lease", "retry", "worker"),
|
||||
"FLEET_SCOPE_ISOLATION": ("tenant", "project", "scope", "conversation", "task", "run"),
|
||||
"FLEET_IDEMPOTENT_DISPATCH": ("idempot", "dispatch", "claim", "lease", "result", "attempt"),
|
||||
"FLEET_FAILURE_RECOVERY": ("checkpoint", "resume", "recover", "retry", "cleanup", "receipt"),
|
||||
}
|
||||
|
||||
MODEL_GOVERNANCE_CONTRACT_TERMS: dict[str, tuple[str, ...]] = {
|
||||
"MODEL_PROVIDER_REGISTRY": ("model", "provider", "registry", "endpoint", "capability"),
|
||||
"MODEL_SELECTION_POLICY": ("model", "select", "route", "policy", "capability"),
|
||||
"MODEL_VERSION_PINNING": ("model", "version", "pin", "revision", "snapshot"),
|
||||
"MODEL_BUDGET_ACCOUNTING": ("budget", "token", "cost", "quota", "usage"),
|
||||
"MODEL_FALLBACK_POLICY": ("fallback", "retry", "route", "provider", "model"),
|
||||
"MODEL_SCOPE_ISOLATION": ("tenant", "project", "scope", "policy", "model"),
|
||||
"MODEL_CREDENTIAL_BOUNDARY": ("credential", "secret", "provider", "tool", "token"),
|
||||
"MODEL_TRACE_AUDIT": ("trace", "audit", "model", "provider", "request"),
|
||||
"MODEL_EVALUATION_GATE": ("eval", "evaluation", "approval", "quality", "gate"),
|
||||
"MODEL_ROUTING_RESILIENCE": ("route", "retry", "fallback", "timeout", "health"),
|
||||
}
|
||||
|
||||
PROFILE_CLASSIFICATION = {
|
||||
"fleet": "P2_INTERFACE_MAP_CAPTURED",
|
||||
"model_governance": "MODEL_GOVERNANCE_INTERFACE_MAP_CAPTURED",
|
||||
}
|
||||
|
||||
MODEL_PROFILE_MARKERS = {
|
||||
"model", "provider", "llm", "inference", "routing", "router", "budget",
|
||||
"token", "cost", "quota", "fallback", "evaluation", "eval", "approval",
|
||||
"version", "pinning",
|
||||
}
|
||||
SENSITIVE_RE = re.compile(
|
||||
r"(?i)(authorization\s*:\s*bearer\s+\S+|"
|
||||
r"(?:password|passwd|secret|token|api[_-]?key)\s*[:=]\s*['\"]?[^'\"\s,;]+)"
|
||||
)
|
||||
SQL_TABLE_RE = re.compile(
|
||||
r"(?is)\bcreate\s+table\s+(?:if\s+not\s+exists\s+)?[\"`[]?([A-Za-z_][A-Za-z0-9_.-]*)"
|
||||
)
|
||||
|
||||
|
||||
class CompactInterfaceMapError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _sha(data: bytes) -> str:
|
||||
return hashlib.sha256(data).hexdigest()
|
||||
|
||||
|
||||
def _bounded_excerpt(text: str, start_line: int, end_line: int) -> str:
|
||||
lines = text.splitlines()
|
||||
start = max(0, start_line - 2)
|
||||
end = min(len(lines), max(end_line + 1, start + 1))
|
||||
excerpt = "\n".join(lines[start:end])
|
||||
excerpt = SENSITIVE_RE.sub("<redacted-sensitive-source-pattern>", excerpt)
|
||||
if len(excerpt) > MAX_EXCERPT_CHARS:
|
||||
excerpt = excerpt[: MAX_EXCERPT_CHARS - 1] + "…"
|
||||
return excerpt
|
||||
|
||||
|
||||
def _signal_terms(signals: Sequence[str]) -> set[str]:
|
||||
terms = set(DEFAULT_TERMS)
|
||||
for signal in signals:
|
||||
if not isinstance(signal, str):
|
||||
continue
|
||||
terms.update(
|
||||
token.lower()
|
||||
for token in re.findall(r"[A-Za-z][A-Za-z0-9_.-]{2,63}", signal)
|
||||
)
|
||||
return set(sorted(terms)[:96])
|
||||
|
||||
|
||||
def _profile(source_signals: Sequence[str]) -> str:
|
||||
tokens: set[str] = set()
|
||||
for signal in source_signals:
|
||||
if not isinstance(signal, str):
|
||||
continue
|
||||
tokens.update(
|
||||
token.lower()
|
||||
for token in re.findall(r"[A-Za-z][A-Za-z0-9_.-]{2,63}", signal)
|
||||
)
|
||||
model_score = len(tokens & MODEL_PROFILE_MARKERS)
|
||||
return "model_governance" if model_score >= 2 else "fleet"
|
||||
|
||||
|
||||
def _contract_terms(profile: str) -> Mapping[str, tuple[str, ...]]:
|
||||
if profile == "model_governance":
|
||||
return MODEL_GOVERNANCE_CONTRACT_TERMS
|
||||
if profile == "fleet":
|
||||
return FLEET_CONTRACT_TERMS
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_PROFILE_INVALID")
|
||||
|
||||
|
||||
def _candidate_files(root: Path, terms: set[str]) -> list[dict[str, Any]]:
|
||||
root = root.resolve()
|
||||
candidates: list[dict[str, Any]] = []
|
||||
for relative_root in ROOTS:
|
||||
directory = (root / relative_root).resolve(strict=False)
|
||||
if directory == root or root not in directory.parents:
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_SCOPE_INVALID")
|
||||
if not directory.is_dir() or directory.is_symlink():
|
||||
continue
|
||||
for path in sorted(directory.rglob("*")):
|
||||
if path.is_symlink() or not path.is_file() or path.suffix.lower() not in SUFFIXES:
|
||||
continue
|
||||
resolved = path.resolve(strict=True)
|
||||
if root not in resolved.parents:
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_PATH_ESCAPE")
|
||||
data = path.read_bytes()
|
||||
if len(data) > MAX_FILE_BYTES:
|
||||
continue
|
||||
relative = path.relative_to(root).as_posix()
|
||||
text = data.decode("utf-8", "ignore")
|
||||
searchable = (relative + "\n" + text[:384000]).lower()
|
||||
matched = sorted(term for term in terms if term in searchable)
|
||||
if not matched:
|
||||
continue
|
||||
path_bonus = sum(3 for term in terms if term in relative.lower())
|
||||
test_bonus = 3 if relative.startswith("tests/") else 0
|
||||
score = min(250, len(matched) + path_bonus + test_bonus)
|
||||
candidates.append(
|
||||
{
|
||||
"path": relative,
|
||||
"sha256": _sha(data),
|
||||
"bytes": len(data),
|
||||
"score": score,
|
||||
"text": text,
|
||||
"suffix": path.suffix.lower(),
|
||||
}
|
||||
)
|
||||
candidates.sort(key=lambda item: (-item["score"], item["path"]))
|
||||
return candidates
|
||||
|
||||
|
||||
def _python_symbols(item: Mapping[str, Any], terms: set[str]) -> list[dict[str, Any]]:
|
||||
try:
|
||||
tree = ast.parse(str(item["text"]), filename=str(item["path"]))
|
||||
except SyntaxError:
|
||||
return []
|
||||
result: list[dict[str, Any]] = []
|
||||
|
||||
class Visitor(ast.NodeVisitor):
|
||||
def __init__(self) -> None:
|
||||
self.stack: list[str] = []
|
||||
|
||||
def _record(self, node: ast.AST, kind: str, name: str) -> None:
|
||||
qualified = ".".join([*self.stack, name]) if self.stack else name
|
||||
haystack = f"{item['path']} {qualified}".lower()
|
||||
score = sum(3 if term in haystack else 0 for term in terms)
|
||||
if name.startswith("test_"):
|
||||
score += 12
|
||||
result.append(
|
||||
{
|
||||
"path": item["path"],
|
||||
"symbol": qualified,
|
||||
"kind": kind,
|
||||
"line": int(getattr(node, "lineno", 1)),
|
||||
"end_line": int(getattr(node, "end_lineno", getattr(node, "lineno", 1))),
|
||||
"score": score + int(item["score"]),
|
||||
}
|
||||
)
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> Any:
|
||||
self._record(node, "class", node.name)
|
||||
self.stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self.stack.pop()
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
|
||||
self._record(node, "function", node.name)
|
||||
self.stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self.stack.pop()
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
|
||||
self._record(node, "async_function", node.name)
|
||||
self.stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self.stack.pop()
|
||||
|
||||
Visitor().visit(tree)
|
||||
return result
|
||||
|
||||
|
||||
def _surface_ref(symbol: Mapping[str, Any]) -> str:
|
||||
return f"{symbol['path']}::{symbol['symbol']}"
|
||||
|
||||
|
||||
def _contract_map(
|
||||
symbols: Sequence[Mapping[str, Any]],
|
||||
modules: Sequence[Mapping[str, Any]],
|
||||
contract_terms: Mapping[str, tuple[str, ...]],
|
||||
) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {}
|
||||
module_paths = [str(item["path"]) for item in modules]
|
||||
for contract, terms in contract_terms.items():
|
||||
scored: list[tuple[int, str]] = []
|
||||
for symbol in symbols:
|
||||
haystack = _surface_ref(symbol).lower()
|
||||
score = sum(1 for term in terms if term in haystack)
|
||||
if score:
|
||||
scored.append((score, _surface_ref(symbol)))
|
||||
scored.sort(key=lambda item: (-item[0], item[1]))
|
||||
surfaces = [value for _, value in scored[:6]]
|
||||
if not surfaces:
|
||||
for path in module_paths:
|
||||
score = sum(1 for term in terms if term in path.lower())
|
||||
if score:
|
||||
surfaces.append(path)
|
||||
if len(surfaces) == 6:
|
||||
break
|
||||
result[contract] = {
|
||||
"status": "PARTIAL" if surfaces else "GAP",
|
||||
"surfaces": surfaces,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def _gap_evidence(
|
||||
symbols: Sequence[Mapping[str, Any]],
|
||||
profile: str,
|
||||
) -> tuple[list[str], list[str], list[str]]:
|
||||
haystack = "\n".join(_surface_ref(item).lower() for item in symbols)
|
||||
|
||||
if profile == "model_governance":
|
||||
invariants = [
|
||||
"Online Task claim/heartbeat remains the authoritative worker ownership boundary.",
|
||||
"Provider/model credentials remain outside uncredentialed planner/executor processes.",
|
||||
"P0-A tenant/project scope and P0-C trace/audit lineage must constrain model routing.",
|
||||
"P1 retry/checkpoint semantics remain authoritative for model-call execution recovery.",
|
||||
]
|
||||
failure_windows: list[str] = []
|
||||
if "provider" not in haystack and "model" not in haystack:
|
||||
failure_windows.append("No first-class model/provider registry surface detected.")
|
||||
if "budget" not in haystack and "cost" not in haystack and "quota" not in haystack:
|
||||
failure_windows.append("No explicit durable model budget/cost/quota accounting surface detected.")
|
||||
if "fallback" not in haystack and "route" not in haystack:
|
||||
failure_windows.append("No explicit model fallback/routing policy surface detected.")
|
||||
if "eval" not in haystack and "approval" not in haystack:
|
||||
failure_windows.append("No explicit evaluation/approval gate surface detected for model changes.")
|
||||
if "version" not in haystack and "pin" not in haystack:
|
||||
failure_windows.append("No explicit model-version pinning surface detected.")
|
||||
desired = (
|
||||
("trace", "P0-C trace/audit lineage"),
|
||||
("scope", "P0-A tenant/project scope isolation"),
|
||||
("credential", "P0-C credential governance"),
|
||||
("checkpoint", "P1 durable execution recovery"),
|
||||
("message", "P0-B MessageEnvelope/handoff"),
|
||||
("retry", "P1 retry/idempotency"),
|
||||
)
|
||||
else:
|
||||
invariants = [
|
||||
"Online Task claim/heartbeat remains the authoritative worker ownership boundary.",
|
||||
"Governed mutation remains separated from read-only analysis and requires WriterLease.",
|
||||
"P1 attempt-aware result idempotency and durable checkpoint recovery are reusable fleet primitives.",
|
||||
"Scope/trace/message lineage must remain attached when work is scheduled or reassigned.",
|
||||
]
|
||||
failure_windows = []
|
||||
if "scheduler" not in haystack:
|
||||
failure_windows.append("No first-class durable fleet scheduler symbol detected; worker-local polling may still be the scheduling mechanism.")
|
||||
if "capacity" not in haystack and "slot" not in haystack:
|
||||
failure_windows.append("No explicit durable worker-capacity/slot primitive detected.")
|
||||
if "fair" not in haystack and "starvation" not in haystack:
|
||||
failure_windows.append("No explicit fairness/starvation policy symbol detected.")
|
||||
if "circuit" not in haystack:
|
||||
failure_windows.append("No explicit circuit-breaker primitive detected for repeated worker/task failures.")
|
||||
if "reassign" not in haystack:
|
||||
failure_windows.append("No explicit worker-reassignment primitive detected beyond lease expiry/retry semantics.")
|
||||
desired = (
|
||||
("checkpoint", "P1 durable checkpoint/resume"),
|
||||
("message", "P0-B MessageEnvelope/handoff"),
|
||||
("trace", "P0-C trace correlation"),
|
||||
("scope", "P0-A scope isolation"),
|
||||
("lease", "claim/WriterLease ownership"),
|
||||
("heartbeat", "worker liveness"),
|
||||
("idempot", "retry-safe dispatch/result"),
|
||||
)
|
||||
|
||||
reuse_points: list[str] = []
|
||||
for token, label in desired:
|
||||
match = next((item for item in symbols if token in _surface_ref(item).lower()), None)
|
||||
if match is not None:
|
||||
reuse_points.append(f"{label}: {_surface_ref(match)}")
|
||||
return invariants, failure_windows, reuse_points
|
||||
|
||||
|
||||
def _trim_to_budget(value: dict[str, Any]) -> dict[str, Any]:
|
||||
def size() -> int:
|
||||
return len(json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8"))
|
||||
|
||||
while size() > TARGET_RESULT_BYTES and value["bounded_excerpts"]:
|
||||
value["bounded_excerpts"].pop()
|
||||
while size() > TARGET_RESULT_BYTES and len(value["implementation_map"]["functions"]) > 20:
|
||||
value["implementation_map"]["functions"].pop()
|
||||
while size() > TARGET_RESULT_BYTES and len(value["implementation_map"]["classes"]) > 15:
|
||||
value["implementation_map"]["classes"].pop()
|
||||
while size() > TARGET_RESULT_BYTES and len(value["implementation_map"]["tests"]) > 15:
|
||||
value["implementation_map"]["tests"].pop()
|
||||
while size() > TARGET_RESULT_BYTES and len(value["implementation_map"]["modules"]) > 20:
|
||||
value["implementation_map"]["modules"].pop()
|
||||
final_size = size()
|
||||
if final_size > MAX_RESULT_BYTES:
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_RESULT_TOO_LARGE")
|
||||
value["result_bytes"] = final_size
|
||||
return value
|
||||
|
||||
|
||||
def build_compact_interface_map(
|
||||
worktree: Path,
|
||||
source_signals: Sequence[str],
|
||||
*,
|
||||
source_commit: str,
|
||||
) -> dict[str, Any]:
|
||||
root = Path(worktree)
|
||||
if not root.is_dir() or root.is_symlink():
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_WORKTREE_INVALID")
|
||||
if re.fullmatch(r"[0-9a-f]{40}", source_commit) is None:
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_SOURCE_COMMIT_INVALID")
|
||||
|
||||
profile = _profile(source_signals)
|
||||
terms = _signal_terms(source_signals)
|
||||
candidates = _candidate_files(root, terms)
|
||||
if not candidates:
|
||||
raise CompactInterfaceMapError("INTERFACE_MAP_EMPTY")
|
||||
|
||||
module_candidates = candidates[:MAX_MODULES]
|
||||
modules = [
|
||||
{
|
||||
"path": item["path"],
|
||||
"sha256": item["sha256"],
|
||||
"bytes": item["bytes"],
|
||||
"score": item["score"],
|
||||
}
|
||||
for item in module_candidates
|
||||
]
|
||||
|
||||
symbols: list[dict[str, Any]] = []
|
||||
tables: list[dict[str, Any]] = []
|
||||
systemd_units: list[dict[str, Any]] = []
|
||||
for item in candidates:
|
||||
if item["suffix"] == ".py":
|
||||
symbols.extend(_python_symbols(item, terms))
|
||||
elif item["suffix"] == ".sql":
|
||||
for table in SQL_TABLE_RE.findall(str(item["text"])):
|
||||
tables.append({"path": item["path"], "table": table, "sha256": item["sha256"]})
|
||||
elif item["suffix"] == ".service":
|
||||
systemd_units.append({"path": item["path"], "unit": Path(str(item["path"])).name, "sha256": item["sha256"]})
|
||||
|
||||
symbols.sort(key=lambda item: (-int(item["score"]), str(item["path"]), str(item["symbol"])))
|
||||
all_symbols = list(symbols)
|
||||
test_symbols = [
|
||||
item for item in all_symbols
|
||||
if str(item["path"]).startswith("tests/") or str(item["symbol"]).split(".")[-1].startswith("test_")
|
||||
][: min(MAX_TESTS, 24)]
|
||||
non_test_symbols = [
|
||||
item for item in all_symbols
|
||||
if item not in test_symbols
|
||||
][: max(0, MAX_SYMBOLS - len(test_symbols))]
|
||||
symbols = [*test_symbols, *non_test_symbols]
|
||||
symbols.sort(key=lambda item: (-int(item["score"]), str(item["path"]), str(item["symbol"])))
|
||||
classes = [
|
||||
{k: item[k] for k in ("path", "symbol", "line", "end_line")}
|
||||
for item in symbols
|
||||
if item["kind"] == "class"
|
||||
]
|
||||
functions = [
|
||||
{k: item[k] for k in ("path", "symbol", "line", "end_line")}
|
||||
for item in symbols
|
||||
if item["kind"] in {"function", "async_function"} and not str(item["symbol"]).split(".")[-1].startswith("test_")
|
||||
]
|
||||
tests = [
|
||||
{k: item[k] for k in ("path", "symbol", "line", "end_line")}
|
||||
for item in symbols
|
||||
if str(item["path"]).startswith("tests/") or str(item["symbol"]).split(".")[-1].startswith("test_")
|
||||
][:MAX_TESTS]
|
||||
|
||||
excerpts: list[dict[str, Any]] = []
|
||||
by_path = {str(item["path"]): item for item in candidates}
|
||||
for symbol in symbols:
|
||||
if len(excerpts) == MAX_EXCERPTS:
|
||||
break
|
||||
item = by_path.get(str(symbol["path"]))
|
||||
if item is None:
|
||||
continue
|
||||
excerpt = _bounded_excerpt(str(item["text"]), int(symbol["line"]), int(symbol["end_line"]))
|
||||
if not excerpt.strip():
|
||||
continue
|
||||
excerpts.append(
|
||||
{
|
||||
"path": symbol["path"],
|
||||
"symbol": symbol["symbol"],
|
||||
"sha256": item["sha256"],
|
||||
"excerpt": excerpt,
|
||||
}
|
||||
)
|
||||
|
||||
invariants, failure_windows, reuse_points = _gap_evidence(all_symbols, profile)
|
||||
result = {
|
||||
"schema": SCHEMA,
|
||||
"profile": profile,
|
||||
"classification": PROFILE_CLASSIFICATION[profile],
|
||||
"source_commit": source_commit,
|
||||
"implementation_map": {
|
||||
"modules": modules,
|
||||
"classes": classes[:MAX_SYMBOLS],
|
||||
"functions": functions[:MAX_SYMBOLS],
|
||||
"tables": tables[:40],
|
||||
"systemd_units": systemd_units[:20],
|
||||
"tests": tests,
|
||||
"current_invariants": invariants,
|
||||
"failure_windows": failure_windows,
|
||||
"reuse_points": reuse_points,
|
||||
},
|
||||
"bounded_excerpts": excerpts,
|
||||
"contracts": _contract_map(all_symbols, modules, _contract_terms(profile)),
|
||||
}
|
||||
return _trim_to_budget(result)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,456 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import datetime as dt
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Mapping
|
||||
|
||||
UTC = dt.timezone.utc
|
||||
_ALLOWED_CONFIDENTIALITY = frozenset({"public", "tenant", "project", "private", "run"})
|
||||
_PRIVILEGED = frozenset({"*", "admin", "system.admin", "message.admin", "task.admin"})
|
||||
|
||||
|
||||
class HandoffError(RuntimeError):
|
||||
code = "HANDOFF_ERROR"
|
||||
|
||||
|
||||
class MessageScopeMismatch(HandoffError):
|
||||
code = "MESSAGE_SCOPE_MISMATCH_DENY"
|
||||
|
||||
|
||||
class ConfidentialityScopeDenied(HandoffError):
|
||||
code = "MESSAGE_CONFIDENTIALITY_SCOPE_ENFORCE"
|
||||
|
||||
|
||||
class MessageExpired(HandoffError):
|
||||
code = "MESSAGE_EXPIRED"
|
||||
|
||||
|
||||
class SequenceGap(HandoffError):
|
||||
code = "MESSAGE_SEQUENCE_GAP_DETECT"
|
||||
|
||||
|
||||
class SequenceConflict(HandoffError):
|
||||
code = "MESSAGE_SEQUENCE_CONFLICT"
|
||||
|
||||
|
||||
class IdempotencyConflict(HandoffError):
|
||||
code = "MESSAGE_IDEMPOTENCY_CONFLICT"
|
||||
|
||||
|
||||
def _utcnow() -> dt.datetime:
|
||||
return dt.datetime.now(UTC)
|
||||
|
||||
|
||||
def _iso(value: dt.datetime | str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00")) if isinstance(value, str) else value
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=UTC)
|
||||
return parsed.astimezone(UTC).isoformat()
|
||||
|
||||
|
||||
def _canonical(value: Any) -> str:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
|
||||
|
||||
|
||||
def _sha(value: Any) -> str:
|
||||
return hashlib.sha256(_canonical(value).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _clean(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _tuple_scope(value: Iterable[str] | str | None) -> tuple[str, ...]:
|
||||
if value is None:
|
||||
return ()
|
||||
if isinstance(value, str):
|
||||
value = value.split(",")
|
||||
return tuple(sorted({str(item).strip() for item in value if str(item).strip()}))
|
||||
|
||||
|
||||
_SECRET_KEY = re.compile(r"(?i)(?:token|secret|password|private[_-]?key|api[_-]?key)")
|
||||
_SECRET_VALUE = re.compile(r"(?i)bearer\s+[A-Za-z0-9._~+/-]{8,}|(?:token|secret|password|api[_-]?key)\s*[:=]\s*\S+")
|
||||
_SAFE_REFERENCE = re.compile(r"^(?:systemd-credential|vault|secret-manager):[A-Za-z0-9][A-Za-z0-9._:/#-]{0,255}$")
|
||||
|
||||
|
||||
def assert_payload_secret_safe(value: Any, *, key: str = "") -> None:
|
||||
if isinstance(value, Mapping):
|
||||
for child_key, child in value.items():
|
||||
assert_payload_secret_safe(child, key=str(child_key))
|
||||
return
|
||||
if isinstance(value, (list, tuple)):
|
||||
for child in value:
|
||||
assert_payload_secret_safe(child, key=key)
|
||||
return
|
||||
if isinstance(value, (bytes, bytearray)):
|
||||
raise ValueError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
if isinstance(value, str):
|
||||
if _SECRET_VALUE.search(value):
|
||||
raise ValueError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
if _SECRET_KEY.search(key) and value and _SAFE_REFERENCE.fullmatch(value) is None:
|
||||
raise ValueError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class DeliveryContext:
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
run_id: str | None = None
|
||||
capabilities: tuple[str, ...] = ()
|
||||
|
||||
@classmethod
|
||||
def from_actor(cls, actor: Any) -> "DeliveryContext":
|
||||
def get(name: str) -> Any:
|
||||
if isinstance(actor, Mapping):
|
||||
return actor.get(name)
|
||||
return getattr(actor, name, None)
|
||||
return cls(
|
||||
tenant_id=_clean(get("tenant_id")),
|
||||
user_id=_clean(get("user_id")),
|
||||
project_id=_clean(get("project_id")),
|
||||
conversation_id=_clean(get("conversation_id")),
|
||||
run_id=_clean(get("run_id")),
|
||||
capabilities=_tuple_scope(get("capabilities") or get("capability_scope")),
|
||||
)
|
||||
|
||||
@property
|
||||
def privileged(self) -> bool:
|
||||
return bool(set(self.capabilities) & _PRIVILEGED)
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class MessageEnvelope:
|
||||
source: str
|
||||
target: str
|
||||
message_type: str
|
||||
payload: Mapping[str, Any]
|
||||
sequence: int
|
||||
idempotency_key: str
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
task_id: str | None = None
|
||||
run_id: str | None = None
|
||||
correlation_id: str | None = None
|
||||
causation_id: str | None = None
|
||||
capability_scope: tuple[str, ...] = ()
|
||||
confidentiality_scope: str = "private"
|
||||
expires_at: str | None = None
|
||||
created_at: str = dataclasses.field(default_factory=lambda: _iso(_utcnow()) or "")
|
||||
envelope_id: str = ""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
assert_payload_secret_safe(self.payload)
|
||||
for field in ("source", "target", "message_type", "idempotency_key", "tenant_id", "user_id",
|
||||
"project_id", "conversation_id", "task_id", "run_id", "correlation_id", "causation_id"):
|
||||
object.__setattr__(self, field, _clean(getattr(self, field)) if field not in {
|
||||
"source", "target", "message_type", "idempotency_key"
|
||||
} else (_clean(getattr(self, field)) or ""))
|
||||
object.__setattr__(self, "capability_scope", _tuple_scope(self.capability_scope))
|
||||
object.__setattr__(self, "confidentiality_scope", (_clean(self.confidentiality_scope) or "private").lower())
|
||||
object.__setattr__(self, "expires_at", _iso(self.expires_at))
|
||||
object.__setattr__(self, "created_at", _iso(self.created_at) or "")
|
||||
if not self.source or not self.target or not self.message_type or not self.idempotency_key:
|
||||
raise ValueError("MESSAGE_ENVELOPE_IDENTITY_REQUIRED")
|
||||
if not isinstance(self.sequence, int) or self.sequence < 1:
|
||||
raise ValueError("MESSAGE_SEQUENCE_MUST_BE_POSITIVE")
|
||||
if self.confidentiality_scope not in _ALLOWED_CONFIDENTIALITY:
|
||||
raise ValueError("MESSAGE_CONFIDENTIALITY_SCOPE_INVALID")
|
||||
if self.confidentiality_scope != "public" and self.tenant_id is None:
|
||||
raise ValueError("MESSAGE_TENANT_REQUIRED")
|
||||
if self.confidentiality_scope in {"project", "private", "run"} and self.project_id is None:
|
||||
raise ValueError("MESSAGE_PROJECT_REQUIRED")
|
||||
if self.confidentiality_scope in {"private", "run"} and self.user_id is None:
|
||||
raise ValueError("MESSAGE_USER_REQUIRED")
|
||||
if self.confidentiality_scope == "run" and self.run_id is None:
|
||||
raise ValueError("MESSAGE_RUN_REQUIRED")
|
||||
expected = stable_event_id(
|
||||
source=self.source, target=self.target, message_type=self.message_type, payload=self.payload,
|
||||
tenant_id=self.tenant_id, user_id=self.user_id, project_id=self.project_id,
|
||||
conversation_id=self.conversation_id, task_id=self.task_id, run_id=self.run_id,
|
||||
correlation_id=self.correlation_id, causation_id=self.causation_id,
|
||||
sequence=self.sequence, idempotency_key=self.idempotency_key,
|
||||
)
|
||||
if self.envelope_id and self.envelope_id != expected:
|
||||
raise ValueError("MESSAGE_ENVELOPE_ID_MISMATCH")
|
||||
object.__setattr__(self, "envelope_id", expected)
|
||||
|
||||
@property
|
||||
def payload_sha256(self) -> str:
|
||||
return _sha(self.payload)
|
||||
|
||||
@property
|
||||
def stream_key(self) -> str:
|
||||
return _sha({
|
||||
"tenant_id": self.tenant_id, "source": self.source, "target": self.target,
|
||||
"conversation_id": self.conversation_id, "task_id": self.task_id, "run_id": self.run_id,
|
||||
})
|
||||
|
||||
def is_expired(self, *, now: dt.datetime | None = None) -> bool:
|
||||
if self.expires_at is None:
|
||||
return False
|
||||
return dt.datetime.fromisoformat(self.expires_at) <= (now or _utcnow()).astimezone(UTC)
|
||||
|
||||
def as_record(self) -> dict[str, Any]:
|
||||
return {
|
||||
"envelope_id": self.envelope_id, "stream_key": self.stream_key, "source": self.source,
|
||||
"target": self.target, "message_type": self.message_type, "payload_json": _canonical(self.payload),
|
||||
"payload_sha256": self.payload_sha256, "tenant_id": self.tenant_id, "user_id": self.user_id,
|
||||
"project_id": self.project_id, "conversation_id": self.conversation_id, "task_id": self.task_id,
|
||||
"run_id": self.run_id, "correlation_id": self.correlation_id,
|
||||
"causation_id": self.causation_id, "sequence": self.sequence,
|
||||
"capability_scope_json": _canonical(self.capability_scope),
|
||||
"confidentiality_scope": self.confidentiality_scope, "idempotency_key": self.idempotency_key,
|
||||
"expires_at": self.expires_at, "created_at": self.created_at,
|
||||
}
|
||||
|
||||
|
||||
def stable_event_id(*, source: Any, target: Any, message_type: Any, payload: Any,
|
||||
tenant_id: Any = None, user_id: Any = None, project_id: Any = None,
|
||||
conversation_id: Any = None, task_id: Any = None, run_id: Any = None,
|
||||
correlation_id: Any = None, causation_id: Any = None,
|
||||
sequence: Any = None, idempotency_key: Any = None) -> str:
|
||||
assert_payload_secret_safe(payload)
|
||||
material = {
|
||||
"source": _clean(source), "target": _clean(target), "message_type": _clean(message_type),
|
||||
"payload_sha256": _sha(payload), "tenant_id": _clean(tenant_id), "user_id": _clean(user_id),
|
||||
"project_id": _clean(project_id), "conversation_id": _clean(conversation_id),
|
||||
"task_id": _clean(task_id), "run_id": _clean(run_id), "causation_id": _clean(causation_id),
|
||||
"sequence": sequence, "idempotency_key": _clean(idempotency_key),
|
||||
}
|
||||
# Preserve historical envelope IDs for legacy rows that predate correlation_id.
|
||||
if _clean(correlation_id) is not None:
|
||||
material["correlation_id"] = _clean(correlation_id)
|
||||
return _sha(material)
|
||||
|
||||
|
||||
def bind_envelope_scope(envelope: MessageEnvelope, actor: Any) -> MessageEnvelope:
|
||||
ctx = DeliveryContext.from_actor(actor)
|
||||
if ctx.privileged:
|
||||
return envelope
|
||||
if ctx.tenant_id is None:
|
||||
if envelope.tenant_id is not None:
|
||||
raise MessageScopeMismatch(MessageScopeMismatch.code)
|
||||
return envelope
|
||||
updates: dict[str, Any] = {}
|
||||
for field in ("tenant_id", "user_id", "project_id", "conversation_id", "run_id"):
|
||||
trusted = getattr(ctx, field)
|
||||
requested = getattr(envelope, field)
|
||||
if requested is not None and requested != trusted:
|
||||
raise MessageScopeMismatch(f"{MessageScopeMismatch.code}:{field}")
|
||||
if requested is None and trusted is not None:
|
||||
updates[field] = trusted
|
||||
if not updates:
|
||||
return envelope
|
||||
data = dataclasses.asdict(envelope)
|
||||
data.update(updates)
|
||||
data["envelope_id"] = ""
|
||||
return MessageEnvelope(**data)
|
||||
|
||||
|
||||
def authorize_delivery(actor: Any, envelope: MessageEnvelope) -> None:
|
||||
ctx = DeliveryContext.from_actor(actor)
|
||||
if ctx.privileged:
|
||||
return
|
||||
if not set(envelope.capability_scope).issubset(set(ctx.capabilities)):
|
||||
raise ConfidentialityScopeDenied(f"{ConfidentialityScopeDenied.code}:capability")
|
||||
if envelope.confidentiality_scope == "public":
|
||||
return
|
||||
if ctx.tenant_id != envelope.tenant_id:
|
||||
raise MessageScopeMismatch(f"{MessageScopeMismatch.code}:tenant_id")
|
||||
if envelope.confidentiality_scope in {"project", "private", "run"} and ctx.project_id != envelope.project_id:
|
||||
raise ConfidentialityScopeDenied(f"{ConfidentialityScopeDenied.code}:project_id")
|
||||
if envelope.confidentiality_scope in {"private", "run"} and ctx.user_id != envelope.user_id:
|
||||
raise ConfidentialityScopeDenied(f"{ConfidentialityScopeDenied.code}:user_id")
|
||||
if envelope.confidentiality_scope == "run" and ctx.run_id != envelope.run_id:
|
||||
raise ConfidentialityScopeDenied(f"{ConfidentialityScopeDenied.code}:run_id")
|
||||
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS message_handoff (
|
||||
envelope_id TEXT PRIMARY KEY,
|
||||
stream_key TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
target TEXT NOT NULL,
|
||||
message_type TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL,
|
||||
payload_sha256 TEXT NOT NULL,
|
||||
tenant_id TEXT,
|
||||
user_id TEXT,
|
||||
project_id TEXT,
|
||||
conversation_id TEXT,
|
||||
task_id TEXT,
|
||||
run_id TEXT,
|
||||
correlation_id TEXT,
|
||||
causation_id TEXT,
|
||||
sequence INTEGER NOT NULL CHECK(sequence >= 1),
|
||||
capability_scope_json TEXT NOT NULL,
|
||||
confidentiality_scope TEXT NOT NULL,
|
||||
idempotency_key TEXT NOT NULL,
|
||||
expires_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
delivered_at TEXT,
|
||||
acked_at TEXT,
|
||||
UNIQUE(stream_key, idempotency_key),
|
||||
UNIQUE(stream_key, sequence)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_message_handoff_target_status
|
||||
ON message_handoff(target, status, created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_message_handoff_scope
|
||||
ON message_handoff(tenant_id, project_id, user_id, conversation_id, run_id);
|
||||
"""
|
||||
|
||||
|
||||
class DurableHandoffStore:
|
||||
def __init__(self, path: str | Path) -> None:
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._connect() as conn:
|
||||
conn.executescript(_SCHEMA)
|
||||
columns = {str(row[1]) for row in conn.execute("PRAGMA table_info(message_handoff)")}
|
||||
if "correlation_id" not in columns:
|
||||
conn.execute("ALTER TABLE message_handoff ADD COLUMN correlation_id TEXT")
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_message_handoff_lineage "
|
||||
"ON message_handoff(correlation_id, causation_id)"
|
||||
)
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(self.path, timeout=30.0)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
conn.execute("PRAGMA busy_timeout=30000")
|
||||
return conn
|
||||
|
||||
@staticmethod
|
||||
def _from_row(row: sqlite3.Row) -> MessageEnvelope:
|
||||
return MessageEnvelope(
|
||||
source=row["source"], target=row["target"], message_type=row["message_type"],
|
||||
payload=json.loads(row["payload_json"]), tenant_id=row["tenant_id"], user_id=row["user_id"],
|
||||
project_id=row["project_id"], conversation_id=row["conversation_id"], task_id=row["task_id"],
|
||||
run_id=row["run_id"], correlation_id=row["correlation_id"],
|
||||
causation_id=row["causation_id"], sequence=int(row["sequence"]),
|
||||
capability_scope=tuple(json.loads(row["capability_scope_json"])),
|
||||
confidentiality_scope=row["confidentiality_scope"], idempotency_key=row["idempotency_key"],
|
||||
expires_at=row["expires_at"], created_at=row["created_at"], envelope_id=row["envelope_id"],
|
||||
)
|
||||
|
||||
def put(self, envelope: MessageEnvelope) -> tuple[str, MessageEnvelope]:
|
||||
if envelope.is_expired():
|
||||
raise MessageExpired(MessageExpired.code)
|
||||
record = envelope.as_record()
|
||||
with self._connect() as conn:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
replay = conn.execute(
|
||||
"SELECT * FROM message_handoff WHERE stream_key=? AND idempotency_key=?",
|
||||
(envelope.stream_key, envelope.idempotency_key),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
existing = self._from_row(replay)
|
||||
if existing.envelope_id != envelope.envelope_id or existing.payload_sha256 != envelope.payload_sha256:
|
||||
raise IdempotencyConflict(IdempotencyConflict.code)
|
||||
conn.commit()
|
||||
return "REPLAY", existing
|
||||
seq = conn.execute(
|
||||
"SELECT * FROM message_handoff WHERE stream_key=? AND sequence=?",
|
||||
(envelope.stream_key, envelope.sequence),
|
||||
).fetchone()
|
||||
if seq is not None:
|
||||
existing = self._from_row(seq)
|
||||
if existing.envelope_id == envelope.envelope_id:
|
||||
conn.commit()
|
||||
return "REPLAY", existing
|
||||
raise SequenceConflict(SequenceConflict.code)
|
||||
maximum = int(conn.execute(
|
||||
"SELECT COALESCE(MAX(sequence),0) FROM message_handoff WHERE stream_key=?",
|
||||
(envelope.stream_key,),
|
||||
).fetchone()[0])
|
||||
expected = maximum + 1
|
||||
if envelope.sequence != expected:
|
||||
raise SequenceGap(f"{SequenceGap.code}:expected={expected}:actual={envelope.sequence}")
|
||||
columns = tuple(record)
|
||||
placeholders = ",".join("?" for _ in columns)
|
||||
conn.execute(
|
||||
f"INSERT INTO message_handoff ({','.join(columns)}) VALUES ({placeholders})",
|
||||
tuple(record[column] for column in columns),
|
||||
)
|
||||
conn.commit()
|
||||
return "CREATED", envelope
|
||||
|
||||
def get(self, envelope_id: str, *, actor: Any) -> MessageEnvelope:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute("SELECT * FROM message_handoff WHERE envelope_id=?", (envelope_id,)).fetchone()
|
||||
if row is None:
|
||||
raise KeyError(envelope_id)
|
||||
envelope = self._from_row(row)
|
||||
if envelope.is_expired():
|
||||
conn.execute("UPDATE message_handoff SET status='EXPIRED' WHERE envelope_id=?", (envelope_id,))
|
||||
conn.commit()
|
||||
raise MessageExpired(MessageExpired.code)
|
||||
authorize_delivery(actor, envelope)
|
||||
return envelope
|
||||
|
||||
def pending(self, *, target: str, actor: Any, limit: int = 100) -> list[MessageEnvelope]:
|
||||
if not 1 <= limit <= 1000:
|
||||
raise ValueError("MESSAGE_PENDING_LIMIT_INVALID")
|
||||
visible: list[MessageEnvelope] = []
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM message_handoff WHERE target=? AND status='PENDING' ORDER BY created_at, sequence LIMIT ?",
|
||||
(target, limit),
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
envelope = self._from_row(row)
|
||||
if envelope.is_expired():
|
||||
conn.execute("UPDATE message_handoff SET status='EXPIRED' WHERE envelope_id=?", (envelope.envelope_id,))
|
||||
continue
|
||||
try:
|
||||
authorize_delivery(actor, envelope)
|
||||
except HandoffError:
|
||||
continue
|
||||
visible.append(envelope)
|
||||
conn.commit()
|
||||
return visible
|
||||
|
||||
def mark_delivered(self, envelope_id: str, *, actor: Any) -> MessageEnvelope:
|
||||
envelope = self.get(envelope_id, actor=actor)
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE message_handoff SET status='DELIVERED', delivered_at=? "
|
||||
"WHERE envelope_id=? AND status='PENDING'",
|
||||
(_iso(_utcnow()), envelope_id),
|
||||
)
|
||||
conn.commit()
|
||||
return envelope
|
||||
|
||||
def acknowledge(self, envelope_id: str, *, actor: Any) -> MessageEnvelope:
|
||||
envelope = self.get(envelope_id, actor=actor)
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE message_handoff SET status='ACKED', acked_at=? WHERE envelope_id=?",
|
||||
(_iso(_utcnow()), envelope_id),
|
||||
)
|
||||
conn.commit()
|
||||
return envelope
|
||||
|
||||
|
||||
class HandoffFabric:
|
||||
def __init__(self, store: DurableHandoffStore) -> None:
|
||||
self.store = store
|
||||
|
||||
def publish(self, envelope: MessageEnvelope, *, actor: Any | None = None) -> tuple[str, MessageEnvelope]:
|
||||
return self.store.put(bind_envelope_scope(envelope, actor) if actor is not None else envelope)
|
||||
@@ -0,0 +1,68 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from message_handoff import MessageEnvelope
|
||||
|
||||
|
||||
def _get(value: Any, *names: str, default: Any = None) -> Any:
|
||||
for name in names:
|
||||
if isinstance(value, Mapping) and name in value:
|
||||
return value[name]
|
||||
if hasattr(value, name):
|
||||
return getattr(value, name)
|
||||
return default
|
||||
|
||||
|
||||
def agent_event_to_envelope(event: Any, *, target: str | None = None, sequence: int = 1,
|
||||
idempotency_key: str | None = None, **scope: Any) -> MessageEnvelope:
|
||||
payload = _get(event, "payload", "data", default={}) or {}
|
||||
target = target or _get(event, "target", "recipient")
|
||||
if not target:
|
||||
raise ValueError("MESSAGE_TARGET_REQUIRED")
|
||||
return MessageEnvelope(
|
||||
source=str(_get(event, "source", "producer", "agent_id", default="agent-event")),
|
||||
target=str(target),
|
||||
message_type=str(_get(event, "event_type", "type", "message_type", default="agent.event")),
|
||||
payload=dict(payload),
|
||||
sequence=sequence,
|
||||
idempotency_key=idempotency_key or str(_get(event, "id", "event_id", default="")) or f"agent-event:{sequence}",
|
||||
causation_id=_get(event, "causation_id"),
|
||||
**scope,
|
||||
)
|
||||
|
||||
|
||||
def taskplane_event_to_envelope(event: Any, *, target: str, sequence: int = 1,
|
||||
idempotency_key: str | None = None, **scope: Any) -> MessageEnvelope:
|
||||
data = _get(event, "data", "payload", default={}) or {}
|
||||
event_id = str(_get(event, "id", "event_id", default=""))
|
||||
return MessageEnvelope(
|
||||
source=str(_get(event, "source", default="task-plane")),
|
||||
target=target,
|
||||
message_type=str(_get(event, "type", "event_type", default="task.event")),
|
||||
payload=dict(data),
|
||||
sequence=sequence,
|
||||
idempotency_key=idempotency_key or event_id or f"task-event:{sequence}",
|
||||
causation_id=event_id or None,
|
||||
**scope,
|
||||
)
|
||||
|
||||
|
||||
def legacy_message_to_envelope(message: Any, *, target: str | None = None, sequence: int = 1,
|
||||
idempotency_key: str | None = None, **scope: Any) -> MessageEnvelope:
|
||||
payload = _get(message, "payload", "data", "body", "content", default={})
|
||||
if not isinstance(payload, Mapping):
|
||||
payload = {"value": payload}
|
||||
target = target or _get(message, "target", "recipient", "receiver")
|
||||
if not target:
|
||||
raise ValueError("MESSAGE_TARGET_REQUIRED")
|
||||
source = _get(message, "source", "sender", "agent_id", default="legacy-message-bus")
|
||||
kind = _get(message, "message_type", "type", "kind", default="legacy.message")
|
||||
identifier = _get(message, "id", "message_id", "event_id")
|
||||
return MessageEnvelope(
|
||||
source=str(source), target=str(target), message_type=str(kind), payload=dict(payload),
|
||||
sequence=sequence,
|
||||
idempotency_key=idempotency_key or (str(identifier) if identifier else f"legacy:{sequence}"),
|
||||
causation_id=str(identifier) if identifier else None,
|
||||
**scope,
|
||||
)
|
||||
@@ -11,6 +11,7 @@ from dataclasses import asdict, dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from .search_gateway_orchestrator import SearchGatewayOrchestrator
|
||||
from .trace_tool_credential_governance import GovernanceActor, TraceContext
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -43,6 +44,9 @@ class SearchGatewayAgentAdapter:
|
||||
domain_profiles: list[str] | None = None,
|
||||
user_scope: str = "public",
|
||||
correlation_id: str | None = None,
|
||||
*,
|
||||
actor: GovernanceActor | dict[str, Any] | None = None,
|
||||
trace: TraceContext | None = None,
|
||||
) -> dict[str, Any]:
|
||||
payload = {
|
||||
"query": query,
|
||||
@@ -59,7 +63,11 @@ class SearchGatewayAgentAdapter:
|
||||
"session_profile": context.session_profile,
|
||||
"correlation_id": correlation_id,
|
||||
}
|
||||
return self.orchestrator.execute(payload)
|
||||
if actor is None and trace is None:
|
||||
# Preserve adapter compatibility; the governed orchestrator itself
|
||||
# still fails closed when trusted context is absent.
|
||||
return self.orchestrator.execute(payload)
|
||||
return self.orchestrator.execute(payload, actor=actor, trace=trace)
|
||||
|
||||
|
||||
class OpenClawSearchClient(SearchGatewayAgentAdapter):
|
||||
|
||||
@@ -13,10 +13,18 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Any, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from .search_profile_resolver import DomainQueryProfileRegistry, SearchProfile, SearchProfileResolver
|
||||
from .trace_tool_credential_governance import (
|
||||
GovernanceActor,
|
||||
GovernanceError,
|
||||
GovernanceStore,
|
||||
TraceContext,
|
||||
)
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
@@ -156,13 +164,25 @@ class SearchPolicyGuard:
|
||||
|
||||
|
||||
class SearchGateway:
|
||||
def __init__(self, providers: list[SearchProvider] | None = None):
|
||||
def __init__(self, providers: list[SearchProvider] | None = None, *, governance_store: GovernanceStore | None = None):
|
||||
self.providers = providers or []
|
||||
self.governance_store = governance_store
|
||||
self.profile_resolver = SearchProfileResolver()
|
||||
self.planner = QueryPlanner()
|
||||
self.guard = SearchPolicyGuard()
|
||||
|
||||
def search(self, request: SearchRequest) -> SearchResponse:
|
||||
def search(self, request: SearchRequest, *, actor: GovernanceActor | dict[str, Any] | None = None, trace: TraceContext | None = None) -> SearchResponse:
|
||||
trusted_actor: GovernanceActor | None = None
|
||||
if self.governance_store is not None:
|
||||
if actor is None or trace is None:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
trusted_actor = GovernanceActor.from_value(actor)
|
||||
if request.requester != trusted_actor.identity:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY", field="requester")
|
||||
if request.correlation_id is not None and request.correlation_id not in {
|
||||
trace.trace_id, trace.correlation_id
|
||||
}:
|
||||
raise GovernanceError("TOOL_SCOPE_PROPAGATION", field="correlation_id")
|
||||
request_id = request.correlation_id or f"search_{uuid4().hex[:12]}"
|
||||
profile = self.profile_resolver.resolve(request.group_profile, request.user_profile, request.session_profile)
|
||||
expanded_queries = self.planner.expand(request, profile)
|
||||
@@ -184,14 +204,60 @@ class SearchGateway:
|
||||
attempts: list[SearchProviderAttempt] = []
|
||||
providers = self._ordered_providers(profile)
|
||||
for provider in providers:
|
||||
if request.authorized_sources and provider.name not in request.authorized_sources:
|
||||
attempts.append(SearchProviderAttempt(provider=provider.name, status="denied", reason="SOURCE_AUTHZ_DENY"))
|
||||
continue
|
||||
if request.user_scope != "public" and provider.name not in request.authorized_sources:
|
||||
attempts.append(SearchProviderAttempt(provider=provider.name, status="denied", reason="SOURCE_AUTHZ_REQUIRED"))
|
||||
continue
|
||||
try:
|
||||
if self.governance_store is not None:
|
||||
assert trusted_actor is not None and trace is not None
|
||||
input_digest = hashlib.sha256(
|
||||
json.dumps(
|
||||
{"query": request.query, "intent": request.intent, "provider": provider.name},
|
||||
sort_keys=True, separators=(",", ":"), ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
).hexdigest()
|
||||
invocation_key = hashlib.sha256(
|
||||
f"{trace.trace_id}:{trace.span_id}:{provider.name}:{input_digest}".encode("utf-8")
|
||||
).hexdigest()
|
||||
authorization = self.governance_store.authorize_tool(
|
||||
actor=trusted_actor,
|
||||
tool_name=f"search.{provider.name}",
|
||||
trace=trace,
|
||||
requested_scope=trace.scope,
|
||||
input_digest=input_digest,
|
||||
invocation_key=invocation_key,
|
||||
)
|
||||
if authorization.get("idempotent"):
|
||||
raise GovernanceError("TOOL_INVOCATION_REPLAY")
|
||||
provider_evidence, attempt = provider.search(request, expanded_queries)
|
||||
if self.governance_store is not None:
|
||||
result_digest = hashlib.sha256(
|
||||
json.dumps(
|
||||
{
|
||||
"evidence": [item.to_dict() for item in provider_evidence],
|
||||
"attempt": attempt.to_dict(),
|
||||
},
|
||||
sort_keys=True, separators=(",", ":"), ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
).hexdigest()
|
||||
self.governance_store.record_tool_receipt(
|
||||
invocation_key=invocation_key,
|
||||
result_digest=result_digest,
|
||||
actor=trusted_actor,
|
||||
trace=trace,
|
||||
)
|
||||
evidence.extend(provider_evidence)
|
||||
attempts.append(attempt)
|
||||
if provider_evidence:
|
||||
break
|
||||
except Exception as exc:
|
||||
attempts.append(SearchProviderAttempt(provider=getattr(provider, "name", "unknown"), status="failed", reason=str(exc)))
|
||||
except GovernanceError:
|
||||
raise
|
||||
except Exception:
|
||||
# Provider exceptions may contain URLs, tokens, or local paths.
|
||||
attempts.append(SearchProviderAttempt(provider=getattr(provider, "name", "unknown"), status="failed", reason="PROVIDER_FAILURE"))
|
||||
|
||||
if evidence:
|
||||
return SearchResponse(
|
||||
|
||||
@@ -7,27 +7,29 @@ try:
|
||||
except Exception:
|
||||
APIRouter = None
|
||||
|
||||
from .search_gateway import SearchGateway, SearchRequest
|
||||
from .search_gateway import SearchRequest
|
||||
from .search_gateway_orchestrator import SearchGatewayOrchestrator
|
||||
from .trace_tool_credential_governance import GovernanceActor, TraceContext
|
||||
from .search_gateway_dashboard import SearchGatewayDashboard
|
||||
|
||||
|
||||
gateway = SearchGateway()
|
||||
orchestrator = SearchGatewayOrchestrator()
|
||||
dashboard = SearchGatewayDashboard()
|
||||
|
||||
|
||||
def search_gateway_request(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
request = SearchRequest(**payload)
|
||||
return gateway.search(request).to_dict()
|
||||
def search_gateway_request(payload: dict[str, Any], *, actor: GovernanceActor | dict[str, Any] | None = None, trace: TraceContext | None = None) -> dict[str, Any]:
|
||||
return orchestrator.execute(payload, actor=actor, trace=trace)
|
||||
|
||||
|
||||
def get_search_gateway_status() -> dict[str, Any]:
|
||||
return {
|
||||
"module": "search-gateway",
|
||||
"status": "ready",
|
||||
"providers_configured": len(gateway.providers),
|
||||
"providers_configured": len(orchestrator.gateway.providers),
|
||||
"policy_guard": "enabled",
|
||||
"profile_customization": "enabled",
|
||||
"fallback_required_when_provider_unavailable": True,
|
||||
"governance": "required",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -3,24 +3,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
import os
|
||||
|
||||
from .search_audit_store import SearchAuditStore
|
||||
from .search_event_bridge import SearchEventBridge
|
||||
from .search_gateway import SearchGateway, SearchRequest
|
||||
from .search_provider_registry import SearchProviderRegistry
|
||||
from .trace_tool_credential_governance import (
|
||||
GovernanceActor,
|
||||
GovernanceScope,
|
||||
GovernanceStore,
|
||||
ToolDescriptor,
|
||||
TraceContext,
|
||||
)
|
||||
|
||||
|
||||
class SearchGatewayOrchestrator:
|
||||
def __init__(self):
|
||||
def __init__(self, *, governance_store: GovernanceStore | None = None):
|
||||
self.registry = SearchProviderRegistry()
|
||||
self.audit = SearchAuditStore()
|
||||
self.events = SearchEventBridge()
|
||||
self.gateway = SearchGateway(self.registry.ordered())
|
||||
self.governance = governance_store or GovernanceStore(
|
||||
os.environ.get("AIOS_SEARCH_GOVERNANCE_DB", "runtime/search-governance.sqlite3")
|
||||
)
|
||||
registrar = GovernanceActor(
|
||||
"search-runtime-bootstrap", GovernanceScope("system"), ("system.admin",)
|
||||
)
|
||||
for provider in self.registry.ordered():
|
||||
self.governance.register_tool(
|
||||
ToolDescriptor(
|
||||
f"search.{provider.name}", "tool", GovernanceScope("global"), ("tool.search",)
|
||||
),
|
||||
registrar,
|
||||
)
|
||||
self.gateway = SearchGateway(self.registry.ordered(), governance_store=self.governance)
|
||||
|
||||
def execute(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
def execute(self, payload: dict[str, Any], *, actor: GovernanceActor | dict[str, Any] | None = None, trace: TraceContext | None = None) -> dict[str, Any]:
|
||||
if actor is None or trace is None:
|
||||
from .trace_tool_credential_governance import GovernanceError
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
self.governance.record_trace(trace, actor)
|
||||
request = SearchRequest(**payload)
|
||||
self.events.search_started(request.to_dict())
|
||||
response = self.gateway.search(request)
|
||||
response = self.gateway.search(request, actor=actor, trace=trace)
|
||||
result = response.to_dict()
|
||||
for attempt in result.get("attempts", []):
|
||||
self.events.provider_attempted(result["request_id"], attempt)
|
||||
@@ -35,6 +60,7 @@ class SearchGatewayOrchestrator:
|
||||
"registry": self.registry.status(),
|
||||
"audit": "enabled",
|
||||
"events": "enabled",
|
||||
"governance": "required",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ from .models import (
|
||||
)
|
||||
from .policy import PolicyEngine
|
||||
from .store import TaskPlaneStore
|
||||
from .scope import bind_task_scope
|
||||
from .worker_registry import WorkerRegistry
|
||||
|
||||
|
||||
@@ -68,28 +69,54 @@ class TaskPlaneService:
|
||||
self.delegation = DelegationEngine(store)
|
||||
self.store.probe_attestation_verifier = self.policy.probe_attestation_verifier
|
||||
|
||||
def create_task(self, body: dict[str, Any], idempotency_key: str | None = None) -> dict[str, Any]:
|
||||
def _authorized_task(self, task_id: str, actor: str, action: str) -> TaskRecord:
|
||||
task = self.store.get_task(task_id, self.policy.actor_context(actor))
|
||||
self.policy.assert_allowed(action, actor, task=task)
|
||||
return task
|
||||
|
||||
def create_task(
|
||||
self,
|
||||
body: dict[str, Any],
|
||||
idempotency_key: str | None = None,
|
||||
*,
|
||||
actor_identity: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if body.get("parent_task_id"):
|
||||
raise ForbiddenError("child tasks must use the delegate endpoint")
|
||||
issuer = body.get("issuer", {})
|
||||
actor = issuer.get("id", "anon") if isinstance(issuer, dict) else "anon"
|
||||
declared_actor = issuer.get("id", "anon") if isinstance(issuer, dict) else "anon"
|
||||
actor = actor_identity or declared_actor
|
||||
self.policy.assert_allowed("create", actor)
|
||||
try:
|
||||
plan = TaskPlan.from_dict(body)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValidationError("invalid task plan") from exc
|
||||
if actor_identity is not None:
|
||||
plan = bind_task_scope(
|
||||
plan, self.policy.actor_context(actor_identity), public_create=True
|
||||
)
|
||||
task = self.store.create_task_with_event(plan, idempotency_key=idempotency_key)
|
||||
return task.to_dict()
|
||||
|
||||
def list_tasks(self, actor: str | None = None) -> list[dict[str, Any]]:
|
||||
context = None
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("read", actor)
|
||||
return [task.to_dict(include_task=False) for task in self.store.list_tasks()]
|
||||
context = self.policy.actor_context(actor)
|
||||
return [
|
||||
task.to_dict(include_task=False)
|
||||
for task in self.store.list_tasks(context)
|
||||
]
|
||||
|
||||
def get_task(self, task_id: str, actor: str | None = None) -> dict[str, Any]:
|
||||
context = None
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("read", actor, task=self.store.get_task(task_id))
|
||||
return self.store.get_task(task_id).to_dict()
|
||||
self.policy.assert_allowed("read", actor)
|
||||
context = self.policy.actor_context(actor)
|
||||
task = self.store.get_task(task_id, context)
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("read", actor, task=task)
|
||||
return task.to_dict()
|
||||
|
||||
def get_worker(self, worker_id: str, actor: str | None = None) -> dict[str, Any]:
|
||||
if actor is not None:
|
||||
@@ -163,9 +190,8 @@ class TaskPlaneService:
|
||||
if not worker_id:
|
||||
raise ValidationError("worker_id required")
|
||||
|
||||
task = self.store.get_task(task_id)
|
||||
task = self._authorized_task(task_id, actor, "claim")
|
||||
self.policy.assert_readonly_worker_binding(actor, str(worker_id))
|
||||
self.policy.assert_allowed("claim", actor, task=task)
|
||||
lease_seconds = _bounded_int(
|
||||
body.get("lease_seconds", 30), "lease_seconds", 1, 3600
|
||||
)
|
||||
@@ -184,10 +210,9 @@ class TaskPlaneService:
|
||||
if not claim_id:
|
||||
raise ValidationError("claim_id required")
|
||||
actor = body.get("actor_identity") or body.get("worker_identity") or "worker"
|
||||
task = self.store.get_task(task_id)
|
||||
task = self._authorized_task(task_id, actor, "heartbeat")
|
||||
claim = self.store.get_claim(claim_id)
|
||||
self.policy.assert_readonly_worker_binding(actor, claim.worker_id)
|
||||
self.policy.assert_allowed("heartbeat", actor, task=task)
|
||||
if claim.worker_identity != actor:
|
||||
raise ForbiddenError("claim identity mismatch")
|
||||
lease_seconds = _bounded_int(
|
||||
@@ -206,10 +231,9 @@ class TaskPlaneService:
|
||||
message = body.get("message")
|
||||
if progress is None:
|
||||
raise ValidationError("progress required")
|
||||
task = self.store.get_task(task_id)
|
||||
task = self._authorized_task(task_id, actor, "progress")
|
||||
claim = self.store.get_claim(claim_id)
|
||||
self.policy.assert_readonly_worker_binding(actor, claim.worker_id)
|
||||
self.policy.assert_allowed("progress", actor, task=task)
|
||||
if task.task_plan.data_classification == "credential-secret":
|
||||
if message not in {None, ""} or contains_secret_text(stage):
|
||||
raise ValidationError(
|
||||
@@ -234,7 +258,7 @@ class TaskPlaneService:
|
||||
if not claim_id:
|
||||
raise ValidationError("claim_id required")
|
||||
actor = body.get("actor_identity") or body.get("worker_identity") or "worker"
|
||||
task = self.store.get_task(task_id)
|
||||
task = self._authorized_task(task_id, actor, "result")
|
||||
|
||||
status = body.get("status", "COMPLETED")
|
||||
outputs = body.get("outputs", [])
|
||||
@@ -247,7 +271,6 @@ class TaskPlaneService:
|
||||
raise ValidationError("evidence items must be objects")
|
||||
claim = self.store.get_claim(claim_id)
|
||||
self.policy.assert_readonly_worker_binding(actor, claim.worker_id)
|
||||
self.policy.assert_allowed("result", actor, task=task)
|
||||
if claim.worker_identity != actor:
|
||||
raise ForbiddenError("claim identity mismatch")
|
||||
if task.task_plan.data_classification == "credential-secret":
|
||||
@@ -326,8 +349,7 @@ class TaskPlaneService:
|
||||
reason = body.get("reason", "failed")
|
||||
actor = body.get("actor_identity") or body.get("worker_identity") or "worker"
|
||||
self.policy.assert_allowed("fail", actor)
|
||||
task = self.store.get_task(task_id)
|
||||
self.policy.assert_allowed("fail", actor, task=task)
|
||||
task = self._authorized_task(task_id, actor, "fail")
|
||||
claim = self.store.get_claim(claim_id)
|
||||
if claim.worker_identity != actor:
|
||||
raise ForbiddenError("claim identity mismatch")
|
||||
@@ -339,8 +361,7 @@ class TaskPlaneService:
|
||||
actor = body.get("actor_identity") or body.get("worker_identity") or "owner"
|
||||
self.policy.assert_allowed("cancel", actor)
|
||||
reason = body.get("reason", "cancelled")
|
||||
task = self.store.get_task(task_id)
|
||||
self.policy.assert_allowed("cancel", actor, task=task)
|
||||
task = self._authorized_task(task_id, actor, "cancel")
|
||||
updated = self.store.cancel_task_with_event(
|
||||
task_id, task.revision, reason=reason
|
||||
)
|
||||
@@ -349,8 +370,7 @@ class TaskPlaneService:
|
||||
def delegate(self, task_id: str, body: dict[str, Any]) -> dict[str, Any]:
|
||||
actor = body.get("actor_identity") or "owner"
|
||||
self.policy.assert_allowed("delegate", actor)
|
||||
parent = self.store.get_task(task_id)
|
||||
self.policy.assert_allowed("delegate", actor, task=parent)
|
||||
parent = self._authorized_task(task_id, actor, "delegate")
|
||||
if "task" in body:
|
||||
try:
|
||||
child_plan = TaskPlan.from_dict(body["task"])
|
||||
@@ -364,25 +384,31 @@ class TaskPlaneService:
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValidationError("invalid child task plan") from exc
|
||||
child_plan.parent_task_id = task_id
|
||||
child_plan = bind_task_scope(
|
||||
child_plan, self.policy.actor_context(actor), public_create=False
|
||||
)
|
||||
child = self.delegation.delegate(task_id, child_plan, caller=actor)
|
||||
return child.to_dict()
|
||||
|
||||
def get_events(self, task_id: str, actor: str | None = None) -> list[dict[str, Any]]:
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("audit", actor, task=self.store.get_task(task_id))
|
||||
self.store.get_task(task_id)
|
||||
self._authorized_task(task_id, actor, "audit")
|
||||
else:
|
||||
self.store.get_task(task_id)
|
||||
return self.store.list_events(task_id)
|
||||
|
||||
def get_evidence(self, task_id: str, actor: str | None = None) -> list[dict[str, Any]]:
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("audit", actor, task=self.store.get_task(task_id))
|
||||
self.store.get_task(task_id)
|
||||
self._authorized_task(task_id, actor, "audit")
|
||||
else:
|
||||
self.store.get_task(task_id)
|
||||
return self.evidence.list(task_id)
|
||||
|
||||
def get_children(self, task_id: str, actor: str | None = None) -> list[str]:
|
||||
if actor is not None:
|
||||
self.policy.assert_allowed("read", actor, task=self.store.get_task(task_id))
|
||||
self.store.get_task(task_id)
|
||||
self._authorized_task(task_id, actor, "read")
|
||||
else:
|
||||
self.store.get_task(task_id)
|
||||
return self.store.list_children(task_id)
|
||||
|
||||
|
||||
@@ -506,7 +532,11 @@ class TaskPlaneRequestHandler(BaseHTTPRequestHandler):
|
||||
):
|
||||
raise ForbiddenError("issuer identity mismatch")
|
||||
body["issuer"] = {"type": "authenticated", "id": identity}
|
||||
return 201, self.service.create_task(body or {}, headers.get("idempotency-key"))
|
||||
return 201, self.service.create_task(
|
||||
body or {},
|
||||
headers.get("idempotency-key"),
|
||||
actor_identity=identity,
|
||||
)
|
||||
if body.get("actor_identity") not in {None, identity}:
|
||||
raise ForbiddenError("actor identity mismatch")
|
||||
body["actor_identity"] = identity
|
||||
|
||||
@@ -106,3 +106,8 @@ class DurableOutboxDispatcher:
|
||||
receipt=str(receipt) if receipt else None,
|
||||
error=None if receipt else f"{adapter.name} receipt unavailable",
|
||||
)
|
||||
|
||||
# R4.2 P0-B compatibility adapter.
|
||||
def to_handoff_envelope(event, *, target, sequence=1, idempotency_key=None, **scope):
|
||||
from message_handoff_adapters import taskplane_event_to_envelope
|
||||
return taskplane_event_to_envelope(event, target=target, sequence=sequence, idempotency_key=idempotency_key, **scope)
|
||||
|
||||
@@ -0,0 +1,911 @@
|
||||
"""Durable multi-agent fleet scheduling for TaskPlane.
|
||||
|
||||
The fleet layer decides *which* eligible worker should receive *which* queued task.
|
||||
It never replaces TaskPlane claim/lease ownership: callers must still perform the
|
||||
authoritative TaskPlane claim CAS after a reservation is returned.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Iterable, Mapping, Sequence
|
||||
|
||||
DEFAULT_RESERVATION_SECONDS = 30
|
||||
DEFAULT_CIRCUIT_THRESHOLD = 3
|
||||
DEFAULT_CIRCUIT_SECONDS = 30
|
||||
MAX_CIRCUIT_SECONDS = 15 * 60
|
||||
MAX_SLOTS = 128
|
||||
|
||||
|
||||
class FleetError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class FleetConflict(FleetError):
|
||||
pass
|
||||
|
||||
|
||||
def _canonical(value: object) -> bytes:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
|
||||
|
||||
def _sha(value: object) -> str:
|
||||
return hashlib.sha256(_canonical(value)).hexdigest()
|
||||
|
||||
|
||||
def _text(value: object, default: str = "") -> str:
|
||||
if value is None:
|
||||
return default
|
||||
return str(value).strip() or default
|
||||
|
||||
|
||||
def _int(value: object, default: int, minimum: int, maximum: int) -> int:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, bool):
|
||||
raise FleetError("FLEET_INTEGER_INVALID")
|
||||
try:
|
||||
result = int(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise FleetError("FLEET_INTEGER_INVALID") from exc
|
||||
if result < minimum or result > maximum:
|
||||
raise FleetError("FLEET_INTEGER_OUT_OF_RANGE")
|
||||
return result
|
||||
|
||||
|
||||
def _float(value: object, default: float, minimum: float, maximum: float) -> float:
|
||||
if value is None:
|
||||
return default
|
||||
try:
|
||||
result = float(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise FleetError("FLEET_NUMBER_INVALID") from exc
|
||||
if not math.isfinite(result) or result < minimum or result > maximum:
|
||||
raise FleetError("FLEET_NUMBER_OUT_OF_RANGE")
|
||||
return result
|
||||
|
||||
|
||||
def _mapping(value: object) -> Mapping[str, Any]:
|
||||
if isinstance(value, Mapping):
|
||||
return value
|
||||
if hasattr(value, "to_dict"):
|
||||
data = value.to_dict()
|
||||
if isinstance(data, Mapping):
|
||||
return data
|
||||
if hasattr(value, "__dict__"):
|
||||
return vars(value)
|
||||
return {}
|
||||
|
||||
|
||||
def _attr(value: object, name: str, default: Any = None) -> Any:
|
||||
if isinstance(value, Mapping):
|
||||
return value.get(name, default)
|
||||
return getattr(value, name, default)
|
||||
|
||||
|
||||
def _task_plan(task: object) -> object:
|
||||
return _attr(task, "task_plan", _attr(task, "plan", task))
|
||||
|
||||
|
||||
def _task_id(task: object) -> str:
|
||||
return _text(_attr(task, "task_id", _attr(_task_plan(task), "task_id", "")))
|
||||
|
||||
|
||||
def _task_status(task: object) -> str:
|
||||
return _text(_attr(task, "status", "")).upper()
|
||||
|
||||
|
||||
def _task_revision(task: object) -> int:
|
||||
return _int(_attr(task, "revision", 0), 0, 0, 2**31 - 1)
|
||||
|
||||
|
||||
def _task_created_at(task: object) -> float:
|
||||
value = _attr(task, "created_at", _attr(task, "updated_at", 0.0))
|
||||
try:
|
||||
return float(value or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _task_required_capabilities(task: object) -> frozenset[str]:
|
||||
plan = _task_plan(task)
|
||||
value = _attr(plan, "required_capabilities", ())
|
||||
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
||||
return frozenset()
|
||||
return frozenset(_text(item) for item in value if _text(item))
|
||||
|
||||
|
||||
def _task_metadata(task: object) -> Mapping[str, Any]:
|
||||
plan = _task_plan(task)
|
||||
value = _attr(plan, "metadata", {})
|
||||
return value if isinstance(value, Mapping) else {}
|
||||
|
||||
|
||||
def _task_scope_key(task: object) -> str:
|
||||
plan = _task_plan(task)
|
||||
scope_type = _text(_attr(plan, "scope_type", "legacy"), "legacy")
|
||||
tenant = _text(_attr(plan, "tenant_id", ""))
|
||||
project = _text(_attr(plan, "project_id", ""))
|
||||
user = _text(_attr(plan, "user_id", ""))
|
||||
if scope_type == "tenant":
|
||||
return "tenant:" + tenant + "|project:" + (project or "*")
|
||||
if scope_type == "system":
|
||||
return "system"
|
||||
return "legacy:" + (user or "*")
|
||||
|
||||
|
||||
def _worker_id(worker: object) -> str:
|
||||
return _text(_attr(worker, "worker_id", ""))
|
||||
|
||||
|
||||
def _worker_capabilities(worker: object) -> frozenset[str]:
|
||||
value = _attr(worker, "capabilities", ())
|
||||
if isinstance(value, Mapping):
|
||||
value = value.keys()
|
||||
if not isinstance(value, Iterable) or isinstance(value, (str, bytes)):
|
||||
return frozenset()
|
||||
return frozenset(_text(item) for item in value if _text(item))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FleetWorkerPolicy:
|
||||
worker_id: str
|
||||
max_slots: int = 1
|
||||
failure_domain: str = "default"
|
||||
weight: float = 1.0
|
||||
labels: Mapping[str, str] | None = None
|
||||
revision: int = 1
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FleetDispatchDecision:
|
||||
reservation_id: str
|
||||
task_id: str
|
||||
worker_id: str
|
||||
task_revision: int
|
||||
scope_key: str
|
||||
failure_domain: str
|
||||
reserved_until: float
|
||||
score: tuple[float, float, float, str]
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"reservation_id": self.reservation_id,
|
||||
"task_id": self.task_id,
|
||||
"worker_id": self.worker_id,
|
||||
"task_revision": self.task_revision,
|
||||
"scope_key": self.scope_key,
|
||||
"failure_domain": self.failure_domain,
|
||||
"reserved_until": self.reserved_until,
|
||||
"score": list(self.score),
|
||||
}
|
||||
|
||||
|
||||
class FleetStateStore:
|
||||
"""Small durable fleet journal colocated in the existing TaskPlane SQLite database."""
|
||||
|
||||
def __init__(self, connection: sqlite3.Connection, lock: threading.RLock | None = None) -> None:
|
||||
self.conn = connection
|
||||
self.lock = lock or threading.RLock()
|
||||
self._savepoint_counter = 0
|
||||
self._ensure_schema()
|
||||
|
||||
@contextmanager
|
||||
def _transaction(self):
|
||||
with self.lock:
|
||||
self._savepoint_counter += 1
|
||||
name = f"fleet_{threading.get_ident()}_{self._savepoint_counter}"
|
||||
self.conn.execute(f"SAVEPOINT {name}")
|
||||
try:
|
||||
yield
|
||||
except Exception:
|
||||
self.conn.execute(f"ROLLBACK TO {name}")
|
||||
self.conn.execute(f"RELEASE {name}")
|
||||
raise
|
||||
else:
|
||||
self.conn.execute(f"RELEASE {name}")
|
||||
|
||||
def _ensure_schema(self) -> None:
|
||||
with self.lock:
|
||||
self.conn.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS fleet_worker_capacity (
|
||||
worker_id TEXT PRIMARY KEY,
|
||||
max_slots INTEGER NOT NULL,
|
||||
failure_domain TEXT NOT NULL,
|
||||
weight REAL NOT NULL,
|
||||
labels_json TEXT NOT NULL,
|
||||
revision INTEGER NOT NULL,
|
||||
updated_at REAL NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS fleet_scope_state (
|
||||
scope_key TEXT PRIMARY KEY,
|
||||
virtual_finish REAL NOT NULL,
|
||||
dispatch_count INTEGER NOT NULL,
|
||||
last_dispatched_at REAL,
|
||||
revision INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS fleet_dispatches (
|
||||
reservation_id TEXT PRIMARY KEY,
|
||||
task_id TEXT NOT NULL UNIQUE,
|
||||
worker_id TEXT NOT NULL,
|
||||
task_revision INTEGER NOT NULL,
|
||||
scope_key TEXT NOT NULL,
|
||||
failure_domain TEXT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
reserved_until REAL NOT NULL,
|
||||
claim_id TEXT,
|
||||
created_at REAL NOT NULL,
|
||||
updated_at REAL NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS ix_fleet_dispatch_worker_state
|
||||
ON fleet_dispatches(worker_id, state, reserved_until);
|
||||
CREATE TABLE IF NOT EXISTS fleet_worker_failures (
|
||||
worker_id TEXT PRIMARY KEY,
|
||||
consecutive_failures INTEGER NOT NULL,
|
||||
circuit_open_until REAL,
|
||||
last_failure_class TEXT,
|
||||
revision INTEGER NOT NULL,
|
||||
updated_at REAL NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS fleet_task_worker_failures (
|
||||
task_id TEXT NOT NULL,
|
||||
worker_id TEXT NOT NULL,
|
||||
failure_count INTEGER NOT NULL,
|
||||
blocked_until REAL,
|
||||
last_failure_class TEXT,
|
||||
revision INTEGER NOT NULL,
|
||||
updated_at REAL NOT NULL,
|
||||
PRIMARY KEY(task_id, worker_id)
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def configure_worker(
|
||||
self,
|
||||
worker_id: str,
|
||||
*,
|
||||
max_slots: int = 1,
|
||||
failure_domain: str = "default",
|
||||
weight: float = 1.0,
|
||||
labels: Mapping[str, str] | None = None,
|
||||
expected_revision: int | None = None,
|
||||
now: float | None = None,
|
||||
) -> FleetWorkerPolicy:
|
||||
worker_id = _text(worker_id)
|
||||
if not worker_id:
|
||||
raise FleetError("FLEET_WORKER_ID_REQUIRED")
|
||||
max_slots = _int(max_slots, 1, 1, MAX_SLOTS)
|
||||
weight = _float(weight, 1.0, 0.05, 100.0)
|
||||
failure_domain = _text(failure_domain, "default")
|
||||
normalized_labels = {
|
||||
_text(k): _text(v)
|
||||
for k, v in dict(labels or {}).items()
|
||||
if _text(k) and _text(v)
|
||||
}
|
||||
now = float(time.time() if now is None else now)
|
||||
with self._transaction():
|
||||
row = self.conn.execute(
|
||||
"SELECT revision FROM fleet_worker_capacity WHERE worker_id=?",
|
||||
(worker_id,),
|
||||
).fetchone()
|
||||
current = int(row[0]) if row else 0
|
||||
if expected_revision is not None and current != int(expected_revision):
|
||||
raise FleetConflict("FLEET_WORKER_POLICY_CAS_MISMATCH")
|
||||
revision = current + 1
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_worker_capacity(
|
||||
worker_id,max_slots,failure_domain,weight,labels_json,revision,updated_at
|
||||
) VALUES (?,?,?,?,?,?,?)
|
||||
ON CONFLICT(worker_id) DO UPDATE SET
|
||||
max_slots=excluded.max_slots,
|
||||
failure_domain=excluded.failure_domain,
|
||||
weight=excluded.weight,
|
||||
labels_json=excluded.labels_json,
|
||||
revision=excluded.revision,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
worker_id,
|
||||
max_slots,
|
||||
failure_domain,
|
||||
weight,
|
||||
json.dumps(normalized_labels, sort_keys=True, separators=(",", ":")),
|
||||
revision,
|
||||
now,
|
||||
),
|
||||
)
|
||||
return FleetWorkerPolicy(
|
||||
worker_id, max_slots, failure_domain, weight, normalized_labels, revision
|
||||
)
|
||||
|
||||
def worker_policy(self, worker_id: str) -> FleetWorkerPolicy:
|
||||
worker_id = _text(worker_id)
|
||||
with self.lock:
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT worker_id,max_slots,failure_domain,weight,labels_json,revision
|
||||
FROM fleet_worker_capacity WHERE worker_id=?
|
||||
""",
|
||||
(worker_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return FleetWorkerPolicy(worker_id=worker_id)
|
||||
return FleetWorkerPolicy(
|
||||
worker_id=str(row[0]),
|
||||
max_slots=int(row[1]),
|
||||
failure_domain=str(row[2]),
|
||||
weight=float(row[3]),
|
||||
labels=json.loads(row[4] or "{}"),
|
||||
revision=int(row[5]),
|
||||
)
|
||||
|
||||
def _cleanup_expired(self, now: float) -> None:
|
||||
self.conn.execute(
|
||||
"""
|
||||
UPDATE fleet_dispatches
|
||||
SET state='EXPIRED', updated_at=?
|
||||
WHERE state='RESERVED' AND reserved_until<=?
|
||||
""",
|
||||
(now, now),
|
||||
)
|
||||
|
||||
def reserved_slots(self, worker_id: str, now: float) -> int:
|
||||
with self.lock:
|
||||
self._cleanup_expired(now)
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT COUNT(*) FROM fleet_dispatches
|
||||
WHERE worker_id=? AND state='RESERVED' AND reserved_until>?
|
||||
""",
|
||||
(worker_id, now),
|
||||
).fetchone()
|
||||
return int(row[0] if row else 0)
|
||||
|
||||
def scope_virtual_finish(self, scope_key: str) -> float:
|
||||
with self.lock:
|
||||
row = self.conn.execute(
|
||||
"SELECT virtual_finish FROM fleet_scope_state WHERE scope_key=?",
|
||||
(scope_key,),
|
||||
).fetchone()
|
||||
return float(row[0]) if row else 0.0
|
||||
|
||||
def reserve(
|
||||
self,
|
||||
*,
|
||||
task_id: str,
|
||||
worker_id: str,
|
||||
task_revision: int,
|
||||
scope_key: str,
|
||||
failure_domain: str,
|
||||
score: tuple[float, float, float, str],
|
||||
reservation_seconds: int,
|
||||
now: float,
|
||||
scope_weight: float,
|
||||
) -> FleetDispatchDecision:
|
||||
reservation_seconds = _int(
|
||||
reservation_seconds, DEFAULT_RESERVATION_SECONDS, 1, 300
|
||||
)
|
||||
reservation_id = "fleet-" + _sha(
|
||||
{
|
||||
"task_id": task_id,
|
||||
"worker_id": worker_id,
|
||||
"task_revision": task_revision,
|
||||
"scope_key": scope_key,
|
||||
}
|
||||
)[:32]
|
||||
reserved_until = now + reservation_seconds
|
||||
with self._transaction():
|
||||
self._cleanup_expired(now)
|
||||
existing = self.conn.execute(
|
||||
"""
|
||||
SELECT reservation_id,worker_id,task_revision,scope_key,failure_domain,
|
||||
reserved_until,state
|
||||
FROM fleet_dispatches WHERE task_id=?
|
||||
""",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
# Capacity is rechecked inside the same durable reservation transaction.
|
||||
# This closes the cross-thread/process race between eligibility scoring and reserve.
|
||||
policy_row = self.conn.execute(
|
||||
"SELECT max_slots FROM fleet_worker_capacity WHERE worker_id=?",
|
||||
(worker_id,),
|
||||
).fetchone()
|
||||
max_slots = int(policy_row[0]) if policy_row else 1
|
||||
if existing and str(existing[6]) in {"RESERVED", "CLAIMED"} and float(existing[5]) > now:
|
||||
if str(existing[1]) != worker_id or int(existing[2]) != task_revision:
|
||||
raise FleetConflict("FLEET_TASK_ALREADY_RESERVED")
|
||||
return FleetDispatchDecision(
|
||||
str(existing[0]), task_id, worker_id, task_revision, str(existing[3]),
|
||||
str(existing[4]), float(existing[5]), score
|
||||
)
|
||||
|
||||
claim_columns = {
|
||||
str(row[1])
|
||||
for row in self.conn.execute("PRAGMA table_info(claims)").fetchall()
|
||||
}
|
||||
active_claims = 0
|
||||
if {"worker_id", "status"}.issubset(claim_columns):
|
||||
sql = "SELECT COUNT(*) FROM claims WHERE worker_id=? AND status='ACTIVE'"
|
||||
params: list[Any] = [worker_id]
|
||||
if "lease_expires_at" in claim_columns:
|
||||
sql += " AND lease_expires_at>?"
|
||||
params.append(now)
|
||||
row = self.conn.execute(sql, tuple(params)).fetchone()
|
||||
active_claims = int(row[0] if row else 0)
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT COUNT(*) FROM fleet_dispatches
|
||||
WHERE worker_id=? AND task_id<>?
|
||||
AND state='RESERVED' AND reserved_until>?
|
||||
""",
|
||||
(worker_id, task_id, now),
|
||||
).fetchone()
|
||||
active_reservations = int(row[0] if row else 0)
|
||||
if active_claims + active_reservations >= max_slots:
|
||||
raise FleetConflict("FLEET_WORKER_CAPACITY_EXHAUSTED")
|
||||
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_dispatches(
|
||||
reservation_id,task_id,worker_id,task_revision,scope_key,failure_domain,
|
||||
state,reserved_until,claim_id,created_at,updated_at
|
||||
) VALUES (?,?,?,?,?,?,'RESERVED',?,NULL,?,?)
|
||||
ON CONFLICT(task_id) DO UPDATE SET
|
||||
reservation_id=excluded.reservation_id,
|
||||
worker_id=excluded.worker_id,
|
||||
task_revision=excluded.task_revision,
|
||||
scope_key=excluded.scope_key,
|
||||
failure_domain=excluded.failure_domain,
|
||||
state='RESERVED',
|
||||
reserved_until=excluded.reserved_until,
|
||||
claim_id=NULL,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
reservation_id, task_id, worker_id, task_revision, scope_key,
|
||||
failure_domain, reserved_until, now, now,
|
||||
),
|
||||
)
|
||||
row = self.conn.execute(
|
||||
"SELECT virtual_finish,dispatch_count,revision FROM fleet_scope_state WHERE scope_key=?",
|
||||
(scope_key,),
|
||||
).fetchone()
|
||||
current_finish = float(row[0]) if row else 0.0
|
||||
dispatch_count = int(row[1]) if row else 0
|
||||
revision = int(row[2]) if row else 0
|
||||
next_finish = max(current_finish, 0.0) + 1.0 / max(scope_weight, 0.05)
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_scope_state(
|
||||
scope_key,virtual_finish,dispatch_count,last_dispatched_at,revision
|
||||
) VALUES (?,?,?,?,?)
|
||||
ON CONFLICT(scope_key) DO UPDATE SET
|
||||
virtual_finish=excluded.virtual_finish,
|
||||
dispatch_count=excluded.dispatch_count,
|
||||
last_dispatched_at=excluded.last_dispatched_at,
|
||||
revision=excluded.revision
|
||||
""",
|
||||
(scope_key, next_finish, dispatch_count + 1, now, revision + 1),
|
||||
)
|
||||
return FleetDispatchDecision(
|
||||
reservation_id, task_id, worker_id, task_revision, scope_key,
|
||||
failure_domain, reserved_until, score
|
||||
)
|
||||
|
||||
def bind_claim(self, reservation_id: str, claim_id: str, *, now: float | None = None) -> None:
|
||||
now = float(time.time() if now is None else now)
|
||||
with self._transaction():
|
||||
row = self.conn.execute(
|
||||
"SELECT state FROM fleet_dispatches WHERE reservation_id=?",
|
||||
(reservation_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise FleetError("FLEET_RESERVATION_NOT_FOUND")
|
||||
if str(row[0]) not in {"RESERVED", "CLAIMED"}:
|
||||
raise FleetConflict("FLEET_RESERVATION_NOT_ACTIVE")
|
||||
self.conn.execute(
|
||||
"""
|
||||
UPDATE fleet_dispatches
|
||||
SET state='CLAIMED', claim_id=?, updated_at=?
|
||||
WHERE reservation_id=?
|
||||
""",
|
||||
(claim_id, now, reservation_id),
|
||||
)
|
||||
|
||||
def release(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
state: str = "RELEASED",
|
||||
now: float | None = None,
|
||||
) -> None:
|
||||
if state not in {"RELEASED", "COMPLETED", "FAILED", "EXPIRED"}:
|
||||
raise FleetError("FLEET_RESERVATION_STATE_INVALID")
|
||||
now = float(time.time() if now is None else now)
|
||||
with self._transaction():
|
||||
self.conn.execute(
|
||||
"UPDATE fleet_dispatches SET state=?,updated_at=? WHERE task_id=?",
|
||||
(state, now, task_id),
|
||||
)
|
||||
|
||||
def worker_circuit_open(self, worker_id: str, now: float) -> bool:
|
||||
with self.lock:
|
||||
row = self.conn.execute(
|
||||
"SELECT circuit_open_until FROM fleet_worker_failures WHERE worker_id=?",
|
||||
(worker_id,),
|
||||
).fetchone()
|
||||
return bool(row and row[0] is not None and float(row[0]) > now)
|
||||
|
||||
def task_worker_blocked(self, task_id: str, worker_id: str, now: float) -> bool:
|
||||
with self.lock:
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT blocked_until FROM fleet_task_worker_failures
|
||||
WHERE task_id=? AND worker_id=?
|
||||
""",
|
||||
(task_id, worker_id),
|
||||
).fetchone()
|
||||
return bool(row and row[0] is not None and float(row[0]) > now)
|
||||
|
||||
def record_worker_failure(
|
||||
self,
|
||||
worker_id: str,
|
||||
failure_class: str,
|
||||
*,
|
||||
threshold: int = DEFAULT_CIRCUIT_THRESHOLD,
|
||||
base_circuit_seconds: int = DEFAULT_CIRCUIT_SECONDS,
|
||||
now: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = float(time.time() if now is None else now)
|
||||
threshold = _int(threshold, DEFAULT_CIRCUIT_THRESHOLD, 1, 20)
|
||||
base_circuit_seconds = _int(base_circuit_seconds, DEFAULT_CIRCUIT_SECONDS, 1, 3600)
|
||||
with self._transaction():
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT consecutive_failures,revision FROM fleet_worker_failures
|
||||
WHERE worker_id=?
|
||||
""",
|
||||
(worker_id,),
|
||||
).fetchone()
|
||||
failures = (int(row[0]) if row else 0) + 1
|
||||
revision = (int(row[1]) if row else 0) + 1
|
||||
open_until = None
|
||||
if failures >= threshold:
|
||||
exponent = min(failures - threshold, 8)
|
||||
open_until = now + min(
|
||||
MAX_CIRCUIT_SECONDS, base_circuit_seconds * (2**exponent)
|
||||
)
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_worker_failures(
|
||||
worker_id,consecutive_failures,circuit_open_until,last_failure_class,
|
||||
revision,updated_at
|
||||
) VALUES (?,?,?,?,?,?)
|
||||
ON CONFLICT(worker_id) DO UPDATE SET
|
||||
consecutive_failures=excluded.consecutive_failures,
|
||||
circuit_open_until=excluded.circuit_open_until,
|
||||
last_failure_class=excluded.last_failure_class,
|
||||
revision=excluded.revision,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(worker_id, failures, open_until, _text(failure_class, "unknown"), revision, now),
|
||||
)
|
||||
return {
|
||||
"worker_id": worker_id,
|
||||
"consecutive_failures": failures,
|
||||
"circuit_open_until": open_until,
|
||||
"revision": revision,
|
||||
}
|
||||
|
||||
def record_worker_success(self, worker_id: str, *, now: float | None = None) -> None:
|
||||
now = float(time.time() if now is None else now)
|
||||
with self._transaction():
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_worker_failures(
|
||||
worker_id,consecutive_failures,circuit_open_until,last_failure_class,
|
||||
revision,updated_at
|
||||
) VALUES (?,0,NULL,NULL,1,?)
|
||||
ON CONFLICT(worker_id) DO UPDATE SET
|
||||
consecutive_failures=0,
|
||||
circuit_open_until=NULL,
|
||||
last_failure_class=NULL,
|
||||
revision=revision+1,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(worker_id, now),
|
||||
)
|
||||
|
||||
def record_task_worker_failure(
|
||||
self,
|
||||
task_id: str,
|
||||
worker_id: str,
|
||||
failure_class: str,
|
||||
*,
|
||||
quarantine_seconds: int = 60,
|
||||
now: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = float(time.time() if now is None else now)
|
||||
quarantine_seconds = _int(quarantine_seconds, 60, 1, 3600)
|
||||
with self._transaction():
|
||||
row = self.conn.execute(
|
||||
"""
|
||||
SELECT failure_count,revision FROM fleet_task_worker_failures
|
||||
WHERE task_id=? AND worker_id=?
|
||||
""",
|
||||
(task_id, worker_id),
|
||||
).fetchone()
|
||||
count = (int(row[0]) if row else 0) + 1
|
||||
revision = (int(row[1]) if row else 0) + 1
|
||||
blocked_until = now + quarantine_seconds
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO fleet_task_worker_failures(
|
||||
task_id,worker_id,failure_count,blocked_until,last_failure_class,
|
||||
revision,updated_at
|
||||
) VALUES (?,?,?,?,?,?,?)
|
||||
ON CONFLICT(task_id,worker_id) DO UPDATE SET
|
||||
failure_count=excluded.failure_count,
|
||||
blocked_until=excluded.blocked_until,
|
||||
last_failure_class=excluded.last_failure_class,
|
||||
revision=excluded.revision,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
task_id, worker_id, count, blocked_until,
|
||||
_text(failure_class, "unknown"), revision, now,
|
||||
),
|
||||
)
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"worker_id": worker_id,
|
||||
"failure_count": count,
|
||||
"blocked_until": blocked_until,
|
||||
"revision": revision,
|
||||
}
|
||||
|
||||
|
||||
class FleetScheduler:
|
||||
"""Deterministic capacity/fairness scheduler over existing TaskPlane records."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Any,
|
||||
workers: Any,
|
||||
*,
|
||||
now_fn: Any = time.time,
|
||||
reservation_seconds: int = DEFAULT_RESERVATION_SECONDS,
|
||||
) -> None:
|
||||
self.store = store
|
||||
self.workers = workers
|
||||
self.now_fn = now_fn
|
||||
self.reservation_seconds = _int(
|
||||
reservation_seconds, DEFAULT_RESERVATION_SECONDS, 1, 300
|
||||
)
|
||||
connection = getattr(store, "conn", None)
|
||||
if not isinstance(connection, sqlite3.Connection):
|
||||
raise FleetError("FLEET_TASKPLANE_SQLITE_CONNECTION_REQUIRED")
|
||||
lock = getattr(store, "_lock", None)
|
||||
self.state = FleetStateStore(connection, lock)
|
||||
|
||||
def configure_worker(self, worker_id: str, **kwargs: Any) -> FleetWorkerPolicy:
|
||||
return self.state.configure_worker(worker_id, **kwargs)
|
||||
|
||||
def _online(self, worker_id: str, now: float) -> bool:
|
||||
method = getattr(self.workers, "is_online", None)
|
||||
if not callable(method):
|
||||
raise FleetError("FLEET_WORKER_REGISTRY_ONLINE_CHECK_REQUIRED")
|
||||
try:
|
||||
return bool(method(worker_id, now=now))
|
||||
except TypeError:
|
||||
return bool(method(worker_id))
|
||||
|
||||
def _worker_record(self, worker_id: str) -> object:
|
||||
getter = getattr(self.store, "get_worker", None)
|
||||
if not callable(getter):
|
||||
raise FleetError("FLEET_WORKER_LOOKUP_REQUIRED")
|
||||
return getter(worker_id)
|
||||
|
||||
def _active_claims(self, worker_id: str, now: float) -> int:
|
||||
# The TaskPlane store owns one SQLite connection in the current runtime.
|
||||
# Fleet reads must share its lock just like fleet writes do, otherwise two
|
||||
# concurrent pullers can issue overlapping operations on the same handle.
|
||||
with self.state.lock:
|
||||
columns = {
|
||||
str(row[1])
|
||||
for row in self.store.conn.execute("PRAGMA table_info(claims)").fetchall()
|
||||
}
|
||||
required = {"worker_id", "status"}
|
||||
if not required.issubset(columns):
|
||||
raise FleetError("FLEET_CLAIMS_SCHEMA_UNSUPPORTED")
|
||||
sql = "SELECT COUNT(*) FROM claims WHERE worker_id=? AND status='ACTIVE'"
|
||||
params: list[Any] = [worker_id]
|
||||
if "lease_expires_at" in columns:
|
||||
sql += " AND lease_expires_at>?"
|
||||
params.append(now)
|
||||
row = self.store.conn.execute(sql, tuple(params)).fetchone()
|
||||
return int(row[0] if row else 0)
|
||||
|
||||
@staticmethod
|
||||
def _labels_match(policy: FleetWorkerPolicy, task: object) -> bool:
|
||||
metadata = _task_metadata(task)
|
||||
affinity = metadata.get("fleet_affinity_labels") or {}
|
||||
if affinity and not isinstance(affinity, Mapping):
|
||||
return False
|
||||
labels = dict(policy.labels or {})
|
||||
for key, value in dict(affinity).items():
|
||||
if labels.get(_text(key)) != _text(value):
|
||||
return False
|
||||
avoid = metadata.get("fleet_avoid_failure_domains") or []
|
||||
if isinstance(avoid, Sequence) and not isinstance(avoid, (str, bytes)):
|
||||
if policy.failure_domain in {_text(item) for item in avoid}:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _eligible_worker(self, worker_id: str, task: object, now: float) -> tuple[bool, int, FleetWorkerPolicy]:
|
||||
if not self._online(worker_id, now):
|
||||
return False, 0, self.state.worker_policy(worker_id)
|
||||
worker = self._worker_record(worker_id)
|
||||
policy = self.state.worker_policy(worker_id)
|
||||
if self.state.worker_circuit_open(worker_id, now):
|
||||
return False, 0, policy
|
||||
if self.state.task_worker_blocked(_task_id(task), worker_id, now):
|
||||
return False, 0, policy
|
||||
if not self._labels_match(policy, task):
|
||||
return False, 0, policy
|
||||
required = _task_required_capabilities(task)
|
||||
capabilities = _worker_capabilities(worker)
|
||||
if required and not required.issubset(capabilities):
|
||||
return False, 0, policy
|
||||
active = self._active_claims(worker_id, now)
|
||||
reserved = self.state.reserved_slots(worker_id, now)
|
||||
load = active + reserved
|
||||
if load >= policy.max_slots:
|
||||
return False, load, policy
|
||||
return True, load, policy
|
||||
|
||||
@staticmethod
|
||||
def _task_priority(task: object, now: float) -> float:
|
||||
metadata = _task_metadata(task)
|
||||
base = _int(metadata.get("fleet_priority"), 0, -100, 100)
|
||||
created = _task_created_at(task)
|
||||
waited = max(0.0, now - created) if created > 0 else 0.0
|
||||
aging = min(100.0, waited / 60.0)
|
||||
return float(base) + aging
|
||||
|
||||
def _task_score(self, task: object, now: float) -> tuple[float, float, float, str]:
|
||||
scope = _task_scope_key(task)
|
||||
virtual_finish = self.state.scope_virtual_finish(scope)
|
||||
priority = self._task_priority(task, now)
|
||||
return (virtual_finish, -priority, _task_created_at(task), _task_id(task))
|
||||
|
||||
def _tasks(self, actor_context: Any = None) -> list[object]:
|
||||
lister = getattr(self.store, "list_tasks", None)
|
||||
if not callable(lister):
|
||||
raise FleetError("FLEET_TASK_LIST_REQUIRED")
|
||||
try:
|
||||
tasks = lister(actor_context)
|
||||
except TypeError:
|
||||
tasks = lister()
|
||||
return [
|
||||
task for task in tasks
|
||||
if _task_status(task) in {"CREATED", "QUEUED"}
|
||||
and _task_id(task)
|
||||
]
|
||||
|
||||
def rank_for_worker(
|
||||
self,
|
||||
worker_id: str,
|
||||
*,
|
||||
actor_context: Any = None,
|
||||
now: float | None = None,
|
||||
) -> list[object]:
|
||||
"""Return fleet-ranked tasks without reserving or claiming anything."""
|
||||
now = float(self.now_fn() if now is None else now)
|
||||
worker_id = _text(worker_id)
|
||||
if not worker_id:
|
||||
return []
|
||||
ranked: list[object] = []
|
||||
for task in self._tasks(actor_context):
|
||||
eligible, _, _ = self._eligible_worker(worker_id, task, now)
|
||||
if eligible:
|
||||
ranked.append(task)
|
||||
return sorted(ranked, key=lambda task: self._task_score(task, now))
|
||||
|
||||
def reserve_for_worker(
|
||||
self,
|
||||
worker_id: str,
|
||||
task: object,
|
||||
*,
|
||||
now: float | None = None,
|
||||
) -> FleetDispatchDecision:
|
||||
"""Reserve exactly one already-ranked task, rechecking eligibility and capacity."""
|
||||
now = float(self.now_fn() if now is None else now)
|
||||
worker_id = _text(worker_id)
|
||||
eligible, _, policy = self._eligible_worker(worker_id, task, now)
|
||||
if not eligible:
|
||||
raise FleetConflict("FLEET_WORKER_NOT_ELIGIBLE")
|
||||
metadata = _task_metadata(task)
|
||||
scope_weight = _float(metadata.get("fleet_scope_weight"), 1.0, 0.05, 100.0)
|
||||
return self.state.reserve(
|
||||
task_id=_task_id(task),
|
||||
worker_id=worker_id,
|
||||
task_revision=_task_revision(task),
|
||||
scope_key=_task_scope_key(task),
|
||||
failure_domain=policy.failure_domain,
|
||||
score=self._task_score(task, now),
|
||||
reservation_seconds=self.reservation_seconds,
|
||||
now=now,
|
||||
scope_weight=scope_weight,
|
||||
)
|
||||
|
||||
def select(
|
||||
self,
|
||||
worker_ids: Sequence[str],
|
||||
*,
|
||||
actor_context: Any = None,
|
||||
limit: int = 1,
|
||||
now: float | None = None,
|
||||
) -> list[FleetDispatchDecision]:
|
||||
now = float(self.now_fn() if now is None else now)
|
||||
limit = _int(limit, 1, 1, 100)
|
||||
worker_ids = tuple(dict.fromkeys(_text(item) for item in worker_ids if _text(item)))
|
||||
tasks = sorted(self._tasks(actor_context), key=lambda task: self._task_score(task, now))
|
||||
decisions: list[FleetDispatchDecision] = []
|
||||
local_load: dict[str, int] = {}
|
||||
for task in tasks:
|
||||
if len(decisions) >= limit:
|
||||
break
|
||||
candidates: list[tuple[int, float, str, FleetWorkerPolicy]] = []
|
||||
for worker_id in worker_ids:
|
||||
eligible, load, policy = self._eligible_worker(worker_id, task, now)
|
||||
if not eligible:
|
||||
continue
|
||||
effective_load = local_load.get(worker_id, load)
|
||||
if effective_load >= policy.max_slots:
|
||||
continue
|
||||
candidates.append(
|
||||
(
|
||||
effective_load,
|
||||
-policy.weight,
|
||||
worker_id,
|
||||
policy,
|
||||
)
|
||||
)
|
||||
if not candidates:
|
||||
continue
|
||||
candidates.sort(key=lambda item: (item[0], item[1], item[2]))
|
||||
_, _, worker_id, policy = candidates[0]
|
||||
score = self._task_score(task, now)
|
||||
metadata = _task_metadata(task)
|
||||
scope_weight = _float(metadata.get("fleet_scope_weight"), 1.0, 0.05, 100.0)
|
||||
try:
|
||||
decision = self.state.reserve(
|
||||
task_id=_task_id(task),
|
||||
worker_id=worker_id,
|
||||
task_revision=_task_revision(task),
|
||||
scope_key=_task_scope_key(task),
|
||||
failure_domain=policy.failure_domain,
|
||||
score=score,
|
||||
reservation_seconds=self.reservation_seconds,
|
||||
now=now,
|
||||
scope_weight=scope_weight,
|
||||
)
|
||||
except FleetConflict:
|
||||
continue
|
||||
local_load[worker_id] = local_load.get(
|
||||
worker_id,
|
||||
self._active_claims(worker_id, now)
|
||||
+ self.state.reserved_slots(worker_id, now)
|
||||
- 1,
|
||||
) + 1
|
||||
decisions.append(decision)
|
||||
return decisions
|
||||
@@ -13,6 +13,8 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from .scope import normalize_task_scope
|
||||
|
||||
|
||||
SENSITIVE_KEYS = {
|
||||
"secret",
|
||||
@@ -66,6 +68,12 @@ TASK_PLAN_KEYS = {
|
||||
"allowed_input_paths",
|
||||
"allowed_output_paths",
|
||||
"output_policy",
|
||||
"scope_type",
|
||||
"tenant_id",
|
||||
"user_id",
|
||||
"project_id",
|
||||
"conversation_id",
|
||||
"run_id",
|
||||
}
|
||||
|
||||
|
||||
@@ -237,6 +245,12 @@ class TaskPlan:
|
||||
schema_version: int = 1
|
||||
task_id: str | None = None
|
||||
task_type: str = "code.change"
|
||||
scope_type: str = "legacy"
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
run_id: str | None = None
|
||||
issuer: dict[str, Any] = field(default_factory=dict)
|
||||
intent: str | None = None
|
||||
required_capabilities: list[str] = field(default_factory=list)
|
||||
@@ -311,10 +325,17 @@ class TaskPlan:
|
||||
raise ValueError("credential-secret requires metadata-only logging")
|
||||
if output_policy != "hash-count-status-only":
|
||||
raise ValueError("credential-secret requires restricted output")
|
||||
scope = normalize_task_scope(data)
|
||||
return cls(
|
||||
schema_version=schema_version,
|
||||
task_id=validate_identifier(data["task_id"], "task_id"),
|
||||
task_type=validate_type(data.get("task_type", "code.change"), "task_type"),
|
||||
scope_type=scope.scope_type,
|
||||
tenant_id=scope.tenant_id,
|
||||
user_id=scope.user_id,
|
||||
project_id=scope.project_id,
|
||||
conversation_id=scope.conversation_id,
|
||||
run_id=scope.run_id,
|
||||
issuer=issuer,
|
||||
intent=data.get("intent"),
|
||||
required_capabilities=_normalize_list(data.get("required_capabilities", [])),
|
||||
@@ -348,6 +369,12 @@ class TaskPlan:
|
||||
"schema_version": self.schema_version,
|
||||
"task_id": self.task_id,
|
||||
"task_type": self.task_type,
|
||||
"scope_type": self.scope_type,
|
||||
"tenant_id": self.tenant_id,
|
||||
"user_id": self.user_id,
|
||||
"project_id": self.project_id,
|
||||
"conversation_id": self.conversation_id,
|
||||
"run_id": self.run_id,
|
||||
"issuer": self.issuer,
|
||||
"intent": self.intent,
|
||||
"required_capabilities": list(self.required_capabilities),
|
||||
|
||||
@@ -3,10 +3,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
from .errors import ForbiddenError, LockedError
|
||||
from .models import ProbeEvidence, ResultEnvelope, TaskRecord
|
||||
from .scope import TaskActorContext, actor_context as build_actor_context
|
||||
|
||||
|
||||
RISK_RANK = {"low": 1, "medium": 2, "high": 3}
|
||||
@@ -93,6 +94,10 @@ class PolicyEngine:
|
||||
result_signature_verifier: Callable[[ResultEnvelope], bool] | None = None,
|
||||
probe_attestation_verifier: Callable[[ProbeEvidence], bool] | None = None,
|
||||
identity_capabilities: dict[str, set[str]] | None = None,
|
||||
identity_actor_contexts: Mapping[
|
||||
str, TaskActorContext | Mapping[str, Any]
|
||||
]
|
||||
| None = None,
|
||||
admin_identities: set[str] | None = None,
|
||||
readonly_worker_execution_enabled: bool = False,
|
||||
readonly_worker_profiles: list[ReadonlyWorkerProfile] | None = None,
|
||||
@@ -109,12 +114,26 @@ class PolicyEngine:
|
||||
identity: set(capabilities)
|
||||
for identity, capabilities in (identity_capabilities or {}).items()
|
||||
}
|
||||
self.identity_actor_contexts = dict(identity_actor_contexts or {})
|
||||
self.admin_identities = set(admin_identities or ())
|
||||
self.readonly_worker_execution_enabled = readonly_worker_execution_enabled
|
||||
self.readonly_worker_profiles = {
|
||||
profile.identity: profile for profile in (readonly_worker_profiles or [])
|
||||
}
|
||||
|
||||
def actor_context(self, identity: str) -> TaskActorContext:
|
||||
configured = self.identity_actor_contexts.get(
|
||||
identity, self.identity_actor_contexts.get("*")
|
||||
)
|
||||
capabilities = set(
|
||||
self.identity_capabilities.get(
|
||||
identity, self.identity_capabilities.get("*", set())
|
||||
)
|
||||
)
|
||||
if identity in self.admin_identities:
|
||||
capabilities.add("task.admin")
|
||||
return build_actor_context(identity, configured, capabilities)
|
||||
|
||||
def _readonly_worker_allows(
|
||||
self, action: str, actor: str, task: TaskRecord | None
|
||||
) -> bool:
|
||||
|
||||
@@ -74,3 +74,203 @@ class TaskScheduler:
|
||||
operation_digest=task.task_plan.operation_digest(),
|
||||
)
|
||||
return claim
|
||||
|
||||
# R4.2 P2 durable multi-agent fleet scheduling public surface.
|
||||
from .fleet import (
|
||||
FleetConflict,
|
||||
FleetDispatchDecision,
|
||||
FleetError,
|
||||
FleetScheduler,
|
||||
FleetStateStore,
|
||||
FleetWorkerPolicy,
|
||||
)
|
||||
# R4.2 P2-B fleet-aware pull scheduling.
|
||||
#
|
||||
# The original TaskScheduler remains the security oracle and claim-CAS implementation.
|
||||
# Fleet ordering only controls the order presented to that oracle.
|
||||
import inspect as _fleet_inspect
|
||||
import threading as _fleet_threading
|
||||
from collections.abc import Mapping as _FleetMapping
|
||||
|
||||
from .worker_registry import WorkerRegistry as _FleetWorkerRegistry
|
||||
from .errors import NotFoundError as _FleetNotFoundError
|
||||
|
||||
_LegacyTaskScheduler = TaskScheduler
|
||||
_LegacyScheduleParameters = _fleet_inspect.signature(
|
||||
_LegacyTaskScheduler.schedule_for_worker
|
||||
).parameters
|
||||
_LegacyNextParameters = _fleet_inspect.signature(
|
||||
_LegacyTaskScheduler.next_queued_task
|
||||
).parameters
|
||||
|
||||
|
||||
class _FleetOrderedStoreProxy:
|
||||
def __init__(self, delegate, tasks):
|
||||
self._delegate = delegate
|
||||
self._tasks = tuple(tasks)
|
||||
|
||||
def list_tasks(self, *args, **kwargs):
|
||||
return list(self._tasks)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._delegate, name)
|
||||
|
||||
|
||||
def _fleet_claim_field(claim, name, default=None):
|
||||
if isinstance(claim, _FleetMapping):
|
||||
return claim.get(name, default)
|
||||
return getattr(claim, name, default)
|
||||
|
||||
|
||||
class TaskScheduler(_LegacyTaskScheduler):
|
||||
"""Fleet-integrated pull scheduler preserving legacy security and claim semantics."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
store = getattr(self, "store", args[0] if args else None)
|
||||
if store is None:
|
||||
raise FleetError("FLEET_SCHEDULER_STORE_REQUIRED")
|
||||
registry = getattr(self, "workers", None)
|
||||
if registry is None:
|
||||
registry = _FleetWorkerRegistry(store)
|
||||
self.fleet_scheduler = FleetScheduler(store, registry)
|
||||
self._fleet_local = _fleet_threading.local()
|
||||
|
||||
def _fleet_legacy_next(self, worker_id, tasks, *args, **kwargs):
|
||||
proxy = _FleetOrderedStoreProxy(self.store, tasks)
|
||||
legacy = _LegacyTaskScheduler(proxy)
|
||||
legacy_kwargs = dict(kwargs)
|
||||
if "actor_context" not in _LegacyNextParameters:
|
||||
legacy_kwargs.pop("actor_context", None)
|
||||
return legacy.next_queued_task(worker_id, *args, **legacy_kwargs)
|
||||
|
||||
def next_queued_task(self, worker_id, *args, **kwargs):
|
||||
actor_context = kwargs.get(
|
||||
"actor_context", getattr(self._fleet_local, "actor_context", None)
|
||||
)
|
||||
|
||||
# Preserve legacy fail-closed worker checks before fleet ordering.
|
||||
probe_kwargs = dict(kwargs)
|
||||
if "actor_context" not in _LegacyNextParameters:
|
||||
probe_kwargs.pop("actor_context", None)
|
||||
legacy_probe = _LegacyTaskScheduler(self.store).next_queued_task(
|
||||
worker_id, *args, **probe_kwargs
|
||||
)
|
||||
if legacy_probe is None:
|
||||
return None
|
||||
|
||||
ranked = self.fleet_scheduler.rank_for_worker(
|
||||
worker_id, actor_context=actor_context
|
||||
)
|
||||
# Preserve FleetScheduler fairness order. The legacy scheduler is used
|
||||
# only as a fail-closed security oracle: present one fleet-ranked task
|
||||
# at a time so its internal sort cannot reorder across scopes.
|
||||
# Keep the legacy priority dimension ahead of Fleet fairness for
|
||||
# compatibility with existing TaskPlane scheduling contracts; the
|
||||
# Fleet order remains authoritative among equal-priority tasks.
|
||||
def _legacy_priority(task):
|
||||
plan = getattr(task, "task_plan", task)
|
||||
metadata = getattr(plan, "metadata", {}) or {}
|
||||
if not isinstance(metadata, _FleetMapping):
|
||||
return 0
|
||||
try:
|
||||
return int(metadata.get("priority", 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
ranked = sorted(ranked, key=lambda task: -_legacy_priority(task))
|
||||
for ranked_candidate in ranked:
|
||||
try:
|
||||
candidate = self._fleet_legacy_next(
|
||||
worker_id, [ranked_candidate], *args, **kwargs
|
||||
)
|
||||
except _FleetNotFoundError:
|
||||
continue
|
||||
if candidate is None:
|
||||
continue
|
||||
|
||||
ranked_id = getattr(
|
||||
ranked_candidate,
|
||||
"task_id",
|
||||
getattr(getattr(ranked_candidate, "task_plan", None), "task_id", None),
|
||||
)
|
||||
candidate_id = getattr(
|
||||
candidate,
|
||||
"task_id",
|
||||
getattr(getattr(candidate, "task_plan", None), "task_id", None),
|
||||
)
|
||||
if candidate_id != ranked_id:
|
||||
raise FleetConflict("FLEET_LEGACY_ORACLE_ORDER_VIOLATION")
|
||||
|
||||
try:
|
||||
decision = self.fleet_scheduler.reserve_for_worker(
|
||||
worker_id, candidate
|
||||
)
|
||||
except FleetConflict:
|
||||
continue
|
||||
self._fleet_local.pending = decision
|
||||
return candidate
|
||||
return None
|
||||
|
||||
def schedule_for_worker(self, worker_id, *args, **kwargs):
|
||||
actor_context = kwargs.get("actor_context")
|
||||
legacy_kwargs = dict(kwargs)
|
||||
if "actor_context" not in _LegacyScheduleParameters:
|
||||
legacy_kwargs.pop("actor_context", None)
|
||||
self._fleet_local.actor_context = actor_context
|
||||
self._fleet_local.pending = None
|
||||
try:
|
||||
claim = _LegacyTaskScheduler.schedule_for_worker(
|
||||
self, worker_id, *args, **legacy_kwargs
|
||||
)
|
||||
decision = getattr(self._fleet_local, "pending", None)
|
||||
if claim is None:
|
||||
if decision is not None:
|
||||
self.fleet_scheduler.state.release(
|
||||
decision.task_id, state="RELEASED"
|
||||
)
|
||||
return None
|
||||
if decision is not None:
|
||||
claim_id = _fleet_claim_field(claim, "claim_id")
|
||||
claim_task_id = _fleet_claim_field(claim, "task_id", decision.task_id)
|
||||
if not claim_id:
|
||||
self.fleet_scheduler.state.release(
|
||||
decision.task_id, state="RELEASED"
|
||||
)
|
||||
raise FleetError("FLEET_CLAIM_ID_REQUIRED")
|
||||
if claim_task_id != decision.task_id:
|
||||
self.fleet_scheduler.state.release(
|
||||
decision.task_id, state="RELEASED"
|
||||
)
|
||||
raise FleetConflict("FLEET_CLAIM_TASK_MISMATCH")
|
||||
self.fleet_scheduler.state.bind_claim(
|
||||
decision.reservation_id, str(claim_id)
|
||||
)
|
||||
return claim
|
||||
except Exception:
|
||||
decision = getattr(self._fleet_local, "pending", None)
|
||||
if decision is not None:
|
||||
self.fleet_scheduler.state.release(
|
||||
decision.task_id, state="RELEASED"
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
self._fleet_local.pending = None
|
||||
self._fleet_local.actor_context = None
|
||||
|
||||
def record_worker_failure(self, worker_id, failure_class, **kwargs):
|
||||
return self.fleet_scheduler.state.record_worker_failure(
|
||||
worker_id, failure_class, **kwargs
|
||||
)
|
||||
|
||||
def record_worker_success(self, worker_id, **kwargs):
|
||||
return self.fleet_scheduler.state.record_worker_success(
|
||||
worker_id, **kwargs
|
||||
)
|
||||
|
||||
def quarantine_task_worker(
|
||||
self, task_id, worker_id, failure_class, **kwargs
|
||||
):
|
||||
return self.fleet_scheduler.state.record_task_worker_failure(
|
||||
task_id, worker_id, failure_class, **kwargs
|
||||
)
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Trusted TaskPlane scope model and SQL authorization helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Any, Mapping
|
||||
|
||||
from .errors import ForbiddenError, ValidationError
|
||||
|
||||
SCOPE_FIELDS = ("tenant_id", "user_id", "project_id", "conversation_id", "run_id")
|
||||
PRIVILEGED_CAPABILITIES = frozenset(
|
||||
{"*", "admin", "system.admin", "task.admin", "task.read.all"}
|
||||
)
|
||||
|
||||
|
||||
def _value(value: Any, field: str) -> Any:
|
||||
if hasattr(value, "task_plan"):
|
||||
value = value.task_plan
|
||||
if isinstance(value, Mapping):
|
||||
return value.get(field)
|
||||
return getattr(value, field, None)
|
||||
|
||||
|
||||
def _text(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
result = str(value).strip()
|
||||
return result or None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TaskScope:
|
||||
scope_type: str
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
run_id: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TaskActorContext:
|
||||
identity: str
|
||||
scope_type: str = "legacy"
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
run_id: str | None = None
|
||||
capabilities: frozenset[str] = frozenset()
|
||||
|
||||
|
||||
def normalize_task_scope(value: Any) -> TaskScope:
|
||||
fields = {field: _text(_value(value, field)) for field in SCOPE_FIELDS}
|
||||
declared = _text(_value(value, "scope_type"))
|
||||
if declared is None:
|
||||
declared = "tenant" if fields["tenant_id"] else "legacy"
|
||||
if declared == "tenant":
|
||||
if fields["tenant_id"] is None:
|
||||
raise ValidationError(
|
||||
"tenant scope requires tenant_id", code="TASK_SCOPE_TENANT_ID_REQUIRED"
|
||||
)
|
||||
elif declared == "legacy":
|
||||
if any(fields.values()):
|
||||
raise ValidationError(
|
||||
"legacy scope cannot contain scoped ownership",
|
||||
code="TASK_SCOPE_LEGACY_WITH_SCOPED_FIELDS",
|
||||
)
|
||||
elif declared == "system":
|
||||
if any(fields[field] for field in ("tenant_id", "user_id", "project_id", "conversation_id")):
|
||||
raise ValidationError(
|
||||
"system scope cannot contain tenant ownership",
|
||||
code="TASK_SCOPE_SYSTEM_WITH_TENANT_FIELDS",
|
||||
)
|
||||
else:
|
||||
raise ValidationError("invalid task scope type", code="TASK_SCOPE_TYPE_INVALID")
|
||||
return TaskScope(declared, **fields)
|
||||
|
||||
|
||||
def actor_is_privileged(actor: TaskActorContext) -> bool:
|
||||
return bool(actor.capabilities & PRIVILEGED_CAPABILITIES)
|
||||
|
||||
|
||||
def actor_context(
|
||||
identity: str,
|
||||
configured: TaskActorContext | Mapping[str, Any] | None,
|
||||
capabilities: set[str] | frozenset[str],
|
||||
) -> TaskActorContext:
|
||||
if isinstance(configured, TaskActorContext):
|
||||
context = configured
|
||||
else:
|
||||
data = dict(configured or {})
|
||||
data.pop("identity", None)
|
||||
data.pop("capabilities", None)
|
||||
context = TaskActorContext(identity=str(identity), **data)
|
||||
scope = normalize_task_scope(context)
|
||||
return replace(
|
||||
context,
|
||||
identity=str(identity),
|
||||
scope_type=scope.scope_type,
|
||||
tenant_id=scope.tenant_id,
|
||||
user_id=scope.user_id,
|
||||
project_id=scope.project_id,
|
||||
conversation_id=scope.conversation_id,
|
||||
run_id=scope.run_id,
|
||||
capabilities=frozenset(str(item) for item in capabilities),
|
||||
)
|
||||
|
||||
|
||||
def bind_task_scope(plan: Any, actor: TaskActorContext, *, public_create: bool) -> Any:
|
||||
requested = normalize_task_scope(plan)
|
||||
if actor_is_privileged(actor):
|
||||
trusted = normalize_task_scope(actor)
|
||||
if requested.scope_type == "legacy" and trusted.scope_type != "legacy":
|
||||
plan.scope_type = trusted.scope_type
|
||||
for field in SCOPE_FIELDS:
|
||||
setattr(plan, field, getattr(actor, field))
|
||||
return plan
|
||||
trusted = normalize_task_scope(actor)
|
||||
if public_create and trusted.scope_type == "legacy":
|
||||
raise ForbiddenError(
|
||||
"authenticated identity has no tenant scope",
|
||||
code="TASK_ACTOR_SCOPE_UNAVAILABLE",
|
||||
)
|
||||
if trusted.scope_type != "tenant":
|
||||
if requested.scope_type == "legacy":
|
||||
plan.scope_type = trusted.scope_type
|
||||
for field in SCOPE_FIELDS:
|
||||
setattr(plan, field, getattr(actor, field))
|
||||
return plan
|
||||
if requested.scope_type != trusted.scope_type:
|
||||
raise ForbiddenError("task scope mismatch", code="TASK_SCOPE_MISMATCH_DENY")
|
||||
return plan
|
||||
if requested.scope_type not in {"legacy", "tenant"}:
|
||||
raise ForbiddenError("task scope mismatch", code="TASK_SCOPE_MISMATCH_DENY")
|
||||
updates: dict[str, Any] = {"scope_type": "tenant"}
|
||||
for field in SCOPE_FIELDS:
|
||||
declared = getattr(requested, field)
|
||||
authoritative = getattr(actor, field)
|
||||
if declared is not None and declared != authoritative:
|
||||
raise ForbiddenError(
|
||||
"task scope mismatch",
|
||||
code="TASK_SCOPE_MISMATCH_DENY",
|
||||
details={"field": field},
|
||||
)
|
||||
updates[field] = authoritative
|
||||
for field, item in updates.items():
|
||||
setattr(plan, field, item)
|
||||
return plan
|
||||
|
||||
|
||||
def actor_scope_sql(actor: TaskActorContext | None) -> tuple[str, tuple[Any, ...]]:
|
||||
if actor is None or actor_is_privileged(actor):
|
||||
return "", ()
|
||||
scope = normalize_task_scope(actor)
|
||||
if scope.scope_type == "legacy":
|
||||
return "scope_type = 'legacy'", ()
|
||||
if scope.scope_type == "system":
|
||||
return "scope_type = 'system'", ()
|
||||
clauses = ["scope_type = 'tenant'", "tenant_id = ?"]
|
||||
params: list[Any] = [scope.tenant_id]
|
||||
for field in ("user_id", "project_id", "conversation_id", "run_id"):
|
||||
value = getattr(scope, field)
|
||||
if value is None:
|
||||
clauses.append(f"{field} IS NULL")
|
||||
else:
|
||||
clauses.append(f"({field} IS NULL OR {field} = ?)")
|
||||
params.append(value)
|
||||
return " AND ".join(clauses), tuple(params)
|
||||
|
||||
|
||||
def task_scope_values(value: Any) -> tuple[str | None, ...]:
|
||||
scope = normalize_task_scope(value)
|
||||
return (
|
||||
scope.scope_type,
|
||||
scope.tenant_id,
|
||||
scope.user_id,
|
||||
scope.project_id,
|
||||
scope.conversation_id,
|
||||
scope.run_id,
|
||||
)
|
||||
@@ -32,6 +32,7 @@ from .models import (
|
||||
redact_value,
|
||||
validate_identifier,
|
||||
)
|
||||
from .scope import TaskActorContext, actor_scope_sql, task_scope_values
|
||||
|
||||
|
||||
EventEnvelope = tuple[str, str, str, str, str, float]
|
||||
@@ -134,6 +135,12 @@ class TaskPlaneStore:
|
||||
status TEXT NOT NULL,
|
||||
task_json TEXT NOT NULL,
|
||||
task_type TEXT NOT NULL,
|
||||
scope_type TEXT NOT NULL DEFAULT 'legacy',
|
||||
tenant_id TEXT,
|
||||
user_id TEXT,
|
||||
project_id TEXT,
|
||||
conversation_id TEXT,
|
||||
run_id TEXT,
|
||||
root_task_id TEXT NOT NULL,
|
||||
parent_task_id TEXT,
|
||||
correlation_id TEXT,
|
||||
@@ -272,6 +279,21 @@ class TaskPlaneStore:
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._ensure_columns(
|
||||
"tasks",
|
||||
{
|
||||
"scope_type": "TEXT NOT NULL DEFAULT 'legacy'",
|
||||
"tenant_id": "TEXT",
|
||||
"user_id": "TEXT",
|
||||
"project_id": "TEXT",
|
||||
"conversation_id": "TEXT",
|
||||
"run_id": "TEXT",
|
||||
},
|
||||
)
|
||||
self.conn.execute(
|
||||
"UPDATE tasks SET scope_type='legacy' "
|
||||
"WHERE scope_type IS NULL OR TRIM(scope_type)=''"
|
||||
)
|
||||
self._ensure_columns(
|
||||
"workers",
|
||||
{
|
||||
@@ -299,6 +321,14 @@ class TaskPlaneStore:
|
||||
)
|
||||
self.conn.execute("CREATE INDEX IF NOT EXISTS ix_tasks_status ON tasks(status)")
|
||||
self.conn.execute("CREATE INDEX IF NOT EXISTS ix_tasks_parent ON tasks(parent_task_id)")
|
||||
self.conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS ix_tasks_scope_tenant_project "
|
||||
"ON tasks(scope_type, tenant_id, project_id)"
|
||||
)
|
||||
self.conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS ix_tasks_scope_owner "
|
||||
"ON tasks(scope_type, tenant_id, user_id, conversation_id, run_id)"
|
||||
)
|
||||
self.conn.execute("CREATE INDEX IF NOT EXISTS ix_workers_status ON workers(status)")
|
||||
self.conn.execute("CREATE INDEX IF NOT EXISTS ix_claims_task ON claims(task_id)")
|
||||
self.conn.execute(
|
||||
@@ -347,6 +377,42 @@ class TaskPlaneStore:
|
||||
raise NotFoundError("task")
|
||||
task_json = json.loads(row["task_json"])
|
||||
task_plan = TaskPlan.from_dict(task_json)
|
||||
stored_scope = tuple(
|
||||
row[field]
|
||||
for field in (
|
||||
"scope_type",
|
||||
"tenant_id",
|
||||
"user_id",
|
||||
"project_id",
|
||||
"conversation_id",
|
||||
"run_id",
|
||||
)
|
||||
)
|
||||
json_has_scope = any(
|
||||
field in task_json
|
||||
for field in (
|
||||
"scope_type",
|
||||
"tenant_id",
|
||||
"user_id",
|
||||
"project_id",
|
||||
"conversation_id",
|
||||
"run_id",
|
||||
)
|
||||
)
|
||||
if json_has_scope and stored_scope != task_scope_values(task_plan):
|
||||
raise ConflictError(
|
||||
"task scope persistence mismatch",
|
||||
code="TASK_SCOPE_PERSISTENCE_MISMATCH",
|
||||
)
|
||||
if not json_has_scope:
|
||||
(
|
||||
task_plan.scope_type,
|
||||
task_plan.tenant_id,
|
||||
task_plan.user_id,
|
||||
task_plan.project_id,
|
||||
task_plan.conversation_id,
|
||||
task_plan.run_id,
|
||||
) = stored_scope
|
||||
task_plan.parent_task_id = row["parent_task_id"]
|
||||
task_plan.correlation_id = row["correlation_id"]
|
||||
return TaskRecord(
|
||||
@@ -389,6 +455,7 @@ class TaskPlaneStore:
|
||||
task_id = task_plan.task_id or f"task-{uuid.uuid4()}"
|
||||
task_plan.task_id = task_id
|
||||
task_json = redact_value(task_plan.to_dict(redact=False))
|
||||
scope_values = task_scope_values(task_plan)
|
||||
now = _now()
|
||||
with self._transaction():
|
||||
if idempotency_key:
|
||||
@@ -397,7 +464,11 @@ class TaskPlaneStore:
|
||||
(idempotency_key,),
|
||||
).fetchone()
|
||||
if existing:
|
||||
if _canonical(json.loads(existing["task_json"])) != _canonical(task_json):
|
||||
existing_plan = TaskPlan.from_dict(
|
||||
json.loads(existing["task_json"])
|
||||
)
|
||||
existing_json = redact_value(existing_plan.to_dict(redact=False))
|
||||
if _canonical(existing_json) != _canonical(task_json):
|
||||
raise ConflictError("idempotency key payload mismatch")
|
||||
return self.get_task(existing["task_id"])
|
||||
root_task_id = task_id
|
||||
@@ -412,14 +483,18 @@ class TaskPlaneStore:
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO tasks
|
||||
(task_id, revision, status, task_json, task_type, root_task_id, parent_task_id, correlation_id, idempotency_key, created_at, updated_at, failure_reason)
|
||||
VALUES (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL)
|
||||
(task_id, revision, status, task_json, task_type,
|
||||
scope_type, tenant_id, user_id, project_id, conversation_id, run_id,
|
||||
root_task_id, parent_task_id, correlation_id, idempotency_key,
|
||||
created_at, updated_at, failure_reason)
|
||||
VALUES (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL)
|
||||
""",
|
||||
(
|
||||
task_id,
|
||||
"CREATED",
|
||||
json.dumps(task_json, ensure_ascii=False),
|
||||
task_plan.task_type,
|
||||
*scope_values,
|
||||
root_task_id,
|
||||
task_plan.parent_task_id,
|
||||
task_plan.correlation_id,
|
||||
@@ -451,21 +526,30 @@ class TaskPlaneStore:
|
||||
)
|
||||
return self.get_task(task.task_id)
|
||||
|
||||
def get_task(self, task_id: str) -> TaskRecord:
|
||||
def get_task(
|
||||
self, task_id: str, actor_context: TaskActorContext | None = None
|
||||
) -> TaskRecord:
|
||||
task_id = _validated_id(task_id, "task_id")
|
||||
predicate, scope_params = actor_scope_sql(actor_context)
|
||||
sql = "SELECT * FROM tasks WHERE task_id=?"
|
||||
if predicate:
|
||||
sql += f" AND ({predicate})"
|
||||
with self._lock:
|
||||
row = self.conn.execute(
|
||||
"SELECT * FROM tasks WHERE task_id=?", (task_id,)
|
||||
).fetchone()
|
||||
row = self.conn.execute(sql, (task_id, *scope_params)).fetchone()
|
||||
if row is None:
|
||||
raise NotFoundError(f"task {task_id} not found")
|
||||
return self._row_to_task(row)
|
||||
|
||||
def list_tasks(self) -> list[TaskRecord]:
|
||||
def list_tasks(
|
||||
self, actor_context: TaskActorContext | None = None
|
||||
) -> list[TaskRecord]:
|
||||
predicate, scope_params = actor_scope_sql(actor_context)
|
||||
sql = "SELECT * FROM tasks"
|
||||
if predicate:
|
||||
sql += f" WHERE ({predicate})"
|
||||
sql += " ORDER BY created_at DESC"
|
||||
with self._lock:
|
||||
rows = self.conn.execute(
|
||||
"SELECT * FROM tasks ORDER BY created_at DESC"
|
||||
).fetchall()
|
||||
rows = self.conn.execute(sql, scope_params).fetchall()
|
||||
return [self._row_to_task(row) for row in rows]
|
||||
|
||||
def get_task_budget(self, task_id: str) -> Budget:
|
||||
|
||||
@@ -0,0 +1,759 @@
|
||||
"""Scope-bound governance for traces, Tool/MCP calls, and credentials.
|
||||
|
||||
This module deliberately stores references and audit metadata only. Secret
|
||||
material remains in the configured credential provider boundary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from contextlib import contextmanager
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Mapping
|
||||
|
||||
|
||||
SCOPE_FIELDS = ("tenant_id", "user_id", "project_id", "conversation_id", "run_id")
|
||||
PRIVILEGED_CAPABILITIES = frozenset({"*", "admin", "system.admin", "governance.admin"})
|
||||
REFERENCE_RE = re.compile(r"^(systemd-credential|vault|secret-manager):[A-Za-z0-9][A-Za-z0-9._:/#-]{0,255}$")
|
||||
IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,179}$")
|
||||
|
||||
|
||||
class GovernanceError(RuntimeError):
|
||||
def __init__(self, code: str, *, field: str | None = None) -> None:
|
||||
super().__init__(code if field is None else f"{code}:{field}")
|
||||
self.code = code
|
||||
self.field = field
|
||||
|
||||
|
||||
def _clean(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _scope_tuple(value: Iterable[str] | str | None) -> tuple[str, ...]:
|
||||
if value is None:
|
||||
return ()
|
||||
if isinstance(value, str):
|
||||
value = value.split(",")
|
||||
return tuple(sorted({_clean(item) for item in value if _clean(item)}))
|
||||
|
||||
|
||||
def _canonical(value: Any) -> str:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
|
||||
|
||||
|
||||
def _sha(value: Any) -> str:
|
||||
return hashlib.sha256(_canonical(value).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _identifier(value: Any, code: str) -> str:
|
||||
text = _clean(value)
|
||||
if text is None or IDENTIFIER_RE.fullmatch(text) is None:
|
||||
raise GovernanceError(code)
|
||||
return text
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GovernanceScope:
|
||||
"""Explicit resource scope; system/global resources never get fake tenants."""
|
||||
|
||||
scope_type: str
|
||||
tenant_id: str | None = None
|
||||
user_id: str | None = None
|
||||
project_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
run_id: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "scope_type", (_clean(self.scope_type) or "").lower())
|
||||
for name in SCOPE_FIELDS:
|
||||
object.__setattr__(self, name, _clean(getattr(self, name)))
|
||||
if self.scope_type == "tenant":
|
||||
if self.tenant_id is None:
|
||||
raise GovernanceError("GOVERNANCE_TENANT_ID_REQUIRED")
|
||||
elif self.scope_type in {"system", "global"}:
|
||||
if any(getattr(self, name) is not None for name in SCOPE_FIELDS):
|
||||
raise GovernanceError("GOVERNANCE_NON_TENANT_SCOPE_INVALID")
|
||||
elif self.scope_type == "legacy":
|
||||
if any(getattr(self, name) is not None for name in SCOPE_FIELDS):
|
||||
raise GovernanceError("GOVERNANCE_LEGACY_SCOPE_INVALID")
|
||||
else:
|
||||
raise GovernanceError("GOVERNANCE_SCOPE_TYPE_INVALID")
|
||||
|
||||
@classmethod
|
||||
def from_value(cls, value: Any) -> "GovernanceScope":
|
||||
def get(name: str) -> Any:
|
||||
return value.get(name) if isinstance(value, Mapping) else getattr(value, name, None)
|
||||
|
||||
declared = _clean(get("scope_type"))
|
||||
tenant_id = _clean(get("tenant_id"))
|
||||
return cls(
|
||||
scope_type=declared or ("tenant" if tenant_id else "legacy"),
|
||||
tenant_id=tenant_id,
|
||||
user_id=get("user_id"),
|
||||
project_id=get("project_id"),
|
||||
conversation_id=get("conversation_id"),
|
||||
run_id=get("run_id"),
|
||||
)
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GovernanceActor:
|
||||
identity: str
|
||||
scope: GovernanceScope
|
||||
capabilities: tuple[str, ...] = ()
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "identity", _identifier(self.identity, "GOVERNANCE_ACTOR_INVALID"))
|
||||
object.__setattr__(self, "capabilities", _scope_tuple(self.capabilities))
|
||||
|
||||
@classmethod
|
||||
def from_value(cls, value: Any) -> "GovernanceActor":
|
||||
if isinstance(value, cls):
|
||||
return value
|
||||
if not isinstance(value, Mapping):
|
||||
raise GovernanceError("GOVERNANCE_ACTOR_INVALID")
|
||||
return cls(
|
||||
identity=value.get("identity") or value.get("actor_id") or value.get("user_id"),
|
||||
scope=GovernanceScope.from_value(value),
|
||||
capabilities=_scope_tuple(value.get("capabilities") or value.get("capability_scope")),
|
||||
)
|
||||
|
||||
@property
|
||||
def privileged(self) -> bool:
|
||||
return bool(set(self.capabilities) & PRIVILEGED_CAPABILITIES)
|
||||
|
||||
|
||||
def authorize_scope(actor: GovernanceActor | Mapping[str, Any], resource: GovernanceScope | Mapping[str, Any]) -> None:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
resource = resource if isinstance(resource, GovernanceScope) else GovernanceScope.from_value(resource)
|
||||
if actor.privileged:
|
||||
return
|
||||
if resource.scope_type in {"system", "global"}:
|
||||
raise GovernanceError("GOVERNANCE_SYSTEM_SCOPE_AUTHZ_DENY")
|
||||
if resource.scope_type == "legacy":
|
||||
if actor.scope.scope_type != "legacy":
|
||||
raise GovernanceError("GOVERNANCE_LEGACY_SCOPE_AUTHZ_DENY")
|
||||
return
|
||||
if actor.scope.scope_type != "tenant":
|
||||
raise GovernanceError("GOVERNANCE_SCOPE_MISMATCH_DENY", field="scope_type")
|
||||
for name in SCOPE_FIELDS:
|
||||
expected = getattr(resource, name)
|
||||
actual = getattr(actor.scope, name)
|
||||
if expected is not None and expected != actual:
|
||||
raise GovernanceError("GOVERNANCE_SCOPE_MISMATCH_DENY", field=name)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TraceContext:
|
||||
trace_id: str
|
||||
span_id: str
|
||||
scope: GovernanceScope
|
||||
parent_span_id: str | None = None
|
||||
correlation_id: str | None = None
|
||||
causation_id: str | None = None
|
||||
task_id: str | None = None
|
||||
claim_id: str | None = None
|
||||
message_id: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "trace_id", _identifier(self.trace_id, "TRACE_ID_INVALID"))
|
||||
object.__setattr__(self, "span_id", _identifier(self.span_id, "TRACE_SPAN_ID_INVALID"))
|
||||
for name in ("parent_span_id", "correlation_id", "causation_id", "task_id", "claim_id", "message_id"):
|
||||
value = _clean(getattr(self, name))
|
||||
if value is not None and IDENTIFIER_RE.fullmatch(value) is None:
|
||||
raise GovernanceError("TRACE_LINEAGE_IDENTIFIER_INVALID", field=name)
|
||||
object.__setattr__(self, name, value)
|
||||
|
||||
def child(self, *, span_id: str, actor: GovernanceActor | Mapping[str, Any], **lineage: Any) -> "TraceContext":
|
||||
authorize_scope(actor, self.scope)
|
||||
for name in ("trace_id", "task_id", "claim_id", "message_id", "correlation_id"):
|
||||
supplied = _clean(lineage.get(name))
|
||||
expected = getattr(self, name)
|
||||
if supplied is not None and expected is not None and supplied != expected:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONTINUITY", field=name)
|
||||
requested_scope = lineage.pop("scope", self.scope)
|
||||
requested_scope = requested_scope if isinstance(requested_scope, GovernanceScope) else GovernanceScope.from_value(requested_scope)
|
||||
if requested_scope != self.scope:
|
||||
raise GovernanceError("TRACE_SCOPE_MISMATCH_DENY")
|
||||
return TraceContext(
|
||||
trace_id=self.trace_id,
|
||||
span_id=span_id,
|
||||
parent_span_id=self.span_id,
|
||||
correlation_id=self.correlation_id,
|
||||
causation_id=_clean(lineage.get("causation_id")) or self.span_id,
|
||||
task_id=self.task_id,
|
||||
claim_id=self.claim_id,
|
||||
message_id=self.message_id,
|
||||
scope=self.scope,
|
||||
)
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"trace_id": self.trace_id,
|
||||
"span_id": self.span_id,
|
||||
"parent_span_id": self.parent_span_id,
|
||||
"correlation_id": self.correlation_id,
|
||||
"causation_id": self.causation_id,
|
||||
"task_id": self.task_id,
|
||||
"claim_id": self.claim_id,
|
||||
"message_id": self.message_id,
|
||||
**self.scope.as_dict(),
|
||||
}
|
||||
|
||||
|
||||
def bind_trace_lineage(
|
||||
*,
|
||||
actor: GovernanceActor | Mapping[str, Any],
|
||||
task: Any,
|
||||
message: Any,
|
||||
trace_id: str,
|
||||
span_id: str,
|
||||
parent: TraceContext | None = None,
|
||||
) -> TraceContext:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
task_scope = GovernanceScope.from_value(task)
|
||||
message_scope = GovernanceScope.from_value(message)
|
||||
if task_scope != message_scope:
|
||||
raise GovernanceError("TRACE_SCOPE_MISMATCH_DENY")
|
||||
authorize_scope(actor, task_scope)
|
||||
|
||||
def get(value: Any, name: str) -> Any:
|
||||
return value.get(name) if isinstance(value, Mapping) else getattr(value, name, None)
|
||||
|
||||
task_id = _clean(get(task, "task_id"))
|
||||
message_task = _clean(get(message, "task_id"))
|
||||
if task_id and message_task and task_id != message_task:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONTINUITY", field="task_id")
|
||||
context = TraceContext(
|
||||
trace_id=trace_id,
|
||||
span_id=span_id,
|
||||
parent_span_id=parent.span_id if parent else None,
|
||||
correlation_id=_clean(get(task, "correlation_id")) or _clean(get(message, "correlation_id")),
|
||||
causation_id=_clean(get(message, "causation_id")) or (parent.span_id if parent else None),
|
||||
task_id=task_id or message_task,
|
||||
claim_id=_clean(get(task, "claim_id")) or _clean(get(message, "claim_id")),
|
||||
message_id=_clean(get(message, "envelope_id")) or _clean(get(message, "id")),
|
||||
scope=task_scope,
|
||||
)
|
||||
if parent is not None:
|
||||
if parent.trace_id != context.trace_id or parent.scope != context.scope:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONTINUITY")
|
||||
if parent.task_id and context.task_id and parent.task_id != context.task_id:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONTINUITY", field="task_id")
|
||||
return context
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolDescriptor:
|
||||
name: str
|
||||
kind: str
|
||||
scope: GovernanceScope
|
||||
required_capabilities: tuple[str, ...] = ()
|
||||
credential_ref: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "name", _identifier(self.name, "TOOL_NAME_INVALID"))
|
||||
object.__setattr__(self, "kind", (_clean(self.kind) or "").lower())
|
||||
if self.kind not in {"tool", "mcp"}:
|
||||
raise GovernanceError("TOOL_KIND_INVALID")
|
||||
object.__setattr__(self, "required_capabilities", _scope_tuple(self.required_capabilities))
|
||||
if self.credential_ref is not None and REFERENCE_RE.fullmatch(self.credential_ref) is None:
|
||||
raise GovernanceError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CredentialGrant:
|
||||
grant_id: str
|
||||
credential_ref: str
|
||||
scope: GovernanceScope
|
||||
operations: tuple[str, ...]
|
||||
task_id: str | None = None
|
||||
tool_name: str | None = None
|
||||
expires_at: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "grant_id", _identifier(self.grant_id, "CREDENTIAL_GRANT_ID_INVALID"))
|
||||
if REFERENCE_RE.fullmatch(self.credential_ref or "") is None:
|
||||
raise GovernanceError("CREDENTIAL_SECRET_EXPOSURE_DENY")
|
||||
object.__setattr__(self, "operations", _scope_tuple(self.operations))
|
||||
if not self.operations:
|
||||
raise GovernanceError("CREDENTIAL_OPERATIONS_REQUIRED")
|
||||
if self.expires_at is not None:
|
||||
parsed = datetime.fromisoformat(self.expires_at.replace("Z", "+00:00"))
|
||||
if parsed.tzinfo is None:
|
||||
raise GovernanceError("CREDENTIAL_EXPIRY_INVALID")
|
||||
|
||||
def authorize(
|
||||
self,
|
||||
actor: GovernanceActor | Mapping[str, Any],
|
||||
*,
|
||||
operation: str,
|
||||
task_id: str | None = None,
|
||||
tool_name: str | None = None,
|
||||
now: datetime | None = None,
|
||||
) -> str:
|
||||
try:
|
||||
authorize_scope(actor, self.scope)
|
||||
except GovernanceError as exc:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field=exc.field) from exc
|
||||
if operation not in self.operations:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field="operation")
|
||||
if self.task_id is not None and self.task_id != _clean(task_id):
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field="task_id")
|
||||
if self.tool_name is not None and self.tool_name != _clean(tool_name):
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field="tool_name")
|
||||
if self.expires_at is not None:
|
||||
current = (now or datetime.now(timezone.utc)).astimezone(timezone.utc)
|
||||
expiry = datetime.fromisoformat(self.expires_at.replace("Z", "+00:00")).astimezone(timezone.utc)
|
||||
if current >= expiry:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field="expires_at")
|
||||
return self.credential_ref
|
||||
|
||||
def public(self) -> dict[str, Any]:
|
||||
return {
|
||||
"grant_id": self.grant_id,
|
||||
"credential_ref": self.credential_ref,
|
||||
"operations": list(self.operations),
|
||||
"task_id": self.task_id,
|
||||
"tool_name": self.tool_name,
|
||||
"expires_at": self.expires_at,
|
||||
**self.scope.as_dict(),
|
||||
}
|
||||
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS governance_schema_metadata (
|
||||
singleton INTEGER PRIMARY KEY CHECK(singleton=1),
|
||||
schema_version INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS governed_tools (
|
||||
name TEXT PRIMARY KEY,
|
||||
kind TEXT NOT NULL,
|
||||
scope_json TEXT NOT NULL,
|
||||
required_capabilities_json TEXT NOT NULL,
|
||||
credential_ref TEXT,
|
||||
descriptor_sha256 TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS governed_trace_spans (
|
||||
trace_id TEXT NOT NULL,
|
||||
span_id TEXT NOT NULL,
|
||||
parent_span_id TEXT,
|
||||
context_json TEXT NOT NULL,
|
||||
context_sha256 TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY(trace_id, span_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS governed_credential_grants (
|
||||
grant_id TEXT PRIMARY KEY,
|
||||
grant_json TEXT NOT NULL,
|
||||
grant_sha256 TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS governed_tool_invocations (
|
||||
invocation_key TEXT PRIMARY KEY,
|
||||
request_sha256 TEXT NOT NULL,
|
||||
tool_name TEXT NOT NULL,
|
||||
task_id TEXT NOT NULL,
|
||||
claim_id TEXT NOT NULL,
|
||||
trace_id TEXT NOT NULL,
|
||||
span_id TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
receipt_sha256 TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS governance_audit_evidence (
|
||||
sequence INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
evidence_id TEXT NOT NULL UNIQUE,
|
||||
action TEXT NOT NULL,
|
||||
actor_identity TEXT NOT NULL,
|
||||
trace_id TEXT,
|
||||
span_id TEXT,
|
||||
resource_id TEXT NOT NULL,
|
||||
outcome TEXT NOT NULL,
|
||||
details_json TEXT NOT NULL,
|
||||
previous_sha256 TEXT NOT NULL,
|
||||
evidence_sha256 TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
class GovernanceStore:
|
||||
def __init__(self, path: str | Path) -> None:
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._connection() as connection:
|
||||
has_metadata = connection.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='governance_schema_metadata'"
|
||||
).fetchone()
|
||||
if has_metadata:
|
||||
row = connection.execute(
|
||||
"SELECT schema_version FROM governance_schema_metadata WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is not None and int(row[0]) != 1:
|
||||
raise GovernanceError("GOVERNANCE_SCHEMA_VERSION_UNSUPPORTED")
|
||||
connection.executescript(_SCHEMA)
|
||||
row = connection.execute(
|
||||
"SELECT schema_version FROM governance_schema_metadata WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is None:
|
||||
connection.execute(
|
||||
"INSERT INTO governance_schema_metadata(singleton,schema_version) VALUES(1,1)"
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.path, timeout=30.0)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys=ON")
|
||||
connection.execute("PRAGMA busy_timeout=30000")
|
||||
return connection
|
||||
|
||||
@contextmanager
|
||||
def _connection(self):
|
||||
connection = self._connect()
|
||||
try:
|
||||
yield connection
|
||||
connection.commit()
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def _audit(
|
||||
self,
|
||||
connection: sqlite3.Connection,
|
||||
*,
|
||||
action: str,
|
||||
actor: GovernanceActor,
|
||||
resource_id: str,
|
||||
outcome: str,
|
||||
trace: TraceContext | None,
|
||||
details: Mapping[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
previous = connection.execute(
|
||||
"SELECT evidence_sha256 FROM governance_audit_evidence ORDER BY sequence DESC LIMIT 1"
|
||||
).fetchone()
|
||||
previous_sha = str(previous[0]) if previous else ""
|
||||
created_at = _now()
|
||||
material = {
|
||||
"action": action,
|
||||
"actor_identity": actor.identity,
|
||||
"trace_id": trace.trace_id if trace else None,
|
||||
"span_id": trace.span_id if trace else None,
|
||||
"resource_id": resource_id,
|
||||
"outcome": outcome,
|
||||
"details": dict(details),
|
||||
"previous_sha256": previous_sha,
|
||||
"created_at": created_at,
|
||||
}
|
||||
digest = _sha(material)
|
||||
evidence_id = "gev-" + digest[:32]
|
||||
connection.execute(
|
||||
"""INSERT INTO governance_audit_evidence
|
||||
(evidence_id,action,actor_identity,trace_id,span_id,resource_id,outcome,
|
||||
details_json,previous_sha256,evidence_sha256,created_at)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
|
||||
(
|
||||
evidence_id, action, actor.identity, material["trace_id"], material["span_id"],
|
||||
resource_id, outcome, _canonical(material["details"]), previous_sha, digest, created_at,
|
||||
),
|
||||
)
|
||||
return {"evidence_id": evidence_id, "evidence_sha256": digest, **material}
|
||||
|
||||
def record_trace(self, trace: TraceContext, actor: GovernanceActor | Mapping[str, Any]) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
authorize_scope(actor, trace.scope)
|
||||
context = trace.as_dict()
|
||||
digest = _sha(context)
|
||||
with self._connection() as connection:
|
||||
existing = connection.execute(
|
||||
"SELECT context_sha256 FROM governed_trace_spans WHERE trace_id=? AND span_id=?",
|
||||
(trace.trace_id, trace.span_id),
|
||||
).fetchone()
|
||||
if existing is not None and existing[0] != digest:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONFLICT")
|
||||
if existing is not None:
|
||||
return {"context": context, "context_sha256": digest, "idempotent": True}
|
||||
if trace.parent_span_id:
|
||||
parent = connection.execute(
|
||||
"SELECT context_json FROM governed_trace_spans WHERE trace_id=? AND span_id=?",
|
||||
(trace.trace_id, trace.parent_span_id),
|
||||
).fetchone()
|
||||
if parent is None:
|
||||
raise GovernanceError("TRACE_LINEAGE_CONTINUITY", field="parent_span_id")
|
||||
parent_context = json.loads(parent[0])
|
||||
if GovernanceScope.from_value(parent_context) != trace.scope:
|
||||
raise GovernanceError("TRACE_SCOPE_MISMATCH_DENY")
|
||||
connection.execute(
|
||||
"INSERT OR IGNORE INTO governed_trace_spans VALUES (?,?,?,?,?,?)",
|
||||
(trace.trace_id, trace.span_id, trace.parent_span_id, _canonical(context), digest, _now()),
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="TRACE_RECORDED", actor=actor,
|
||||
resource_id=f"{trace.trace_id}:{trace.span_id}", outcome="ALLOW", trace=trace,
|
||||
details={"context_sha256": digest},
|
||||
)
|
||||
return {"context": context, "context_sha256": digest, "evidence": evidence}
|
||||
|
||||
def register_tool(self, descriptor: ToolDescriptor, actor: GovernanceActor | Mapping[str, Any]) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
if not actor.privileged:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
public = {
|
||||
"name": descriptor.name,
|
||||
"kind": descriptor.kind,
|
||||
"required_capabilities": list(descriptor.required_capabilities),
|
||||
"credential_ref": descriptor.credential_ref,
|
||||
**descriptor.scope.as_dict(),
|
||||
}
|
||||
digest = _sha(public)
|
||||
with self._connection() as connection:
|
||||
existing = connection.execute(
|
||||
"SELECT descriptor_sha256 FROM governed_tools WHERE name=?", (descriptor.name,)
|
||||
).fetchone()
|
||||
if existing is not None and existing[0] != digest:
|
||||
raise GovernanceError("TOOL_DESCRIPTOR_CONFLICT")
|
||||
if existing is not None:
|
||||
return {"descriptor": public, "descriptor_sha256": digest, "idempotent": True}
|
||||
connection.execute(
|
||||
"INSERT OR IGNORE INTO governed_tools VALUES (?,?,?,?,?,?)",
|
||||
(descriptor.name, descriptor.kind, _canonical(descriptor.scope.as_dict()),
|
||||
_canonical(descriptor.required_capabilities), descriptor.credential_ref, digest),
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="TOOL_REGISTERED", actor=actor, resource_id=descriptor.name,
|
||||
outcome="ALLOW", trace=None, details={"descriptor_sha256": digest},
|
||||
)
|
||||
return {"descriptor": public, "descriptor_sha256": digest, "evidence": evidence}
|
||||
|
||||
def authorize_tool(
|
||||
self,
|
||||
*,
|
||||
actor: GovernanceActor | Mapping[str, Any],
|
||||
tool_name: str,
|
||||
trace: TraceContext,
|
||||
requested_scope: GovernanceScope | Mapping[str, Any],
|
||||
input_digest: str,
|
||||
invocation_key: str,
|
||||
) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
invocation_key = _identifier(invocation_key, "TOOL_INVOCATION_KEY_INVALID")
|
||||
if not re.fullmatch(r"[0-9a-f]{64}", input_digest):
|
||||
raise GovernanceError("TOOL_INPUT_DIGEST_INVALID")
|
||||
if trace.task_id is None or trace.claim_id is None:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY", field="active_claim")
|
||||
requested_scope = requested_scope if isinstance(requested_scope, GovernanceScope) else GovernanceScope.from_value(requested_scope)
|
||||
if requested_scope != trace.scope:
|
||||
raise GovernanceError("TOOL_SCOPE_PROPAGATION")
|
||||
authorize_scope(actor, requested_scope)
|
||||
with self._connection() as connection:
|
||||
persisted_trace = connection.execute(
|
||||
"SELECT context_sha256 FROM governed_trace_spans WHERE trace_id=? AND span_id=?",
|
||||
(trace.trace_id, trace.span_id),
|
||||
).fetchone()
|
||||
if persisted_trace is None or persisted_trace[0] != _sha(trace.as_dict()):
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY", field="trace")
|
||||
row = connection.execute("SELECT * FROM governed_tools WHERE name=?", (tool_name,)).fetchone()
|
||||
if row is None:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
descriptor_scope = GovernanceScope.from_value(json.loads(row["scope_json"]))
|
||||
if descriptor_scope.scope_type == "tenant" and descriptor_scope != requested_scope:
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
required = set(json.loads(row["required_capabilities_json"]))
|
||||
if not required.issubset(set(actor.capabilities)):
|
||||
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
||||
request_sha = _sha({
|
||||
"tool_name": tool_name, "input_digest": input_digest,
|
||||
"task_id": trace.task_id, "claim_id": trace.claim_id,
|
||||
"trace_id": trace.trace_id, "span_id": trace.span_id,
|
||||
"scope": requested_scope.as_dict(),
|
||||
})
|
||||
replay = connection.execute(
|
||||
"SELECT * FROM governed_tool_invocations WHERE invocation_key=?",
|
||||
(invocation_key,),
|
||||
).fetchone()
|
||||
if replay is not None:
|
||||
if replay["request_sha256"] != request_sha:
|
||||
raise GovernanceError("TOOL_INVOCATION_IDEMPOTENCY_CONFLICT")
|
||||
if replay["status"] != "VERIFIED":
|
||||
raise GovernanceError("TOOL_INVOCATION_RECONCILIATION_REQUIRED")
|
||||
return {
|
||||
"tool_name": tool_name, "trace_id": trace.trace_id,
|
||||
"span_id": trace.span_id, "invocation_key": invocation_key,
|
||||
"status": "VERIFIED", "receipt_sha256": replay["receipt_sha256"],
|
||||
"idempotent": True,
|
||||
}
|
||||
created_at = _now()
|
||||
connection.execute(
|
||||
"""INSERT INTO governed_tool_invocations
|
||||
(invocation_key,request_sha256,tool_name,task_id,claim_id,trace_id,span_id,
|
||||
status,receipt_sha256,created_at,updated_at)
|
||||
VALUES(?,?,?,?,?,?,?,?,?,?,?)""",
|
||||
(invocation_key, request_sha, tool_name, trace.task_id, trace.claim_id,
|
||||
trace.trace_id, trace.span_id, "AUTHORIZED", None, created_at, created_at),
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="TOOL_AUTHORIZED", actor=actor, resource_id=tool_name,
|
||||
outcome="ALLOW", trace=trace,
|
||||
details={"input_sha256": input_digest, "scope_sha256": _sha(requested_scope.as_dict()),
|
||||
"invocation_key": invocation_key},
|
||||
)
|
||||
return {
|
||||
"tool_name": tool_name,
|
||||
"trace_id": trace.trace_id,
|
||||
"span_id": trace.span_id,
|
||||
"invocation_key": invocation_key,
|
||||
"status": "AUTHORIZED",
|
||||
"credential_required": row["credential_ref"] is not None,
|
||||
"evidence": evidence,
|
||||
}
|
||||
|
||||
def record_tool_receipt(
|
||||
self,
|
||||
*,
|
||||
invocation_key: str,
|
||||
result_digest: str,
|
||||
actor: GovernanceActor | Mapping[str, Any],
|
||||
trace: TraceContext,
|
||||
) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
invocation_key = _identifier(invocation_key, "TOOL_INVOCATION_KEY_INVALID")
|
||||
if not re.fullmatch(r"[0-9a-f]{64}", result_digest):
|
||||
raise GovernanceError("TOOL_RECEIPT_DIGEST_INVALID")
|
||||
authorize_scope(actor, trace.scope)
|
||||
with self._connection() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT * FROM governed_tool_invocations WHERE invocation_key=?",
|
||||
(invocation_key,),
|
||||
).fetchone()
|
||||
if row is None or row["trace_id"] != trace.trace_id or row["span_id"] != trace.span_id:
|
||||
raise GovernanceError("TOOL_RECEIPT_BINDING_INVALID")
|
||||
if row["status"] == "VERIFIED":
|
||||
if row["receipt_sha256"] != result_digest:
|
||||
raise GovernanceError("TOOL_RECEIPT_CONFLICT")
|
||||
return {"invocation_key": invocation_key, "status": "VERIFIED", "idempotent": True}
|
||||
if row["status"] != "AUTHORIZED":
|
||||
raise GovernanceError("TOOL_INVOCATION_RECONCILIATION_REQUIRED")
|
||||
connection.execute(
|
||||
"UPDATE governed_tool_invocations SET status='VERIFIED',receipt_sha256=?,updated_at=? "
|
||||
"WHERE invocation_key=? AND status='AUTHORIZED'",
|
||||
(result_digest, _now(), invocation_key),
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="TOOL_RESULT_VERIFIED", actor=actor,
|
||||
resource_id=row["tool_name"], outcome="ALLOW", trace=trace,
|
||||
details={"invocation_key": invocation_key, "result_sha256": result_digest},
|
||||
)
|
||||
return {"invocation_key": invocation_key, "status": "VERIFIED", "evidence": evidence}
|
||||
|
||||
def store_grant(self, grant: CredentialGrant, actor: GovernanceActor | Mapping[str, Any]) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
if not actor.privileged:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY")
|
||||
public = grant.public()
|
||||
digest = _sha(public)
|
||||
with self._connection() as connection:
|
||||
existing = connection.execute(
|
||||
"SELECT grant_sha256 FROM governed_credential_grants WHERE grant_id=?", (grant.grant_id,)
|
||||
).fetchone()
|
||||
if existing is not None and existing[0] != digest:
|
||||
raise GovernanceError("CREDENTIAL_GRANT_CONFLICT")
|
||||
connection.execute(
|
||||
"INSERT OR IGNORE INTO governed_credential_grants VALUES (?,?,?,?)",
|
||||
(grant.grant_id, _canonical(public), digest, _now()),
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="CREDENTIAL_GRANT_STORED", actor=actor,
|
||||
resource_id=grant.grant_id, outcome="ALLOW", trace=None,
|
||||
details={"grant_sha256": digest},
|
||||
)
|
||||
return {"grant": public, "grant_sha256": digest, "evidence": evidence}
|
||||
|
||||
def authorize_credential(
|
||||
self,
|
||||
*,
|
||||
grant_id: str,
|
||||
actor: GovernanceActor | Mapping[str, Any],
|
||||
trace: TraceContext,
|
||||
operation: str,
|
||||
task_id: str | None,
|
||||
tool_name: str | None,
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
actor = GovernanceActor.from_value(actor)
|
||||
with self._connection() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT grant_json FROM governed_credential_grants WHERE grant_id=?", (grant_id,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY")
|
||||
raw = json.loads(row[0])
|
||||
grant = CredentialGrant(
|
||||
grant_id=raw["grant_id"], credential_ref=raw["credential_ref"],
|
||||
scope=GovernanceScope.from_value(raw), operations=tuple(raw["operations"]),
|
||||
task_id=raw.get("task_id"), tool_name=raw.get("tool_name"),
|
||||
expires_at=raw.get("expires_at"),
|
||||
)
|
||||
if trace.scope != grant.scope:
|
||||
raise GovernanceError("CREDENTIAL_SCOPE_DENY", field="trace_scope")
|
||||
reference = grant.authorize(
|
||||
actor, operation=operation, task_id=task_id, tool_name=tool_name, now=now
|
||||
)
|
||||
evidence = self._audit(
|
||||
connection, action="CREDENTIAL_AUTHORIZED", actor=actor,
|
||||
resource_id=grant_id, outcome="ALLOW", trace=trace,
|
||||
details={"operation": operation, "task_id": task_id, "tool_name": tool_name,
|
||||
"credential_ref_sha256": _sha(reference)},
|
||||
)
|
||||
return {"credential_ref": reference, "evidence": evidence}
|
||||
|
||||
def verify_audit_chain(self, required_actions: Iterable[str] = ()) -> bool:
|
||||
ordered_actions: list[str] = []
|
||||
previous = ""
|
||||
with self._connection() as connection:
|
||||
rows = connection.execute("SELECT * FROM governance_audit_evidence ORDER BY sequence").fetchall()
|
||||
for row in rows:
|
||||
details = json.loads(row["details_json"])
|
||||
material = {
|
||||
"action": row["action"], "actor_identity": row["actor_identity"],
|
||||
"trace_id": row["trace_id"], "span_id": row["span_id"],
|
||||
"resource_id": row["resource_id"], "outcome": row["outcome"],
|
||||
"details": details, "previous_sha256": row["previous_sha256"],
|
||||
"created_at": row["created_at"],
|
||||
}
|
||||
if row["previous_sha256"] != previous or row["evidence_sha256"] != _sha(material):
|
||||
return False
|
||||
previous = row["evidence_sha256"]
|
||||
ordered_actions.append(row["action"])
|
||||
position = 0
|
||||
for required in required_actions:
|
||||
try:
|
||||
position = ordered_actions.index(required, position) + 1
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CredentialGrant", "GovernanceActor", "GovernanceError", "GovernanceScope",
|
||||
"GovernanceStore", "ToolDescriptor", "TraceContext", "authorize_scope",
|
||||
"bind_trace_lineage",
|
||||
]
|
||||
@@ -0,0 +1,36 @@
|
||||
-- R4.2 P0-B durable Message/Handoff Fabric.
|
||||
-- Forward-only. The existing 20260707 message_outbox remains a compatibility source.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS message_handoff (
|
||||
envelope_id TEXT PRIMARY KEY,
|
||||
stream_key TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
target TEXT NOT NULL,
|
||||
message_type TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL,
|
||||
payload_sha256 TEXT NOT NULL,
|
||||
tenant_id TEXT,
|
||||
user_id TEXT,
|
||||
project_id TEXT,
|
||||
conversation_id TEXT,
|
||||
task_id TEXT,
|
||||
run_id TEXT,
|
||||
causation_id TEXT,
|
||||
sequence INTEGER NOT NULL CHECK(sequence >= 1),
|
||||
capability_scope_json TEXT NOT NULL,
|
||||
confidentiality_scope TEXT NOT NULL,
|
||||
idempotency_key TEXT NOT NULL,
|
||||
expires_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
delivered_at TEXT,
|
||||
acked_at TEXT,
|
||||
UNIQUE(stream_key, idempotency_key),
|
||||
UNIQUE(stream_key, sequence)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_message_handoff_target_status
|
||||
ON message_handoff(target, status, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_message_handoff_scope
|
||||
ON message_handoff(tenant_id, project_id, user_id, conversation_id, run_id);
|
||||
@@ -0,0 +1,37 @@
|
||||
[Unit]
|
||||
Description=AIOS owner-controlled Task execution-rebind API
|
||||
After=aios-agent-task-api.service
|
||||
Wants=aios-agent-task-api.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
User=aios-agent-task
|
||||
Group=aios-agent-task
|
||||
LoadCredential=agent_task_control_token:/etc/aios-gitops/agent-task-control.token
|
||||
EnvironmentFile=/etc/aios-gitops/agent-task-execution-rebind.env
|
||||
Environment=AIOS_AGENT_TASK_DB=/var/lib/aios-agent-task/state.sqlite3
|
||||
Environment=AIOS_AGENT_TASK_CONTROL_TOKEN_FILE=%d/agent_task_control_token
|
||||
Environment=AIOS_GOVERNED_WRITE_API_URL=http://127.0.0.1:18892
|
||||
Environment=AIOS_TASK_EXECUTION_REBIND_HOST=127.0.0.1
|
||||
Environment=AIOS_TASK_EXECUTION_REBIND_PORT=18942
|
||||
ExecStart=/usr/bin/python3 -B /opt/aios/current/active-runtime/app/aios_agent_task_execution_rebind.py
|
||||
Restart=on-failure
|
||||
RestartSec=2
|
||||
NoNewPrivileges=true
|
||||
PrivateTmp=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
ProtectKernelTunables=true
|
||||
ProtectKernelModules=true
|
||||
ProtectControlGroups=true
|
||||
LockPersonality=true
|
||||
RestrictSUIDSGID=true
|
||||
RestrictRealtime=true
|
||||
RestrictNamespaces=true
|
||||
MemoryDenyWriteExecute=true
|
||||
RestrictAddressFamilies=AF_UNIX AF_INET
|
||||
ReadWritePaths=/var/lib/aios-agent-task
|
||||
UMask=0077
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,36 @@
|
||||
[Unit]
|
||||
Description=AIOS owner-controlled Task retry-extension API
|
||||
After=aios-agent-task-api.service
|
||||
Wants=aios-agent-task-api.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
User=aios-agent-task
|
||||
Group=aios-agent-task
|
||||
LoadCredential=agent_task_control_token:/etc/aios-gitops/agent-task-control.token
|
||||
EnvironmentFile=/etc/aios-gitops/agent-task-retry-extension.env
|
||||
Environment=AIOS_AGENT_TASK_DB=/var/lib/aios-agent-task/state.sqlite3
|
||||
Environment=AIOS_AGENT_TASK_CONTROL_TOKEN_FILE=%d/agent_task_control_token
|
||||
Environment=AIOS_TASK_RETRY_HOST=127.0.0.1
|
||||
Environment=AIOS_TASK_RETRY_PORT=18941
|
||||
ExecStart=/usr/bin/python3 -B /opt/aios/current/active-runtime/app/aios_agent_task_retry_extension.py
|
||||
Restart=on-failure
|
||||
RestartSec=2
|
||||
NoNewPrivileges=true
|
||||
PrivateTmp=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
ProtectKernelTunables=true
|
||||
ProtectKernelModules=true
|
||||
ProtectControlGroups=true
|
||||
LockPersonality=true
|
||||
RestrictSUIDSGID=true
|
||||
RestrictRealtime=true
|
||||
RestrictNamespaces=true
|
||||
MemoryDenyWriteExecute=true
|
||||
RestrictAddressFamilies=AF_UNIX AF_INET
|
||||
ReadWritePaths=/var/lib/aios-agent-task
|
||||
UMask=0077
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -1,5 +0,0 @@
|
||||
[Service]
|
||||
LoadCredential=agent_task_control_token:/etc/aios-gitops/agent-task-control.token
|
||||
Environment=AIOS_AGENT_TASK_API_URL=http://127.0.0.1:18940
|
||||
Environment=AIOS_ONLINE_TASK_PROJECTION_ENABLED=1
|
||||
Environment=AIOS_AGENT_TASK_CONTROL_TOKEN_FILE=%d/agent_task_control_token
|
||||
@@ -14,11 +14,12 @@ Environment=AIOS_AGENT_TASK_API_URL=http://127.0.0.1:18940
|
||||
Environment=AIOS_GOVERNED_WRITE_API_URL=http://127.0.0.1:18892
|
||||
Environment=AIOS_AGENT_ID=online-engineering-worker
|
||||
Environment=AIOS_AGENT_TYPE=trusted-online-engineering-worker
|
||||
Environment=AIOS_AGENT_VERSION=2026-08-17.1-dynamic
|
||||
Environment=AIOS_AGENT_VERSION=2026-08-17.2-durable
|
||||
Environment=AIOS_WORKER_CAPABILITIES=git.read,git.test,governed.write.branch,governed.write.file
|
||||
Environment=AIOS_AUTHORITY_GIT_DIR=/data/aios-gitops/repo.git
|
||||
Environment=AIOS_ISOLATED_EXECUTOR_ROOT=/var/lib/aios-online-agent-executor
|
||||
Environment=AIOS_WORKER_STATE_ROOT=/var/lib/aios-online-agent-worker
|
||||
Environment=AIOS_WORKER_CHECKPOINT_DB=/var/lib/aios-online-agent-worker/checkpoints.sqlite3
|
||||
Environment=AIOS_WORKER_WORKSPACE_ROOT=/var/lib/aios-online-agent-worker/workspaces
|
||||
ExecStart=/usr/bin/python3 -B /opt/aios/current/active-runtime/app/aios_online_agent_worker.py
|
||||
Restart=always
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
[Unit]
|
||||
Description=AIOS trusted dynamic engineering consumer
|
||||
After=aios-agent-task-api.service aios-chatgpt-bridge-v2-staging.service
|
||||
Wants=aios-agent-task-api.service aios-chatgpt-bridge-v2-staging.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
User=aios-engineering-worker
|
||||
Group=aios-engineering-worker
|
||||
SupplementaryGroups=gitops aios-executor-exchange
|
||||
LoadCredentialEncrypted=agent_token:/etc/credstore.encrypted/aios-online-engineering-20260815/agent_token
|
||||
Environment=PYTHONUNBUFFERED=1
|
||||
Environment=AIOS_AGENT_TASK_API_URL=http://127.0.0.1:18940
|
||||
Environment=AIOS_GOVERNED_WRITE_API_URL=http://127.0.0.1:18892
|
||||
Environment=AIOS_AGENT_ID=online-engineering-consumer
|
||||
Environment=AIOS_AGENT_TYPE=trusted-online-engineering-worker
|
||||
Environment=AIOS_AGENT_VERSION=2026-08-17.1-dynamic
|
||||
Environment=AIOS_WORKER_CAPABILITIES=git.read,git.test,governed.write.branch,governed.write.file
|
||||
Environment=AIOS_AUTHORITY_GIT_DIR=/data/aios-gitops/repo.git
|
||||
Environment=AIOS_ISOLATED_EXECUTOR_ROOT=/var/lib/aios-online-agent-executor
|
||||
Environment=AIOS_WORKER_STATE_ROOT=/var/lib/aios-online-engineering-consumer
|
||||
Environment=AIOS_WORKER_WORKSPACE_ROOT=/var/lib/aios-online-engineering-consumer/workspaces
|
||||
ExecStart=/usr/bin/python3 -B /opt/aios/current/active-runtime/app/aios_online_agent_worker.py
|
||||
Restart=always
|
||||
RestartSec=5
|
||||
TimeoutStopSec=20
|
||||
NoNewPrivileges=true
|
||||
PrivateTmp=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
ProtectKernelTunables=true
|
||||
ProtectKernelModules=true
|
||||
ProtectControlGroups=true
|
||||
LockPersonality=true
|
||||
RestrictSUIDSGID=true
|
||||
RestrictRealtime=true
|
||||
RestrictNamespaces=true
|
||||
MemoryDenyWriteExecute=true
|
||||
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
|
||||
ReadOnlyPaths=/data/aios-gitops/repo.git
|
||||
ReadWritePaths=/var/lib/aios-online-engineering-consumer
|
||||
UMask=0027
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -790,91 +790,3 @@ paths:
|
||||
'409': {description: Lease, expected-head, idempotency, or repository conflict}
|
||||
'503': {description: Reconciliation required}
|
||||
x-openai-isConsequential: true
|
||||
|
||||
/v2/execution-envelope:
|
||||
post:
|
||||
operationId: projectOnlineAgentExecution
|
||||
summary: Project the canonical execution envelope for an existing governed task
|
||||
description: Returns the durable task binding and canonical source_capture or deterministic projection.
|
||||
security: [{bearerAuth: []}]
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
additionalProperties: false
|
||||
required: [task_id]
|
||||
properties:
|
||||
task_id: {type: string, minLength: 3}
|
||||
responses:
|
||||
'200':
|
||||
description: Canonical execution envelope
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
required: [schema, task_id, job_id, action, execution_digest]
|
||||
properties:
|
||||
schema: {type: string, enum: [aios.online-agent.execution.v2]}
|
||||
task_id: {type: string}
|
||||
job_id: {type: string}
|
||||
action: {type: string, enum: [source_capture, deterministic]}
|
||||
execution_digest: {type: string, pattern: '^[0-9a-f]{64}$'}
|
||||
execution: {type: object, additionalProperties: true}
|
||||
'400': {description: Missing or invalid task binding}
|
||||
'409': {description: Canonical projection conflict}
|
||||
x-openai-isConsequential: true
|
||||
/v2/execution-plan:
|
||||
post:
|
||||
operationId: registerOnlineAgentExecutionPlan
|
||||
summary: Register a deterministic plan for an existing governed task
|
||||
description: Registers a non-empty deterministic change set and bounded verification plan against the exact source head.
|
||||
security: [{bearerAuth: []}]
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
additionalProperties: false
|
||||
required: [schema, task_id, job_id, source_head_oid, generation, idempotency_key, execution]
|
||||
properties:
|
||||
schema: {type: string, enum: [aios.online-agent.plan.v1]}
|
||||
task_id: {type: string, minLength: 3}
|
||||
job_id: {type: string, minLength: 8}
|
||||
source_head_oid: {type: string, pattern: '^[0-9a-f]{40}$'}
|
||||
generation: {type: integer, minimum: 1}
|
||||
idempotency_key: {type: string, minLength: 8}
|
||||
expected_previous_plan_digest: {type: string, pattern: '^[0-9a-f]{64}$'}
|
||||
execution:
|
||||
type: object
|
||||
additionalProperties: false
|
||||
required: [action, change_set, verification]
|
||||
properties:
|
||||
action: {type: string, const: deterministic}
|
||||
adapter: {type: string}
|
||||
capability_scope: {type: array, items: {type: string}}
|
||||
change_set: {type: array, minItems: 1, maxItems: 128, items: {type: object, additionalProperties: true}}
|
||||
verification: {type: array, maxItems: 128, items: {type: object, additionalProperties: true}}
|
||||
responses:
|
||||
'200':
|
||||
description: Deterministic plan registered or returned idempotently
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
required: [schema, task_id, job_id, source_head_oid, generation, idempotency_key, plan_digest, execution]
|
||||
properties:
|
||||
schema: {type: string, enum: [aios-online-agent.plan-record.v1]}
|
||||
task_id: {type: string}
|
||||
job_id: {type: string}
|
||||
source_head_oid: {type: string, pattern: '^[0-9a-f]{40}$'}
|
||||
generation: {type: integer}
|
||||
idempotency_key: {type: string}
|
||||
plan_digest: {type: string, pattern: '^[0-9a-f]{64}$'}
|
||||
execution: {type: object, additionalProperties: true}
|
||||
execution_digest: {type: string, pattern: '^[0-9a-f]{64}$'}
|
||||
'400': {description: Invalid deterministic plan or source head}
|
||||
'409': {description: Binding, idempotency, or replacement CAS conflict}
|
||||
x-openai-isConsequential: true
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# R4.2 P0-A TaskPlane Scope Isolation Design
|
||||
|
||||
Status: approved candidate design. This document is not evidence of a protected-main merge or production deployment.
|
||||
|
||||
## Boundary
|
||||
|
||||
- Task: `r42-p0a-taskplane-scope-isolation-20260816`
|
||||
- Governed Job: `job-3fd743b51971c6007923699c1233f84c`
|
||||
- Repository: `aios/aios`
|
||||
- Base: `0eb6e1a864edfdaeef37b438991707a110e21558`
|
||||
- Target branch: `agent/r42-p0a-taskplane-scope-isolation-20260816`
|
||||
- Generation: `1`
|
||||
|
||||
The implementation reuses this exact boundary. It does not create a replacement Task, Job, branch, generation, or idempotency key.
|
||||
|
||||
## Model and invariants
|
||||
|
||||
`TaskPlan` gains first-class `scope_type`, `tenant_id`, `user_id`, `project_id`, `conversation_id`, and `run_id` fields. The fields participate in accepted-key validation, parsing, serialization, persistence, and authorization.
|
||||
|
||||
Historical rows without scope data normalize to `scope_type=legacy`; no tenant is inferred. Tenant-scoped tasks require a non-empty `tenant_id`. System/global resources use an explicit `scope_type=system` invariant and do not receive a fabricated tenant. A task may add user, project, conversation, and run ownership within its tenant.
|
||||
|
||||
## Trusted ActorContext
|
||||
|
||||
TaskPlane defines a small immutable `TaskActorContext` used only for authorization. The HTTP authenticator returns an identity. `PolicyEngine` resolves that identity through server-owned policy configuration into the effective context and capabilities. Request fields such as `tenant_id`, `user_id`, `group_id`, and `role` are declarations only and never expand authority.
|
||||
|
||||
Public create, list, and get paths always use the resolved context. Internal compatibility calls may omit a context and retain explicit legacy behavior, but the HTTP adapter never uses that path.
|
||||
|
||||
## Persistence
|
||||
|
||||
SQLite remains the current runtime. `_create_schema()` performs an idempotent forward migration that adds the six scope columns and indexes, then backfills only missing `scope_type` values to `legacy`.
|
||||
|
||||
New writes persist scope in both the canonical `task_json` representation and first-class columns. Reads accept historical JSON without scope, using legacy columns, but reject a new-row mismatch between JSON and columns. SQL list/get predicates are parameterized and applied before row hydration.
|
||||
|
||||
PostgreSQL remains a separate migration target and is not claimed by this change.
|
||||
|
||||
## Authorization flow
|
||||
|
||||
1. HTTP bearer authentication yields an identity.
|
||||
2. `PolicyEngine` verifies capability and resolves the identity to `TaskActorContext`.
|
||||
3. Create binds the requested scope to the trusted context and rejects mismatches.
|
||||
4. List/get pass the trusted context to `TaskPlaneStore`.
|
||||
5. Store applies tenant and ownership predicates in SQL; privileged capabilities may explicitly cross scope.
|
||||
|
||||
Tenant actors cannot read legacy tasks, another tenant, or mismatched owned resources. Legacy actors can only read legacy tasks. System actors require explicit privileged capability for cross-scope access.
|
||||
|
||||
## Failure behavior
|
||||
|
||||
Missing tenant scope, forged scope, unknown scope types, JSON/column divergence, and unauthorized list/get operations fail closed with canonical TaskPlane errors. Public get does not reveal whether an out-of-scope task exists.
|
||||
|
||||
## Verification
|
||||
|
||||
Tests exercise real model round trips, SQLite migration/restart, Service/Policy/Store wiring, tenant A versus tenant B list/get denial, ownership filtering, forged create payload rejection, legacy isolation, privileged access, and SQL parameterization. Existing TaskPlane API, RBAC, recovery/security, and security-contract suites remain mandatory.
|
||||
|
||||
The candidate is complete only after real `TASK_CLAIMED`, `TASK_HEARTBEAT`, exact WriterLease, governed commit, OID convergence, result submission, and `TASK_VERIFIED`. Protected main and production remain unchanged.
|
||||
@@ -1,134 +0,0 @@
|
||||
# Online Agent Automatic Claim Infrastructure
|
||||
|
||||
## Status and fixed boundary
|
||||
|
||||
This candidate targets `p0-claim-online-agent-auto-execution-20260817` on governed Job
|
||||
`job-2031a9e5682a4238f98e3ac7aee8a593`, base
|
||||
`a642b045775b9154a00ae78543037fbef4fb447b`, target branch
|
||||
`agent/p0-claim-online-agent-auto-execution-20260817`, generation 1.
|
||||
|
||||
It does not authorize protected-main, Production, tag, baseline, raw Git, Task DB, or
|
||||
credential mutations.
|
||||
|
||||
## Considered approaches
|
||||
|
||||
1. Require every caller to write `metadata.execution`. This is simple but preserves the
|
||||
current operational gap: a valid governed Job can still project an unexecutable Online
|
||||
Task.
|
||||
2. Let workers infer repository and WriterLease coordinates from natural-language Task
|
||||
fields. This removes the skip but violates fail-closed lineage and can bind the wrong Job,
|
||||
branch, base, or generation.
|
||||
3. Use a canonical server-side projector plus a task-agnostic worker resolver. This is the
|
||||
selected approach. The governed binding remains the only authority for repository, Job,
|
||||
branch, base, expected head, generation, and idempotency. Producers can persist the
|
||||
projected envelope when creating the Online Task. A worker may recover an omitted envelope
|
||||
only from the same canonical projector and durable binding; it never guesses coordinates.
|
||||
|
||||
## Components
|
||||
|
||||
### Canonical execution contract
|
||||
|
||||
`aios_online_agent_execution.py` owns schema `aios.online-agent.execution.v2`. It validates
|
||||
the complete governed lineage and supports two explicit actions:
|
||||
|
||||
- `deterministic`: a hash-bound change set, verification commands, and governed commit;
|
||||
- `source_capture`: a read-only first attempt that returns `NEEDS_REWORK_SOURCE_CAPTURE` and
|
||||
cannot acquire WriterLease or mutate Git.
|
||||
|
||||
Wrong task, Job, repository, base, expected head, branch, generation, capability scope, or
|
||||
governance policy fails closed with a canonical diagnostic.
|
||||
|
||||
### Governed Job to Online Task projection
|
||||
|
||||
`GovernedTaskBootstrap` exposes an exact durable binding lookup and builds Online Task
|
||||
projection data through the canonical contract. When the Bridge is configured with the
|
||||
official Agent Task control credential reference, governed Job creation submits that complete
|
||||
projection through the official Task HTTP API. The Bridge also exposes a loopback envelope
|
||||
projection operation used for recovery. The persisted projection includes the versioned
|
||||
envelope rather than relying on manual metadata edits.
|
||||
|
||||
Projection delivery is a durable outbox, not a best-effort tail call. The outbox record is
|
||||
written before HTTP delivery and a Bridge reconciler retries it after transient API failure or
|
||||
process restart. A response is accepted only when its complete execution envelope equals the
|
||||
submitted projection. Requests that explicitly require projection are rejected before governed
|
||||
bootstrap when the projection principal is unavailable.
|
||||
|
||||
Loopback is not treated as an identity boundary. Every POST route requires an exact
|
||||
runtime-principal bearer credential supplied by systemd credential projection; unauthenticated
|
||||
local callers are denied before request dispatch. The control-plane credential used to create
|
||||
Online Tasks remains a separate Bridge-only reference and is never projected into the worker.
|
||||
|
||||
For legacy Online Tasks with an omitted envelope, the generic worker asks the same Bridge
|
||||
operation for a read-only `source_capture` envelope. Projection succeeds only when the exact
|
||||
durable binding exists. An omitted envelope is therefore recoverable; an invalid envelope is
|
||||
never replaced or weakened.
|
||||
|
||||
Persisted envelopes are also reprojected through the Bridge before claim and compared by exact
|
||||
digest and Job lineage. A self-consistent client-generated envelope is therefore not authority.
|
||||
|
||||
### Dynamic worker and trusted consumer
|
||||
|
||||
The repository gains a deployable generic worker and systemd unit definitions as governed
|
||||
source files. The worker validates or
|
||||
resolves an envelope before claim, revalidates after atomic claim, starts heartbeat, and then
|
||||
executes the selected action. The trusted engineering consumer uses this same task-agnostic
|
||||
entrypoint. No compiled worker contains a Task ID, base OID, branch, generation, or plan path.
|
||||
|
||||
Skip observability distinguishes:
|
||||
|
||||
- `TASK_SKIPPED_PROJECTION_MISSING`;
|
||||
- `TASK_SKIPPED_ENVELOPE_INVALID`;
|
||||
- `TASK_SKIPPED_CAPABILITY_MISMATCH`;
|
||||
- `TASK_SKIPPED_POLICY_DENY`;
|
||||
- `TASK_POLL_FAILED` / worker unavailable.
|
||||
|
||||
## State flow
|
||||
|
||||
```text
|
||||
Governed Job BOUND
|
||||
-> canonical durable binding
|
||||
-> Online Task projection with execution.v2
|
||||
-> capability-filtered poll
|
||||
-> envelope and lineage validation
|
||||
-> atomic claim
|
||||
-> initial heartbeat/start
|
||||
-> refetch Task and re-resolve envelope
|
||||
-> compare execution digest and Job
|
||||
-> source capture OR deterministic governed execution
|
||||
-> result construction
|
||||
-> result submission
|
||||
-> external server-side verification gate
|
||||
```
|
||||
|
||||
Legacy omission uses `binding -> source_capture projection` between poll and validation. It
|
||||
does not write Task metadata and does not convert a malformed envelope into a valid one.
|
||||
|
||||
## Recovery and safety
|
||||
|
||||
- Projection is deterministic and digest-bound.
|
||||
- Governed Task creation is serialized by a cross-process, Task-ID-digest lock. Durable paths
|
||||
use collision-resistant exact Task keys while retaining read compatibility with verified
|
||||
legacy bindings.
|
||||
- A new Task cannot adopt an existing branch from caller-selected generation data. It must name
|
||||
an exact durable parent Task/Job whose single ACTIVE lease owns the branch head; generation is
|
||||
then derived from that lease.
|
||||
- Claim is attempted only after validation; claimed payload drift is denied.
|
||||
- Source capture has no WriterLease or commit path.
|
||||
- Deterministic execution retains detached worktree, exact WriterLease, governed commit,
|
||||
durable current-head CAS, and receipt convergence.
|
||||
- Local deterministic recovery is keyed by Task and execution digest; live acceptance must
|
||||
independently reconverge Job, lease, receipt, and branch OIDs before verification.
|
||||
- An ambiguous commit transport or server outcome becomes `RECONCILIATION_REQUIRED`; the
|
||||
worker does not submit a canonical failure until exact Job receipt and head evidence is
|
||||
observed.
|
||||
- A completed read-only `source_capture` state is ignored when a later attempt has a changed
|
||||
digest and no governed state exists; all other digest drift is denied.
|
||||
- Protected refs are never worker targets.
|
||||
|
||||
## Acceptance
|
||||
|
||||
The implementation must pass the named projection, schema, claimability, dynamic-binding,
|
||||
wrong-lineage, capability, P0-A, and P0-B tests. Runtime acceptance additionally requires a
|
||||
real Online Task to be claimed and heartbeated without per-task worker or systemd plan
|
||||
injection, plus governed receipt/OID/CAS convergence for this implementation Task.
|
||||
Unit tests and deployable unit definitions do not by themselves satisfy runtime acceptance.
|
||||
@@ -1,56 +0,0 @@
|
||||
# P0-CLAIM Scheme A follow-up
|
||||
|
||||
This follow-up continues the existing candidate `71523aef6e48b1e2fa516c456ca189b7a039104b`.
|
||||
It does not create a replacement Task, Job, branch or generation.
|
||||
|
||||
## Retry extension
|
||||
|
||||
A versioned Owner-controlled loopback API provides
|
||||
`POST /v1/tasks/{task_id}/retry-extension`.
|
||||
|
||||
The request binds FAILED status, revision, exhausted attempt/max_attempts, next max attempts,
|
||||
previous verification SHA256, Owner authorization ID and idempotency key. The service performs one
|
||||
SQLite `BEGIN IMMEDIATE` transaction containing the CAS update, durable retry-extension record and
|
||||
Task event. Each authorization extends the boundary by exactly one attempt.
|
||||
|
||||
## Same-Job continuation
|
||||
|
||||
The durable GovernedTask binding remains the lineage anchor. After the first CREATE_NEW commit,
|
||||
`execution_binding()` may normalize that SAME Job to `CONTINUE_EXISTING` only when exactly one
|
||||
ACTIVE WriterLease for the Job owns the current authority branch head, generation and lineage
|
||||
anchor. The lineage `base_oid` stays unchanged; only `expected_branch_oid` becomes the current head.
|
||||
|
||||
GovernedTask creation is also serialized by a branch-digest lock, preventing sibling continuation
|
||||
tasks from racing generation admission.
|
||||
|
||||
## Projection and recovery
|
||||
|
||||
Natural-language source signals are normalized server-side to at most 16 unique 512-byte values.
|
||||
Malformed persisted envelopes still fail closed.
|
||||
|
||||
Fresh commits and recovery require a COMPLETED Job, `remote_verified=true`, exact
|
||||
authority/API/remote receipt OIDs, and exactly one ACTIVE WriterLease whose durable current head
|
||||
matches the commit. Delayed local ref visibility becomes `RECONCILIATION_REQUIRED`, not a terminal
|
||||
canonical failure.
|
||||
|
||||
## Credential isolation
|
||||
|
||||
The credentialed Coordinator performs only Task/Bridge coordination and governed commit. All
|
||||
deterministic repository modification and test commands run in
|
||||
`aios-online-agent-isolated-executor` under a separate system user with:
|
||||
|
||||
- no Agent/Owner credential projection;
|
||||
- no gitops supplementary group;
|
||||
- `PrivateNetwork=true`;
|
||||
- `RestrictAddressFamilies=AF_UNIX AF_INET`, with `IPAddressDeny=any` and
|
||||
`IPAddressAllow=localhost`, so regression tests may use an isolated loopback
|
||||
HTTP server without gaining host or external network reachability;
|
||||
- `InaccessiblePaths=/run/credentials`.
|
||||
|
||||
The executor receives a hash-bound source archive and deterministic change set through a filesystem
|
||||
exchange and returns only changed file bytes/hashes plus command evidence. The Coordinator validates
|
||||
that response before materializing exact allowed files and calling the governed Bridge.
|
||||
|
||||
Codex/Hermes mutation adapters remain disabled until a separately audited uncredentialed planner
|
||||
backend exists. Under the current operating model, implementation code is authored before Online
|
||||
Agent execution and is carried as a deterministic plan.
|
||||
@@ -0,0 +1,90 @@
|
||||
# R4.2 P0-C Trace / Tool-MCP / Credential Governance
|
||||
|
||||
## Status boundary
|
||||
|
||||
This change is a candidate implementation on the governed P0-C branch. It does
|
||||
not promote protected `main`, deploy Production, or assert that all R4.2 target
|
||||
architecture is complete.
|
||||
|
||||
## Scope model
|
||||
|
||||
Governance resources use an explicit `scope_type`:
|
||||
|
||||
- `tenant` requires `tenant_id` and may narrow to user, project,
|
||||
conversation, and run;
|
||||
- `system` and `global` prohibit tenant ownership fields;
|
||||
- `legacy` preserves unscoped historical records without inventing ownership.
|
||||
|
||||
System and global tools therefore never receive a synthetic tenant merely to
|
||||
satisfy a tenant-oriented schema.
|
||||
|
||||
## Trace lineage
|
||||
|
||||
`TaskPlan` and `MessageEnvelope` scope is bound to a `TraceContext`. The trace
|
||||
records `trace_id`, `span_id`, parent span, correlation, causation, task,
|
||||
message, and the canonical P0-A scope. Child spans must keep the parent's trace,
|
||||
task, and scope. Privileged visibility does not permit cross-scope trace
|
||||
splicing.
|
||||
|
||||
`MessageEnvelope.correlation_id` is persisted and identity-bound. Existing
|
||||
handoff databases are upgraded in place with an idempotent column migration;
|
||||
legacy rows retain their prior event IDs because a null correlation field is
|
||||
not added to the historical identity material.
|
||||
|
||||
Observability keeps its legacy presentation fields, but new governed records
|
||||
also carry immutable trace and scope fields. Scoped records require a trusted
|
||||
server actor for queries. A repeated execution ID is an idempotent replay only
|
||||
when the complete record matches; conflicting overwrites are denied.
|
||||
|
||||
## Tool and MCP governance
|
||||
|
||||
`ToolDescriptor` binds a tool or MCP server to:
|
||||
|
||||
- an explicit resource scope;
|
||||
- required actor capabilities;
|
||||
- an optional credential reference.
|
||||
|
||||
Invocation authorization requires exact propagation of the trace scope and a
|
||||
capability superset. Unknown tools, cross-tenant calls, scope substitution, and
|
||||
non-privileged access to system/global resources fail closed. Every successful
|
||||
registration and invocation appends hash-chained audit evidence.
|
||||
|
||||
## Credential governance
|
||||
|
||||
Credentials are represented only by references whose schemes are
|
||||
`systemd-credential`, `vault`, or `secret-manager`. Raw bearer tokens, passwords,
|
||||
and arbitrary values are rejected. Grants are bound to scope, operation, task,
|
||||
tool, and expiry. Public/audit records contain a reference hash where disclosure
|
||||
is unnecessary; secret material is never persisted by this module.
|
||||
|
||||
The existing runtime credential broker remains backward compatible for system
|
||||
leases, adds explicit tenant/system/global invariants, and rejects providers
|
||||
that attempt to return raw token/secret/password environment values instead of
|
||||
file/reference indirection.
|
||||
|
||||
## Durable evidence
|
||||
|
||||
The SQLite governance store is restart-safe and idempotently initializes four
|
||||
tables: governed tools, trace spans, credential grants, and audit evidence.
|
||||
Audit entries form a SHA-256 chain. Acceptance requires the chain to verify and
|
||||
contain the required actions; missing evidence is not treated as success.
|
||||
|
||||
## Acceptance tests
|
||||
|
||||
The candidate must pass:
|
||||
|
||||
- `TRACE_LINEAGE_CONTINUITY`;
|
||||
- `TRACE_SCOPE_MISMATCH_DENY`;
|
||||
- `TOOL_MCP_AUTHZ_DENY`;
|
||||
- `TOOL_SCOPE_PROPAGATION`;
|
||||
- `CREDENTIAL_SCOPE_DENY`;
|
||||
- `CREDENTIAL_SECRET_EXPOSURE_DENY`;
|
||||
- `AUDIT_EVIDENCE_COMPLETENESS`;
|
||||
- P0-A scope and P0-B handoff regression suites;
|
||||
- Python compilation, `git diff --check`, and a diff-scoped secret scan.
|
||||
|
||||
The repository-wide baseline currently has seven unrelated failures in the
|
||||
isolated host environment (five stale continuation-contract fixtures and two
|
||||
host-permission fixtures). The same seven failures reproduce at the untouched
|
||||
`a642b045775b9154a00ae78543037fbef4fb447b` base; they are not suppressed or
|
||||
modified by P0-C.
|
||||
@@ -0,0 +1,75 @@
|
||||
# R4.2 P1 — Durable Online Agent Execution
|
||||
|
||||
Status: candidate implementation. This document does not imply protected-main promotion or Production deployment.
|
||||
|
||||
## State model
|
||||
|
||||
The generic worker uses two local durable layers:
|
||||
|
||||
1. `tasks/<task_id>.json`: atomic fast-state cache for compatibility and quick restart.
|
||||
2. `checkpoints.sqlite3`: append-only attempt/phase journal with `synchronous=FULL`.
|
||||
|
||||
Checkpoint identity is bound to:
|
||||
|
||||
- task_id
|
||||
- attempt
|
||||
- execution_digest
|
||||
- job_id
|
||||
- repository
|
||||
- base_oid
|
||||
- expected_branch_oid
|
||||
- target_branch
|
||||
- generation
|
||||
|
||||
A checkpoint from another repository/branch/generation is never reused.
|
||||
|
||||
Phases are monotonic within an attempt:
|
||||
|
||||
`claimed → workspace_prepared → verification_passed → source_capture_done|commit_done → result_pending → result_submitted`
|
||||
|
||||
`reconciliation_required` is an explicit recovery phase.
|
||||
|
||||
## Attempt-aware result idempotency
|
||||
|
||||
Result submission idempotency keys include the official Task attempt and the execution digest. The same
|
||||
attempt reuses the same key after a timeout/restart; a later official attempt cannot collide with an
|
||||
earlier result.
|
||||
|
||||
Before submitting a result the worker durably writes `result_pending`, including the result hash and
|
||||
idempotency key. After a successful or exact-idempotency-reconciled API response it writes
|
||||
`result_submitted`.
|
||||
|
||||
A 409 is accepted only if the Task API already contains the exact same result body and the Task is
|
||||
VERIFYING/SUCCEEDED. A conflicting result remains a hard failure.
|
||||
|
||||
## Resume rules
|
||||
|
||||
- `source_capture_done` and `result_pending` may be reused across later official attempts only when the
|
||||
execution/scope binding is identical. This avoids repeating read-only source capture.
|
||||
- A deterministic task that crashed before governed commit replays from a new clean detached worktree.
|
||||
This is safe because there is no external side effect yet.
|
||||
- `commit_done` is never recommitted. Recovery requires exact Job receipt OID convergence plus exactly
|
||||
one ACTIVE WriterLease whose `current_head_oid` equals the governed commit.
|
||||
- delayed Git visibility remains `RECONCILIATION_REQUIRED`, preserving an explicit recovery entry.
|
||||
|
||||
## Crash cleanup
|
||||
|
||||
On worker startup, stale `task-*` workspaces under the configured workspace root are removed before
|
||||
polling begins. The path boundary is verified before deletion.
|
||||
|
||||
## Contracts
|
||||
|
||||
P1 gates include:
|
||||
|
||||
- CHECKPOINT_DURABLE_WRITE
|
||||
- RESUME_FROM_CHECKPOINT
|
||||
- RETRY_IDEMPOTENT_NO_DUPLICATE_SIDE_EFFECT
|
||||
- CRASH_RECOVERY_CONVERGENCE
|
||||
- WORKER_RESTART_RESUME
|
||||
- CHECKPOINT_SCOPE_ISOLATION
|
||||
- RESULT_SUBMISSION_REPLAY_SAFE
|
||||
- WORKSPACE_CLEANUP_RECOVERY
|
||||
- LEASE_RENEWAL_RECOVERY
|
||||
|
||||
WriterLease and governed Job authority remain in the existing Bridge/GitOps durable stores; the worker
|
||||
checkpoint DB is not a competing authority for Git state.
|
||||
@@ -0,0 +1,18 @@
|
||||
# R4.2 P2 Agent Fleet / Multi-Agent Scheduling / Fault Isolation
|
||||
|
||||
P2 adds a durable fleet decision layer inside TaskPlane without replacing TaskPlane claim/lease
|
||||
ownership. Worker registration/heartbeat remains authoritative for liveness. TaskPlane claim CAS
|
||||
remains authoritative for ownership. Probe-before-reassignment remains mandatory after an expired
|
||||
lease. P1 checkpoint/result replay remains authoritative after worker restart.
|
||||
|
||||
`FleetScheduler` adds durable worker capacity slots, weighted scope fairness, starvation aging,
|
||||
affinity/anti-affinity labels, worker circuit breakers, per-task/worker quarantine and idempotent
|
||||
dispatch reservations. Fleet state is colocated in the existing TaskPlane SQLite database and is
|
||||
created idempotently; no new daemon or database is introduced.
|
||||
|
||||
A fleet reservation is only a scheduling decision. It does not claim a Task. Callers must perform
|
||||
the normal TaskPlane claim and may then bind the returned claim ID to the reservation. This preserves
|
||||
the existing ownership and recovery invariants.
|
||||
|
||||
Failure recovery never bypasses probe evidence. When TaskPlane recovery places work back in a
|
||||
claimable state, P2 may quarantine the failed worker and prefer another worker/failure domain.
|
||||
@@ -0,0 +1,79 @@
|
||||
# R4.2 P2 Compact Interface Map Runtime
|
||||
|
||||
## Purpose
|
||||
|
||||
The persistent Online Agent already supports the canonical read-only `source_capture` action, but
|
||||
its `SourceCaptureExecutor` always returns an `aios.engineering.source-snapshot.v2` tar.gz+base64
|
||||
payload. That is correct for implementation source transfer, but too large for control-plane
|
||||
consumers that only need a bounded implementation map.
|
||||
|
||||
P2 adds a new **adapter**, not a new execution action:
|
||||
|
||||
```text
|
||||
action=source_capture
|
||||
adapter=compact_interface_map
|
||||
```
|
||||
|
||||
This preserves the existing execution-v2 schema and Bridge policy:
|
||||
|
||||
- read-only repository;
|
||||
- no change set;
|
||||
- no verification commands;
|
||||
- `writer_lease=false`;
|
||||
- no branch or commit;
|
||||
- existing capability and execution-digest validation.
|
||||
|
||||
## Compact map contract
|
||||
|
||||
`aios_online_agent_compact_map.py` scans only approved repository roots and source suffixes. It
|
||||
returns:
|
||||
|
||||
- at most 40 relevant modules;
|
||||
- at most 80 selected Python symbols in total;
|
||||
- at most 40 test symbols;
|
||||
- at most 25 excerpts;
|
||||
- at most 1200 characters per excerpt;
|
||||
- SQL table names;
|
||||
- systemd unit names;
|
||||
- current invariants, failure windows and reuse points;
|
||||
- P2 fleet contract-to-surface mappings.
|
||||
|
||||
The result is hard-bounded to 120 KiB and contains no source archive or `payload_b64`.
|
||||
|
||||
Source excerpts apply a conservative sensitive-pattern redaction before result construction.
|
||||
|
||||
## P2 contracts
|
||||
|
||||
The map supplies evidence surfaces for:
|
||||
|
||||
- `FLEET_AGENT_REGISTRATION`
|
||||
- `FLEET_CAPACITY_SCHEDULING`
|
||||
- `FLEET_FAIRNESS`
|
||||
- `FLEET_FAULT_ISOLATION`
|
||||
- `FLEET_WORKER_REASSIGNMENT`
|
||||
- `FLEET_SCOPE_ISOLATION`
|
||||
- `FLEET_IDEMPOTENT_DISPATCH`
|
||||
- `FLEET_FAILURE_RECOVERY`
|
||||
|
||||
Missing first-class scheduling/fleet primitives are reported as gaps rather than invented.
|
||||
|
||||
## Result
|
||||
|
||||
When the adapter is selected, the worker returns:
|
||||
|
||||
```text
|
||||
status=SUCCEEDED
|
||||
classification=P2_INTERFACE_MAP_CAPTURED
|
||||
execution_mode=compact_interface_map
|
||||
writer_lease=NONE
|
||||
```
|
||||
|
||||
The existing `source_capture_done` / P1 checkpoint/retry path is reused, so the compact analysis
|
||||
inherits attempt-aware result idempotency and restart-safe result submission.
|
||||
|
||||
## Governance
|
||||
|
||||
This capability is runtime infrastructure only. The implementation candidate may use a governed
|
||||
branch/WriterLease for its own code change, but any task executed through `compact_interface_map`
|
||||
is read-only and must not create a WriterLease, branch, commit, protected-ref write, or Production
|
||||
side effect.
|
||||
@@ -0,0 +1,17 @@
|
||||
# R4.2 P2-B Fleet Order Preservation Repair
|
||||
|
||||
P2-C live acceptance exposed a real integration defect in the verified P2-B candidate. The fleet
|
||||
scheduler correctly produced durable cross-scope fairness order, but the integration passed the
|
||||
entire ranked list to the legacy scheduler. The legacy scheduler then applied its own priority /
|
||||
created_at / task_id sort and silently discarded fleet ordering.
|
||||
|
||||
The repair preserves the legacy scheduler as the fail-closed security oracle while changing how it
|
||||
is called: one fleet-ranked candidate is presented at a time. A candidate rejected by legacy
|
||||
security/policy advances to the next fleet candidate. A candidate returned by the oracle must have
|
||||
the same task ID as the candidate presented, otherwise execution fails closed.
|
||||
|
||||
Task ownership remains unchanged: the original schedule_for_worker/claim_task CAS remains
|
||||
authoritative, and a fleet reservation is still only bound after the claim succeeds.
|
||||
|
||||
This repair does not modify FleetScheduler state semantics, capacity, WriterLease, checkpoint,
|
||||
protected refs or Production behavior.
|
||||
@@ -0,0 +1,23 @@
|
||||
# R4.2 P2-B Fleet Runtime Integration
|
||||
|
||||
P2-B connects the verified P2 fleet primitives to the existing TaskPlane pull scheduler without
|
||||
replacing its security or ownership logic.
|
||||
|
||||
The integration uses two stages. `FleetScheduler.rank_for_worker()` performs only read-only
|
||||
capacity/fairness/fault-isolation ranking. The original `TaskScheduler.next_queued_task()` is then
|
||||
run over that ordered list and remains the fail-closed security oracle for worker liveness,
|
||||
capabilities, data classification, execution zone and other existing policy checks. Only a task
|
||||
accepted by the original scheduler receives a short durable fleet reservation.
|
||||
|
||||
The original `TaskScheduler.schedule_for_worker()` remains the sole claim path and continues to call
|
||||
the existing TaskPlane `claim_task` CAS. After a successful claim, the claim ID is bound to the
|
||||
fleet reservation. A claim error releases the reservation. A fleet reservation is never treated as
|
||||
task ownership.
|
||||
|
||||
P2-B also makes the reservation capacity check atomic inside the fleet SQLite transaction so two
|
||||
concurrent pullers cannot oversubscribe the same worker slot. Circuit and task/worker quarantine
|
||||
state are enforced on the real pull path. Explicit execution-failure signals can be recorded through
|
||||
the integrated scheduler's failure/success/quarantine methods; policy/claim failures are not
|
||||
automatically mislabeled as worker crashes.
|
||||
|
||||
No new daemon, database or production write path is introduced.
|
||||
@@ -0,0 +1,18 @@
|
||||
# Generic Compact Interface-Map Profiles
|
||||
|
||||
The compact interface-map runtime is a reusable read-only control-plane capability. It keeps the
|
||||
canonical `source_capture` action and `compact_interface_map` adapter, but selects a bounded semantic
|
||||
profile from the task's source signals.
|
||||
|
||||
The existing `fleet` profile remains backward compatible and returns
|
||||
`P2_INTERFACE_MAP_CAPTURED` with the `FLEET_*` contracts.
|
||||
|
||||
The new `model_governance` profile is selected when model/provider/routing/budget/evaluation terms
|
||||
are present. It returns `MODEL_GOVERNANCE_INTERFACE_MAP_CAPTURED` and model-specific contracts for
|
||||
provider registry, selection policy, version pinning, budget accounting, fallback, scope isolation,
|
||||
credential boundary, trace/audit, evaluation gates and routing resilience.
|
||||
|
||||
Profile and classification are emitted by the compact-map module and revalidated by the worker
|
||||
before result submission. Unknown or mismatched profile/classification pairs fail closed.
|
||||
|
||||
No new Bridge action, daemon, database, WriterLease, mutation path or credential surface is added.
|
||||
@@ -0,0 +1,36 @@
|
||||
# Owner-Controlled Same-Binding Execution Rebind
|
||||
|
||||
## Purpose
|
||||
|
||||
`aios_agent_task_execution_rebind.py` closes a projection/plan race without replacing the Task,
|
||||
Governed Job, target branch, generation, or idempotency lineage.
|
||||
|
||||
The only supported transition is:
|
||||
|
||||
```text
|
||||
official NEEDS_REWORK
|
||||
→ Task QUEUED with no active lease
|
||||
→ persisted source_capture envelope
|
||||
→ canonical Bridge resolves an already-registered deterministic plan
|
||||
→ CAS-bound metadata.execution + execution_projection replacement
|
||||
→ audit event
|
||||
→ normal worker claim
|
||||
```
|
||||
|
||||
The service rejects arbitrary metadata patches. It derives the new envelope from canonical Bridge
|
||||
`/v2/execution-envelope`, validates every governed binding coordinate against the persisted envelope,
|
||||
requires `source_capture → deterministic`, requires exact old/new execution digests supplied by the
|
||||
Owner recovery request, and increments the Task revision under `BEGIN IMMEDIATE`.
|
||||
|
||||
It records `TASK_EXECUTION_REBIND_AUTHORIZED` and an idempotent `task_execution_rebinds` receipt.
|
||||
No Task/Job/branch replacement is created and the worker still performs the real claim/heartbeat,
|
||||
WriterLease and governed commit.
|
||||
|
||||
The endpoint is loopback-only:
|
||||
|
||||
`POST /v1/tasks/{task_id}/execution-rebind`
|
||||
|
||||
and requires the existing Owner control credential plus
|
||||
`AIOS_AGENT_TASK_EXECUTION_REBIND_AUTHORIZATION_ID`.
|
||||
|
||||
Direct/manual Task database edits remain forbidden; the service is the audited mutation boundary.
|
||||
@@ -0,0 +1,9 @@
|
||||
# R4.2 P0-A TaskPlane Scope Isolation
|
||||
|
||||
Status: candidate implementation. This is not proof of protected-main promotion or production deployment.
|
||||
|
||||
Task scope is now represented by `scope_type`, `tenant_id`, `user_id`, `project_id`, `conversation_id`, and `run_id` in the model and SQLite persistence. Historical rows become `legacy`; their tenant is never guessed. Tenant tasks require `tenant_id`. System resources use `scope_type=system` without a fabricated tenant.
|
||||
|
||||
The HTTP adapter uses the authenticated identity. `PolicyEngine` resolves that identity through server-owned `identity_actor_contexts`; request payload fields do not grant scope. Public create binds the trusted scope, and public list/get apply parameterized scope predicates in `TaskPlaneStore`.
|
||||
|
||||
SQLite is the current runtime database. PostgreSQL remains a separate migration and runtime-acceptance gate.
|
||||
@@ -0,0 +1,38 @@
|
||||
# R4.2 P0-B — Durable Message / Handoff Fabric
|
||||
|
||||
Status: candidate implementation; not evidence of protected-main merge or Production deployment.
|
||||
|
||||
`MessageEnvelope` is the single durable handoff identity and carries source/target/type, tenant/user/
|
||||
project/conversation/task/run scope, causation, sequence, capability scope, confidentiality scope,
|
||||
idempotency key, expiry and delivery state.
|
||||
|
||||
Identity includes **target + scope + sequence + idempotency key + payload hash**, fixing the existing
|
||||
AgentEventBus target collision. Server ActorContext is authoritative; caller scope may narrow but may
|
||||
not widen/forge trusted tenant/user/project/conversation/run fields.
|
||||
|
||||
`DurableHandoffStore` provides SQLite durability in the current runtime. Each logical stream has
|
||||
unique `(stream_key, idempotency_key)` and `(stream_key, sequence)`. Exact retry returns `REPLAY`;
|
||||
changed payload under the same key is rejected; sequence gaps are rejected; expired messages are not
|
||||
delivered.
|
||||
|
||||
Confidentiality scopes:
|
||||
|
||||
- public: no tenant requirement
|
||||
- tenant: same tenant
|
||||
- project: same tenant + project
|
||||
- private: same tenant + project + user
|
||||
- run: same tenant + project + user + run
|
||||
|
||||
Required `capability_scope` must be present on the authenticated delivery actor unless the actor is
|
||||
explicitly administrative.
|
||||
|
||||
Compatibility adapters convert AgentEventBus events, TaskPlane CloudEvents and legacy Agent OS
|
||||
MessageBus messages into the durable envelope. Existing message_outbox remains intact as a migration
|
||||
source; P0-B does not delete or silently rewrite it.
|
||||
|
||||
Required gates:
|
||||
`MESSAGE_SCOPE_MISMATCH_DENY`, `MESSAGE_TARGET_COLLISION_DENY`, `MESSAGE_IDEMPOTENCY_REPLAY`,
|
||||
`MESSAGE_SEQUENCE_GAP_DETECT`, `MESSAGE_CONFIDENTIALITY_SCOPE_ENFORCE`,
|
||||
`HANDOFF_CROSS_TENANT_DENY`, plus existing P0-A TaskPlane scope/security regressions.
|
||||
|
||||
SQLite remains current runtime persistence. PostgreSQL remains a separate migration gate.
|
||||
@@ -48,3 +48,8 @@ class MessageBus:
|
||||
messages = list(self.queue)
|
||||
self.queue.clear()
|
||||
return messages
|
||||
|
||||
# R4.2 P0-B compatibility adapter.
|
||||
def to_handoff_envelope(message, *, target=None, sequence=1, idempotency_key=None, **scope):
|
||||
from message_handoff_adapters import legacy_message_to_envelope
|
||||
return legacy_message_to_envelope(message, target=target, sequence=sequence, idempotency_key=idempotency_key, **scope)
|
||||
|
||||
@@ -45,6 +45,9 @@ class ShadowClientRBACPublicInterfaceTests(unittest.TestCase):
|
||||
identity_verifier=self.IdentityVerifier({"shadow-client"}),
|
||||
execution_enabled=False,
|
||||
identity_capabilities={"shadow-client": set(capabilities)},
|
||||
identity_actor_contexts={
|
||||
"shadow-client": {"scope_type": "system"}
|
||||
},
|
||||
)
|
||||
service = self.TaskPlaneService(store, policy=policy)
|
||||
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
ACTIVE_RUNTIME = ROOT / "active-runtime"
|
||||
if str(ACTIVE_RUNTIME) not in sys.path:
|
||||
sys.path.insert(0, str(ACTIVE_RUNTIME))
|
||||
|
||||
import sqlite3
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from app.task_plane.fleet import FleetConflict, FleetScheduler
|
||||
|
||||
|
||||
@dataclass
|
||||
class Plan:
|
||||
task_id: str
|
||||
required_capabilities: list[str] = field(default_factory=lambda:["code.read"])
|
||||
metadata: dict = field(default_factory=dict)
|
||||
scope_type: str = "tenant"
|
||||
tenant_id: str = "tenant-a"
|
||||
project_id: str | None = "project-a"
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
task_id: str
|
||||
revision: int = 1
|
||||
status: str = "QUEUED"
|
||||
created_at: float = 1.0
|
||||
task_plan: Plan | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.task_plan is None:
|
||||
self.task_plan=Plan(self.task_id)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Worker:
|
||||
worker_id: str
|
||||
capabilities: list[str]=field(default_factory=lambda:["code.read"])
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self, path=":memory:"):
|
||||
self.conn=sqlite3.connect(path,check_same_thread=False)
|
||||
self._lock=threading.RLock()
|
||||
self.conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS claims(
|
||||
claim_id TEXT PRIMARY KEY,task_id TEXT,worker_id TEXT,status TEXT,lease_expires_at REAL
|
||||
);
|
||||
""")
|
||||
self.tasks=[]
|
||||
self.workers={}
|
||||
|
||||
def list_tasks(self,_ctx=None):
|
||||
return list(self.tasks)
|
||||
|
||||
def get_worker(self,worker_id):
|
||||
return self.workers[worker_id]
|
||||
|
||||
|
||||
class Registry:
|
||||
def __init__(self):
|
||||
self.online=set()
|
||||
def is_online(self,worker_id,now=None):
|
||||
return worker_id in self.online
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fleet():
|
||||
store=Store(); registry=Registry()
|
||||
for wid in ("w-a","w-b","w-c"):
|
||||
store.workers[wid]=Worker(wid)
|
||||
registry.online.add(wid)
|
||||
scheduler=FleetScheduler(store,registry,now_fn=lambda:1000.0)
|
||||
scheduler.configure_worker("w-a",max_slots=1,failure_domain="zone-a",weight=1)
|
||||
scheduler.configure_worker("w-b",max_slots=2,failure_domain="zone-b",weight=1)
|
||||
scheduler.configure_worker("w-c",max_slots=1,failure_domain="zone-c",weight=2)
|
||||
return store,registry,scheduler
|
||||
|
||||
|
||||
def test_FLEET_AGENT_REGISTRATION_capacity_policy_is_durable(fleet):
|
||||
store,_,scheduler=fleet
|
||||
policy=scheduler.state.worker_policy("w-b")
|
||||
assert policy.max_slots==2 and policy.failure_domain=="zone-b"
|
||||
reloaded=type(scheduler)(store, scheduler.workers, now_fn=lambda:1000.0)
|
||||
assert reloaded.state.worker_policy("w-b").max_slots==2
|
||||
|
||||
|
||||
def test_FLEET_CAPACITY_SCHEDULING_does_not_oversubscribe_worker(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[Task("task-1"),Task("task-2"),Task("task-3")]
|
||||
decisions=scheduler.select(["w-a"],limit=3,now=1000.0)
|
||||
assert [d.worker_id for d in decisions]==["w-a"]
|
||||
|
||||
|
||||
def test_FLEET_CAPACITY_SCHEDULING_accounts_for_active_claims(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.conn.execute(
|
||||
"INSERT INTO claims VALUES ('claim-1','other','w-a','ACTIVE',2000)"
|
||||
)
|
||||
store.tasks=[Task("task-1")]
|
||||
assert scheduler.select(["w-a"],limit=1,now=1000.0)==[]
|
||||
|
||||
|
||||
def test_FLEET_FAIRNESS_scope_virtual_finish_prevents_hot_scope_monopoly(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[
|
||||
Task("a1",created_at=1,task_plan=Plan("a1",tenant_id="tenant-a",project_id="p-a")),
|
||||
Task("a2",created_at=2,task_plan=Plan("a2",tenant_id="tenant-a",project_id="p-a")),
|
||||
Task("b1",created_at=3,task_plan=Plan("b1",tenant_id="tenant-b",project_id="p-b")),
|
||||
]
|
||||
first=scheduler.select(["w-b"],limit=1,now=1000.0)[0]
|
||||
assert first.task_id=="a1"
|
||||
scheduler.state.release("a1",state="COMPLETED",now=1000.1)
|
||||
store.tasks=[task for task in store.tasks if task.task_id!="a1"]
|
||||
second=scheduler.select(["w-b"],limit=1,now=1000.2)[0]
|
||||
assert second.task_id=="b1"
|
||||
|
||||
|
||||
def test_FLEET_FAULT_ISOLATION_opens_worker_circuit_after_threshold(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[Task("task-1")]
|
||||
for _ in range(3):
|
||||
scheduler.state.record_worker_failure("w-a","runtime_crash",now=1000.0)
|
||||
assert scheduler.state.worker_circuit_open("w-a",1001.0)
|
||||
assert scheduler.select(["w-a"],limit=1,now=1001.0)==[]
|
||||
|
||||
|
||||
def test_FLEET_WORKER_REASSIGNMENT_avoids_quarantined_worker(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[Task("task-1")]
|
||||
scheduler.state.record_task_worker_failure(
|
||||
"task-1","w-a","lease_expired",quarantine_seconds=60,now=1000.0
|
||||
)
|
||||
decision=scheduler.select(["w-a","w-b"],limit=1,now=1001.0)[0]
|
||||
assert decision.worker_id=="w-b"
|
||||
|
||||
|
||||
def test_FLEET_SCOPE_ISOLATION_uses_tenant_project_fairness_keys(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[
|
||||
Task("a",task_plan=Plan("a",tenant_id="tenant-a",project_id="p")),
|
||||
Task("b",task_plan=Plan("b",tenant_id="tenant-b",project_id="p")),
|
||||
]
|
||||
decisions=scheduler.select(["w-b"],limit=2,now=1000.0)
|
||||
assert {d.scope_key for d in decisions}=={
|
||||
"tenant:tenant-a|project:p","tenant:tenant-b|project:p"
|
||||
}
|
||||
|
||||
|
||||
def test_FLEET_IDEMPOTENT_DISPATCH_replay_same_task_worker(fleet):
|
||||
store,_,scheduler=fleet
|
||||
store.tasks=[Task("task-1")]
|
||||
one=scheduler.select(["w-b"],limit=1,now=1000.0)[0]
|
||||
# Reservation remains authoritative; a conflicting second worker cannot steal it.
|
||||
with pytest.raises(FleetConflict):
|
||||
scheduler.state.reserve(
|
||||
task_id="task-1",worker_id="w-c",task_revision=1,
|
||||
scope_key=one.scope_key,failure_domain="zone-c",score=one.score,
|
||||
reservation_seconds=30,now=1001.0,scope_weight=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_FLEET_FAILURE_RECOVERY_worker_success_closes_circuit(fleet):
|
||||
store,_,scheduler=fleet
|
||||
for _ in range(3):
|
||||
scheduler.state.record_worker_failure("w-a","runtime_crash",now=1000.0)
|
||||
scheduler.state.record_worker_success("w-a",now=1001.0)
|
||||
assert not scheduler.state.worker_circuit_open("w-a",1001.1)
|
||||
|
||||
|
||||
def test_affinity_and_failure_domain_anti_affinity(fleet):
|
||||
store,_,scheduler=fleet
|
||||
scheduler.configure_worker("w-b",max_slots=2,failure_domain="zone-b",labels={"gpu":"false"})
|
||||
scheduler.configure_worker("w-c",max_slots=1,failure_domain="zone-c",labels={"gpu":"true"})
|
||||
task=Task("gpu-task")
|
||||
task.task_plan.metadata={
|
||||
"fleet_affinity_labels":{"gpu":"true"},
|
||||
"fleet_avoid_failure_domains":["zone-b"],
|
||||
}
|
||||
store.tasks=[task]
|
||||
decision=scheduler.select(["w-b","w-c"],limit=1,now=1000.0)[0]
|
||||
assert decision.worker_id=="w-c"
|
||||
|
||||
|
||||
def test_FLEET_FAILURE_RECOVERY_state_survives_sqlite_restart(tmp_path):
|
||||
path=tmp_path/"task-plane.sqlite3"
|
||||
first=Store(str(path)); registry=Registry()
|
||||
first.workers["w-a"]=Worker("w-a"); registry.online.add("w-a")
|
||||
scheduler=FleetScheduler(first,registry,now_fn=lambda:1000.0)
|
||||
scheduler.configure_worker("w-a",max_slots=3,failure_domain="zone-a")
|
||||
scheduler.state.record_worker_failure("w-a","runtime_crash",threshold=1,now=1000.0)
|
||||
first.conn.close()
|
||||
|
||||
second=Store(str(path)); second.workers["w-a"]=Worker("w-a")
|
||||
resumed=FleetScheduler(second,registry,now_fn=lambda:1001.0)
|
||||
assert resumed.state.worker_policy("w-a").max_slots==3
|
||||
assert resumed.state.worker_circuit_open("w-a",1001.0)
|
||||
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
ACTIVE_RUNTIME = ROOT / "active-runtime"
|
||||
if str(ACTIVE_RUNTIME) not in sys.path:
|
||||
sys.path.insert(0, str(ACTIVE_RUNTIME))
|
||||
|
||||
from app.task_plane.scheduler import TaskScheduler
|
||||
|
||||
|
||||
def _security_contract():
|
||||
path = Path(__file__).with_name("test_task_plane_security_contract.py")
|
||||
spec = importlib.util.spec_from_file_location("p2b_order_security_contract", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _create_task(service, contract, task_id, tenant, project):
|
||||
return service.create_task(
|
||||
{
|
||||
"task_id": task_id,
|
||||
"issuer": {"type": "agent", "id": contract.ADMIN_IDENTITY},
|
||||
"required_capabilities": ["code.read"],
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": tenant,
|
||||
"project_id": project,
|
||||
"metadata": {"fleet_priority": 0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _task_id(claim):
|
||||
return getattr(
|
||||
claim,
|
||||
"task_id",
|
||||
claim.get("task_id") if isinstance(claim, dict) else None,
|
||||
)
|
||||
|
||||
|
||||
def test_P2B_FLEET_ORDER_is_preserved_through_legacy_claim_oracle():
|
||||
contract = _security_contract()
|
||||
service = contract._service()
|
||||
service.register_worker(contract._worker_payload(max_concurrency=3))
|
||||
worker_id = contract.WORKER_ID
|
||||
|
||||
scheduler = TaskScheduler(service.store)
|
||||
scheduler.fleet_scheduler.configure_worker(
|
||||
worker_id, max_slots=3, failure_domain="zone-a"
|
||||
)
|
||||
|
||||
_create_task(service, contract, "p2b-order-a1", "tenant-a", "project-a")
|
||||
_create_task(service, contract, "p2b-order-a2", "tenant-a", "project-a")
|
||||
_create_task(service, contract, "p2b-order-b1", "tenant-b", "project-b")
|
||||
|
||||
first = scheduler.schedule_for_worker(worker_id)
|
||||
second = scheduler.schedule_for_worker(worker_id)
|
||||
|
||||
assert _task_id(first) == "p2b-order-a1"
|
||||
assert _task_id(second) == "p2b-order-b1"
|
||||
|
||||
rows = service.store.conn.execute(
|
||||
"SELECT scope_key,dispatch_count FROM fleet_scope_state ORDER BY scope_key"
|
||||
).fetchall()
|
||||
assert len(rows) >= 2
|
||||
@@ -0,0 +1,181 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
ACTIVE_RUNTIME = ROOT / "active-runtime"
|
||||
if str(ACTIVE_RUNTIME) not in sys.path:
|
||||
sys.path.insert(0, str(ACTIVE_RUNTIME))
|
||||
|
||||
from app.task_plane.fleet import FleetConflict, FleetScheduler
|
||||
from app.task_plane.errors import NotFoundError
|
||||
from app.task_plane.scheduler import TaskScheduler
|
||||
|
||||
|
||||
def _security_contract():
|
||||
path = Path(__file__).with_name("test_task_plane_security_contract.py")
|
||||
spec = importlib.util.spec_from_file_location("p2b_security_contract", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_P2B_REAL_DISPATCH_claim_binds_fleet_reservation_and_circuit():
|
||||
contract = _security_contract()
|
||||
service = contract._service()
|
||||
service.register_worker(contract._worker_payload(max_concurrency=2))
|
||||
worker_id = contract.WORKER_ID
|
||||
|
||||
service.create_task(
|
||||
{
|
||||
"task_id": "task-p2b-real-1",
|
||||
"issuer": {"type": "agent", "id": contract.ADMIN_IDENTITY},
|
||||
"required_capabilities": ["code.read"],
|
||||
}
|
||||
)
|
||||
|
||||
scheduler = TaskScheduler(service.store)
|
||||
scheduler.fleet_scheduler.configure_worker(
|
||||
worker_id,
|
||||
max_slots=2,
|
||||
failure_domain="zone-real",
|
||||
labels={"runtime": "task-plane"},
|
||||
)
|
||||
first = scheduler.schedule_for_worker(worker_id)
|
||||
assert first is not None
|
||||
assert first.task_id == "task-p2b-real-1"
|
||||
|
||||
row = service.store.conn.execute(
|
||||
"SELECT state,claim_id,worker_id FROM fleet_dispatches WHERE task_id=?",
|
||||
("task-p2b-real-1",),
|
||||
).fetchone()
|
||||
assert tuple(row) == ("CLAIMED", first.claim_id, worker_id)
|
||||
|
||||
# A real circuit blocks the actual scheduler even though one capacity slot remains.
|
||||
failure_now = time.time()
|
||||
scheduler.record_worker_failure(
|
||||
worker_id, "runtime_crash", threshold=1, now=failure_now
|
||||
)
|
||||
service.create_task(
|
||||
{
|
||||
"task_id": "task-p2b-real-2",
|
||||
"issuer": {"type": "agent", "id": contract.ADMIN_IDENTITY},
|
||||
"required_capabilities": ["code.read"],
|
||||
}
|
||||
)
|
||||
with pytest.raises(NotFoundError, match="no queued tasks available"):
|
||||
scheduler.schedule_for_worker(worker_id)
|
||||
|
||||
scheduler.record_worker_success(worker_id, now=failure_now + 1.0)
|
||||
second = scheduler.schedule_for_worker(worker_id)
|
||||
assert second is not None
|
||||
assert second.task_id == "task-p2b-real-2"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Plan:
|
||||
task_id: str
|
||||
required_capabilities: list[str] = field(default_factory=lambda: ["code.read"])
|
||||
metadata: dict = field(default_factory=dict)
|
||||
scope_type: str = "tenant"
|
||||
tenant_id: str = "tenant-a"
|
||||
project_id: str | None = "project-a"
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
task_id: str
|
||||
revision: int = 1
|
||||
status: str = "QUEUED"
|
||||
created_at: float = 1.0
|
||||
task_plan: Plan | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.task_plan is None:
|
||||
self.task_plan = Plan(self.task_id)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Worker:
|
||||
worker_id: str
|
||||
capabilities: list[str] = field(default_factory=lambda: ["code.read"])
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self):
|
||||
self.conn = sqlite3.connect(":memory:", check_same_thread=False)
|
||||
self._lock = threading.RLock()
|
||||
self.conn.execute(
|
||||
"CREATE TABLE claims("
|
||||
"claim_id TEXT PRIMARY KEY,task_id TEXT,worker_id TEXT,status TEXT,"
|
||||
"lease_expires_at REAL)"
|
||||
)
|
||||
self.tasks = []
|
||||
self.workers = {}
|
||||
|
||||
def list_tasks(self, _ctx=None):
|
||||
return list(self.tasks)
|
||||
|
||||
def get_worker(self, worker_id):
|
||||
return self.workers[worker_id]
|
||||
|
||||
|
||||
class Registry:
|
||||
def __init__(self):
|
||||
self.online = set()
|
||||
|
||||
def is_online(self, worker_id, now=None):
|
||||
return worker_id in self.online
|
||||
|
||||
|
||||
def _fleet():
|
||||
store = Store()
|
||||
registry = Registry()
|
||||
store.workers["w-a"] = Worker("w-a")
|
||||
registry.online.add("w-a")
|
||||
fleet = FleetScheduler(store, registry, now_fn=lambda: 1000.0)
|
||||
fleet.configure_worker("w-a", max_slots=1, failure_domain="zone-a")
|
||||
return store, fleet
|
||||
|
||||
|
||||
def test_P2B_CAPACITY_RESERVATION_is_atomic_under_concurrent_pullers():
|
||||
store, fleet = _fleet()
|
||||
tasks = [Task("task-a"), Task("task-b")]
|
||||
store.tasks = tasks
|
||||
barrier = threading.Barrier(2)
|
||||
outcomes = []
|
||||
|
||||
def reserve(task):
|
||||
barrier.wait()
|
||||
try:
|
||||
fleet.reserve_for_worker("w-a", task, now=1000.0)
|
||||
outcomes.append("reserved")
|
||||
except FleetConflict:
|
||||
outcomes.append("conflict")
|
||||
|
||||
threads = [threading.Thread(target=reserve, args=(task,)) for task in tasks]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
assert sorted(outcomes) == ["conflict", "reserved"]
|
||||
assert fleet.state.reserved_slots("w-a", 1000.0) == 1
|
||||
|
||||
|
||||
def test_P2B_RANK_is_read_only_until_legacy_security_accepts_candidate():
|
||||
store, fleet = _fleet()
|
||||
store.tasks = [Task("task-a"), Task("task-b")]
|
||||
ranked = fleet.rank_for_worker("w-a", now=1000.0)
|
||||
assert [task.task_id for task in ranked] == ["task-a", "task-b"]
|
||||
assert fleet.state.reserved_slots("w-a", 1000.0) == 0
|
||||
@@ -0,0 +1,298 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from http.server import ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = os.getcwd()
|
||||
if os.path.join(ROOT, "active-runtime") not in sys.path:
|
||||
sys.path.append(os.path.join(ROOT, "active-runtime"))
|
||||
|
||||
from app.task_plane.api import TaskPlaneRequestHandler, TaskPlaneService
|
||||
from app.task_plane.errors import ForbiddenError, NotFoundError, ValidationError
|
||||
from app.task_plane.models import TaskPlan
|
||||
from app.task_plane.policy import PolicyEngine
|
||||
from app.task_plane.store import TaskPlaneStore
|
||||
|
||||
|
||||
def _policy() -> PolicyEngine:
|
||||
return PolicyEngine(
|
||||
identity_capabilities={
|
||||
"actor-a": {"task.creator", "task.read"},
|
||||
"actor-b": {"task.creator", "task.read"},
|
||||
"reader-a": {"task.read"},
|
||||
"reader-b": {"task.read"},
|
||||
"owner": {"task.creator", "task.read", "task.read.all"},
|
||||
"unbound": {"task.creator", "task.read"},
|
||||
},
|
||||
identity_actor_contexts={
|
||||
"actor-a": {
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-a",
|
||||
"user_id": "user-a",
|
||||
"project_id": "project-a",
|
||||
"conversation_id": "conversation-a",
|
||||
"run_id": "run-a",
|
||||
},
|
||||
"reader-a": {
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-a",
|
||||
"user_id": "user-a",
|
||||
"project_id": "project-a",
|
||||
"conversation_id": "conversation-a",
|
||||
"run_id": "run-a",
|
||||
},
|
||||
"actor-b": {
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-b",
|
||||
"user_id": "user-b",
|
||||
"project_id": "project-b",
|
||||
},
|
||||
"reader-b": {
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-b",
|
||||
"user_id": "user-b",
|
||||
"project_id": "project-b",
|
||||
},
|
||||
"owner": {"scope_type": "system"},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _service() -> TaskPlaneService:
|
||||
return TaskPlaneService(TaskPlaneStore(":memory:"), policy=_policy())
|
||||
|
||||
|
||||
def test_taskplan_scope_fields_round_trip_as_first_class_values() -> None:
|
||||
plan = TaskPlan.from_dict(
|
||||
{
|
||||
"task_id": "task-scope-roundtrip",
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-a",
|
||||
"user_id": "user-a",
|
||||
"project_id": "project-a",
|
||||
"conversation_id": "conversation-a",
|
||||
"run_id": "run-a",
|
||||
}
|
||||
)
|
||||
payload = plan.to_dict()
|
||||
assert TaskPlan.from_dict(payload).to_dict() == payload
|
||||
|
||||
|
||||
def test_tenant_scope_requires_tenant_id_and_system_scope_does_not() -> None:
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
TaskPlan.from_dict({"task_id": "task-missing-tenant", "scope_type": "tenant"})
|
||||
assert exc.value.error_code == "TASK_SCOPE_TENANT_ID_REQUIRED"
|
||||
system = TaskPlan.from_dict(
|
||||
{"task_id": "task-system", "scope_type": "system", "run_id": "system-run"}
|
||||
)
|
||||
assert system.tenant_id is None
|
||||
|
||||
|
||||
def test_ACTOR_BINDING_AUTHZ_MERGE_uses_server_scope() -> None:
|
||||
service = _service()
|
||||
created = service.create_task(
|
||||
{"task_id": "task-a", "task_type": "code.change"},
|
||||
actor_identity="actor-a",
|
||||
)
|
||||
assert created["plan"]["tenant_id"] == "tenant-a"
|
||||
assert created["plan"]["user_id"] == "user-a"
|
||||
assert created["plan"]["project_id"] == "project-a"
|
||||
declared = service.create_task(
|
||||
{
|
||||
"task_id": "task-a-declared",
|
||||
"tenant_id": "tenant-a",
|
||||
"user_id": "user-a",
|
||||
"project_id": "project-a",
|
||||
"conversation_id": "conversation-a",
|
||||
"run_id": "run-a",
|
||||
},
|
||||
actor_identity="actor-a",
|
||||
)
|
||||
assert declared["plan"]["scope_type"] == "tenant"
|
||||
|
||||
|
||||
def test_TASK_SCOPE_MISMATCH_DENY_rejects_forged_client_scope() -> None:
|
||||
service = _service()
|
||||
with pytest.raises(ForbiddenError) as exc:
|
||||
service.create_task(
|
||||
{
|
||||
"task_id": "task-forged",
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-b",
|
||||
"user_id": "user-b",
|
||||
"project_id": "project-b",
|
||||
},
|
||||
actor_identity="actor-a",
|
||||
)
|
||||
assert exc.value.error_code == "TASK_SCOPE_MISMATCH_DENY"
|
||||
|
||||
|
||||
def test_public_create_fails_closed_without_server_actor_scope() -> None:
|
||||
service = _service()
|
||||
with pytest.raises(ForbiddenError) as exc:
|
||||
service.create_task(
|
||||
{"task_id": "task-unbound"}, actor_identity="unbound"
|
||||
)
|
||||
assert exc.value.error_code == "TASK_ACTOR_SCOPE_UNAVAILABLE"
|
||||
|
||||
|
||||
def test_TASKPLANE_LIST_AUTHZ_DENY_and_TASK_OWNERSHIP_FILTER() -> None:
|
||||
service = _service()
|
||||
service.create_task({"task_id": "task-a"}, actor_identity="actor-a")
|
||||
service.create_task({"task_id": "task-b"}, actor_identity="actor-b")
|
||||
assert [item["task_id"] for item in service.list_tasks("reader-a")] == ["task-a"]
|
||||
assert [item["task_id"] for item in service.list_tasks("reader-b")] == ["task-b"]
|
||||
|
||||
|
||||
def test_TASKPLANE_GET_AUTHZ_DENY_does_not_disclose_cross_tenant_task() -> None:
|
||||
service = _service()
|
||||
service.create_task({"task_id": "task-a"}, actor_identity="actor-a")
|
||||
with pytest.raises(NotFoundError):
|
||||
service.get_task("task-a", "reader-b")
|
||||
|
||||
|
||||
def test_http_authentication_drives_server_scope_and_cross_tenant_404() -> None:
|
||||
service = _service()
|
||||
|
||||
class Handler(TaskPlaneRequestHandler):
|
||||
pass
|
||||
|
||||
Handler.service = service
|
||||
Handler.authenticator = lambda token: {
|
||||
"token-a": "actor-a",
|
||||
"token-b": "reader-b",
|
||||
}.get(token)
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
base = f"http://127.0.0.1:{server.server_port}"
|
||||
try:
|
||||
request = urllib.request.Request(
|
||||
base + "/tasks",
|
||||
data=json.dumps({"task_id": "task-http-a"}).encode(),
|
||||
headers={
|
||||
"Authorization": "Bearer token-a",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=3) as response:
|
||||
created = json.loads(response.read())
|
||||
assert created["plan"]["tenant_id"] == "tenant-a"
|
||||
|
||||
denied = urllib.request.Request(
|
||||
base + "/tasks/task-http-a",
|
||||
headers={"Authorization": "Bearer token-b"},
|
||||
)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc:
|
||||
urllib.request.urlopen(denied, timeout=3)
|
||||
assert exc.value.code == 404
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=3)
|
||||
|
||||
|
||||
def test_mutating_endpoint_is_scoped_before_cancel() -> None:
|
||||
policy = PolicyEngine(
|
||||
execution_enabled=True,
|
||||
identity_capabilities={
|
||||
"actor-a": {"task.creator", "task.read"},
|
||||
"actor-b": {"task.cancel"},
|
||||
},
|
||||
identity_actor_contexts={
|
||||
"actor-a": {"scope_type": "tenant", "tenant_id": "tenant-a"},
|
||||
"actor-b": {"scope_type": "tenant", "tenant_id": "tenant-b"},
|
||||
},
|
||||
)
|
||||
service = TaskPlaneService(TaskPlaneStore(":memory:"), policy=policy)
|
||||
service.create_task({"task_id": "task-a"}, actor_identity="actor-a")
|
||||
with pytest.raises(NotFoundError):
|
||||
service.cancel("task-a", {"actor_identity": "actor-b"})
|
||||
|
||||
|
||||
def test_privileged_server_actor_can_cross_scope() -> None:
|
||||
service = _service()
|
||||
service.create_task({"task_id": "task-a"}, actor_identity="actor-a")
|
||||
service.create_task({"task_id": "task-b"}, actor_identity="actor-b")
|
||||
assert {item["task_id"] for item in service.list_tasks("owner")} == {
|
||||
"task-a",
|
||||
"task-b",
|
||||
}
|
||||
|
||||
|
||||
def test_scope_sql_parameters_do_not_interpolate_tenant_values() -> None:
|
||||
service = _service()
|
||||
service.create_task({"task_id": "task-a"}, actor_identity="actor-a")
|
||||
policy = PolicyEngine(
|
||||
identity_capabilities={"attacker": {"task.read"}},
|
||||
identity_actor_contexts={
|
||||
"attacker": {
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-a' OR 1=1 --",
|
||||
}
|
||||
},
|
||||
)
|
||||
attacker_service = TaskPlaneService(service.store, policy=policy)
|
||||
assert attacker_service.list_tasks("attacker") == []
|
||||
|
||||
|
||||
def test_existing_sqlite_rows_migrate_to_legacy_without_tenant_guess(tmp_path) -> None:
|
||||
path = tmp_path / "legacy.sqlite3"
|
||||
connection = sqlite3.connect(path)
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE tasks (
|
||||
task_id TEXT PRIMARY KEY, revision INTEGER NOT NULL, status TEXT NOT NULL,
|
||||
task_json TEXT NOT NULL, task_type TEXT NOT NULL, root_task_id TEXT NOT NULL,
|
||||
parent_task_id TEXT, correlation_id TEXT, idempotency_key TEXT UNIQUE,
|
||||
created_at REAL NOT NULL, updated_at REAL NOT NULL, failure_reason TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.execute(
|
||||
"INSERT INTO tasks VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
|
||||
(
|
||||
"task-legacy",
|
||||
1,
|
||||
"CREATED",
|
||||
json.dumps({"task_id": "task-legacy"}),
|
||||
"code.change",
|
||||
"task-legacy",
|
||||
None,
|
||||
None,
|
||||
"legacy-key",
|
||||
1.0,
|
||||
1.0,
|
||||
None,
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
store = TaskPlaneStore(str(path))
|
||||
task = store.get_task("task-legacy")
|
||||
assert task.task_plan.scope_type == "legacy"
|
||||
assert task.task_plan.tenant_id is None
|
||||
replay = store.create_task(
|
||||
TaskPlan.from_dict({"task_id": "task-legacy"}),
|
||||
idempotency_key="legacy-key",
|
||||
)
|
||||
assert replay.task_id == "task-legacy"
|
||||
assert tuple(
|
||||
store.conn.execute(
|
||||
"SELECT scope_type,tenant_id FROM tasks WHERE task_id='task-legacy'"
|
||||
).fetchone()
|
||||
) == ("legacy", None)
|
||||
tenant_context = _policy().actor_context("reader-a")
|
||||
with pytest.raises(NotFoundError):
|
||||
store.get_task("task-legacy", tenant_context)
|
||||
store.close()
|
||||
@@ -0,0 +1,293 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0,str(ROOT/"active-runtime/app"))
|
||||
|
||||
from aios_agent_task_execution_rebind import (
|
||||
ExecutionRebindError,
|
||||
ExecutionRebindRequest,
|
||||
ExecutionRebindStore,
|
||||
)
|
||||
from aios_online_agent_execution import (
|
||||
GovernedExecutionBinding,
|
||||
project_execution_envelope,
|
||||
)
|
||||
|
||||
AUTHORIZATION="owner-execution-rebind-20260818"
|
||||
TASK_ID="task-rebind-a"
|
||||
JOB_ID="job-rebind-a"
|
||||
BASE="a"*40
|
||||
BRANCH="agent/task-rebind-a"
|
||||
IDEMPOTENCY="task-rebind-a-binding-v1"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _auth(monkeypatch):
|
||||
monkeypatch.setenv(
|
||||
"AIOS_AGENT_TASK_EXECUTION_REBIND_AUTHORIZATION_ID",
|
||||
AUTHORIZATION,
|
||||
)
|
||||
|
||||
|
||||
def canon(value):
|
||||
return hashlib.sha256(
|
||||
json.dumps(value,sort_keys=True,separators=(",",":")).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def binding():
|
||||
return GovernedExecutionBinding.from_mapping({
|
||||
"task_id":TASK_ID,
|
||||
"job_id":JOB_ID,
|
||||
"repository":"aios/aios",
|
||||
"base_oid":BASE,
|
||||
"expected_branch_oid":BASE,
|
||||
"target_branch":BRANCH,
|
||||
"branch_mode":"CREATE_NEW",
|
||||
"generation":1,
|
||||
"idempotency_key":IDEMPOTENCY,
|
||||
})
|
||||
|
||||
|
||||
def old_envelope():
|
||||
return project_execution_envelope(
|
||||
binding(),
|
||||
{
|
||||
"action":"source_capture",
|
||||
"adapter":"compact_interface_map",
|
||||
"capability_scope":[
|
||||
"git.read","git.test","governed.write.branch","governed.write.file"
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def new_envelope():
|
||||
return project_execution_envelope(
|
||||
binding(),
|
||||
{
|
||||
"action":"deterministic",
|
||||
"adapter":"deterministic",
|
||||
"capability_scope":[
|
||||
"git.read","git.test","governed.write.branch","governed.write.file"
|
||||
],
|
||||
"change_set":[
|
||||
{
|
||||
"path":"active-runtime/app/example.py",
|
||||
"operation":"write_file",
|
||||
"content":"VALUE=1\\n",
|
||||
}
|
||||
],
|
||||
"verification":[
|
||||
{
|
||||
"label":"unit",
|
||||
"argv":["/usr/bin/python3","-m","pytest","-q"],
|
||||
"cwd":".",
|
||||
"timeout_seconds":120,
|
||||
}
|
||||
],
|
||||
"governance":{
|
||||
"commit_idempotency_key":"task-rebind-a-commit-v1",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def make_db(tmp_path, *, status="QUEUED", revision=5, attempt=2, lease=False):
|
||||
path=tmp_path/"state.sqlite3"
|
||||
con=sqlite3.connect(path)
|
||||
con.executescript("""
|
||||
CREATE TABLE tasks(
|
||||
task_id TEXT PRIMARY KEY,
|
||||
status TEXT NOT NULL,
|
||||
attempt INTEGER NOT NULL,
|
||||
max_attempts INTEGER NOT NULL,
|
||||
revision INTEGER NOT NULL,
|
||||
claimed_by TEXT,
|
||||
lease_token_hash TEXT,
|
||||
lease_expires_at TEXT,
|
||||
verification_json TEXT,
|
||||
result_json TEXT,
|
||||
metadata_json TEXT,
|
||||
required_capabilities_json TEXT
|
||||
);
|
||||
CREATE TABLE task_events(
|
||||
event_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id TEXT,
|
||||
event_type TEXT,
|
||||
agent_id TEXT,
|
||||
payload_json TEXT,
|
||||
created_at TEXT
|
||||
);
|
||||
""")
|
||||
verification={"decision":"NEEDS_REWORK","reason":"WRONG_ENVELOPE_BEFORE_PLAN"}
|
||||
old=old_envelope()
|
||||
metadata={
|
||||
"execution":old,
|
||||
"execution_projection":{
|
||||
"schema":"aios.online-agent.task-projection.v1",
|
||||
"task_id":TASK_ID,
|
||||
"job_id":JOB_ID,
|
||||
"execution_digest":old["execution_digest"],
|
||||
},
|
||||
"preserved":{"value":"unchanged"},
|
||||
}
|
||||
con.execute(
|
||||
"INSERT INTO tasks VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
|
||||
(
|
||||
TASK_ID,status,attempt,3,revision,
|
||||
"previous-worker" if not lease else "active-worker",
|
||||
"lease-hash" if lease else None,
|
||||
"2099-01-01T00:00:00Z" if lease else None,
|
||||
json.dumps(verification),
|
||||
json.dumps({"status":"SUCCEEDED"}),
|
||||
json.dumps(metadata),
|
||||
json.dumps([
|
||||
"git.read","git.test","governed.write.branch","governed.write.file"
|
||||
]),
|
||||
),
|
||||
)
|
||||
con.commit();con.close()
|
||||
return path,verification
|
||||
|
||||
|
||||
def request(verification, *, revision=5, old=None, new=None, key="rebind-task-a-v1"):
|
||||
old=old or old_envelope()
|
||||
new=new or new_envelope()
|
||||
return ExecutionRebindRequest.from_mapping({
|
||||
"expected_status":"QUEUED",
|
||||
"expected_revision":revision,
|
||||
"expected_attempt":2,
|
||||
"previous_verification_sha256":canon(verification),
|
||||
"expected_old_execution_digest":old["execution_digest"],
|
||||
"expected_new_execution_digest":new["execution_digest"],
|
||||
"authorization_id":AUTHORIZATION,
|
||||
"idempotency_key":key,
|
||||
"reason":"projection-before-plan race recovery",
|
||||
})
|
||||
|
||||
|
||||
def resolver(_task_id,_caps):
|
||||
return new_envelope()
|
||||
|
||||
|
||||
def test_rebind_cas_updates_only_execution_projection_and_audits(tmp_path):
|
||||
path,verification=make_db(tmp_path)
|
||||
result=ExecutionRebindStore(path,resolver).rebind(
|
||||
TASK_ID,request(verification)
|
||||
)
|
||||
assert result["status"]=="QUEUED"
|
||||
assert result["revision"]==6
|
||||
assert result["idempotent"] is False
|
||||
|
||||
con=sqlite3.connect(path)
|
||||
row=con.execute(
|
||||
"SELECT status,attempt,revision,claimed_by,metadata_json FROM tasks"
|
||||
).fetchone()
|
||||
assert row[:4]==("QUEUED",2,6,None)
|
||||
metadata=json.loads(row[4])
|
||||
assert metadata["preserved"]=={"value":"unchanged"}
|
||||
assert metadata["execution"]["action"]=="deterministic"
|
||||
assert metadata["execution"]["execution_digest"]==new_envelope()["execution_digest"]
|
||||
assert metadata["execution_projection"]["execution_digest"]==new_envelope()["execution_digest"]
|
||||
|
||||
event=con.execute(
|
||||
"SELECT event_type,payload_json FROM task_events"
|
||||
).fetchone()
|
||||
assert event[0]=="TASK_EXECUTION_REBIND_AUTHORIZED"
|
||||
payload=json.loads(event[1])
|
||||
assert payload["old_action"]=="source_capture"
|
||||
assert payload["new_action"]=="deterministic"
|
||||
assert payload["stale_claim_normalized"] is True
|
||||
assert con.execute("SELECT COUNT(*) FROM task_execution_rebinds").fetchone()[0]==1
|
||||
|
||||
|
||||
def test_rebind_idempotent_replay(tmp_path):
|
||||
path,verification=make_db(tmp_path)
|
||||
store=ExecutionRebindStore(path,resolver)
|
||||
req=request(verification)
|
||||
first=store.rebind(TASK_ID,req)
|
||||
second=store.rebind(TASK_ID,req)
|
||||
assert first["idempotent"] is False
|
||||
assert second["idempotent"] is True
|
||||
assert second["new_execution_digest"]==first["new_execution_digest"]
|
||||
|
||||
|
||||
def test_rebind_rejects_active_lease(tmp_path):
|
||||
path,verification=make_db(tmp_path,lease=True)
|
||||
with pytest.raises(ExecutionRebindError,match="EXECUTION_REBIND_ACTIVE_LEASE_DENY"):
|
||||
ExecutionRebindStore(path,resolver).rebind(TASK_ID,request(verification))
|
||||
|
||||
|
||||
def test_rebind_rejects_verification_drift(tmp_path):
|
||||
path,_=make_db(tmp_path)
|
||||
with pytest.raises(ExecutionRebindError,match="EXECUTION_REBIND_VERIFICATION_DRIFT"):
|
||||
ExecutionRebindStore(path,resolver).rebind(
|
||||
TASK_ID,request({"decision":"NEEDS_REWORK","reason":"other"})
|
||||
)
|
||||
|
||||
|
||||
def test_rebind_rejects_nonqueued_task(tmp_path):
|
||||
path,verification=make_db(tmp_path,status="VERIFYING")
|
||||
with pytest.raises(ExecutionRebindError,match="EXECUTION_REBIND_CAS_MISMATCH"):
|
||||
ExecutionRebindStore(path,resolver).rebind(TASK_ID,request(verification))
|
||||
|
||||
|
||||
def test_rebind_rejects_non_deterministic_canonical_resolution(tmp_path):
|
||||
path,verification=make_db(tmp_path)
|
||||
def wrong(_task,_caps):
|
||||
return old_envelope()
|
||||
with pytest.raises(ExecutionRebindError,match="EXECUTION_REBIND_DETERMINISTIC_PLAN_REQUIRED"):
|
||||
ExecutionRebindStore(path,wrong).rebind(TASK_ID,request(verification))
|
||||
|
||||
|
||||
def test_rebind_rejects_binding_drift(tmp_path):
|
||||
path,verification=make_db(tmp_path)
|
||||
different=project_execution_envelope(
|
||||
GovernedExecutionBinding.from_mapping({
|
||||
"task_id":TASK_ID,
|
||||
"job_id":"job-other",
|
||||
"repository":"aios/aios",
|
||||
"base_oid":BASE,
|
||||
"expected_branch_oid":BASE,
|
||||
"target_branch":BRANCH,
|
||||
"branch_mode":"CREATE_NEW",
|
||||
"generation":1,
|
||||
"idempotency_key":IDEMPOTENCY,
|
||||
}),
|
||||
{
|
||||
"action":"deterministic",
|
||||
"capability_scope":["git.read"],
|
||||
"change_set":[{"path":"x.py","operation":"write_file","content":"x=1\\n"}],
|
||||
},
|
||||
)
|
||||
def wrong(_task,_caps):
|
||||
return different
|
||||
with pytest.raises(Exception):
|
||||
ExecutionRebindStore(path,wrong).rebind(TASK_ID,request(verification))
|
||||
|
||||
|
||||
def test_rebind_request_is_narrow_and_digest_bound():
|
||||
verification={"decision":"NEEDS_REWORK"}
|
||||
data=request(verification).public()
|
||||
data["arbitrary_metadata"]={"x":1}
|
||||
with pytest.raises(ExecutionRebindError,match="EXECUTION_REBIND_FIELDS_INVALID"):
|
||||
ExecutionRebindRequest.from_mapping(data)
|
||||
|
||||
|
||||
def test_execution_rebind_unit_is_unprivileged_and_owner_credential_bound():
|
||||
unit=(ROOT/"deploy/systemd/aios-agent-task-execution-rebind.service").read_text()
|
||||
assert "User=aios-agent-task" in unit
|
||||
assert "Group=aios-agent-task" in unit
|
||||
assert "LoadCredential=agent_task_control_token:" in unit
|
||||
assert "Environment=AIOS_GOVERNED_WRITE_API_URL=http://127.0.0.1:18892" in unit
|
||||
assert "RestrictAddressFamilies=AF_UNIX AF_INET" in unit
|
||||
assert "User=root" not in unit
|
||||
@@ -0,0 +1,117 @@
|
||||
from pathlib import Path
|
||||
import hashlib, json, sqlite3, sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0,str(ROOT/"active-runtime/app"))
|
||||
from aios_agent_task_retry_extension import RetryExtensionRequest, RetryExtensionStore, RetryExtensionError
|
||||
|
||||
AUTHORIZATION_ID="owner-scheme-a-20260817"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _authorization_id(monkeypatch):
|
||||
monkeypatch.setenv("AIOS_AGENT_TASK_AUTHORIZATION_ID", AUTHORIZATION_ID)
|
||||
|
||||
|
||||
def canon(v):
|
||||
return hashlib.sha256(json.dumps(v,sort_keys=True,separators=(",",":")).encode()).hexdigest()
|
||||
|
||||
|
||||
def db(tmp_path):
|
||||
p=tmp_path/"state.sqlite3"
|
||||
c=sqlite3.connect(p)
|
||||
c.executescript("""
|
||||
CREATE TABLE tasks(
|
||||
task_id TEXT PRIMARY KEY,status TEXT,attempt INTEGER,max_attempts INTEGER,revision INTEGER,
|
||||
claimed_by TEXT,lease_token_hash TEXT,lease_expires_at TEXT,verification_json TEXT,result_json TEXT
|
||||
);
|
||||
CREATE TABLE task_events(
|
||||
event_id INTEGER PRIMARY KEY AUTOINCREMENT,task_id TEXT,event_type TEXT,agent_id TEXT,
|
||||
payload_json TEXT,created_at TEXT
|
||||
);
|
||||
""")
|
||||
verification={"decision":"NEEDS_REWORK","reason":"BASE_OID_BINDING_INVALID"}
|
||||
c.execute("INSERT INTO tasks VALUES (?,?,?,?,?,?,?,?,?,?)",
|
||||
("task-a","FAILED",10,10,41,None,None,None,json.dumps(verification),json.dumps({"status":"FAILED"})))
|
||||
c.commit(); c.close()
|
||||
return p,verification
|
||||
|
||||
|
||||
def request(v):
|
||||
return RetryExtensionRequest.from_mapping({
|
||||
"expected_status":"FAILED","expected_revision":41,"expected_attempt":10,"expected_max_attempts":10,
|
||||
"next_max_attempts":11,"previous_verification_sha256":canon(v),
|
||||
"authorization_id":AUTHORIZATION_ID,"idempotency_key":"retry-task-a-attempt11",
|
||||
})
|
||||
|
||||
|
||||
def test_retry_extension_cas_and_audit(tmp_path):
|
||||
p,v=db(tmp_path)
|
||||
out=RetryExtensionStore(p).extend("task-a",request(v))
|
||||
assert out["status"]=="QUEUED" and out["max_attempts"]==11 and out["revision"]==42
|
||||
c=sqlite3.connect(p)
|
||||
assert c.execute("select status,attempt,max_attempts,revision from tasks").fetchone()==("QUEUED",10,11,42)
|
||||
event=c.execute("select event_type,payload_json from task_events").fetchone()
|
||||
assert event[0]=="TASK_RETRY_EXTENSION_AUTHORIZED"
|
||||
assert json.loads(event[1])["authorization_id"]=="owner-scheme-a-20260817"
|
||||
|
||||
|
||||
def test_retry_extension_idempotent_replay(tmp_path):
|
||||
p,v=db(tmp_path); store=RetryExtensionStore(p); req=request(v)
|
||||
first=store.extend("task-a",req); second=store.extend("task-a",req)
|
||||
assert first["idempotent"] is False and second["idempotent"] is True
|
||||
c=sqlite3.connect(p)
|
||||
assert c.execute("select count(*) from task_retry_extensions").fetchone()[0]==1
|
||||
|
||||
|
||||
def test_retry_extension_normalizes_stale_terminal_claim_without_active_lease(tmp_path):
|
||||
p,v=db(tmp_path)
|
||||
c=sqlite3.connect(p)
|
||||
c.execute("update tasks set claimed_by='previous-worker'")
|
||||
c.commit(); c.close()
|
||||
out=RetryExtensionStore(p).extend("task-a",request(v))
|
||||
assert out["status"]=="QUEUED"
|
||||
c=sqlite3.connect(p)
|
||||
assert c.execute(
|
||||
"select claimed_by,lease_token_hash,lease_expires_at from tasks"
|
||||
).fetchone()==(None,None,None)
|
||||
payload=json.loads(c.execute("select payload_json from task_events").fetchone()[0])
|
||||
assert payload["previous_claimed_by"]=="previous-worker"
|
||||
assert payload["stale_terminal_claim_normalized"] is True
|
||||
assert c.execute("select count(*) from task_events").fetchone()[0]==1
|
||||
|
||||
|
||||
def test_retry_extension_verification_drift_denied(tmp_path):
|
||||
p,v=db(tmp_path)
|
||||
bad=request({"decision":"NEEDS_REWORK","reason":"other"})
|
||||
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_VERIFICATION_DRIFT"):
|
||||
RetryExtensionStore(p).extend("task-a",bad)
|
||||
|
||||
|
||||
def test_retry_extension_authorization_id_mismatch_rejected(tmp_path, monkeypatch):
|
||||
p,v=db(tmp_path)
|
||||
monkeypatch.setenv("AIOS_AGENT_TASK_AUTHORIZATION_ID", "other-owner-id")
|
||||
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_AUTHORIZATION_ID_MISMATCH"):
|
||||
RetryExtensionStore(p).extend("task-a",request(v))
|
||||
|
||||
|
||||
def test_retry_extension_bounded_one_step():
|
||||
v={"decision":"NEEDS_REWORK"}
|
||||
data={
|
||||
"expected_status":"FAILED","expected_revision":41,"expected_attempt":10,"expected_max_attempts":10,
|
||||
"next_max_attempts":12,"previous_verification_sha256":canon(v),
|
||||
"authorization_id":"owner-scheme-a-20260817","idempotency_key":"retry-task-a-attempt11",
|
||||
}
|
||||
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_MUST_BE_BOUNDED_ONE_STEP"):
|
||||
RetryExtensionRequest.from_mapping(data)
|
||||
|
||||
|
||||
def test_retry_extension_unit_is_unprivileged_and_requires_owner_binding():
|
||||
unit=(ROOT/"deploy/systemd/aios-agent-task-retry-extension.service").read_text()
|
||||
assert "User=aios-agent-task" in unit
|
||||
assert "Group=aios-agent-task" in unit
|
||||
assert "LoadCredential=agent_task_control_token:" in unit
|
||||
assert "EnvironmentFile=/etc/aios-gitops/agent-task-retry-extension.env" in unit
|
||||
assert "User=root" not in unit
|
||||
@@ -77,7 +77,7 @@ def test_slash_branch_uses_paginated_collection():
|
||||
|
||||
def test_http_route_contract_is_complete():
|
||||
expected = {
|
||||
"/v2/runtime", "/v2/capabilities", "/v2/governed-task", "/v2/execution-envelope", "/v2/execution-plan",
|
||||
"/v2/runtime", "/v2/capabilities", "/v2/governed-task",
|
||||
"/v2/writer-lease", "/v2/writer-lease/bootstrap", "/v2/writer-lease/renew", "/v2/commit-file", "/v2/commit-files",
|
||||
"/v2/job", "/v2/job/resume", "/v2/job/reconcile",
|
||||
"/v2/tag/approval-request",
|
||||
@@ -203,38 +203,9 @@ def test_http_allows_omitted_job_id_for_durable_server_resolution():
|
||||
)
|
||||
|
||||
|
||||
def test_http_mutation_routes_require_exact_runtime_principal():
|
||||
handler = BridgeHandler.__new__(BridgeHandler)
|
||||
runtime_principal = "runtime-" "agent-principal-token-1234"
|
||||
handler.control_token = runtime_principal
|
||||
handler.agent_task_api_url = "http://127.0.0.1:1"
|
||||
handler.headers = {"Authorization": f"Bearer {runtime_principal}"}
|
||||
assert handler._authorized("createGovernedTask", {}) is True
|
||||
handler.headers = {"Authorization": "Bearer wrong-runtime-principal-token"}
|
||||
assert handler._authorized("createGovernedTask", {}) is False
|
||||
handler.headers = {}
|
||||
assert handler._authorized("createGovernedTask", {}) is False
|
||||
|
||||
|
||||
def test_capability_contracts_return_exact_job_id_lineage():
|
||||
contracts = BridgeV2Adapter(runtime_status={}).getGovernedWriteCapabilities()["operation_contracts"]
|
||||
for operation in ("acquireWriterLease", "bootstrapWriterLease", "renewWriterLease"):
|
||||
assert "job_id" in contracts[operation]["response_fields"]
|
||||
for operation in ("commitFileGoverned", "commitFilesGoverned"):
|
||||
assert "job_id" in contracts[operation]["response_fields"]
|
||||
|
||||
|
||||
def test_online_agent_execution_routes_and_action_projection_are_exposed():
|
||||
from aios_chatgpt_bridge_v2 import BridgeV2Adapter, OPERATION_REQUIRED_FIELDS
|
||||
from aios_chatgpt_bridge_v2_server import ROUTES
|
||||
assert ROUTES["/v2/execution-envelope"] == "projectOnlineAgentExecution"
|
||||
assert ROUTES["/v2/execution-plan"] == "registerOnlineAgentExecutionPlan"
|
||||
adapter = BridgeV2Adapter(runtime_status={})
|
||||
contracts = adapter.getGovernedWriteCapabilities()["operation_contracts"]
|
||||
assert tuple(contracts["projectOnlineAgentExecution"]["required_fields"]) == OPERATION_REQUIRED_FIELDS["projectOnlineAgentExecution"]
|
||||
assert tuple(contracts["registerOnlineAgentExecutionPlan"]["required_fields"]) == OPERATION_REQUIRED_FIELDS["registerOnlineAgentExecutionPlan"]
|
||||
schema = Path("docs/gitops/openapi-chatgpt-actions.yaml").read_text(encoding="utf-8")
|
||||
assert "/v2/execution-envelope:" in schema
|
||||
assert "operationId: projectOnlineAgentExecution" in schema
|
||||
assert "/v2/execution-plan:" in schema
|
||||
assert "operationId: registerOnlineAgentExecutionPlan" in schema
|
||||
|
||||
@@ -248,9 +248,6 @@ def test_create_new_executor_resumes_executing_and_uses_absence_bound_writers()
|
||||
self.advanced.append(state.value)
|
||||
return SimpleNamespace(as_dict=lambda: {"state": state.value})
|
||||
|
||||
def complete_continuation_cycle(self, *_args: object, **_kwargs: object) -> None:
|
||||
return None
|
||||
|
||||
class Workspace:
|
||||
def commit_files(self, _files: object, _message: str) -> str:
|
||||
return new_oid
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
APP = pathlib.Path(__file__).resolve().parents[1] / "active-runtime" / "app"
|
||||
if str(APP) not in sys.path:
|
||||
sys.path.insert(0, str(APP))
|
||||
|
||||
from message_handoff import (
|
||||
ConfidentialityScopeDenied,
|
||||
DeliveryContext,
|
||||
DurableHandoffStore,
|
||||
HandoffFabric,
|
||||
IdempotencyConflict,
|
||||
MessageEnvelope,
|
||||
MessageExpired,
|
||||
MessageScopeMismatch,
|
||||
SequenceGap,
|
||||
authorize_delivery,
|
||||
bind_envelope_scope,
|
||||
stable_event_id,
|
||||
)
|
||||
|
||||
|
||||
def actor(**overrides):
|
||||
values = dict(
|
||||
tenant_id="tenant-a", user_id="user-a", project_id="project-a",
|
||||
conversation_id="conversation-a", run_id="run-a",
|
||||
capabilities=("task.read", "message.read"),
|
||||
)
|
||||
values.update(overrides)
|
||||
return DeliveryContext(**values)
|
||||
|
||||
|
||||
def envelope(**overrides):
|
||||
values = dict(
|
||||
source="agent-a", target="agent-b", message_type="handoff.request", payload={"value": 1},
|
||||
tenant_id="tenant-a", user_id="user-a", project_id="project-a",
|
||||
conversation_id="conversation-a", task_id="task-a", run_id="run-a",
|
||||
sequence=1, idempotency_key="idem-1", capability_scope=("message.read",),
|
||||
confidentiality_scope="private",
|
||||
)
|
||||
values.update(overrides)
|
||||
return MessageEnvelope(**values)
|
||||
|
||||
|
||||
def test_MESSAGE_TARGET_COLLISION_DENY():
|
||||
a = stable_event_id(source="a", target="user-a", message_type="event", payload={"x": 1},
|
||||
tenant_id="tenant-a", sequence=1, idempotency_key="same")
|
||||
b = stable_event_id(source="a", target="user-b", message_type="event", payload={"x": 1},
|
||||
tenant_id="tenant-a", sequence=1, idempotency_key="same")
|
||||
assert a != b
|
||||
|
||||
|
||||
def test_MESSAGE_SCOPE_MISMATCH_DENY():
|
||||
with pytest.raises(MessageScopeMismatch, match="MESSAGE_SCOPE_MISMATCH_DENY"):
|
||||
bind_envelope_scope(envelope(tenant_id="tenant-b"), actor())
|
||||
|
||||
|
||||
def test_HANDOFF_CROSS_TENANT_DENY(tmp_path):
|
||||
store = DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
env = envelope()
|
||||
store.put(env)
|
||||
with pytest.raises(MessageScopeMismatch, match="MESSAGE_SCOPE_MISMATCH_DENY"):
|
||||
store.get(env.envelope_id, actor=actor(tenant_id="tenant-b"))
|
||||
|
||||
|
||||
def test_MESSAGE_CONFIDENTIALITY_SCOPE_ENFORCE():
|
||||
with pytest.raises(ConfidentialityScopeDenied, match="MESSAGE_CONFIDENTIALITY_SCOPE_ENFORCE"):
|
||||
authorize_delivery(actor(project_id="project-b"), envelope(confidentiality_scope="project"))
|
||||
with pytest.raises(ConfidentialityScopeDenied, match="MESSAGE_CONFIDENTIALITY_SCOPE_ENFORCE"):
|
||||
authorize_delivery(actor(), envelope(capability_scope=("message.read", "secret.read")))
|
||||
|
||||
|
||||
def test_MESSAGE_IDEMPOTENCY_REPLAY(tmp_path):
|
||||
store = DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
env = envelope()
|
||||
assert store.put(env)[0] == "CREATED"
|
||||
status, replay = store.put(env)
|
||||
assert status == "REPLAY"
|
||||
assert replay.envelope_id == env.envelope_id
|
||||
with pytest.raises(IdempotencyConflict, match="MESSAGE_IDEMPOTENCY_CONFLICT"):
|
||||
store.put(envelope(payload={"value": 2}))
|
||||
|
||||
|
||||
def test_MESSAGE_SEQUENCE_GAP_DETECT(tmp_path):
|
||||
store = DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
store.put(envelope(sequence=1, idempotency_key="one"))
|
||||
with pytest.raises(SequenceGap, match="MESSAGE_SEQUENCE_GAP_DETECT"):
|
||||
store.put(envelope(sequence=3, idempotency_key="three"))
|
||||
|
||||
|
||||
def test_sequence_advances_and_scope_filters_pending(tmp_path):
|
||||
store = DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
store.put(envelope(sequence=1, idempotency_key="one"))
|
||||
store.put(envelope(sequence=2, idempotency_key="two", payload={"value": 2}))
|
||||
assert [item.sequence for item in store.pending(target="agent-b", actor=actor())] == [1, 2]
|
||||
assert store.pending(target="agent-b", actor=actor(tenant_id="tenant-b")) == []
|
||||
|
||||
|
||||
def test_expired_message_rejected(tmp_path):
|
||||
expired = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(seconds=1)).isoformat()
|
||||
with pytest.raises(MessageExpired, match="MESSAGE_EXPIRED"):
|
||||
DurableHandoffStore(tmp_path / "handoff.sqlite3").put(envelope(expires_at=expired))
|
||||
|
||||
|
||||
def test_server_actor_binding_fills_scope_and_recomputes_identity():
|
||||
public = MessageEnvelope(
|
||||
source="agent-a", target="agent-b", message_type="handoff.request", payload={"value": 1},
|
||||
sequence=1, idempotency_key="bind-1", confidentiality_scope="public",
|
||||
)
|
||||
bound = bind_envelope_scope(public, actor())
|
||||
assert (bound.tenant_id, bound.user_id, bound.project_id, bound.run_id) == (
|
||||
"tenant-a", "user-a", "project-a", "run-a"
|
||||
)
|
||||
assert bound.envelope_id != public.envelope_id
|
||||
|
||||
|
||||
def test_handoff_fabric_publish_is_durable(tmp_path):
|
||||
store = DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
status, env = HandoffFabric(store).publish(envelope(), actor=actor())
|
||||
assert status == "CREATED"
|
||||
assert store.get(env.envelope_id, actor=actor()).payload == {"value": 1}
|
||||
@@ -1,524 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
APP = ROOT / "active-runtime" / "app"
|
||||
sys.path.insert(0, str(APP))
|
||||
|
||||
from aios_chatgpt_bridge_v2 import ( # noqa: E402
|
||||
BRIDGE_OPERATION_CONTRACTS,
|
||||
BridgeV2Adapter,
|
||||
)
|
||||
from aios_chatgpt_bridge_v2_server import GovernedRuntime # noqa: E402
|
||||
from aios_governed_task_bootstrap import BootstrapError, GovernedTaskBinding, GovernedTaskBootstrap, _task_key # noqa: E402
|
||||
from aios_online_agent_execution import ( # noqa: E402
|
||||
ExecutionProjectionError,
|
||||
GovernedExecutionBinding,
|
||||
OnlineTaskProjectionClient,
|
||||
project_execution_envelope,
|
||||
project_online_task,
|
||||
validate_execution_envelope,
|
||||
)
|
||||
import aios_online_agent_worker as worker_module # noqa: E402
|
||||
from aios_online_agent_worker import ExecutionEnvelope, WorkerError, WorkerSupervisor, governed_commit, load_task_state # noqa: E402
|
||||
|
||||
|
||||
BASE = "a642b045775b9154a00ae78543037fbef4fb447b"
|
||||
TASK = "p0-claim-online-agent-auto-execution-20260817"
|
||||
JOB = "job-2031a9e5682a4238f98e3ac7aee8a593"
|
||||
BRANCH = "agent/p0-claim-online-agent-auto-execution-20260817"
|
||||
|
||||
|
||||
def binding() -> GovernedExecutionBinding:
|
||||
return GovernedExecutionBinding.from_mapping(
|
||||
{
|
||||
"task_id": TASK,
|
||||
"job_id": JOB,
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"expected_branch_oid": BASE,
|
||||
"target_branch": BRANCH,
|
||||
"branch_mode": "CREATE_NEW",
|
||||
"generation": 1,
|
||||
"idempotency_key": "p0-claim-online-agent-auto-execution-v1",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def task_without_envelope() -> dict:
|
||||
return {
|
||||
"task_id": TASK,
|
||||
"goal": "make governed Online Tasks automatically executable",
|
||||
"instructions": "project, claim, heartbeat, and preserve exact lineage",
|
||||
"required_capabilities": ["git.read", "git.test"],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
|
||||
def test_project_online_task_builds_execution_metadata_from_binding() -> None:
|
||||
projected = project_online_task(binding(), task_without_envelope())
|
||||
execution = projected["metadata"]["execution"]
|
||||
assert execution["schema"] == "aios.online-agent.execution.v2"
|
||||
assert execution["job_id"] == JOB
|
||||
assert execution["task_id"] == TASK
|
||||
assert execution["repository"] == "aios/aios"
|
||||
assert execution["base_oid"] == BASE
|
||||
assert execution["expected_branch_oid"] == BASE
|
||||
assert execution["target_branch"] == BRANCH
|
||||
assert execution["generation"] == 1
|
||||
assert projected["metadata"]["execution_projection"]["execution_digest"] == execution["execution_digest"]
|
||||
|
||||
|
||||
def test_job_to_task_execution_projection(tmp_path: Path) -> None:
|
||||
exact = binding()
|
||||
durable = GovernedTaskBinding(
|
||||
task_id=exact.task_id,
|
||||
job_id=exact.job_id,
|
||||
repository=exact.repository,
|
||||
base_oid=exact.base_oid,
|
||||
expected_branch_oid=exact.expected_branch_oid,
|
||||
target_branch=exact.target_branch,
|
||||
branch_mode=exact.branch_mode,
|
||||
generation=exact.generation,
|
||||
idempotency_key=exact.idempotency_key,
|
||||
risk="engineering",
|
||||
description="automatic claim infrastructure",
|
||||
)
|
||||
|
||||
class Bootstrap:
|
||||
def create(self, _request: dict) -> dict:
|
||||
return durable.response()
|
||||
|
||||
def online_task_projection(self, _task_id: str, online: dict, execution: dict | None, **_kwargs: object) -> dict:
|
||||
return project_online_task(exact, online, execution)
|
||||
|
||||
def get_binding(self, _task_id: str) -> GovernedTaskBinding:
|
||||
return durable
|
||||
|
||||
def execution_binding(self, _binding: GovernedTaskBinding) -> GovernedExecutionBinding:
|
||||
return exact
|
||||
|
||||
class Projector:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def create(self, projected: dict, *, binding: GovernedExecutionBinding) -> dict:
|
||||
self.calls += 1
|
||||
assert binding == exact
|
||||
return {**projected, "status": "QUEUED"}
|
||||
|
||||
runtime = GovernedRuntime.__new__(GovernedRuntime)
|
||||
runtime.bootstrap = Bootstrap()
|
||||
runtime.bindings = {}
|
||||
runtime.online_task_projector = Projector()
|
||||
runtime.projection_outbox_root = tmp_path / "projection-outbox"
|
||||
runtime.projection_outbox_root.mkdir()
|
||||
response = GovernedRuntime.create_task(
|
||||
runtime,
|
||||
{
|
||||
"task_id": TASK,
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"target_branch": BRANCH,
|
||||
"risk": "engineering",
|
||||
"description": "automatic claim infrastructure",
|
||||
"idempotency_key": exact.idempotency_key,
|
||||
"online_task": task_without_envelope(),
|
||||
},
|
||||
)
|
||||
assert response["online_task_projection"]["status"] == "PROJECTED"
|
||||
assert response["online_task_projection"]["task_status"] == "QUEUED"
|
||||
assert runtime.online_task_projector.calls == 1
|
||||
|
||||
|
||||
def test_execution_envelope_schema_validation() -> None:
|
||||
envelope = project_execution_envelope(binding())
|
||||
validate_execution_envelope(envelope, binding=binding())
|
||||
malformed = dict(envelope)
|
||||
malformed["schema"] = "aios.online-agent.execution.v1"
|
||||
with pytest.raises(ExecutionProjectionError, match="EXECUTION_ENVELOPE_SCHEMA_VALIDATION"):
|
||||
validate_execution_envelope(malformed, binding=binding())
|
||||
|
||||
|
||||
def test_projection_response_must_equal_submitted_execution(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
submitted = project_online_task(binding(), task_without_envelope(), {"action": "source_capture", "source_signals": ["submitted"]})
|
||||
returned = project_online_task(binding(), task_without_envelope(), {"action": "source_capture", "source_signals": ["different"]})
|
||||
returned.update({"status": "QUEUED"})
|
||||
|
||||
class Response:
|
||||
status = 201
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_args: object) -> None:
|
||||
return None
|
||||
|
||||
def read(self) -> bytes:
|
||||
return json.dumps(returned).encode("utf-8")
|
||||
|
||||
monkeypatch.setattr("aios_online_agent_execution.urllib.request.urlopen", lambda *_args, **_kwargs: Response())
|
||||
token = tmp_path / "control.token"
|
||||
token.write_text("x" * 40, encoding="utf-8")
|
||||
with pytest.raises(ExecutionProjectionError, match="ONLINE_TASK_PROJECTION_RESPONSE_BINDING_INVALID"):
|
||||
OnlineTaskProjectionClient("http://127.0.0.1:1", token).create(submitted, binding=binding())
|
||||
|
||||
|
||||
class _ProjectedBridge:
|
||||
def __init__(self, value: dict) -> None:
|
||||
self.value = value
|
||||
self.calls = 0
|
||||
|
||||
def project_execution(self, _task: dict, _execution: dict | None = None) -> tuple[int, dict]:
|
||||
self.calls += 1
|
||||
return 200, self.value
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
def __init__(self, tasks: list[dict]) -> None:
|
||||
self.tasks = tasks
|
||||
self.claimed: list[str] = []
|
||||
|
||||
def poll(self) -> list[dict]:
|
||||
return self.tasks
|
||||
|
||||
def claim(self, task_id: str) -> tuple[int, dict]:
|
||||
self.claimed.append(task_id)
|
||||
return 200, {"task_id": task_id, "lease_token": "x" * 32, "attempt": 1, "revision": 2}
|
||||
|
||||
|
||||
def test_task_visible_to_worker_claimable(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(worker_module, "STATE_ROOT", tmp_path)
|
||||
supervisor = WorkerSupervisor("runtime-token-not-secret")
|
||||
supervisor.bridge = _ProjectedBridge(project_execution_envelope(binding()))
|
||||
envelope, source = supervisor.resolve_envelope(task_without_envelope())
|
||||
assert source == "recovered_projection"
|
||||
assert envelope.task_id == TASK
|
||||
assert envelope.job_id == JOB
|
||||
assert envelope.action == "source_capture"
|
||||
|
||||
|
||||
def test_task_claim_after_projection(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(worker_module, "STATE_ROOT", tmp_path)
|
||||
supervisor = WorkerSupervisor("runtime-token-not-secret")
|
||||
fake_agent = _FakeAgent([task_without_envelope()])
|
||||
supervisor.agent = fake_agent
|
||||
supervisor.bridge = _ProjectedBridge(project_execution_envelope(binding()))
|
||||
processed: list[tuple[str, str]] = []
|
||||
supervisor.process_claim = lambda task, claim, envelope: processed.append((task["task_id"], envelope.execution_digest if hasattr(envelope, "execution_digest") else envelope.digest)) # type: ignore[method-assign]
|
||||
assert supervisor.once() == 0
|
||||
assert fake_agent.claimed == [TASK]
|
||||
assert processed and processed[0][0] == TASK
|
||||
|
||||
|
||||
def test_claim_heartbeat_source_capture_and_result_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
execution = project_execution_envelope(binding(), {"action": "source_capture", "source_signals": ["agent"]})
|
||||
task = task_without_envelope()
|
||||
task["metadata"] = {"execution": execution}
|
||||
|
||||
class Agent:
|
||||
def __init__(self) -> None:
|
||||
self.heartbeats: list[str] = []
|
||||
self.results: list[dict] = []
|
||||
|
||||
def heartbeat(self, _task_id: str, _lease_token: str, stage: str) -> tuple[int, dict]:
|
||||
self.heartbeats.append(stage)
|
||||
return 200, {}
|
||||
|
||||
def get(self, _task_id: str) -> tuple[int, dict]:
|
||||
return 200, task
|
||||
|
||||
def result(self, _task_id: str, _lease_token: str, _key: str, result: dict) -> tuple[int, dict]:
|
||||
self.results.append(result)
|
||||
return 200, {}
|
||||
|
||||
worktree = tmp_path / "worktree"
|
||||
source = worktree / "active-runtime" / "app" / "agent_auto_fixture.py"
|
||||
source.parent.mkdir(parents=True)
|
||||
source.write_text("VALUE = 1\n", encoding="utf-8")
|
||||
monkeypatch.setattr(worker_module, "STATE_ROOT", tmp_path / "state")
|
||||
monkeypatch.setattr(worker_module, "new_run_root", lambda: tmp_path / "run")
|
||||
monkeypatch.setattr(worker_module, "prepare_worktree", lambda *_args, **_kwargs: worktree)
|
||||
monkeypatch.setattr(worker_module, "cleanup_run_root", lambda *_args, **_kwargs: None)
|
||||
agent = Agent()
|
||||
supervisor = WorkerSupervisor("runtime-token-not-secret")
|
||||
supervisor.agent = agent
|
||||
supervisor.bridge = _ProjectedBridge(execution)
|
||||
envelope, _ = supervisor.resolve_envelope(task)
|
||||
supervisor.process_claim(task, {"task_id": TASK, "lease_token": "l" * 32, "attempt": 2, "revision": 8}, envelope)
|
||||
assert agent.heartbeats and agent.heartbeats[0] == "claimed"
|
||||
assert len(agent.results) == 1
|
||||
assert agent.results[0]["status"] == "SUCCEEDED"
|
||||
assert agent.results[0]["classification"] == "NEEDS_REWORK_SOURCE_CAPTURE"
|
||||
|
||||
|
||||
def test_dynamic_task_binding_no_hardcoded_task_id() -> None:
|
||||
worker = (APP / "aios_online_agent_worker.py").read_text(encoding="utf-8")
|
||||
assert TASK not in worker
|
||||
assert BASE not in worker
|
||||
for name in ("aios-online-agent-worker.service", "aios-online-engineering-consumer.service"):
|
||||
unit = (ROOT / "deploy" / "systemd" / name).read_text(encoding="utf-8")
|
||||
assert TASK not in unit
|
||||
assert JOB not in unit
|
||||
assert BASE not in unit
|
||||
assert BRANCH not in unit
|
||||
assert "AIOS_ENGINEERING_PLAN_FILE" not in unit
|
||||
assert "aios_online_agent_worker.py" in unit
|
||||
projection = (ROOT / "deploy" / "systemd" / "aios-chatgpt-bridge-v2-agent-task-projection.conf").read_text(encoding="utf-8")
|
||||
assert "AIOS_ONLINE_TASK_PROJECTION_ENABLED=1" in projection
|
||||
assert "LoadCredential=agent_task_control_token:" in projection
|
||||
assert "LoadCredentialEncrypted=" not in projection
|
||||
assert "Bearer " not in projection
|
||||
|
||||
|
||||
def test_durable_task_key_has_collision_resistant_exact_identity() -> None:
|
||||
assert _task_key("task:a") != _task_key("task-a")
|
||||
assert _task_key("x" * 140 + "a") != _task_key("x" * 140 + "b")
|
||||
|
||||
|
||||
def test_existing_binding_rejects_idempotency_head_and_generation_drift(tmp_path: Path) -> None:
|
||||
class Authority:
|
||||
@staticmethod
|
||||
def commit_exists(_oid: str) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def read_ref(_ref: str) -> None:
|
||||
return None
|
||||
|
||||
bootstrap = GovernedTaskBootstrap(Authority(), tmp_path, remote_name="origin", token_path=tmp_path / "token", ssh_key_path=tmp_path / "key")
|
||||
request = {
|
||||
"task_id": TASK,
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"target_branch": BRANCH,
|
||||
"risk": "engineering",
|
||||
"description": "automatic claim infrastructure",
|
||||
"idempotency_key": "p0-claim-online-agent-auto-execution-v1",
|
||||
"expected_branch_oid": BASE,
|
||||
"generation": 1,
|
||||
}
|
||||
bootstrap.create(request)
|
||||
for drift in (
|
||||
{"idempotency_key": "p0-claim-online-agent-auto-execution-v2"},
|
||||
{"expected_branch_oid": "b" * 40},
|
||||
{"generation": 2},
|
||||
):
|
||||
with pytest.raises(BootstrapError, match="TASK_BINDING_CONFLICT"):
|
||||
bootstrap.create({**request, **drift})
|
||||
|
||||
|
||||
def test_existing_branch_without_durable_parent_ownership_is_denied(tmp_path: Path) -> None:
|
||||
class Authority:
|
||||
@staticmethod
|
||||
def commit_exists(_oid: str) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def read_ref(_ref: str) -> str:
|
||||
return BASE
|
||||
|
||||
bootstrap = GovernedTaskBootstrap(Authority(), tmp_path, remote_name="origin", token_path=tmp_path / "token", ssh_key_path=tmp_path / "key")
|
||||
with pytest.raises(BootstrapError, match="EXISTING_BRANCH_OWNERSHIP_REQUIRED"):
|
||||
bootstrap.create({
|
||||
"task_id": TASK,
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"target_branch": BRANCH,
|
||||
"risk": "engineering",
|
||||
"description": "automatic claim infrastructure",
|
||||
"idempotency_key": "p0-claim-online-agent-auto-execution-v1",
|
||||
"expected_branch_oid": BASE,
|
||||
"generation": 2,
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "error"),
|
||||
[
|
||||
("repository", "other/repo", "WRONG_REPOSITORY_ENVELOPE_DENY"),
|
||||
("target_branch", "agent/wrong-branch", "WRONG_BRANCH_ENVELOPE_DENY"),
|
||||
("generation", 2, "WRONG_GENERATION_ENVELOPE_DENY"),
|
||||
],
|
||||
)
|
||||
def test_wrong_lineage_envelope_deny(field: str, value: object, error: str) -> None:
|
||||
envelope = project_execution_envelope(binding())
|
||||
envelope[field] = value
|
||||
with pytest.raises(ExecutionProjectionError, match=error):
|
||||
validate_execution_envelope(envelope, binding=binding())
|
||||
|
||||
|
||||
def test_capability_mismatch_deny() -> None:
|
||||
envelope = project_execution_envelope(
|
||||
binding(),
|
||||
{"action": "source_capture", "capability_scope": ["git.read", "governed.write.file"]},
|
||||
)
|
||||
with pytest.raises(ExecutionProjectionError, match="CAPABILITY_MISMATCH_DENY"):
|
||||
validate_execution_envelope(envelope, binding=binding(), worker_capabilities=["git.read"])
|
||||
|
||||
|
||||
def test_governance_lineage_override_deny() -> None:
|
||||
envelope = project_execution_envelope(binding())
|
||||
envelope["governance"] = {**envelope["governance"], "job_id": "job-wrong-lineage"}
|
||||
with pytest.raises(ExecutionProjectionError, match="JOB_BINDING_CONFLICT"):
|
||||
validate_execution_envelope(envelope, binding=binding())
|
||||
|
||||
|
||||
def test_malformed_persisted_envelope_does_not_fallback() -> None:
|
||||
task = task_without_envelope()
|
||||
task["metadata"] = {"execution": {"schema": "broken"}}
|
||||
bridge = _ProjectedBridge(project_execution_envelope(binding()))
|
||||
supervisor = WorkerSupervisor("runtime-token-not-secret")
|
||||
supervisor.bridge = bridge
|
||||
with pytest.raises(WorkerError, match="EXECUTION_ENVELOPE_SCHEMA_VALIDATION"):
|
||||
supervisor.resolve_envelope(task)
|
||||
assert bridge.calls == 0
|
||||
|
||||
|
||||
def test_self_consistent_persisted_wrong_lineage_is_denied_by_authority() -> None:
|
||||
task = task_without_envelope()
|
||||
wrong = binding()
|
||||
wrong = GovernedExecutionBinding.from_mapping({**wrong.public(), "job_id": "job-00000000000000000000000000000000"})
|
||||
task["metadata"] = {"execution": project_execution_envelope(wrong)}
|
||||
supervisor = WorkerSupervisor("runtime-token-not-secret")
|
||||
supervisor.bridge = _ProjectedBridge(project_execution_envelope(binding()))
|
||||
with pytest.raises(WorkerError, match="TASK_PERSISTED_BINDING_CONFLICT"):
|
||||
supervisor.resolve_envelope(task)
|
||||
|
||||
|
||||
def test_projection_control_missing_denies_before_governed_side_effect() -> None:
|
||||
class Bootstrap:
|
||||
called = False
|
||||
|
||||
def create(self, _request: dict) -> dict:
|
||||
self.called = True
|
||||
raise AssertionError("must not create governed state")
|
||||
|
||||
runtime = GovernedRuntime.__new__(GovernedRuntime)
|
||||
runtime.bootstrap = Bootstrap()
|
||||
runtime.online_task_projector = None
|
||||
with pytest.raises(RuntimeError, match="ONLINE_TASK_PROJECTION_CONTROL_UNAVAILABLE"):
|
||||
GovernedRuntime.create_task(runtime, {"task_id": TASK, "require_online_task_projection": True})
|
||||
assert runtime.bootstrap.called is False
|
||||
|
||||
|
||||
def test_source_capture_projection_cannot_write() -> None:
|
||||
envelope = project_execution_envelope(binding())
|
||||
assert envelope["action"] == "source_capture"
|
||||
assert envelope["change_set"] == []
|
||||
assert envelope["verification"] == []
|
||||
assert envelope["governance"]["writer_lease"] is False
|
||||
assert all(envelope["governance"][field] is False for field in ("main_write", "production_write", "tag_write", "baseline_write", "force_push", "protected_ref_bypass"))
|
||||
|
||||
|
||||
def test_source_capture_state_allows_next_attempt_deterministic_digest(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(worker_module, "STATE_ROOT", tmp_path)
|
||||
state_path = tmp_path / "tasks" / f"{TASK}.json"
|
||||
state_path.parent.mkdir(parents=True)
|
||||
state_path.write_text(json.dumps({"task_id": TASK, "execution_digest": "a" * 64, "phase": "source_capture_done"}), encoding="utf-8")
|
||||
assert load_task_state(TASK, "b" * 64) is None
|
||||
|
||||
|
||||
def test_bridge_exposes_controlled_projection_operation() -> None:
|
||||
seen: list[dict] = []
|
||||
adapter = BridgeV2Adapter(
|
||||
runtime_status={},
|
||||
handlers={"projectOnlineAgentExecution": lambda request: seen.append(dict(request)) or project_execution_envelope(binding())},
|
||||
)
|
||||
response = adapter.projectOnlineAgentExecution({"task_id": TASK})
|
||||
assert response["job_id"] == JOB
|
||||
assert seen == [{"task_id": TASK}]
|
||||
assert BRIDGE_OPERATION_CONTRACTS["projectOnlineAgentExecution"]["path"] == "/v2/execution-envelope"
|
||||
|
||||
|
||||
def test_execution_digest_is_canonical() -> None:
|
||||
envelope = project_execution_envelope(binding())
|
||||
material = {key: value for key, value in envelope.items() if key != "execution_digest"}
|
||||
expected = hashlib.sha256(json.dumps(material, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
|
||||
assert envelope["execution_digest"] == expected
|
||||
|
||||
|
||||
def test_governed_commit_uses_existing_exact_job_without_recreating(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
execution = project_execution_envelope(
|
||||
binding(),
|
||||
{
|
||||
"action": "deterministic",
|
||||
"change_set": [{"operation": "write_file", "path": "result.txt", "content": "ok\n"}],
|
||||
"governance": {"allowed_files": ["result.txt"]},
|
||||
},
|
||||
)
|
||||
envelope = ExecutionEnvelope.from_task({"task_id": TASK, "metadata": {"execution": execution}})
|
||||
(tmp_path / "result.txt").write_text("ok\n", encoding="utf-8")
|
||||
commit_oid = "b" * 40
|
||||
|
||||
class Bridge:
|
||||
def __init__(self) -> None:
|
||||
self.routes: list[str] = []
|
||||
|
||||
def call(self, route: str, body: dict) -> tuple[int, dict]:
|
||||
self.routes.append(route)
|
||||
if route == "/v2/job":
|
||||
assert body == {"job_id": JOB}
|
||||
else:
|
||||
assert body["task_id"] == TASK
|
||||
assert body["job_id"] == JOB
|
||||
if route == "/v2/writer-lease":
|
||||
return 200, {
|
||||
"task_id": TASK,
|
||||
"job_id": JOB,
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"expected_branch_oid": BASE,
|
||||
"branch_mode": "CREATE_NEW",
|
||||
"generation": 1,
|
||||
"target_branch": BRANCH,
|
||||
}
|
||||
if route == "/v2/commit-files":
|
||||
return 200, {
|
||||
"job_id": JOB,
|
||||
"target_branch": BRANCH,
|
||||
"generation": 1,
|
||||
"branch_mode": "CONTINUE_EXISTING",
|
||||
"commit_oid": commit_oid,
|
||||
"commit_visible": True,
|
||||
"receipt": {
|
||||
"remote_verified": True,
|
||||
"authority_oid": commit_oid,
|
||||
"api_branch_oid": commit_oid,
|
||||
"git_remote_oid": commit_oid,
|
||||
},
|
||||
}
|
||||
if route == "/v2/job":
|
||||
return 200, {
|
||||
"state": "COMPLETED",
|
||||
"receipt": {
|
||||
"remote_verified": True,
|
||||
"authority_oid": commit_oid,
|
||||
"api_branch_oid": commit_oid,
|
||||
"git_remote_oid": commit_oid,
|
||||
},
|
||||
"branch_leases": [{
|
||||
"job_id": JOB,
|
||||
"branch": BRANCH,
|
||||
"state": "ACTIVE",
|
||||
"generation": 1,
|
||||
"current_head_oid": commit_oid,
|
||||
}],
|
||||
}
|
||||
raise AssertionError(route)
|
||||
|
||||
bridge = Bridge()
|
||||
monkeypatch.setattr(worker_module.subprocess, "run", lambda *args, **kwargs: SimpleNamespace(returncode=0, stdout=commit_oid + "\n", stderr=""))
|
||||
heartbeat = SimpleNamespace(ensure=lambda: None)
|
||||
result = governed_commit(envelope, tmp_path, bridge, heartbeat)
|
||||
assert bridge.routes == ["/v2/writer-lease", "/v2/commit-files", "/v2/job"]
|
||||
assert result["job_id"] == JOB
|
||||
assert result["commit_oid"] == commit_oid
|
||||
@@ -0,0 +1,244 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
APP = ROOT / "active-runtime/app"
|
||||
sys.path.insert(0, str(APP))
|
||||
|
||||
from aios_online_agent_compact_map import (
|
||||
MAX_RESULT_BYTES,
|
||||
build_compact_interface_map,
|
||||
)
|
||||
from aios_online_agent_execution import (
|
||||
ExecutionProjectionError,
|
||||
GovernedExecutionBinding,
|
||||
project_execution_envelope,
|
||||
validate_execution_envelope,
|
||||
)
|
||||
import aios_online_agent_worker as worker
|
||||
from aios_online_agent_worker import (
|
||||
CompactInterfaceMapExecutor,
|
||||
ExecutionEnvelope,
|
||||
build_compact_interface_map_result,
|
||||
)
|
||||
|
||||
|
||||
BASE = "68b95fea0d9af3b52269e79b7cf878254f191b61"
|
||||
|
||||
|
||||
def make_repo(tmp_path: Path) -> Path:
|
||||
root = tmp_path / "repo"
|
||||
(root / "active-runtime/app").mkdir(parents=True)
|
||||
(root / "tests").mkdir(parents=True)
|
||||
(root / "db/migrations").mkdir(parents=True)
|
||||
(root / "deploy/systemd").mkdir(parents=True)
|
||||
(root / "active-runtime/app/fleet.py").write_text(
|
||||
"""
|
||||
class WorkerRegistry:
|
||||
def register_worker(self, worker_id, capabilities):
|
||||
return {"worker_id": worker_id, "capabilities": capabilities}
|
||||
|
||||
class TaskScheduler:
|
||||
def claim_next(self, tenant_id, capacity):
|
||||
return None
|
||||
|
||||
def recover_assignment(task_id, worker_id):
|
||||
return task_id, worker_id
|
||||
|
||||
TOKEN = "fixturetok"
|
||||
""".strip() + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(root / "active-runtime/app/checkpoint.py").write_text(
|
||||
"""
|
||||
class CheckpointStore:
|
||||
def resume(self, task_id, attempt):
|
||||
return None
|
||||
""".strip() + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(root / "tests/test_fleet.py").write_text(
|
||||
"""
|
||||
def test_scheduler_fairness():
|
||||
assert True
|
||||
|
||||
def test_worker_reassignment():
|
||||
assert True
|
||||
""".strip() + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(root / "db/migrations/20260817_fleet.sql").write_text(
|
||||
"CREATE TABLE IF NOT EXISTS worker_slots (worker_id TEXT, capacity INTEGER);\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(root / "deploy/systemd/aios-fleet-worker.service").write_text(
|
||||
"[Service]\nExecStart=/usr/bin/true\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return root
|
||||
|
||||
|
||||
def binding() -> GovernedExecutionBinding:
|
||||
return GovernedExecutionBinding.from_mapping(
|
||||
{
|
||||
"task_id": "r42-p2-agent-fleet-interface-map-20260817",
|
||||
"job_id": "job-2c451360f78ef595e41663e3b444a8fb",
|
||||
"repository": "aios/aios",
|
||||
"base_oid": BASE,
|
||||
"expected_branch_oid": BASE,
|
||||
"target_branch": "agent/r42-p2-agent-fleet-interface-map-20260817",
|
||||
"branch_mode": "CREATE_NEW",
|
||||
"generation": 1,
|
||||
"idempotency_key": "r42-p2-agent-fleet-interface-map-68b95fe-v1",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def envelope_mapping():
|
||||
return project_execution_envelope(
|
||||
binding(),
|
||||
{
|
||||
"action": "source_capture",
|
||||
"adapter": "compact_interface_map",
|
||||
"capability_scope": ["git.read", "git.test"],
|
||||
"source_signals": [
|
||||
"worker registry scheduler capacity fairness fault isolation",
|
||||
"checkpoint resume retry scope trace message",
|
||||
],
|
||||
},
|
||||
required_capabilities=["git.read", "git.test"],
|
||||
)
|
||||
|
||||
|
||||
def task_for(envelope):
|
||||
return {
|
||||
"task_id": envelope["task_id"],
|
||||
"metadata": {"execution": envelope},
|
||||
"required_capabilities": ["git.read", "git.test"],
|
||||
}
|
||||
|
||||
|
||||
def test_compact_map_bounds_and_surfaces(tmp_path):
|
||||
root = make_repo(tmp_path)
|
||||
result = build_compact_interface_map(
|
||||
root,
|
||||
["worker scheduler capacity fairness checkpoint recovery"],
|
||||
source_commit=BASE,
|
||||
)
|
||||
encoded = json.dumps(result, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()
|
||||
assert result["schema"] == "aios.engineering.interface-map.v1"
|
||||
assert len(encoded) <= MAX_RESULT_BYTES
|
||||
assert len(result["implementation_map"]["modules"]) <= 40
|
||||
assert len(result["bounded_excerpts"]) <= 25
|
||||
assert len(result["implementation_map"]["tests"]) <= 40
|
||||
assert any(item["symbol"] == "WorkerRegistry" for item in result["implementation_map"]["classes"])
|
||||
assert any(item["table"] == "worker_slots" for item in result["implementation_map"]["tables"])
|
||||
assert any(item["unit"] == "aios-fleet-worker.service" for item in result["implementation_map"]["systemd_units"])
|
||||
|
||||
|
||||
def test_compact_map_redacts_sensitive_source_patterns(tmp_path):
|
||||
root = make_repo(tmp_path)
|
||||
result = build_compact_interface_map(root, ["token fleet"], source_commit=BASE)
|
||||
excerpts = json.dumps(result["bounded_excerpts"], ensure_ascii=False)
|
||||
assert "fixturetok" not in excerpts
|
||||
|
||||
|
||||
def test_compact_execution_projection_is_read_only():
|
||||
envelope = envelope_mapping()
|
||||
assert envelope["action"] == "source_capture"
|
||||
assert envelope["adapter"] == "compact_interface_map"
|
||||
assert envelope["governance"]["writer_lease"] is False
|
||||
assert envelope["change_set"] == []
|
||||
assert envelope["verification"] == []
|
||||
validate_execution_envelope(envelope, binding=binding(), worker_capabilities=["git.read", "git.test"])
|
||||
|
||||
|
||||
def test_compact_execution_rejects_mutation():
|
||||
with pytest.raises(ExecutionProjectionError, match="SOURCE_CAPTURE_MUTATION_DENY"):
|
||||
project_execution_envelope(
|
||||
binding(),
|
||||
{
|
||||
"action": "source_capture",
|
||||
"adapter": "compact_interface_map",
|
||||
"change_set": [{"path": "x.py", "operation": "write_file", "content": "x=1\n"}],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_worker_accepts_compact_adapter_and_builds_result(tmp_path, monkeypatch):
|
||||
root = make_repo(tmp_path)
|
||||
monkeypatch.setattr(worker, "audit", lambda *args, **kwargs: None)
|
||||
mapped = envelope_mapping()
|
||||
envelope = ExecutionEnvelope.from_task(
|
||||
task_for(mapped),
|
||||
worker_capabilities=("git.read", "git.test"),
|
||||
)
|
||||
execution = CompactInterfaceMapExecutor().execute(envelope, root)
|
||||
result = build_compact_interface_map_result(envelope.task_id, envelope, execution)
|
||||
assert result["status"] == "SUCCEEDED"
|
||||
assert result["classification"] == "P2_INTERFACE_MAP_CAPTURED"
|
||||
assert result["execution_mode"] == "compact_interface_map"
|
||||
assert result["writer_lease"] == "NONE"
|
||||
assert "source_snapshot" not in result
|
||||
assert result["result_bytes"] <= 120 * 1024
|
||||
assert result["source_commit"] == BASE
|
||||
|
||||
|
||||
def test_compact_contract_keys_are_complete(tmp_path):
|
||||
root = make_repo(tmp_path)
|
||||
result = build_compact_interface_map(root, ["fleet worker scheduler retry"], source_commit=BASE)
|
||||
assert set(result["contracts"]) == {
|
||||
"FLEET_AGENT_REGISTRATION",
|
||||
"FLEET_CAPACITY_SCHEDULING",
|
||||
"FLEET_FAIRNESS",
|
||||
"FLEET_FAULT_ISOLATION",
|
||||
"FLEET_WORKER_REASSIGNMENT",
|
||||
"FLEET_SCOPE_ISOLATION",
|
||||
"FLEET_IDEMPOTENT_DISPATCH",
|
||||
"FLEET_FAILURE_RECOVERY",
|
||||
}
|
||||
|
||||
|
||||
def test_compact_result_has_no_source_archive(tmp_path):
|
||||
root = make_repo(tmp_path)
|
||||
result = build_compact_interface_map(root, ["fleet"], source_commit=BASE)
|
||||
serialized = json.dumps(result)
|
||||
assert "payload_b64" not in serialized
|
||||
assert "tar.gz+base64" not in serialized
|
||||
|
||||
|
||||
def test_create_new_retries_discard_known_policy_reconciliation_marker(tmp_path, monkeypatch):
|
||||
task_id = "r42-p2-compact-interface-map-runtime-20260817"
|
||||
job_id = "job-74d0982f616a7f2653ef346144139044"
|
||||
monkeypatch.setattr(worker, "STATE_ROOT", tmp_path / "worker-state")
|
||||
state_path = worker.task_state_path(task_id)
|
||||
state_path.parent.mkdir(parents=True)
|
||||
state_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"task_id": task_id,
|
||||
"execution_digest": "old-digest",
|
||||
"phase": "reconciliation_required",
|
||||
"attempt": 16,
|
||||
"job_id": job_id,
|
||||
"details": {
|
||||
"api_status": 503,
|
||||
"error": "SECRET_LIKE_PAYLOAD_REJECTED",
|
||||
},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert worker.load_task_state(
|
||||
task_id,
|
||||
"new-digest",
|
||||
attempt=18,
|
||||
continuation_head_oid=None,
|
||||
continuation_job_id=job_id,
|
||||
commit_idempotency_key="r42-p2-compact-interface-map-runtime-commit-v1",
|
||||
) is None
|
||||
@@ -0,0 +1,139 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
APP = ROOT / "active-runtime/app"
|
||||
sys.path.insert(0, str(APP))
|
||||
|
||||
from aios_online_agent_compact_map import build_compact_interface_map
|
||||
import aios_online_agent_worker as worker
|
||||
from aios_online_agent_worker import build_compact_interface_map_result
|
||||
|
||||
|
||||
def _repo(tmp_path: Path) -> Path:
|
||||
root = tmp_path / "repo"
|
||||
(root / "active-runtime/app").mkdir(parents=True)
|
||||
(root / "tests").mkdir(parents=True)
|
||||
(root / "active-runtime/app/model_router.py").write_text(
|
||||
"""
|
||||
class ProviderRegistry:
|
||||
def register_provider(self, provider_id, model_version):
|
||||
return provider_id, model_version
|
||||
|
||||
class ModelRouter:
|
||||
def select_model(self, tenant_id, budget_tokens, fallback_provider):
|
||||
return fallback_provider
|
||||
|
||||
def record_model_cost(trace_id, token_count, cost_usd):
|
||||
return trace_id, token_count, cost_usd
|
||||
|
||||
def evaluate_model_version(model_id, version, approval_gate):
|
||||
return approval_gate
|
||||
""".strip() + "\\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(root / "tests/test_model_router.py").write_text(
|
||||
"""
|
||||
def test_model_budget_and_fallback():
|
||||
assert True
|
||||
""".strip() + "\\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return root
|
||||
|
||||
|
||||
def _envelope(base_oid: str):
|
||||
return SimpleNamespace(
|
||||
base_oid=base_oid,
|
||||
job_id="job-model-governance",
|
||||
repository="aios/aios",
|
||||
expected_branch_oid=base_oid,
|
||||
target_branch="agent/model-governance-map",
|
||||
generation=1,
|
||||
digest="e" * 64,
|
||||
)
|
||||
|
||||
|
||||
def test_model_governance_profile_contracts_and_classification(tmp_path):
|
||||
root = _repo(tmp_path)
|
||||
value = build_compact_interface_map(
|
||||
root,
|
||||
[
|
||||
"model provider registry routing budget token cost fallback",
|
||||
"evaluation approval version pinning tenant project credential trace",
|
||||
],
|
||||
source_commit="a" * 40,
|
||||
)
|
||||
assert value["profile"] == "model_governance"
|
||||
assert value["classification"] == "MODEL_GOVERNANCE_INTERFACE_MAP_CAPTURED"
|
||||
assert set(value["contracts"]) == {
|
||||
"MODEL_PROVIDER_REGISTRY",
|
||||
"MODEL_SELECTION_POLICY",
|
||||
"MODEL_VERSION_PINNING",
|
||||
"MODEL_BUDGET_ACCOUNTING",
|
||||
"MODEL_FALLBACK_POLICY",
|
||||
"MODEL_SCOPE_ISOLATION",
|
||||
"MODEL_CREDENTIAL_BOUNDARY",
|
||||
"MODEL_TRACE_AUDIT",
|
||||
"MODEL_EVALUATION_GATE",
|
||||
"MODEL_ROUTING_RESILIENCE",
|
||||
}
|
||||
assert "FLEET_FAIRNESS" not in value["contracts"]
|
||||
|
||||
|
||||
def test_fleet_profile_is_backward_compatible(tmp_path):
|
||||
root = _repo(tmp_path)
|
||||
value = build_compact_interface_map(
|
||||
root,
|
||||
["agent fleet scheduler worker capacity fairness"],
|
||||
source_commit="b" * 40,
|
||||
)
|
||||
assert value["profile"] == "fleet"
|
||||
assert value["classification"] == "P2_INTERFACE_MAP_CAPTURED"
|
||||
assert "FLEET_AGENT_REGISTRATION" in value["contracts"]
|
||||
|
||||
|
||||
def test_worker_binds_model_profile_to_model_classification():
|
||||
interface_map = {
|
||||
"schema": "aios.engineering.interface-map.v1",
|
||||
"profile": "model_governance",
|
||||
"classification": "MODEL_GOVERNANCE_INTERFACE_MAP_CAPTURED",
|
||||
"source_commit": "c" * 40,
|
||||
"implementation_map": {"modules": [], "classes": [], "functions": [], "tests": []},
|
||||
"bounded_excerpts": [],
|
||||
"contracts": {"MODEL_PROVIDER_REGISTRY": {"status": "GAP", "surfaces": []}},
|
||||
"result_bytes": 512,
|
||||
}
|
||||
result = build_compact_interface_map_result(
|
||||
"r42-model-governance-interface-map-20260818",
|
||||
_envelope("c" * 40),
|
||||
{"interface_map": interface_map},
|
||||
)
|
||||
assert result["classification"] == "MODEL_GOVERNANCE_INTERFACE_MAP_CAPTURED"
|
||||
assert result["interface_map_profile"] == "model_governance"
|
||||
assert result["writer_lease"] == "NONE"
|
||||
|
||||
|
||||
def test_worker_rejects_profile_classification_mismatch():
|
||||
interface_map = {
|
||||
"schema": "aios.engineering.interface-map.v1",
|
||||
"profile": "model_governance",
|
||||
"classification": "P2_INTERFACE_MAP_CAPTURED",
|
||||
"source_commit": "c" * 40,
|
||||
"implementation_map": {"modules": [], "classes": [], "functions": [], "tests": []},
|
||||
"bounded_excerpts": [],
|
||||
"contracts": {},
|
||||
"result_bytes": 512,
|
||||
}
|
||||
with pytest.raises(worker.WorkerError, match="COMPACT_INTERFACE_MAP_PROFILE_BINDING_INVALID"):
|
||||
build_compact_interface_map_result(
|
||||
"r42-model-governance-interface-map-20260818",
|
||||
_envelope("c" * 40),
|
||||
{"interface_map": interface_map},
|
||||
)
|
||||
@@ -0,0 +1,216 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
APP=ROOT/"active-runtime/app"
|
||||
sys.path.insert(0,str(APP))
|
||||
|
||||
from aios_online_agent_checkpoint import (
|
||||
CheckpointBinding,
|
||||
CheckpointError,
|
||||
CheckpointStore,
|
||||
)
|
||||
import aios_online_agent_worker as worker
|
||||
|
||||
|
||||
def binding(attempt: int = 1, *, tenant_head: str = "b"*40):
|
||||
return CheckpointBinding(
|
||||
task_id="task-p1-a",
|
||||
attempt=attempt,
|
||||
execution_digest="d"*64,
|
||||
job_id="job-p1-a",
|
||||
repository="aios/aios",
|
||||
base_oid="a"*40,
|
||||
expected_branch_oid=tenant_head,
|
||||
target_branch="agent/task-p1-a",
|
||||
generation=1,
|
||||
)
|
||||
|
||||
|
||||
def test_CHECKPOINT_DURABLE_WRITE(tmp_path):
|
||||
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
||||
first=store.append(binding(), "claimed", {"phase":"claimed","attempt":1})
|
||||
second=store.append(binding(), "workspace_prepared", {"phase":"workspace_prepared","attempt":1})
|
||||
assert first["sequence"] == 1
|
||||
assert second["sequence"] == 2
|
||||
latest=store.latest("task-p1-a")
|
||||
assert latest["phase"]=="workspace_prepared"
|
||||
assert latest["payload"]["attempt"]==1
|
||||
|
||||
|
||||
def test_CHECKPOINT_SCOPE_ISOLATION(tmp_path):
|
||||
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
||||
store.append(binding(), "claimed", {"phase":"claimed"})
|
||||
changed=CheckpointBinding(
|
||||
task_id="task-p1-a",attempt=1,execution_digest="d"*64,job_id="job-p1-a",
|
||||
repository="aios/other",base_oid="a"*40,expected_branch_oid="b"*40,
|
||||
target_branch="agent/task-p1-a",generation=1,
|
||||
)
|
||||
with pytest.raises(CheckpointError,match="CHECKPOINT_SCOPE_MISMATCH_DENY"):
|
||||
store.append(changed,"workspace_prepared",{"phase":"workspace_prepared"})
|
||||
|
||||
|
||||
def test_RESUME_FROM_CHECKPOINT_across_attempt_same_binding(tmp_path):
|
||||
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
||||
store.append(binding(1), "source_capture_done", {
|
||||
"phase":"source_capture_done",
|
||||
"attempt":1,
|
||||
"result":{"status":"SUCCEEDED","classification":"NEEDS_REWORK_SOURCE_CAPTURE"},
|
||||
})
|
||||
resumed=store.verify_latest_binding(binding(2))
|
||||
assert resumed is not None
|
||||
assert resumed["attempt"]==1
|
||||
assert resumed["phase"]=="source_capture_done"
|
||||
|
||||
|
||||
def test_checkpoint_new_attempt_can_start_from_claimed(tmp_path):
|
||||
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
||||
store.append(binding(1),"result_submitted",{"phase":"result_submitted"})
|
||||
cp=store.append(binding(2),"claimed",{"phase":"claimed"})
|
||||
assert cp["attempt"]==2
|
||||
assert cp["sequence"]==1
|
||||
|
||||
|
||||
def envelope():
|
||||
return worker.ExecutionEnvelope(
|
||||
task_id="task-p1-a",
|
||||
job_id="job-p1-a",
|
||||
repository="aios/aios",
|
||||
base_oid="a"*40,
|
||||
expected_branch_oid="b"*40,
|
||||
target_branch="agent/task-p1-a",
|
||||
branch_mode="CONTINUE_EXISTING",
|
||||
generation=1,
|
||||
idempotency_key="governed-p1-a",
|
||||
action="source_capture",
|
||||
adapter="source_capture",
|
||||
capability_scope=("git.read","git.test"),
|
||||
source_signals=(),
|
||||
change_set=(),
|
||||
verification=(),
|
||||
governance={"task_id":"task-p1-a","job_id":"job-p1-a"},
|
||||
workspace={},
|
||||
digest="d"*64,
|
||||
)
|
||||
|
||||
|
||||
def test_RESULT_SUBMISSION_REPLAY_SAFE_attempt_aware_key():
|
||||
env=envelope()
|
||||
result={"status":"SUCCEEDED","execution_digest":env.digest}
|
||||
one=worker.result_idempotency_key(env.task_id,1,env,result)
|
||||
two=worker.result_idempotency_key(env.task_id,2,env,result)
|
||||
assert one != two
|
||||
assert "-a1-result-" in one
|
||||
assert "-a2-result-" in two
|
||||
assert one == worker.result_idempotency_key(env.task_id,1,env,result)
|
||||
|
||||
|
||||
class RetryClient:
|
||||
def __init__(self):
|
||||
self.calls=[]
|
||||
def result(self,task_id,lease_token,key,result):
|
||||
self.calls.append((task_id,key,dict(result)))
|
||||
return (599,{"error":"SERVICE_UNAVAILABLE"}) if len(self.calls)==1 else (201,{"status":"VERIFYING"})
|
||||
def get(self,task_id):
|
||||
return 500,{}
|
||||
|
||||
|
||||
def test_RETRY_IDEMPOTENT_NO_DUPLICATE_SIDE_EFFECT_reuses_same_key(monkeypatch):
|
||||
client=RetryClient()
|
||||
monkeypatch.setattr(worker.time,"sleep",lambda _:None)
|
||||
status,response=worker.submit_result_replay_safe(
|
||||
client,"task-p1-a","lease-token-"+"x"*32,"stable-key",{"status":"SUCCEEDED"}
|
||||
)
|
||||
assert status==201
|
||||
assert [item[1] for item in client.calls]==["stable-key","stable-key"]
|
||||
|
||||
|
||||
class ConflictClient:
|
||||
def __init__(self,result):
|
||||
self.result_value=result
|
||||
def result(self,*args):
|
||||
return 409,{"error":"IDEMPOTENCY_CONFLICT"}
|
||||
def get(self,task_id):
|
||||
return 200,{"status":"VERIFYING","result":dict(self.result_value)}
|
||||
|
||||
|
||||
def test_RESULT_SUBMISSION_REPLAY_SAFE_reconciles_exact_result():
|
||||
result={"status":"SUCCEEDED","value":1}
|
||||
status,response=worker.submit_result_replay_safe(
|
||||
ConflictClient(result),"task-p1-a","lease-token-"+"x"*32,"stable-key",result
|
||||
)
|
||||
assert status==200
|
||||
assert response["idempotent"] is True
|
||||
assert response["reconciled"] is True
|
||||
|
||||
|
||||
def test_RESULT_SUBMISSION_REPLAY_SAFE_rejects_conflicting_result():
|
||||
status,response=worker.submit_result_replay_safe(
|
||||
ConflictClient({"status":"SUCCEEDED","value":2}),
|
||||
"task-p1-a","lease-token-"+"x"*32,"stable-key",
|
||||
{"status":"SUCCEEDED","value":1},
|
||||
)
|
||||
assert status==409
|
||||
assert response["error"]=="IDEMPOTENCY_CONFLICT"
|
||||
|
||||
|
||||
def test_WORKSPACE_CLEANUP_RECOVERY(tmp_path,monkeypatch):
|
||||
root=tmp_path/"workspaces"; root.mkdir()
|
||||
stale=root/"task-old"; stale.mkdir(); (stale/"x").write_text("x")
|
||||
monkeypatch.setattr(worker,"WORKSPACE_ROOT",root)
|
||||
monkeypatch.setattr(worker,"STATE_ROOT",tmp_path/"state")
|
||||
removed=worker.recover_stale_workspaces()
|
||||
assert removed==1
|
||||
assert not stale.exists()
|
||||
|
||||
|
||||
class LeaseBridge:
|
||||
def __init__(self):
|
||||
self.calls=0
|
||||
def call(self,route,body):
|
||||
assert route=="/v2/writer-lease"
|
||||
self.calls+=1
|
||||
if self.calls==1:
|
||||
return 599,{"error":"SERVICE_UNAVAILABLE"}
|
||||
return 200,{
|
||||
"lease_id":"lease-p1-a",
|
||||
"job_id":"job-p1-a",
|
||||
"task_id":"task-p1-a",
|
||||
"repository":"aios/aios",
|
||||
"base_oid":"a"*40,
|
||||
"expected_branch_oid":"b"*40,
|
||||
"branch_mode":"CONTINUE_EXISTING",
|
||||
"generation":1,
|
||||
"target_branch":"agent/task-p1-a",
|
||||
"state":"ACTIVE",
|
||||
}
|
||||
|
||||
|
||||
class LeaseHeartbeatStub:
|
||||
def ensure(self):
|
||||
return None
|
||||
|
||||
|
||||
def test_LEASE_RENEWAL_RECOVERY_retries_same_binding(monkeypatch):
|
||||
env=envelope()
|
||||
env=worker.ExecutionEnvelope(
|
||||
task_id=env.task_id,job_id=env.job_id,repository=env.repository,
|
||||
base_oid=env.base_oid,expected_branch_oid=env.expected_branch_oid,
|
||||
target_branch=env.target_branch,branch_mode=env.branch_mode,generation=env.generation,
|
||||
idempotency_key=env.idempotency_key,action=env.action,adapter=env.adapter,
|
||||
capability_scope=env.capability_scope,source_signals=env.source_signals,
|
||||
change_set=env.change_set,verification=env.verification,
|
||||
governance={"task_id":env.task_id,"job_id":env.job_id,"lease_seconds":900},
|
||||
workspace=env.workspace,digest=env.digest,
|
||||
)
|
||||
bridge=LeaseBridge()
|
||||
monkeypatch.setattr(worker.time,"sleep",lambda _:None)
|
||||
monkeypatch.setattr(worker,"STATE_ROOT",Path("/tmp")/"aios-p1-test-state")
|
||||
response=worker.acquire_writer_lease_replay_safe(env,bridge,LeaseHeartbeatStub())
|
||||
assert response["lease_id"]=="lease-p1-a"
|
||||
assert bridge.calls==2
|
||||
@@ -1,245 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
APP=ROOT/"active-runtime/app"
|
||||
sys.path.insert(0,str(APP))
|
||||
|
||||
from aios_online_agent_execution import GovernedExecutionBinding, project_execution_envelope
|
||||
from aios_governed_task_bootstrap import GovernedTaskBinding, GovernedTaskBootstrap
|
||||
from aios_gitops_job_controller import GitOpsJobController, JobState, JobStore
|
||||
from aios_continuation_executor import ContinuationExecutionError, DetachedGitWorkspace
|
||||
from aios_online_agent_worker import (
|
||||
BridgeClient,
|
||||
IsolatedDeterministicExecutor,
|
||||
load_task_state,
|
||||
_claim_order,
|
||||
_isolated_executor_wait_seconds,
|
||||
_verify_converged_job,
|
||||
WorkerError,
|
||||
)
|
||||
|
||||
|
||||
BASE="a642b045775b9154a00ae78543037fbef4fb447b"
|
||||
HEAD="71523aef6e48b1e2fa516c456ca189b7a039104b"
|
||||
TASK="p0-claim-online-agent-auto-execution-20260817"
|
||||
JOB="job-2031a9e5682a4238f98e3ac7aee8a593"
|
||||
BRANCH="agent/p0-claim-online-agent-auto-execution-20260817"
|
||||
|
||||
|
||||
def test_source_signals_are_bounded_by_projection():
|
||||
binding=GovernedExecutionBinding.from_mapping({
|
||||
"task_id":TASK,"job_id":JOB,"repository":"aios/aios","base_oid":BASE,
|
||||
"expected_branch_oid":BASE,"target_branch":BRANCH,"branch_mode":"CREATE_NEW",
|
||||
"generation":1,"idempotency_key":"p0-claim-online-agent-auto-execution-a642b04-v1",
|
||||
})
|
||||
execution=project_execution_envelope(binding,{
|
||||
"action":"source_capture",
|
||||
"source_signals":["x"*5000,"x"*5000,"trace credential tool "*100],
|
||||
})
|
||||
assert len(execution["source_signals"]) <= 16
|
||||
assert all(0 < len(item) <= 512 for item in execution["source_signals"])
|
||||
|
||||
|
||||
def test_same_job_create_new_binding_normalizes_to_current_head():
|
||||
binding=GovernedTaskBinding(
|
||||
task_id=TASK,job_id=JOB,repository="aios/aios",base_oid=BASE,
|
||||
expected_branch_oid=BASE,target_branch=BRANCH,branch_mode="CREATE_NEW",
|
||||
generation=1,idempotency_key="p0-claim-online-agent-auto-execution-a642b04-v1",
|
||||
risk="engineering",description="automatic claim",
|
||||
)
|
||||
bootstrap=GovernedTaskBootstrap.__new__(GovernedTaskBootstrap)
|
||||
bootstrap.authority=SimpleNamespace(read_ref=lambda ref: HEAD)
|
||||
class Connection:
|
||||
def __enter__(self): return self
|
||||
def __exit__(self,*a): return False
|
||||
def execute(self,*a,**k):
|
||||
return SimpleNamespace(fetchall=lambda:[{
|
||||
"current_head_oid":HEAD,"generation":1,"state":"ACTIVE","lineage_anchor_oid":BASE
|
||||
}])
|
||||
bootstrap._store=lambda _task:(SimpleNamespace(_connect=lambda:Connection()),None)
|
||||
normalized=bootstrap.execution_binding(binding)
|
||||
assert normalized.base_oid == BASE
|
||||
assert normalized.expected_branch_oid == HEAD
|
||||
assert normalized.branch_mode == "CONTINUE_EXISTING"
|
||||
assert normalized.generation == 1
|
||||
|
||||
|
||||
def test_recovery_receipt_requires_exact_oid_and_active_lease():
|
||||
envelope=SimpleNamespace(job_id=JOB,target_branch=BRANCH,generation=1)
|
||||
response={
|
||||
"state":"COMPLETED",
|
||||
"receipt":{"remote_verified":True,"authority_oid":HEAD,"api_branch_oid":HEAD,"git_remote_oid":HEAD},
|
||||
"branch_leases":[{
|
||||
"job_id":JOB,"branch":BRANCH,"state":"ACTIVE","generation":1,"current_head_oid":HEAD
|
||||
}],
|
||||
}
|
||||
_verify_converged_job(response,envelope,HEAD)
|
||||
bad={**response,"receipt":{**response["receipt"],"git_remote_oid":BASE}}
|
||||
with pytest.raises(WorkerError,match="GOVERNED_RECEIPT_OID_MISMATCH"):
|
||||
_verify_converged_job(bad,envelope,HEAD)
|
||||
|
||||
|
||||
def test_executor_unit_is_uncredentialed():
|
||||
unit=(ROOT/"deploy/systemd/aios-online-agent-isolated-executor.service").read_text()
|
||||
assert "PrivateNetwork=true" in unit
|
||||
assert "RestrictAddressFamilies=AF_UNIX AF_INET" in unit
|
||||
assert "IPAddressDeny=any" in unit
|
||||
assert "IPAddressAllow=localhost" in unit
|
||||
assert "InaccessiblePaths=/run/credentials" in unit
|
||||
assert "LoadCredential" not in unit
|
||||
assert "SupplementaryGroups=gitops" not in unit
|
||||
|
||||
|
||||
def test_isolated_executor_wait_budget_includes_safety_margin():
|
||||
assert _isolated_executor_wait_seconds(({"timeout_seconds":1200},)) == 1230
|
||||
|
||||
|
||||
def test_prepared_deterministic_plan_precedes_source_capture_at_same_priority():
|
||||
deterministic=SimpleNamespace(action="deterministic",task_id="task-deterministic")
|
||||
source_capture=SimpleNamespace(action="source_capture",task_id="task-source-capture")
|
||||
task={"priority":100,"created_at":"2026-08-17T00:00:00Z"}
|
||||
ordered=sorted(((task,source_capture),(task,deterministic)),key=_claim_order)
|
||||
assert ordered[0][1] is deterministic
|
||||
|
||||
|
||||
def test_plan_registry_gets_precedence_when_task_has_no_execution(monkeypatch):
|
||||
client=BridgeClient("agent-token-placeholder")
|
||||
seen=[]
|
||||
monkeypatch.setattr(client,"call",lambda route,body:(seen.append((route,body)) or (200,{})))
|
||||
task={"task_id":TASK,"required_capabilities":["git.read"]}
|
||||
client.project_execution(task)
|
||||
assert seen == [("/v2/execution-envelope",{"task_id":TASK,"required_capabilities":["git.read"]})]
|
||||
client.project_execution(task,{"action":"source_capture","source_signals":["trace"]})
|
||||
assert seen[-1][1]["execution"]["action"] == "source_capture"
|
||||
|
||||
|
||||
def test_executor_exchange_files_are_group_readable_and_directories_setgid(tmp_path):
|
||||
request=tmp_path/"request.json"
|
||||
IsolatedDeterministicExecutor._atomic(request,b"{}")
|
||||
assert request.stat().st_mode & 0o777 == 0o640
|
||||
tmpfiles=(ROOT/"deploy/tmpfiles.d/aios-online-agent-executor.conf").read_text()
|
||||
assert tmpfiles.count(" 2770 ") == 5
|
||||
|
||||
|
||||
def test_precommit_state_recovery_requires_a_later_official_attempt(monkeypatch,tmp_path):
|
||||
import aios_online_agent_worker as worker_module
|
||||
monkeypatch.setattr(worker_module,"STATE_ROOT",tmp_path)
|
||||
state_path=tmp_path/"tasks"/f"{TASK}.json"
|
||||
state_path.parent.mkdir(parents=True)
|
||||
state_path.write_text(json.dumps({"task_id":TASK,"execution_digest":"a"*64,"phase":"workspace_prepared"}))
|
||||
assert load_task_state(TASK,"b"*64) is None
|
||||
state_path.write_text(json.dumps({"task_id":TASK,"execution_digest":"a"*64,"phase":"verification_passed","attempt":16}))
|
||||
with pytest.raises(WorkerError,match="TASK_PAYLOAD_DRIFT"):
|
||||
load_task_state(TASK,"b"*64,attempt=16)
|
||||
assert load_task_state(TASK,"b"*64,attempt=17) is None
|
||||
|
||||
|
||||
def test_reconciliation_state_recovery_requires_later_attempt_and_same_job(monkeypatch,tmp_path):
|
||||
import aios_online_agent_worker as worker_module
|
||||
monkeypatch.setattr(worker_module,"STATE_ROOT",tmp_path)
|
||||
state_path=tmp_path/"tasks"/f"{TASK}.json"
|
||||
state_path.parent.mkdir(parents=True)
|
||||
state_path.write_text(json.dumps({
|
||||
"task_id":TASK,
|
||||
"execution_digest":"a"*64,
|
||||
"phase":"reconciliation_required",
|
||||
"attempt":18,
|
||||
"job_id":JOB,
|
||||
"details":{"api_status":503,"error":"WORKTREE_DIRTY_SCOPE_MISMATCH"},
|
||||
}))
|
||||
with pytest.raises(WorkerError,match="TASK_PAYLOAD_DRIFT"):
|
||||
load_task_state(TASK,"b"*64,attempt=18,continuation_job_id=JOB,continuation_head_oid=HEAD)
|
||||
assert load_task_state(TASK,"b"*64,attempt=19,continuation_job_id=JOB,continuation_head_oid=HEAD) is None
|
||||
with pytest.raises(WorkerError,match="TASK_PAYLOAD_DRIFT"):
|
||||
load_task_state(TASK,"b"*64,attempt=19,continuation_job_id="job-other",continuation_head_oid=HEAD)
|
||||
|
||||
|
||||
def test_legacy_dirty_scope_reconciliation_state_is_exactly_bounded(monkeypatch,tmp_path):
|
||||
import aios_online_agent_worker as worker_module
|
||||
monkeypatch.setattr(worker_module,"STATE_ROOT",tmp_path)
|
||||
state_path=tmp_path/"tasks"/f"{TASK}.json"
|
||||
state_path.parent.mkdir(parents=True)
|
||||
state_path.write_text(json.dumps({
|
||||
"task_id":TASK,
|
||||
"execution_digest":"a"*64,
|
||||
"phase":"reconciliation_required",
|
||||
"job_id":JOB,
|
||||
"details":{"api_status":503,"error":"WORKTREE_DIRTY_SCOPE_MISMATCH"},
|
||||
}))
|
||||
assert load_task_state(TASK,"b"*64,attempt=19,continuation_job_id=JOB,continuation_head_oid=HEAD) is None
|
||||
with pytest.raises(WorkerError,match="TASK_PAYLOAD_DRIFT"):
|
||||
load_task_state(TASK,"b"*64,attempt=19,continuation_job_id=JOB,continuation_head_oid="not-an-oid")
|
||||
|
||||
|
||||
def test_completed_same_job_begins_one_idempotent_continuation_cycle(tmp_path):
|
||||
controller=GitOpsJobController(JobStore(tmp_path/"jobs",payload_root=tmp_path/"payloads"))
|
||||
submitted=controller.submit({
|
||||
"idempotency_key":"same-job-continuation-test",
|
||||
"campaign_id":TASK,
|
||||
"task_id":TASK,
|
||||
"repository":"aios/aios",
|
||||
"operation":"GOVERNED_TASK_CREATE",
|
||||
"branch":BRANCH,
|
||||
"expected_branch_oid":HEAD,
|
||||
"payload_digest":"a"*64,
|
||||
})
|
||||
job_id=submitted.job_id
|
||||
controller.advance(job_id,JobState.AUTHORIZED,side_effect_state="AUTHORIZED",expected_revision=0)
|
||||
controller.lease_job(
|
||||
job_id,
|
||||
lease_id="lease-"+job_id,
|
||||
lineage_anchor_oid=BASE,
|
||||
current_head_oid=HEAD,
|
||||
generation=1,
|
||||
lease_expires_at="2099-01-01T00:00:00+00:00",
|
||||
expected_revision=1,
|
||||
)
|
||||
controller.advance(job_id,JobState.WORKTREE_READY,side_effect_state="WORKTREE_READY",expected_revision=2)
|
||||
controller.advance(job_id,JobState.EXECUTING,side_effect_state="IN_PROGRESS",expected_revision=3)
|
||||
controller.advance(job_id,JobState.LOCAL_VERIFIED,side_effect_state="LOCAL_VERIFIED",receipt={"authority_oid":HEAD},expected_revision=4)
|
||||
controller.advance(job_id,JobState.REMOTE_SYNCING,side_effect_state="IN_PROGRESS",expected_revision=5)
|
||||
receipt={
|
||||
"remote_verified":True,
|
||||
"commit_visible":True,
|
||||
"authority_oid":HEAD,
|
||||
"api_branch_oid":HEAD,
|
||||
"git_remote_oid":HEAD,
|
||||
}
|
||||
controller.advance(job_id,JobState.REMOTE_VERIFIED,side_effect_state="REMOTE_VERIFIED",receipt=receipt,expected_revision=6)
|
||||
controller.advance(job_id,JobState.COMPLETED,side_effect_state="REMOTE_VERIFIED",receipt=receipt,expected_revision=7)
|
||||
first=controller.begin_continuation_cycle(
|
||||
job_id,
|
||||
idempotency_key="same-job-continuation-commit-v1",
|
||||
expected_head_oid=HEAD,
|
||||
generation=1,
|
||||
expected_revision=8,
|
||||
)
|
||||
assert first == {"completed":False,"commit_oid":None,"idempotent":False}
|
||||
assert controller.store.read(job_id)["state"] == "LEASED"
|
||||
replay=controller.begin_continuation_cycle(
|
||||
job_id,
|
||||
idempotency_key="same-job-continuation-commit-v1",
|
||||
expected_head_oid=HEAD,
|
||||
generation=1,
|
||||
expected_revision=8,
|
||||
)
|
||||
assert replay == {"completed":False,"commit_oid":None,"idempotent":True}
|
||||
|
||||
|
||||
def test_deterministic_materialization_recovers_only_an_in_scope_dirty_subset(monkeypatch,tmp_path):
|
||||
workspace=object.__new__(DetachedGitWorkspace)
|
||||
workspace.path=tmp_path
|
||||
monkeypatch.setattr(DetachedGitWorkspace,"_dirty_paths",lambda _self:{"old.py"})
|
||||
workspace.materialize_files({"old.py":b"new-old\n","new.py":b"new\n"})
|
||||
assert (tmp_path/"old.py").read_bytes() == b"new-old\n"
|
||||
assert (tmp_path/"new.py").read_bytes() == b"new\n"
|
||||
monkeypatch.setattr(DetachedGitWorkspace,"_dirty_paths",lambda _self:{"outside.py"})
|
||||
with pytest.raises(ContinuationExecutionError,match="WORKTREE_DIRTY_SCOPE_MISMATCH"):
|
||||
workspace.materialize_files({"old.py":b"again\n"})
|
||||
@@ -0,0 +1,487 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from active_runtime.app.aios_credential_broker import CredentialBroker
|
||||
from active_runtime.app.agent_observability_hooks import AgentObservabilityHooks
|
||||
from active_runtime.app.message_handoff import DurableHandoffStore, MessageEnvelope
|
||||
from active_runtime.app.search_gateway import (
|
||||
SearchEvidence,
|
||||
SearchGateway,
|
||||
SearchProviderAttempt,
|
||||
SearchRequest,
|
||||
)
|
||||
from active_runtime.app.search_gateway_orchestrator import SearchGatewayOrchestrator
|
||||
from active_runtime.app.trace_tool_credential_governance import (
|
||||
CredentialGrant,
|
||||
GovernanceActor,
|
||||
GovernanceError,
|
||||
GovernanceScope,
|
||||
GovernanceStore,
|
||||
ToolDescriptor,
|
||||
TraceContext,
|
||||
bind_trace_lineage,
|
||||
)
|
||||
|
||||
|
||||
TENANT_SCOPE = GovernanceScope(
|
||||
"tenant",
|
||||
tenant_id="tenant-a",
|
||||
user_id="user-a",
|
||||
project_id="project-a",
|
||||
conversation_id="conversation-a",
|
||||
run_id="run-a",
|
||||
)
|
||||
|
||||
|
||||
def actor(*capabilities: str, tenant_id: str = "tenant-a") -> GovernanceActor:
|
||||
return GovernanceActor(
|
||||
"actor-a",
|
||||
GovernanceScope(
|
||||
"tenant",
|
||||
tenant_id=tenant_id,
|
||||
user_id="user-a",
|
||||
project_id="project-a",
|
||||
conversation_id="conversation-a",
|
||||
run_id="run-a",
|
||||
),
|
||||
tuple(capabilities),
|
||||
)
|
||||
|
||||
|
||||
def owner() -> GovernanceActor:
|
||||
return GovernanceActor("owner-control", GovernanceScope("system"), ("system.admin",))
|
||||
|
||||
|
||||
def message(**changes):
|
||||
values = dict(
|
||||
source="task-plane",
|
||||
target="worker-a",
|
||||
message_type="tool.invoke",
|
||||
payload={"query": "safe"},
|
||||
sequence=1,
|
||||
idempotency_key="message-a-1",
|
||||
tenant_id="tenant-a",
|
||||
user_id="user-a",
|
||||
project_id="project-a",
|
||||
conversation_id="conversation-a",
|
||||
task_id="task-a",
|
||||
run_id="run-a",
|
||||
confidentiality_scope="run",
|
||||
)
|
||||
values.update(changes)
|
||||
return MessageEnvelope(**values)
|
||||
|
||||
|
||||
def task(**changes):
|
||||
value = {
|
||||
"task_id": "task-a",
|
||||
"claim_id": "claim-a",
|
||||
"correlation_id": "correlation-a",
|
||||
"scope_type": "tenant",
|
||||
"tenant_id": "tenant-a",
|
||||
"user_id": "user-a",
|
||||
"project_id": "project-a",
|
||||
"conversation_id": "conversation-a",
|
||||
"run_id": "run-a",
|
||||
}
|
||||
value.update(changes)
|
||||
return value
|
||||
|
||||
|
||||
def root_trace():
|
||||
return bind_trace_lineage(
|
||||
actor=actor("tool.search"),
|
||||
task=task(),
|
||||
message=message(),
|
||||
trace_id="trace-a",
|
||||
span_id="span-root",
|
||||
)
|
||||
|
||||
|
||||
def test_trace_lineage_continuity_and_durable_restart(tmp_path: Path) -> None:
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
root = root_trace()
|
||||
child = root.child(span_id="span-tool", actor=actor("tool.search"))
|
||||
store.record_trace(root, actor("tool.search"))
|
||||
store.record_trace(child, actor("tool.search"))
|
||||
|
||||
reopened = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
assert reopened.verify_audit_chain(("TRACE_RECORDED",))
|
||||
assert child.trace_id == root.trace_id
|
||||
assert child.parent_span_id == root.span_id
|
||||
assert child.task_id == "task-a"
|
||||
assert child.claim_id == "claim-a"
|
||||
|
||||
|
||||
def test_trace_lineage_continuity_rejects_task_change() -> None:
|
||||
with pytest.raises(GovernanceError, match="TRACE_LINEAGE_CONTINUITY"):
|
||||
bind_trace_lineage(
|
||||
actor=actor("tool.search"),
|
||||
task=task(task_id="task-b"),
|
||||
message=message(),
|
||||
trace_id="trace-a",
|
||||
span_id="span-root",
|
||||
)
|
||||
|
||||
|
||||
def test_trace_scope_mismatch_deny() -> None:
|
||||
with pytest.raises(GovernanceError, match="TRACE_SCOPE_MISMATCH_DENY"):
|
||||
bind_trace_lineage(
|
||||
actor=actor("tool.search"),
|
||||
task=task(),
|
||||
message=message(tenant_id="tenant-b"),
|
||||
trace_id="trace-a",
|
||||
span_id="span-root",
|
||||
)
|
||||
|
||||
|
||||
def test_message_correlation_is_identity_bound() -> None:
|
||||
first = message(correlation_id="correlation-a")
|
||||
second = message(correlation_id="correlation-b")
|
||||
assert first.envelope_id != second.envelope_id
|
||||
assert first.as_record()["correlation_id"] == "correlation-a"
|
||||
|
||||
|
||||
def test_message_handoff_legacy_schema_migrates_without_changing_id(tmp_path: Path) -> None:
|
||||
import sqlite3
|
||||
|
||||
path = tmp_path / "legacy.sqlite3"
|
||||
legacy = message()
|
||||
record = legacy.as_record()
|
||||
record.pop("correlation_id")
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE message_handoff (
|
||||
envelope_id TEXT PRIMARY KEY, stream_key TEXT NOT NULL, source TEXT NOT NULL,
|
||||
target TEXT NOT NULL, message_type TEXT NOT NULL, payload_json TEXT NOT NULL,
|
||||
payload_sha256 TEXT NOT NULL, tenant_id TEXT, user_id TEXT, project_id TEXT,
|
||||
conversation_id TEXT, task_id TEXT, run_id TEXT, causation_id TEXT,
|
||||
sequence INTEGER NOT NULL, capability_scope_json TEXT NOT NULL,
|
||||
confidentiality_scope TEXT NOT NULL, idempotency_key TEXT NOT NULL,
|
||||
expires_at TEXT, created_at TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
delivered_at TEXT, acked_at TEXT,
|
||||
UNIQUE(stream_key,idempotency_key), UNIQUE(stream_key,sequence)
|
||||
);
|
||||
"""
|
||||
)
|
||||
columns = tuple(record)
|
||||
connection.execute(
|
||||
f"INSERT INTO message_handoff ({','.join(columns)}) VALUES ({','.join('?' for _ in columns)})",
|
||||
tuple(record[name] for name in columns),
|
||||
)
|
||||
store = DurableHandoffStore(path)
|
||||
loaded = store.get(
|
||||
legacy.envelope_id,
|
||||
actor={
|
||||
"tenant_id": "tenant-a", "user_id": "user-a", "project_id": "project-a",
|
||||
"conversation_id": "conversation-a", "run_id": "run-a",
|
||||
"capabilities": [],
|
||||
},
|
||||
)
|
||||
assert loaded.envelope_id == legacy.envelope_id
|
||||
assert loaded.correlation_id is None
|
||||
|
||||
|
||||
def test_message_handoff_rejects_raw_credential_material(tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError, match="CREDENTIAL_SECRET_EXPOSURE_DENY"):
|
||||
message(payload={"value": "Bearer " + "x" * 24})
|
||||
DurableHandoffStore(tmp_path / "handoff.sqlite3")
|
||||
import sqlite3
|
||||
with sqlite3.connect(tmp_path / "handoff.sqlite3") as connection:
|
||||
assert connection.execute("SELECT count(*) FROM message_handoff").fetchone()[0] == 0
|
||||
|
||||
|
||||
def test_tool_mcp_authz_deny_and_scope_propagation(tmp_path: Path) -> None:
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
store.register_tool(
|
||||
ToolDescriptor("mcp.search", "mcp", TENANT_SCOPE, ("tool.search",)),
|
||||
owner(),
|
||||
)
|
||||
trace = root_trace()
|
||||
store.record_trace(trace, actor("tool.search"))
|
||||
|
||||
with pytest.raises(GovernanceError, match="TOOL_MCP_AUTHZ_DENY"):
|
||||
store.authorize_tool(
|
||||
actor=actor(), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-deny-capability",
|
||||
)
|
||||
with pytest.raises(GovernanceError, match="TOOL_SCOPE_PROPAGATION"):
|
||||
store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=GovernanceScope("tenant", tenant_id="tenant-b"),
|
||||
input_digest="a" * 64, invocation_key="invoke-deny-scope",
|
||||
)
|
||||
|
||||
invocation = store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-search-a",
|
||||
)
|
||||
assert invocation["trace_id"] == "trace-a"
|
||||
assert invocation["evidence"]["action"] == "TOOL_AUTHORIZED"
|
||||
|
||||
|
||||
def test_tool_requires_recorded_trace_and_receipt_is_idempotent(tmp_path: Path) -> None:
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
store.register_tool(
|
||||
ToolDescriptor(
|
||||
"mcp.search", "mcp", TENANT_SCOPE, ("tool.search",),
|
||||
credential_ref="systemd-credential:search-provider-a",
|
||||
),
|
||||
owner(),
|
||||
)
|
||||
trace = root_trace()
|
||||
with pytest.raises(GovernanceError, match="TOOL_MCP_AUTHZ_DENY:trace"):
|
||||
store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-unrecorded",
|
||||
)
|
||||
store.record_trace(trace, actor("tool.search"))
|
||||
authorized = store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-replay-a",
|
||||
)
|
||||
assert authorized["credential_required"] is True
|
||||
assert "credential_ref" not in authorized
|
||||
receipt = store.record_tool_receipt(
|
||||
invocation_key="invoke-replay-a", result_digest="b" * 64,
|
||||
actor=actor("tool.search"), trace=trace,
|
||||
)
|
||||
assert receipt["status"] == "VERIFIED"
|
||||
replay = store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-replay-a",
|
||||
)
|
||||
assert replay["idempotent"] is True
|
||||
assert replay["receipt_sha256"] == "b" * 64
|
||||
assert store.verify_audit_chain(
|
||||
("TRACE_RECORDED", "TOOL_AUTHORIZED", "TOOL_RESULT_VERIFIED")
|
||||
)
|
||||
|
||||
|
||||
def test_system_and_global_resources_do_not_use_fake_tenants(tmp_path: Path) -> None:
|
||||
with pytest.raises(GovernanceError, match="GOVERNANCE_NON_TENANT_SCOPE_INVALID"):
|
||||
GovernanceScope("system", tenant_id="fake-tenant")
|
||||
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
descriptor = ToolDescriptor("system.health", "tool", GovernanceScope("global"), ("system.admin",))
|
||||
registered = store.register_tool(descriptor, owner())
|
||||
assert registered["descriptor"]["scope_type"] == "global"
|
||||
assert registered["descriptor"]["tenant_id"] is None
|
||||
|
||||
|
||||
def test_credential_scope_and_secret_exposure_deny(tmp_path: Path) -> None:
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
grant = CredentialGrant(
|
||||
"grant-search-a",
|
||||
"systemd-credential:search-provider-a",
|
||||
TENANT_SCOPE,
|
||||
("search.invoke",),
|
||||
task_id="task-a",
|
||||
tool_name="mcp.search",
|
||||
expires_at=(datetime.now(timezone.utc) + timedelta(minutes=5)).isoformat(),
|
||||
)
|
||||
stored = store.store_grant(grant, owner())
|
||||
assert "secret" not in stored["grant"]
|
||||
trace = root_trace()
|
||||
store.record_trace(trace, actor("tool.search"))
|
||||
|
||||
with pytest.raises(GovernanceError, match="CREDENTIAL_SCOPE_DENY"):
|
||||
store.authorize_credential(
|
||||
grant_id=grant.grant_id,
|
||||
actor=actor("tool.search", tenant_id="tenant-b"),
|
||||
trace=trace,
|
||||
operation="search.invoke",
|
||||
task_id="task-a",
|
||||
tool_name="mcp.search",
|
||||
)
|
||||
with pytest.raises(GovernanceError, match="CREDENTIAL_SECRET_EXPOSURE_DENY"):
|
||||
CredentialGrant("grant-raw", "Bearer raw-secret", TENANT_SCOPE, ("search.invoke",))
|
||||
|
||||
authorized = store.authorize_credential(
|
||||
grant_id=grant.grant_id,
|
||||
actor=actor("tool.search"),
|
||||
trace=trace,
|
||||
operation="search.invoke",
|
||||
task_id="task-a",
|
||||
tool_name="mcp.search",
|
||||
)
|
||||
assert authorized["credential_ref"] == "systemd-credential:search-provider-a"
|
||||
assert "credential_ref" not in authorized["evidence"]["details"]
|
||||
|
||||
|
||||
def test_audit_evidence_completeness_and_tamper_detection(tmp_path: Path) -> None:
|
||||
path = tmp_path / "governance.sqlite3"
|
||||
store = GovernanceStore(path)
|
||||
trace = root_trace()
|
||||
store.record_trace(trace, actor("tool.search"))
|
||||
store.register_tool(
|
||||
ToolDescriptor("mcp.search", "mcp", TENANT_SCOPE, ("tool.search",)), owner()
|
||||
)
|
||||
store.authorize_tool(
|
||||
actor=actor("tool.search"), tool_name="mcp.search", trace=trace,
|
||||
requested_scope=TENANT_SCOPE, input_digest="a" * 64,
|
||||
invocation_key="invoke-audit-a",
|
||||
)
|
||||
assert store.verify_audit_chain(("TRACE_RECORDED", "TOOL_REGISTERED", "TOOL_AUTHORIZED"))
|
||||
assert not store.verify_audit_chain(("CREDENTIAL_AUTHORIZED",))
|
||||
|
||||
import sqlite3
|
||||
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.execute(
|
||||
"UPDATE governance_audit_evidence SET outcome='DENY' WHERE sequence=1"
|
||||
)
|
||||
assert not GovernanceStore(path).verify_audit_chain()
|
||||
|
||||
|
||||
def test_observability_record_carries_governed_trace(tmp_path: Path) -> None:
|
||||
from active_runtime.app.agent_observability_store import AgentObservabilityStore
|
||||
|
||||
hooks = AgentObservabilityHooks(AgentObservabilityStore(tmp_path / "observability"))
|
||||
trace = root_trace()
|
||||
started = hooks.start(
|
||||
"user-a", "project-a", "openclaw",
|
||||
actor=actor("tool.search"), trace=trace,
|
||||
)
|
||||
record = hooks.finish(started["execution_id"])["record"]
|
||||
assert record["trace_id"] == "trace-a"
|
||||
assert record["tenant_id"] == "tenant-a"
|
||||
assert record["task_id"] == "task-a"
|
||||
assert record["message_id"] == message().envelope_id
|
||||
|
||||
|
||||
def test_active_runtime_credential_broker_scope_and_raw_secret_deny() -> None:
|
||||
class Provider:
|
||||
def environment(self, lease):
|
||||
return {"ACCESS_TOKEN": "raw-secret"}
|
||||
|
||||
broker = CredentialBroker(Provider())
|
||||
lease = broker.issue(
|
||||
task_id="task-a", repository="aios/aios", target_branch="agent/task-a",
|
||||
generation=1, scopes=("tool.invoke",), scope_type="tenant",
|
||||
tenant_id="tenant-a", user_id="user-a", project_id="project-a", run_id="run-a",
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="CREDENTIAL_SCOPE_DENY"):
|
||||
broker.authorize(
|
||||
lease["lease_id"], task_id="task-a", operation="tool.invoke",
|
||||
scope_type="tenant", tenant_id="tenant-b", user_id="user-a",
|
||||
project_id="project-a", run_id="run-a",
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="CREDENTIAL_SECRET_EXPOSURE_DENY"):
|
||||
broker.authorize(
|
||||
lease["lease_id"], task_id="task-a", operation="tool.invoke",
|
||||
scope_type="tenant", tenant_id="tenant-a", user_id="user-a",
|
||||
project_id="project-a", run_id="run-a",
|
||||
)
|
||||
|
||||
|
||||
def test_search_gateway_uses_governed_tool_boundary(tmp_path: Path) -> None:
|
||||
class Provider:
|
||||
name = "official_site"
|
||||
called = False
|
||||
|
||||
def search(self, request, queries):
|
||||
self.called = True
|
||||
return (
|
||||
[SearchEvidence("Official", "https://example.test", self.name, "safe")],
|
||||
SearchProviderAttempt(self.name, "ok", 1),
|
||||
)
|
||||
|
||||
store = GovernanceStore(tmp_path / "governance.sqlite3")
|
||||
store.register_tool(
|
||||
ToolDescriptor("search.official_site", "tool", TENANT_SCOPE, ("tool.search",)),
|
||||
owner(),
|
||||
)
|
||||
provider = Provider()
|
||||
gateway = SearchGateway([provider], governance_store=store)
|
||||
trace = root_trace()
|
||||
store.record_trace(trace, actor("tool.search"))
|
||||
response = gateway.search(
|
||||
SearchRequest(
|
||||
query="AIOS", requester="actor-a", correlation_id="trace-a",
|
||||
authorized_sources=["official_site"],
|
||||
),
|
||||
actor=actor("tool.search"),
|
||||
trace=trace,
|
||||
)
|
||||
assert response.status == "found"
|
||||
assert provider.called is True
|
||||
assert store.verify_audit_chain(
|
||||
("TOOL_REGISTERED", "TRACE_RECORDED", "TOOL_AUTHORIZED", "TOOL_RESULT_VERIFIED")
|
||||
)
|
||||
|
||||
with pytest.raises(GovernanceError, match="TOOL_MCP_AUTHZ_DENY"):
|
||||
gateway.search(
|
||||
SearchRequest(query="AIOS", requester="forged", correlation_id="trace-a"),
|
||||
actor=actor("tool.search"),
|
||||
trace=trace,
|
||||
)
|
||||
|
||||
|
||||
def test_search_gateway_redacts_provider_exception() -> None:
|
||||
class Provider:
|
||||
name = "general_web"
|
||||
|
||||
def search(self, request, queries):
|
||||
raise RuntimeError("Bearer should-not-leak")
|
||||
|
||||
response = SearchGateway([Provider()]).search(SearchRequest(query="AIOS", requester="legacy"))
|
||||
assert response.attempts[0].reason == "PROVIDER_FAILURE"
|
||||
assert "should-not-leak" not in str(response.to_dict())
|
||||
|
||||
|
||||
def test_search_runtime_default_path_requires_governance(tmp_path: Path) -> None:
|
||||
orchestrator = SearchGatewayOrchestrator(
|
||||
governance_store=GovernanceStore(tmp_path / "search-governance.sqlite3")
|
||||
)
|
||||
assert orchestrator.status()["governance"] == "required"
|
||||
with pytest.raises(GovernanceError, match="TOOL_MCP_AUTHZ_DENY"):
|
||||
orchestrator.execute({"query": "AIOS", "requester": "actor-a"})
|
||||
|
||||
|
||||
def test_governance_newer_schema_fails_closed(tmp_path: Path) -> None:
|
||||
import sqlite3
|
||||
path = tmp_path / "newer.sqlite3"
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.execute(
|
||||
"CREATE TABLE governance_schema_metadata(singleton INTEGER PRIMARY KEY, schema_version INTEGER NOT NULL)"
|
||||
)
|
||||
connection.execute(
|
||||
"INSERT INTO governance_schema_metadata(singleton,schema_version) VALUES(1,999)"
|
||||
)
|
||||
with pytest.raises(GovernanceError, match="GOVERNANCE_SCHEMA_VERSION_UNSUPPORTED"):
|
||||
GovernanceStore(path)
|
||||
|
||||
|
||||
def test_changed_production_files_have_no_embedded_secret_material() -> None:
|
||||
import re
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
paths = (
|
||||
"active-runtime/app/agent_event_bus.py",
|
||||
"active-runtime/app/agent_event_observability_bridge.py",
|
||||
"active-runtime/app/agent_observability.py",
|
||||
"active-runtime/app/agent_observability_api.py",
|
||||
"active-runtime/app/agent_observability_hooks.py",
|
||||
"active-runtime/app/agent_observability_store.py",
|
||||
"active-runtime/app/aios_credential_broker.py",
|
||||
"active-runtime/app/message_handoff.py",
|
||||
"active-runtime/app/search_agent_adapter.py",
|
||||
"active-runtime/app/search_gateway.py",
|
||||
"active-runtime/app/search_gateway_api.py",
|
||||
"active-runtime/app/search_gateway_orchestrator.py",
|
||||
"active-runtime/app/trace_tool_credential_governance.py",
|
||||
)
|
||||
patterns = (
|
||||
re.compile(r"BEGIN (?:RSA|OPENSSH|EC) PRIVATE KEY"),
|
||||
re.compile(r"Bearer\s+[A-Za-z0-9._~+/-]{24,}"),
|
||||
)
|
||||
hits = [path for path in paths for pattern in patterns if pattern.search((root / path).read_text(encoding="utf-8"))]
|
||||
assert hits == []
|
||||
@@ -1,73 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
def die(msg: str) -> None:
|
||||
print(f"ERROR: {msg}", file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
|
||||
def backup(path: pathlib.Path) -> pathlib.Path:
|
||||
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
out = path.with_name(path.name + f".bak.{stamp}")
|
||||
shutil.copy2(path, out)
|
||||
return out
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Repair misplaced subprocess import in execution-rebind runtime."
|
||||
)
|
||||
ap.add_argument("--rebind", required=True)
|
||||
args = ap.parse_args()
|
||||
path = pathlib.Path(args.rebind).resolve()
|
||||
if not path.is_file():
|
||||
die(f"missing rebind source: {path}")
|
||||
|
||||
text = path.read_text(encoding="utf-8")
|
||||
marker = "from __future__ import annotations\n"
|
||||
if marker not in text:
|
||||
die("future import marker missing")
|
||||
|
||||
lines = text.splitlines(keepends=True)
|
||||
filtered = [line for line in lines if line.strip() != "import subprocess"]
|
||||
text2 = "".join(filtered)
|
||||
|
||||
future_index = text2.find(marker)
|
||||
if future_index < 0:
|
||||
die("future import marker missing after normalization")
|
||||
|
||||
insert_at = future_index + len(marker)
|
||||
normalized = text2[:insert_at] + "\nimport subprocess\n" + text2[insert_at:]
|
||||
|
||||
if normalized == text:
|
||||
print("OK import order already normalized")
|
||||
else:
|
||||
b = backup(path)
|
||||
path.write_text(normalized, encoding="utf-8")
|
||||
print(f"PATCHED import order; backup={b}")
|
||||
|
||||
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
print(f"OK syntax: {path}")
|
||||
|
||||
current = path.read_text(encoding="utf-8")
|
||||
required_markers = (
|
||||
"EXECUTION_REBIND_NON_FAST_FORWARD_DENY",
|
||||
"old_expected_branch_oid",
|
||||
"new_expected_branch_oid",
|
||||
)
|
||||
missing = [m for m in required_markers if m not in current]
|
||||
if missing:
|
||||
die("re-anchor patch markers missing: " + ",".join(missing))
|
||||
|
||||
print("RESULT: IMPORT_ORDER_REPAIRED")
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,172 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
from datetime import datetime, timezone
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
def die(msg: str) -> None:
|
||||
print(f"ERROR: {msg}", file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
|
||||
def backup(path: pathlib.Path) -> pathlib.Path:
|
||||
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
out = path.with_name(path.name + f".bak.{stamp}")
|
||||
shutil.copy2(path, out)
|
||||
return out
|
||||
|
||||
def ensure_import(path: pathlib.Path, import_line: str) -> None:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
if import_line + "\n" in text:
|
||||
return
|
||||
lines = text.splitlines(keepends=True)
|
||||
at = 0
|
||||
if lines and lines[0].startswith("#!"):
|
||||
at = 1
|
||||
if at < len(lines) and lines[at].startswith("from __future__ import"):
|
||||
at += 1
|
||||
lines.insert(at, import_line + "\n")
|
||||
path.write_text("".join(lines), encoding="utf-8")
|
||||
|
||||
def replace_once(path: pathlib.Path, old: str, new: str, label: str) -> None:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
if new in text:
|
||||
print(f"OK already patched: {label}")
|
||||
return
|
||||
count = text.count(old)
|
||||
if count != 1:
|
||||
die(f"{label}: expected exactly one match, found {count}")
|
||||
b = backup(path)
|
||||
path.write_text(text.replace(old, new, 1), encoding="utf-8")
|
||||
print(f"PATCHED {label}; backup={b}")
|
||||
|
||||
def syntax_check(path: pathlib.Path) -> None:
|
||||
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
print(f"OK syntax: {path}")
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description="Add guarded same-job execution re-anchor support.")
|
||||
ap.add_argument("--rebind", required=True)
|
||||
args = ap.parse_args()
|
||||
path = pathlib.Path(args.rebind).resolve()
|
||||
if not path.is_file():
|
||||
die(f"missing rebind source: {path}")
|
||||
|
||||
ensure_import(path, "import subprocess")
|
||||
|
||||
old = ''' binding = GovernedExecutionBinding.from_mapping(old_execution)
|
||||
resolved = dict(self.resolver(task_id, capabilities))
|
||||
new_execution = validate_execution_envelope(resolved, binding=binding)
|
||||
if new_execution.get("action") != "deterministic":
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DETERMINISTIC_PLAN_REQUIRED", 409)
|
||||
if new_execution["execution_digest"] != request.expected_new_execution_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_NEW_DIGEST_DRIFT", 409)
|
||||
if new_execution["execution_digest"] == old_execution["execution_digest"]:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DIGEST_UNCHANGED", 409)
|
||||
'''
|
||||
new = ''' binding = GovernedExecutionBinding.from_mapping(old_execution)
|
||||
resolved = dict(self.resolver(task_id, capabilities))
|
||||
candidate_binding = GovernedExecutionBinding.from_mapping(resolved)
|
||||
|
||||
# Same-job continuation may advance only the branch head (and may
|
||||
# transition CREATE_NEW -> CONTINUE_EXISTING after the branch exists).
|
||||
immutable_fields = (
|
||||
"task_id",
|
||||
"job_id",
|
||||
"repository",
|
||||
"base_oid",
|
||||
"target_branch",
|
||||
"generation",
|
||||
"idempotency_key",
|
||||
)
|
||||
for field in immutable_fields:
|
||||
if getattr(candidate_binding, field) != getattr(binding, field):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_BINDING_CONFLICT", 409)
|
||||
|
||||
head_advanced = candidate_binding.expected_branch_oid != binding.expected_branch_oid
|
||||
if head_advanced:
|
||||
allowed_modes = {
|
||||
(binding.branch_mode, binding.branch_mode),
|
||||
("CREATE_NEW", "CONTINUE_EXISTING"),
|
||||
}
|
||||
if (binding.branch_mode, candidate_binding.branch_mode) not in allowed_modes:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_BRANCH_MODE_CONFLICT", 409)
|
||||
|
||||
authority_git_dir = "/data/aios-gitops/repo.git"
|
||||
authority_ref = f"refs/heads/{binding.target_branch}"
|
||||
try:
|
||||
authority_head = subprocess.check_output(
|
||||
[
|
||||
"/usr/bin/git",
|
||||
f"--git-dir={authority_git_dir}",
|
||||
"rev-parse",
|
||||
authority_ref,
|
||||
],
|
||||
text=True,
|
||||
timeout=5,
|
||||
).strip()
|
||||
ancestry = subprocess.run(
|
||||
[
|
||||
"/usr/bin/git",
|
||||
f"--git-dir={authority_git_dir}",
|
||||
"merge-base",
|
||||
"--is-ancestor",
|
||||
binding.expected_branch_oid,
|
||||
candidate_binding.expected_branch_oid,
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=5,
|
||||
check=False,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_AUTHORITY_GIT_UNAVAILABLE", 503) from exc
|
||||
|
||||
if authority_head != candidate_binding.expected_branch_oid:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_CURRENT_HEAD_MISMATCH", 409)
|
||||
if ancestry.returncode != 0:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_NON_FAST_FORWARD_DENY", 409)
|
||||
elif candidate_binding.branch_mode != binding.branch_mode:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_BRANCH_MODE_CONFLICT", 409)
|
||||
|
||||
# Validate the canonical envelope on its own after the continuation
|
||||
# binding proof above. Worker-side strict equality remains unchanged.
|
||||
new_execution = validate_execution_envelope(resolved)
|
||||
|
||||
# Re-anchor may legitimately remain source_capture when no deterministic
|
||||
# plan is registered yet. A non-head-changing rebind must still advance
|
||||
# source_capture -> deterministic.
|
||||
if new_execution.get("action") not in {"source_capture", "deterministic"}:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTION_INVALID", 409)
|
||||
if not head_advanced and new_execution.get("action") != "deterministic":
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DETERMINISTIC_PLAN_REQUIRED", 409)
|
||||
if new_execution["execution_digest"] != request.expected_new_execution_digest:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_NEW_DIGEST_DRIFT", 409)
|
||||
if new_execution["execution_digest"] == old_execution["execution_digest"]:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_DIGEST_UNCHANGED", 409)
|
||||
'''
|
||||
replace_once(path, old, new, "same-job fast-forward execution re-anchor")
|
||||
|
||||
old_event = ''' "old_action": old_execution.get("action"),
|
||||
"new_action": new_execution.get("action"),
|
||||
"previous_revision": request.expected_revision,
|
||||
'''
|
||||
new_event = ''' "old_action": old_execution.get("action"),
|
||||
"new_action": new_execution.get("action"),
|
||||
"old_expected_branch_oid": binding.expected_branch_oid,
|
||||
"new_expected_branch_oid": candidate_binding.expected_branch_oid,
|
||||
"head_advanced": head_advanced,
|
||||
"previous_revision": request.expected_revision,
|
||||
'''
|
||||
replace_once(path, old_event, new_event, "re-anchor audit evidence")
|
||||
syntax_check(path)
|
||||
print("RESULT: PATCH_OK")
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,141 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
from datetime import datetime, timezone
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
def die(msg: str) -> None:
|
||||
print(f"ERROR: {msg}", file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
|
||||
def backup(path: pathlib.Path) -> pathlib.Path:
|
||||
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
out = path.with_name(path.name + f".bak.{stamp}")
|
||||
shutil.copy2(path, out)
|
||||
return out
|
||||
|
||||
def ensure_datetime_import(path: pathlib.Path) -> bool:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
target = "from datetime import datetime, timezone\n"
|
||||
if target in text:
|
||||
return False
|
||||
lines = text.splitlines(keepends=True)
|
||||
at = 0
|
||||
if lines and lines[0].startswith("#!"):
|
||||
at = 1
|
||||
if at < len(lines) and lines[at].startswith("from __future__ import"):
|
||||
at += 1
|
||||
b = backup(path)
|
||||
lines.insert(at, target)
|
||||
path.write_text("".join(lines), encoding="utf-8")
|
||||
print(f"PATCHED datetime import: {path}; backup={b}")
|
||||
return True
|
||||
|
||||
def replace_once(path: pathlib.Path, old: str, new: str, label: str) -> bool:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
if new in text:
|
||||
print(f"OK already patched: {label}: {path}")
|
||||
return False
|
||||
n = text.count(old)
|
||||
if n != 1:
|
||||
die(f"{label}: expected exactly 1 match, found {n}: {path}")
|
||||
b = backup(path)
|
||||
path.write_text(text.replace(old, new, 1), encoding="utf-8")
|
||||
print(f"PATCHED {label}: {path}; backup={b}")
|
||||
return True
|
||||
|
||||
def syntax_check(path: pathlib.Path) -> None:
|
||||
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
print(f"OK syntax: {path}")
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--rebind", required=True)
|
||||
ap.add_argument("--worker", required=True)
|
||||
ap.add_argument("--bridge", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
rebind = pathlib.Path(args.rebind).resolve()
|
||||
worker = pathlib.Path(args.worker).resolve()
|
||||
bridge = pathlib.Path(args.bridge).resolve()
|
||||
for p in (rebind, worker, bridge):
|
||||
if not p.is_file():
|
||||
die(f"missing file: {p}")
|
||||
|
||||
old_rebind = ''' if task["lease_token_hash"] not in (None, "") or task["lease_expires_at"] not in (None, ""):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
|
||||
'''
|
||||
new_rebind = ''' # A persisted lease token is not by itself proof of an active lease.
|
||||
lease_expires_at = task["lease_expires_at"]
|
||||
if lease_expires_at not in (None, ""):
|
||||
try:
|
||||
parsed_lease_expires_at = datetime.fromisoformat(
|
||||
str(lease_expires_at).replace("Z", "+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_LEASE_EXPIRY_INVALID", 409) from exc
|
||||
if parsed_lease_expires_at > datetime.now(timezone.utc):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
|
||||
'''
|
||||
|
||||
old_worker = ''' active = [
|
||||
item for item in leases
|
||||
if isinstance(item, Mapping)
|
||||
and item.get("state") == "ACTIVE"
|
||||
and item.get("job_id") == envelope.job_id
|
||||
and item.get("branch") == envelope.target_branch
|
||||
]
|
||||
'''
|
||||
new_worker = ''' now = datetime.now(timezone.utc)
|
||||
active = []
|
||||
for item in leases:
|
||||
if (
|
||||
not isinstance(item, Mapping)
|
||||
or item.get("state") != "ACTIVE"
|
||||
or item.get("job_id") != envelope.job_id
|
||||
or item.get("branch") != envelope.target_branch
|
||||
):
|
||||
continue
|
||||
try:
|
||||
expires_at = datetime.fromisoformat(
|
||||
str(item.get("lease_expires_at") or "").replace("Z", "+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if expires_at > now:
|
||||
active.append(item)
|
||||
'''
|
||||
|
||||
old_bridge = ''' if durable_expires_at <= datetime.now(timezone.utc):
|
||||
return self.renew_writer_lease(request)
|
||||
'''
|
||||
new_bridge = ''' # Expired durable lease is recoverable on the same task/job/generation.
|
||||
if durable_expires_at <= datetime.now(timezone.utc):
|
||||
return self.renew_writer_lease(request)
|
||||
'''
|
||||
|
||||
changed = []
|
||||
for p in (rebind, worker):
|
||||
if ensure_datetime_import(p):
|
||||
changed.append(p)
|
||||
if replace_once(rebind, old_rebind, new_rebind, "execution-rebind effective lease expiry"):
|
||||
changed.append(rebind)
|
||||
if replace_once(worker, old_worker, new_worker, "worker effective lease expiry"):
|
||||
changed.append(worker)
|
||||
if replace_once(bridge, old_bridge, new_bridge, "bridge actual-runtime marker"):
|
||||
changed.append(bridge)
|
||||
|
||||
for p in (rebind, worker, bridge):
|
||||
syntax_check(p)
|
||||
print("RESULT: PATCH_OK")
|
||||
print("CHANGED:", ",".join(dict.fromkeys(str(p) for p in changed)) if changed else "none")
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,63 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
from datetime import datetime, timezone
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
def die(msg: str) -> None:
|
||||
print(f"ERROR: {msg}", file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
|
||||
def backup(path: pathlib.Path) -> pathlib.Path:
|
||||
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
out = path.with_name(path.name + f".bak.{stamp}")
|
||||
shutil.copy2(path, out)
|
||||
return out
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description="Allow pre-claim execution rebind at attempt=0.")
|
||||
ap.add_argument("--rebind", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
path = pathlib.Path(args.rebind).resolve()
|
||||
if not path.is_file():
|
||||
die(f"missing rebind source: {path}")
|
||||
|
||||
text = path.read_text(encoding="utf-8")
|
||||
old = " or attempt < 1\n"
|
||||
new = " or attempt < 0\n"
|
||||
|
||||
if new in text:
|
||||
print("OK already patched: pre-claim attempt=0")
|
||||
else:
|
||||
count = text.count(old)
|
||||
if count != 1:
|
||||
die(f"expected exactly one attempt<1 guard, found {count}")
|
||||
b = backup(path)
|
||||
path.write_text(text.replace(old, new, 1), encoding="utf-8")
|
||||
print(f"PATCHED pre-claim attempt=0; backup={b}")
|
||||
|
||||
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
|
||||
current = path.read_text(encoding="utf-8")
|
||||
for marker in (
|
||||
"or attempt < 0",
|
||||
"EXECUTION_REBIND_NON_FAST_FORWARD_DENY",
|
||||
"old_expected_branch_oid",
|
||||
"new_expected_branch_oid",
|
||||
):
|
||||
if marker not in current:
|
||||
die(f"required marker missing: {marker}")
|
||||
|
||||
print(f"OK syntax: {path}")
|
||||
print("RESULT: PRECLAIM_REBIND_PATCH_OK")
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,145 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import datetime as dt
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
UTC = dt.timezone.utc
|
||||
|
||||
def die(msg: str) -> None:
|
||||
print(f"ERROR: {msg}", file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
|
||||
def backup(path: pathlib.Path) -> pathlib.Path:
|
||||
stamp = dt.datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
out = path.with_suffix(path.suffix + f".bak.{stamp}")
|
||||
shutil.copy2(path, out)
|
||||
return out
|
||||
|
||||
def replace_once(path: pathlib.Path, old: str, new: str, label: str) -> bool:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
if new in text:
|
||||
print(f"OK already patched: {label}")
|
||||
return False
|
||||
count = text.count(old)
|
||||
if count != 1:
|
||||
die(f"{label}: expected exactly 1 match, found {count} in {path}")
|
||||
backup_path = backup(path)
|
||||
path.write_text(text.replace(old, new, 1), encoding="utf-8")
|
||||
print(f"PATCHED {label}; backup={backup_path}")
|
||||
return True
|
||||
|
||||
def ensure_datetime_import(path: pathlib.Path) -> bool:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
target = "from datetime import datetime, timezone\n"
|
||||
if target in text:
|
||||
return False
|
||||
lines = text.splitlines(keepends=True)
|
||||
insert_at = 0
|
||||
if lines and lines[0].startswith("#!"):
|
||||
insert_at = 1
|
||||
if insert_at < len(lines) and lines[insert_at].startswith("from __future__ import"):
|
||||
insert_at += 1
|
||||
lines.insert(insert_at, target)
|
||||
backup_path = backup(path)
|
||||
path.write_text("".join(lines), encoding="utf-8")
|
||||
print(f"PATCHED datetime import: {path}; backup={backup_path}")
|
||||
return True
|
||||
|
||||
def syntax_check(path: pathlib.Path) -> None:
|
||||
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
print(f"OK syntax: {path}")
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description="Repair stale governed writer lease handling.")
|
||||
ap.add_argument("--app-dir", default="/opt/aios/current/active-runtime/app")
|
||||
args = ap.parse_args()
|
||||
app = pathlib.Path(args.app_dir)
|
||||
|
||||
rebind = app / "aios_agent_task_execution_rebind.py"
|
||||
worker = app / "aios_online_agent_worker.py"
|
||||
bridge = app / "aios_chatgpt_bridge_v2_server.py"
|
||||
for p in (rebind, worker, bridge):
|
||||
if not p.is_file():
|
||||
die(f"missing file: {p}")
|
||||
|
||||
old_rebind = ''' if task["lease_token_hash"] not in (None, "") or task["lease_expires_at"] not in (None, ""):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
|
||||
'''
|
||||
new_rebind = ''' # A persisted lease token is not by itself proof of an active lease.
|
||||
# Treat an unexpired lease as active; stale lease projection is recoverable.
|
||||
lease_expires_at = task["lease_expires_at"]
|
||||
if lease_expires_at not in (None, ""):
|
||||
try:
|
||||
parsed_lease_expires_at = datetime.fromisoformat(
|
||||
str(lease_expires_at).replace("Z", "+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_LEASE_EXPIRY_INVALID", 409) from exc
|
||||
if parsed_lease_expires_at > datetime.now(timezone.utc):
|
||||
raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
|
||||
'''
|
||||
|
||||
old_worker = ''' active = [
|
||||
item for item in leases
|
||||
if isinstance(item, Mapping)
|
||||
and item.get("state") == "ACTIVE"
|
||||
and item.get("job_id") == envelope.job_id
|
||||
and item.get("branch") == envelope.target_branch
|
||||
]
|
||||
'''
|
||||
new_worker = ''' now = datetime.now(timezone.utc)
|
||||
active = []
|
||||
for item in leases:
|
||||
if (
|
||||
not isinstance(item, Mapping)
|
||||
or item.get("state") != "ACTIVE"
|
||||
or item.get("job_id") != envelope.job_id
|
||||
or item.get("branch") != envelope.target_branch
|
||||
):
|
||||
continue
|
||||
try:
|
||||
expires_at = datetime.fromisoformat(
|
||||
str(item.get("lease_expires_at") or "").replace("Z", "+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if expires_at > now:
|
||||
active.append(item)
|
||||
'''
|
||||
|
||||
old_bridge = ''' if durable_expires_at <= datetime.now(timezone.utc):
|
||||
return self.renew_writer_lease(request)
|
||||
'''
|
||||
new_bridge = ''' # An expired durable lease is recoverable on the same governed
|
||||
# task/job/generation. Renew it instead of treating stale ACTIVE state
|
||||
# as a durable-binding failure.
|
||||
if durable_expires_at <= datetime.now(timezone.utc):
|
||||
return self.renew_writer_lease(request)
|
||||
'''
|
||||
|
||||
changed = []
|
||||
if ensure_datetime_import(rebind):
|
||||
changed.append(rebind)
|
||||
if ensure_datetime_import(worker):
|
||||
changed.append(worker)
|
||||
if replace_once(rebind, old_rebind, new_rebind, "execution-rebind effective lease expiry"):
|
||||
changed.append(rebind)
|
||||
if replace_once(worker, old_worker, new_worker, "worker converged-job effective lease expiry"):
|
||||
changed.append(worker)
|
||||
if replace_once(bridge, old_bridge, new_bridge, "bridge expired-lease recovery contract"):
|
||||
changed.append(bridge)
|
||||
|
||||
for p in (rebind, worker, bridge):
|
||||
syntax_check(p)
|
||||
|
||||
print("RESULT: PATCH_OK")
|
||||
print("CHANGED:", ",".join(dict.fromkeys(str(p) for p in changed)) if changed else "none")
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user