Improve Foundry auth and task flows (#240)

This commit is contained in:
Nathan Flurry 2026-03-11 18:13:31 -07:00 committed by GitHub
parent d75e8c31d1
commit dbc2ff0682
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 621 additions and 137 deletions

View file

@ -142,10 +142,14 @@ export const task = actor({
async provision(c, cmd: InitializeCommand): Promise<{ ok: true }> {
const self = selfTask(c);
await self.send(taskWorkflowQueueName("task.command.provision"), cmd ?? {}, {
const result = await self.send(taskWorkflowQueueName("task.command.provision"), cmd ?? {}, {
wait: true,
timeout: 30 * 60_000,
});
const response = expectQueueResponse<{ ok: boolean; error?: string }>(result);
if (!response.ok) {
throw new Error(response.error ?? "task provisioning failed");
}
return { ok: true };
},

View file

@ -2,7 +2,8 @@
import { basename } from "node:path";
import { asc, eq } from "drizzle-orm";
import { getActorRuntimeContext } from "../context.js";
import { getOrCreateTaskStatusSync, getOrCreateProject, getOrCreateWorkspace, getSandboxInstance } from "../handles.js";
import { getOrCreateTaskStatusSync, getOrCreateProject, getOrCreateWorkspace, getSandboxInstance, selfTask } from "../handles.js";
import { resolveWorkspaceGithubAuth } from "../../services/github-auth.js";
import { task as taskTable, taskRuntime, taskWorkbenchSessions } from "./db/schema.js";
import { getCurrentRecord } from "./workflow/common.js";
@ -547,7 +548,26 @@ export async function renameWorkbenchBranch(c: any, value: string): Promise<void
}
export async function createWorkbenchSession(c: any, model?: string): Promise<{ tabId: string }> {
const record = await ensureWorkbenchSeeded(c);
let record = await ensureWorkbenchSeeded(c);
if (!record.activeSandboxId) {
const providerId = record.providerId ?? c.state.providerId ?? getActorRuntimeContext().providers.defaultProviderId();
await selfTask(c).provision({ providerId });
record = await ensureWorkbenchSeeded(c);
}
if (record.activeSessionId) {
const existingSessions = await listSessionMetaRows(c);
if (existingSessions.length === 0) {
await ensureSessionMeta(c, {
sessionId: record.activeSessionId,
model: model ?? defaultModelForAgent(record.agentType),
sessionName: "Session 1",
});
await notifyWorkbenchUpdated(c);
return { tabId: record.activeSessionId };
}
}
if (!record.activeSandboxId) {
throw new Error("cannot create session without an active sandbox");
}
@ -783,7 +803,10 @@ export async function publishWorkbenchPr(c: any): Promise<void> {
throw new Error("cannot publish PR without a branch");
}
const { driver } = getActorRuntimeContext();
const created = await driver.github.createPr(c.state.repoLocalPath, record.branchName, record.title ?? c.state.task);
const auth = await resolveWorkspaceGithubAuth(c, c.state.workspaceId);
const created = await driver.github.createPr(c.state.repoLocalPath, record.branchName, record.title ?? c.state.task, undefined, {
githubToken: auth?.githubToken ?? null,
});
await c.db
.update(taskTable)
.set({

View file

@ -104,7 +104,10 @@ const commandHandlers: Record<TaskQueueName, WorkflowHandler> = {
await msg.complete({ ok: true });
} catch (error) {
await loopCtx.step("init-failed-v2", async () => initFailedActivity(loopCtx, error));
await msg.complete({ ok: false });
await msg.complete({
ok: false,
error: resolveErrorMessage(error),
});
}
},

View file

@ -1,6 +1,7 @@
// @ts-nocheck
import { desc, eq } from "drizzle-orm";
import { resolveCreateFlowDecision } from "../../../services/create-flow.js";
import { resolveWorkspaceGithubAuth } from "../../../services/github-auth.js";
import { getActorRuntimeContext } from "../../context.js";
import { getOrCreateTaskStatusSync, getOrCreateHistory, getOrCreateProject, getOrCreateSandboxInstance, getSandboxInstance, selfTask } from "../../handles.js";
import { logActorWarning, resolveErrorMessage } from "../../logging.js";
@ -150,8 +151,9 @@ export async function initEnsureNameActivity(loopCtx: any): Promise<void> {
}
const { driver } = getActorRuntimeContext();
const auth = await resolveWorkspaceGithubAuth(loopCtx, loopCtx.state.workspaceId);
try {
await driver.git.fetch(loopCtx.state.repoLocalPath);
await driver.git.fetch(loopCtx.state.repoLocalPath, { githubToken: auth?.githubToken ?? null });
} catch (error) {
logActorWarning("task.init", "fetch before naming failed", {
workspaceId: loopCtx.state.workspaceId,
@ -160,7 +162,9 @@ export async function initEnsureNameActivity(loopCtx: any): Promise<void> {
error: resolveErrorMessage(error),
});
}
const remoteBranches = (await driver.git.listRemoteBranches(loopCtx.state.repoLocalPath)).map((branch: any) => branch.branchName);
const remoteBranches = (await driver.git.listRemoteBranches(loopCtx.state.repoLocalPath, { githubToken: auth?.githubToken ?? null })).map(
(branch: any) => branch.branchName,
);
const project = await getOrCreateProject(loopCtx, loopCtx.state.workspaceId, loopCtx.state.repoId, loopCtx.state.repoRemote);
const reservedBranches = await project.listReservedBranches({});
@ -274,6 +278,7 @@ export async function initCreateSandboxActivity(loopCtx: any, body: any): Promis
});
try {
const auth = await resolveWorkspaceGithubAuth(loopCtx, loopCtx.state.workspaceId);
const sandbox = await withActivityTimeout(timeoutMs, "createSandbox", async () =>
provider.createSandbox({
workspaceId: loopCtx.state.workspaceId,
@ -281,6 +286,7 @@ export async function initCreateSandboxActivity(loopCtx: any, body: any): Promis
repoRemote: loopCtx.state.repoRemote,
branchName: loopCtx.state.branchName,
taskId: loopCtx.state.taskId,
githubToken: auth?.githubToken ?? null,
debug: (message, context) => debugInit(loopCtx, message, context),
}),
);

View file

@ -2,6 +2,7 @@
import { eq } from "drizzle-orm";
import { getActorRuntimeContext } from "../../context.js";
import { logActorWarning, resolveErrorMessage } from "../../logging.js";
import { resolveWorkspaceGithubAuth } from "../../../services/github-auth.js";
import { task as taskTable, taskRuntime, taskSandboxes } from "../db/schema.js";
import { TASK_ROW_ID, appendHistory, resolveErrorDetail } from "./common.js";
import { pushActiveBranchActivity } from "./push.js";
@ -77,8 +78,10 @@ export async function idleSubmitPrActivity(loopCtx: any): Promise<void> {
if (self && self.prSubmitted) return;
const auth = await resolveWorkspaceGithubAuth(loopCtx, loopCtx.state.workspaceId);
try {
await driver.git.fetch(loopCtx.state.repoLocalPath);
await driver.git.fetch(loopCtx.state.repoLocalPath, { githubToken: auth?.githubToken ?? null });
} catch (error) {
logActorWarning("task.status-sync", "fetch before PR submit failed", {
workspaceId: loopCtx.state.workspaceId,
@ -98,7 +101,9 @@ export async function idleSubmitPrActivity(loopCtx: any): Promise<void> {
historyKind: "task.push.auto",
});
const pr = await driver.github.createPr(loopCtx.state.repoLocalPath, loopCtx.state.branchName, loopCtx.state.title);
const pr = await driver.github.createPr(loopCtx.state.repoLocalPath, loopCtx.state.branchName, loopCtx.state.title, undefined, {
githubToken: auth?.githubToken ?? null,
});
await db.update(taskTable).set({ prSubmitted: 1, updatedAt: Date.now() }).where(eq(taskTable.id, TASK_ROW_ID)).run();