diff --git a/doc/connections/AI-CONNECTION-ROUTERS.md b/doc/connections/AI-CONNECTION-ROUTERS.md index 4bf4f4d78b7..f95f15077de 100644 --- a/doc/connections/AI-CONNECTION-ROUTERS.md +++ b/doc/connections/AI-CONNECTION-ROUTERS.md @@ -11,6 +11,8 @@ observations to `onRouteAiConnection`. The plugin proposes an opaque member ID; it cannot receive credentials or expand authorization. Pure round robin does not probe usage. Usage-aware selection has a shared 15-second probe budget, 60-second freshness cache and ordinary grant/secret freshness invalidation. +Concurrent starts share an in-flight probe for the same grant and credential +freshness. Each caller retains its selection deadline. One cursor spans all agents in a pool. An allocation transaction locks the cursor, checks config and cursor revisions, rechecks authorization, and writes @@ -23,6 +25,9 @@ Pins snapshot the concrete binding and member profile. Removing or editing a member affects future allocations. Composer changes can change a supported model or effort, with a note on fallback, but never the account or harness. Session reset and compaction retain the pin. Revocation requires operator repair. +Authentication repair uses the failed run’s durable concrete allocation and +current member permissions; reconnect-and-continue preserves the agent’s pool +binding. Changing the agent’s pool invalidates the old repair card. A pre-existing managed session can adopt its saved account when it is an eligible member; otherwise the operator must explicitly reset the session. diff --git a/packages/db/src/connections-v3-schema-core-migration.test.ts b/packages/db/src/connections-v3-schema-core-migration.test.ts index a1ee50090cf..c476de7f104 100644 --- a/packages/db/src/connections-v3-schema-core-migration.test.ts +++ b/packages/db/src/connections-v3-schema-core-migration.test.ts @@ -38,6 +38,11 @@ describeEmbeddedPostgres("connections v3 schema core migration", () => { await sql`DELETE FROM "drizzle"."__drizzle_migrations" WHERE "hash" = ${await migrationHash()}`; // AI defaults arrive in 0273/0277 and depend on the composite grant key // from 0232. Rewind those tables before recreating the 0182 grant schema. + // Router pins/cursors depend on pools, whose composite connection FK also + // arrives after 0182. Rewind these empty fixture tables in dependency order. + await sql`DROP TABLE IF EXISTS "ai_connection_task_pins"`; + await sql`DROP TABLE IF EXISTS "ai_connection_router_cursors"`; + await sql`DROP TABLE IF EXISTS "ai_connection_pools"`; await sql`DROP TABLE IF EXISTS "ai_provider_defaults"`; await sql`DROP TABLE IF EXISTS "ai_connection_defaults"`; await sql`DROP TABLE IF EXISTS "connection_grant_delegations"`; diff --git a/server/src/__tests__/ai-connection-router.test.ts b/server/src/__tests__/ai-connection-router.test.ts index 9bd7e7cd6f4..0ce44befd19 100644 --- a/server/src/__tests__/ai-connection-router.test.ts +++ b/server/src/__tests__/ai-connection-router.test.ts @@ -4,7 +4,7 @@ import { mkdtemp, rm } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { and, eq } from "drizzle-orm"; -import { createDb, companies, agents, companyMemberships, connectionGrants, plugins, aiConnectionRouterCursors, aiConnectionTaskPins, toolConnections, companySecrets } from "@paperclipai/db"; +import { createDb, companies, agents, companyMemberships, connectionGrants, plugins, aiConnectionRouterCursors, aiConnectionTaskPins, toolConnections, companySecrets, heartbeatRuns, issues } from "@paperclipai/db"; import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres"; import type { AiConnectionPoolMember, AiConnectionRouterRequest, AiConnectionRouterResult, PaperclipPluginManifestV1 } from "@paperclipai/shared"; import { aiConnectionRouterService, AiConnectionPoolExhausted, poolMemberRuntimeConfig, applyAiConnectionRouterTaskSettings } from "../services/ai-connection-router.js"; @@ -18,6 +18,10 @@ import express from "express"; import request from "supertest"; import { aiConnectionRoutes } from "../routes/ai-connections.js"; import { errorHandler } from "../middleware/index.js"; +import { heartbeatService } from "../services/heartbeat.js"; +import { connectionIntentService } from "../services/connection-intents.js"; +import { legacyExecutionNeedsReconciliation } from "../services/legacy-execution-recovery.js"; +import { getServerAdapter, registerServerAdapter, unregisterServerAdapter } from "../adapters/index.js"; let database: Awaited>; let db: ReturnType; @@ -33,6 +37,7 @@ const proposal = vi.fn(async (_id: string, _method: string, request: AiConnectio return selected ? { kind: "selected", memberId: selected } : { kind: "unavailable", skipped: {} }; }); const worker = { isRunning: () => true, call: proposal } as unknown as PluginWorkerManager; +const defaultProposal = proposal.getMockImplementation()!; const service = () => aiConnectionRouterService(db, worker); const resolve = (poolId: string, taskKey: string, extra = {}) => service().resolve({ companyId, poolId, agentId, userId: "alice", adapterType: "paperclip_runner", taskKey, ...extra }); const makePool = async (config = {}) => service().save(pluginKey, { companyId, config: { name: randomUUID(), enabled: true, mode: "round_robin", thresholdPercent: 90, members, ...config } }, "alice"); @@ -176,6 +181,20 @@ describe("durable, authorized connection routing", () => { const pins = await db.select().from(aiConnectionTaskPins).where(eq(aiConnectionTaskPins.poolId, aware.id)); expect(pins).toHaveLength(1); } finally { spy.mockRestore(); } }); + it("shares an in-flight usage observation across concurrent allocations", async () => { + const [grant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, members[1]!.binding.grantId)); + const ref = grant!.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + await secretService(db).rotate(ref.secretId, { value: "fixture-shared-probe", preserveAiSessionEpoch: true }, { userId: "alice" }); + const fetch = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + await new Promise(resolve => setTimeout(resolve, 100)); + return new Response(JSON.stringify({ five_hour: { utilization: 12 } })); + }); + try { + const pool = await makePool({ mode: "usage_aware", members: [members[1]!] }); + await Promise.all(Array.from({ length: 6 }, (_, index) => resolve(pool.id, `shared-probe-${index}`))); + expect(fetch).toHaveBeenCalledTimes(1); + } finally { fetch.mockRestore(); } + }); it("expires a usage observation at a reported reset", async () => { const [grant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, members[1]!.binding.grantId)); const ref = grant!.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; @@ -194,6 +213,59 @@ describe("durable, authorized connection routing", () => { try { const pool = await makePool({ mode: "usage_aware", members: [members[1]!] }); const start = Date.now(); expect((await resolve(pool.id, "slow-probe")).memberId).toBe(members[1]!.id); expect(Date.now() - start).toBeLessThan(20_000); expect(proposal.mock.calls.at(-1)?.[2].candidates[0]?.usage).toBeUndefined(); } finally { spy.mockRestore(); } }, 25_000); + it("defers a pinned task before provider work and schedules a retry with its affinity intact", async () => { + await instanceSettingsService(db).updateExperimental({ enableAiConnectionRouters: true, enableWorktreeRunExecution: true }); + const pool = await makePool({ members: [members[1]!] }); + const [agent] = await db.insert(agents).values({ companyId, name: "Quota fixture", adapterType: "claude_local", adapterConfig: { cwd: home, engine: "cli" }, runtimeConfig: { aiConnection: { mode: "router", connectionId: pool.id }, heartbeat: { enabled: false } } }).returning(); + const [issue] = await db.insert(issues).values({ companyId, title: "Quota fixture", status: "todo", assigneeAgentId: agent!.id, responsibleUserId: "alice", createdByUserId: "alice" }).returning(); + const pinned = await resolve(pool.id, issue!.id, { agentId: agent!.id, adapterType: "claude_local" }); + await service().save(pluginKey, { companyId, id: pool.id, expectedRevision: pool.revision, config: { name: pool.name, enabled: true, mode: "usage_aware", thresholdPercent: 90, members: [members[1]!] } }, "alice"); + const execute = vi.fn(); registerServerAdapter({ ...getServerAdapter("claude_local"), execute }); + const fetch = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ five_hour: { utilization: 100 } }))); + proposal.mockResolvedValueOnce({ kind: "exhausted", retryAt: new Date(Date.now() + 60_000).toISOString(), skipped: { [members[1]!.id]: "Exhausted" } }); + const heartbeat = heartbeatService(db, { pluginWorkerManager: worker, runtimeEnv: {} }); + try { + const run = await heartbeat.wakeup(agent!.id, { source: "assignment", reason: "issue_assigned", payload: { issueId: issue!.id }, contextSnapshot: { issueId: issue!.id, wakeReason: "issue_assigned", responsibleUserId: "alice" }, triggerDetail: "system" }); + expect(run).not.toBeNull(); + await expect.poll(async () => (await heartbeat.getRun(run!.id))?.status, { timeout: 20_000 }).toBe("cancelled"); + const cancelled = await heartbeat.getRun(run!.id); expect(cancelled?.errorCode).toBe("ai_connection_pool_exhausted"); + expect(legacyExecutionNeedsReconciliation(cancelled!)).toBe(false); + expect(legacyExecutionNeedsReconciliation({ ...cancelled!, resultJson: { executionRecovery: { kind: "ai_connection_wait", providerWorkStarted: true } } })).toBe(true); + await expect.poll(async () => (await db.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.agentId, agent!.id), eq(heartbeatRuns.scheduledRetryReason, "ai_connection_pool_wait")))).length, { timeout: 5000 }).toBe(1); + const [retry] = await db.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.agentId, agent!.id), eq(heartbeatRuns.scheduledRetryReason, "ai_connection_pool_wait"))); + expect(retry?.contextSnapshot?.aiRouterTaskKey).toBe(issue!.id); expect(retry?.contextSnapshot?.executionRetryAccounting).toMatchObject({ failureRetries: 0 }); expect(execute).not.toHaveBeenCalled(); + const [pin] = await db.select().from(aiConnectionTaskPins).where(and(eq(aiConnectionTaskPins.poolId, pool.id), eq(aiConnectionTaskPins.agentId, agent!.id))); expect(pin?.selection).toEqual(pinned); + const [cursor] = await db.select().from(aiConnectionRouterCursors).where(eq(aiConnectionRouterCursors.poolId, pool.id)); expect(cursor?.version).toBe(1); + await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.agentId, agent!.id)); + } finally { fetch.mockRestore(); proposal.mockReset(); proposal.mockImplementation(defaultProposal); unregisterServerAdapter("claude_local"); await heartbeat.drainActiveRunExecutions(); } + }, 30_000); + it("repairs the failed pool account without replacing the pool or adopting the configured provider", async () => { + const pool = await makePool({ members: [members[1]!] }); + const [agent] = await db.insert(agents).values({ companyId, name: "Pool repair", adapterType: "paperclip_runner", adapterConfig: { provider: "codex", model: "gpt-5.6-sol" }, runtimeConfig: { aiConnection: { mode: "router", connectionId: pool.id } } }).returning(); + const [issue] = await db.insert(issues).values({ companyId, title: "Pool repair", status: "in_progress", assigneeAgentId: agent!.id, responsibleUserId: "alice", createdByUserId: "alice" }).returning(); + const selection = await resolve(pool.id, issue!.id, { agentId: agent!.id }); + const runtime = await prepareManagedAiRuntime(db, { companyId, agentId: agent!.id, responsibleUserId: "alice", adapterType: "paperclip_runner", config: selection.runtimeConfig, binding: selection.binding }); + const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId: agent!.id, status: "failed", errorCode: "acpx_auth_required", responsibleUserId: "alice", startedAt: new Date(), contextSnapshot: { issueId: issue!.id, aiRouterTaskKey: issue!.id, aiRouterSelection: selection, aiConnection: { ...runtime!.attribution, identity: runtime!.identity } } }).returning(); + try { + const intents = connectionIntentService(db); + const card = await intents.requestForRunAuthFailure(run!.id); + expect(card?.service).toBe("anthropic"); expect(card?.state).toBe("needs_user_action"); + const options = await intents.setupOptions(card!.interactionId!); + expect(options.aiConnection).toEqual(selection.binding); expect(options.aiConnectionRequiresAdoption).toBeUndefined(); + const [failedGrant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, members[1]!.binding.grantId)); expect(failedGrant?.status).toBe("needs_reauthorization"); + await db.update(connectionGrants).set({ status: "active" }).where(eq(connectionGrants.id, members[1]!.binding.grantId)); + await db.update(toolConnections).set({ healthStatus: "ok" }).where(eq(toolConnections.id, members[1]!.binding.connectionId)); + expect(await intents.complete(card!.interactionId!, members[1]!.binding.connectionId, "alice")).toMatchObject({ status: "accepted" }); + const [savedAgent] = await db.select().from(agents).where(eq(agents.id, agent!.id)); expect(savedAgent?.runtimeConfig.aiConnection).toEqual({ mode: "router", connectionId: pool.id }); + expect((await resolve(pool.id, issue!.id, { agentId: agent!.id })).binding).toEqual(selection.binding); + await db.update(agents).set({ runtimeConfig: { aiConnection: { mode: "router", connectionId: randomUUID() } } }).where(eq(agents.id, agent!.id)); + await expect(intents.setupOptions(card!.interactionId!)).rejects.toThrow("configuration changed"); + } finally { + await runtime?.cleanup(); + await db.update(connectionGrants).set({ status: "active" }).where(eq(connectionGrants.id, members[1]!.binding.grantId)); + await db.update(toolConnections).set({ healthStatus: "ok" }).where(eq(toolConnections.id, members[1]!.binding.connectionId)); + } + }); it("rejects disabled routing but recovers server-owned native evidence after disable/uninstall", async () => { const pool = await makePool(); const selected = await resolve(pool.id, "recovery"); await instanceSettingsService(db).updateExperimental({ enableAiConnectionRouters: false }); diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index f37ba2eba83..88a5ceabacb 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -7,6 +7,7 @@ import { Router } from "express"; import { z } from "zod"; import { createAiConnectionSchema, + aiConnectionPoolConfigSchema, aiConnectionLoginIntentSchema, localAiConnectionSchema, localAiLoginStartSchema, @@ -10373,6 +10374,31 @@ registerCurrentRoute({ // --- AI runtime connections ------------------------------------------------- +registerCurrentRoute({ + method: "get", + path: "/api/companies/{companyId}/ai-connection-pools", + tags: ["ai-connections"], + summary: "List company AI connection pools for connection managers", + responses: { 200: r.ok(), 401: r.unauthorized, 403: r.forbidden }, +}); + +registerCurrentRoute({ + method: "post", + path: "/api/companies/{companyId}/ai-connection-pools", + tags: ["ai-connections"], + summary: "Create or revise an experimental plugin-owned connection pool", + body: z.object({ pluginKey: z.string().min(1), id: z.string().uuid().optional(), expectedRevision: z.number().int().positive().optional(), config: aiConnectionPoolConfigSchema }).strict(), + responses: { 200: r.ok(), 400: r.badRequest, 401: r.unauthorized, 403: r.forbidden, 404: r.notFound, 409: r.conflict, 422: r.unprocessable }, +}); + +registerCurrentRoute({ + method: "get", + path: "/api/companies/{companyId}/ai-connection-pools/{poolId}/inspection", + tags: ["ai-connections"], + summary: "Inspect authorized pool members and fresh cached usage without probing", + responses: { 200: r.ok(), 401: r.unauthorized, 403: r.forbidden, 404: r.notFound }, +}); + registerCurrentRoute({ method: "get", path: "/api/companies/{companyId}/ai-connections/{connectionId}/usage", diff --git a/server/src/services/ai-connection-router.ts b/server/src/services/ai-connection-router.ts index 18d85a3a1c5..67708d42e76 100644 --- a/server/src/services/ai-connection-router.ts +++ b/server/src/services/ai-connection-router.ts @@ -13,6 +13,7 @@ import type { PluginWorkerManager } from "./plugin-worker-manager.js"; import { resolvePaperclipRunnerProviderProfile, resolvePaperclipRunnerNativeProviderInput } from "./native-runtime/provider-profile.js"; const usageCache = new Map(); +const usageInFlight = new Map>(); const routedRuntimeKeys = new Set(["provider", "acpxAgent", "model", "modelReasoningEffort", "reasoningEffort", "effort", "variant"]); export function applyAiConnectionRouterTaskSettings(config: Record, selection: AiConnectionRouterSelection) { return { ...Object.fromEntries(Object.entries(config).filter(([key]) => !routedRuntimeKeys.has(key))), ...selection.runtimeConfig }; @@ -148,20 +149,31 @@ export function aiConnectionRouterService(db: Db, workerManager?: PluginWorkerMa if (!key) return undefined; const cached = usageCache.get(key); if (cached && cached.expires > Date.now()) return cached.value; + let pending = usageInFlight.get(key); + if (!pending) { + let probeTimer: ReturnType | undefined; + pending = Promise.race([ + aiConnectionService(db).probeUsage(companyId, userId, member.binding.connectionId, member.binding.grantId).catch(() => undefined), + new Promise(resolve => { probeTimer = setTimeout(() => resolve(undefined), 15_000); }), + ]).then(value => { + if (value) { + const resets = value.limits.map(v => Date.parse(v.resetsAt ?? "")).filter(v => Number.isFinite(v) && v > Date.now()); + if (usageCache.size > 1000) usageCache.clear(); + usageCache.set(key, { value, expires: Math.min(Date.now() + 60_000, ...resets) }); + } + return value; + }).finally(() => { if (probeTimer) clearTimeout(probeTimer); usageInFlight.delete(key); }); + usageInFlight.set(key, pending); + } let timer: ReturnType | undefined; try { - const value = await Promise.race([ - aiConnectionService(db).probeUsage(companyId, userId, member.binding.connectionId, member.binding.grantId).catch(() => undefined), - new Promise((resolve) => { timer = setTimeout(() => resolve(undefined), Math.max(0, deadline - Date.now())); }), + return await Promise.race([ + pending, + new Promise(resolve => { timer = setTimeout(() => resolve(undefined), Math.max(0, deadline - Date.now())); }), ]); - if (value) { - const resets = value.limits.map((v) => Date.parse(v.resetsAt ?? "")).filter((v) => Number.isFinite(v) && v > Date.now()); - if (usageCache.size > 1000) usageCache.clear(); - usageCache.set(key, { value, expires: Math.min(Date.now() + 60_000, ...resets) }); - } - return value; } finally { if (timer) clearTimeout(timer); } } + async function resolve(input: { companyId: string; poolId: string; agentId: string; userId: string | null; adapterType: string; taskKey: string; overrides?: Record; existingGrantId?: string; requireExisting?: boolean; persisted?: AiConnectionRouterSelection }) { const pinWhere = and(eq(aiConnectionTaskPins.companyId, input.companyId), eq(aiConnectionTaskPins.poolId, input.poolId), eq(aiConnectionTaskPins.agentId, input.agentId), eq(aiConnectionTaskPins.taskKey, input.taskKey)); const validateSelection = async (selection: AiConnectionRouterSelection, client = db) => aiConnectionService(client).select({ companyId: input.companyId, agentId: input.agentId, userId: input.userId, adapterType: input.adapterType, binding: selection.binding, model: selection.runtimeConfig.model, runnerProvider: selection.runtimeConfig.provider, acpxAgent: selection.runtimeConfig.acpxAgent }); diff --git a/server/src/services/connection-intents.ts b/server/src/services/connection-intents.ts index 34cec954f7f..ccef3280af4 100644 --- a/server/src/services/connection-intents.ts +++ b/server/src/services/connection-intents.ts @@ -3,12 +3,13 @@ import { emailConnectionService } from "./email-connections.js"; import { agentService } from "./agents.js"; import { logActivity } from "./activity-log.js"; import { aiConnectionService } from "./ai-connections.js"; -import { aiConnectionBindingSchema } from "@paperclipai/shared"; +import { aiConnectionBindingSchema, aiConnectionRouterBindingSchema } from "@paperclipai/shared"; import { aiBindingForAuthRecovery, isAiAuthenticationFailure } from "./ai-auth-failure.js"; import { and, eq, desc, isNull, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { agents, + aiConnectionTaskPins, companies, toolConnections, toolCatalogEntries, @@ -227,11 +228,26 @@ export function connectionIntentService(db: Db) { }; } - async function managedAgent(companyId: string, agentId: string, serviceSlug: string, fallback?: AiConnectionBinding) { - const [agent] = await db.select().from(agents).where(and(eq(agents.companyId, companyId), eq(agents.id, agentId))); - const saved = aiConnectionBindingSchema.safeParse(agent?.runtimeConfig?.aiConnection).data; - const binding = saved ?? fallback; - return agent && binding?.provider === serviceSlug ? { agent, binding, requiresAdoption: !saved } : null; + async function managedAgent(companyId: string, agentId: string, serviceSlug: string, options: { fallback?: AiConnectionBinding; sourceRunId?: string | null } = {}, client = db) { + const [agent] = await client.select().from(agents).where(and(eq(agents.companyId, companyId), eq(agents.id, agentId))); + if (!agent) return null; + const router = aiConnectionRouterBindingSchema.safeParse(agent.runtimeConfig.aiConnection).data; + if (router) { + // Repair the durable task allocation, without adopting a personal default + // or replacing the agent's pool. A changed pool invalidates the old card. + if (!options.sourceRunId) return null; + const [source] = await client.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, companyId), eq(heartbeatRuns.agentId, agentId), eq(heartbeatRuns.id, options.sourceRunId))); + const taskKey = text(source?.contextSnapshot?.aiRouterTaskKey); + const evidence = record(source?.contextSnapshot?.aiRouterSelection); + if (!source || !isAiAuthenticationFailure(source.errorCode) || source.status !== "failed" || evidence?.poolId !== router.connectionId || !taskKey) return null; + const [pin] = await client.select().from(aiConnectionTaskPins).where(and(eq(aiConnectionTaskPins.companyId, companyId), eq(aiConnectionTaskPins.poolId, router.connectionId), eq(aiConnectionTaskPins.agentId, agentId), eq(aiConnectionTaskPins.taskKey, taskKey))); + const binding = pin?.selection.binding; + if (!pin || binding?.provider !== serviceSlug || JSON.stringify(evidence.binding) !== JSON.stringify(binding)) return null; + return { agent, binding, adapterConfig: { ...agent.adapterConfig, ...pin.selection.runtimeConfig }, requiresAdoption: false }; + } + const saved = aiConnectionBindingSchema.safeParse(agent.runtimeConfig.aiConnection).data; + const binding = saved ?? options.fallback; + return binding?.provider === serviceSlug ? { agent, binding, adapterConfig: agent.adapterConfig, requiresAdoption: !saved } : null; } async function usableConnectionForAgent(input: { @@ -240,12 +256,13 @@ export function connectionIntentService(db: Db) { responsibleUserId: string; serviceSlug: string; purpose?: "ai" | "channel"; + sourceRunId?: string | null; inventory?: Awaited>; }) { - const managed = input.purpose === "ai" ? await managedAgent(input.companyId, input.agentId, input.serviceSlug) : null; + const managed = input.purpose === "ai" ? await managedAgent(input.companyId, input.agentId, input.serviceSlug, { sourceRunId: input.sourceRunId }) : null; if (managed) { try { - const selected = await aiConnectionService(db).select({ companyId: input.companyId, agentId: input.agentId, userId: input.responsibleUserId, adapterType: managed.agent.adapterType, model: managed.agent.adapterConfig.model, runnerProvider: managed.agent.adapterConfig.provider, acpxAgent: managed.agent.adapterConfig.acpxAgent, binding: managed.binding }); + const selected = await aiConnectionService(db).select({ companyId: input.companyId, agentId: input.agentId, userId: input.responsibleUserId, adapterType: managed.agent.adapterType, model: managed.adapterConfig.model, runnerProvider: managed.adapterConfig.provider, acpxAgent: managed.adapterConfig.acpxAgent, binding: managed.binding }); return access.getConnection(selected.connection.id, input.companyId); } catch (error) { if ([403, 404, 422].includes((error as { status?: number }).status ?? 0)) return null; throw error; } } @@ -690,6 +707,7 @@ export function connectionIntentService(db: Db) { responsibleUserId: context.run.responsibleUserId!, serviceSlug: app.slug, purpose: options.purpose, + sourceRunId: context.run.id, }); if (ready) { return { @@ -784,7 +802,7 @@ export function connectionIntentService(db: Db) { readyConnectionId: inboxes.find(inbox => inbox.id === interactionId)?.connectionId ?? null }, }; } - let managed = payload.purpose === "ai" ? await managedAgent(loaded.issue.companyId, payload.requestingAgentId, app.slug) : null; + let managed = payload.purpose === "ai" ? await managedAgent(loaded.issue.companyId, payload.requestingAgentId, app.slug, { sourceRunId: loaded.interaction.sourceRunId }) : null; if (payload.purpose === "ai" && !managed) { const [source] = await db.select().from(heartbeatRuns).where(and( eq(heartbeatRuns.id, loaded.interaction.sourceRunId!), eq(heartbeatRuns.companyId, loaded.issue.companyId), @@ -792,21 +810,21 @@ export function connectionIntentService(db: Db) { )); const [agent] = await db.select().from(agents).where(and(eq(agents.id, payload.requestingAgentId), eq(agents.companyId, loaded.issue.companyId))); if (source?.status === "failed" && isAiAuthenticationFailure(source.errorCode) && !source.contextSnapshot?.aiConnection && agent && !agent.runtimeConfig.aiConnection) { - managed = await managedAgent(loaded.issue.companyId, agent.id, app.slug, aiBindingForAuthRecovery(agent.adapterType, agent.adapterConfig)); + managed = await managedAgent(loaded.issue.companyId, agent.id, app.slug, { fallback: aiBindingForAuthRecovery(agent.adapterType, agent.adapterConfig) }); } } if (payload.purpose === "ai" && !managed) throw conflict("The agent’s AI configuration changed. Start a new execution."); const inventory = await connectionInventory(loaded.issue.companyId); let usableAiConnection = managed ? await usableConnectionForAgent({ companyId: loaded.issue.companyId, agentId: payload.requestingAgentId, - responsibleUserId: loaded.interaction.addresseeUserId!, serviceSlug: app.slug, purpose: "ai", + responsibleUserId: loaded.interaction.addresseeUserId!, serviceSlug: app.slug, purpose: "ai", sourceRunId: loaded.interaction.sourceRunId, }) : null; if (managed?.requiresAdoption) { try { const selected = await aiConnectionService(db).select({ companyId: loaded.issue.companyId, agentId: managed.agent.id, userId: loaded.interaction.addresseeUserId!, - adapterType: managed.agent.adapterType, model: managed.agent.adapterConfig.model, - runnerProvider: managed.agent.adapterConfig.provider, acpxAgent: managed.agent.adapterConfig.acpxAgent, + adapterType: managed.agent.adapterType, model: managed.adapterConfig.model, + runnerProvider: managed.adapterConfig.provider, acpxAgent: managed.adapterConfig.acpxAgent, binding: managed.binding, allowUninstalledPersonal: true, }); usableAiConnection = await access.getConnection(selected.connection.id, loaded.issue.companyId); @@ -955,19 +973,19 @@ export function connectionIntentService(db: Db) { action: "agent.updated", entityType: "agent", entityId: agent.id, details: { connectionIntentId: interactionId, aiConnectionAdopted: true }, }); - managed = { agent: updated, binding, requiresAdoption: false }; + managed = { agent: updated, binding, adapterConfig: updated.adapterConfig, requiresAdoption: false }; } else { - managed = await managedAgent(loaded.issue.companyId, payload.requestingAgentId, payload.serviceSlug); + managed = await managedAgent(loaded.issue.companyId, payload.requestingAgentId, payload.serviceSlug, { sourceRunId: loaded.interaction.sourceRunId }, txDb); } if (!managed) throw conflict("Configure the agent’s AI connection before using this account"); const service = aiConnectionService(txDb); if (managed.binding.mode === "responsible_user") { - const selected = await service.select({ companyId: loaded.issue.companyId, agentId: payload.requestingAgentId, userId, adapterType: managed.agent.adapterType, model: managed.agent.adapterConfig.model, runnerProvider: managed.agent.adapterConfig.provider, acpxAgent: managed.agent.adapterConfig.acpxAgent, binding: managed.binding, allowUninstalledPersonal: true }); + const selected = await service.select({ companyId: loaded.issue.companyId, agentId: payload.requestingAgentId, userId, adapterType: managed.agent.adapterType, model: managed.adapterConfig.model, runnerProvider: managed.adapterConfig.provider, acpxAgent: managed.adapterConfig.acpxAgent, binding: managed.binding, allowUninstalledPersonal: true }); if (selected.connection.id !== selectedConnection.id) throw conflict("Choose this account as your personal default in Connections first"); const installs = await txAccess.listConnectionInstalls(selectedConnection.id, loaded.issue.companyId); await txAccess.putConnectionInstalls(selectedConnection.id, { installs: [...installs, { targetType: "agent", targetId: payload.requestingAgentId }] }, { actorType: "user", actorId: userId }); } - const selected = await service.select({ companyId: loaded.issue.companyId, agentId: payload.requestingAgentId, userId, adapterType: managed.agent.adapterType, model: managed.agent.adapterConfig.model, runnerProvider: managed.agent.adapterConfig.provider, acpxAgent: managed.agent.adapterConfig.acpxAgent, binding: managed.binding }); + const selected = await service.select({ companyId: loaded.issue.companyId, agentId: payload.requestingAgentId, userId, adapterType: managed.agent.adapterType, model: managed.adapterConfig.model, runnerProvider: managed.adapterConfig.provider, acpxAgent: managed.adapterConfig.acpxAgent, binding: managed.binding }); if (selected.connection.id !== selectedConnection.id) throw conflict("This is not the account selected for the agent"); return txInteractions.resolveConnectionIntent(loaded.issue, interactionId, { version: 1, outcome: "connected", connectionId: selected.connection.id }, { userId }); } @@ -1094,11 +1112,14 @@ export function connectionIntentService(db: Db) { const [agent] = await db.select().from(agents).where(and(eq(agents.id, run.agentId), eq(agents.companyId, run.companyId))); if (!agent) return null; const saved = aiConnectionBindingSchema.safeParse(agent.runtimeConfig.aiConnection).data; - const binding = saved ?? aiBindingForAuthRecovery(agent.adapterType, agent.adapterConfig); + const router = aiConnectionRouterBindingSchema.safeParse(agent.runtimeConfig.aiConnection).data; + const selected = router ? record(run.contextSnapshot?.aiRouterSelection) : null; + const managed = router && selected ? await managedAgent(run.companyId, run.agentId, String(record(selected.binding)?.provider), { sourceRunId: run.id }) : null; + const binding = router ? managed?.binding : saved ?? aiBindingForAuthRecovery(agent.adapterType, agent.adapterConfig); if (!binding) return null; const attribution = record(run.contextSnapshot?.aiConnection); if (attribution && attribution.provider !== binding.provider) return null; - if (saved && attribution && typeof attribution.identity === "string" && typeof attribution.connectionId === "string" && typeof attribution.grantId === "string") { + if ((saved || managed) && attribution && typeof attribution.identity === "string" && typeof attribution.connectionId === "string" && typeof attribution.grantId === "string") { await aiConnectionService(db).markAuthenticationFailed({ companyId: run.companyId, runId: run.id, agentId: run.agentId, runStartedAt: run.startedAt ?? run.createdAt, attribution: attribution as unknown as AiConnectionAttribution & { identity: string } }); diff --git a/server/src/services/legacy-execution-recovery.ts b/server/src/services/legacy-execution-recovery.ts index 98c4c03f7e5..bcdf6eb53ce 100644 --- a/server/src/services/legacy-execution-recovery.ts +++ b/server/src/services/legacy-execution-recovery.ts @@ -38,7 +38,7 @@ export function legacyExecutionNeedsReconciliation( && (run.resultJson?.executionCancellation as Record | undefined)?.state === "acknowledged") return false; // Waiting for a subscription or workspace precedes provider execution. It is // a resource wait, not a failed provider attempt or permission to replay work. - if (run.status === "cancelled" && run.errorCode === "ai_connection_busy" && + if (run.status === "cancelled" && (run.errorCode === "ai_connection_busy" || run.errorCode === "ai_connection_pool_exhausted") && evidence?.kind === "ai_connection_wait" && evidence.providerWorkStarted === false) return false; if (run.status === "cancelled" && run.errorCode === "workspace_busy" && evidence?.kind === "workspace_wait" && evidence.providerWorkStarted === false) return false; diff --git a/ui/src/pages/InstanceExperimentalSettings.tsx b/ui/src/pages/InstanceExperimentalSettings.tsx index a9e9e47db1e..7f4018c5bde 100644 --- a/ui/src/pages/InstanceExperimentalSettings.tsx +++ b/ui/src/pages/InstanceExperimentalSettings.tsx @@ -305,6 +305,17 @@ export function InstanceExperimentalSettings() { ariaLabel="Toggle agent chat experimental setting" /> + toggleMutation.mutate({ enableAiConnectionRouters: checked })} + disabled={toggleMutation.isPending} + settingKey="enableAiConnectionRouters" + managed={managedKeys.enableAiConnectionRouters} + ariaLabel="Toggle AI connection routers experimental setting" + /> + - toggleMutation.mutate({ enableAiConnectionRouters: checked })} - disabled={toggleMutation.isPending} - settingKey="enableAiConnectionRouters" - managed={managedKeys.enableAiConnectionRouters} - ariaLabel="Toggle AI connection routers experimental setting" - /> -