diff --git a/src/app/api/projects/[projectId]/runs/route.ts b/src/app/api/projects/[projectId]/runs/route.ts index 505ac09c4..35bd18f60 100644 --- a/src/app/api/projects/[projectId]/runs/route.ts +++ b/src/app/api/projects/[projectId]/runs/route.ts @@ -109,16 +109,20 @@ function normalizeSchedulerDrainRequest(value: unknown) { value && typeof value === "object" && !Array.isArray(value) ? (value as Record) : {}; + const researchRunId = readOptionalString(body.researchRunId); + const executionProfileId = readOptionalString(body.executionProfileId); + const threadId = readOptionalString(body.threadId); + if (!researchRunId || !executionProfileId || !threadId) { + throw new Error( + "Background run drain requires researchRunId, executionProfileId, and threadId.", + ); + } return { - ...(readOptionalString(body.threadId) ? { threadId: readOptionalString(body.threadId) } : {}), + researchRunId, + executionProfileId, + threadId, ...(readOptionalString(body.owner) ? { owner: readOptionalString(body.owner) } : {}), ...(readOptionalString(body.goal) ? { goal: readOptionalString(body.goal) } : {}), - ...(readOptionalStringArray(body.approvedToolIds) - ? { approvedToolIds: readOptionalStringArray(body.approvedToolIds) } - : {}), - ...(readOptionalStringArray(body.commandAllowlist) - ? { commandAllowlist: readOptionalStringArray(body.commandAllowlist) } - : {}), ...(readOptionalPositiveInteger(body.concurrency) ? { concurrency: readOptionalPositiveInteger(body.concurrency) } : {}), @@ -135,16 +139,6 @@ function readOptionalString(value: unknown) { return typeof value === "string" && value.trim() ? value.trim() : undefined; } -function readOptionalStringArray(value: unknown) { - if (!Array.isArray(value)) { - return undefined; - } - const items = value - .filter((item): item is string => typeof item === "string" && item.trim().length > 0) - .map((item) => item.trim()); - return items.length > 0 ? items : undefined; -} - function readOptionalPositiveInteger(value: unknown) { if (value === undefined) { return undefined; diff --git a/src/mastra/agent-controller/agent-controller.ts b/src/mastra/agent-controller/agent-controller.ts index c42c6ad02..748bf04ff 100644 --- a/src/mastra/agent-controller/agent-controller.ts +++ b/src/mastra/agent-controller/agent-controller.ts @@ -8,10 +8,8 @@ import { type ToolCategory, } from "@mastra/core/agent-controller"; import { MastraModelGateway, type MastraModelGatewayInterface } from "@mastra/core/llm"; -import { parse } from "llm-strings"; import { z } from "zod"; -import { asModelConnectionString, modelProviderFromHost } from "../../lib/models"; import { compactResearchStateSchema, emptyCompactResearchState } from "../../lib/research-terminal"; import { classifySecurityCapability, LOCAL_TOOL_IDS } from "../../lib/tools/catalog"; import { @@ -21,7 +19,11 @@ import { } from "../agents/security-research"; import { securityResearchAgent } from "../agents/security-research-agent"; import { resolveSecurityResearchControllerMemory } from "../config/memory"; -import { resolveSecurityResearchMastraModelUri } from "../config/model"; +import { + getMastraChatModel, + getMastraModelRuntimeOptions, + resolveSecurityResearchMastraModelUri, +} from "../config/model"; import { mastraStorage } from "../config/storage"; import { securityResearchWorkspace } from "../config/workspace"; import { selectSecurityResearchTools } from "../tools"; @@ -43,6 +45,14 @@ class SecurityResearchNativeModelGateway extends MastraModelGateway { gateway: this.id, npm: "@ai-sdk/anthropic", }, + profile: { + apiKeyEnvVar: "MODEL_DEFAULT", + name: "Pinned ExploitHunter profile", + models: [], + docUrl: "https://ExploitHunter.app/", + gateway: this.id, + npm: "ai", + }, }; } @@ -51,11 +61,7 @@ class SecurityResearchNativeModelGateway extends MastraModelGateway { } async getApiKey() { - const apiKey = process.env.ANTHROPIC_API_KEY?.trim(); - if (!apiKey) { - throw new Error("ANTHROPIC_API_KEY is required for native Anthropic AgentController models."); - } - return apiKey; + return process.env.ANTHROPIC_API_KEY?.trim() || "pinned-profile"; } resolveLanguageModel({ @@ -69,6 +75,13 @@ class SecurityResearchNativeModelGateway extends MastraModelGateway { apiKey: string; headers?: Record; }) { + if (providerId === "profile") { + const modelUri = Buffer.from(modelId, "base64url").toString("utf8"); + return applyPinnedControllerModelOptions( + modelUri, + getMastraChatModel(modelUri) as Record, + ) as never; + } if (providerId !== "anthropic") { throw new Error(`Unsupported native AgentController provider: ${providerId}`); } @@ -76,6 +89,30 @@ class SecurityResearchNativeModelGateway extends MastraModelGateway { } } +export function applyPinnedControllerModelOptions( + modelUri: string, + model: Record, +) { + const runtime = getMastraModelRuntimeOptions(modelUri); + return new Proxy(model, { + get(target, property, receiver) { + const value = Reflect.get(target, property, receiver); + if ((property === "doGenerate" || property === "doStream") && typeof value === "function") { + return (options: Record) => + value.call(target, { + ...options, + ...runtime.modelSettings, + providerOptions: { + ...((options.providerOptions as Record | undefined) ?? {}), + ...runtime.providerOptions, + }, + }); + } + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + export const securityResearchControllerGateways: MastraModelGatewayInterface[] = [ new SecurityResearchNativeModelGateway(), ]; @@ -119,16 +156,7 @@ const DEFAULT_SECURITY_RESEARCH_MODEL_URI = resolveSecurityResearchMastraModelUr ); export function toMastraGatewayModelId(modelUri: string) { - const parsed = parse(asModelConnectionString(modelUri)); - const provider = - modelProviderFromHost(parsed.hostAlias) ?? - modelProviderFromHost(parsed.host) ?? - parsed.hostAlias ?? - parsed.host; - if (provider === "anthropic" && parsed.model === "claude-opus-5") { - return `${SECURITY_RESEARCH_NATIVE_GATEWAY_ID}/${provider}/${parsed.model}`; - } - return provider ? `${provider}/${parsed.model}` : parsed.model; + return `${SECURITY_RESEARCH_NATIVE_GATEWAY_ID}/profile/${Buffer.from(modelUri).toString("base64url")}`; } const DEFAULT_SECURITY_RESEARCH_MODEL = toMastraGatewayModelId(DEFAULT_SECURITY_RESEARCH_MODEL_URI); @@ -206,8 +234,8 @@ export const securityResearchControllerSubagents: AgentControllerSubagent[] = id: stage, name: securityResearchControllerSubagentStageNames[stage], description: securityResearchControllerSubagentDescriptions[stage], - instructions: async () => - String(await securityResearchStageAgents[agentKey].getInstructions()), + instructions: async ({ requestContext }) => + String(await securityResearchStageAgents[agentKey].getInstructions({ requestContext })), allowedControllerTools: [...getSecurityResearchStageToolIds(stage)], defaultModelId: DEFAULT_SECURITY_RESEARCH_MODEL, maxSteps: stage === "hunt" || stage === "trace" ? 24 : 12, diff --git a/src/mastra/agents/security-research/stage-agents.ts b/src/mastra/agents/security-research/stage-agents.ts index b001c93ae..044b2a73f 100644 --- a/src/mastra/agents/security-research/stage-agents.ts +++ b/src/mastra/agents/security-research/stage-agents.ts @@ -288,11 +288,16 @@ function stageModel(definition: StageDefinition) { async function stageInstructions( definition: StageDefinition, capabilities: readonly string[] = [], + pinnedSkills: readonly { id: string; revision: string; detail?: string }[] = [], + pinnedProfile = false, ) { - const skillInstructions = await formatSecurityResearchSkillInstructions( - definition.skillHints, - { capabilities }, - ); + const skillInstructions = pinnedProfile + ? pinnedSkills + .filter((skill) => definition.skillHints.includes(skill.id)) + .filter((skill) => skill.detail) + .map((skill) => `## ${skill.id} [${skill.revision}]\n${skill.detail}`) + .join("\n\n") + : await formatSecurityResearchSkillInstructions(definition.skillHints, { capabilities }); return `${commonStageInstructions(definition.stage)} Stage role: @@ -320,6 +325,8 @@ function createStageAgent(definition: StageDefinition) { stageInstructions( definition, readRuntimeSkillCapabilities(requestContext.get("runtimeSkillCapabilities")), + readPinnedSkillRefs(requestContext.get("selectedSkillRefs")), + typeof requestContext.get("researchExecutionProfileId") === "string", ), model: stageModel(definition), memory: createSecurityResearchStageMemory( @@ -345,6 +352,23 @@ function createStageAgent(definition: StageDefinition) { }); } +function readPinnedSkillRefs(value: unknown) { + if (!Array.isArray(value)) return []; + return value.flatMap((item) => { + if (!item || typeof item !== "object" || Array.isArray(item)) return []; + const skill = item as Record; + return typeof skill.id === "string" && typeof skill.revision === "string" + ? [ + { + id: skill.id, + revision: skill.revision, + ...(typeof skill.detail === "string" ? { detail: skill.detail } : {}), + }, + ] + : []; + }); +} + function readRuntimeSkillCapabilities(value: unknown) { if (!Array.isArray(value)) return []; return value.filter((item): item is string => typeof item === "string"); diff --git a/src/mastra/config/skill-search.ts b/src/mastra/config/skill-search.ts index ef012c30b..0be14ad6b 100644 --- a/src/mastra/config/skill-search.ts +++ b/src/mastra/config/skill-search.ts @@ -1,39 +1,34 @@ import { - BaseProcessor, - type ProcessInputStepArgs, - type ProcessInputStepResult, + BaseProcessor, + type ProcessInputStepArgs, + type ProcessInputStepResult, } from "@mastra/core/processors"; import { createTool } from "@mastra/core/tools"; import { z } from "zod"; import { readSecurityResearchRuntimeToolProfile } from "../../lib/security-chat/runtime-tool-profile"; import { - createProductSkillRegistry, - formatProductSkillDirectory, - type ProductSkillRegistry, + createProductSkillRegistry, + formatProductSkillDirectory, + type ProductSkillRegistry, } from "../../server/skills/product-skill-registry"; -import { - SECURITY_RESEARCH_SKILLS_ROOT, - securityResearchWorkspace, -} from "./workspace"; +import { SECURITY_RESEARCH_SKILLS_ROOT, securityResearchWorkspace } from "./workspace"; -const readRequestContext = ( - requestContext: ProcessInputStepArgs["requestContext"], - key: string, -) => requestContext?.get?.(key); +const readRequestContext = (requestContext: ProcessInputStepArgs["requestContext"], key: string) => + requestContext?.get?.(key); /** A foreground-command-only run deliberately exposes no overlapping discovery tools. */ export const shouldUseSecurityResearchSkillSearch = ( - requestContext: ProcessInputStepArgs["requestContext"], + requestContext: ProcessInputStepArgs["requestContext"], ) => - readSecurityResearchRuntimeToolProfile( - readRequestContext(requestContext, "runtimeToolProfile"), - ) !== "foreground-command-only"; + readSecurityResearchRuntimeToolProfile( + readRequestContext(requestContext, "runtimeToolProfile"), + ) !== "foreground-command-only"; type LoadedSkill = { revision: string; instructions: string }; type ThreadSkillState = { - snapshotKey: string; - loaded: Map; + snapshotKey: string; + loaded: Map; }; /** @@ -41,190 +36,208 @@ type ThreadSkillState = { * both re-read request-scoped capabilities and revalidate through Workspace. */ export class RegistryBackedSkillSearchProcessor extends BaseProcessor<"security-research-skill-search"> { - readonly id = "security-research-skill-search" as const; - readonly name = "Security Research Skill Search"; - readonly providesSkillDiscovery = "on-demand" as const; - - private readonly threadState = new Map(); - - constructor(private readonly registry: ProductSkillRegistry) { - super(); - } - - async processInputStep( - args: ProcessInputStepArgs, - ): Promise { - if (!shouldUseSecurityResearchSkillSearch(args.requestContext)) { - return { tools: args.tools }; - } - - const capabilitiesForRequest = () => - readRuntimeSkillCapabilities( - readRequestContext(args.requestContext, "runtimeSkillCapabilities"), - ); - const capabilities = capabilitiesForRequest(); - const snapshot = await this.registry.list({ capabilities }); - const snapshotKey = [ - snapshot.status, - snapshot.revision ?? "incomplete", - capabilityKey(capabilities), - ].join(":"); - const threadId = readThreadId(args); - let state = this.threadState.get(threadId); - if (!state || state.snapshotKey !== snapshotKey) { - state = { snapshotKey, loaded: new Map() }; - this.threadState.set(threadId, state); - } - - args.messageList.addSystem( - `Product skill directory (Workspace-owned). Full bodies are available only through on-demand loading.\n${formatProductSkillDirectory(snapshot)}`, - ); - args.messageList.addSystem( - "To discover an applicable product procedure, call search_skills. Load only the procedure needed for the current task with load_skill.", - ); - for (const [skillName, skill] of state.loaded) { - args.messageList.addSystem( - `[Skill: ${skillName}; revision: ${skill.revision}]\n\n${skill.instructions}`, - ); - } - - const searchSkills = createTool({ - id: "search_skills", - description: - "Search reviewed product research skills applicable to the current runtime capabilities.", - inputSchema: z.object({ query: z.string() }), - outputSchema: z.object({ - results: z.array( - z.object({ - name: z.string(), - description: z.string(), - revision: z.string(), - }), - ), - status: z.enum(["complete", "incomplete"]), - message: z.string(), - }), - execute: async ({ query }) => { - const normalizedQuery = query.trim(); - if (!normalizedQuery) { - return { - results: [], - status: "complete" as const, - message: "A non-empty skill search query is required.", - }; - } - const current = await this.registry.list({ - query: normalizedQuery, - limit: 5, - capabilities: capabilitiesForRequest(), - }); - return { - results: current.skills.map((skill) => ({ - name: skill.id, - description: skill.description.slice(0, 150), - revision: skill.revision, - })), - status: current.status, - message: - current.status === "incomplete" - ? `Skill discovery is incomplete: ${current.diagnostics.join("; ")}` - : `Found ${current.skills.length} applicable skill(s).`, - }; - }, - }); - - const loadSkill = createTool({ - id: "load_skill", - description: - "Load one reviewed product skill after revalidating current visibility and capabilities.", - inputSchema: z.object({ skillName: z.string() }), - outputSchema: z.object({ - success: z.boolean(), - status: z.enum(["complete", "incomplete"]), - message: z.string(), - skillName: z.string().optional(), - revision: z.string().optional(), - }), - execute: async ({ skillName }) => { - const normalizedName = skillName.trim(); - const current = await this.registry.list({ - query: normalizedName, - includeDetails: true, - limit: 100, - capabilities: capabilitiesForRequest(), - }); - const skill = current.skills.find( - (entry) => entry.id === normalizedName, - ); - if (current.status !== "complete" || !skill) { - state.loaded.delete(normalizedName); - return { - success: false, - status: current.status, - message: - current.status === "incomplete" - ? `Skill discovery is incomplete: ${current.diagnostics.join("; ")}` - : `Skill "${normalizedName}" is unavailable or not applicable to the current runtime capabilities.`, - }; - } - - state.loaded.set(skill.id, { - revision: skill.revision, - instructions: skill.detail, - }); - return { - success: true, - status: "complete" as const, - message: `Skill "${skill.id}" loaded for the current snapshot.`, - skillName: skill.id, - revision: skill.revision, - }; - }, - }); - - return { - tools: { - ...(args.tools ?? {}), - search_skills: searchSkills, - load_skill: loadSkill, - }, - }; - } + readonly id = "security-research-skill-search" as const; + readonly name = "Security Research Skill Search"; + readonly providesSkillDiscovery = "on-demand" as const; + + private readonly threadState = new Map(); + + constructor(private readonly registry: ProductSkillRegistry) { + super(); + } + + async processInputStep(args: ProcessInputStepArgs): Promise { + if (!shouldUseSecurityResearchSkillSearch(args.requestContext)) { + return { tools: args.tools }; + } + const pinnedSkills = readPinnedSkills( + readRequestContext(args.requestContext, "selectedSkillRefs"), + ); + const pinnedProfileId = readRequestContext(args.requestContext, "researchExecutionProfileId"); + if (typeof pinnedProfileId === "string") { + args.messageList.addSystem( + `Pinned product skill snapshot (${pinnedProfileId}):\n${ + pinnedSkills + .map((skill) => `[Skill: ${skill.id}; revision: ${skill.revision}]\n\n${skill.detail}`) + .join("\n\n") || "No product skills were selected for this profile." + }`, + ); + return { tools: args.tools }; + } + + const capabilitiesForRequest = () => + readRuntimeSkillCapabilities( + readRequestContext(args.requestContext, "runtimeSkillCapabilities"), + ); + const capabilities = capabilitiesForRequest(); + const snapshot = await this.registry.list({ capabilities }); + const snapshotKey = [ + snapshot.status, + snapshot.revision ?? "incomplete", + capabilityKey(capabilities), + ].join(":"); + const threadId = readThreadId(args); + let state = this.threadState.get(threadId); + if (!state || state.snapshotKey !== snapshotKey) { + state = { snapshotKey, loaded: new Map() }; + this.threadState.set(threadId, state); + } + + args.messageList.addSystem( + `Product skill directory (Workspace-owned). Full bodies are available only through on-demand loading.\n${formatProductSkillDirectory(snapshot)}`, + ); + args.messageList.addSystem( + "To discover an applicable product procedure, call search_skills. Load only the procedure needed for the current task with load_skill.", + ); + for (const [skillName, skill] of state.loaded) { + args.messageList.addSystem( + `[Skill: ${skillName}; revision: ${skill.revision}]\n\n${skill.instructions}`, + ); + } + + const searchSkills = createTool({ + id: "search_skills", + description: + "Search reviewed product research skills applicable to the current runtime capabilities.", + inputSchema: z.object({ query: z.string() }), + outputSchema: z.object({ + results: z.array( + z.object({ + name: z.string(), + description: z.string(), + revision: z.string(), + }), + ), + status: z.enum(["complete", "incomplete"]), + message: z.string(), + }), + execute: async ({ query }) => { + const normalizedQuery = query.trim(); + if (!normalizedQuery) { + return { + results: [], + status: "complete" as const, + message: "A non-empty skill search query is required.", + }; + } + const current = await this.registry.list({ + query: normalizedQuery, + limit: 5, + capabilities: capabilitiesForRequest(), + }); + return { + results: current.skills.map((skill) => ({ + name: skill.id, + description: skill.description.slice(0, 150), + revision: skill.revision, + })), + status: current.status, + message: + current.status === "incomplete" + ? `Skill discovery is incomplete: ${current.diagnostics.join("; ")}` + : `Found ${current.skills.length} applicable skill(s).`, + }; + }, + }); + + const loadSkill = createTool({ + id: "load_skill", + description: + "Load one reviewed product skill after revalidating current visibility and capabilities.", + inputSchema: z.object({ skillName: z.string() }), + outputSchema: z.object({ + success: z.boolean(), + status: z.enum(["complete", "incomplete"]), + message: z.string(), + skillName: z.string().optional(), + revision: z.string().optional(), + }), + execute: async ({ skillName }) => { + const normalizedName = skillName.trim(); + const current = await this.registry.list({ + query: normalizedName, + includeDetails: true, + limit: 100, + capabilities: capabilitiesForRequest(), + }); + const skill = current.skills.find((entry) => entry.id === normalizedName); + if (current.status !== "complete" || !skill) { + state.loaded.delete(normalizedName); + return { + success: false, + status: current.status, + message: + current.status === "incomplete" + ? `Skill discovery is incomplete: ${current.diagnostics.join("; ")}` + : `Skill "${normalizedName}" is unavailable or not applicable to the current runtime capabilities.`, + }; + } + + state.loaded.set(skill.id, { + revision: skill.revision, + instructions: skill.detail, + }); + return { + success: true, + status: "complete" as const, + message: `Skill "${skill.id}" loaded for the current snapshot.`, + skillName: skill.id, + revision: skill.revision, + }; + }, + }); + + return { + tools: { + ...(args.tools ?? {}), + search_skills: searchSkills, + load_skill: loadSkill, + }, + }; + } +} + +function readPinnedSkills(value: unknown) { + if (!Array.isArray(value)) return []; + return value.flatMap((item) => { + if (!item || typeof item !== "object" || Array.isArray(item)) return []; + const skill = item as Record; + return typeof skill.id === "string" && + typeof skill.revision === "string" && + typeof skill.detail === "string" + ? [{ id: skill.id, revision: skill.revision, detail: skill.detail }] + : []; + }); } function readRuntimeSkillCapabilities(value: unknown) { - if (!Array.isArray(value)) return []; - return [ - ...new Set( - value.filter((item): item is string => typeof item === "string"), - ), - ].sort(); + if (!Array.isArray(value)) return []; + return [...new Set(value.filter((item): item is string => typeof item === "string"))].sort(); } function capabilityKey(capabilities: readonly string[]) { - return capabilities.join(","); + return capabilities.join(","); } function readThreadId(args: ProcessInputStepArgs) { - for (const key of ["mastra__threadId", "threadId", "workspaceThreadId"]) { - const value = readRequestContext(args.requestContext, key); - if (typeof value === "string" && value.trim()) return value; - } - return "default"; + for (const key of ["mastra__threadId", "threadId", "workspaceThreadId"]) { + const value = readRequestContext(args.requestContext, key); + if (typeof value === "string" && value.trim()) return value; + } + return "default"; } const workspaceSkills = securityResearchWorkspace.skills; if (!workspaceSkills) { - throw new Error("Security research Workspace must configure product skills."); + throw new Error("Security research Workspace must configure product skills."); } const productSkillRegistry = createProductSkillRegistry( - workspaceSkills, - SECURITY_RESEARCH_SKILLS_ROOT, + workspaceSkills, + SECURITY_RESEARCH_SKILLS_ROOT, ); -export const securityResearchSkillSearchProcessor = - new RegistryBackedSkillSearchProcessor(productSkillRegistry); +export const securityResearchSkillSearchProcessor = new RegistryBackedSkillSearchProcessor( + productSkillRegistry, +); -export const securityResearchSkillInputProcessors = [ - securityResearchSkillSearchProcessor, -]; +export const securityResearchSkillInputProcessors = [securityResearchSkillSearchProcessor]; diff --git a/src/mastra/workflows/scheduler-backed-research.ts b/src/mastra/workflows/scheduler-backed-research.ts index 33a92917b..57e90b81d 100644 --- a/src/mastra/workflows/scheduler-backed-research.ts +++ b/src/mastra/workflows/scheduler-backed-research.ts @@ -4,6 +4,13 @@ import { buildSecurityResearchStageRequestContext, buildSecurityResearchTaskWorkspaceContext, } from "../../server/chat/security-research-runtime-context"; +import type { ResearchExecutionBackgroundProfile } from "../../server/research/execution-profile"; +import { + persistStageHandoff, + readStageHandoff, + type StageHandoff, + type StageHandoffDraft, +} from "../../server/research/stage-handoff"; import { getDefaultSchedulerService, type SchedulerService, @@ -27,12 +34,6 @@ import { } from "../agents/security-research"; import { getSecurityResearchConcurrencyConfig } from "../config/concurrency"; import type { SecurityResearchStageId } from "../config/memory"; -import { - persistStageHandoff, - readStageHandoff, - type StageHandoff, - type StageHandoffDraft, -} from "../../server/research/stage-handoff"; export const SCHEDULER_RESEARCH_LOCK_TTL_MS = 5 * 60 * 1000; const SCHEDULER_RESEARCH_MAX_HEARTBEAT_INTERVAL_MS = 30 * 1000; @@ -51,6 +52,8 @@ export type SchedulerResearchTask = { export type SchedulerResearchEnqueueInput = { projectId: string; + researchRunId?: string; + executionProfileId?: string; threadId?: string; goal: string; tasks: SchedulerResearchTask[]; @@ -64,6 +67,8 @@ export type SchedulerResearchEnqueueInput = { export type SchedulerResearchRunInput = { projectId: string; + researchRunId?: string; + executionProfileId?: string; threadId?: string; owner: string; concurrency: number; @@ -72,6 +77,7 @@ export type SchedulerResearchRunInput = { lockTtlMs?: number; goal?: string; isolateTaskWorkspaces?: boolean; + executionProfile?: ResearchExecutionBackgroundProfile; }; export type SchedulerResearchRunResult = { @@ -111,6 +117,8 @@ export type SchedulerResearchWorkerLoopSummary = { export type SchedulerResearchExecuteTaskContext = { projectId: string; + researchRunId?: string; + executionProfileId?: string; threadId?: string; stage?: SecurityResearchStageId; approvedToolIds: string[]; @@ -119,6 +127,7 @@ export type SchedulerResearchExecuteTaskContext = { isolateTaskWorkspaces: boolean; taskWorkspaceId?: string; parentHandoffs: StageHandoff[]; + executionProfile?: ResearchExecutionBackgroundProfile; }; export type SchedulerResearchExecuteTaskInput = { @@ -155,16 +164,6 @@ export type SchedulerResearchWorkerLoopDeps = { sleep?: (ms: number, signal?: AbortSignal) => Promise; }; -function firstParagraph(value: string) { - return ( - value - .split(/\n\s*\n/) - .map((part) => part.trim()) - .find(Boolean) - ?.slice(0, 600) ?? "" - ); -} - function readAgentText(response: unknown) { if ( response && @@ -237,6 +236,8 @@ function payloadForTask(task: SchedulerTask) { agentId?: string; shellSessionId?: string; parentHandoffIds?: string[]; + researchRunId?: string; + executionProfileId?: string; budgetUsage?: StageHandoff["budgetUsage"]; resourceClaims?: StageHandoff["resourceClaims"]; }; @@ -304,9 +305,14 @@ async function runStageAgentForTask(input: SchedulerResearchExecuteTaskInput): P schedulerTaskId: task.id, targetIds: task.targetIds, approvalClass, - enabledToolIds: context.approvedToolIds, + enabledToolIds: context.executionProfile + ? context.approvedToolIds.filter((id) => + context.executionProfile?.effective.capabilityIds.includes(id), + ) + : context.approvedToolIds, commandAllowPatterns: context.commandAllowlist, commandDefaultAction: context.commandAllowlist.length > 0 ? "approval" : "deny", + ...(context.executionProfile ? { executionProfile: context.executionProfile } : {}), ...(workspace ? { workspace } : {}), }); const response = await agentForStage(task.stage).generate(buildStagePrompt(task, context.goal), { @@ -353,6 +359,8 @@ export function createSchedulerBackedResearchRunner( expectedEvidence: task.expectedEvidence, source: "harness-seed", approvalClass: task.active ? "active" : "passive", + ...(input.researchRunId ? { researchRunId: input.researchRunId } : {}), + ...(input.executionProfileId ? { executionProfileId: input.executionProfileId } : {}), ...(input.isolateTaskWorkspaces ? { workspaceIsolation: "task", @@ -380,12 +388,15 @@ export function createSchedulerBackedResearchRunner( const concurrency = clampSchedulerResearchConcurrency(input.concurrency); const context: SchedulerResearchExecuteTaskContext = { projectId: input.projectId, + ...(input.researchRunId ? { researchRunId: input.researchRunId } : {}), + ...(input.executionProfileId ? { executionProfileId: input.executionProfileId } : {}), ...(input.threadId ? { threadId: input.threadId } : {}), approvedToolIds: input.approvedToolIds ?? [], commandAllowlist: input.commandAllowlist ?? [], goal: input.goal ?? "", isolateTaskWorkspaces: input.isolateTaskWorkspaces ?? false, parentHandoffs: [], + ...(input.executionProfile ? { executionProfile: input.executionProfile } : {}), }; while (true) { @@ -673,6 +684,14 @@ async function runOneTask( ...(isolateTaskWorkspaces ? { taskWorkspaceId: taskWorkspaceIdFor(task) } : {}), parentHandoffs: await loadParentHandoffs(task, readHandoff), }; + const taskPayload = payloadForTask(task); + if ( + context.executionProfile && + (taskPayload.researchRunId !== context.researchRunId || + taskPayload.executionProfileId !== context.executionProfile.profileId) + ) { + throw new Error(`Scheduler task ${task.id} is not bound to this Research Execution Profile.`); + } output = await executeTask({ task, context: taskContext }); } catch (error) { const message = error instanceof Error ? error.message : String(error); @@ -700,17 +719,14 @@ async function runOneTask( const persisted = await persistHandoff({ projectId: task.projectId, ...(task.threadId ? { threadId: task.threadId } : {}), + ...(context.researchRunId ? { researchRunId: context.researchRunId } : {}), + ...(context.executionProfileId ? { executionProfileId: context.executionProfileId } : {}), schedulerTaskId: task.id, stage: task.stage, taskId: task.taskId ?? task.id, attempt: task.attempt, producerAgentId: payload.agentId ?? task.owner ?? `stage:${task.stage}`, - executionProfile: { - approvedToolIds: context.approvedToolIds, - commandAllowlist: context.commandAllowlist, - isolateTaskWorkspaces: context.isolateTaskWorkspaces, - targetIds: task.targetIds, - }, + executionProfile: context.executionProfile ?? {}, draft: { ...draft, attribution: { diff --git a/src/server/chat/mastra-background-tasks.ts b/src/server/chat/mastra-background-tasks.ts index 20beed018..a40e04920 100644 --- a/src/server/chat/mastra-background-tasks.ts +++ b/src/server/chat/mastra-background-tasks.ts @@ -8,6 +8,7 @@ import { type SchedulerBackedResearchRunner, type SchedulerResearchRunInput, } from "../../mastra/workflows"; +import { readResearchExecutionBackgroundProfile } from "../research/execution-profile"; export const SCHEDULER_DRAIN_BACKGROUND_TOOL_NAME = "schedulerDrainTool"; @@ -15,6 +16,9 @@ export type StartSchedulerDrainBackgroundTaskInput = Omit< SchedulerResearchRunInput, "owner" | "concurrency" > & { + researchRunId: string; + executionProfileId: string; + threadId: string; owner?: string; concurrency?: number; timeoutMs?: number; @@ -31,6 +35,7 @@ export type StartSchedulerDrainBackgroundTaskDeps = { runner?: SchedulerBackedResearchRunner; runId?: () => string; toolCallId?: () => string; + resolveExecutionProfile?: typeof readResearchExecutionBackgroundProfile; }; export async function listProjectMastraBackgroundTasks( @@ -59,7 +64,24 @@ export async function startSchedulerDrainBackgroundTask( throw new Error("Mastra background tasks are not enabled."); } const runner = deps.runner ?? createSchedulerBackedResearchRunner(); - const args = normalizeSchedulerDrainArgs(input); + const inheritedProfile = await ( + deps.resolveExecutionProfile ?? readResearchExecutionBackgroundProfile + )({ + projectId: input.projectId, + threadId: input.threadId, + researchRunId: input.researchRunId, + profileId: input.executionProfileId, + }); + if (inheritedProfile.profileId !== input.executionProfileId) { + throw new Error("Background task execution profile identity mismatch."); + } + if (input.executionProfile && input.executionProfile.profileId !== inheritedProfile.profileId) { + throw new Error("Background task supplied profile does not match its Research Run."); + } + const args = normalizeSchedulerDrainArgs({ + ...input, + ...(inheritedProfile ? { executionProfile: inheritedProfile } : {}), + }); const backgroundTask = createBackgroundTask(manager, { runId: deps.runId?.() ?? createId("activity"), toolName: SCHEDULER_DRAIN_BACKGROUND_TOOL_NAME, @@ -75,6 +97,8 @@ export async function startSchedulerDrainBackgroundTask( execute: async () => runner.runEnqueuedTasks({ projectId: args.projectId, + researchRunId: args.researchRunId, + executionProfileId: args.executionProfileId, ...(args.threadId ? { threadId: args.threadId } : {}), owner: args.owner, concurrency: args.concurrency, @@ -83,6 +107,7 @@ export async function startSchedulerDrainBackgroundTask( lockTtlMs: args.lockTtlMs, goal: args.goal, isolateTaskWorkspaces: args.isolateTaskWorkspaces, + ...(args.executionProfile ? { executionProfile: args.executionProfile } : {}), }), }, }, @@ -96,6 +121,8 @@ function normalizeSchedulerDrainArgs( ): SchedulerResearchRunInput { return { projectId: input.projectId, + researchRunId: input.researchRunId, + executionProfileId: input.executionProfileId, ...(input.threadId ? { threadId: input.threadId } : {}), owner: input.owner ?? "mastra-background-scheduler-drain", concurrency: input.concurrency ?? 1, @@ -104,5 +131,6 @@ function normalizeSchedulerDrainArgs( lockTtlMs: input.lockTtlMs ?? SCHEDULER_RESEARCH_LOCK_TTL_MS, goal: input.goal ?? "", isolateTaskWorkspaces: input.isolateTaskWorkspaces ?? false, + ...(input.executionProfile ? { executionProfile: input.executionProfile } : {}), }; } diff --git a/src/server/chat/security-research-run.ts b/src/server/chat/security-research-run.ts index 40f1183a5..6d8b5e4cb 100644 --- a/src/server/chat/security-research-run.ts +++ b/src/server/chat/security-research-run.ts @@ -60,6 +60,7 @@ import { SECURITY_RESEARCH_CONTROLLER_ID, toMastraGatewayModelId, } from "../../mastra/agent-controller/agent-controller"; +import { SECURITY_RESEARCH_STAGE_IDS } from "../../mastra/agents/security-research"; import { securityResearchLangfusePromptReference } from "../../mastra/config/langfuse-prompt"; import { createLocalOpenAiUsageCapture, @@ -90,8 +91,8 @@ import { loadControllerToolInputRepeatGuardSnapshot, saveControllerToolInputRepeatGuardSnapshot, } from "../research/turn-ledger"; -import { readThreadTargetConfig, type ThreadTargetConfig } from "../workspaces/target-mode"; import { consumeTargetAuthorization } from "../targets"; +import { readThreadTargetConfig, type ThreadTargetConfig } from "../workspaces/target-mode"; import { type AgentControllerBackgroundLifecycleEvent, continueAgentControllerAfterBackgroundTasks, @@ -735,6 +736,7 @@ async function runSecurityResearchAgentController({ ...(runtimeMaxSteps ? { runtimeMaxSteps } : {}), ...(runtimeMaxToolCalls ? { runtimeMaxToolCalls } : {}), }); + await configurePinnedControllerSubagentModels(session, requestContext); const activityEvents: ChatActivityEvent[] = []; const labCommandEvents: LabCommandStreamEvent[] = []; const targetAuthorizationEvents: TargetAuthorizationStreamEvent[] = []; @@ -1027,12 +1029,16 @@ async function runSecurityResearchAgentController({ const repeatGuardSnapshot = researchRunId ? await loadControllerToolInputRepeatGuardSnapshot(researchRunId) : undefined; - const toolInputRepeatGuard = new ControllerToolInputRepeatGuard(maxRepeatedToolCalls, { - // Schema-invalid calls never reach tool_start/tool_end, so the durable - // terminal repair protocol cannot count them itself. Bound all generated - // terminal inputs to the same initial-attempt-plus-repairs allowance. - finish_research: (terminalRepairLimit ?? 2) + 1, - }, repeatGuardSnapshot); + const toolInputRepeatGuard = new ControllerToolInputRepeatGuard( + maxRepeatedToolCalls, + { + // Schema-invalid calls never reach tool_start/tool_end, so the durable + // terminal repair protocol cannot count them itself. Bound all generated + // terminal inputs to the same initial-attempt-plus-repairs allowance. + finish_research: (terminalRepairLimit ?? 2) + 1, + }, + repeatGuardSnapshot, + ); let repeatGuardPersistence = Promise.resolve(); const persistRepeatGuard = () => { if (!researchRunId) return; @@ -1631,8 +1637,7 @@ async function runSecurityResearchAgentController({ requestContext, tracingOptions, }), - continueAfterBackgroundTasks: () => - continueAfterBackgroundTasks(truncationRecoveryStartedAt), + continueAfterBackgroundTasks: () => continueAfterBackgroundTasks(truncationRecoveryStartedAt), repairMissingTerminal: repairMissingTerminalTransitions, }); }; @@ -1750,10 +1755,7 @@ async function runSecurityResearchAgentController({ return; } await continueAfterBackgroundTasks(controllerActionStartedAt); - if ( - controllerOutputTruncated && - !outputTruncationTracker.attempted - ) { + if (controllerOutputTruncated && !outputTruncationTracker.attempted) { await recoverControllerOutputTruncation(); } if ( @@ -2207,6 +2209,36 @@ async function runSecurityResearchAgentController({ }; } +export async function configurePinnedControllerSubagentModels( + session: ControllerSessionLike, + requestContext: unknown, +) { + const modelSelection = session.subagents?.model; + if (!modelSelection?.set) return; + const modelOverrides = readRuntimeModelOverrides( + readRequestContextValue(requestContext, securityResearchRuntimeContextKeys.modelOverrides), + ); + const coordinatorModelUri = readRequestContextValue( + requestContext, + securityResearchRuntimeContextKeys.coordinatorModelUri, + ); + for (const stage of SECURITY_RESEARCH_STAGE_IDS) { + const modelUri = isModelOverrideTarget(stage) + ? readModelOverride(modelOverrides, stage) + : typeof coordinatorModelUri === "string" + ? coordinatorModelUri + : undefined; + if (!modelUri) continue; + await modelSelection.set({ agentType: stage, modelId: toMastraGatewayModelId(modelUri) }); + } +} + +function readRequestContextValue(requestContext: unknown, key: string) { + if (!requestContext || typeof requestContext !== "object") return undefined; + const get = (requestContext as { get?: (name: string) => unknown }).get; + return typeof get === "function" ? get.call(requestContext, key) : undefined; +} + /** * Count only model-issued tool executions. Observational-memory compression and * reflection events are ancillary lifecycle diagnostics, not candidate tool @@ -2545,7 +2577,8 @@ function attachDurableApprovalToRequestContext(requestContext: unknown, durableA set?: (key: string, value: unknown) => void; }; if (typeof candidate.set !== "function") return; - const current = typeof candidate.get === "function" ? candidate.get("chatPermissions") : undefined; + const current = + typeof candidate.get === "function" ? candidate.get("chatPermissions") : undefined; const permissions = current && typeof current === "object" && !Array.isArray(current) ? (current as Record) @@ -3722,6 +3755,13 @@ export const runSecurityResearchAgent = async ( ? { evaluationArm: input.metadata.evaluationArm } : {}), terminalProtocolMode, + researchExecutionProfileId: readStringValue(input.metadata?.researchExecutionProfileId), + capabilityManifestRevision: readStringValue(input.metadata?.capabilityManifestRevision), + skillRegistryRevision: readStringValue(input.metadata?.skillRegistryRevision), + selectedSkillRefs: Array.isArray(input.metadata?.selectedSkillRefs) + ? (input.metadata.selectedSkillRefs as never[]) + : undefined, + runtimeSkillCapabilities: enabledToolIds, ...(input.metadata?.terminalCleanupRequiredBeforeFinish === true ? { terminalCleanupRequiredBeforeFinish: true } : {}), @@ -3838,9 +3878,7 @@ export const runSecurityResearchAgent = async ( effectiveToolIds: effectiveEnabledToolIds, infrastructureCost, researchRunId: - typeof input.metadata?.researchRunId === "string" - ? input.metadata.researchRunId - : undefined, + typeof input.metadata?.researchRunId === "string" ? input.metadata.researchRunId : undefined, }); // Paid eval gates query Langfuse while this production process remains // alive. Flush the completed candidate trace at that boundary so exporter @@ -4062,10 +4100,7 @@ function stableSerializeForFingerprint(value: unknown, seen = new WeakSet 0) { return Math.min( MAX_AGENT_CONTROLLER_TIMEOUT_MS, diff --git a/src/server/chat/security-research-runtime-context.ts b/src/server/chat/security-research-runtime-context.ts index a97515a41..8fb1c8ee3 100644 --- a/src/server/chat/security-research-runtime-context.ts +++ b/src/server/chat/security-research-runtime-context.ts @@ -8,6 +8,7 @@ import { } from "../../lib/security-chat/chat-approvals"; import type { SecurityResearchRuntimeToolProfile } from "../../lib/security-chat/runtime-tool-profile"; import type { ProjectWorkspaceConfig } from "../../lib/workspace-config"; +import type { ResearchExecutionBackgroundProfile } from "../research/execution-profile"; import type { ContextCompressionMode, ProjectApprovalMode } from "./types"; type RequestContextEntry = readonly [string, unknown]; @@ -77,6 +78,15 @@ export const securityResearchRuntimeContextKeys = { evaluationArm: "evaluationArm", terminalProtocolMode: "terminalProtocolMode", terminalCleanupRequiredBeforeFinish: "terminalCleanupRequiredBeforeFinish", + researchExecutionProfileId: "researchExecutionProfileId", + capabilityManifestRevision: "capabilityManifestRevision", + skillRegistryRevision: "skillRegistryRevision", + selectedSkillRefs: "selectedSkillRefs", + runtimeSkillCapabilities: "runtimeSkillCapabilities", + researchExecutionProfile: "researchExecutionProfile", + networkProfile: "networkProfile", + renderJsonCollection: "renderJsonCollection", + policyVersions: "policyVersions", } as const; export function createSecurityResearchRequestContext( @@ -127,6 +137,11 @@ export function buildSecurityResearchChatRequestContext(input: { evaluationArm?: string; terminalProtocolMode?: "legacy" | "typed" | "gated"; terminalCleanupRequiredBeforeFinish?: boolean; + researchExecutionProfileId?: string; + capabilityManifestRevision?: string; + skillRegistryRevision?: string; + selectedSkillRefs?: readonly { id: string; revision: string; detail?: string }[]; + runtimeSkillCapabilities?: readonly string[]; }) { const keys = securityResearchRuntimeContextKeys; const resourceId = buildSecurityResearchMemoryResourceId( @@ -168,6 +183,11 @@ export function buildSecurityResearchChatRequestContext(input: { [keys.evaluationArm, input.evaluationArm], [keys.terminalProtocolMode, input.terminalProtocolMode], [keys.terminalCleanupRequiredBeforeFinish, input.terminalCleanupRequiredBeforeFinish], + [keys.researchExecutionProfileId, input.researchExecutionProfileId], + [keys.capabilityManifestRevision, input.capabilityManifestRevision], + [keys.skillRegistryRevision, input.skillRegistryRevision], + [keys.selectedSkillRefs, input.selectedSkillRefs], + [keys.runtimeSkillCapabilities, input.runtimeSkillCapabilities], [keys.workspaceConfig, input.workspaceConfig], [keys.commandAllowPatterns, input.commandAllowPatterns], [keys.commandBlockPatterns, input.commandBlockPatterns], @@ -241,6 +261,7 @@ export function buildSecurityResearchStageRequestContext(input: { enabledToolIds: readonly string[]; commandAllowPatterns: readonly string[]; commandDefaultAction: "approval" | "deny" | string; + executionProfile?: ResearchExecutionBackgroundProfile; }) { const keys = securityResearchRuntimeContextKeys; const baseEntries = input.baseContext ? Array.from(input.baseContext.entries()) : []; @@ -257,6 +278,31 @@ export function buildSecurityResearchStageRequestContext(input: { [keys.enabledToolIds, input.enabledToolIds], [keys.commandAllowPatterns, input.commandAllowPatterns], [keys.commandDefaultAction, input.commandDefaultAction], + [keys.researchExecutionProfileId, input.executionProfile?.profileId], + [keys.capabilityManifestRevision, input.executionProfile?.effective.capabilityManifestRevision], + [keys.skillRegistryRevision, input.executionProfile?.effective.skillRegistryRevision], + [keys.selectedSkillRefs, input.executionProfile?.effective.selectedSkills], + [keys.runtimeSkillCapabilities, input.executionProfile?.effective.capabilityIds], + [keys.modelUri, input.executionProfile?.effective.models.coordinator.modelUri], + [keys.coordinatorModelUri, input.executionProfile?.effective.models.coordinator.modelUri], + [ + keys.modelOverrides, + input.executionProfile + ? Object.fromEntries( + Object.entries(input.executionProfile.effective.models).map(([stage, model]) => [ + stage, + model.modelUri, + ]), + ) + : undefined, + ], + [keys.targetMode, input.executionProfile?.effective.targetMode], + [keys.networkProfile, input.executionProfile?.effective.networkPolicy], + [keys.renderJsonCollection, input.executionProfile?.effective.uiCollection], + [keys.policyVersions, input.executionProfile?.effective.policies], + [keys.researchExecutionProfile, input.executionProfile], + [keys.runtimeMaxToolCalls, input.executionProfile?.effective.budgets.maxToolCalls], + [keys.autonomyMode, input.executionProfile?.effective.budgets], ]); } diff --git a/src/server/chat/securityResearchTurn.ts b/src/server/chat/securityResearchTurn.ts index 2b5854b69..9e12f3935 100644 --- a/src/server/chat/securityResearchTurn.ts +++ b/src/server/chat/securityResearchTurn.ts @@ -11,6 +11,7 @@ import { import { deriveThreadName } from "../../lib/security-chat/naming"; import { getSecurityCapability, + SECURITY_CAPABILITY_MANIFEST, SECURITY_CAPABILITY_MANIFEST_REVISION, sanitizeRuntimeToolIds, } from "../../lib/tools/catalog"; @@ -24,9 +25,15 @@ import { import { createArtifactService } from "../evidence/artifact-service"; import { runCompletedTurnPassivePolicyShadow } from "../policy/passive-policy-shadow-runtime"; import { isPlaceholderThreadTitle } from "../projects/naming"; +import { + projectResearchExecutionProfileToRuntime, + readResearchExecutionProfileForRun, + researchExecutionProfileForensics, + resolveResearchExecutionProfile, +} from "../research/execution-profile"; import { recoverResearchTurn, shouldRecoverResearchResult } from "../research/recovery-runtime"; import { withResearchRunContext } from "../research/run-context"; -import { getSecurityResearchSkillRevision } from "../research/skill-revision"; +import { getSecurityResearchSkillSnapshot } from "../research/skill-revision"; import { beginResearchTurnLedger, finalizeResearchTurnLedger, @@ -210,6 +217,71 @@ export const runSecurityResearchTurn = async ( agentInput.content, readOptionalString(agentInput.metadata?.researchRunId) ?? undefined, ); + const requestedEnvelope = requestedModelEnvelope(agentInput.metadata); + const resumedExecutionProfile = resumeResearchRunId + ? await readResearchExecutionProfileForRun({ + projectId, + threadId: message.threadId, + researchRunId: resumeResearchRunId, + }) + : undefined; + const configuredCapabilityIds = resumedExecutionProfile + ? resumedExecutionProfile.effective.capabilityIds + : sanitizeRuntimeToolIds(agentInput.metadata?.enabledToolIds); + const skillSnapshot = resumedExecutionProfile + ? { + revision: resumedExecutionProfile.effective.skillRegistryRevision, + skills: resumedExecutionProfile.effective.selectedSkills, + } + : await getSecurityResearchSkillSnapshot(configuredCapabilityIds); + const skillRegistryRevision = skillSnapshot.revision; + const targetConfig = readTargetConfigForExecutionProfile(agentInput.metadata); + const executionProfile = + resumedExecutionProfile ?? + resolveResearchExecutionProfile( + { + modelUri: readOptionalString(agentInput.metadata?.modelUri) ?? undefined, + runtimeModelUri: readOptionalString(agentInput.metadata?.runtimeModelUri) ?? undefined, + modelOverrides: readModelOverrides(agentInput.metadata?.modelOverrides), + contextWindowTokens: + readOptionalNumber(agentInput.metadata?.contextWindowTokens) ?? undefined, + maxOutputTokens: readOptionalNumber(agentInput.metadata?.maxOutputTokens) ?? undefined, + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: configuredCapabilityIds, + skillRegistryRevision, + selectedSkills: skillSnapshot.skills, + targetMode: targetConfig.targetMode, + computeTargetId: + targetConfig.targetMode === "remote" ? targetConfig.computeTargetId : undefined, + networkPolicy: + readOptionalString(agentInput.metadata?.networkProfile) ?? + (targetConfig.targetMode === "none" ? "none" : "approved-targets"), + maxToolCalls: readOptionalNumber(agentInput.metadata?.runtimeMaxToolCalls) ?? undefined, + maxRuntimeMs: isRecord(agentInput.metadata?.autonomyMode) + ? (readOptionalNumber(agentInput.metadata.autonomyMode.maxRuntimeMs) ?? undefined) + : undefined, + maxTurns: isRecord(agentInput.metadata?.autonomyMode) + ? (readOptionalNumber(agentInput.metadata.autonomyMode.maxTurns) ?? undefined) + : undefined, + maxCostUsd: isRecord(agentInput.metadata?.autonomyMode) + ? (readOptionalNumber(agentInput.metadata.autonomyMode.maxCostUsd) ?? undefined) + : undefined, + uiCollection: readExecutionProfileUiCollection(agentInput.metadata?.renderJsonCollection), + policies: { + approval: "durable-intent-v1", + terminal: + readOptionalString(agentInput.metadata?.terminalProtocolMode) ?? "gated-terminal-v1", + runtimeToolProfile: + readOptionalString(agentInput.metadata?.runtimeToolProfile) ?? "default-v1", + }, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: SECURITY_CAPABILITY_MANIFEST.map((capability) => capability.id), + skillRegistryRevision, + selectedSkills: skillSnapshot.skills, + }, + ); const ledger = await beginResearchTurnLedger({ projectId, threadId: message.threadId, @@ -219,15 +291,16 @@ export const runSecurityResearchTurn = async ( ...extractArtifactSourceIds(agentInput.content), ...extractPassivePolicySourceIds(agentInput.metadata), ], - requestedModelEnvelope: requestedModelEnvelope(agentInput.metadata), - capabilityIds: sanitizeRuntimeToolIds(agentInput.metadata?.enabledToolIds), - capabilitySchemaRefs: sanitizeRuntimeToolIds(agentInput.metadata?.enabledToolIds).map((id) => ({ + requestedModelEnvelope: requestedEnvelope, + capabilityIds: configuredCapabilityIds, + capabilitySchemaRefs: configuredCapabilityIds.map((id) => ({ capabilityId: id, manifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, schemaRevision: getSecurityCapability(id)?.schemaRevision ?? null, schemaAvailability: getSecurityCapability(id) ? "manifest-revision" : "unavailable", })), - skillRevision: await getSecurityResearchSkillRevision(), + skillRevision: skillRegistryRevision, + executionProfile, metadata: { autonomyBudgetEnvelope: isRecord(agentInput.metadata?.autonomyMode) ? { @@ -264,10 +337,14 @@ export const runSecurityResearchTurn = async ( { source: "workspace-skill-projection", reason: "not-yet-addressed-by-durable-source-id" }, ], }); + const pinnedExecutionProfile = ledger.executionProfile ?? executionProfile; + const pinnedRuntime = projectResearchExecutionProfileToRuntime(pinnedExecutionProfile); const ledgerAgentInput = { ...agentInput, metadata: { ...(agentInput.metadata ?? {}), + ...pinnedRuntime, + researchExecutionProfile: researchExecutionProfileForensics(pinnedExecutionProfile), researchRunId: ledger.researchRunId, researchTurnLedgerId: ledger.turnLedgerId, capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, @@ -818,6 +895,28 @@ function readOptionalNumber(value: unknown) { return typeof value === "number" && Number.isFinite(value) ? value : null; } +function readTargetConfigForExecutionProfile(metadata: Record | undefined) { + const workspace = metadata?.threadWorkspace; + if (isRecord(workspace) && workspace.available === true) { + return readThreadTargetConfig(workspace); + } + return readThreadTargetConfig(metadata); +} + +function readExecutionProfileUiCollection(value: unknown) { + switch (value) { + case "onboarding": + case "planning": + case "evidence": + case "approvals": + case "blockers": + case "all": + return value; + default: + return "all"; + } +} + function toJsonObject(value: Record) { return JSON.parse(JSON.stringify(value)) as Record; } diff --git a/src/server/research/execution-profile.ts b/src/server/research/execution-profile.ts new file mode 100644 index 000000000..d9d55dccf --- /dev/null +++ b/src/server/research/execution-profile.ts @@ -0,0 +1,517 @@ +import { createHash } from "node:crypto"; + +import { normalize, parse } from "llm-strings"; + +import { + asModelConnectionString, + findModelEntryForRef, + MODEL_OVERRIDE_TARGETS, + type ModelOverrideMap, + type ModelOverrideTarget, + modelProviderFromHost, + readModelOverride, +} from "../../lib/models"; +import type { SecurityResearchComponentCollection } from "../../lib/render-json/security-research-catalog"; +import { + getMastraModelRuntimeOptions, + resolveSecurityResearchMastraModelUri, +} from "../../mastra/config/model"; +import { withDatabase } from "../db/client"; +import type { JsonObject } from "../db/types"; + +export const RESEARCH_EXECUTION_PROFILE_SCHEMA_VERSION = "research-execution-profile-v1"; +export const RESEARCH_EXECUTION_PROFILE_METADATA_KEY = "executionProfile"; + +export type ResearchExecutionTargetMode = "none" | "container" | "remote"; + +export type RequestedResearchExecutionProfile = { + modelUri?: string; + runtimeModelUri?: string; + modelOverrides?: ModelOverrideMap; + contextWindowTokens?: number; + maxOutputTokens?: number; + capabilityManifestRevision: string; + capabilityIds: readonly string[]; + skillRegistryRevision: string; + selectedSkills?: readonly { id: string; revision: string; detail?: string }[]; + targetMode: ResearchExecutionTargetMode; + computeTargetId?: string; + networkPolicy: string; + maxToolCalls?: number; + maxRuntimeMs?: number; + maxTurns?: number; + maxCostUsd?: number; + uiCollection: SecurityResearchComponentCollection; + policies: Readonly>; +}; + +export type ResearchExecutionModel = { + modelUri: string; + provider: string; + model: string; + contextWindowTokens: number | null; + maxOutputTokens: number | null; + modelSettings: JsonObject; + providerOptions: JsonObject; +}; + +export type ResearchExecutionProfile = { + schemaVersion: typeof RESEARCH_EXECUTION_PROFILE_SCHEMA_VERSION; + profileId: string; + requested: JsonObject; + effective: { + models: Record; + capabilityManifestRevision: string; + capabilityIds: string[]; + skillRegistryRevision: string; + selectedSkills: { id: string; revision: string; detail?: string }[]; + targetMode: ResearchExecutionTargetMode; + computeTargetId: string | null; + networkPolicy: string; + budgets: { + maxToolCalls: number | null; + maxRuntimeMs: number | null; + maxTurns: number | null; + maxCostUsd: number | null; + }; + uiCollection: SecurityResearchComponentCollection; + policies: Record; + }; +}; + +export type ResearchExecutionProfileRuntimeProjection = { + modelUri: string; + modelOverrides: ModelOverrideMap; + contextWindowTokens?: number; + maxOutputTokens?: number; + enabledToolIds: string[]; + runtimeMaxToolCalls?: number; + autonomyMode?: { + maxRuntimeMs?: number; + maxTurns?: number; + maxCostUsd?: number; + }; + targetMode: ResearchExecutionTargetMode; + computeTargetId?: string; + networkProfile: string; + renderJsonCollection: SecurityResearchComponentCollection; + capabilityManifestRevision: string; + skillRegistryRevision: string; + selectedSkillRefs: { id: string; revision: string; detail?: string }[]; + researchExecutionProfileId: string; +}; + +/** Complete authority-free projection used by durable stage/background work. */ +export type ResearchExecutionBackgroundProfile = ResearchExecutionProfile; + +export function resolveResearchExecutionProfile( + input: RequestedResearchExecutionProfile, + current: { + capabilityManifestRevision: string; + capabilityIds: readonly string[]; + skillRegistryRevision: string; + selectedSkills: readonly { id: string; revision: string; detail?: string }[]; + }, +): ResearchExecutionProfile { + requireCurrentRevision( + "capability manifest", + input.capabilityManifestRevision, + current.capabilityManifestRevision, + ); + requireCurrentRevision( + "skill registry", + input.skillRegistryRevision, + current.skillRegistryRevision, + ); + const selectedSkills = canonicalSelectedSkills(input.selectedSkills ?? []); + if (Buffer.byteLength(JSON.stringify(selectedSkills), "utf8") > 256 * 1024) { + throw new Error("Research Execution Profile selected skill payload exceeds 256 KiB."); + } + const currentCapabilities = new Set(uniqueStrings(current.capabilityIds)); + for (const capabilityId of uniqueStrings(input.capabilityIds)) { + if (!currentCapabilities.has(capabilityId)) { + throw new Error( + `Research Execution Profile capability ${capabilityId} is unavailable in manifest ${current.capabilityManifestRevision}.`, + ); + } + } + const currentSkills = new Map( + canonicalSelectedSkills(current.selectedSkills).map((skill) => [skill.id, skill.revision]), + ); + for (const skill of selectedSkills) { + if (!skill.revision.trim()) { + throw new Error(`Research Execution Profile skill revision is missing for ${skill.id}.`); + } + const currentRevision = currentSkills.get(skill.id); + if (currentRevision !== skill.revision) { + throw new Error( + `Research Execution Profile skill ${skill.id} revision ${skill.revision} is unavailable; current revision is ${currentRevision ?? "missing"}.`, + ); + } + } + + const requestedModelUri = cleanString(input.modelUri); + const runtimeModelUri = cleanString(input.runtimeModelUri); + const requestedOverrides = canonicalModelOverrides(input.modelOverrides); + const coordinatorRequest = + runtimeModelUri ?? readModelOverride(requestedOverrides, "coordinator") ?? requestedModelUri; + const coordinatorUri = resolveSecurityResearchMastraModelUri(coordinatorRequest); + const models = Object.fromEntries( + MODEL_OVERRIDE_TARGETS.map((target) => { + const targetRequest = runtimeModelUri + ? runtimeModelUri + : target === "coordinator" + ? coordinatorRequest + : (readModelOverride(requestedOverrides, target) ?? coordinatorUri); + return [ + target, + resolveExecutionModel(targetRequest, { + contextWindowTokens: target === "coordinator" ? input.contextWindowTokens : undefined, + maxOutputTokens: target === "coordinator" ? input.maxOutputTokens : undefined, + }), + ]; + }), + ) as Record; + + const requested = toJsonObject({ + modelUri: requestedModelUri ?? null, + runtimeModelUri: runtimeModelUri ?? null, + modelOverrides: requestedOverrides, + contextWindowTokens: positiveIntegerOrNull(input.contextWindowTokens), + maxOutputTokens: positiveIntegerOrNull(input.maxOutputTokens), + capabilityManifestRevision: input.capabilityManifestRevision, + capabilityIds: uniqueStrings(input.capabilityIds), + skillRegistryRevision: input.skillRegistryRevision, + selectedSkills, + targetMode: input.targetMode, + computeTargetId: cleanString(input.computeTargetId) ?? null, + networkPolicy: input.networkPolicy, + maxToolCalls: positiveIntegerOrNull(input.maxToolCalls), + maxRuntimeMs: positiveIntegerOrNull(input.maxRuntimeMs), + maxTurns: positiveIntegerOrNull(input.maxTurns), + maxCostUsd: nonNegativeNumberOrNull(input.maxCostUsd), + uiCollection: input.uiCollection, + policies: canonicalStringRecord(input.policies), + }); + const body = { + schemaVersion: RESEARCH_EXECUTION_PROFILE_SCHEMA_VERSION, + requested, + effective: { + models, + capabilityManifestRevision: current.capabilityManifestRevision, + capabilityIds: uniqueStrings(input.capabilityIds), + skillRegistryRevision: current.skillRegistryRevision, + selectedSkills, + targetMode: input.targetMode, + computeTargetId: cleanString(input.computeTargetId) ?? null, + networkPolicy: input.networkPolicy, + budgets: { + maxToolCalls: positiveIntegerOrNull(input.maxToolCalls), + maxRuntimeMs: positiveIntegerOrNull(input.maxRuntimeMs), + maxTurns: positiveIntegerOrNull(input.maxTurns), + maxCostUsd: nonNegativeNumberOrNull(input.maxCostUsd), + }, + uiCollection: input.uiCollection, + policies: canonicalStringRecord(input.policies), + }, + } as const; + return deepFreeze({ + ...body, + profileId: `research-profile-sha256:${sha256(stableStringify(body))}`, + }); +} + +/** Projection safe for stage/background inheritance; deliberately excludes approval grants. */ +export function projectResearchExecutionProfileToBackground( + profile: ResearchExecutionProfile, +): ResearchExecutionBackgroundProfile { + assertResearchExecutionProfile(profile); + return readResearchExecutionProfile(JSON.parse(JSON.stringify(profile))); +} + +export async function readResearchExecutionBackgroundProfile(input: { + projectId: string; + threadId: string; + researchRunId: string; + profileId: string; +}) { + return withDatabase(async (db) => { + const result = await db.query<{ metadata: JsonObject }>( + `SELECT metadata FROM research_runs + WHERE project_id = $1 AND thread_id = $2 AND id = $3`, + [input.projectId, input.threadId, input.researchRunId], + ); + const stored = result.rows[0]?.metadata?.[RESEARCH_EXECUTION_PROFILE_METADATA_KEY]; + if (stored === undefined) { + throw new Error(`Research Run ${input.researchRunId} has no pinned execution profile.`); + } + const profile = readResearchExecutionProfile(stored); + if (profile.profileId !== input.profileId) { + throw new Error( + `Research Run ${input.researchRunId} does not match execution profile ${input.profileId}.`, + ); + } + return projectResearchExecutionProfileToBackground(profile); + }); +} + +export async function readResearchExecutionProfileForRun(input: { + projectId: string; + threadId: string; + researchRunId: string; +}) { + return withDatabase(async (db) => { + const row = ( + await db.query<{ metadata: JsonObject }>( + `SELECT metadata FROM research_runs WHERE project_id = $1 AND thread_id = $2 AND id = $3`, + [input.projectId, input.threadId, input.researchRunId], + ) + ).rows[0]; + const stored = row?.metadata?.[RESEARCH_EXECUTION_PROFILE_METADATA_KEY]; + if (stored === undefined) { + throw new Error(`Research Run ${input.researchRunId} has no pinned execution profile.`); + } + return readResearchExecutionProfile(stored); + }); +} + +export function projectResearchExecutionProfileToRuntime( + profile: ResearchExecutionProfile, +): ResearchExecutionProfileRuntimeProjection { + assertResearchExecutionProfile(profile); + const coordinator = profile.effective.models.coordinator; + const modelOverrides = Object.fromEntries( + MODEL_OVERRIDE_TARGETS.filter((target) => target !== "coordinator").map((target) => [ + target, + profile.effective.models[target].modelUri, + ]), + ) as ModelOverrideMap; + const budgets = profile.effective.budgets; + return { + modelUri: coordinator.modelUri, + modelOverrides, + ...(coordinator.contextWindowTokens + ? { contextWindowTokens: coordinator.contextWindowTokens } + : {}), + ...(coordinator.maxOutputTokens ? { maxOutputTokens: coordinator.maxOutputTokens } : {}), + enabledToolIds: [...profile.effective.capabilityIds], + ...(budgets.maxToolCalls ? { runtimeMaxToolCalls: budgets.maxToolCalls } : {}), + ...(budgets.maxRuntimeMs || budgets.maxTurns || budgets.maxCostUsd !== null + ? { + autonomyMode: { + ...(budgets.maxRuntimeMs ? { maxRuntimeMs: budgets.maxRuntimeMs } : {}), + ...(budgets.maxTurns ? { maxTurns: budgets.maxTurns } : {}), + ...(budgets.maxCostUsd !== null ? { maxCostUsd: budgets.maxCostUsd } : {}), + }, + } + : {}), + targetMode: profile.effective.targetMode, + ...(profile.effective.computeTargetId + ? { computeTargetId: profile.effective.computeTargetId } + : {}), + networkProfile: profile.effective.networkPolicy, + renderJsonCollection: profile.effective.uiCollection, + capabilityManifestRevision: profile.effective.capabilityManifestRevision, + skillRegistryRevision: profile.effective.skillRegistryRevision, + selectedSkillRefs: profile.effective.selectedSkills.map((skill) => ({ ...skill })), + researchExecutionProfileId: profile.profileId, + }; +} + +export function researchExecutionProfileForensics(profile: ResearchExecutionProfile): JsonObject { + assertResearchExecutionProfile(profile); + return toJsonObject(profile); +} + +export function readResearchExecutionProfile(value: unknown): ResearchExecutionProfile { + assertResearchExecutionProfile(value); + return deepFreeze(value as ResearchExecutionProfile); +} + +export function assertResearchExecutionProfile( + value: unknown, +): asserts value is ResearchExecutionProfile { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("Research Execution Profile is missing or malformed."); + } + const candidate = value as Partial; + if ( + candidate.schemaVersion !== RESEARCH_EXECUTION_PROFILE_SCHEMA_VERSION || + typeof candidate.profileId !== "string" || + !candidate.requested || + !candidate.effective + ) { + throw new Error("Research Execution Profile revision is missing or unsupported."); + } + const body = { + schemaVersion: candidate.schemaVersion, + requested: candidate.requested, + effective: candidate.effective, + }; + const expectedId = `research-profile-sha256:${sha256(stableStringify(body))}`; + if (candidate.profileId !== expectedId) { + throw new Error("Research Execution Profile content does not match its immutable profile id."); + } + const models = candidate.effective.models; + if (!models || MODEL_OVERRIDE_TARGETS.some((target) => !models[target]?.modelUri)) { + throw new Error("Research Execution Profile is missing an effective model revision."); + } + requireNonEmpty("capability manifest", candidate.effective.capabilityManifestRevision); + requireNonEmpty("skill registry", candidate.effective.skillRegistryRevision); +} + +function resolveExecutionModel( + requestedModelUri: string | undefined, + overrides: { contextWindowTokens?: number; maxOutputTokens?: number }, +): ResearchExecutionModel { + const modelUri = withModelEnvelope( + resolveSecurityResearchMastraModelUri(requestedModelUri), + overrides, + ); + const parsed = parse(asModelConnectionString(modelUri)); + const normalized = normalize(parsed); + const registryEntry = findModelEntryForRef(modelUri); + const runtimeOptions = getMastraModelRuntimeOptions(modelUri); + const provider = + normalized.provider ?? + modelProviderFromHost(parsed.hostAlias) ?? + modelProviderFromHost(parsed.host) ?? + parsed.hostAlias ?? + parsed.host; + const supportedProviders = new Set([ + "openrouter", + "ollama", + "lmstudio", + "vllm", + "anthropic", + "moonshotai", + "deepseek", + "alibaba", + ]); + if (!supportedProviders.has(provider) && !parsed.params.baseUrl && !parsed.params.base_url) { + throw new Error( + `Research Execution Profile model provider ${provider} is unsupported without an explicit baseUrl route.`, + ); + } + const normalizedMaxOutput = positiveIntegerOrNull(normalized.config.params.max_tokens); + return { + modelUri, + provider, + model: parsed.model, + contextWindowTokens: + positiveIntegerOrNull(overrides.contextWindowTokens) ?? registryEntry?.contextWindow ?? null, + maxOutputTokens: + positiveIntegerOrNull(overrides.maxOutputTokens) ?? + normalizedMaxOutput ?? + registryEntry?.maxOutputTokens ?? + null, + modelSettings: toJsonObject(runtimeOptions.modelSettings), + providerOptions: toJsonObject(runtimeOptions.providerOptions), + }; +} + +function withModelEnvelope( + modelUri: string, + overrides: { contextWindowTokens?: number; maxOutputTokens?: number }, +) { + const url = new URL(modelUri); + if (positiveIntegerOrNull(overrides.contextWindowTokens)) { + url.searchParams.set("contextWindow", String(overrides.contextWindowTokens)); + } + if (positiveIntegerOrNull(overrides.maxOutputTokens)) { + url.searchParams.set("maxTokens", String(overrides.maxOutputTokens)); + } + return url.toString(); +} + +function canonicalModelOverrides(value: ModelOverrideMap | undefined): ModelOverrideMap { + return Object.fromEntries( + MODEL_OVERRIDE_TARGETS.flatMap((target) => { + const modelUri = readModelOverride(value, target); + return modelUri ? [[target, modelUri]] : []; + }), + ); +} + +function canonicalSelectedSkills( + value: readonly { id: string; revision: string; detail?: string }[], +) { + return [...value] + .map((skill) => ({ + id: skill.id.trim(), + revision: skill.revision.trim(), + ...(skill.detail?.trim() ? { detail: skill.detail.trim() } : {}), + })) + .filter((skill) => skill.id) + .sort((left, right) => left.id.localeCompare(right.id)); +} + +function canonicalStringRecord(value: Readonly>) { + return Object.fromEntries( + Object.entries(value) + .map(([key, item]) => [key.trim(), item.trim()] as const) + .filter(([key, item]) => key && item) + .sort(([left], [right]) => left.localeCompare(right)), + ); +} + +function uniqueStrings(values: readonly string[]) { + return [...new Set(values.map((value) => value.trim()).filter(Boolean))].sort(); +} + +function requireCurrentRevision(label: string, requested: string, current: string) { + requireNonEmpty(label, requested); + requireNonEmpty(label, current); + if (requested !== current) { + throw new Error( + `Research Execution Profile ${label} revision ${requested} is unavailable; current revision is ${current}.`, + ); + } +} + +function requireNonEmpty(label: string, value: unknown): asserts value is string { + if (typeof value !== "string" || !value.trim()) { + throw new Error(`Research Execution Profile ${label} revision is missing.`); + } +} + +function cleanString(value: unknown) { + return typeof value === "string" && value.trim() ? value.trim() : undefined; +} + +function positiveIntegerOrNull(value: unknown) { + return typeof value === "number" && Number.isSafeInteger(value) && value > 0 ? value : null; +} + +function nonNegativeNumberOrNull(value: unknown) { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : null; +} + +function toJsonObject(value: unknown): JsonObject { + return JSON.parse( + JSON.stringify(value, (_key, item) => (item === undefined ? null : item)), + ) as JsonObject; +} + +function stableStringify(value: unknown): string { + if (Array.isArray(value)) return `[${value.map(stableStringify).join(",")}]`; + if (value && typeof value === "object") { + return `{${Object.entries(value as Record) + .sort(([left], [right]) => left.localeCompare(right)) + .map(([key, item]) => `${JSON.stringify(key)}:${stableStringify(item)}`) + .join(",")}}`; + } + return JSON.stringify(value); +} + +function sha256(value: string) { + return createHash("sha256").update(value).digest("hex"); +} + +function deepFreeze(value: T): T { + if (value && typeof value === "object" && !Object.isFrozen(value)) { + Object.freeze(value); + for (const nested of Object.values(value as Record)) deepFreeze(nested); + } + return value; +} diff --git a/src/server/research/skill-revision.ts b/src/server/research/skill-revision.ts index f1b922127..eb7b62776 100644 --- a/src/server/research/skill-revision.ts +++ b/src/server/research/skill-revision.ts @@ -1,31 +1,50 @@ import { - SECURITY_RESEARCH_SKILLS_ROOT, - securityResearchWorkspace, + SECURITY_RESEARCH_SKILLS_ROOT, + securityResearchWorkspace, } from "../../mastra/config/workspace"; import { - createProductSkillRegistry, - type ProductSkillRegistry, + createProductSkillRegistry, + type ProductSkillRegistry, } from "../skills/product-skill-registry"; const workspaceSkills = securityResearchWorkspace.skills; if (!workspaceSkills) { - throw new Error("Security research Workspace must configure product skills."); + throw new Error("Security research Workspace must configure product skills."); } const executionProfileSkillRegistry = createProductSkillRegistry( - workspaceSkills, - SECURITY_RESEARCH_SKILLS_ROOT, + workspaceSkills, + SECURITY_RESEARCH_SKILLS_ROOT, ); /** Resolve the current reviewed Workspace snapshot; never pin an empty or stale catalog. */ export async function getSecurityResearchSkillRevision( - registry: ProductSkillRegistry = executionProfileSkillRegistry, + registry: ProductSkillRegistry = executionProfileSkillRegistry, ) { - const snapshot = await registry.list(); - if (snapshot.status !== "complete" || !snapshot.revision) { - throw new Error( - `Product skill discovery is incomplete: ${snapshot.diagnostics.join("; ") || "unknown discovery failure"}`, - ); - } - return snapshot.revision; + return (await getSecurityResearchSkillSnapshot([], registry)).revision; +} + +/** Resolve the exact reviewed skill references visible to one runtime capability set. */ +export async function getSecurityResearchSkillSnapshot( + capabilities: readonly string[] = [], + registry: ProductSkillRegistry = executionProfileSkillRegistry, +) { + const snapshot = await registry.list({ + capabilities, + includeDetails: true, + includeFullDetails: true, + }); + if (snapshot.status !== "complete" || !snapshot.revision) { + throw new Error( + `Product skill discovery is incomplete: ${snapshot.diagnostics.join("; ") || "unknown discovery failure"}`, + ); + } + return { + revision: snapshot.revision, + skills: snapshot.skills.map((skill) => ({ + id: skill.id, + revision: skill.revision, + detail: skill.detail, + })), + }; } diff --git a/src/server/research/stage-handoff.ts b/src/server/research/stage-handoff.ts index 8a9f804df..e00f34872 100644 --- a/src/server/research/stage-handoff.ts +++ b/src/server/research/stage-handoff.ts @@ -1,5 +1,5 @@ import { type Queryable, withDatabase } from "../db/client"; -import { getArtifactService, type ArtifactServiceInstance } from "../evidence"; +import { type ArtifactServiceInstance, getArtifactService } from "../evidence"; export type StageHandoffOutcome = "completed" | "partial" | "blocked" | "failed" | "cancelled"; @@ -7,6 +7,8 @@ export type StageHandoff = { schemaVersion: 1; projectId: string; threadId?: string; + researchRunId?: string; + executionProfileId?: string; schedulerTaskId: string; stage: string; taskId: string; @@ -40,13 +42,22 @@ export type StageHandoff = { export type StageHandoffDraft = Partial< Omit< StageHandoff, - "schemaVersion" | "projectId" | "schedulerTaskId" | "stage" | "taskId" | "attempt" | "createdAt" + | "schemaVersion" + | "projectId" + | "schedulerTaskId" + | "stage" + | "taskId" + | "attempt" + | "createdAt" + | "executionProfile" > > & { output: string }; type PersistInput = { projectId: string; threadId?: string; + researchRunId?: string; + executionProfileId?: string; schedulerTaskId: string; stage: string; taskId: string; @@ -68,6 +79,8 @@ export async function persistStageHandoff( schemaVersion: 1, projectId: input.projectId, ...(input.threadId ? { threadId: input.threadId } : {}), + ...(input.researchRunId ? { researchRunId: input.researchRunId } : {}), + ...(input.executionProfileId ? { executionProfileId: input.executionProfileId } : {}), schedulerTaskId: input.schedulerTaskId, stage: input.stage, taskId: input.taskId, @@ -85,7 +98,7 @@ export async function persistStageHandoff( blockers: input.draft.blockers ?? [], proposedTaskIds: input.draft.proposedTaskIds ?? [], approvalNeeds: input.draft.approvalNeeds ?? [], - executionProfile: input.draft.executionProfile ?? input.executionProfile, + executionProfile: input.executionProfile, attribution: { producerAgentId: input.draft.attribution?.producerAgentId ?? input.producerAgentId, parentHandoffIds: input.draft.attribution?.parentHandoffIds ?? [], diff --git a/src/server/research/turn-ledger.ts b/src/server/research/turn-ledger.ts index 4586f1461..ea8702d16 100644 --- a/src/server/research/turn-ledger.ts +++ b/src/server/research/turn-ledger.ts @@ -3,6 +3,12 @@ import { SECURITY_CAPABILITY_MANIFEST_REVISION } from "../../lib/tools/catalog"; import { mergeJsonObject, withDatabase, withTransaction } from "../db/client"; import type { JsonObject, JsonValue, ResearchRunRow, ResearchTurnLedgerRow } from "../db/types"; import type { ControllerToolInputRepeatGuardSnapshot } from "../chat/controller-tool-input-repeat-guard"; +import { + readResearchExecutionProfile, + RESEARCH_EXECUTION_PROFILE_METADATA_KEY, + researchExecutionProfileForensics, + type ResearchExecutionProfile, +} from "./execution-profile"; export const RESEARCH_TURN_LEDGER_VERSION = "research-turn-ledger-v1"; export const UNAVAILABLE_SKILL_REVISION = "unavailable:not-versioned"; @@ -22,12 +28,14 @@ export type BeginResearchTurnLedgerInput = { skillRefs?: JsonValue[]; metadata?: JsonObject; resumeResearchRunId?: string; + executionProfile?: ResearchExecutionProfile; }; export type ResearchTurnLedgerHandle = { researchRunId: string; turnLedgerId: string; startedAt: string; + executionProfile?: ResearchExecutionProfile; }; const CONTROLLER_REPEAT_GUARD_METADATA_KEY = "controllerToolInputRepeatGuard"; @@ -109,6 +117,7 @@ export async function beginResearchTurnLedger( const startedAt = new Date().toISOString(); const userMessageSourceId = `message:${input.userMessageId}`; const sourceIds = uniqueIds(input.modelVisibleSourceIds ?? [userMessageSourceId]); + let executionProfile = input.executionProfile; if (!sourceIds.includes(userMessageSourceId)) { throw new Error( "Research Turn Ledger requires the durable user message as a model-visible source.", @@ -119,8 +128,8 @@ export async function beginResearchTurnLedger( withTransaction(db, async (tx) => { await validateModelVisibleSources(tx, input.projectId, input.threadId, sourceIds); const resumable = input.resumeResearchRunId - ? await tx.query<{ id: string }>( - `SELECT id FROM research_runs + ? await tx.query<{ id: string; metadata: JsonObject }>( + `SELECT id, metadata FROM research_runs WHERE id = $1 AND project_id = $2 AND thread_id = $3 AND status IN ('running', 'blocked')`, [input.resumeResearchRunId, input.projectId, input.threadId], @@ -128,6 +137,16 @@ export async function beginResearchTurnLedger( : { rows: [] }; if (resumable.rows[0]) { researchRunId = resumable.rows[0].id; + const persistedProfile = resumable.rows[0].metadata?.[ + RESEARCH_EXECUTION_PROFILE_METADATA_KEY + ]; + if (persistedProfile !== undefined) { + executionProfile = readResearchExecutionProfile(persistedProfile); + } else if (input.executionProfile) { + throw new Error( + "Cannot resume a research run without its original Research Execution Profile.", + ); + } await tx.query( "UPDATE research_runs SET status = 'running', finished_at = NULL, updated_at = now() WHERE id = $1", [researchRunId], @@ -146,7 +165,17 @@ export async function beginResearchTurnLedger( SECURITY_CAPABILITY_MANIFEST_REVISION, input.skillRevision ?? UNAVAILABLE_SKILL_REVISION, startedAt, - JSON.stringify({ version: RESEARCH_TURN_LEDGER_VERSION, ...(input.metadata ?? {}) }), + JSON.stringify({ + version: RESEARCH_TURN_LEDGER_VERSION, + ...(input.metadata ?? {}), + ...(input.executionProfile + ? { + [RESEARCH_EXECUTION_PROFILE_METADATA_KEY]: researchExecutionProfileForensics( + input.executionProfile, + ), + } + : {}), + }), ], ); } @@ -180,7 +209,7 @@ export async function beginResearchTurnLedger( ); }), ); - return { researchRunId, turnLedgerId, startedAt }; + return { researchRunId, turnLedgerId, startedAt, ...(executionProfile ? { executionProfile } : {}) }; } export async function finalizeResearchTurnLedger(input: FinalizeResearchTurnLedgerInput) { diff --git a/src/server/skills/product-skill-registry.ts b/src/server/skills/product-skill-registry.ts index 314bf9c67..c0bf1aa34 100644 --- a/src/server/skills/product-skill-registry.ts +++ b/src/server/skills/product-skill-registry.ts @@ -1,67 +1,61 @@ import { createHash } from "node:crypto"; -import type { - Skill, - SkillMetadata, - WorkspaceSkills, -} from "@mastra/core/workspace"; +import type { Skill, SkillMetadata, WorkspaceSkills } from "@mastra/core/workspace"; export type ProductSkillReviewStatus = "reviewed"; export type ProductSkillRegistryStatus = "complete" | "incomplete"; export type ProductSkillRegistryEntry = { - id: string; - name: string; - description: string; - path: string; - revision: string; - reviewStatus: ProductSkillReviewStatus; - reviewedProvenance: "mastra-workspace"; - applicability: { - userInvocable: boolean; - requiredCapabilities: string[]; - }; - searchTags: string[]; - detail: string; - detailRevision?: string; + id: string; + name: string; + description: string; + path: string; + revision: string; + reviewStatus: ProductSkillReviewStatus; + reviewedProvenance: "mastra-workspace"; + applicability: { + userInvocable: boolean; + requiredCapabilities: string[]; + }; + searchTags: string[]; + detail: string; + detailRevision?: string; }; export type ListProductSkillsInput = { - query?: string; - includeDetails?: boolean; - limit?: number; - capabilities?: readonly string[]; + query?: string; + includeDetails?: boolean; + limit?: number; + capabilities?: readonly string[]; + /** Internal custody path for immutable execution profiles; never expose directly to model tools. */ + includeFullDetails?: boolean; }; export type ListProductSkillsResult = { - root: string; - status: ProductSkillRegistryStatus; - revision: string | null; - diagnostics: string[]; - skillCount: number; - skills: ProductSkillRegistryEntry[]; + root: string; + status: ProductSkillRegistryStatus; + revision: string | null; + diagnostics: string[]; + skillCount: number; + skills: ProductSkillRegistryEntry[]; }; export type ProductSkillRegistry = { - list(input?: ListProductSkillsInput): Promise; + list(input?: ListProductSkillsInput): Promise; }; export function formatProductSkillDirectory(snapshot: ListProductSkillsResult) { - if (snapshot.status === "incomplete") { - return [ - "Product skill discovery status: incomplete.", - ...[...new Set(snapshot.diagnostics)].map( - (diagnostic) => `- ${diagnostic}`, - ), - ].join("\n"); - } - - return [ - `Product skill snapshot: ${snapshot.revision}`, - ...snapshot.skills.map( - (skill) => `- ${skill.id} [${skill.revision}]: ${skill.description}`, - ), - ].join("\n"); + if (snapshot.status === "incomplete") { + return [ + "Product skill discovery status: incomplete.", + ...[...new Set(snapshot.diagnostics)].map((diagnostic) => `- ${diagnostic}`), + ].join("\n"); + } + + return [ + `Product skill snapshot: ${snapshot.revision}`, + ...snapshot.skills.map((skill) => `- ${skill.id} [${skill.revision}]: ${skill.description}`), + ].join("\n"); } const DEFAULT_SKILLS_ROOT = "sandbox/skills"; @@ -74,316 +68,290 @@ const MAX_SKILL_DETAIL_LENGTH = 1_600; * remains responsible for discovery, refresh, validation, and full-body loading. */ export function createProductSkillRegistry( - workspaceSkills: WorkspaceSkills, - root = DEFAULT_SKILLS_ROOT, + workspaceSkills: WorkspaceSkills, + root = DEFAULT_SKILLS_ROOT, ): ProductSkillRegistry { - return { - async list(input: ListProductSkillsInput = {}) { - const capabilities = new Set(input.capabilities ?? []); - const diagnostics: string[] = []; - let metadata: SkillMetadata[]; - - try { - // Publication must invalidate Mastra's in-memory discovery state in every - // runtime mode. The reviewed product corpus is deliberately small, so a - // native refresh is preferable to an indefinitely stale process cache. - await refreshWorkspaceSkills(workspaceSkills); - metadata = await workspaceSkills.list(); - } catch (error) { - return incompleteResult(root, discoveryDiagnostic(error)); - } - - const runtimeSkills = metadata.filter(isRuntimeProductSkill); - - if (runtimeSkills.length === 0) { - return incompleteResult( - root, - "Mastra Workspace discovery returned no visible product research skills.", - ); - } - const visible = runtimeSkills.filter((skill) => - hasRequiredCapabilities(skill, capabilities), - ); - - const query = normalizeQuery(input.query); - const allEntries = await loadSnapshotEntries({ - workspaceSkills, - metadata: visible, - diagnostics, - }); - const revision = catalogRevision(allEntries); - const matchingEntries = allEntries.filter( - (skill) => !query || matchesQuery(skill, query), - ); - const selectedEntries = matchingEntries.slice( - 0, - normalizeLimit(input.limit), - ); - - const skills = input.includeDetails - ? await loadCurrentDetails({ - workspaceSkills, - entries: selectedEntries, - capabilities, - diagnostics, - }) - : selectedEntries; - - return { - root, - status: diagnostics.length > 0 ? "incomplete" : "complete", - revision, - diagnostics, - skillCount: matchingEntries.length, - skills, - }; - }, - }; + return { + async list(input: ListProductSkillsInput = {}) { + const capabilities = new Set(input.capabilities ?? []); + const diagnostics: string[] = []; + let metadata: SkillMetadata[]; + + try { + // Publication must invalidate Mastra's in-memory discovery state in every + // runtime mode. The reviewed product corpus is deliberately small, so a + // native refresh is preferable to an indefinitely stale process cache. + await refreshWorkspaceSkills(workspaceSkills); + metadata = await workspaceSkills.list(); + } catch (error) { + return incompleteResult(root, discoveryDiagnostic(error)); + } + + const runtimeSkills = metadata.filter(isRuntimeProductSkill); + + if (runtimeSkills.length === 0) { + return incompleteResult( + root, + "Mastra Workspace discovery returned no visible product research skills.", + ); + } + const visible = runtimeSkills.filter((skill) => hasRequiredCapabilities(skill, capabilities)); + + const query = normalizeQuery(input.query); + const allEntries = await loadSnapshotEntries({ + workspaceSkills, + metadata: visible, + diagnostics, + }); + const revision = catalogRevision(allEntries); + const matchingEntries = allEntries.filter((skill) => !query || matchesQuery(skill, query)); + const selectedEntries = matchingEntries.slice(0, normalizeLimit(input.limit)); + + const skills = input.includeDetails + ? await loadCurrentDetails({ + workspaceSkills, + entries: selectedEntries, + capabilities, + diagnostics, + includeFullDetails: input.includeFullDetails === true, + }) + : selectedEntries; + + return { + root, + status: diagnostics.length > 0 ? "incomplete" : "complete", + revision, + diagnostics, + skillCount: matchingEntries.length, + skills, + }; + }, + }; } -function incompleteResult( - root: string, - diagnostic: string, -): ListProductSkillsResult { - return { - root, - status: "incomplete", - revision: null, - diagnostics: [diagnostic], - skillCount: 0, - skills: [], - }; +function incompleteResult(root: string, diagnostic: string): ListProductSkillsResult { + return { + root, + status: "incomplete", + revision: null, + diagnostics: [diagnostic], + skillCount: 0, + skills: [], + }; } async function loadCurrentDetails(input: { - workspaceSkills: WorkspaceSkills; - entries: ProductSkillRegistryEntry[]; - capabilities: ReadonlySet; - diagnostics: string[]; + workspaceSkills: WorkspaceSkills; + entries: ProductSkillRegistryEntry[]; + capabilities: ReadonlySet; + diagnostics: string[]; + includeFullDetails: boolean; }) { - const loaded: ProductSkillRegistryEntry[] = []; - let currentSkills: SkillMetadata[]; - try { - await refreshWorkspaceSkills(input.workspaceSkills); - currentSkills = await input.workspaceSkills.list(); - } catch (error) { - input.diagnostics.push( - `Skill bodies could not be revalidated against Mastra Workspace: ${errorMessage(error)}`, - ); - return loaded; - } - - for (const entry of input.entries) { - try { - // Revalidate both discovery and applicability immediately before loading. - const currentMetadata = currentSkills.find( - (candidate) => - candidate.name === entry.id && candidate.path === entry.path, - ); - if ( - !currentMetadata || - !isRuntimeProductSkill(currentMetadata) || - !hasRequiredCapabilities(currentMetadata, input.capabilities) - ) { - input.diagnostics.push( - `Skill ${entry.id} was no longer visible or applicable when its body was requested.`, - ); - continue; - } - - const skill = await input.workspaceSkills.get(currentMetadata.path); - if (!skill || !isRuntimeProductSkill(skill)) { - input.diagnostics.push( - `Skill ${entry.id} could not be loaded from Mastra Workspace.`, - ); - continue; - } - - const currentEntry = toRegistryEntry(skill); - if (currentEntry.revision !== entry.revision) { - input.diagnostics.push( - `Skill ${entry.id} changed after the registry snapshot; retry the lookup against the new revision.`, - ); - continue; - } - - loaded.push(withDetail(currentEntry, skill)); - } catch (error) { - input.diagnostics.push( - `Skill ${entry.id} could not be revalidated: ${errorMessage(error)}`, - ); - } - } - - return loaded; + const loaded: ProductSkillRegistryEntry[] = []; + let currentSkills: SkillMetadata[]; + try { + await refreshWorkspaceSkills(input.workspaceSkills); + currentSkills = await input.workspaceSkills.list(); + } catch (error) { + input.diagnostics.push( + `Skill bodies could not be revalidated against Mastra Workspace: ${errorMessage(error)}`, + ); + return loaded; + } + + for (const entry of input.entries) { + try { + // Revalidate both discovery and applicability immediately before loading. + const currentMetadata = currentSkills.find( + (candidate) => candidate.name === entry.id && candidate.path === entry.path, + ); + if ( + !currentMetadata || + !isRuntimeProductSkill(currentMetadata) || + !hasRequiredCapabilities(currentMetadata, input.capabilities) + ) { + input.diagnostics.push( + `Skill ${entry.id} was no longer visible or applicable when its body was requested.`, + ); + continue; + } + + const skill = await input.workspaceSkills.get(currentMetadata.path); + if (!skill || !isRuntimeProductSkill(skill)) { + input.diagnostics.push(`Skill ${entry.id} could not be loaded from Mastra Workspace.`); + continue; + } + + const currentEntry = toRegistryEntry(skill); + if (currentEntry.revision !== entry.revision) { + input.diagnostics.push( + `Skill ${entry.id} changed after the registry snapshot; retry the lookup against the new revision.`, + ); + continue; + } + + loaded.push(withDetail(currentEntry, skill, input.includeFullDetails)); + } catch (error) { + input.diagnostics.push(`Skill ${entry.id} could not be revalidated: ${errorMessage(error)}`); + } + } + + return loaded; } async function refreshWorkspaceSkills(workspaceSkills: WorkspaceSkills) { - // Mastra's current refresh() expects dynamic skill paths to have been resolved - // by initialization. maybeRefresh() performs that initialization first. - await workspaceSkills.maybeRefresh(); - await workspaceSkills.refresh(); + // Mastra's current refresh() expects dynamic skill paths to have been resolved + // by initialization. maybeRefresh() performs that initialization first. + await workspaceSkills.maybeRefresh(); + await workspaceSkills.refresh(); } async function loadSnapshotEntries(input: { - workspaceSkills: WorkspaceSkills; - metadata: SkillMetadata[]; - diagnostics: string[]; + workspaceSkills: WorkspaceSkills; + metadata: SkillMetadata[]; + diagnostics: string[]; }) { - const entries: ProductSkillRegistryEntry[] = []; - for (const metadata of input.metadata) { - try { - const skill = await input.workspaceSkills.get(metadata.path); - if (!skill || !isRuntimeProductSkill(skill)) { - input.diagnostics.push( - `Skill ${metadata.name} could not be resolved from the current Mastra Workspace snapshot.`, - ); - continue; - } - entries.push(toRegistryEntry(skill)); - } catch (error) { - input.diagnostics.push( - `Skill ${metadata.name} could not be fingerprinted: ${errorMessage(error)}`, - ); - } - } - return entries.sort((left, right) => left.id.localeCompare(right.id)); + const entries: ProductSkillRegistryEntry[] = []; + for (const metadata of input.metadata) { + try { + const skill = await input.workspaceSkills.get(metadata.path); + if (!skill || !isRuntimeProductSkill(skill)) { + input.diagnostics.push( + `Skill ${metadata.name} could not be resolved from the current Mastra Workspace snapshot.`, + ); + continue; + } + entries.push(toRegistryEntry(skill)); + } catch (error) { + input.diagnostics.push( + `Skill ${metadata.name} could not be fingerprinted: ${errorMessage(error)}`, + ); + } + } + return entries.sort((left, right) => left.id.localeCompare(right.id)); } function toRegistryEntry(skill: SkillMetadata): ProductSkillRegistryEntry { - const requiredCapabilities = readRequiredCapabilities(skill); - const userInvocable = skill["user-invocable"] !== false; - const revision = hashJson({ - name: skill.name, - path: skill.path, - description: skill.description, - userInvocable, - requiredCapabilities, - metadata: skill.metadata ?? null, - instructions: "instructions" in skill ? skill.instructions : null, - }); - - return { - id: skill.name, - name: skill.name, - description: skill.description, - path: skill.path, - revision, - reviewStatus: "reviewed", - reviewedProvenance: "mastra-workspace", - applicability: { - userInvocable, - requiredCapabilities, - }, - searchTags: buildSearchTags( - skill.name, - skill.description, - ...requiredCapabilities, - ), - detail: "", - detailRevision: revision, - }; + const requiredCapabilities = readRequiredCapabilities(skill); + const userInvocable = skill["user-invocable"] !== false; + const revision = hashJson({ + name: skill.name, + path: skill.path, + description: skill.description, + userInvocable, + requiredCapabilities, + metadata: skill.metadata ?? null, + instructions: "instructions" in skill ? skill.instructions : null, + }); + + return { + id: skill.name, + name: skill.name, + description: skill.description, + path: skill.path, + revision, + reviewStatus: "reviewed", + reviewedProvenance: "mastra-workspace", + applicability: { + userInvocable, + requiredCapabilities, + }, + searchTags: buildSearchTags(skill.name, skill.description, ...requiredCapabilities), + detail: "", + detailRevision: revision, + }; } function withDetail( - entry: ProductSkillRegistryEntry, - skill: Skill, + entry: ProductSkillRegistryEntry, + skill: Skill, + includeFullDetails = false, ): ProductSkillRegistryEntry { - const detail = skill.instructions.slice(0, MAX_SKILL_DETAIL_LENGTH).trim(); - return { - ...entry, - detail, - detailRevision: entry.revision, - }; + const detail = ( + includeFullDetails ? skill.instructions : skill.instructions.slice(0, MAX_SKILL_DETAIL_LENGTH) + ).trim(); + return { + ...entry, + detail, + detailRevision: entry.revision, + }; } function isRuntimeProductSkill(skill: SkillMetadata) { - const normalizedPath = skill.path.replaceAll("\\", "/"); - const pathSegments = normalizedPath.split("/"); - const visibility = readMetadataString(skill.metadata, "visibility"); - return !pathSegments.includes(".agents") && visibility !== "maintainer"; + const normalizedPath = skill.path.replaceAll("\\", "/"); + const pathSegments = normalizedPath.split("/"); + const visibility = readMetadataString(skill.metadata, "visibility"); + return !pathSegments.includes(".agents") && visibility !== "maintainer"; } -function hasRequiredCapabilities( - skill: SkillMetadata, - availableCapabilities: ReadonlySet, -) { - return readRequiredCapabilities(skill).every((capability) => - availableCapabilities.has(capability), - ); +function hasRequiredCapabilities(skill: SkillMetadata, availableCapabilities: ReadonlySet) { + return readRequiredCapabilities(skill).every((capability) => + availableCapabilities.has(capability), + ); } function readRequiredCapabilities(skill: SkillMetadata) { - const raw = skill.metadata?.requiredCapabilities; - if (!Array.isArray(raw)) return []; - return [ - ...new Set( - raw.filter((value): value is string => typeof value === "string"), - ), - ].sort(); + const raw = skill.metadata?.requiredCapabilities; + if (!Array.isArray(raw)) return []; + return [...new Set(raw.filter((value): value is string => typeof value === "string"))].sort(); } -function readMetadataString( - metadata: Record | undefined, - key: string, -) { - const value = metadata?.[key]; - return typeof value === "string" ? value : undefined; +function readMetadataString(metadata: Record | undefined, key: string) { + const value = metadata?.[key]; + return typeof value === "string" ? value : undefined; } function discoveryDiagnostic(error: unknown) { - return `Mastra Workspace skill discovery is incomplete: ${errorMessage(error)}`; + return `Mastra Workspace skill discovery is incomplete: ${errorMessage(error)}`; } function errorMessage(error: unknown) { - return error instanceof Error ? error.message : String(error); + return error instanceof Error ? error.message : String(error); } function catalogRevision(skills: ProductSkillRegistryEntry[]) { - return hashJson( - skills.map((skill) => ({ id: skill.id, revision: skill.revision })), - ); + return hashJson(skills.map((skill) => ({ id: skill.id, revision: skill.revision }))); } function hashJson(value: unknown) { - return hashText(JSON.stringify(value)); + return hashText(JSON.stringify(value)); } function hashText(value: string) { - return createHash("sha256").update(value).digest("hex"); + return createHash("sha256").update(value).digest("hex"); } function buildSearchTags(...values: string[]) { - const tags = new Set(); - for (const value of values) { - for (const token of value.toLowerCase().split(/[^a-z0-9]+/)) { - if (token.length >= 3) tags.add(token); - } - } - return [...tags].slice(0, 24); + const tags = new Set(); + for (const value of values) { + for (const token of value.toLowerCase().split(/[^a-z0-9]+/)) { + if (token.length >= 3) tags.add(token); + } + } + return [...tags].slice(0, 24); } function normalizeQuery(value: string | undefined) { - const query = value?.trim().toLowerCase(); - return query || undefined; + const query = value?.trim().toLowerCase(); + return query || undefined; } function matchesQuery(skill: ProductSkillRegistryEntry, query: string) { - return [ - skill.id, - skill.name, - skill.description, - skill.path, - skill.reviewStatus, - ...skill.applicability.requiredCapabilities, - ...skill.searchTags, - ] - .join(" ") - .toLowerCase() - .includes(query); + return [ + skill.id, + skill.name, + skill.description, + skill.path, + skill.reviewStatus, + ...skill.applicability.requiredCapabilities, + ...skill.searchTags, + ] + .join(" ") + .toLowerCase() + .includes(query); } function normalizeLimit(value: number | undefined) { - if (typeof value !== "number" || !Number.isFinite(value)) return 50; - return Math.max(1, Math.min(Math.floor(value), 100)); + if (typeof value !== "number" || !Number.isFinite(value)) return 50; + return Math.max(1, Math.min(Math.floor(value), 100)); } diff --git a/tests/integration/research-execution-profile.test.ts b/tests/integration/research-execution-profile.test.ts new file mode 100644 index 000000000..6ffa3f82b --- /dev/null +++ b/tests/integration/research-execution-profile.test.ts @@ -0,0 +1,390 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { RequestContext } from "@mastra/core/request-context"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { SECURITY_CAPABILITY_MANIFEST_REVISION } from "../../src/lib/tools/catalog"; +import { + applyPinnedControllerModelOptions, + toMastraGatewayModelId, +} from "../../src/mastra/agent-controller/agent-controller"; +import { getMastraModelRuntimeOptions } from "../../src/mastra/config/model"; +import { getProjectStore } from "../../src/server/chat/projectAdapter"; +import { configurePinnedControllerSubagentModels } from "../../src/server/chat/security-research-run"; +import { runSecurityResearchTurn } from "../../src/server/chat/securityResearchTurn"; +import { withDatabase } from "../../src/server/db/client"; +import { + projectResearchExecutionProfileToRuntime, + readResearchExecutionBackgroundProfile, + readResearchExecutionProfile, + resolveResearchExecutionProfile, +} from "../../src/server/research/execution-profile"; +import { beginResearchTurnLedger } from "../../src/server/research/turn-ledger"; + +const skillRegistryRevision = "skill-registry-test-revision"; + +describe("Research Execution Profile", () => { + let databaseRoot: string; + let previousDatabaseUrl: string | undefined; + + beforeEach(async () => { + previousDatabaseUrl = process.env.EH_APP_DB_URL; + databaseRoot = await mkdtemp(join(tmpdir(), "exploit-hunter-execution-profile-")); + process.env.EH_APP_DB_URL = `sqlite://${join(databaseRoot, "app.sqlite")}`; + }); + + afterEach(async () => { + if (previousDatabaseUrl === undefined) delete process.env.EH_APP_DB_URL; + else process.env.EH_APP_DB_URL = previousDatabaseUrl; + await rm(databaseRoot, { recursive: true, force: true }); + }); + + it("pins the first run profile across continuations and rejects unavailable revisions", async () => { + const store = await getProjectStore(); + const project = await store.createProject({ + name: "Pinned execution profile", + }); + const thread = await store.createThread(project.id, { + title: "Profile continuity", + }); + const firstMessage = await store.addMessage(project.id, { + threadId: thread.id, + role: "user", + content: "Start passive review.", + }); + const firstProfile = resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/qwen/qwen3-30b-a3b?maxTokens=4096", + modelOverrides: { + hunt: "llm://lmstudio/hunt/strong-hunt?maxTokens=8192", + report: "llm://lmstudio/report/report-model?maxTokens=2048", + }, + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + targetMode: "none", + networkPolicy: "none", + maxToolCalls: 12, + maxRuntimeMs: 30_000, + maxTurns: 8, + uiCollection: "evidence", + policies: { approval: "durable-intent-v1", terminal: "gated-v1" }, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + selectedSkills: [], + }, + ); + const first = await beginResearchTurnLedger({ + projectId: project.id, + threadId: thread.id, + userMessageId: firstMessage.id, + requestedModelEnvelope: { + modelUri: firstProfile.requested.modelUri ?? null, + }, + executionProfile: firstProfile, + }); + + const secondMessage = await store.addMessage(project.id, { + threadId: thread.id, + role: "user", + content: "Continue with the existing run.", + }); + const changedProfile = resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/google/gemma-3-27b-it?maxTokens=1024", + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: [], + skillRegistryRevision, + targetMode: "container", + networkPolicy: "approved-targets", + maxToolCalls: 2, + uiCollection: "planning", + policies: { approval: "durable-intent-v1", terminal: "gated-v1" }, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + selectedSkills: [], + }, + ); + const resumed = await beginResearchTurnLedger({ + projectId: project.id, + threadId: thread.id, + userMessageId: secondMessage.id, + requestedModelEnvelope: { + modelUri: changedProfile.requested.modelUri ?? null, + }, + executionProfile: changedProfile, + resumeResearchRunId: first.researchRunId, + }); + + const resumedProfile = required(resumed.executionProfile, "resumed execution profile"); + expect(resumedProfile.profileId).toBe(firstProfile.profileId); + expect(resumedProfile.profileId).not.toBe(changedProfile.profileId); + const resumedRuntime = projectResearchExecutionProfileToRuntime(resumedProfile); + expect(resumedRuntime).toMatchObject({ + modelUri: firstProfile.effective.models.coordinator.modelUri, + modelOverrides: { + hunt: firstProfile.effective.models.hunt.modelUri, + report: firstProfile.effective.models.report.modelUri, + }, + enabledToolIds: ["tool:artifactAccessTool"], + runtimeMaxToolCalls: 12, + }); + const selectedModels = new Map(); + await configurePinnedControllerSubagentModels( + { + subagents: { + model: { + set: async ({ agentType, modelId }: { agentType?: string; modelId: string }) => { + selectedModels.set(required(agentType, "stage id"), modelId); + }, + }, + }, + } as never, + new RequestContext([ + ["coordinatorModelUri", resumedRuntime.modelUri], + ["modelOverrides", resumedRuntime.modelOverrides], + ]), + ); + expect(selectedModels.get("hunt")).toBe( + toMastraGatewayModelId(firstProfile.effective.models.hunt.modelUri), + ); + expect(selectedModels.get("report")).toBe( + toMastraGatewayModelId(firstProfile.effective.models.report.modelUri), + ); + expect(selectedModels.get("composition")).toBe( + toMastraGatewayModelId(firstProfile.effective.models.coordinator.modelUri), + ); + await withDatabase(async (db) => { + const row = required( + ( + await db.query<{ metadata: Record }>( + "SELECT metadata FROM research_runs WHERE id = $1", + [first.researchRunId], + ) + ).rows[0], + "persisted research run", + ); + expect(readResearchExecutionProfile(row.metadata.executionProfile).profileId).toBe( + firstProfile.profileId, + ); + }); + await expect( + readResearchExecutionBackgroundProfile({ + projectId: project.id, + threadId: thread.id, + researchRunId: first.researchRunId, + profileId: firstProfile.profileId, + }), + ).resolves.toMatchObject({ + profileId: firstProfile.profileId, + effective: { + models: { + hunt: { modelUri: firstProfile.effective.models.hunt.modelUri }, + report: { modelUri: firstProfile.effective.models.report.modelUri }, + }, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + }, + }); + await expect( + readResearchExecutionBackgroundProfile({ + projectId: project.id, + threadId: "wrong-thread", + researchRunId: first.researchRunId, + profileId: firstProfile.profileId, + }), + ).rejects.toThrow(/has no pinned execution profile/); + + expect(() => + resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/qwen/qwen3-30b-a3b", + capabilityManifestRevision: "removed-capability-revision", + capabilityIds: [], + skillRegistryRevision, + targetMode: "none", + networkPolicy: "none", + uiCollection: "all", + policies: {}, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + selectedSkills: [], + }, + ), + ).toThrow(/revision .* is unavailable/); + + expect(() => + resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/qwen/qwen3-30b-a3b", + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:missingTool"], + skillRegistryRevision, + selectedSkills: [{ id: "missing-skill", revision: "stale" }], + targetMode: "none", + networkPolicy: "none", + uiCollection: "all", + policies: {}, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision, + selectedSkills: [], + }, + ), + ).toThrow(/capability tool:missingTool is unavailable/); + + expect(() => + resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/qwen/qwen3-30b-a3b", + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: [], + skillRegistryRevision, + selectedSkills: [{ id: "review", revision: "stale" }], + targetMode: "none", + networkPolicy: "none", + uiCollection: "all", + policies: {}, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: [], + skillRegistryRevision, + selectedSkills: [{ id: "review", revision: "current" }], + }, + ), + ).toThrow(/skill review revision stale is unavailable/); + expect(() => + resolveResearchExecutionProfile( + { + modelUri: "llm://unsupported-provider/no-model", + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: [], + skillRegistryRevision, + targetMode: "none", + networkPolicy: "none", + uiCollection: "all", + policies: {}, + }, + { + capabilityManifestRevision: SECURITY_CAPABILITY_MANIFEST_REVISION, + capabilityIds: [], + skillRegistryRevision, + selectedSkills: [], + }, + ), + ).toThrow(/provider unsupported-provider is unsupported/); + }, 15_000); + + it("applies pinned URI options at the controller gateway seam", async () => { + const observed: Record[] = []; + const modelUri = "llm://openrouter/qwen/qwen3?maxTokens=321&temperature=0.2"; + const model = applyPinnedControllerModelOptions(modelUri, { + specificationVersion: "v2", + provider: "test", + modelId: "pinned", + doGenerate: async (options: Record) => { + observed.push(options); + return {}; + }, + }); + await (model.doGenerate as (input: unknown) => Promise)({ prompt: [] }); + expect(observed[0]).toMatchObject(getMastraModelRuntimeOptions(modelUri).modelSettings); + expect(toMastraGatewayModelId(modelUri)).toMatch(/^exploit-hunter-native\/profile\//); + }); + + it("publishes the same effective model values to the turn runner and forensic run record", async () => { + const store = await getProjectStore(); + const project = await store.createProject({ + name: "Profile runtime projection", + }); + const thread = await store.createThread(project.id, { + title: "Runtime projection", + }); + const runAgent = vi.fn(async (_projectId, input, persistedMessage) => ({ + message: { + id: "profile-agent-message", + projectId: project.id, + threadId: persistedMessage.threadId, + role: "assistant" as const, + content: "Profile observed.", + createdAt: new Date().toISOString(), + metadata: { + researchTurnProvenance: { + effectiveModelEnvelope: { + modelUri: input.metadata?.modelUri, + contextWindowTokens: input.metadata?.contextWindowTokens, + maxOutputTokens: input.metadata?.maxOutputTokens, + }, + }, + researchRunStatus: { lifecycle: "completed" }, + }, + }, + })); + + const result = await runSecurityResearchTurn( + project.id, + { + threadId: thread.id, + role: "user", + content: "Review the saved evidence.", + metadata: { + modelUri: "llm://lmstudio/qwen/qwen3-30b-a3b?maxTokens=4096", + modelOverrides: { + hunt: "llm://lmstudio/qwen/qwen3-30b-a3b?maxTokens=8192", + }, + enabledToolIds: ["tool:artifactAccessTool"], + runtimeMaxToolCalls: 9, + }, + }, + { + store, + runAgent, + resolveWorkspace: async () => null, + generateThreadTitle: async () => null, + }, + ); + + const runtimeCall = required(runAgent.mock.calls[0], "turn runner invocation"); + const runtimeMetadata = required(runtimeCall[1].metadata, "turn runtime metadata"); + const runtimeProfile = readResearchExecutionProfile(runtimeMetadata.researchExecutionProfile); + expect(runtimeMetadata).toMatchObject({ + modelUri: runtimeProfile.effective.models.coordinator.modelUri, + modelOverrides: { hunt: runtimeProfile.effective.models.hunt.modelUri }, + researchExecutionProfileId: runtimeProfile.profileId, + runtimeMaxToolCalls: 9, + }); + await withDatabase(async (db) => { + const row = required( + ( + await db.query<{ metadata: Record }>( + "SELECT metadata FROM research_runs WHERE id = $1", + [result.agent?.message?.metadata?.researchRunId], + ) + ).rows[0], + "persisted research run", + ); + expect(readResearchExecutionProfile(row.metadata.executionProfile).profileId).toBe( + runtimeProfile.profileId, + ); + }); + }, 15_000); +}); + +function required(value: T | null | undefined, label: string): T { + if (value === null || value === undefined) throw new Error(`Missing ${label}.`); + return value; +} diff --git a/tests/integration/scheduler-backed-research.test.ts b/tests/integration/scheduler-backed-research.test.ts index b69807444..cf4557214 100644 --- a/tests/integration/scheduler-backed-research.test.ts +++ b/tests/integration/scheduler-backed-research.test.ts @@ -8,6 +8,7 @@ import { type SchedulerResearchTask, } from "../../src/mastra/workflows"; import { startSchedulerDrainBackgroundTask } from "../../src/server/chat/mastra-background-tasks"; +import { resolveResearchExecutionProfile } from "../../src/server/research/execution-profile"; import { createInMemorySchedulerRepository, createSchedulerService, @@ -34,6 +35,29 @@ function makeService(): SchedulerService { return createSchedulerService(createInMemorySchedulerRepository()); } +function makeExecutionProfile() { + return resolveResearchExecutionProfile( + { + modelUri: "llm://lmstudio/coordinator/pinned", + modelOverrides: { hunt: "llm://lmstudio/hunt/pinned" }, + capabilityManifestRevision: "builtin-v1", + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision: "skills-pinned", + selectedSkills: [{ id: "runtime-review", revision: "skill-pinned" }], + targetMode: "none", + networkPolicy: "none", + uiCollection: "all", + policies: {}, + }, + { + capabilityManifestRevision: "builtin-v1", + capabilityIds: ["tool:artifactAccessTool"], + skillRegistryRevision: "skills-pinned", + selectedSkills: [{ id: "runtime-review", revision: "skill-pinned" }], + }, + ); +} + function wrapScheduler( base: SchedulerService, overrides: Partial, @@ -55,41 +79,46 @@ function makeRunner( executeTask: (input: SchedulerResearchExecuteTaskInput) => Promise, workspaceService: Record = unlockedWorkspaceService(), blockerService?: Record, + onPersist?: (input: { executionProfile: Record }) => void, ): SchedulerBackedResearchRunner { return createSchedulerBackedResearchRunner({ scheduler, executeTask, workspaceService: workspaceService as never, - persistHandoff: async (input) => ({ - artifactId: `artifact-${input.schedulerTaskId}`, - handoff: { - schemaVersion: 1, - projectId: input.projectId, - schedulerTaskId: input.schedulerTaskId, - stage: input.stage, - taskId: input.taskId, - attempt: input.attempt, - outcome: "completed", - summary: typeof input.draft.summary === "string" ? input.draft.summary : input.draft.output, - coverage: [], - negativeCoverage: [], - assertionIds: [], - hypothesisIds: [], - artifactIds: [], - findingIds: [], - completedBranches: [], - openBranches: [], - blockers: [], - proposedTaskIds: [], - approvalNeeds: [], - executionProfile: input.executionProfile, - attribution: { producerAgentId: input.producerAgentId, parentHandoffIds: [] }, - budgetUsage: {}, - resourceClaims: [], - output: input.draft.output, - createdAt: new Date(0).toISOString(), - }, - }), + persistHandoff: async (input) => { + onPersist?.(input); + return { + artifactId: `artifact-${input.schedulerTaskId}`, + handoff: { + schemaVersion: 1, + projectId: input.projectId, + schedulerTaskId: input.schedulerTaskId, + stage: input.stage, + taskId: input.taskId, + attempt: input.attempt, + outcome: "completed", + summary: + typeof input.draft.summary === "string" ? input.draft.summary : input.draft.output, + coverage: [], + negativeCoverage: [], + assertionIds: [], + hypothesisIds: [], + artifactIds: [], + findingIds: [], + completedBranches: [], + openBranches: [], + blockers: [], + proposedTaskIds: [], + approvalNeeds: [], + executionProfile: input.executionProfile, + attribution: { producerAgentId: input.producerAgentId, parentHandoffIds: [] }, + budgetUsage: {}, + resourceClaims: [], + output: input.draft.output, + createdAt: new Date(0).toISOString(), + }, + }; + }, readHandoff: async () => { throw new Error("unexpected parent handoff read"); }, @@ -229,6 +258,52 @@ describe("scheduler-backed research runner", () => { expect(all.every((task) => task.handoffSummary?.startsWith("Hunt output"))).toBe(true); }); + it("carries the pinned profile into background work and handoffs without approval grants", async () => { + const scheduler = makeService(); + const executed: SchedulerResearchExecuteTaskInput[] = []; + const persisted: Array<{ + researchRunId?: string; + executionProfileId?: string; + executionProfile: Record; + }> = []; + const runner = makeRunner( + scheduler, + async (input) => { + executed.push(input); + return "profile-bound output"; + }, + unlockedWorkspaceService(), + undefined, + (input) => persisted.push(input), + ); + const profile = makeExecutionProfile(); + + await runner.enqueueResearchTasks({ + ...enqueueInput("project-1", [makeTask({ id: "t-profile" })]), + researchRunId: "research-run-profile", + executionProfileId: profile.profileId, + }); + await runner.runEnqueuedTasks({ + projectId: "project-1", + researchRunId: "research-run-profile", + executionProfileId: profile.profileId, + owner: "worker-profile", + concurrency: 1, + approvedToolIds: ["tool:artifactAccessTool"], + commandAllowlist: ["curl *"], + executionProfile: profile, + }); + + expect(executed[0]?.context.executionProfile).toEqual(profile); + expect(executed[0]?.context.executionProfile).not.toHaveProperty("approvedToolIds"); + expect(executed[0]?.context.executionProfile).not.toHaveProperty("commandAllowlist"); + expect(persisted[0]).toMatchObject({ + researchRunId: "research-run-profile", + executionProfileId: profile.profileId, + }); + expect(persisted[0]?.executionProfile).toEqual(profile); + }, 15_000); + it("runs a long-lived worker loop that polls again after idle ticks", async () => { const scheduler = makeService(); const runner = makeRunner(scheduler, async ({ task }) => `loop output for ${task.taskId}`); @@ -245,6 +320,8 @@ describe("scheduler-backed research runner", () => { const summary = await runSchedulerResearchWorkerLoop( { projectId: "project-1", + researchRunId: "research-run-1", + executionProfileId: "profile-1", owner: "worker-loop-1", concurrency: 1, idlePollIntervalMs: 123, @@ -864,9 +941,12 @@ describe("scheduler-backed research runner", () => { }), }; + const executionProfile = makeExecutionProfile(); const result = await startSchedulerDrainBackgroundTask( { projectId: "project-1", + researchRunId: "research-run-1", + executionProfileId: executionProfile.profileId, threadId: "thread-1", owner: "worker-bg", concurrency: 2, @@ -880,6 +960,7 @@ describe("scheduler-backed research runner", () => { runner: { runEnqueuedTasks } as never, runId: () => "run-drain-1", toolCallId: () => "call-drain-1", + resolveExecutionProfile: async () => executionProfile, }, ); @@ -924,5 +1005,5 @@ describe("scheduler-backed research runner", () => { isolateTaskWorkspaces: true, }), ); - }); + }, 15_000); });