flock-core 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl

flock/runtime.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from typing import Any
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from flock.artifacts import Artifact
 
@@ -28,7 +28,9 @@ class EvalInputs(BaseModel):
 
         Example:
             >>> class TaskProcessor(EngineComponent):
-            ...     async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:
+            ...     async def evaluate(
+            ...         self, agent, ctx, inputs: EvalInputs
+            ...     ) -> EvalResult:
             ...         task = inputs.first_as(Task)
             ...         if not task:
             ...             return EvalResult.empty()
@@ -88,9 +90,13 @@ class EvalResult(BaseModel):
 
         Example:
             >>> class TaskProcessor(EngineComponent):
-            ...     async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:
+            ...     async def evaluate(
+            ...         self, agent, ctx, inputs: EvalInputs
+            ...     ) -> EvalResult:
             ...         task = inputs.first_as(Task)
-            ...         processed = Task(name=f"Done: {task.name}", priority=task.priority)
+            ...         processed = Task(
+            ...             name=f"Done: {task.name}", priority=task.priority
+            ...         )
             ...         return EvalResult.from_object(processed, agent=agent)
         """
         from flock.artifacts import Artifact
@@ -136,14 +142,16 @@ class EvalResult(BaseModel):
 
         Example:
             >>> class MovieEngine(EngineComponent):
-            ...     async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:
+            ...     async def evaluate(
+            ...         self, agent, ctx, inputs: EvalInputs
+            ...     ) -> EvalResult:
             ...         idea = inputs.first_as(Idea)
-            ...         movie = Movie(title=idea.topic.upper(), runtime=240, synopsis="...")
+            ...         movie = Movie(
+            ...             title=idea.topic.upper(), runtime=240, synopsis="..."
+            ...         )
             ...         tagline = Tagline(line="Don't miss it!")
             ...         return EvalResult.from_objects(
-            ...             movie, tagline,
-            ...             agent=agent,
-            ...             metrics={"confidence": 0.9}
+            ...             movie, tagline, agent=agent, metrics={"confidence": 0.9}
             ...         )
         """
         from flock.artifacts import Artifact
@@ -190,7 +198,9 @@ class EvalResult(BaseModel):
 
         Example:
             >>> class ConditionalProcessor(EngineComponent):
-            ...     async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:
+            ...     async def evaluate(
+            ...         self, agent, ctx, inputs: EvalInputs
+            ...     ) -> EvalResult:
             ...         task = inputs.first_as(Task)
             ...         if task.priority < 3:
             ...             return EvalResult.empty()  # Skip low priority
@@ -229,12 +239,15 @@ class EvalResult(BaseModel):
 
         Example:
             >>> class ValidationAgent(EngineComponent):
-            ...     async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:
+            ...     async def evaluate(
+            ...         self, agent, ctx, inputs: EvalInputs
+            ...     ) -> EvalResult:
             ...         task = inputs.first_as(Task)
             ...         is_valid = task.priority >= 1
-            ...         return EvalResult.with_state(
-            ...             {"validation_passed": is_valid, "validator": "priority_check"}
-            ...         )
+            ...         return EvalResult.with_state({
+            ...             "validation_passed": is_valid,
+            ...             "validator": "priority_check",
+            ...         })
         """
         return cls(
             artifacts=[],
@@ -245,13 +258,51 @@ class EvalResult(BaseModel):
 
 
 class Context(BaseModel):
-    board: Any
-    orchestrator: Any
-    correlation_id: UUID | None = None  # NEW!
+    """Runtime context for agent execution.
+
+    SECURITY FIX (2025-10-17): Simplified to data-only design.
+    Context is now just pre-filtered data with ZERO capabilities.
+
+    Vulnerabilities fixed:
+    - Vulnerability #1 (READ): Agents could bypass visibility via ctx.board.list()
+    - Vulnerability #2 (WRITE): Agents could bypass validation via ctx.board.publish()
+    - Vulnerability #3 (GOD MODE): Agents had unlimited ctx.orchestrator access
+    - Vulnerability #4 (STORE ACCESS): Agents could access ctx.store or ctx.provider._store
+
+    Solution: Orchestrator evaluates context BEFORE creating Context.
+    Engines receive only pre-filtered artifact data via ctx.artifacts.
+    No provider, no store, no capabilities - just immutable data.
+
+    Design Philosophy: Engines are pure functions (input + context → output).
+    They don't query, they don't mutate - they only transform data.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    # ❌ REMOVED: board: Any (security vulnerability)
+    # ❌ REMOVED: orchestrator: Any (security vulnerability)
+    # ❌ REMOVED: provider: Any (security vulnerability - engines could call provider methods)
+    # ❌ REMOVED: store: Any (security vulnerability - direct store access)
+
+    # ✅ FINAL SOLUTION: Pre-filtered artifacts (evaluated by orchestrator)
+    # Engines can only read this list - they cannot query for more data
+    artifacts: list[Artifact] = Field(
+        default_factory=list,
+        description="Pre-filtered conversation context artifacts (evaluated by orchestrator using context provider)",
+    )
+
+    # ✅ Agent identity (informational only - used by orchestrator for logging/tracing)
+    agent_identity: Any = Field(
+        default=None,
+        description="Agent identity (informational) - engines cannot use this to query data",
+    )
+
+    correlation_id: UUID | None = None
     task_id: str
     state: dict[str, Any] = Field(default_factory=dict)
     is_batch: bool = Field(
-        default=False, description="True if this execution is processing a BatchSpec accumulation"
+        default=False,
+        description="True if this execution is processing a BatchSpec accumulation",
     )
 
     def get_variable(self, key: str, default: Any = None) -> Any:

flock/service.py

@@ -40,7 +40,9 @@ class BlackboardHTTPService:
                 "visibility": artifact.visibility.model_dump(mode="json"),
                 "visibility_kind": getattr(artifact.visibility, "kind", "Unknown"),
                 "created_at": artifact.created_at.isoformat(),
-                "correlation_id": str(artifact.correlation_id) if artifact.correlation_id else None,
+                "correlation_id": str(artifact.correlation_id)
+                if artifact.correlation_id
+                else None,
                 "partition_key": artifact.partition_key,
                 "tags": sorted(artifact.tags),
                 "version": artifact.version,
@@ -56,7 +58,9 @@ class BlackboardHTTPService:
                     }
                     for record in consumptions
                 ]
-                data["consumed_by"] = sorted({record.consumer for record in consumptions})
+                data["consumed_by"] = sorted({
+                    record.consumer for record in consumptions
+                })
             return data
 
         def _parse_datetime(value: str | None, label: str) -> datetime | None:
@@ -65,7 +69,9 @@ class BlackboardHTTPService:
             try:
                 return datetime.fromisoformat(value)
             except ValueError as exc:  # pragma: no cover - FastAPI converts
-                raise HTTPException(status_code=400, detail=f"Invalid {label}: {value}") from exc
+                raise HTTPException(
+                    status_code=400, detail=f"Invalid {label}: {value}"
+                ) from exc
 
         def _make_filter_config(
             type_names: list[str] | None,
@@ -129,7 +135,9 @@ class BlackboardHTTPService:
             items: list[dict[str, Any]] = []
             for artifact in artifacts:
                 if isinstance(artifact, ArtifactEnvelope):
-                    items.append(_serialize_artifact(artifact.artifact, artifact.consumptions))
+                    items.append(
+                        _serialize_artifact(artifact.artifact, artifact.consumptions)
+                    )
                 else:
                     items.append(_serialize_artifact(artifact))
             return {
@@ -179,7 +187,9 @@ class BlackboardHTTPService:
                 type_name = item.get("type")
                 payload = item.get("payload") or {}
                 if not type_name:
-                    raise HTTPException(status_code=400, detail="Each input requires 'type'.")
+                    raise HTTPException(
+                        status_code=400, detail="Each input requires 'type'."
+                    )
                 model = type_registry.resolve(type_name)
                 instance = model(**payload)
                 inputs.append(instance)
@@ -253,7 +263,10 @@ class BlackboardHTTPService:
 
         @app.get("/metrics")
         async def metrics() -> PlainTextResponse:
-            lines = [f"blackboard_{key} {value}" for key, value in orchestrator.metrics.items()]
+            lines = [
+                f"blackboard_{key} {value}"
+                for key, value in orchestrator.metrics.items()
+            ]
             return PlainTextResponse("\n".join(lines))
 
     def run(

flock/store.py

@@ -227,7 +227,9 @@ class InMemoryBlackboardStore(BlackboardStore):
         self._lock = Lock()
         self._by_id: dict[UUID, Artifact] = {}
         self._by_type: dict[str, list[Artifact]] = defaultdict(list)
-        self._consumptions_by_artifact: dict[UUID, list[ConsumptionRecord]] = defaultdict(list)
+        self._consumptions_by_artifact: dict[UUID, list[ConsumptionRecord]] = (
+            defaultdict(list)
+        )
         self._agent_snapshots: dict[str, AgentSnapshotRecord] = {}
 
     async def publish(self, artifact: Artifact) -> None:
@@ -254,7 +256,9 @@ class InMemoryBlackboardStore(BlackboardStore):
             artifacts = self._by_type.get(canonical, [])
             return [artifact_type(**artifact.payload) for artifact in artifacts]  # type: ignore
 
-    async def extend(self, artifacts: Iterable[Artifact]) -> None:  # pragma: no cover - helper
+    async def extend(
+        self, artifacts: Iterable[Artifact]
+    ) -> None:  # pragma: no cover - helper
         for artifact in artifacts:
             await self.publish(artifact)
 
@@ -280,7 +284,9 @@ class InMemoryBlackboardStore(BlackboardStore):
         filters = filters or FilterConfig()
         canonical: set[str] | None = None
         if filters.type_names:
-            canonical = {type_registry.resolve_name(name) for name in filters.type_names}
+            canonical = {
+                type_registry.resolve_name(name) for name in filters.type_names
+            }
 
         visibility_filter = filters.visibility or set()
 
@@ -347,7 +353,9 @@ class InMemoryBlackboardStore(BlackboardStore):
             if not isinstance(artifact, Artifact):
                 raise TypeError("Expected Artifact instance")
             by_type[artifact.type] = by_type.get(artifact.type, 0) + 1
-            by_producer[artifact.produced_by] = by_producer.get(artifact.produced_by, 0) + 1
+            by_producer[artifact.produced_by] = (
+                by_producer.get(artifact.produced_by, 0) + 1
+            )
             kind = getattr(artifact.visibility, "kind", "Unknown")
             by_visibility[kind] = by_visibility.get(kind, 0) + 1
             for tag in artifact.tags:
@@ -476,7 +484,9 @@ class SQLiteBlackboardStore(BlackboardStore):
                 "version": artifact.version,
                 "visibility": visibility_json,
                 "tags": tags_json,
-                "correlation_id": str(artifact.correlation_id) if artifact.correlation_id else None,
+                "correlation_id": str(artifact.correlation_id)
+                if artifact.correlation_id
+                else None,
                 "partition_key": artifact.partition_key,
                 "created_at": created_at,
             }
@@ -816,7 +826,8 @@ class SQLiteBlackboardStore(BlackboardStore):
         by_visibility_rows = await cursor.fetchall()
         await cursor.close()
         by_visibility = {
-            (row["visibility_kind"] or "Unknown"): row["count"] for row in by_visibility_rows
+            (row["visibility_kind"] or "Unknown"): row["count"]
+            for row in by_visibility_rows
         }
 
         tag_query = f"""
@@ -839,7 +850,9 @@ class SQLiteBlackboardStore(BlackboardStore):
         cursor = await conn.execute(range_query, params_tuple)
         range_row = await cursor.fetchone()
         await cursor.close()
-        earliest = range_row["earliest"] if range_row and range_row["earliest"] else None
+        earliest = (
+            range_row["earliest"] if range_row and range_row["earliest"] else None
+        )
         latest = range_row["latest"] if range_row and range_row["latest"] else None
 
         return {
@@ -902,7 +915,9 @@ class SQLiteBlackboardStore(BlackboardStore):
         consumption_rows = await cursor.fetchall()
         await cursor.close()
 
-        consumed_by_type = {row["canonical_type"]: row["count"] for row in consumption_rows}
+        consumed_by_type = {
+            row["canonical_type"]: row["count"] for row in consumption_rows
+        }
         consumed_total = sum(consumed_by_type.values())
 
         return {
@@ -1156,7 +1171,9 @@ class SQLiteBlackboardStore(BlackboardStore):
         params: list[Any] = []
 
         if filters.type_names:
-            canonical = {type_registry.resolve_name(name) for name in filters.type_names}
+            canonical = {
+                type_registry.resolve_name(name) for name in filters.type_names
+            }
             placeholders = ", ".join("?" for _ in canonical)
             conditions.append(f"{prefix}canonical_type IN ({placeholders})")
             params.extend(sorted(canonical))
@@ -1172,7 +1189,9 @@ class SQLiteBlackboardStore(BlackboardStore):
 
         if filters.visibility:
             placeholders = ", ".join("?" for _ in filters.visibility)
-            conditions.append(f"json_extract({prefix}visibility, '$.kind') IN ({placeholders})")
+            conditions.append(
+                f"json_extract({prefix}visibility, '$.kind') IN ({placeholders})"
+            )
             params.extend(sorted(filters.visibility))
 
         if filters.start is not None:

flock/subscription.py

@@ -102,7 +102,7 @@ class Subscription:
         where: Sequence[Predicate] | None = None,
         text_predicates: Sequence[TextPredicate] | None = None,
         from_agents: Iterable[str] | None = None,
-        channels: Iterable[str] | None = None,
+        tags: Iterable[str] | None = None,
         join: JoinSpec | None = None,
         batch: BatchSpec | None = None,
         delivery: str = "exclusive",
@@ -116,7 +116,9 @@ class Subscription:
 
         # Register all types and build counts (supports duplicates for count-based AND gates)
         type_name_list = [type_registry.register(t) for t in types]
-        self.type_names: set[str] = set(type_name_list)  # Unique type names (for matching)
+        self.type_names: set[str] = set(
+            type_name_list
+        )  # Unique type names (for matching)
 
         # Count-based AND gate: Track how many of each type are required
         # Example: .consumes(A, A, B) → {"TypeA": 2, "TypeB": 1}
@@ -127,7 +129,7 @@ class Subscription:
         self.where = list(where or [])
         self.text_predicates = list(text_predicates or [])
         self.from_agents = set(from_agents or [])
-        self.channels = set(channels or [])
+        self.tags = set(tags or [])
         self.join = join
         self.batch = batch
         self.delivery = delivery
@@ -145,7 +147,7 @@ class Subscription:
             return False
         if self.from_agents and artifact.produced_by not in self.from_agents:
             return False
-        if self.channels and not artifact.tags.intersection(self.channels):
+        if self.tags and not artifact.tags.intersection(self.tags):
             return False
 
         # Evaluate where predicates on typed payloads

flock/utilities.py

@@ -31,7 +31,9 @@ class MetricsUtility(AgentComponent):
 
     name: str | None = "metrics"
 
-    async def on_pre_evaluate(self, agent, ctx: Context, inputs: EvalInputs) -> EvalInputs:
+    async def on_pre_evaluate(
+        self, agent, ctx: Context, inputs: EvalInputs
+    ) -> EvalInputs:
         ctx.state.setdefault("metrics", {})[f"{agent.name}:start"] = time.perf_counter()
         return inputs
 
@@ -42,7 +44,9 @@ class MetricsUtility(AgentComponent):
         start = metrics.get(f"{agent.name}:start")
         if start:
             metrics[f"{agent.name}:duration_ms"] = (time.perf_counter() - start) * 1000
-        result.metrics.update({k: v for k, v in metrics.items() if k.endswith("duration_ms")})
+        result.metrics.update({
+            k: v for k, v in metrics.items() if k.endswith("duration_ms")
+        })
         return result
 
 
@@ -78,11 +82,15 @@ class LoggingUtility(AgentComponent):
 
     async def on_pre_consume(self, agent, ctx: Context, inputs: list[Any]):
         summary = ", ".join(self._summarize_artifact(art) for art in inputs) or "<none>"
-        self._console.log(f"[{agent.name}] consume n={len(inputs)} artifacts -> {summary}")
+        self._console.log(
+            f"[{agent.name}] consume n={len(inputs)} artifacts -> {summary}"
+        )
         self._render_artifacts(agent.name, inputs, role="input")
         return await super().on_pre_consume(agent, ctx, inputs)
 
-    async def on_pre_evaluate(self, agent, ctx: Context, inputs: EvalInputs) -> EvalInputs:
+    async def on_pre_evaluate(
+        self, agent, ctx: Context, inputs: EvalInputs
+    ) -> EvalInputs:
         if self._stream_tokens:
             self._maybe_start_stream(agent, ctx)
         return await super().on_pre_evaluate(agent, ctx, inputs)
@@ -91,7 +99,9 @@ class LoggingUtility(AgentComponent):
         self, agent, ctx: Context, inputs: EvalInputs, result: EvalResult
     ) -> EvalResult:
         self._render_metrics(agent.name, result.metrics)
-        self._render_artifacts(agent.name, result.artifacts or inputs.artifacts, role="output")
+        self._render_artifacts(
+            agent.name, result.artifacts or inputs.artifacts, role="output"
+        )
         if result.logs:
             self._render_logs(agent.name, result.logs)
         awaited = await super().on_post_evaluate(agent, ctx, inputs, result)
@@ -102,7 +112,9 @@ class LoggingUtility(AgentComponent):
     async def on_post_publish(self, agent, ctx: Context, artifact):
         visibility = getattr(artifact.visibility, "kind", "Public")
         subtitle = f"visibility={visibility}"
-        panel = self._build_artifact_panel(artifact, role="published", subtitle=subtitle)
+        panel = self._build_artifact_panel(
+            artifact, role="published", subtitle=subtitle
+        )
         self._console.print(panel)
         await super().on_post_publish(agent, ctx, artifact)
 
@@ -121,7 +133,9 @@ class LoggingUtility(AgentComponent):
     # ------------------------------------------------------------------
     # Rendering helpers
 
-    def _render_artifacts(self, agent_name: str, artifacts: Sequence[Any], *, role: str) -> None:
+    def _render_artifacts(
+        self, agent_name: str, artifacts: Sequence[Any], *, role: str
+    ) -> None:
         for artifact in artifacts:
             panel = self._build_artifact_panel(artifact, role=role)
             self._console.print(panel)
@@ -189,7 +203,9 @@ class LoggingUtility(AgentComponent):
             else:
                 textual.append(line)
         for payload in json_sections:
-            panel = Panel(payload, title=f"{agent_name} ▸ dspy.output", border_style="green")
+            panel = Panel(
+                payload, title=f"{agent_name} ▸ dspy.output", border_style="green"
+            )
             self._console.print(panel)
         if textual:
             body = Text("\n".join(textual) + "\n")
@@ -244,7 +260,9 @@ class LoggingUtility(AgentComponent):
         with contextlib.suppress(asyncio.CancelledError):
             await task
 
-    async def _consume_stream(self, agent_name: str, stream_key: str, queue: asyncio.Queue) -> None:
+    async def _consume_stream(
+        self, agent_name: str, stream_key: str, queue: asyncio.Queue
+    ) -> None:
         body = Text()
         live: Live | None = None
         try:
@@ -254,7 +272,9 @@ class LoggingUtility(AgentComponent):
                     break
                 kind = event.get("kind")
                 if live is None:
-                    live_panel = Panel(body, title=f"{agent_name} ▸ streaming", border_style="cyan")
+                    live_panel = Panel(
+                        body, title=f"{agent_name} ▸ streaming", border_style="cyan"
+                    )
                     live = Live(
                         live_panel,
                         console=self._console,
@@ -274,19 +294,27 @@ class LoggingUtility(AgentComponent):
                     message = event.get("message") or ""
                     body.append(f"\n⚠ {message}\n", style="bold red")
                 if live is not None:
-                    live.update(Panel(body, title=f"{agent_name} ▸ streaming", border_style="cyan"))
+                    live.update(
+                        Panel(
+                            body, title=f"{agent_name} ▸ streaming", border_style="cyan"
+                        )
+                    )
         finally:
             if live is not None:
                 live.__exit__(None, None, None)
         if body.plain:
             self._console.print(
-                Panel(body, title=f"{agent_name} ▸ stream transcript", border_style="cyan")
+                Panel(
+                    body, title=f"{agent_name} ▸ stream transcript", border_style="cyan"
+                )
             )
 
     def _stream_key(self, agent, ctx: Context) -> str:
         return f"{ctx.task_id}:{agent.name}"
 
-    def _attach_stream_queue(self, state: MutableMapping[str, Any], queue: asyncio.Queue) -> None:
+    def _attach_stream_queue(
+        self, state: MutableMapping[str, Any], queue: asyncio.Queue
+    ) -> None:
         state.setdefault("_logging", {})["stream_queue"] = queue
 
     def _detach_stream_queue(self, state: MutableMapping[str, Any]) -> None:

flock/utility/output_utility_component.py

@@ -28,8 +28,12 @@ class OutputUtilityConfig(AgentComponentConfig):
     theme: OutputTheme = Field(
         default=OutputTheme.catppuccin_mocha, description="Theme for output formatting"
     )
-    render_table: bool = Field(default=True, description="Whether to render output as a table")
-    max_length: int = Field(default=1000, description="Maximum length for displayed output")
+    render_table: bool = Field(
+        default=True, description="Whether to render output as a table"
+    )
+    max_length: int = Field(
+        default=1000, description="Maximum length for displayed output"
+    )
     truncate_long_values: bool = Field(
         default=True, description="Whether to truncate long values in display"
     )
@@ -61,7 +65,9 @@ class OutputUtilityComponent(AgentComponent):
         default_factory=OutputUtilityConfig, description="Output configuration"
     )
 
-    def __init__(self, name: str = "output", config: OutputUtilityConfig | None = None, **data):
+    def __init__(
+        self, name: str = "output", config: OutputUtilityConfig | None = None, **data
+    ):
         if config is None:
             config = OutputUtilityConfig()
         super().__init__(name=name, config=config, **data)
@@ -96,7 +102,11 @@ class OutputUtilityComponent(AgentComponent):
         items = []
         prefix = "  " * indent
         for key, value in d.items():
-            if self.config.truncate_long_values and isinstance(value, str) and len(value) > 100:
+            if (
+                self.config.truncate_long_values
+                and isinstance(value, str)
+                and len(value) > 100
+            ):
                 value = value[:97] + "..."
             formatted_value = self._format_value(value, key)
             items.append(f"{prefix}  {key}: {formatted_value}")
@@ -125,7 +135,9 @@ class OutputUtilityComponent(AgentComponent):
             return f"[CODE:{language}]\n{code}\n[/CODE]"
 
         # Replace markdown-style code blocks
-        return re.sub(r"```(\w+)?\n(.*?)\n```", replace_code_block, text, flags=re.DOTALL)
+        return re.sub(
+            r"```(\w+)?\n(.*?)\n```", replace_code_block, text, flags=re.DOTALL
+        )
 
     async def on_post_evaluate(
         self, agent: "Agent", ctx: Context, inputs: EvalInputs, result: EvalResult
@@ -138,7 +150,9 @@ class OutputUtilityComponent(AgentComponent):
         streamed_artifact_id = None
 
         if ctx:
-            streaming_live_handled = bool(ctx.get_variable("_flock_stream_live_active", False))
+            streaming_live_handled = bool(
+                ctx.get_variable("_flock_stream_live_active", False)
+            )
             output_queued = bool(ctx.get_variable("_flock_output_queued", False))
             streamed_artifact_id = ctx.get_variable("_flock_streamed_artifact_id")
 
@@ -162,20 +176,24 @@ class OutputUtilityComponent(AgentComponent):
 
         # Skip output if streaming already handled it (and no ID to update)
         if streaming_live_handled:
-            logger.debug("Skipping static table because streaming rendered live output.")
+            logger.debug(
+                "Skipping static table because streaming rendered live output."
+            )
             return result
 
         # If output was queued due to concurrent stream, wait and then display
         if output_queued:
             # Wait for active streams to complete
-            orchestrator = getattr(ctx, "orchestrator", None)
-            if orchestrator:
+            # Phase 6+7 Security Fix: Use Agent class variable instead of ctx.state
+            if ctx:
                 import asyncio
 
+                from flock.agent import Agent
+
                 # Wait until no streams are active
                 max_wait = 30  # seconds
                 waited = 0
-                while getattr(orchestrator, "_active_streams", 0) > 0 and waited < max_wait:
+                while Agent._streaming_counter > 0 and waited < max_wait:
                     await asyncio.sleep(0.1)
                     waited += 0.1
                 logger.debug(
@@ -189,7 +207,9 @@ class OutputUtilityComponent(AgentComponent):
             try:
                 # Create a copy or select relevant parts to avoid modifying original result dict directly
                 display_result = result.copy()
-                display_result["context_snapshot"] = ctx.to_dict()  # Potential performance hit
+                display_result["context_snapshot"] = (
+                    ctx.to_dict()
+                )  # Potential performance hit
             except Exception:
                 display_result = result.copy()
                 display_result["context_snapshot"] = "[Error serializing context]"