Add Pro prompt redaction integration

flocks/session/goal.py

@@ -11,6 +11,12 @@
 
 from flocks.provider.options import build_provider_options
 from flocks.provider.provider import ChatMessage, Provider
+from flocks.session.llm_hook_utils import (
+    apply_hook_request_output,
+    restore_text_with_replacements,
+    serialize_chat_message,
+    stream_text_replacements_from_hook_output,
+)
 from flocks.storage.storage import Storage
 from flocks.utils.log import Log
 
@@ -174,6 +180,7 @@ async def judge_goal_with_model(
     *,
     provider_id: str,
     model_id: str,
+    session_id: Optional[str] = None,
     initial_clarification: Optional[GoalClarification] = None,
 ) -> tuple[GoalVerdict, str]:
     """Hermes-style model judge using the active session provider/model."""
@@ -183,26 +190,84 @@ async def judge_goal_with_model(
 
     provider_options = build_provider_options(provider_id, model_id)
     provider_options.pop("max_tokens", None)
+    messages = [
+        ChatMessage(role="system", content=_MODEL_JUDGE_SYSTEM_PROMPT),
+        ChatMessage(
+            role="user",
+            content=(
+                f"Goal:\n{_format_goal_context(objective, initial_clarification)}\n\n"
+                "Latest assistant final response (truncated to the last 4KB):\n"
+                f"{_judge_input(last_response)}"
+            ),
+        ),
+    ]
+    replacements: list[tuple[str, str]] = []
+    if session_id:
+        from flocks.hooks.pipeline import HookPipeline, HookStage
+
+        hook_metadata = {
+            "sessionID": session_id,
+            "agent": "goal_judge",
+            "model": {
+                "providerID": provider_id,
+                "modelID": model_id,
+            },
+            "purpose": "goal_judge",
+        }
+        try:
+            llm_before_enabled = await HookPipeline.has_stage_handlers(
+                HookStage.LLM_BEFORE,
+                hook_metadata,
+            )
+        except Exception as exc:
+            log.error("goal.model_judge.hook_stage_probe_failed", {
+                "session_id": session_id,
+                "provider_id": provider_id,
+                "model_id": model_id,
+                "error": str(exc),
+            })
+            raise RuntimeError("Goal judge before-hook stage probe failed; request was not sent") from exc
+
+        if llm_before_enabled:
+            try:
+                hook_ctx = await HookPipeline.run_llm_before({
+                    **hook_metadata,
+                    "request": {
+                        "messageCount": len(messages),
+                        "messages": [serialize_chat_message(message) for message in messages],
+                        "toolCount": 0,
+                        "tools": [],
+                        "providerOptions": dict(provider_options),
+                        "providerToolsEnabled": False,
+                    },
+                })
+                hook_output = getattr(hook_ctx, "output", None) or {}
+                replacements = stream_text_replacements_from_hook_output(hook_output)
+                messages, provider_options = apply_hook_request_output(
+                    messages,
+                    provider_options,
+                    hook_output,
+                )
+                provider_options.pop("max_tokens", None)
+            except Exception as exc:
+                log.error("goal.model_judge.llm_before_failed", {
+                    "session_id": session_id,
+                    "provider_id": provider_id,
+                    "model_id": model_id,
+                    "error": str(exc),
+                })
+                raise RuntimeError("Goal judge before-hook failed; request was not sent") from exc
 
     response = await provider.chat(
         model_id=model_id,
-        messages=[
-            ChatMessage(role="system", content=_MODEL_JUDGE_SYSTEM_PROMPT),
-            ChatMessage(
-                role="user",
-                content=(
-                    f"Goal:\n{_format_goal_context(objective, initial_clarification)}\n\n"
-                    "Latest assistant final response (truncated to the last 4KB):\n"
-                    f"{_judge_input(last_response)}"
-                ),
-            ),
-        ],
+        messages=messages,
         **provider_options,
         max_tokens=JUDGE_MAX_TOKENS,
         temperature=0,
     )
 
-    payload = _extract_json_object(response.content)
+    response_content = restore_text_with_replacements(response.content, replacements)
+    payload = _extract_json_object(response_content)
     verdict = str(payload.get("verdict") or "").strip().lower()
     reason = _trim_reason(str(payload.get("reason") or ""))
     if verdict not in {"complete", "blocked", "waiting", "continue"}:
@@ -352,6 +417,7 @@ async def evaluate_after_turn(
                     last_response,
                     provider_id=provider_id,
                     model_id=model_id,
+                    session_id=session_id,
                     initial_clarification=state.initial_clarification,
                 )
             except Exception as exc:

flocks/session/lifecycle/compaction/compaction.py

@@ -1190,6 +1190,7 @@ async def process(
                         focus_instruction=focus_instruction,
                         previous_summary=previous_summary,
                         chat_messages=chat_messages,
+                        session_id=session_id,
                     )
                 except RuntimeError as _e:
                     # No provider configured — long cooldown (hermes: 600s)

flocks/session/lifecycle/compaction/summary.py

@@ -13,8 +13,16 @@
 # here in the future).
 ProgressCallback = Callable[[str, Dict[str, Any]], Awaitable[None]]
 
+from flocks.provider.provider import ChatMessage
 from flocks.utils.log import Log
 from flocks.session.prompt import SessionPrompt
+from flocks.session.llm_hook_utils import (
+    apply_hook_request_output,
+    restore_text_with_replacements,
+    serialize_chat_message,
+    stream_text_replacements_from_hook_output,
+    strip_think_blocks,
+)
 from .models import DEFAULT_COMPACTION_PROMPT_WITH_PREVIOUS
 
 log = Log.create(service="session.compaction.summarization")
@@ -159,16 +167,72 @@ async def _llm_chat_with_timeout(
     messages: list,
     max_tokens: int,
     timeout: int = COMPACTION_TIMEOUT_SECONDS,
+    session_id: Optional[str] = None,
+    purpose: str = "compaction_summary",
 ) -> Any:
     """Call provider_client.chat with a timeout guard."""
-    return await asyncio.wait_for(
+    provider_options: Dict[str, Any] = {"max_tokens": max_tokens}
+    replacements: list[tuple[str, str]] = []
+    if session_id:
+        try:
+            from flocks.hooks.pipeline import HookPipeline, HookStage
+
+            provider_id = (
+                getattr(provider_client, "provider_id", None)
+                or getattr(provider_client, "id", None)
+                or provider_client.__class__.__name__
+            )
+            llm_hook_metadata = {
+                "sessionID": session_id,
+                "agent": "session.compaction",
+                "step": None,
+                "model": {
+                    "providerID": provider_id,
+                    "modelID": model_id,
+                },
+                "purpose": purpose,
+            }
+            if await HookPipeline.has_stage_handlers(HookStage.LLM_BEFORE, llm_hook_metadata):
+                llm_before_ctx = await HookPipeline.run_llm_before({
+                    **llm_hook_metadata,
+                    "request": {
+                        "messageCount": len(messages),
+                        "messages": [serialize_chat_message(message) for message in messages],
+                        "toolCount": 0,
+                        "tools": [],
+                        "providerOptions": dict(provider_options),
+                        "providerToolsEnabled": False,
+                    },
+                })
+                hook_output = getattr(llm_before_ctx, "output", None) or {}
+                messages, provider_options = apply_hook_request_output(
+                    messages,
+                    provider_options,
+                    hook_output,
+                )
+                replacements = stream_text_replacements_from_hook_output(hook_output)
+        except Exception as hook_err:
+            log.error("compaction.llm_before_hook.error", {
+                "session_id": session_id,
+                "purpose": purpose,
+                "error": str(hook_err),
+            })
+            raise RuntimeError("compaction llm_before hook failed; request was not sent") from hook_err
+
+    provider_options.setdefault("max_tokens", max_tokens)
+    response = await asyncio.wait_for(
         provider_client.chat(
             model_id=model_id,
             messages=messages,
-            max_tokens=max_tokens,
+            **provider_options,
         ),
         timeout=timeout,
     )
+    if replacements and response is not None and isinstance(getattr(response, "content", None), str):
+        response.content = restore_text_with_replacements(response.content, replacements)
+    if response is not None and isinstance(getattr(response, "content", None), str):
+        response.content = strip_think_blocks(response.content)
+    return response
 
 
 async def summarize_single_pass(
@@ -181,6 +245,7 @@ async def summarize_single_pass(
     focus_instruction: Optional[str] = None,
     previous_summary: Optional[str] = None,
     chat_messages: Optional[list] = None,
+    session_id: Optional[str] = None,
 ) -> Optional[str]:
     """Generate summary in a single LLM call.
 
@@ -201,8 +266,6 @@ async def summarize_single_pass(
     as "merge new turns into the prior summary" rather than compressing from
     scratch.
     """
-    from flocks.provider.provider import ChatMessage
-
     if chat_messages:
         # Per-message truncation path (hermes-style): every turn contributes
         # a capped fragment (head + tail per message), so early decisions
@@ -236,6 +299,8 @@ async def summarize_single_pass(
             model_id=model_id,
             messages=[ChatMessage(role="user", content=request)],
             max_tokens=max_tokens,
+            session_id=session_id,
+            purpose="compaction_summary",
         )
     except asyncio.TimeoutError:
         log.error("compaction.single_pass.timeout", {
@@ -452,6 +517,8 @@ async def summarize_chunked_iterative(
                 messages=[ChatMessage(role="user", content=chunk_prompt)],
                 max_tokens=chunk_max_tokens,
                 timeout=COMPACTION_TIMEOUT_SECONDS,
+                session_id=session_id,
+                purpose="compaction_summary_chunk",
             )
             duration_ms = (time.perf_counter() - started) * 1000
             if resp and resp.content:

flocks/session/lifecycle/title.py

@@ -12,6 +12,12 @@
 
 from flocks.utils.log import Log
 from flocks.provider.provider import ChatMessage
+from flocks.session.llm_hook_utils import (
+    apply_hook_request_output,
+    restore_text_with_replacements,
+    serialize_chat_message,
+    stream_text_replacements_from_hook_output,
+)
 
 EventPublishCallback = Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]]
 
@@ -132,16 +138,61 @@ async def generate_title_after_first_message(
             # Send PROMPT_TITLE as system instruction, user question as user message
             title = ""
             try:
+                title_messages = [
+                    ChatMessage(role="system", content=_CANONICAL_TITLE_PROMPT),
+                    ChatMessage(role="user", content=question),
+                ]
+                provider_options: Dict[str, Any] = {"max_tokens": 50}
+                replacements: list[tuple[str, str]] = []
+                try:
+                    from flocks.hooks.pipeline import HookPipeline, HookStage
+
+                    llm_hook_metadata = {
+                        "sessionID": session_id,
+                        "agent": "session.title",
+                        "step": None,
+                        "model": {
+                            "providerID": provider_id,
+                            "modelID": model_id,
+                        },
+                        "purpose": "title_generation",
+                    }
+                    if await HookPipeline.has_stage_handlers(HookStage.LLM_BEFORE, llm_hook_metadata):
+                        llm_before_ctx = await HookPipeline.run_llm_before({
+                            **llm_hook_metadata,
+                            "request": {
+                                "messageCount": len(title_messages),
+                                "messages": [serialize_chat_message(message) for message in title_messages],
+                                "toolCount": 0,
+                                "tools": [],
+                                "providerOptions": dict(provider_options),
+                                "providerToolsEnabled": False,
+                            },
+                        })
+                        hook_output = getattr(llm_before_ctx, "output", None) or {}
+                        title_messages, provider_options = apply_hook_request_output(
+                            title_messages,
+                            provider_options,
+                            hook_output,
+                        )
+                        replacements = stream_text_replacements_from_hook_output(hook_output)
+                except Exception as hook_err:
+                    log.error("title.llm_before_hook.error", {
+                        "session_id": session_id,
+                        "error": str(hook_err),
+                    })
+                    raise RuntimeError("title llm_before hook failed; request was not sent") from hook_err
+
+                provider_options.setdefault("max_tokens", 50)
                 async for chunk in provider.chat_stream(
                     model_id,
-                    [
-                        ChatMessage(role="system", content=_CANONICAL_TITLE_PROMPT),
-                        ChatMessage(role="user", content=question),
-                    ],
-                    max_tokens=50,
+                    title_messages,
+                    **provider_options,
                 ):
                     if hasattr(chunk, 'delta') and chunk.delta:
                         title += chunk.delta
+                if replacements:
+                    title = restore_text_with_replacements(title, replacements)
             except Exception as llm_err:
                 log.warn("title.llm_failed", {
                     "session_id": session_id,