Files
aios/tests/test_trace_tool_credential_governance.py
T

488 lines
18 KiB
Python

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 == []