imprint-mcp 0.3.0 → 0.4.0

package/src/imprint/teach.ts

@@ -82,7 +82,9 @@ import {
   isExistingTeachFile as isExistingFile,
   loadTeachState,
   nextTeachStep as nextStep,
+  pruneStalePendingTeachWorkflows,
   resolveTeachStatePath,
+  resolveWorkflowTriagedPath,
   saveTeachState,
   toRelativeTeachStatePath as toRelative,
 } from './teach-state.ts';
@@ -98,7 +100,11 @@ import { setSpanAttributes, traced } from './tracing.ts';
 import { CronConfigSchema, SessionSchema, WorkflowSchema } from './types.ts';
 import type { CronConfig, Playbook, Session, Workflow } from './types.ts';
 
-export { buildTeachStateFromSession, resolveTeachStatePath } from './teach-state.ts';
+export {
+  buildTeachStateFromSession,
+  resolveTeachStatePath,
+  resolveWorkflowTriagedPath,
+} from './teach-state.ts';
 
 /**
  * How many compile agents run in parallel when more than one tool is selected.
@@ -396,6 +402,9 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
 
   const completedWorkflows = discoverCompletedWorkflows(site);
   const completedSet = new Set(completedWorkflows);
+  if (pruneStalePendingTeachWorkflows(site, state)) {
+    saveTeachState(site, state);
+  }
   const incompleteWorkflows = Object.entries(state.workflows).filter(
     ([name]) => !completedSet.has(name),
   );
@@ -725,7 +734,6 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
   // Run them in parallel so the user can select tools while replay runs.
   let siteClassifications: ClassifiedValue[] | undefined;
   let triageResult: TriageResult | undefined;
-  let triagedPath: string | null = null;
   let plans: CandidateCompilePlan[];
 
   let needsReplay = startIdx <= STEPS.indexOf('replay-and-diff') && !opts.skipReplay;
@@ -783,7 +791,7 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
 
       // Branch B: triage → detect-candidates → user selection (fast, ~30s)
       type CandidateChainResult = {
-        triageRes?: { result: TriageResult; sessionPath: string };
+        triageResult?: TriageResult;
         plans: CandidateCompilePlan[];
       };
       const candidatePromise = (async (): Promise<CandidateChainResult> => {
@@ -836,8 +844,11 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
           );
         } else {
           const ws = state.workflows[workflowKey];
-          if (ws?.triagedPath) {
-            localTriagedPath = resolveTeachStatePath(site, ws.triagedPath);
+          localTriagedPath = resolveWorkflowTriagedPath(site, ws);
+          if (ws && localTriagedPath && !ws.triagedPath) {
+            updateCheckpoint(site, state, workflowKey, 'triage', {
+              triagedPath: toRelative(site, localTriagedPath),
+            });
           }
         }
 
@@ -875,6 +886,9 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
           kind: 'raw',
         });
         const baseState = buildTeachStateFromSession(site, rawSessionPath, redactedPath);
+        if (localTriagedPath) {
+          baseState.triagedPath = toRelative(site, localTriagedPath);
+        }
         const candidatePlans = selected.map((candidate) => {
           checkpoint(site, state, candidate.toolName, {
             ...baseState,
@@ -896,9 +910,7 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
         }
 
         return {
-          triageRes: localTriageResult
-            ? { result: localTriageResult, sessionPath: replaySessionPath }
-            : undefined,
+          triageResult: localTriageResult,
           plans: candidatePlans,
         };
       })();
@@ -907,13 +919,7 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
       const candidateResult = await candidatePromise;
       plans = candidateResult.plans;
 
-      if (candidateResult.triageRes) {
-        triageResult = candidateResult.triageRes.result;
-        triagedPath = candidateResult.triageRes.sessionPath.replace(
-          /\.redacted\.json$/,
-          '.triaged.json',
-        );
-      }
+      triageResult = candidateResult.triageResult;
 
       // Wait for replay — may already be done, or show progress while waiting
       let replaySettled = false;
@@ -938,18 +944,19 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
       mp.clear();
 
       // Checkpoints — write sequentially after both complete
-      if (needsReplay) {
-        updateCheckpoint(site, state, workflowKey, 'replay-and-diff', {
-          classificationsPath: siteClassifications
-            ? toRelative(site, pathJoin(localSiteDir(site), '.classifications.json'))
-            : undefined,
-        });
-      }
-      if (candidateResult.triageRes && triagedPath) {
-        updateCheckpoint(site, state, workflowKey, 'triage', {
-          triagedPath: toRelative(site, triagedPath),
-        });
-      }
+      updateCandidateStageCheckpoints({
+        site,
+        state,
+        plans,
+        fallbackWorkflowKey: workflowKey,
+        replay: needsReplay
+          ? {
+              classificationsPath: siteClassifications
+                ? toRelative(site, pathJoin(localSiteDir(site), '.classifications.json'))
+                : undefined,
+            }
+          : undefined,
+      });
     } finally {
       unmuteLog();
     }
@@ -964,8 +971,11 @@ export async function teach(opts: TeachOptions): Promise<TeachResult> {
       }
     }
     const ws = state.workflows[workflowKey];
-    if (ws?.triagedPath) {
-      triagedPath = resolveTeachStatePath(site, ws.triagedPath);
+    const resolvedTriagedPath = resolveWorkflowTriagedPath(site, ws);
+    if (ws && resolvedTriagedPath && !ws.triagedPath) {
+      updateCheckpoint(site, state, workflowKey, 'triage', {
+        triagedPath: toRelative(site, resolvedTriagedPath),
+      });
     }
     plans = [
       {
@@ -1251,6 +1261,29 @@ export interface CandidateCompilePlan {
   sharedContext?: SharedCompileContext;
 }
 
+function candidateStageCheckpointKeys(
+  plans: CandidateCompilePlan[],
+  fallbackWorkflowKey: string,
+): string[] {
+  const keys = plans.map((plan) => plan.workflowKey).filter((key) => key.length > 0);
+  return [...new Set(keys.length > 0 ? keys : [fallbackWorkflowKey])];
+}
+
+export function updateCandidateStageCheckpoints(opts: {
+  site: string;
+  state: TeachState;
+  plans: CandidateCompilePlan[];
+  fallbackWorkflowKey: string;
+  replay?: Partial<WorkflowState>;
+  triage?: Partial<WorkflowState>;
+}): void {
+  const keys = candidateStageCheckpointKeys(opts.plans, opts.fallbackWorkflowKey);
+  for (const key of keys) {
+    if (opts.replay) updateCheckpoint(opts.site, opts.state, key, 'replay-and-diff', opts.replay);
+    if (opts.triage) updateCheckpoint(opts.site, opts.state, key, 'triage', opts.triage);
+  }
+}
+
 async function detectTeachCandidates(opts: {
   sessionPath: string;
   providerName: ProviderName;