multica/server/migrations/073_task_usage_daily_rollup.up.sql

-- Daily rollup table for `task_usage`. Background: the dashboard query
-- ListRuntimeUsage runs `SUM() GROUP BY DATE(created_at), provider, model`
-- against the raw event stream and is called once per runtime row on the
-- runtimes list (plus once per detail page load), so it dominates DB load
-- as event volume grows. We materialise the day-bucketed aggregate here
-- so reads scan O(days × providers × models) rows instead of O(events).
--
-- All query dimensions are denormalised into the table so reads never
-- need to join `agent_task_queue`. The PK doubles as the upsert key for
-- the rollup worker.
CREATE TABLE task_usage_daily (
    bucket_date         DATE        NOT NULL,
    workspace_id        UUID        NOT NULL,
    runtime_id          UUID        NOT NULL,
    provider            TEXT        NOT NULL,
    model               TEXT        NOT NULL,
    input_tokens        BIGINT      NOT NULL DEFAULT 0,
    output_tokens       BIGINT      NOT NULL DEFAULT 0,
    cache_read_tokens   BIGINT      NOT NULL DEFAULT 0,
    cache_write_tokens  BIGINT      NOT NULL DEFAULT 0,
    event_count         BIGINT      NOT NULL DEFAULT 0,
    updated_at          TIMESTAMPTZ NOT NULL DEFAULT now(),
    PRIMARY KEY (bucket_date, workspace_id, runtime_id, provider, model)
);

-- Primary read path: runtime detail page + runtimes-list cost cell, both
-- filter by runtime_id and order by date DESC. bucket_date DESC in the
-- index lets the query avoid an extra sort.
CREATE INDEX idx_task_usage_daily_runtime_date
    ON task_usage_daily (runtime_id, bucket_date DESC);

-- Workspace-wide aggregations hit this index instead of fanning out per
-- runtime.
CREATE INDEX idx_task_usage_daily_workspace_date
    ON task_usage_daily (workspace_id, bucket_date DESC);

-- Single-row state table tracking how far the rollup worker has consumed.
CREATE TABLE task_usage_rollup_state (
    id                    SMALLINT    PRIMARY KEY DEFAULT 1 CHECK (id = 1),
    watermark_at          TIMESTAMPTZ NOT NULL DEFAULT '1970-01-01 00:00:00+00',
    last_run_started_at   TIMESTAMPTZ,
    last_run_finished_at  TIMESTAMPTZ,
    last_run_rows         BIGINT      NOT NULL DEFAULT 0,
    last_error            TEXT
);
INSERT INTO task_usage_rollup_state (id) VALUES (1) ON CONFLICT DO NOTHING;

-- Window-based aggregation primitive. Used by both the cron-driven
-- watermark advancer and the offline backfill command, so they stay
-- byte-identical in their semantics. Returns the number of output rows
-- touched.
--
-- IDEMPOTENCY CONTRACT (this is the important bit):
--   For every (bucket_date, workspace_id, runtime_id, provider, model)
--   key that has at least one task_usage row whose `updated_at` falls in
--   [p_from, p_to), this function REPLACES the corresponding daily row
--   with the SUM of *all* task_usage rows for that key (regardless of
--   their updated_at). It does NOT add a delta.
--
-- Consequences:
--   * Replaying the same window is safe — the row is rebuilt from raw
--     each time, so the result converges.
--   * Two callers (cron + backfill) processing overlapping windows is
--     safe — both write the same value.
--   * `UpsertTaskUsage` corrections that overwrite token counts are
--     captured: the corrected row's updated_at gets bumped, the next
--     window picks up its bucket key, and the bucket is recomputed
--     from current truth.
--
-- Cost: the recompute reads ALL task_usage rows for each dirty bucket,
-- not just the windowed slice. In steady state only "today" buckets are
-- dirty (a handful of keys per active runtime), so this stays cheap.
-- During backfill the entire history's bucket keys become dirty once;
-- the backfill walks history in monthly slices to bound the working
-- set per call.
CREATE OR REPLACE FUNCTION rollup_task_usage_daily_window(
    p_from TIMESTAMPTZ,
    p_to   TIMESTAMPTZ
)
RETURNS BIGINT
LANGUAGE plpgsql
AS $$
DECLARE
    v_rows BIGINT;
BEGIN
    IF p_from >= p_to THEN
        RETURN 0;
    END IF;

    WITH dirty_keys AS (
        SELECT DISTINCT
            DATE(tu.created_at) AS bucket_date,
            i.workspace_id      AS workspace_id,
            atq.runtime_id      AS runtime_id,
            tu.provider         AS provider,
            tu.model            AS model
        FROM task_usage tu
        JOIN agent_task_queue atq ON atq.id      = tu.task_id
        JOIN issue            i   ON i.id        = atq.issue_id
        WHERE atq.runtime_id IS NOT NULL
          AND (
              -- Steady state: rows updated within the watermark window.
              -- Hits idx_task_usage_updated_at directly.
              (tu.updated_at >= p_from AND tu.updated_at < p_to)
              -- Legacy rows from before migration 072 (updated_at IS NULL)
              -- — discoverable via created_at + the partial index added
              -- in 077. Steady-state windows after backfill never include
              -- historical dates, so this branch is a no-op once the
              -- backfill has swept history.
              OR (tu.updated_at IS NULL
                  AND tu.created_at >= p_from
                  AND tu.created_at <  p_to)
          )
    ),
    recomputed AS (
        SELECT
            dk.bucket_date,
            dk.workspace_id,
            dk.runtime_id,
            dk.provider,
            dk.model,
            SUM(tu.input_tokens)::bigint        AS input_tokens,
            SUM(tu.output_tokens)::bigint       AS output_tokens,
            SUM(tu.cache_read_tokens)::bigint   AS cache_read_tokens,
            SUM(tu.cache_write_tokens)::bigint  AS cache_write_tokens,
            COUNT(*)::bigint                    AS event_count
        FROM dirty_keys dk
        JOIN agent_task_queue atq ON atq.runtime_id = dk.runtime_id
        JOIN issue            i   ON i.id           = atq.issue_id
                                  AND i.workspace_id = dk.workspace_id
        JOIN task_usage       tu  ON tu.task_id     = atq.id
                                  AND tu.provider   = dk.provider
                                  AND tu.model      = dk.model
                                  AND DATE(tu.created_at) = dk.bucket_date
        GROUP BY 1, 2, 3, 4, 5
    )
    INSERT INTO task_usage_daily AS d (
        bucket_date, workspace_id, runtime_id, provider, model,
        input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
        event_count
    )
    SELECT
        bucket_date, workspace_id, runtime_id, provider, model,
        input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
        event_count
    FROM recomputed
    ON CONFLICT (bucket_date, workspace_id, runtime_id, provider, model) DO UPDATE
        SET input_tokens       = EXCLUDED.input_tokens,
            output_tokens      = EXCLUDED.output_tokens,
            cache_read_tokens  = EXCLUDED.cache_read_tokens,
            cache_write_tokens = EXCLUDED.cache_write_tokens,
            event_count        = EXCLUDED.event_count,
            updated_at         = now();

    GET DIAGNOSTICS v_rows = ROW_COUNT;
    RETURN v_rows;
END;
$$;

-- Cron entry point. Advances the watermark by one window each call.
--
-- Invariants:
--  * `pg_try_advisory_lock(4242)` serialises overlapping ticks.
--  * The window upper bound is `now() - 5 minutes`. The lag exists
--    because `task_usage` rows are written from a separate transaction;
--    a row with updated_at = T can become visible to this snapshot at
--    some t > T. 5 minutes is a generous bound on that visibility delay
--    and keeps the dashboard "today" bucket at most ~10 min stale
--    (5 min lag + 5 min cron period).
--  * On error we record `last_error` and re-raise; the watermark is NOT
--    advanced because the UPDATE that advances it only runs after the
--    upsert succeeds.
--  * SAFE TO RUN CONCURRENTLY WITH BACKFILL: the window primitive is
--    idempotent (see contract above), so even if cron fires while the
--    offline backfill is also walking history, the worst case is some
--    bucket gets written twice with the same value.
CREATE OR REPLACE FUNCTION rollup_task_usage_daily()
RETURNS BIGINT
LANGUAGE plpgsql
AS $$
DECLARE
    v_lock_ok BOOLEAN;
    v_from    TIMESTAMPTZ;
    v_to      TIMESTAMPTZ;
    v_rows    BIGINT := 0;
BEGIN
    SELECT pg_try_advisory_lock(4242) INTO v_lock_ok;
    IF NOT v_lock_ok THEN
        RETURN 0;
    END IF;

    BEGIN
        UPDATE task_usage_rollup_state
           SET last_run_started_at = now(),
               last_error          = NULL
         WHERE id = 1
        RETURNING watermark_at INTO v_from;

        v_to := now() - INTERVAL '5 minutes';

        IF v_from < v_to THEN
            v_rows := rollup_task_usage_daily_window(v_from, v_to);

            UPDATE task_usage_rollup_state
               SET watermark_at         = v_to,
                   last_run_finished_at = now(),
                   last_run_rows        = v_rows
             WHERE id = 1;
        ELSE
            UPDATE task_usage_rollup_state
               SET last_run_finished_at = now(),
                   last_run_rows        = 0
             WHERE id = 1;
        END IF;

        PERFORM pg_advisory_unlock(4242);
        RETURN v_rows;
    EXCEPTION WHEN OTHERS THEN
        UPDATE task_usage_rollup_state
           SET last_error           = SQLERRM,
               last_run_finished_at = now()
         WHERE id = 1;
        PERFORM pg_advisory_unlock(4242);
        RAISE;
    END;
END;
$$;