From 9a8d18eae79ef6afaa41720a71f44230df42ccb9 Mon Sep 17 00:00:00 2001 From: Camiel van Schoonhoven Date: Thu, 22 Oct 2026 11:50:45 -0700 Subject: [PATCH] feat(runs): submit and manage runs as a team A run looks up its secrets in the account of whoever created it, so a run that uses a team's secrets has to be created as that team. Once it is, the backend only lets the run's creator change it. These methods now take an optional team id and send the impersonation header when given one. Every caller still passes none, so nothing changes yet. - `submitPipelineRun` takes `teamId` and passes it to `createPipelineRun`. - `saveRunAnnotations` and `updateRunAnnotation` take it too. The notes and tags written right after submit are separate requests, and without the header they would be refused for a run the team created. - `cancelPipelineRun` takes it, for cancelling a run the team created. - `useTeamForRunAnnotations` resolves the team of the projects a run is filed under. That one rule serves both a run about to be submitted (the annotations it will carry) and a run already stored, so a rerun or cancel from a run page outside any project still finds the team. - `referencedSecretNames` and `missingSecretNames` find the secrets a run needs and which of them an account lacks, for a check before submit. A secret reference holds only a name, so a run submitted from the wrong context would otherwise fail at launch or quietly use a same-named secret from another account. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../secretReferences.test.ts | 87 +++++++++++++++++++ .../SecretsManagement/secretReferences.ts | 33 +++++++ .../Tangle/saveRunAnnotations.test.ts | 59 +++++++++++++ .../Submitters/Tangle/saveRunAnnotations.ts | 5 +- src/services/pipelineRunService.test.ts | 69 ++++++++++++++- src/services/pipelineRunService.ts | 12 ++- .../teams/useTeamForRunAnnotations.test.ts | 33 +++++++ .../teams/useTeamForRunAnnotations.ts | 10 +++ src/utils/submitPipeline.test.ts | 30 +++++++ src/utils/submitPipeline.ts | 2 + 10 files changed, 337 insertions(+), 3 deletions(-) create mode 100644 src/components/shared/SecretsManagement/secretReferences.test.ts create mode 100644 src/components/shared/SecretsManagement/secretReferences.ts create mode 100644 src/components/shared/Submitters/Tangle/saveRunAnnotations.test.ts create mode 100644 src/services/teams/useTeamForRunAnnotations.test.ts create mode 100644 src/services/teams/useTeamForRunAnnotations.ts diff --git a/src/components/shared/SecretsManagement/secretReferences.test.ts b/src/components/shared/SecretsManagement/secretReferences.test.ts new file mode 100644 index 0000000000..5095d950e9 --- /dev/null +++ b/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( + [], + ); + }); +}); diff --git a/src/components/shared/SecretsManagement/secretReferences.ts b/src/components/shared/SecretsManagement/secretReferences.ts new file mode 100644 index 0000000000..2433b8c7d4 --- /dev/null +++ b/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) { + 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(); + for (const source of sources) collectSecretNames(source, names); + return [...names]; +} + +export function missingSecretNames( + referenced: readonly string[], + available: readonly Pick[], +): string[] { + const availableNames = new Set(available.map((secret) => secret.name)); + return referenced.filter((name) => !availableNames.has(name)); +} diff --git a/src/components/shared/Submitters/Tangle/saveRunAnnotations.test.ts b/src/components/shared/Submitters/Tangle/saveRunAnnotations.test.ts new file mode 100644 index 0000000000..60c732eda9 --- /dev/null +++ b/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, + ); + }); +}); diff --git a/src/components/shared/Submitters/Tangle/saveRunAnnotations.ts b/src/components/shared/Submitters/Tangle/saveRunAnnotations.ts index 0a62bf7930..1399fdef7b 100644 --- a/src/components/shared/Submitters/Tangle/saveRunAnnotations.ts +++ b/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"); diff --git a/src/services/pipelineRunService.test.ts b/src/services/pipelineRunService.test.ts index bfb13668ca..5fc666bbc8 100644 --- a/src/services/pipelineRunService.test.ts +++ b/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 }); + }); +}); diff --git a/src/services/pipelineRunService.ts b/src/services/pipelineRunService.ts index a8f4e15ae1..e387beffa6 100644 --- a/src/services/pipelineRunService.ts +++ b/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), }); }; diff --git a/src/services/teams/useTeamForRunAnnotations.test.ts b/src/services/teams/useTeamForRunAnnotations.test.ts new file mode 100644 index 0000000000..9253d1c544 --- /dev/null +++ b/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([]); + }); +}); diff --git a/src/services/teams/useTeamForRunAnnotations.ts b/src/services/teams/useTeamForRunAnnotations.ts new file mode 100644 index 0000000000..e2062bbc19 --- /dev/null +++ b/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 | null | undefined, +): TeamResolution { + return useTeamForProjects(projectIdsFromAnnotations(annotations)); +} diff --git a/src/utils/submitPipeline.test.ts b/src/utils/submitPipeline.test.ts index a54dedb0f0..16edcd39c7 100644 --- a/src/utils/submitPipeline.test.ts +++ b/src/utils/submitPipeline.test.ts @@ -159,6 +159,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, mockAuthToken, + undefined, ); expect(pipelineRunService.savePipelineRun).toHaveBeenCalledWith( mockPipelineRun, @@ -198,6 +199,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -216,6 +218,25 @@ describe("submitPipelineRun", () => { expect.any(Object), mockBackendUrl, undefined, + undefined, + ); + }); + + it("submits the run as the team it is given", async () => { + const componentSpec: ComponentSpec = { + name: "team-component", + implementation: { container: { image: "test:latest" } }, + }; + + await submitPipelineRun(componentSpec, mockBackendUrl, { + teamId: "ml-team@project.iam.gserviceaccount.com", + }); + + expect(pipelineRunService.createPipelineRun).toHaveBeenCalledWith( + expect.any(Object), + mockBackendUrl, + undefined, + "ml-team@project.iam.gserviceaccount.com", ); }); @@ -448,6 +469,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -491,6 +513,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -533,6 +556,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -572,6 +596,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -604,6 +629,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -697,6 +723,7 @@ describe("submitPipelineRun", () => { }, mockBackendUrl, undefined, + undefined, ); }); @@ -861,6 +888,7 @@ describe("submitPipelineRun", () => { }), mockBackendUrl, undefined, + undefined, ); }); }); @@ -1108,6 +1136,7 @@ describe("submitPipelineRun", () => { }), mockBackendUrl, undefined, + undefined, ); }); @@ -1143,6 +1172,7 @@ describe("submitPipelineRun", () => { }), mockBackendUrl, undefined, + undefined, ); }); diff --git a/src/utils/submitPipeline.ts b/src/utils/submitPipeline.ts index 6af85773ea..ec2a093800 100644 --- a/src/utils/submitPipeline.ts +++ b/src/utils/submitPipeline.ts @@ -30,6 +30,7 @@ export async function submitPipelineRun( options?: { taskArguments?: Record; authorizationToken?: string; + teamId?: string; canonicalName?: string; runAnnotations?: Record; runSource?: string; @@ -116,6 +117,7 @@ export async function submitPipelineRun( payload as BodyCreateApiPipelineRunsPost, backendUrl, options?.authorizationToken, + options?.teamId, ); if (responseData.id) {