feat: add .after test chain

packages/shortest/index.d.ts

@@ -24,6 +24,7 @@ declare module '@antiwork/shortest' {
     expect(description: string): TestChain;
     expect(description: string, fn?: (context: TestContextProps) => Promise<void>): TestChain;
     expect(description: string, payload?: any, fn?: (context: TestContextProps) => Promise<void>): TestChain;
+    after(fn: (context: TestContextProps) => void | Promise<void>): TestChain;
   };
 
   export type TestAPI = {

packages/shortest/src/core/runner/index.ts

@@ -194,30 +194,44 @@ export class TestRunner {
     ].filter(Boolean).join('\n');
 
     // Execute test with enhanced prompt
-    const result = await aiClient.processAction(prompt, browserTool, (content: Anthropic.Beta.Messages.BetaContentBlockParam) => {
-      if (content.type === 'text') {
-        // this.logger.reportStatus(`🤖 ${(content as Anthropic.Beta.Messages.BetaTextBlock).text}`);
-      }
-    });
+    const result = await aiClient.processAction(prompt, browserTool);
 
     if (!result) {
       throw new Error('AI processing failed: no result returned');
     }
 
+    // Parse AI result first
     const finalMessage = result.finalResponse.content.find(block => 
       block.type === 'text' && 
       (block as Anthropic.Beta.Messages.BetaTextBlock).text.includes('"result":')
     );
 
-    if (finalMessage && finalMessage.type === 'text') {
-      const jsonMatch = (finalMessage as Anthropic.Beta.Messages.BetaTextBlock).text.match(/{[\s\S]*}/);
-      if (jsonMatch) {
-        const testResult = JSON.parse(jsonMatch[0]) as TestResult;
-        return testResult;
+    if (!finalMessage || finalMessage.type !== 'text') {
+      throw new Error('No test result found in AI response');
+    }
+
+    const jsonMatch = (finalMessage as Anthropic.Beta.Messages.BetaTextBlock).text.match(/{[\s\S]*}/);
+    if (!jsonMatch) {
+      throw new Error('Invalid test result format');
+    }
+
+    const aiResult = JSON.parse(jsonMatch[0]) as TestResult;
+
+    // Execute after function if present
+    if (test.afterFn) {
+      try {
+        await test.afterFn(testContext);
+      } catch (error) {
+        return {
+          result: 'fail' as const,
+          reason: aiResult.result === 'fail'
+            ? `AI: ${aiResult.reason}, After: ${error instanceof Error ? error.message : String(error)}`
+            : error instanceof Error ? error.message : String(error)
+        };
       }
     }
 
-    throw new Error('No test result found in AI response');
+    return aiResult;
   }
 
   private async executeTestFile(file: string) {

packages/shortest/src/index.ts

@@ -128,6 +128,9 @@ function createTestChain(
     return {
       expect: () => {
         throw new Error('expect() cannot be called on direct execution test');
+      },
+      after: () => {
+        throw new Error('after() cannot be called on direct execution test');
       }
     };
   }
@@ -167,6 +170,10 @@ function createTestChain(
         fn: typeof payloadOrFn === 'function' ? payloadOrFn : fn
       });
       return chain;
+    },
+    after(fn: (context: TestContext) => void | Promise<void>) {
+      test.afterFn = (context) => Promise.resolve(fn(context));
+      return chain;
     }
   };
 

packages/shortest/src/types/test.ts

@@ -49,13 +49,15 @@ export interface TestFunction {
     directExecution?: boolean;
   }[];
   directExecution?: boolean;
+  afterFn?: (context: TestContext) => void | Promise<void>;
 }
 
 export type TestChain = {
   expect(fn: (context: TestContext) => Promise<void>): TestChain;
   expect(description: string): TestChain;
   expect(description: string, fn?: (context: TestContext) => Promise<void>): TestChain;
   expect(description: string, payload?: any, fn?: (context: TestContext) => Promise<void>): TestChain;
+  after(fn: (context: TestContext) => void | Promise<void>): TestChain;
 };
 
 export type TestAPI = {