feat: introduce method to fetch test durations from the Flakiness.io

README.md

@@ -102,11 +102,31 @@ Use this entry point when you need to process or manipulate reports in browser-b
 
 ### Working with Reports
 - **`readReport()`** - Read a Flakiness report and its attachments from disk
+- **`fetchTestDurations()`** - Fetch historical test durations from Flakiness.io and return a report enriched with timings
 - **`showReport()`** - Start a local server and open the report in your browser
 - **`showReportCommand()`** - Build a shell command for opening the report later with the Flakiness CLI
 - **`uploadReport()`** - Upload reports and attachments to Flakiness.io
 - **`writeReport()`** - Write reports to disk in the standard Flakiness report format
 
+## Fetching Test Durations
+
+`fetchTestDurations()` sends a report to Flakiness.io and returns a copy enriched
+with historical test durations. Test runners can use these timings to split tests
+into balanced shards.
+
+```typescript
+import { fetchTestDurations } from '@flakiness/sdk';
+
+const reportWithDurations = await fetchTestDurations(report, {
+  flakinessAccessToken: 'your-token',
+});
+```
+
+Authentication follows the same priority order as `uploadReport()`:
+
+1. **Access token** — pass `flakinessAccessToken` option or set the `FLAKINESS_ACCESS_TOKEN` environment variable.
+2. **GitHub Actions OIDC** — when running inside GitHub Actions and the report has `flakinessProject` set.
+
 ## Uploading Reports
 
 `uploadReport()` authenticates using one of the following methods (in order of priority):

src/fetchTestDurations.ts

@@ -0,0 +1,177 @@
+import { FlakinessReport } from '@flakiness/flakiness-report';
+import { URL } from 'url';
+import { compressTextAsync, retryWithBackoff, sha1Text } from './_internalUtils.js';
+import { GithubOIDC } from './githubOIDC.js';
+
+type TestDurationsFetcherOptions = {
+  flakinessEndpoint: string;
+  flakinessAccessToken: string;
+}
+
+/**
+ * Options for {@link fetchTestDurations}.
+ */
+export type FetchTestDurationsOptions = {
+  /**
+   * Custom Flakiness.io endpoint URL.
+   *
+   * Defaults to the `FLAKINESS_ENDPOINT` environment variable, or 'https://flakiness.io'
+   * if the environment variable is not set.
+   *
+   * @example 'https://custom.flakiness.io'
+   */
+  flakinessEndpoint?: string;
+
+  /**
+   * Access token for authenticating with the Flakiness.io platform.
+   *
+   * Defaults to the `FLAKINESS_ACCESS_TOKEN` environment variable. If no token is provided
+   * through this option or the environment variable, the function will attempt to authenticate
+   * via GitHub Actions OIDC when running in GitHub Actions (requires `report.flakinessProject`
+   * to be set and the project to be bound to the repository).
+   *
+   * @example 'flakiness-io-1234567890abcdef...'
+   */
+  flakinessAccessToken?: string;
+}
+
+// Retry schedule for uploading the report to the presigned URL.
+const HTTP_BACKOFF = [100, 500, 1000, 1000, 1000, 1000];
+
+// The computed historical durations are not ready immediately after submit, so
+// we poll the download URL for up to ~90 seconds before giving up.
+const DOWNLOAD_BACKOFF = [Array(10).fill(1000), Array(10).fill(2000), Array(20).fill(3000)].flat();
+
+/**
+ * Fetches historical test durations for a report from the Flakiness.io platform.
+ *
+ * This is used to compute "balanced shards" — by knowing how long each test took
+ * historically, a test runner can balance tests across shards so that every shard
+ * finishes at roughly the same time.
+ *
+ * The function performs the following steps:
+ * 1. Authenticates using an access token or GitHub Actions OIDC.
+ * 2. Computes a shard-group key from the report so that all shards of the same run
+ *    fetch an identical set of timings.
+ * 3. Uploads the (compressed) report so the platform knows which tests to time.
+ * 4. Submits the request and polls the resulting download URL until the computed
+ *    durations are ready (up to ~90 seconds).
+ *
+ * ## Authentication
+ *
+ * Authentication follows the same priority order as {@link uploadReport}:
+ * 1. **Access token** — provided via `flakinessAccessToken` option or `FLAKINESS_ACCESS_TOKEN` env var.
+ * 2. **GitHub Actions OIDC** — when running in GitHub Actions with no access token. This requires
+ *    `report.flakinessProject` to be set and the project to be bound to the GitHub repository.
+ *
+ * @param {FlakinessReport.Report} report - The report describing the tests to fetch durations for.
+ * @param {FetchTestDurationsOptions} options - Optional configuration object.
+ *
+ * @returns {Promise<FlakinessReport.Report>} A report enriched with historical test durations.
+ *
+ * @throws {Error} If no access token is available, any API call fails, or the durations are not
+ *   ready within the polling timeout.
+ *
+ * @example
+ * ```typescript
+ * const reportWithDurations = await fetchTestDurations(report);
+ * ```
+ */
+export async function fetchTestDurations(
+  report: FlakinessReport.Report,
+  options?: FetchTestDurationsOptions,
+): Promise<FlakinessReport.Report> {
+  let flakinessAccessToken = options?.flakinessAccessToken ?? process.env['FLAKINESS_ACCESS_TOKEN'];
+
+  const githubOIDC = GithubOIDC.initializeFromEnv();
+  if (!flakinessAccessToken && githubOIDC && report.flakinessProject)
+    flakinessAccessToken = await githubOIDC.createFlakinessAccessToken(report.flakinessProject);
+
+  if (!flakinessAccessToken)
+    throw new Error('No Flakiness access token available (set FLAKINESS_ACCESS_TOKEN, pass `flakinessAccessToken`, or run in GitHub Actions with `id-token: write` and a configured `flakinessProject`)');
+
+  const flakinessEndpoint = options?.flakinessEndpoint ?? process.env['FLAKINESS_ENDPOINT'] ?? 'https://flakiness.io';
+  const fetcher = new TestDurationsFetcher(report, { flakinessAccessToken, flakinessEndpoint });
+  return await fetcher.fetch();
+}
+
+class TestDurationsFetcher {
+  private _report: FlakinessReport.Report;
+  private _options: TestDurationsFetcherOptions;
+
+  constructor(report: FlakinessReport.Report, options: TestDurationsFetcherOptions) {
+    this._report = report;
+    this._options = options;
+  }
+
+  private async _api<OUTPUT>(pathname: string, token: string, body?: any): Promise<OUTPUT> {
+    const url = new URL(this._options.flakinessEndpoint);
+    url.pathname = pathname;
+    return await retryWithBackoff(async () => {
+      const response = await fetch(url, {
+        method: 'POST',
+        headers: {
+          'Authorization': `Bearer ${token}`,
+          'Content-Type': 'application/json',
+        },
+        body: body ? JSON.stringify(body) : undefined,
+      });
+      if (!response.ok)
+        throw new Error(`Failed to fetch ${url}: ${response.status} ${response.statusText}`)
+      return response;
+    }, HTTP_BACKOFF).then(async response => await response.json());
+  }
+
+  async fetch(): Promise<FlakinessReport.Report> {
+    // Shard group key makes sure that all shards fetch the same timings.
+    // We do NOT use envs as a shard group since environments might 
+    // fluctuate between runs: the shard jobs might run on different versions of the 
+    // OS, i.e. when Github Actions do a gradual rollout of a new VM Image.
+    const shardGroupKey = sha1Text(JSON.stringify({
+      commitId: this._report.commitId,
+      category: this._report.category,
+      testRunnerName: this._report.testRunner?.name ?? 'unknown',
+      testRunnerVersion: this._report.testRunner?.version ?? 'unknown',
+    }));
+
+    const createResponse = await this._api<{ testDurationsToken: string, uploadUrl: string }>(
+      '/api/testDurations/create',
+      this._options.flakinessAccessToken,
+      { commitId: this._report.commitId, shardGroupKey },
+    );
+
+    await this._uploadReport(JSON.stringify(this._report), createResponse.uploadUrl);
+
+    const submitResponse = await this._api<{ downloadUrl: string }>(
+      '/api/testDurations/submit',
+      createResponse.testDurationsToken,
+    );
+
+    return await retryWithBackoff(async () => {
+      const response = await fetch(submitResponse.downloadUrl);
+      if (!response.ok)
+        throw new Error(`Request to ${submitResponse.downloadUrl} failed with ${response.status}`);
+      return response;
+    }, DOWNLOAD_BACKOFF).then(async response => await response.json() as FlakinessReport.Report);
+  }
+
+  private async _uploadReport(data: string, uploadUrl: string) {
+    const compressed = await compressTextAsync(data);
+    const headers = {
+      'Content-Type': 'application/json',
+      'Content-Length': Buffer.byteLength(compressed) + '',
+      'Content-Encoding': 'br',
+    };
+    await retryWithBackoff(async () => {
+      const response = await fetch(uploadUrl, {
+        method: 'PUT',
+        headers,
+        body: Buffer.from(compressed),
+      });
+      if (!response.ok)
+        throw new Error(`Request to ${uploadUrl} failed with ${response.status}`);
+      // Read response to ensure it completes
+      await response.arrayBuffer();
+    }, HTTP_BACKOFF);
+  }
+}

src/index.ts

@@ -7,6 +7,7 @@ export { GithubOIDC } from './githubOIDC.js';
 export * as ReportUtils from './reportUtils.js';
 
 // Working with reports
+export { fetchTestDurations, type FetchTestDurationsOptions } from './fetchTestDurations.js';
 export { readReport } from './readReport.js';
 export { showReport } from './showReport.js';
 export { showReportCommand } from './showReportCommand.js';