diff --git a/server/src/services/chat-github-tools.test.ts b/server/src/services/chat-github-tools.test.ts new file mode 100644 index 00000000000..67165c1bb8c --- /dev/null +++ b/server/src/services/chat-github-tools.test.ts @@ -0,0 +1,97 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { + agents, + chatConversations, + chatDeliveries, + chatEndpoints, + chatMessageLinks, + companies, + createDb, + heartbeatRuns, + issueComments, + issues, + runIdentityContexts, + toolApplications, + toolConnections, +} from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { githubGuestBotConnectionForSession } from "./chat-github-tools.js"; + +const support = await getEmbeddedPostgresTestSupport(); +(support.supported ? describe : describe.skip)("GitHub guest bot connection discovery", () => { + let database: Awaited>; + let db: ReturnType; + + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("paperclip-github-guest-"); + db = createDb(database.connectionString); + }, 30_000); + afterAll(async () => { await database?.cleanup(); }, 60_000); + + async function fixture() { + const [company] = await db.insert(companies).values({ name: "Guest discovery", issuePrefix: randomUUID().slice(0, 8) }).returning(); + const [first, second] = await db.insert(agents).values([ + { companyId: company.id, name: "First agent", role: "general", adapterType: "codex_local" }, + { companyId: company.id, name: "Second agent", role: "general", adapterType: "codex_local" }, + ]).returning(); + const [issue, other] = await db.insert(issues).values([ + { companyId: company.id, title: "Reassigned task", assigneeAgentId: second.id }, + { companyId: company.id, title: "Other task", assigneeAgentId: second.id }, + ]).returning(); + return { company, first, second, issue, other }; + } + + async function continuation(f: Awaited>, parent: { agentId: string; issueId: string }) { + const [source, continued] = await db.insert(heartbeatRuns).values([ + { companyId: f.company.id, agentId: parent.agentId, status: "succeeded", nativeIssueId: parent.issueId, contextSnapshot: {}, responsibleUserId: "requester" }, + { companyId: f.company.id, agentId: f.second.id, status: "running", nativeIssueId: f.issue.id, contextSnapshot: {}, responsibleUserId: "requester" }, + ]).returning(); + const [origin] = await db.insert(runIdentityContexts).values({ companyId: f.company.id, runId: source.id, revision: 1, responsibleUserId: "requester", cause: "dispatch", correlationId: "dispatch" }).returning(); + const [context] = await db.insert(runIdentityContexts).values({ companyId: f.company.id, runId: continued.id, revision: 1, responsibleUserId: "requester", parentContextId: origin.id, cause: "dispatch", correlationId: "dispatch" }).returning(); + await db.update(heartbeatRuns).set({ activeIdentityContextId: context.id }).where(eq(heartbeatRuns.id, continued.id)); + return { companyId: f.company.id, agentId: f.second.id, runId: continued.id, issueId: f.issue.id }; + } + + async function guestSource(f: Awaited>, guest: boolean) { + const endpointId = randomUUID(); + const [application] = await db.insert(toolApplications).values({ companyId: f.company.id, applicationKey: `chat:github:${endpointId}`, name: `GitHub ${endpointId}`, type: "chat", status: "active" }).returning(); + const [connection] = await db.insert(toolConnections).values({ companyId: f.company.id, applicationId: application.id, name: "GitHub", uid: `chat-github-${endpointId}`, connectionPurpose: "channel", transport: "chat_sdk", status: "active", enabled: true }).returning(); + await db.insert(chatEndpoints).values({ id: endpointId, companyId: f.company.id, connectionId: connection.id, provider: "github", publicId: randomUUID(), assignedAgentId: f.second.id, status: "active" }); + const [conversation] = await db.insert(chatConversations).values({ companyId: f.company.id, endpointId, issueId: f.issue.id, externalConversationId: `github:owner/repo:${endpointId}`, externalLabel: "Guest thread" }).returning(); + const [comment] = await db.insert(issueComments).values({ companyId: f.company.id, issueId: f.issue.id, authorType: "user", authorUserId: "requester", body: "Please review" }).returning(); + const [delivery] = await db.insert(chatDeliveries).values({ companyId: f.company.id, endpointId, conversationId: conversation.id, providerEventId: randomUUID(), deduplicationKey: randomUUID(), eventKind: "message", normalizedEvent: { githubAuthority: { guest } }, state: "processed" }).returning(); + await db.insert(chatMessageLinks).values({ companyId: f.company.id, endpointId, conversationId: conversation.id, deliveryId: delivery.id, commentId: comment.id, providerMessageId: randomUUID(), direction: "inbound" }); + return { connectionId: connection.id, commentId: comment.id }; + } + + async function sponsoredRun(f: Awaited>, commentId: string) { + const [run] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.second.id, status: "running", nativeIssueId: f.issue.id, contextSnapshot: {}, responsibleUserId: "requester" }).returning(); + const [context] = await db.insert(runIdentityContexts).values({ companyId: f.company.id, runId: run.id, revision: 1, responsibleUserId: "requester", messageId: commentId, cause: "instruction", correlationId: "message" }).returning(); + await db.update(heartbeatRuns).set({ activeIdentityContextId: context.id }).where(eq(heartbeatRuns.id, run.id)); + return { companyId: f.company.id, agentId: f.second.id, runId: run.id, issueId: f.issue.id }; + } + + it.each(["agent", "task"])("finds no guest sponsor when the run chain continues another %s", async (fault) => { + const f = await fixture(); + const session = await continuation(f, fault === "agent" + ? { agentId: f.first.id, issueId: f.issue.id } + : { agentId: f.second.id, issueId: f.other.id }); + expect(await githubGuestBotConnectionForSession(db, session)).toBeNull(); + }); + + it("still refuses a guest-sponsored task whose run chain continues another agent", async () => { + const f = await fixture(); + await guestSource(f, true); + const session = await continuation(f, { agentId: f.first.id, issueId: f.issue.id }); + await expect(githubGuestBotConnectionForSession(db, session)).rejects.toMatchObject({ status: 403 }); + }); + + it.each([true, false])("resolves the bot connection only for a run sponsored by a GitHub guest (guest=%s)", async (guest) => { + const f = await fixture(); + const source = await guestSource(f, guest); + const session = await sponsoredRun(f, source.commentId); + expect(await githubGuestBotConnectionForSession(db, session)).toBe(guest ? source.connectionId : null); + }); +}); diff --git a/server/src/services/chat-github-tools.ts b/server/src/services/chat-github-tools.ts index f87a9f1c663..a221016caf4 100644 --- a/server/src/services/chat-github-tools.ts +++ b/server/src/services/chat-github-tools.ts @@ -385,6 +385,33 @@ export async function githubGuestBotConnectionForSession( >, ): Promise { if (!session.runId || !session.agentId || !session.issueId) return null; + const guestSource = (commentId?: string) => db + .select({ connectionId: chatEndpoints.connectionId }) + .from(chatMessageLinks) + .innerJoin( + chatDeliveries, + eq(chatDeliveries.id, chatMessageLinks.deliveryId), + ) + .innerJoin( + chatConversations, + eq(chatConversations.id, chatMessageLinks.conversationId), + ) + .innerJoin( + chatEndpoints, + eq(chatEndpoints.id, chatConversations.endpointId), + ) + .where( + and( + eq(chatMessageLinks.companyId, session.companyId), + ...(commentId ? [eq(chatMessageLinks.commentId, commentId)] : []), + eq(chatConversations.issueId, session.issueId!), + eq(chatEndpoints.assignedAgentId, session.agentId!), + eq(chatEndpoints.provider, "github"), + sql`${chatDeliveries.normalizedEvent} -> 'githubAuthority' ->> 'guest' = 'true'`, + ), + ) + .limit(1); + if (!(await guestSource()).length) return null; const [run] = await db .select({ id: heartbeatRuns.id, @@ -406,35 +433,6 @@ export async function githubGuestBotConnectionForSession( ); const commentId = run ? await chatRunSourceMessageId(db, run, { allowLegacyWakeComment: true, provider: "github" }) : null; if (typeof commentId !== "string") return null; - const [source] = await db - .select({ - event: chatDeliveries.normalizedEvent, - connectionId: chatEndpoints.connectionId, - }) - .from(chatMessageLinks) - .innerJoin( - chatDeliveries, - eq(chatDeliveries.id, chatMessageLinks.deliveryId), - ) - .innerJoin( - chatConversations, - eq(chatConversations.id, chatMessageLinks.conversationId), - ) - .innerJoin( - chatEndpoints, - eq(chatEndpoints.id, chatConversations.endpointId), - ) - .where( - and( - eq(chatMessageLinks.companyId, session.companyId), - eq(chatMessageLinks.commentId, commentId), - eq(chatConversations.issueId, session.issueId), - eq(chatEndpoints.assignedAgentId, session.agentId), - eq(chatEndpoints.provider, "github"), - ), - ); - return (source?.event.githubAuthority as { guest?: boolean } | undefined) - ?.guest === true - ? source!.connectionId - : null; + const [source] = await guestSource(commentId); + return source?.connectionId ?? null; }