diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index d7bc02d754..502d92c5a8 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -35,6 +35,7 @@ import { type StoredLogEntry, sendableQueuePrefixLength, sessionSupportsNativeSteer, + type TaskRunArtifact, type TaskRunStatus, } from "@posthog/shared"; import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; @@ -1599,7 +1600,7 @@ export class SessionService { /** Deduplicates concurrent manifest reads when a message renders many images. */ private cloudAttachmentManifestRequests = new Map< string, - Promise> + Promise >(); private idleKilledSubscription: { unsubscribe: () => void } | null = null; /** @@ -4673,6 +4674,7 @@ export class SessionService { cloudStatus: run.status, cloudStage: run.stage ?? null, cloudOutput: run.output ?? null, + cloudArtifacts: run.artifacts ?? [], cloudErrorMessage: run.error_message, logUrl: run.log_url ?? session.logUrl, }); @@ -7293,22 +7295,34 @@ export class SessionService { } } + async getCloudRunArtifacts( + taskId: string, + runId: string, + ): Promise { + const authStatus = await this.getAuthCredentialsStatus(); + if (authStatus.kind !== "ready") return []; + + return this.getCloudAttachmentManifest( + authStatus.auth.client, + `${authStatus.auth.apiHost}:${authStatus.auth.projectId}`, + taskId, + runId, + ); + } + private getCloudAttachmentManifest( client: AuthClient, authIdentity: string, taskId: string, runId: string, - ): Promise> { + ): Promise { const key = `${authIdentity}:${taskId}:${runId}`; const existing = this.cloudAttachmentManifestRequests.get(key); if (existing) return existing; const request = client .getTaskRun(taskId, runId) - .then( - (run: { artifacts?: Array<{ id?: string; storage_path?: string }> }) => - run.artifacts ?? [], - ); + .then((run: { artifacts?: TaskRunArtifact[] }) => run.artifacts ?? []); this.cloudAttachmentManifestRequests.set(key, request); const clear = () => { diff --git a/packages/shared/src/domain-types.ts b/packages/shared/src/domain-types.ts index 273bf9a87b..e070d38788 100644 --- a/packages/shared/src/domain-types.ts +++ b/packages/shared/src/domain-types.ts @@ -3,6 +3,7 @@ import type { Adapter } from "./adapter"; import type { AgentRuntime } from "./agent-runtime"; import type { DismissalReasonOptionValue } from "./dismissal-reasons"; import type { StoredLogEntry } from "./session-events"; +import type { TaskRunArtifact } from "./task"; // Execution mode schema and type - shared between main and renderer export const executionModeSchema = z.enum([ @@ -180,6 +181,7 @@ export interface TaskRun { error_message: string | null; output: Record | null; // Structured output (PR URL, commit SHA, etc.) state: Record; // Intermediate run state (defaults to {}, never null) + artifacts?: TaskRunArtifact[]; created_at: string; updated_at: string; completed_at: string | null; diff --git a/packages/shared/src/sessions.ts b/packages/shared/src/sessions.ts index ca64dcb01e..76a49cd736 100644 --- a/packages/shared/src/sessions.ts +++ b/packages/shared/src/sessions.ts @@ -10,7 +10,7 @@ import type { Adapter } from "./adapter"; import type { SkillButtonId } from "./analytics-events"; import type { ExecutionMode } from "./exec-types"; import type { AcpMessage } from "./session-events"; -import type { TaskRunStatus } from "./task"; +import type { TaskRunArtifact, TaskRunStatus } from "./task"; export type { Adapter }; @@ -94,6 +94,7 @@ export interface AgentSession { cloudStatus?: TaskRunStatus; cloudStage?: string | null; cloudOutput?: Record | null; + cloudArtifacts?: TaskRunArtifact[]; cloudErrorMessage?: string | null; initialPrompt?: ContentBlock[]; cloudBranch?: string | null; diff --git a/packages/ui/src/features/sessions/components/CloudArtifactDownloads.test.tsx b/packages/ui/src/features/sessions/components/CloudArtifactDownloads.test.tsx new file mode 100644 index 0000000000..c7cc9f8581 --- /dev/null +++ b/packages/ui/src/features/sessions/components/CloudArtifactDownloads.test.tsx @@ -0,0 +1,88 @@ +import { Theme } from "@radix-ui/themes"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { CloudArtifactDownloads } from "./CloudArtifactDownloads"; + +const getCloudAttachmentPreviewUrl = vi.fn(); +const fetchedArtifacts = [ + { + id: "output-1", + name: "report.pdf", + type: "output", + size: 12_000, + storage_path: "tasks/run-1/report.pdf", + }, + { + id: "internal-1", + name: "handoff.pack", + type: "artifact", + storage_path: "tasks/run-1/handoff.pack", + }, +]; + +vi.mock("@posthog/core/sessions/sessionService", () => ({ + SESSION_SERVICE: Symbol("SESSION_SERVICE"), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: () => ({ getCloudAttachmentPreviewUrl }), +})); + +vi.mock("@posthog/ui/features/sessions/sessionStore", () => ({ + useSessionSelector: () => undefined, +})); + +vi.mock("@posthog/ui/features/auth/store", () => ({ + getAuthIdentity: () => "auth-1", + useAuthStateValue: () => "auth-1", +})); + +vi.mock("@tanstack/react-query", () => ({ + useQuery: () => ({ data: fetchedArtifacts }), +})); + +const task = { + id: "task-1", + latest_run: { + id: "run-1", + status: "completed", + }, +} as never; + +describe("CloudArtifactDownloads", () => { + beforeEach(() => { + getCloudAttachmentPreviewUrl.mockReset(); + }); + + it("shows output artifacts and opens their download URL", async () => { + getCloudAttachmentPreviewUrl.mockResolvedValue( + "https://files.example/report.pdf", + ); + const open = vi.spyOn(window, "open").mockImplementation(() => null); + + render( + + + , + ); + + expect(screen.getByText("report.pdf")).toBeInTheDocument(); + expect(screen.getByText("12 KB")).toBeInTheDocument(); + expect(screen.queryByText("handoff.pack")).not.toBeInTheDocument(); + + fireEvent.click(screen.getByRole("button", { name: "Download" })); + + await waitFor(() => + expect(open).toHaveBeenCalledWith( + "https://files.example/report.pdf", + "_blank", + "noopener,noreferrer", + ), + ); + expect(getCloudAttachmentPreviewUrl).toHaveBeenCalledWith( + "task-1", + "run-1", + "output-1", + ); + }); +}); diff --git a/packages/ui/src/features/sessions/components/CloudArtifactDownloads.tsx b/packages/ui/src/features/sessions/components/CloudArtifactDownloads.tsx new file mode 100644 index 0000000000..364629846c --- /dev/null +++ b/packages/ui/src/features/sessions/components/CloudArtifactDownloads.tsx @@ -0,0 +1,135 @@ +import { DownloadSimple } from "@phosphor-icons/react"; +import { + SESSION_SERVICE, + type SessionService, +} from "@posthog/core/sessions/sessionService"; +import { useService } from "@posthog/di/react"; +import { Button } from "@posthog/quill"; +import type { TaskRunArtifact } from "@posthog/shared"; +import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; +import { + getAuthIdentity, + useAuthStateValue, +} from "@posthog/ui/features/auth/store"; +import { useSessionSelector } from "@posthog/ui/features/sessions/sessionStore"; +import { FileIcon } from "@posthog/ui/primitives/FileIcon"; +import { toast } from "@posthog/ui/primitives/toast"; +import { Box, Flex, Text } from "@radix-ui/themes"; +import { useQuery } from "@tanstack/react-query"; +import { useCallback, useMemo, useState } from "react"; + +function formatFileSize(size: number | undefined): string | null { + if (size === undefined) return null; + if (size < 1_000) return `${size} B`; + if (size < 1_000_000) return `${Math.round(size / 1_000)} KB`; + return `${(size / 1_000_000).toFixed(1)} MB`; +} + +export function CloudArtifactDownloads({ + taskId, + task, +}: { + taskId: string | undefined; + task: Task | undefined; +}) { + const sessionService = useService(SESSION_SERVICE); + const sessionArtifacts = useSessionSelector( + taskId, + (session) => session?.cloudArtifacts, + ); + const cloudStatus = useSessionSelector( + taskId, + (session) => session?.cloudStatus, + ); + const authIdentity = useAuthStateValue(getAuthIdentity); + const [downloadingId, setDownloadingId] = useState(null); + const runId = task?.latest_run?.id; + const { data: fetchedArtifacts } = useQuery({ + queryKey: ["cloudRunArtifacts", authIdentity, taskId, runId], + queryFn: () => + sessionService.getCloudRunArtifacts(taskId ?? "", runId ?? ""), + enabled: + authIdentity !== null && + taskId !== undefined && + runId !== undefined && + isTerminalStatus(cloudStatus ?? task?.latest_run?.status), + retry: false, + staleTime: Infinity, + }); + const artifacts = useMemo( + () => + ( + fetchedArtifacts ?? + sessionArtifacts ?? + task?.latest_run?.artifacts ?? + [] + ).filter((artifact) => artifact.type === "output"), + [fetchedArtifacts, sessionArtifacts, task?.latest_run?.artifacts], + ); + + const downloadArtifact = useCallback( + async (artifact: TaskRunArtifact): Promise => { + if (!taskId || !runId || !artifact.id) return; + setDownloadingId(artifact.id); + try { + const url = await sessionService.getCloudAttachmentPreviewUrl( + taskId, + runId, + artifact.id, + ); + if (!url) { + toast.error("This file is no longer available"); + return; + } + window.open(url, "_blank", "noopener,noreferrer"); + } catch { + toast.error("Couldn't download file"); + } finally { + setDownloadingId(null); + } + }, + [runId, sessionService, taskId], + ); + + if (!runId || artifacts.length === 0) return null; + + return ( + + Files + + {artifacts.map((artifact) => { + const size = formatFileSize(artifact.size); + const canDownload = Boolean(artifact.id); + return ( + + + + {artifact.name} + {size !== null && ( + + {size} + + )} + + + + ); + })} + + + ); +} diff --git a/packages/ui/src/features/sessions/components/ConversationView.tsx b/packages/ui/src/features/sessions/components/ConversationView.tsx index 5a049e930a..9f6ab47eab 100644 --- a/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -15,6 +15,7 @@ import type { ConversationItem, TurnContext, } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import { CloudArtifactDownloads } from "@posthog/ui/features/sessions/components/CloudArtifactDownloads"; import { ConversationSearchBar } from "@posthog/ui/features/sessions/components/ConversationSearchBar"; import { PROMPT_RECALL_HINT_KEY, @@ -452,6 +453,7 @@ export function ConversationView({ const footer = (
+ + <> + + + } />