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

"""Fence WP09 approvals, callable functions, and reservation settlement."""

from alembic import op

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


def upgrade() -> None:
    op.execute(r"""
    ALTER TABLE public.agent_runtime_reservations
      ADD COLUMN IF NOT EXISTS agent_uid uuid REFERENCES public.governed_agents(uid),
      ADD COLUMN IF NOT EXISTS principal_id uuid REFERENCES public.users(id),
      ADD COLUMN IF NOT EXISTS worker_id text,
      ADD COLUMN IF NOT EXISTS lease_expires_at timestamptz,
      ADD COLUMN IF NOT EXISTS settlement_digest char(64);
    UPDATE public.agent_runtime_reservations
       SET lease_expires_at=COALESCE(lease_expires_at,created_at+interval '5 minutes')
     WHERE lease_expires_at IS NULL;
    ALTER TABLE public.agent_runtime_reservations
      ALTER COLUMN lease_expires_at SET NOT NULL;

    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;
    REVOKE ALL ON FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint) FROM PUBLIC;
    REVOKE ALL ON FUNCTION public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) FROM PUBLIC;

    -- Retain the 497 helper overload for a true downgrade; the stricter 498
    -- overload below is the only one called by the replacement authorizer.
    CREATE FUNCTION public.agent_runtime_approval_valid(
      p_task uuid,p_agent uuid,p_principal 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 $approval$
      SELECT EXISTS(
        SELECT 1 FROM public.governance_tasks t
        WHERE t.uid=p_task AND t.status='approved' AND t.subject_uid=p_agent::text
          AND t.context->>'agent_uid'=p_agent::text AND t.context->>'principal_id'=p_principal::text
          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_principal)
      )
    $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 $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->>'agent_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'grant_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'principal_id') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'business_domain_uid') !~ '^[0-9a-f-]{36}$' OR (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') OR (p_payload->>'approval_task_uid') !~ '^[0-9a-f-]{36}$' THEN RAISE EXCEPTION 'agent runtime approval is missing'; END IF;
        v_approval:=(p_payload->>'approval_task_uid')::uuid;
        IF NOT EXISTS(SELECT 1 FROM public.governance_tasks t WHERE t.uid=v_approval AND t.subject_uid=(p_payload->>'agent_uid') AND t.context->>'agent_uid'=p_payload->>'agent_uid') THEN RAISE EXCEPTION 'approval agent scope denied'; END IF;
        IF NOT EXISTS(SELECT 1 FROM public.governance_tasks t WHERE t.uid=v_approval AND t.context->>'principal_id'=p_payload->>'principal_id') THEN RAISE EXCEPTION 'approval principal scope denied'; END IF;
      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->>'agent_uid')::uuid,(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,agent_uid,principal_id,lease_expires_at) 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',(p_payload->>'agent_uid')::uuid,(p_payload->>'principal_id')::uuid,v_now+interval '5 minutes');
      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; $authorize$;

    CREATE TABLE IF NOT EXISTS public.agent_runtime_settlement_audits (
      reservation_id uuid PRIMARY KEY REFERENCES public.agent_runtime_reservations(reservation_id),
      tenant_id text NOT NULL,
      idempotency_key text NOT NULL,
      outcome text NOT NULL,
      actual_tokens bigint NOT NULL,
      actual_cost_micros bigint NOT NULL,
      actual_tools integer NOT NULL,
      actual_time_ms bigint NOT NULL,
      lease_fence bigint NOT NULL,
      settled_at timestamptz NOT NULL DEFAULT clock_timestamp(),
      UNIQUE (tenant_id, idempotency_key)
    );
    ALTER TABLE public.agent_runtime_settlement_audits OWNER TO dataops_agent_runtime_owner;

    CREATE FUNCTION public.agent_runtime_settle(p_payload jsonb)
    RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $settle$
    DECLARE r record; b record; digest text; refund_tokens bigint; refund_cost bigint; refund_tools integer; refund_time bigint;
    BEGIN
      IF jsonb_typeof(p_payload)<>'object' OR NOT (p_payload ?& ARRAY['agent_uid','principal_id','worker_id','idempotency_key','lease_fence','outcome','actual_tokens','actual_cost_micros','actual_tools','actual_time_ms']) OR p_payload-ARRAY['agent_uid','principal_id','worker_id','idempotency_key','lease_fence','outcome','actual_tokens','actual_cost_micros','actual_tools','actual_time_ms']<>'{}'::jsonb THEN RAISE EXCEPTION 'agent runtime settlement is closed'; END IF;
      IF (p_payload->>'agent_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'principal_id') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'worker_id') !~ '^agent-runtime-http:[0-9a-f-]{36}$' OR (p_payload->>'lease_fence') !~ '^[1-9][0-9]*$' OR p_payload->>'outcome' NOT IN ('success','failed','cancelled') OR (p_payload->>'actual_tokens') !~ '^[0-9]{1,6}$' OR (p_payload->>'actual_cost_micros') !~ '^[0-9]{1,9}$' OR (p_payload->>'actual_tools') !~ '^[0-9]$' OR (p_payload->>'actual_time_ms') !~ '^[0-9]{1,6}$' THEN RAISE EXCEPTION 'agent runtime settlement value invalid'; END IF;
      SELECT * INTO r FROM public.agent_runtime_reservations WHERE tenant_id IN (SELECT tenant_id FROM public.agent_invocation_audits WHERE idempotency_key=p_payload->>'idempotency_key') AND idempotency_key=p_payload->>'idempotency_key' AND agent_uid=(p_payload->>'agent_uid')::uuid AND principal_id=(p_payload->>'principal_id')::uuid FOR UPDATE;
      IF NOT FOUND THEN RAISE EXCEPTION 'agent runtime settlement scope denied'; END IF;
      digest:=encode(sha256(convert_to(p_payload::text,'utf8')),'hex');
      IF r.status IN ('settled','failed','cancelled') THEN IF r.settlement_digest<>digest THEN RAISE EXCEPTION 'agent runtime settlement conflict'; END IF; RETURN jsonb_build_object('replay',true,'rejected',r.settlement='lease_expired','status',r.status,'lease_fence',r.lease_fence); END IF;
      IF r.lease_fence<>(p_payload->>'lease_fence')::bigint THEN RAISE EXCEPTION 'settlement fence conflict'; END IF;
      IF r.worker_id IS NULL THEN UPDATE public.agent_runtime_reservations SET worker_id=p_payload->>'worker_id' WHERE reservation_id=r.reservation_id; ELSIF r.worker_id<>p_payload->>'worker_id' THEN RAISE EXCEPTION 'settlement worker conflict'; END IF;
      IF r.lease_expires_at<=clock_timestamp() THEN
        UPDATE public.agent_runtime_budgets SET token_remaining=token_remaining+r.token_count,cost_remaining_micros=cost_remaining_micros+r.cost_micros,tool_remaining=tool_remaining+r.tool_count,time_remaining_ms=time_remaining_ms+r.time_ms,concurrency_remaining=concurrency_remaining+1,lease_fence=lease_fence+1 WHERE tenant_id=r.tenant_id;
        UPDATE public.agent_runtime_model_budgets SET remaining=remaining+1 WHERE tenant_id=r.tenant_id AND model=r.model;
        UPDATE public.agent_runtime_reservations SET status='failed',settlement='lease_expired',settled_at=clock_timestamp(),settlement_digest=digest WHERE reservation_id=r.reservation_id;
        INSERT INTO public.agent_runtime_settlement_audits(reservation_id,tenant_id,idempotency_key,outcome,actual_tokens,actual_cost_micros,actual_tools,actual_time_ms,lease_fence) VALUES(r.reservation_id,r.tenant_id,r.idempotency_key,'lease_expired',0,0,0,0,r.lease_fence);
        RETURN jsonb_build_object('replay',false,'rejected',true,'status','failed','lease_fence',r.lease_fence);
      END IF;
      IF (p_payload->>'actual_tokens')::bigint>r.token_count OR (p_payload->>'actual_cost_micros')::bigint>r.cost_micros OR (p_payload->>'actual_tools')::integer>r.tool_count OR (p_payload->>'actual_time_ms')::bigint>r.time_ms THEN RAISE EXCEPTION 'agent runtime settlement exceeds reservation'; END IF;
      IF p_payload->>'outcome'='success' THEN refund_tokens:=r.token_count-(p_payload->>'actual_tokens')::bigint; refund_cost:=r.cost_micros-(p_payload->>'actual_cost_micros')::bigint; refund_tools:=r.tool_count-(p_payload->>'actual_tools')::integer; refund_time:=r.time_ms-(p_payload->>'actual_time_ms')::bigint; ELSE refund_tokens:=r.token_count; refund_cost:=r.cost_micros; refund_tools:=r.tool_count; refund_time:=r.time_ms; END IF;
      UPDATE public.agent_runtime_budgets SET token_remaining=token_remaining+refund_tokens,cost_remaining_micros=cost_remaining_micros+refund_cost,tool_remaining=tool_remaining+refund_tools,time_remaining_ms=time_remaining_ms+refund_time,concurrency_remaining=concurrency_remaining+1,lease_fence=lease_fence+1 WHERE tenant_id=r.tenant_id;
      UPDATE public.agent_runtime_model_budgets SET remaining=remaining+1 WHERE tenant_id=r.tenant_id AND model=r.model;
      UPDATE public.agent_runtime_reservations SET status=CASE WHEN p_payload->>'outcome'='success' THEN 'settled' ELSE p_payload->>'outcome' END,settlement=p_payload->>'outcome',settled_at=clock_timestamp(),settlement_digest=digest WHERE reservation_id=r.reservation_id;
      INSERT INTO public.agent_runtime_settlement_audits(reservation_id,tenant_id,idempotency_key,outcome,actual_tokens,actual_cost_micros,actual_tools,actual_time_ms,lease_fence) VALUES(r.reservation_id,r.tenant_id,r.idempotency_key,p_payload->>'outcome',(p_payload->>'actual_tokens')::bigint,(p_payload->>'actual_cost_micros')::bigint,(p_payload->>'actual_tools')::integer,(p_payload->>'actual_time_ms')::bigint,r.lease_fence);
      RETURN jsonb_build_object('replay',false,'rejected',false,'status',CASE WHEN p_payload->>'outcome'='success' THEN 'settled' ELSE p_payload->>'outcome' END,'lease_fence',r.lease_fence);
    END; $settle$;
    REVOKE ALL ON FUNCTION public.agent_runtime_approval_valid(uuid,uuid,uuid,uuid,text,text,text,text) FROM PUBLIC;
    REVOKE ALL ON FUNCTION public.agent_runtime_settle(jsonb) FROM PUBLIC;
    REVOKE ALL ON TABLE public.agent_runtime_settlement_audits FROM PUBLIC,dataops_agent_runtime;
    GRANT SELECT,INSERT,UPDATE ON TABLE public.agent_runtime_settlement_audits,
      public.agent_runtime_states,public.agent_generation_canaries TO dataops_agent_runtime_owner;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_authorize(jsonb) TO dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_settle(jsonb) TO dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint) TO dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) TO dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_authorize(jsonb),
      public.agent_runtime_settle(jsonb),
      public.agent_runtime_transition(text,text,uuid,uuid,bigint),
      public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text)
      TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_authorize(jsonb) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_approval_valid(uuid,uuid,uuid,uuid,text,text,text,text) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_settle(jsonb) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) OWNER TO dataops_agent_runtime_owner;
    """)


def downgrade() -> None:
    op.execute("""
    REVOKE ALL ON FUNCTION public.agent_runtime_settle(jsonb) FROM PUBLIC;
    DROP FUNCTION IF EXISTS public.agent_runtime_settle(jsonb);
    DROP FUNCTION IF EXISTS public.agent_runtime_approval_valid(uuid,uuid,uuid,uuid,text,text,text,text);
    DROP TABLE IF EXISTS public.agent_runtime_settlement_audits;
    ALTER TABLE public.agent_runtime_reservations DROP COLUMN IF EXISTS settlement_digest, DROP COLUMN IF EXISTS lease_expires_at, DROP COLUMN IF EXISTS worker_id, DROP COLUMN IF EXISTS principal_id, DROP COLUMN IF EXISTS agent_uid;
    """)