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

"""Harden WP09 runtime authorization into one fenced database transaction."""

from alembic import op

revision = "20260813_497"
down_revision = "20260813_496"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute(r"""
    DO $$ BEGIN
      IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname='dataops_agent_runtime_owner')
         OR NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname='dataops_agent_runtime') THEN
        RAISE EXCEPTION 'agent runtime role-init is required before migration 20260813_497';
      END IF;
      IF NOT pg_has_role(current_user, 'dataops_agent_runtime_owner', 'MEMBER') THEN
        RAISE EXCEPTION 'migrator must be pre-granted agent runtime owner role';
      END IF;
    END $$;
    ALTER TABLE public.agent_invocation_audits
      ADD COLUMN IF NOT EXISTS request_digest char(64),
      ADD COLUMN IF NOT EXISTS approval_task_uid uuid REFERENCES public.governance_tasks(uid),
      ADD COLUMN IF NOT EXISTS incident_uid uuid REFERENCES public.data_incidents(uid);
    UPDATE public.agent_invocation_audits SET request_digest=input_hash WHERE request_digest IS NULL;
    ALTER TABLE public.agent_invocation_audits ALTER COLUMN request_digest SET NOT NULL;
    CREATE UNIQUE INDEX IF NOT EXISTS uq_agent_runtime_request_digest
      ON public.agent_invocation_audits(tenant_id,idempotency_key,request_digest);
    ALTER TABLE public.agent_runtime_reservations
      ADD COLUMN IF NOT EXISTS settled_at timestamptz,
      ADD COLUMN IF NOT EXISTS settlement text;

    CREATE OR REPLACE FUNCTION public.agent_runtime_approval_valid(
      p_task uuid,p_actor uuid,p_domain uuid,p_environment text,p_tool text,p_action text,p_digest text
    ) RETURNS boolean LANGUAGE sql SECURITY DEFINER SET search_path=pg_catalog,public AS $agent_runtime_approval$
      SELECT EXISTS(
        SELECT 1 FROM public.governance_tasks t
        WHERE t.uid=p_task AND t.status='approved'
          AND t.business_domain_uid=p_domain::text AND t.context->>'environment'=p_environment
          AND t.context->>'tool_name'=p_tool AND t.context->>'action'=p_action
          AND t.context->>'request_digest'=p_digest AND t.due_at>clock_timestamp()
          AND t.context->>'automatic_execution_allowed'='false'
          AND EXISTS(SELECT 1 FROM public.governance_task_reviews r
            WHERE r.task_uid=t.uid AND r.decision='approve' AND r.reviewer_uid<>p_actor)
      )
    $agent_runtime_approval$;

    CREATE OR REPLACE FUNCTION public.agent_runtime_authorize(p_payload jsonb)
    RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER
    SET search_path=pg_catalog,public AS $agent_runtime_authorize$
    DECLARE v_budget record; v_model record; v_grant record; v_route record;
      v_audit record; v_fence bigint; v_now timestamptz; v_approval uuid;
    BEGIN
      IF jsonb_typeof(p_payload) <> 'object'
         OR NOT (p_payload ?& ARRAY['agent_uid','grant_uid','tenant_id','principal_id',
           'business_domain_uid','environment','provider','model','prompt_version',
           'generation','interface_type','tool_name','action','risk_level','idempotency_key',
           'request_digest','input_hash','evidence_digests','estimated_tokens',
           'estimated_cost_micros','requested_time_ms'])
         OR p_payload - ARRAY['agent_uid','grant_uid','tenant_id','principal_id',
           'business_domain_uid','environment','provider','model','prompt_version',
           'generation','interface_type','tool_name','action','risk_level','idempotency_key',
           'request_digest','input_hash','evidence_digests','estimated_tokens',
           'estimated_cost_micros','requested_time_ms','approval_task_uid'] <> '{}'::jsonb THEN
        RAISE EXCEPTION 'agent runtime request is closed';
      END IF;
      IF (p_payload->>'request_digest') !~ '^[0-9a-f]{64}$'
         OR (p_payload->>'input_hash') !~ '^[0-9a-f]{64}$'
         OR jsonb_typeof(p_payload->'evidence_digests') <> 'array'
         OR (p_payload->>'estimated_tokens') !~ '^[1-9][0-9]{0,5}$'
         OR (p_payload->>'estimated_cost_micros') !~ '^[1-9][0-9]{0,8}$'
         OR (p_payload->>'requested_time_ms') !~ '^[1-9][0-9]{0,5}$' THEN
        RAISE EXCEPTION 'agent runtime budget value is invalid';
      END IF;
      SELECT g.*,a.owner_uid INTO v_grant FROM public.agent_tool_grants g
      JOIN public.governed_agents a ON a.uid=g.agent_uid
      WHERE g.uid=(p_payload->>'grant_uid')::uuid
        AND g.agent_uid=(p_payload->>'agent_uid')::uuid
        AND g.status='active' AND a.status='active'
        AND a.owner_uid=(p_payload->>'principal_id')::uuid
        AND g.business_domain_uid=(p_payload->>'business_domain_uid')::uuid
        AND g.environment=p_payload->>'environment'
        AND g.interface_type=p_payload->>'interface_type'
        AND g.tool_name=p_payload->>'tool_name' AND g.action=p_payload->>'action'
        AND g.risk_level=p_payload->>'risk_level' FOR UPDATE;
      IF NOT FOUND THEN
        RAISE EXCEPTION 'agent runtime grant scope denied';
      END IF;
      IF v_grant.action='execute' OR v_grant.risk_level <> 'low' OR v_grant.requires_approval THEN
        IF NOT (p_payload ? 'approval_task_uid') THEN RAISE EXCEPTION 'agent runtime approval is missing'; END IF;
        v_approval := (p_payload->>'approval_task_uid')::uuid;
      END IF;
      SELECT * INTO v_route FROM public.model_gateway_routes
      WHERE tenant_id=p_payload->>'tenant_id'
        AND principal_id=(p_payload->>'principal_id')::uuid
        AND business_domain_uid=(p_payload->>'business_domain_uid')::uuid
        AND environment=p_payload->>'environment' AND provider=p_payload->>'provider'
        AND model=p_payload->>'model' AND prompt_version=p_payload->>'prompt_version'
        AND generation=p_payload->>'generation' AND canary_status='approved' FOR UPDATE;
      IF NOT FOUND THEN RAISE EXCEPTION 'agent runtime route scope denied'; END IF;
      SELECT * INTO v_budget FROM public.agent_runtime_budgets
        WHERE tenant_id=p_payload->>'tenant_id' FOR UPDATE;
      IF NOT FOUND THEN RAISE EXCEPTION 'agent runtime budget missing'; END IF;
      SELECT * INTO v_audit FROM public.agent_invocation_audits
        WHERE tenant_id=p_payload->>'tenant_id' AND idempotency_key=p_payload->>'idempotency_key' FOR UPDATE;
      IF FOUND THEN
        IF v_audit.request_digest <> p_payload->>'request_digest' THEN
          RAISE EXCEPTION 'agent runtime idempotency conflict';
        END IF;
        RETURN jsonb_build_object('decision',v_audit.decision,'invocation_id',v_audit.invocation_id::text,
          'lease_fence',v_audit.lease_fence,'replay',true);
      END IF;
      IF v_approval IS NOT NULL AND NOT public.agent_runtime_approval_valid(
        v_approval,(p_payload->>'principal_id')::uuid,(p_payload->>'business_domain_uid')::uuid,
        p_payload->>'environment',p_payload->>'tool_name',p_payload->>'action',p_payload->>'input_hash') THEN
        RAISE EXCEPTION 'agent runtime approval is expired, cross-scope, or consumed';
      END IF;
      SELECT * INTO v_model FROM public.agent_runtime_model_budgets
        WHERE tenant_id=v_budget.tenant_id AND model=p_payload->>'model' FOR UPDATE;
      IF NOT FOUND OR v_budget.token_remaining < (p_payload->>'estimated_tokens')::bigint
        OR v_budget.cost_remaining_micros < (p_payload->>'estimated_cost_micros')::bigint
        OR v_budget.tool_remaining < 1 OR v_budget.time_remaining_ms < (p_payload->>'requested_time_ms')::bigint
        OR v_budget.concurrency_remaining < 1 OR v_model.remaining < 1 THEN
        RAISE EXCEPTION 'agent runtime budget exhausted';
      END IF;
      IF v_approval IS NOT NULL THEN
        UPDATE public.governance_tasks SET status='closed',closed_at=clock_timestamp(),
          updated_at=clock_timestamp()
          WHERE uid=v_approval AND status='approved' RETURNING uid INTO v_approval;
        IF NOT FOUND THEN RAISE EXCEPTION 'agent runtime approval was already consumed'; END IF;
      END IF;
      v_now:=clock_timestamp();
      UPDATE public.agent_runtime_budgets SET token_remaining=token_remaining-(p_payload->>'estimated_tokens')::bigint,
        cost_remaining_micros=cost_remaining_micros-(p_payload->>'estimated_cost_micros')::bigint,
        tool_remaining=tool_remaining-1,time_remaining_ms=time_remaining_ms-(p_payload->>'requested_time_ms')::bigint,
        concurrency_remaining=concurrency_remaining-1,lease_fence=lease_fence+1,updated_at=v_now
        WHERE tenant_id=v_budget.tenant_id RETURNING lease_fence INTO v_fence;
      UPDATE public.agent_runtime_model_budgets SET remaining=remaining-1
        WHERE tenant_id=v_budget.tenant_id AND model=v_model.model;
      INSERT INTO public.agent_runtime_reservations(reservation_id,tenant_id,idempotency_key,model,token_count,cost_micros,tool_count,time_ms,lease_fence,status)
        VALUES(gen_random_uuid(),v_budget.tenant_id,p_payload->>'idempotency_key',v_model.model,
          (p_payload->>'estimated_tokens')::bigint,(p_payload->>'estimated_cost_micros')::bigint,1,
          (p_payload->>'requested_time_ms')::bigint,v_fence,'reserved');
      INSERT INTO public.agent_invocation_audits(invocation_id,tenant_id,idempotency_key,route_id,generation,lease_fence,input_hash,evidence_digests,decision,request_digest,approval_task_uid)
        VALUES(gen_random_uuid(),v_budget.tenant_id,p_payload->>'idempotency_key',v_route.route_id,
          v_route.generation,v_fence,p_payload->>'input_hash',p_payload->'evidence_digests','authorized',p_payload->>'request_digest',v_approval)
        RETURNING * INTO v_audit;
      RETURN jsonb_build_object('decision','authorized','invocation_id',v_audit.invocation_id::text,
        'lease_fence',v_fence,'replay',false);
    END; $agent_runtime_authorize$;
    REVOKE ALL ON TABLE public.agent_invocation_audits,public.agent_runtime_budgets,
      public.agent_runtime_model_budgets,public.agent_runtime_reservations FROM PUBLIC,dataops_agent_runtime;
    REVOKE ALL ON FUNCTION public.agent_runtime_authorize(jsonb) FROM PUBLIC;
    REVOKE ALL ON FUNCTION public.agent_runtime_approval_valid(uuid,uuid,uuid,text,text,text,text) FROM PUBLIC;
    GRANT SELECT,INSERT,UPDATE ON TABLE public.agent_invocation_audits,public.agent_runtime_budgets,
      public.agent_runtime_model_budgets,public.agent_runtime_reservations,
      public.agent_tool_grants,public.governed_agents,public.model_gateway_routes,
      public.governance_tasks,public.governance_task_reviews TO dataops_agent_runtime_owner;
    -- Role-init, not this low-privilege migration, grants dataops_agent_runtime to dataops.
    -- Schema usage is pre-provisioned by role-init for the runtime role.
    GRANT EXECUTE ON FUNCTION public.agent_runtime_authorize(jsonb) TO dataops_agent_runtime,dataops;
    ALTER FUNCTION public.agent_runtime_authorize(jsonb) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_approval_valid(uuid,uuid,uuid,text,text,text,text) OWNER TO dataops_agent_runtime_owner;
    CREATE OR REPLACE FUNCTION public.agent_runtime_transition(
      p_tenant text,p_state text,p_incident uuid,p_approval uuid,p_expected_fence bigint
    ) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $agent_runtime_transition$
    DECLARE v_fence bigint;
    BEGIN
      IF p_state NOT IN ('degraded','paused','active') THEN RAISE EXCEPTION 'runtime state invalid'; END IF;
      IF p_state='active' AND (p_approval IS NULL OR NOT EXISTS(SELECT 1 FROM governance_tasks WHERE uid=p_approval AND status='approved' AND due_at>clock_timestamp())) THEN RAISE EXCEPTION 'runtime recovery approval required'; END IF;
      IF p_state IN ('degraded','paused') AND p_incident IS NULL THEN RAISE EXCEPTION 'runtime incident required'; END IF;
      UPDATE agent_runtime_states SET state=p_state,incident_ref=COALESCE(p_incident::text,incident_ref),approval_ref=p_approval::text,lease_fence=lease_fence+1,updated_at=clock_timestamp()
       WHERE tenant_id=p_tenant AND lease_fence=p_expected_fence RETURNING lease_fence INTO v_fence;
      IF NOT FOUND THEN RAISE EXCEPTION 'runtime state fence conflict'; END IF;
      RETURN v_fence;
    END; $agent_runtime_transition$;
    CREATE OR REPLACE FUNCTION public.agent_runtime_canary_transition(
      p_tenant text,p_generation text,p_status text,p_approval uuid,p_expected_fence bigint,p_metrics jsonb,p_dataset_version text
    ) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $agent_runtime_canary_transition$
    DECLARE v_fence bigint;
    BEGIN
      IF p_status NOT IN ('approved','rolled_back') OR p_dataset_version !~ '^[A-Za-z0-9._-]{1,120}$' OR jsonb_typeof(p_metrics)<>'object' THEN RAISE EXCEPTION 'canary evidence invalid'; END IF;
      IF p_approval IS NULL OR NOT EXISTS(SELECT 1 FROM governance_tasks WHERE uid=p_approval AND status='approved' AND due_at>clock_timestamp()) THEN RAISE EXCEPTION 'canary approval required'; END IF;
      UPDATE agent_generation_canaries SET status=p_status,approval_ref=p_approval::text,lease_fence=lease_fence+1
       WHERE tenant_id=p_tenant AND generation=p_generation AND lease_fence=p_expected_fence RETURNING lease_fence INTO v_fence;
      IF NOT FOUND THEN RAISE EXCEPTION 'generation fence conflict'; END IF;
      RETURN v_fence;
    END; $agent_runtime_canary_transition$;
    """)


def downgrade() -> None:
    op.execute("""
    REVOKE dataops_agent_runtime FROM dataops;
    REVOKE EXECUTE ON FUNCTION public.agent_runtime_authorize(jsonb) FROM dataops_agent_runtime,dataops;
    DROP FUNCTION public.agent_runtime_authorize(jsonb);
    DROP FUNCTION public.agent_runtime_approval_valid(uuid,uuid,uuid,text,text,text,text);
    DROP FUNCTION IF EXISTS public.agent_runtime_transition(text,text,uuid,uuid,bigint);
    DROP FUNCTION IF EXISTS public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text);
    DROP INDEX IF EXISTS public.uq_agent_runtime_request_digest;
    ALTER TABLE public.agent_runtime_reservations DROP COLUMN IF EXISTS settlement, DROP COLUMN IF EXISTS settled_at;
    ALTER TABLE public.agent_invocation_audits DROP COLUMN IF EXISTS incident_uid, DROP COLUMN IF EXISTS approval_task_uid, DROP COLUMN IF EXISTS request_digest;
    """)