diff --git a/src/mgr/postgres-store.ts b/src/mgr/postgres-store.ts index d0d56d5..f27382e 100644 --- a/src/mgr/postgres-store.ts +++ b/src/mgr/postgres-store.ts @@ -1190,6 +1190,17 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( return result.rows[0] ? sessionFromRow(result.rows[0]) : null; } + async getLatestSessionResourceRun(sessionId: string): Promise { + const result = await this.pool.query( + `SELECT * FROM agentrun_runs + WHERE session_ref->>'sessionId' = $1 AND jsonb_typeof(resource_bundle_ref) = 'object' + ORDER BY created_at DESC, id DESC + LIMIT 1`, + [sessionId], + ); + return result.rows[0] ? runFromRow(result.rows[0]) : null; + } + async getSessionSummary(sessionId: string, readerId: string | null = null): Promise { const session = await this.getSession(sessionId); if (!session) throw new AgentRunError("schema-invalid", `session ${sessionId} was not found`, { httpStatus: 404 }); diff --git a/src/mgr/server.ts b/src/mgr/server.ts index 97a9b70..569b391 100644 --- a/src/mgr/server.ts +++ b/src/mgr/server.ts @@ -1168,7 +1168,8 @@ async function sendToSession(input: { store: AgentRunStore; sessionId: string; b } const runRecord = asRecord(record.run ?? record.runBase ?? null, "sessionSend.run"); - const runBody = sessionSendRunBody(input.sessionId, runRecord); + const resourceRun = existing ? await input.store.getLatestSessionResourceRun(input.sessionId) : null; + const runBody = inheritSessionRunResources(sessionSendRunBody(input.sessionId, runRecord), resourceRun); const commandBody: JsonRecord = { type: "turn", payload }; if (commandIdempotencyKey) commandBody.idempotencyKey = commandIdempotencyKey; const runnerJobBody = validateSessionRunnerJobInput(record.runnerJob === undefined || record.runnerJob === null ? {} : record.runnerJob); @@ -1269,6 +1270,29 @@ function sessionSendRunBody(sessionId: string, runRecord: JsonRecord): JsonRecor return { ...runRecord, sessionRef: { ...sessionRef, sessionId, metadata } }; } +function inheritSessionRunResources(runBody: JsonRecord, resourceRun: RunRecord | null): JsonRecord { + if (runBody.resourceBundleRef !== undefined && runBody.resourceBundleRef !== null) return runBody; + if (!resourceRun?.resourceBundleRef) return runBody; + const executionPolicy = asJsonRecord(runBody.executionPolicy) ?? {}; + const secretScope = asJsonRecord(executionPolicy.secretScope) ?? {}; + const sourceToolCredentials = resourceRun.executionPolicy.secretScope.toolCredentials; + const toolCredentials = Array.isArray(secretScope.toolCredentials) && secretScope.toolCredentials.length > 0 + ? secretScope.toolCredentials + : sourceToolCredentials; + return { + ...runBody, + workspaceRef: resourceRun.workspaceRef, + resourceBundleRef: resourceRun.resourceBundleRef, + executionPolicy: { + ...executionPolicy, + secretScope: { + ...secretScope, + ...(toolCredentials && toolCredentials.length > 0 ? { toolCredentials } : {}), + }, + }, + }; +} + function sessionSendPlan(sessionId: string, decision: "steer" | "turn", active: Awaited>, request: JsonRecord, runBody: JsonRecord | null, reusedIdleRun: JsonRecord | null = null): JsonRecord { return { action: "session-send-plan", diff --git a/src/mgr/store.ts b/src/mgr/store.ts index 7442ed9..e6217a2 100644 --- a/src/mgr/store.ts +++ b/src/mgr/store.ts @@ -98,6 +98,7 @@ export interface AgentRunStore { cancelRun(runId: string, reason?: string): MaybePromise; cancelCommand(commandId: string, reason?: string): MaybePromise; getSession(sessionId: string): MaybePromise; + getLatestSessionResourceRun(sessionId: string): MaybePromise; getSessionSummary(sessionId: string, readerId?: string | null): MaybePromise; listSessions(input: ListSessionsInput): MaybePromise; listSessionTrace(sessionId: string, input: SessionEventPageInput): MaybePromise; @@ -776,6 +777,12 @@ export class MemoryAgentRunStore implements AgentRunStore { return this.sessions.get(sessionId) ?? null; } + getLatestSessionResourceRun(sessionId: string): RunRecord | null { + return Array.from(this.runs.values()) + .filter((run) => run.sessionRef?.sessionId === sessionId && run.resourceBundleRef !== null) + .sort((a, b) => b.createdAt.localeCompare(a.createdAt) || b.id.localeCompare(a.id))[0] ?? null; + } + getSessionSummary(sessionId: string, readerId: string | null = null): SessionSummary { const session = this.getSession(sessionId); if (!session) throw new AgentRunError("schema-invalid", `session ${sessionId} was not found`, { httpStatus: 404 }); diff --git a/src/selftest/cases/66-session-turn-admission.ts b/src/selftest/cases/66-session-turn-admission.ts index 5374cb5..92ac0ad 100644 --- a/src/selftest/cases/66-session-turn-admission.ts +++ b/src/selftest/cases/66-session-turn-admission.ts @@ -27,6 +27,7 @@ const selfTest: SelfTestCase = async (context) => { await assertKubernetesObservationFailureTerminalizes(fakeKubectl); await assertKubernetesCreateFailureTerminalizes(fakeKubectl); await assertSuccessfulSameSessionContinuity(fakeKubectl); + await assertSessionFollowUpInheritsAipodResources(fakeKubectl); } finally { restoreEnv("AGENTRUN_SELFTEST_KUBECTL_MODE", previousMode); } @@ -42,6 +43,7 @@ const selfTest: SelfTestCase = async (context) => { "kubectl-observation-failure-terminalizes-and-releases-session", "kubectl-create-failure-terminalizes-and-releases-session", "successful-same-session-continuity", + "session-follow-up-inherits-aipod-workspace-bundles-and-tool-credentials", ], }; }; @@ -279,6 +281,44 @@ async function assertSuccessfulSameSessionContinuity(kubectlCommand: string): Pr assert.equal(session?.lastCommandId, second.command.id); } +async function assertSessionFollowUpInheritsAipodResources(kubectlCommand: string): Promise { + const store = new MemoryAgentRunStore(); + const sessionId = "ses_session_admission_aipod_resources"; + const initial = store.createRun(artificerRunInput(sessionId)); + const initialCommand = store.createCommand(initial.id, validateCreateCommand({ type: "turn", payload: { prompt: "initial artificer turn" } })); + store.finishCommand(initialCommand.id, { terminalStatus: "completed", failureKind: null, failureMessage: null }); + store.finishRun(initial.id, { terminalStatus: "completed", failureKind: null, failureMessage: null }); + + const regressed = store.createRun(runInput(sessionId)); + const regressedCommand = store.createCommand(regressed.id, validateCreateCommand({ type: "turn", payload: { prompt: "regressed follow-up" } })); + store.finishCommand(regressedCommand.id, { terminalStatus: "cancelled", failureKind: "cancelled", failureMessage: "wrong workspace" }); + store.finishRun(regressed.id, { terminalStatus: "cancelled", failureKind: "cancelled", failureMessage: "wrong workspace" }); + assert.equal(store.getSession(sessionId)?.lastRunId, regressed.id); + assert.equal(store.getLatestSessionResourceRun(sessionId)?.id, initial.id); + + const server = await startManagerServer({ + store, + sourceCommit: "selftest", + runnerJobDefaults: defaults(kubectlCommand), + runnerDispatcherOptions: { enabled: false }, + kafkaOutboxRelayOptions: { enabled: false }, + }); + try { + const client = new ManagerClient(server.baseUrl); + const response = await client.post(`/api/v1/sessions/${sessionId}/send`, { + ...sessionSendBody(sessionId, "continue in artificer workspace", "session-admission-aipod-resource-follow-up"), + createRunnerJob: false, + }) as JsonRecord; + const run = response.run as RunRecord; + assert.deepEqual(run.workspaceRef, initial.workspaceRef); + assert.deepEqual(run.resourceBundleRef, initial.resourceBundleRef); + assert.deepEqual(run.executionPolicy.secretScope.toolCredentials, initial.executionPolicy.secretScope.toolCredentials); + assert.deepEqual(run.executionPolicy.secretScope.providerCredentials, runInput(sessionId).executionPolicy.secretScope.providerCredentials); + } finally { + await closeServer(server.server); + } +} + function assertTerminalizedAdmission(store: MemoryAgentRunStore, runId: string, commandId: string, sessionId: string, reason: string): void { const intent = store.getRunnerDispatchIntent(commandId); assert.equal(intent?.state, "failed"); @@ -354,6 +394,35 @@ function runInput(sessionId: string, resourceCommit?: string): CreateRunInput { }; } +function artificerRunInput(sessionId: string): CreateRunInput { + const input = runInput(sessionId, "1111111111111111111111111111111111111111"); + return validateCreateRun({ + ...input, + resourceBundleRef: { + kind: "gitbundle", + repoUrl: "git@github.com:pikasTech/unidesk.git", + ref: "master", + bundles: [ + { name: "unidesk-skills", subpath: ".agents/skills", targetPath: ".agents/skills" }, + { name: "agentrun-runner-tools", repoUrl: "git@github.com:pikasTech/agentrun.git", ref: "v0.2", subpath: "tools", targetPath: "tools" }, + ], + submodules: false, + lfs: false, + }, + executionPolicy: { + ...input.executionPolicy, + secretScope: { + ...input.executionPolicy.secretScope, + toolCredentials: [ + { tool: "github", purpose: "github-pr", secretRef: { name: "agentrun-v01-tool-github-pr", keys: ["GH_TOKEN"] }, projection: { kind: "env", envName: "GH_TOKEN", secretKey: "GH_TOKEN" } }, + { tool: "github", purpose: "github-ssh", secretRef: { name: "agentrun-v01-tool-github-ssh", keys: ["id_ed25519", "known_hosts"] }, projection: { kind: "volume", mountPath: "/home/agentrun/.ssh" } }, + ], + }, + }, + traceSink: { kind: "queue", taskId: "qt_artificer_selftest", queue: "commander", lane: "v0.2" }, + }); +} + function defaults(kubectlCommand: string, withRetention = false): RunnerJobDefaults { return { namespace: "agentrun-v02", diff --git a/src/selftest/integration/postgres-session-turn-admission.ts b/src/selftest/integration/postgres-session-turn-admission.ts index 51c7065..2fc79e9 100644 --- a/src/selftest/integration/postgres-session-turn-admission.ts +++ b/src/selftest/integration/postgres-session-turn-admission.ts @@ -12,6 +12,7 @@ const suffix = `${Date.now()}_${process.pid}`; const concurrentSessionId = `ses_pg_admission_${suffix}`; const legacySessionId = `ses_pg_admission_legacy_${suffix}`; const rollbackSessionId = `ses_pg_admission_rollback_${suffix}`; +const resourceSessionId = `ses_pg_admission_resource_${suffix}`; const eventOutbox = { enabled: true, topic: "agentrun.event.v1", source: "session-admission-selftest" } as const; const primary = await createPostgresAgentRunStore({ connectionString, poolMax: 4, eventOutbox }); const concurrent = await createPostgresAgentRunStore({ connectionString, poolMax: 4, eventOutbox }); @@ -22,6 +23,7 @@ try { await assertConcurrentAdmissionAndRestartRecovery(); await assertLegacyHalfCommitRecoveryAfterRestart(); await assertAdmissionTransactionRollback(); + await assertLatestSessionResourceRunSkipsJsonNull(); console.log(JSON.stringify({ ok: true, tests: [ @@ -31,6 +33,7 @@ try { "postgres-session-turn-same-session-continuity", "postgres-legacy-half-commit-recovery-after-restart", "postgres-session-turn-transaction-rollback", + "postgres-session-resource-run-skips-json-null", ], valuesPrinted: false, })); @@ -39,6 +42,7 @@ try { await cleanupSession(concurrentSessionId).catch(() => undefined); await cleanupSession(legacySessionId).catch(() => undefined); await cleanupSession(rollbackSessionId).catch(() => undefined); + await cleanupSession(resourceSessionId).catch(() => undefined); await Promise.all(reopenedStores.map(async (store) => await store.close())); await Promise.all([primary.close(), concurrent.close(), control.end()]); } @@ -134,6 +138,17 @@ async function assertAdmissionTransactionRollback(): Promise { assert.deepEqual(counts.rows[0], { session_count: 0, run_count: 0, command_count: 0, event_count: 0, intent_count: 0, outbox_count: 0 }); } +async function assertLatestSessionResourceRunSkipsJsonNull(): Promise { + const initial = await primary.createRun(resourceRunInput(resourceSessionId)); + const regressed = await primary.createRun(runInput(resourceSessionId)); + const storedTypes = await control.query( + "SELECT id, jsonb_typeof(resource_bundle_ref) AS resource_type FROM agentrun_runs WHERE id = ANY($1::text[]) ORDER BY id", + [[initial.id, regressed.id]], + ); + assert.deepEqual(new Set(storedTypes.rows.map((row) => row.resource_type)), new Set(["object", "null"])); + assert.equal((await primary.getLatestSessionResourceRun(resourceSessionId))?.id, initial.id); +} + function admissionInput(sessionId: string, prompt: string, idempotencyKey: string): SessionTurnAdmissionInput { return { sessionId, @@ -165,6 +180,21 @@ function runInput(sessionId: string): CreateRunInput { }; } +function resourceRunInput(sessionId: string): CreateRunInput { + return validateCreateRun({ + ...runInput(sessionId), + workspaceRef: { kind: "opaque", path: "." }, + resourceBundleRef: { + kind: "gitbundle", + repoUrl: "git@github.com:pikasTech/unidesk.git", + ref: "master", + bundles: [{ name: "unidesk-skills", subpath: ".agents/skills", targetPath: ".agents/skills" }], + submodules: false, + lfs: false, + }, + }); +} + async function installRollbackTrigger(): Promise { await dropRollbackTrigger(); await control.query(`