feat(runs): submit and manage runs as a team

src/components/shared/SecretsManagement/secretReferences.test.ts

@@ -0,0 +1,87 @@
+import { describe, expect, it } from "vitest";
+
+import type { ComponentSpec } from "@/utils/componentSpec";
+
+import { missingSecretNames, referencedSecretNames } from "./secretReferences";
+import { createSecretArgument } from "./types";
+
+const pipelineWithSecrets: ComponentSpec = {
+  name: "pipeline",
+  implementation: {
+    graph: {
+      tasks: {
+        train: {
+          componentRef: {
+            spec: {
+              name: "train",
+              implementation: {
+                graph: {
+                  tasks: {
+                    fetch: {
+                      componentRef: {},
+                      arguments: {
+                        token: createSecretArgument("HF_TOKEN"),
+                      },
+                    },
+                  },
+                },
+              },
+            },
+          },
+          arguments: {
+            api_key: createSecretArgument("OPENAI_API_KEY"),
+            epochs: "3",
+          },
+        },
+        evaluate: {
+          componentRef: {},
+          arguments: { api_key: createSecretArgument("OPENAI_API_KEY") },
+        },
+      },
+    },
+  },
+};
+
+describe("referencedSecretNames", () => {
+  it("finds every secret a pipeline's tasks take, nested subgraphs included, once each", () => {
+    expect(referencedSecretNames(pipelineWithSecrets)).toEqual([
+      "HF_TOKEN",
+      "OPENAI_API_KEY",
+    ]);
+  });
+
+  it("includes the secrets passed as run arguments", () => {
+    expect(
+      referencedSecretNames(pipelineWithSecrets, {
+        db_password: createSecretArgument("DB_PASSWORD"),
+        region: "us-east1",
+      }),
+    ).toEqual(["HF_TOKEN", "OPENAI_API_KEY", "DB_PASSWORD"]);
+  });
+
+  it("is empty for a pipeline that takes no secrets", () => {
+    expect(
+      referencedSecretNames({
+        name: "plain",
+        implementation: { container: { image: "test:latest" } },
+      }),
+    ).toEqual([]);
+  });
+});
+
+describe("missingSecretNames", () => {
+  it("lists the referenced secrets the account does not have", () => {
+    expect(
+      missingSecretNames(
+        ["HF_TOKEN", "OPENAI_API_KEY"],
+        [{ name: "OPENAI_API_KEY" }, { name: "UNUSED" }],
+      ),
+    ).toEqual(["HF_TOKEN"]);
+  });
+
+  it("is empty when every referenced secret exists", () => {
+    expect(missingSecretNames(["HF_TOKEN"], [{ name: "HF_TOKEN" }])).toEqual(
+      [],
+    );
+  });
+});

src/components/shared/SecretsManagement/secretReferences.ts

@@ -0,0 +1,33 @@
+import { isSecretArgument } from "@/utils/componentSpec";
+import { isRecord } from "@/utils/typeGuards";
+
+import { extractSecretName, type Secret } from "./types";
+
+function collectSecretNames(value: unknown, names: Set<string>) {
+  if (isSecretArgument(value)) {
+    const name = extractSecretName(value);
+    if (name) names.add(name);
+    return;
+  }
+  if (Array.isArray(value)) {
+    for (const item of value) collectSecretNames(item, names);
+    return;
+  }
+  if (isRecord(value)) {
+    for (const item of Object.values(value)) collectSecretNames(item, names);
+  }
+}
+
+export function referencedSecretNames(...sources: unknown[]): string[] {
+  const names = new Set<string>();
+  for (const source of sources) collectSecretNames(source, names);
+  return [...names];
+}
+
+export function missingSecretNames(
+  referenced: readonly string[],
+  available: readonly Pick<Secret, "name">[],
+): string[] {
+  const availableNames = new Set(available.map((secret) => secret.name));
+  return referenced.filter((name) => !availableNames.has(name));
+}

src/components/shared/Submitters/Tangle/saveRunAnnotations.test.ts

@@ -0,0 +1,59 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@/services/pipelineRunService", () => ({
+  updateRunAnnotation: vi.fn(),
+}));
+
+import { updateRunAnnotation } from "@/services/pipelineRunService";
+import type { ComponentSpec } from "@/utils/componentSpec";
+
+import { saveRunAnnotations } from "./saveRunAnnotations";
+
+const BACKEND_URL = "https://api.example.com";
+const TEAM_ID = "ml-team@project.iam.gserviceaccount.com";
+
+const taggedSpec: ComponentSpec = {
+  name: "tagged",
+  metadata: { annotations: { tags: "nightly,training" } },
+  implementation: { container: { image: "test:latest" } },
+};
+
+describe("saveRunAnnotations", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(updateRunAnnotation).mockResolvedValue(undefined);
+  });
+
+  it("writes the notes and tags as the team that started the run", async () => {
+    await saveRunAnnotations(
+      "42",
+      BACKEND_URL,
+      { notes: "Baseline", componentSpec: taggedSpec },
+      TEAM_ID,
+    );
+
+    expect(updateRunAnnotation).toHaveBeenCalledWith(
+      "42",
+      BACKEND_URL,
+      { key: "notes", value: "Baseline" },
+      TEAM_ID,
+    );
+    expect(updateRunAnnotation).toHaveBeenCalledWith(
+      "42",
+      BACKEND_URL,
+      { key: "tags", value: "nightly,training" },
+      TEAM_ID,
+    );
+  });
+
+  it("writes them as the user without a team", async () => {
+    await saveRunAnnotations("42", BACKEND_URL, { notes: "Baseline" });
+
+    expect(updateRunAnnotation).toHaveBeenCalledWith(
+      "42",
+      BACKEND_URL,
+      { key: "notes", value: "Baseline" },
+      undefined,
+    );
+  });
+});

src/components/shared/Submitters/Tangle/saveRunAnnotations.ts

@@ -22,6 +22,7 @@ export async function saveRunAnnotations(
   runId: string,
   backendUrl: string,
   { notes, componentSpec }: RunAnnotationsFromSubmission,
+  teamId?: string,
 ) {
   const writes: { key: string; value: string }[] = [];
 
@@ -37,7 +38,9 @@ export async function saveRunAnnotations(
   // Neither annotation depends on the other, so one being refused must not take
   // the other down with it.
   const settled = await Promise.allSettled(
-    writes.map((write) => updateRunAnnotation(runId, backendUrl, write)),
+    writes.map((write) =>
+      updateRunAnnotation(runId, backendUrl, write, teamId),
+    ),
   );
 
   const failed = settled.filter((result) => result.status === "rejected");

src/services/pipelineRunService.test.ts

@@ -10,7 +10,11 @@ import {
 
 import type { BodyCreateApiPipelineRunsPost } from "@/api/types.gen";
 
-import { createPipelineRun } from "./pipelineRunService";
+import {
+  cancelPipelineRun,
+  createPipelineRun,
+  updateRunAnnotation,
+} from "./pipelineRunService";
 
 const BACKEND_URL = "https://api.example.com";
 const PAYLOAD = {} as BodyCreateApiPipelineRunsPost;
@@ -90,3 +94,66 @@ describe("createPipelineRun", () => {
     );
   });
 });
+
+describe("acting as a team", () => {
+  const TEAM_ID = "ml-team@project.iam.gserviceaccount.com";
+  const HEADER_NAME = "x-tangle-impersonate-team-id";
+  let mockFetch: Mock;
+
+  const sentHeaders = () => {
+    const [, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+    return init.headers;
+  };
+
+  beforeEach(() => {
+    mockFetch = vi.fn().mockResolvedValue(jsonResponse(200, { id: 123 }));
+    global.fetch = mockFetch;
+    window.__TANGLE_TEAM_IMPERSONATION__ = { headerName: HEADER_NAME };
+  });
+
+  afterEach(() => {
+    delete window.__TANGLE_TEAM_IMPERSONATION__;
+    vi.restoreAllMocks();
+  });
+
+  it("creates the run as the team", async () => {
+    await createPipelineRun(PAYLOAD, BACKEND_URL, "token", TEAM_ID);
+
+    expect(sentHeaders()).toEqual({
+      "Content-Type": "application/json",
+      Authorization: "Bearer token",
+      [HEADER_NAME]: TEAM_ID,
+    });
+  });
+
+  it("creates the run as the user without a team", async () => {
+    await createPipelineRun(PAYLOAD, BACKEND_URL);
+
+    expect(sentHeaders()).toEqual({ "Content-Type": "application/json" });
+  });
+
+  it("cancels the run as the team", async () => {
+    await cancelPipelineRun("42", BACKEND_URL, TEAM_ID);
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      `${BACKEND_URL}/api/pipeline_runs/42/cancel`,
+      expect.objectContaining({ method: "POST" }),
+    );
+    expect(sentHeaders()).toEqual({ [HEADER_NAME]: TEAM_ID });
+  });
+
+  it("sets an annotation as the team", async () => {
+    await updateRunAnnotation(
+      "42",
+      BACKEND_URL,
+      { key: "notes", value: "Rerun with more memory" },
+      TEAM_ID,
+    );
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      `${BACKEND_URL}/api/pipeline_runs/42/annotations/notes?value=Rerun+with+more+memory`,
+      expect.objectContaining({ method: "PUT" }),
+    );
+    expect(sentHeaders()).toEqual({ [HEADER_NAME]: TEAM_ID });
+  });
+});

src/services/pipelineRunService.ts

@@ -14,6 +14,7 @@ import {
   getPipelineCreationBackendUrl,
 } from "@/services/pipelineStorage/pipelineCreation";
 import type { PipelineStorageService } from "@/services/pipelineStorage/PipelineStorageService";
+import { teamRequestHeaders } from "@/services/teams/teamRequestHeaders";
 import type { PipelineRun } from "@/types/pipelineRun";
 import { EDITOR_FLOW_DIRECTION_ANNOTATION } from "@/utils/annotations";
 import { removeCachingStrategyFromSpec } from "@/utils/cache";
@@ -31,6 +32,7 @@ export const createPipelineRun = async (
   payload: BodyCreateApiPipelineRunsPost,
   backendUrl: string,
   authorizationToken?: string,
+  teamId?: string,
 ) => {
   const authorizationHeader = authorizationToken
     ? { Authorization: `Bearer ${authorizationToken}` }
@@ -41,6 +43,7 @@ export const createPipelineRun = async (
     headers: {
       "Content-Type": "application/json",
       ...authorizationHeader,
+      ...teamRequestHeaders(teamId),
     },
     body: JSON.stringify(payload),
   });
@@ -200,11 +203,16 @@ export const fetchPipelineRuns = async (
   }
 };
 
-export const cancelPipelineRun = async (runId: string, backendUrl: string) => {
+export const cancelPipelineRun = async (
+  runId: string,
+  backendUrl: string,
+  teamId?: string,
+) => {
   await fetchWithErrorHandling(
     `${backendUrl}/api/pipeline_runs/${runId}/cancel`,
     {
       method: "POST",
+      headers: teamRequestHeaders(teamId),
     },
   );
 };
@@ -224,6 +232,7 @@ export const updateRunAnnotation = async (
     key: string;
     value: string;
   },
+  teamId?: string,
 ) => {
   const url = new URL(
     `${backendUrl}/api/pipeline_runs/${runId}/annotations/${annotation.key}`,
@@ -232,5 +241,6 @@ export const updateRunAnnotation = async (
 
   await fetchWithErrorHandling(url.toString(), {
     method: "PUT",
+    headers: teamRequestHeaders(teamId),
   });
 };

src/services/teams/useTeamForRunAnnotations.test.ts

@@ -0,0 +1,33 @@
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+
+vi.mock("./useTeamForProjects", () => ({
+  useTeamForProjects: vi.fn(() => ({ status: "none" })),
+}));
+
+import { projectRunAnnotations } from "@/utils/projectRunAnnotation";
+
+import { useTeamForProjects } from "./useTeamForProjects";
+import { useTeamForRunAnnotations } from "./useTeamForRunAnnotations";
+
+describe("useTeamForRunAnnotations", () => {
+  it("resolves the team of the projects a run is filed under", () => {
+    const annotations = {
+      ...projectRunAnnotations(["research-project", "data-project"]),
+      source: "web-app",
+    };
+
+    renderHook(() => useTeamForRunAnnotations(annotations));
+
+    expect(useTeamForProjects).toHaveBeenCalledWith([
+      "research-project",
+      "data-project",
+    ]);
+  });
+
+  it("resolves no projects for a run filed under none", () => {
+    renderHook(() => useTeamForRunAnnotations(null));
+
+    expect(useTeamForProjects).toHaveBeenLastCalledWith([]);
+  });
+});

src/services/teams/useTeamForRunAnnotations.ts

@@ -0,0 +1,10 @@
+import { projectIdsFromAnnotations } from "@/utils/projectRunAnnotation";
+
+import type { TeamResolution } from "./types";
+import { useTeamForProjects } from "./useTeamForProjects";
+
+export function useTeamForRunAnnotations(
+  annotations: Record<string, unknown> | null | undefined,
+): TeamResolution {
+  return useTeamForProjects(projectIdsFromAnnotations(annotations));
+}