From 501a344789ad0097c2c00468cafc47874ad8f10b Mon Sep 17 00:00:00 2001 From: Tycen McCann Date: Wed, 13 May 2026 11:42:37 -0700 Subject: [PATCH 1/3] feat(blueprint): add tool profiles schema for dynamic per-task tool selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces named Tool Profiles — deploy-time configurations that define which MCP servers, skills, and Cedar policies are available to the agent on a per-task basis. This is PR 1 of 2: schema, validation, and CLI flag only (no runtime resolution yet). - Add ToolProfile interface and toolProfiles prop to Blueprint construct - Store profiles as JSON in RepoConfig DynamoDB table - Add tool_profile field to CreateTaskRequest, TaskRecord, TaskDetail - Validate profile name format (lowercase alphanumeric + hyphens, 1-64 chars) - Validate profile exists in repo's Blueprint at task admission - Add --tool-profile flag to CLI submit command - Add tool_profile field to agent TaskConfig model - Mirror types across CDK ↔ CLI sync boundary Co-Authored-By: Claude Opus 4.6 --- agent/src/config.py | 3 + agent/src/models.py | 3 + agent/tests/test_config.py | 19 +++ agent/tests/test_models.py | 17 +++ cdk/src/constructs/blueprint.ts | 88 ++++++++++++++ cdk/src/handlers/shared/create-task-core.ts | 22 +++- cdk/src/handlers/shared/repo-config.ts | 44 +++++++ cdk/src/handlers/shared/types.ts | 11 ++ cdk/src/handlers/shared/validation.ts | 21 ++++ cdk/test/constructs/blueprint.test.ts | 111 +++++++++++++++++- .../handlers/shared/create-task-core.test.ts | 110 +++++++++++++++++ cdk/test/handlers/shared/repo-config.test.ts | 48 +++++++- cdk/test/handlers/shared/validation.test.ts | 57 +++++++++ cli/src/commands/submit.ts | 10 ++ cli/src/types.ts | 11 ++ 15 files changed, 571 insertions(+), 4 deletions(-) diff --git a/agent/src/config.py b/agent/src/config.py index 0e9e4958..337b1517 100644 --- a/agent/src/config.py +++ b/agent/src/config.py @@ -92,6 +92,7 @@ def build_config( channel_metadata: dict[str, str] | None = None, trace: bool = False, user_id: str = "", + tool_profile: str = "", ) -> TaskConfig: """Build and validate configuration from explicit parameters. @@ -146,6 +147,7 @@ def build_config( channel_metadata=channel_metadata or {}, trace=trace, user_id=user_id, + tool_profile=tool_profile, ) @@ -170,6 +172,7 @@ def get_config() -> TaskConfig: # an unreachable ``traces//`` key. trace=os.environ.get("TRACE", "").lower() in ("1", "true", "yes"), user_id=os.environ.get("USER_ID", ""), + tool_profile=os.environ.get("TOOL_PROFILE", ""), ) except ValueError as e: print(f"ERROR: {e}", file=sys.stderr) diff --git a/agent/src/models.py b/agent/src/models.py index 255c18f9..137f312b 100644 --- a/agent/src/models.py +++ b/agent/src/models.py @@ -127,6 +127,9 @@ class TaskConfig(BaseModel): # previews live, so dropping ``trace`` here silently no-ops the # feature for the fields that matter. trace: bool = False + # Tool profile selected at task submission (from Blueprint.toolProfiles). + # Empty string means legacy single-tier behavior (no profile selected). + tool_profile: str = "" # Enriched mid-flight by pipeline.py: cedar_policies: list[str] = [] issue: GitHubIssue | None = None diff --git a/agent/tests/test_config.py b/agent/tests/test_config.py index d9e32c84..5b9d22c9 100644 --- a/agent/tests/test_config.py +++ b/agent/tests/test_config.py @@ -85,3 +85,22 @@ def test_auto_generated_task_id(self): ) assert config.task_id assert len(config.task_id) == 12 + + def test_tool_profile_defaults_to_empty(self): + config = build_config( + repo_url="owner/repo", + task_description="fix bug", + github_token="ghp_test", + aws_region="us-east-1", + ) + assert config.tool_profile == "" + + def test_tool_profile_passed_through(self): + config = build_config( + repo_url="owner/repo", + task_description="fix bug", + github_token="ghp_test", + aws_region="us-east-1", + tool_profile="frontend", + ) + assert config.tool_profile == "frontend" diff --git a/agent/tests/test_models.py b/agent/tests/test_models.py index 91236115..f20ea6e2 100644 --- a/agent/tests/test_models.py +++ b/agent/tests/test_models.py @@ -314,6 +314,23 @@ def test_trace_false_allows_empty_user_id(self): assert config.trace is False assert config.user_id == "" + def test_tool_profile_defaults_to_empty_string(self): + config = TaskConfig( + repo_url="owner/repo", + github_token="ghp_test", + aws_region="us-east-1", + ) + assert config.tool_profile == "" + + def test_tool_profile_accepts_valid_name(self): + config = TaskConfig( + repo_url="owner/repo", + github_token="ghp_test", + aws_region="us-east-1", + tool_profile="frontend", + ) + assert config.tool_profile == "frontend" + class TestRepoSetup: def test_construction(self): diff --git a/cdk/src/constructs/blueprint.ts b/cdk/src/constructs/blueprint.ts index 453bf069..ac465332 100644 --- a/cdk/src/constructs/blueprint.ts +++ b/cdk/src/constructs/blueprint.ts @@ -25,6 +25,48 @@ import { Construct, IValidation } from 'constructs'; const REPO_PATTERN = /^[a-zA-Z0-9._-]+\/[a-zA-Z0-9._-]+$/; const DOMAIN_PATTERN = /^(\*\.)?[a-z0-9]([a-z0-9-]*[a-z0-9])?(\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)*$/; +const TOOL_PROFILE_NAME_PATTERN = /^[a-z0-9][a-z0-9-]{0,62}[a-z0-9]$/; + +/** + * A named tool profile that defines which tools, MCP servers, skills, + * and Cedar policies are available to the agent for a given task scope. + * + * Profiles are deploy-time artifacts — defined in CDK source, deployed + * via CloudFormation, and stored in DynamoDB. At task submission time, + * only the profile *name* is user-controlled; the definition itself + * is trusted (same trust level as Blueprint.security.cedarPolicies). + */ +export interface ToolProfile { + /** + * Tool capability tier for this profile. + * @default 'default' + */ + readonly capabilityTier?: 'default' | 'extended'; + + /** + * MCP server identifiers activated for this profile. + * These must correspond to MCP servers registered with the platform. + */ + readonly mcpServers?: readonly string[]; + + /** + * Skill identifiers activated for this profile. + * References deploy-time skill definitions (SKILL.md directories) + * bundled into the agent runtime image or fetched at session start. + */ + readonly skills?: readonly string[]; + + /** + * Additional Cedar policy statements for this profile. + * Appended to baseline deny-list policies during evaluation. + */ + readonly cedarPolicies?: readonly string[]; + + /** + * Human-readable description of the profile's purpose. + */ + readonly description?: string; +} /** * Properties for the Blueprint construct. @@ -119,6 +161,17 @@ export interface BlueprintProps { */ readonly egressAllowlist?: string[]; }; + + /** + * Named tool profiles defining per-task tool configurations. + * Keys are profile names (lowercase alphanumeric + hyphens, 2-64 chars). + * At task submission, the user selects a profile by name; the platform + * activates only the tools/skills/policies defined in that profile. + * + * If omitted, the repo uses legacy single-tier behavior based on + * security.cedarPolicies alone. + */ + readonly toolProfiles?: Readonly>; } /** @@ -145,15 +198,22 @@ export class Blueprint extends Construct { */ public readonly cedarPolicies: readonly string[]; + /** + * Tool profiles from the toolProfiles prop, exposed for inspection. + */ + public readonly toolProfiles: Readonly>; + constructor(scope: Construct, id: string, props: BlueprintProps) { super(scope, id); this.egressAllowlist = [...(props.networking?.egressAllowlist ?? [])]; this.cedarPolicies = [...(props.security?.cedarPolicies ?? [])]; + this.toolProfiles = props.toolProfiles ?? {}; // Validate repo format at construct time this.node.addValidation(new RepoFormatValidation(props.repo)); this.node.addValidation(new DomainFormatValidation(this.egressAllowlist)); + this.node.addValidation(new ToolProfileNameValidation(this.toolProfiles)); const now = new Date().toISOString(); @@ -192,6 +252,9 @@ export class Blueprint extends Construct { if (this.cedarPolicies.length > 0) { item.cedar_policies = { L: this.cedarPolicies.map(p => ({ S: p })) }; } + if (Object.keys(this.toolProfiles).length > 0) { + item.tool_profiles = { S: JSON.stringify(this.toolProfiles) }; + } new cr.AwsCustomResource(this, 'RepoConfigCR', { timeout: Duration.minutes(5), @@ -263,6 +326,7 @@ export class Blueprint extends Construct { if (props.pipeline?.pollIntervalMs !== undefined) fields.push(', #poll_interval_ms = :poll_interval_ms'); if (this.egressAllowlist.length > 0) fields.push(', #egress_allowlist = :egress_allowlist'); if (this.cedarPolicies.length > 0) fields.push(', #cedar_policies = :cedar_policies'); + if (Object.keys(this.toolProfiles).length > 0) fields.push(', #tool_profiles = :tool_profiles'); return fields.join(''); } @@ -277,6 +341,7 @@ export class Blueprint extends Construct { if (props.pipeline?.pollIntervalMs !== undefined) names['#poll_interval_ms'] = 'poll_interval_ms'; if (this.egressAllowlist.length > 0) names['#egress_allowlist'] = 'egress_allowlist'; if (this.cedarPolicies.length > 0) names['#cedar_policies'] = 'cedar_policies'; + if (Object.keys(this.toolProfiles).length > 0) names['#tool_profiles'] = 'tool_profiles'; return names; } @@ -291,6 +356,7 @@ export class Blueprint extends Construct { if (props.pipeline?.pollIntervalMs !== undefined) values[':poll_interval_ms'] = { N: String(props.pipeline.pollIntervalMs) }; if (this.egressAllowlist.length > 0) values[':egress_allowlist'] = { L: this.egressAllowlist.map(d => ({ S: d })) }; if (this.cedarPolicies.length > 0) values[':cedar_policies'] = { L: this.cedarPolicies.map(p => ({ S: p })) }; + if (Object.keys(this.toolProfiles).length > 0) values[':tool_profiles'] = { S: JSON.stringify(this.toolProfiles) }; return values; } } @@ -325,3 +391,25 @@ class DomainFormatValidation implements IValidation { return errors; } } + +/** + * Validates tool profile names (lowercase alphanumeric + hyphens, 2-64 chars). + * Single-char profile names are allowed if they match [a-z0-9]. + */ +class ToolProfileNameValidation implements IValidation { + constructor(private readonly profiles: Readonly>) {} + + public validate(): string[] { + const errors: string[] = []; + for (const name of Object.keys(this.profiles)) { + if (name.length === 1) { + if (!/^[a-z0-9]$/.test(name)) { + errors.push(`Invalid tool profile name: '${name}'. Expected lowercase alphanumeric and hyphens (2-64 chars).`); + } + } else if (!TOOL_PROFILE_NAME_PATTERN.test(name)) { + errors.push(`Invalid tool profile name: '${name}'. Expected lowercase alphanumeric and hyphens (2-64 chars).`); + } + } + return errors; + } +} diff --git a/cdk/src/handlers/shared/create-task-core.ts b/cdk/src/handlers/shared/create-task-core.ts index f8d5c69d..db0cee19 100644 --- a/cdk/src/handlers/shared/create-task-core.ts +++ b/cdk/src/handlers/shared/create-task-core.ts @@ -29,10 +29,10 @@ import type { APIGatewayProxyResult } from 'aws-lambda'; import { ulid } from 'ulid'; import { generateBranchName } from './gateway'; import { logger } from './logger'; -import { checkRepoOnboarded } from './repo-config'; +import { checkRepoOnboarded, loadRepoConfig, parseToolProfiles, isValidToolProfile } from './repo-config'; import { ErrorCode, errorResponse, successResponse } from './response'; import { type ChannelSource, type CreateTaskRequest, isPrTaskType, type TaskRecord, type TaskType, toTaskDetail } from './types'; -import { computeTtlEpoch, DEFAULT_MAX_TURNS, hasTaskSpec, isValidIdempotencyKey, isValidRepo, isValidTaskDescriptionLength, isValidTaskType, MAX_TASK_DESCRIPTION_LENGTH, validateMaxBudgetUsd, validateMaxTurns, validatePrNumber } from './validation'; +import { computeTtlEpoch, DEFAULT_MAX_TURNS, hasTaskSpec, isValidIdempotencyKey, isValidRepo, isValidTaskDescriptionLength, isValidTaskType, MAX_TASK_DESCRIPTION_LENGTH, validateMaxBudgetUsd, validateMaxTurns, validatePrNumber, validateToolProfile } from './validation'; import { TaskStatus } from '../../constructs/task-status'; /** @@ -132,6 +132,23 @@ export async function createTaskCore( } const userTrace = body.trace === true; + // Validate tool_profile format (if provided) + const toolProfileResult = validateToolProfile(body.tool_profile); + if (toolProfileResult === null) { + return errorResponse(400, ErrorCode.VALIDATION_ERROR, 'Invalid tool_profile. Must be a lowercase alphanumeric string with hyphens (1-64 chars).', requestId); + } + + // If a tool_profile is specified, validate it exists in the repo's Blueprint + if (toolProfileResult !== undefined) { + const repoConfig = await loadRepoConfig(body.repo); + if (repoConfig) { + const profiles = parseToolProfiles(repoConfig.tool_profiles); + if (!isValidToolProfile(toolProfileResult, profiles)) { + return errorResponse(422, ErrorCode.VALIDATION_ERROR, 'Invalid tool_profile. The specified profile does not exist for this repository.', requestId); + } + } + } + // 2. Screen task description with Bedrock Guardrail (fail-closed: unscreened content // must not reach the agent — a Bedrock outage blocks task submissions) if (bedrockClient && body.task_description) { @@ -233,6 +250,7 @@ export async function createTaskCore( ...(userMaxTurns !== undefined && { max_turns: userMaxTurns }), ...(userMaxBudgetUsd !== undefined && { max_budget_usd: userMaxBudgetUsd }), ...(userTrace && { trace: true }), + ...(toolProfileResult !== undefined && { tool_profile: toolProfileResult }), ...(context.idempotencyKey && { idempotency_key: context.idempotencyKey }), channel_source: context.channelSource, channel_metadata: context.channelMetadata, diff --git a/cdk/src/handlers/shared/repo-config.ts b/cdk/src/handlers/shared/repo-config.ts index 60dd23aa..ca884ff0 100644 --- a/cdk/src/handlers/shared/repo-config.ts +++ b/cdk/src/handlers/shared/repo-config.ts @@ -27,6 +27,15 @@ import { logger } from './logger'; */ export type ComputeType = 'agentcore' | 'ecs'; +/** Runtime representation of a tool profile stored in RepoConfig. */ +export interface StoredToolProfile { + readonly capabilityTier?: 'default' | 'extended'; + readonly mcpServers?: readonly string[]; + readonly skills?: readonly string[]; + readonly cedarPolicies?: readonly string[]; + readonly description?: string; +} + export interface RepoConfig { readonly repo: string; readonly status: 'active' | 'removed'; @@ -42,6 +51,8 @@ export interface RepoConfig { readonly poll_interval_ms?: number; readonly egress_allowlist?: string[]; readonly cedar_policies?: string[]; + /** JSON-serialized map of profile name → ToolProfile, written by Blueprint. */ + readonly tool_profiles?: string; } /** @@ -59,6 +70,8 @@ export interface BlueprintConfig { readonly poll_interval_ms?: number; readonly egress_allowlist?: string[]; readonly cedar_policies?: string[]; + /** Parsed tool profiles map (profile name → definition). */ + readonly tool_profiles?: Readonly>; } const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); @@ -138,3 +151,34 @@ export async function loadRepoConfig(repo: string): Promise { throw new Error(`Unable to load repo config for '${repo}': ${String(err)}`); } } + +/** + * Parse the tool_profiles JSON string from a RepoConfig into a typed map. + * Returns an empty object if the field is absent or unparseable. + */ +export function parseToolProfiles(raw: string | undefined): Readonly> { + if (!raw) return {}; + try { + const parsed = JSON.parse(raw); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + logger.warn('tool_profiles is not a valid object, ignoring', { raw_type: typeof parsed }); + return {}; + } + return parsed as Record; + } catch (err) { + logger.warn('Failed to parse tool_profiles JSON', { + error: err instanceof Error ? err.message : String(err), + }); + return {}; + } +} + +/** + * Validate that a tool profile name exists in the given profiles map. + * @param profileName - the profile name from the task request. + * @param profiles - the parsed tool profiles from RepoConfig. + * @returns true if the profile exists. + */ +export function isValidToolProfile(profileName: string, profiles: Readonly>): boolean { + return Object.prototype.hasOwnProperty.call(profiles, profileName); +} diff --git a/cdk/src/handlers/shared/types.ts b/cdk/src/handlers/shared/types.ts index 9aff1918..d02eb498 100644 --- a/cdk/src/handlers/shared/types.ts +++ b/cdk/src/handlers/shared/types.ts @@ -106,6 +106,8 @@ export interface TaskRecord { readonly memory_written?: boolean; readonly compute_type?: ComputeType; readonly compute_metadata?: Record; + /** Tool profile name selected at task submission (from Blueprint.toolProfiles). */ + readonly tool_profile?: string; readonly ttl?: number; /** * Optional per-task override for the FanOutConsumer's channel filters @@ -198,6 +200,8 @@ export interface TaskDetail { * the field being present; CLI download resolves this via the * ``get-trace-url`` handler rather than hitting S3 directly. */ readonly trace_s3_uri: string | null; + /** Tool profile selected at submission, or ``null`` for legacy single-tier tasks. */ + readonly tool_profile: string | null; } /** @@ -275,6 +279,12 @@ export interface CreateTaskRequest { readonly attachments?: Attachment[]; /** Enable 4 KB debug previews (design §10.1, opt-in per task). */ readonly trace?: boolean; + /** + * Named tool profile to activate for this task. Must reference a profile + * defined in the repo's Blueprint.toolProfiles. When omitted, the repo's + * legacy single-tier behavior applies. + */ + readonly tool_profile?: string; } /** @@ -333,6 +343,7 @@ export function toTaskDetail(record: TaskRecord): TaskDetail { prompt_version: record.prompt_version ?? null, trace: record.trace === true, trace_s3_uri: record.trace_s3_uri ?? null, + tool_profile: record.tool_profile ?? null, }; } diff --git a/cdk/src/handlers/shared/validation.ts b/cdk/src/handlers/shared/validation.ts index 11398c58..0dd9cf8a 100644 --- a/cdk/src/handlers/shared/validation.ts +++ b/cdk/src/handlers/shared/validation.ts @@ -210,6 +210,27 @@ export function computeTtlEpoch(retentionDays: number): number { return Math.floor(Date.now() / 1000) + retentionDays * 86400; } +/** Maximum allowed length for a tool profile name. */ +export const MAX_TOOL_PROFILE_NAME_LENGTH = 64; +const TOOL_PROFILE_NAME_PATTERN = /^[a-z0-9][a-z0-9-]*[a-z0-9]$/; + +/** + * Validate a tool_profile value from a request body. + * @param value - the raw value from the request. + * @returns the valid string, null if invalid (caller should return 400), or undefined if absent. + */ +export function validateToolProfile(value: unknown): string | null | undefined { + if (value === undefined || value === null) return undefined; + if (typeof value !== 'string') return null; + if (value.length === 0 || value.length > MAX_TOOL_PROFILE_NAME_LENGTH) return null; + // Single-char profile names: must be [a-z0-9] + if (value.length === 1) { + return /^[a-z0-9]$/.test(value) ? value : null; + } + if (!TOOL_PROFILE_NAME_PATTERN.test(value)) return null; + return value; +} + /** Valid task type values. Compile-time check ensures this stays in sync with TaskType. */ const TASK_TYPE_LIST = ['new_task', 'pr_iteration', 'pr_review'] as const satisfies readonly TaskType[]; type _AssertExhaustive = Exclude extends never ? true : never; diff --git a/cdk/test/constructs/blueprint.test.ts b/cdk/test/constructs/blueprint.test.ts index 82c9b745..d4ed48e1 100644 --- a/cdk/test/constructs/blueprint.test.ts +++ b/cdk/test/constructs/blueprint.test.ts @@ -20,7 +20,7 @@ import { App, Stack } from 'aws-cdk-lib'; import { Template, Match } from 'aws-cdk-lib/assertions'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; -import { Blueprint, type BlueprintProps } from '../../src/constructs/blueprint'; +import { Blueprint, type BlueprintProps, type ToolProfile } from '../../src/constructs/blueprint'; function createStack(props?: Partial): { stack: Stack; template: Template } { const app = new App(); @@ -407,4 +407,113 @@ describe('Blueprint validation', () => { expect(() => app.synth()).not.toThrow(); }); + + test('rejects invalid tool profile name (uppercase)', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + + new Blueprint(stack, 'Blueprint', { + repo: 'my-org/my-repo', + repoTable, + toolProfiles: { 'INVALID': { capabilityTier: 'default' } }, + }); + + expect(() => app.synth()).toThrow(/Invalid tool profile name/); + }); + + test('accepts valid tool profile names', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + + new Blueprint(stack, 'Blueprint', { + repo: 'my-org/my-repo', + repoTable, + toolProfiles: { + frontend: { capabilityTier: 'extended', mcpServers: ['eslint-mcp'], skills: ['react-patterns'] }, + backend: { capabilityTier: 'default' }, + 'my-infra': { mcpServers: ['aws-cdk-mcp'] }, + }, + }); + + expect(() => app.synth()).not.toThrow(); + }); +}); + +describe('Blueprint toolProfiles', () => { + test('maps tool_profiles to DynamoDB JSON string', () => { + const profiles = { + frontend: { capabilityTier: 'extended' as const, mcpServers: ['eslint-mcp'], skills: ['react-patterns'] }, + backend: { capabilityTier: 'default' as const }, + }; + const { template } = createStack({ toolProfiles: profiles }); + const parts = getCreateJoinParts(template); + const serialized = parts.join(''); + expect(serialized).toContain('"tool_profiles":{"S":'); + expect(serialized).toContain('frontend'); + expect(serialized).toContain('eslint-mcp'); + expect(serialized).toContain('react-patterns'); + }); + + test('omits tool_profiles when not provided', () => { + const { template } = createStack(); + const parts = getCreateJoinParts(template); + const serialized = parts.join(''); + expect(serialized).not.toContain('tool_profiles'); + }); + + test('omits tool_profiles when empty object', () => { + const { template } = createStack({ toolProfiles: {} }); + const parts = getCreateJoinParts(template); + const serialized = parts.join(''); + expect(serialized).not.toContain('tool_profiles'); + }); + + test('onUpdate includes tool_profiles in UpdateExpression', () => { + const { template } = createStack({ + toolProfiles: { frontend: { capabilityTier: 'extended' as const } }, + }); + const parts = getUpdateJoinParts(template); + const serialized = parts.join(''); + expect(serialized).toContain('#tool_profiles'); + }); + + test('exposes toolProfiles as public property', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + + const profiles = { frontend: { capabilityTier: 'extended' as const } }; + const blueprint = new Blueprint(stack, 'Blueprint', { + repo: 'org/my-repo', + repoTable, + toolProfiles: profiles, + }); + + expect(blueprint.toolProfiles).toEqual(profiles); + }); + + test('toolProfiles defaults to empty object', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + + const blueprint = new Blueprint(stack, 'Blueprint', { + repo: 'org/my-repo', + repoTable, + }); + + expect(blueprint.toolProfiles).toEqual({}); + }); }); diff --git a/cdk/test/handlers/shared/create-task-core.test.ts b/cdk/test/handlers/shared/create-task-core.test.ts index e03323de..0e16fce4 100644 --- a/cdk/test/handlers/shared/create-task-core.test.ts +++ b/cdk/test/handlers/shared/create-task-core.test.ts @@ -40,8 +40,12 @@ jest.mock('@aws-sdk/client-bedrock-runtime', () => ({ })); const mockCheckRepoOnboarded = jest.fn(); +const mockLoadRepoConfig = jest.fn(); jest.mock('../../../src/handlers/shared/repo-config', () => ({ checkRepoOnboarded: mockCheckRepoOnboarded, + loadRepoConfig: mockLoadRepoConfig, + parseToolProfiles: jest.requireActual('../../../src/handlers/shared/repo-config').parseToolProfiles, + isValidToolProfile: jest.requireActual('../../../src/handlers/shared/repo-config').isValidToolProfile, })); let ulidCounter = 0; @@ -73,6 +77,7 @@ beforeEach(() => { mockLambdaSend.mockResolvedValue({}); mockBedrockSend.mockResolvedValue({ action: 'NONE' }); mockCheckRepoOnboarded.mockResolvedValue({ onboarded: true }); + mockLoadRepoConfig.mockResolvedValue(null); }); describe('createTaskCore', () => { @@ -551,4 +556,109 @@ describe('createTaskCore', () => { expect(result.statusCode).toBe(400); expect(JSON.parse(result.body).error.message).toContain('trace'); }); + + // -- tool_profile (dynamic tool selection) ---------------------------- + + test('tool_profile omitted creates task without tool_profile field', async () => { + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug' }, + makeContext(), + 'req-tp-1', + ); + expect(result.statusCode).toBe(201); + const body = JSON.parse(result.body); + expect(body.data.tool_profile).toBeNull(); + }); + + test('tool_profile persists on task record and surfaces in response', async () => { + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + tool_profiles: JSON.stringify({ frontend: { capabilityTier: 'extended' } }), + }); + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 'frontend' }, + makeContext(), + 'req-tp-2', + ); + expect(result.statusCode).toBe(201); + const body = JSON.parse(result.body); + expect(body.data.tool_profile).toBe('frontend'); + }); + + test('returns 400 for invalid tool_profile format (uppercase)', async () => { + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 'FRONTEND' } as any, + makeContext(), + 'req-tp-3', + ); + expect(result.statusCode).toBe(400); + expect(JSON.parse(result.body).error.message).toContain('tool_profile'); + }); + + test('returns 400 for invalid tool_profile format (special chars)', async () => { + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 'front end!' } as any, + makeContext(), + 'req-tp-4', + ); + expect(result.statusCode).toBe(400); + expect(JSON.parse(result.body).error.message).toContain('tool_profile'); + }); + + test('returns 422 when tool_profile does not exist in repo config', async () => { + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + tool_profiles: JSON.stringify({ backend: { capabilityTier: 'default' } }), + }); + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 'nonexistent' }, + makeContext(), + 'req-tp-5', + ); + expect(result.statusCode).toBe(422); + expect(JSON.parse(result.body).error.message).toContain('does not exist'); + }); + + test('tool_profile passes when profile exists in repo config', async () => { + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + tool_profiles: JSON.stringify({ infra: { capabilityTier: 'extended', mcpServers: ['aws-cdk-mcp'] } }), + }); + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Deploy infra', tool_profile: 'infra' }, + makeContext(), + 'req-tp-6', + ); + expect(result.statusCode).toBe(201); + const body = JSON.parse(result.body); + expect(body.data.tool_profile).toBe('infra'); + }); + + test('tool_profile skips existence check when repo has no tool_profiles', async () => { + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + // no tool_profiles field + }); + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 'frontend' }, + makeContext(), + 'req-tp-7', + ); + // When repo has no profiles defined, specifying a profile that doesn't exist should 422 + expect(result.statusCode).toBe(422); + }); + + test('returns 400 for non-string tool_profile', async () => { + const result = await createTaskCore( + { repo: 'org/repo', task_description: 'Fix bug', tool_profile: 123 } as any, + makeContext(), + 'req-tp-8', + ); + expect(result.statusCode).toBe(400); + expect(JSON.parse(result.body).error.message).toContain('tool_profile'); + }); }); diff --git a/cdk/test/handlers/shared/repo-config.test.ts b/cdk/test/handlers/shared/repo-config.test.ts index bc79f8d3..c5085661 100644 --- a/cdk/test/handlers/shared/repo-config.test.ts +++ b/cdk/test/handlers/shared/repo-config.test.ts @@ -30,7 +30,7 @@ jest.mock('../../../src/handlers/shared/logger', () => ({ process.env.REPO_TABLE_NAME = 'RepoConfig'; -import { checkRepoOnboarded, loadRepoConfig } from '../../../src/handlers/shared/repo-config'; +import { checkRepoOnboarded, loadRepoConfig, parseToolProfiles, isValidToolProfile } from '../../../src/handlers/shared/repo-config'; beforeEach(() => { jest.clearAllMocks(); @@ -147,3 +147,49 @@ describe('loadRepoConfig', () => { ); }); }); + +describe('parseToolProfiles', () => { + test('returns empty object for undefined', () => { + expect(parseToolProfiles(undefined)).toEqual({}); + }); + + test('returns empty object for empty string', () => { + expect(parseToolProfiles('')).toEqual({}); + }); + + test('parses valid JSON profiles', () => { + const profiles = { + frontend: { capabilityTier: 'extended', mcpServers: ['eslint-mcp'], skills: ['react-patterns'] }, + backend: { capabilityTier: 'default' }, + }; + expect(parseToolProfiles(JSON.stringify(profiles))).toEqual(profiles); + }); + + test('returns empty object for invalid JSON', () => { + expect(parseToolProfiles('not json')).toEqual({}); + }); + + test('returns empty object for JSON array', () => { + expect(parseToolProfiles('["not", "an", "object"]')).toEqual({}); + }); + + test('returns empty object for JSON null', () => { + expect(parseToolProfiles('null')).toEqual({}); + }); +}); + +describe('isValidToolProfile', () => { + test('returns true for existing profile', () => { + const profiles = { frontend: { capabilityTier: 'extended' as const }, backend: {} }; + expect(isValidToolProfile('frontend', profiles)).toBe(true); + }); + + test('returns false for non-existent profile', () => { + const profiles = { frontend: { capabilityTier: 'extended' as const } }; + expect(isValidToolProfile('backend', profiles)).toBe(false); + }); + + test('returns false for empty profiles map', () => { + expect(isValidToolProfile('any', {})).toBe(false); + }); +}); diff --git a/cdk/test/handlers/shared/validation.test.ts b/cdk/test/handlers/shared/validation.test.ts index 4aecf4e5..a39ae453 100644 --- a/cdk/test/handlers/shared/validation.test.ts +++ b/cdk/test/handlers/shared/validation.test.ts @@ -35,6 +35,7 @@ import { VALID_TASK_TYPES, validateMaxTurns, validatePrNumber, + validateToolProfile, } from '../../../src/handlers/shared/validation'; describe('parseBody', () => { @@ -407,3 +408,59 @@ describe('validatePrNumber', () => { expect(validatePrNumber(true)).toBeNull(); }); }); + +describe('validateToolProfile', () => { + test('returns undefined for absent values', () => { + expect(validateToolProfile(undefined)).toBeUndefined(); + expect(validateToolProfile(null)).toBeUndefined(); + }); + + test('returns the string for valid profile names', () => { + expect(validateToolProfile('frontend')).toBe('frontend'); + expect(validateToolProfile('backend')).toBe('backend'); + expect(validateToolProfile('my-infra')).toBe('my-infra'); + expect(validateToolProfile('a1')).toBe('a1'); + }); + + test('accepts single-char alphanumeric names', () => { + expect(validateToolProfile('a')).toBe('a'); + expect(validateToolProfile('9')).toBe('9'); + }); + + test('returns null for uppercase names', () => { + expect(validateToolProfile('FRONTEND')).toBeNull(); + expect(validateToolProfile('Frontend')).toBeNull(); + }); + + test('returns null for names with special characters', () => { + expect(validateToolProfile('front end')).toBeNull(); + expect(validateToolProfile('front_end')).toBeNull(); + expect(validateToolProfile('front.end')).toBeNull(); + expect(validateToolProfile('front/end')).toBeNull(); + }); + + test('returns null for names starting or ending with hyphen', () => { + expect(validateToolProfile('-frontend')).toBeNull(); + expect(validateToolProfile('frontend-')).toBeNull(); + }); + + test('returns null for empty string', () => { + expect(validateToolProfile('')).toBeNull(); + }); + + test('returns null for names exceeding 64 chars', () => { + expect(validateToolProfile('a'.repeat(65))).toBeNull(); + }); + + test('accepts name at exactly 64 chars', () => { + const name = 'a'.repeat(64); + expect(validateToolProfile(name)).toBe(name); + }); + + test('returns null for non-string types', () => { + expect(validateToolProfile(123)).toBeNull(); + expect(validateToolProfile(true)).toBeNull(); + expect(validateToolProfile({})).toBeNull(); + expect(validateToolProfile([])).toBeNull(); + }); +}); diff --git a/cli/src/commands/submit.ts b/cli/src/commands/submit.ts index 3ebc4e82..8618187b 100644 --- a/cli/src/commands/submit.ts +++ b/cli/src/commands/submit.ts @@ -36,6 +36,7 @@ export function makeSubmitCommand(): Command { .option('--review-pr ', 'PR number to review (sets task_type to pr_review)', parseInt) .option('--idempotency-key ', 'Idempotency key for deduplication') .option('--trace', 'Capture 4 KB debug previews (design §10.1). Opt-in per task; not routine observability.') + .option('--tool-profile ', 'Tool profile to activate (must be defined in repo Blueprint)') .option('--wait', 'Wait for task to complete') .option('--output ', 'Output format (text or json)', 'text') .action(async (opts) => { @@ -64,6 +65,14 @@ export function makeSubmitCommand(): Command { throw new CliError('--max-budget must be a number between 0.01 and 100.'); } } + if (opts.toolProfile !== undefined) { + if (!/^[a-z0-9][a-z0-9-]*[a-z0-9]$/.test(opts.toolProfile) && !/^[a-z0-9]$/.test(opts.toolProfile)) { + throw new CliError('--tool-profile must be lowercase alphanumeric with hyphens (1-64 chars).'); + } + if (opts.toolProfile.length > 64) { + throw new CliError('--tool-profile must be lowercase alphanumeric with hyphens (1-64 chars).'); + } + } const client = new ApiClient(); const body: CreateTaskRequest = { @@ -76,6 +85,7 @@ export function makeSubmitCommand(): Command { ...(opts.pr !== undefined && { task_type: 'pr_iteration' as const, pr_number: opts.pr }), ...(opts.reviewPr !== undefined && { task_type: 'pr_review' as const, pr_number: opts.reviewPr }), ...(opts.trace && { trace: true }), + ...(opts.toolProfile && { tool_profile: opts.toolProfile }), }; const task = await client.createTask(body, opts.idempotencyKey); diff --git a/cli/src/types.ts b/cli/src/types.ts index 628d4562..cf9349b0 100644 --- a/cli/src/types.ts +++ b/cli/src/types.ts @@ -96,6 +96,9 @@ export interface TaskDetail { * the URI in ``status --output json`` lets users / scripts detect * completion without an extra round trip. */ readonly trace_s3_uri: string | null; + /** Tool profile selected at submission, or ``null`` for legacy + * single-tier tasks. Mirrors ``cdk/src/handlers/shared/types.ts``. */ + readonly tool_profile: string | null; } /** Response body of ``GET /v1/tasks/{task_id}/trace`` (design §10.1). */ @@ -165,6 +168,14 @@ export interface CreateTaskRequest { * ``bgagent watch`` / notifications. */ readonly trace?: boolean; + /** + * Named tool profile to activate for this task. Must reference a profile + * defined in the repo's Blueprint.toolProfiles. When omitted, the repo's + * legacy single-tier behavior applies. + * + * Mirrors ``cdk/src/handlers/shared/types.ts::CreateTaskRequest``. + */ + readonly tool_profile?: string; } /** From ce35114dc1a73939b234648a52e83b5ebd66da4d Mon Sep 17 00:00:00 2001 From: bgagent Date: Wed, 13 May 2026 11:56:16 -0700 Subject: [PATCH 2/3] feat(orchestrator): resolve tool profiles and activate MCP servers at runtime The orchestrator now resolves the task's tool_profile against the Blueprint's stored profiles, merging profile cedar policies with base policies and including profile MCP servers and skills in the agent payload. The agent reads these fields, writes profile MCP server entries to .mcp.json (convention-based URL via env vars), and logs skills for future activation. Co-Authored-By: Claude Opus 4.6 --- agent/src/channel_mcp.py | 64 ++++++++++++++ agent/src/pipeline.py | 13 +++ agent/src/server.py | 12 +++ agent/tests/test_channel_mcp.py | 53 ++++++++++++ cdk/src/handlers/shared/orchestrator.ts | 66 ++++++++++++++- cdk/test/handlers/orchestrate-task.test.ts | 97 ++++++++++++++++++++++ 6 files changed, 303 insertions(+), 2 deletions(-) diff --git a/agent/src/channel_mcp.py b/agent/src/channel_mcp.py index f9c51c03..08e9b02f 100644 --- a/agent/src/channel_mcp.py +++ b/agent/src/channel_mcp.py @@ -110,3 +110,67 @@ def configure_channel_mcp(repo_dir: str, channel_source: str) -> bool: f"Linear MCP configured at {mcp_path} (server key: {LINEAR_MCP_SERVER_KEY})", ) return True + + +def configure_profile_mcp(repo_dir: str, mcp_servers: list[str]) -> bool: + """Write tool-profile MCP server entries into ``.mcp.json``. + + Each entry in ``mcp_servers`` is a server identifier (e.g. ``"eslint-mcp"``). + The server entries are written as Streamable HTTP stubs — the actual URL + resolution is expected to be handled by the MCP server registry or + environment-based configuration. + + For v1, profile MCP server entries use a convention-based URL pattern: + ``MCP_SERVER__URL`` env var, falling back to a placeholder. + The Claude Agent SDK will attempt to connect; if the URL is invalid or + unreachable the server simply won't contribute tools. + + Args: + repo_dir: the cloned-repo working directory. + mcp_servers: list of MCP server identifiers from the tool profile. + + Returns: + True if at least one server entry was written, False otherwise. + """ + if not mcp_servers: + return False + + if not repo_dir or not os.path.isdir(repo_dir): + log("WARN", f"configure_profile_mcp: repo_dir missing or not a directory: {repo_dir!r}") + return False + + mcp_path = os.path.join(repo_dir, ".mcp.json") + config = _read_existing_mcp_config(mcp_path) + + servers = config.get("mcpServers") + if not isinstance(servers, dict): + servers = {} + + for server_name in mcp_servers: + # Convention: look for MCP_SERVER__URL env var for the endpoint + env_key = f"MCP_SERVER_{server_name.upper().replace('-', '_')}_URL" + url = os.environ.get(env_key, "") + if not url: + log("WARN", f"No URL found for profile MCP server '{server_name}' (checked ${env_key}), skipping") + continue + servers[server_name] = { + "type": "http", + "url": url, + } + + config["mcpServers"] = servers + + try: + with open(mcp_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + f.write("\n") + except OSError as e: + log("ERROR", f"Failed to write profile MCP config to {mcp_path}: {e}") + return False + + written = [s for s in mcp_servers if s in servers] + log( + "TASK", + f"Profile MCP servers configured at {mcp_path}: {written}", + ) + return len(written) > 0 diff --git a/agent/src/pipeline.py b/agent/src/pipeline.py index 9a20afe9..7d8e5003 100644 --- a/agent/src/pipeline.py +++ b/agent/src/pipeline.py @@ -244,6 +244,9 @@ def run_task( branch_name: str = "", pr_number: str = "", cedar_policies: list[str] | None = None, + tool_profile: str = "", + profile_mcp_servers: list[str] | None = None, + profile_skills: list[str] | None = None, channel_source: str = "", channel_metadata: dict[str, str] | None = None, trace: bool = False, @@ -282,6 +285,7 @@ def run_task( channel_metadata=channel_metadata, trace=trace, user_id=user_id, + tool_profile=tool_profile, ) # Inject Cedar policies into config for the PolicyEngine in runner.py @@ -418,6 +422,15 @@ def _on_trace_truncated(max_bytes: int, first_dropped: int) -> None: resolve_linear_api_token() configure_channel_mcp(setup.repo_dir, config.channel_source) + # Tool-profile MCP wiring. Write profile MCP server entries into + # .mcp.json so the SDK picks them up via setting_sources=["project"]. + if profile_mcp_servers: + from channel_mcp import configure_profile_mcp + + configure_profile_mcp(setup.repo_dir, profile_mcp_servers) + if profile_skills: + log("TASK", f"Tool profile skills: {profile_skills}") + # 👀 on the Linear issue — acknowledges the task is picked up. # No-op for non-Linear tasks. Best-effort; failures are logged # but do not block the pipeline. Capture the reaction id so we diff --git a/agent/src/server.py b/agent/src/server.py index aa547a2d..742ccca3 100644 --- a/agent/src/server.py +++ b/agent/src/server.py @@ -273,6 +273,9 @@ def _run_task_background( branch_name: str = "", pr_number: str = "", cedar_policies: list[str] | None = None, + tool_profile: str = "", + profile_mcp_servers: list[str] | None = None, + profile_skills: list[str] | None = None, channel_source: str = "", channel_metadata: dict[str, str] | None = None, trace: bool = False, @@ -322,6 +325,9 @@ def _run_task_background( branch_name=branch_name, pr_number=pr_number, cedar_policies=cedar_policies, + tool_profile=tool_profile, + profile_mcp_servers=profile_mcp_servers, + profile_skills=profile_skills, channel_source=channel_source, channel_metadata=channel_metadata, trace=trace, @@ -371,6 +377,9 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: branch_name = inp.get("branch_name", "") pr_number = str(inp.get("pr_number", "")) cedar_policies = inp.get("cedar_policies") or [] + tool_profile = inp.get("tool_profile", "") or "" + profile_mcp_servers = inp.get("profile_mcp_servers") or [] + profile_skills = inp.get("profile_skills") or [] channel_source = inp.get("channel_source", "") or "" channel_metadata = inp.get("channel_metadata") or {} # ``trace`` is strictly opt-in (design §10.1). Accept only real @@ -418,6 +427,9 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: "branch_name": branch_name, "pr_number": pr_number, "cedar_policies": cedar_policies, + "tool_profile": tool_profile, + "profile_mcp_servers": profile_mcp_servers, + "profile_skills": profile_skills, "channel_source": channel_source, "channel_metadata": channel_metadata, "trace": trace, diff --git a/agent/tests/test_channel_mcp.py b/agent/tests/test_channel_mcp.py index 9ef4c221..7aad2f42 100644 --- a/agent/tests/test_channel_mcp.py +++ b/agent/tests/test_channel_mcp.py @@ -10,6 +10,7 @@ LINEAR_MCP_SERVER_KEY, LINEAR_MCP_URL, configure_channel_mcp, + configure_profile_mcp, ) @@ -139,3 +140,55 @@ def test_missing_repo_dir(self, tmp_path): def test_empty_repo_dir_string(self): wrote = configure_channel_mcp("", "linear") assert wrote is False + + +class TestConfigureProfileMcp: + """Tests for configure_profile_mcp — writing profile MCP server entries.""" + + def test_no_op_for_empty_servers(self, tmp_path): + wrote = configure_profile_mcp(str(tmp_path), []) + assert wrote is False + + def test_no_op_for_missing_repo_dir(self, tmp_path): + wrote = configure_profile_mcp(str(tmp_path / "nope"), ["eslint-mcp"]) + assert wrote is False + + def test_writes_server_entry_when_env_var_present(self, tmp_path, monkeypatch): + monkeypatch.setenv("MCP_SERVER_ESLINT_MCP_URL", "https://eslint.example/mcp") + wrote = configure_profile_mcp(str(tmp_path), ["eslint-mcp"]) + assert wrote is True + config = _read_mcp(str(tmp_path)) + assert "eslint-mcp" in config["mcpServers"] + assert config["mcpServers"]["eslint-mcp"]["url"] == "https://eslint.example/mcp" + assert config["mcpServers"]["eslint-mcp"]["type"] == "http" + + def test_skips_server_without_env_var(self, tmp_path, monkeypatch): + monkeypatch.delenv("MCP_SERVER_MYSERVER_URL", raising=False) + wrote = configure_profile_mcp(str(tmp_path), ["myserver"]) + assert wrote is False + + def test_multiple_servers_partial_env(self, tmp_path, monkeypatch): + monkeypatch.setenv("MCP_SERVER_SERVER_A_URL", "https://a.example/mcp") + monkeypatch.delenv("MCP_SERVER_SERVER_B_URL", raising=False) + wrote = configure_profile_mcp(str(tmp_path), ["server-a", "server-b"]) + assert wrote is True + config = _read_mcp(str(tmp_path)) + assert "server-a" in config["mcpServers"] + assert "server-b" not in config["mcpServers"] + + def test_merges_with_existing_mcp_config(self, tmp_path, monkeypatch): + monkeypatch.setenv("MCP_SERVER_NEW_MCP_URL", "https://new.example/mcp") + existing = {"mcpServers": {"existing-server": {"type": "http", "url": "https://existing.example"}}} + (tmp_path / ".mcp.json").write_text(json.dumps(existing)) + configure_profile_mcp(str(tmp_path), ["new-mcp"]) + config = _read_mcp(str(tmp_path)) + assert "existing-server" in config["mcpServers"] + assert "new-mcp" in config["mcpServers"] + + def test_coexists_with_channel_mcp(self, tmp_path, monkeypatch): + monkeypatch.setenv("MCP_SERVER_ESLINT_MCP_URL", "https://eslint.example/mcp") + configure_channel_mcp(str(tmp_path), "linear") + configure_profile_mcp(str(tmp_path), ["eslint-mcp"]) + config = _read_mcp(str(tmp_path)) + assert LINEAR_MCP_SERVER_KEY in config["mcpServers"] + assert "eslint-mcp" in config["mcpServers"] diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts index f3c21467..0543da25 100644 --- a/cdk/src/handlers/shared/orchestrator.ts +++ b/cdk/src/handlers/shared/orchestrator.ts @@ -25,7 +25,7 @@ import { logger } from './logger'; import { writeMinimalEpisode } from './memory'; import { coerceNumericOrNull } from './numeric'; import { computePromptVersion } from './prompt-version'; -import { loadRepoConfig, type BlueprintConfig, type ComputeType } from './repo-config'; +import { loadRepoConfig, parseToolProfiles, type BlueprintConfig, type ComputeType } from './repo-config'; import type { TaskRecord } from './types'; import { computeTtlEpoch, DEFAULT_MAX_TURNS } from './validation'; import { TaskStatus, TERMINAL_STATUSES, VALID_TRANSITIONS, type TaskStatusType } from '../../constructs/task-status'; @@ -241,9 +241,70 @@ export async function loadBlueprintConfig(task: TaskRecord): Promise { + const basePolicies = blueprintConfig?.cedar_policies ?? []; + const profileName = task.tool_profile; + let profilePolicies: readonly string[] = []; + + if (profileName && blueprintConfig?.tool_profiles) { + const profile = blueprintConfig.tool_profiles[profileName]; + if (profile?.cedarPolicies && profile.cedarPolicies.length > 0) { + profilePolicies = profile.cedarPolicies; + } + } + + const merged = [...basePolicies, ...profilePolicies]; + if (merged.length === 0) return {}; + return { cedar_policies: merged }; +} + +/** + * Build the tool profile payload fields (mcp_servers, skills, tool_profile name) + * resolved from the task's selected profile. + */ +function buildToolProfilePayload(task: TaskRecord, blueprintConfig?: BlueprintConfig): Record { + const profileName = task.tool_profile; + if (!profileName) return {}; + + const result: Record = { tool_profile: profileName }; + + if (!blueprintConfig?.tool_profiles) return result; + + const profile = blueprintConfig.tool_profiles[profileName]; + if (!profile) { + logger.warn('Task references unknown tool_profile — passing name only', { + task_id: task.task_id, + tool_profile: profileName, + }); + return result; + } + + if (profile.mcpServers && profile.mcpServers.length > 0) { + result.profile_mcp_servers = [...profile.mcpServers]; + } + if (profile.skills && profile.skills.length > 0) { + result.profile_skills = [...profile.skills]; + } + + logger.info('Resolved tool profile for payload', { + task_id: task.task_id, + tool_profile: profileName, + mcp_servers_count: profile.mcpServers?.length ?? 0, + skills_count: profile.skills?.length ?? 0, + cedar_policies_count: profile.cedarPolicies?.length ?? 0, + }); + + return result; +} + /** * Transition task to HYDRATING and assemble the invocation payload. * @param task - the task record. @@ -347,7 +408,8 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B ...(task.trace === true && { trace: true }), ...(blueprintConfig?.model_id && { model_id: blueprintConfig.model_id }), ...(blueprintConfig?.system_prompt_overrides && { system_prompt_overrides: blueprintConfig.system_prompt_overrides }), - ...(blueprintConfig?.cedar_policies && blueprintConfig.cedar_policies.length > 0 && { cedar_policies: blueprintConfig.cedar_policies }), + ...(buildCedarPoliciesPayload(task, blueprintConfig)), + ...(buildToolProfilePayload(task, blueprintConfig)), prompt_version: promptVersion, ...(MEMORY_ID && { memory_id: MEMORY_ID }), hydrated_context: hydratedContext, diff --git a/cdk/test/handlers/orchestrate-task.test.ts b/cdk/test/handlers/orchestrate-task.test.ts index 0151fe74..2c9d970f 100644 --- a/cdk/test/handlers/orchestrate-task.test.ts +++ b/cdk/test/handlers/orchestrate-task.test.ts @@ -53,6 +53,14 @@ const mockLoadRepoConfig = jest.fn(); jest.mock('../../src/handlers/shared/repo-config', () => ({ loadRepoConfig: mockLoadRepoConfig, checkRepoOnboarded: jest.fn(), + parseToolProfiles: jest.fn((raw: string | undefined) => { + if (!raw) return {}; + try { + const parsed = JSON.parse(raw); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) return {}; + return parsed; + } catch { return {}; } + }), })); let ulidCounter = 0; @@ -506,6 +514,30 @@ describe('loadBlueprintConfig', () => { const config = await loadBlueprintConfig(baseTask as any); expect(config.cedar_policies).toBeUndefined(); }); + + test('parses tool_profiles from repo config JSON string', async () => { + const profiles = { frontend: { mcpServers: ['eslint-mcp'], capabilityTier: 'extended' as const } }; + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + onboarded_at: '2024-01-01T00:00:00Z', + updated_at: '2024-01-01T00:00:00Z', + tool_profiles: JSON.stringify(profiles), + }); + const config = await loadBlueprintConfig(baseTask as any); + expect(config.tool_profiles).toEqual(profiles); + }); + + test('returns empty tool_profiles when repo config has none', async () => { + mockLoadRepoConfig.mockResolvedValueOnce({ + repo: 'org/repo', + status: 'active', + onboarded_at: '2024-01-01T00:00:00Z', + updated_at: '2024-01-01T00:00:00Z', + }); + const config = await loadBlueprintConfig(baseTask as any); + expect(config.tool_profiles).toEqual({}); + }); }); describe('hydrateAndTransition with blueprint config', () => { @@ -619,6 +651,71 @@ describe('hydrateAndTransition with blueprint config', () => { }); expect(payload.cedar_policies).toBeUndefined(); }); + + test('includes tool_profile and resolved profile fields in payload', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + const taskWithProfile = { ...baseTask, tool_profile: 'frontend' }; + const payload = await hydrateAndTransition(taskWithProfile as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + tool_profiles: { + frontend: { + mcpServers: ['eslint-mcp'], + skills: ['react-patterns'], + cedarPolicies: ['permit (principal, action, resource);'], + }, + }, + }); + expect(payload.tool_profile).toBe('frontend'); + expect(payload.profile_mcp_servers).toEqual(['eslint-mcp']); + expect(payload.profile_skills).toEqual(['react-patterns']); + expect(payload.cedar_policies).toEqual(['permit (principal, action, resource);']); + }); + + test('merges blueprint and profile cedar policies', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + const taskWithProfile = { ...baseTask, tool_profile: 'backend' }; + const payload = await hydrateAndTransition(taskWithProfile as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + cedar_policies: ['base-policy'], + tool_profiles: { + backend: { cedarPolicies: ['profile-policy'] }, + }, + }); + expect(payload.cedar_policies).toEqual(['base-policy', 'profile-policy']); + }); + + test('omits profile fields when no tool_profile on task', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + const payload = await hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + tool_profiles: { + frontend: { mcpServers: ['eslint-mcp'] }, + }, + }); + expect(payload.tool_profile).toBeUndefined(); + expect(payload.profile_mcp_servers).toBeUndefined(); + expect(payload.profile_skills).toBeUndefined(); + }); + + test('passes tool_profile name even for unknown profile', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + const taskWithProfile = { ...baseTask, tool_profile: 'nonexistent' }; + const payload = await hydrateAndTransition(taskWithProfile as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + tool_profiles: { frontend: { mcpServers: ['eslint-mcp'] } }, + }); + expect(payload.tool_profile).toBe('nonexistent'); + expect(payload.profile_mcp_servers).toBeUndefined(); + expect(payload.profile_skills).toBeUndefined(); + }); }); describe('finalizeTask', () => { From 4f7665c3364235ee0c6afea347cbef96f2b518ec Mon Sep 17 00:00:00 2001 From: bgagent Date: Wed, 13 May 2026 12:12:28 -0700 Subject: [PATCH 3/3] test: add unit tests for tool profile payload extraction and pipeline activation Adds server.py tests for tool_profile/profile_mcp_servers/profile_skills extraction from orchestrator payload, and pipeline integration tests verifying configure_profile_mcp is called only when profile_mcp_servers is non-empty. Co-Authored-By: Claude Opus 4.6 --- agent/tests/test_pipeline.py | 123 +++++++++++++++++++++++++++++++++++ agent/tests/test_server.py | 57 ++++++++++++++++ 2 files changed, 180 insertions(+) diff --git a/agent/tests/test_pipeline.py b/agent/tests/test_pipeline.py index ff5c49c0..9d7f54b6 100644 --- a/agent/tests/test_pipeline.py +++ b/agent/tests/test_pipeline.py @@ -143,6 +143,129 @@ async def fake_run_agent(_prompt, _system_prompt, config, cwd=None, trajectory=N assert captured_config.cedar_policies == [] +class TestProfileMcpActivation: + @patch("pipeline.run_agent") + @patch("pipeline.build_system_prompt") + @patch("pipeline.discover_project_config") + @patch("repo.setup_repo") + @patch("pipeline.task_span") + @patch("pipeline.task_state") + def test_configure_profile_mcp_called_when_servers_provided( + self, + _mock_task_state, + mock_task_span, + mock_setup_repo, + _mock_discover, + _mock_build_prompt, + mock_run_agent, + monkeypatch, + ): + """When profile_mcp_servers is non-empty, configure_profile_mcp is called.""" + monkeypatch.setenv("GITHUB_TOKEN", "ghp_test") + monkeypatch.setenv("AWS_REGION", "us-east-1") + + mock_setup_repo.return_value = RepoSetup( + repo_dir="/workspace/repo", + branch="bgagent/test/branch", + build_before=True, + ) + + async def fake_run_agent(_prompt, _system_prompt, config, cwd=None, trajectory=None): + return AgentResult(status="success", turns=1, cost_usd=0.01, num_turns=1) + + mock_run_agent.side_effect = fake_run_agent + + mock_span = MagicMock() + mock_span.__enter__ = MagicMock(return_value=mock_span) + mock_span.__exit__ = MagicMock(return_value=False) + mock_task_span.return_value = mock_span + + mock_configure_profile = MagicMock(return_value=True) + + with ( + patch("pipeline.ensure_committed", return_value=False), + patch("pipeline.verify_build", return_value=True), + patch("pipeline.verify_lint", return_value=True), + patch("pipeline.ensure_pr", return_value="https://github.com/org/repo/pull/1"), + patch("pipeline.get_disk_usage", return_value=0), + patch("pipeline.print_metrics"), + patch("channel_mcp.configure_profile_mcp", mock_configure_profile), + ): + from pipeline import run_task + + run_task( + repo_url="owner/repo", + task_description="fix bug", + github_token="ghp_test", + aws_region="us-east-1", + task_id="test-id", + profile_mcp_servers=["eslint-mcp", "prettier-mcp"], + ) + + mock_configure_profile.assert_called_once_with( + "/workspace/repo", ["eslint-mcp", "prettier-mcp"] + ) + + @patch("pipeline.run_agent") + @patch("pipeline.build_system_prompt") + @patch("pipeline.discover_project_config") + @patch("repo.setup_repo") + @patch("pipeline.task_span") + @patch("pipeline.task_state") + def test_configure_profile_mcp_not_called_when_empty( + self, + _mock_task_state, + mock_task_span, + mock_setup_repo, + _mock_discover, + _mock_build_prompt, + mock_run_agent, + monkeypatch, + ): + """When profile_mcp_servers is empty, configure_profile_mcp is NOT called.""" + monkeypatch.setenv("GITHUB_TOKEN", "ghp_test") + monkeypatch.setenv("AWS_REGION", "us-east-1") + + mock_setup_repo.return_value = RepoSetup( + repo_dir="/workspace/repo", + branch="bgagent/test/branch", + build_before=True, + ) + + async def fake_run_agent(_prompt, _system_prompt, config, cwd=None, trajectory=None): + return AgentResult(status="success", turns=1, cost_usd=0.01, num_turns=1) + + mock_run_agent.side_effect = fake_run_agent + + mock_span = MagicMock() + mock_span.__enter__ = MagicMock(return_value=mock_span) + mock_span.__exit__ = MagicMock(return_value=False) + mock_task_span.return_value = mock_span + + mock_configure_profile = MagicMock(return_value=False) + + with ( + patch("pipeline.ensure_committed", return_value=False), + patch("pipeline.verify_build", return_value=True), + patch("pipeline.verify_lint", return_value=True), + patch("pipeline.ensure_pr", return_value="https://github.com/org/repo/pull/1"), + patch("pipeline.get_disk_usage", return_value=0), + patch("pipeline.print_metrics"), + patch("channel_mcp.configure_profile_mcp", mock_configure_profile), + ): + from pipeline import run_task + + run_task( + repo_url="owner/repo", + task_description="fix bug", + github_token="ghp_test", + aws_region="us-east-1", + task_id="test-id", + ) + + mock_configure_profile.assert_not_called() + + class TestChainPriorAgentError: def test_none_agent_result_returns_exception_only(self): exc = RuntimeError("post-hook crash") diff --git a/agent/tests/test_server.py b/agent/tests/test_server.py index 619702ee..351290ff 100644 --- a/agent/tests/test_server.py +++ b/agent/tests/test_server.py @@ -511,3 +511,60 @@ def test_user_id_non_string_logs_warn(self, capsys): assert "user_id payload field is not a string" in captured.out assert "type=int" in captured.out assert "'t-warn'" in captured.out + + +class TestExtractToolProfile: + """Tests for tool_profile, profile_mcp_servers, profile_skills extraction.""" + + def _fake_req(self) -> Any: + return _FakeRequest() + + def _base_payload(self, **overrides) -> dict: + base: dict[str, Any] = { + "task_id": "t-profile", + "repo_url": "o/r", + "prompt": "fix", + "github_token": "ghp_x", + "aws_region": "us-east-1", + } + base.update(overrides) + return base + + def test_tool_profile_defaults_to_empty(self): + params = server._extract_invocation_params( + self._base_payload(), + self._fake_req(), + ) + assert params["tool_profile"] == "" + assert params["profile_mcp_servers"] == [] + assert params["profile_skills"] == [] + + def test_tool_profile_extracted(self): + params = server._extract_invocation_params( + self._base_payload(tool_profile="frontend"), + self._fake_req(), + ) + assert params["tool_profile"] == "frontend" + + def test_profile_mcp_servers_extracted(self): + params = server._extract_invocation_params( + self._base_payload(profile_mcp_servers=["eslint-mcp", "prettier-mcp"]), + self._fake_req(), + ) + assert params["profile_mcp_servers"] == ["eslint-mcp", "prettier-mcp"] + + def test_profile_skills_extracted(self): + params = server._extract_invocation_params( + self._base_payload(profile_skills=["react-patterns"]), + self._fake_req(), + ) + assert params["profile_skills"] == ["react-patterns"] + + def test_null_profile_fields_default_to_empty(self): + params = server._extract_invocation_params( + self._base_payload(tool_profile=None, profile_mcp_servers=None, profile_skills=None), + self._fake_req(), + ) + assert params["tool_profile"] == "" + assert params["profile_mcp_servers"] == [] + assert params["profile_skills"] == []