From d3c97a0d3d36f95d7073997281e4e654252ea987 Mon Sep 17 00:00:00 2001 From: lordspline <74811063+lordspline@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:53:46 +0000 Subject: [PATCH 1/2] feat(aio): hydrate screenshot attachments in sandbox agents --- .../agent/src/server/agent-server.test.ts | 197 ++++++++++++++++++ .../packages/agent/src/server/agent-server.ts | 28 ++- .../agent/src/server/pi-agent-server.test.ts | 124 +++++++++++ .../agent/src/server/pi-agent-server.ts | 21 +- .../src/server/resolve-user-artifacts.ts | 65 ++++++ .../packages/agent/src/server/schemas.test.ts | 34 ++- .../packages/agent/src/server/schemas.ts | 10 +- 7 files changed, 468 insertions(+), 11 deletions(-) create mode 100644 products/desktop/packages/agent/src/server/resolve-user-artifacts.ts diff --git a/products/desktop/packages/agent/src/server/agent-server.test.ts b/products/desktop/packages/agent/src/server/agent-server.test.ts index 838c43e8cd0f..702befbdc184 100644 --- a/products/desktop/packages/agent/src/server/agent-server.test.ts +++ b/products/desktop/packages/agent/src/server/agent-server.test.ts @@ -2373,6 +2373,203 @@ describe("AgentServer HTTP Mode", () => { }); }, 20000); + it("resolves screenshot IDs into native ACP image blocks", async () => { + const pngAttachmentId = "11111111-1111-4111-8111-111111111111"; + const jpegAttachmentId = "22222222-2222-4222-8222-222222222222"; + const s = createServer(); + await s.start(); + const prompt = vi.fn(async () => ({ stopReason: "cancelled" })); + const getTaskRun = vi.fn(async () => + createTaskRun({ + artifacts: [ + { + id: pngAttachmentId, + name: "screen.png", + type: "context", + source: "user_attachment", + size: 3, + content_type: "image/png", + storage_path: "artifacts/screen.png", + uploaded_by: "user", + uploaded_by_user_id: 1, + }, + { + id: jpegAttachmentId, + name: "photo.jpg", + type: "context", + source: "user_attachment", + size: 4, + content_type: "image/jpeg", + storage_path: "artifacts/photo.jpg", + uploaded_by: "user", + uploaded_by_user_id: 1, + }, + ], + }), + ); + const downloadArtifact = vi.fn( + async (_taskId: string, _runId: string, storagePath: string) => + exactArrayBuffer( + new TextEncoder().encode( + storagePath.endsWith("screen.png") ? "png" : "jpeg", + ), + ), + ); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } }; + posthogAPI: { + getTaskRun: typeof getTaskRun; + downloadArtifact: typeof downloadArtifact; + }; + }; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.posthogAPI.getTaskRun = getTaskRun; + serverInternals.posthogAPI.downloadArtifact = downloadArtifact; + + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${createToken()}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "screenshots", + method: "user_message", + params: { + content: "Compare these screenshots", + artifact_ids: [pngAttachmentId, jpegAttachmentId], + }, + }), + }); + + expect(response.status).toBe(200); + expect(getTaskRun).toHaveBeenCalledWith("test-task-id", "test-run-id"); + expect(prompt).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: [ + { type: "text", text: "Compare these screenshots" }, + { + type: "image", + data: Buffer.from("png").toString("base64"), + mimeType: "image/png", + }, + { + type: "image", + data: Buffer.from("jpeg").toString("base64"), + mimeType: "image/jpeg", + }, + ], + }), + ); + }, 20000); + + it("rejects screenshot IDs missing from the current ACP run manifest", async () => { + const missingAttachmentId = "11111111-1111-4111-8111-111111111111"; + const s = createServer(); + await s.start(); + const prompt = vi.fn(async () => ({ stopReason: "cancelled" })); + const getTaskRun = vi.fn(async () => createTaskRun({ artifacts: [] })); + const downloadArtifact = vi.fn(); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } }; + posthogAPI: { + getTaskRun: typeof getTaskRun; + downloadArtifact: typeof downloadArtifact; + }; + }; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.posthogAPI.getTaskRun = getTaskRun; + serverInternals.posthogAPI.downloadArtifact = downloadArtifact; + + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${createToken()}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "missing-screenshot", + method: "user_message", + params: { + content: "Inspect this screenshot", + artifact_ids: [missingAttachmentId], + }, + }), + }); + const body = (await response.json()) as { + error?: { code?: number; message?: string }; + }; + + expect(body.error).toEqual({ + code: -32000, + message: "Screenshot attachments are unavailable for this task run", + }); + expect(downloadArtifact).not.toHaveBeenCalled(); + expect(prompt).not.toHaveBeenCalled(); + }, 20000); + + it("rejects IDs for non-screenshot artifacts in the current ACP run manifest", async () => { + const skillArtifactId = "11111111-1111-4111-8111-111111111111"; + const s = createServer(); + await s.start(); + const prompt = vi.fn(async () => ({ stopReason: "cancelled" })); + const getTaskRun = vi.fn(async () => + createTaskRun({ + artifacts: [ + { + id: skillArtifactId, + name: "bundle.zip", + type: "skill_bundle", + source: "posthog_code_skill", + size: 100, + content_type: "application/zip", + storage_path: "artifacts/bundle.zip", + uploaded_by: "agent", + }, + ], + }), + ); + const downloadArtifact = vi.fn(); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } }; + posthogAPI: { + getTaskRun: typeof getTaskRun; + downloadArtifact: typeof downloadArtifact; + }; + }; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.posthogAPI.getTaskRun = getTaskRun; + serverInternals.posthogAPI.downloadArtifact = downloadArtifact; + + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${createToken()}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "non-screenshot", + method: "user_message", + params: { + artifact_ids: [skillArtifactId], + }, + }), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toMatchObject({ + error: { + code: -32000, + message: "Screenshot attachments are unavailable for this task run", + }, + }); + expect(downloadArtifact).not.toHaveBeenCalled(); + expect(prompt).not.toHaveBeenCalled(); + }, 20000); + it("rewrites a bundled local skill slash command before sending the prompt", async () => { const skillDefinition = [ "---", diff --git a/products/desktop/packages/agent/src/server/agent-server.ts b/products/desktop/packages/agent/src/server/agent-server.ts index 866f7c284341..a583d1ae3b38 100644 --- a/products/desktop/packages/agent/src/server/agent-server.ts +++ b/products/desktop/packages/agent/src/server/agent-server.ts @@ -95,7 +95,7 @@ import type { TaskRun, TaskRunArtifact, } from "../types"; -import { resourceLink } from "../utils/acp-content"; +import { image, resourceLink } from "../utils/acp-content"; import { AsyncMutex } from "../utils/async-mutex"; import { withTimeout } from "../utils/common"; import { resolveGatewayProduct, resolveGatewayTarget } from "../utils/gateway"; @@ -113,6 +113,7 @@ import { checkoutExistingPullRequest, type ExistingPrCheckoutResult, } from "./pr-checkout"; +import { resolveUserArtifactsById } from "./resolve-user-artifacts"; import { resolveRtkSavings } from "./rtk-savings"; import { RunUsageAccumulator } from "./run-usage"; import { @@ -1180,12 +1181,26 @@ export class AgentServer { artifactCount: Array.isArray(params.artifacts) ? params.artifacts.length : 0, + artifactIdCount: Array.isArray(params.artifact_ids) + ? params.artifact_ids.length + : 0, }); + const resolvedArtifacts = await resolveUserArtifactsById( + this.posthogAPI, + commandSession.payload.task_id, + commandSession.payload.run_id, + Array.isArray(params.artifact_ids) + ? (params.artifact_ids as string[]) + : [], + ); const builtPrompt = await this.buildPromptFromContentAndArtifacts({ content: params.content as string | ContentBlock[] | undefined, - artifacts: Array.isArray(params.artifacts) - ? (params.artifacts as TaskRunArtifact[]) - : [], + artifacts: [ + ...(Array.isArray(params.artifacts) + ? (params.artifacts as TaskRunArtifact[]) + : []), + ...resolvedArtifacts, + ], taskId: commandSession.payload.task_id, runId: commandSession.payload.run_id, }); @@ -3331,6 +3346,11 @@ export class AgentServer { throw new Error(`Failed to download artifact ${artifact.name}`); } + const mimeType = artifact.content_type?.toLowerCase(); + if (mimeType === "image/png" || mimeType === "image/jpeg") { + return image(Buffer.from(data).toString("base64"), mimeType); + } + const safeName = this.getSafeArtifactName(artifact.name); const artifactDir = join( this.config.repositoryPath ?? "/tmp/workspace", diff --git a/products/desktop/packages/agent/src/server/pi-agent-server.test.ts b/products/desktop/packages/agent/src/server/pi-agent-server.test.ts index ad68a77f922d..c73c12539156 100644 --- a/products/desktop/packages/agent/src/server/pi-agent-server.test.ts +++ b/products/desktop/packages/agent/src/server/pi-agent-server.test.ts @@ -275,13 +275,16 @@ describe("PiAgentServer", () => { const sendCommand = vi.fn( async (_command: Record) => ({}), ); + const getTaskRun = vi.fn(); const server = new PiAgentServer(config()) as unknown as { + posthogAPI: { getTaskRun: typeof getTaskRun }; session: unknown; executeCommand( method: string, params: Record, ): Promise; }; + server.posthogAPI.getTaskRun = getTaskRun; server.session = { runtime: { client: { @@ -302,6 +305,7 @@ describe("PiAgentServer", () => { message: "hello", images: [], }); + expect(getTaskRun).not.toHaveBeenCalled(); }); it("preserves the native Pi user prompt when auto-publish is enabled", async () => { @@ -397,6 +401,126 @@ describe("PiAgentServer", () => { await rm(repositoryPath, { recursive: true }); }); + it("resolves screenshot IDs into native Pi image inputs", async () => { + const pngAttachmentId = "11111111-1111-4111-8111-111111111111"; + const jpegAttachmentId = "22222222-2222-4222-8222-222222222222"; + const sendCommand = vi.fn( + async (_command: Record) => ({}), + ); + const getTaskRun = vi.fn(async () => ({ + artifacts: [ + { + id: pngAttachmentId, + name: "screen.png", + type: "context", + source: "user_attachment", + size: 3, + content_type: "image/png", + storage_path: "artifacts/screen.png", + uploaded_by: "user", + uploaded_by_user_id: 1, + }, + { + id: jpegAttachmentId, + name: "photo.jpg", + type: "context", + source: "user_attachment", + size: 4, + content_type: "image/jpeg", + storage_path: "artifacts/photo.jpg", + uploaded_by: "user", + uploaded_by_user_id: 1, + }, + ], + })); + const downloadArtifact = vi.fn( + async (_taskId: string, _runId: string, storagePath: string) => + Buffer.from(storagePath.endsWith("screen.png") ? "png" : "jpeg"), + ); + const server = new PiAgentServer(config()) as unknown as { + posthogAPI: { + getTaskRun: typeof getTaskRun; + downloadArtifact: typeof downloadArtifact; + }; + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.posthogAPI.getTaskRun = getTaskRun; + server.posthogAPI.downloadArtifact = downloadArtifact; + server.session = { + runtime: { + client: { getState: vi.fn(async () => ({ isStreaming: false })) }, + sendCommand, + }, + }; + + await server.executeCommand("user_message", { + content: "Compare these screenshots", + artifact_ids: [pngAttachmentId, jpegAttachmentId], + }); + + expect(getTaskRun).toHaveBeenCalledWith("task-1", "run-1"); + expect(sendCommand).toHaveBeenCalledWith({ + id: expect.any(String), + type: "prompt", + message: "Compare these screenshots", + images: [ + { + type: "image", + data: Buffer.from("png").toString("base64"), + mimeType: "image/png", + fileName: "screen.png", + }, + { + type: "image", + data: Buffer.from("jpeg").toString("base64"), + mimeType: "image/jpeg", + fileName: "photo.jpg", + }, + ], + }); + }); + + it("rejects screenshot IDs missing from the current Pi run manifest", async () => { + const missingAttachmentId = "11111111-1111-4111-8111-111111111111"; + const getTaskRun = vi.fn(async () => ({ artifacts: [] })); + const downloadArtifact = vi.fn(); + const sendCommand = vi.fn(); + const server = new PiAgentServer(config()) as unknown as { + posthogAPI: { + getTaskRun: typeof getTaskRun; + downloadArtifact: typeof downloadArtifact; + }; + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.posthogAPI.getTaskRun = getTaskRun; + server.posthogAPI.downloadArtifact = downloadArtifact; + server.session = { + runtime: { + client: { getState: vi.fn(async () => ({ isStreaming: false })) }, + sendCommand, + }, + }; + + await expect( + server.executeCommand("user_message", { + content: "Inspect this", + artifact_ids: [missingAttachmentId], + }), + ).rejects.toThrow( + "Screenshot attachments are unavailable for this task run", + ); + expect(downloadArtifact).not.toHaveBeenCalled(); + expect(sendCommand).not.toHaveBeenCalled(); + }); + it("aborts the streaming run and re-prompts when a steer arrives", async () => { const sendCommand = vi.fn(async (_command: Record) => ({ success: true, diff --git a/products/desktop/packages/agent/src/server/pi-agent-server.ts b/products/desktop/packages/agent/src/server/pi-agent-server.ts index f34b90705502..ba3ee640ce63 100644 --- a/products/desktop/packages/agent/src/server/pi-agent-server.ts +++ b/products/desktop/packages/agent/src/server/pi-agent-server.ts @@ -33,7 +33,8 @@ import { resolveLlmGatewayUrl } from "../utils/gateway"; import { Logger } from "../utils/logger"; import { TaskRunEventStreamSender } from "./event-stream-sender"; import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; -import { jsonRpcRequestSchema } from "./schemas"; +import { resolveUserArtifactsById } from "./resolve-user-artifacts"; +import { jsonRpcRequestSchema, userAttachmentIdsSchema } from "./schemas"; import type { AgentServerConfig } from "./types"; interface SseController { @@ -57,11 +58,15 @@ const userMessageCommandSchema = z .object({ content: z.string().min(1).optional(), artifacts: z.array(z.record(z.string(), z.unknown())).optional(), + artifact_ids: userAttachmentIdsSchema.optional(), messageId: z.string().min(1).optional(), steer: z.boolean().optional(), }) .refine( - (params) => params.content || (params.artifacts?.length ?? 0) > 0, + (params) => + params.content || + (params.artifacts?.length ?? 0) > 0 || + (params.artifact_ids?.length ?? 0) > 0, "Either content or artifacts are required", ); @@ -705,12 +710,20 @@ export class PiAgentServer { runtime: PiRuntime, params: Record, ): Promise { - const artifacts = Array.isArray(params.artifacts) + const legacyArtifacts = Array.isArray(params.artifacts) ? (params.artifacts as TaskRunArtifact[]) : []; + const resolvedArtifacts = await resolveUserArtifactsById( + this.posthogAPI, + this.config.taskId, + this.config.runId, + Array.isArray(params.artifact_ids) + ? (params.artifact_ids as string[]) + : [], + ); const message = await this.prepareUserMessage( typeof params.content === "string" ? params.content : "", - artifacts, + [...legacyArtifacts, ...resolvedArtifacts], ); const result = await this.dispatchUserMessage( runtime, diff --git a/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts b/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts new file mode 100644 index 000000000000..927cd61c0ecd --- /dev/null +++ b/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts @@ -0,0 +1,65 @@ +import type { PostHogAPIClient } from "../posthog-api"; +import type { TaskRunArtifact } from "../types"; + +const MAX_USER_ATTACHMENT_BYTES = 4 * 1024 * 1024; + +function isScreenshotAttachment( + artifact: TaskRunArtifact, +): artifact is TaskRunArtifact & { + id: string; + content_type: "image/png" | "image/jpeg"; + size: number; + storage_path: string; +} { + return ( + artifact.type === "context" && + artifact.source === "user_attachment" && + artifact.uploaded_by === "user" && + typeof artifact.uploaded_by_user_id === "number" && + Number.isInteger(artifact.uploaded_by_user_id) && + artifact.uploaded_by_user_id > 0 && + (artifact.content_type === "image/png" || + artifact.content_type === "image/jpeg") && + typeof artifact.size === "number" && + Number.isInteger(artifact.size) && + artifact.size > 0 && + artifact.size <= MAX_USER_ATTACHMENT_BYTES && + typeof artifact.storage_path === "string" && + artifact.storage_path.length > 0 + ); +} + +export async function resolveUserArtifactsById( + posthogAPI: Pick, + taskId: string, + runId: string, + artifactIds: string[], +): Promise { + if (artifactIds.length === 0) { + return []; + } + + const taskRun = await posthogAPI.getTaskRun(taskId, runId); + const artifactsById = new Map( + (taskRun.artifacts ?? []).flatMap((artifact) => + artifact.id && isScreenshotAttachment(artifact) + ? [[artifact.id, artifact] as const] + : [], + ), + ); + const resolvedArtifacts: TaskRunArtifact[] = []; + const missingArtifactIds = new Set(); + for (const artifactId of artifactIds) { + const artifact = artifactsById.get(artifactId); + if (artifact) { + resolvedArtifacts.push(artifact); + } else { + missingArtifactIds.add(artifactId); + } + } + if (missingArtifactIds.size > 0) { + throw new Error("Screenshot attachments are unavailable for this task run"); + } + + return resolvedArtifacts; +} diff --git a/products/desktop/packages/agent/src/server/schemas.test.ts b/products/desktop/packages/agent/src/server/schemas.test.ts index a31eef90a1b5..eab06cdeee9d 100644 --- a/products/desktop/packages/agent/src/server/schemas.test.ts +++ b/products/desktop/packages/agent/src/server/schemas.test.ts @@ -146,11 +146,16 @@ describe("validateCommandParams", () => { expect(result.success).toBe(true); }); - it("accepts artifact-only user_message payloads", () => { - const result = validateCommandParams("user_message", { + it.each([ + { artifacts: [ { id: "artifact-1", storage_path: "tasks/artifacts/file.pdf" }, ], + }, + { artifact_ids: ["11111111-1111-4111-8111-111111111111"] }, + ])("accepts artifact-only user_message payloads", (attachments) => { + const result = validateCommandParams("user_message", { + ...attachments, }); expect(result.success).toBe(true); @@ -164,6 +169,31 @@ describe("validateCommandParams", () => { expect(result.success).toBe(false); }); + it.each([ + { artifactIds: ["not-a-uuid"] }, + { + artifactIds: [ + "11111111-1111-4111-8111-111111111111", + "11111111-1111-4111-8111-111111111111", + ], + }, + { + artifactIds: [ + "11111111-1111-4111-8111-111111111111", + "22222222-2222-4222-8222-222222222222", + "33333333-3333-4333-8333-333333333333", + "44444444-4444-4444-8444-444444444444", + "55555555-5555-4555-8555-555555555555", + ], + }, + ])("rejects invalid screenshot attachment IDs", ({ artifactIds }) => { + const result = validateCommandParams("user_message", { + artifact_ids: artifactIds, + }); + + expect(result.success).toBe(false); + }); + it("accepts valid permission_response", () => { const result = validateCommandParams("permission_response", { requestId: "abc-123", diff --git a/products/desktop/packages/agent/src/server/schemas.ts b/products/desktop/packages/agent/src/server/schemas.ts index d494d39a83b6..47c760f889cf 100644 --- a/products/desktop/packages/agent/src/server/schemas.ts +++ b/products/desktop/packages/agent/src/server/schemas.ts @@ -55,6 +55,11 @@ export const jsonRpcRequestSchema = z.object({ export type JsonRpcRequest = z.infer; +export const userAttachmentIdsSchema = z + .array(z.uuid()) + .max(4) + .refine((artifactIds) => new Set(artifactIds).size === artifactIds.length); + export const userMessageParamsSchema = z .object({ content: z @@ -66,6 +71,7 @@ export const userMessageParamsSchema = z ]) .optional(), artifacts: z.array(z.record(z.string(), z.unknown())).optional(), + artifact_ids: userAttachmentIdsSchema.optional(), messageId: z.string().min(1).optional(), steer: z.boolean().optional(), }) @@ -77,8 +83,10 @@ export const userMessageParamsSchema = z : Array.isArray(params.content) && params.content.length > 0; const hasArtifacts = Array.isArray(params.artifacts) && params.artifacts.length > 0; + const hasArtifactIds = + Array.isArray(params.artifact_ids) && params.artifact_ids.length > 0; - return hasContent || hasArtifacts; + return hasContent || hasArtifacts || hasArtifactIds; }, { error: "Either content or artifacts are required" }, ); From 33279344701847c6b42ee4c88d06ff909ff0636c Mon Sep 17 00:00:00 2001 From: lordspline <74811063+lordspline@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:22:04 +0000 Subject: [PATCH 2/2] feat(aio): add image attachments to sandbox conversations --- .../engineering/ai/sandboxed-agents.md | 8 + ee/api/conversation.py | 55 +- ee/api/tests/test_conversation.py | 1 + .../scenes/max/components/HandsFreeButton.tsx | 6 +- .../max/components/QuestionInput.test.tsx | 128 +++ .../scenes/max/components/QuestionInput.tsx | 137 ++- .../src/scenes/max/maxThreadLogic.test.ts | 194 +++- frontend/src/scenes/max/maxThreadLogic.tsx | 453 +++++---- posthog/settings/web.py | 1 + .../frontend/generated/api.schemas.ts | 12 +- .../frontend/generated/api.zod.ts | 12 +- .../packages/agent/src/posthog-api.test.ts | 53 ++ .../desktop/packages/agent/src/posthog-api.ts | 41 +- .../agent/src/server/agent-server.test.ts | 167 ++-- .../packages/agent/src/server/agent-server.ts | 59 +- .../agent/src/server/pi-agent-server.test.ts | 40 +- .../agent/src/server/pi-agent-server.ts | 41 +- .../src/server/resolve-user-artifacts.test.ts | 94 ++ .../src/server/resolve-user-artifacts.ts | 70 +- .../packages/shared/src/domain-types.ts | 12 +- products/posthog_ai/backend/api/__init__.py | 3 +- .../posthog_ai/backend/api/attachments.py | 228 +++++ products/posthog_ai/backend/attachments.py | 893 ++++++++++++++++++ .../posthog_ai/backend/message_routing.py | 253 +++-- products/posthog_ai/backend/routes.py | 8 +- .../backend/tests/test_attachment_api.py | 189 ++++ .../backend/tests/test_attachments.py | 615 ++++++++++++ .../backend/tests/test_message_routing.py | 152 +++ products/posthog_ai/frontend/api/logics.ts | 10 + .../posthog_ai/frontend/api/primitives.ts | 4 + .../composer/ImageAttachmentButton.tsx | 45 + .../composer/ImageAttachmentPreviewList.tsx | 83 ++ .../frontend/generated/api.schemas.ts | 102 ++ products/posthog_ai/frontend/generated/api.ts | 74 +- .../posthog_ai/frontend/generated/api.zod.ts | 55 ++ .../logics/assistantAttachmentsLogic.test.ts | 253 +++++ .../logics/assistantAttachmentsLogic.ts | 504 ++++++++++ products/tasks/backend/facade/api.py | 265 +++++- products/tasks/backend/facade/contracts.py | 8 + .../tasks/backend/presentation/serializers.py | 47 +- .../tasks/backend/presentation/views/api.py | 29 +- products/tasks/backend/tests/test_api.py | 25 + .../test_posthog_ai_attachment_promotion.py | 294 ++++++ .../tasks/frontend/generated/api.schemas.ts | 19 +- products/tasks/frontend/generated/api.zod.ts | 371 +++++--- services/mcp/src/api/generated.ts | 133 ++- 46 files changed, 5629 insertions(+), 617 deletions(-) create mode 100644 products/desktop/packages/agent/src/server/resolve-user-artifacts.test.ts create mode 100644 products/posthog_ai/backend/api/attachments.py create mode 100644 products/posthog_ai/backend/attachments.py create mode 100644 products/posthog_ai/backend/tests/test_attachment_api.py create mode 100644 products/posthog_ai/backend/tests/test_attachments.py create mode 100644 products/posthog_ai/frontend/components/composer/ImageAttachmentButton.tsx create mode 100644 products/posthog_ai/frontend/components/composer/ImageAttachmentPreviewList.tsx create mode 100644 products/posthog_ai/frontend/logics/assistantAttachmentsLogic.test.ts create mode 100644 products/posthog_ai/frontend/logics/assistantAttachmentsLogic.ts create mode 100644 products/tasks/backend/tests/test_posthog_ai_attachment_promotion.py diff --git a/docs/published/handbook/engineering/ai/sandboxed-agents.md b/docs/published/handbook/engineering/ai/sandboxed-agents.md index 9d7c21ed000e..b88184be4472 100644 --- a/docs/published/handbook/engineering/ai/sandboxed-agents.md +++ b/docs/published/handbook/engineering/ai/sandboxed-agents.md @@ -51,6 +51,14 @@ The agent inside the sandbox gets: - Access to the **PostHog MCP server** for querying data - **Code execution** capabilities within the sandbox +### PostHog AI screenshot context + +PostHog AI sandbox conversations can include up to four PNG or JPEG screenshots with a nonblank message. Each image can be up to 4 MiB, and the images in one message can total up to 10 MiB. + +The browser uploads each image directly to object storage through a 15-minute signed form. The API validates and normalizes the image before it promotes a separate copy into the task run's artifact manifest. Only opaque attachment IDs pass through Temporal and the agent command protocol, so image bytes and storage credentials never enter workflow payloads. + +Staged uploads are scoped to one team, user, and conversation. The agent server resolves promoted IDs only from the current task run and accepts normalized PNG or JPEG context artifacts uploaded by a user. + ## Creating a sandboxed agent Use `Task.create_and_run()` to launch a sandboxed agent from your product code: diff --git a/ee/api/conversation.py b/ee/api/conversation.py index 770a8c495ee9..b9d35f909028 100644 --- a/ee/api/conversation.py +++ b/ee/api/conversation.py @@ -2,7 +2,7 @@ import uuid import asyncio from collections.abc import AsyncGenerator, Iterable -from typing import cast +from typing import Any, cast from django.conf import settings from django.core.exceptions import ValidationError @@ -228,19 +228,26 @@ def _validate_sandbox_task(task_id: uuid.UUID, team_id: int, user_id: int | None class SandboxOpenSerializer(serializers.Serializer): - """Request body for `POST /conversations/{id}/open/`. A string `content` processes a turn; a - null/absent `content` warms a sandbox that idles awaiting the first message.""" + """Request body for `POST /conversations/{id}/open/`. Nonblank `content` processes a turn and may include + `attachment_ids`; null or absent `content` with no attachments warms a sandbox awaiting its first message.""" content = serializers.CharField( required=False, allow_null=True, allow_blank=True, max_length=40000, - help_text="The user's message text. Omit or null to warm a sandbox (boot + idle) ahead of the first message.", + help_text="The user's message text. Omit or null to warm a sandbox ahead of the first message, unless attachment_ids are provided.", ) trace_id = serializers.UUIDField( required=False, help_text="Client-generated trace id correlated with the resulting Run's SSE stream." ) + attachment_ids = serializers.ListField( + child=serializers.UUIDField(), + required=False, + min_length=1, + max_length=4, + help_text="Finalized sandbox image attachment IDs to send with this message.", + ) # Deprecated with the legacy Max bridge (see SandboxAttachedContextItemSerializer) — do not extend. attached_context = serializers.ListField( required=False, @@ -264,6 +271,15 @@ class SandboxOpenSerializer(serializers.Serializer): ), ) + def validate(self, attrs: dict[str, Any]) -> dict[str, Any]: + attachment_ids = attrs.get("attachment_ids") or [] + if len(set(attachment_ids)) != len(attachment_ids): + raise serializers.ValidationError({"attachment_ids": "Attachment IDs must be unique."}) + content = attrs.get("content") or "" + if attachment_ids and not content.strip(): + raise serializers.ValidationError({"attachment_ids": "Attachment IDs require a message send."}) + return attrs + def validate_task_id(self, value: uuid.UUID) -> uuid.UUID: """Resolve the Task to bind, scoped to the team and the requesting user's visibility. @@ -729,20 +745,27 @@ def open(self, request: Request, *args, **kwargs): if conversation.task_id is not None: _validate_sandbox_task(conversation.task_id, self.team.id, request.user.id) - has_content = bool(serializer.validated_data.get("content")) - convert_to_acp, resumed_context = self._compute_sandbox_conversion(request, conversation, has_content) + content = serializer.validated_data.get("content") or "" + attachment_ids = serializer.validated_data.get("attachment_ids") or [] + has_message = bool(content.strip()) or bool(attachment_ids) + convert_to_acp, resumed_context = self._compute_sandbox_conversion(request, conversation, has_message) # Sandbox-only endpoint. A converting LangGraph thread is still LANGGRAPH here (the flip happens # inside the routing service), so allow it through; reject any other non-sandbox conversation. if conversation.agent_runtime != Conversation.AgentRuntime.SANDBOX and not convert_to_acp: raise exceptions.ValidationError("This conversation is not on the sandbox runtime.") - if has_content and conversation.title is None: - conversation.title = serializer.validated_data["content"][:80] + if has_message and conversation.title is None: + conversation.title = content[:80] if content.strip() else "Image attachment" conversation.save(update_fields=["title"]) return self._route_sandbox_message( - request, conversation, resumed_context=resumed_context, convert_to_acp=convert_to_acp, created=created + request, + conversation, + payload=serializer.validated_data, + resumed_context=resumed_context, + convert_to_acp=convert_to_acp, + created=created, ) def _get_or_create_sandbox_conversation( @@ -815,7 +838,7 @@ def _compute_sandbox_conversion( resumed_context = None return True, resumed_context - def _auto_route_repository(self, request: Request, conversation: Conversation, user: User) -> str | None: + def _auto_route_repository(self, payload: dict[str, Any], conversation: Conversation, user: User) -> str | None: """Auto-select the repository a sandbox conversation's first message is about. Runs only on a first message — no backing Task yet (`task_id is None`) and real content. @@ -825,7 +848,7 @@ def _auto_route_repository(self, request: Request, conversation: Conversation, u """ if conversation.task_id is not None: return None - content = request.data.get("content") + content = payload.get("content") if not isinstance(content, str) or not content.strip(): return None return asgi_async_to_sync(tasks_facade.select_repository_for_message)( @@ -837,14 +860,15 @@ def _route_sandbox_message( request: Request, conversation: Conversation, *, + payload: dict[str, Any], resumed_context: str | None = None, convert_to_acp: bool = False, created: bool = False, ) -> Response: user = cast(User, request.user) - repository = self._auto_route_repository(request, conversation, user) + repository = self._auto_route_repository(payload, conversation, user) result = SandboxSession(conversation, user).open( - request.data, resumed_context=resumed_context, convert_to_acp=convert_to_acp, repository=repository + payload, resumed_context=resumed_context, convert_to_acp=convert_to_acp, repository=repository ) if result is None: # Warm intent that provisioned nothing (pool full / released) — no run to open. Drop the @@ -852,8 +876,9 @@ def _route_sandbox_message( if created: conversation.delete() return Response(status=status.HTTP_204_NO_CONTENT) - content = request.data.get("content") - if isinstance(content, str) and content.strip(): + content = payload.get("content") + attachment_ids = payload.get("attachment_ids") or [] + if (isinstance(content, str) and content.strip()) or attachment_ids: report_user_action( user, "prompt sent", diff --git a/ee/api/tests/test_conversation.py b/ee/api/tests/test_conversation.py index b1d884e27fed..f9894a887087 100644 --- a/ee/api/tests/test_conversation.py +++ b/ee/api/tests/test_conversation.py @@ -1767,6 +1767,7 @@ def test_open_validates_request_body(self): {"content": "x" * 40001}, # over the content length cap {"content": "hello", "trace_id": "not-a-uuid"}, # malformed trace id {"content": "hello", "initial_permission_mode": "full-access"}, # Codex-only mode, not valid for Claude + {"content": " ", "attachment_ids": [str(uuid.uuid4())]}, ] for payload in bad_payloads: with patch("ee.api.conversation.SandboxSession") as m_session: diff --git a/frontend/src/scenes/max/components/HandsFreeButton.tsx b/frontend/src/scenes/max/components/HandsFreeButton.tsx index 822f7985a519..a17b616de7c6 100644 --- a/frontend/src/scenes/max/components/HandsFreeButton.tsx +++ b/frontend/src/scenes/max/components/HandsFreeButton.tsx @@ -11,9 +11,10 @@ import { handsFreeLogic } from '../handsFreeLogic' interface HandsFreeButtonProps { panelId?: string + disabledReason?: string } -export function HandsFreeButton({ panelId }: HandsFreeButtonProps): JSX.Element | null { +export function HandsFreeButton({ panelId, disabledReason }: HandsFreeButtonProps): JSX.Element | null { const flagEnabled = useFeatureFlag('MAX_HANDS_FREE') const { status, canUseHandsFree } = useValues(handsFreeLogic({ panelId })) const { toggleHandsFree } = useActions(handsFreeLogic({ panelId })) @@ -30,8 +31,9 @@ export function HandsFreeButton({ panelId }: HandsFreeButtonProps): JSX.Element type="tertiary" icon={} onClick={toggleHandsFree} - tooltip="Enter hands-free" + tooltip={disabledReason || 'Enter hands-free'} aria-label="Enter hands-free" + disabledReason={disabledReason} /> ) diff --git a/frontend/src/scenes/max/components/QuestionInput.test.tsx b/frontend/src/scenes/max/components/QuestionInput.test.tsx index 5503e6f6d63a..3c91de2a3264 100644 --- a/frontend/src/scenes/max/components/QuestionInput.test.tsx +++ b/frontend/src/scenes/max/components/QuestionInput.test.tsx @@ -3,9 +3,20 @@ import '@testing-library/jest-dom' import { cleanup, fireEvent, render, screen, waitFor } from '@testing-library/react' import { BindLogic, Provider } from 'kea' +import { FEATURE_FLAGS } from 'lib/constants' +import { featureFlagLogic } from 'lib/logic/featureFlagLogic' +import { projectLogic } from 'scenes/projectLogic' + import { useMocks } from '~/mocks/jest' import { initKeaTests } from '~/test/init' +import { + assistantAttachmentsDeleteCreate, + assistantAttachmentsFinalizeCreate, + assistantAttachmentsPrepareCreate, +} from 'products/posthog_ai/frontend/generated/api' + +import { handsFreeLogic } from '../handsFreeLogic' import { maxGlobalLogic } from '../maxGlobalLogic' import { maxLogic } from '../maxLogic' import { maxThreadLogic } from '../maxThreadLogic' @@ -21,13 +32,53 @@ jest.mock( { virtual: true } ) +jest.mock('products/posthog_ai/frontend/generated/api', () => ({ + assistantAttachmentsPrepareCreate: jest.fn(), + assistantAttachmentsFinalizeCreate: jest.fn(), + assistantAttachmentsDeleteCreate: jest.fn(), +})) + describe('QuestionInput', () => { let maxLogicInstance: ReturnType + let projectLogicInstance: ReturnType let threadLogicInstance: ReturnType beforeEach(() => { useMocks(maxMocks) initKeaTests() + Object.defineProperty(URL, 'createObjectURL', { + configurable: true, + value: jest.fn((file: File) => `blob:${file.name}`), + }) + Object.defineProperty(URL, 'revokeObjectURL', { configurable: true, value: jest.fn() }) + const baseFetch = global.fetch.bind(global) + global.fetch = jest.fn((input: RequestInfo | URL, init?: RequestInit) => { + if (typeof input === 'string' && input.startsWith('https://upload.test')) { + return Promise.resolve({ ok: true } as Response) + } + return baseFetch(input, init) + }) as any + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ + attachments: [ + { + id: 'attachment-1', + file_name: 'upload.png', + content_type: 'image/png', + size: 1, + upload_url: 'https://upload.test', + upload_fields: { key: 'value' }, + }, + ], + }) + ;(assistantAttachmentsFinalizeCreate as jest.Mock).mockResolvedValue({ + id: 'attachment-1', + file_name: 'upload.png', + content_type: 'image/png', + size: 1, + width: 1, + height: 1, + }) + ;(assistantAttachmentsDeleteCreate as jest.Mock).mockResolvedValue(undefined) const maxGlobalLogicInstance = maxGlobalLogic() maxGlobalLogicInstance.mount() @@ -36,6 +87,10 @@ describe('QuestionInput', () => { maxLogicInstance = maxLogic({ panelId: 'test' }) maxLogicInstance.mount() + projectLogicInstance = projectLogic() + projectLogicInstance.mount() + projectLogicInstance.actions.loadCurrentProjectSuccess({ id: 1, name: 'Test project' } as any) + const threadProps = { panelId: 'test', conversationId: maxLogicInstance.values.frontendConversationId } threadLogicInstance = maxThreadLogic(threadProps) threadLogicInstance.mount() @@ -56,6 +111,7 @@ describe('QuestionInput', () => { threadLogicInstance?.unmount() maxLogicInstance?.cache.eventSourceController?.abort() maxLogicInstance?.unmount() + projectLogicInstance?.unmount() jest.restoreAllMocks() }) @@ -109,6 +165,78 @@ describe('QuestionInput', () => { await waitFor(() => expect(slashCommandItem()).toBeInTheDocument()) }) + it('keeps send disabled for a blank message even after an attachment is ready', async () => { + threadLogicInstance.actions.setIsSandboxMode(true) + await waitFor(() => expect(screen.getByLabelText('Choose PNG or JPEG images')).toBeInTheDocument()) + const fileInput = screen.getByLabelText('Choose PNG or JPEG images') as HTMLInputElement + const sendButton = document.querySelector('[data-attr="max-send-message"]') as HTMLElement + + fireEvent.change(fileInput, { + target: { files: [new File(['a'], 'upload.png', { type: 'image/png' })] }, + }) + + await waitFor(() => expect(screen.getByText('upload.png')).toBeInTheDocument()) + await waitFor(() => expect(screen.getByText('Ready')).toBeInTheDocument()) + expect(sendButton).toHaveAttribute('aria-disabled', 'true') + }) + + it('disables hands-free mode while an image is staged', async () => { + featureFlagLogic.actions.setFeatureFlags([FEATURE_FLAGS.MAX_HANDS_FREE], { + [FEATURE_FLAGS.MAX_HANDS_FREE]: true, + }) + handsFreeLogic({ panelId: 'test' }).actions.setSdkAvailable(true) + threadLogicInstance.actions.setIsSandboxMode(true) + await waitFor(() => expect(screen.getByLabelText('Choose PNG or JPEG images')).toBeInTheDocument()) + + fireEvent.change(screen.getByLabelText('Choose PNG or JPEG images'), { + target: { files: [new File(['a'], 'upload.png', { type: 'image/png' })] }, + }) + + await waitFor(() => expect(screen.getByText('Ready')).toBeInTheDocument()) + expect(screen.getByLabelText('Enter hands-free')).toHaveAttribute('aria-disabled', 'true') + }) + + it.each([ + [ + 'drop', + (input: HTMLElement, file: File) => + fireEvent.drop(input.closest('label') as HTMLElement, { + dataTransfer: { files: [file], types: ['Files'] }, + }), + ], + [ + 'paste', + (input: HTMLElement, file: File) => + fireEvent.paste(input, { + clipboardData: { + items: [ + { + kind: 'file', + getAsFile: () => file, + }, + ], + }, + }), + ], + ])('adds sandbox attachments via %s', async (method, addFile) => { + threadLogicInstance.actions.setIsSandboxMode(true) + await waitFor(() => expect(screen.getByLabelText('Choose PNG or JPEG images')).toBeInTheDocument()) + const input = screen.getByRole('textbox') + const composer = input.closest('label') as HTMLElement + const file = new File(['a'], `from-${method}.png`, { type: 'image/png' }) + + if (method === 'drop') { + fireEvent.dragEnter(composer, { dataTransfer: { files: [file], types: ['Files'] } }) + await waitFor(() => expect(composer.className).toContain('bg-accent-highlight-secondary/20')) + } + + addFile(input, file) + + await waitFor(() => expect(screen.getByText(`from-${method}.png`)).toBeInTheDocument()) + await waitFor(() => expect(screen.getByText('Ready')).toBeInTheDocument()) + await waitFor(() => expect(composer.className).not.toContain('bg-accent-highlight-secondary/20')) + }) + describe('stop button cancel state', () => { const sendButton = (): HTMLElement | null => document.querySelector('[data-attr="max-send-message"]') const stopButton = (): HTMLElement | null => document.querySelector('[data-attr="max-stop-generation"]') diff --git a/frontend/src/scenes/max/components/QuestionInput.tsx b/frontend/src/scenes/max/components/QuestionInput.tsx index 97c41d3cf2e9..bda118b741db 100644 --- a/frontend/src/scenes/max/components/QuestionInput.tsx +++ b/frontend/src/scenes/max/components/QuestionInput.tsx @@ -18,6 +18,9 @@ import { userLogic } from 'scenes/userLogic' import { AgentMode } from '~/queries/schema/schema-assistant-messages' import { ConversationQueueMessage } from '~/types' +import { assistantAttachmentsLogic } from 'products/posthog_ai/frontend/api/logics' +import { ImageAttachmentButton, ImageAttachmentPreviewList } from 'products/posthog_ai/frontend/api/primitives' + import { ContextDisplay } from '../Context' import { handsFreeLogic } from '../handsFreeLogic' import { maxGlobalLogic } from '../maxGlobalLogic' @@ -166,6 +169,8 @@ export const QuestionInput = React.forwardRef 0 + const attachmentsEnabled = isSandboxMode && !isSharedThread && !handsFreeActive + const submit = (prompt: string): void => { + if (attachmentsAreSending) { + return + } // askMax reads the prompt arg directly and clears `question` afterwards, so drop any // pending debounce to stop it from re-populating the just-sent text. debouncedSetQuestion.cancel() @@ -238,19 +253,47 @@ export const QuestionInput = React.forwardRef 0 + const handleSelectedFiles = (files: File[]): void => { + if (!attachmentsEnabled || attachmentsAreSending || files.length === 0) { + return + } + addFiles(files) + } + + const handlePastedFiles = (event: React.ClipboardEvent): void => { + if (!attachmentsEnabled) { + return + } + const files = Array.from(event.clipboardData.items) + .filter((item) => item.kind === 'file') + .flatMap((item) => { + const file = item.getAsFile() + return file ? [file] : [] + }) + if (files.length === 0) { + return + } + event.preventDefault() + handleSelectedFiles(files) + } + + const attachmentsBlockedByQueue = queueingEnabled && threadLoading && attachments.length > 0 // A fill-in suggestion typed its prefix in and is waiting for the user to complete it. const showFillInHint = !!fillInHint - const isQueueingSubmission = queueingEnabled && threadLoading && hasQuestion + const isQueueingSubmission = queueingEnabled && threadLoading && hasQuestion && !attachmentsBlockedByQueue const showStopButton = threadLoading && !isQueueingSubmission && !cancelLoading // Mirrors maxThreadLogic's `submissionDisabledReason` selector, but using the local input // value so the submit guard stays correct while the debounced sync to kea is still pending. const submissionDisabledReason = contextDisabledReason ? contextDisabledReason - : !inputValue + : !hasQuestion ? 'I need some input first' - : queueDisabledReason + : attachmentSubmissionDisabledReason + ? attachmentSubmissionDisabledReason + : attachmentsBlockedByQueue + ? 'Wait for the current reply before sending images' + : queueDisabledReason // Update autocomplete visibility when the input changes useEffect(() => { @@ -268,7 +311,9 @@ export const QuestionInput = React.forwardRef { + if (!attachmentsEnabled || !Array.from(event.dataTransfer?.types || []).includes('Files')) { + return + } + event.preventDefault() + setDragActive(true) + }} + onDragLeave={(event) => { + if (!attachmentsEnabled) { + return + } + event.preventDefault() + setDragActive(false) + }} + onDragOver={(event) => { + if (!attachmentsEnabled || !Array.from(event.dataTransfer?.types || []).includes('Files')) { + return + } + event.preventDefault() + }} + onDrop={(event) => { + if (!attachmentsEnabled) { + return + } + event.preventDefault() + setDragActive(false) + handleSelectedFiles(Array.from(event.dataTransfer.files || [])) + }} > {handsFreeActive ? ( @@ -353,7 +427,7 @@ export const QuestionInput = React.forwardRef -
+
{!inputValue && (
)} - {!isSharedThread && - !handsFreeActive && ( - // When the hands-free flag is on, reserve ~80px (pr-20) so the chip - // row doesn't wrap under the absolutely-positioned mic + send pair. - // Without the flag the row only has send and the legacy pr-12 is - // enough — keep it so non-flagged users see the original layout. + {!isSharedThread && !handsFreeActive && ( + <> + + {selectionError && ( +
+ {selectionError} +
+ )} + {/* When the hands-free flag is on, reserve ~80px (pr-20) so the chip + row doesn't wrap under the absolutely-positioned mic + send pair. + Without the flag the row only has send and the legacy pr-12 is + enough — keep it so non-flagged users see the original layout. */}
{!isThreadVisible ? (
+ {attachmentsEnabled && ( + + )} {topActions}
@@ -489,7 +582,8 @@ export const QuestionInput = React.forwardRef )}
- )} + + )}
- + {attachmentsEnabled && isThreadVisible && ( + + )} + 0 ? 'Remove images before entering hands-free' : undefined + } + /> {!handsFreeActive && ( ) } - loading={threadLoading && !dataProcessingAccepted} + loading={attachmentsAreSending || (threadLoading && !dataProcessingAccepted)} disabledReason={disabledReason} className={disabledReason ? 'opacity-[0.5]' : ''} size="small" diff --git a/frontend/src/scenes/max/maxThreadLogic.test.ts b/frontend/src/scenes/max/maxThreadLogic.test.ts index faa83e5bcd1c..10e0fe022b3f 100644 --- a/frontend/src/scenes/max/maxThreadLogic.test.ts +++ b/frontend/src/scenes/max/maxThreadLogic.test.ts @@ -12,6 +12,7 @@ import { lemonToast } from 'lib/lemon-ui/LemonToast' import { featureFlagLogic } from 'lib/logic/featureFlagLogic' import { notebookLogic } from 'scenes/notebooks/Notebook/notebookLogic' import { NotebookTarget } from 'scenes/notebooks/types' +import { projectLogic } from 'scenes/projectLogic' import { sceneLogic } from 'scenes/sceneLogic' import { Scene } from 'scenes/sceneTypes' import { urls } from 'scenes/urls' @@ -30,7 +31,9 @@ import { import { initKeaTests } from '~/test/init' import { Conversation, ConversationDetail, ConversationStatus, ConversationType } from '~/types' +import * as conversationsFrontendApi from 'products/conversations/frontend/generated/api' import { attachedContextLogic, runStreamLogic } from 'products/posthog_ai/frontend/api/logics' +import { assistantAttachmentsLogic } from 'products/posthog_ai/frontend/logics/assistantAttachmentsLogic' import { RuntimeEnumApi } from 'products/tasks/frontend/generated/api.schemas' import { EnhancedToolCall, TOOL_DEFINITIONS } from './max-constants' @@ -60,6 +63,7 @@ jest.mock( describe('maxThreadLogic', () => { let logic: ReturnType let maxLogicInstance: ReturnType + let projectLogicInstance: ReturnType beforeEach(() => { useMocks({ @@ -87,6 +91,10 @@ describe('maxThreadLogic', () => { maxLogicInstance.mount() maxLogicInstance.actions.setConversationId(MOCK_CONVERSATION_ID) + projectLogicInstance = projectLogic() + projectLogicInstance.mount() + projectLogicInstance.actions.loadCurrentProjectSuccess({ id: 997, name: 'Test project' } as any) + logic = maxThreadLogic({ conversationId: MOCK_CONVERSATION_ID, panelId: 'test' }) logic.mount() }) @@ -105,6 +113,7 @@ describe('maxThreadLogic', () => { maxLogicInstance.cache.eventSourceController?.abort() maxLogicInstance.unmount() } + projectLogicInstance?.unmount() // Clean up any remaining mocks jest.restoreAllMocks() @@ -1816,12 +1825,12 @@ describe('maxThreadLogic', () => { // Warm = open with content: null. Keep the POST in flight so the abandon races it. let resolvePrewarm: (handle: typeof warmHandle) => void = () => {} const openSpy = jest - .spyOn(api.conversations, 'open') + .spyOn(conversationsFrontendApi, 'conversationsOpenCreate') .mockReturnValue(new Promise((resolve) => (resolvePrewarm = resolve)) as any) logic.actions.prewarmSandbox() await flush() - expect(openSpy).toHaveBeenCalledWith(MOCK_CONVERSATION_ID, { + expect(openSpy).toHaveBeenCalledWith('997', MOCK_CONVERSATION_ID, { content: null, initial_permission_mode: 'auto', }) @@ -1843,7 +1852,7 @@ describe('maxThreadLogic', () => { it('does not release a warm that resolved with no pending abandon', async () => { await mountIdleSandbox() - jest.spyOn(api.conversations, 'open').mockResolvedValue(warmHandle) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockResolvedValue(warmHandle) await expectLogic(logic, () => { logic.actions.prewarmSandbox() @@ -3574,7 +3583,7 @@ describe('maxThreadLogic', () => { const taskSpy = jest .spyOn(api.tasks, 'get') .mockResolvedValue({ id: 'pi-task', runtime: RuntimeEnumApi.Pi } as any) - const openSpy = jest.spyOn(api.conversations, 'open') + const openSpy = jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate') logic.actions.askMax('hello') await Promise.resolve() @@ -3591,7 +3600,7 @@ describe('maxThreadLogic', () => { const taskSpy = jest .spyOn(api.tasks, 'get') .mockResolvedValue({ id: 'acp-task', runtime: RuntimeEnumApi.Acp } as any) - jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockResolvedValue(sandboxRunResponse) await expectLogic(logic, () => { logic.actions.askMax('hello') @@ -3602,7 +3611,9 @@ describe('maxThreadLogic', () => { }) it('holds the streaming lock until the sandbox turn completes and releases exactly once', async () => { - const openSpy = jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + const openSpy = jest + .spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + .mockResolvedValue(sandboxRunResponse) await expectLogic(logic, () => { logic.actions.streamConversation( @@ -3612,12 +3623,14 @@ describe('maxThreadLogic', () => { }).toDispatchActions(['openSandboxSse']) expect(openSpy).toHaveBeenCalledWith( + '997', MOCK_CONVERSATION_ID, expect.objectContaining({ content: 'hello', initial_permission_mode: 'auto', }) ) + expect(openSpy.mock.calls[0]?.[2]).not.toHaveProperty('attachment_ids') // The POST finished, but the turn is still streaming — the lock must still be held expect(maxLogicInstance.values.activeStreamingThreads).toEqual(1) @@ -3637,7 +3650,7 @@ describe('maxThreadLogic', () => { }) it('does not release the lock on a non-terminal task_run_state frame', async () => { - jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockResolvedValue(sandboxRunResponse) await expectLogic(logic, () => { logic.actions.streamConversation( @@ -3660,7 +3673,7 @@ describe('maxThreadLogic', () => { }) it('lights the optimistic boot indicator before the open POST and clears it once the SSE opens', async () => { - jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockResolvedValue(sandboxRunResponse) await expectLogic(logic, () => { logic.actions.streamConversation( @@ -3674,7 +3687,7 @@ describe('maxThreadLogic', () => { }) it('releases the lock immediately and surfaces an error when the send POST fails', async () => { - jest.spyOn(api.conversations, 'open').mockRejectedValue(new Error('boom')) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockRejectedValue(new Error('boom')) await expectLogic(logic, () => { logic.actions.streamConversation( @@ -3694,8 +3707,28 @@ describe('maxThreadLogic', () => { ).toEqual(true) }) + it('fails a sandbox send if the current project has not loaded yet', async () => { + projectLogicInstance.actions.loadCurrentProjectSuccess(null as any) + const openSpy = jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + + await expectLogic(logic, () => { + logic.actions.streamConversation( + { agent_mode: null, is_sandbox: true, content: 'hello', conversation: MOCK_CONVERSATION_ID }, + 0 + ) + }).toDispatchActions(['pushSandboxError', 'decrActiveStreamingThreads']) + + expect(openSpy).not.toHaveBeenCalled() + expect(maxLogicInstance.values.activeStreamingThreads).toEqual(0) + expect( + runStreamLogic({ streamKey: MOCK_CONVERSATION_ID }).values.threadItems.some( + (item) => item.type === 'human_message' && item.text === 'hello' + ) + ).toEqual(false) + }) + it('releases the lock immediately when no run was started', async () => { - const openSpy = jest.spyOn(api.conversations, 'open') + const openSpy = jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate') await expectLogic(logic, () => { logic.actions.streamConversation( @@ -3708,8 +3741,144 @@ describe('maxThreadLogic', () => { expect(maxLogicInstance.values.activeStreamingThreads).toEqual(0) }) + it('sends ready attachment ids only on sandbox turns', async () => { + const attachmentLogicInstance = assistantAttachmentsLogic({ conversationId: MOCK_CONVERSATION_ID }) + attachmentLogicInstance.mount() + attachmentLogicInstance.actions.addAttachment({ + localId: 'local-1', + file: new File(['a'], 'a.png', { type: 'image/png', lastModified: 1 }), + fileKey: 'a.png:1:1:image/png', + previewUrl: 'blob:a.png', + status: 'ready', + }) + attachmentLogicInstance.actions.setAttachmentReady('local-1', 'attachment-1') + logic.actions.setIsSandboxMode(true) + + const openSpy = jest + .spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + .mockResolvedValue(sandboxRunResponse) + + await expectLogic(logic, () => { + logic.actions.askMax('hello') + }).toDispatchActions(['openSandboxSse']) + + expect(openSpy).toHaveBeenCalledWith( + '997', + MOCK_CONVERSATION_ID, + expect.objectContaining({ + content: 'hello', + attachment_ids: ['attachment-1'], + }) + ) + + attachmentLogicInstance.unmount() + }) + + it('keeps the submitted attachment snapshot stable while routing the sandbox turn', async () => { + const attachmentLogicInstance = assistantAttachmentsLogic({ conversationId: MOCK_CONVERSATION_ID }) + attachmentLogicInstance.mount() + attachmentLogicInstance.actions.addAttachment({ + localId: 'local-first', + file: new File(['a'], 'first.png', { type: 'image/png', lastModified: 1 }), + fileKey: 'first.png:1:1:image/png', + previewUrl: 'blob:first.png', + status: 'ready', + }) + attachmentLogicInstance.actions.setAttachmentReady('local-first', 'attachment-first') + logic.actions.setIsSandboxMode(true) + maxLogicInstance.actions.setPendingBindTaskId('acp-task') + + let resolveTask!: (task: { id: string; runtime: RuntimeEnumApi }) => void + jest.spyOn(api.tasks, 'get').mockImplementation( + () => new Promise((resolve) => (resolveTask = resolve as typeof resolveTask)) + ) + const openSpy = jest + .spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + .mockResolvedValue(sandboxRunResponse) + + logic.actions.askMax('first message') + expect(attachmentLogicInstance.values.attachmentsAreSending).toBe(true) + + attachmentLogicInstance.actions.addAttachment({ + localId: 'local-second', + file: new File(['b'], 'second.png', { type: 'image/png', lastModified: 2 }), + fileKey: 'second.png:1:2:image/png', + previewUrl: 'blob:second.png', + status: 'ready', + }) + attachmentLogicInstance.actions.setAttachmentReady('local-second', 'attachment-second') + logic.actions.askMax('duplicate message') + + await expectLogic(logic, () => { + resolveTask({ id: 'acp-task', runtime: RuntimeEnumApi.Acp }) + }).toDispatchActions(['openSandboxSse']) + + expect(openSpy).toHaveBeenCalledTimes(1) + expect(openSpy).toHaveBeenCalledWith( + '997', + MOCK_CONVERSATION_ID, + expect.objectContaining({ + content: 'first message', + attachment_ids: ['attachment-first'], + }) + ) + expect(attachmentLogicInstance.values.readyAttachmentIds).toEqual(['attachment-second']) + + attachmentLogicInstance.unmount() + }) + + it('keeps ready attachments after the sandbox open request fails', async () => { + const attachmentLogicInstance = assistantAttachmentsLogic({ conversationId: MOCK_CONVERSATION_ID }) + attachmentLogicInstance.mount() + attachmentLogicInstance.actions.addAttachment({ + localId: 'local-failed', + file: new File(['a'], 'a.png', { type: 'image/png', lastModified: 1 }), + fileKey: 'a.png:1:1:image/png', + previewUrl: 'blob:a.png', + status: 'ready', + }) + attachmentLogicInstance.actions.setAttachmentReady('local-failed', 'attachment-failed') + logic.actions.setIsSandboxMode(true) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockRejectedValue(new Error('open failed')) + + await expectLogic(logic, () => { + logic.actions.askMax('hello') + }).toDispatchActions(['markSandboxAttachmentsSendFailed']) + + expect(attachmentLogicInstance.values.attachmentsAreSending).toBe(false) + expect(attachmentLogicInstance.values.readyAttachmentIds).toEqual(['attachment-failed']) + + attachmentLogicInstance.unmount() + }) + + it('does not open an attachment-only sandbox turn', async () => { + const attachmentLogicInstance = assistantAttachmentsLogic({ conversationId: MOCK_CONVERSATION_ID }) + attachmentLogicInstance.mount() + attachmentLogicInstance.actions.addAttachment({ + localId: 'local-1', + file: new File(['a'], 'a.png', { type: 'image/png', lastModified: 1 }), + fileKey: 'a.png:1:1:image/png', + previewUrl: 'blob:a.png', + status: 'ready', + }) + attachmentLogicInstance.actions.setAttachmentReady('local-1', 'attachment-1') + logic.actions.setIsSandboxMode(true) + + const openSpy = jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + logic.actions.askMax(' ') + await Promise.resolve() + await Promise.resolve() + + expect(openSpy).not.toHaveBeenCalled() + expect(attachmentLogicInstance.values.attachments[0]?.status).toEqual('ready') + + attachmentLogicInstance.unmount() + }) + it('degrades a keyed non-allowlisted context item to a text attachment instead of dropping it', async () => { - const openSpy = jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + const openSpy = jest + .spyOn(conversationsFrontendApi, 'conversationsOpenCreate') + .mockResolvedValue(sandboxRunResponse) // initKeaTests() in beforeEach resets the kea context, so no explicit unmount is needed attachedContextLogic.mount() attachedContextLogic.actions.registerContext('test-provider', [ @@ -3724,6 +3893,7 @@ describe('maxThreadLogic', () => { }).toDispatchActions(['openSandboxSse']) expect(openSpy).toHaveBeenCalledWith( + '997', MOCK_CONVERSATION_ID, expect.objectContaining({ attached_context: expect.arrayContaining([{ type: 'text', value: 'trace 0189-abc ("LLM trace")' }]), @@ -3748,7 +3918,7 @@ describe('maxThreadLogic', () => { addEventListener(): void {} close(): void {} } - jest.spyOn(api.conversations, 'open').mockResolvedValue(sandboxRunResponse) + jest.spyOn(conversationsFrontendApi, 'conversationsOpenCreate').mockResolvedValue(sandboxRunResponse) }) afterEach(() => { diff --git a/frontend/src/scenes/max/maxThreadLogic.tsx b/frontend/src/scenes/max/maxThreadLogic.tsx index 4a1af9f0be04..fbb047ebd401 100644 --- a/frontend/src/scenes/max/maxThreadLogic.tsx +++ b/frontend/src/scenes/max/maxThreadLogic.tsx @@ -29,6 +29,7 @@ import { uuid } from 'lib/utils/dom' import { maxContextLogic } from 'scenes/max/maxContextLogic' import { notebookLogic } from 'scenes/notebooks/Notebook/notebookLogic' import { NotebookTarget } from 'scenes/notebooks/types' +import { projectLogic } from 'scenes/projectLogic' import { sceneLogic } from 'scenes/sceneLogic' import { Scene } from 'scenes/sceneTypes' import { urls } from 'scenes/urls' @@ -65,11 +66,13 @@ import { SidePanelTab, } from '~/types' +import { conversationsOpenCreate } from 'products/conversations/frontend/generated/api' +import type { InitialPermissionModeEnumApi } from 'products/conversations/frontend/generated/api.schemas' import { + assistantAttachmentsLogic, attachedContextLogic, getRandomThinkingMessage, isTerminalRunStatus, - INITIAL_PERMISSION_MODE, runStreamLogic, } from 'products/posthog_ai/frontend/api/logics' import { LogEntry, parseLogEvent } from 'products/posthog_ai/frontend/lib/parse-logs' @@ -111,6 +114,8 @@ import { export const MAX_DASHBOARD_CONTEXT_WAIT_MS = 8000 const DASHBOARD_CONTEXT_POLL_INTERVAL_MS = 100 +const SANDBOX_INITIAL_PERMISSION_MODE: InitialPermissionModeEnumApi = 'auto' + export type MessageStatus = 'loading' | 'completed' | 'error' export type ThreadMessage = RootAssistantMessage & { @@ -147,6 +152,8 @@ async function shouldBlockPendingPiTask(pendingBindTaskId: string): Promise { + attachmentIds: string[] + } // assistantAttachmentsLogic + markSandboxAttachmentsSendFailed: (attachmentIds: string[]) => { + attachmentIds: string[] + } // assistantAttachmentsLogic + markSandboxAttachmentsSent: (attachmentIds: string[]) => { + attachmentIds: string[] + } // assistantAttachmentsLogic loadConversation: (conversationId: string) => string // maxGlobalLogic askMax: ( prompt: string | null, @@ -547,6 +564,7 @@ export interface maxThreadLogicActions { conversation?: string is_sandbox?: boolean resume_payload?: ResumePayload | null + sandbox_attachment_ids?: string[] ui_context?: any }, generationAttempt: number, @@ -561,6 +579,7 @@ export interface maxThreadLogicActions { conversation?: string | undefined is_sandbox?: boolean | undefined resume_payload?: ResumePayload | null | undefined + sandbox_attachment_ids?: string[] | undefined ui_context?: any } } @@ -743,6 +762,8 @@ export const maxThreadLogic = kea([ ['billingContext'], featureFlagLogic, ['featureFlags'], + projectLogic, + ['currentProjectId'], sceneLogic, ['sceneId'], // Mounts this conversation's sandbox attachment store so its `attachments` are readable @@ -750,6 +771,8 @@ export const maxThreadLogic = kea([ // mounts it itself). posthogAiContextLogic({ conversationId }), ['attachments as sandboxAttachments'], + assistantAttachmentsLogic({ conversationId }), + ['readyAttachmentIds as sandboxAttachmentIds', 'attachmentsAreSending as sandboxAttachmentsAreSending'], // Surfaces the sandbox stream's input-area state to components outside ThreadView's // BindLogic subtree (the input area renders for LangGraph conversations too, so they // can't bind the keyed stream logic themselves). @@ -788,6 +811,12 @@ export const maxThreadLogic = kea([ ], posthogAiContextLogic({ conversationId }), ['clearAttachments as clearSandboxAttachments'], + assistantAttachmentsLogic({ conversationId }), + [ + 'beginAttachmentsSend as beginSandboxAttachmentsSend', + 'markAttachmentsSendFailed as markSandboxAttachmentsSendFailed', + 'markAttachmentsSent as markSandboxAttachmentsSent', + ], ], })), @@ -803,6 +832,7 @@ export const maxThreadLogic = kea([ contextual_tools?: Record ui_context?: any resume_payload?: ResumePayload | null + sandbox_attachment_ids?: string[] }, generationAttempt: number, addToThread: boolean = true @@ -1223,7 +1253,12 @@ export const maxThreadLogic = kea([ listeners((logic) => ({ streamConversation: async ( { - streamData: { agent_mode: agentMode, is_sandbox: isSandbox, ...streamData }, + streamData: { + agent_mode: agentMode, + is_sandbox: isSandbox, + sandbox_attachment_ids: sandboxAttachmentIds = [], + ...streamData + }, generationAttempt, addToThread = true, }, @@ -1271,18 +1306,22 @@ export const maxThreadLogic = kea([ } if (isSandboxConversation) { - // ThreadView renders runStreamLogic's threadItems, not this logic's thread, - // so the human message must be echoed there to show up in the UI. - if (generationAttempt === 0 && streamData.content && addToThread) { - actions.pushSandboxHumanMessage(streamData.content) - // Pull the current scene's `maxContext` into the sandbox attachments at send - // (consumption) time so on-scene entities flow into `sandboxAttachments` below. - // Nothing else dispatches this, so without it scene context never auto-attaches. - posthogAiContextLogic({ conversationId: props.conversationId }).actions.syncSceneAttachments() - } + const sandboxContent = streamData.content?.trim() ? streamData.content : null try { + if (sandboxContent && values.currentProjectId == null) { + throw new Error('Current project has not loaded yet') + } + if (generationAttempt === 0 && sandboxContent && addToThread) { + // ThreadView renders runStreamLogic's threadItems, not this logic's thread, + // so the human message must be echoed there to show up in the UI. + actions.pushSandboxHumanMessage(sandboxContent) + // Pull the current scene's `maxContext` into the sandbox attachments at send + // (consumption) time so on-scene entities flow into `sandboxAttachments` below. + // Nothing else dispatches this, so without it scene context never auto-attaches. + posthogAiContextLogic({ conversationId: props.conversationId }).actions.syncSceneAttachments() + } const conversationId = values.conversation?.id || values.conversationId - if (conversationId && streamData.content) { + if (conversationId && sandboxContent) { // The sandbox runtime has no agent modes — they're a legacy LangGraph concept. If the // user still picked one, carry it through as a context note so the agent can acknowledge it. const attachedContext: AttachedContext[] = [...values.sandboxAttachments] @@ -1341,16 +1380,18 @@ export const maxThreadLogic = kea([ // Single create-or-resume opener: it creates the conversation row on first use, // starts/continues the Run, and returns the (task, run) handle. A message always // provisions a run (a null handle only happens on a warm with a full pool). - const handle = await api.conversations.open(conversationId, { - content: streamData.content, + const handle = await conversationsOpenCreate(String(values.currentProjectId), conversationId, { + content: sandboxContent, trace_id: traceId, attached_context: attachedContext, - initial_permission_mode: INITIAL_PERMISSION_MODE, + ...(sandboxAttachmentIds.length > 0 ? { attachment_ids: sandboxAttachmentIds } : {}), + initial_permission_mode: SANDBOX_INITIAL_PERMISSION_MODE, // Bind a brand-new conversation to an existing Task (inbox "Open task") so the // backend resumes that Task's run. Only the first message carries it. ...(values.pendingBindTaskId ? { task_id: values.pendingBindTaskId } : {}), }) if (handle) { + actions.markSandboxAttachmentsSent(sandboxAttachmentIds) // The sent message consumes any in-flight warm — it's now the active run, so // drop the release handle to avoid cancelling the run out from under it. cache.warmRun = null @@ -1384,6 +1425,7 @@ export const maxThreadLogic = kea([ posthog.captureException(e) actions.pushSandboxError('Failed to send your message. Please try again.') } + actions.markSandboxAttachmentsSendFailed(sandboxAttachmentIds) // The POST failed or no run was started — nothing will stream. Drop the optimistic boot // indicator and release the lock now. actions.setSandboxRunOpening(false) @@ -1725,9 +1767,12 @@ export const maxThreadLogic = kea([ // Warm = open with no message: boots a Run that idles awaiting the first message. The // returned handle lets a later release cancel exactly that Run via the relay (a full // pool returns null — nothing to release). - const warm = await api.conversations.open(values.conversationId, { + if (values.currentProjectId == null) { + return + } + const warm = await conversationsOpenCreate(String(values.currentProjectId), values.conversationId, { content: null, - initial_permission_mode: INITIAL_PERMISSION_MODE, + initial_permission_mode: SANDBOX_INITIAL_PERMISSION_MODE, }) cache.warmRun = warm ? { taskId: warm.task_id, runId: warm.run_id } : null cache.prewarmed = true @@ -1901,195 +1946,225 @@ export const maxThreadLogic = kea([ if (isPiTaskRuntime(values.conversation?.task?.runtime)) { return } - if ( - !values.conversation?.task && - values.pendingBindTaskId && - (await shouldBlockPendingPiTask(values.pendingBindTaskId)) - ) { - return - } - - // A sent message consumes any sandbox pre-warm: the warm Run is the in-progress run the - // sandbox routing follows up on, so cancel pending timers and clear the flag WITHOUT - // issuing a release/DELETE. - cache.disposables.dispose('prewarm-debounce') - cache.disposables.dispose('prewarm-release') - cache.prewarmed = false - cache.warmRun = null - // A sent message consumes the warm — drop any in-flight release intent so the - // run the message follows up on isn't cancelled out from under it. - cache.pendingRelease = false - // Wait for the open dashboard to finish loading before collecting context (see the - // constants above for why). The scene is re-read every tick, so the gate releases the - // moment the dashboard's metadata lands — and immediately if the user navigates away - // mid-wait (no longer on a dashboard, or onto a different one that's already loaded). - const isDashboardSceneLoading = (): boolean => { - if (sceneLogic.values.activeSceneId !== Scene.Dashboard) { - return false - } - const activeSceneLogic = sceneLogic.values.activeSceneLogic - if (!activeSceneLogic) { - // No dashboard scene logic to wait on — its key hasn't resolved or it can't be - // built. Nothing will land, so don't block: send now rather than stalling for the - // full cap and shipping without context anyway. - return false - } - if (!activeSceneLogicHasMaxContext(activeSceneLogic)) { - // The logic exists but isn't mounted yet — building, or briefly unmounted mid - // dashboard→dashboard navigation. Keep waiting (bounded by the cap) so context - // collection picks up the dashboard once it mounts. - return true - } - return !(activeSceneLogic.values as { dashboard?: unknown }).dashboard - } - // Measure real elapsed time, not tick count: breakpoint() only guarantees a *minimum* - // delay, so a busy event loop would make a tick counter under-report the wait — letting - // it run past the cap and skewing the telemetry below. performance.now() is monotonic. - const dashboardWaitStart = performance.now() - while ( - isDashboardSceneLoading() && - performance.now() - dashboardWaitStart < MAX_DASHBOARD_CONTEXT_WAIT_MS - ) { - await breakpoint(DASHBOARD_CONTEXT_POLL_INTERVAL_MS) - } - if (isDashboardSceneLoading()) { - // We hit the wait cap while the dashboard was still loading, so the message ships - // without dashboard context (the original "Max can't see this dashboard" symptom). - // Capture it so we can tell whether the cap is ever the binding constraint in prod. - const activeLoadedScene = sceneLogic.values.activeLoadedScene - const sceneProps = activeLoadedScene?.paramsToProps?.(activeLoadedScene?.sceneParams) || {} - posthog.capture('max dashboard context wait timed out', { - waited_ms: Math.round(performance.now() - dashboardWaitStart), - dashboard_id: (sceneProps as { id?: number | string }).id, - conversation_id: values.conversation?.id || values.conversationId, - }) - } - const contextualTools = Object.fromEntries(values.tools.map((tool) => [tool.identifier, tool.context])) - // Always send voice_mode as an explicit boolean when handsFreeLogic is mounted, - // not just when active. Otherwise a typed turn following a spoken one inherits - // the earlier system instruction from conversation history and - // keeps formatting for speech (no markdown, spelled-out numbers). - const handsFree = handsFreeLogic.findMounted({ panelId: props.panelId }) - const voiceMode = handsFree ? { voice_mode: handsFree.values.isActive } : undefined - const mergedUiContext = - uiContext || voiceMode - ? { ...values.compiledContext, ...uiContext, ...voiceMode } - : values.compiledContext || undefined - const billingContext = - values.billingContext && values.featureFlags[FEATURE_FLAGS.MAX_BILLING_CONTEXT] - ? values.billingContext - : undefined - if ( + const isSandboxSubmission = + values.isSandboxMode || + values.conversation?.agent_runtime === 'sandbox' || + values.pendingBindTaskId != null + const sandboxAttachmentIds = isSandboxSubmission ? [...values.sandboxAttachmentIds] : [] + const isQueueingSubmission = values.queueingEnabled && values.threadLoading && addToThread && typeof prompt === 'string' && prompt.trim() !== '' - ) { - if (values.queueIsFull) { - lemonToast.error('You can only queue two messages at a time.') + const shouldLockSandboxAttachments = + sandboxAttachmentIds.length > 0 && values.dataProcessingAccepted && !isQueueingSubmission + if (shouldLockSandboxAttachments && values.sandboxAttachmentsAreSending) { + return + } + if (shouldLockSandboxAttachments) { + actions.beginSandboxAttachmentsSend(sandboxAttachmentIds) + } + let submissionHandedOff = false + try { + if ( + !values.conversation?.task && + values.pendingBindTaskId && + (await shouldBlockPendingPiTask(values.pendingBindTaskId)) + ) { return } - actions.enqueueQueuedMessage({ - content: prompt, - contextualTools, - uiContext: mergedUiContext, - billingContext, - agentMode: values.agentMode, - }) + + // A sent message consumes any sandbox pre-warm: the warm Run is the in-progress run the + // sandbox routing follows up on, so cancel pending timers and clear the flag WITHOUT + // issuing a release/DELETE. + cache.disposables.dispose('prewarm-debounce') + cache.disposables.dispose('prewarm-release') + cache.prewarmed = false + cache.warmRun = null + // A sent message consumes the warm — drop any in-flight release intent so the + // run the message follows up on isn't cancelled out from under it. + cache.pendingRelease = false + // Wait for the open dashboard to finish loading before collecting context (see the + // constants above for why). The scene is re-read every tick, so the gate releases the + // moment the dashboard's metadata lands — and immediately if the user navigates away + // mid-wait (no longer on a dashboard, or onto a different one that's already loaded). + const isDashboardSceneLoading = (): boolean => { + if (sceneLogic.values.activeSceneId !== Scene.Dashboard) { + return false + } + const activeSceneLogic = sceneLogic.values.activeSceneLogic + if (!activeSceneLogic) { + // No dashboard scene logic to wait on — its key hasn't resolved or it can't be + // built. Nothing will land, so don't block: send now rather than stalling for the + // full cap and shipping without context anyway. + return false + } + if (!activeSceneLogicHasMaxContext(activeSceneLogic)) { + // The logic exists but isn't mounted yet — building, or briefly unmounted mid + // dashboard→dashboard navigation. Keep waiting (bounded by the cap) so context + // collection picks up the dashboard once it mounts. + return true + } + return !(activeSceneLogic.values as { dashboard?: unknown }).dashboard + } + // Measure real elapsed time, not tick count: breakpoint() only guarantees a *minimum* + // delay, so a busy event loop would make a tick counter under-report the wait — letting + // it run past the cap and skewing the telemetry below. performance.now() is monotonic. + const dashboardWaitStart = performance.now() + while ( + isDashboardSceneLoading() && + performance.now() - dashboardWaitStart < MAX_DASHBOARD_CONTEXT_WAIT_MS + ) { + await breakpoint(DASHBOARD_CONTEXT_POLL_INTERVAL_MS) + } + if (isDashboardSceneLoading()) { + // We hit the wait cap while the dashboard was still loading, so the message ships + // without dashboard context (the original "Max can't see this dashboard" symptom). + // Capture it so we can tell whether the cap is ever the binding constraint in prod. + const activeLoadedScene = sceneLogic.values.activeLoadedScene + const sceneProps = activeLoadedScene?.paramsToProps?.(activeLoadedScene?.sceneParams) || {} + posthog.capture('max dashboard context wait timed out', { + waited_ms: Math.round(performance.now() - dashboardWaitStart), + dashboard_id: (sceneProps as { id?: number | string }).id, + conversation_id: values.conversation?.id || values.conversationId, + }) + } + const contextualTools = Object.fromEntries(values.tools.map((tool) => [tool.identifier, tool.context])) + // Always send voice_mode as an explicit boolean when handsFreeLogic is mounted, + // not just when active. Otherwise a typed turn following a spoken one inherits + // the earlier system instruction from conversation history and + // keeps formatting for speech (no markdown, spelled-out numbers). + const handsFree = handsFreeLogic.findMounted({ panelId: props.panelId }) + const voiceMode = handsFree ? { voice_mode: handsFree.values.isActive } : undefined + const mergedUiContext = + uiContext || voiceMode + ? { ...values.compiledContext, ...uiContext, ...voiceMode } + : values.compiledContext || undefined + const billingContext = + values.billingContext && values.featureFlags[FEATURE_FLAGS.MAX_BILLING_CONTEXT] + ? values.billingContext + : undefined + + if ( + values.queueingEnabled && + values.threadLoading && + addToThread && + typeof prompt === 'string' && + prompt.trim() !== '' + ) { + if (values.queueIsFull) { + lemonToast.error('You can only queue two messages at a time.') + return + } + actions.enqueueQueuedMessage({ + content: prompt, + contextualTools, + uiContext: mergedUiContext, + billingContext, + agentMode: values.agentMode, + }) + actions.setQuestion('') + if (props.panelId === SIDE_PANEL_PANEL_ID && sidePanelStateLogic.isMounted()) { + sidePanelStateLogic.actions.setSidePanelOptions(null) + } + return + } + if (!values.dataProcessingAccepted) { + // Persist prompt to sessionStorage in case of OAuth redirect during consent flow + if (prompt) { + try { + sessionStorage.setItem( + PENDING_AI_PROMPT_KEY, + JSON.stringify({ + prompt, + timestamp: Date.now(), + }) + ) + } catch { + // sessionStorage might be unavailable + } + } + return // Skip - this will be re-fired by the `onApprove` on `AIConsentPopoverWrapper` + } + + // Clear any stored prompt since we're proceeding with submission + try { + sessionStorage.removeItem(PENDING_AI_PROMPT_KEY) + } catch { + // sessionStorage might be unavailable + } + + // Build auto-rejection payload if there's a pending approval that hasn't already been resolved + // (pendingApprovalProposalId might get re-set during streaming even after user approved/rejected) + let autoRejectPayload: { action: 'reject'; proposal_id: string; feedback?: string } | undefined = + undefined + const pendingProposalId = values.pendingApprovalProposalId + const alreadyResolved = pendingProposalId + ? !!values.resolvedApprovalStatuses[pendingProposalId]?.status + : false + if (pendingProposalId && !alreadyResolved) { + autoRejectPayload = { + action: 'reject', + proposal_id: pendingProposalId, + feedback: prompt ?? undefined, + } + actions.clearPendingApproval() + actions.setResolvedApprovalStatus(pendingProposalId, 'auto_rejected') + } + // A pending task-bind (inbox "Open task") makes this first message a sandbox task-resume: + // the conversation is created bound to the Task and its run is resumed. Force sandbox mode + // so the message routes through the sandbox `open` endpoint (which carries the task_id). + // Reducers apply synchronously, so the `is_sandbox` derivation below sees the new value. + if (values.pendingBindTaskId && !values.isSandboxMode) { + actions.setIsSandboxMode(true) + } + const agentMode = values.agentMode + + // Clear the question actions.setQuestion('') + // Drop #panel=max:… options so reload doesn't re-run auto-send from the hash if (props.panelId === SIDE_PANEL_PANEL_ID && sidePanelStateLogic.isMounted()) { sidePanelStateLogic.actions.setSidePanelOptions(null) } - return - } - if (!values.dataProcessingAccepted) { - // Persist prompt to sessionStorage in case of OAuth redirect during consent flow - if (prompt) { - try { - sessionStorage.setItem( - PENDING_AI_PROMPT_KEY, - JSON.stringify({ - prompt, - timestamp: Date.now(), - }) - ) - } catch { - // sessionStorage might be unavailable + // For a new conversations, set the frontend conversation ID + if (!values.conversation) { + actions.setConversationId(values.conversationId) + } else { + const updatedConversation = { + ...values.conversation, + agent_mode: agentMode || values.conversation?.agent_mode, + status: ConversationStatus.InProgress, + updated_at: dayjs().toISOString(), } + // Update the current status + actions.setConversation(updatedConversation) + // Update the global conversation cache + actions.updateGlobalConversationCache(updatedConversation) } - return // Skip - this will be re-fired by the `onApprove` on `AIConsentPopoverWrapper` - } - // Clear any stored prompt since we're proceeding with submission - try { - sessionStorage.removeItem(PENDING_AI_PROMPT_KEY) - } catch { - // sessionStorage might be unavailable - } - - // Build auto-rejection payload if there's a pending approval that hasn't already been resolved - // (pendingApprovalProposalId might get re-set during streaming even after user approved/rejected) - let autoRejectPayload: { action: 'reject'; proposal_id: string; feedback?: string } | undefined = undefined - const pendingProposalId = values.pendingApprovalProposalId - const alreadyResolved = pendingProposalId - ? !!values.resolvedApprovalStatuses[pendingProposalId]?.status - : false - if (pendingProposalId && !alreadyResolved) { - autoRejectPayload = { - action: 'reject', - proposal_id: pendingProposalId, - feedback: prompt ?? undefined, - } - actions.clearPendingApproval() - actions.setResolvedApprovalStatus(pendingProposalId, 'auto_rejected') - } - // A pending task-bind (inbox "Open task") makes this first message a sandbox task-resume: - // the conversation is created bound to the Task and its run is resumed. Force sandbox mode - // so the message routes through the sandbox `open` endpoint (which carries the task_id). - // Reducers apply synchronously, so the `is_sandbox` derivation below sees the new value. - if (values.pendingBindTaskId && !values.isSandboxMode) { - actions.setIsSandboxMode(true) - } - const agentMode = values.agentMode - - // Clear the question - actions.setQuestion('') - // Drop #panel=max:… options so reload doesn't re-run auto-send from the hash - if (props.panelId === SIDE_PANEL_PANEL_ID && sidePanelStateLogic.isMounted()) { - sidePanelStateLogic.actions.setSidePanelOptions(null) - } - // For a new conversations, set the frontend conversation ID - if (!values.conversation) { - actions.setConversationId(values.conversationId) - } else { - const updatedConversation = { - ...values.conversation, - agent_mode: agentMode || values.conversation?.agent_mode, - status: ConversationStatus.InProgress, - updated_at: dayjs().toISOString(), + submissionHandedOff = true + actions.streamConversation( + { + agent_mode: values.isSandboxMode ? null : agentMode, + is_sandbox: values.isSandboxMode || undefined, + content: prompt, + contextual_tools: contextualTools, + ui_context: mergedUiContext, + conversation: values.conversation?.id || values.conversationId, + // Include auto-rejection payload if there was a pending approval + resume_payload: autoRejectPayload, + ...(isSandboxSubmission ? { sandbox_attachment_ids: sandboxAttachmentIds } : {}), + }, + 0, + addToThread + ) + } finally { + if (shouldLockSandboxAttachments && !submissionHandedOff) { + actions.markSandboxAttachmentsSendFailed(sandboxAttachmentIds) } - // Update the current status - actions.setConversation(updatedConversation) - // Update the global conversation cache - actions.updateGlobalConversationCache(updatedConversation) } - - actions.streamConversation( - { - agent_mode: values.isSandboxMode ? null : agentMode, - is_sandbox: values.isSandboxMode || undefined, - content: prompt, - contextual_tools: contextualTools, - ui_context: mergedUiContext, - conversation: values.conversation?.id || values.conversationId, - // Include auto-rejection payload if there was a pending approval - resume_payload: autoRejectPayload, - }, - 0, - addToThread - ) }, stopGeneration: async () => { if (!values.conversation?.id) { diff --git a/posthog/settings/web.py b/posthog/settings/web.py index 4b458300444b..26aadff6181d 100644 --- a/posthog/settings/web.py +++ b/posthog/settings/web.py @@ -639,6 +639,7 @@ def static_varies_origin(headers, path, url): "NotebookSQLV2NodeTypeEnum": ["hogql", "python"], "NotebookSQLV2RefKindEnum": ["hogql", "local"], "TaskRunStatusEnum": "products.tasks.backend.models.TaskRun.Status", + "AssistantAttachmentContentTypeEnum": ["image/png", "image/jpeg"], # Inline-choices variant of TaskRun.Status (labels == values), shared by # TaskRunUpdate.status and ExperimentFlagCleanupTask.run_status. "RunStatusEnum": ["not_started", "queued", "in_progress", "completed", "failed", "cancelled"], diff --git a/products/conversations/frontend/generated/api.schemas.ts b/products/conversations/frontend/generated/api.schemas.ts index 79349969bf09..86f125398956 100644 --- a/products/conversations/frontend/generated/api.schemas.ts +++ b/products/conversations/frontend/generated/api.schemas.ts @@ -526,18 +526,24 @@ export const InitialPermissionModeEnumApi = { } as const /** - * Request body for `POST /conversations/{id}/open/`. A string `content` processes a turn; a - * null/absent `content` warms a sandbox that idles awaiting the first message. + * Request body for `POST /conversations/{id}/open/`. Nonblank `content` processes a turn and may include + * `attachment_ids`; null or absent `content` with no attachments warms a sandbox awaiting its first message. */ export interface SandboxOpenApi { /** - * The user's message text. Omit or null to warm a sandbox (boot + idle) ahead of the first message. + * The user's message text. Omit or null to warm a sandbox ahead of the first message, unless attachment_ids are provided. * @maxLength 40000 * @nullable */ content?: string | null /** Client-generated trace id correlated with the resulting Run's SSE stream. */ trace_id?: string + /** + * Finalized sandbox image attachment IDs to send with this message. + * @minItems 1 + * @maxItems 4 + */ + attachment_ids?: string[] /** Typed PostHog entities (and free text) attached to this message. */ attached_context?: SandboxAttachedContextItemApi[] /** Initial permission mode for the sandbox agent session. Defaults to `auto`, which allows safe tool use while preserving explicit confirmations. diff --git a/products/conversations/frontend/generated/api.zod.ts b/products/conversations/frontend/generated/api.zod.ts index 0d5c3175becf..2429e50dcb88 100644 --- a/products/conversations/frontend/generated/api.zod.ts +++ b/products/conversations/frontend/generated/api.zod.ts @@ -76,6 +76,8 @@ export const ConversationsCancelPartialUpdateBody = /* @__PURE__ */ zod.looseObj */ export const conversationsOpenCreateBodyContentMax = 40000 +export const conversationsOpenCreateBodyAttachmentIdsMax = 4 + export const ConversationsOpenCreateBody = /* @__PURE__ */ zod .object({ content: zod @@ -83,12 +85,18 @@ export const ConversationsOpenCreateBody = /* @__PURE__ */ zod .max(conversationsOpenCreateBodyContentMax) .nullish() .describe( - "The user's message text. Omit or null to warm a sandbox (boot + idle) ahead of the first message." + "The user's message text. Omit or null to warm a sandbox ahead of the first message, unless attachment_ids are provided." ), trace_id: zod .uuid() .optional() .describe("Client-generated trace id correlated with the resulting Run's SSE stream."), + attachment_ids: zod + .array(zod.uuid()) + .min(1) + .max(conversationsOpenCreateBodyAttachmentIdsMax) + .optional() + .describe('Finalized sandbox image attachment IDs to send with this message.'), attached_context: zod .array( zod @@ -145,7 +153,7 @@ export const ConversationsOpenCreateBody = /* @__PURE__ */ zod ), }) .describe( - 'Request body for `POST \/conversations\/{id}\/open\/`. A string `content` processes a turn; a\nnull\/absent `content` warms a sandbox that idles awaiting the first message.' + 'Request body for `POST \/conversations\/{id}\/open\/`. Nonblank `content` processes a turn and may include\n`attachment_ids`; null or absent `content` with no attachments warms a sandbox awaiting its first message.' ) export const ConversationsQueueCreateBody = /* @__PURE__ */ zod.looseObject({}) diff --git a/products/desktop/packages/agent/src/posthog-api.test.ts b/products/desktop/packages/agent/src/posthog-api.test.ts index 54d543f045ae..b6194fe7d558 100644 --- a/products/desktop/packages/agent/src/posthog-api.test.ts +++ b/products/desktop/packages/agent/src/posthog-api.test.ts @@ -154,6 +154,59 @@ describe("PostHogAPIClient", () => { ); }); + it("rejects an oversized artifact before reading its response body", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + const arrayBuffer = vi.fn(); + mockFetch.mockResolvedValueOnce({ + ok: true, + headers: new Headers({ "content-length": "5" }), + arrayBuffer, + }); + + const artifact = await client.downloadArtifact( + "task-1", + "run-1", + "tasks/artifacts/image.png", + 4, + ); + + expect(artifact).toBeNull(); + expect(arrayBuffer).not.toHaveBeenCalled(); + }); + + it("cancels a streamed artifact as soon as it exceeds the byte limit", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + const cancel = vi.fn(); + const read = vi + .fn() + .mockResolvedValueOnce({ done: false, value: new Uint8Array([1, 2, 3]) }) + .mockResolvedValueOnce({ done: false, value: new Uint8Array([4, 5, 6]) }); + mockFetch.mockResolvedValueOnce({ + ok: true, + headers: new Headers(), + body: { getReader: () => ({ read, cancel }) }, + }); + + const artifact = await client.downloadArtifact( + "task-1", + "run-1", + "tasks/artifacts/image.png", + 4, + ); + + expect(artifact).toBeNull(); + expect(read).toHaveBeenCalledTimes(2); + expect(cancel).toHaveBeenCalledOnce(); + }); + it.each([ [ "includes message_id and text_parts when provided", diff --git a/products/desktop/packages/agent/src/posthog-api.ts b/products/desktop/packages/agent/src/posthog-api.ts index 8872cbc851e0..33c9163226f0 100644 --- a/products/desktop/packages/agent/src/posthog-api.ts +++ b/products/desktop/packages/agent/src/posthog-api.ts @@ -540,6 +540,7 @@ export class PostHogAPIClient { taskId: string, runId: string, storagePath: string, + maxBytes?: number, ): Promise { const teamId = this.getTeamId(); @@ -554,7 +555,45 @@ export class PostHogAPIClient { if (!response.ok) { throw new Error(`Failed to download artifact: ${response.status}`); } - return response.arrayBuffer(); + if (maxBytes === undefined) { + return response.arrayBuffer(); + } + + const contentLength = Number(response.headers.get("content-length")); + if (Number.isFinite(contentLength) && contentLength > maxBytes) { + throw new Error("Artifact exceeds the download size limit"); + } + if (!response.body) { + const data = await response.arrayBuffer(); + if (data.byteLength > maxBytes) { + throw new Error("Artifact exceeds the download size limit"); + } + return data; + } + + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let totalBytes = 0; + for (;;) { + const { done, value } = await reader.read(); + if (done) { + break; + } + totalBytes += value.byteLength; + if (totalBytes > maxBytes) { + await reader.cancel(); + throw new Error("Artifact exceeds the download size limit"); + } + chunks.push(value); + } + + const data = new Uint8Array(totalBytes); + let offset = 0; + for (const chunk of chunks) { + data.set(chunk, offset); + offset += chunk.byteLength; + } + return data.buffer; } catch { return null; } diff --git a/products/desktop/packages/agent/src/server/agent-server.test.ts b/products/desktop/packages/agent/src/server/agent-server.test.ts index 702befbdc184..dde5aeb39683 100644 --- a/products/desktop/packages/agent/src/server/agent-server.test.ts +++ b/products/desktop/packages/agent/src/server/agent-server.test.ts @@ -38,7 +38,7 @@ import { type TestRepo, } from "../test/fixtures/api"; import { createPostHogHandlers } from "../test/mocks/msw-handlers"; -import type { StoredEntry, TaskRun } from "../types"; +import type { StoredEntry, TaskRun, TaskRunArtifact } from "../types"; import { AgentServer, isTurnCompleteNotification, @@ -2585,6 +2585,21 @@ describe("AgentServer HTTP Mode", () => { const checksum = createHash("sha256") .update(Buffer.from(bundle)) .digest("hex"); + const skillArtifact = { + id: "skill-artifact-1", + name: "local-test-skill.zip", + type: "skill_bundle", + source: "posthog_code_skill", + storage_path: "tasks/artifacts/local-test-skill.zip", + content_type: "application/zip", + metadata: { + skill_name: "local-test-skill", + skill_source: "user", + content_sha256: checksum, + bundle_format: "zip", + schema_version: 1, + }, + } as TaskRunArtifact; const s = createServer(); await s.start(); @@ -2597,10 +2612,16 @@ describe("AgentServer HTTP Mode", () => { const downloadArtifact = vi.fn(async () => exactArrayBuffer(bundle)); const serverInternals = s as unknown as { session: { clientConnection: { prompt: typeof prompt } }; - posthogAPI: { downloadArtifact: typeof downloadArtifact }; + posthogAPI: { + downloadArtifact: typeof downloadArtifact; + getTaskRun: ReturnType; + }; }; serverInternals.session.clientConnection.prompt = prompt; serverInternals.posthogAPI.downloadArtifact = downloadArtifact; + serverInternals.posthogAPI.getTaskRun = vi.fn(async () => + createTaskRun({ artifacts: [skillArtifact] }), + ); const token = createToken(); const response = await fetch(`http://localhost:${port}/command`, { @@ -2615,23 +2636,7 @@ describe("AgentServer HTTP Mode", () => { method: "user_message", params: { content: "/local-test-skill with context", - artifacts: [ - { - id: "skill-artifact-1", - name: "local-test-skill.zip", - type: "skill_bundle", - source: "posthog_code_skill", - storage_path: "tasks/artifacts/local-test-skill.zip", - content_type: "application/zip", - metadata: { - skill_name: "local-test-skill", - skill_source: "user", - content_sha256: checksum, - bundle_format: "zip", - schema_version: 1, - }, - }, - ], + artifacts: [skillArtifact], }, }), }); @@ -2685,6 +2690,25 @@ describe("AgentServer HTTP Mode", () => { const depBundle = makeBundle("dep-skill", "Dependency instructions."); const checksumOf = (bundle: Uint8Array) => createHash("sha256").update(Buffer.from(bundle)).digest("hex"); + const makeArtifact = ( + id: string, + name: string, + bundle: Uint8Array, + ): TaskRunArtifact => ({ + id, + name: `${name}.zip`, + type: "skill_bundle", + source: "posthog_code_skill", + storage_path: `tasks/artifacts/${name}.zip`, + content_type: "application/zip", + metadata: { + skill_name: name, + skill_source: "user", + content_sha256: checksumOf(bundle), + bundle_format: "zip", + schema_version: 1, + }, + }); const s = createServer(); await s.start(); @@ -2700,32 +2724,22 @@ describe("AgentServer HTTP Mode", () => { storagePath.includes("dep-skill") ? depBundle : invokedBundle, ), ); + const artifacts = [ + makeArtifact("skill-artifact-parent", "parent-skill", invokedBundle), + makeArtifact("skill-artifact-dep", "dep-skill", depBundle), + ]; const serverInternals = s as unknown as { session: { clientConnection: { prompt: typeof prompt } }; - posthogAPI: { downloadArtifact: typeof downloadArtifact }; + posthogAPI: { + downloadArtifact: typeof downloadArtifact; + getTaskRun: ReturnType; + }; }; serverInternals.session.clientConnection.prompt = prompt; serverInternals.posthogAPI.downloadArtifact = downloadArtifact; - - const makeArtifact = ( - id: string, - name: string, - bundle: Uint8Array, - ): Record => ({ - id, - name: `${name}.zip`, - type: "skill_bundle", - source: "posthog_code_skill", - storage_path: `tasks/artifacts/${name}.zip`, - content_type: "application/zip", - metadata: { - skill_name: name, - skill_source: "user", - content_sha256: checksumOf(bundle), - bundle_format: "zip", - schema_version: 1, - }, - }); + serverInternals.posthogAPI.getTaskRun = vi.fn(async () => + createTaskRun({ artifacts }), + ); const token = createToken(); const response = await fetch(`http://localhost:${port}/command`, { @@ -2740,14 +2754,7 @@ describe("AgentServer HTTP Mode", () => { method: "user_message", params: { content: "/parent-skill run it", - artifacts: [ - makeArtifact( - "skill-artifact-parent", - "parent-skill", - invokedBundle, - ), - makeArtifact("skill-artifact-dep", "dep-skill", depBundle), - ], + artifacts, }, }), }); @@ -2784,6 +2791,25 @@ describe("AgentServer HTTP Mode", () => { const prefixBundle = makeBundle("mentioned", "PREFIX_SKILL_MARKER body."); const checksumOf = (bundle: Uint8Array) => createHash("sha256").update(Buffer.from(bundle)).digest("hex"); + const makeArtifact = ( + id: string, + name: string, + bundle: Uint8Array, + ): TaskRunArtifact => ({ + id, + name: `${name}.zip`, + type: "skill_bundle", + source: "posthog_code_skill", + storage_path: `tasks/artifacts/${name}.zip`, + content_type: "application/zip", + metadata: { + skill_name: name, + skill_source: "user", + content_sha256: checksumOf(bundle), + bundle_format: "zip", + schema_version: 1, + }, + }); const s = createServer(); await s.start(); @@ -2801,32 +2827,26 @@ describe("AgentServer HTTP Mode", () => { : prefixBundle, ), ); + const artifacts = [ + makeArtifact( + "skill-artifact-mentioned", + "mentioned-skill", + mentionedBundle, + ), + makeArtifact("skill-artifact-prefix", "mentioned", prefixBundle), + ]; const serverInternals = s as unknown as { session: { clientConnection: { prompt: typeof prompt } }; - posthogAPI: { downloadArtifact: typeof downloadArtifact }; + posthogAPI: { + downloadArtifact: typeof downloadArtifact; + getTaskRun: ReturnType; + }; }; serverInternals.session.clientConnection.prompt = prompt; serverInternals.posthogAPI.downloadArtifact = downloadArtifact; - - const makeArtifact = ( - id: string, - name: string, - bundle: Uint8Array, - ): Record => ({ - id, - name: `${name}.zip`, - type: "skill_bundle", - source: "posthog_code_skill", - storage_path: `tasks/artifacts/${name}.zip`, - content_type: "application/zip", - metadata: { - skill_name: name, - skill_source: "user", - content_sha256: checksumOf(bundle), - bundle_format: "zip", - schema_version: 1, - }, - }); + serverInternals.posthogAPI.getTaskRun = vi.fn(async () => + createTaskRun({ artifacts }), + ); const token = createToken(); const response = await fetch(`http://localhost:${port}/command`, { @@ -2841,14 +2861,7 @@ describe("AgentServer HTTP Mode", () => { method: "user_message", params: { content: "please use /mentioned-skill on the diff", - artifacts: [ - makeArtifact( - "skill-artifact-mentioned", - "mentioned-skill", - mentionedBundle, - ), - makeArtifact("skill-artifact-prefix", "mentioned", prefixBundle), - ], + artifacts, }, }), }); diff --git a/products/desktop/packages/agent/src/server/agent-server.ts b/products/desktop/packages/agent/src/server/agent-server.ts index a583d1ae3b38..8da17f6ae926 100644 --- a/products/desktop/packages/agent/src/server/agent-server.ts +++ b/products/desktop/packages/agent/src/server/agent-server.ts @@ -113,7 +113,12 @@ import { checkoutExistingPullRequest, type ExistingPrCheckoutResult, } from "./pr-checkout"; -import { resolveUserArtifactsById } from "./resolve-user-artifacts"; +import { + MAX_USER_ATTACHMENT_BYTES, + resolveLegacyArtifactsFromManifest, + resolveUserArtifactsById, + validateUserArtifactData, +} from "./resolve-user-artifacts"; import { resolveRtkSavings } from "./rtk-savings"; import { RunUsageAccumulator } from "./run-usage"; import { @@ -1185,6 +1190,12 @@ export class AgentServer { ? params.artifact_ids.length : 0, }); + const legacyArtifacts = await resolveLegacyArtifactsFromManifest( + this.posthogAPI, + commandSession.payload.task_id, + commandSession.payload.run_id, + Array.isArray(params.artifacts) ? params.artifacts : [], + ); const resolvedArtifacts = await resolveUserArtifactsById( this.posthogAPI, commandSession.payload.task_id, @@ -1195,12 +1206,7 @@ export class AgentServer { ); const builtPrompt = await this.buildPromptFromContentAndArtifacts({ content: params.content as string | ContentBlock[] | undefined, - artifacts: [ - ...(Array.isArray(params.artifacts) - ? (params.artifacts as TaskRunArtifact[]) - : []), - ...resolvedArtifacts, - ], + artifacts: [...legacyArtifacts, ...resolvedArtifacts], taskId: commandSession.payload.task_id, runId: commandSession.payload.run_id, }); @@ -2992,7 +2998,9 @@ export class AgentServer { const hasMatchingArtifact = artifacts.some( (artifact) => artifact.type === "skill_bundle" && - artifact.metadata?.skill_name === invocation.skillName, + artifact.metadata && + "skill_name" in artifact.metadata && + artifact.metadata.skill_name === invocation.skillName, ); const installedSkill = hasMatchingArtifact ? this.installedSkillBundleInfo.get( @@ -3034,7 +3042,11 @@ export class AgentServer { ): LocalSkillPromptContext | null { const installed = artifacts .filter((artifact) => artifact.type === "skill_bundle") - .map((artifact) => artifact.metadata?.skill_name) + .map((artifact) => + artifact.metadata && "skill_name" in artifact.metadata + ? artifact.metadata.skill_name + : undefined, + ) .filter((name): name is string => typeof name === "string") .map((name) => this.installedSkillBundleInfo.get( @@ -3155,8 +3167,12 @@ export class AgentServer { artifact: TaskRunArtifact, ): Promise { const metadata = artifact.metadata; - const skillName = metadata?.skill_name; - const expectedSha256 = metadata?.content_sha256; + const skillName = + metadata && "skill_name" in metadata ? metadata.skill_name : undefined; + const expectedSha256 = + metadata && "content_sha256" in metadata + ? metadata.content_sha256 + : undefined; if (!artifact.storage_path || !skillName || !expectedSha256) { throw new Error( @@ -3337,14 +3353,23 @@ export class AgentServer { return null; } - const data = await this.posthogAPI.downloadArtifact( - taskId, - runId, - artifact.storage_path, - ); + const data = + artifact.source === "user_attachment" + ? await this.posthogAPI.downloadArtifact( + taskId, + runId, + artifact.storage_path, + MAX_USER_ATTACHMENT_BYTES, + ) + : await this.posthogAPI.downloadArtifact( + taskId, + runId, + artifact.storage_path, + ); if (!data) { throw new Error(`Failed to download artifact ${artifact.name}`); } + validateUserArtifactData(artifact, data); const mimeType = artifact.content_type?.toLowerCase(); if (mimeType === "image/png" || mimeType === "image/jpeg") { @@ -3357,7 +3382,7 @@ export class AgentServer { ".posthog", "attachments", runId, - artifact.id ?? safeName, + this.getSafeArtifactName(artifact.id ?? safeName), ); await mkdir(artifactDir, { recursive: true }); diff --git a/products/desktop/packages/agent/src/server/pi-agent-server.test.ts b/products/desktop/packages/agent/src/server/pi-agent-server.test.ts index c73c12539156..eb987b3e1026 100644 --- a/products/desktop/packages/agent/src/server/pi-agent-server.test.ts +++ b/products/desktop/packages/agent/src/server/pi-agent-server.test.ts @@ -342,8 +342,28 @@ describe("PiAgentServer", () => { .fn() .mockResolvedValueOnce(Buffer.from("notes")) .mockResolvedValueOnce(Buffer.from("image")); + const artifacts = [ + { + id: "file-1", + name: "notes.txt", + type: "user_attachment", + content_type: "text/plain", + storage_path: "artifacts/notes.txt", + }, + { + id: "image-1", + name: "image.png", + type: "user_attachment", + content_type: "image/png", + storage_path: "artifacts/image.png", + }, + ]; + const getTaskRun = vi.fn(async () => ({ artifacts })); const server = new PiAgentServer(config({ repositoryPath })) as unknown as { - posthogAPI: { downloadArtifact: typeof downloadArtifact }; + posthogAPI: { + downloadArtifact: typeof downloadArtifact; + getTaskRun: typeof getTaskRun; + }; session: unknown; executeCommand( method: string, @@ -351,6 +371,7 @@ describe("PiAgentServer", () => { ): Promise; }; server.posthogAPI.downloadArtifact = downloadArtifact; + server.posthogAPI.getTaskRun = getTaskRun; server.session = { runtime: { client: { @@ -362,22 +383,7 @@ describe("PiAgentServer", () => { await server.executeCommand("user_message", { content: "Read these", - artifacts: [ - { - id: "file-1", - name: "notes.txt", - type: "user_attachment", - content_type: "text/plain", - storage_path: "artifacts/notes.txt", - }, - { - id: "image-1", - name: "image.png", - type: "user_attachment", - content_type: "image/png", - storage_path: "artifacts/image.png", - }, - ], + artifacts, }); const command = sendCommand.mock.calls[0][0]; diff --git a/products/desktop/packages/agent/src/server/pi-agent-server.ts b/products/desktop/packages/agent/src/server/pi-agent-server.ts index ba3ee640ce63..3f96242dcddc 100644 --- a/products/desktop/packages/agent/src/server/pi-agent-server.ts +++ b/products/desktop/packages/agent/src/server/pi-agent-server.ts @@ -33,7 +33,12 @@ import { resolveLlmGatewayUrl } from "../utils/gateway"; import { Logger } from "../utils/logger"; import { TaskRunEventStreamSender } from "./event-stream-sender"; import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; -import { resolveUserArtifactsById } from "./resolve-user-artifacts"; +import { + MAX_USER_ATTACHMENT_BYTES, + resolveLegacyArtifactsFromManifest, + resolveUserArtifactsById, + validateUserArtifactData, +} from "./resolve-user-artifacts"; import { jsonRpcRequestSchema, userAttachmentIdsSchema } from "./schemas"; import type { AgentServerConfig } from "./types"; @@ -710,9 +715,12 @@ export class PiAgentServer { runtime: PiRuntime, params: Record, ): Promise { - const legacyArtifacts = Array.isArray(params.artifacts) - ? (params.artifacts as TaskRunArtifact[]) - : []; + const legacyArtifacts = await resolveLegacyArtifactsFromManifest( + this.posthogAPI, + this.config.taskId, + this.config.runId, + Array.isArray(params.artifacts) ? params.artifacts : [], + ); const resolvedArtifacts = await resolveUserArtifactsById( this.posthogAPI, this.config.taskId, @@ -754,14 +762,23 @@ export class PiAgentServer { if (!artifact.storage_path) { continue; } - const data = await this.posthogAPI.downloadArtifact( - this.config.taskId, - this.config.runId, - artifact.storage_path, - ); + const data = + artifact.source === "user_attachment" + ? await this.posthogAPI.downloadArtifact( + this.config.taskId, + this.config.runId, + artifact.storage_path, + MAX_USER_ATTACHMENT_BYTES, + ) + : await this.posthogAPI.downloadArtifact( + this.config.taskId, + this.config.runId, + artifact.storage_path, + ); if (!data) { throw new Error(`Failed to download attachment: ${artifact.name}`); } + validateUserArtifactData(artifact, data); const mimeType = artifact.content_type ?? "application/octet-stream"; if (mimeType.startsWith("image/")) { @@ -775,7 +792,11 @@ export class PiAgentServer { } await mkdir(attachmentDirectory, { recursive: true }); - const fileName = `${artifact.id}-${basename(artifact.name)}`; + const safeArtifactId = basename(artifact.id ?? "attachment").replace( + /[^\w.-]/g, + "_", + ); + const fileName = `${safeArtifactId}-${basename(artifact.name)}`; const filePath = join(attachmentDirectory, fileName); await writeFile(filePath, Buffer.from(data)); filePaths.push(filePath); diff --git a/products/desktop/packages/agent/src/server/resolve-user-artifacts.test.ts b/products/desktop/packages/agent/src/server/resolve-user-artifacts.test.ts new file mode 100644 index 000000000000..c70a0402f72a --- /dev/null +++ b/products/desktop/packages/agent/src/server/resolve-user-artifacts.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it, vi } from "vitest"; + +import type { TaskRunArtifact } from "../types"; +import { + resolveLegacyArtifactsFromManifest, + validateUserArtifactData, +} from "./resolve-user-artifacts"; + +describe("validateUserArtifactData", () => { + const artifact: TaskRunArtifact = { + id: "11111111-1111-4111-8111-111111111111", + name: "screen.png", + type: "context", + source: "user_attachment", + size: 3, + content_type: "image/png", + storage_path: "artifacts/screen.png", + uploaded_by: "user", + uploaded_by_user_id: 1, + }; + + it.each([ + ["empty", new ArrayBuffer(0)], + ["different from the manifest size", new Uint8Array([1, 2]).buffer], + ["larger than 4 MiB", new ArrayBuffer(4 * 1024 * 1024 + 1)], + ])("rejects %s screenshot content", (_name, data) => { + expect(() => validateUserArtifactData(artifact, data)).toThrow( + "Screenshot attachment content is invalid", + ); + }); + + it("accepts screenshot content matching the validated manifest", () => { + expect(() => + validateUserArtifactData(artifact, new Uint8Array([1, 2, 3]).buffer), + ).not.toThrow(); + }); +}); + +describe("resolveLegacyArtifactsFromManifest", () => { + const manifestArtifact: TaskRunArtifact = { + id: "artifact-1", + name: "report.txt", + type: "context", + source: "agent_output", + size: 6, + content_type: "text/plain", + storage_path: "tasks/artifacts/report.txt", + uploaded_by: "agent", + }; + + it("returns canonical manifest metadata instead of caller metadata", async () => { + const getTaskRun = vi.fn(async () => ({ + artifacts: [manifestArtifact], + })); + + const resolved = await resolveLegacyArtifactsFromManifest( + { getTaskRun } as never, + "task-1", + "run-1", + [ + { + ...manifestArtifact, + name: "../../escaped.txt", + source: "user_attachment", + content_type: "image/png", + }, + ], + ); + + expect(resolved).toEqual([manifestArtifact]); + }); + + it.each([ + ["unknown id", { id: "../../target" }], + [ + "mismatched storage path", + { id: manifestArtifact.id, storage_path: "tasks/artifacts/other.txt" }, + ], + ["missing selector", { name: "report.txt" }], + ])("rejects a legacy artifact with %s", async (_name, suppliedArtifact) => { + const getTaskRun = vi.fn(async () => ({ + artifacts: [manifestArtifact], + })); + + await expect( + resolveLegacyArtifactsFromManifest( + { getTaskRun } as never, + "task-1", + "run-1", + [suppliedArtifact], + ), + ).rejects.toThrow("Artifacts are unavailable for this task run"); + }); +}); diff --git a/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts b/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts index 927cd61c0ecd..55169d8b27b2 100644 --- a/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts +++ b/products/desktop/packages/agent/src/server/resolve-user-artifacts.ts @@ -1,7 +1,23 @@ import type { PostHogAPIClient } from "../posthog-api"; import type { TaskRunArtifact } from "../types"; -const MAX_USER_ATTACHMENT_BYTES = 4 * 1024 * 1024; +export const MAX_USER_ATTACHMENT_BYTES = 4 * 1024 * 1024; + +export function validateUserArtifactData( + artifact: TaskRunArtifact, + data: ArrayBuffer, +): void { + if (artifact.source !== "user_attachment") { + return; + } + if ( + data.byteLength === 0 || + data.byteLength > MAX_USER_ATTACHMENT_BYTES || + data.byteLength !== artifact.size + ) { + throw new Error("Screenshot attachment content is invalid"); + } +} function isScreenshotAttachment( artifact: TaskRunArtifact, @@ -63,3 +79,55 @@ export async function resolveUserArtifactsById( return resolvedArtifacts; } + +export async function resolveLegacyArtifactsFromManifest( + posthogAPI: Pick, + taskId: string, + runId: string, + suppliedArtifacts: unknown[], +): Promise { + if (suppliedArtifacts.length === 0) { + return []; + } + + const taskRun = await posthogAPI.getTaskRun(taskId, runId); + const manifest = taskRun.artifacts ?? []; + const resolvedArtifacts: TaskRunArtifact[] = []; + const resolvedKeys = new Set(); + + for (const suppliedArtifact of suppliedArtifacts) { + if (!suppliedArtifact || typeof suppliedArtifact !== "object") { + throw new Error("Artifacts are unavailable for this task run"); + } + const supplied = suppliedArtifact as Record; + const suppliedId = + typeof supplied.id === "string" && supplied.id.length > 0 + ? supplied.id + : null; + const suppliedStoragePath = + typeof supplied.storage_path === "string" && + supplied.storage_path.length > 0 + ? supplied.storage_path + : null; + const artifact = manifest.find((candidate) => + suppliedId + ? candidate.id === suppliedId && + (!suppliedStoragePath || + candidate.storage_path === suppliedStoragePath) + : suppliedStoragePath + ? candidate.storage_path === suppliedStoragePath + : false, + ); + if (!artifact) { + throw new Error("Artifacts are unavailable for this task run"); + } + + const key = artifact.id ?? artifact.storage_path; + if (key && !resolvedKeys.has(key)) { + resolvedKeys.add(key); + resolvedArtifacts.push(artifact); + } + } + + return resolvedArtifacts; +} diff --git a/products/desktop/packages/shared/src/domain-types.ts b/products/desktop/packages/shared/src/domain-types.ts index 414227afb114..4efe1c79b180 100644 --- a/products/desktop/packages/shared/src/domain-types.ts +++ b/products/desktop/packages/shared/src/domain-types.ts @@ -288,7 +288,7 @@ export type ArtifactSource = | "user_attachment" | "posthog_code_skill"; -export interface TaskRunArtifactMetadata { +export interface TaskRunSkillBundleMetadata { skill_name: string; skill_source: UploadableSkillSource; content_sha256: string; @@ -296,6 +296,16 @@ export interface TaskRunArtifactMetadata { schema_version: number; } +export interface TaskRunImageAttachmentMetadata { + conversation_id: string; + width: number; + height: number; +} + +export type TaskRunArtifactMetadata = + | TaskRunSkillBundleMetadata + | TaskRunImageAttachmentMetadata; + export interface TaskRunArtifact { id?: string; name: string; diff --git a/products/posthog_ai/backend/api/__init__.py b/products/posthog_ai/backend/api/__init__.py index b96e41966634..8380c79fed5c 100644 --- a/products/posthog_ai/backend/api/__init__.py +++ b/products/posthog_ai/backend/api/__init__.py @@ -1,3 +1,4 @@ +from .attachments import AssistantAttachmentsViewSet from .mcp_tools import MCPToolsViewSet -__all__ = ["MCPToolsViewSet"] +__all__ = ["AssistantAttachmentsViewSet", "MCPToolsViewSet"] diff --git a/products/posthog_ai/backend/api/attachments.py b/products/posthog_ai/backend/api/attachments.py new file mode 100644 index 000000000000..c5a0ee00eee9 --- /dev/null +++ b/products/posthog_ai/backend/api/attachments.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import uuid +from typing import Any, cast + +from drf_spectacular.utils import OpenApiResponse +from rest_framework import serializers +from rest_framework.decorators import action +from rest_framework.exceptions import NotFound, PermissionDenied, ValidationError +from rest_framework.request import Request +from rest_framework.response import Response +from rest_framework.viewsets import GenericViewSet + +from posthog.api.mixins import ValidatedRequest, validated_request +from posthog.api.routing import TeamAndOrgViewSetMixin +from posthog.models.user import User + +from products.posthog_ai.backend.attachments import ( + MAX_ATTACHMENT_BYTES, + MAX_ATTACHMENT_COUNT, + AttachmentNotFoundError, + AttachmentStorageError, + AttachmentValidationError, + delete_attachment, + finalize_attachments, + prepare_attachments, +) +from products.posthog_ai.backend.models.assistant import Conversation + +from ee.hogai.utils.feature_flags import has_sandbox_mode_feature_flag + + +class AttachmentPrepareItemSerializer(serializers.Serializer): + file_name = serializers.CharField(max_length=255, help_text="File name to associate with the uploaded image.") + size = serializers.IntegerField( + min_value=1, + max_value=MAX_ATTACHMENT_BYTES, + help_text=f"Expected upload size in bytes. Each image must be {MAX_ATTACHMENT_BYTES} bytes or smaller.", + ) + content_type = serializers.ChoiceField( + choices=["image/png", "image/jpeg"], + help_text="Exact MIME type for the direct upload. Only PNG and JPEG are supported.", + ) + + +class AttachmentPrepareRequestSerializer(serializers.Serializer): + conversation_id = serializers.UUIDField(help_text="Conversation UUID the staged image attachments belong to.") + attachments = serializers.ListField( + child=AttachmentPrepareItemSerializer(), + min_length=1, + max_length=MAX_ATTACHMENT_COUNT, + help_text="Images to stage for the next sandbox message.", + ) + + +class AttachmentPrepareResponseItemSerializer(serializers.Serializer): + id = serializers.UUIDField(help_text="Opaque attachment ID for finalize, delete, and sandbox message sends.") + file_name = serializers.CharField(help_text="Sanitized file name recorded for the image attachment.") + content_type = serializers.ChoiceField( + choices=["image/png", "image/jpeg"], + help_text="Signed MIME type for the upload.", + ) + size = serializers.IntegerField(help_text="Expected upload size in bytes.") + upload_url = serializers.URLField(help_text="Signed direct-upload URL for the image attachment.") + upload_fields = serializers.DictField( + child=serializers.CharField(), + help_text="Signed S3-compatible form fields to include with the upload request.", + ) + + +class AttachmentPrepareResponseSerializer(serializers.Serializer): + attachments = AttachmentPrepareResponseItemSerializer(many=True, help_text="Prepared image uploads.") + + +class AttachmentFinalizeRequestSerializer(serializers.Serializer): + conversation_id = serializers.UUIDField(help_text="Conversation UUID the uploaded image attachment belongs to.") + attachment_id = serializers.UUIDField(help_text="Prepared attachment ID to validate and finalize for sandbox use.") + + +class AttachmentFinalizeResponseSerializer(serializers.Serializer): + id = serializers.UUIDField(help_text="Opaque attachment ID for sandbox message sends.") + file_name = serializers.CharField(help_text="Sanitized file name recorded for the image attachment.") + content_type = serializers.ChoiceField( + choices=["image/png", "image/jpeg"], + help_text="Validated MIME type for the normalized image.", + ) + size = serializers.IntegerField(help_text="Normalized image size in bytes.") + width = serializers.IntegerField(help_text="Normalized image width in pixels.") + height = serializers.IntegerField(help_text="Normalized image height in pixels.") + + +class AttachmentDeleteRequestSerializer(serializers.Serializer): + conversation_id = serializers.UUIDField(help_text="Conversation UUID the image attachment belongs to.") + attachment_id = serializers.UUIDField(help_text="Single attachment ID to delete from staging.") + + +class AssistantAttachmentsViewSet(TeamAndOrgViewSetMixin, GenericViewSet): + scope_object = "conversation" + serializer_class = AttachmentPrepareRequestSerializer + + def _resolve_conversation(self, request: Request, conversation_id: uuid.UUID) -> None: + user = cast(User, request.user) + conversation = ( + Conversation.objects.filter(id=conversation_id) + .only("team_id", "user_id", "deleted", "agent_runtime") + .first() + ) + if conversation is None: + if not has_sandbox_mode_feature_flag(self.team, user): + raise ValidationError("This conversation is not on the sandbox runtime.") + return + if conversation.deleted: + raise NotFound("Conversation not found.") + if conversation.team_id != self.team.id or conversation.user_id != user.id: + raise PermissionDenied("Cannot access other users' conversations") + if conversation.agent_runtime != Conversation.AgentRuntime.SANDBOX: + raise ValidationError("This conversation is not on the sandbox runtime.") + + @validated_request( + request_serializer=AttachmentPrepareRequestSerializer, + responses={ + 200: OpenApiResponse(response=AttachmentPrepareResponseSerializer, description="Prepared image uploads."), + 400: OpenApiResponse(description="Invalid image attachment payload."), + }, + operation_id="assistantAttachmentsPrepareCreate", + summary="Prepare sandbox image uploads", + description="Reserve direct-upload image attachment slots for a sandbox conversation and return signed POST fields.", + strict_request_validation=True, + ) + @action(detail=False, methods=["post"], url_path="prepare", required_scopes=["project:write"]) + def prepare(self, request: ValidatedRequest, *args: Any, **kwargs: Any) -> Response: + data = cast(dict[str, Any], request.validated_data) + conversation_id = data["conversation_id"] + self._resolve_conversation(request, conversation_id) + try: + prepared = prepare_attachments( + team_id=self.team.id, + user_id=cast(User, request.user).id, + conversation_id=conversation_id, + attachments=data["attachments"], + ) + except AttachmentValidationError as error: + raise ValidationError(str(error)) from error + except AttachmentStorageError as error: + raise ValidationError(str(error)) from error + serializer = AttachmentPrepareResponseSerializer( + { + "attachments": [ + { + "id": item.id, + "file_name": item.file_name, + "content_type": item.content_type, + "size": item.size, + "upload_url": item.upload_url, + "upload_fields": item.upload_fields, + } + for item in prepared + ] + } + ) + return Response(serializer.data) + + @validated_request( + request_serializer=AttachmentFinalizeRequestSerializer, + responses={ + 200: OpenApiResponse( + response=AttachmentFinalizeResponseSerializer, description="Finalized image attachment." + ), + 400: OpenApiResponse(description="Invalid image attachment payload."), + }, + operation_id="assistantAttachmentsFinalizeCreate", + summary="Finalize sandbox image uploads", + description="Validate, normalize, and finalize one staged image upload for sandbox messages.", + strict_request_validation=True, + ) + @action(detail=False, methods=["post"], url_path="finalize", required_scopes=["project:write"]) + def finalize(self, request: ValidatedRequest, *args: Any, **kwargs: Any) -> Response: + data = cast(dict[str, Any], request.validated_data) + conversation_id = data["conversation_id"] + self._resolve_conversation(request, conversation_id) + try: + finalized = finalize_attachments( + team_id=self.team.id, + user_id=cast(User, request.user).id, + conversation_id=conversation_id, + attachment_id=str(data["attachment_id"]), + ) + except AttachmentNotFoundError as error: + raise NotFound(str(error)) from error + except AttachmentValidationError as error: + raise ValidationError(str(error)) from error + except AttachmentStorageError as error: + raise ValidationError(str(error)) from error + serializer = AttachmentFinalizeResponseSerializer( + { + "id": finalized.id, + "file_name": finalized.file_name, + "content_type": finalized.content_type, + "size": finalized.size, + "width": finalized.width, + "height": finalized.height, + } + ) + return Response(serializer.data) + + @validated_request( + request_serializer=AttachmentDeleteRequestSerializer, + responses={204: OpenApiResponse(description="Attachment deleted.")}, + operation_id="assistantAttachmentsDeleteCreate", + summary="Delete one sandbox image attachment", + description="Idempotently delete a single staged or finalized sandbox image attachment for one conversation.", + strict_request_validation=True, + ) + @action(detail=False, methods=["post"], url_path="delete", required_scopes=["project:write"]) + def delete(self, request: ValidatedRequest, *args: Any, **kwargs: Any) -> Response: + data = cast(dict[str, Any], request.validated_data) + conversation_id = data["conversation_id"] + self._resolve_conversation(request, conversation_id) + try: + delete_attachment( + team_id=self.team.id, + user_id=cast(User, request.user).id, + conversation_id=conversation_id, + attachment_id=str(data["attachment_id"]), + ) + except AttachmentNotFoundError: + pass + return Response(status=204) diff --git a/products/posthog_ai/backend/attachments.py b/products/posthog_ai/backend/attachments.py new file mode 100644 index 000000000000..e6ff204d51b0 --- /dev/null +++ b/products/posthog_ai/backend/attachments.py @@ -0,0 +1,893 @@ +from __future__ import annotations + +import os +import time +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from datetime import timedelta +from io import BytesIO +from typing import Any +from uuid import UUID, uuid4 + +from django.utils import timezone + +import structlog +from PIL import Image, ImageOps, UnidentifiedImageError + +from posthog.dataclasses import frozen +from posthog.storage import object_storage + +from products.tasks.backend.facade import api as tasks_facade +from products.tasks.backend.redis import get_tasks_cache + +MAX_ATTACHMENT_COUNT = 4 +MAX_ATTACHMENT_BYTES = 4 * 1024 * 1024 +MAX_MESSAGE_ATTACHMENT_BYTES = 10 * 1024 * 1024 +MAX_USER_STAGED_ATTACHMENT_COUNT = 20 +MAX_USER_STAGED_ATTACHMENT_BYTES = 50 * 1024 * 1024 +MAX_DECODED_PIXELS = 20_000_000 +ATTACHMENT_CACHE_TTL_SECONDS = 24 * 60 * 60 +ATTACHMENT_UPLOAD_EXPIRATION_SECONDS = 15 * 60 +ATTACHMENT_STAGING_TTL_DAYS = "1" +SUPPORTED_CONTENT_TYPES = ("image/png", "image/jpeg") +_FORMAT_BY_CONTENT_TYPE = {"image/png": "PNG", "image/jpeg": "JPEG"} +_EXTENSION_BY_CONTENT_TYPE = {"image/png": "png", "image/jpeg": "jpg"} +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +_JPEG_SIGNATURE = b"\xff\xd8\xff" +_CACHE_PREFIX = "posthog_ai:assistant_attachment" +_LOCK_TTL_SECONDS = 3 * 60 +_LOCK_WAIT_SECONDS = 15 + +logger = structlog.get_logger(__name__) + + +class AttachmentError(Exception): + pass + + +class AttachmentValidationError(AttachmentError): + pass + + +class AttachmentNotFoundError(AttachmentError): + pass + + +class AttachmentStorageError(AttachmentError): + pass + + +@frozen +class PreparedAttachment: + id: str + file_name: str + content_type: str + size: int + upload_url: str + upload_fields: dict[str, str] + + +@frozen +class FinalizedAttachment: + id: str + file_name: str + content_type: str + size: int + width: int + height: int + + +@frozen +class AttachmentPromotion: + attachment_ids: list[str] + newly_attached_ids: list[str] + task_id: str + run_id: str + records: list[dict[str, Any]] + + +@frozen +class _PreparedAttachmentSpec: + id: str + file_name: str + content_type: str + size: int + raw_storage_path: str + upload_url: str + upload_fields: dict[str, str] + + +def _cache() -> Any: + return get_tasks_cache() + + +def _record_key(attachment_id: str) -> str: + return f"{_CACHE_PREFIX}:record:{attachment_id}" + + +def _conversation_index_key(team_id: int, user_id: int, conversation_id: UUID | str) -> str: + return f"{_CACHE_PREFIX}:conversation:{team_id}:{user_id}:{conversation_id}" + + +def _user_index_key(team_id: int, user_id: int) -> str: + return f"{_CACHE_PREFIX}:user:{team_id}:{user_id}" + + +def _lock_key(name: str) -> str: + return f"{_CACHE_PREFIX}:lock:{name}" + + +def _user_lock_key(team_id: int, user_id: int) -> str: + return _lock_key(f"user:{team_id}:{user_id}") + + +def _attachment_lock_key(attachment_id: str) -> str: + return _lock_key(f"attachment:{attachment_id}") + + +@contextmanager +def _cache_locks( + keys: Sequence[str], *, ttl: int = _LOCK_TTL_SECONDS, wait: int = _LOCK_WAIT_SECONDS +) -> Iterator[None]: + cache = _cache() + ordered_keys = list(dict.fromkeys(sorted(keys))) + acquired: list[tuple[str, str]] = [] + deadline = time.monotonic() + wait + + try: + for key in ordered_keys: + token = str(uuid4()) + while True: + if cache.add(key, token, timeout=ttl): + acquired.append((key, token)) + break + if time.monotonic() >= deadline: + raise AttachmentStorageError("Could not acquire the attachment lock. Try again.") + time.sleep(0.05) + yield + finally: + for key, token in reversed(acquired): + try: + if cache.get(key) == token: + cache.delete(key) + except Exception: + logger.warning("assistant_attachment_lock_release_failed", lock_key=key) + + +def _safe_name(name: str, content_type: str) -> str: + base_name = os.path.basename(name).strip() or "image" + stem = os.path.splitext(base_name)[0].strip() or "image" + extension = _EXTENSION_BY_CONTENT_TYPE[content_type] + return f"{stem[:240]}.{extension}" + + +def _raw_storage_path(team_id: int, user_id: int, conversation_id: UUID, attachment_id: str) -> str: + return ( + f"posthog_ai/assistant_attachments/team_{team_id}/user_{user_id}/" + f"conversation_{conversation_id}/uploads/{attachment_id}" + ) + + +def _normalized_storage_path( + team_id: int, user_id: int, conversation_id: UUID, attachment_id: str, content_type: str +) -> str: + extension = _EXTENSION_BY_CONTENT_TYPE[content_type] + return ( + f"posthog_ai/assistant_attachments/team_{team_id}/user_{user_id}/" + f"conversation_{conversation_id}/normalized/{attachment_id}.{extension}" + ) + + +def _get_index(key: str) -> list[str]: + value = _cache().get(key) + if not isinstance(value, list): + return [] + return [item for item in value if isinstance(item, str)] + + +def _get_record(attachment_id: str) -> dict[str, Any] | None: + value = _cache().get(_record_key(attachment_id)) + return value if isinstance(value, dict) else None + + +def _write_records_and_indexes( + *, + records: Sequence[dict[str, Any]], + team_id: int, + user_id: int, + conversation_id: UUID, + conversation_ids: list[str], + user_ids: list[str], +) -> None: + conversation_key = _conversation_index_key(team_id, user_id, conversation_id) + user_key = _user_index_key(team_id, user_id) + values = { + conversation_key: conversation_ids, + user_key: user_ids, + **{_record_key(str(record["id"])): record for record in records}, + } + _cache().set_many(values, timeout=ATTACHMENT_CACHE_TTL_SECONDS) + + +def _active_records(index_ids: list[str]) -> tuple[list[dict[str, Any]], list[str]]: + records: list[dict[str, Any]] = [] + active_ids: list[str] = [] + for attachment_id in index_ids: + record = _get_record(attachment_id) + if record is None or record.get("status") not in {"prepared", "finalized"}: + continue + records.append(record) + active_ids.append(attachment_id) + return records, active_ids + + +def _quota_records(index_ids: list[str]) -> tuple[list[dict[str, Any]], list[str]]: + records: list[dict[str, Any]] = [] + active_ids: list[str] = [] + now = timezone.now().timestamp() + for attachment_id in index_ids: + record = _get_record(attachment_id) + if record is None: + continue + status = record.get("status") + if status not in {"prepared", "finalized"} and not ( + status == "deleted" and float(record.get("upload_expires_at", 0)) > now + ): + continue + records.append(record) + active_ids.append(attachment_id) + return records, active_ids + + +def _record_quota_size(record: dict[str, Any]) -> int: + declared_size = record.get("declared_size") + normalized_size = record.get("normalized_size") + sizes = [size for size in (declared_size, normalized_size) if isinstance(size, int)] + return max(sizes, default=0) + + +def _as_finalized_record(record: dict[str, Any]) -> dict[str, Any]: + restored = {key: value for key, value in record.items() if key not in {"promoted_task_id", "promoted_run_id"}} + restored["status"] = "finalized" + return restored + + +def prepare_attachments( + *, + team_id: int, + user_id: int, + conversation_id: UUID, + attachments: Sequence[dict[str, Any]], +) -> list[PreparedAttachment]: + if not attachments or len(attachments) > MAX_ATTACHMENT_COUNT: + raise AttachmentValidationError("Attach between 1 and 4 images.") + + requested_bytes = sum(int(attachment["size"]) for attachment in attachments) + if requested_bytes > MAX_MESSAGE_ATTACHMENT_BYTES: + raise AttachmentValidationError("Images must be 10 MiB or smaller in total.") + + tagging = f"ttl_days={ATTACHMENT_STAGING_TTL_DAYS}&team_id={team_id}" + upload_expires_at = (timezone.now() + timedelta(seconds=ATTACHMENT_UPLOAD_EXPIRATION_SECONDS)).timestamp() + prepared_specs: list[_PreparedAttachmentSpec] = [] + for attachment in attachments: + content_type = str(attachment["content_type"]) + if content_type not in SUPPORTED_CONTENT_TYPES: + raise AttachmentValidationError("Only PNG and JPEG images are supported.") + attachment_id = str(uuid4()) + raw_path = _raw_storage_path(team_id, user_id, conversation_id, attachment_id) + presigned_post = object_storage.get_presigned_post( + raw_path, + conditions=[ + ["content-length-range", int(attachment["size"]), int(attachment["size"])], + {"Content-Type": content_type}, + {"x-amz-tagging": tagging}, + ], + expiration=ATTACHMENT_UPLOAD_EXPIRATION_SECONDS, + ) + if not presigned_post: + raise AttachmentStorageError("Could not prepare the image upload. Try again.") + fields = {str(key): str(value) for key, value in dict(presigned_post.get("fields") or {}).items()} + fields["Content-Type"] = content_type + fields["x-amz-tagging"] = tagging + prepared_specs.append( + _PreparedAttachmentSpec( + id=attachment_id, + file_name=_safe_name(str(attachment["file_name"]), content_type), + content_type=content_type, + size=int(attachment["size"]), + raw_storage_path=raw_path, + upload_url=str(presigned_post["url"]), + upload_fields=fields, + ) + ) + + with _cache_locks([_user_lock_key(team_id, user_id)]): + user_key = _user_index_key(team_id, user_id) + conversation_key = _conversation_index_key(team_id, user_id, conversation_id) + user_records, active_user_ids = _quota_records(_get_index(user_key)) + conversation_records, active_conversation_ids = _active_records(_get_index(conversation_key)) + + if len(active_conversation_ids) + len(prepared_specs) > MAX_ATTACHMENT_COUNT: + raise AttachmentValidationError("You can attach up to 4 images to one message.") + if ( + sum(_record_quota_size(record) for record in conversation_records) + requested_bytes + > MAX_MESSAGE_ATTACHMENT_BYTES + ): + raise AttachmentValidationError("Images must be 10 MiB or smaller in total.") + if len(active_user_ids) + len(prepared_specs) > MAX_USER_STAGED_ATTACHMENT_COUNT: + raise AttachmentValidationError("You can have up to 20 images waiting to be sent.") + if ( + sum(_record_quota_size(record) for record in user_records) + requested_bytes + > MAX_USER_STAGED_ATTACHMENT_BYTES + ): + raise AttachmentValidationError("Images waiting to be sent must be 50 MiB or smaller in total.") + + records = [ + { + "id": spec.id, + "team_id": team_id, + "user_id": user_id, + "conversation_id": str(conversation_id), + "file_name": spec.file_name, + "content_type": spec.content_type, + "declared_size": spec.size, + "raw_storage_path": spec.raw_storage_path, + "status": "prepared", + "created_at": timezone.now().isoformat(), + "upload_expires_at": upload_expires_at, + } + for spec in prepared_specs + ] + _write_records_and_indexes( + records=records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=[*active_conversation_ids, *(str(record["id"]) for record in records)], + user_ids=[*active_user_ids, *(str(record["id"]) for record in records)], + ) + + return [ + PreparedAttachment( + id=spec.id, + file_name=spec.file_name, + content_type=spec.content_type, + size=spec.size, + upload_url=spec.upload_url, + upload_fields=spec.upload_fields, + ) + for spec in prepared_specs + ] + + +def _validate_and_normalize_image(content: bytes, content_type: str) -> tuple[bytes, int, int]: + if content_type == "image/png" and not content.startswith(_PNG_SIGNATURE): + raise AttachmentValidationError("The image contents do not match the selected PNG type.") + if content_type == "image/jpeg" and not content.startswith(_JPEG_SIGNATURE): + raise AttachmentValidationError("The image contents do not match the selected JPEG type.") + + try: + with Image.open(BytesIO(content)) as image: + if image.format != _FORMAT_BY_CONTENT_TYPE[content_type]: + raise AttachmentValidationError("The image contents do not match the selected file type.") + if bool(getattr(image, "is_animated", False)) or int(getattr(image, "n_frames", 1)) != 1: + raise AttachmentValidationError("Animated images are not supported.") + width, height = image.size + if width <= 0 or height <= 0 or width * height >= MAX_DECODED_PIXELS: + raise AttachmentValidationError("The image must contain fewer than 20 million pixels.") + image.verify() + + with Image.open(BytesIO(content)) as image: + image.load() + transposed = ImageOps.exif_transpose(image) + width, height = transposed.size + if width * height >= MAX_DECODED_PIXELS: + raise AttachmentValidationError("The image must contain fewer than 20 million pixels.") + output = BytesIO() + if content_type == "image/jpeg": + clean = transposed.convert("RGB") + clean.save(output, format="JPEG", quality=90, optimize=True) + else: + has_transparency = "A" in transposed.getbands() or "transparency" in transposed.info + clean = transposed.convert("RGBA" if has_transparency else "RGB") + clean.save(output, format="PNG", optimize=True) + except AttachmentValidationError: + raise + except (Image.DecompressionBombError, UnidentifiedImageError, OSError, ValueError) as error: + raise AttachmentValidationError("The uploaded file is not a valid PNG or JPEG image.") from error + + normalized = output.getvalue() + if len(normalized) > MAX_ATTACHMENT_BYTES: + raise AttachmentValidationError("The normalized image is larger than 4 MiB.") + return normalized, width, height + + +def _scoped_records( + *, team_id: int, user_id: int, conversation_id: UUID, attachment_ids: Sequence[str] +) -> list[dict[str, Any]]: + if ( + not attachment_ids + or len(attachment_ids) > MAX_ATTACHMENT_COUNT + or len(set(attachment_ids)) != len(attachment_ids) + ): + raise AttachmentValidationError("Provide between 1 and 4 unique attachment IDs.") + records: list[dict[str, Any]] = [] + for attachment_id in attachment_ids: + record = _get_record(attachment_id) + if ( + record is None + or record.get("team_id") != team_id + or record.get("user_id") != user_id + or record.get("conversation_id") != str(conversation_id) + ): + raise AttachmentNotFoundError("One or more image attachments were not found.") + records.append(record) + return records + + +def _delete_objects(paths: Sequence[str]) -> None: + for path in paths: + try: + object_storage.delete(path) + except Exception: + pass + + +def _remove_records_locked( + *, team_id: int, user_id: int, conversation_id: UUID, records: Sequence[dict[str, Any]] +) -> None: + removed_ids = {str(record["id"]) for record in records} + conversation_key = _conversation_index_key(team_id, user_id, conversation_id) + user_key = _user_index_key(team_id, user_id) + conversation_ids = [item for item in _get_index(conversation_key) if item not in removed_ids] + user_ids = [item for item in _get_index(user_key) if item not in removed_ids] + now = timezone.now().timestamp() + tombstones = [ + {**record, "status": "deleted"} for record in records if float(record.get("upload_expires_at", 0)) > now + ] + user_ids.extend(str(record["id"]) for record in tombstones if str(record["id"]) not in user_ids) + _write_records_and_indexes( + records=tombstones, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=conversation_ids, + user_ids=user_ids, + ) + tombstone_ids = {str(record["id"]) for record in tombstones} + expired_ids = removed_ids - tombstone_ids + if expired_ids: + _cache().delete_many([_record_key(attachment_id) for attachment_id in expired_ids]) + + +def finalize_attachments( + *, team_id: int, user_id: int, conversation_id: UUID, attachment_id: str +) -> FinalizedAttachment: + lock_keys = [_user_lock_key(team_id, user_id), _attachment_lock_key(attachment_id)] + with _cache_locks(lock_keys): + record = _scoped_records( + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + attachment_ids=[attachment_id], + )[0] + if record.get("status") == "finalized": + return FinalizedAttachment( + id=str(record["id"]), + file_name=str(record["file_name"]), + content_type=str(record["content_type"]), + size=int(record["normalized_size"]), + width=int(record["width"]), + height=int(record["height"]), + ) + if record.get("status") != "prepared": + raise AttachmentValidationError("This image attachment cannot be finalized.") + if float(record.get("upload_expires_at", 0)) <= timezone.now().timestamp(): + _remove_records_locked(team_id=team_id, user_id=user_id, conversation_id=conversation_id, records=[record]) + raise AttachmentValidationError("This image upload has expired. Upload it again.") + + normalized_paths: list[str] = [] + try: + raw_path = str(record["raw_storage_path"]) + stored_object = object_storage.read_object(raw_path, missing_ok=True) + if stored_object is None: + raise AttachmentValidationError("An uploaded image could not be found. Upload it again.") + content, stored_content_type = stored_object + content_length = len(content) + if content_length != record["declared_size"]: + raise AttachmentValidationError("The uploaded image size does not match the signed upload size.") + if stored_content_type != record["content_type"]: + raise AttachmentValidationError("The uploaded image type does not match the signed upload type.") + normalized, width, height = _validate_and_normalize_image(content, str(record["content_type"])) + normalized_path = _normalized_storage_path( + team_id, user_id, conversation_id, str(record["id"]), str(record["content_type"]) + ) + object_storage.write(normalized_path, normalized, extras={"ContentType": str(record["content_type"])}) + normalized_paths.append(normalized_path) + object_storage.tag( + normalized_path, + {"ttl_days": ATTACHMENT_STAGING_TTL_DAYS, "team_id": str(team_id)}, + ) + except AttachmentValidationError: + _delete_objects([str(record["raw_storage_path"]), *normalized_paths]) + _remove_records_locked(team_id=team_id, user_id=user_id, conversation_id=conversation_id, records=[record]) + raise + except Exception as error: + _delete_objects(normalized_paths) + raise AttachmentStorageError("Could not process the image. Try uploading it again.") from error + + current_record = _scoped_records( + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + attachment_ids=[attachment_id], + )[0] + if current_record.get("status") == "finalized": + return FinalizedAttachment( + id=str(current_record["id"]), + file_name=str(current_record["file_name"]), + content_type=str(current_record["content_type"]), + size=int(current_record["normalized_size"]), + width=int(current_record["width"]), + height=int(current_record["height"]), + ) + if current_record.get("status") != "prepared": + _delete_objects(normalized_paths) + raise AttachmentValidationError("This image attachment changed while finalizing.") + + user_records, active_user_ids = _quota_records(_get_index(_user_index_key(team_id, user_id))) + conversation_records, active_conversation_ids = _active_records( + _get_index(_conversation_index_key(team_id, user_id, conversation_id)) + ) + normalized_size = len(normalized) + other_user_bytes = sum(_record_quota_size(r) for r in user_records if str(r["id"]) != attachment_id) + other_conversation_bytes = sum( + _record_quota_size(r) for r in conversation_records if str(r["id"]) != attachment_id + ) + if other_conversation_bytes + normalized_size > MAX_MESSAGE_ATTACHMENT_BYTES: + _delete_objects(normalized_paths) + raise AttachmentValidationError("Images must be 10 MiB or smaller in total after processing.") + if other_user_bytes + normalized_size > MAX_USER_STAGED_ATTACHMENT_BYTES: + _delete_objects(normalized_paths) + raise AttachmentValidationError( + "Images waiting to be sent must be 50 MiB or smaller in total after processing." + ) + + updated_record = { + **current_record, + "status": "finalized", + "normalized_storage_path": normalized_path, + "normalized_size": normalized_size, + "width": width, + "height": height, + } + _write_records_and_indexes( + records=[updated_record], + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=active_conversation_ids, + user_ids=active_user_ids, + ) + + _delete_objects([str(record["raw_storage_path"])]) + return FinalizedAttachment( + id=str(updated_record["id"]), + file_name=str(updated_record["file_name"]), + content_type=str(updated_record["content_type"]), + size=int(updated_record["normalized_size"]), + width=int(updated_record["width"]), + height=int(updated_record["height"]), + ) + + +def delete_attachment(*, team_id: int, user_id: int, conversation_id: UUID, attachment_id: str) -> None: + lock_keys = [_user_lock_key(team_id, user_id), _attachment_lock_key(attachment_id)] + with _cache_locks(lock_keys): + records = _scoped_records( + team_id=team_id, user_id=user_id, conversation_id=conversation_id, attachment_ids=[attachment_id] + ) + record = records[0] + paths = [str(record["raw_storage_path"])] + if record.get("normalized_storage_path"): + paths.append(str(record["normalized_storage_path"])) + _delete_objects(paths) + _remove_records_locked(team_id=team_id, user_id=user_id, conversation_id=conversation_id, records=records) + + +def validate_new_run_attachments( + *, team_id: int, user_id: int, conversation_id: UUID, attachment_ids: Sequence[str] +) -> None: + if not attachment_ids: + return + + lock_keys = [ + _user_lock_key(team_id, user_id), + *(_attachment_lock_key(attachment_id) for attachment_id in attachment_ids), + ] + with _cache_locks(lock_keys): + records = _scoped_records( + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + attachment_ids=attachment_ids, + ) + if any(record.get("status") == "promoted" for record in records): + raise AttachmentValidationError("An image attachment has already been sent in another message.") + if any( + record.get("status") not in {"finalized", "promotion_pending", "rollback_pending"} for record in records + ): + raise AttachmentValidationError("Finalize every image before sending the message.") + + +def promote_attachments( + *, + team_id: int, + user_id: int, + conversation_id: UUID, + task_id: str, + run_id: str, + attachment_ids: Sequence[str], +) -> AttachmentPromotion | None: + if not attachment_ids: + return None + + lock_keys = [ + _user_lock_key(team_id, user_id), + *(_attachment_lock_key(attachment_id) for attachment_id in attachment_ids), + ] + with _cache_locks(lock_keys): + records = _scoped_records( + team_id=team_id, user_id=user_id, conversation_id=conversation_id, attachment_ids=attachment_ids + ) + pending_records = [ + record for record in records if record.get("status") in {"promotion_pending", "rollback_pending"} + ] + if pending_records: + pending_by_run: dict[tuple[str, str], list[str]] = {} + for record in pending_records: + pending_key = (str(record["promoted_task_id"]), str(record["promoted_run_id"])) + pending_by_run.setdefault(pending_key, []).append(str(record["id"])) + for (pending_task_id, pending_run_id), pending_ids in pending_by_run.items(): + tasks_facade.rollback_posthog_ai_attachments( + task_id=pending_task_id, + run_id=pending_run_id, + team_id=team_id, + attachment_ids=pending_ids, + ) + pending_id_set = {str(record["id"]) for record in pending_records} + records = [ + _as_finalized_record(record) if str(record["id"]) in pending_id_set else record for record in records + ] + conversation_ids = _get_index(_conversation_index_key(team_id, user_id, conversation_id)) + user_ids = _get_index(_user_index_key(team_id, user_id)) + for pending_record in pending_records: + pending_id = str(pending_record["id"]) + if pending_id not in conversation_ids: + conversation_ids.append(pending_id) + if pending_id not in user_ids: + user_ids.append(pending_id) + _write_records_and_indexes( + records=records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=conversation_ids, + user_ids=user_ids, + ) + if any(record.get("status") not in {"finalized", "promoted"} for record in records): + raise AttachmentValidationError("Finalize every image before sending the message.") + for record in records: + if record.get("status") == "promoted" and record.get("promoted_run_id") != run_id: + raise AttachmentValidationError("An image attachment has already been sent in another message.") + + finalized_records = [record for record in records if record.get("status") == "finalized"] + if not finalized_records: + return AttachmentPromotion( + attachment_ids=[str(record["id"]) for record in records], + newly_attached_ids=[], + task_id=task_id, + run_id=run_id, + records=records, + ) + + finalized_ids = {str(record["id"]) for record in finalized_records} + checkpoint_records = [ + { + **record, + "status": "promotion_pending", + "promoted_task_id": task_id, + "promoted_run_id": run_id, + } + for record in finalized_records + ] + conversation_ids = [ + item + for item in _get_index(_conversation_index_key(team_id, user_id, conversation_id)) + if item not in finalized_ids + ] + user_ids = [item for item in _get_index(_user_index_key(team_id, user_id)) if item not in finalized_ids] + _write_records_and_indexes( + records=checkpoint_records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=conversation_ids, + user_ids=user_ids, + ) + + try: + result = tasks_facade.promote_posthog_ai_attachments( + task_id=task_id, + run_id=run_id, + team_id=team_id, + user_id=user_id, + conversation_id=str(conversation_id), + attachments=records, + ) + if result is None: + raise AttachmentNotFoundError("The sandbox run for these image attachments was not found.") + except Exception: + try: + restored_conversation_ids = _get_index(_conversation_index_key(team_id, user_id, conversation_id)) + restored_user_ids = _get_index(_user_index_key(team_id, user_id)) + for finalized_id in finalized_ids: + if finalized_id not in restored_conversation_ids: + restored_conversation_ids.append(finalized_id) + if finalized_id not in restored_user_ids: + restored_user_ids.append(finalized_id) + _write_records_and_indexes( + records=finalized_records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=restored_conversation_ids, + user_ids=restored_user_ids, + ) + except Exception as error: + logger.warning( + "assistant_attachment_promotion_checkpoint_restore_failed", + task_id=task_id, + run_id=run_id, + error=str(error), + ) + raise + + promoted_records = [ + {**record, "status": "promoted", "promoted_task_id": task_id, "promoted_run_id": run_id} + for record in records + ] + try: + _write_records_and_indexes( + records=promoted_records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=conversation_ids, + user_ids=user_ids, + ) + except Exception: + try: + tasks_facade.rollback_posthog_ai_attachments( + task_id=task_id, + run_id=run_id, + team_id=team_id, + attachment_ids=result.newly_attached_ids, + ) + except Exception as error: + logger.warning( + "assistant_attachment_promotion_rollback_failed", + task_id=task_id, + run_id=run_id, + error=str(error), + ) + raise + + return AttachmentPromotion( + attachment_ids=result.attached_ids, + newly_attached_ids=result.newly_attached_ids, + task_id=task_id, + run_id=run_id, + records=records, + ) + + +def rollback_attachment_promotion(promotion: AttachmentPromotion | None) -> None: + if promotion is None or not promotion.newly_attached_ids: + return + + newly_attached_ids = set(promotion.newly_attached_ids) + records = [record for record in promotion.records if str(record["id"]) in newly_attached_ids] + first = records[0] + team_id = int(first["team_id"]) + user_id = int(first["user_id"]) + conversation_id = UUID(str(first["conversation_id"])) + try: + with _cache_locks( + [ + _user_lock_key(team_id, user_id), + *(_attachment_lock_key(attachment_id) for attachment_id in promotion.newly_attached_ids), + ] + ): + pending_records = [ + { + **record, + "status": "rollback_pending", + "promoted_task_id": promotion.task_id, + "promoted_run_id": promotion.run_id, + } + for record in records + ] + _write_records_and_indexes( + records=pending_records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=_get_index(_conversation_index_key(team_id, user_id, conversation_id)), + user_ids=_get_index(_user_index_key(team_id, user_id)), + ) + except Exception as error: + logger.warning( + "assistant_attachment_rollback_checkpoint_failed", + task_id=promotion.task_id, + run_id=promotion.run_id, + error=str(error), + ) + + try: + tasks_facade.rollback_posthog_ai_attachments( + task_id=promotion.task_id, + run_id=promotion.run_id, + team_id=team_id, + attachment_ids=promotion.newly_attached_ids, + ) + except Exception as error: + logger.warning( + "assistant_attachment_detach_failed", + task_id=promotion.task_id, + run_id=promotion.run_id, + error=str(error), + ) + return + + try: + with _cache_locks( + [ + _user_lock_key(team_id, user_id), + *(_attachment_lock_key(attachment_id) for attachment_id in promotion.newly_attached_ids), + ] + ): + conversation_ids = _get_index(_conversation_index_key(team_id, user_id, conversation_id)) + user_ids = _get_index(_user_index_key(team_id, user_id)) + for attachment_id in promotion.newly_attached_ids: + if attachment_id not in conversation_ids: + conversation_ids.append(attachment_id) + if attachment_id not in user_ids: + user_ids.append(attachment_id) + _write_records_and_indexes( + records=records, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + conversation_ids=conversation_ids, + user_ids=user_ids, + ) + except Exception as error: + logger.warning( + "assistant_attachment_cache_restore_failed", + task_id=promotion.task_id, + run_id=promotion.run_id, + error=str(error), + ) + + +def commit_attachment_promotion(promotion: AttachmentPromotion | None) -> None: + if promotion is None or not promotion.newly_attached_ids: + return + first = promotion.records[0] + tasks_facade.commit_posthog_ai_attachments( + task_id=promotion.task_id, + run_id=promotion.run_id, + team_id=int(first["team_id"]), + attachment_ids=promotion.newly_attached_ids, + ) diff --git a/products/posthog_ai/backend/message_routing.py b/products/posthog_ai/backend/message_routing.py index 9ddccf1277d7..641fa2acf788 100644 --- a/products/posthog_ai/backend/message_routing.py +++ b/products/posthog_ai/backend/message_routing.py @@ -23,6 +23,15 @@ from posthog.models.user import User from posthog.storage import object_storage +from products.posthog_ai.backend.attachments import ( + AttachmentNotFoundError, + AttachmentPromotion, + AttachmentValidationError, + commit_attachment_promotion, + promote_attachments, + rollback_attachment_promotion, + validate_new_run_attachments, +) from products.posthog_ai.backend.context_wrapper import ( ALLOWED_TYPES, MAX_ATTACHED_ITEMS, @@ -128,20 +137,22 @@ def open( ) -> SandboxRouteResult | None: initial_permission_mode = self._initial_permission_mode(data.get("initial_permission_mode")) content = data.get("content") - if not isinstance(content, str) or not content.strip(): - # No message — warm intent: boot a Run that idles awaiting the first `user_message`. - # Returns the warm handle, or None when the pool is full and nothing was provisioned. - return self._warm(trace_id=data.get("trace_id"), initial_permission_mode=initial_permission_mode) + trace_value = data.get("trace_id") + trace_id = str(trace_value) if trace_value is not None else None + attachment_ids = [str(attachment_id) for attachment_id in (data.get("attachment_ids") or [])] + if attachment_ids and not (isinstance(content, str) and content.strip()): + raise exceptions.ValidationError("Attachment IDs require a message send.") + has_message = isinstance(content, str) and bool(content.strip()) + if not has_message: + return self._warm(trace_id=trace_id, initial_permission_mode=initial_permission_mode) - trace_id = data.get("trace_id") attached_context = self._validate_attached_context(data.get("attached_context")) + normalized_content = content if isinstance(content, str) else "" if self.conversation.task_id is None: - # `resumed_context` / `convert_to_acp` only apply to the conversion event, which is - # always a first message (the gate requires `task_id is None`). `repository` is the - # auto-routed repo for this first message — followups/resumes reuse the existing Task. return self._handle_first_message( - content=content, + content=normalized_content, + attachment_ids=attachment_ids, trace_id=trace_id, attached_context=attached_context, initial_permission_mode=initial_permission_mode, @@ -160,12 +171,13 @@ def open( context_service = ContextService() prior_seen = self._collect_seen_entity_refs(current_run) deduped = context_service.prune_repeated_entity_refs(attached_context, prior=prior_seen) - wrapped = context_service.wrap_user_message(content, deduped) + wrapped = context_service.wrap_user_message(normalized_content, deduped) if current_run.status in self._IN_PROGRESS_STATUSES: return self._handle_in_progress_followup( run=current_run, wrapped=wrapped, + attachment_ids=attachment_ids, trace_id=trace_id, attached_context=attached_context, ) @@ -173,6 +185,7 @@ def open( return self._handle_terminal_resume( run=current_run, wrapped=wrapped, + attachment_ids=attachment_ids, trace_id=trace_id, attached_context=attached_context, initial_permission_mode=initial_permission_mode, @@ -256,6 +269,7 @@ def _handle_first_message( self, *, content: str, + attachment_ids: list[str], trace_id: str | None, attached_context: list[AttachedContext], initial_permission_mode: InitialPermissionMode, @@ -272,12 +286,25 @@ def _handle_first_message( # sandbox agent has continuity, then the user's own attachments + message. wrapped = f"{resumed_context}\n\n{wrapped}" + if attachment_ids: + try: + validate_new_run_attachments( + team_id=self.team.id, + user_id=self.user.pk, + conversation_id=self.conversation.id, + attachment_ids=attachment_ids, + ) + except (AttachmentNotFoundError, AttachmentValidationError) as error: + raise exceptions.ValidationError(str(error)) from error + system_prompt = PromptService(self.team, self.user).build() + task_title = content[:80] if content.strip() else "Image attachment" + task_description = content or "Image attachment" created = tasks_facade.create_and_run_task( team=self.team, - title=content[:80], - description=content, + title=task_title, + description=task_description, origin_product=tasks_facade.TaskOriginProduct.POSTHOG_AI, user_id=self.user.pk, repository=repository, @@ -293,46 +320,47 @@ def _handle_first_message( if run_dto is None: raise exceptions.ValidationError("Failed to create sandbox task run.") - # Seed the PostHog AI per-Run state keys. `attached_context` keeps the full, - # undeduped list. These aren't `create_and_run` arguments, so merge them into the - # run state here, before the workflow starts and reads it. `exclude_unset` keeps the - # merge limited to exactly these keys — model defaults must not leak into the bag. - ph_state = PostHogAIRunState( - system_prompt=system_prompt, - attached_context=attached_context, - initial_permission_mode=initial_permission_mode, - interaction_origin=POSTHOG_AI_INTERACTION_ORIGIN, - pending_user_message=wrapped, - ) - state_updates = ph_state.model_dump(mode="json", by_alias=True, exclude_unset=True) - # Persist the enriched run state and conversation linkage together, under the row lock so a - # concurrent first message / conversion in another tab can't double-link. Re-check - # `task_id is None` inside the lock; a half-write would orphan the run (enriched state, but - # conversation.task still NULL) and the next retry would look like a fresh first message. On - # a conversion, the runtime flip to sandbox happens here too, atomically with the link. - with lock_conversation_for_followup(str(self.conversation.id), self.team.pk) as locked: - if locked.task_id is not None: - raise Conflict("This conversation was just resumed in another tab. Please try again.") - tasks_facade.update_task_run_state(run_dto.id, updates=state_updates) - locked.task_id = created.task_id - update_fields = ["task", "updated_at"] - if convert_to_acp: - locked.agent_runtime = Conversation.AgentRuntime.SANDBOX - update_fields = ["task", "agent_runtime", "updated_at"] - locked.save(update_fields=update_fields) - - # Mirror the committed writes onto the in-memory instance for the response + the rollback below. - self.conversation.task_id = created.task_id - if convert_to_acp: - self.conversation.agent_runtime = Conversation.AgentRuntime.SANDBOX - - # Start the run after the commit. `posthog_mcp_scopes="full"` mirrors the legacy - # first-message path: the agent creates insights, dashboards, and notebooks, so it - # needs write scopes (the workflow client otherwise defaults to read-only). If the - # start fails, un-link the conversation so the user's retry is a fresh first message - # rather than a follow-up onto a run that never started; on a conversion, also revert the - # runtime flip so the user is left on a clean idle LangGraph conversation. + promotion = None + linked_conversation = False try: + if attachment_ids: + try: + promotion = promote_attachments( + team_id=self.team.id, + user_id=self.user.pk, + conversation_id=self.conversation.id, + task_id=str(created.task_id), + run_id=str(run_dto.id), + attachment_ids=attachment_ids, + ) + except (AttachmentNotFoundError, AttachmentValidationError) as error: + raise exceptions.ValidationError(str(error)) from error + + ph_state = PostHogAIRunState( + system_prompt=system_prompt, + attached_context=attached_context, + initial_permission_mode=initial_permission_mode, + interaction_origin=POSTHOG_AI_INTERACTION_ORIGIN, + pending_user_message=wrapped, + pending_user_artifact_ids=promotion.attachment_ids if promotion is not None else None, + ) + state_updates = ph_state.model_dump(mode="json", by_alias=True, exclude_unset=True) + with lock_conversation_for_followup(str(self.conversation.id), self.team.pk) as locked: + if locked.task_id is not None: + raise Conflict("This conversation was just resumed in another tab. Please try again.") + tasks_facade.update_task_run_state(run_dto.id, updates=state_updates) + locked.task_id = created.task_id + update_fields = ["task", "updated_at"] + if convert_to_acp: + locked.agent_runtime = Conversation.AgentRuntime.SANDBOX + update_fields = ["task", "agent_runtime", "updated_at"] + locked.save(update_fields=update_fields) + linked_conversation = True + + self.conversation.task_id = created.task_id + if convert_to_acp: + self.conversation.agent_runtime = Conversation.AgentRuntime.SANDBOX + execute_task_processing_workflow( task_id=str(created.task_id), run_id=str(run_dto.id), @@ -341,13 +369,22 @@ def _handle_first_message( create_pr=False, posthog_mcp_scopes="full", ) + self._commit_attachment_promotion(promotion, run_id=str(run_dto.id)) except Exception: - self.conversation.task_id = None - revert_fields = ["task", "updated_at"] - if convert_to_acp: - self.conversation.agent_runtime = Conversation.AgentRuntime.LANGGRAPH - revert_fields = ["task", "agent_runtime", "updated_at"] - self.conversation.save(update_fields=revert_fields) + rollback_attachment_promotion(promotion) + tasks_facade.claim_and_fail_stale_run(run_dto.id, "Failed to start the sandbox run.") + if linked_conversation: + with lock_conversation_for_followup(str(self.conversation.id), self.team.pk) as locked: + if locked.task_id == created.task_id: + locked.task_id = None + update_fields = ["task", "updated_at"] + if convert_to_acp: + locked.agent_runtime = Conversation.AgentRuntime.LANGGRAPH + update_fields = ["task", "agent_runtime", "updated_at"] + locked.save(update_fields=update_fields) + self.conversation.task_id = None + if convert_to_acp: + self.conversation.agent_runtime = Conversation.AgentRuntime.LANGGRAPH raise return SandboxRouteResult( @@ -364,6 +401,7 @@ def _handle_in_progress_followup( *, run: "TaskRun", wrapped: str, + attachment_ids: list[str], trace_id: str | None, attached_context: list[AttachedContext], ) -> SandboxRouteResult: @@ -373,9 +411,28 @@ def _handle_in_progress_followup( The full undeduped context is recorded on the logged message's `_meta.attached_context` so the persisted ACP log stays a complete record. """ + promotion = None + if attachment_ids: + try: + promotion = promote_attachments( + team_id=self.team.id, + user_id=self.user.pk, + conversation_id=self.conversation.id, + task_id=str(run.task_id), + run_id=str(run.id), + attachment_ids=attachment_ids, + ) + except (AttachmentNotFoundError, AttachmentValidationError) as error: + raise exceptions.ValidationError(str(error)) from error try: - signal_task_followup_message(run.workflow_id, wrapped, artifact_ids=[]) + signal_task_followup_message( + run.workflow_id, + wrapped, + artifact_ids=attachment_ids, + message_id=trace_id, + ) except Exception as e: + rollback_attachment_promotion(promotion) # Status race: the run still reads in-progress but its workflow has already finished or # was terminated, so the signal can't be delivered. Surface a recoverable conflict # instead of an unhandled 500, and don't log the turn — a retry then routes cleanly @@ -388,6 +445,8 @@ def _handle_in_progress_followup( ) raise Conflict("The sandbox run is no longer accepting messages. Please try again.") from e + self._commit_attachment_promotion(promotion, run_id=str(run.id)) + # The Run has received its first human message, so it is no longer speculative — drop the # warm flag so the warm-pool cap stops counting it (it's now an active Run governed by AI # credits). Best-effort: a failure only over-counts the warm pool until the Run terminates, @@ -400,7 +459,7 @@ def _handle_in_progress_followup( try: # Persist only after the signal was accepted, so the log never records a message the # agent never received. - self._log_user_message(run, wrapped, attached_context) + self._log_user_message(run, wrapped, attached_context, attachment_ids) except Exception as e: # The agent already has the message; a log-append failure must not fail the request — # it only degrades context-dedup on the next follow-up. @@ -420,6 +479,7 @@ def _handle_terminal_resume( *, run: "TaskRun", wrapped: str, + attachment_ids: list[str], trace_id: str | None, attached_context: list[AttachedContext], initial_permission_mode: InitialPermissionMode, @@ -441,6 +501,17 @@ def _handle_terminal_resume( if task is None: raise exceptions.ValidationError("This conversation has no backing task to resume.") + if attachment_ids: + try: + validate_new_run_attachments( + team_id=self.team.id, + user_id=self.user.pk, + conversation_id=self.conversation.id, + attachment_ids=attachment_ids, + ) + except (AttachmentNotFoundError, AttachmentValidationError) as error: + raise exceptions.ValidationError(str(error)) from error + system_prompt = PromptService(self.team, self.user).build() with lock_conversation_for_followup(str(self.conversation.id), self.team.pk) as locked: @@ -466,16 +537,40 @@ def _handle_terminal_resume( new_run = task.create_run(mode="interactive", extra_state=extra_state) - # Same write scopes as the first message — the resumed agent keeps creating - # insights/dashboards/notebooks on follow-up turns. - execute_task_processing_workflow( - task_id=str(task.id), - run_id=str(new_run.id), - team_id=self.team.id, - user_id=self.user.pk, - create_pr=False, - posthog_mcp_scopes="full", - ) + promotion = None + try: + if attachment_ids: + try: + promotion = promote_attachments( + team_id=self.team.id, + user_id=self.user.pk, + conversation_id=self.conversation.id, + task_id=str(task.id), + run_id=str(new_run.id), + attachment_ids=attachment_ids, + ) + except (AttachmentNotFoundError, AttachmentValidationError) as error: + raise exceptions.ValidationError(str(error)) from error + if promotion is None: + raise exceptions.ValidationError("Failed to attach images to the sandbox run.") + tasks_facade.update_task_run_state( + new_run.id, + updates={"pending_user_artifact_ids": promotion.attachment_ids}, + ) + + execute_task_processing_workflow( + task_id=str(task.id), + run_id=str(new_run.id), + team_id=self.team.id, + user_id=self.user.pk, + create_pr=False, + posthog_mcp_scopes="full", + ) + self._commit_attachment_promotion(promotion, run_id=str(new_run.id)) + except Exception: + rollback_attachment_promotion(promotion) + tasks_facade.claim_and_fail_stale_run(new_run.id, "Failed to resume the sandbox run.") + raise return SandboxRouteResult( task_id=str(task.id), @@ -486,6 +581,12 @@ def _handle_terminal_resume( attached_context_count=len(attached_context), ) + def _commit_attachment_promotion(self, promotion: AttachmentPromotion | None, *, run_id: str) -> None: + try: + commit_attachment_promotion(promotion) + except Exception as error: + logger.warning("sandbox_attachment_delivery_commit_failed", run_id=run_id, error=str(error)) + def _collect_seen_entity_refs(self, run: "TaskRun") -> list[tuple[str, str | int]]: """Collect `(type, id)` pairs for entities already named in the conversation. @@ -532,18 +633,24 @@ def _absorb(items: Any) -> None: return seen - def _log_user_message(self, run: "TaskRun", wrapped: str, attached_context: list[AttachedContext]) -> None: + def _log_user_message( + self, + run: "TaskRun", + wrapped: str, + attached_context: list[AttachedContext], + artifact_ids: list[str], + ) -> None: """Append a `_posthog/user_message` entry to the Run's ACP log. - Records the wrapped content plus the full undeduped `attached_context` under - `_meta` so the persisted log is a complete per-message record that - `_collect_seen_entity_refs` reads on later follow-ups. + Records the wrapped content, attachment ids, and the full undeduped + `attached_context` so the persisted log is a complete per-message record. """ entry = { "notification": { "method": "_posthog/user_message", "params": { "content": wrapped, + "artifact_ids": artifact_ids, "_meta": {"attached_context": attached_context}, }, } diff --git a/products/posthog_ai/backend/routes.py b/products/posthog_ai/backend/routes.py index a68a90185a02..41341fb5fb1a 100644 --- a/products/posthog_ai/backend/routes.py +++ b/products/posthog_ai/backend/routes.py @@ -1,6 +1,6 @@ from posthog.api.routing import RouterRegistry -from products.posthog_ai.backend.api import MCPToolsViewSet +from products.posthog_ai.backend.api import AssistantAttachmentsViewSet, MCPToolsViewSet def register_routes(routers: RouterRegistry) -> None: @@ -10,3 +10,9 @@ def register_routes(routers: RouterRegistry) -> None: "project_mcp_tools", ["team_id"], ) + routers.projects.register( + r"assistant_attachments", + AssistantAttachmentsViewSet, + "project_assistant_attachments", + ["team_id"], + ) diff --git a/products/posthog_ai/backend/tests/test_attachment_api.py b/products/posthog_ai/backend/tests/test_attachment_api.py new file mode 100644 index 000000000000..2d0e7ff0237c --- /dev/null +++ b/products/posthog_ai/backend/tests/test_attachment_api.py @@ -0,0 +1,189 @@ +import uuid + +from posthog.test.base import APIBaseTest +from unittest.mock import patch + +from rest_framework import status + +from posthog.models import Team + +from products.posthog_ai.backend.attachments import AttachmentNotFoundError, FinalizedAttachment, PreparedAttachment +from products.posthog_ai.backend.models.assistant import Conversation + + +class TestAssistantAttachmentsViewSet(APIBaseTest): + def _url(self, action: str) -> str: + return f"/api/projects/{self.team.id}/assistant_attachments/{action}/" + + def _conversation(self, *, runtime: str = Conversation.AgentRuntime.SANDBOX) -> Conversation: + return Conversation.objects.create( + user=self.user, + team=self.team, + agent_runtime=runtime, + type=Conversation.Type.ASSISTANT, + ) + + def test_prepare_uses_flat_upload_contract(self) -> None: + conversation = self._conversation() + prepared = PreparedAttachment( + id=str(uuid.uuid4()), + file_name="image.png", + content_type="image/png", + size=123, + upload_url="https://example.com/upload", + upload_fields={"key": "value"}, + ) + with patch("products.posthog_ai.backend.api.attachments.prepare_attachments", return_value=[prepared]): + response = self.client.post( + self._url("prepare"), + { + "conversation_id": str(conversation.id), + "attachments": [{"file_name": "image.png", "content_type": "image/png", "size": 123}], + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual( + response.json(), + { + "attachments": [ + { + "id": prepared.id, + "file_name": "image.png", + "content_type": "image/png", + "size": 123, + "upload_url": "https://example.com/upload", + "upload_fields": {"key": "value"}, + } + ] + }, + ) + + def test_finalize_uses_single_attachment_contract(self) -> None: + conversation = self._conversation() + finalized = FinalizedAttachment( + id=str(uuid.uuid4()), + file_name="image.png", + content_type="image/png", + size=120, + width=10, + height=12, + ) + with patch( + "products.posthog_ai.backend.api.attachments.finalize_attachments", return_value=finalized + ) as m_finalize: + response = self.client.post( + self._url("finalize"), + { + "conversation_id": str(conversation.id), + "attachment_id": finalized.id, + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual( + response.json(), + { + "id": finalized.id, + "file_name": "image.png", + "content_type": "image/png", + "size": 120, + "width": 10, + "height": 12, + }, + ) + self.assertEqual(m_finalize.call_args.kwargs["attachment_id"], finalized.id) + + def test_delete_is_idempotent_when_attachment_is_already_absent(self) -> None: + conversation = self._conversation() + with patch( + "products.posthog_ai.backend.api.attachments.delete_attachment", + side_effect=AttachmentNotFoundError("Attachment not found"), + ): + response = self.client.post( + self._url("delete"), + { + "conversation_id": str(conversation.id), + "attachment_id": str(uuid.uuid4()), + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) + + def test_existing_langgraph_conversation_is_rejected(self) -> None: + conversation = self._conversation(runtime=Conversation.AgentRuntime.LANGGRAPH) + response = self.client.post( + self._url("prepare"), + { + "conversation_id": str(conversation.id), + "attachments": [{"file_name": "image.png", "content_type": "image/png", "size": 123}], + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(response.json()["detail"], "This conversation is not on the sandbox runtime.") + + def test_other_users_conversation_is_rejected_before_preparing_upload(self) -> None: + conversation = Conversation.objects.create( + user=self._create_user("other-user@posthog.com"), + team=self.team, + agent_runtime=Conversation.AgentRuntime.SANDBOX, + type=Conversation.Type.ASSISTANT, + ) + + with patch("products.posthog_ai.backend.api.attachments.prepare_attachments") as prepare: + response = self.client.post( + self._url("prepare"), + { + "conversation_id": str(conversation.id), + "attachments": [{"file_name": "image.png", "content_type": "image/png", "size": 123}], + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + prepare.assert_not_called() + + def test_other_teams_conversation_is_rejected_before_preparing_upload(self) -> None: + other_team = Team.objects.create(organization=self.organization, name="Other project") + conversation = Conversation.objects.create( + user=self.user, + team=other_team, + agent_runtime=Conversation.AgentRuntime.SANDBOX, + type=Conversation.Type.ASSISTANT, + ) + + with patch("products.posthog_ai.backend.api.attachments.prepare_attachments") as prepare: + response = self.client.post( + self._url("prepare"), + { + "conversation_id": str(conversation.id), + "attachments": [{"file_name": "image.png", "content_type": "image/png", "size": 123}], + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + prepare.assert_not_called() + + def test_deleted_conversation_is_rejected_before_preparing_upload(self) -> None: + conversation = self._conversation() + conversation.deleted = True + conversation.save(update_fields=["deleted"]) + + with patch("products.posthog_ai.backend.api.attachments.prepare_attachments") as prepare: + response = self.client.post( + self._url("prepare"), + { + "conversation_id": str(conversation.id), + "attachments": [{"file_name": "image.png", "content_type": "image/png", "size": 123}], + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + prepare.assert_not_called() diff --git a/products/posthog_ai/backend/tests/test_attachments.py b/products/posthog_ai/backend/tests/test_attachments.py new file mode 100644 index 000000000000..76ec3d701071 --- /dev/null +++ b/products/posthog_ai/backend/tests/test_attachments.py @@ -0,0 +1,615 @@ +import time +import uuid +import threading +from concurrent.futures import ThreadPoolExecutor +from io import BytesIO +from types import SimpleNamespace +from typing import Any + +from posthog.test.base import APIBaseTest +from unittest.mock import patch + +from parameterized import parameterized +from PIL import Image, PngImagePlugin + +from products.posthog_ai.backend.attachments import ( + ATTACHMENT_UPLOAD_EXPIRATION_SECONDS, + AttachmentPromotion, + AttachmentStorageError, + AttachmentValidationError, + _conversation_index_key, + _get_index, + _get_record, + _normalized_storage_path, + _raw_storage_path, + _user_index_key, + _validate_and_normalize_image, + _write_records_and_indexes, + delete_attachment, + finalize_attachments, + prepare_attachments, + promote_attachments, + rollback_attachment_promotion, +) + + +class TestAssistantAttachments(APIBaseTest): + def setUp(self) -> None: + super().setUp() + self.conversation_id = uuid.uuid4() + + def _png_bytes(self) -> bytes: + image = Image.new("RGB", (4, 4), color="red") + buffer = BytesIO() + image.save(buffer, format="PNG") + return buffer.getvalue() + + def _jpeg_bytes(self) -> bytes: + image = Image.new("RGB", (4, 4), color="red") + buffer = BytesIO() + image.save(buffer, format="JPEG") + return buffer.getvalue() + + def _seed_prepared_record(self, attachment_id: str, *, size: int) -> None: + record = { + "id": attachment_id, + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": str(self.conversation_id), + "file_name": "image.png", + "content_type": "image/png", + "declared_size": size, + "raw_storage_path": _raw_storage_path(self.team.id, self.user.id, self.conversation_id, attachment_id), + "status": "prepared", + "created_at": "2026-01-01T00:00:00+00:00", + "upload_expires_at": time.time() + ATTACHMENT_UPLOAD_EXPIRATION_SECONDS, + } + _write_records_and_indexes( + records=[record], + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + conversation_ids=[attachment_id], + user_ids=[attachment_id], + ) + + def test_prepare_attachments_uses_fifteen_minute_expiry(self) -> None: + with patch("products.posthog_ai.backend.attachments.object_storage.get_presigned_post") as m_post: + m_post.return_value = {"url": "https://example.com/upload", "fields": {"policy": "abc"}} + prepared = prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachments=[{"file_name": "image.png", "content_type": "image/png", "size": 10}], + ) + + self.assertEqual(m_post.call_args.kwargs["expiration"], ATTACHMENT_UPLOAD_EXPIRATION_SECONDS) + self.assertIn(["content-length-range", 10, 10], m_post.call_args.kwargs["conditions"]) + self.assertEqual(ATTACHMENT_UPLOAD_EXPIRATION_SECONDS, 15 * 60) + self.assertEqual(prepared[0].upload_fields["policy"], "abc") + self.assertEqual(prepared[0].upload_fields["Content-Type"], "image/png") + self.assertEqual( + prepared[0].upload_fields["x-amz-tagging"], + f"ttl_days=1&team_id={self.team.id}", + ) + + def test_prepare_enforces_per_user_staging_count_across_conversations(self) -> None: + upload = {"file_name": "image.png", "content_type": "image/png", "size": 1} + presigned_post = {"url": "https://example.com/upload", "fields": {"policy": "abc"}} + + with patch( + "products.posthog_ai.backend.attachments.object_storage.get_presigned_post", + return_value=presigned_post, + ): + for _ in range(5): + prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=uuid.uuid4(), + attachments=[upload] * 4, + ) + + with self.assertRaisesRegex(AttachmentValidationError, "up to 20 images"): + prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=uuid.uuid4(), + attachments=[upload], + ) + + def test_prepare_enforces_per_user_staging_bytes_across_conversations(self) -> None: + four_mib_upload = { + "file_name": "image.png", + "content_type": "image/png", + "size": 4 * 1024 * 1024, + } + three_mib_upload = { + "file_name": "image.png", + "content_type": "image/png", + "size": 3 * 1024 * 1024, + } + presigned_post = {"url": "https://example.com/upload", "fields": {"policy": "abc"}} + + with patch( + "products.posthog_ai.backend.attachments.object_storage.get_presigned_post", + return_value=presigned_post, + ): + for _ in range(6): + prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=uuid.uuid4(), + attachments=[four_mib_upload] * 2, + ) + + with self.assertRaisesRegex(AttachmentValidationError, "50 MiB"): + prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=uuid.uuid4(), + attachments=[three_mib_upload], + ) + + def test_deleted_upload_forms_still_count_toward_user_quota_until_expiry(self) -> None: + upload = {"file_name": "image.png", "content_type": "image/png", "size": 1} + presigned_post = {"url": "https://example.com/upload", "fields": {"policy": "abc"}} + + with ( + patch( + "products.posthog_ai.backend.attachments.object_storage.get_presigned_post", + return_value=presigned_post, + ), + patch("products.posthog_ai.backend.attachments.object_storage.delete"), + ): + for _ in range(20): + conversation_id = uuid.uuid4() + prepared = prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[upload], + ) + delete_attachment( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachment_id=prepared[0].id, + ) + + with self.assertRaisesRegex(AttachmentValidationError, "up to 20 images"): + prepare_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=uuid.uuid4(), + attachments=[upload], + ) + + def test_finalize_rejects_upload_with_different_size_than_signed(self) -> None: + content = self._png_bytes() + attachment_id = "att-size-mismatch" + self._seed_prepared_record(attachment_id, size=len(content) + 1) + + with ( + patch( + "products.posthog_ai.backend.attachments.object_storage.read_object", + return_value=(content, "image/png"), + ), + patch("products.posthog_ai.backend.attachments.object_storage.delete"), + ): + with self.assertRaisesRegex(AttachmentValidationError, "signed upload size"): + finalize_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachment_id=attachment_id, + ) + + def test_finalize_waits_and_returns_idempotently(self) -> None: + content = self._png_bytes() + attachment_id = "att-1" + self._seed_prepared_record(attachment_id, size=len(content)) + + write_started = threading.Event() + write_calls = 0 + write_lock = threading.Lock() + + def slow_write(*args, **kwargs): + nonlocal write_calls + with write_lock: + write_calls += 1 + write_started.set() + time.sleep(0.2) + + with ( + patch( + "products.posthog_ai.backend.attachments.object_storage.read_object", + return_value=(content, "image/png"), + ), + patch("products.posthog_ai.backend.attachments.object_storage.write", side_effect=slow_write), + patch("products.posthog_ai.backend.attachments.object_storage.tag"), + patch("products.posthog_ai.backend.attachments.object_storage.delete"), + ): + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit( + finalize_attachments, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachment_id=attachment_id, + ) + self.assertTrue(write_started.wait(timeout=2)) + second = pool.submit( + finalize_attachments, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachment_id=attachment_id, + ) + first_result = first.result(timeout=5) + second_result = second.result(timeout=5) + + self.assertEqual(write_calls, 1) + self.assertEqual(first_result.id, attachment_id) + self.assertEqual(second_result.id, attachment_id) + self.assertEqual(second_result.size, first_result.size) + + def test_finalize_tag_failure_deletes_normalized_object(self) -> None: + content = self._png_bytes() + attachment_id = "att-2" + self._seed_prepared_record(attachment_id, size=len(content)) + normalized_path = _normalized_storage_path( + self.team.id, self.user.id, self.conversation_id, attachment_id, "image/png" + ) + + with ( + patch( + "products.posthog_ai.backend.attachments.object_storage.read_object", + return_value=(content, "image/png"), + ), + patch("products.posthog_ai.backend.attachments.object_storage.write"), + patch("products.posthog_ai.backend.attachments.object_storage.tag", side_effect=RuntimeError("tag failed")), + patch("products.posthog_ai.backend.attachments.object_storage.delete") as m_delete, + ): + with self.assertRaises(AttachmentStorageError): + finalize_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachment_id=attachment_id, + ) + + record = _get_record(attachment_id) + assert record is not None + self.assertEqual(record["status"], "prepared") + deleted_paths = [call.args[0] for call in m_delete.call_args_list] + self.assertIn(normalized_path, deleted_paths) + + @parameterized.expand( + [ + ("png_declared_for_jpeg", "image/png", "jpeg"), + ("jpeg_declared_for_png", "image/jpeg", "png"), + ] + ) + def test_image_validation_rejects_spoofed_content(self, _name: str, content_type: str, source_type: str) -> None: + content = self._jpeg_bytes() if source_type == "jpeg" else self._png_bytes() + + with self.assertRaises(AttachmentValidationError): + _validate_and_normalize_image(content, content_type) + + def test_image_validation_rejects_animated_png(self) -> None: + first = Image.new("RGB", (2, 2), color="red") + second = Image.new("RGB", (2, 2), color="blue") + buffer = BytesIO() + first.save(buffer, format="PNG", save_all=True, append_images=[second], duration=100) + + with self.assertRaisesRegex(AttachmentValidationError, "Animated images"): + _validate_and_normalize_image(buffer.getvalue(), "image/png") + + def test_image_validation_rejects_twenty_megapixels(self) -> None: + image = Image.new("1", (5_000, 4_000)) + buffer = BytesIO() + image.save(buffer, format="PNG") + + with self.assertRaisesRegex(AttachmentValidationError, "fewer than 20 million pixels"): + _validate_and_normalize_image(buffer.getvalue(), "image/png") + + def test_png_normalization_preserves_transparency_and_strips_metadata(self) -> None: + image = Image.new("RGBA", (2, 3), color=(255, 0, 0, 64)) + metadata = PngImagePlugin.PngInfo() + metadata.add_text("private", "remove me") + buffer = BytesIO() + image.save(buffer, format="PNG", pnginfo=metadata) + + normalized, width, height = _validate_and_normalize_image(buffer.getvalue(), "image/png") + + with Image.open(BytesIO(normalized)) as clean: + self.assertEqual((width, height), (2, 3)) + self.assertEqual(clean.mode, "RGBA") + self.assertEqual(clean.getpixel((0, 0)), (255, 0, 0, 64)) + self.assertNotIn("private", clean.info) + + def test_jpeg_normalization_applies_orientation_and_strips_metadata(self) -> None: + image = Image.new("RGB", (2, 3), color="red") + exif = Image.Exif() + exif[274] = 6 + exif[315] = "private" + buffer = BytesIO() + image.save(buffer, format="JPEG", exif=exif) + + normalized, width, height = _validate_and_normalize_image(buffer.getvalue(), "image/jpeg") + + with Image.open(BytesIO(normalized)) as clean: + self.assertEqual((width, height), (3, 2)) + self.assertEqual(clean.size, (3, 2)) + self.assertEqual(len(clean.getexif()), 0) + + def test_invalid_image_cleanup_removes_raw_object_and_staging_record(self) -> None: + attachment_id = "att-invalid" + content = self._jpeg_bytes() + self._seed_prepared_record(attachment_id, size=len(content)) + raw_path = _raw_storage_path(self.team.id, self.user.id, self.conversation_id, attachment_id) + + with ( + patch( + "products.posthog_ai.backend.attachments.object_storage.read_object", + return_value=(content, "image/png"), + ), + patch("products.posthog_ai.backend.attachments.object_storage.delete") as delete, + ): + with self.assertRaises(AttachmentValidationError): + finalize_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + attachment_id=attachment_id, + ) + + delete.assert_called_once_with(raw_path) + record = _get_record(attachment_id) + assert record is not None + self.assertEqual(record["status"], "deleted") + self.assertNotIn( + attachment_id, + _get_index(_conversation_index_key(self.team.id, self.user.id, self.conversation_id)), + ) + self.assertIn(attachment_id, _get_index(_user_index_key(self.team.id, self.user.id))) + + def _promotion(self, attachment_id: str) -> AttachmentPromotion: + return AttachmentPromotion( + attachment_ids=[attachment_id], + newly_attached_ids=[attachment_id], + task_id="task-1", + run_id="run-1", + records=[ + { + "id": attachment_id, + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": str(self.conversation_id), + "file_name": "image.png", + "content_type": "image/png", + "declared_size": 10, + "raw_storage_path": _raw_storage_path( + self.team.id, self.user.id, self.conversation_id, attachment_id + ), + "status": "finalized", + } + ], + ) + + def test_rollback_keeps_retryable_checkpoint_when_detach_fails(self) -> None: + attachment_id = "att-3" + self._seed_prepared_record(attachment_id, size=10) + promotion = self._promotion(attachment_id) + _write_records_and_indexes( + records=[ + {**promotion.records[0], "status": "promoted", "promoted_task_id": "task-1", "promoted_run_id": "run-1"} + ], + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + conversation_ids=[], + user_ids=[], + ) + + with patch( + "products.posthog_ai.backend.attachments.tasks_facade.rollback_posthog_ai_attachments", + side_effect=RuntimeError("detach failed"), + ): + rollback_attachment_promotion(promotion) + + self.assertNotIn( + attachment_id, + _get_index(_conversation_index_key(self.team.id, self.user.id, self.conversation_id)), + ) + self.assertNotIn(attachment_id, _get_index(_user_index_key(self.team.id, self.user.id))) + record = _get_record(attachment_id) + assert record is not None + self.assertEqual(record["status"], "rollback_pending") + + def test_rollback_checkpoint_failure_still_detaches_manifest(self) -> None: + promotion = self._promotion("att-4") + with ( + patch("products.posthog_ai.backend.attachments.tasks_facade.rollback_posthog_ai_attachments") as m_detach, + patch( + "products.posthog_ai.backend.attachments._cache_locks", side_effect=RuntimeError("cache lock failed") + ), + patch("products.posthog_ai.backend.attachments.logger.warning") as m_warning, + ): + rollback_attachment_promotion(promotion) + + m_detach.assert_called_once_with( + task_id="task-1", + run_id="run-1", + team_id=self.team.id, + attachment_ids=["att-4"], + ) + warning_codes = [call.args[0] for call in m_warning.call_args_list] + self.assertIn("assistant_attachment_rollback_checkpoint_failed", warning_codes) + + def test_promotion_cache_failure_preserves_original_error_when_detach_fails(self) -> None: + attachment_id = "att-5" + record = { + **self._promotion(attachment_id).records[0], + "normalized_storage_path": "normalized/image.png", + "normalized_size": 10, + "width": 2, + "height": 5, + } + _write_records_and_indexes( + records=[record], + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + conversation_ids=[attachment_id], + user_ids=[attachment_id], + ) + cache_write_count = 0 + + def fail_final_cache_write(**kwargs: Any) -> None: + nonlocal cache_write_count + cache_write_count += 1 + if cache_write_count == 1: + _write_records_and_indexes(**kwargs) + return + raise RuntimeError("cache transition failed") + + with ( + patch( + "products.posthog_ai.backend.attachments.tasks_facade.promote_posthog_ai_attachments", + return_value=SimpleNamespace(attached_ids=[attachment_id], newly_attached_ids=[attachment_id]), + ), + patch( + "products.posthog_ai.backend.attachments._write_records_and_indexes", + side_effect=fail_final_cache_write, + ), + patch( + "products.posthog_ai.backend.attachments.tasks_facade.rollback_posthog_ai_attachments", + side_effect=RuntimeError("detach failed"), + ), + patch("products.posthog_ai.backend.attachments.logger.warning") as warning, + ): + with self.assertRaisesRegex(RuntimeError, "cache transition failed"): + promote_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + task_id="task-1", + run_id="run-1", + attachment_ids=[attachment_id], + ) + + warning.assert_any_call( + "assistant_attachment_promotion_rollback_failed", + task_id="task-1", + run_id="run-1", + error="detach failed", + ) + stored_record = _get_record(attachment_id) + assert stored_record is not None + self.assertEqual(stored_record["status"], "promotion_pending") + self.assertNotIn(attachment_id, _get_index(_user_index_key(self.team.id, self.user.id))) + + def test_promotion_resumes_pending_rollback_before_reusing_attachment(self) -> None: + attachment_id = str(uuid.uuid4()) + record = { + **self._promotion(attachment_id).records[0], + "normalized_storage_path": "normalized/image.png", + "normalized_size": 10, + "width": 2, + "height": 5, + "status": "rollback_pending", + "promoted_task_id": "old-task", + "promoted_run_id": "old-run", + } + _write_records_and_indexes( + records=[record], + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + conversation_ids=[], + user_ids=[], + ) + + with ( + patch("products.posthog_ai.backend.attachments.tasks_facade.rollback_posthog_ai_attachments") as rollback, + patch( + "products.posthog_ai.backend.attachments.tasks_facade.promote_posthog_ai_attachments", + return_value=SimpleNamespace(attached_ids=[attachment_id], newly_attached_ids=[attachment_id]), + ), + ): + promotion = promote_attachments( + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + task_id="new-task", + run_id="new-run", + attachment_ids=[attachment_id], + ) + + rollback.assert_called_once_with( + task_id="old-task", + run_id="old-run", + team_id=self.team.id, + attachment_ids=[attachment_id], + ) + assert promotion is not None + self.assertEqual(promotion.run_id, "new-run") + + def test_concurrent_promotion_cannot_send_one_attachment_to_two_runs(self) -> None: + attachment_id = str(uuid.uuid4()) + record = { + **self._promotion(attachment_id).records[0], + "normalized_storage_path": "normalized/image.png", + "normalized_size": 10, + "width": 2, + "height": 5, + } + _write_records_and_indexes( + records=[record], + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + conversation_ids=[attachment_id], + user_ids=[attachment_id], + ) + promotion_started = threading.Event() + release_promotion = threading.Event() + + def slow_promotion(**_kwargs: object) -> SimpleNamespace: + promotion_started.set() + release_promotion.wait(timeout=5) + return SimpleNamespace(attached_ids=[attachment_id], newly_attached_ids=[attachment_id]) + + with patch( + "products.posthog_ai.backend.attachments.tasks_facade.promote_posthog_ai_attachments", + side_effect=slow_promotion, + ) as promote: + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit( + promote_attachments, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + task_id="task-1", + run_id="run-1", + attachment_ids=[attachment_id], + ) + self.assertTrue(promotion_started.wait(timeout=2)) + second = pool.submit( + promote_attachments, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=self.conversation_id, + task_id="task-1", + run_id="run-2", + attachment_ids=[attachment_id], + ) + release_promotion.set() + first_promotion = first.result(timeout=5) + assert first_promotion is not None + self.assertEqual(first_promotion.run_id, "run-1") + with self.assertRaisesRegex(AttachmentValidationError, "another message"): + second.result(timeout=5) + + promote.assert_called_once() diff --git a/products/posthog_ai/backend/tests/test_message_routing.py b/products/posthog_ai/backend/tests/test_message_routing.py index 8bdb871812de..0a91591c9d88 100644 --- a/products/posthog_ai/backend/tests/test_message_routing.py +++ b/products/posthog_ai/backend/tests/test_message_routing.py @@ -7,6 +7,7 @@ from posthog.exceptions import Conflict, QuotaLimitExceeded +from products.posthog_ai.backend.attachments import AttachmentPromotion from products.posthog_ai.backend.context_wrapper import MAX_ATTACHED_ITEMS, MAX_TEXT_LENGTH from products.posthog_ai.backend.message_routing import ( POSTHOG_AI_INTERACTION_ORIGIN, @@ -136,6 +137,37 @@ def test_first_message_without_context_forwards_bare_content(self): assert run.state["pending_user_message"] == "Hello" assert run.state["attached_context"] == [] + def test_first_message_conflict_does_not_clear_winner_task_link(self): + task, orphan_run = self._stub_task() + winning_task, _ = self._stub_task() + Conversation.objects.filter(id=self.conversation.id).update(task_id=winning_task.id) + + car, workflow, sysprompt = self._patches(task) + with car, workflow, sysprompt, self.assertRaises(Conflict): + self._service().open({"content": "Hello", "trace_id": "t"}) + + self.conversation.refresh_from_db() + orphan_run.refresh_from_db() + assert self.conversation.task_id == winning_task.id + assert orphan_run.status == TaskRun.Status.FAILED + + def test_attachment_ids_require_nonblank_message(self): + with self.assertRaises(exceptions.ValidationError): + self._service().open({"content": " ", "attachment_ids": ["att-1"]}) + + def test_first_message_rejects_invalid_attachment_before_creating_task(self): + with ( + patch( + f"{ROUTING}.validate_new_run_attachments", + side_effect=exceptions.ValidationError("not found"), + ), + patch.object(Task, "create_and_run") as create_and_run, + self.assertRaises(exceptions.ValidationError), + ): + self._service().open({"content": "inspect", "attachment_ids": ["att-1"]}) + + create_and_run.assert_not_called() + def test_unknown_attached_context_type_raises(self): with self.assertRaises(exceptions.ValidationError): self._service().open({"content": "x", "attached_context": [{"type": "bogus", "id": 1}]}) @@ -190,6 +222,7 @@ def test_in_progress_followup_signals_existing_run(self): signal_args, _ = m_signal.call_args assert signal_args[0] == run.workflow_id assert "and the mobile funnel?" in signal_args[1] + assert m_signal.call_args.kwargs["message_id"] == "trace-2" assert task.runs.count() == 1 # The follow-up is logged with the full undeduped attached_context on _meta. @@ -198,6 +231,36 @@ def test_in_progress_followup_signals_existing_run(self): meta = logged_entries[0]["notification"]["params"]["_meta"] assert meta["attached_context"] == [{"type": "insight", "id": "abc"}] + def test_in_progress_followup_logs_artifact_ids(self): + task, run = self._stub_task() + run.status = TaskRun.Status.IN_PROGRESS + run.save(update_fields=["status"]) + self._attach_task(task) + promotion = AttachmentPromotion( + attachment_ids=["att-1"], + newly_attached_ids=["att-1"], + task_id=str(task.id), + run_id=str(run.id), + records=[ + { + "id": "att-1", + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": str(self.conversation.id), + } + ], + ) + + with ( + patch(f"{ROUTING}.promote_attachments", return_value=promotion), + patch(f"{ROUTING}.signal_task_followup_message"), + patch.object(TaskRun, "append_log") as m_append, + ): + self._service().open({"content": "follow up", "attachment_ids": ["att-1"], "attached_context": []}) + + logged_entries = m_append.call_args[0][0] + assert logged_entries[0]["notification"]["params"]["artifact_ids"] == ["att-1"] + def test_in_progress_followup_clears_warm_flag(self): # A warm Run that receives its first human message is no longer speculative — the warm flag is # cleared so the warm-pool cap stops counting it (it becomes an active, AI-credit-governed Run). @@ -291,6 +354,95 @@ def test_terminal_followup_creates_new_run_with_resume(self): # The resumed agent keeps the same write scopes as the first message. assert m_workflow.call_args.kwargs["posthog_mcp_scopes"] == "full" + def test_terminal_followup_rejects_invalid_attachment_before_creating_run(self): + task, run = self._stub_task() + run.status = TaskRun.Status.COMPLETED + run.save(update_fields=["status"]) + self._attach_task(task) + + with ( + patch( + f"{ROUTING}.validate_new_run_attachments", + side_effect=exceptions.ValidationError("not found"), + ), + self.assertRaises(exceptions.ValidationError), + ): + self._service().open({"content": "inspect", "attachment_ids": ["att-1"]}) + + assert task.runs.count() == 1 + + def test_terminal_followup_sets_pending_artifact_ids_after_promotion(self): + task, run = self._stub_task() + run.status = TaskRun.Status.COMPLETED + run.save(update_fields=["status"]) + self._attach_task(task) + promotion = AttachmentPromotion( + attachment_ids=["att-1"], + newly_attached_ids=["att-1"], + task_id=str(task.id), + run_id="ignored", + records=[ + { + "id": "att-1", + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": str(self.conversation.id), + } + ], + ) + + with ( + patch.object(PromptService, "build", return_value=SYS_PROMPT), + patch(f"{ROUTING}.validate_new_run_attachments"), + patch(f"{ROUTING}.promote_attachments", return_value=promotion), + patch(f"{ROUTING}.execute_task_processing_workflow"), + patch.object(Task, "create_run", wraps=task.create_run) as m_create_run, + ): + result = self._service().open({"content": "resume please", "attachment_ids": ["att-1"]}) + + assert result is not None + extra_state = m_create_run.call_args.kwargs["extra_state"] + assert "pending_user_artifact_ids" not in extra_state + new_run = task.runs.order_by("-created_at").first() + assert new_run is not None + assert new_run.state["pending_user_artifact_ids"] == ["att-1"] + + def test_terminal_followup_workflow_failure_rolls_back_and_fails_run(self): + task, run = self._stub_task() + run.status = TaskRun.Status.COMPLETED + run.save(update_fields=["status"]) + self._attach_task(task) + promotion = AttachmentPromotion( + attachment_ids=["att-1"], + newly_attached_ids=["att-1"], + task_id=str(task.id), + run_id="ignored", + records=[ + { + "id": "att-1", + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": str(self.conversation.id), + } + ], + ) + + with ( + patch.object(PromptService, "build", return_value=SYS_PROMPT), + patch(f"{ROUTING}.validate_new_run_attachments"), + patch(f"{ROUTING}.promote_attachments", return_value=promotion), + patch(f"{ROUTING}.rollback_attachment_promotion") as m_rollback, + patch(f"{ROUTING}.execute_task_processing_workflow", side_effect=RuntimeError("boom")), + ): + with self.assertRaises(RuntimeError): + self._service().open({"content": "resume please", "attachment_ids": ["att-1"]}) + + new_run = task.runs.order_by("-created_at").first() + assert new_run is not None + new_run.refresh_from_db() + m_rollback.assert_called_once() + assert new_run.status == TaskRun.Status.FAILED + def test_dedupes_entities_named_in_prior_run_state(self): task, run = self._stub_task() run.status = TaskRun.Status.COMPLETED diff --git a/products/posthog_ai/frontend/api/logics.ts b/products/posthog_ai/frontend/api/logics.ts index 7d9606242b08..c7a2c860f403 100644 --- a/products/posthog_ai/frontend/api/logics.ts +++ b/products/posthog_ai/frontend/api/logics.ts @@ -75,3 +75,13 @@ export type { RunnerPanelLogicProps, ActiveCreation } from '../logics/runnerPane // auto-submitting it); the paired `taskTrackerSceneLogic` consumes it on mount or when it arrives. export { composerSeedLogic } from '../logics/composerSeedLogic' export type { ComposerSeed, ComposerSeedLogicProps } from '../logics/composerSeedLogic' + +// --- Assistant image attachments --- +export { + ASSISTANT_ATTACHMENT_ACCEPT, + MAX_ASSISTANT_ATTACHMENTS, + MAX_ASSISTANT_ATTACHMENT_BYTES, + MAX_ASSISTANT_ATTACHMENTS_TOTAL_BYTES, + assistantAttachmentsLogic, +} from '../logics/assistantAttachmentsLogic' +export type { AssistantAttachment, AssistantAttachmentsLogicProps } from '../logics/assistantAttachmentsLogic' diff --git a/products/posthog_ai/frontend/api/primitives.ts b/products/posthog_ai/frontend/api/primitives.ts index f3b8d3d9d418..6975e7e324d3 100644 --- a/products/posthog_ai/frontend/api/primitives.ts +++ b/products/posthog_ai/frontend/api/primitives.ts @@ -30,6 +30,10 @@ export type { ComposerModelEffortPickersProps } from '../components/composer/Com // The composer's context affordance: @-picker (TaxonomicPopover) + removable chips over the // attached-context store. Drop into `Composer.Header`; headless half is in api/logics. export { AttachedContextBar } from '../components/composer/AttachedContextBar' +export { ImageAttachmentButton } from '../components/composer/ImageAttachmentButton' +export type { ImageAttachmentButtonProps } from '../components/composer/ImageAttachmentButton' +export { ImageAttachmentPreviewList } from '../components/composer/ImageAttachmentPreviewList' +export type { ImageAttachmentPreviewListProps } from '../components/composer/ImageAttachmentPreviewList' // Welcome header (logomark + headline + subheadline) and its overridable default headlines. export { Welcome } from '../components/welcome/Welcome' diff --git a/products/posthog_ai/frontend/components/composer/ImageAttachmentButton.tsx b/products/posthog_ai/frontend/components/composer/ImageAttachmentButton.tsx new file mode 100644 index 000000000000..ccbeb626ac23 --- /dev/null +++ b/products/posthog_ai/frontend/components/composer/ImageAttachmentButton.tsx @@ -0,0 +1,45 @@ +import { useRef } from 'react' + +import { IconImage } from '@posthog/icons' +import { LemonButton } from '@posthog/lemon-ui' + +import { ASSISTANT_ATTACHMENT_ACCEPT } from '../../logics/assistantAttachmentsLogic' + +export interface ImageAttachmentButtonProps { + disabled?: boolean + onFilesSelected: (files: File[]) => void +} + +export function ImageAttachmentButton({ disabled, onFilesSelected }: ImageAttachmentButtonProps): JSX.Element { + const fileInputRef = useRef(null) + + return ( + <> + { + const files = Array.from(event.target.files || []) + if (files.length > 0) { + onFilesSelected(files) + } + event.target.value = '' + }} + ref={fileInputRef} + type="file" + /> + } + noPadding + onClick={() => fileInputRef.current?.click()} + size="small" + type="tertiary" + /> + + ) +} diff --git a/products/posthog_ai/frontend/components/composer/ImageAttachmentPreviewList.tsx b/products/posthog_ai/frontend/components/composer/ImageAttachmentPreviewList.tsx new file mode 100644 index 000000000000..5c594a8f3a3d --- /dev/null +++ b/products/posthog_ai/frontend/components/composer/ImageAttachmentPreviewList.tsx @@ -0,0 +1,83 @@ +import { IconImage, IconRefresh, IconTrash } from '@posthog/icons' +import { LemonButton, Spinner } from '@posthog/lemon-ui' + +import type { AssistantAttachment } from '../../logics/assistantAttachmentsLogic' + +export interface ImageAttachmentPreviewListProps { + attachments: AssistantAttachment[] + disabled?: boolean + onRemove: (localId: string) => void + onRetry: (localId: string) => void +} + +export function ImageAttachmentPreviewList({ + attachments, + disabled = false, + onRemove, + onRetry, +}: ImageAttachmentPreviewListProps): JSX.Element | null { + if (attachments.length === 0) { + return null + } + + return ( +
+ {attachments.map((attachment) => ( +
+
+ {attachment.previewUrl ? ( + {attachment.file.name} + ) : ( + + )} +
+
+
+ {attachment.file.name} +
+ {attachment.status === 'uploading' ? ( + + Uploading + + ) : attachment.status === 'error' ? ( + {attachment.error || 'Upload failed'} + ) : ( + Ready + )} +
+
+ {attachment.status === 'error' && ( + } + noPadding + onClick={() => onRetry(attachment.localId)} + size="xsmall" + type="tertiary" + /> + )} + } + noPadding + onClick={() => onRemove(attachment.localId)} + size="xsmall" + type="tertiary" + /> +
+
+ ))} +
+ ) +} diff --git a/products/posthog_ai/frontend/generated/api.schemas.ts b/products/posthog_ai/frontend/generated/api.schemas.ts index 63af694ab95d..c819412a0966 100644 --- a/products/posthog_ai/frontend/generated/api.schemas.ts +++ b/products/posthog_ai/frontend/generated/api.schemas.ts @@ -7,6 +7,108 @@ * PostHog API - generated * OpenAPI spec version: 1.0.0 */ +export interface AttachmentDeleteRequestApi { + /** Conversation UUID the image attachment belongs to. */ + conversation_id: string + /** Single attachment ID to delete from staging. */ + attachment_id: string +} + +export interface AttachmentFinalizeRequestApi { + /** Conversation UUID the uploaded image attachment belongs to. */ + conversation_id: string + /** Prepared attachment ID to validate and finalize for sandbox use. */ + attachment_id: string +} + +/** + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg + */ +export type AssistantAttachmentContentTypeEnumApi = + (typeof AssistantAttachmentContentTypeEnumApi)[keyof typeof AssistantAttachmentContentTypeEnumApi] + +export const AssistantAttachmentContentTypeEnumApi = { + ImagePng: 'image/png', + ImageJpeg: 'image/jpeg', +} as const + +export interface AttachmentFinalizeResponseApi { + /** Opaque attachment ID for sandbox message sends. */ + id: string + /** Sanitized file name recorded for the image attachment. */ + file_name: string + /** Validated MIME type for the normalized image. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnumApi + /** Normalized image size in bytes. */ + size: number + /** Normalized image width in pixels. */ + width: number + /** Normalized image height in pixels. */ + height: number +} + +export interface AttachmentPrepareItemApi { + /** + * File name to associate with the uploaded image. + * @maxLength 255 + */ + file_name: string + /** + * Expected upload size in bytes. Each image must be 4194304 bytes or smaller. + * @minimum 1 + * @maximum 4194304 + */ + size: number + /** Exact MIME type for the direct upload. Only PNG and JPEG are supported. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnumApi +} + +export interface AttachmentPrepareRequestApi { + /** Conversation UUID the staged image attachments belong to. */ + conversation_id: string + /** + * Images to stage for the next sandbox message. + * @minItems 1 + * @maxItems 4 + */ + attachments: AttachmentPrepareItemApi[] +} + +/** + * Signed S3-compatible form fields to include with the upload request. + */ +export type AttachmentPrepareResponseItemApiUploadFields = { [key: string]: string } + +export interface AttachmentPrepareResponseItemApi { + /** Opaque attachment ID for finalize, delete, and sandbox message sends. */ + id: string + /** Sanitized file name recorded for the image attachment. */ + file_name: string + /** Signed MIME type for the upload. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnumApi + /** Expected upload size in bytes. */ + size: number + /** Signed direct-upload URL for the image attachment. */ + upload_url: string + /** Signed S3-compatible form fields to include with the upload request. */ + upload_fields: AttachmentPrepareResponseItemApiUploadFields +} + +export interface AttachmentPrepareResponseApi { + /** Prepared image uploads. */ + attachments: AttachmentPrepareResponseItemApi[] +} + export interface DocsSearchRequestApi { /** Natural-language description of what to find in the PostHog documentation. Inkeep performs hybrid (semantic + full-text) RAG, so phrase the query the way a user would ask the question. */ query: string diff --git a/products/posthog_ai/frontend/generated/api.ts b/products/posthog_ai/frontend/generated/api.ts index 0a62bc0b2929..e37ae45190bf 100644 --- a/products/posthog_ai/frontend/generated/api.ts +++ b/products/posthog_ai/frontend/generated/api.ts @@ -8,7 +8,79 @@ import { apiMutator } from '../../../../frontend/src/lib/api-orval-mutator' * PostHog API - generated * OpenAPI spec version: 1.0.0 */ -import type { DocsSearchRequestApi, DocsSearchResponseApi, McpToolsCreate200 } from './api.schemas' +import type { + AttachmentDeleteRequestApi, + AttachmentFinalizeRequestApi, + AttachmentFinalizeResponseApi, + AttachmentPrepareRequestApi, + AttachmentPrepareResponseApi, + DocsSearchRequestApi, + DocsSearchResponseApi, + McpToolsCreate200, +} from './api.schemas' + +export const getAssistantAttachmentsDeleteCreateUrl = (projectId: string) => { + return `/api/projects/${projectId}/assistant_attachments/delete/` +} + +/** + * Idempotently delete a single staged or finalized sandbox image attachment for one conversation. + * @summary Delete one sandbox image attachment + */ +export const assistantAttachmentsDeleteCreate = async ( + projectId: string, + attachmentDeleteRequestApi: AttachmentDeleteRequestApi, + options?: RequestInit +): Promise => { + return apiMutator(getAssistantAttachmentsDeleteCreateUrl(projectId), { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/json', ...options?.headers }, + body: JSON.stringify(attachmentDeleteRequestApi), + }) +} + +export const getAssistantAttachmentsFinalizeCreateUrl = (projectId: string) => { + return `/api/projects/${projectId}/assistant_attachments/finalize/` +} + +/** + * Validate, normalize, and finalize one staged image upload for sandbox messages. + * @summary Finalize sandbox image uploads + */ +export const assistantAttachmentsFinalizeCreate = async ( + projectId: string, + attachmentFinalizeRequestApi: AttachmentFinalizeRequestApi, + options?: RequestInit +): Promise => { + return apiMutator(getAssistantAttachmentsFinalizeCreateUrl(projectId), { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/json', ...options?.headers }, + body: JSON.stringify(attachmentFinalizeRequestApi), + }) +} + +export const getAssistantAttachmentsPrepareCreateUrl = (projectId: string) => { + return `/api/projects/${projectId}/assistant_attachments/prepare/` +} + +/** + * Reserve direct-upload image attachment slots for a sandbox conversation and return signed POST fields. + * @summary Prepare sandbox image uploads + */ +export const assistantAttachmentsPrepareCreate = async ( + projectId: string, + attachmentPrepareRequestApi: AttachmentPrepareRequestApi, + options?: RequestInit +): Promise => { + return apiMutator(getAssistantAttachmentsPrepareCreateUrl(projectId), { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/json', ...options?.headers }, + body: JSON.stringify(attachmentPrepareRequestApi), + }) +} export const getMcpToolsCreateUrl = (projectId: string, toolName: string) => { return `/api/projects/${projectId}/mcp_tools/${toolName}/` diff --git a/products/posthog_ai/frontend/generated/api.zod.ts b/products/posthog_ai/frontend/generated/api.zod.ts index 971cfc372160..550e296d68a8 100644 --- a/products/posthog_ai/frontend/generated/api.zod.ts +++ b/products/posthog_ai/frontend/generated/api.zod.ts @@ -9,6 +9,61 @@ */ import * as zod from 'zod' +/** + * Idempotently delete a single staged or finalized sandbox image attachment for one conversation. + * @summary Delete one sandbox image attachment + */ +export const AssistantAttachmentsDeleteCreateBody = /* @__PURE__ */ zod.object({ + conversation_id: zod.uuid().describe('Conversation UUID the image attachment belongs to.'), + attachment_id: zod.uuid().describe('Single attachment ID to delete from staging.'), +}) + +/** + * Validate, normalize, and finalize one staged image upload for sandbox messages. + * @summary Finalize sandbox image uploads + */ +export const AssistantAttachmentsFinalizeCreateBody = /* @__PURE__ */ zod.object({ + conversation_id: zod.uuid().describe('Conversation UUID the uploaded image attachment belongs to.'), + attachment_id: zod.uuid().describe('Prepared attachment ID to validate and finalize for sandbox use.'), +}) + +/** + * Reserve direct-upload image attachment slots for a sandbox conversation and return signed POST fields. + * @summary Prepare sandbox image uploads + */ +export const assistantAttachmentsPrepareCreateBodyAttachmentsItemFileNameMax = 255 + +export const assistantAttachmentsPrepareCreateBodyAttachmentsItemSizeMax = 4194304 + +export const assistantAttachmentsPrepareCreateBodyAttachmentsMax = 4 + +export const AssistantAttachmentsPrepareCreateBody = /* @__PURE__ */ zod.object({ + conversation_id: zod.uuid().describe('Conversation UUID the staged image attachments belong to.'), + attachments: zod + .array( + zod.object({ + file_name: zod + .string() + .max(assistantAttachmentsPrepareCreateBodyAttachmentsItemFileNameMax) + .describe('File name to associate with the uploaded image.'), + size: zod + .number() + .min(1) + .max(assistantAttachmentsPrepareCreateBodyAttachmentsItemSizeMax) + .describe('Expected upload size in bytes. Each image must be 4194304 bytes or smaller.'), + content_type: zod + .enum(['image/png', 'image/jpeg']) + .describe('\* `image\/png` - image\/png\n\* `image\/jpeg` - image\/jpeg') + .describe( + 'Exact MIME type for the direct upload. Only PNG and JPEG are supported.\n\n\* `image\/png` - image\/png\n\* `image\/jpeg` - image\/jpeg' + ), + }) + ) + .min(1) + .max(assistantAttachmentsPrepareCreateBodyAttachmentsMax) + .describe('Images to stage for the next sandbox message.'), +}) + /** * Run a hybrid (semantic + full-text) RAG search over the PostHog documentation via Inkeep. Returns a markdown body with title, URL, and excerpt for each match for the agent to cite back to the user. * @summary Search PostHog documentation diff --git a/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.test.ts b/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.test.ts new file mode 100644 index 000000000000..3865146ebe9c --- /dev/null +++ b/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.test.ts @@ -0,0 +1,253 @@ +import { expectLogic } from 'kea-test-utils' + +import { projectLogic } from 'scenes/projectLogic' + +import { useMocks } from '~/mocks/jest' +import { initKeaTests } from '~/test/init' + +import { + assistantAttachmentsDeleteCreate, + assistantAttachmentsFinalizeCreate, + assistantAttachmentsPrepareCreate, +} from '../generated/api' +import { MAX_ASSISTANT_ATTACHMENTS } from './assistantAttachmentsLogic' +import { assistantAttachmentsLogic } from './assistantAttachmentsLogic' + +jest.mock('../generated/api', () => ({ + assistantAttachmentsPrepareCreate: jest.fn(), + assistantAttachmentsFinalizeCreate: jest.fn(), + assistantAttachmentsDeleteCreate: jest.fn(), +})) + +function deferred(): { promise: Promise; reject: (error: unknown) => void; resolve: (value: T) => void } { + let resolve!: (value: T) => void + let reject!: (error: unknown) => void + const promise = new Promise((innerResolve, innerReject) => { + resolve = innerResolve + reject = innerReject + }) + return { promise, resolve, reject } +} + +describe('assistantAttachmentsLogic', () => { + let logic: ReturnType + let projectLogicInstance: ReturnType + + beforeEach(() => { + initKeaTests() + useMocks({ + get: { + '/_preflight/': {}, + }, + }) + Object.defineProperty(URL, 'createObjectURL', { + configurable: true, + value: jest.fn((file: File) => `blob:${file.name}`), + }) + Object.defineProperty(URL, 'revokeObjectURL', { configurable: true, value: jest.fn() }) + const baseFetch = global.fetch.bind(global) + global.fetch = jest.fn((input: RequestInfo | URL, init?: RequestInit) => { + if (typeof input === 'string' && input.startsWith('https://upload.test')) { + return Promise.resolve({ ok: true } as Response) + } + return baseFetch(input, init) + }) as any + + projectLogicInstance = projectLogic() + projectLogicInstance.mount() + projectLogicInstance.actions.loadCurrentProjectSuccess({ id: 1, name: 'Test project' } as any) + + logic = assistantAttachmentsLogic({ conversationId: 'conv-1' }) + logic.mount() + }) + + afterEach(() => { + logic?.unmount() + projectLogicInstance?.unmount() + jest.restoreAllMocks() + }) + + it('validates type, count, size, and uniqueness', async () => { + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ attachments: [] }) + + const validFiles = Array.from( + { length: MAX_ASSISTANT_ATTACHMENTS }, + (_, index) => new File([String(index)], `${index}.png`, { type: 'image/png', lastModified: index }) + ) + const duplicate = new File(['0'], '0.png', { type: 'image/png', lastModified: 0 }) + const unsupported = new File(['svg'], 'a.svg', { type: 'image/svg+xml', lastModified: 10 }) + const empty = new File([], 'empty.png', { type: 'image/png', lastModified: 13 }) + const oversized = new File([new Uint8Array(4 * 1024 * 1024 + 1)], 'big.png', { + type: 'image/png', + lastModified: 11, + }) + const extra = new File(['x'], 'extra.png', { type: 'image/png', lastModified: 12 }) + + logic.actions.addFiles([...validFiles, duplicate, unsupported, empty, oversized, extra]) + + await expectLogic(logic).toMatchValues({ attachments: expect.any(Array) }) + expect(logic.values.attachments).toHaveLength(MAX_ASSISTANT_ATTACHMENTS) + expect(logic.values.selectionError).toContain('That image is already attached.') + expect(logic.values.selectionError).toContain('Choose a PNG or JPEG image.') + expect(logic.values.selectionError).toContain('Choose a non-empty image.') + expect(logic.values.selectionError).toContain('Each image must be 4 MiB or smaller.') + expect(logic.values.selectionError).toContain('You can attach up to 4 images.') + }) + + it('rejects files when the selected total would exceed 10 MiB', async () => { + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ attachments: [] }) + + logic.actions.addFiles([ + new File([new Uint8Array(4 * 1024 * 1024)], 'a.png', { type: 'image/png', lastModified: 1 }), + new File([new Uint8Array(4 * 1024 * 1024)], 'b.png', { type: 'image/png', lastModified: 2 }), + new File([new Uint8Array(3 * 1024 * 1024)], 'c.png', { type: 'image/png', lastModified: 3 }), + ]) + + await expectLogic(logic).toMatchValues({ attachments: expect.any(Array) }) + expect(logic.values.attachments).toHaveLength(2) + expect(logic.values.selectionError).toContain('Selected images must total 10 MiB or less.') + }) + + it('deletes a prepared attachment once if upload resolves after removal', async () => { + const uploadDeferred = deferred() + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ + attachments: [ + { + id: 'prepared-1', + file_name: 'a.png', + content_type: 'image/png', + size: 1, + upload_url: 'https://upload.test', + upload_fields: { key: 'value' }, + }, + ], + }) + ;(global.fetch as jest.Mock).mockReturnValue(uploadDeferred.promise) + ;(assistantAttachmentsFinalizeCreate as jest.Mock).mockResolvedValue({ + id: 'prepared-1', + file_name: 'a.png', + content_type: 'image/png', + size: 1, + width: 1, + height: 1, + }) + ;(assistantAttachmentsDeleteCreate as jest.Mock).mockResolvedValue(undefined) + + logic.actions.addFiles([new File(['a'], 'a.png', { type: 'image/png' })]) + await expectLogic(logic).toDispatchActions(['uploadAttachment']) + + const localId = logic.values.attachments[0]?.localId + expect(localId).toBeTruthy() + logic.actions.removeAttachment(localId as string) + uploadDeferred.resolve({ ok: true } as Response) + + await expectLogic(logic).toFinishAllListeners() + expect(assistantAttachmentsDeleteCreate).toHaveBeenCalledTimes(1) + expect(assistantAttachmentsDeleteCreate).toHaveBeenCalledWith('1', { + conversation_id: 'conv-1', + attachment_id: 'prepared-1', + }) + }) + + it('deletes a prepared attachment once if finalize resolves after unmount', async () => { + const finalizeDeferred = deferred() + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ + attachments: [ + { + id: 'prepared-2', + file_name: 'b.png', + content_type: 'image/png', + size: 1, + upload_url: 'https://upload.test', + upload_fields: { key: 'value' }, + }, + ], + }) + ;(global.fetch as jest.Mock).mockResolvedValue({ ok: true } as Response) + ;(assistantAttachmentsFinalizeCreate as jest.Mock).mockReturnValue(finalizeDeferred.promise) + ;(assistantAttachmentsDeleteCreate as jest.Mock).mockResolvedValue(undefined) + + logic.actions.addFiles([new File(['b'], 'b.png', { type: 'image/png' })]) + await expectLogic(logic).toDispatchActions(['uploadAttachment']) + + logic.unmount() + finalizeDeferred.resolve({ + id: 'prepared-2', + file_name: 'b.png', + content_type: 'image/png', + size: 1, + width: 1, + height: 1, + }) + + await Promise.resolve() + await Promise.resolve() + + expect(assistantAttachmentsDeleteCreate).toHaveBeenCalledTimes(1) + expect(assistantAttachmentsDeleteCreate).toHaveBeenCalledWith('1', { + conversation_id: 'conv-1', + attachment_id: 'prepared-2', + }) + }) + + it('keeps attachments while a send is pending and clears them only after success', () => { + logic.actions.addAttachment({ + localId: 'local-send', + file: new File(['a'], 'a.png', { type: 'image/png' }), + fileKey: 'a.png:image/png:1:0', + previewUrl: 'blob:a.png', + status: 'ready', + attachmentId: 'prepared-send', + }) + + logic.actions.beginAttachmentsSend(['prepared-send']) + expect(logic.values.attachmentsAreSending).toBe(true) + expect(logic.values.attachmentSubmissionDisabledReason).toBe('Wait for images to finish sending') + expect(logic.values.attachments).toHaveLength(1) + + logic.actions.markAttachmentsSendFailed(['prepared-send']) + expect(logic.values.attachmentsAreSending).toBe(false) + expect(logic.values.attachments).toHaveLength(1) + + logic.actions.beginAttachmentsSend(['prepared-send']) + logic.actions.markAttachmentsSent(['prepared-send']) + expect(logic.values.attachmentsAreSending).toBe(false) + expect(logic.values.attachments).toHaveLength(0) + }) + + it('deletes an attachment through the project that prepared it', async () => { + ;(assistantAttachmentsPrepareCreate as jest.Mock).mockResolvedValue({ + attachments: [ + { + id: 'prepared-project', + file_name: 'project.png', + content_type: 'image/png', + size: 1, + upload_url: 'https://upload.test', + upload_fields: { key: 'value' }, + }, + ], + }) + ;(assistantAttachmentsFinalizeCreate as jest.Mock).mockResolvedValue({ + id: 'prepared-project', + file_name: 'project.png', + content_type: 'image/png', + size: 1, + width: 1, + height: 1, + }) + ;(assistantAttachmentsDeleteCreate as jest.Mock).mockResolvedValue(undefined) + + logic.actions.addFiles([new File(['a'], 'project.png', { type: 'image/png' })]) + await expectLogic(logic).toFinishAllListeners() + projectLogicInstance.actions.loadCurrentProjectSuccess({ id: 2, name: 'Other project' } as any) + + logic.actions.removeAttachment(logic.values.attachments[0].localId) + await expectLogic(logic).toFinishAllListeners() + + expect(assistantAttachmentsDeleteCreate).toHaveBeenCalledWith('1', { + conversation_id: 'conv-1', + attachment_id: 'prepared-project', + }) + }) +}) diff --git a/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.ts b/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.ts new file mode 100644 index 000000000000..faf4e0ec4792 --- /dev/null +++ b/products/posthog_ai/frontend/logics/assistantAttachmentsLogic.ts @@ -0,0 +1,504 @@ +import { + MakeLogicType, + actions, + afterMount, + beforeUnmount, + connect, + kea, + key, + listeners, + path, + props, + reducers, + selectors, +} from 'kea' +import posthog from 'posthog-js' + +import { lemonToast } from '@posthog/lemon-ui' + +import { uuid } from 'lib/utils/dom' +import { projectLogic } from 'scenes/projectLogic' + +import { + assistantAttachmentsDeleteCreate, + assistantAttachmentsFinalizeCreate, + assistantAttachmentsPrepareCreate, +} from '../generated/api' +import type { AssistantAttachmentContentTypeEnumApi } from '../generated/api.schemas' + +export const MAX_ASSISTANT_ATTACHMENTS = 4 +export const MAX_ASSISTANT_ATTACHMENT_BYTES = 4 * 1024 * 1024 +export const MAX_ASSISTANT_ATTACHMENTS_TOTAL_BYTES = 10 * 1024 * 1024 +export const ASSISTANT_ATTACHMENT_ACCEPT = 'image/png,image/jpeg' + +const ACCEPTED_IMAGE_TYPES = new Set(['image/png', 'image/jpeg']) + +type PreparedAttachmentResponse = Awaited>['attachments'][number] +type FinalizedAttachmentResponse = Awaited> + +export interface AssistantAttachmentsLogicProps { + conversationId: string +} + +export interface AssistantAttachment { + localId: string + file: File + fileKey: string + previewUrl: string + status: 'uploading' | 'ready' | 'error' + attachmentId?: string + error?: string +} + +function getFileKey(file: File): string { + return `${file.name}:${file.type}:${file.size}:${file.lastModified}` +} + +function getUploadErrorMessage(): string { + return "Couldn't upload this image. Retry the upload or remove it." +} + +async function deletePreparedAttachment( + projectId: string | undefined, + conversationId: string, + attachmentId: string | undefined +): Promise { + if (!attachmentId || !projectId) { + return + } + let lastError: unknown + for (let attempt = 0; attempt < 3; attempt++) { + try { + await assistantAttachmentsDeleteCreate(projectId, { + conversation_id: conversationId, + attachment_id: attachmentId, + }) + return + } catch (error) { + lastError = error + if (attempt < 2) { + await new Promise((resolve) => window.setTimeout(resolve, 250 * 2 ** attempt)) + } + } + } + posthog.captureException(lastError, { feature: 'posthog_ai_assistant_attachment_cleanup' }) + throw lastError +} + +// Generated by kea-typegen. Update if you're an agent, ignore if you're human. +export interface assistantAttachmentsLogicValues { + currentProjectId: number | null // projectLogic + attachmentSubmissionDisabledReason: string | undefined + attachments: AssistantAttachment[] + attachmentsAreSending: boolean + isDragActive: boolean + readyAttachmentIds: string[] + selectionError: string | null + totalBytes: number +} + +// Generated by kea-typegen. Update if you're an agent, ignore if you're human. +export interface assistantAttachmentsLogicActions { + addAttachment: (attachment: AssistantAttachment) => { + attachment: AssistantAttachment + } + addFiles: (files: File[]) => { + files: File[] + } + beginAttachmentsSend: (attachmentIds: string[]) => { + attachmentIds: string[] + } + markAttachmentsSendFailed: (attachmentIds: string[]) => { + attachmentIds: string[] + } + markAttachmentsSent: (attachmentIds: string[]) => { + attachmentIds: string[] + } + removeAttachment: (localId: string) => { + localId: string + } + retryAttachment: (localId: string) => { + localId: string + } + setAttachmentError: ( + localId: string, + error: string + ) => { + error: string + localId: string + } + setAttachmentReady: ( + localId: string, + attachmentId: string + ) => { + attachmentId: string + localId: string + } + setAttachmentUploading: (localId: string) => { + localId: string + } + setDragActive: (isActive: boolean) => { + isActive: boolean + } + setSelectionError: (error: string | null) => { + error: string | null + } + uploadAttachment: (localId: string) => { + localId: string + } +} + +// Generated by kea-typegen. Update if you're an agent, ignore if you're human. +export interface assistantAttachmentsLogicMeta { + key: string + __keaTypeGenInternalSelectorTypes: { + readyAttachmentIds: (attachments: AssistantAttachment[]) => string[] + attachmentSubmissionDisabledReason: ( + attachments: AssistantAttachment[], + attachmentsAreSending: boolean + ) => string | undefined + totalBytes: (attachments: AssistantAttachment[]) => number + } +} + +export type assistantAttachmentsLogicType = MakeLogicType< + assistantAttachmentsLogicValues, + assistantAttachmentsLogicActions, + AssistantAttachmentsLogicProps, + assistantAttachmentsLogicMeta +> + +export const assistantAttachmentsLogic = kea([ + props({} as AssistantAttachmentsLogicProps), + key((props) => props.conversationId), + path((key) => ['products', 'posthog_ai', 'frontend', 'logics', 'assistantAttachmentsLogic', key]), + + connect(() => ({ + values: [projectLogic, ['currentProjectId']], + })), + + actions({ + addFiles: (files: File[]) => ({ files }), + addAttachment: (attachment: AssistantAttachment) => ({ attachment }), + uploadAttachment: (localId: string) => ({ localId }), + retryAttachment: (localId: string) => ({ localId }), + removeAttachment: (localId: string) => ({ localId }), + setAttachmentUploading: (localId: string) => ({ localId }), + setAttachmentReady: (localId: string, attachmentId: string) => ({ localId, attachmentId }), + setAttachmentError: (localId: string, error: string) => ({ localId, error }), + setSelectionError: (error: string | null) => ({ error }), + setDragActive: (isActive: boolean) => ({ isActive }), + beginAttachmentsSend: (attachmentIds: string[]) => ({ attachmentIds }), + markAttachmentsSendFailed: (attachmentIds: string[]) => ({ attachmentIds }), + markAttachmentsSent: (attachmentIds: string[]) => ({ attachmentIds }), + }), + + reducers({ + attachments: [ + [] as AssistantAttachment[], + { + addAttachment: (state, { attachment }) => [...state, attachment], + removeAttachment: (state, { localId }) => state.filter((attachment) => attachment.localId !== localId), + setAttachmentUploading: (state, { localId }) => + state.map((attachment) => + attachment.localId === localId + ? { ...attachment, status: 'uploading', error: undefined, attachmentId: undefined } + : attachment + ), + setAttachmentReady: (state, { localId, attachmentId }) => + state.map((attachment) => + attachment.localId === localId + ? { ...attachment, status: 'ready', attachmentId, error: undefined } + : attachment + ), + setAttachmentError: (state, { localId, error }) => + state.map((attachment) => + attachment.localId === localId + ? { ...attachment, status: 'error', error, attachmentId: undefined } + : attachment + ), + markAttachmentsSent: (state, { attachmentIds }) => + state.filter( + (attachment) => !attachment.attachmentId || !attachmentIds.includes(attachment.attachmentId) + ), + }, + ], + selectionError: [null as string | null, { setSelectionError: (_, { error }) => error }], + isDragActive: [false, { setDragActive: (_, { isActive }) => isActive, addFiles: () => false }], + attachmentsAreSending: [ + false, + { + beginAttachmentsSend: () => true, + markAttachmentsSendFailed: () => false, + markAttachmentsSent: () => false, + }, + ], + }), + + selectors({ + readyAttachmentIds: [ + (s) => [s.attachments], + (attachments: AssistantAttachment[]): string[] => + attachments.flatMap((attachment) => + attachment.status === 'ready' && attachment.attachmentId ? [attachment.attachmentId] : [] + ), + ], + attachmentSubmissionDisabledReason: [ + (s) => [s.attachments, s.attachmentsAreSending], + (attachments: AssistantAttachment[], attachmentsAreSending: boolean): string | undefined => { + if (attachmentsAreSending) { + return 'Wait for images to finish sending' + } + if (attachments.some((attachment) => attachment.status === 'uploading')) { + return 'Wait for images to finish uploading' + } + if (attachments.some((attachment) => attachment.status === 'error')) { + return 'Retry or remove failed images before sending' + } + return undefined + }, + ], + totalBytes: [ + (s) => [s.attachments], + (attachments: AssistantAttachment[]): number => + attachments.reduce((total, attachment) => total + attachment.file.size, 0), + ], + }), + + listeners(({ actions, values, props, cache }) => { + const cleanupPreparedAttachment = ( + attachmentId: string | undefined, + projectId: string | undefined, + notifyFailure?: boolean + ): void => { + if (!attachmentId || cache.cleanupAttachmentIds.has(attachmentId)) { + return + } + cache.cleanupAttachmentIds.add(attachmentId) + void deletePreparedAttachment(projectId, props.conversationId, attachmentId).catch(() => { + if (notifyFailure) { + lemonToast.error("Couldn't remove the image from storage. Please try again.") + } + }) + } + + return { + addFiles: ({ files }) => { + actions.setSelectionError(null) + const existingKeys = new Set(values.attachments.map((attachment) => attachment.fileKey)) + let selectedCount = values.attachments.length + let selectedBytes = values.totalBytes + const validationErrors = new Set() + + for (const file of files) { + const fileKey = getFileKey(file) + if (existingKeys.has(fileKey)) { + validationErrors.add('That image is already attached.') + continue + } + if (!ACCEPTED_IMAGE_TYPES.has(file.type)) { + validationErrors.add('Choose a PNG or JPEG image.') + continue + } + if (file.size === 0) { + validationErrors.add('Choose a non-empty image.') + continue + } + if (file.size > MAX_ASSISTANT_ATTACHMENT_BYTES) { + validationErrors.add('Each image must be 4 MiB or smaller.') + continue + } + if (selectedCount >= MAX_ASSISTANT_ATTACHMENTS) { + validationErrors.add(`You can attach up to ${MAX_ASSISTANT_ATTACHMENTS} images.`) + continue + } + if (selectedBytes + file.size > MAX_ASSISTANT_ATTACHMENTS_TOTAL_BYTES) { + validationErrors.add('Selected images must total 10 MiB or less.') + continue + } + + const localId = uuid() + const previewUrl = URL.createObjectURL(file) + cache.previewUrls.set(localId, previewUrl) + existingKeys.add(fileKey) + selectedCount += 1 + selectedBytes += file.size + actions.addAttachment({ + localId, + file, + fileKey, + previewUrl, + status: 'uploading', + }) + actions.uploadAttachment(localId) + } + + if (validationErrors.size > 0) { + actions.setSelectionError([...validationErrors].join(' ')) + } + }, + uploadAttachment: async ({ localId }) => { + const attachment = values.attachments.find((candidate) => candidate.localId === localId) + if (!attachment || cache.uploadControllers.has(localId)) { + return + } + if (values.currentProjectId == null) { + actions.setAttachmentError(localId, getUploadErrorMessage()) + return + } + + const abortController = new AbortController() + cache.uploadControllers.set(localId, abortController) + let preparedAttachmentId: string | undefined + const uploadProjectId = String(values.currentProjectId) + + try { + const preparedResponse = await assistantAttachmentsPrepareCreate(uploadProjectId, { + conversation_id: props.conversationId, + attachments: [ + { + file_name: attachment.file.name, + content_type: attachment.file.type as AssistantAttachmentContentTypeEnumApi, + size: attachment.file.size, + }, + ], + }) + const preparedAttachment: PreparedAttachmentResponse | undefined = preparedResponse.attachments[0] + if (!preparedAttachment) { + throw new Error('Attachment preparation returned no upload target') + } + preparedAttachmentId = preparedAttachment.id + + if (cache.isUnmounted || abortController.signal.aborted) { + cleanupPreparedAttachment(preparedAttachmentId, uploadProjectId) + return + } + + cache.serverAttachmentIds.set(localId, preparedAttachmentId) + cache.serverAttachmentProjectIds.set(localId, uploadProjectId) + const uploadBody = new FormData() + for (const [field, value] of Object.entries(preparedAttachment.upload_fields)) { + uploadBody.append(field, value) + } + uploadBody.append('file', attachment.file) + const uploadResponse = await fetch(preparedAttachment.upload_url, { + method: 'POST', + body: uploadBody, + signal: abortController.signal, + }) + if (!uploadResponse.ok) { + throw new Error('Attachment object upload failed') + } + + if (cache.isUnmounted || abortController.signal.aborted) { + cleanupPreparedAttachment(preparedAttachmentId, uploadProjectId) + return + } + + const finalizedAttachment: FinalizedAttachmentResponse = await assistantAttachmentsFinalizeCreate( + uploadProjectId, + { + conversation_id: props.conversationId, + attachment_id: preparedAttachmentId, + } + ) + + if (cache.isUnmounted || abortController.signal.aborted) { + cleanupPreparedAttachment(preparedAttachmentId, uploadProjectId) + return + } + + actions.setAttachmentReady(localId, finalizedAttachment.id) + } catch (error) { + if (preparedAttachmentId && (cache.isUnmounted || abortController.signal.aborted)) { + cleanupPreparedAttachment(preparedAttachmentId, uploadProjectId) + return + } + posthog.captureException(error, { feature: 'posthog_ai_assistant_attachment' }) + actions.setAttachmentError(localId, getUploadErrorMessage()) + } finally { + if (cache.uploadControllers.get(localId) === abortController) { + cache.uploadControllers.delete(localId) + } + } + }, + retryAttachment: ({ localId }) => { + const attachment = values.attachments.find((candidate) => candidate.localId === localId) + if (!attachment || attachment.status !== 'error' || cache.uploadControllers.has(localId)) { + return + } + + const previousServerAttachmentId = cache.serverAttachmentIds.get(localId) + const previousProjectId = cache.serverAttachmentProjectIds.get(localId) + cache.serverAttachmentIds.delete(localId) + cache.serverAttachmentProjectIds.delete(localId) + cleanupPreparedAttachment(previousServerAttachmentId, previousProjectId) + actions.setAttachmentUploading(localId) + actions.uploadAttachment(localId) + }, + removeAttachment: ({ localId }) => { + cache.uploadControllers.get(localId)?.abort() + cache.uploadControllers.delete(localId) + + const previewUrl = cache.previewUrls.get(localId) + if (previewUrl) { + URL.revokeObjectURL(previewUrl) + cache.previewUrls.delete(localId) + } + + const attachmentId = cache.serverAttachmentIds.get(localId) + const projectId = cache.serverAttachmentProjectIds.get(localId) + cache.serverAttachmentIds.delete(localId) + cache.serverAttachmentProjectIds.delete(localId) + actions.setSelectionError(null) + cleanupPreparedAttachment(attachmentId, projectId, true) + }, + markAttachmentsSent: ({ attachmentIds }) => { + for (const [localId, attachmentId] of cache.serverAttachmentIds.entries()) { + if (!attachmentIds.includes(attachmentId)) { + continue + } + cache.serverAttachmentIds.delete(localId) + cache.serverAttachmentProjectIds.delete(localId) + const previewUrl = cache.previewUrls.get(localId) + if (previewUrl) { + URL.revokeObjectURL(previewUrl) + cache.previewUrls.delete(localId) + } + } + }, + } + }), + + afterMount(({ cache }) => { + cache.isUnmounted = false + cache.previewUrls = new Map() + cache.serverAttachmentIds = new Map() + cache.serverAttachmentProjectIds = new Map() + cache.cleanupAttachmentIds = new Set() + cache.uploadControllers = new Map() + }), + + beforeUnmount(({ props, cache }) => { + cache.isUnmounted = true + for (const abortController of cache.uploadControllers.values()) { + abortController.abort() + } + cache.uploadControllers.clear() + + for (const [localId, attachmentId] of cache.serverAttachmentIds.entries()) { + const projectId = cache.serverAttachmentProjectIds.get(localId) + if (projectId && !cache.cleanupAttachmentIds.has(attachmentId)) { + cache.cleanupAttachmentIds.add(attachmentId) + void deletePreparedAttachment(projectId, props.conversationId, attachmentId).catch(() => undefined) + } + } + cache.serverAttachmentIds.clear() + cache.serverAttachmentProjectIds.clear() + + for (const previewUrl of cache.previewUrls.values()) { + URL.revokeObjectURL(previewUrl) + } + cache.previewUrls.clear() + }), +]) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 5b415589c9d7..e59ee3e3f5f4 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -236,6 +236,8 @@ "list_tasks", "pi_cloud_runtime_enabled", "prepare_task_run_artifact_uploads", + "promote_posthog_ai_attachments", + "rollback_posthog_ai_attachments", "prepare_task_staged_artifacts", "presign_task_run_artifact", "presign_task_run_artifact_download", @@ -326,7 +328,7 @@ def _task_run_to_dto(run: TaskRun, *, task: Task | None = None) -> contracts.Tas error_message=run.error_message, output=run.output, state=run.state or {}, - artifacts=run.artifacts or [], + artifacts=_task_run_artifacts(run, include_pending=False), created_at=run.created_at, updated_at=run.updated_at, completed_at=run.completed_at, @@ -430,7 +432,18 @@ def _task_run_log_url(run: TaskRun) -> str | None: return presigned_url -def _task_run_detail_to_dto(run: TaskRun) -> contracts.TaskRunDetailDTO: +_POSTHOG_AI_PENDING_DELIVERY_KEY = "_posthog_ai_pending_delivery" + + +def _task_run_artifacts(run: TaskRun, *, include_pending: bool) -> list[dict]: + return [ + {key: value for key, value in entry.items() if key != _POSTHOG_AI_PENDING_DELIVERY_KEY} + for entry in (run.artifacts or []) + if include_pending or not entry.get(_POSTHOG_AI_PENDING_DELIVERY_KEY) + ] + + +def _task_run_detail_to_dto(run: TaskRun, *, include_pending_artifacts: bool = False) -> contracts.TaskRunDetailDTO: """Map a ``TaskRun`` to its HTTP detail DTO. Reproduces the SMF-derived fields ``TaskRunDetailSerializer`` computed: ``log_url`` does @@ -457,7 +470,7 @@ def _task_run_detail_to_dto(run: TaskRun) -> contracts.TaskRunDetailDTO: error_message=run.error_message, output=run.output, state=_public_task_run_state(run.state), - artifacts=run.artifacts or [], + artifacts=_task_run_artifacts(run, include_pending=include_pending_artifacts), created_at=run.created_at, updated_at=run.updated_at, completed_at=run.completed_at, @@ -2290,10 +2303,18 @@ def list_task_runs(task_id: str | UUID, team_id: int) -> list[contracts.TaskRunD return [_task_run_detail_to_dto(run) for run in runs] -def get_task_run_detail(run_id: str | UUID, task_id: str | UUID, team_id: int) -> contracts.TaskRunDetailDTO | None: +def get_task_run_detail( + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + include_pending_artifacts: bool = False, +) -> contracts.TaskRunDetailDTO | None: """A single run as a detail DTO, scoped to its task + team.""" run = _get_visible_run(run_id, task_id, team_id) - return _task_run_detail_to_dto(run) if run is not None else None + return ( + _task_run_detail_to_dto(run, include_pending_artifacts=include_pending_artifacts) if run is not None else None + ) def get_task_run_stream_info( @@ -3082,7 +3103,7 @@ def upload_task_run_artifacts( {**entry, "url": absolute_uri(_build_artifact_download_path(run, entry["id"]))} if entry.get("id") else dict(entry) - for entry in manifest + for entry in _task_run_artifacts(run, include_pending=uploaded_by == "agent") ] return uploaded, response_manifest @@ -3185,6 +3206,8 @@ def finalize_task_run_artifact_uploads( existing_entry = _find_artifact_manifest_entry(manifest, artifact_id, storage_path) if existing_entry is not None: + if existing_entry.get(_POSTHOG_AI_PENDING_DELIVERY_KEY) and uploaded_by != "agent": + return None, "Artifact upload not found in object storage" finalized_entries.append(existing_entry) continue @@ -3255,6 +3278,196 @@ def finalize_task_run_artifact_uploads( return response_entries, None +def _validate_posthog_ai_attachment( + attachment: dict[str, Any], *, team_id: int, user_id: int, conversation_id: str +) -> None: + attachment_id = str(attachment.get("id", "")) + try: + UUID(attachment_id) + except (TypeError, ValueError) as error: + raise ValueError("PostHog AI attachment IDs must be UUIDs.") from error + + content_type = attachment.get("content_type") + extension = {"image/png": "png", "image/jpeg": "jpg"}.get(content_type) if isinstance(content_type, str) else None + if extension is None: + raise ValueError("PostHog AI attachments must be PNG or JPEG images.") + if ( + attachment.get("team_id") != team_id + or attachment.get("user_id") != user_id + or attachment.get("conversation_id") != conversation_id + or attachment.get("status") not in {"finalized", "promoted"} + ): + raise ValueError("PostHog AI attachment scope does not match the sandbox run.") + + normalized_size = attachment.get("normalized_size") + width = attachment.get("width") + height = attachment.get("height") + if ( + isinstance(normalized_size, bool) + or not isinstance(normalized_size, int) + or not 0 < normalized_size <= 4 * 1024 * 1024 + or isinstance(width, bool) + or not isinstance(width, int) + or width <= 0 + or isinstance(height, bool) + or not isinstance(height, int) + or height <= 0 + or width * height >= 20_000_000 + ): + raise ValueError("PostHog AI attachment dimensions or size are invalid.") + + expected_storage_path = ( + f"posthog_ai/assistant_attachments/team_{team_id}/user_{user_id}/" + f"conversation_{conversation_id}/normalized/{attachment_id}.{extension}" + ) + if attachment.get("normalized_storage_path") != expected_storage_path: + raise ValueError("PostHog AI attachment storage path is invalid.") + + +def promote_posthog_ai_attachments( + *, + task_id: str | UUID, + run_id: str | UUID, + team_id: int, + user_id: int | None, + conversation_id: str, + attachments: list[dict[str, Any]], +) -> contracts.PostHogAIAttachmentPromotionResult | None: + from posthog.storage import object_storage # noqa: PLC0415 — keeps storage deps off facade import path + + run = _get_visible_run(run_id, task_id, team_id) + if run is None or user_id is None or user_id <= 0: + return None + if not 0 < len(attachments) <= 4 or len({str(attachment.get("id")) for attachment in attachments}) != len( + attachments + ): + raise ValueError("Provide between 1 and 4 unique PostHog AI attachments.") + for attachment in attachments: + _validate_posthog_ai_attachment( + attachment, + team_id=team_id, + user_id=user_id, + conversation_id=conversation_id, + ) + if sum(int(attachment["normalized_size"]) for attachment in attachments) > 10 * 1024 * 1024: + raise ValueError("PostHog AI attachments exceed the per-message size limit.") + + prepared_entries: list[dict[str, Any]] = [] + copied_targets_by_id: dict[str, str] = {} + existing_ids = {str(entry.get("id")) for entry in (run.artifacts or []) if entry.get("id")} + for attachment in attachments: + artifact_id = str(attachment["id"]) + if artifact_id in existing_ids: + continue + source_storage_path = str(attachment["normalized_storage_path"]) + safe_name, target_storage_path = _build_artifact_storage_path(run, artifact_id, str(attachment["file_name"])) + try: + object_storage.copy(source_storage_path, target_storage_path) + copied_targets_by_id[artifact_id] = target_storage_path + object_storage.tag(target_storage_path, {"ttl_days": "30", "team_id": str(run.team_id)}) + except Exception as error: + for copied_target in copied_targets_by_id.values(): + try: + object_storage.delete(copied_target) + except Exception: + pass + raise RuntimeError("Could not copy the image attachment into the sandbox run.") from error + entry = _build_artifact_manifest_entry( + artifact_id=artifact_id, + name=safe_name, + artifact_type="context", + source="user_attachment", + size=int(attachment["normalized_size"]), + content_type=str(attachment["content_type"]), + storage_path=target_storage_path, + uploaded_at=django_timezone.now().isoformat(), + metadata={ + "conversation_id": conversation_id, + "width": int(attachment["width"]), + "height": int(attachment["height"]), + }, + ) + entry[_POSTHOG_AI_PENDING_DELIVERY_KEY] = True + entry["uploaded_by"] = "user" + if user_id is not None: + entry["uploaded_by_user_id"] = user_id + prepared_entries.append(entry) + + newly_attached_ids: list[str] = [] + try: + if prepared_entries: + with transaction.atomic(): + locked_run = TaskRun.objects.select_for_update().get(pk=run.pk) + existing_ids = {str(entry.get("id")) for entry in (locked_run.artifacts or []) if entry.get("id")} + entries_to_append = [entry for entry in prepared_entries if str(entry["id"]) not in existing_ids] + newly_attached_ids = [str(entry["id"]) for entry in entries_to_append] + if entries_to_append: + merged = list(locked_run.artifacts or []) + merged.extend(entries_to_append) + _save_artifact_manifest(locked_run, merged) + except Exception: + for copied_target in copied_targets_by_id.values(): + try: + object_storage.delete(copied_target) + except Exception: + pass + raise + + return contracts.PostHogAIAttachmentPromotionResult( + attached_ids=[str(attachment["id"]) for attachment in attachments], + newly_attached_ids=newly_attached_ids, + ) + + +def rollback_posthog_ai_attachments( + *, task_id: str | UUID, run_id: str | UUID, team_id: int, attachment_ids: list[str] +) -> None: + from posthog.storage import object_storage # noqa: PLC0415 — keeps storage deps off facade import path + + run = _get_visible_run(run_id, task_id, team_id) + if run is None or not attachment_ids: + return + attachment_id_set = set(attachment_ids) + with transaction.atomic(): + locked_run = TaskRun.objects.select_for_update().get(pk=run.pk) + manifest = list(locked_run.artifacts or []) + removed = [ + entry + for entry in manifest + if entry.get("id") in attachment_id_set and entry.get(_POSTHOG_AI_PENDING_DELIVERY_KEY) + ] + kept = [entry for entry in manifest if entry not in removed] + _save_artifact_manifest(locked_run, kept) + for entry in removed: + storage_path = entry.get("storage_path") + if not storage_path: + continue + try: + object_storage.delete(str(storage_path)) + except Exception: + pass + + +def commit_posthog_ai_attachments( + *, task_id: str | UUID, run_id: str | UUID, team_id: int, attachment_ids: list[str] +) -> None: + run = _get_visible_run(run_id, task_id, team_id) + if run is None or not attachment_ids: + return + attachment_id_set = set(attachment_ids) + with transaction.atomic(): + locked_run = TaskRun.objects.select_for_update().get(pk=run.pk) + manifest = [ + ( + {key: value for key, value in entry.items() if key != _POSTHOG_AI_PENDING_DELIVERY_KEY} + if entry.get("id") in attachment_id_set + else entry + ) + for entry in (locked_run.artifacts or []) + ] + _save_artifact_manifest(locked_run, manifest) + + def list_task_run_living_artifacts(run_id: str | UUID, task_id: str | UUID, team_id: int) -> list[dict] | None: from products.tasks.backend.logic.services.living_artifacts import ( # noqa: PLC0415 — keep storage deps off the api import path get_task_artifacts_for_run, @@ -3363,7 +3576,12 @@ def edit_task_run_living_artifact( def presign_task_run_artifact( - run_id: str | UUID, task_id: str | UUID, team_id: int, *, storage_path: str + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + storage_path: str, + include_pending_artifacts: bool = False, ) -> tuple[str | None, str | None]: """Presign a download URL for an artifact on the run. @@ -3376,7 +3594,7 @@ def presign_task_run_artifact( if run is None: return None, None - artifacts = run.artifacts or [] + artifacts = _task_run_artifacts(run, include_pending=include_pending_artifacts) if not any(artifact.get("storage_path") == storage_path for artifact in artifacts): return None, "not_found" @@ -3409,7 +3627,8 @@ def set_task_run_artifacts_dismissed( locked_run = TaskRun.objects.select_for_update().get(pk=run.pk) manifest = list(locked_run.artifacts or []) requested = set(artifact_ids) - if not requested.issubset({entry.get("id") for entry in manifest}): + visible_ids = {entry.get("id") for entry in manifest if not entry.get(_POSTHOG_AI_PENDING_DELIVERY_KEY)} + if not requested.issubset(visible_ids): return None, "not_found" # Restoring drops the key rather than nulling it, so a manifest entry only ever carries @@ -3425,11 +3644,16 @@ def set_task_run_artifacts_dismissed( ] _save_artifact_manifest(locked_run, manifest) - return manifest, None + return _task_run_artifacts(locked_run, include_pending=False), None def presign_task_run_artifact_download( - run_id: str | UUID, task_id: str | UUID, team_id: int, *, artifact_id: str + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + artifact_id: str, + include_pending_artifacts: bool = False, ) -> tuple[str | None, str | None]: """Presign a download URL for an artifact addressed by its manifest id. @@ -3442,7 +3666,14 @@ def presign_task_run_artifact_download( if run is None: return None, None - entry = next((a for a in run.artifacts or [] if a.get("id") == artifact_id), None) + entry = next( + ( + artifact + for artifact in _task_run_artifacts(run, include_pending=include_pending_artifacts) + if artifact.get("id") == artifact_id + ), + None, + ) if entry is None: return None, "not_found" @@ -3458,7 +3689,12 @@ def presign_task_run_artifact_download( def read_task_run_artifact( - run_id: str | UUID, task_id: str | UUID, team_id: int, *, storage_path: str + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + storage_path: str, + include_pending_artifacts: bool = False, ) -> tuple[bytes | None, dict | None, str | None]: """Read artifact bytes for download, walking the resume chain. @@ -3473,8 +3709,9 @@ def read_task_run_artifact( return None, None, None artifact = run.find_artifact_in_resume_chain(storage_path) - if artifact is None: + if artifact is None or (artifact.get(_POSTHOG_AI_PENDING_DELIVERY_KEY) and not include_pending_artifacts): return None, None, "not_found" + artifact = {key: value for key, value in artifact.items() if key != _POSTHOG_AI_PENDING_DELIVERY_KEY} try: content = object_storage.read_bytes(storage_path, missing_ok=True) diff --git a/products/tasks/backend/facade/contracts.py b/products/tasks/backend/facade/contracts.py index 2c715e200bee..769144281c18 100644 --- a/products/tasks/backend/facade/contracts.py +++ b/products/tasks/backend/facade/contracts.py @@ -435,6 +435,14 @@ class StagedArtifactFinalizeResult: error: str | None = None +@dataclass(frozen=True) +class PostHogAIAttachmentPromotionResult: + """Outcome of promoting sandbox image attachments into a task run's artifact manifest.""" + + attached_ids: list[str] + newly_attached_ids: list[str] + + @dataclass(frozen=True) class SlackThreadContextRepoResearchDTO: """The internal sandbox run the discovery agent used to pick a run's repo.""" diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 4bb59c291151..20f5a020f7c8 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -242,7 +242,7 @@ class TaskRunUpdateSerializer(serializers.Serializer): ) -class TaskRunArtifactMetadataSerializer(serializers.Serializer): +class TaskRunSkillBundleMetadataSerializer(serializers.Serializer): skill_name = serializers.CharField( allow_blank=False, max_length=255, @@ -266,16 +266,45 @@ class TaskRunArtifactMetadataSerializer(serializers.Serializer): ) +class TaskRunImageAttachmentMetadataSerializer(serializers.Serializer): + conversation_id = serializers.UUIDField( + help_text="PostHog AI conversation that supplied a user attachment.", + ) + width = serializers.IntegerField( + min_value=1, + help_text="Normalized image width in pixels for a user attachment.", + ) + height = serializers.IntegerField( + min_value=1, + help_text="Normalized image height in pixels for a user attachment.", + ) + + +@extend_schema_field( + PolymorphicProxySerializer( + component_name="TaskRunArtifactMetadata", + serializers=[TaskRunSkillBundleMetadataSerializer, TaskRunImageAttachmentMetadataSerializer], + resource_type_field_name=None, + ) +) +class TaskRunArtifactMetadataField(serializers.JSONField): + pass + + def validate_task_run_artifact_metadata(attrs: dict[str, Any]) -> dict[str, Any]: artifact_type = attrs.get("type") metadata = attrs.get("metadata") if artifact_type != "skill_bundle": + if metadata: + TaskRunImageAttachmentMetadataSerializer(data=metadata).is_valid(raise_exception=True) return attrs if not metadata: raise serializers.ValidationError({"metadata": "Skill bundle artifacts require metadata"}) + TaskRunSkillBundleMetadataSerializer(data=metadata).is_valid(raise_exception=True) + return attrs @@ -290,7 +319,7 @@ class TaskRunArtifactResponseSerializer(serializers.Serializer): ) size = serializers.IntegerField(required=False, help_text="Artifact size in bytes") content_type = serializers.CharField(required=False, allow_blank=True, help_text="Optional MIME type") - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -914,7 +943,7 @@ class TaskRunArtifactUploadSerializer(serializers.Serializer): allow_blank=True, help_text="Optional MIME type for the artifact", ) - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1201,7 +1230,7 @@ class TaskRunArtifactPrepareUploadSerializer(serializers.Serializer): allow_blank=True, help_text="Optional MIME type for the artifact upload", ) - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1248,7 +1277,7 @@ class TaskRunArtifactPrepareUploadResponseSerializer(serializers.Serializer): ) size = serializers.IntegerField(help_text="Expected upload size in bytes") content_type = serializers.CharField(required=False, allow_blank=True, help_text="Optional MIME type") - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1281,7 +1310,7 @@ class TaskRunArtifactFinalizeUploadSerializer(serializers.Serializer): allow_blank=True, help_text="Optional MIME type recorded for the artifact", ) - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1324,7 +1353,7 @@ class TaskStagedArtifactPrepareUploadSerializer(serializers.Serializer): allow_blank=True, help_text="Optional MIME type for the artifact upload", ) - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1365,7 +1394,7 @@ class TaskStagedArtifactPrepareUploadResponseSerializer(serializers.Serializer): ) size = serializers.IntegerField(help_text="Expected upload size in bytes") content_type = serializers.CharField(required=False, allow_blank=True, help_text="Optional MIME type") - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) @@ -1404,7 +1433,7 @@ class TaskStagedArtifactFinalizeUploadSerializer(serializers.Serializer): allow_blank=True, help_text="Optional MIME type recorded for the artifact", ) - metadata = TaskRunArtifactMetadataSerializer( + metadata = TaskRunArtifactMetadataField( required=False, help_text="Optional structured metadata for special artifact types, such as skill bundles.", ) diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 41cdf45f0900..65e5fa9fa0eb 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -1218,7 +1218,12 @@ def _ensure_task_accessible(self) -> str: def _get_run_or_404(self, pk) -> tasks_contracts.TaskRunDetailDTO: task_id = self._ensure_task_accessible() - run = tasks_facade.get_task_run_detail(pk, task_id, self.team_id) + run = tasks_facade.get_task_run_detail( + pk, + task_id, + self.team_id, + include_pending_artifacts=self._is_sandbox_agent_request(task_id), + ) if run is None: raise NotFound() return run @@ -1848,7 +1853,13 @@ def artifacts_finalize_upload(self, request, pk=None, **kwargs): def artifacts_presign(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() storage_path = request.validated_data["storage_path"] - url, error = tasks_facade.presign_task_run_artifact(pk, task_id, self.team_id, storage_path=storage_path) + url, error = tasks_facade.presign_task_run_artifact( + pk, + task_id, + self.team_id, + storage_path=storage_path, + include_pending_artifacts=self._is_sandbox_agent_request(task_id), + ) if url is None and error is None: raise NotFound() if error == "not_found": @@ -1930,7 +1941,11 @@ def artifacts_download(self, request, pk=None, **kwargs): # Walk the resume chain so cloud→cloud resume runs can fetch the git checkpoint # pack/index that lives on the prior run they were forked from. content, artifact, error = tasks_facade.read_task_run_artifact( - pk, task_id, self.team_id, storage_path=storage_path + pk, + task_id, + self.team_id, + storage_path=storage_path, + include_pending_artifacts=self._is_sandbox_agent_request(task_id), ) if content is None and artifact is None and error is None: raise NotFound() @@ -1992,7 +2007,13 @@ def artifacts_download(self, request, pk=None, **kwargs): ) def artifacts_download_by_id(self, request, pk=None, artifact_id=None, **kwargs): task_id = self._ensure_task_accessible() - url, error = tasks_facade.presign_task_run_artifact_download(pk, task_id, self.team_id, artifact_id=artifact_id) + url, error = tasks_facade.presign_task_run_artifact_download( + pk, + task_id, + self.team_id, + artifact_id=artifact_id, + include_pending_artifacts=self._is_sandbox_agent_request(task_id), + ) if error == "unavailable": return Response( TaskRunErrorResponseSerializer({"error": "Unable to generate download URL"}).data, diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index f69846fb53f3..2ca37ac42179 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -6331,6 +6331,31 @@ def test_upload_artifacts_rejects_skill_bundle_without_metadata(self): self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("metadata", json.dumps(response.json())) + def test_upload_artifacts_rejects_partial_skill_bundle_metadata(self): + task = self.create_task() + run = TaskRun.objects.create(task=task, team=self.team, status=TaskRun.Status.IN_PROGRESS) + + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/artifacts/", + { + "artifacts": [ + { + "name": "local-skill.zip", + "type": "skill_bundle", + "source": "posthog_code_skill", + "content": "c2tpbGw=", + "content_encoding": "base64", + "content_type": "application/zip", + "metadata": {"skill_name": "partial-skill"}, + } + ] + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertIn("skill_source", json.dumps(response.json())) + def test_upload_artifacts_rejects_invalid_base64_content(self): task = self.create_task() run = TaskRun.objects.create(task=task, team=self.team, status=TaskRun.Status.IN_PROGRESS) diff --git a/products/tasks/backend/tests/test_posthog_ai_attachment_promotion.py b/products/tasks/backend/tests/test_posthog_ai_attachment_promotion.py new file mode 100644 index 000000000000..40bd4d319e00 --- /dev/null +++ b/products/tasks/backend/tests/test_posthog_ai_attachment_promotion.py @@ -0,0 +1,294 @@ +import uuid +from typing import Any + +from posthog.test.base import APIBaseTest +from unittest.mock import patch + +from parameterized import parameterized + +from posthog.models.user import User + +from products.tasks.backend.facade import api as tasks_facade +from products.tasks.backend.models import Task, TaskRun + + +class TestPosthogAIAttachmentPromotionFacade(APIBaseTest): + def _task_and_run(self) -> tuple[Task, TaskRun]: + task = Task.objects.create( + team=self.team, + title="t", + description="d", + origin_product=Task.OriginProduct.POSTHOG_AI, + created_by=self.user, + ) + run = task.create_run(mode="interactive") + return task, run + + def _attachment(self, *, attachment_id: str, conversation_id: str) -> dict[str, Any]: + return { + "id": attachment_id, + "team_id": self.team.id, + "user_id": self.user.id, + "conversation_id": conversation_id, + "file_name": "image.png", + "normalized_storage_path": ( + f"posthog_ai/assistant_attachments/team_{self.team.id}/user_{self.user.id}/" + f"conversation_{conversation_id}/normalized/{attachment_id}.png" + ), + "normalized_size": 123, + "content_type": "image/png", + "width": 10, + "height": 12, + "status": "finalized", + } + + def test_tag_failure_leaves_no_manifest_entry(self) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + + with ( + patch("posthog.storage.object_storage.copy"), + patch("posthog.storage.object_storage.tag", side_effect=RuntimeError("tag failed")), + patch("posthog.storage.object_storage.delete") as m_delete, + ): + with self.assertRaises(RuntimeError): + tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + run.refresh_from_db() + self.assertEqual(run.artifacts, []) + self.assertEqual(m_delete.call_count, 1) + + def test_existing_artifact_id_is_not_duplicated(self) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + run.artifacts = [ + { + "id": attachment_id, + "name": "image.png", + "type": "context", + "source": "user_attachment", + "size": 123, + "content_type": "image/png", + "storage_path": "existing/path.png", + "uploaded_at": "2026-01-01T00:00:00+00:00", + "metadata": {"conversation_id": conversation_id, "width": 10, "height": 12}, + } + ] + run.save(update_fields=["artifacts", "updated_at"]) + + with ( + patch("posthog.storage.object_storage.copy") as copy, + patch("posthog.storage.object_storage.tag") as tag, + ): + result = tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + run.refresh_from_db() + assert result is not None + self.assertEqual(result.attached_ids, [attachment_id]) + self.assertEqual(result.newly_attached_ids, []) + self.assertEqual([entry["id"] for entry in run.artifacts], [attachment_id]) + copy.assert_not_called() + tag.assert_not_called() + + def test_team_member_can_promote_attachment_to_task_created_by_another_user(self) -> None: + creator = User.objects.create_user(email="creator@example.com", first_name="Creator", password="password") + task = Task.objects.create( + team=self.team, + title="t", + description="d", + origin_product=Task.OriginProduct.POSTHOG_AI, + created_by=creator, + ) + run = task.create_run(mode="interactive") + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + + with ( + patch("posthog.storage.object_storage.copy"), + patch("posthog.storage.object_storage.tag"), + ): + result = tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + assert result is not None + self.assertEqual(result.attached_ids, [attachment_id]) + + def test_pending_attachment_is_hidden_until_delivery_is_committed(self) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + + with ( + patch("posthog.storage.object_storage.copy"), + patch("posthog.storage.object_storage.tag"), + ): + result = tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + response = self.client.get(f"/api/projects/@current/tasks/{task.id}/") + + self.assertEqual(response.status_code, 200) + self.assertEqual(response.json()["latest_run"]["artifacts"], []) + agent_run = tasks_facade.get_task_run_detail( + run.id, + task.id, + self.team.id, + include_pending_artifacts=True, + ) + assert agent_run is not None + self.assertEqual([artifact["id"] for artifact in agent_run.artifacts], [attachment_id]) + self.assertNotIn("_posthog_ai_pending_delivery", agent_run.artifacts[0]) + run.refresh_from_db() + storage_path = run.artifacts[0]["storage_path"] + with patch("posthog.storage.object_storage.get_presigned_url", return_value="https://example.com/image"): + self.assertEqual( + tasks_facade.presign_task_run_artifact( + run.id, + task.id, + self.team.id, + storage_path=storage_path, + ), + (None, "not_found"), + ) + self.assertEqual( + tasks_facade.presign_task_run_artifact( + run.id, + task.id, + self.team.id, + storage_path=storage_path, + include_pending_artifacts=True, + ), + ("https://example.com/image", None), + ) + + assert result is not None + tasks_facade.commit_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + attachment_ids=result.newly_attached_ids, + ) + response = self.client.get(f"/api/projects/@current/tasks/{task.id}/") + + self.assertEqual(response.status_code, 200) + artifact = response.json()["latest_run"]["artifacts"][0] + self.assertEqual(artifact["id"], attachment_id) + self.assertEqual( + artifact["metadata"], + {"conversation_id": conversation_id, "width": 10, "height": 12}, + ) + + def test_failed_rollback_leaves_pending_attachment_hidden(self) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + + with ( + patch("posthog.storage.object_storage.copy"), + patch("posthog.storage.object_storage.tag"), + ): + tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + with patch( + "products.tasks.backend.facade.api._save_artifact_manifest", side_effect=RuntimeError("save failed") + ): + with self.assertRaisesRegex(RuntimeError, "save failed"): + tasks_facade.rollback_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + attachment_ids=[attachment_id], + ) + + run.refresh_from_db() + self.assertTrue(run.artifacts[0]["_posthog_ai_pending_delivery"]) + public_run = tasks_facade.get_task_run_detail(run.id, task.id, self.team.id) + assert public_run is not None + self.assertEqual(public_run.artifacts, []) + + def test_manifest_failure_deletes_every_copied_object(self) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + + with ( + patch("posthog.storage.object_storage.copy"), + patch("posthog.storage.object_storage.tag"), + patch("posthog.storage.object_storage.delete") as delete, + patch("products.tasks.backend.facade.api._save_artifact_manifest", side_effect=RuntimeError("save failed")), + ): + with self.assertRaisesRegex(RuntimeError, "save failed"): + tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[self._attachment(attachment_id=attachment_id, conversation_id=conversation_id)], + ) + + delete.assert_called_once() + + @parameterized.expand( + [ + ("wrong_team", "team_id", -1), + ("unsupported_type", "content_type", "text/plain"), + ("oversized", "normalized_size", 4 * 1024 * 1024 + 1), + ("too_many_pixels", "width", 20_000_000), + ("untrusted_path", "normalized_storage_path", "https://example.com/image.png"), + ] + ) + def test_rejects_unsafe_manifest_fields(self, _name: str, field: str, value: Any) -> None: + task, run = self._task_and_run() + attachment_id = str(uuid.uuid4()) + conversation_id = str(uuid.uuid4()) + attachment = self._attachment(attachment_id=attachment_id, conversation_id=conversation_id) + attachment[field] = value + + with patch("posthog.storage.object_storage.copy") as copy: + with self.assertRaises(ValueError): + tasks_facade.promote_posthog_ai_attachments( + task_id=task.id, + run_id=run.id, + team_id=self.team.id, + user_id=self.user.id, + conversation_id=conversation_id, + attachments=[attachment], + ) + + copy.assert_not_called() diff --git a/products/tasks/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index 0da815226e2a..2f5d22e1837b 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -1418,7 +1418,7 @@ export const TaskRunReasoningEffortEnumApi = { Ultracode: 'ultracode', } as const -export interface TaskRunArtifactMetadataApi { +export interface TaskRunSkillBundleMetadataApi { /** * Name of the local skill included in a skill_bundle artifact. * @maxLength 255 @@ -1447,6 +1447,23 @@ export interface TaskRunArtifactMetadataApi { schema_version: number } +export interface TaskRunImageAttachmentMetadataApi { + /** PostHog AI conversation that supplied a user attachment. */ + conversation_id: string + /** + * Normalized image width in pixels for a user attachment. + * @minimum 1 + */ + width: number + /** + * Normalized image height in pixels for a user attachment. + * @minimum 1 + */ + height: number +} + +export type TaskRunArtifactMetadataApi = TaskRunSkillBundleMetadataApi | TaskRunImageAttachmentMetadataApi + /** * * `agent` - agent * * `user` - user diff --git a/products/tasks/frontend/generated/api.zod.ts b/products/tasks/frontend/generated/api.zod.ts index 06eeb2db49d8..ed6126162a37 100644 --- a/products/tasks/frontend/generated/api.zod.ts +++ b/products/tasks/frontend/generated/api.zod.ts @@ -2050,9 +2050,9 @@ export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemStoragePat export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemContentTypeMax = 255 -export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneSkillNameMax = 255 +export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax = 255 -export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp = new RegExp( +export const tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp = new RegExp( '^[a-f0-9]{64}$' ) @@ -2097,34 +2097,51 @@ export const TasksStagedArtifactsFinalizeUploadCreateBody = /* @__PURE__ */ zod. .optional() .describe('Optional MIME type recorded for the artifact'), metadata: zod - .object({ - skill_name: zod - .string() - .max(tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneSkillNameMax) - .describe('Name of the local skill included in a skill_bundle artifact.'), - skill_source: zod - .enum(['user', 'repo', 'marketplace', 'codex']) - .describe( - '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ) - .describe( - 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ), - content_sha256: zod - .string() - .regex( - tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp - ) - .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), - bundle_format: zod - .enum(['zip']) - .describe('\* `zip` - zip') - .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), - schema_version: zod - .number() - .min(1) - .describe('Version of the local skill bundle metadata schema.'), - }) + .union([ + zod.object({ + skill_name: zod + .string() + .max( + tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax + ) + .describe('Name of the local skill included in a skill_bundle artifact.'), + skill_source: zod + .enum(['user', 'repo', 'marketplace', 'codex']) + .describe( + '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ) + .describe( + 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ), + content_sha256: zod + .string() + .regex( + tasksStagedArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp + ) + .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), + bundle_format: zod + .enum(['zip']) + .describe('\* `zip` - zip') + .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), + schema_version: zod + .number() + .min(1) + .describe('Version of the local skill bundle metadata schema.'), + }), + zod.object({ + conversation_id: zod + .uuid() + .describe('PostHog AI conversation that supplied a user attachment.'), + width: zod + .number() + .min(1) + .describe('Normalized image width in pixels for a user attachment.'), + height: zod + .number() + .min(1) + .describe('Normalized image height in pixels for a user attachment.'), + }), + ]) .optional() .describe('Optional structured metadata for special artifact types, such as skill bundles.'), }) @@ -2145,9 +2162,9 @@ export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemSizeMax = 3 export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemContentTypeMax = 255 -export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneSkillNameMax = 255 +export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax = 255 -export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp = new RegExp( +export const tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp = new RegExp( '^[a-f0-9]{64}$' ) @@ -2192,34 +2209,49 @@ export const TasksStagedArtifactsPrepareUploadCreateBody = /* @__PURE__ */ zod.o .optional() .describe('Optional MIME type for the artifact upload'), metadata: zod - .object({ - skill_name: zod - .string() - .max(tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneSkillNameMax) - .describe('Name of the local skill included in a skill_bundle artifact.'), - skill_source: zod - .enum(['user', 'repo', 'marketplace', 'codex']) - .describe( - '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ) - .describe( - 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ), - content_sha256: zod - .string() - .regex( - tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp - ) - .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), - bundle_format: zod - .enum(['zip']) - .describe('\* `zip` - zip') - .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), - schema_version: zod - .number() - .min(1) - .describe('Version of the local skill bundle metadata schema.'), - }) + .union([ + zod.object({ + skill_name: zod + .string() + .max(tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax) + .describe('Name of the local skill included in a skill_bundle artifact.'), + skill_source: zod + .enum(['user', 'repo', 'marketplace', 'codex']) + .describe( + '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ) + .describe( + 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ), + content_sha256: zod + .string() + .regex( + tasksStagedArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp + ) + .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), + bundle_format: zod + .enum(['zip']) + .describe('\* `zip` - zip') + .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), + schema_version: zod + .number() + .min(1) + .describe('Version of the local skill bundle metadata schema.'), + }), + zod.object({ + conversation_id: zod + .uuid() + .describe('PostHog AI conversation that supplied a user attachment.'), + width: zod + .number() + .min(1) + .describe('Normalized image width in pixels for a user attachment.'), + height: zod + .number() + .min(1) + .describe('Normalized image height in pixels for a user attachment.'), + }), + ]) .optional() .describe('Optional structured metadata for special artifact types, such as skill bundles.'), }) @@ -2433,9 +2465,9 @@ export const tasksRunsArtifactsCreateBodyArtifactsItemSourceMax = 64 export const tasksRunsArtifactsCreateBodyArtifactsItemContentEncodingDefault = `utf-8` export const tasksRunsArtifactsCreateBodyArtifactsItemContentTypeMax = 255 -export const tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneSkillNameMax = 255 +export const tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneOneSkillNameMax = 255 -export const tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneContentSha256RegExp = new RegExp('^[a-f0-9]{64}$') +export const tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp = new RegExp('^[a-f0-9]{64}$') export const TasksRunsArtifactsCreateBody = /* @__PURE__ */ zod.object({ artifacts: zod @@ -2481,32 +2513,47 @@ export const TasksRunsArtifactsCreateBody = /* @__PURE__ */ zod.object({ .optional() .describe('Optional MIME type for the artifact'), metadata: zod - .object({ - skill_name: zod - .string() - .max(tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneSkillNameMax) - .describe('Name of the local skill included in a skill_bundle artifact.'), - skill_source: zod - .enum(['user', 'repo', 'marketplace', 'codex']) - .describe( - '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ) - .describe( - 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ), - content_sha256: zod - .string() - .regex(tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneContentSha256RegExp) - .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), - bundle_format: zod - .enum(['zip']) - .describe('\* `zip` - zip') - .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), - schema_version: zod - .number() - .min(1) - .describe('Version of the local skill bundle metadata schema.'), - }) + .union([ + zod.object({ + skill_name: zod + .string() + .max(tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneOneSkillNameMax) + .describe('Name of the local skill included in a skill_bundle artifact.'), + skill_source: zod + .enum(['user', 'repo', 'marketplace', 'codex']) + .describe( + '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ) + .describe( + 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ), + content_sha256: zod + .string() + .regex(tasksRunsArtifactsCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp) + .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), + bundle_format: zod + .enum(['zip']) + .describe('\* `zip` - zip') + .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), + schema_version: zod + .number() + .min(1) + .describe('Version of the local skill bundle metadata schema.'), + }), + zod.object({ + conversation_id: zod + .uuid() + .describe('PostHog AI conversation that supplied a user attachment.'), + width: zod + .number() + .min(1) + .describe('Normalized image width in pixels for a user attachment.'), + height: zod + .number() + .min(1) + .describe('Normalized image height in pixels for a user attachment.'), + }), + ]) .optional() .describe('Optional structured metadata for special artifact types, such as skill bundles.'), }) @@ -2563,9 +2610,9 @@ export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemStoragePathM export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemContentTypeMax = 255 -export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneSkillNameMax = 255 +export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax = 255 -export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp = new RegExp( +export const tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp = new RegExp( '^[a-f0-9]{64}$' ) @@ -2610,34 +2657,49 @@ export const TasksRunsArtifactsFinalizeUploadCreateBody = /* @__PURE__ */ zod.ob .optional() .describe('Optional MIME type recorded for the artifact'), metadata: zod - .object({ - skill_name: zod - .string() - .max(tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneSkillNameMax) - .describe('Name of the local skill included in a skill_bundle artifact.'), - skill_source: zod - .enum(['user', 'repo', 'marketplace', 'codex']) - .describe( - '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ) - .describe( - 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ), - content_sha256: zod - .string() - .regex( - tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp - ) - .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), - bundle_format: zod - .enum(['zip']) - .describe('\* `zip` - zip') - .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), - schema_version: zod - .number() - .min(1) - .describe('Version of the local skill bundle metadata schema.'), - }) + .union([ + zod.object({ + skill_name: zod + .string() + .max(tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax) + .describe('Name of the local skill included in a skill_bundle artifact.'), + skill_source: zod + .enum(['user', 'repo', 'marketplace', 'codex']) + .describe( + '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ) + .describe( + 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ), + content_sha256: zod + .string() + .regex( + tasksRunsArtifactsFinalizeUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp + ) + .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), + bundle_format: zod + .enum(['zip']) + .describe('\* `zip` - zip') + .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), + schema_version: zod + .number() + .min(1) + .describe('Version of the local skill bundle metadata schema.'), + }), + zod.object({ + conversation_id: zod + .uuid() + .describe('PostHog AI conversation that supplied a user attachment.'), + width: zod + .number() + .min(1) + .describe('Normalized image width in pixels for a user attachment.'), + height: zod + .number() + .min(1) + .describe('Normalized image height in pixels for a user attachment.'), + }), + ]) .optional() .describe('Optional structured metadata for special artifact types, such as skill bundles.'), }) @@ -2658,9 +2720,9 @@ export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemSizeMax = 314 export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemContentTypeMax = 255 -export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneSkillNameMax = 255 +export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax = 255 -export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp = new RegExp( +export const tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp = new RegExp( '^[a-f0-9]{64}$' ) @@ -2705,32 +2767,49 @@ export const TasksRunsArtifactsPrepareUploadCreateBody = /* @__PURE__ */ zod.obj .optional() .describe('Optional MIME type for the artifact upload'), metadata: zod - .object({ - skill_name: zod - .string() - .max(tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneSkillNameMax) - .describe('Name of the local skill included in a skill_bundle artifact.'), - skill_source: zod - .enum(['user', 'repo', 'marketplace', 'codex']) - .describe( - '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ) - .describe( - 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' - ), - content_sha256: zod - .string() - .regex(tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneContentSha256RegExp) - .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), - bundle_format: zod - .enum(['zip']) - .describe('\* `zip` - zip') - .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), - schema_version: zod - .number() - .min(1) - .describe('Version of the local skill bundle metadata schema.'), - }) + .union([ + zod.object({ + skill_name: zod + .string() + .max(tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneSkillNameMax) + .describe('Name of the local skill included in a skill_bundle artifact.'), + skill_source: zod + .enum(['user', 'repo', 'marketplace', 'codex']) + .describe( + '\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ) + .describe( + 'Local source for the uploaded skill bundle, such as user or repo.\n\n\* `user` - user\n\* `repo` - repo\n\* `marketplace` - marketplace\n\* `codex` - codex' + ), + content_sha256: zod + .string() + .regex( + tasksRunsArtifactsPrepareUploadCreateBodyArtifactsItemMetadataOneOneContentSha256RegExp + ) + .describe('SHA-256 hex digest of the uploaded skill bundle bytes.'), + bundle_format: zod + .enum(['zip']) + .describe('\* `zip` - zip') + .describe('Archive format used for the local skill bundle.\n\n\* `zip` - zip'), + schema_version: zod + .number() + .min(1) + .describe('Version of the local skill bundle metadata schema.'), + }), + zod.object({ + conversation_id: zod + .uuid() + .describe('PostHog AI conversation that supplied a user attachment.'), + width: zod + .number() + .min(1) + .describe('Normalized image width in pixels for a user attachment.'), + height: zod + .number() + .min(1) + .describe('Normalized image height in pixels for a user attachment.'), + }), + ]) .optional() .describe('Optional structured metadata for special artifact types, such as skill bundles.'), }) diff --git a/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index 1340f786de4c..fcd28d8a90eb 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -10202,6 +10202,18 @@ export namespace Schemas { Role: 'role', } as const; + /** + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg + */ + export type AssistantAttachmentContentTypeEnum = typeof AssistantAttachmentContentTypeEnum[keyof typeof AssistantAttachmentContentTypeEnum]; + + + export const AssistantAttachmentContentTypeEnum = { + ImagePng: 'image/png', + ImageJpeg: 'image/jpeg', + } as const; + export interface AsyncDeletionStatus { /** The UUID of the person whose events are queued for deletion. */ person_uuid: string; @@ -10216,6 +10228,96 @@ export namespace Schemas { delete_verified_at: string | null; } + export interface AttachmentDeleteRequest { + /** Conversation UUID the image attachment belongs to. */ + conversation_id: string; + /** Single attachment ID to delete from staging. */ + attachment_id: string; + } + + export interface AttachmentFinalizeRequest { + /** Conversation UUID the uploaded image attachment belongs to. */ + conversation_id: string; + /** Prepared attachment ID to validate and finalize for sandbox use. */ + attachment_id: string; + } + + export interface AttachmentFinalizeResponse { + /** Opaque attachment ID for sandbox message sends. */ + id: string; + /** Sanitized file name recorded for the image attachment. */ + file_name: string; + /** Validated MIME type for the normalized image. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnum; + /** Normalized image size in bytes. */ + size: number; + /** Normalized image width in pixels. */ + width: number; + /** Normalized image height in pixels. */ + height: number; + } + + export interface AttachmentPrepareItem { + /** + * File name to associate with the uploaded image. + * @maxLength 255 + */ + file_name: string; + /** + * Expected upload size in bytes. Each image must be 4194304 bytes or smaller. + * @minimum 1 + * @maximum 4194304 + */ + size: number; + /** Exact MIME type for the direct upload. Only PNG and JPEG are supported. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnum; + } + + export interface AttachmentPrepareRequest { + /** Conversation UUID the staged image attachments belong to. */ + conversation_id: string; + /** + * Images to stage for the next sandbox message. + * @minItems 1 + * @maxItems 4 + */ + attachments: AttachmentPrepareItem[]; + } + + /** + * Signed S3-compatible form fields to include with the upload request. + */ + export type AttachmentPrepareResponseItemUploadFields = {[key: string]: string}; + + export interface AttachmentPrepareResponseItem { + /** Opaque attachment ID for finalize, delete, and sandbox message sends. */ + id: string; + /** Sanitized file name recorded for the image attachment. */ + file_name: string; + /** Signed MIME type for the upload. + * + * * `image/png` - image/png + * * `image/jpeg` - image/jpeg */ + content_type: AssistantAttachmentContentTypeEnum; + /** Expected upload size in bytes. */ + size: number; + /** Signed direct-upload URL for the image attachment. */ + upload_url: string; + /** Signed S3-compatible form fields to include with the upload request. */ + upload_fields: AttachmentPrepareResponseItemUploadFields; + } + + export interface AttachmentPrepareResponse { + /** Prepared image uploads. */ + attachments: AttachmentPrepareResponseItem[]; + } + export interface AttributeBreakdownRow { count: number; error_count: number; @@ -52597,7 +52699,7 @@ export namespace Schemas { Ultracode: 'ultracode', } as const; - export interface TaskRunArtifactMetadata { + export interface TaskRunSkillBundleMetadata { /** * Name of the local skill included in a skill_bundle artifact. * @maxLength 255 @@ -52626,6 +52728,23 @@ export namespace Schemas { schema_version: number; } + export interface TaskRunImageAttachmentMetadata { + /** PostHog AI conversation that supplied a user attachment. */ + conversation_id: string; + /** + * Normalized image width in pixels for a user attachment. + * @minimum 1 + */ + width: number; + /** + * Normalized image height in pixels for a user attachment. + * @minimum 1 + */ + height: number; + } + + export type TaskRunArtifactMetadata = TaskRunSkillBundleMetadata | TaskRunImageAttachmentMetadata; + /** * * `agent` - agent * * `user` - user @@ -69509,18 +69628,24 @@ export namespace Schemas { } /** - * Request body for `POST /conversations/{id}/open/`. A string `content` processes a turn; a - * null/absent `content` warms a sandbox that idles awaiting the first message. + * Request body for `POST /conversations/{id}/open/`. Nonblank `content` processes a turn and may include + * `attachment_ids`; null or absent `content` with no attachments warms a sandbox awaiting its first message. */ export interface SandboxOpen { /** - * The user's message text. Omit or null to warm a sandbox (boot + idle) ahead of the first message. + * The user's message text. Omit or null to warm a sandbox ahead of the first message, unless attachment_ids are provided. * @maxLength 40000 * @nullable */ content?: string | null; /** Client-generated trace id correlated with the resulting Run's SSE stream. */ trace_id?: string; + /** + * Finalized sandbox image attachment IDs to send with this message. + * @minItems 1 + * @maxItems 4 + */ + attachment_ids?: string[]; /** Typed PostHog entities (and free text) attached to this message. */ attached_context?: SandboxAttachedContextItem[]; /** Initial permission mode for the sandbox agent session. Defaults to `auto`, which allows safe tool use while preserving explicit confirmations.