Merge pull request #27 from dsp-testing/fix/forge-terminal-and-query-volume

Harden Forge interrupted extraction handling

Author: GolaraJ

plugins/repo-dreamer/skills/repository-skill-forge/SKILL.md

@@ -233,6 +233,19 @@ already found. Continue invoking the controller while its state is `running`;
 the agent must not independently declare the run blocked because remaining work
 is slow or numerous.
 
+Before leaving extraction, updating state, or reporting any terminal outcome,
+require the controller to confirm that extraction is terminal:
+
+```bash
+python3 "$SKILL_DIR/scripts/extraction-controller.py" assert-terminal \
+  --state "$RUN_DIR/extraction-state.json"
+```
+
+This command fails while state is `running` and identifies any issued actions
+that still require outcomes. A run is blocked only when this command returns
+`status: blocked` with the controller-recorded blocker. Never describe
+incomplete `running` work as blocked.
+
 Generate the next bounded action batch with:
 
 ```bash

plugins/repo-dreamer/skills/repository-skill-forge/scripts/extraction-controller.py

@@ -202,6 +202,36 @@ def finalize_extraction(state: dict[str, Any]) -> None:
     state["status"] = "partial" if state["omittedUnits"] else "complete"
 
 
+def terminal_summary(state: dict[str, Any]) -> dict[str, Any]:
+    validate_state_invariants(state)
+    status = str(state.get("status"))
+    if status == "running":
+        pending = [
+            str(action.get("actionId"))
+            for action in state.get("issuedActions", [])
+            if action.get("actionId")
+        ]
+        detail = f"; pending actions: {', '.join(pending)}" if pending else ""
+        raise ValueError(
+            "extraction is not terminal: status is running"
+            f"{detail}; continue invoking next and recording every outcome"
+        )
+    if status not in {"complete", "partial", "blocked"}:
+        raise ValueError(f"extraction has unsupported terminal status: {status}")
+    summary: dict[str, Any] = {
+        "kind": "terminal",
+        "status": status,
+        "terminal": True,
+    }
+    if state.get("coverage") is not None:
+        summary["coverage"] = state["coverage"]
+    if status == "partial":
+        summary["omittedUnits"] = state["omittedUnits"]
+    if status == "blocked":
+        summary["blocker"] = state["blockers"][-1]
+    return summary
+
+
 def next_actions(
     state: dict[str, Any],
     max_actions: int,
@@ -604,7 +634,6 @@ def record_success(
         elif action["kind"] == "files":
             batch["filesCursor"] = {
                 "sessionId": last["session_id"],
-                "turnIndex": last["turn_index"],
                 "filePath": last["file_path"],
                 "toolName": last["tool_name"],
             }
@@ -1142,6 +1171,10 @@ def main() -> None:
     next_parser.add_argument("--out")
     next_parser.add_argument("--parallel", action="store_true")
 
+    terminal = subparsers.add_parser("assert-terminal")
+    terminal.add_argument("--state", required=True)
+    terminal.add_argument("--out")
+
     success = subparsers.add_parser("record-success")
     success.add_argument("--state", required=True)
     success.add_argument("--action", required=True)
@@ -1193,22 +1226,38 @@ def main() -> None:
             )
             actions = next_actions(state, max_actions)
             write_json(args.state, state)
-            done = {"kind": "done", "status": state["status"]}
+            done = {
+                "kind": "done",
+                "status": state["status"],
+                "terminal": True,
+            }
             if state["status"] == "blocked":
                 done["blocker"] = state["blockers"][-1]
             if state["coverage"] is not None:
                 done["coverage"] = state["coverage"]
             if state["status"] == "partial":
                 done["omittedUnits"] = state["omittedUnits"]
             if args.parallel and actions:
-                payload = {"kind": "action-batch", "actions": actions}
+                payload = {
+                    "kind": "action-batch",
+                    "status": state["status"],
+                    "terminal": False,
+                    "actions": actions,
+                }
             else:
                 payload = actions[0] if actions else done
             if args.out:
                 write_json(args.out, payload)
             else:
                 print(json.dumps(payload, indent=2))
             return
+        if args.command == "assert-terminal":
+            summary = terminal_summary(state)
+            if args.out:
+                write_json(args.out, summary)
+            else:
+                print(json.dumps(summary, indent=2))
+            return
         action = load_action(args.action, state)
         if args.command == "record-success":
             try:

plugins/repo-dreamer/skills/repository-skill-forge/scripts/materialize-session-query.py

@@ -27,8 +27,45 @@ class QueryHandoffMismatch(ValueError):
     """The action ID matched, but the submitted SQL differed."""
 
 
-def normalize_sql(sql: str) -> str:
-    return " ".join(sql.split())
+def normalize_sql(sql: str) -> tuple[str, ...]:
+    tokens: list[str] = []
+    index = 0
+    while index < len(sql):
+        character = sql[index]
+        if character.isspace():
+            index += 1
+            continue
+        if character in {"'", '"'}:
+            quote = character
+            end = index + 1
+            while end < len(sql):
+                if sql[end] == quote:
+                    if end + 1 < len(sql) and sql[end + 1] == quote:
+                        end += 2
+                        continue
+                    end += 1
+                    break
+                end += 1
+            tokens.append(sql[index:end])
+            index = end
+            continue
+        if character.isalnum() or character in {"_", "$"}:
+            end = index + 1
+            while end < len(sql) and (
+                sql[end].isalnum() or sql[end] in {"_", "$"}
+            ):
+                end += 1
+            tokens.append(sql[index:end])
+            index = end
+            continue
+        operator = sql[index : index + 2]
+        if operator in {">=", "<=", "<>", "!=", "||", "::"}:
+            tokens.append(operator)
+            index += 2
+            continue
+        tokens.append(character)
+        index += 1
+    return tuple(tokens)
 
 
 def read_events(path: Path) -> Iterator[dict[str, Any]]:

plugins/repo-dreamer/skills/repository-skill-forge/scripts/session_queries.py

@@ -111,16 +111,22 @@ def build_files_query(
     after = ""
     if cursor:
         after = (
-            "\n  AND (session_id, turn_index, file_path, tool_name) > "
-            f"('{sql_literal(str(cursor['sessionId']))}', {int(cursor['turnIndex'])}, "
-            f"'{sql_literal(str(cursor['filePath']))}', '{sql_literal(str(cursor['toolName']))}')"
+            "\nWHERE (session_id, file_path, tool_name) > "
+            f"('{sql_literal(str(cursor['sessionId']))}', "
+            f"'{sql_literal(str(cursor['filePath']))}', "
+            f"'{sql_literal(str(cursor['toolName']))}')"
         )
-    return f"""SELECT session_id, file_path, tool_name, turn_index
-FROM session_files
-WHERE first_seen_at >= TIMESTAMP '{sql_literal(start)}'
-  AND first_seen_at < TIMESTAMP '{sql_literal(end)}'
-  AND session_id IN ({ids}){after}
-ORDER BY session_id, turn_index, file_path, tool_name
+    return f"""WITH selected_files AS (
+    SELECT session_id, file_path, tool_name, min(turn_index) AS turn_index
+    FROM session_files
+    WHERE first_seen_at >= TIMESTAMP '{sql_literal(start)}'
+      AND first_seen_at < TIMESTAMP '{sql_literal(end)}'
+      AND session_id IN ({ids})
+    GROUP BY session_id, file_path, tool_name
+)
+SELECT session_id, file_path, tool_name, turn_index
+FROM selected_files{after}
+ORDER BY session_id, file_path, tool_name
 LIMIT {limit + 1}"""
 
 

plugins/repo-dreamer/skills/repository-skill-forge/tests/test_extraction_controller.py

@@ -16,7 +16,7 @@
 SCRIPTS_DIR = SKILL_DIR / "scripts"
 sys.path.insert(0, str(SCRIPTS_DIR))
 
-from session_queries import build_discovery_query
+from session_queries import build_discovery_query, build_files_query
 
 CONTROLLER_SPEC = importlib.util.spec_from_file_location(
     "extraction_controller",
@@ -133,6 +133,118 @@ def test_cli_defaults_to_twenty_five_session_batches(self) -> None:
                 ],
             )
 
+    def test_assert_terminal_rejects_running_state(self) -> None:
+        with tempfile.TemporaryDirectory() as run_dir:
+            state_path = Path(run_dir) / "state.json"
+            state = controller.initialize(arguments(run_dir))
+            action = controller.next_action(state)
+            assert action is not None
+            controller.write_json(str(state_path), state)
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    str(SCRIPTS_DIR / "extraction-controller.py"),
+                    "assert-terminal",
+                    "--state",
+                    str(state_path),
+                ],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            self.assertNotEqual(0, result.returncode)
+            self.assertIn("status is running", result.stderr)
+            self.assertIn(action["actionId"], result.stderr)
+
+    def test_assert_terminal_accepts_complete_state(self) -> None:
+        with tempfile.TemporaryDirectory() as run_dir:
+            state_path = Path(run_dir) / "state.json"
+            state = controller.initialize(arguments(run_dir))
+            state["status"] = "complete"
+            state["coverage"] = 1.0
+            controller.write_json(str(state_path), state)
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    str(SCRIPTS_DIR / "extraction-controller.py"),
+                    "assert-terminal",
+                    "--state",
+                    str(state_path),
+                ],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            self.assertEqual("", result.stderr)
+            self.assertEqual(0, result.returncode)
+            self.assertEqual(
+                {
+                    "kind": "terminal",
+                    "status": "complete",
+                    "terminal": True,
+                    "coverage": 1.0,
+                },
+                json.loads(result.stdout),
+            )
+
+    def test_parallel_action_manifest_is_explicitly_nonterminal(self) -> None:
+        with tempfile.TemporaryDirectory() as run_dir:
+            state_path = Path(run_dir) / "state.json"
+            actions_path = Path(run_dir) / "actions.json"
+            controller.write_json(
+                str(state_path),
+                controller.initialize(arguments(run_dir)),
+            )
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    str(SCRIPTS_DIR / "extraction-controller.py"),
+                    "next",
+                    "--state",
+                    str(state_path),
+                    "--parallel",
+                    "--out",
+                    str(actions_path),
+                ],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            self.assertEqual("", result.stderr)
+            self.assertEqual(0, result.returncode)
+            manifest = json.loads(actions_path.read_text(encoding="utf-8"))
+            self.assertEqual("action-batch", manifest["kind"])
+            self.assertEqual("running", manifest["status"])
+            self.assertFalse(manifest["terminal"])
+
+    def test_file_query_deduplicates_before_pagination(self) -> None:
+        query = build_files_query(
+            session_ids=["session-1"],
+            start="2026-08-01T00:00:00Z",
+            end="2026-08-08T00:00:00Z",
+            limit=500,
+            cursor={
+                "sessionId": "session-1",
+                "filePath": "src/example.py",
+                "toolName": "edit",
+            },
+        )
+
+        self.assertIn("min(turn_index) AS turn_index", query)
+        self.assertIn("GROUP BY session_id, file_path, tool_name", query)
+        self.assertIn(
+            "WHERE (session_id, file_path, tool_name) > "
+            "('session-1', 'src/example.py', 'edit')",
+            query,
+        )
+        self.assertIn("ORDER BY session_id, file_path, tool_name", query)
+
     def test_discovery_uses_ordered_keyset_query(self) -> None:
         query = build_discovery_query(
             repository="owner/repository",

plugins/repo-dreamer/skills/repository-skill-forge/tests/test_materialize_session_query.py

@@ -342,6 +342,60 @@ def test_matches_whitespace_normalized_sql(self) -> None:
                 materializer.result_content(events_root, sql, "discovery-1"),
             )
 
+    def test_matches_punctuation_spacing_normalized_sql(self) -> None:
+        with tempfile.TemporaryDirectory() as temporary:
+            events_root = Path(temporary)
+            event_dir = events_root / "session-1"
+            event_dir.mkdir()
+            sql = (
+                "SELECT session_id, tool_call_id FROM tool_requests "
+                "WHERE session_id IN ('one', 'two') AND name = 'bash'"
+            )
+            submitted = (
+                "SELECT session_id,tool_call_id FROM tool_requests "
+                "WHERE session_id IN('one','two') AND name='bash'"
+            )
+            content = "Query returned 0 rows."
+            events = [
+                {
+                    "type": "tool.execution_start",
+                    "data": {
+                        "toolCallId": "call-1",
+                        "toolName": "session_store_sql",
+                        "arguments": {
+                            "description": "tools-batch-1",
+                            "query": submitted,
+                        },
+                    },
+                },
+                {
+                    "type": "tool.execution_complete",
+                    "data": {
+                        "toolCallId": "call-1",
+                        "success": True,
+                        "result": {
+                            "content": content,
+                            "detailedContent": "SQL result omitted",
+                        },
+                    },
+                },
+            ]
+            (event_dir / "events.jsonl").write_text(
+                "".join(json.dumps(event) + "\n" for event in events),
+                encoding="utf-8",
+            )
+
+            self.assertEqual(
+                content,
+                materializer.result_content(events_root, sql, "tools-batch-1"),
+            )
+
+    def test_sql_normalization_preserves_literal_contents(self) -> None:
+        self.assertNotEqual(
+            materializer.normalize_sql("SELECT 'one two'"),
+            materializer.normalize_sql("SELECT 'onetwo'"),
+        )
+
     def test_ignores_matching_description_from_other_tools(self) -> None:
         with tempfile.TemporaryDirectory() as temporary:
             events_root = Path(temporary)