From 9535f12d48a8d12dc9f115ed1940d39efc2a96b1 Mon Sep 17 00:00:00 2001 From: Gabrielle Oliveira Date: Tue, 13 Oct 2026 22:35:40 -0300 Subject: [PATCH] fix(pi): use leverage's client field and start sessions with the first prompt --- src/index.ts | 17 +++++++++++--- src/render.ts | 2 +- src/session.ts | 32 +++++++++++++++---------- src/workspace/api.ts | 18 ++++++++------ src/workspace/schema.ts | 10 ++++---- src/workspace/ui.ts | 2 +- src/workspace/view.ts | 4 ++-- tests/extension.test.ts | 28 ++++++++++++---------- tests/view.test.ts | 8 +++---- tests/workspace/api.test.ts | 43 ++++++++++++++++++++-------------- tests/workspace/fixture.ts | 27 +++++++++++++++++++-- tests/workspace/socket.test.ts | 4 ++-- 12 files changed, 126 insertions(+), 69 deletions(-) diff --git a/src/index.ts b/src/index.ts index 81ba4aa..f505f28 100644 --- a/src/index.ts +++ b/src/index.ts @@ -42,6 +42,7 @@ import { SessionView, sameSession, sessionLink, + uploadFiles, viewRemoteHistory, } from "./session"; import { @@ -518,7 +519,7 @@ export default function leverage(pi: ExtensionAPI): void { await createAndSend(prompt, ctx); } }; - // The first prompt creates the empty session, opens it, then sends the prompt. + // The first prompt creates the session as its first message, then opens it. const createAndSend = async (prompt: Prompt, ctx: ExtensionContext) => { if (creating) { throw new Error("Session setup is already in progress."); @@ -547,7 +548,18 @@ export default function leverage(pi: ExtensionAPI): void { pendingCreation = prompt; status(ctx); try { - const id = await workspace.create(draft, lifetime.signal); + // A created draft needs no upload: its session already has the message. + const attachmentIds = draft.sessionId + ? [] + : await uploadFiles(prompt, lifetime.signal); + const id = await workspace.create( + draft, + { + content: prompt.text, + ...(attachmentIds.length ? { attachmentIds } : {}), + }, + lifetime.signal, + ); rememberModel(api.connection, { providerFamily: draft.providerFamily, model: draft.model, @@ -560,7 +572,6 @@ export default function leverage(pi: ExtensionAPI): void { if (!view?.ready || view.session.id !== session.id) { await attach(session, ctx); } - await requireSession(true).send(prompt); if (opening === generation) { draft = undefined; pendingCreation = undefined; diff --git a/src/render.ts b/src/render.ts index e125fea..f649c94 100644 --- a/src/render.ts +++ b/src/render.ts @@ -444,7 +444,7 @@ function person(theme: Theme, shown: Shown): Component { const author = clean(authorOf(data, members, viewerId) ?? "User"); const own = author.endsWith(" (you)"); const [state, tone] = DELIVERY[data.status ?? ""] ?? []; - const app = SENT_FROM.get(data.harness ?? ""); + const app = SENT_FROM.get(data.client ?? ""); card.addChild( new Text( [ diff --git a/src/session.ts b/src/session.ts index 6f6fe67..a730a42 100644 --- a/src/session.ts +++ b/src/session.ts @@ -40,6 +40,25 @@ export type Prompt = { delivery: "send" | "queue"; }; +/** Uploads a prompt's files and returns their attachment IDs. */ +export function uploadFiles( + prompt: Prompt, + signal: AbortSignal, +): Promise { + return Promise.all( + (prompt.files ?? []).map((file) => + workspace.upload( + { + filename: file.filename, + contentType: file.contentType, + bytes: Buffer.from(file.data, "base64"), + }, + signal, + ), + ), + ); +} + export interface ViewHooks { // The status lines and footer need a redraw. changed(): void; @@ -479,18 +498,7 @@ export class SessionView { } this.sending = true; try { - const attachmentIds = await Promise.all( - (prompt.files ?? []).map((file) => - workspace.upload( - { - filename: file.filename, - contentType: file.contentType, - bytes: Buffer.from(file.data, "base64"), - }, - this.signal, - ), - ), - ); + const attachmentIds = await uploadFiles(prompt, this.signal); const next = this.nextModel; const input = await workspace.send( this.session.id, diff --git a/src/workspace/api.ts b/src/workspace/api.ts index d21547d..153dc1b 100644 --- a/src/workspace/api.ts +++ b/src/workspace/api.ts @@ -473,11 +473,15 @@ export class WorkspaceClient { } /** - * Creates the empty session for a draft. The draft keeps the created ID, so - * a retry after a lost acknowledgement uses the same session instead of - * creating a second one. + * Creates a draft's session with its first message. The draft keeps the + * created ID, so a retry after a lost acknowledgement uses the same session + * instead of creating a second one. */ - create(draft: SessionDraft, signal: AbortSignal): Promise { + create( + draft: SessionDraft, + first: Pick, + signal: AbortSignal, + ): Promise { const program = Effect.gen({ self: this }, function* () { if (draft.sessionId) { return draft.sessionId; @@ -493,14 +497,14 @@ export class WorkspaceClient { const request = sessionCreateSchema.parse({ type: "session.create", clientRequestId: draft.requestId, - prompt: "", + ...first, ...(title?.success ? { title: title.data } : {}), context: draft.context, providerFamily: draft.providerFamily, model: draft.model, reasoningEffort: draft.reasoningEffort, checkoutStrategy: "clone", - harness: "pi", + client: "pi", }); // The socket owns the "not confirmed" outcome, so the signal goes to it. const reply = yield* Effect.tryPromise({ @@ -537,7 +541,7 @@ export class WorkspaceClient { type: "session.message", sessionId, ...message, - harness: "pi", + client: "pi", }); const reply = await socket.request( frame, diff --git a/src/workspace/schema.ts b/src/workspace/schema.ts index 1f67cef..90ad6c1 100644 --- a/src/workspace/schema.ts +++ b/src/workspace/schema.ts @@ -82,7 +82,7 @@ export const inputSchema = z.object({ sessionId: z.string(), authorId: z.string().nullable(), authorName: z.string().nullable(), - harness: z.string().nullable(), + client: z.string().nullable(), content: z.string(), status: z.string(), createdAt: z.string(), @@ -478,14 +478,16 @@ export type WorkspaceEvent = z.infer; export const sessionCreateSchema = z.object({ type: z.literal("session.create"), clientRequestId: z.string().uuid(), - prompt: z.string(), + // The first message. Leverage needs its text or an attachment. + content: z.string(), + attachmentIds: z.array(z.string()).optional(), title: z.string().optional(), context: sessionContextSchema, providerFamily: familySchema.optional(), model: z.string().optional(), reasoningEffort: z.string().optional(), checkoutStrategy: z.literal("clone"), - harness: z.literal("pi"), + client: z.literal("pi"), }); export const sessionMessageSchema = z.object({ @@ -497,7 +499,7 @@ export const sessionMessageSchema = z.object({ attachmentIds: z.array(z.string()).optional(), model: z.string().optional(), reasoningEffort: z.string().optional(), - harness: z.literal("pi"), + client: z.literal("pi"), }); export const clientFrameSchema = z.discriminatedUnion("type", [ diff --git a/src/workspace/ui.ts b/src/workspace/ui.ts index 43926b0..6f902b1 100644 --- a/src/workspace/ui.ts +++ b/src/workspace/ui.ts @@ -70,7 +70,7 @@ export async function editDraft( await textDrawer(ctx, { title: "Session setup is waiting", read: () => - "The empty session already exists. Retry the original prompt to finish its setup. Start a new session to use different settings.", + "The session already exists. Retry the original prompt to finish its setup. Start a new session to use different settings.", signal, }); return; diff --git a/src/workspace/view.ts b/src/workspace/view.ts index f940873..7fe3b02 100644 --- a/src/workspace/view.ts +++ b/src/workspace/view.ts @@ -120,7 +120,7 @@ export interface EntryData { // A person's message: how far it got, and the app it was sent from. inputUuid?: string; status?: string; - harness?: string | null; + client?: string | null; attachments?: Attachment[]; // A sub-agent's work, one line per step. steps?: string[]; @@ -462,7 +462,7 @@ export class EntryFold { : {}), ...(input?.authorName ? { authorName: input.authorName } : {}), ...(input?.status ? { status: input.status } : {}), - ...(input?.harness !== undefined ? { harness: input.harness } : {}), + ...(input?.client !== undefined ? { client: input.client } : {}), ...(attachments.length ? { attachments } : {}), }, [user], diff --git a/tests/extension.test.ts b/tests/extension.test.ts index 123e979..7e38029 100644 --- a/tests/extension.test.ts +++ b/tests/extension.test.ts @@ -543,7 +543,7 @@ describe("Pi hosted frontend", () => { ).toBeUndefined(); }); - test("the first prompt creates the draft's session with its title, then sends the prompt", async () => { + test("the first prompt and its image create the draft's session with its title", async () => { const f = leverage(); const runner = await load(); f.configure(runner); @@ -554,24 +554,26 @@ describe("Pi hosted frontend", () => { }, }); await command(runner, "new Explore the project"); - await runner.emitInput("Map the modules", undefined, "interactive"); + await runner.emitInput( + "Map the modules", + [{ type: "image", data: "aGVsbG8=", mimeType: "image/png" }], + "interactive", + ); expect(f.state.createCount).toBe(1); - expect(f.sent("session.create")).toEqual([ - expect.objectContaining({ title: "Explore the project", prompt: "" }), + // The image is uploaded first, then the create names its attachment. + expect(f.uploads).toEqual([ + { id: expect.any(String), bytes: 5, completed: true }, ]); - expect(f.sent("session.message")).toEqual([ + expect(f.sent("session.create")).toEqual([ expect.objectContaining({ - sessionId: SESSION, + title: "Explore the project", content: "Map the modules", - delivery: "queue", + attachmentIds: [f.uploads[0]?.id], + client: "pi", }), ]); - // The create frame goes out before the message that needs its session. - expect( - f.frames.findIndex((frame) => frame.type === "session.create"), - ).toBeLessThan( - f.frames.findIndex((frame) => frame.type === "session.message"), - ); + // The session starts with the prompt, so it is not sent again. + expect(f.sent("session.message")).toEqual([]); expect( sessionLink(runner.createContext().sessionManager.getBranch())?.sessionId, ).toBe(SESSION); diff --git a/tests/view.test.ts b/tests/view.test.ts index 42a6821..cddbb78 100644 --- a/tests/view.test.ts +++ b/tests/view.test.ts @@ -187,7 +187,7 @@ describe("Session view entries", () => { const message = input("Canonical content", { authorId: "bob", authorName: "Old name", - harness: "codex", + client: "codex", status: "consumed", }); const people = { members: [{ id: "bob", name: "Bob" }], viewerId: "alice" }; @@ -206,10 +206,10 @@ describe("Session view entries", () => { const text = shown(f.shown(canonical!, people)); expect(text).toContain("Bob"); expect(text).toContain("Sent from Codex"); - const sentFrom = (harness: string) => + const sentFrom = (client: string) => shown( f.shown( - { ...canonical!, data: { ...entryData(canonical!), harness } }, + { ...canonical!, data: { ...entryData(canonical!), client } }, people, ), ); @@ -223,7 +223,7 @@ describe("Session view entries", () => { expect(text).not.toContain("Old name"); expect(text).not.toContain("(you)"); expect(entryData(unknown!).authorId).toBeUndefined(); - expect(entryData(unknown!).harness).toBeUndefined(); + expect(entryData(unknown!).client).toBeUndefined(); expect(f.text(unknown!, people)).toStartWith("User\n"); }); diff --git a/tests/workspace/api.test.ts b/tests/workspace/api.test.ts index ff2f8f1..d6924d6 100644 --- a/tests/workspace/api.test.ts +++ b/tests/workspace/api.test.ts @@ -8,6 +8,9 @@ import { workspaceFixture, } from "./fixture"; +// Leverage starts a session with its first message. +const FIRST = { content: "Map the modules" }; + const fixtures: ReturnType[] = []; afterEach(async () => { for (const fixture of fixtures.splice(0)) { @@ -21,19 +24,23 @@ function fixture(extra?: Parameters[0]) { return value; } -test("creating a session sends its title with the empty first message", async () => { +test("creating a session sends its title with its first message", async () => { const f = fixture(); const api = f.client(); const draft = api.draft(); draft.title = " Named draft "; - expect(await api.create(draft, new AbortController().signal)).toBe(SESSION); + expect(await api.create(draft, FIRST, new AbortController().signal)).toBe( + SESSION, + ); const create = f.frames.find((frame) => frame.type === "session.create")!; - expect(create.prompt).toBe(""); + expect(create.content).toBe("Map the modules"); expect(create.title).toBe("Named draft"); - expect(create.harness).toBe("pi"); + expect(create.client).toBe("pi"); expect(f.session.title).toBe("Named draft"); // A created draft reuses its session instead of asking for another. - expect(await api.create(draft, new AbortController().signal)).toBe(SESSION); + expect(await api.create(draft, FIRST, new AbortController().signal)).toBe( + SESSION, + ); expect(f.state.createCount).toBe(1); }); @@ -44,9 +51,9 @@ test("a default-model session is created even when the model catalog is missing" : undefined, ); const api = f.client(); - expect(await api.create(api.draft(), new AbortController().signal)).toBe( - SESSION, - ); + expect( + await api.create(api.draft(), FIRST, new AbortController().signal), + ).toBe(SESSION); expect(f.state.createCount).toBe(1); expect(f.requests.some((r) => r.path.endsWith("/provider-models"))).toBe( false, @@ -87,13 +94,13 @@ test("a lost creation acknowledgement retries the stable request ID without dupl const draft = api.draft(); f.state.dropCreated = true; const life = new AbortController(); - const attempt = api.create(draft, life.signal); + const attempt = api.create(draft, FIRST, life.signal); const failure = rejects(attempt, /not confirmed/); await eventually(() => f.state.createCount === 1); life.abort(); await failure; expect(draft.sessionId).toBeUndefined(); - await api.create(draft, new AbortController().signal); + await api.create(draft, FIRST, new AbortController().signal); expect(f.state.createCount).toBe(1); const ids = f.frames .filter((e) => e.type === "session.create") @@ -107,20 +114,20 @@ test("revalidates removed channels, unavailable providers, and unsupported reaso const draft = api.draft(); const signal = new AbortController().signal; draft.context = { type: "channel", channelId: "removed" }; - await rejects(api.create(draft, signal), /channel/); + await rejects(api.create(draft, FIRST, signal), /channel/); draft.context = { type: "none" }; draft.providerFamily = "codex"; - await rejects(api.create(draft, signal), /unavailable/); + await rejects(api.create(draft, FIRST, signal), /unavailable/); draft.providerFamily = "claude_code"; draft.model = "hosted-model"; draft.reasoningEffort = "impossible"; - await rejects(api.create(draft, signal), /reasoning/); + await rejects(api.create(draft, FIRST, signal), /reasoning/); draft.model = "retired-model"; - await rejects(api.create(draft, signal), /no longer available/); + await rejects(api.create(draft, FIRST, signal), /no longer available/); expect(f.state.createCount).toBe(0); draft.model = "hosted-model"; draft.reasoningEffort = "high"; - await api.create(draft, signal); + await api.create(draft, FIRST, signal); expect(f.state.createCount).toBe(1); }); @@ -146,11 +153,11 @@ test("a workspace with no subscription starts sessions on Leverage's own models" } // A catalog with no model on credits leaves nothing to run. f.state.catalog = [{ ...model, sources: ["pool"] }]; - await rejects(api.create(api.draft(), signal), /unavailable/); + await rejects(api.create(api.draft(), FIRST, signal), /unavailable/); expect(f.state.createCount).toBe(0); f.state.catalog = [{ ...model, sources: ["pool", "leverage"] }]; const draft = api.draft(); - expect(await api.create(draft, signal)).toBe(SESSION); + expect(await api.create(draft, FIRST, signal)).toBe(SESSION); expect(draft.providerFamily).toBe(model.family); }); @@ -168,7 +175,7 @@ test("a message resolves once Leverage stores it, and an unconfirmed one fails", expect(stored.content).toBe("Run the tests"); const frame = f.frames.find((one) => one.type === "session.message")!; expect(frame.delivery).toBe("send"); - expect(frame.harness).toBe("pi"); + expect(frame.client).toBe("pi"); f.state.dropMessageAck = true; const life = new AbortController(); const lost = api.send( diff --git a/tests/workspace/fixture.ts b/tests/workspace/fixture.ts index d6df1fd..31fa9cc 100644 --- a/tests/workspace/fixture.ts +++ b/tests/workspace/fixture.ts @@ -49,7 +49,7 @@ export function input( sessionId: SESSION, authorId: "owner", authorName: "Alice", - harness: null, + client: null, content, status: "received", createdAt: new Date().toISOString(), @@ -576,13 +576,36 @@ export function workspaceFixture(extra?: Extra) { watched(ws, sessionId, true, state); return; } + // Like Leverage, a session starts with its first message. case "session.create": { + const attachmentIds = Array.isArray(frame.attachmentIds) + ? frame.attachmentIds + : []; + if (!String(frame.content ?? "").trim() && !attachmentIds.length) { + reply({ + type: "error", + message: "Invalid event", + clientRequestId: frame.clientRequestId, + }); + return; + } if (!created.has(String(frame.clientRequestId))) { created.add(String(frame.clientRequestId)); state.createCount++; if (typeof frame.title === "string") { session.title = frame.title; } + state.nativeMessages = [ + ...state.nativeMessages, + input(String(frame.content ?? ""), { + uuid: String(frame.clientRequestId), + authorId: ws.data.user, + authorName: ws.data.user === "owner" ? "Alice" : "Bob", + client: + typeof frame.client === "string" ? frame.client : null, + }), + ]; + state.version++; } if (state.dropCreated) { state.dropCreated = false; @@ -604,7 +627,7 @@ export function workspaceFixture(extra?: Extra) { sessionId: String(frame.sessionId), authorId: ws.data.user, authorName: ws.data.user === "owner" ? "Alice" : "Bob", - harness: String(frame.harness), + client: String(frame.client), status: frame.delivery === "queue" ? "queued" : "received", }); state.nativeMessages = [ diff --git a/tests/workspace/socket.test.ts b/tests/workspace/socket.test.ts index 8d922e3..382b447 100644 --- a/tests/workspace/socket.test.ts +++ b/tests/workspace/socket.test.ts @@ -46,7 +46,7 @@ const bob = input("From Bob", { uuid: "22222222-2222-4222-8222-222222222222", authorId: MEMBER, authorName: "Bob", - harness: "codex", + client: "codex", status: "queued", }); @@ -227,7 +227,7 @@ test("a stale read cannot overwrite a newer message update", async () => { release(); await refreshing; expect(shared.input(bob.uuid)?.content).toBe("From Bob"); - expect(shared.input(bob.uuid)?.harness).toBe("codex"); + expect(shared.input(bob.uuid)?.client).toBe("codex"); shared.close(); f.publish({ type: "session.messages.updated",