diff --git a/packages/ui/src/features/task-detail/components/CustomImageBadge.tsx b/packages/ui/src/features/task-detail/components/CustomImageBadge.tsx
new file mode 100644
index 0000000000..ac7b5b1dd1
--- /dev/null
+++ b/packages/ui/src/features/task-detail/components/CustomImageBadge.tsx
@@ -0,0 +1,76 @@
+import { Cube } from "@phosphor-icons/react";
+import type { Task } from "@posthog/shared/domain-types";
+import { Badge } from "../../../primitives/Badge";
+import { Tooltip } from "../../../primitives/Tooltip";
+import { openSettings } from "../../settings/hooks/useOpenSettings";
+import { useSandboxCustomImages } from "../../settings/sections/environments/useSandboxCustomImages";
+import { useSandboxEnvironments } from "../../settings/sections/environments/useSandboxEnvironments";
+
+export function CustomImageBadge({ task }: { task: Task }) {
+ const run = task.latest_run;
+ const state = run?.state as
+ | { custom_image_id?: unknown; sandbox_environment_id?: unknown }
+ | undefined;
+ const customImageId =
+ typeof state?.custom_image_id === "string" ? state.custom_image_id : null;
+ const sandboxEnvironmentId =
+ typeof state?.sandbox_environment_id === "string"
+ ? state.sandbox_environment_id
+ : null;
+
+ // Only mount the data-fetching part when the run could have used a custom image.
+ if (
+ run?.environment !== "cloud" ||
+ (!customImageId && !sandboxEnvironmentId)
+ ) {
+ return null;
+ }
+ return (
+
+ );
+}
+
+function ResolvedCustomImageBadge({
+ customImageId,
+ sandboxEnvironmentId,
+}: {
+ customImageId: string | null;
+ sandboxEnvironmentId: string | null;
+}) {
+ const { images } = useSandboxCustomImages();
+ const { environments } = useSandboxEnvironments();
+
+ const environment = sandboxEnvironmentId
+ ? environments.find((env) => env.id === sandboxEnvironmentId)
+ : undefined;
+ const imageId = customImageId ?? environment?.custom_image_id ?? null;
+ if (!imageId) return null;
+
+ const imageName =
+ images.find((image) => image.id === imageId)?.name ??
+ environment?.custom_image_name ??
+ null;
+
+ return (
+
+
+
+ );
+}
diff --git a/packages/ui/src/features/task-detail/components/TaskDetail.tsx b/packages/ui/src/features/task-detail/components/TaskDetail.tsx
index 707ca47443..4d62a700af 100644
--- a/packages/ui/src/features/task-detail/components/TaskDetail.tsx
+++ b/packages/ui/src/features/task-detail/components/TaskDetail.tsx
@@ -25,6 +25,7 @@ import { useWorkspace } from "../../workspace/useWorkspace";
import { useWorkspaceEvents } from "../../workspace/useWorkspaceEvents";
import { HeaderTitleEditor } from "../HeaderTitleEditor";
import { useTaskData } from "../hooks/useTaskData";
+import { CustomImageBadge } from "./CustomImageBadge";
import { ExternalAppsOpener } from "./ExternalAppsOpener";
import { WorkspaceModeBadge } from "./WorkspaceModeBadge";
@@ -147,12 +148,13 @@ export function TaskDetail({
channelName={channelName}
channelId={channelId}
leafIcon={
- workspaceMode ? (
+
- ) : undefined
+
+
}
leafLabel={task.title}
onRename={handleTitleEditSubmit}
@@ -172,6 +174,7 @@ export function TaskDetail({
mode={workspaceMode}
checkoutPath={effectiveRepoPath}
/>
+