implement resume request batching

Author: eyw520

servers/fai-lambda-deploy/scripts/fai-scribe-stack.ts

@@ -126,8 +126,6 @@ export class FaiScribeStack extends Stack {
             filesystem: lambda.FileSystem.fromEfsAccessPoint(accessPoint, "/mnt/efs")
         });
 
-        // Grant SQS permissions for editing session queues
-        // These queues are created dynamically by the FAI server per editing session
         lambdaFunction.addToRolePolicy(
             new iam.PolicyStatement({
                 effect: iam.Effect.ALLOW,
@@ -138,10 +136,7 @@ export class FaiScribeStack extends Stack {
                     "sqs:GetQueueUrl",
                     "sqs:ChangeMessageVisibility"
                 ],
-                resources: [
-                    // Allow access to all editing-session queues in this region
-                    `arn:aws:sqs:${this.region}:${this.account}:editing-session-*.fifo`
-                ]
+                resources: [`arn:aws:sqs:${this.region}:${this.account}:editing-session-*.fifo`]
             })
         );
 

servers/fai-lambda/fai-scribe/src/handler.py

@@ -109,7 +109,6 @@ async def handle_editing_request(
             else:
                 LOGGER.info(f"Successfully updated editing session: {editing_id}")
 
-            # Clean up the SQS queue after successful completion
             LOGGER.info(f"Cleaning up SQS queue for session: {editing_id}")
             cleanup_response = await client.delete(f"{SETTINGS.FAI_API_URL}/editing-sessions/{editing_id}/queue")
             if cleanup_response.status_code == 200:

servers/fai-lambda/fai-scribe/src/utils/agent.py

@@ -186,6 +186,49 @@ async def run_editing_session(
     if session_id is None:
         raise RuntimeError("Failed to obtain session ID from Claude")
 
+    LOGGER.info(f"Checking for RESUME messages in queue for session {editing_id}")
+    resume_messages = sqs_client.get_resume_messages()
+
+    if resume_messages:
+        LOGGER.info(f"Found {len(resume_messages)} RESUME messages, batching prompts")
+
+        all_prompts = []
+        for resume_msg in resume_messages:
+            body = resume_msg["body"]
+            prompts = body.get("prompts", [])
+            if prompts:
+                all_prompts.extend(prompts)
+            sqs_client.delete_message(resume_msg["receipt_handle"])
+
+        if all_prompts:
+            batched_prompt = "\n\n".join([f"Request {i+1}: {p}" for i, p in enumerate(all_prompts)])
+            LOGGER.info(f"Processing {len(all_prompts)} batched prompts: {batched_prompt[:200]}...")
+
+            async with ClaudeSDKClient(options=options) as client:
+                await client.query(batched_prompt)
+
+                async for msg in client.receive_response():
+                    is_interrupted, receipt_handle = check_if_interrupted_via_sqs(sqs_client)
+                    if is_interrupted:
+                        LOGGER.warning(f"Session interrupted while processing batched prompts: {editing_id}")
+                        if receipt_handle:
+                            sqs_client.delete_message(receipt_handle)
+                        await update_session_status(editing_id, "waiting")
+                        raise SessionInterruptedError(f"Editing session {editing_id} was interrupted")
+
+                    if isinstance(msg, AssistantMessage):
+                        for block in msg.content:
+                            if isinstance(block, TextBlock):
+                                for line in block.text.split("\n"):
+                                    if "PR_URL:" in line:
+                                        extracted_text = line.split("PR_URL:", 1)[1].strip()
+                                        if extracted_text:
+                                            match = GITHUB_PR_URL_PATTERN.search(extracted_text)
+                                            if match:
+                                                pr_url = match.group(0)
+                                                LOGGER.info(f"Updated PR URL: {pr_url}")
+                                                await update_session_metadata(editing_id, pr_url=pr_url)
+
     await update_session_status(editing_id, "waiting")
 
     return session_id, pr_url

servers/fai/deploy/src/deploy-stack.ts

@@ -159,7 +159,6 @@ export class FernAiDeployStack extends Stack {
             );
         }
 
-        // Grant SQS permissions for creating and managing editing session queues
         fargateService.taskDefinition.taskRole.addToPrincipalPolicy(
             new iam.PolicyStatement({
                 effect: iam.Effect.ALLOW,
@@ -172,10 +171,7 @@ export class FernAiDeployStack extends Stack {
                     "sqs:SetQueueAttributes",
                     "sqs:TagQueue"
                 ],
-                resources: [
-                    // Allow creating and managing editing-session queues
-                    `arn:aws:sqs:us-east-1:985111089818:editing-session-*.fifo`
-                ]
+                resources: [`arn:aws:sqs:us-east-1:985111089818:editing-session-*.fifo`]
             })
         );