v3 run bar: hold the time estimate inside the API's reported stage

app/src/api/societyWideCalculation.ts

@@ -29,6 +29,12 @@ export interface SocietyWideCalculationResponse {
   result: SocietyWideReportOutput | null;
   error?: string;
   policyengine_bundle?: PolicyEngineBundle | null;
+  /** Optional stage fields while computing (policyengine-sim-api#781 via the API). */
+  stage?: string;
+  progress?: number;
+  progress_max?: number;
+  stage_started_at?: string;
+  segments?: { completed: number; total: number };
 }
 
 function mergePolicyEngineBundle(

app/src/components/flagship/report/RunProgress.tsx

@@ -1,7 +1,8 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { Progress, Spinner, Stack, Text } from '@/components/ui';
 import { colors, spacing, typography } from '@/designTokens';
 import { formatElapsed } from '@/libs/flagship/reportStages';
+import type { RunStage } from '@/types/calculation';
 
 /** Milliseconds since the component mounted, ticking each second. */
 function useElapsed(): number {
@@ -53,16 +54,52 @@ export function RunLoading({ label }: { label: string }) {
   );
 }
 
+/** One line per stage the API reports; an unknown stage gets no line. */
+export function runStageCaption(runStage?: RunStage): string | undefined {
+  switch (runStage?.stage) {
+    case 'queued':
+      return 'Waiting for a worker';
+    case 'loading_data':
+      return 'Loading the data';
+    case 'baseline':
+      return 'Running current law';
+    case 'reform':
+      return 'Running the reform';
+    case 'simulating':
+      return runStage.segments
+        ? `Running the regions, ${runStage.segments.completed} of ${runStage.segments.total} done`
+        : 'Running the regions';
+    case 'outputs':
+      return 'Calculating the results';
+    default:
+      return undefined;
+  }
+}
+
+/**
+ * The bar's value: the time estimate, held inside the stage the API reports
+ * (from where the stage starts to just before the next one), else the
+ * estimate alone.
+ */
+export function runProgressValue(estimate: number, runStage?: RunStage): number {
+  if (typeof runStage?.floor !== 'number' || typeof runStage.max !== 'number') {
+    return estimate;
+  }
+  return Math.min(Math.max(estimate, runStage.floor), runStage.max);
+}
+
 /**
  * A society-wide run in progress, as a bar: what is running, how long it
- * has taken, and how far along a typical run of this size would be. The
- * API reports no stages for a single-year run, so the bar follows the
- * run's typical duration (the calculation strategy's estimate, held short
- * of the end until the results arrive).
+ * has taken, and how far along it is. The bar follows the run's typical
+ * duration (the calculation strategy's estimate, held short of the end
+ * until the results arrive). When the API reports the run's stage, the
+ * estimate is held inside that stage, and the stage names the line under
+ * the title. The bar never moves backwards.
  */
 export default function RunProgress({
   title,
   progress,
+  runStage,
   detail,
   expectedMs,
   note,
@@ -71,6 +108,8 @@ export default function RunProgress({
   title: string;
   /** 0–100, from the calculation status; absent, the bar follows the clock. */
   progress?: number;
+  /** The run's stage from the API, when it reports one. */
+  runStage?: RunStage;
   /** One line under the title, e.g. a queue position. */
   detail?: string;
   /** How long a run like this one typically takes. */
@@ -79,9 +118,15 @@ export default function RunProgress({
   note?: string;
 }) {
   const elapsed = useElapsed();
-  const value =
+  const estimate =
     typeof progress === 'number' ? progress : Math.min((elapsed / expectedMs) * 100, 95);
+  // The highest value shown so far: a stage's range can start lower than an
+  // earlier estimate (a queued run's path is not known yet).
+  const shown = useRef(0);
+  shown.current = Math.max(shown.current, runProgressValue(estimate, runStage));
+  const value = shown.current;
   const percent = Math.max(0, Math.min(Math.round(value), 99));
+  const line = detail ?? runStageCaption(runStage);
 
   return (
     <Stack role="status" aria-live="polite" style={CARD_STYLE}>
@@ -110,9 +155,9 @@ export default function RunProgress({
         </Text>
       </Stack>
       <Progress value={value} aria-label={`${title}, about ${percent}%`} className="tw:h-2" />
-      {detail && (
+      {line && (
         <Text style={{ margin: 0, fontSize: typography.fontSize.xs, color: colors.text.primary }}>
-          {detail}
+          {line}
         </Text>
       )}
       {note && (

app/src/hooks/useBillEconomy.ts

@@ -7,7 +7,7 @@ import { useCurrentCountry } from '@/hooks/useCurrentCountry';
 import { SocietyWideCalcStrategy } from '@/libs/calculations/economy/SocietyWideCalcStrategy';
 import { createReformPolicy } from '@/libs/flagship/runReport';
 import type { RootState } from '@/store';
-import type { CalcStatus } from '@/types/calculation';
+import type { CalcStatus, RunStage } from '@/types/calculation';
 
 const strategy = new SocietyWideCalcStrategy();
 
@@ -18,6 +18,8 @@ export interface BillEconomy {
   queuePosition?: number;
   /** 0–100 while pending: the calculation strategy's estimate from typical run times. */
   progress?: number;
+  /** The run's stage from the API while pending, when it reports one. */
+  runStage?: RunStage;
   reformPolicyId?: string;
   baselinePolicyId?: string;
   year: string;
@@ -118,6 +120,7 @@ export function useBillEconomy(
         : economy.data?.message,
     queuePosition: economy.data?.queuePosition,
     progress: economy.data?.progress,
+    runStage: economy.data?.runStage,
     reformPolicyId,
     baselinePolicyId,
     year,

app/src/hooks/useCalculationStatus.ts

@@ -121,6 +121,7 @@ function useSingleCalculationStatus(calcId: string, targetType: 'report' | 'simu
 
     // Keep server data available
     queuePosition: status?.queuePosition,
+    runStage: status?.runStage,
     result: status?.result,
     error: status?.error,
     metadata: status?.metadata,

app/src/libs/calculations/economy/SocietyWideCalcStrategy.ts

@@ -5,7 +5,7 @@ import {
   SocietyWideCalculationResponse,
 } from '@/api/societyWideCalculation';
 import { getDurationForCountry, getDurationForRegion } from '@/constants/calculationDurations';
-import { CalcMetadata, CalcParams, CalcStatus } from '@/types/calculation';
+import { CalcMetadata, CalcParams, CalcStatus, RunStage } from '@/types/calculation';
 import { CalcExecutionStrategy, RefetchConfig } from '../strategies/types';
 
 /** How a caller wants the run's duration estimated. */
@@ -105,10 +105,13 @@ export class SocietyWideCalcStrategy implements CalcExecutionStrategy {
       const estimatedDuration = this.getEstimatedDuration(countryId, estimate);
       const progress = Math.min((elapsed / estimatedDuration) * 100, 95); // Cap at 95%
 
+      const runStage = estimate.estimateByRegion ? this.getRunStage(response) : undefined;
+
       return {
         status: 'pending',
         progress,
         ...(estimate.estimateByRegion && { estimatedDurationMs: estimatedDuration }),
+        ...(runStage && { runStage }),
         queuePosition: response.queue_position,
         message: this.getComputingMessage(response.queue_position),
         metadata,
@@ -149,6 +152,24 @@ export class SocietyWideCalcStrategy implements CalcExecutionStrategy {
     return this.transformResponseWithMetadata(apiResponse, metadata);
   }
 
+  /**
+   * The API's stage fields, when it reports any
+   */
+  private getRunStage(response: SocietyWideCalculationResponse): RunStage | undefined {
+    const runStage: RunStage = {};
+    if (typeof response.stage === 'string') {
+      runStage.stage = response.stage;
+    }
+    if (typeof response.progress === 'number' && typeof response.progress_max === 'number') {
+      runStage.floor = response.progress;
+      runStage.max = Math.max(response.progress, response.progress_max);
+    }
+    if (response.segments) {
+      runStage.segments = response.segments;
+    }
+    return Object.keys(runStage).length > 0 ? runStage : undefined;
+  }
+
   /**
    * Get human-readable message for computing state
    */

app/src/pages/flagship/BillReport.page.tsx

@@ -15,7 +15,7 @@ import StoredImpactCharts from '@/components/flagship/report/StoredImpactCharts'
 import { ErrorState } from '@/components/flagship/states';
 import { BillExternalEstimates, useModelTrackRecord } from '@/components/flagship/ValidationPanel';
 import { Button, Stack, Text } from '@/components/ui';
-import { getDurationForCountry } from '@/constants/calculationDurations';
+import { getDurationForRegion } from '@/constants/calculationDurations';
 import { CongressionalDistrictDataProvider } from '@/contexts/CongressionalDistrictDataContext';
 import { colors, spacing, typography } from '@/designTokens';
 import { useBillEconomy, type BillEconomy } from '@/hooks/useBillEconomy';
@@ -203,8 +203,9 @@ function FullResultsStatus({
     <RunProgress
       title="Calculating the full results"
       progress={economy.progress}
+      runStage={economy.runStage}
       detail={queued}
-      expectedMs={getDurationForCountry(countryId)}
+      expectedMs={getDurationForRegion(countryId, economy.region)}
       note="Later visits load from the cache."
     />
   );

app/src/pages/flagship/Report.page.tsx

@@ -13,7 +13,7 @@ import { ReportUnresolvable } from '@/components/flagship/ReportComputing';
 import { ErrorState } from '@/components/flagship/states';
 import { trackRecordPrograms, useModelTrackRecord } from '@/components/flagship/ValidationPanel';
 import { Stack, Title } from '@/components/ui';
-import { getDurationForCountry } from '@/constants/calculationDurations';
+import { getDurationForRegion } from '@/constants/calculationDurations';
 import { colors, spacing } from '@/designTokens';
 import { useCalculationStatus } from '@/hooks/useCalculationStatus';
 import { useFlagshipMetadataStatus } from '@/hooks/useFlagshipMetadataStatus';
@@ -266,12 +266,13 @@ export default function FlagshipReportPage({ userReportId: propId }: FlagshipRep
                 : 'Calculating the nationwide results'
             }
             progress={calcStatus.progress}
+            runStage={'runStage' in calcStatus ? calcStatus.runStage : undefined}
             detail={
               typeof calcStatus.queuePosition === 'number' && calcStatus.queuePosition > 0
                 ? `In queue, position ${calcStatus.queuePosition}`
                 : undefined
             }
-            expectedMs={getDurationForCountry(report?.countryId ?? 'us')}
+            expectedMs={getDurationForRegion(report?.countryId ?? 'us', region)}
           />
         ) : (
           <RunLoading label="Loading report" />

app/src/tests/unit/components/flagship/report/RunProgress.test.tsx

@@ -0,0 +1,89 @@
+import { render, screen } from '@test-utils';
+import { describe, expect, test } from 'vitest';
+import RunProgress, {
+  runProgressValue,
+  runStageCaption,
+} from '@/components/flagship/report/RunProgress';
+
+describe('runProgressValue', () => {
+  test('given no run stage then the time estimate is the value', () => {
+    expect(runProgressValue(40)).toBe(40);
+  });
+
+  test('given an estimate behind the stage then the stage start holds the bar', () => {
+    expect(runProgressValue(30, { stage: 'reform', floor: 67, max: 96 })).toBe(67);
+  });
+
+  test('given an estimate past the stage then the bar waits below the next stage', () => {
+    expect(runProgressValue(99, { stage: 'reform', floor: 67, max: 96 })).toBe(96);
+  });
+
+  test('given an estimate inside the stage then the estimate is the value', () => {
+    expect(runProgressValue(80, { stage: 'reform', floor: 67, max: 96 })).toBe(80);
+  });
+});
+
+describe('runStageCaption', () => {
+  test('given a segmented run then the caption counts the regions', () => {
+    expect(runStageCaption({ stage: 'simulating', segments: { completed: 19, total: 20 } })).toBe(
+      'Running the regions, 19 of 20 done'
+    );
+  });
+
+  test('given a stage the app does not know then there is no caption', () => {
+    expect(runStageCaption({ stage: 'a_future_stage' })).toBeUndefined();
+  });
+});
+
+describe('RunProgress', () => {
+  test('given a run stage then the stage names the line under the title', () => {
+    render(
+      <RunProgress
+        title="Calculating the nationwide results"
+        expectedMs={360000}
+        progress={10}
+        runStage={{ stage: 'reform', floor: 67, max: 96 }}
+      />
+    );
+
+    expect(screen.getByText('Running the reform')).toBeInTheDocument();
+    expect(screen.getByText(/67%/)).toBeInTheDocument();
+  });
+
+  test('given a queue detail then it is shown in place of the stage', () => {
+    render(
+      <RunProgress
+        title="Calculating the nationwide results"
+        expectedMs={360000}
+        progress={1}
+        detail="In queue, position 2"
+        runStage={{ stage: 'queued', floor: 0, max: 24 }}
+      />
+    );
+
+    expect(screen.getByText('In queue, position 2')).toBeInTheDocument();
+    expect(screen.queryByText('Waiting for a worker')).not.toBeInTheDocument();
+  });
+
+  test('given a later stage range that starts lower then the bar does not move backwards', () => {
+    const { rerender } = render(
+      <RunProgress
+        title="Calculating the nationwide results"
+        expectedMs={360000}
+        progress={20}
+        runStage={{ stage: 'queued', floor: 0, max: 24 }}
+      />
+    );
+
+    rerender(
+      <RunProgress
+        title="Calculating the nationwide results"
+        expectedMs={360000}
+        progress={21}
+        runStage={{ stage: 'loading_data', floor: 3, max: 4 }}
+      />
+    );
+
+    expect(screen.getByText(/· 20%/)).toBeInTheDocument();
+  });
+});

app/src/tests/unit/libs/calculations/strategies/SocietyWideCalcStrategy.test.ts

@@ -319,6 +319,55 @@ describe('SocietyWideCalcStrategy', () => {
       expect(result.progress).toBeCloseTo(50, 0);
       expect(result).not.toHaveProperty('estimatedDurationMs');
     });
+
+    it('given estimateByRegion and API stage fields then passes the run stage on', () => {
+      // When
+      const result = strategy.transformResponseWithMetadata(
+        {
+          ...mockSocietyWideComputingResponse(),
+          stage: 'simulating',
+          progress: 5,
+          progress_max: 96,
+          segments: { completed: 19, total: 20 },
+        },
+        startedAgo(1000),
+        'us',
+        { estimateByRegion: true }
+      );
+
+      // Then
+      expect(result.runStage).toEqual({
+        stage: 'simulating',
+        floor: 5,
+        max: 96,
+        segments: { completed: 19, total: 20 },
+      });
+    });
+
+    it('given API stage fields without estimateByRegion then the v2 status is unchanged', () => {
+      // When
+      const result = strategy.transformResponseWithMetadata(
+        { ...mockSocietyWideComputingResponse(), stage: 'reform', progress: 67, progress_max: 96 },
+        startedAgo(1000),
+        'us'
+      );
+
+      // Then
+      expect(result).not.toHaveProperty('runStage');
+    });
+
+    it('given estimateByRegion and no API stage fields then there is no run stage', () => {
+      // When
+      const result = strategy.transformResponseWithMetadata(
+        mockSocietyWideComputingResponse(),
+        startedAgo(1000),
+        'us',
+        { estimateByRegion: true }
+      );
+
+      // Then
+      expect(result).not.toHaveProperty('runStage');
+    });
   });
 
   describe('getRefetchConfig', () => {