diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 9eb76b1cfe..42a36681cb 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -451,6 +451,50 @@ describe("PostHogAPIClient", () => { ); }); + it("presigns a task run artifact for preview", async () => { + const fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + url: "https://s3.example.com/screenshot.png?signature=abc", + expires_in: 3600, + }), + }); + const client = new PostHogAPIClient( + "http://localhost:8000", + async () => "token", + async () => "token", + 123, + ); + + ( + client as unknown as { + api: { baseUrl: string; fetcher: { fetch: typeof fetch } }; + } + ).api = { + baseUrl: "http://localhost:8000", + fetcher: { fetch }, + }; + + await expect( + client.presignTaskRunArtifact( + "task-123", + "run-123", + "tasks/run-123/artifacts/screenshot.png", + ), + ).resolves.toBe("https://s3.example.com/screenshot.png?signature=abc"); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ + method: "post", + path: "/api/projects/123/tasks/task-123/runs/run-123/artifacts/presign/", + overrides: { + body: JSON.stringify({ + storage_path: "tasks/run-123/artifacts/screenshot.png", + }), + }, + }), + ); + }); + it("returns the redirect URL when authorizing an MCP installation", async () => { const fetch = vi.fn().mockResolvedValue({ ok: true, diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 6468f58ace..9d1907d4ec 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -2900,6 +2900,34 @@ export class PostHogAPIClient { return data.artifacts ?? []; } + async presignTaskRunArtifact( + taskId: string, + runId: string, + storagePath: string, + ): Promise { + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/presign/`, + ); + const response = await this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/presign/`, + overrides: { + body: JSON.stringify({ storage_path: storagePath }), + }, + }); + + if (!response.ok) { + throw new Error( + `Failed to generate artifact preview URL: ${response.statusText}`, + ); + } + + const data = (await response.json()) as { url: string }; + return data.url; + } + async resumeRunInCloud(taskId: string, runId: string): Promise { const teamId = await this.getTeamId(); const url = new URL( diff --git a/packages/core/src/sessions/promptContent.test.ts b/packages/core/src/sessions/promptContent.test.ts index 7bd174e974..6cfe9dee15 100644 --- a/packages/core/src/sessions/promptContent.test.ts +++ b/packages/core/src/sessions/promptContent.test.ts @@ -49,7 +49,8 @@ describe("promptContent", () => { }); it("extracts cloud resource_link attachments from file URIs", () => { - const fileUri = "file:///tmp/workspace/attachments/Receipt-2264-0277.pdf"; + const fileUri = + "file:///tmp/workspace/.posthog/attachments/run-123/artifact-456/Receipt-2264-0277.pdf"; const result = extractPromptDisplayContent([ { type: "text", text: "what is this about?" }, @@ -62,7 +63,23 @@ describe("promptContent", () => { expect(result.text).toBe("what is this about?"); expect(result.attachments).toEqual([ - { id: fileUri, label: "Receipt-2264-0277.pdf" }, + { + id: fileUri, + label: "Receipt-2264-0277.pdf", + cloudArtifact: { runId: "run-123", artifactId: "artifact-456" }, + }, + ]); + }); + + it("does not mark ordinary file URIs as cloud artifacts", () => { + const fileUri = "file:///tmp/screenshot.png"; + + const result = extractPromptDisplayContent([ + { type: "resource_link", uri: fileUri, name: "screenshot.png" }, + ]); + + expect(result.attachments).toEqual([ + { id: fileUri, label: "screenshot.png" }, ]); }); }); diff --git a/packages/core/src/sessions/promptContent.ts b/packages/core/src/sessions/promptContent.ts index 5754d7f4e1..fb19f6296b 100644 --- a/packages/core/src/sessions/promptContent.ts +++ b/packages/core/src/sessions/promptContent.ts @@ -23,6 +23,30 @@ export function makeAttachmentUri(filePath: string): string { export interface AttachmentRef { id: string; label: string; + cloudArtifact?: CloudArtifactRef; +} + +export interface CloudArtifactRef { + runId: string; + artifactId: string; +} + +function parseCloudArtifactRef(pathname: string): CloudArtifactRef | undefined { + const segments = pathname.split("/").filter(Boolean); + const posthogIndex = segments.lastIndexOf(".posthog"); + if ( + posthogIndex < 0 || + segments[posthogIndex + 1] !== "attachments" || + !segments[posthogIndex + 2] || + !segments[posthogIndex + 3] + ) { + return undefined; + } + + return { + runId: segments[posthogIndex + 2], + artifactId: segments[posthogIndex + 3], + }; } export function parseAttachmentUri(uri: string): AttachmentRef | null { @@ -56,7 +80,8 @@ function parseFileUri( const pathname = decodeURIComponent(new URL(uri).pathname); const label = fallbackLabel?.trim() || getFileName(pathname) || "attachment"; - return { id: uri, label }; + const cloudArtifact = parseCloudArtifactRef(pathname); + return { id: uri, label, ...(cloudArtifact ? { cloudArtifact } : {}) }; } catch { const label = fallbackLabel?.trim() || getFileName(uri) || "attachment"; return { id: uri, label }; diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 65830d943c..38a1ede091 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1595,6 +1595,11 @@ export class SessionService { string, Promise >(); + /** Deduplicates concurrent manifest reads when a message renders many images. */ + private cloudAttachmentManifestRequests = new Map< + string, + Promise> + >(); private idleKilledSubscription: { unsubscribe: () => void } | null = null; /** * Cached preview-config-options responses keyed by `${apiHost}::${adapter}`. @@ -7088,6 +7093,69 @@ export class SessionService { } } + async getCloudAttachmentPreviewUrl( + taskId: string, + runId: string, + artifactId: string, + ): Promise { + const authStatus = await this.getAuthCredentialsStatus(); + if (authStatus.kind !== "ready") return null; + + try { + const artifacts = await this.getCloudAttachmentManifest( + authStatus.auth.client, + `${authStatus.auth.apiHost}:${authStatus.auth.projectId}`, + taskId, + runId, + ); + const artifact = artifacts.find( + (candidate) => candidate.id === artifactId, + ); + if (!artifact?.storage_path) return null; + + return await authStatus.auth.client.presignTaskRunArtifact( + taskId, + runId, + artifact.storage_path, + ); + } catch (error) { + this.d.log.warn("Failed to resolve cloud attachment preview", { + taskId, + runId, + artifactId, + error: String(error), + }); + return null; + } + } + + private getCloudAttachmentManifest( + client: AuthClient, + authIdentity: string, + taskId: string, + runId: string, + ): 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 ?? [], + ); + this.cloudAttachmentManifestRequests.set(key, request); + + const clear = () => { + if (this.cloudAttachmentManifestRequests.get(key) === request) { + this.cloudAttachmentManifestRequests.delete(key); + } + }; + void request.then(clear, clear); + return request; + } + // --- Helper Methods --- private async resolveCloudPrompt( diff --git a/packages/ui/src/features/sessions/components/UserMessageAttachments.tsx b/packages/ui/src/features/sessions/components/UserMessageAttachments.tsx new file mode 100644 index 0000000000..7b4b4d45b4 --- /dev/null +++ b/packages/ui/src/features/sessions/components/UserMessageAttachments.tsx @@ -0,0 +1,140 @@ +import { File } from "@phosphor-icons/react"; +import { + SESSION_SERVICE, + type SessionService, +} from "@posthog/core/sessions/sessionService"; +import { useService } from "@posthog/di/react"; +import { isRasterImageFile, parseImageDataUrl } from "@posthog/shared"; +import { + getAuthIdentity, + useAuthStateValue, +} from "@posthog/ui/features/auth/store"; +import { readFileAsDataUrl } from "@posthog/ui/features/message-editor/hostApi"; +import { MentionChip } from "@posthog/ui/features/sessions/components/session-update/parseFileMentions"; +import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes"; +import { useSessionTaskId } from "@posthog/ui/features/sessions/useSessionTaskId"; +import { SafeImagePreview } from "@posthog/ui/primitives/SafeImagePreview"; +import { Dialog, Text } from "@radix-ui/themes"; +import { useQuery } from "@tanstack/react-query"; + +function attachmentFilePath(id: string): string | null { + if (id.startsWith("/") || /^[A-Za-z]:[\\/]/.test(id)) return id; + if (!id.startsWith("file://")) return null; + + try { + return decodeURIComponent(new URL(id).pathname); + } catch { + return null; + } +} + +function ImageAttachment({ + attachment, +}: { + attachment: UserMessageAttachment; +}) { + const filePath = attachmentFilePath(attachment.id); + const taskId = useSessionTaskId(); + const authIdentity = useAuthStateValue(getAuthIdentity); + const sessionService = useService(SESSION_SERVICE); + const cloudArtifact = attachment.cloudArtifact; + const { data: previewUrl } = useQuery({ + queryKey: cloudArtifact + ? [ + "cloudArtifactPreview", + authIdentity, + taskId, + cloudArtifact.runId, + cloudArtifact.artifactId, + ] + : ["os", "readFileAsDataUrl", filePath], + queryFn: () => { + if (cloudArtifact && taskId) { + return sessionService.getCloudAttachmentPreviewUrl( + taskId, + cloudArtifact.runId, + cloudArtifact.artifactId, + ); + } + return readFileAsDataUrl({ filePath: filePath ?? "" }); + }, + enabled: cloudArtifact + ? taskId !== null && authIdentity !== null + : filePath !== null, + retry: false, + staleTime: cloudArtifact ? 50 * 60 * 1000 : Infinity, + }); + const parsedImage = previewUrl?.startsWith("data:") + ? parseImageDataUrl(previewUrl) + : null; + + if (!previewUrl) { + return } label={attachment.label} />; + } + + return ( + + + + + + + {attachment.label} + + {parsedImage ? ( + + ) : previewUrl.startsWith("data:") ? ( + + Unable to load image preview + + ) : ( + {attachment.label} + )} + + + ); +} + +export function UserMessageAttachments({ + attachments, +}: { + attachments: UserMessageAttachment[]; +}) { + return ( +
+ {attachments.map((attachment) => + isRasterImageFile(attachment.label) ? ( + + ) : ( + } + label={attachment.label} + /> + ), + )} +
+ ); +} diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 75a79c707a..5e1fc63c67 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -67,6 +67,7 @@ import { } from "@posthog/ui/features/sessions/components/session-update/parseFileMentions"; import { SessionUpdateView } from "@posthog/ui/features/sessions/components/session-update/SessionUpdateView"; import { UserShellExecuteView } from "@posthog/ui/features/sessions/components/session-update/UserShellExecuteView"; +import { UserMessageAttachments } from "@posthog/ui/features/sessions/components/UserMessageAttachments"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; import { useAgentConversationItems } from "@posthog/ui/features/sessions/hooks/useAgentConversationItems"; @@ -406,14 +407,8 @@ function UserBubble({ )} {attachments.length > 0 && !containsFileMentions && ( -
- {attachments.map((attachment) => ( - } - label={attachment.label} - /> - ))} +
+
)} {isOverflowing && ( diff --git a/packages/ui/src/features/sessions/components/session-update/UserMessage.tsx b/packages/ui/src/features/sessions/components/session-update/UserMessage.tsx index e6478da0e3..00cf6ff966 100644 --- a/packages/ui/src/features/sessions/components/session-update/UserMessage.tsx +++ b/packages/ui/src/features/sessions/components/session-update/UserMessage.tsx @@ -1,7 +1,6 @@ import { Check, Copy, - File, FileText, Scroll, SlackLogo, @@ -15,6 +14,7 @@ import { MarkdownRenderer } from "../../../editor/components/MarkdownRenderer"; import { useFeatureFlag } from "../../../feature-flags/useFeatureFlag"; import { usePanelLayoutStore } from "../../../panels/panelLayoutStore"; import type { UserMessageAttachment } from "../../userMessageTypes"; +import { UserMessageAttachments } from "../UserMessageAttachments"; import { CollapsibleMessageContent } from "./CollapsibleMessageContent"; import { extractCanvasInstructions } from "./canvasInstructions"; import { extractChannelContext } from "./channelContext"; @@ -179,19 +179,9 @@ export const UserMessage = memo(function UserMessage({ )} {showAttachmentChips && ( - - {attachments.map((attachment) => ( - } - label={attachment.label} - /> - ))} - +
+ +
)} {sourceUrl && ( diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index 88deaf7e04..90a90ca460 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -133,6 +133,7 @@ const mockAuthenticatedClient = vi.hoisted(() => ({ finalizeTaskRunArtifactUploads: vi.fn(), prepareTaskStagedArtifactUploads: vi.fn(), finalizeTaskStagedArtifactUploads: vi.fn(), + presignTaskRunArtifact: vi.fn(), startGithubUserIntegrationConnect: vi.fn(), getTaskRunSessionLogs: vi.fn(), getTaskRunSessionLogsResult: vi.fn(), @@ -577,6 +578,169 @@ describe("SessionService", () => { ); }); + describe("cloud attachment previews", () => { + it("deduplicates concurrent manifest requests for the same run", async () => { + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ + artifacts: [ + { + id: "artifact-1", + storage_path: "tasks/run-123/artifacts/one.png", + }, + { + id: "artifact-2", + storage_path: "tasks/run-123/artifacts/two.png", + }, + ], + }); + mockAuthenticatedClient.presignTaskRunArtifact.mockImplementation( + async (_taskId, _runId, storagePath) => + `https://s3.example.com/${storagePath}`, + ); + const service = getSessionService(); + + await Promise.all([ + service.getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-1", + ), + service.getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-2", + ), + ]); + + expect(mockAuthenticatedClient.getTaskRun).toHaveBeenCalledTimes(1); + expect( + mockAuthenticatedClient.presignTaskRunArtifact, + ).toHaveBeenCalledTimes(2); + }); + + it("does not share manifest requests across projects", async () => { + const resolvers: Array<(value: { artifacts: unknown[] }) => void> = []; + mockAuthenticatedClient.getTaskRun.mockImplementation( + () => + new Promise((resolve) => { + resolvers.push(resolve); + }), + ); + const service = getSessionService(); + const first = service.getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-1", + ); + await vi.waitFor(() => + expect(mockAuthenticatedClient.getTaskRun).toHaveBeenCalledTimes(1), + ); + + mockAuth.fetchAuthState.mockResolvedValue({ + status: "authenticated", + bootstrapComplete: true, + cloudRegion: "us", + orgProjectsMap: {}, + currentOrgId: "org-2", + currentProjectId: 456, + hasCodeAccess: true, + needsScopeReauth: false, + }); + const second = service.getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-1", + ); + await vi.waitFor(() => + expect(mockAuthenticatedClient.getTaskRun).toHaveBeenCalledTimes(2), + ); + + for (const resolve of resolvers) resolve({ artifacts: [] }); + await Promise.all([first, second]); + }); + + it("resolves an artifact id to a presigned URL", async () => { + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ + artifacts: [ + { + id: "artifact-456", + storage_path: "tasks/run-123/artifacts/screenshot.png", + }, + ], + }); + mockAuthenticatedClient.presignTaskRunArtifact.mockResolvedValue( + "https://s3.example.com/screenshot.png?signature=abc", + ); + + await expect( + getSessionService().getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-456", + ), + ).resolves.toBe("https://s3.example.com/screenshot.png?signature=abc"); + expect( + mockAuthenticatedClient.presignTaskRunArtifact, + ).toHaveBeenCalledWith( + "task-123", + "run-123", + "tasks/run-123/artifacts/screenshot.png", + ); + }); + + it("returns null when the artifact is absent", async () => { + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ artifacts: [] }); + + await expect( + getSessionService().getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "missing", + ), + ).resolves.toBeNull(); + expect( + mockAuthenticatedClient.presignTaskRunArtifact, + ).not.toHaveBeenCalled(); + }); + + it("returns null when authentication is unavailable", async () => { + mockAuth.fetchAuthState.mockResolvedValue({ + status: "anonymous", + bootstrapComplete: true, + }); + + await expect( + getSessionService().getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-456", + ), + ).resolves.toBeNull(); + expect(mockAuthenticatedClient.getTaskRun).not.toHaveBeenCalled(); + }); + + it("returns null when presigning fails", async () => { + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ + artifacts: [ + { + id: "artifact-456", + storage_path: "tasks/run-123/artifacts/screenshot.png", + }, + ], + }); + mockAuthenticatedClient.presignTaskRunArtifact.mockRejectedValue( + new Error("presign unavailable"), + ); + + await expect( + getSessionService().getCloudAttachmentPreviewUrl( + "task-123", + "run-123", + "artifact-456", + ), + ).resolves.toBeNull(); + }); + }); + describe("connectToTask", () => { it("skips local connection for cloud runs", async () => { const service = getSessionService(); diff --git a/packages/ui/src/features/sessions/userMessageTypes.ts b/packages/ui/src/features/sessions/userMessageTypes.ts index 1209a353a3..339c788106 100644 --- a/packages/ui/src/features/sessions/userMessageTypes.ts +++ b/packages/ui/src/features/sessions/userMessageTypes.ts @@ -1,4 +1,3 @@ -export interface UserMessageAttachment { - id: string; - label: string; -} +import type { AttachmentRef } from "@posthog/core/sessions/promptContent"; + +export type UserMessageAttachment = AttachmentRef;