diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxModels.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxModels.ts index a71f41139a..8b0c89563a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxModels.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxModels.ts @@ -1,19 +1,42 @@ import { ExecuteCommandParams } from 'vscode-languageserver' -import { TransformationPlan } from '@amzn/codewhisperer-runtime' - -// ATX Job Status enum (matches client-side C# definition) -export enum AtxJobStatus { - CREATED = 'CREATED', - STARTING = 'STARTING', - ASSESSING = 'ASSESSING', - PLANNING = 'PLANNING', - PLANNED = 'PLANNED', - EXECUTING = 'EXECUTING', - AWAITING_HUMAN_INPUT = 'AWAITING_HUMAN_INPUT', - COMPLETED = 'COMPLETED', - FAILED = 'FAILED', - STOPPING = 'STOPPING', - STOPPED = 'STOPPED', +import { PlanStepStatus } from '@amazon/elastic-gumby-frontend-client' + +// Re-export for convenience +export { PlanStepStatus } + +/** + * Represents a step in an ATX transformation plan tree structure. + * Matches C# AtxPlanStep class. + */ +export interface AtxPlanStep { + StepId: string + ParentStepId: string | null + StepName: string + Description: string + Status: PlanStepStatus + Children: AtxPlanStep[] +} + +/** + * Tree structure for transformation plan. + * Matches C# AtxTransformationPlan class. + */ +export interface AtxTransformationPlan { + Root: AtxPlanStep +} + +/** + * Creates an empty root node for the transformation plan tree. + */ +export function createEmptyRootNode(): AtxPlanStep { + return { + StepId: 'root', + ParentStepId: null, + StepName: 'Root', + Description: '', + Status: 'NOT_STARTED', + Children: [], + } } // ATX Workspace Models @@ -28,11 +51,11 @@ export interface AtxCreatedWorkspaceInfo { WorkspaceName: string } -// ATX Transformation Job (matches client-side C# definition) +// ATX Transformation Job export interface AtxTransformationJob { WorkspaceId: string JobId: string - Status: AtxJobStatus + Status: string FailureReason?: string } @@ -50,6 +73,7 @@ export interface AtxListOrCreateWorkspaceResponse { export interface AtxStartTransformRequest extends ExecuteCommandParams { WorkspaceId: string JobName?: string + InteractiveMode?: boolean StartTransformRequest: object // Original RTS-style request for ZIP creation } @@ -71,9 +95,22 @@ export interface AtxGetTransformInfoResponse { PlanPath?: string | null ReportPath?: string | null WorklogPath?: string | null - TransformationPlan?: TransformationPlan | null + TransformationPlan?: AtxTransformationPlan | null ArtifactPath?: string | null ErrorString?: string | null + StepInformation?: AtxStepInformation | null +} + +/** + * Information about a step during execution phase HITL. + */ +export interface AtxStepInformation { + StepId: string + DiffArtifactPath: string + RetryInstruction?: string + IsInvalid?: boolean + InvalidInstruction?: string + InvalidReason?: string } // ATX Stop Job request @@ -93,3 +130,31 @@ export interface AtxUploadPlanResponse { PlanPath?: string ReportPath?: string } + +// ATX Set Checkpoints request/response (interactive mode) +export interface AtxSetCheckpointsRequest extends ExecuteCommandParams { + TransformationJobId: string + WorkspaceId: string + SolutionRootPath: string + Checkpoints: Record +} + +export interface AtxSetCheckpointsResponse { + Success: boolean + Error?: string +} + +// ATX Checkpoint Action request/response +export interface AtxCheckpointActionRequest extends ExecuteCommandParams { + Action: string // "APPLY" or "RETRY" + NewInstruction?: string // Only used when Action is "RETRY" + StepId: string + TransformationJobId: string + WorkspaceId: string + SolutionRootPath: string +} + +export interface AtxCheckpointActionResponse { + Success: boolean + Error?: string | null +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxNetTransformServer.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxNetTransformServer.ts index 76f6073412..2c058830be 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxNetTransformServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxNetTransformServer.ts @@ -12,6 +12,8 @@ import { AtxGetTransformInfoRequest, AtxStopJobRequest, AtxUploadPlanRequest, + AtxSetCheckpointsRequest, + AtxCheckpointActionRequest, } from './atxModels' // ATX FES Commands - Consolidated APIs @@ -20,6 +22,8 @@ const AtxStartTransformCommand = 'aws/atxTransform/startTransform' const AtxGetTransformInfoCommand = 'aws/atxTransform/getTransformInfo' const AtxStopJobCommand = 'aws/atxTransform/stopJob' const AtxUploadPlanCommand = 'aws/atxTransform/uploadPlan' +const AtxSetCheckpointsCommand = 'aws/atxTransform/setCheckpoints' +const AtxCheckpointActionCommand = 'aws/atxTransform/checkpointAction' export const AtxNetTransformServerToken = (): Server => @@ -36,7 +40,8 @@ export const AtxNetTransformServerToken = return result } case AtxStartTransformCommand: { - const { WorkspaceId, JobName, StartTransformRequest } = params as AtxStartTransformRequest + const { WorkspaceId, JobName, InteractiveMode, StartTransformRequest } = + params as AtxStartTransformRequest if (!WorkspaceId) { throw new Error('WorkspaceId is required for startTransform') @@ -45,6 +50,7 @@ export const AtxNetTransformServerToken = const result = await atxTransformHandler.startTransform({ workspaceId: WorkspaceId, jobName: JobName, + interactiveMode: InteractiveMode, startTransformRequest: StartTransformRequest, }) @@ -79,6 +85,30 @@ export const AtxNetTransformServerToken = const result = await atxTransformHandler.stopJob(WorkspaceId, JobId) return { Status: result } } + case AtxSetCheckpointsCommand: { + const { WorkspaceId, TransformationJobId, SolutionRootPath, Checkpoints } = + params as AtxSetCheckpointsRequest + + return await atxTransformHandler.setCheckpoints( + WorkspaceId, + TransformationJobId, + SolutionRootPath, + Checkpoints || {} + ) + } + case AtxCheckpointActionCommand: { + const { WorkspaceId, TransformationJobId, StepId, Action, NewInstruction, SolutionRootPath } = + params as AtxCheckpointActionRequest + + return await atxTransformHandler.checkpointAction( + WorkspaceId, + TransformationJobId, + StepId, + Action, + SolutionRootPath, + NewInstruction + ) + } default: { throw new Error(`Unknown ATX FES command: ${params.command}`) } @@ -106,6 +136,8 @@ export const AtxNetTransformServerToken = AtxGetTransformInfoCommand, AtxUploadPlanCommand, AtxStopJobCommand, + AtxSetCheckpointsCommand, + AtxCheckpointActionCommand, ], }, }, diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxTransformHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxTransformHandler.ts index 7408b67197..a86455205a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxTransformHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/atxTransformHandler.ts @@ -15,6 +15,8 @@ import { GetHitlTaskCommand, ListHitlTasksCommand, SubmitCriticalHitlTaskCommand, + SubmitStandardHitlTaskCommand, + UpdateHitlTaskCommand, GetJobCommand, ListJobPlanStepsCommand, ListWorklogsCommand, @@ -36,10 +38,16 @@ import { AtxTransformationJob, AtxUploadPlanRequest, AtxUploadPlanResponse, + AtxSetCheckpointsResponse, + AtxTransformationPlan, + AtxPlanStep, + PlanStepStatus, + createEmptyRootNode, + AtxStepInformation, + AtxCheckpointActionResponse, } from './atxModels' import { v4 as uuidv4 } from 'uuid' import { request } from 'http' -import { TransformationPlan } from '@amzn/codewhisperer-runtime' import { Utils, workspaceFolderName } from './utils' @@ -54,6 +62,7 @@ export class ATXTransformHandler { private runtime: Runtime private atxClient: ElasticGumbyFrontendClient | null = null private cachedHitl: string | null = null + private cachedStepHitl: string | null = null constructor(serviceManager: AtxTokenServiceManager, workspace: Workspace, logging: Logging, runtime: Runtime) { this.serviceManager = serviceManager @@ -339,6 +348,7 @@ export class ATXTransformHandler { workspaceId: string jobName?: string targetFramework?: string + interactiveMode?: boolean }): Promise<{ jobId: string; status: string } | null> { try { this.logging.log(`ATX: Starting CreateJob for workspace: ${request.workspaceId}`) @@ -348,9 +358,15 @@ export class ATXTransformHandler { throw new Error('ATX FES client not initialized') } + // Build objective object with target_framework and optionally interactive_mode + const objective: any = { + target_framework: request.targetFramework || 'net10.0', + interactive_mode: request.interactiveMode || false, + } + const command = new CreateJobCommand({ workspaceId: request.workspaceId, - objective: JSON.stringify({ target_framework: request.targetFramework || 'net10.0' }), + objective: JSON.stringify(objective), jobType: 'DOTNET_IDE' as any, jobName: request.jobName || `transform-job-${Date.now()}`, intent: 'LANGUAGE_UPGRADE', @@ -522,6 +538,7 @@ export class ATXTransformHandler { async startTransform(request: { workspaceId: string jobName?: string + interactiveMode?: boolean startTransformRequest: object }): Promise<{ TransformationJobId: string; ArtifactPath: string; UploadId: string } | null> { try { @@ -532,6 +549,7 @@ export class ATXTransformHandler { workspaceId: request.workspaceId, jobName: request.jobName || 'Transform Job', targetFramework: (request.startTransformRequest as any).TargetFramework, + interactiveMode: request.interactiveMode, }) if (!createJobResponse?.jobId) { @@ -735,6 +753,35 @@ export class ATXTransformHandler { } } + async updateHitl(workspaceId: string, jobId: string, taskId: string, humanArtifactId: string): Promise { + try { + this.logging.log(`ATX: Starting UpdateHitl for task: ${taskId}`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + this.logging.error('ATX: Failed to initialize client for UpdateHitl') + return null + } + + const command = new UpdateHitlTaskCommand({ + workspaceId: workspaceId, + jobId: jobId, + taskId: taskId, + humanArtifact: { + artifactId: humanArtifactId, + }, + }) + + await this.addAuthToCommand(command) + const result = await this.atxClient!.send(command) + + this.logging.log(`ATX: UpdateHitl completed successfully with status ${result.status}`) + return result + } catch (error) { + this.logging.error(`ATX: UpdateHitl error: ${String(error)}`) + return null + } + } + async getHitl(workspaceId: string, jobId: string, taskId: string): Promise { try { this.logging.log(`ATX: Getting Hitl task: ${taskId}`) @@ -842,7 +889,7 @@ export class ATXTransformHandler { } /** - * Get transform info - dummy implementation + * Get transform info */ async getTransformInfo(request: AtxGetTransformInfoRequest): Promise { try { @@ -914,11 +961,9 @@ export class ATXTransformHandler { TransformationPlan: plan, } as AtxGetTransformInfoResponse } else if (jobStatus === 'AWAITING_HUMAN_INPUT') { - const response = await this.getHitlAgentArtifact( - request.WorkspaceId, - request.TransformationJobId, - request.SolutionRootPath - ) + return await this.handleAwaitingHumanInput(request) + } else { + await this.listWorklogs(request.WorkspaceId, request.TransformationJobId, request.SolutionRootPath) return { TransformationJob: { @@ -926,26 +971,297 @@ export class ATXTransformHandler { JobId: request.TransformationJobId, Status: jobStatus, } as AtxTransformationJob, - PlanPath: response?.PlanPath, - ReportPath: response?.ReportPath, } as AtxGetTransformInfoResponse - } else { - await this.listWorklogs(request.WorkspaceId, request.TransformationJobId, request.SolutionRootPath) + } + } catch (error) { + this.logging.error(`ATX: GetTransformInfo error: ${String(error)}`) + return null + } + } + + /** + * Handles AWAITING_HUMAN_INPUT status. + * Two scenarios: + * 1. Planning phase: No plan exists yet, need to get HITL artifacts (plan.md, report.md) + * 2. Execution phase: Plan exists, HITL raised for a specific step + */ + private async handleAwaitingHumanInput(request: AtxGetTransformInfoRequest): Promise { + // Try to get the transformation plan first + const plan = await this.getTransformationPlan( + request.WorkspaceId, + request.TransformationJobId, + request.SolutionRootPath + ) + + const hasPlan = plan.Root.Children[0].Children.length > 0 + + if (hasPlan) { + // Execution phase: Plan exists, HITL raised during transformation + return await this.handleExecutionPhaseHitl(request, plan) + } else { + // Planning phase: No plan yet, get HITL artifacts for plan review + return await this.handlePlanningPhaseHitl(request) + } + } + + /** + * Handles HITL during planning phase - downloads plan.md and report.md for user review. + */ + private async handlePlanningPhaseHitl(request: AtxGetTransformInfoRequest): Promise { + const response = await this.getHitlAgentArtifact( + request.WorkspaceId, + request.TransformationJobId, + request.SolutionRootPath + ) + + return { + TransformationJob: { + WorkspaceId: request.WorkspaceId, + JobId: request.TransformationJobId, + Status: 'AWAITING_HUMAN_INPUT', + } as AtxTransformationJob, + PlanPath: response?.PlanPath, + ReportPath: response?.ReportPath, + } as AtxGetTransformInfoResponse + } + /** + * Handles HITL during execution phase - plan exists, step-level HITL raised. + */ + private async handleExecutionPhaseHitl( + request: AtxGetTransformInfoRequest, + plan: AtxTransformationPlan + ): Promise { + this.logging.log(`ATX: Execution phase HITL - plan has ${plan.Root.Children.length} steps`) + + try { + // Find the step with PENDING_HUMAN_INPUT status in the plan + const pendingStep = this.findPendingHumanInputStep(plan.Root) + + if (!pendingStep) { + this.logging.log('ATX: No step with PENDING_HUMAN_INPUT status found in plan') return { TransformationJob: { WorkspaceId: request.WorkspaceId, JobId: request.TransformationJobId, - Status: jobStatus, + Status: 'AWAITING_HUMAN_INPUT', } as AtxTransformationJob, + TransformationPlan: plan, } as AtxGetTransformInfoResponse } + + this.logging.log(`ATX: Found pending step: ${pendingStep.StepId}`) + + // Find the step-level HITL using tag: {stepId}-review + const stepHitl = await this.findStepLevelHitl( + request.WorkspaceId, + request.TransformationJobId, + pendingStep.StepId + ) + + if (!stepHitl) { + this.logging.log('ATX: No step-level HITL found, returning plan only') + return { + TransformationJob: { + WorkspaceId: request.WorkspaceId, + JobId: request.TransformationJobId, + Status: 'AWAITING_HUMAN_INPUT', + } as AtxTransformationJob, + TransformationPlan: plan, + } as AtxGetTransformInfoResponse + } + + // Download and parse the agent artifact JSON + const stepInformation = await this.downloadAndParseStepHitlArtifact( + request.WorkspaceId, + request.TransformationJobId, + stepHitl, + pendingStep.StepId, + request.SolutionRootPath + ) + + return { + TransformationJob: { + WorkspaceId: request.WorkspaceId, + JobId: request.TransformationJobId, + Status: 'AWAITING_HUMAN_INPUT', + } as AtxTransformationJob, + TransformationPlan: plan, + StepInformation: stepInformation, + } as AtxGetTransformInfoResponse } catch (error) { - this.logging.error(`ATX: GetTransformInfo error: ${String(error)}`) + this.logging.error(`ATX: handleExecutionPhaseHitl error: ${String(error)}`) + return { + TransformationJob: { + WorkspaceId: request.WorkspaceId, + JobId: request.TransformationJobId, + Status: 'AWAITING_HUMAN_INPUT', + } as AtxTransformationJob, + TransformationPlan: plan, + } as AtxGetTransformInfoResponse + } + } + + /** + * Recursively finds a step with PENDING_HUMAN_INPUT status in the plan tree. + */ + private findPendingHumanInputStep(step: AtxPlanStep): AtxPlanStep | null { + if (step.Status === 'PENDING_HUMAN_INPUT') { + return step + } + + for (const child of step.Children) { + const found = this.findPendingHumanInputStep(child) + if (found) { + return found + } + } + + return null + } + + /** + * Finds the step-level HITL task by filtering for the "{stepId}-review" tag. + */ + private async findStepLevelHitl(workspaceId: string, jobId: string, stepId: string): Promise { + try { + this.logging.log(`ATX: Finding step-level HITL for job: ${jobId}, step: ${stepId}`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + this.logging.error('ATX: Failed to initialize client for findStepLevelHitl') + return null + } + + // List HITLs with "{stepId}-review" tag + const command = new ListHitlTasksCommand({ + workspaceId: workspaceId, + jobId: jobId, + taskType: 'NORMAL', + taskFilter: { + taskStatuses: ['AWAITING_HUMAN_INPUT'], + tag: `${stepId}-review`, + }, + }) + + await this.addAuthToCommand(command) + const result = await this.atxClient!.send(command) + + if (!result.hitlTasks || result.hitlTasks.length === 0) { + this.logging.log(`ATX: No step-level HITL found with tag ${stepId}-review`) + return null + } + + const stepHitl = result.hitlTasks[0] + this.logging.log(`ATX: Found step-level HITL: ${stepHitl.taskId}`) + + // Cache the step HITL task ID for later use in checkpointAction + this.cachedStepHitl = stepHitl.taskId || null + + return stepHitl + } catch (error) { + this.logging.error(`ATX: findStepLevelHitl error: ${String(error)}`) + return null + } + } + + /** + * Downloads and parses the step HITL agent artifact JSON, and extracts the diff artifact. + */ + private async downloadAndParseStepHitlArtifact( + workspaceId: string, + jobId: string, + hitlTask: any, + stepId: string, + solutionRootPath: string + ): Promise { + try { + const taskId = hitlTask.taskId + const agentArtifactId = hitlTask.agentArtifact?.artifactId + + if (!agentArtifactId) { + this.logging.error('ATX: Step HITL has no agent artifact') + return null + } + + this.logging.log(`ATX: Downloading step HITL artifact: ${agentArtifactId}`) + + // Download the agent artifact JSON + const downloadInfo = await this.createArtifactDownloadUrl(workspaceId, jobId, agentArtifactId) + if (!downloadInfo) { + throw new Error('Failed to get download URL for step HITL artifact') + } + + // Download the JSON content + const response = await got.get(downloadInfo.s3PresignedUrl, { + headers: downloadInfo.requestHeaders || {}, + responseType: 'text', + }) + + const artifactJson = JSON.parse(response.body) + this.logging.log(`ATX: Parsed step HITL artifact JSON`) + + // Extract diff artifact if present + let diffArtifactPath = '' + if (artifactJson.diffArtifactId) { + diffArtifactPath = await this.downloadDiffArtifact( + workspaceId, + jobId, + artifactJson.diffArtifactId, + solutionRootPath, + stepId + ) + } + + return { + StepId: stepId, + DiffArtifactPath: diffArtifactPath, + ...(artifactJson.retryInstruction && { RetryInstruction: artifactJson.retryInstruction }), + ...(artifactJson.isInvalid !== undefined && { IsInvalid: artifactJson.isInvalid }), + ...(artifactJson.invalidInstruction && { InvalidInstruction: artifactJson.invalidInstruction }), + ...(artifactJson.invalidReason && { InvalidReason: artifactJson.invalidReason }), + } + } catch (error) { + this.logging.error(`ATX: downloadAndParseStepHitlArtifact error: ${String(error)}`) return null } } + /** + * Downloads and extracts the diff artifact ZIP. + */ + private async downloadDiffArtifact( + workspaceId: string, + jobId: string, + diffArtifactId: string, + solutionRootPath: string, + stepId: string + ): Promise { + try { + this.logging.log(`ATX: Downloading diff artifact: ${diffArtifactId}`) + + const downloadInfo = await this.createArtifactDownloadUrl(workspaceId, jobId, diffArtifactId) + if (!downloadInfo) { + throw new Error('Failed to get download URL for diff artifact') + } + + const pathToDownload = path.join(solutionRootPath, workspaceFolderName, jobId, 'checkpoints', stepId) + + await Utils.downloadAndExtractArchive( + downloadInfo.s3PresignedUrl, + downloadInfo.requestHeaders, + pathToDownload, + `${stepId}.zip`, + this.logging + ) + + this.logging.log(`ATX: Diff artifact extracted to: ${pathToDownload}`) + return pathToDownload + } catch (error) { + this.logging.error(`ATX: downloadDiffArtifact error: ${String(error)}`) + return '' + } + } + async uploadPlan(request: AtxUploadPlanRequest): Promise { this.logging.log('ATX: Starting upload plan') @@ -1110,266 +1426,165 @@ export class ATXTransformHandler { workspaceId: string, jobId: string, solutionRootPath: string - ): Promise { + ): Promise { try { - // Get real plan steps from ATX FES (only if job status >= PLANNED) - const planSteps = await this.getATXFESJobPlanSteps(workspaceId, jobId) - - if (planSteps) { - this.logging.log(`ATX FES: Found ${planSteps.length} transformation steps`) + const plan = await this.fetchPlanTree(workspaceId, jobId) - // Sort steps by score (primary) and startTime (tiebreaker) to match RTS ordering - planSteps.sort((a: any, b: any) => { - const scoreDiff = (a.score || 0) - (b.score || 0) - if (scoreDiff !== 0) return scoreDiff + // Fetch worklogs in parallel (fire and forget, don't block plan return) + this.fetchWorklogs(workspaceId, jobId, solutionRootPath).catch(e => { + this.logging.log(`ATX: Could not get worklogs for workspace: ${workspaceId}, job: ${jobId}`) + }) - // Tiebreaker for identical scores: sort by startTime - const timeA = a.startTime ? new Date(a.startTime).getTime() : 0 - const timeB = b.startTime ? new Date(b.startTime).getTime() : 0 - return timeA - timeB - }) + this.logging.log(`ATX: Successfully built plan tree with ${plan.Root.Children.length} root steps`) + return plan + } catch (error) { + this.logging.error(`ATX: getTransformationPlan error: ${String(error)}`) + return { Root: createEmptyRootNode() } + } + } - // Return in exact same format as RTS with all required fields - const transformationPlan = { - transformationSteps: planSteps.map((step: any, index: number) => { - try { - // Map substeps to ProgressUpdates for IDE display - const progressUpdates = (step.substeps || []).map((substep: any) => { - // Map ATX substep status to IDE TransformationProgressUpdateStatus enum values - let substepStatus = 'IN_PROGRESS' // Default - no NOT_STARTED in this enum - switch (substep.status) { - case 'SUCCEEDED': - case 'COMPLETED': - substepStatus = 'COMPLETED' - break - case 'IN_PROGRESS': - case 'RUNNING': - substepStatus = 'IN_PROGRESS' - break - case 'FAILED': - substepStatus = 'FAILED' - break - case 'SKIPPED': - substepStatus = 'SKIPPED' - break - case 'NOT_STARTED': - case 'CREATED': - default: - substepStatus = 'IN_PROGRESS' // No NOT_STARTED option in ProgressUpdate enum - break - } - - // Map nested progress updates (3rd level) - const nestedProgressUpdates = (substep.substeps || []).map((nestedUpdate: any) => { - let nestedStatus = 'IN_PROGRESS' - switch (nestedUpdate.status) { - case 'SUCCEEDED': - case 'COMPLETED': - nestedStatus = 'COMPLETED' - break - case 'IN_PROGRESS': - case 'RUNNING': - nestedStatus = 'IN_PROGRESS' - break - case 'FAILED': - nestedStatus = 'FAILED' - break - case 'SKIPPED': - nestedStatus = 'SKIPPED' - break - default: - nestedStatus = 'IN_PROGRESS' - break - } - return { - name: nestedUpdate.stepName || 'Unknown Nested Update', - description: nestedUpdate.description || '', - status: nestedStatus, - stepId: nestedUpdate.stepId ?? undefined, - } - }) - - return { - name: substep.stepName || 'Unknown Substep', - description: substep.description || '', - status: substepStatus, - startTime: substep.startTime ? new Date(substep.startTime) : undefined, - endTime: substep.endTime ? new Date(substep.endTime) : undefined, - stepId: substep.stepId ?? undefined, - progressUpdates: nestedProgressUpdates, - } - }) - - // Use ATX status directly - IDE supports most values, minimal mapping needed - let mappedStatus = step.status || 'NOT_STARTED' - // Only map the few values IDE doesn't have - if (mappedStatus === 'SUCCEEDED') { - mappedStatus = 'COMPLETED' - } else if (mappedStatus === 'RUNNING') { - mappedStatus = 'IN_PROGRESS' - } else if (mappedStatus === 'CREATED') { - mappedStatus = 'NOT_STARTED' - } - - // Use ATX step data directly without hardcoded ordering - const stepNumber = index + 1 - const stepName = `Step ${stepNumber} - ${step.stepName || 'Unknown Step'}` - - return { - id: step.stepId || `step-${stepNumber}`, - name: stepName, - description: step.description || '', - status: mappedStatus, - progressUpdates: progressUpdates, - startTime: step.startTime ? new Date(step.startTime) : undefined, - endTime: step.endTime ? new Date(step.endTime) : undefined, - } - } catch (error) { - this.logging.error(`ATX FES: Error mapping step ${index}: ${String(error)}`) - // Return a safe fallback step - const stepNumber = index + 1 - return { - id: step.stepId || `fallback-${stepNumber}`, - name: `Step ${stepNumber} - ${step.stepName || `Step ${stepNumber}`}`, - description: step.description || '', - status: 'NOT_STARTED', - progressUpdates: [], - startTime: undefined, - endTime: undefined, - } - } - }), - } as TransformationPlan - try { - await this.listWorklogs(workspaceId, jobId, solutionRootPath) - } catch (e) { - this.logging.log(`ATX FES: Could not get worklog for workspaces: ${workspaceId}, job id: ${jobId}`) - } + /** + * Fetches all plan steps in a single API call and builds the tree locally. + */ + private async fetchPlanTree(workspaceId: string, jobId: string): Promise { + const root = createEmptyRootNode() - this.logging.log( - `ATX FES: Successfully mapped ${transformationPlan.transformationSteps?.length || 0} steps` - ) + if (!this.atxClient && !(await this.initializeAtxClient())) { + this.logging.error('ATX: Failed to initialize client for fetchPlanTree') + return { Root: root } + } - return transformationPlan - } else { - this.logging.log('ATX FES: No plan steps available yet - returning empty plan') - return { - transformationSteps: [] as any, - } as TransformationPlan - } - } catch (error) { - this.logging.error(`ATX FES getTransformationPlan error: ${String(error)}`) - // Return empty plan on error - return { - transformationSteps: [] as any, - } as TransformationPlan + // Fetch ALL steps in a single call (no parentStepId = returns all steps) + const allSteps = await this.fetchAllSteps(workspaceId, jobId) + if (!allSteps || allSteps.length === 0) { + this.logging.log('ATX: No plan steps available yet') + return { Root: root } } + + // Build tree from flat list + root.Children = this.buildTreeFromFlatList(allSteps) + + this.logging.log( + `ATX: fetchPlanTree completed - Built tree with ${root.Children.length} root steps from ${allSteps.length} total steps` + ) + return { Root: root } } - private async getATXFESJobPlanSteps(workspaceId: string, jobId: string): Promise { + /** + * Fetches all steps in a single API call (no parentStepId filter). + */ + private async fetchAllSteps(workspaceId: string, jobId: string): Promise { + const allSteps: any[] = [] + let nextToken: string | undefined + try { - this.logging.log(`ATX FES: getting plan steps with substeps...`) - const result = await this.listJobPlanSteps(workspaceId, jobId) - if (result) { - const steps = result || [] - this.logging.log(`ListJobPlanSteps: SUCCESS - Found ${steps.length} plan steps with substeps`) - return steps - } - return null + do { + const command = new ListJobPlanStepsCommand({ + workspaceId: workspaceId, + jobId: jobId, + maxResults: 100, + ...(nextToken && { nextToken }), + }) + + await this.addAuthToCommand(command) + const result = await this.atxClient!.send(command) + + if (result?.steps) { + allSteps.push(...result.steps) + } + nextToken = result?.nextToken + } while (nextToken) + + this.logging.log(`ATX: Fetched ${allSteps.length} total steps`) + return allSteps } catch (error) { - this.logging.error(`ListJobPlanSteps error: ${error instanceof Error ? error.message : 'Unknown error'}`) - return null + this.logging.error(`ATX: Error fetching all steps: ${String(error)}`) + return [] } } /** - * Lists job plan steps using FES client with recursive substep fetching + * Builds a tree structure from a flat list of steps using ParentStepId relationships. */ - private async listJobPlanSteps(workspaceId: string, jobId: string): Promise { - try { - this.logging.log(`ATX: Starting ListJobPlanSteps for job: ${jobId}`) - - if (!this.atxClient && !(await this.initializeAtxClient())) { - this.logging.error('ATX: Failed to initialize client for ListJobPlanSteps') - return null + private buildTreeFromFlatList(flatSteps: any[]): AtxPlanStep[] { + // Create a map of StepId -> AtxPlanStep for quick lookup + const stepMap = new Map() + + // First pass: convert all API steps to AtxPlanStep objects + for (const apiStep of flatSteps) { + const step = this.mapApiStepToNode(apiStep) + if (step.StepId) { + stepMap.set(step.StepId, step) } + } - // Get root steps first - const rootSteps = await this.getStepsRecursive(workspaceId, jobId, 'root') - - if (rootSteps && rootSteps.length > 0) { - // For each root step, get its substeps - for (const step of rootSteps) { - const substeps = await this.getStepsRecursive(workspaceId, jobId, step.stepId) - step.substeps = substeps || [] - - // Sort substeps by score (primary) and startTime (tiebreaker) to match RTS ordering - if (step.substeps.length > 0) { - step.substeps.sort((a: any, b: any) => { - const scoreDiff = (a.score || 0) - (b.score || 0) - if (scoreDiff !== 0) return scoreDiff - - // Tiebreaker for identical scores: sort by startTime - const timeA = a.startTime ? new Date(a.startTime).getTime() : 0 - const timeB = b.startTime ? new Date(b.startTime).getTime() : 0 - return timeA - timeB - }) - for (const substep of step.substeps) { - const superSubsteps = await this.getStepsRecursive(workspaceId, jobId, substep.stepId) - substep.substeps = superSubsteps || [] - - // Sort substeps by score (primary) and startTime (tiebreaker) to match RTS ordering - if (substep.substeps.length > 0) { - substep.substeps.sort((a: any, b: any) => { - const scoreDiff = (a.score || 0) - (b.score || 0) - if (scoreDiff !== 0) return scoreDiff - - // Tiebreaker for identical scores: sort by startTime - const timeA = a.startTime ? new Date(a.startTime).getTime() : 0 - const timeB = b.startTime ? new Date(b.startTime).getTime() : 0 - return timeA - timeB - }) - } - } - } - } + // Second pass: build parent-child relationships + const rootChildren: AtxPlanStep[] = [] - this.logging.log(`ATX: ListJobPlanSteps completed - Found ${rootSteps.length} steps with substeps`) - return rootSteps + for (const step of stepMap.values()) { + if (step.ParentStepId === 'root' || !step.ParentStepId) { + rootChildren.push(step) + } else { + const parent = stepMap.get(step.ParentStepId) + if (parent) { + parent.Children.push(step) + } else { + // Orphan step - treat as root level + rootChildren.push(step) + } } + } - this.logging.log('ATX: ListJobPlanSteps - No root steps found') - return null - } catch (error) { - this.logging.error(`ATX: ListJobPlanSteps error: ${String(error)}`) - return null + // Sort all children arrays by score + this.sortStepsByScore(rootChildren) + for (const step of stepMap.values()) { + if (step.Children.length > 0) { + this.sortStepsByScore(step.Children) + } } + + return rootChildren } /** - * Recursively gets steps for a given parent step ID + * Maps an API step response to AtxPlanStep. + * Converts from FES camelCase to C#-compatible PascalCase. */ - private async getStepsRecursive(workspaceId: string, jobId: string, parentStepId: string): Promise { - try { - const command = new ListJobPlanStepsCommand({ - workspaceId: workspaceId, - jobId: jobId, - parentStepId: parentStepId, - maxResults: 100, - }) + private mapApiStepToNode(apiStep: any): AtxPlanStep & { score?: number } { + return { + StepId: apiStep.stepId || '', + ParentStepId: apiStep.parentStepId === 'root' ? null : apiStep.parentStepId || null, + StepName: apiStep.stepName || '', + Description: apiStep.description || '', + Status: this.mapApiStatus(apiStep.status), + Children: [], + // Keep score for sorting (not sent to C#) + score: apiStep.score || 0, + } + } - await this.addAuthToCommand(command) - const result = await this.atxClient!.send(command) + /** + * Maps API status string to PlanStepStatus. + * Returns the status directly if valid, otherwise defaults to NOT_STARTED. + */ + private mapApiStatus(status: string | undefined): PlanStepStatus { + if (!status) return 'NOT_STARTED' + // The API returns valid PlanStepStatus values directly + return status as PlanStepStatus + } - if (result && result.steps && result.steps.length > 0) { - return result.steps - } + /** + * Sorts steps by score (primary). + */ + private sortStepsByScore(steps: (AtxPlanStep & { score?: number })[]): void { + steps.sort((a, b) => (a.score || 0) - (b.score || 0)) + } - return null - } catch (error) { - this.logging.error(`Error getting steps for parent ${parentStepId}: ${String(error)}`) - return null - } + /** + * Fetches worklogs for a job and saves them to disk. + */ + private async fetchWorklogs(workspaceId: string, jobId: string, solutionRootPath: string): Promise { + await this.listWorklogs(workspaceId, jobId, solutionRootPath) } /** @@ -1408,7 +1623,7 @@ export class ATXTransformHandler { } /** - * Lists artifacts using FES client with CUSTOMER_OUTPUT filtering + * Lists worklogs for a job and saves them to disk grouped by step ID. */ private async listWorklogs( workspaceId: string, @@ -1451,4 +1666,260 @@ export class ATXTransformHandler { return null } } + + /** + * Set checkpoints for interactive mode transformation. + * Lists HITLs with "checkpoint-settings" tag, uploads checkpoints as JSON artifact, + * and updates the HITL task with the new artifact ID. + */ + async setCheckpoints( + workspaceId: string, + jobId: string, + solutionRootPath: string, + checkpoints: Record + ): Promise { + try { + this.logging.log(`ATX: Starting setCheckpoints for job: ${jobId}`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + return { Success: false, Error: 'ATX FES client not initialized' } + } + + // Step 1: List HITLs with "checkpoint-settings" tag + const hitlTask = await this.findCheckpointSettingsHitl(workspaceId, jobId) + + if (!hitlTask) { + return { Success: false, Error: 'No HITL task found with checkpoint-settings tag' } + } + + this.logging.log(`ATX: Found checkpoint-settings HITL task: ${hitlTask.taskId}`) + + // Step 2: Create JSON file with checkpoints mapping in checkpoints folder + const artifactDir = path.join(solutionRootPath, workspaceFolderName, jobId, 'checkpoints') + if (!fs.existsSync(artifactDir)) { + fs.mkdirSync(artifactDir, { recursive: true }) + } + + const jsonFilePath = path.join(artifactDir, 'checkpoint-settings.json') + fs.writeFileSync(jsonFilePath, JSON.stringify(checkpoints, null, 2)) + + // Step 3: Upload the JSON artifact + const uploadInfo = await this.createArtifactUploadUrl( + workspaceId, + jobId, + jsonFilePath, + CategoryType.HITL_FROM_USER, + FileType.JSON + ) + + if (!uploadInfo) { + return { Success: false, Error: 'Failed to create artifact upload URL' } + } + + const uploadSuccess = await Utils.uploadArtifact( + uploadInfo.uploadUrl, + jsonFilePath, + uploadInfo.requestHeaders, + this.logging + ) + + if (!uploadSuccess) { + return { Success: false, Error: 'Failed to upload checkpoints artifact to S3' } + } + + // Step 4: Complete artifact upload + const completeResponse = await this.completeArtifactUpload(workspaceId, jobId, uploadInfo.uploadId) + + if (!completeResponse?.success) { + return { Success: false, Error: 'Failed to complete artifact upload' } + } + + // Step 5: Update HITL task with the new artifact ID + const updateResult = await this.updateHitl(workspaceId, jobId, hitlTask.taskId, uploadInfo.uploadId) + + if (!updateResult) { + return { Success: false, Error: 'Failed to update HITL task with checkpoints artifact' } + } + + this.logging.log(`ATX: setCheckpoints completed successfully`) + return { Success: true } + } catch (error) { + this.logging.error(`ATX: setCheckpoints error: ${String(error)}`) + return { Success: false, Error: String(error) } + } + } + + /** + * Find HITL task with "checkpoint-settings" tag + */ + private async findCheckpointSettingsHitl(workspaceId: string, jobId: string): Promise { + try { + this.logging.log(`ATX: Looking for HITL task with checkpoint-settings tag`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + this.logging.error('ATX: Failed to initialize client for findCheckpointSettingsHitl') + return null + } + + const command = new ListHitlTasksCommand({ + workspaceId: workspaceId, + jobId: jobId, + taskType: 'NORMAL', + taskFilter: { + taskStatuses: ['AWAITING_HUMAN_INPUT', 'IN_PROGRESS'], + tag: `${jobId}-checkpoint`, + }, + }) + + await this.addAuthToCommand(command) + const result = await this.atxClient!.send(command) + + if (result.hitlTasks && result.hitlTasks.length > 0) { + this.logging.log(`ATX: Found ${result.hitlTasks.length} HITL task(s) with checkpoint-settings tag`) + return result.hitlTasks[0] + } + + this.logging.log('ATX: No HITL task found with checkpoint-settings tag') + return null + } catch (error) { + this.logging.error(`ATX: findCheckpointSettingsHitl error: ${String(error)}`) + return null + } + } + + /** + * Handle checkpoint action (APPLY or RETRY) for a step-level HITL. + */ + async checkpointAction( + workspaceId: string, + jobId: string, + stepId: string, + action: string, + solutionRootPath: string, + newInstruction?: string + ): Promise { + try { + this.logging.log(`ATX: Starting checkpointAction for job: ${jobId}, step: ${stepId}, action: ${action}`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + return { Success: false, Error: 'ATX FES client not initialized' } + } + + // Get the cached step HITL task ID, or query for it if not cached + let taskId: string | null = this.cachedStepHitl + if (!taskId) { + this.logging.log('ATX: No cached step HITL, querying for active step HITL') + const stepHitl = await this.findStepLevelHitl(workspaceId, jobId, stepId) + if (!stepHitl || !stepHitl.taskId) { + return { Success: false, Error: 'No active step HITL found' } + } + taskId = stepHitl.taskId + } + + // At this point taskId is guaranteed to be a string + const validTaskId = taskId as string + + // Create the human artifact JSON + const artifactContent: any = { + action: action, + } + if (action === 'RETRY' && newInstruction) { + artifactContent.newInstruction = newInstruction + } + + // Create the JSON file at {solutionRootPath}/{workspaceFolderName}/{jobId}/checkpoints/checkpoint-action.json + const artifactDir = path.join(solutionRootPath, workspaceFolderName, jobId, 'checkpoints') + if (!fs.existsSync(artifactDir)) { + fs.mkdirSync(artifactDir, { recursive: true }) + } + + const jsonFilePath = path.join(artifactDir, 'checkpoint-action.json') + fs.writeFileSync(jsonFilePath, JSON.stringify(artifactContent, null, 2)) + + // Upload the JSON artifact + const uploadInfo = await this.createArtifactUploadUrl( + workspaceId, + jobId, + jsonFilePath, + CategoryType.HITL_FROM_USER, + FileType.JSON + ) + + if (!uploadInfo) { + return { Success: false, Error: 'Failed to create artifact upload URL' } + } + + const uploadSuccess = await Utils.uploadArtifact( + uploadInfo.uploadUrl, + jsonFilePath, + uploadInfo.requestHeaders, + this.logging + ) + + if (!uploadSuccess) { + return { Success: false, Error: 'Failed to upload checkpoint action artifact to S3' } + } + + // Complete artifact upload + const completeResponse = await this.completeArtifactUpload(workspaceId, jobId, uploadInfo.uploadId) + + if (!completeResponse?.success) { + return { Success: false, Error: 'Failed to complete artifact upload' } + } + + // Submit the standard HITL task with the human artifact + const submitResult = await this.submitStandardHitl(workspaceId, jobId, validTaskId, uploadInfo.uploadId) + + if (!submitResult) { + return { Success: false, Error: 'Failed to submit checkpoint action' } + } + + // Clear the cached step HITL after successful submission + this.cachedStepHitl = null + + this.logging.log(`ATX: checkpointAction completed successfully`) + return { Success: true } + } catch (error) { + this.logging.error(`ATX: checkpointAction error: ${String(error)}`) + return { Success: false, Error: String(error) } + } + } + + /** + * Submit a standard HITL task with a human artifact. + */ + private async submitStandardHitl( + workspaceId: string, + jobId: string, + taskId: string, + humanArtifactId: string + ): Promise { + try { + this.logging.log(`ATX: Starting SubmitStandardHitl for task: ${taskId}`) + + if (!this.atxClient && !(await this.initializeAtxClient())) { + this.logging.error('ATX: Failed to initialize client for SubmitStandardHitl') + return null + } + + const command = new SubmitStandardHitlTaskCommand({ + workspaceId: workspaceId, + jobId: jobId, + taskId: taskId, + action: 'APPROVE', + humanArtifact: { + artifactId: humanArtifactId, + }, + }) + + await this.addAuthToCommand(command) + const result = await this.atxClient!.send(command) + + this.logging.log(`ATX: SubmitStandardHitl completed - task status: ${result.status || 'UNKNOWN'}`) + return result + } catch (error) { + this.logging.error(`ATX: SubmitStandardHitl error: ${String(error)}`) + return null + } + } }