Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion apps/code/src/renderer/di/bindings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,11 @@ import {
GITHUB_CONNECT_CLIENT,
type GithubConnectClient,
} from "@posthog/core/onboarding/identifiers";
import { CLOUD_ARTIFACT_READ_FILE_AS_BASE64 } from "@posthog/core/sessions/cloudArtifactIdentifiers";
import {
type BundleLocalSkill,
CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
CLOUD_ARTIFACT_READ_FILE_AS_BASE64,
} from "@posthog/core/sessions/cloudArtifactIdentifiers";
import {
LOCAL_HANDOFF_DIALOG,
LOCAL_HANDOFF_HOST,
Expand Down Expand Up @@ -262,6 +266,7 @@ export interface RendererBindings {
[CODE_REVIEW_WORKSPACE_CLIENT]: CodeReviewWorkspaceClient;
[REVERT_HUNK_SERVICE]: RevertHunkService;
[SKILLS_WORKSPACE_CLIENT]: SkillsWorkspaceClient;
[CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL]: BundleLocalSkill;
[CLOUD_ARTIFACT_READ_FILE_AS_BASE64]: ReadFileAsBase64;
[LLM_GATEWAY_SERVICE]: LlmGatewayService;
[TITLE_GENERATOR_FILE_READ_CLIENT]: FileReadClient;
Expand Down
10 changes: 9 additions & 1 deletion apps/code/src/renderer/di/container.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,10 @@ import {
import { LLM_GATEWAY_SERVICE } from "@posthog/core/llm-gateway/identifiers";
import type { LlmGatewayService } from "@posthog/core/llm-gateway/llm-gateway";
import type { LlmMessage } from "@posthog/core/llm-gateway/schemas";
import { CLOUD_ARTIFACT_READ_FILE_AS_BASE64 } from "@posthog/core/sessions/cloudArtifactIdentifiers";
import {
CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
CLOUD_ARTIFACT_READ_FILE_AS_BASE64,
} from "@posthog/core/sessions/cloudArtifactIdentifiers";
import {
LOCAL_HANDOFF_DIALOG,
LOCAL_HANDOFF_HOST,
Expand Down Expand Up @@ -365,6 +368,11 @@ container
.toConstantValue((filePath: string) =>
trpcClient.fs.readFileAsBase64.query({ filePath }),
);
container
.bind(CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL)
.toConstantValue((skillBundleRef) =>
hostTrpcClient.skills.bundleLocal.query(skillBundleRef),
);
container.bind(LLM_GATEWAY_SERVICE).toConstantValue({
prompt: (
messages: LlmMessage[],
Expand Down
92 changes: 86 additions & 6 deletions packages/core/src/sessions/cloudArtifactService.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
import { describe, expect, it, vi } from "vitest";
import type { CloudArtifactClient } from "./cloudArtifactIdentifiers";
import type {
BundleLocalSkill,
CloudArtifactClient,
} from "./cloudArtifactIdentifiers";
import {
CLOUD_ATTACHMENT_MAX_SIZE_BYTES,
CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES,
Expand All @@ -15,9 +18,23 @@ function makeClient(): CloudArtifactClient {
};
}

const bundleLocalSkill: BundleLocalSkill = vi.fn(async (skillBundleRef) => {
const contentBase64 = btoa("skill-bundle");
return {
name: skillBundleRef.name,
source: skillBundleRef.source,
fileName: `${skillBundleRef.name}.zip`,
contentType: "application/zip" as const,
contentBase64,
contentSha256:
"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
size: 12,
};
});

describe("CloudArtifactService", () => {
it("returns empty ids when no file paths are provided", async () => {
const service = new CloudArtifactService(vi.fn());
const service = new CloudArtifactService(vi.fn(), bundleLocalSkill);
expect(
await service.uploadRunAttachments(makeClient(), "t", "r", []),
).toEqual([]);
Expand All @@ -26,7 +43,10 @@ describe("CloudArtifactService", () => {
it("rejects attachments that exceed the max size", async () => {
const oversized = CLOUD_ATTACHMENT_MAX_SIZE_BYTES + 1;
const base64 = btoa("a".repeat(oversized));
const service = new CloudArtifactService(vi.fn().mockResolvedValue(base64));
const service = new CloudArtifactService(
vi.fn().mockResolvedValue(base64),
bundleLocalSkill,
);

await expect(
service.uploadRunAttachments(makeClient(), "task-1", "run-1", [
Expand All @@ -38,7 +58,10 @@ describe("CloudArtifactService", () => {
it("rejects PDFs that exceed the stricter cloud limit", async () => {
const oversized = CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES + 1;
const base64 = btoa("a".repeat(oversized));
const service = new CloudArtifactService(vi.fn().mockResolvedValue(base64));
const service = new CloudArtifactService(
vi.fn().mockResolvedValue(base64),
bundleLocalSkill,
);

await expect(
service.uploadRunAttachments(makeClient(), "task-1", "run-1", [
Expand All @@ -50,7 +73,10 @@ describe("CloudArtifactService", () => {
});

it("throws when a file cannot be read", async () => {
const service = new CloudArtifactService(vi.fn().mockResolvedValue(null));
const service = new CloudArtifactService(
vi.fn().mockResolvedValue(null),
bundleLocalSkill,
);

await expect(
service.uploadRunAttachments(makeClient(), "task-1", "run-1", [
Expand All @@ -64,7 +90,10 @@ describe("CloudArtifactService", () => {
.spyOn(globalThis, "fetch")
.mockResolvedValue({ ok: true } as Response);
const base64 = btoa("hello");
const service = new CloudArtifactService(vi.fn().mockResolvedValue(base64));
const service = new CloudArtifactService(
vi.fn().mockResolvedValue(base64),
bundleLocalSkill,
);

const client = makeClient();
(
Expand Down Expand Up @@ -93,4 +122,55 @@ describe("CloudArtifactService", () => {
);
fetchMock.mockRestore();
});

it("uploads local skill bundles as skill bundle artifacts", async () => {
const fetchMock = vi
.spyOn(globalThis, "fetch")
.mockResolvedValue({ ok: true } as Response);
const service = new CloudArtifactService(vi.fn(), bundleLocalSkill);
const client = makeClient();

(
client.prepareTaskRunArtifactUploads as ReturnType<typeof vi.fn>
).mockResolvedValue([
{
id: "prep-1",
name: "local-skill.zip",
type: "skill_bundle",
size: 12,
presigned_post: { url: "https://s3/upload", fields: { key: "k" } },
},
]);
(
client.finalizeTaskRunArtifactUploads as ReturnType<typeof vi.fn>
).mockResolvedValue([{ id: "skill-artifact-1" }]);

const ids = await service.uploadRunAttachments(
client,
"task-1",
"run-1",
[],
[{ name: "local-skill", source: "user", path: "/tmp/local-skill" }],
);

expect(ids).toEqual(["skill-artifact-1"]);
expect(client.prepareTaskRunArtifactUploads).toHaveBeenCalledWith(
"task-1",
"run-1",
[
expect.objectContaining({
name: "local-skill.zip",
type: "skill_bundle",
content_type: "application/zip",
metadata: expect.objectContaining({
skill_name: "local-skill",
skill_source: "user",
bundle_format: "zip",
schema_version: 1,
}),
}),
],
);
fetchMock.mockRestore();
});
});
63 changes: 59 additions & 4 deletions packages/core/src/sessions/cloudArtifactService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,20 @@ import type { ReadFileAsBase64 } from "@posthog/core/editor/cloud-prompt";
import { getFileName } from "@posthog/shared";
import { inject, injectable } from "inversify";
import {
type BundleLocalSkill,
CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
CLOUD_ARTIFACT_READ_FILE_AS_BASE64,
type CloudArtifactClient,
type CloudArtifactUploadRequest,
type CloudSkillBundleRef,
type FinalizedCloudArtifact,
type PreparedCloudArtifact,
} from "./cloudArtifactIdentifiers";

const ATTACHMENT_SOURCE = "posthog_code";
const SKILL_BUNDLE_SOURCE = "posthog_code_skill";
const DEFAULT_CONTENT_TYPE = "application/octet-stream";
const SKILL_BUNDLE_CONTENT_TYPE = "application/zip";
export const CLOUD_ATTACHMENT_MAX_SIZE_BYTES = 30 * 1024 * 1024;
export const CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES = 10 * 1024 * 1024;

Expand Down Expand Up @@ -116,18 +121,24 @@ export class CloudArtifactService {
constructor(
@inject(CLOUD_ARTIFACT_READ_FILE_AS_BASE64)
private readonly readFileAsBase64: ReadFileAsBase64,
@inject(CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL)
private readonly bundleLocalSkill: BundleLocalSkill,
) {}

async uploadTaskStagedAttachments(
client: CloudArtifactClient,
taskId: string,
filePaths: string[],
skillBundles: CloudSkillBundleRef[] = [],
): Promise<string[]> {
if (!filePaths.length) {
if (!filePaths.length && !skillBundles.length) {
return [];
}

const attachments = await this.loadCloudAttachments(filePaths);
const attachments = [
...(await this.loadCloudAttachments(filePaths)),
...(await this.loadCloudSkillBundles(skillBundles)),
];
const preparedArtifacts = await client.prepareTaskStagedArtifactUploads(
taskId,
attachments.map((attachment) => attachment.upload),
Expand All @@ -148,12 +159,16 @@ export class CloudArtifactService {
taskId: string,
runId: string,
filePaths: string[],
skillBundles: CloudSkillBundleRef[] = [],
): Promise<string[]> {
if (!filePaths.length) {
if (!filePaths.length && !skillBundles.length) {
return [];
}

const attachments = await this.loadCloudAttachments(filePaths);
const attachments = [
...(await this.loadCloudAttachments(filePaths)),
...(await this.loadCloudSkillBundles(skillBundles)),
];
const preparedArtifacts = await client.prepareTaskRunArtifactUploads(
taskId,
runId,
Expand Down Expand Up @@ -207,6 +222,46 @@ export class CloudArtifactService {
);
}

private async loadCloudSkillBundles(
skillBundleRefs: CloudSkillBundleRef[],
): Promise<LoadedCloudAttachment[]> {
return Promise.all(
skillBundleRefs.map(async (skillBundleRef) => {
const bundle = await this.bundleLocalSkill(skillBundleRef);
const bytes = base64ToUint8Array(bundle.contentBase64);
if (bytes.byteLength !== bundle.size) {
throw new Error(
`Unable to prepare local skill ${skillBundleRef.name}`,
);
}
if (bytes.byteLength > CLOUD_ATTACHMENT_MAX_SIZE_BYTES) {
throw new Error(
`${bundle.fileName} exceeds the 30MB attachment limit`,
);
}

return {
filePath: skillBundleRef.path,
bytes,
upload: {
name: bundle.fileName,
type: "skill_bundle",
source: SKILL_BUNDLE_SOURCE,
size: bytes.byteLength,
content_type: SKILL_BUNDLE_CONTENT_TYPE,
metadata: {
skill_name: bundle.name,
skill_source: bundle.source,
content_sha256: bundle.contentSha256,
bundle_format: "zip",
schema_version: 1,
},
},
};
}),
);
}

private async uploadPreparedArtifacts(
attachments: LoadedCloudAttachment[],
preparedArtifacts: PreparedCloudArtifact[],
Expand Down
40 changes: 40 additions & 0 deletions packages/ui/src/features/message-editor/commands.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
import { describe, expect, it } from "vitest";
import { rewriteLocalSkillCommandPrompt } from "./commands";
import type { EditorAvailableCommand } from "./types";

const commands: EditorAvailableCommand[] = [
{
name: "local-test-skill",
description: "Local user skill",
localSkill: {
name: "local-test-skill",
source: "user",
path: "/Users/example/.claude/skills/local-test-skill",
},
},
];

describe("message editor commands", () => {
it("rewrites local skill slash commands to skill tags", () => {
expect(rewriteLocalSkillCommandPrompt("/local-test-skill", commands)).toBe(
'<skill name="local-test-skill" source="user" path="/Users/example/.claude/skills/local-test-skill" />',
);
});

it("preserves local skill arguments after the skill tag", () => {
expect(
rewriteLocalSkillCommandPrompt(
"/local-test-skill with context",
commands,
),
).toBe(
'<skill name="local-test-skill" source="user" path="/Users/example/.claude/skills/local-test-skill" /> with context',
);
});

it("does not rewrite unknown commands", () => {
expect(
rewriteLocalSkillCommandPrompt("/feedback looks good", commands),
).toBe(null);
});
});
Loading
Loading