DataOps-platform/migrations/versions/20260813_496_model_gateway_agent_runtime.py

"""Add durable default-deny model gateway runtime governance.

Revision ID: 20260813_496
Revises: 20260811_495
"""

from alembic import op

revision = "20260813_496"
down_revision = "20260811_495"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.model_gateway_routes (
            route_id text PRIMARY KEY, tenant_id text NOT NULL,
            principal_id uuid NOT NULL, business_domain_uid uuid NOT NULL,
            environment text NOT NULL, provider text NOT NULL, model text NOT NULL,
            prompt_version text NOT NULL, generation text NOT NULL,
            canary_status text NOT NULL, created_at timestamptz NOT NULL DEFAULT now(),
            UNIQUE (tenant_id, principal_id, business_domain_uid, environment,
                    provider, model, prompt_version, generation)
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_runtime_budgets (
            tenant_id text PRIMARY KEY, token_remaining bigint NOT NULL CHECK (token_remaining >= 0),
            cost_remaining_micros bigint NOT NULL CHECK (cost_remaining_micros >= 0),
            tool_remaining integer NOT NULL CHECK (tool_remaining >= 0),
            time_remaining_ms bigint NOT NULL CHECK (time_remaining_ms >= 0),
            concurrency_remaining integer NOT NULL CHECK (concurrency_remaining >= 0),
            lease_fence bigint NOT NULL DEFAULT 0, updated_at timestamptz NOT NULL DEFAULT now()
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_runtime_model_budgets (
            tenant_id text NOT NULL REFERENCES public.agent_runtime_budgets(tenant_id) ON DELETE CASCADE,
            model text NOT NULL, remaining integer NOT NULL CHECK (remaining >= 0),
            PRIMARY KEY (tenant_id, model)
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_runtime_reservations (
            reservation_id uuid PRIMARY KEY, tenant_id text NOT NULL,
            idempotency_key text NOT NULL, model text NOT NULL, token_count bigint NOT NULL,
            cost_micros bigint NOT NULL, tool_count integer NOT NULL, time_ms bigint NOT NULL,
            lease_fence bigint NOT NULL, status text NOT NULL, created_at timestamptz NOT NULL DEFAULT now(),
            UNIQUE (tenant_id, idempotency_key)
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_runtime_capabilities (
            capability_id uuid PRIMARY KEY, tenant_id text NOT NULL,
            business_domain_uid uuid NOT NULL, environment text NOT NULL,
            interface_type text NOT NULL, tool_name text NOT NULL, action text NOT NULL,
            capability_set jsonb NOT NULL, risk_level text NOT NULL,
            enabled boolean NOT NULL DEFAULT false, created_at timestamptz NOT NULL DEFAULT now(),
            UNIQUE (tenant_id, business_domain_uid, environment, interface_type, tool_name, action)
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_runtime_states (
            tenant_id text PRIMARY KEY, state text NOT NULL, incident_ref text,
            approval_ref text, lease_fence bigint NOT NULL DEFAULT 0,
            updated_at timestamptz NOT NULL DEFAULT now()
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_generation_canaries (
            tenant_id text NOT NULL, generation text NOT NULL, parent_generation text,
            status text NOT NULL, lease_fence bigint NOT NULL DEFAULT 0,
            approval_ref text, created_at timestamptz NOT NULL DEFAULT now(),
            PRIMARY KEY (tenant_id, generation)
        )
    """)
    op.execute("""
        CREATE TABLE IF NOT EXISTS public.agent_invocation_audits (
            invocation_id uuid PRIMARY KEY, tenant_id text NOT NULL,
            idempotency_key text NOT NULL, route_id text NOT NULL, generation text NOT NULL,
            lease_fence bigint NOT NULL, input_hash text NOT NULL,
            evidence_digests jsonb NOT NULL, decision text NOT NULL,
            created_at timestamptz NOT NULL DEFAULT now(),
            UNIQUE (tenant_id, idempotency_key)
        )
    """)


def downgrade() -> None:
    op.execute("DROP TABLE IF EXISTS public.agent_invocation_audits")
    op.execute("DROP TABLE IF EXISTS public.agent_generation_canaries")
    op.execute("DROP TABLE IF EXISTS public.agent_runtime_states")
    op.execute("DROP TABLE IF EXISTS public.agent_runtime_capabilities")
    op.execute("DROP TABLE IF EXISTS public.agent_runtime_reservations")
    op.execute("DROP TABLE IF EXISTS public.agent_runtime_model_budgets")
    op.execute("DROP TABLE IF EXISTS public.agent_runtime_budgets")
    op.execute("DROP TABLE IF EXISTS public.model_gateway_routes")