mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-11 13:24:55 +01:00
feat(connections): add experimental task-pinned AI routing
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
1 parent
806af230b6
commit
c2efe3e92a
51 files changed
+1763
-60
No files matched your search
@@ -0,0 +1,238 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { and, eq, sql } from "drizzle-orm";
|
||||
import { aiConnectionPools, aiConnectionRouterCursors, aiConnectionTaskPins, companySecrets, connectionGrants, instanceSettings, plugins, toolApplications, toolConnections, toolConnectionInstalls, type Db } from "@paperclipai/db";
|
||||
import { aiConnectionPoolConfigSchema, isAiConnectionCompatible, type AiConnectionPool, type AiConnectionPoolMember, type AiConnectionRouterRequest, type AiConnectionRouterSelection, type AiConnectionPoolSaveInput, type AiConnectionUsage } from "@paperclipai/shared";
|
||||
import { isCodexLocalKnownModel, codexLocalReasoningEffortsForModel } from "@paperclipai/adapter-codex-local";
|
||||
import { models as claudeModels, claudeLocalReasoningEffortsForModel } from "@paperclipai/adapter-claude-local";
|
||||
import { models as grokModels, grokLocalReasoningEffortsForModel } from "@paperclipai/adapter-grok-local";
|
||||
import { aiConnectionService } from "./ai-connections.js";
|
||||
import { instanceSettingsService } from "./instance-settings.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { forbidden, notFound, conflict, unprocessable } from "../errors.js";
|
||||
import type { PluginWorkerManager } from "./plugin-worker-manager.js";
|
||||
import { resolvePaperclipRunnerProviderProfile, resolvePaperclipRunnerNativeProviderInput } from "./native-runtime/provider-profile.js";
|
||||
|
||||
const usageCache = new Map<string, { expires: number; value: AiConnectionUsage }>();
|
||||
const routedRuntimeKeys = new Set(["provider", "acpxAgent", "model", "modelReasoningEffort", "reasoningEffort", "effort", "variant"]);
|
||||
export function applyAiConnectionRouterTaskSettings(config: Record<string, unknown>, selection: AiConnectionRouterSelection) {
|
||||
return { ...Object.fromEntries(Object.entries(config).filter(([key]) => !routedRuntimeKeys.has(key))), ...selection.runtimeConfig };
|
||||
}
|
||||
export class AiConnectionPoolExhausted extends Error {
|
||||
readonly code = "ai_connection_pool_exhausted";
|
||||
constructor(readonly retryAt: string, readonly poolId: string) { super("All eligible pool accounts are over their usage limit. Waiting for usage to reset."); }
|
||||
}
|
||||
|
||||
/** Routing can override provider/model only, never environment, permissions, or instructions. */
|
||||
export function poolMemberRuntimeConfig(member: AiConnectionPoolMember, adapterType: string, overrides: Record<string, unknown> = {}) {
|
||||
const notes: string[] = [];
|
||||
const provider = member.binding.provider;
|
||||
const profile = member.profile;
|
||||
const expected = provider === "openai" ? profile.provider === "codex"
|
||||
: provider === "anthropic" ? profile.provider === "acpx" && profile.acpxAgent === "claude"
|
||||
: provider === "xai" ? profile.provider === "acpx" && profile.acpxAgent === "grok"
|
||||
: profile.provider === "opencode";
|
||||
if (!expected) throw unprocessable("Pool member harness does not match its account");
|
||||
const supportsModel = (value: string) => provider === "openai" ? isCodexLocalKnownModel(value)
|
||||
: provider === "anthropic" ? claudeModels.some((m) => m.id === value)
|
||||
: provider === "xai" ? grokModels.some((m) => m.id === value) : value.startsWith("openrouter/");
|
||||
const config: Record<string, unknown> = { provider: profile.provider, ...(profile.acpxAgent ? { acpxAgent: profile.acpxAgent } : {}), model: profile.model };
|
||||
const validate = () => {
|
||||
if (!isAiConnectionCompatible(member.binding, adapterType, config.model, config.provider, config.acpxAgent)) throw unprocessable("Pool member is incompatible with this harness");
|
||||
if (adapterType === "paperclip_runner") {
|
||||
const resolved = resolvePaperclipRunnerProviderProfile(config);
|
||||
resolvePaperclipRunnerNativeProviderInput({ backend: resolved.backend, adapterConfig: config });
|
||||
}
|
||||
};
|
||||
validate();
|
||||
if (typeof overrides.model === "string" && overrides.model !== profile.model) {
|
||||
if (supportsModel(overrides.model)) {
|
||||
config.model = overrides.model;
|
||||
try { validate(); } catch { config.model = profile.model; notes.push("Model override is unavailable for this member; using its default."); }
|
||||
} else notes.push("Model override is unavailable for this member; using its default.");
|
||||
}
|
||||
const model = String(config.model);
|
||||
const efforts = provider === "openai" ? codexLocalReasoningEffortsForModel(model)
|
||||
: provider === "anthropic" ? claudeLocalReasoningEffortsForModel(model)
|
||||
: provider === "xai" ? grokLocalReasoningEffortsForModel(model) : [];
|
||||
// ACPX currently has no qualified effort override; keep that host limitation explicit.
|
||||
const effortKey = adapterType === "paperclip_runner" ? profile.provider === "codex" ? "modelReasoningEffort" : null
|
||||
: provider === "anthropic" ? "effort" : provider === "xai" ? "reasoningEffort" : provider === "openai" ? "modelReasoningEffort" : "variant";
|
||||
const requested = overrides.modelReasoningEffort ?? overrides.reasoningEffort ?? overrides.effort ?? overrides.variant;
|
||||
const effort = typeof requested === "string" && efforts.includes(requested) && effortKey ? requested
|
||||
: profile.effort && efforts.includes(profile.effort) && effortKey ? profile.effort : null;
|
||||
if (requested && requested !== effort) notes.push("Effort override is unavailable for this member; using its default.");
|
||||
if (effort && effortKey) config[effortKey] = effort;
|
||||
validate();
|
||||
return { config, notes };
|
||||
}
|
||||
|
||||
export function aiConnectionRouterService(db: Db, workerManager?: PluginWorkerManager) {
|
||||
async function owner(pluginKey: string) {
|
||||
const [plugin] = await db.select().from(plugins).where(eq(plugins.pluginKey, pluginKey));
|
||||
if (!plugin || plugin.status !== "ready" || !plugin.manifestJson.capabilities.includes("ai.connections.route")) throw unprocessable("Enable a compatible AI connection router plugin");
|
||||
return plugin;
|
||||
}
|
||||
async function enabled() {
|
||||
if (!(await instanceSettingsService(db).getExperimental()).enableAiConnectionRouters) throw unprocessable("Enable AI connection routers in Experimental settings", { code: "ai_connection_router_disabled" });
|
||||
}
|
||||
async function list(companyId: string) {
|
||||
const rows = await db.select().from(aiConnectionPools).where(eq(aiConnectionPools.companyId, companyId));
|
||||
return rows.map((row): AiConnectionPool => ({ ...row.config, id: row.id, companyId, pluginKey: row.pluginKey, revision: row.revision }));
|
||||
}
|
||||
async function selectable(companyId: string, userId: string) {
|
||||
if (!(await instanceSettingsService(db).getExperimental()).enableAiConnectionRouters) return [];
|
||||
const accounts = await aiConnectionService(db).list(companyId, userId);
|
||||
const result: AiConnectionPool[] = [];
|
||||
for (const pool of await list(companyId)) {
|
||||
if (!pool.enabled) continue;
|
||||
try { await owner(pool.pluginKey); } catch { continue; }
|
||||
const members = pool.members.filter(member => accounts.some(account => account.id === member.binding.connectionId && account.grantId === member.binding.grantId));
|
||||
if (members.length) result.push({ ...pool, members });
|
||||
}
|
||||
return result;
|
||||
}
|
||||
async function inspect(companyId: string, poolId: string, userId: string) {
|
||||
const pool = (await list(companyId)).find(pool => pool.id === poolId);
|
||||
if (!pool) throw notFound("Connection pool not found");
|
||||
const visible = await aiConnectionService(db).list(companyId, userId);
|
||||
const result: Record<string, { reason: string; checkedAt: string | null; usage?: AiConnectionUsage }> = {};
|
||||
for (const member of pool.members) {
|
||||
const account = visible.find(account => account.id === member.binding.connectionId && account.grantId === member.binding.grantId);
|
||||
const key = account ? await usageCacheKey(member, companyId) : null;
|
||||
const cached = key ? usageCache.get(key) : undefined;
|
||||
const observation = cached && cached.expires > Date.now() ? cached.value : undefined;
|
||||
result[member.id] = { reason: account?.status === "connected" ? "Eligible if agent access permits" : account?.status ?? "Not permitted", checkedAt: account && observation ? observation.checkedAt : null, ...(account && observation ? { usage: observation } : {}) };
|
||||
}
|
||||
return result;
|
||||
}
|
||||
async function save(pluginKey: string, input: AiConnectionPoolSaveInput, userId: string) {
|
||||
await enabled();
|
||||
const plugin = await owner(pluginKey);
|
||||
const config = aiConnectionPoolConfigSchema.parse(input.config);
|
||||
const visible = await aiConnectionService(db).list(input.companyId, userId);
|
||||
for (const member of config.members) {
|
||||
if (!visible.some((v) => v.id === member.binding.connectionId && v.grantId === member.binding.grantId)) throw forbidden("Choose an account you may use in this company");
|
||||
poolMemberRuntimeConfig(member, "paperclip_runner");
|
||||
}
|
||||
const id = input.id ?? randomUUID();
|
||||
await db.transaction(async (tx) => {
|
||||
if (input.id) {
|
||||
const [row] = await tx.select().from(aiConnectionPools).where(and(eq(aiConnectionPools.id, id), eq(aiConnectionPools.companyId, input.companyId), eq(aiConnectionPools.pluginKey, pluginKey))).for("update");
|
||||
if (!row) throw notFound("Connection pool not found");
|
||||
if (input.expectedRevision !== row.revision) throw conflict("Pool changed; reload before saving");
|
||||
await tx.update(aiConnectionPools).set({ config, revision: row.revision + 1, updatedAt: new Date() }).where(eq(aiConnectionPools.id, id));
|
||||
await tx.update(toolConnections).set({ name: config.name, enabled: config.enabled, updatedAt: new Date() }).where(and(eq(toolConnections.id, id), eq(toolConnections.companyId, input.companyId)));
|
||||
} else {
|
||||
const key = `plugin:${pluginKey}:ai-router`;
|
||||
await tx.insert(toolApplications).values({ companyId: input.companyId, applicationKey: key, name: `${plugin.manifestJson.displayName} connections`, type: "paperclip_plugin", pluginId: plugin.id }).onConflictDoNothing();
|
||||
const [application] = await tx.select().from(toolApplications).where(and(eq(toolApplications.companyId, input.companyId), eq(toolApplications.applicationKey, key)));
|
||||
if (!application) throw new Error("Router application could not be created");
|
||||
await tx.insert(toolConnections).values({ id, companyId: input.companyId, applicationId: application.id, name: config.name, uid: id, connectionPurpose: "ai", transport: "runtime_auth", enabled: config.enabled, status: "active", healthStatus: "ok", config: { aiRouter: { pluginKey } }, createdByUserId: userId });
|
||||
await tx.insert(aiConnectionPools).values({ id, companyId: input.companyId, pluginKey, config });
|
||||
await tx.insert(aiConnectionRouterCursors).values({ poolId: id, companyId: input.companyId });
|
||||
await tx.insert(toolConnectionInstalls).values({ companyId: input.companyId, connectionId: id, targetType: "company", targetId: input.companyId, createdByUserId: userId });
|
||||
}
|
||||
await logActivity(tx as unknown as Db, { companyId: input.companyId, actorType: "user", actorId: userId, action: "ai_connection.pool_saved", entityType: "tool_connection", entityId: id, details: { enabled: config.enabled, memberCount: config.members.length, mode: config.mode } });
|
||||
});
|
||||
return (await list(input.companyId)).find((p) => p.id === id)!;
|
||||
}
|
||||
async function usageCacheKey(member: AiConnectionPoolMember, companyId: string) {
|
||||
const [grant] = await db.select().from(connectionGrants).where(and(eq(connectionGrants.companyId, companyId), eq(connectionGrants.id, member.binding.grantId)));
|
||||
if (!grant) return null;
|
||||
const ref = grant.credentialSecretRefs.find((r) => r.configPath === "ai.credential");
|
||||
const [secret] = ref ? await db.select({ version: companySecrets.latestVersion }).from(companySecrets).where(and(eq(companySecrets.companyId, companyId), eq(companySecrets.id, ref.secretId))) : [];
|
||||
return `${companyId}:${grant.id}:${secret?.version}:${grant.updatedAt.toISOString()}`;
|
||||
}
|
||||
async function probe(member: AiConnectionPoolMember, companyId: string, userId: string, deadline: number) {
|
||||
const key = await usageCacheKey(member, companyId);
|
||||
if (!key) return undefined;
|
||||
const cached = usageCache.get(key);
|
||||
if (cached && cached.expires > Date.now()) return cached.value;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
const value = await Promise.race([
|
||||
aiConnectionService(db).probeUsage(companyId, userId, member.binding.connectionId, member.binding.grantId).catch(() => undefined),
|
||||
new Promise<undefined>((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<string, unknown>; 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 });
|
||||
if (input.persisted) {
|
||||
if (input.persisted.poolId !== input.poolId) throw unprocessable("Persisted pool selection does not match this run");
|
||||
const [pin] = await db.select().from(aiConnectionTaskPins).where(pinWhere);
|
||||
// Uninstall can remove the virtual connection. Server-owned native recovery
|
||||
// evidence still authorizes the original concrete selection, never a new one.
|
||||
if (pin && (pin.selection.memberId !== input.persisted.memberId || JSON.stringify(pin.selection.binding) !== JSON.stringify(input.persisted.binding))) throw unprocessable("Persisted pool allocation is inconsistent");
|
||||
await validateSelection(input.persisted);
|
||||
return input.persisted;
|
||||
}
|
||||
await enabled();
|
||||
for (let attempt = 0; attempt < 20; attempt++) {
|
||||
const [row] = await db.select().from(aiConnectionPools).where(and(eq(aiConnectionPools.id, input.poolId), eq(aiConnectionPools.companyId, input.companyId)));
|
||||
if (!row || !row.config.enabled) throw unprocessable("Enable the selected connection pool");
|
||||
const plugin = await owner(row.pluginKey);
|
||||
if (!workerManager?.isRunning(plugin.id)) throw unprocessable("The connection router worker is unavailable; retry when it is ready");
|
||||
const [pin] = await db.select().from(aiConnectionTaskPins).where(pinWhere);
|
||||
const [cursor] = await db.select().from(aiConnectionRouterCursors).where(and(eq(aiConnectionRouterCursors.poolId, row.id), eq(aiConnectionRouterCursors.companyId, input.companyId)));
|
||||
if (!cursor) throw unprocessable("The connection pool cursor is unavailable");
|
||||
const pool: AiConnectionPool = { ...row.config, id: row.id, companyId: row.companyId, pluginKey: row.pluginKey, revision: row.revision };
|
||||
const deadline = Date.now() + 15_000;
|
||||
const candidates: AiConnectionRouterRequest["candidates"] = [];
|
||||
for (const member of pin ? [pin.member] : pool.members) {
|
||||
try {
|
||||
const effective = poolMemberRuntimeConfig(member, input.adapterType, input.overrides);
|
||||
const selection: AiConnectionRouterSelection = { poolId: row.id, memberId: member.id, binding: member.binding, runtimeConfig: effective.config, notes: effective.notes };
|
||||
await validateSelection(selection);
|
||||
candidates.push({ member, runtimeConfig: effective.config, notes: effective.notes });
|
||||
} catch (error) { if (pin) throw error; }
|
||||
}
|
||||
if (!pin && input.requireExisting) {
|
||||
const existing = candidates.find((candidate) => candidate.member.binding.grantId === input.existingGrantId);
|
||||
if (!existing) throw unprocessable("Reset this task session before adopting the connection pool");
|
||||
candidates.splice(0, candidates.length, existing);
|
||||
}
|
||||
if (pool.mode === "usage_aware" && input.userId) {
|
||||
// At most four provider calls at once; the shared deadline bounds the whole pool.
|
||||
for (let i = 0; i < candidates.length && Date.now() < deadline; i += 4) await Promise.all(candidates.slice(i, i + 4).map(async (candidate) => { candidate.usage = await probe(candidate.member, input.companyId, input.userId!, deadline); }));
|
||||
}
|
||||
const authorizedPool = { ...pool, members: candidates.map((candidate) => candidate.member) };
|
||||
const result = await workerManager.call(plugin.id, "routeAiConnection", { companyId: input.companyId, pool: authorizedPool, memberOrder: pool.members.map((member) => member.id), agentId: input.agentId, taskKey: input.taskKey, lastMemberId: cursor.lastMemberId, cursorVersion: cursor.version, ...(pin ? { pinnedMemberId: pin.member.id } : {}), candidates, now: new Date().toISOString() }, 5000);
|
||||
if (result.kind === "exhausted") {
|
||||
const retry = Date.parse(result.retryAt);
|
||||
throw new AiConnectionPoolExhausted(new Date(Number.isFinite(retry) ? Math.max(Date.now() + 1000, retry) : Date.now() + 60_000).toISOString(), pool.id);
|
||||
}
|
||||
if (result.kind !== "selected") throw unprocessable("No pool account is permitted and compatible for this task");
|
||||
const candidate = candidates.find((v) => v.member.id === result.memberId);
|
||||
if (!candidate || (pin && candidate.member.id !== pin.member.id)) throw unprocessable("The router returned an unauthorized member");
|
||||
const selected: AiConnectionRouterSelection = { poolId: row.id, memberId: candidate.member.id, binding: candidate.member.binding, runtimeConfig: candidate.runtimeConfig, notes: candidate.notes };
|
||||
if (pin) return selected;
|
||||
const committed = await db.transaction(async (tx) => {
|
||||
const [locked] = await tx.select().from(aiConnectionRouterCursors).where(and(eq(aiConnectionRouterCursors.poolId, row.id), eq(aiConnectionRouterCursors.companyId, input.companyId))).for("update");
|
||||
const [existing] = await tx.select().from(aiConnectionTaskPins).where(pinWhere);
|
||||
if (existing) return { retry: true };
|
||||
const [current] = await tx.select().from(aiConnectionPools).where(and(eq(aiConnectionPools.id, row.id), eq(aiConnectionPools.companyId, input.companyId))).for("share");
|
||||
if (!locked || locked.version !== cursor.version || !current || current.revision !== row.revision) return { retry: true };
|
||||
const [currentPlugin] = await tx.select().from(plugins).where(eq(plugins.id, plugin.id)).for("share");
|
||||
if (currentPlugin?.status !== "ready") throw unprocessable("The router plugin was disabled during allocation");
|
||||
await tx.select({ id: instanceSettings.id }).from(instanceSettings).where(eq(instanceSettings.singletonKey, "default")).for("share");
|
||||
if (!(await instanceSettingsService(tx as unknown as Db).getExperimental()).enableAiConnectionRouters) throw unprocessable("AI connection routers were disabled during allocation");
|
||||
await validateSelection(selected, tx as unknown as Db);
|
||||
await tx.insert(aiConnectionTaskPins).values({ companyId: input.companyId, poolId: row.id, agentId: input.agentId, taskKey: input.taskKey, selection: selected, member: candidate.member });
|
||||
await tx.update(aiConnectionRouterCursors).set({ lastMemberId: candidate.member.id, version: sql`${aiConnectionRouterCursors.version} + 1` }).where(eq(aiConnectionRouterCursors.poolId, row.id));
|
||||
await logActivity(tx as unknown as Db, { companyId: input.companyId, actorType: "system", actorId: "system", agentId: input.agentId, action: "ai_connection.pool_selected", entityType: "tool_connection", entityId: row.id, details: { memberId: candidate.member.id, connectionId: candidate.member.binding.connectionId, taskKey: input.taskKey } });
|
||||
return { retry: false };
|
||||
});
|
||||
if (!committed.retry) return selected;
|
||||
}
|
||||
throw conflict("Concurrent pool selections changed the cursor; retry this task");
|
||||
}
|
||||
return { list, selectable, inspect, save, resolve };
|
||||
}
|
||||
@@ -189,7 +189,13 @@ export function managedAiSessionFingerprintConfig(
|
||||
for (const [key, value] of Object.entries(managedAiHomeEnvironment(managedHome))) {
|
||||
if (env[key] === value) env[key] = stable[key];
|
||||
}
|
||||
return { ...config, env };
|
||||
const managed = config.managedAiConnection as Record<string, unknown> | undefined;
|
||||
if (managed?.sessionIdentity) {
|
||||
for (const key of ["ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "CLAUDE_CODE_OAUTH_TOKEN", "OPENAI_API_KEY", "CODEX_API_KEY", "OPENROUTER_API_KEY", "XAI_API_KEY", "GROK_API_KEY", "OPENCODE_AUTH_JSON", "OPENCODE_CONFIG_CONTENT"]) {
|
||||
if (env[key]) env[key] = "<managed-ai-credential>";
|
||||
}
|
||||
}
|
||||
return { ...config, env, ...(managed?.sessionIdentity ? { managedAiConnection: { ...managed, identity: managed.sessionIdentity } } : {}) };
|
||||
}
|
||||
|
||||
export async function prepareManagedAiRuntime(
|
||||
@@ -251,7 +257,14 @@ export async function prepareManagedAiRuntime(
|
||||
throw unprocessable(
|
||||
"The selected default changed. Retry this execution.",
|
||||
);
|
||||
const credentialRef = selection.grant.credentialSecretRefs.find((ref) => ref.configPath === "ai.credential");
|
||||
if (!credentialRef) throw unprocessable("The selected AI credential is unavailable");
|
||||
const readFreshness = async () => (await db.select({ epoch: companySecrets.aiSessionEpoch, version: companySecrets.latestVersion })
|
||||
.from(companySecrets).where(and(eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, credentialRef.secretId))).limit(1))[0];
|
||||
const freshness = await readFreshness();
|
||||
const value = await service.credential(selection);
|
||||
const afterRead = await readFreshness();
|
||||
if (!freshness || freshness.version !== afterRead?.version || freshness.epoch !== afterRead.epoch) throw unprocessable("The AI credential changed during preparation; retry this execution");
|
||||
home = await mkdtemp(
|
||||
path.join(
|
||||
os.tmpdir(),
|
||||
@@ -298,11 +311,13 @@ export async function prepareManagedAiRuntime(
|
||||
.digest("hex")
|
||||
.slice(0, 16);
|
||||
const identity = `${selection.grant.id}:${input.responsibleUserId ?? "shared"}:${generation}`;
|
||||
const sessionIdentity = `${selection.grant.id}:${input.responsibleUserId ?? "shared"}:${freshness.epoch}`;
|
||||
return {
|
||||
sessionIdentity,
|
||||
config: {
|
||||
...input.config,
|
||||
env,
|
||||
managedAiConnection: { ...selection.attribution, identity },
|
||||
managedAiConnection: { ...selection.attribution, identity, sessionIdentity },
|
||||
},
|
||||
attribution: selection.attribution,
|
||||
accountName: selection.connection.name,
|
||||
@@ -367,7 +382,7 @@ export async function prepareManagedAiRuntime(
|
||||
if (decision !== 10) return;
|
||||
await secretService(tx).rotate(
|
||||
ref.secretId,
|
||||
{ value: refreshed },
|
||||
{ value: refreshed, preserveAiSessionEpoch: true },
|
||||
{ userId: grant.subjectUserId },
|
||||
);
|
||||
await tx
|
||||
|
||||
@@ -27,7 +27,7 @@ describe("failure attempts across resource waits", () => {
|
||||
});
|
||||
|
||||
describe("persisted independent accounting", () => {
|
||||
it.each(["max_turns_continuation", "issue_disposition_repair", "workspace_busy", "ai_connection_busy"])("%s cannot erase prior infrastructure debits or spend more", scheduledRetryReason => {
|
||||
it.each(["max_turns_continuation", "issue_disposition_repair", "workspace_busy", "ai_connection_busy", "ai_connection_pool_wait"])("%s cannot erase prior infrastructure debits or spend more", scheduledRetryReason => {
|
||||
expect(executionFailureRetryCount({ scheduledRetryReason, scheduledRetryAttempt: 20,
|
||||
contextSnapshot: { executionRetryAccounting: { version: 1, failureRetries: 2, maxTurnContinuations: 1 } },
|
||||
})).toBe(2);
|
||||
|
||||
@@ -28,7 +28,7 @@ function savedAccounting(run: RetryRun): ExecutionRetryAccounting | null {
|
||||
|
||||
function historicalFailureCount(run: RetryRun): number {
|
||||
if (run.scheduledRetryReason === "max_turns_continuation" || run.scheduledRetryReason === "issue_disposition_repair") return 0;
|
||||
if (run.scheduledRetryReason === "ai_connection_busy") {
|
||||
if (run.scheduledRetryReason === "ai_connection_busy" || run.scheduledRetryReason === "ai_connection_pool_wait") {
|
||||
const saved = count(run.contextSnapshot?.failureRetriesBeforeAiConnectionWait);
|
||||
if (saved !== null) return saved;
|
||||
}
|
||||
@@ -42,7 +42,7 @@ function historicalFailureCount(run: RetryRun): number {
|
||||
|
||||
export function executionRetryAccounting(run: RetryRun): ExecutionRetryAccounting {
|
||||
const saved = savedAccounting(run);
|
||||
const nonFailureLane = ["max_turns_continuation", "issue_disposition_repair", "workspace_busy", "ai_connection_busy"].includes(run.scheduledRetryReason ?? "");
|
||||
const nonFailureLane = ["max_turns_continuation", "issue_disposition_repair", "workspace_busy", "ai_connection_busy", "ai_connection_pool_wait"].includes(run.scheduledRetryReason ?? "");
|
||||
return {
|
||||
version: 1,
|
||||
failureRetries: Math.max(saved?.failureRetries ?? 0, saved && nonFailureLane ? 0 : historicalFailureCount(run)),
|
||||
@@ -57,7 +57,7 @@ export function executionFailureRetryCount(run: RetryRun): number {
|
||||
}
|
||||
|
||||
export function executionRetryAttemptCount(run: RetryRun, reason: string): number {
|
||||
if (reason === "workspace_busy" || reason === "ai_connection_busy") {
|
||||
if (reason === "workspace_busy" || reason === "ai_connection_busy" || reason === "ai_connection_pool_wait") {
|
||||
return run.scheduledRetryReason === reason ? count(run.scheduledRetryAttempt) ?? 0 : 0;
|
||||
}
|
||||
const accounting = executionRetryAccounting(run);
|
||||
@@ -67,6 +67,6 @@ export function executionRetryAttemptCount(run: RetryRun, reason: string): numbe
|
||||
export function accountingForScheduledRetry(run: RetryRun, reason: string, attempt: number): ExecutionRetryAccounting {
|
||||
const accounting = executionRetryAccounting(run);
|
||||
if (reason === "max_turns_continuation") accounting.maxTurnContinuations = attempt;
|
||||
else if (reason !== "workspace_busy" && reason !== "ai_connection_busy") accounting.failureRetries = attempt;
|
||||
else if (reason !== "workspace_busy" && reason !== "ai_connection_busy" && reason !== "ai_connection_pool_wait") accounting.failureRetries = attempt;
|
||||
return accounting;
|
||||
}
|
||||
@@ -34,7 +34,8 @@ import { applyConnectorSkills, prepareConnectorSkillDelivery, resolveConnectorAs
|
||||
import { admitExplicitNativeContinuation, undeliveredLegacyUserCommentIds } from "./explicit-native-continuation.js";
|
||||
import { connectionIntentService } from "./connection-intents.js";
|
||||
import { managedAiSessionFingerprintConfig, prepareManagedAiRuntime, assertManagedAiProjectAuth, stripAiAuthBindings, isAiConnectionBusy, AI_AUTH_ENV_KEYS } from "./ai-connection-runtime.js";
|
||||
import { aiConnectionBindingSchema } from "@paperclipai/shared";
|
||||
import { aiConnectionBindingSchema, aiRuntimeConnectionBindingSchema, type AiConnectionRouterSelection } from "@paperclipai/shared";
|
||||
import { aiConnectionRouterService, AiConnectionPoolExhausted, applyAiConnectionRouterTaskSettings } from "./ai-connection-router.js";
|
||||
import { executionBlockerPredicate, getExecutionBlocker } from "./execution-blocker.js";
|
||||
import { CONVERSATION_CONTINUATION_POLICY, claimedAdapterType, runUsedConversationAdapter, hasConversationContinuationPolicy, isConversationAdapter } from "./conversation-continuation.js";
|
||||
import { recordExecutionWait } from "./execution-wait.js";
|
||||
@@ -535,6 +536,7 @@ import {
|
||||
MAX_TURN_CONTINUATION_RETRY_REASON,
|
||||
WORKSPACE_BUSY_RETRY_REASON,
|
||||
AI_CONNECTION_BUSY_RETRY_REASON,
|
||||
AI_CONNECTION_POOL_WAIT_RETRY_REASON,
|
||||
INTERACTION_CONTINUATION_INFRA_RETRY_REASON,
|
||||
INTERACTION_CONTINUATION_INFRA_WAKE_REASON,
|
||||
WAKE_COMMENT_IDS_KEY,
|
||||
@@ -15542,7 +15544,7 @@ export function heartbeatService(
|
||||
const requiresIssueGate =
|
||||
isTransientWorkspaceGitScanCode(run.errorCode) ||
|
||||
hasConversationContinuationPolicy(run.resultJson) ||
|
||||
retryReason === AI_CONNECTION_BUSY_RETRY_REASON ||
|
||||
(retryReason === AI_CONNECTION_BUSY_RETRY_REASON || retryReason === AI_CONNECTION_POOL_WAIT_RETRY_REASON) ||
|
||||
retryReason === MAX_TURN_CONTINUATION_RETRY_REASON ||
|
||||
retryReason === INTERACTION_CONTINUATION_INFRA_RETRY_REASON;
|
||||
if (requiresIssueGate) {
|
||||
@@ -15612,7 +15614,7 @@ export function heartbeatService(
|
||||
executionFailureRetryCount(run),
|
||||
}
|
||||
: {}),
|
||||
...(retryReason === AI_CONNECTION_BUSY_RETRY_REASON
|
||||
...((retryReason === AI_CONNECTION_BUSY_RETRY_REASON || retryReason === AI_CONNECTION_POOL_WAIT_RETRY_REASON)
|
||||
? { failureRetriesBeforeAiConnectionWait: executionFailureRetryCount(run) }
|
||||
: {}),
|
||||
...(shouldQuarantineWorkspaceForRetry
|
||||
@@ -15938,7 +15940,7 @@ export function heartbeatService(
|
||||
}
|
||||
|
||||
if (
|
||||
retryReason === AI_CONNECTION_BUSY_RETRY_REASON && issueId &&
|
||||
(retryReason === AI_CONNECTION_BUSY_RETRY_REASON || retryReason === AI_CONNECTION_POOL_WAIT_RETRY_REASON) && issueId &&
|
||||
!isNonAssigneeWorkspaceBusyRetry(retryReason, contextSnapshot)
|
||||
) {
|
||||
// The issue row is locked above. Recheck after the preflight gate so
|
||||
@@ -20433,7 +20435,7 @@ export function heartbeatService(
|
||||
let readFailureReportSecrets: () => string[] = () => [];
|
||||
|
||||
try {
|
||||
const agent = await getAgent(run.agentId);
|
||||
let agent = await getAgent(run.agentId);
|
||||
if (!agent) {
|
||||
await setRunStatus(runId, "failed", {
|
||||
error: "Agent not found",
|
||||
@@ -20684,7 +20686,7 @@ export function heartbeatService(
|
||||
: row;
|
||||
})
|
||||
: null;
|
||||
const issueAssigneeOverrides =
|
||||
let issueAssigneeOverrides =
|
||||
issueContext && issueContext.assigneeAgentId === agent.id
|
||||
? parseIssueAssigneeAdapterOverrides(
|
||||
issueContext.assigneeAdapterOverrides,
|
||||
@@ -20890,7 +20892,7 @@ export function heartbeatService(
|
||||
// durable for authorization and operation-time credential resolution.
|
||||
context.executionPolicy = retainedTrust.executionPolicy;
|
||||
}
|
||||
const config = parseObject(agent.adapterConfig);
|
||||
let config = parseObject(agent.adapterConfig);
|
||||
const taskSession = taskKey
|
||||
? await getTaskSession(
|
||||
agent.companyId,
|
||||
@@ -20899,6 +20901,46 @@ export function heartbeatService(
|
||||
taskKey,
|
||||
)
|
||||
: null;
|
||||
const routerHasPersistedInput = Object.keys(parseObject(parseObject(run.runnerProfileJson).nativeExecutionInput)).length > 0;
|
||||
const persistedRouterPoolId = routerHasPersistedInput ? readNonEmptyString(parseObject(context.aiRouterSelection).poolId) : null;
|
||||
const requestedAiBinding = persistedRouterPoolId ? { mode: "router" as const, connectionId: persistedRouterPoolId } : agent.runtimeConfig?.aiConnection ? aiRuntimeConnectionBindingSchema.parse(agent.runtimeConfig.aiConnection) : undefined;
|
||||
let aiBinding = requestedAiBinding?.mode === "router" ? undefined : requestedAiBinding;
|
||||
if (requestedAiBinding?.mode === "router") {
|
||||
try {
|
||||
context.aiRouterTaskKey = readNonEmptyString(context.aiRouterTaskKey) ?? taskKey ?? run.id;
|
||||
const savedIdentity = taskSession?.sessionParamsJson?.paperclipAiCredentialIdentity;
|
||||
const selection = await aiConnectionRouterService(db, options.pluginWorkerManager).resolve({
|
||||
companyId: agent.companyId, poolId: requestedAiBinding.connectionId, agentId: agent.id,
|
||||
userId: responsibleUserId, adapterType: agent.adapterType, taskKey: String(context.aiRouterTaskKey),
|
||||
overrides: issueAssigneeOverrides?.adapterConfig ?? {},
|
||||
existingGrantId: typeof savedIdentity === "string" ? savedIdentity.split(":")[0] : undefined,
|
||||
requireExisting: Boolean(taskSession?.sessionDisplayId) && !shouldResetTaskSessionForWake(context),
|
||||
persisted: routerHasPersistedInput ? context.aiRouterSelection as AiConnectionRouterSelection | undefined : undefined,
|
||||
});
|
||||
aiBinding = selection.binding;
|
||||
// Provider-specific fields from the configured harness must not leak into the selected member.
|
||||
config = applyAiConnectionRouterTaskSettings(config, selection);
|
||||
if (issueAssigneeOverrides) issueAssigneeOverrides = { ...issueAssigneeOverrides, adapterConfig: applyAiConnectionRouterTaskSettings(issueAssigneeOverrides.adapterConfig ?? {}, selection) };
|
||||
agent = { ...agent, adapterConfig: config };
|
||||
context.aiRouterSelection = selection;
|
||||
await db.update(heartbeatRuns).set({ contextSnapshot: sql`coalesce(${heartbeatRuns.contextSnapshot}, '{}'::jsonb) || ${JSON.stringify({ aiRouterTaskKey: context.aiRouterTaskKey, aiRouterSelection: selection })}::jsonb` }).where(eq(heartbeatRuns.id, run.id));
|
||||
await appendRunEvent(run, { eventType: "lifecycle", stream: "system", level: "info", message: "Using task-pinned pool account", payload: { poolId: selection.poolId, memberId: selection.memberId, provider: aiBinding.provider, model: config.model, notes: selection.notes } });
|
||||
} catch (error) {
|
||||
if (error instanceof AiConnectionPoolExhausted && !routerHasPersistedInput) {
|
||||
const now = new Date();
|
||||
context.aiConnectionBusyDeferredWhileAssignee = issueContext?.assigneeAgentId === agent.id;
|
||||
const cancelled = await setRunStatusIfRunning(run.id, "cancelled", { error: error.message, errorCode: error.code, finishedAt: now, resultJson: { executionRecovery: { kind: "ai_connection_wait", providerWorkStarted: false }, retryAt: error.retryAt }, contextSnapshot: context });
|
||||
if (cancelled.updated) {
|
||||
await setWakeupStatus(run.wakeupRequestId, "cancelled", { finishedAt: now, error: error.message });
|
||||
const retry = await scheduleBoundedRetryForRun(cancelled.run ?? run, agent, { now, retryReason: AI_CONNECTION_POOL_WAIT_RETRY_REASON, wakeReason: "ai_connection_pool_retry", maxAttempts: (run.scheduledRetryAttempt ?? 0) + 1, delayMs: Math.max(1000, Date.parse(error.retryAt) - Date.now()) });
|
||||
if (retry.outcome !== "scheduled") await releaseIssueExecutionAndPromote(cancelled.run ?? run);
|
||||
await finalizeAgentStatus(run.agentId, "cancelled");
|
||||
}
|
||||
return;
|
||||
}
|
||||
throw new ConfigurationIncompleteFailure(error instanceof Error ? error.message : "Configure this connection pool", { configurationIncomplete: { reason: "ai_connection_unavailable", companyId: agent.companyId, agentId: agent.id, responsibleUserId, actionUrl: `/agents/${agent.id}/runtime`, fingerprint: `ai-router:${requestedAiBinding.connectionId}` } });
|
||||
}
|
||||
}
|
||||
if (isConversation(issueContext)) {
|
||||
delete context.resumeSessionParams;
|
||||
delete context.resumeSessionDisplayId;
|
||||
@@ -21569,7 +21611,8 @@ export function heartbeatService(
|
||||
});
|
||||
const mergedConfig = {
|
||||
...workspaceManagedConfig,
|
||||
...(issueAssigneeOverrides?.adapterConfig ?? {}),
|
||||
...Object.fromEntries(Object.entries(issueAssigneeOverrides?.adapterConfig ?? {}).filter(([key]) => requestedAiBinding?.mode !== "router" || !["provider", "acpxAgent", "model", "modelReasoningEffort", "reasoningEffort", "effort", "variant"].includes(key))),
|
||||
...(requestedAiBinding?.mode === "router" ? parseObject(parseObject(context.aiRouterSelection).runtimeConfig) : {}),
|
||||
// The base below is already task-owned. Keep directory transport while
|
||||
// preserving isolated mode and the mandatory sandbox preflight.
|
||||
...(useIsolatedTaskDirectory ? { workspaceStrategy: { type: "project_primary" } } : {}),
|
||||
@@ -21606,7 +21649,7 @@ export function heartbeatService(
|
||||
["local", "ssh"].includes(
|
||||
selectedEnvironmentForConfig?.driver ?? "local",
|
||||
);
|
||||
const aiBinding = agent.runtimeConfig?.aiConnection ? aiConnectionBindingSchema.parse(agent.runtimeConfig.aiConnection) : undefined;
|
||||
|
||||
const { resolvedConfig, secretKeys, secretManifest } =
|
||||
await resolveExecutionRunAdapterConfig({
|
||||
managedAiCredentials: Boolean(aiBinding),
|
||||
@@ -21656,13 +21699,16 @@ export function heartbeatService(
|
||||
fingerprint: `ai:${agent.id}:${responsibleUserId}:${JSON.stringify(aiBinding)}` },
|
||||
});
|
||||
}
|
||||
if (persistedNativeExecutionInput && parseObject(run.contextSnapshot?.aiConnection).identity !== managedAiRuntime.identity) {
|
||||
const savedAiAccount = parseObject(run.contextSnapshot?.aiConnection);
|
||||
if (persistedNativeExecutionInput && (savedAiAccount.sessionIdentity
|
||||
? savedAiAccount.sessionIdentity !== managedAiRuntime.sessionIdentity
|
||||
: savedAiAccount.identity !== managedAiRuntime.identity)) {
|
||||
throw new ConfigurationIncompleteFailure("The AI account changed while this native run was suspended. Start a new execution.", { configurationIncomplete: { reason: "ai_connection_changed", actionUrl: `/agents/${agent.id}/runtime` } });
|
||||
}
|
||||
Object.assign(resolvedConfig, managedAiRuntime.config);
|
||||
for (const key of AI_AUTH_ENV_KEYS) secretKeys.add(key);
|
||||
context.aiConnection = { ...managedAiRuntime.attribution, identity: managedAiRuntime.identity };
|
||||
await db.update(heartbeatRuns).set({ contextSnapshot: sql`coalesce(${heartbeatRuns.contextSnapshot}, '{}'::jsonb) || ${JSON.stringify({ aiConnection: context.aiConnection })}::jsonb` }).where(eq(heartbeatRuns.id, run.id));
|
||||
context.aiConnection = { ...managedAiRuntime.attribution, accountName: managedAiRuntime.accountName, identity: managedAiRuntime.identity, sessionIdentity: managedAiRuntime.sessionIdentity };
|
||||
await db.update(heartbeatRuns).set({ contextSnapshot: sql`coalesce(${heartbeatRuns.contextSnapshot}, '{}'::jsonb) || ${JSON.stringify({ aiConnection: context.aiConnection, ...(context.aiRouterTaskKey ? { aiRouterTaskKey: context.aiRouterTaskKey } : {}), ...(context.aiRouterSelection ? { aiRouterSelection: context.aiRouterSelection } : {}) })}::jsonb` }).where(eq(heartbeatRuns.id, run.id));
|
||||
}
|
||||
if (secretManifest.length > 0) {
|
||||
context.paperclipSecrets = {
|
||||
@@ -23044,10 +23090,10 @@ export function heartbeatService(
|
||||
|
||||
const taskSessionCredentialCompatible = isTaskSessionCredentialCompatible(
|
||||
taskSession?.sessionParamsJson,
|
||||
managedAiRuntime?.identity,
|
||||
);
|
||||
managedAiRuntime?.sessionIdentity,
|
||||
) || Boolean(managedAiRuntime && isTaskSessionCredentialCompatible(taskSession?.sessionParamsJson, managedAiRuntime.identity));
|
||||
if (managedAiRuntime) {
|
||||
sessionConfigMetadata.aiCredentialIdentity = managedAiRuntime.identity;
|
||||
sessionConfigMetadata.aiCredentialIdentity = managedAiRuntime.sessionIdentity;
|
||||
if (!taskSessionCredentialCompatible) {
|
||||
runtimeSessionIdForAdapter = null;
|
||||
runtimeSessionParamsForAdapter = null;
|
||||
|
||||
@@ -223,6 +223,7 @@ export function normalizeExperimentalSettings(raw: unknown): InstanceExperimenta
|
||||
return {
|
||||
enableEnvironments: parsed.data.enableEnvironments ?? false,
|
||||
enableNativeRunner: parsed.data.enableNativeRunner ?? true,
|
||||
enableAiConnectionRouters: parsed.data.enableAiConnectionRouters ?? false,
|
||||
enableManagedSandboxOnly: parsed.data.enableManagedSandboxOnly ?? false,
|
||||
enableIsolatedWorkspaces: parsed.data.enableIsolatedWorkspaces ?? false,
|
||||
enableIsolatedWorkspacesByDefault: parsed.data.enableIsolatedWorkspacesByDefault ?? false,
|
||||
@@ -268,6 +269,7 @@ export function normalizeExperimentalSettings(raw: unknown): InstanceExperimenta
|
||||
return {
|
||||
enableEnvironments: false,
|
||||
enableNativeRunner: true,
|
||||
enableAiConnectionRouters: false,
|
||||
enableManagedSandboxOnly: false,
|
||||
enableIsolatedWorkspaces: false,
|
||||
enableIsolatedWorkspacesByDefault: false,
|
||||
|
||||
@@ -2397,6 +2397,8 @@ export function secretService(db: Db | DbTransaction) {
|
||||
// update matches the latest version, so a concurrent rotation between the
|
||||
// read and the write cannot pass. A mismatch throws a 409 conflict.
|
||||
expectedLatestVersion?: number;
|
||||
/** Internal authenticated refresh write-back only; never accepted by HTTP schemas. */
|
||||
preserveAiSessionEpoch?: boolean;
|
||||
},
|
||||
actor?: { userId?: string | null; agentId?: string | null },
|
||||
) {
|
||||
@@ -2534,6 +2536,7 @@ export function secretService(db: Db | DbTransaction) {
|
||||
.update(companySecrets)
|
||||
.set({
|
||||
latestVersion: nextVersion,
|
||||
aiSessionEpoch: input.preserveAiSessionEpoch ? secret.aiSessionEpoch : sql`${companySecrets.aiSessionEpoch} + 1`,
|
||||
externalRef: prepared.externalRef,
|
||||
providerConfigId,
|
||||
lastRotatedAt: new Date(),
|
||||
@@ -4589,6 +4592,8 @@ export function secretService(db: Db | DbTransaction) {
|
||||
providerVersionRef?: string | null;
|
||||
providerConfigId?: string | null;
|
||||
expectedLatestVersion?: number;
|
||||
/** Internal authenticated refresh write-back only; never accepted by HTTP schemas. */
|
||||
preserveAiSessionEpoch?: boolean;
|
||||
},
|
||||
actor?: { userId?: string | null; agentId?: string | null },
|
||||
) => {
|
||||
|
||||
Reference in new issue
Block a user